VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning

1 INTRODUCTION . Self-supervised representation learning has made significant progress over the last years , almost reaching the performance of supervised baselines on many downstream tasks Bachman et al . ( 2019 ) ; Misra & Maaten ( 2020 ) ; He et al . ( 2020 ) ; Tian et al . ( 2020 ) ; Caron et al . ( 2020 ) ; Grill et al . ( 2020 ) ; Chen & He ( 2020 ) ; Gidaris et al . ( 2021 ) ; Zbontar et al . ( 2021 ) . Several recent approaches rely on a joint embedding architecture in which two networks are trained to produce similar embeddings for different views of the same image . A popular instance is the Siamese network architecture Bromley et al . ( 1994 ) , where the two networks share the same weights . The main challenge with joint embedding architectures is to prevent a collapse in which the two branches ignore the inputs and produce identical and constant output vectors . There are two main approaches to preventing collapse : contrastive methods and information maximization methods . Contrastive Bromley et al . ( 1994 ) ; Chopra et al . ( 2005 ) ; He et al . ( 2020 ) ; Hjelm et al . ( 2019 ) ; Chen et al . ( 2020a ) methods tend to be costly , require large batch sizes or memory banks , and use a loss that explicitly pushes the embeddings of dissimilar images away from each other . They often require a mining procedure to search for offending dissimilar samples from a memory bank He et al . ( 2020 ) or from the current batch Chen et al . ( 2020a ) . Quantization-based approaches Caron et al . ( 2020 ; 2018 ) force the embeddings of different samples to belong to different clusters on the unit sphere . Collapse is prevented by ensuring that the assignment of samples to clusters is as uniform as possible . A similarity term encourages the cluster assignment score vectors from the two branches to be similar . More recently , a few methods have appeared that do not rely on contrastive samples or vector quantization , yet produce high-quality representations , for example BYOL Grill et al . ( 2020 ) and SimSiam Chen & He ( 2020 ) . They exploit several tricks : batch-wise or feature-wise normalization , a `` momentum encoder '' in which the parameter vector of one branch is a low-pass-filtered version of the parameter vector of the other branch Grill et al . ( 2020 ) ; Richemond et al . ( 2020 ) , or a stop-gradient operation in one of the branches Chen & He ( 2020 ) . The dynamics of learning in these methods , and how they avoid collapse , is not fully understood , although theoretical and empirical studies point to the crucial importance of batch-wise or feature-wise normalization Richemond et al . ( 2020 ) ; Tian et al . ( 2021 ) . Finally , an alternative class of collapse prevention methods relies on maximizing the information content of the embedding Zbontar et al . ( 2021 ) ; Ermolov et al . ( 2021 ) . These methods prevent informational collapse by decorrelating every pair of variables of the embedding vectors . This indirectly maximizes the information content of the embedding vectors . The Barlow Twins method drives the normalized cross-correlation matrix of the two embeddings towards the identity Zbontar et al . ( 2021 ) , while the Whitening-MSE method whitens and spreads out the embedding vectors on the unit sphere Ermolov et al . ( 2021 ) . 2 VICREG : INTUITION . We introduce VICReg ( Variance-Invariance-Covariance Regularization ) , a self-supervised method for training joint embedding architectures based on the principle of preserving the information content of the embeddings . The basic idea is to use a loss function with three terms : • Invariance : the mean square distance between the embedding vectors . • Variance : a hinge loss to maintain the standard deviation ( over a batch ) of each variable of the embedding above a given threshold . This term forces the embedding vectors of samples within a batch to be different . • Covariance : a term that attracts the covariances ( over a batch ) between every pair of ( centered ) embedding variables towards zero . This term decorrelates the variables of each embedding and prevents an informational collapse in which the variables would vary together or be highly correlated . Variance and Covariance terms are applied to both branches of the architecture separately , thereby preserving the information content of each embedding at a certain level and preventing informational collapse independently for the two branches . The main contribution of this paper is the Variance preservation term , which explicitly prevents a collapse due to a shrinkage of the embedding vectors towards zero . The Covariance criterion is borrowed from the Barlow Twins method and prevents informational collapse due to redundancy between the embedding variables Zbontar et al . ( 2021 ) . VICReg is more generally applicable than most of the aforementioned methods because of fewer constraints on the architecture . In particular , VICReg : • does not require that the weights of the two branches be shared , not that the architectures be identical , nor that the inputs be of the same nature ; • does not require a memory bank , nor contrastive samples , nor a large batch size ; • does not require batch-wise nor feature-wise normalization ; and • does not require vector quantization nor a predictor module . Other methods require asymmetric stop gradient operations , as in SimSiam Chen & He ( 2020 ) , weight sharing between the two branches as in classical Siamese nets , or weight sharing through exponential moving average dampening with stop gradient in one branch , as in BYOL and MoCo He et al . ( 2020 ) ; Grill et al . ( 2020 ) ; Chen et al . ( 2020c ) , large batches of contrastive samples , as in SimCLR Chen et al . ( 2020a ) , or batch-wise and/or feature-wise normalization Caron et al . ( 2020 ) ; Grill et al . ( 2020 ) ; Chen & He ( 2020 ) ; Zbontar et al . ( 2021 ) ; Ermolov et al . ( 2021 ) . One of the most interesting feature of VICReg is the fact that the two branches are not required to share the same parameters , architecture , or input modality . This opens the door to the use of non-contrastive self-supervised joint-embedding for multi-modal signals , such as video and audio . We demonstrate the effectiveness of the proposed approach by evaluating the representations learned with VICReg on several downstream image recognition tasks including linear head and semi-supervised evaluation protocols for image classification on ImageNet Deng et al . ( 2009 ) , and other classification , detection , instance segmentation , and retrieval tasks . Furthermore , we show that incorporating variance preservation into other self-supervised joint-embedding methods yields better training stability and performance improvement on downstream tasks . More generally , we show that VICReg is an explicit and effective , yet simple method for preventing collapse in self-supervised joint-embedding learning . 3 RELATED WORK . Contrastive learning . In contrastive SSL methods applied to joint embedding architectures , the output embeddings for a sample and its distorted version are brought close to each other , while other samples and their distortions are pushed away . The method is most often applied to Siamese architectures in which the two branches have identical architectures and share weights Misra & Maaten ( 2020 ) ; He et al . ( 2020 ) ; Bromley et al . ( 1994 ) ; Hjelm et al . ( 2019 ) ; Chen et al . ( 2020a ; c ) ; Hadsell et al . ( 2006 ) ; Ye et al . ( 2019 ) ; Wu et al . ( 2018 ) ; van den Oord et al . ( 2018 ) ; Chen et al . ( 2020b ) . Many authors use the InfoNCE loss van den Oord et al . ( 2018 ) in which the repulsive force is larger for contrastive samples that are closer to the reference . While these methods yield good performance , they require large amounts of contrastive pairs in order to work well . These contrastive pairs can be sampled from a memory bank as in MoCo He et al . ( 2020 ) , or given by the current batch of data as in SimCLR Chen et al . ( 2020a ) , with a significant memory footprint . This downside of contrastive methods motivates a search for alternatives . Clustering methods . Instead of viewing each sample as its own class , clustering-based methods group them into clusters based on some similarity measure Caron et al . ( 2020 ; 2018 ) ; Bautista et al . ( 2016 ) ; Yang et al . ( 2016 ) ; Xie et al . ( 2016 ) ; Huang et al . ( 2019 ) ; Zhuang et al . ( 2019 ) ; Caron et al . ( 2019 ) ; Asano et al . ( 2020 ) ; Yan et al . ( 2020 ) . DeepCluster Caron et al . ( 2018 ) uses k-means assignments of representations from previous iterations as pseudo-labels for the new representations , which requires an expensive clustering phase done asynchronously , and makes the method hard to scale up . SwAV Caron et al . ( 2020 ) mitigates this issue by learning the clusters online while maintaining a balanced partition of the assignments through the Sinkhorn-Knopp transform Cuturi ( 2013 ) . These clustering approaches can be viewed as contrastive learning at the level of clusters which still requires a lot of negative comparisons to work well . Distillation methods . Recent proposals such as BYOL , SimSiam , OBoW and variants Grill et al . ( 2020 ) ; Chen & He ( 2020 ) ; Gidaris et al . ( 2021 ) ; Richemond et al . ( 2020 ) ; Gidaris et al . ( 2020 ) have shown that collapse can be avoided by using architectural tricks inspired by knowledge distillation Hinton et al . ( 2015 ) . These methods train a student network to predict the representations of a teacher network , for which the weights are a running average of the student network ’ s weights Grill et al . ( 2020 ) , or are shared with the student network , but no gradient is back-propagated through the teacher Chen & He ( 2020 ) . These methods are effective , but there is no clear understanding of why and how they avoid collapse . Alternatively , the images can be represented as bags of word over a dictionary of visual features , which effectively prevents collapse . In OBoW Gidaris et al . ( 2020 ) and Gidaris et al . ( 2021 ) the dictionary is obtained by off-line or on-line clustering . By contrast , our method explicitly prevents collapse in the two branches independently , which removes the requirement for shared weights and identical architecture , opening the door to the application of joint-embedding SSL to multi-modal signals . Information maximization methods . A principle to prevent collapse is to maximize the information content of the embeddings . Two such methods were recently proposed : W-MSE Ermolov et al . ( 2021 ) and Barlow Twins Zbontar et al . ( 2021 ) . In W-MSE , an extra module transforms the embeddings into the eigenspace of their covariance matrix ( whitening or Karhunen-Loève transform ) , and forces the vectors thereby obtained to be uniformly distributed on the unit sphere . In Barlow Twins , a loss term attempts to make the normalized cross-correlation matrix of the embedding vectors from the two branches to be close to the identity . Both methods attempt to produce embedding variables that are decorrelated from each other , thus preventing an informational collapse in which the variables carry redundant information . Because all variables are normalized over a batch , there is no incentive for them to shrink nor expand . This seems to sufficient to prevent collapse . Our method borrows the decorrelation mechanism of Barlow Twins . But it includes an explicit variance-preservation term for each variable of the two embeddings and thus does not require any normalization .

Summary: This paper proposes a Variance-Invariance-Covariance regularization (VICReg) for self-supervised learning. The total loss function is the weighted sum of three terms: 1. distance between two views of the same input batch; 2. hinge loss on the variance of the embedding variables, encouraging them to stay away from zero; 3. loss on the covariance between different embedding variables, encouraging them to be close to zero, and hence preventing a representation collapse. The proposed model does not require the embedding branches to be identical, nor does it require mining of contrastive pairs. Experiments are conducted to demonstrate the effectiveness of the proposed regularization.