Towards Demystifying Representation Learning with Non-contrastive Self-supervision

Non-contrastive methods of self-supervised learning ( such as BYOL and SimSiam ) learn representations by minimizing the distance between two views of the same image . These approaches have achieved remarkable performance in practice , but it is not well understood how the representation is learned based on the augmentation process . Tian et al . ( 2021 ) explained why the representation does not collapse to zero and proposed DirectPred that sets the predictor directly . In our work , we analyze a generalized version of DirectPred , called DirectSet ( α ) . We show that in a simple linear network , DirectSet ( α ) provably learns a desirable projection matrix and also reduces the sample complexity on downstream tasks . Our analysis suggests that weight decay acts as an implicit threshold that discard the features with high variance under augmentation , and keep the features with low variance . Inspired by our theory , we simplify DirectPred by removing the expensive eigen-decomposition step . On CIFAR-10 , CIFAR-100 , STL-10 and ImageNet , DirectCopy , our simpler and more computationally efficient algorithm , rivals or even outperforms DirectPred . 1 INTRODUCTION . Self-supervised learning recently emerges as a promising direction to learn representations without manual labels . While contrastive learning ( Oord et al. , 2018 ; Tian et al. , 2019 ; Bachman et al. , 2019 ; He et al. , 2020 ; Chen et al. , 2020a ) minimizes the distance of representation between positive pairs , and maximizes such distances between negative pairs , recently , non-contrastive selfsupervised learning ( abbreviated as nc-SSL ) is able to learn nontrivial representation with only positive pairs , using an extra predictor and a stop-gradient operation . Furthermore , the learned representation shows comparable ( or even better ) performance for downstream tasks ( e.g. , image classification ) ( Grill et al. , 2020 ; Chen & He , 2020 ) . This brings about two fundamental questions : ( 1 ) why the learned representation does not collapse to trivial ( i.e. , constant ) solutions , and ( 2 ) without negative pairs , what representation nc-SSL learns from the training and how the learned representation reduces the sample complexity in downstream tasks . While many theoretical results on contrastive SSL ( Arora et al. , 2019 ; Lee et al. , 2020 ; Tosh et al. , 2020 ; Wen & Li , 2021 ) do exist , similar study on nc-SSL has been very rare . As one of the first work towards this direction , Tian et al . ( 2021 ) show that while the global optimum of the noncontrastive loss is indeed a trivial one , following gradient direction in nc-SSL , one can find a local optimum that admits a nontrivial representation . Based on their theoretical findings on gradientbased methods , they proposed a new approach , DirectPred , that directly sets the predictor using the eigen-decomposition of the correlation matrix of input before the predictor , rather than updating it with gradient methods . As a method for nc-SSL , DirectPred shows comparable or better performance in multiple datasets , including CIFAR-10 ( Krizhevsky et al. , 2009 ) , STL-10 ( Coates et al. , 2011 ) and ImageNet ( Deng et al. , 2009 ) , compared to BYOL ( Grill et al. , 2020 ) and SimSiam ( Chen & He , 2020 ) that optimize the predictor using gradient descent . While Tian et al . ( 2021 ) address the first question , i.e. , why the learned representation does not collapse , they do not address the second question , i.e. , what representation is learned in nc-SSL and how the learned representation is related to the data distribution and augmentation process and in turn whether it reduces the sample complexity in downstream tasks . Main Contributions . In this paper , we make a first attempt towards the second question , by studying a family of algorithms named DirectSet ( α ) , in which the DirectPred algorithm proposed by Tian et al . ( 2021 ) is a special case with α = 1/2 . Our contribution is two-folds : First , we perform a theoretical analysis on DirectSet ( α ) with linear networks . Our analysis shows that there exists an implicit threshold , determined by weight decay parameter η , that governs which features are learned and which are discarded . More specifically , the threshold is applied to the variance of the feature across different data augmentations ( or “ views ” ) of the same instance : nuisance features ( features with high variances under augmentation ) are discarded , while invariant features ( i.e. , with low variances ) are kept . We further make a formal statement on the sample complexity of the learning process and performance guarantees of the downstream tasks , in the linear setting similar to Tian et al . ( 2021 ) . To our knowledge , this is the first sample complexity result in nc-SSL . Second , we show that DirectCopy , a special case of DirectSet ( α ) when α = 1 , performs comparably with ( or even outperforms ) DirectPred in downstream tasks in CIFAR-10 , CIFAR-100 , STL-10 and ImageNet . In DirectCopy , the predictor can be set without the expensive eigen-decomposition operation , which makes DirectCopy much simpler and more efficient than DirectPred . Related works . In nc-SSL , different techniques are proposed to avoid collapsing . BYOL and SimSiam use an extra predictor and stop gradient operation . Beyond these , BatchNorm ( including its variants ( Richemond et al. , 2020 ) ) , de-correlation ( Zbontar et al. , 2021 ; Bardes et al. , 2021 ; Hua et al. , 2021 ) , whitening ( Ermolov et al. , 2021 ) , centering ( Caron et al. , 2021 ) , and online clustering ( Caron et al. , 2020 ) are all effective ways to enforce implicit contrastive constraints among samples for collapsing prevention . We study BYOL and SimSiam as representative nc-SSL methods . Organization . The paper is organized as follows . Section 2-3 introduce DirectSet ( α ) , prove it learns a projection matrix onto the invariant features , and the learned representation reduces sample complexity in downstream tasks . Section 4 demonstrates that DirectCopy achieves comparable or even better performance than the original DirectPred algorithm in various datasets , and Section 5 shows ablation experiments . Finally , limitation and future works are discussed in Section 6-7 . 2 PRELIMINARIES . 2.1 NOTATIONS . We use Id to denote the d×d identity matrix and simply write I when the dimension is clear . For any linear subspace S in Rd , we use PS ∈ Rd×d to denote the projection matrix on S. More precisely , the projection matrix PS equals UU > , where the columns of U constitute a set of orthonormal bases for subspace S. We useN ( µ , Σ ) to denote the Gaussian distribution with mean µ and covariance Σ . We use ‖·‖ to denote spectral norm for a matrix , or ` 2 norm for a vector and use ‖·‖F to denote Frobenius norm for a matrix . For a real symmetric matrix A ∈ Rd×d whose eigen-decomposition is∑d i=1 λiuiu > i , we use |A| to denote ∑d i=1 |λi|uiu > i . If A is also positive semi-definite , we use Aα to denote ∑d i=1 λ α i uiu > i for any positive α ∈ R. 2.2 DIRECTSET ( α ) AND DIRECTCOPY In nc-SSL , recent methods as BYOL ( Grill et al. , 2020 ) and SimSiam ( Chen & He , 2020 ) employ a dual pair of Siamese networks ( Bromley et al. , 1994 ) : one side is a composition of an online network ( including a projector ) and a predictor network , the other side is a target network ( see Figure 1 for a simple example ) . The target network has the same architecture as the online network , but has potentially different weights . Given an input x , two augmented views x1 , x2 are generated , and the network is trained to match the representation of x1 ( through the online network and the predictor network ) and the representation of x2 ( through the target network ) . More precisely , suppose the online network and the target network are two mappings fθ , fθa : Rd 7→ Rh and the predictor network is a mapping gθp : Rh 7→ Rh , the network is trained to minimize the following loss : L ( θ , θp , θa ) : = 1 2 Ex1 , x2 ∥∥∥∥∥ gθp ( fθ ( x1 ) ) ∥∥gθp ( fθ ( x1 ) ) ∥∥ − StopGrad ( fθa ( x2 ) ‖fθa ( x2 ) ‖ ) ∥∥∥∥∥ 2 . In BYOL and SimSiam , the online network and the target network are trained by running gradient methods on L. The target network is not trained by gradient methods ; instead , it is directly set with the weights in the online network ( Chen & He , 2020 ) or an exponential moving average ( EMA ) of the online network ( Grill et al. , 2020 ; He et al. , 2020 ; Chen et al. , 2020b ) . The DirectSet ( α ) algorithm , as shown in Figure 1 , directly sets the predictor based on the correlation matrix F of the predictor inputs : Wp = Fα ‖Fα‖ + I , where F = Ex1fθ ( x1 ) fθ ( x1 ) > . In practice , F is estimated by a moving average over batches . That is , F̂ = µF̂ + ( 1− µ ) EB [ fθ ( x1 ) fθ ( x1 ) > ] , where EB is the expectation over one batch . In the original DirectPred proposed by Tian et al . ( 2021 ) , α is fixed at 1/2 . To compute F̂ 1/2 , one needs to first compute the eigen-decomposition of F̂ , and then taking the root of each eigenvalue . This step of eigen-decomposition can be expensive especially when the representation dimension h is high . To avoid the eigen-decomposition step , we propose DirectCopy ( α = 1 ) , in which the predictor Wp is a direct copy of the F̂ ( with normalization and regularization ) 1 . As we shall see , DirectCopy enjoys both theoretical guarantees and strong empirical performance . 3 THEORETICAL ANALYSIS OF DIRECTSET ( α ) Deep linear networks have been widely used as a tractable theoretical model for studying nonconvex loss landscapes ( Kawaguchi , 2016 ; Du & Hu , 2019 ; Laurent & Brecht , 2018 ) and nonlinear learning dynamics ( Saxe et al. , 2013 ; 2019 ; Lampinen & Ganguli , 2018 ; Arora et al. , 2018 ) . However , most of them are for supervised learning setting . Tian et al . ( 2021 ) analyzed nc-SSL on a linear network , but did not analyze their proposed approach DirectPred . Here , we analyze the representation learning process of DirectSet ( α ) on a minimal setting where the online network fθ is a single linear layer . We also verify DirectSet ( α ) works for practical nonlinear deep models and realistic datasets . 3.1 SETUP . In this subsection , we define the network model , data distribution and simplify DirectSet ( α ) algorithm for our theoretical analysis . We consider the following network model ( see Figure 1 ) , Assumption 1 ( Linear network model ) . The online , predictor and target network are all single-layer linear network without bias , with weight matrices denoted as W , Wp , Wa ∈ Rd×d respectively . For the data distribution , we assume the input space is a direct sum of a invariant feature subspace and a nuisance feature subspace . Specifically , we assume Assumption 2 ( Data distribution ) . The input x is sampled from N ( 0 , Id ) , and its augmented view x1 , x2 are independently sampled from N ( x , σ2PB ) , where B is a ( d − r ) -dimensional subspace . We denote S as the orthogonal subspace of B in Rd . 1Computing the spectral norm of F̂ is much faster than computing the eigen-decomposition of F̂ , because the former only needs the top eigen-vector of F̂ . Table 4 shows that the spectral norm can also be replaced by Frobenius norm or no normalization , and similar performance can be achieved . In this simple data distribution , subspace S corresponds to the features that are invariant to augmentations and its orthogonal subspaceB is the nuisance subspace which the augmentation changes . We will prove that DirectSet ( α ) can learn the projection matrix onto S subspace . Note in the previous work ( Tian et al. , 2021 ) , they assumed the covariance of the augmentation distribution to be σ2I and did not study what representation is learned . For the convenience of analysis , we consider a simplified version of DirectSet ( α ) . We compute the loss function without normalizing the two representations , so the population loss is L ( W , Wa , Wp ) : = 1 2 Ex1 , x2 ‖WpWx1 − StopGrad ( Wax2 ) ‖ 2 , ( 1 ) and the empirical loss is L̂ ( W , Wp , Wa ) : = 1 2n n∑ i=1 ∥∥∥WpWx ( i ) 1 − StopGrad ( Wax ( i ) 2 ) ∥∥∥2 , ( 2 ) where x ( i ) ’ s are independently sampled from N ( 0 , I ) , and augmented views x ( i ) 1 and x ( i ) 2 are independently sampled from N ( x ( i ) , σ2PB ) . To train our model , first , we initialize W as δI with δ a positive real number . We run gradient flow or gradient descent on online network W with weight decay η , and set the the target network Wa = W. For clarity of presentation , when training on the population loss , we set Wp as ( WExxx > W > ) α = ( WW > ) α instead of ( WEx1x1x > 1 W > ) α as in practice ; when training on the empirical loss , we set Wp as ( W 1n ∑n i=1 x ( i ) [ x ( i ) ] > W > ) α . Here , we set the predictor regularization = 0 and its influence will be studied in Section 5 . In the following , DirectSet ( α ) is shown to recover the projection matrix PS with polynomial number of samples . Furthermore , given that the learned matrix is close to PS , the sample complexity on downstream tasks is reduced .

Summary: In this paper, the authors make theoretical progress on understanding non-contrastive self-supervised learning (ncSSL).  ncSSL has previously demonstrated strong empirical performance, even outperforming contrastive learning, but the theory behind it is still unclear.  In this work, the authors build off of prior analysis by [1], and showcase the role weight decay has in learning a desirable representation; it acts as a threshold that discards noisy features with high variance introduced by the data augmentation, and keeps stable features with low variance.    [1] Tian et al., Understanding Self-Supervised Learning Dynamics without Contrastive Pairs, 2019.