rustytorch/crates/training/rtx-transformers/src/ssl/vicreg_tests.rs

//! VICReg Test Suite
//!
//! Comprehensive tests for VICReg (Variance-Invariance-Covariance Regularization) implementation.
//! Follows strict TDD - tests are written first, then implementation.
//!
//! VICReg is a self-supervised learning method that explicitly avoids collapse through:
//! - 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 cov)

#[cfg(all(test, feature = "disabled_tests"))]
mod tests {
    use super::super::byol::Backbone;
    use super::super::vicreg::*;
    use crate::prelude::*;

    /// Test VICRegConfig creation with default hyperparameters
    #[test]
    fn test_vicreg_config_default() {
        let config = VICRegConfig::default();

        // Default hyperparameters from VICReg paper
        assert_eq!(config.sim_coeff, 25.0); // λ - invariance weight
        assert_eq!(config.std_coeff, 25.0); // μ - variance weight  
        assert_eq!(config.cov_coeff, 1.0); // ν - covariance weight
        assert_eq!(config.variance_target, 1.0); // γ - target standard deviation
        assert_eq!(config.epsilon, 1e-4); // ε - numerical stability
        assert_eq!(config.expander_dims, vec![8192, 8192, 8192]);
    }

    /// Test VICRegConfig creation with custom parameters
    #[test]
    fn test_vicreg_config_custom() -> Result<()> {
        let config = VICRegConfig::new(2048, vec![4096, 4096, 2048])
            .with_sim_coeff(20.0)
            .with_std_coeff(30.0)
            .with_cov_coeff(2.0)
            .with_variance_target(0.5)
            .with_epsilon(1e-5);

        assert_eq!(config.backbone_dim, 2048);
        assert_eq!(config.expander_dims, vec![4096, 4096, 2048]);
        assert_eq!(config.sim_coeff, 20.0);
        assert_eq!(config.std_coeff, 30.0);
        assert_eq!(config.cov_coeff, 2.0);
        assert_eq!(config.variance_target, 0.5);
        assert_eq!(config.epsilon, 1e-5);

        Ok(())
    }

    /// Test expander network initialization  
    #[test]
    fn test_expander_network_init() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let input_dim = 2048;
        let dims = vec![4096, 4096, 2048];

        let expander = ExpanderNetwork::new(input_dim, dims.clone(), &device)?;

        assert_eq!(expander.input_dim(), input_dim);
        assert_eq!(expander.output_dim(), *dims.last().unwrap());
        assert_eq!(expander.layer_dims(), &dims);

        Ok(())
    }

    /// Test expander forward pass dimensions
    #[test]
    fn test_expander_forward_dimensions() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 32;
        let input_dim = 2048;
        let output_dim = 8192;
        let dims = vec![4096, 4096, output_dim];

        let expander = ExpanderNetwork::new(input_dim, dims, &device)?;
        let input = Tensor::randn(vec![batch_size, input_dim], DType::F32, &device)?;

        let output = expander.forward(&input)?;

        assert_eq!(output.shape(), &[batch_size, output_dim]);

        Ok(())
    }

    /// Test invariance loss computation (MSE between representations)
    #[test]
    fn test_invariance_loss() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 64;
        let feature_dim = 8192;

        // Create two representations
        let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
        let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;

        let inv_loss = compute_invariance_loss(&y1, &y2)?;

        // Invariance loss should be a scalar (MSE)
        assert_eq!(inv_loss.shape(), &[]);

        // Loss should be positive
        let loss_val = inv_loss.to_scalar::<f32>()?;
        assert!(loss_val >= 0.0);

        Ok(())
    }

    /// Test variance loss computation (hinge loss to maintain std ≥ γ)
    #[test]
    fn test_variance_loss() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 64;
        let feature_dim = 8192;
        let gamma = 1.0;
        let epsilon = 1e-4;

        // Create representation
        let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;

        let var_loss = compute_variance_loss(&y, gamma, epsilon)?;

        // Variance loss should be a scalar
        assert_eq!(var_loss.shape(), &[]);

        // Loss should be non-negative (hinge loss)
        let loss_val = var_loss.to_scalar::<f32>()?;
        assert!(loss_val >= 0.0);

        Ok(())
    }

    /// Test covariance loss computation (decorrelation of features)
    #[test]
    fn test_covariance_loss() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 64;
        let feature_dim = 128; // Smaller for computational efficiency in tests
        let epsilon = 1e-4;

        // Create representation
        let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;

        let cov_loss = compute_covariance_loss(&y, epsilon)?;

        // Covariance loss should be a scalar
        assert_eq!(cov_loss.shape(), &[]);

        // Loss should be non-negative (sum of squared off-diagonal elements)
        let loss_val = cov_loss.to_scalar::<f32>()?;
        assert!(loss_val >= 0.0);

        Ok(())
    }

    /// Test complete VICReg loss computation
    #[test]
    fn test_vicreg_loss_computation() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 64;
        let feature_dim = 8192;

        let config = VICRegConfig::default();

        // Create two representations (from augmented views)
        let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
        let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;

        let loss_result = compute_vicreg_loss(&y1, &y2, &config)?;

        // Check that all loss components are computed
        assert_eq!(loss_result.total_loss.shape(), &[]);
        assert_eq!(loss_result.invariance_loss.shape(), &[]);
        assert_eq!(loss_result.variance_loss_y1.shape(), &[]);
        assert_eq!(loss_result.variance_loss_y2.shape(), &[]);
        assert_eq!(loss_result.covariance_loss_y1.shape(), &[]);
        assert_eq!(loss_result.covariance_loss_y2.shape(), &[]);

        // All losses should be non-negative
        let total = loss_result.total_loss.to_scalar::<f32>()?;
        let inv = loss_result.invariance_loss.to_scalar::<f32>()?;
        let var1 = loss_result.variance_loss_y1.to_scalar::<f32>()?;
        let var2 = loss_result.variance_loss_y2.to_scalar::<f32>()?;
        let cov1 = loss_result.covariance_loss_y1.to_scalar::<f32>()?;
        let cov2 = loss_result.covariance_loss_y2.to_scalar::<f32>()?;

        assert!(total >= 0.0);
        assert!(inv >= 0.0);
        assert!(var1 >= 0.0);
        assert!(var2 >= 0.0);
        assert!(cov1 >= 0.0);
        assert!(cov2 >= 0.0);

        // Total loss should be weighted sum of components
        let expected_total = config.sim_coeff * inv
            + config.std_coeff * (var1 + var2)
            + config.cov_coeff * (cov1 + cov2);

        assert!((total - expected_total).abs() < 1e-3);

        Ok(())
    }

    /// Test VICReg trainer initialization
    #[test]
    fn test_vicreg_trainer_init() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = VICRegConfig::default();

        // Create simple backbone mock
        struct SimpleBackbone {
            output_dim: usize,
            device: Device,
        }

        impl Backbone for SimpleBackbone {
            fn forward(&self, x: &Tensor) -> Result<Tensor> {
                let batch_size = x.shape()[0];
                Ok(Tensor::randn(
                    vec![batch_size, self.output_dim],
                    DType::F32,
                    &self.device,
                )?)
            }

            fn output_dim(&self) -> usize {
                self.output_dim
            }
        }

        let backbone = SimpleBackbone {
            output_dim: 2048,
            device: device.clone(),
        };

        let trainer = VICRegTrainer::new(backbone, config, &device)?;

        assert_eq!(trainer.config().backbone_dim, 2048);
        assert_eq!(trainer.config().expander_dims, vec![8192, 8192, 8192]);

        Ok(())
    }

    /// Test VICReg trainer forward pass
    #[test]
    fn test_vicreg_trainer_forward() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 32;
        let config = VICRegConfig::default();

        struct SimpleBackbone {
            output_dim: usize,
            device: Device,
        }

        impl Backbone for SimpleBackbone {
            fn forward(&self, x: &Tensor) -> Result<Tensor> {
                let batch_size = x.shape()[0];
                Ok(Tensor::randn(
                    vec![batch_size, self.output_dim],
                    DType::F32,
                    &self.device,
                )?)
            }

            fn output_dim(&self) -> usize {
                self.output_dim
            }
        }

        let backbone = SimpleBackbone {
            output_dim: 2048,
            device: device.clone(),
        };

        let trainer = VICRegTrainer::new(backbone, config.clone(), &device)?;

        // Create two augmented views
        let x1 = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;
        let x2 = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;

        let result = trainer.forward(&x1, &x2)?;

        // Should return VICReg loss components
        assert_eq!(result.total_loss.shape(), &[]);
        assert!(result.total_loss.to_scalar::<f32>()? >= 0.0);

        Ok(())
    }

    /// Test training step with gradient updates
    #[test]
    fn test_vicreg_training_step() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 16; // Small batch for test
        let config = VICRegConfig::default();

        struct SimpleBackbone {
            output_dim: usize,
            device: Device,
        }

        impl Backbone for SimpleBackbone {
            fn forward(&self, x: &Tensor) -> Result<Tensor> {
                let batch_size = x.shape()[0];
                Ok(Tensor::randn(
                    vec![batch_size, self.output_dim],
                    DType::F32,
                    &self.device,
                )?)
            }

            fn output_dim(&self) -> usize {
                self.output_dim
            }
        }

        let backbone = SimpleBackbone {
            output_dim: 2048,
            device: device.clone(),
        };

        let mut trainer = VICRegTrainer::new(backbone, config, &device)?;

        // Create batch of images (normally these would be augmented views)
        let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;

        // Perform training step
        let result = trainer.train_step(&images)?;

        // Should return training metrics
        assert!(result.total_loss >= 0.0);
        assert!(result.invariance_loss >= 0.0);
        assert!(result.variance_loss >= 0.0);
        assert!(result.covariance_loss >= 0.0);

        Ok(())
    }

    /// Test identical inputs produce zero invariance loss
    #[test]
    fn test_identical_inputs_zero_invariance() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 32;
        let feature_dim = 256;

        // Identical representations
        let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;

        let inv_loss = compute_invariance_loss(&y, &y)?;
        let loss_val = inv_loss.to_scalar::<f32>()?;

        // Invariance loss should be very close to zero for identical inputs
        assert!(loss_val < 1e-6);

        Ok(())
    }

    /// Test high variance inputs have low variance loss
    #[test]
    fn test_high_variance_low_loss() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 100;
        let feature_dim = 256;
        let gamma = 1.0;
        let epsilon = 1e-4;

        // Create high variance input (scale up random values)
        let y_base = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
        let scale = Tensor::full(y_base.shape(), 3.0, DType::F32, &device)?;
        let y = y_base.mul(&scale)?;

        let var_loss = compute_variance_loss(&y, gamma, epsilon)?;
        let loss_val = var_loss.to_scalar::<f32>()?;

        // High variance should result in low variance loss (hinge = max(0, γ - std))
        assert!(loss_val < 0.1);

        Ok(())
    }

    /// Test uncorrelated features have low covariance loss
    #[test]
    fn test_uncorrelated_features_low_covariance_loss() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 100;
        let feature_dim = 10; // Small for controlled test
        let epsilon = 1e-4;

        // Create orthogonal/uncorrelated features manually
        let mut y_data = vec![0.0; batch_size * feature_dim];

        // Fill each feature dimension independently
        for i in 0..feature_dim {
            for b in 0..batch_size {
                y_data[b * feature_dim + i] = (b as f32 + i as f32 * 100.0).sin();
            }
        }

        let y = Tensor::from_slice(&y_data, vec![batch_size, feature_dim], DType::F32, &device)?;

        let cov_loss = compute_covariance_loss(&y, epsilon)?;
        let loss_val = cov_loss.to_scalar::<f32>()?;

        // Uncorrelated features should have relatively low covariance loss
        assert!(loss_val >= 0.0); // Should be non-negative

        Ok(())
    }

    /// Test batch size consistency across loss computations
    #[test]
    fn test_batch_size_consistency() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let feature_dim = 512;
        let epsilon = 1e-4;
        let gamma = 1.0;

        // Test different batch sizes
        for batch_size in [1, 8, 32, 128] {
            let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
            let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;

            // All loss computations should work regardless of batch size
            let inv_loss = compute_invariance_loss(&y1, &y2)?;
            let var_loss = compute_variance_loss(&y1, gamma, epsilon)?;
            let cov_loss = compute_covariance_loss(&y1, epsilon)?;

            assert_eq!(inv_loss.shape(), &[]);
            assert_eq!(var_loss.shape(), &[]);
            assert_eq!(cov_loss.shape(), &[]);

            assert!(inv_loss.to_scalar::<f32>()? >= 0.0);
            assert!(var_loss.to_scalar::<f32>()? >= 0.0);
            assert!(cov_loss.to_scalar::<f32>()? >= 0.0);
        }

        Ok(())
    }
}