Understanding Dimensional Collapse in Contrastive Self-supervised Learning

1 INTRODUCTION . Self-supervised learning aims to learn useful representations of the input data without relying on human annotations . Recent advances in self-supervised visual representation learning based on joint embedding methods ( Misra & Maaten , 2020b ; He et al. , 2020 ; Chen et al. , 2020a ; Chen & He , 2020 ; Grill et al. , 2020 ; Zbontar et al. , 2021 ; Bardes et al. , 2021 ; Chen et al. , 2020b ; Dwibedi et al. , 2021 ; Li et al. , 2021 ; Misra & Maaten , 2020a ; HaoChen et al. , 2021 ; Assran et al. , 2021 ; Caron et al. , 2021 ) show that self-supervised representations have competitive performances compared with supervised ones . These methods generally aim to learn representations invariant to data augmentations by maximizing the agreement between embedding vectors from different distortions of the same images . As there are trivial solutions where the model maps all input to the same constant vector , known as the collapsing problem , various methods have been proposed to solve this problem that rely on different mechanisms . Contrastive methods like Chen et al . ( 2020a ) and He et al . ( 2016 ) define ‘ positive ’ and ‘ negative ’ sample pairs which are treated differently in the loss function . Non-contrastive methbods like Grill et al . ( 2020 ) and Chen & He ( 2020 ) use stop-gradient , and an extra predictor to prevent collapse without negative pairs ; Caron et al . ( 2018 ; 2020 ) use an additional clustering step ; and Zbontar et al . ( 2021 ) minimize the redundant information between two branches . These self-supervised learning methods are successful in preventing complete collapse whereby all representation vectors shrink into a single point . However , it has been observed empirically in noncontrastive learning methods ( Hua et al. , 2021 ; Tian et al. , 2021 ) that while embedding vectors do not completely collapse ; they collapse along certain dimensions . This is known as dimensional collapse ( Hua et al. , 2021 ) , whereby the embedding vectors only span a lower-dimensional subspace . In contrastive methods that explicitly use positive and negative pairs in the loss function , it seems intuitive to speculate that the repulsive effect of negative examples should prevent this kind of dimensional collapse and make full use of all dimensions . However , contrary to intuition , contrastive learning methods still suffer from dimensional collapse ( See Fig . 7 ) . In this work , we theoretically study the dynamics behind this phenomenon . We show there are two different mechanisms that cause collapsing : ( 1 ) along the feature direction where the variance caused by the data augmentation is larger than the variance caused by the data distribution , the weight collapses . Moreover , ( 2 ) even if the covariance of data augmentation has a smaller magnitude than the data variance along all dimensions , the weight will still collapse due to the interplay of weight matrices at different layers known as implicit regularization . This kind of collapsing happens only in networks where the network has more than one layer . Inspired by our theory , we propose a novel contrastive learning method , called DirectCLR , which directly optimizes the encoder ( i.e. , representation space ) without relying on a trainable projector . DirectCLR outperforms SimCLR with a linear trainable projector on ImageNet . We summarize our contributions as follows : • We empirically show that contrastive self-supervised learning suffers from dimensional collapse whereby all the embedding vectors fall into a lower-dimensional subspace instead of the entire available embedding space . • We showed that there are two mechanisms causing the dimensional collapse in contrastive learning : ( 1 ) strong augmentation along feature dimensions ( 2 ) implicit regularization driving models toward low-rank solutions . • We propose DirectCLR , a novel contrastive learning method that directly optimizes the representation space without relying on a trainable projector . DirectCLR outperforms SimCLR with a linear trainable projector . 2 RELATED WORKS . 2.1 SELF-SUPERVISED LEARNING METHODS . Joint embedding methods are a promising approach in self-supervised learning , whose principle is to match the embedding vectors of augmented views of a training instance . Contrastive methods ( Chen et al. , 2020a ; He et al. , 2016 ) directly compare training samples by effectively viewing each sample as its own class , typically based on the InfoNCE contrastive loss ( van den Oord et al. , 2018 ) which encourages representations from positive pairs of examples to be close in the embedding space while representations from negative pairs are pushed away from each other . In practice , contrastive methods are known to require a large number of negative samples . Non-contrastive methods do not directly rely on explicit negative samples . These include clustering-based methods ( Caron et al. , 2018 ; 2020 ) , redundancy reduction methods ( Zbontar et al. , 2021 ; Bardes et al. , 2021 ) and methods using special architecture design ( Grill et al. , 2020 ; Chen & He , 2020 ) . 2.2 THEORETICAL UNDERSTANDING OF SELF-SUPERVISED LEARNING . Although self-supervised learning models have shown success in learning useful representations and have outperformed their supervised counterpart in several downstream transfer learning benchmarks ( Chen et al. , 2020a ) , the underlying dynamics of these methods remains somewhat mysterious and poorly understood . Several theoretical works have attempted to understand it . Arora et al . ( 2019b ) ; Lee et al . ( 2020 ) ; Tosh et al . ( 2021 ) theoretically proved that the learned representations via contrastive learning are useful for downstream tasks . Tian et al . ( 2021 ) explained why non-contrastive learning methods like BYOL ( Grill et al. , 2020 ) and SimSiam ( Chen & He , 2020 ) work : the dynamics of the alignment of eigenspaces between the predictor and its input correlation matrix play a key role in preventing complete collapse . 2.3 IMPLICIT REGULARIZATION . It has been theoretically explained that gradient descent will drive adjacent matrices aligned in a linear neural network setting ( Ji & Telgarsky , 2019 ) . Under the aligned matrix assumption , Gunasekar et al . ( 2018 ) prove that gradient descent can derive minimal nuclear norm solution . Arora et al . ( 2019a ) extend this concept to the deep linear network case by theoretically and empirically demonstrating that a deep linear network can derive low-rank solutions . In general , over-parametrized neural networks tend to find flatter local minima ( Saxe et al. , 2019 ; Neyshabur et al. , 2019 ; Soudry et al. , 2018 ; Barrett & Dherin , 2021 ) . 3 DIMENSIONAL COLLAPSE . Self-supervised learning methods learn useful representation by minimizing the distances between embedding vectors from augmented images ( Figure 1a ) . On its own , this would result in a collapsed solution where the produced representation becomes constant ( Figure 1b ) . Contrastive methods prevent complete collapse via the negative term that pushes embedding vectors of different input images away from each other . In this section , we show that while they prevent complete collapse , contrastive methods still experience a dimensional collapse in which the embedding vectors occupy a lower-dimensional subspace than their dimension ( Figure 1c ) . We train a SimCLR model ( Chen et al . ( 2020a ) ) with a two-layer MLP projector . We followed the standard recipe and trained the model on ImageNet for 100 epoch . We evaluate the dimensionality by collecting the embedding vectors on the validation set . Each embedding vector has a size of d = 128 . We compute the covariance matrix C ∈ Rd×d of the embedding layer ( here z̄ : = ∑N i=1 zi/N and N is the total number of samples ) : C = 1 N N∑ i=1 ( zi − z̄ ) ( zi − z̄ ) T ( 1 ) Figure 2 shows singular value decomposition on this matrix ( C = USV T , S = diag ( σk ) ) . in sorted order and logarithmic scale ( { log ( σk ) } ) . We observe that a number of singular values collapse to zero , thus representing collapsed dimensions . 4 DIMENSIONAL COLLAPSE CAUSED BY STRONG AUGMENTATION . 4.1 LINEAR MODEL . In this section , we explain one scenario for contrastive learning to have collapsed embedding dimensions , where the augmentation surpasses the input information . We focus on a simple linear network setting . We denote the input vector as x and the augmentation is an additive noise . The network is a single linear layer with weight matrix is W . Hence , the embedding vector is z = Wx . We focus on a typical contrastive loss , InfoNCE ( van den Oord et al. , 2018 ) : L = − N∑ i=1 log exp ( −|zi − z′i|2/2 ) ∑ j 6=i exp ( −|zi − zj |2/2 ) + exp ( −|zi − z′i|2/2 ) ( 2 ) where zi and z′i are a pair of embedding vectors from the two branches , zj indicates the negative samples within the minibatch . When all zi and z′i are normalized to be unit vector , the negative distance −|zi − z′i|2/2 can be replaced by inner products zTi z′i . The model is trained with a basic stochastic gradient descent without momentum or weight decay . 4.2 GRADIENT FLOW DYNAMICS . We study the dynamics via gradient flow , i.e. , gradient descent with an infinitesimally small learning rate . Lemma 1 . The weight matrix in a linear contrastive self-supervised learning model evolves by : Ẇ = −G ( 3 ) where G = ∑ i ( gzix T i + gz′ix ′T i ) , and gzi is the gradient on the embedding vector zi ( similarly gz′i ) . This can be easily proven based on the chain rule . See proof in Appendix B.1 . For InfoNCE loss defined in Eqn 2 , the gradient of the embedding vector for each branch can be written as gzi = ∑ j 6=i αij ( zj − z′i ) + ∑ j 6=i αji ( zj − zi ) , gz′i = ∑ j 6=i αij ( z′i − zi ) ( 4 ) where { αij } are the softmax of similarity of between zi and { zj } , defined by αij = exp ( −|zi − zj |2/2 ) /Zi , αii = exp ( −|zi−z′i|2/2 ) /Zi , andZi = ∑ j 6=i exp ( −|zi−zj |2/2 ) +exp ( −|zi−z′i|2/2 ) . Hence , ∑ j αij = 1 . Since zi = Wxi , we have G = −WX ( 5 ) where X : = − ∑ i ∑ j 6=i αij ( x′i − xj ) + ∑ j 6=i αji ( xi − xj ) xTi −∑ i ( 1− αii ) ( x′i − xi ) x′i T ( 6 ) Lemma 2 . X is a difference of two PSD matrices : X = Σ̂0 − Σ̂1 ( 7 ) Here Σ̂0 = ∑ i , j αij ( xi − xj ) ( xi − xj ) T is a weighted data distribution covariance matrix and Σ̂1 = ∑ i ( 1− αii ) ( x′i − xi ) ( x′i − xi ) T is a weighted augmentation distribution covariance matrix . See proof in Appendix B.2 . Therefore , the amplitude of augmentation determines whether X is a positive definite matrix . Next , we look into the dynamics of weight matrix W given property of X . Theorem 1 . With fixed matrix X ( defined in Eqn 6 ) and strong augmentation such that X has negative eigenvalues , the weight matrix W has vanishing singular values . See proof in Appendix B.3 . Corollary 1 ( Dimensional Collapse Caused by Strong Augmentation ) . With strong augmentation , the embedding space covariance matrix becomes low-rank . The embedding space is identified by the singular value spectrum of the covariance matrix on the embedding ( Eqn . 1 ) , C = ∑ i ( zi− z̄ ) ( zi− z̄ ) T /N = ∑ iW ( xi− x̄ ) ( xi− x̄ ) TWT /N . Since W has vanishing singular values , C is also low-rank , indicating collapsed dimensions . We experimentally verified the theory in Appendix C. Our theory in this section is limited to linear network settings . For more complex nonlinear networks , the collapsing condition will still depend on “ strong augmentation ” but interpreted differently . A strong augmentation will be determined by more complicated properties of the augmentation ( higher-order statistics of augmentation , manifold property of augmentation vs. data distribution ) conditioned on the capacity of the networks .

Summary: This paper investigates the collapsing problem of contrastive learning. It attempts to attribute the collapsing phenomena to strong augmentation and implicit regularization, using simple linear network models. Based on the above analysis, the paper then propose a simple sub-vector based CL method called DirectCLR.