405 lines
14 KiB
Rust
405 lines
14 KiB
Rust
//! VICReg (Variance-Invariance-Covariance Regularization) Implementation
|
||
//!
|
||
//! Self-supervised learning method that explicitly avoids collapse through three loss terms:
|
||
//! - Invariance: Similar representations for augmented views (MSE loss)
|
||
//! - Variance: Maintains variance ≥ γ in each dimension (hinge loss)
|
||
//! - Covariance: Decorrelates different dimensions (Frobenius norm of off-diagonal covariance)
|
||
//!
|
||
//! Based on "VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning"
|
||
//! (Bardes et al., 2021).
|
||
//!
|
||
//! Key features:
|
||
//! - No momentum encoders or large batches required
|
||
//! - Explicit variance and covariance regularization prevents collapse
|
||
//! - Expander network (projector) transforms representations
|
||
//! - Support for different backbone architectures
|
||
//! - Symmetric loss computation on both augmented views
|
||
|
||
use crate::prelude::*;
|
||
use super::byol::Backbone;
|
||
use std::sync::Arc;
|
||
use parking_lot::RwLock;
|
||
|
||
/// VICReg configuration parameters
|
||
#[derive(Debug, Clone)]
|
||
pub struct VICRegConfig {
|
||
/// Dimension of backbone output features
|
||
pub backbone_dim: usize,
|
||
/// Expander network layer dimensions
|
||
pub expander_dims: Vec<usize>,
|
||
/// λ - Coefficient for invariance loss (similarity between views)
|
||
pub sim_coeff: f32,
|
||
/// μ - Coefficient for variance loss (prevent collapse)
|
||
pub std_coeff: f32,
|
||
/// ν - Coefficient for covariance loss (decorrelate features)
|
||
pub cov_coeff: f32,
|
||
/// γ - Target standard deviation (variance regularization)
|
||
pub variance_target: f32,
|
||
/// ε - Epsilon for numerical stability
|
||
pub epsilon: f32,
|
||
}
|
||
|
||
impl Default for VICRegConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
backbone_dim: 2048,
|
||
expander_dims: vec![8192, 8192, 8192],
|
||
sim_coeff: 25.0,
|
||
std_coeff: 25.0,
|
||
cov_coeff: 1.0,
|
||
variance_target: 1.0,
|
||
epsilon: 1e-4,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl VICRegConfig {
|
||
/// Create new VICReg configuration
|
||
pub fn new(backbone_dim: usize, expander_dims: Vec<usize>) -> Self {
|
||
Self {
|
||
backbone_dim,
|
||
expander_dims,
|
||
..Default::default()
|
||
}
|
||
}
|
||
|
||
/// Set backbone dimension
|
||
pub fn with_backbone_dim(mut self, backbone_dim: usize) -> Self {
|
||
self.backbone_dim = backbone_dim;
|
||
self
|
||
}
|
||
|
||
/// Set invariance loss coefficient (λ)
|
||
pub fn with_sim_coeff(mut self, sim_coeff: f32) -> Self {
|
||
self.sim_coeff = sim_coeff;
|
||
self
|
||
}
|
||
|
||
/// Set variance loss coefficient (μ)
|
||
pub fn with_std_coeff(mut self, std_coeff: f32) -> Self {
|
||
self.std_coeff = std_coeff;
|
||
self
|
||
}
|
||
|
||
/// Set covariance loss coefficient (ν)
|
||
pub fn with_cov_coeff(mut self, cov_coeff: f32) -> Self {
|
||
self.cov_coeff = cov_coeff;
|
||
self
|
||
}
|
||
|
||
/// Set variance target (γ)
|
||
pub fn with_variance_target(mut self, variance_target: f32) -> Self {
|
||
self.variance_target = variance_target;
|
||
self
|
||
}
|
||
|
||
/// Set epsilon for numerical stability
|
||
pub fn with_epsilon(mut self, epsilon: f32) -> Self {
|
||
self.epsilon = epsilon;
|
||
self
|
||
}
|
||
}
|
||
|
||
/// Expander network that maps backbone features to embedding space
|
||
/// Similar to projector but terminology from VICReg paper
|
||
#[derive(Debug)]
|
||
pub struct ExpanderNetwork {
|
||
layers: Vec<LinearLayer>,
|
||
input_dim: usize,
|
||
layer_dims: Vec<usize>,
|
||
device: Device,
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
struct LinearLayer {
|
||
weight: Arc<RwLock<Tensor>>,
|
||
bias: Arc<RwLock<Tensor>>,
|
||
}
|
||
|
||
impl ExpanderNetwork {
|
||
/// Create new expander network
|
||
pub fn new(
|
||
input_dim: usize,
|
||
layer_dims: Vec<usize>,
|
||
device: &Device,
|
||
) -> Result<Self> {
|
||
let mut layers = Vec::new();
|
||
let mut current_dim = input_dim;
|
||
|
||
for &output_dim in layer_dims.iter() {
|
||
// Initialize weights with Xavier/Glorot normal initialization
|
||
let scale = (2.0 / (current_dim + output_dim) as f32).sqrt();
|
||
let weight = Tensor::randn(vec![current_dim, output_dim], DType::F32, device)?
|
||
.mul(&Tensor::full(vec![current_dim, output_dim], scale, DType::F32, device)?)?;
|
||
let bias = Tensor::zeros(vec![output_dim], device)?;
|
||
|
||
layers.push(LinearLayer {
|
||
weight: Arc::new(RwLock::new(weight)),
|
||
bias: Arc::new(RwLock::new(bias)),
|
||
});
|
||
|
||
current_dim = output_dim;
|
||
}
|
||
|
||
Ok(Self {
|
||
layers,
|
||
input_dim,
|
||
layer_dims,
|
||
device: device.clone(),
|
||
})
|
||
}
|
||
|
||
/// Forward pass through expander network
|
||
pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
||
let mut x = input.clone();
|
||
|
||
for (i, layer) in self.layers.iter().enumerate() {
|
||
// Linear transformation
|
||
let weight = layer.weight.read();
|
||
let bias = layer.bias.read();
|
||
x = x.matmul(&*weight)?.add(&*bias)?;
|
||
|
||
// ReLU activation (except for the last layer)
|
||
if i < self.layers.len() - 1 {
|
||
x = x.relu()?;
|
||
}
|
||
}
|
||
|
||
Ok(x)
|
||
}
|
||
|
||
pub fn input_dim(&self) -> usize {
|
||
self.input_dim
|
||
}
|
||
|
||
pub fn output_dim(&self) -> usize {
|
||
*self.layer_dims.last().unwrap_or(&self.input_dim)
|
||
}
|
||
|
||
pub fn layer_dims(&self) -> &[usize] {
|
||
&self.layer_dims
|
||
}
|
||
}
|
||
|
||
/// Result of VICReg loss computation
|
||
#[derive(Debug, Clone)]
|
||
pub struct VICRegLossResult {
|
||
pub total_loss: Tensor,
|
||
pub invariance_loss: Tensor,
|
||
pub variance_loss_y1: Tensor,
|
||
pub variance_loss_y2: Tensor,
|
||
pub covariance_loss_y1: Tensor,
|
||
pub covariance_loss_y2: Tensor,
|
||
}
|
||
|
||
/// Training result from VICReg training step
|
||
#[derive(Debug, Clone)]
|
||
pub struct VICRegTrainingResult {
|
||
pub total_loss: f32,
|
||
pub invariance_loss: f32,
|
||
pub variance_loss: f32,
|
||
pub covariance_loss: f32,
|
||
}
|
||
|
||
/// VICReg trainer combining backbone and expander networks
|
||
pub struct VICRegTrainer<B: Backbone> {
|
||
backbone: B,
|
||
expander: ExpanderNetwork,
|
||
config: VICRegConfig,
|
||
device: Device,
|
||
}
|
||
|
||
impl<B: Backbone> VICRegTrainer<B> {
|
||
/// Create new VICReg trainer
|
||
pub fn new(
|
||
backbone: B,
|
||
config: VICRegConfig,
|
||
device: &Device,
|
||
) -> Result<Self> {
|
||
let backbone_dim = backbone.output_dim();
|
||
let mut config = config;
|
||
config.backbone_dim = backbone_dim;
|
||
|
||
let expander = ExpanderNetwork::new(
|
||
backbone_dim,
|
||
config.expander_dims.clone(),
|
||
device,
|
||
)?;
|
||
|
||
Ok(Self {
|
||
backbone,
|
||
expander,
|
||
config,
|
||
device: device.clone(),
|
||
})
|
||
}
|
||
|
||
/// Forward pass through both networks
|
||
pub fn forward(&self, x1: &Tensor, x2: &Tensor) -> Result<VICRegLossResult> {
|
||
// Get representations from backbone
|
||
let z1 = self.backbone.forward(x1)?;
|
||
let z2 = self.backbone.forward(x2)?;
|
||
|
||
// Transform through expander network
|
||
let y1 = self.expander.forward(&z1)?;
|
||
let y2 = self.expander.forward(&z2)?;
|
||
|
||
// Compute VICReg loss
|
||
compute_vicreg_loss(&y1, &y2, &self.config)
|
||
}
|
||
|
||
/// Training step with augmentation (simplified for test)
|
||
pub fn train_step(&mut self, images: &Tensor) -> Result<VICRegTrainingResult> {
|
||
// In practice, you would apply different augmentations here
|
||
// For testing, we'll use the same image with small noise as "augmentation"
|
||
let noise = Tensor::randn(images.shape().to_vec(), &self.device)?
|
||
.mul(&Tensor::full(images.shape().to_vec(), 0.01, DType::F32, &self.device)?)?;
|
||
let x2 = images.add(&noise)?;
|
||
|
||
let loss_result = self.forward(images, &x2)?;
|
||
|
||
// In real implementation, you would:
|
||
// 1. Compute gradients
|
||
// 2. Update parameters
|
||
// 3. Zero gradients
|
||
|
||
Ok(VICRegTrainingResult {
|
||
total_loss: loss_result.total_loss.to_scalar::<f32>()?,
|
||
invariance_loss: loss_result.invariance_loss.to_scalar::<f32>()?,
|
||
variance_loss: (loss_result.variance_loss_y1.to_scalar::<f32>()? +
|
||
loss_result.variance_loss_y2.to_scalar::<f32>()?) * 0.5,
|
||
covariance_loss: (loss_result.covariance_loss_y1.to_scalar::<f32>()? +
|
||
loss_result.covariance_loss_y2.to_scalar::<f32>()?) * 0.5,
|
||
})
|
||
}
|
||
|
||
pub fn config(&self) -> &VICRegConfig {
|
||
&self.config
|
||
}
|
||
}
|
||
|
||
/// Compute invariance loss: MSE between representations
|
||
pub fn compute_invariance_loss(y1: &Tensor, y2: &Tensor) -> Result<Tensor> {
|
||
let diff = y1.sub(y2)?;
|
||
let squared_diff = diff.mul(&diff)?;
|
||
|
||
// Mean over all dimensions
|
||
let total_elements = squared_diff.shape().iter().product::<usize>() as f32;
|
||
let sum = squared_diff.sum(None)?;
|
||
sum.div(&Tensor::full(vec![], total_elements, DType::F32, squared_diff.device())?)
|
||
}
|
||
|
||
/// Compute variance loss: Hinge loss to maintain std ≥ γ
|
||
pub fn compute_variance_loss(y: &Tensor, gamma: f32, epsilon: f32) -> Result<Tensor> {
|
||
let batch_size = y.shape()[0] as f32;
|
||
let feature_dim = y.shape()[1];
|
||
|
||
// Compute mean along batch dimension
|
||
let mean = y.sum_axis(0)?.div(&Tensor::full(vec![feature_dim], batch_size, DType::F32, y.device())?)?;
|
||
|
||
// Center the features: y - mean
|
||
let centered = y.sub(&mean.unsqueeze(0)?)?;
|
||
|
||
// Compute variance: E[(y - mean)²]
|
||
let squared_diff = centered.mul(¢ered)?;
|
||
let variance = squared_diff.sum_axis(0)?.div(&Tensor::full(vec![feature_dim], batch_size, DType::F32, y.device())?)?;
|
||
|
||
// Compute standard deviation with epsilon for stability
|
||
let eps_tensor = Tensor::full(vec![feature_dim], epsilon, DType::F32, y.device())?;
|
||
let variance_eps = variance.add(&eps_tensor)?;
|
||
let std = variance_eps.sqrt()?;
|
||
|
||
// Hinge loss: ReLU(γ - std)
|
||
let gamma_tensor = Tensor::full(vec![feature_dim], gamma, DType::F32, y.device())?;
|
||
let hinge = gamma_tensor.sub(&std)?;
|
||
let hinge_loss = hinge.relu()?;
|
||
|
||
// Mean over all features
|
||
hinge_loss.mean_all()
|
||
}
|
||
|
||
/// Compute covariance loss: Frobenius norm of off-diagonal covariance matrix
|
||
pub fn compute_covariance_loss(y: &Tensor, epsilon: f32) -> Result<Tensor> {
|
||
let batch_size = y.shape()[0] as f32;
|
||
let feature_dim = y.shape()[1];
|
||
|
||
// Center the features
|
||
let mean = y.sum_axis(0)?.div(&Tensor::full(vec![feature_dim], batch_size, DType::F32, y.device())?)?;
|
||
let centered = y.sub(&mean.unsqueeze(0)?)?;
|
||
|
||
// Compute covariance matrix: (1/N) * X^T * X
|
||
let covariance = centered.transpose(0, 1)?.matmul(¢ered)?
|
||
.div(&Tensor::full(vec![feature_dim, feature_dim], batch_size, DType::F32, y.device())?)?;
|
||
|
||
// Add epsilon to diagonal for numerical stability
|
||
let eps_tensor = Tensor::full(vec![feature_dim], epsilon, DType::F32, y.device())?;
|
||
let eye = create_identity_matrix(feature_dim, y.device())?;
|
||
let eps_eye = eye.mul(&eps_tensor.unsqueeze(1)?)?;
|
||
let stable_cov = covariance.add(&eps_eye)?;
|
||
|
||
// Zero out diagonal elements (we only want off-diagonal)
|
||
let off_diagonal_mask = create_off_diagonal_mask(feature_dim, y.device())?;
|
||
let off_diagonal_cov = stable_cov.mul(&off_diagonal_mask)?;
|
||
|
||
// Compute Frobenius norm of off-diagonal elements
|
||
let squared = off_diagonal_cov.mul(&off_diagonal_cov)?;
|
||
squared.sum(None)
|
||
}
|
||
|
||
/// Compute complete VICReg loss with all three terms
|
||
pub fn compute_vicreg_loss(
|
||
y1: &Tensor,
|
||
y2: &Tensor,
|
||
config: &VICRegConfig
|
||
) -> Result<VICRegLossResult> {
|
||
// Invariance loss: similarity between representations
|
||
let invariance_loss = compute_invariance_loss(y1, y2)?;
|
||
|
||
// Variance losses: prevent dimensional collapse
|
||
let variance_loss_y1 = compute_variance_loss(y1, config.variance_target, config.epsilon)?;
|
||
let variance_loss_y2 = compute_variance_loss(y2, config.variance_target, config.epsilon)?;
|
||
|
||
// Covariance losses: decorrelate features
|
||
let covariance_loss_y1 = compute_covariance_loss(y1, config.epsilon)?;
|
||
let covariance_loss_y2 = compute_covariance_loss(y2, config.epsilon)?;
|
||
|
||
// Weighted total loss
|
||
let sim_term = invariance_loss.mul(&Tensor::full(vec![], config.sim_coeff, DType::F32, y1.device())?)?;
|
||
let std_term1 = variance_loss_y1.mul(&Tensor::full(vec![], config.std_coeff, DType::F32, y1.device())?)?;
|
||
let std_term2 = variance_loss_y2.mul(&Tensor::full(vec![], config.std_coeff, DType::F32, y1.device())?)?;
|
||
let cov_term1 = covariance_loss_y1.mul(&Tensor::full(vec![], config.cov_coeff, DType::F32, y1.device())?)?;
|
||
let cov_term2 = covariance_loss_y2.mul(&Tensor::full(vec![], config.cov_coeff, DType::F32, y1.device())?)?;
|
||
|
||
let total_loss = sim_term
|
||
.add(&std_term1)?
|
||
.add(&std_term2)?
|
||
.add(&cov_term1)?
|
||
.add(&cov_term2)?;
|
||
|
||
Ok(VICRegLossResult {
|
||
total_loss,
|
||
invariance_loss,
|
||
variance_loss_y1,
|
||
variance_loss_y2,
|
||
covariance_loss_y1,
|
||
covariance_loss_y2,
|
||
})
|
||
}
|
||
|
||
/// Helper function to create identity matrix
|
||
fn create_identity_matrix(size: usize, device: &Device) -> Result<Tensor> {
|
||
let mut data = vec![0.0f32; size * size];
|
||
for i in 0..size {
|
||
data[i * size + i] = 1.0;
|
||
}
|
||
Tensor::from_slice(&data, vec![size, size], DType::F32, device)
|
||
}
|
||
|
||
/// Helper function to create off-diagonal mask (1s off-diagonal, 0s on diagonal)
|
||
fn create_off_diagonal_mask(size: usize, device: &Device) -> Result<Tensor> {
|
||
let mut data = vec![1.0f32; size * size];
|
||
for i in 0..size {
|
||
data[i * size + i] = 0.0;
|
||
}
|
||
Tensor::from_slice(&data, vec![size, size], DType::F32, device)
|
||
} |