Year: 2020

Abstract:
Recently, a newly proposed self-supervised framework Bootstrap Your Own Latent (BYOL) seriously challenges the necessity of negative samples in contrastive learning frameworks. BYOL works like a charm despite the fact that it discards the negative samples completely and there is no measure to prevent collapse in its training objective. In this paper, we suggest understanding BYOL from the view of our newly proposed interpretable self-supervised learning framework, Run Away From your Teacher (RAFT). RAFT optimizes two objectives at the same time: (i) aligning two views of the same data to similar representations and (ii) running away from the model’s Mean Teacher (MT, the exponential moving average of the history models) instead of BYOL’s running towards it. The second term of RAFT explicitly prevents the representation collapse and thus makes RAFT a more conceptually reliable framework. We provide basic benchmarks of RAFT on CIFAR10 to validate the effectiveness of our method. Furthermore, we prove that BYOL is equivalent to RAFT under certain conditions, providing solid reasoning for BYOL’s counter-intuitive success.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND AND RELATED WORK

heading: 2.1 TWO METRICS OPTIMIZED IN CONTRASTIVE LEARNING
text: 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.

heading: 2.2 BYOL: BIZARRE ALTERNATIVE OF CONTRASTIVE
text: 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.

heading: 3 ON-AND-OFF BYOL: FAILURE WITHOUT PREDICTOR
text: 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.

heading: 4 RUN AWAY FROM YOUR TEACHER: MORE EFFECTIVE REGULARIZER

heading: 4.1 DISENTANGLE THE BYOL LOSS BY UPPER BOUNDING
text: Analyzing LBYOL is hard, since it only has one single mean squared error term and there are many factors entangled within it, e.g., two augmented views of the same data, predictor, and the EMA updating rule. Inspired by the Bias-Variance decomposition on squared loss (Geman et al., 1992), we extract the alignment loss by subtracting and adding the same term qw(fθ(t2(x))) and further yield the upper bound of LBYOL. For details, please refer to Appendix G. Definition 4.1 (Cross-model loss) The cross-model loss Lcross-model(f, g;X ) of the function f and g over the data distribution X is defined as
Lcross-model(f, g;X ) , Ex∼X [∥∥f(x)− g(x)∥∥2
2
] . (5)
Definition 4.2 (BYOL′ loss) The BYOL′ loss LBYOL′ is defined as
LBYOL′ , αLalign(qw ◦ fθ;Ppos) + βLcross-model(qw ◦ fθ, fξ;X2) (6)
where α, β > 0 are constants, Ppos is defined in Eq. 1 and X2 = T2(X ) is the distribution of the augmented data. For the sake of simplicity, we use Lalign(qw ◦ fθ) to denote Lalign(qw ◦ fθ;Ppos) in the following content. For the sake of symmetry, we use Lcross-model(qw ◦ fθ) to denote (1/2)[Lcross-model(qw ◦ fθ, fξ,X1) + Lcross-model(qw ◦ fθ, fξ,X2)] to compute the cross-model loss.
Theorem 4.1 (LBYOL′ is an upper bound of LBYOL) LBYOL′ is an upper bound of LBYOL if we ignore the scalar multiplication. Concretely speaking, for any given constants α, β > 0, we have
LBYOL ≤ ( 1
α +
1 β )LBYOL′ . (7)
Proof Please refer to Appendix G.
Ideally, minimizing LBYOL′ would yield similar performance as minimizing LBYOL. We exemplify the legitimacy of LBYOL′ by setting (α, β) = (1, 1). In Table 1, the performance of BYOL and BYOL′ are close to each other with respect to three metrics: alignment, uniformity, and downstream linear evaluation protocol, regardless of the form of predictors. When the predictor is linear mapping, the performance differences between them are subtle. Besides, when the predictor is removed, the representation collapse also happens to BYOL′. So we conclude that optimizing LBYOL′ is almost equivalent to LBYOL. In spite of the performance similarity, LBYOL′ is of a more disentangled form than LBYOL and therefore we focus on studying the former instead of the latter. The new objective consists of two terms: the first term Lalign minimizes the representation distance between samples from a positive pair and has already been shown crucial to the successful contrastive methods (Wang & Isola, 2020). Intuitively, it provides the motive power to concentrate similar data in the representation space. Based on the form of BYOL′, we conclude that MT is used to regularize the alignment loss. This perspective of two terms regularizing each other is crucial to our analysis and improvement of the original BYOL framework. Understanding why BYOL works without collapse is approximately equivalent to understanding how minimizing Lcross-model(qw ◦ fθ, fξ) effectively regularizes the alignment loss, or even actively optimizes the uniformity.

heading: 4.2 RAFT: RUN AWAY FROM YOUR TEACHER
text: The major difficulty of correlating Lcross-model with Luniform is that their optimization intentions are not only irrelevant, but somewhat opposite. Minimizing the cross-model loss asks the network to produce close representations for certain inputs, while optimizing the uniformity loss requires it to produce varying representations. The disparity residing in the form pushes us to question the original motivation of BYOL: do we really want the online network to approach to the Mean Teacher? To test our suspicion, we minimize [Lalign(qw ◦ fθ) − Lcross-model(qw ◦ fθ, fξ)] instead of [Lalign(qw ◦ fθ) + Lcross-model(qw ◦ fθ, fξ)], and we find it works as well. This bizarre phenomenon will be explained in Section 5. Removing the predictor, we observe that although minimizing [Lalign(fθ)−Lcross-model(fθ, fξ)] fails to yield better representation than the random baseline,
it prevents the overly-optimized alignment loss, i.e. it works as an effective regularizer for the alignment loss, while minimizing Lcross-model(fθ, fξ) does not. Based on the conclusion above and law of Occam’s Razor, we propose a new self-supervised learning framework, Run Away From your Teacher (RAFT), which optimizes two learning objectives simultaneously: (i) minimize the alignment loss of two samples from a positive pair and (ii) maximize the distance between the online network and its MT (refer to Figure 1 and Algorithm 1). Definition 4.3 (RAFT loss) The RAFT loss LRAFT is defined as
LRAFT , αLalign(qw ◦ fθ;Ppos)− βLcross-model(qw ◦ fθ, fξ;X2), (8)
where α, β > 0 are constants and other components follows the Definition 4.2.
Compared to BYOL and BYOL′, RAFT better conforms to our knowledge and is a conceptually non-collapsing algorithm. There has been a lot of work demonstrating that weight averaging is roughly equal to sample averaging (Tarvainen & Valpola, 2017), thus if two samples’ representations are close to each other at the beginning and their initial updating directions are opposite, then RAFT consistently separates them in the representation space. All the forms of loss terms could be classified into three categories: uniformity optimizer, effective regularizer for alignment loss, and others (refer to Figure 2b). According to our experiments, when the predictor is removed, running away from MT remains an effective regularizer for the alignment loss while BYOL’s running towards MT fails to do so, thus RAFT is of more unified and consistent form. In summary, our proposed learning framework RAFT is completely based on the intention of solving the inconsistency of the predictor in BYOL, and it’s better than BYOL in threefold:
• Consistency. Compared to BYOL, our newly proposed method has an effective regularizer for the alignment loss regardless of the presence of predictor.
• Interpretability. Mean teacher uses the technique of weight averaging and thus could be considered as an approximate ensemble of the previous versions of the model. Running away from the mean teacher intuitively encourages the diversity of the representation, which is positively correlated to the uniformity.
• Disentanglement. The learning objective is decoupled into aligning two augmented views and running away from the mean teacher, and hence could be independently studied.
We will discuss the relationship between RAFT and BYOL′ in the next section, and we find BYOL′ is a special form of RAFT under certain conditions, which makes the performance of BYOL′ a guarantee of the effectiveness of RAFT. We provide benchmarks of alignment, uniformity, and downstream linear evaluation performance on CIFAR10 (Table 3). We discover that balancing the alignment loss and the cross-model loss is not an easy job with the predictor taken away. The imbalance between the alignment loss and the cross-model loss would lead to representation collapse or over-regularized alignment where every data is randomly projected. One interesting research direction is to study the efficacy of the predictor. The reason why it helps the two terms to achieve an equilibrium is left to be answered.

heading: 5 UNDERSTANDING BYOL VIA RAFT
text: In Section 4.1 we derive an upper bound LBYOL′ of LBYOL and explicitly extract two terms Lalign and Lcross-model. In BYOL′, two terms are simultaneously minimized, while in RAFT, we minimize Lalign but maximize Lcross-model instead. To clearly distinguish the difference between the two objectives, we rewrite them as following:
LBYOL′ = αLalign(qw ◦ fθ) + βLcross-model(qw ◦ fθ, fξ), (9) LRAFT = αLalign(qw ◦ fθ)− βLcross-model(qw ◦ fθ, fξ), (10)
where α, β > 0 are constants.
In form, LRAFT and LBYOL′ seem to evolve in opposite optimizing direction on the second term, but the empirical study has shown that both of them work. How can two opposite optimization goals produce similar effect? Since RAFT is a conceptually working method, we analyze the mechanism of BYOL′ by establishing the equivalence between the parameters of BYOL′ and RAFT under mild conditions. Theorem 5.1 (One-to-one correspondence between BYOL′ and RAFT) There is a one-to-one correspondence between parameter trajectories of BYOL′ and RAFT when the following three conditions hold:
i. the representation space is a hypersphere;
ii. the predictor is a linear transformation, i.e. qw(·) = W (·);
iii. only the tangential component of the gradient on the hypersphere is preserved.
Proof We prove the theorem by construction. For the detail, please refer to Appendix H.
Remark The third condition conforms to the property of the hypersphere representation space and is easy to achieve. One can preserve only the tangential gradient by slightly modifying the loss. For example, suppose the representation of the MT is z̄ and the representation of the input is z which are both normalized, the cross-model loss ∥∥z̄ − z∥∥2 2 can be revised as ∥∥z̄ − λz∥∥2 2 /λ, where λ = sg(〈z, z̄〉) stands for stopping gradient of the inner product 〈z, z̄〉. Our experiments in Table 1 demonstrates that the condition of the tangential component of the gradient doesn’t turn any of the algorithms including BYOL, BYOL′ and RAFT into a collapsed one.
In Theorem 1, we show that optimizing LBYOL′ with initial parameters (θ(0),W (0)) is equivalent to optimizing LRAFT with initial parameters (θ(0),−W (0)) when the aforementioned three conditions are satisfied. This equivalence demonstrates that the final encoder network fθ and fθ′ equal to each other. Therefore we conclude that, as representation learning framework, BYOL′ is equivalent to our newly proposed RAFT.
From a geometric point of view, the optimization process is the data points moving in the representation space under the guidance of the training loss. The loss function measures the potential energy of the parameters, and the gradient with regard to the data points is the motive force. If the representation space is a hypersphere as in BYOL, then the tangential force, i.e. the tangential component of the gradient, is the only key to scattering or concentrating the data points in the representation space. By the central symmetry of the hypersphere, clockwise and counterclockwise moving directions are equivalent to some extent, for example, pushing a point by π/2 and pulling it by π/2 on the 2-dimensional sphere causes the same effect.
The equivalence between BYOL′ and RAFT offers us a direct way to understand some strange phenomena we observe which are also reported in the original BYOL paper. Firstly, the non-collapse of BYOL is explained, since the RAFT is an intuitively and practically working algorithm. The equivalence of BYOL′ and RAFT when predictor is linear helps us understand why BYOL is an effective self-supervised learning algorithm. It also explains our initial question why BYOL fails to avoid representation collapse without the predictor: removing the predictor means fixing W = I , which breaks the RAFT’s designing principle of running away from the MT. Secondly, though the BYOL’s optimization procedure is of the form that two models approaching to each other, there has been no report of convergence in the original paper. The established equivalence perfectly explains it. RAFT incorporates the MT in an extremely dynamic way since it continuously varies from the history models, thus there would be no convergence of the data points. So does the parameters.

heading: 6 CONCLUSION AND FUTURE WORK
text: In this paper, we address the problem of why the newly proposed self-supervised learning framework Bootstrap Your Own Latent (BYOL) works without negative samples. By decomposing, upper bounding and approximating the original loss of BYOL, we establish another interpretable selfsupervised learning method, Run Away From your Teacher. We show that RAFT contains an explicit term that prevents the representation collapse and we also empirically validate the effectiveness of RAFT. By constructing a one-to-one correspondence from RAFT to BYOL′ (variant of BYOL), we successfully explain the mechanism behind BYOL that makes it work and therefore implies the huge potential of our proposed RAFT. Based on the observation and the conclusion, here we have several suggestions for future work:
Theoretical guarantees of RAFT. Though we have intuitively explained why running away from the MT is an effective regularizer, we don’t provide theoretical guarantees why optimizing RAFT would be favorable with respect to the representation learning. In future, one can try to relate RAFT to the theory of Mutual Information (MI) maximization (Belghazi et al., 2018; Hjelm et al., 2018; Tschannen et al., 2019), as the training objective of contrastive learning InfoNCE has been proven to be a lower bound of the MI (Poole et al., 2019). One detail should be noticed when attempting to correlate RAFT with MI maximization. Even though RAFT is an effective regularizer, it fails to yield good-quality representations when the predictor is removed, thus any theoretical proof on the effectiveness of RAFT should well explain the mechanism behind this extra predictor.
On the efficacy of the predictor. It has become a popular and almost standardized method to add an extra MLP on top of the network in contrastive learning methods (Chen et al., 2020a;b; Grill et al., 2020), while most of the work adopts this method as a special trick without considering the effect this MLP brings to the algorithm. In this paper, however, we find that this extra MLP may bring some unexpected properties to the original training objective: although the representations are optimized by disparate motivations (in our paper, BYOL′ running towards MT and RAFT running away from MT), the encoder network is trained to be exactly the same. This observation indicates that the mechanism of the extra MLP to the network needs to be further studied.

heading: A MAIN THREAD OF PAPER
text: The proposal of RAFT is based on a series of theoretical derivation and empirical approximation. Therefore the logic chain of our paper is fundamental to the legitimacy of our explanation on BYOL and the superiority of our newly proposed RAFT. Here we organize our main thread in the same order as the sections, to provide a clear view with readers.
In Section 3,
• As a learning framework, BYOL does not consistently work. It heavily relies on the existence of the predictor. We want to understand why this inconsistency exists.
• The architecture of the predictor doesn’t affect the collapse of BYOL, the fact that the linear predictor qw(·) = W (·) prevents collapse will be used as a crucial condition in Section 5.
In Section 4.1,
• A new disentangled objective LBYOL′ = αLalign(qw ◦ fθ) + βLcross-model(qw ◦ fθ, fξ) is established by upper bounding.
• We showcase that minimizing LBYOL′ is close to minimizing LBYOL in terms of alignment, uniformity, and linear evaluation protocol, which indicates that understanding the behavior of optimizing BYOL’s upper bound is approximately equivalent to understanding BYOL.
In Section 4.2,
• We find minimizing [Lalign(qw ◦ fθ) − Lcross-model(qw ◦ fθ, fξ)] works as well, which has the exact opposite way of incorporating the cross-model loss to BYOL′.
• Based on the observation above, we propose a new self-supervised learning approach Run Away From your Teacher, which regularizes Lalign(qw ◦fθ) by maximizing Lcross-model(qw ◦ fθ, fξ). Compared with BYOL, RAFT accords more with our common understanding. • Additional experiments show that without predictor, BYOL′ fails to regularize Lalign(fθ), let alone optimizing uniformity. On the contrary, although not able to actively optimize uniformity either, RAFT’s maximizing Lcross-model(fθ, fξ) continues to be an effective regularizer for Lalign(fθ), which makes it more favorable (Figure 2b).
In Section 5,
• We prove that when the predictor is linear (qw = W ) and the representation space is a hypersphere where only the tangential component of gradient is preserved during training, minimizing Lcross-model(W ◦ fθ, fξ) and maximizing it obtain the same encoder fθ. • Based on the equivalence above, we conclude that BYOL′ is a special case of RAFT un-
der conditions above. The equivalence established helps understanding several counterintuitive behaviors of BYOL.

heading: B BACKGROUND AND RELATED WORK
text: B.1 CONTRASTIVE LEARNING
Contrastive methods relies on the assumption that two views of the same data point share the information, and thus creates a positive pair. By separating the positives and the negatives, the neural network trained by the algorithm learns to extract the most useful information from the data and performs better on the downstream tasks. Typically, the algorithm uses the InfoNCE objective:
Lcontrast(h,K) = E (x,x+)∼Ppos {x−i } K i=1∼X K
[ − log e h(x,x+)
eh(x,x+) + ∑K i=1 e h(x,x−i )
] , (11)
where (x, x+) are sampled from the positive pair distribution Ppos, which is built by a series of data augmentation functions [ref]. The negative samples {x−i }K are i.i.d sampled for K times from the data distribution X ; function h(x, y) measures the similarity between two input data (x, y). Empirically for the sake of symmetry, the measurement function h(x, y) = d(f(x), f(y)) has an encoder f(·) and a similarity metric d(·, ·) evaluating how close the two representations are.
B.2 MEAN TEACHER
There is one type of semi-supervised learning method that BYOL constantly reminds people of, Mean Teacher (MT) (Tarvainen & Valpola, 2017; Laine & Aila, 2016). Like BYOL, MT is also of Teacher-Student (T-S) framework, where the teacher network is also the EMA of the student network. The additional consistency loss between the teacher and student is applied to the supervised signals. There has been a lot of work demonstrating the efficacy of MT(Athiwaratkun et al., 2019; Novak et al., 2018; Chaudhari et al., 2019), among which the major conclusion states that the consistency loss between the student and its MT acts as a regularizer for better generalization. The proven properties of MT might lead us to focus on how the online network’s learning from MT effectively regularizes Lalign in BYOL. In this paper, however, we propose the opposite way of leveraging MT in contrastive methods.

heading: C EXPERIMENTAL SETUP
text: Dataset Our main goal is to unravel the mystery why BYOL doesn’t collapse during training and to solve the predictor-inconsistency. The most important metric is whether the algorithm collapses or not, and we don’t target on developing a more powerful self-supervised learning algorithm that surpasses SOTA on large dataset. In this repsect, we limit our experiments to the scope of the CIFAR10 dataset. Each image is resized from 32× 32 to 96× 96. This change is the consequence of the tradeoff between the effect of the data augmentation and batch size: larger size of the image would allow more subtle and informative data augmentation scheme while it will reduce the training batch size, which has already been empirically shown is harmful to the model performance.
Model architecture In our experiments, the model is composed of three stages: an encoder fθ that adopts the ResNet18 architecture (without the classifier on top); a projector gθ that is comprised of a linear layer with output size 512, batch normalization, rectified linear units (ReLU), and a final linear layer with output size 128; a predictor qw that is comprised of the same architecture as the projector but without the batch normalization.
Training We adopt the same data augmentation scheme that is used in Chen et al. (2020a) and Grill et al. (2020) and train the BYOL on the training set for 300 epochs with batch size 128 on 3 random seeds. The objective of training is specified accordingly and the model is trained on the Adam optimizer with learning rate 3 × 10−4 (Kingma & Ba, 2014). Unless stated otherwise, we update the target network with the EMA rate 4× 10−3 without the cosine smoothing trick. Evaluation After training, we evaluate the encoder’s performance on the widely adopted linear evaluation protocol: we fix the parameter of the encoder and we train another linear classifier on top of it using all the training labels for 100 epochs with learning rate 5 × 10−4. The final classification accuracy indicates to what degree the representations of the same class concentrate and the representations of the different class separate, and thus tells the quality of the representation.

heading: D TABLES

heading: E ALGORITHMS
text: Algorithm 1: RAFT: Run Away From your Teacher Inputs : X , T1, and T2 set of images and distributions of transformations θ and fθ model parameters and encoder w and qw predictor parameters and predictor ξ and fξ MT parameters and MT optimizer optimizer, updates online parameters using the loss gradient K and N total number of optimization steps and batch size {τk}Kk=1 and {ηk}Kk=1 target network update schedule and learning rate schedule
1 for k = 1 to K do 2 B ← {xi}Ni=1 ∼ XN // sample a batch of N images 3 for xi ∈ B do 4 t1 ∼ T1 and t2 ∼ T2 // sample image transformations 5 z1 ← qw(fθ(t1(xi))) and z2 ← qw(fθ(t2(xi))) // reps for model 6 z′1 ← fξ(t1(xi)) and z′2 ← fξ(t2(xi)) // reps for MT 7 li = ∥∥ z1 ‖z1‖2 − z2 ‖z2‖2 ∥∥2 2 // loss for alignment
8 l′i = − 12 (∥∥ z1 ‖z1‖2 − z′1 ‖z′1‖2 ∥∥2 2 + ∥∥ z2 ‖z2‖2 − z′2 ‖z′2‖2 ∥∥2 2 ) // loss for cross-model 9 end
10 δθ ← 1N ( N∑ i=1 ∂θli + ∂θl ′ i ) // compute the loss gradient w.r.t. θ 11 θ ← optimizer(θ, δθ, ηk) // update trainable parameters 12 ξ ← τkξ + (1− τk)θ // update target parameters 13 end
Output: encoder fθ
F VISUALIZATION OF TRAINING EVOLUTION
F.1 REPRESENTATION DISTRIBUTION EVOLUTION OF DIFFERENT LEARNING ALGORITHMS
F.2 EVOLUTION OF NUMERICAL METRIC IN DIFFERENT LEARNING ALGORITHMS

heading: G PROOF OF BYOL UPPER BOUNDING
text: In this section, we provide how we derive the upper bound of LBYOL. For the sake of simplicity, without loss of rigor, we use t1 = t1(x) to represent the transformed input x.
LBYOL = Ex∼X ,t1∼T1,t2∼T2 [ ‖qw(fθ(t1(x)))− fξ(t2(x))‖22 ] = E [ ‖qw(fθ(x1))− qw(fθ(x2)) + qw(fθ(x2))− fξ(x2)‖22 ] . (12)
By applying the Cauchy-Schwarz’s inequality to Eq. 12, we yield:
LBYOL ≤ (1 + 1 λ ) ( E [∥∥qw(fθ(x1))− qw(fθ(x2))∥∥22]+ λE [∥∥qw(fθ(x2))− fξ(x2)∥∥22])
= (1 + 1
λ ) (Lalign(qw ◦ fθ;Ppos) + λLcross-model(qw ◦ fθ, fξ,X )) (13)
which stands for any λ > 0, where the positive-pair distribution Ppos is modeled by the chain rule of the conditional probability:
Ppos(x1, x2) = X (x) · T1(t1|x) · T2(t2|x).
For any given pair α, β > 0, we let λ = β/α and substitute it back to Eq. 13, yielding
LBYOL ≤ (1 + α
β )
[ Lalign(qw ◦ fθ;Ppos) + β
α Lcross-model(qw ◦ fθ, fξ,X ) ] = ( 1
α +
1 β ) [αLalign(qw ◦ fθ;Ppos) + βLcross-model(qw ◦ fθ, fξ,X )]
= ( 1
α +
1 β )LBYOL′, (14)
and as an optimization objective, we have
min( 1
α +
1 β )LBYOL′ ⇔ minLBYOL′ (15)
Therefore we have proven that LBYOL′ as optimization objective is the upper bound of LBYOL. To note here, one can subtract and add a different term fξ(x1) to form the alignment loss on the side of MT fξ,
LBYOL = E [∥∥qw(fθ(x1))− fξ(x1) + fξ(x1)− fξ(x2)∥∥22] , (16)
while it doesn’t help to solve the problem since the alignment constraint on the side of MT doesn’t generate gradients.

heading: H PROOF OF ONE-TO-ONE CORRESPONDENCE BETWEEN BYOL′ AND RAFT
text: Theorem (One-to-one correspondence between BYOL′ and RAFT) There is a one-to-one correspondence between parameter trajectories of BYOL′ and RAFT when the following three conditions hold:
i. the representation space is a hypersphere;
ii. the predictor is a linear transformation, i.e. qw(·) = W (·);
iii. only the tangential component of the gradient on the hypersphere is preserved.
Without losing generality, suppose that x1 = t1(x), x2 = t2(x) where x is an arbitrary input and batch size is 1, and (α, β) = (1, 1). We set BYOL′ and RAFT with initial parameters (θ′,W ′) = (θ(0),W (0)) and (θ,W ) = (θ(0),−W (0)) respectively. For convenience, we assumes the dot product “·” ignores the row layout or column layout in the chain rule of derivatives and we define the following symbols:
z2 = fξ(x2), (17) z′1 = W ′fθ′(x1), (18) z′2 = W ′fθ′(x2), (19)
z1 = Wfθ(x1), (20) z2 = Wfθ(x2). (21)
Based on the notations defined, we rewrite the loss terms of BYOL′ and RAFT as follows:
LBYOL ′ align = ∥∥z′1 − z′2∥∥22, (22)
LBYOL ′ cross-model = ∥∥z′2 − z2∥∥22, (23)
LRAFTalign = ∥∥z1 − z2∥∥22, (24)
LRAFTcross-model = − ∥∥z2 − z2∥∥22. (25)
The two objectives are following:
LBYOL′ = LBYOL ′ align + LBYOL ′ cross-model, (26)
LRAFT = LRAFTalign + LRAFTcross-model. (27)
We claim that under the third condition, the following equations hold:[∂LBYOL′ ∂θ′ ] ‖ = [∂LRAFT ∂θ ] ‖ , [∂LBYOL′ ∂W ′ ] ‖ = − [∂LRAFT ∂W ] ‖ , (28)
where subscript ‖ denotes the tangential component of the gradient.
Firstly we show the equivalence with respect to θ. Differentiate LBYOL′align ,LRAFTalign with respect to θ′ij , θij respectively, we obtain
∂LBYOL′align ∂θ′ij
= 2 [ (z′1 − z′2)‖ + (z ′ 1 − z′2)⊥ ] · ( ∂z′1 ∂θ′ij − ∂z ′ 2 ∂θ′ij ) , (29)
∂LRAFTalign ∂θij
= 2 [ (z1 − z2)‖ + (z1 − z2)⊥ ] · ( ∂z1 ∂θij − ∂z2 ∂θij ) , (30)
[∂LBYOL′align ∂θ′ij ] ‖ = 2 (z′1 − z′2)‖ · ( ∂z′1 ∂θ′ij − ∂z ′ 2 ∂θ′ij ) , (31)
[∂LRAFTalign ∂θij ] ‖ = 2 (z1 − z2)‖ · ( ∂z1 ∂θij − ∂z2 ∂θij ) , (32)
where (z′1 − z′2) , (z1 − z2) are vectors at the points z′2 and z2 on the hypersphere and we decompose the vector into the tangential (denoted by ‖) and normal component (denoted by ⊥):
(z′1 − z′2) = [ (z′1 − z′2)‖ + (z ′ 1 − z′2)⊥ ] , (z1 − z2) = [ (z1 − z2)‖ + (z1 − z2)⊥ ] , (33)
Generally, suppose z is a unit vector starting at the origin point, which is perpendicular to the unit hypersphere at the point z, for any vector v starting at the point z, we have
v⊥ = 〈v, z〉 · z, v‖ = v − v⊥ = v − 〈v, z〉 · z. (34)
Then we can compute the tangential component of the gradient:
(z′1 − z′2)‖ = (z ′ 1 − z′2)− 〈z′1 − z′2, z′2〉 · z′2 =z′1 − 〈z′2, z′1〉 · z′2, (z1 − z2)‖ = (z1 − z2)− 〈z1 − z2, z2〉 · z2
=z1 − 〈z2, z1〉 · z2. (35)
Because of the initialization, z′1 = −z1, z2 = −z′2, therefore we have
(z′1 − z′2)‖ = − (z1 − z2)‖ , (36)( ∂z′1 ∂θ′ij − ∂z ′ 2 ∂θ′ij ) = − ( ∂z1 ∂θij − ∂z2 ∂θij ) . (37)
So we show that [∂LBYOL′align ∂θ′ij ] ‖ = [∂LRAFTalign ∂θij ] ‖ . (38)
We differentiate LBYOL′cross-model, LRAFTcross-model with respect to θ′ij , θij respectively, we obtain that[∂LBYOL′cross-model ∂θ′ij ] ‖ = −2 (z2 − z′2)‖ ·W ′ · ∂fθ ′(x2) ∂θ′ij , (39)
[∂LRAFTcross-model ∂θij ] ‖ = 2 (z2 − z2)‖ ·W · ∂fθ(x2) ∂θij . (40)
Similar to Eq. 35, we derive that
(z2 − z′2)‖ = (z2 − z2)‖ (41)
Since θ′ = θ, ∂fθ′(x2)/∂θ′ij = ∂fθ(x2)/∂θij and W ′ = −W , we have that[∂LBYOL′cross-model ∂θ′ij ] ‖ = [∂LRAFTcross-model ∂θij ] ‖ . (42)
Therefore by Eq. 38 and Eq. 42, RAFT’s updating of the parameter θ is equal to BYOL′:[∂LBYOL′ ∂θ′ ] ‖ = [∂LRAFT ∂θ ] ‖ . (43)
Also we differentiate LBYOL′align ,LRAFTalign with respect to W ′ij , Wij respectively, we obtain that[∂LBYOL′align ∂W ′ij ] ‖ = 2 (z′1 − z′2)‖ · ( ∂z′1 ∂W ′ij − ∂z ′ 2 ∂W ′ij ) , (44)
[∂LRAFTalign ∂Wij ] ‖ = 2 (z1 − z2)‖ · ( ∂z1 ∂Wij − ∂z2 ∂Wij ) . (45)
Note that z1 = −z′1, z2 = −z′2 and similar to Eq. 35, easy to show that
(z′1 − z′2)‖ = −(z1 − z2)‖. (46)
Also,
∂z′1 ∂W ′ij = ∂z1 ∂Wij , (47)
∂z′2 ∂W ′ij = ∂z2 ∂Wij . (48)
So we have [∂LBYOL′align ∂W ′ij ] ‖ = − [∂LRAFTalign ∂Wij ] ‖ . (49)
Differentiate LBYOL′cross-model,LRAFTcross-model with respect to W ′ij ,Wij respectively, we obtain that[∂LBYOL′cross-model ∂W ′ij ] ‖ = −2 (z2 − z′2)‖ · ∂z′2 ∂W ′ij , (50)
[∂LRAFTcross-model ∂Wij ] ‖ = 2 (z2 − z2)‖ · ∂z2 ∂Wij . (51)
Then we have that [∂LBYOL′cross-model ∂W ′ij ] ‖ = − [∂LRAFTcross-model ∂Wij ] ‖ . (52)
By Eq. 49 and Eq. 52, we prove that the cross-model loss of BYOL′ generates the opposite gradient to RAFT, namely, [∂LBYOL′
∂W ′ ] ‖ = − [∂LRAFT ∂W ] ‖ . (53)
Therefore by the two main conclusions Eq. 43 and Eq. 53, for BYOL′ with parameters (θ′,W ′) = (θ(0),W (0)) and RAFT with parameters (θ,W ) = (θ(0),−W (0)) respectively, we have
BYOL: ( θ′(1) = θ′(0) − η [∂LBYOL′ ∂θ′ ] ‖ ∣∣∣∣ θ′=θ′(0) ,W ′(1) = W ′(0) − η [∂LBYOL′ ∂W ′ ] ‖ ∣∣∣∣ W ′=W ′(0) ) ,
(54) RAFT: ( θ(1) = θ(0) − η [∂LRAFT ∂θ ] ‖ ∣∣∣∣ θ=θ(0) ,W (1) = W (0) − η [∂LRAFT ∂W ] ‖ ∣∣∣∣ W=W (0) ) . (55)
We derive that θ(1) = θ′(1),W (1) = −W ′(1), and furthermore, θ(k) = θ′(k),W (k) = −W ′(k) at any iteration k. In this way, we establish an one-to-one correspondence between the parameter trajectories of BYOL′ and RAFT in training, referred to asH:
H : RAFT(θ,W ) 7→ BYOL′(θ,−W ) (56)

heading: I NON-TRIVIAL SOLUTIONS CREATED BY PREDICTOR
text: Suppose inputs x1 = t1(x) and x2 = t2(x) is n-dimensional. And in linear model, fθ, fξ and qw is parameterized by matrices (θij)n×n, (ξij)n×n and (Wij)m×m respectively.
The objective is LBYOL = E(x,t1,t2)∼(X ,T1,T2) [∥∥qw(fθ(t1(x)))− fξ(t2(x))∥∥22]
= E(x,t1,t2)∼(X ,T1,T2) [ ‖Wθx1 − ξx2‖22 ] (57)
Differentiate ‖Wθx1 − ξx2‖22 with respect to θij and Wij , we have ∂ [ ‖Wθx1 − ξx2‖22 ] ∂θij = m∑ k=1 ∂ [Wk,:(θx1)− ξk,:x2]2 ∂θij
= m∑ k=1 2 [Wk,:(θx1)− ξk,:x2] ∂ [Wk,:(θx1)− ξk,:x2] θij
= m∑ k=1 2TkWk,:(x1)j
= 2 [( W>T ) x>1 ] ij , (58)
∂ [ ‖Wθx1 − ξx2‖22 ] ∂Wij = m∑ k=1 ∂ [Wk,:(θx1)− ξk,:x2]2 ∂Wij
= m∑ k=1 2 [Wk,:(θx1)− ξk,:x2] ∂ [Wk,:(θx1)− ξk,:x2] ∂Wij
= m∑ k=1 2Tk(θx1)j1{k=i}
= 2Tiθj,:x1 = 2 [ T (θx1) >] ij , (59)
where Tk = [Wk,:(θx1)− ξk,:x2], T = (T1, T2, . . . , Tm)> = W (θx1)− ξx2, and Wk,:, ξk,: are the k−th row of W and ξ respectively. So
∂ [ ‖Wθx1 − ξx2‖22 ] ∂θ = 2 [( W>T ) x> ] , ∂ [ ‖Wθx1 − ξx2‖22 ] ∂W = 2 [ T (θx)> ] (60)
Let
∂LBYOL ∂θ = 0, ∂LBYOL ∂W = 0, (61)
we have that ∂E(x,t1,t2)∼(X ,T1,T2) [ ‖Wθx1 − ξx2‖22 ] ∂θ = E(x,t1,t2)∼(X ,T1,T2)[ ∂‖Wθx1 − ξx2‖22 ∂θ ] = 0 (62)
∂E(x,t1,t2)∼(X ,T1,T2) [ ‖Wθx1 − ξx2‖22 ] ∂W = E(x,t1,t2)∼(X ,T1,T2)[ ∂‖Wθx1 − ξx2‖22 ∂W ] = 0 (63)
When the weight of target ξ converge, we have ξ(k) = ξ(k+1) in the updating rule,
ξ(k+1) = τkξ k + (1− τk)θ(k)
θ(k) = ξ(k+1) = ξ(k) (64)
Substituting ξ by θ, we obtain E(x,t1,t2)∼(X ,T1,T2) [ W>(Wθx1 − θx2)x>1 ] = 0 (65)
E(x,t1,t2)∼(X ,T1,T2) [ (Wθx1 − θx2)x>1 θ> ] = 0 (66)
Let E(x,t1,t2)∼(X ,T1,T2)(x1x>1 ) = A, E(x,t1,t2)∼(X ,T1,T2)(x2x>1 ) = B, we have that
eq : sylW>(WθA) = W>θB
WθAθ> = θBθ>
⇒Wθ − θBA−1 = 0 (67)
To solve Eq. ?? (which is called Sylvester’s equation), we using the Kronecker product notation and the vectorization operator vec, we can rewrite the equation in the form(
Im ⊗W − (BA−1)T ⊗ In ) vec θ = vec0 (68)
So it has a non-trivial solution θ if and only if ( Im ⊗W − (BA−1)T ⊗ In ) has a non-trivial null space. An equivalent condition to having a non-trivial null space is having zero as an eigenvalue. Let W has eigenvalues in common with BA−1, then we have a non-trivial solution of θ, which is exactly the prevention for collapse.

Review summaries:
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.

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.