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 combines three objective functions for the self-supervised visual pre-training on ImageNet.  (1) The alignment between the two different views of an identical image, which is very common for existing methods;  (2) The covariance term to bring the off-diagonal coefficients of the features' covariance matrix to zero, which is modified from the Barlow Twins;  (3) The variance term that defines a hinge function on the standard deviation of embeddings along the batch dimension for every specific dimension of the feature projections . To the best of the reviewer's knowledge, such objective function is firstly applied for the visual pre-training in this paper, although the same measure has been used to analyze the model collapse problem (e.g., in the paper of SimSiam), but not be designed as a specific pre-trained loss function.