Initial commit

crates/training/rtx-transformers/src/ssl/barlow_twins.rs

@@ -0,0 +1,434 @@
+//! Barlow Twins Implementation
+//!
+//! Self-supervised learning method using redundancy reduction to learn representations.
+//! Based on "Barlow Twins: Self-Supervised Learning via Redundancy Reduction" (Zbontar et al., 2021).
+//!
+//! Key features:
+//! - Cross-correlation matrix computation between twin representations
+//! - Redundancy reduction loss (making correlation matrix close to identity) 
+//! - Projector network for embedding transformation
+//! - Support for different backbone architectures
+//! - Batch normalization in projector
+//! - Symmetric loss computation
+
+use crate::prelude::*;
+use super::byol::Backbone;  // Reuse the Backbone trait
+use std::sync::Arc;
+use parking_lot::RwLock;
+
+/// Barlow Twins configuration parameters
+#[derive(Debug, Clone)]
+pub struct BarlowTwinsConfig {
+    /// Dimension of backbone output features
+    pub backbone_dim: usize,
+    /// Projector network layer dimensions  
+    pub projector_dims: Vec<usize>,
+    /// Lambda coefficient for redundancy reduction term
+    pub lambda_coeff: f32,
+    /// Whether to use batch normalization in projector
+    pub batch_norm: bool,
+    /// Scale factor for numerical stability
+    pub scale_loss: f32,
+}
+
+impl Default for BarlowTwinsConfig {
+    fn default() -> Self {
+        Self {
+            backbone_dim: 2048,  // Will be overridden based on backbone
+            projector_dims: vec![8192, 8192, 8192],
+            lambda_coeff: 0.005,
+            batch_norm: true,
+            scale_loss: 1.0 / 32.0,  // For numerical stability
+        }
+    }
+}
+
+impl BarlowTwinsConfig {
+    /// Create new Barlow Twins configuration
+    pub fn new(backbone_dim: usize, projector_dims: Vec<usize>) -> Self {
+        Self {
+            backbone_dim,
+            projector_dims,
+            ..Default::default()
+        }
+    }
+
+    /// Set backbone dimension
+    pub fn with_backbone_dim(mut self, backbone_dim: usize) -> Self {
+        self.backbone_dim = backbone_dim;
+        self
+    }
+
+    /// Set lambda coefficient for redundancy reduction
+    pub fn with_lambda_coeff(mut self, lambda_coeff: f32) -> Self {
+        self.lambda_coeff = lambda_coeff;
+        self
+    }
+
+    /// Set whether to use batch normalization
+    pub fn with_batch_norm(mut self, batch_norm: bool) -> Self {
+        self.batch_norm = batch_norm;
+        self
+    }
+
+    /// Set loss scale factor
+    pub fn with_scale_loss(mut self, scale_loss: f32) -> Self {
+        self.scale_loss = scale_loss;
+        self
+    }
+}
+
+/// Projector network that maps backbone features to embedding space
+#[derive(Debug)]
+pub struct ProjectorNetwork {
+    layers: Vec<LinearLayer>,
+    batch_norms: Vec<Option<BatchNorm>>,
+    input_dim: usize,
+    layer_dims: Vec<usize>,
+    use_batch_norm: bool,
+    device: Device,
+}
+
+#[derive(Debug)]
+struct LinearLayer {
+    weight: Arc<RwLock<Tensor>>,
+    bias: Arc<RwLock<Tensor>>,
+}
+
+#[derive(Debug)]
+struct BatchNorm {
+    weight: Arc<RwLock<Tensor>>,
+    bias: Arc<RwLock<Tensor>>,
+    running_mean: Arc<RwLock<Tensor>>,
+    running_var: Arc<RwLock<Tensor>>,
+    eps: f32,
+    momentum: f32,
+}
+
+impl ProjectorNetwork {
+    /// Create new projector network
+    pub fn new(
+        input_dim: usize,
+        layer_dims: Vec<usize>,
+        use_batch_norm: bool,
+        device: &Device,
+    ) -> Result<Self> {
+        let mut layers = Vec::new();
+        let mut batch_norms = Vec::new();
+        
+        let mut current_dim = input_dim;
+        
+        for (i, &output_dim) in layer_dims.iter().enumerate() {
+            // Linear layer
+            let weight = Tensor::randn(vec![current_dim, output_dim], 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)),
+            });
+            
+            // Batch normalization (except for the last layer)
+            if use_batch_norm && i < layer_dims.len() - 1 {
+                let bn_weight = Tensor::ones(vec![output_dim], device)?;
+                let bn_bias = Tensor::zeros(vec![output_dim], device)?;
+                let running_mean = Tensor::zeros(vec![output_dim], device)?;
+                let running_var = Tensor::ones(vec![output_dim], device)?;
+                
+                batch_norms.push(Some(BatchNorm {
+                    weight: Arc::new(RwLock::new(bn_weight)),
+                    bias: Arc::new(RwLock::new(bn_bias)),
+                    running_mean: Arc::new(RwLock::new(running_mean)),
+                    running_var: Arc::new(RwLock::new(running_var)),
+                    eps: 1e-5,
+                    momentum: 0.1,
+                }));
+            } else {
+                batch_norms.push(None);
+            }
+            
+            current_dim = output_dim;
+        }
+        
+        Ok(Self {
+            layers,
+            batch_norms,
+            input_dim,
+            layer_dims,
+            use_batch_norm,
+            device: device.clone(),
+        })
+    }
+
+    /// Forward pass through projector 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)?;
+            
+            // Batch normalization (if enabled and not last layer)
+            if let Some(ref batch_norm) = self.batch_norms[i] {
+                x = self.apply_batch_norm(&x, batch_norm)?;
+            }
+            
+            // ReLU activation (except for the last layer)
+            if i < self.layers.len() - 1 {
+                x = x.relu()?;
+            }
+        }
+        
+        Ok(x)
+    }
+
+    fn apply_batch_norm(&self, x: &Tensor, batch_norm: &BatchNorm) -> Result<Tensor> {
+        // Simplified batch normalization implementation
+        let weight = batch_norm.weight.read();
+        let bias = batch_norm.bias.read();
+        let running_mean = batch_norm.running_mean.read();
+        let running_var = batch_norm.running_var.read();
+        
+        // Normalize: (x - mean) / sqrt(var + eps)
+        let eps_tensor = Tensor::full(running_var.shape(), batch_norm.eps, DType::F32, &self.device)?;
+        let var_eps = running_var.add(&eps_tensor)?;
+        let std = var_eps.sqrt()?;
+        
+        let normalized = x.sub(&*running_mean)?.div(&std)?;
+        let scaled = normalized.mul(&*weight)?;
+        let output = scaled.add(&*bias)?;
+        
+        Ok(output)
+    }
+
+    /// Get input dimension
+    pub fn input_dim(&self) -> usize {
+        self.input_dim
+    }
+
+    /// Get output dimension
+    pub fn output_dim(&self) -> usize {
+        self.layer_dims.last().copied().unwrap_or(self.input_dim)
+    }
+
+    /// Get layer dimensions
+    pub fn layer_dims(&self) -> &[usize] {
+        &self.layer_dims
+    }
+
+    /// Check if using batch normalization
+    pub fn has_batch_norm(&self) -> bool {
+        self.use_batch_norm
+    }
+}
+
+/// Normalize features to have zero mean and unit variance per feature dimension
+pub fn normalize_features(features: &Tensor) -> Result<Tensor> {
+    let mean = features.mean(&[0])?;  // Mean across batch dimension
+    let centered = features.sub(&mean)?;
+    
+    let variance = centered.pow_scalar(2.0)?.mean(&[0])?;
+    let eps = Tensor::full(variance.shape(), 1e-8, DType::F32, variance.device())?;
+    let std = variance.add(&eps)?.sqrt()?;
+    
+    centered.div(&std)
+}
+
+/// Compute cross-correlation matrix between two normalized embeddings
+pub fn compute_cross_correlation_matrix(y1: &Tensor, y2: &Tensor) -> Result<Tensor> {
+    let batch_size = y1.shape()[0] as f32;
+    
+    // Normalize features
+    let y1_norm = normalize_features(y1)?;
+    let y2_norm = normalize_features(y2)?;
+    
+    // Compute cross-correlation: C[i,j] = sum(y1_norm[i] * y2_norm[j]) / batch_size
+    let y1_t = y1_norm.transpose(0, 1)?;  // [feature_dim, batch_size]
+    let cross_corr = y1_t.matmul(&y2_norm)?;  // [feature_dim, feature_dim]
+    
+    let batch_size_tensor = Tensor::full(&[], batch_size, DType::F32, cross_corr.device())?;
+    cross_corr.div(&batch_size_tensor)
+}
+
+/// Extract diagonal elements from a square matrix
+pub fn extract_diagonal(matrix: &Tensor) -> Result<Tensor> {
+    let shape = matrix.shape();
+    let size = shape[0];
+    
+    let mut diag_values = Vec::with_capacity(size);
+    let matrix_data = matrix.to_vec::<f32>()?;
+    
+    for i in 0..size {
+        let idx = i * size + i;  // Diagonal index in flattened matrix
+        diag_values.push(matrix_data[idx]);
+    }
+    
+    Tensor::from_vec(diag_values, vec![size], matrix.device())
+}
+
+/// Result of Barlow Twins loss computation
+#[derive(Debug)]
+pub struct BarlowTwinsLossResult {
+    /// Total loss
+    pub loss: Tensor,
+    /// Invariance loss (diagonal term)
+    pub invariance_loss: f32,
+    /// Redundancy reduction loss (off-diagonal term)
+    pub redundancy_loss: f32,
+    /// Cross-correlation matrix
+    pub cross_correlation: Tensor,
+}
+
+/// Compute Barlow Twins loss
+pub fn compute_barlow_twins_loss(
+    y1: &Tensor,
+    y2: &Tensor,
+    lambda_coeff: f32,
+    scale_loss: f32,
+) -> Result<BarlowTwinsLossResult> {
+    // Compute cross-correlation matrix
+    let cross_corr = compute_cross_correlation_matrix(y1, y2)?;
+    let feature_dim = cross_corr.shape()[0];
+    
+    // Create identity matrix
+    let identity = Tensor::eye(feature_dim, DType::F32, cross_corr.device())?;
+    
+    // Invariance loss: sum((1 - C[i,i])^2) - diagonal should be 1
+    let diag_diff = identity.sub(&cross_corr)?;
+    let diag_squared = diag_diff.pow_scalar(2.0)?;
+    
+    // Extract diagonal elements for invariance loss
+    let diagonal = extract_diagonal(&cross_corr)?;
+    let ones = Tensor::ones(diagonal.shape(), diagonal.device())?;
+    let diag_loss_vec = ones.sub(&diagonal)?.pow_scalar(2.0)?;
+    let invariance_loss_tensor = diag_loss_vec.sum(None)?;
+    let invariance_loss = invariance_loss_tensor.to_vec::<f32>()?[0];
+    
+    // Redundancy reduction loss: sum(C[i,j]^2 for i≠j) - off-diagonal should be 0
+    let cross_corr_squared = cross_corr.pow_scalar(2.0)?;
+    let total_squared = cross_corr_squared.sum(None)?;
+    let diag_squared_sum = diag_squared.sum(None)?;
+    let redundancy_loss_tensor = total_squared.sub(&diag_squared_sum)?;
+    let redundancy_loss = redundancy_loss_tensor.to_vec::<f32>()?[0];
+    
+    // Total loss = invariance_loss + lambda * redundancy_loss
+    let lambda_tensor = Tensor::full(&[], lambda_coeff, DType::F32, cross_corr.device())?;
+    let weighted_redundancy = redundancy_loss_tensor.mul(&lambda_tensor)?;
+    let total_loss = invariance_loss_tensor.add(&weighted_redundancy)?;
+    
+    // Apply scaling
+    let scale_tensor = Tensor::full(&[], scale_loss, DType::F32, cross_corr.device())?;
+    let scaled_loss = total_loss.mul(&scale_tensor)?;
+    
+    Ok(BarlowTwinsLossResult {
+        loss: scaled_loss,
+        invariance_loss,
+        redundancy_loss,
+        cross_correlation: cross_corr,
+    })
+}
+
+/// Barlow Twins trainer
+pub struct BarlowTwinsTrainer {
+    backbone: Box<dyn Backbone>,
+    projector: ProjectorNetwork,
+    config: BarlowTwinsConfig,
+    device: Device,
+    training: bool,
+}
+
+/// Result of a Barlow Twins training step
+#[derive(Debug)]
+pub struct BarlowTwinsTrainingResult {
+    /// Total loss
+    pub loss: Tensor,
+    /// Invariance loss component
+    pub invariance_loss: f32,
+    /// Redundancy reduction loss component
+    pub redundancy_loss: f32,
+    /// Cross-correlation matrix
+    pub cross_correlation: Tensor,
+}
+
+impl BarlowTwinsTrainer {
+    /// Create new Barlow Twins trainer
+    pub fn new(
+        backbone: Box<dyn Backbone>,
+        config: BarlowTwinsConfig,
+        device: &Device,
+    ) -> Result<Self> {
+        let backbone_dim = backbone.output_dim();
+        
+        // Update config with actual backbone dimension
+        let mut config = config;
+        config.backbone_dim = backbone_dim;
+        
+        let projector = ProjectorNetwork::new(
+            backbone_dim,
+            config.projector_dims.clone(),
+            config.batch_norm,
+            device,
+        )?;
+        
+        Ok(Self {
+            backbone,
+            projector,
+            config,
+            device: device.clone(),
+            training: true,
+        })
+    }
+
+    /// Perform one training step with two augmented views
+    pub fn train_step(&mut self, images: &Tensor, _seed: Option<u64>) -> Result<BarlowTwinsTrainingResult> {
+        if !self.training {
+            return Err(TransformerError::ConfigurationError("Trainer must be in training mode".to_string()));
+        }
+        
+        // For now, create two simple "augmented" views by adding noise
+        // In practice, this would use proper augmentation pipeline
+        let noise1 = Tensor::randn(images.shape(), &self.device)?.mul_scalar(0.01)?;
+        let noise2 = Tensor::randn(images.shape(), &self.device)?.mul_scalar(0.01)?;
+        let view1 = images.add(&noise1)?;
+        let view2 = images.add(&noise2)?;
+        
+        // Forward pass through backbone and projector
+        let z1 = self.backbone.forward(&view1)?;
+        let z2 = self.backbone.forward(&view2)?;
+        
+        let y1 = self.projector.forward(&z1)?;
+        let y2 = self.projector.forward(&z2)?;
+        
+        // Compute Barlow Twins loss
+        let loss_result = compute_barlow_twins_loss(&y1, &y2, self.config.lambda_coeff, self.config.scale_loss)?;
+        
+        Ok(BarlowTwinsTrainingResult {
+            loss: loss_result.loss,
+            invariance_loss: loss_result.invariance_loss,
+            redundancy_loss: loss_result.redundancy_loss,
+            cross_correlation: loss_result.cross_correlation,
+        })
+    }
+
+    /// Extract features for evaluation (using backbone only)
+    pub fn extract_features(&self, images: &Tensor) -> Result<Tensor> {
+        self.backbone.forward(images)
+    }
+
+    /// Set to training mode
+    pub fn train(&mut self) {
+        self.training = true;
+    }
+
+    /// Set to evaluation mode
+    pub fn eval(&mut self) {
+        self.training = false;
+    }
+
+    /// Get configuration
+    pub fn config(&self) -> &BarlowTwinsConfig {
+        &self.config
+    }
+}