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: This paper attempts to investigate the reasons behind why non-contrastive SSL methods such as BYOL and Sim-Siam do not collapse to trivial solutions, how they learn representations that are related to the data distribution and augmentation process, and how these reduce the sample-complexity of downstream tasks. It heavily draws upon DirectPred (Tian et al) and generalizes the method for directly setting the weights of the predictor network via a parameter that seems to be tied to the strength and distinguishability of the learnt features.   Tian et al: https://arxiv.org/abs/2102.06810