Run Away From your Teacher: a New Self-Supervised Approach Solving the Puzzle of BYOL

1 INTRODUCTION . Recently the performance gap between self-supervised learning and supervised learning has been narrowed thanks to the development of contrastive learning ( Chen et al. , 2020b ; a ; Tian et al. , 2019 ; Chen et al. , 2020b ; Sohn , 2016 ; Zhuang et al. , 2019 ; He et al. , 2020 ; Oord et al. , 2018 ; Hadsell et al. , 2006 ) . Contrastive learning distinguishes positive pairs of data from the negative . It has been shown that when the representation space is l2-normalized , i.e . a hypersphere , optimizing the contrastive loss is approximately equivalent to optimizing the alignment of positive pairs and the uniformity of the representation distribution at the same time ( Wang & Isola , 2020 ) . This equivalence conforms to our intuitive understanding . One can easily imagine a failed method when we only optimize either of the properties : aligning the positive pairs without uniformity constraint causes representation collapse , mapping different data all to the same point ; scattering the data uniformly in the representation space without aligning similar ones yields no more meaningful representation than random . The proposal of Bootstrap Your Own Latent ( BYOL ) fiercely challenges our consensus that negative samples are necessary to contrastive methods ( Grill et al. , 2020 ) . BYOL trains the model ( online network ) to predict its Mean Teacher ( moving average of the online , refer to Appendix B.2 ) on two augmented views of the same data ( Tarvainen & Valpola , 2017 ) . There is no explicit constraint on uniformity in BYOL , while the expected collapse never happens , what ’ s more , it reaches the SOTA performance on the downstream tasks . Although BYOL has been empirically proven to be an effective self-supervised learning approach , the mechanism that keeps it from collapse remains unrevealed . Without disclosing this mystery , it would be disturbing for us to adapt BYOL to other problems , let alone further improve it . Therefore solving the puzzle of BYOL is an urgent task . In this paper , we explain how BYOL works through another interpretable learning framework which leverages the MT in the exact opposite way . Based on a series of theoretical derivation and empirical approximation , we build a new self-supervised learning framework , Run Away From your Teacher ( RAFT ) , which optimizes two objectives at the same time : ( i ) minimize the representation distance between two samples from a positive pair and ( ii ) maximize the representation distance between the online and its MT . The second objective of RAFT incorporates the MT in a way exactly opposite to BYOL , and it explicitly prevents the representation collapse by encouraging the online to be different from its history ( Figure 2a ) . Moreover , we empirically show that the second objective of RAFT is a more effective and consistent regularizer for the first objective , which makes RAFT more favorable than BYOL . Finally , we solve the puzzle of BYOL by theoretically proving that BYOL is a special form of RAFT when certain conditions and approximation hold . This proof explains why collapse does not happen in BYOL , and also makes the performance of BYOL′ an approximate guarantee of the effectiveness of RAFT . The main body of the paper is organized in the same order of how we explore the properties of BYOL and establish RAFT based on them ( refer to Appendix A for more details ) . In section 3 , we investigate the phenomenon that BYOL fails to work when the predictor is removed . In section 4 , we establish two meaningful objectives out of BYOL by upper bounding . Based on that , we propose RAFT due to its stronger regularization effect and its accordance with our knowledge . In section 5 , we prove that , as a representation learning framework , BYOL is a special form of RAFT under certain achievable conditions . In summary , our contributions are listed as follows : • We present a new self-supervised learning framework RAFT that minimizes the alignment and maximizes the distance between the online network and its MT . The motivation of RAFT conforms to our understanding of balancing alignment and uniformity of the representation space , and thus could be easily extended and adapted to future problems . • We equate two seemingly opposite ways of incorporating MT in contrastive methods under certain conditions . By doing so , we unravel the puzzle of how BYOL avoids representation collapse . 2 BACKGROUND AND RELATED WORK . 2.1 TWO METRICS OPTIMIZED IN CONTRASTIVE LEARNING . Optimizing contrastive learning objective has been empirically proven to have positive correlations with the downstream task performance ( Chen et al. , 2020b ; a ; Tian et al. , 2019 ; Chen et al. , 2020b ; Sohn , 2016 ; Zhuang et al. , 2019 ; He et al. , 2020 ; Oord et al. , 2018 ) . Wang & Isola ( 2020 ) puts the contrastive learning under the context of hypersphere and formally showcases that optimizing the contrastive loss ( for preliminary of contrastive learning , refer to Appendix B.1 ) is equivalent to optimizing two metrics of the encoder network when the size of negative samplesK is sufficiently large : the alignment of the two augmented views of the same data and the uniformity of the representation population . We introduce the alignment objective and uniformity objective as follows . Definition 2.1 ( Alignment loss ) The alignment loss Lalign ( f , Ppos ) of the function f over positivepair distribution Ppos is defined as : Lalign ( f ; Ppos ) , E ( x1 , x2 ) ∼Ppos [ ‖f ( x1 ) − f ( x2 ) ‖22 ] , ( 1 ) where the positive pair ( x1 , x2 ) are two augmented views of the same input data x ∼ X , i.e . ( x1 , x2 ) = ( t1 ( x ) , t2 ( x ) ) and t1 ∼ T1 , t2 ∼ T2 are two augmentations . For the sake of simplicity , we omit Ppos and use Lalign ( f ) in the following content . Definition 2.2 ( Uniformity loss ) The loss of uniformity Luniform ( f ; X ) of the encoder function f over data distribution X is defined as Luniform ( f ; X ) , logE ( x , y ) ∼X 2 [ e−t‖f ( x ) −f ( y ) ‖ 2 2 ] , ( 2 ) where t > 0 is a fixed parameter and is empirically set to t = 2 . To note here , the vectors in the representation space are automatically l2-normalized , i.e . f ( x ) , f ( x ) /‖f ( x ) ‖2 , as we limit the representation space to a hypersphere following Wang & Isola ( 2020 ) and Grill et al . ( 2020 ) and the representation vectors in the following context are also automatically l2-normalized , unless specified otherwise . Wang & Isola ( 2020 ) has empirically demonstrated that the balance of the alignment loss and the uniformity loss is necessary when learning representations through contrastive method . The rationale behind it is straightforward : Lalign provides the motive power that concentrates the similar data , and Luniform prevents it from mapping all the data to the same meaningless point . 2.2 BYOL : BIZARRE ALTERNATIVE OF CONTRASTIVE . A recently proposed self-supervised representation learning algorithm BYOL hugely challenges the common understanding , that the alignment should be balanced by negative samples during the contrastive learning . It establishes two networks , online and target , approaching to each other during training . The online is trained to predict the target ’ s representations and the target is the Exponential Moving Average ( EMA ) of the parameters of the online . The loss of BYOL at every iteration could be written as LBYOL , E ( x , t1 , t2 ) ∼ ( X , T1 , T2 ) [ ∥∥qw ( fθ ( t1 ( x ) ) ) − fξ ( t2 ( x ) ) ∥∥22 ] , ( 3 ) where two vectors in representation space are automatically l2-normalized . fθ is the online encoder network parameterized by θ and qw is the predictor network parameterized by w. x ∼ X is the input sampled from the data distribution X , and t1 ( x ) , t2 ( x ) are two augmented views of x where t1 ∼ T1 , t2 ∼ T2 are two data augmentations . The target network fξ is of the same architecture as fθ and is updated by EMA with τ controlling to what degree the target network preserves its history ξ ← τξ + ( 1− τ ) θ . ( 4 ) From the scheme of BYOL training , it seems like there is no constraint on the uniformity , and thus most frequently asked question about BYOL is how it prevents the representation collapse . Theoretically , we would expect that when the final convergence of the online and target is reached , LBYOL degenerates to Lalign and therefore causes representation collapse , while this speculation never happens in reality . Despite the perfect SOTA performance of BYOL , there is one inconsistency not to be neglected : it fails with representation collapse when the predictor is removed , which means qw ( x ) = x for any given x . This inconsistent behavior of BYOL weakens its reliability and further poses questions on future adaptation of the algorithm . The motivation of understanding and even solving this inconsistency is the start point of this paper . 3 ON-AND-OFF BYOL : FAILURE WITHOUT PREDICTOR . We start by presenting a dissatisfactory property of BYOL : its success heavily relies on the existence of the predictor qw . The experimental setup of this paper is listed in Appendix C. The performance of BYOL original model , whose predictor qw is a two-layer MLP with batch normalization , evaluated on the linear evaluation protocol ( Kolesnikov et al. , 2019 ; Kornblith et al. , 2019 ; Chen et al. , 2020a ; He et al. , 2020 ; Grill et al. , 2020 ) reaches 68.08 ± 0.84 % . When the predictor is removed , the performance degenerates to 20.92±1.29 % , which is even lower than the random baseline ’ s 42.74± 0.41 % . We examine the speculation that the performance drop is caused by the representation collapse both visually ( refer to Appendix F.1 ) and numerically . Inspired by Wang & Isola ( 2020 ) , we use Luniform ( fθ ; X ) to evaluate to what degree the representations are spread on the hypersphere and Lalign ( qw ◦ fθ ) to evaluate how the similar samples are aligned in the representation space . The results in Table 1 show that with the predictor , BYOL optimizes the uniformity of the representation distribution . On the contrary , when taken away the predictor , the alignment of two augmented views is overly optimized and the uniformity of the representation deteriorates ( Figure 4 ) , therefore we conclude the predictor is essential to the collapse prevention in BYOL . One reasonable follow-up explanation on the efficacy of the predictor may consider its specially designed architecture or some good properties brought by the weight initialization , which makes it hard to understand the mechanism behind it . Fortunately , after replacing the current predictor , two-layer MLP with batch normalization ( Ioffe & Szegedy , 2015 ) , with different network architectures and weight initializations , we find that there is no significant change either on linear evaluation protocol or on the model behavior during training ( Table 1 , for detailed training trajectory , refer to Figure 4 ) . We first replace the complex structure with linear mapping qw ( · ) = W ( · ) . This replacement provides a naive solution to representation collapse : W = I , while it never converges to this apparent collapse . Surprisingly enough when we go harsher on this linear predictor by initializing W with the apparent collapse solution I , the model itself seems to have a self-recovering mechanism even though it starts off at a poor position : the loss quickly approaches to 0 and the uniformity deteriorates for 10-20 epochs and suddenly it deflects from the collapse and keeps on the right track . We conduct a theoretical proof that a randomly initialized linear predictor prevents the ( more strict form of ) representation collapse by creating infinite non-trivial solutions when the convergence is achieved ( refer to Appendix I ) , while we fail to correlate the consistently optimized uniformity with the presence of the predictor , which indicates that a deeper rationale needs to be found .

Summary: the paper aims to explain the success of BYOL, a recently proposed contrastive method that mysteriously avoids the trivial constant solution without requiring negative samples. The paper proposes a new loss named RAFT. Compared to BYOL, RAFT is more general since it subsumes a variation of BYOL as its special case, and contains a cross-model term to be maximized which regularizes the alignment loss and encourages the online encoder to "run away" from the mean teacher.