VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning — part 2/3

Cleaned-text code points [14847, 29669). Bibliography omitted here; source.md includes it.

We describe here the variance, invariance and covariance terms that compose our loss function. The images are processed in batches, and we denote Z = [ z 1 , … , z n ] Z=[z_{1},\dots,z_{n}] and Z ′ = [ z 1 ′ , … , z n ′ ] Z^{\prime}=[z_{1}^{\prime},\dots,z_{n}^{\prime}] the two batches composed of n n vectors of dimension d d , of embeddings coming out of the two branches of the siamese architecture. We denote by z j z^{j} the vector composed of each value at dimension j j in all vectors in Z Z . We define the variance regularization term v v as a hinge function on the standard deviation of the embeddings along the batch dimension: v ⁡ ( Z ) = 1 d ​ ∑ j = 1 d max ⁡ ( 0 , γ − S ⁡ ( z j , ϵ ) ) , v(Z)=\frac{1}{d}\sum_{j=1}^{d}\max(0,\gamma-S(z^{j},\epsilon)), (1) where S S is the regularized standard deviation defined by: S ⁡ ( x , ϵ ) = Var ⁡ ( x ) + ϵ , S(x,\epsilon)=\sqrt{\mathrm{Var}(x)+\epsilon}, (2) γ \gamma is a constant target value for the standard deviation, fixed to 1 1 in our experiments, ϵ \epsilon is a small scalar preventing numerical instabilities. This criterion encourages the variance inside the current batch to be equal to γ \gamma along each dimension, preventing collapse with all the inputs mapped on the same vector. Using the standard deviation and not directly the variance is crucial. Indeed, if we take S ⁡ ( x ) = Var ⁡ ( x ) S(x)=\mathrm{Var}(x) in the hinge function, the gradient of S S with respect to x x becomes close to 0 when x x is close to x ¯ \bar{x} . In this case, the gradient of v v also becomes close to 0 and the embeddings collapse. We define the covariance matrix of Z Z as: C ⁡ ( Z ) = 1 n − 1 ​ ∑ i = 1 n ( z i − z ¯ ) ​ ( z i − z ¯ ) T , where z ¯ = 1 n ​ ∑ i = 1 n z i . C(Z)=\frac{1}{n-1}\sum_{i=1}^{n}(z_{i}-\bar{z})(z_{i}-\bar{z})^{T},\ \ \ \textrm{where}\ \ \ \bar{z}=\frac{1}{n}\sum_{i=1}^{n}z_{i}. (3) Inspired by Barlow Twins Zbontar et al. (2021) , we can then define the covariance regularization term c c as the sum of the squared off-diagonal coefficients of C ⁡ ( Z ) C(Z) , with a factor 1 / d 1/d that scales the criterion as a function of the dimension: c ⁡ ( Z ) = 1 d ​ ∑ i ≠ j [ C ⁡ ( Z ) ] i , j 2 . c(Z)=\frac{1}{d}\sum_{i\neq j}[C(Z)]_{i,j}^{2}. (4) This term encourages the off-diagonal coefficients of C ⁡ ( Z ) C(Z) to be close to 0 0 , decorrelating the different dimensions of the embeddings and preventing them from encoding similar information. Decorrelation at the embedding level ultimately has a decorrelation effect at the representation level, which is a non trivial phenomenon that we study in Appendix D . We finally define the invariance criterion s s between Z Z and Z ′ Z^{\prime} as the mean-squared euclidean distance between each pair of vectors, without any normalization: s ⁡ ( Z , Z ′ ) = 1 n ​ ∑ i ‖ z i − z i ′ ‖ 2 2 . s(Z,Z^{\prime})=\frac{1}{n}\sum_{i}\|z_{i}-z^{\prime}_{i}\|_{2}^{2}. (5) The overall loss function is a weighted average of the invariance, variance and covariance terms: ℓ ⁡ ( Z , Z ′ ) = λ ​ s ​ ( Z , Z ′ ) + μ ⁡ [ v ⁡ ( Z ) + v ⁡ ( Z ′ ) ] + ν ⁡ [ c ⁡ ( Z ) + c ⁡ ( Z ′ ) ] , \ell(Z,Z^{\prime})=\ \lambda s(Z,Z^{\prime})+\mu[v(Z)+v(Z^{\prime})]+\nu[c(Z)+c(Z^{\prime})], (6) where λ \lambda , μ \mu and ν \nu are hyper-parameters controlling the importance of each term in the loss. In our experiments, we set ν = 1 \nu=1 and perform a grid search on the values of λ \lambda and μ \mu with the base condition λ = μ > 1 \lambda=\mu>1 . The overall objective function taken on all images over an unlabelled dataset 𝒟 \mathcal{D} is given by: ℒ = ∑ I ∈ 𝒟 ∑ t , t ′ ∼ 𝒯 ℓ ⁡ ( Z I , Z ′ I ) , \mathcal{L}=\sum_{I\in\mathcal{D}}\sum_{t,t^{\prime}\sim\mathcal{T}}\ell(Z^{I},Z^{\prime I}), (7) where Z I Z^{I} and Z ′ I Z^{\prime I} are the batches of embeddings corresponding to the batch of images I I transformed by t t and t ′ t^{\prime} . The objective is minimized for several epochs, over the encoder parameters θ \theta and expander parameters ϕ \phi . We illustrate the architecture and loss function of VICReg in Figure 1 .

### 4.2 Implementation details

Implementation details for pretraining with VICReg on the 1000-classes ImagetNet dataset without labels are as follows. Coefficients λ \lambda and μ \mu are 25 25 and ν \nu is 1 in Eq. ( 6 ), and ϵ \epsilon is 0.0001 0.0001 in Eq. ( 1 ). We give more details on how we choose the coefficients of the loss function in Appendix D.4 . The encoder network f θ f_{\theta} is a standard ResNet-50 backbone He et al. (2016) with 2048 output units. The expander h ϕ h_{\phi} is composed of two fully-connected layers with batch normalization (BN) Ioffe & Szegedy (2015) and ReLU, and a third linear layer. The sizes of all 3 layers were set to 8192. As with Barlow Twins, performance improves when the size of the expander layers is larger than the dimension of the representation. The impact of the expander dimension on performance is studied in Appendix D . The training protocol follows those of BYOL and Barlow Twins: LARS optimizer You et al. (2017) ; Goyal et al. (2017) run for 1000 epochs with a weight decay of 10 − 6 10^{-6} and a learning rate l ​ r = b ​ a ​ t ​ c ​ h ​ _ ​ s ​ i ​ z ​ e / 256 × b ​ a ​ s ​ e ​ _ ​ l ​ r lr=batch\_size/256\times base\_lr , where b ​ a ​ t ​ c ​ h ​ _ ​ s ​ i ​ z ​ e batch\_size is set to 2048 2048 by default and b ​ a ​ s ​ e ​ _ ​ l ​ r base\_lr is a base learning rate set to 0.2 0.2 . The learning rate follows a cosine decay schedule Loshchilov & Hutter (2017) , starting from 0 0 with 10 10 warmup epochs and with final value of 0.002 0.002 .

## 5 Results

In this section, we evaluate the representations obtained after self-supervised pretraining of a ResNet-50 He et al. (2016) backbone with VICReg during 1000 epochs, on the training set of ImageNet, using the training protocol described in section 4 . We also pretrain on pairs of image and text data and evaluate on retrieval tasks on the MS-COCO dataset.

### 5.1 Evaluation on ImageNet

Following the ImageNet Deng et al. (2009) linear evaluation protocol, we train a linear classifier on top of the frozen representations of the ResNet-50 backbone pretrained with VICReg. We also evaluate the performance of the backbone when fine-tuned with a linear classifier on a subset of ImageNet’s training set using 1% or 10% of the labels, using the split of Chen et al. (2020a) . We give implementation details about the optimization procedure for these tasks in Appendix C . We have applied the training procedure described in section 4 with three different random initialization. The numbers reported in Table 1 for VICReg are the mean scores, and we have observed that the difference between worse and best run is lower than 0.1% accuracy for linear classification, which shows that VICReg is a very stable algorithm. Lack of time has prevented us from doing the same for the semi-supervised classification experiments, and the experiments of section 5.2 and 6 , but we expect similar conclusion to hold. We compare in Table 1 our results on both tasks against other methods on the validation set of ImageNet. The performance of VICReg is on par with the state of the art without using the negative pairs of SimCLR, the clusters of SwAV, the bag-of-words representations of OBoW, or any asymmetric networks architectural tricks such as the momentum encoder of BYOL and the stop-gradient operation of SimSiam. The performance is comparable to that of Barlow Twins, which shows that VICReg’s more explicit way of constraining the variance and comparing views has the same power than maximizing cross-correlations between pairs of twin dimensions. The main advantage of VICReg is the modularity of its objective function and the applicability to multi-modal setups.

### 5.2 Transfer to other downstream tasks

Following the setup from Misra & Maaten (2020) , we train a linear classifier on top of the frozen representations learnt by our pretrained ResNet-50 backbone on a variety of different datasets: the Places205 Zhou et al. (2014) scene classification dataset, the VOC07 Everingham et al. (2010) multi-label image classification dataset and the iNaturalist2018 Horn et al. (2018) fine-grained image classification dataset. We then evaluate the quality of the representations by transferring to other vision tasks including VOC07+12 Everingham et al. (2010) object detection using Faster R-CNN Ren et al. (2015) with a R50-C4 backbone, and COCO Lin et al. (2014) instance segmentation using Mask-R-CNN He et al. (2017) with a R50-FPN backbone. We report the performance in Table 2 , VICReg performs on par with most concurrent methods, and better than Barlow Twins, across all classification tasks, but is slightly behind the top-3 on detection tasks.

### 5.3 Multi-modal pretraining on MS-COCO

One fundamental difference of VICReg compared to Barlow Twins is the way the branches are regularized. In VICReg, both branches are regularized independently, as the covariance term is applied on each branch separately, which works better in the scenarios where the branches are completely different, have different types of architecture and process different types of data. Indeed, the statistics of the output of the two branches can be very different, and the amount of regularization required for each may vary a lot. In Barlow Twins, the regularization is applied on the cross-correlation matrix, which favors the scenarios where the branches produce outputs with similar statistics. We demonstrate the capabilities of VICReg in a multi-modal experiment where we pretrain on pairs of images and corresponding captions on the MS-COCO dataset. We regularize each branch with a different coefficient, which is not possible with Barlow Twins, and we show that VICReg outperforms Barlow Twins on image and text retrieval downstream tasks. Table 3 reports the performance of VICReg against the contrastive loss proposed by VSE++ Faghri et al. (2018) , and against Barlow Twins, in the identical setting proposed in Faghri et al. (2018) . VICReg outperforms the two by a significant margin.

## 6 Analysis

In this section we study how the different components of our method contribute to its performance, as well as how they interact with components from other self-supervised methods. We also evaluate different scenarios where the branches have different weights and architecture. All reported results are obtained on the linear evaluation protocol, using a ResNet-50 backbone if not mentioned otherwise, and 100 epochs of pretraining, which gives results consistent with those obtained with 1000 epochs of pretraining. The optimization setting used for each experiment is described in Appendix C .

Asymmetric networks. We study the impact of different components used in asymmetric architectures and the effects of adding variance and covariance regularization, in terms of performance and training stability. Starting from a simple symmetric architecture with an encoder and an expander without batch normalization, which correspond to VICReg without batch normalization in the expander, we progressively add batch normalization in the inner layers of the expander, a predictor, a stop-gradient operation and a momentum encoder. We use the training protocol and architecture of SimSiam Chen & He (2020) when a stop-gradient is used and the training protocol and architecture of BYOL Grill et al. (2020) when a momentum encoder is used. The predictor as used in SimSiam and BYOL is a learnable module g ψ g_{\psi} that predicts the embedding of a view given the embedding of the other view of the same image. If z z and z ′ z^{\prime} are the embeddings of two views of an image, then p = g ψ ​ ( z ) p=g_{\psi}(z) and p ′ = g ψ ​ ( z ′ ) p^{\prime}=g_{\psi}(z^{\prime}) are the predictions of each view. The invariance loss function of Eq. ( 5 ) is now computed between a batch of embeddings Z = [ z 1 , … , z n ] Z=[z_{1},\dots,z_{n}] and the corresponding batch of predictions P = [ p 1 ′ , … , p n ′ ] P=[p_{1}^{\prime},\dots,p_{n}^{\prime}] , then symmetrized: s ⁡ ( Z , Z ′ , P , P ′ ) = 1 2 ​ n ​ ∑ i D ⁡ ( z i − p i ′ ) + 1 2 ​ n ​ ∑ i D ⁡ ( z i ′ − p i ) , s(Z,Z^{\prime},P,P^{\prime})=\frac{1}{2n}\sum_{i}D(z_{i}-p^{\prime}_{i})+\frac{1}{2n}\sum_{i}D(z^{\prime}_{i}-p_{i}), (8) where D D is a distance function that depends on the method used. BYOL uses the mean square error between l 2 l_{2} -normalized vectors, SimSiam uses the negative cosine similarity loss and VICReg uses the mean square error without l 2 l_{2} -normalization. The variance and covariance terms are regularizing the output Z Z and Z ′ Z^{\prime} of the expander, which we empirically found to work better than regularizing the output of the predictor. We compare different settings in Table 4 , based on the default data augmentation, optimization and architecture settings of the original BYOL, SimSiam and VICReg methods. In all settings, the absence of BN indicates that BN is also removed in the predictor when one is used.

We analyse first the impact of variance regularization (VR) in the different settings. When using VR, adding a predictor (PR) to VICReg does not lead to a significant change of the performance, which indicates that PR is redundant with VR. In comparison, without VR, the representations collapse, and both stop-gradient (SG) and PR are necessary. Batch normalization in the inner layers of the expander (BN) in VICReg leads to a 1.0% increase in the performance, which is not a big improvement considering that SG and PR without BN is performing very poorly at 35.1%.

Finally, incorporating VR with SG or ME further improves the performance by small margins of respectively 0.2% and 0.9%, which might be explained by the fact that these architectural tricks that prevent collapse are not perfectly maintaining the variance of the representations, i.e. very slow collapse is happening with these methods. We explain this intuition by studying the evolution of the standard deviation of the representations during pretraining for BYOL and SimSiam in Appendix D . We then analyse the impact of adding additional covariance regularization (CR) in the different settings, along with variance regularization. We found that optimization with SG and CR is hard, even if our analysis of the average correlation coefficient of the representations during pretraining in Appendix D shows that both fulfill the same objective.

The performance of BYOL and SimSiam slightly drops compared to VR only, except when PR is removed, where SG becomes useless. BN is still useful and improves the performance by 1.3%. Finally with CR, PR does not harm the performance and even improves it by a very small margin. VICReg+PR with 1000 epochs of pretraining exactly matches the score of VICReg (73.2% on linear classification).