Aligning Feature Distributions in VICReg Using Maximum Mean Discrepancy for Enhanced Manifold Awareness in Self-Supervised Representation Learning

Abstract

Self-supervised learning (SSL) methods like VICReg have shown considerable success in generating ro- bust data representation by promoting invariance across augmented views. However, VICReg’s focus on pairwise alignment between augmentations lim- its its capacity to ensure broader consistency across entire batches of diverse transformations. In this paper, we enhance VICReg by integrating a Maxi- mum Mean Discrepancy (MMD) term, which aligns feature distributions across the entire batch in a Reproducing Kernel Hilbert Space (RKHS), thereby promoting batch-level invariance. By enforcing a unified feature distribution across a batch, MMD enables the model to capture higher-order depen- dencies and reduce variability among augmented views. We have evaluated our approach on MNIST, CIFAR-10, and STL-10, where the results demon- strate improved representation quality, as evidenced by clustering accuracy and linear classification per- formance. The results highlight the effectiveness of incorporating MMD term into VICReg in enhancing the representation quality