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 provides a new perspective on the BYOL self-supervised learning method. First, the paper introduces an upper-bound objective, BYOL', that is easier to analyze than BYOL because it is composed of two well understood losses: an alignment loss and cross-model loss. Further, it shows empirically that optimizing BYOL' is similar to optimizing BYOL. Second, the paper introduces the RAFT method which maximizes the alignment loss instead of minimizing it. The paper proves that under some assumptions, such as a linear predictor function, optimizing BYOL' is equivalent to RAFT. Based on this analysis, the paper explains why the predictor function is essential for BYOL and why it is hard to achieve convergence.