Title: Implicit variance regularization in non-contrastive SSL
Venue: NeurIPS 2023
Accepted: yes
Keywords: Self-supervised learning

Non-contrastive learning

Learning dynamics
Abstract: Non-contrastive SSL methods like BYOL and SimSiam rely on asymmetric predictor networks to avoid representational collapse without negative samples. Yet, how predictor networks facilitate stable learning is not fully understood. While previous theoretical analyses assumed Euclidean losses, most practical implementations rely on cosine similarity. To gain further theoretical insight into non-contrastive SSL, we analytically study learning dynamics in conjunction with Euclidean and cosine similarity in the eigenspace of closed-form linear predictor networks. We show that both avoid collapse through implicit variance regularization albeit through different dynamical mechanisms. Moreover, we find that the eigenvalues act as effective learning rate multipliers and propose a family of isotropic loss functions (IsoLoss) that equalize convergence rates across eigenmodes. Empirically, IsoLoss speeds up the initial learning dynamics and increases robustness, thereby allowing us to dispense with the EMA target network typically used with non-contrastive methods. Our analysis sheds light on the variance regularization mechanisms of non-contrastive SSL and lays the theoretical grounds for crafting novel loss functions that shape the learning dynamics of the predictor's spectrum.

Implicit variance regularization in non-contrastive SSL

Abstract

Non-contrastive self-supervised learning (SSL) methods like BYOL and SimSiam rely on asymmetric predictor networks to avoid representational collapse without negative samples. Yet, how predictor networks facilitate stable learning is not fully understood. While previous theoretical analyses assumed Euclidean losses, most practical implementations rely on cosine similarity. To gain further theoretical insight into non-contrastive SSL, we analytically study learning dynamics in conjunction with Euclidean and cosine similarity in the eigenspace of closed-form linear predictor networks. We show that both avoid collapse through implicit variance regularization albeit through different dynamical mechanisms. Moreover, we find that the eigenvalues act as effective learning rate multipliers and propose a family of isotropic loss functions (IsoLoss) that equalize convergence rates across eigenmodes. Empirically, IsoLoss speeds up the initial learning dynamics and increases robustness, thereby allowing us to dispense with the exponential moving average (EMA) target network typically used with non-contrastive methods. Our analysis sheds light on the variance regularization mechanisms of non-contrastive SSL and lays the theoretical grounds for crafting novel loss functions that shape the learning dynamics of the predictor's spectrum.

1 Introduction

SSL has emerged as a powerful method to learn useful representations from vast quantities of unlabeled data [1-8]. In SSL, the network's objective is to 'pull' together its outputs for two differently augmented versions of the same input, so that they learn representations that are predictive across randomized transformations [9]. To avoid the trivial solution whereby the network output becomes constant, also called representational collapse, SSL methods use either a contrastive objective to 'push' apart representations of unrelated images [2, 3, 10-13] or other non-contrastive strategies. Non-contrastive methods comprise explicit variance regularization techniques [6, 7, 14], whitening approaches [15, 16], and asymmetric losses as in Bootstrap Your Own Latent (BYOL) [1] and SimSiam [5]. Asymmetric losses break symmetry between the two branches by passing one of the representations through a predictor network and stopping gradients from flowing through the other 'target' branch. How this architectural modification prevents representational collapse is not obvious and has been the focus of several theoretical [17-21] and empirical studies [22-24]. A significant advance was provided by Tian et al. [17] who showed that linear predictors align with the correlation matrix of the embeddings, and proposed the closed-form predictor DirectPred based on this insight. However, previous analyses assumed a Euclidean loss at the output [17-19, 21] except [20], whereas practical implementations typically use the cosine loss [1, 5] which yields superior performance on downstream tasks. This difference raises the question whether analysis based on the Euclidean loss provides an accurate account of the learning dynamics under the cosine loss.

In this work, we provide a comparative analysis of the learning dynamics for the Euclidean and cosine-based asymmetric losses in the eigenspace of the closed-form predictor DirectPred. Our analysis shows how both losses implicitly regularize the variance of the representations, revealing a connection between asymmetric losses and explicit variance regularization in VICReg [7]. Yet, the learning dynamics induced by the two losses are markedly different. While the learning dynamics of different eigenmodes decouple in the Euclidean case, dynamics remain coupled for the cosine loss.

Moreover, our analysis shows that for both losses, the predictor's eigenvalues act as learning rate multipliers, thereby slowing down learning for modes with small eigenvalues. Based on our analysis, we craft an isotropic loss function (IsoLoss) for each case that resolves this problem and speeds up the initial learning dynamics. Furthermore, IsoLoss works without an EMA target network possibly because it boosts small eigenvalues, the purported role of the EMA in DirectPred [17]. In summary, our main contributions are the following:

We analyze the SSL dynamics in the eigenspace of closed-form linear predictors for asymmetric Euclidean and cosine losses and show that both perform implicit variance regularization, but with markedly different learning dynamics.

Our analysis shows that predictor eigenvalues act as learning rate multipliers which slows down learning for small eigenvalues.

We propose isotropic loss functions for both cases that equalize the dynamics across eigenmodes and improve robustness, thereby allowing to learn without an EMA target network.

2 Eigenspace analysis of the learning dynamics

To gain a better analytic understanding of the SSL dynamics underlying non-contrastive methods such as BYOL and SimSiam [1, 5], we analyze them in the predictor's eigenspace. Specifically we proceed in three steps. First, building on DirectPred, we invoke the neural tangent kernel (NTK) to derive simple dynamic expressions of the predictor's eigenmodes for Euclidean and cosine loss. This formulation uncovers the implicit variance regularization mechanisms that prevent representational collapse. Using the eigenspace framework, we illustrate how removing the predictor or the stopgradient results in collapse or run-away dynamics. Finally, we find that predictor eigenvalues act as learning rate multipliers for their associated mode, thereby slowing down learning for small eigenvalues. We derive a modified isotropic loss function (IsoLoss) that provides more equalized learning dynamics across modes, which showcases how our analytic insights help to design novel loss functions that actively shape the predictor spectrum. However, before we start our analysis, we will briefly review DirectPred [17] and the NTK [25], a powerful theoretical tool linking representational changes and parameter updates. We will rely on both concepts for our analysis.

2.1

Background and problem setup

We begin by reviewing DirectPred [17] and defining our notation. In the following, we consider a Siamese neural network z = f ( x ; θ ) with output z ∈ R M , input x ∈ R N , and parameters θ . We further assume a linear predictor network W P ∈ R M × M and use the same parameters for the online and target branches as in SimSiam [5]. We denote pairs of representations as z (1) , z (2) corresponding to pairs of inputs x (1) , x (2) related through augmentation and implicitly assume that all losses are averaged over many augmented pairs. The asymmetric loss function (Fig. 1a), introduced in BYOL [1], is then given by:

where SG denotes the stop-gradient operation, and d is either the Euclidean distance metric d ( a , b ) = 1 2 ∥ a -b ∥ 2 or the cosine distance metric d ( a , b ) = -a ⊤ b ∥ a ∥∥ b ∥ . We refer to the corresponding loss functions as L euc and L cos respectively.

DirectPred. Tian et al. [17] showed that a linear predictor in the BYOL setting aligns during learning with the correlation matrix of representations C z := E x [ zz ⊤ ] , where the expectation is taken over the data distribution. Since the correlation matrix is a real symmetric matrix, one can diagonalize it over R : C z = UD C U ⊤ , where U is an orthogonal matrix whose columns are the eigenvectors of C z and D C is the real-valued diagonal matrix of the eigenvalues s m with m ∈ [1 , M ] .

a

Loss

Stop-grad

x

Predictor

Online net

Target net

L

=

d

(

W

P

z

(1)

,

SG(

z

(2)

)

)

b

Euclidean distance

Cosine distance

0

π

2

π

c

Given this eigendecomposition, the authors proposed DirectPred, in which the predictor is not learned via gradient descent but directly set to:

where α is a positive constant exponent applied element-wise to D C . The eigenvalues λ m of the predictor matrix W P are then λ m = s α m . We use D to denote the diagonal matrix containing the eigenvalues λ m . While DirectPred used α = 0 . 5 , the follow-up study DirectCopy [18] showed that α = 1 is also effective while avoiding the expensive diagonalization step. While Tian et al. [17] based their analysis on the Euclidean loss L euc , most practical models, including Tian et al.'s large-scale experiments, relied on the cosine similarity loss L cos . This discrepancy raises the question to what extent setting the predictor to the above expression is justified for the cosine loss. Empirically, we find that a trainable linear predictor does align its eigenspace with that of the representation correlation matrix also for the cosine loss (see Fig. 4 in Appendix A).

Neural tangent kernel (NTK). The NTK is a powerful analytical tool characterizing the learning dynamics of neural networks [25, 26]. Here, we recall the definition of the empirical NTK [26] corresponding to a single instantiation of the network's parameters θ . If | D | denotes the size of the training dataset, L : R M → R an arbitrary loss function, X , the training data concatenated into one vector of size N | D | , and Z = z ( X ) , the concatenated output of size M | D | , then the empirical NTK is the ( M | D | × M | D | ) -sized matrix:

and the continuous-time gradient-descent dynamics [26] of the representations z are given by:

In other words, the empirical NTK links the representational dynamics d z d t under gradient descent on the parameters θ , and the 'representational gradient' ∇ Z L .

2.2 Implicit variance regularization in non-contrastive SSL

As a starting point for our analysis, we first express the relevant loss functions in the eigenbasis of the predictor network. We do this using a closed-form linear predictor as prescribed by DirectPred. In the following, we use ˆ z = U ⊤ z to denote the representation expressed in the eigenbasis.

Lemma 1. (Euclidean and cosine loss in the predictor eigenspace) Let W P be a linear predictor set according to DirectPred with eigenvalues λ m , and ˆ z the representations expressed in the predictor's eigenbasis. Then the asymmetric Euclidean loss L euc and cosine loss L cos can be expressed as:

for which we defer the simple proof to Appendix B. Rewriting the losses in the eigenbasis makes it clear that the asymmetric loss with DirectPred can be viewed as an implicit loss function in the predictor's eigenspace, where the variance of each mode naturally appears through the λ m terms. In the following analysis, we will show how the learning dynamics implicitly regularize these variances λ m . From Eq. (3) we directly see that L euc is a sum of M terms, one for each eigenmode, which decouples the learning dynamics, a fact first noted by Tian et al. [17]. In contrast, the form of L cos yields coupled dynamics due to the ∥ D ˆ z (1) ∥ = √ ∑ k ( λ k ˆ z (1) k ) 2 term in the denominator. This coupling arises from the normalization of the representation vectors to the unit hypersphere when calculating the cosine distance. The normalization effectively removes one degree of freedom and, in the process, adds a dependence between all the representation dimensions (Fig. 1b and 1c).

To get an analytic handle on the evolution of the eigen-representations ˆ z as the encoder learns, we first note that if training were to update the representations directly, instead of indirectly through updating the weights θ , they would evolve along the following 'representational gradients':

In practice, however, representations of different samples do not evolve independently along these gradients, but influence each other through parameter changes in θ . This interdependence of representations and parameters are captured by the empirical NTK Θ t ( X , X ) (cf. Eq. (2)). Because the NTK is positive semi-definite, loosely speaking, gradient descent on the parameters changes representations 'in the direction' of the above representational gradients.

To see this link more formally, we express the NTK in the eigenbasis as ˆ Θ t ( X , X ) = ∇ θ ˆ Z∇ θ ˆ Z ⊤ where ˆ Z = ˆ z t ( X ) = U ⊤ z t ( X ) . Since we are concerned with the learning dynamics in this rotated basis, we will rewrite Eq. (2) for continuous-time gradient descent for a generic loss function L as:

Note, that structurally these dynamics are the same as the embedding space dynamics in Eq. (2) but merely expressed in the predictor eigenbasis (see Lemma 2 in Appendix B for a derivation). Although ˆ Θ t changes over time and is generally intractable in finite-width networks, it is positive semidefinite. This property guarantees that the cosine angle between the representational training dynamics under the parameter-space optimization of a neural network d d t ˆ Z ∝ ˆ Θ t ∇ ˆ Z L and the dynamics that would result from optimizing the representations d d t ˆ Z ∝ -∇ ˆ Z L is non-negative:

In other words, the representational updates due to network training lie within a 180-degree cone of the dynamics prescribed by Eqs. (5) and (6). This guarantee makes it possible to draw qualitative conclusions about asymptotic collective behavior, e.g., whether a network is bound to collapse or not, from analyzing the more tractable dynamics that follow the representational gradients d d t ˆ Z ∝ -∇ ˆ Z L of the transformed BYOL/SimSiam loss. For ease of analysis, we now consider linear networks with Gaussian i.i.d inputs, an important limiting case amenable for theoretical analysis [27]. In this setting

the empirical NTK becomes the identity and the simplified representational dynamics are exact, allowing us to fully characterize the representational dynamics for L euc and L cos in the following two theorems. In the proofs for these theorems, we show that the assumption of Gaussian inputs can be relaxed further.

Theorem 1. (Representational dynamics under L euc ) For a linear network with i.i.d Gaussian inputs learning with L euc , the representational dynamics of each mode m independently follow the gradient of the loss -∇ ˆ z L euc . More specifically, the dynamics uncouple and follow M independent differential equations:

which, after taking the expectation over augmentations yields the dynamics:

We provide the proof in Appendix B and appreciate that d d t ˆ z m has the same sign as ˆ z m whenever λ m < 1 and the opposite sign whenever λ m > 1 . These dynamics are convergent and approach an eigenvalue λ m of one, thereby preventing collapse of mode m . Since the eigenmodes are orthogonal and uncorrelated, and the condition simultaneously holds for all modes, this ultimately prevents both representational and dimensional collapse [28]. Since the eigenvalues also correspond to the variance of the representations, the underlying mechanism constitutes an implicit form of variance regularization. Finally, we note that the above decoupling of the dynamics for the Euclidean loss has been described previously in Tian et al. [17].

Nevertheless, the representational dynamics are different for the commonly used cosine loss L cos .

Theorem 2. (Representational dynamics under L cos ) For a linear network with i.i.d Gaussian inputs trained with L cos , the dynamics follow a system of M coupled differential equations:

̸

and reach a regime in which the eigenvalues are comparable in magnitude. In this regime, the expected update over augmentations is well approximated by:

̸

̸

where we have assumed averages over augmentations. See Appendix B for the proof. Theorem 2 states that d d t ˆ z m has the same or different sign as ˆ z m depending on the sign of the aggregate sum ∑ k = m λ k ( λ k -λ m ) . This relation suggests that a steady state is only reached through mutual agreement when the non-zero eigenvalues are all equal. In contrast to the Euclidean case, there is no pre-specified target value (see Fig. 5 in Appendix A). Thus, the cosine loss also induces implicit variance regularization, but through a markedly different mechanism in which eigenmodes cooperate.

2.3 Stop-grad and predictor network are essential for implicit variance regularization.

We now extend our analysis to explain the known failure modes due to ablating the predictor or the stop-gradient for each distance metric. When we omit the stop-grad operator from L euc , we have:

so that d d t ˆ z m and ˆ z m always have opposite signs (see Appendix C for the derivation). This drives the representations toward zero with exponentially decaying eigenvalues, causing the notorious representational collapse [5]. Omitting the stop-grad operator from L cos yields a nontrivial expression for the dynamics causing the largest eigenmode to diverge (see Appendix C) . Interestingly, this is different from the collapse to zero inferred for the Euclidean distance.

Similarly, when removing the predictor network in the Euclidean loss case, the dynamics read:

meaning that no learning updates occur. When the predictor is removed in the cosine loss case, the dynamics are:

̸

As we show in Appendix C, these dynamics also avoid collapse. However, the effective learning rates become impractically small without the eigenvalue factors from Eq. (11). We summarized the predicted dynamics of all settings in Table 1. Thus, our analysis provides mechanistic explanations for why stop-grad and predictor networks are required for avoiding collapse in non-contrastive SSL.

̸

̸

2.4 Isotropic losses that equalize convergence across eigenmodes

In Eqs. (9) and (11) the eigenvalues appear as multiplicative learning rate modifiers in front of the difference terms that determine the fixed point. Hence, modes with larger eigenvalues converge faster than modes with smaller eigenvalues, reminiscent of previous theoretical work on supervised learning [27]. We hypothesized that the anisotropy in learning dynamics could lead to slow convergence for small eigenvalue modes or instability for large eigenvalues. To alleviate this issue, we designed alternative isotropic loss functions that equalize relaxation dynamics for all eigenmodes by exploiting the stop-grad function. Put simply, this involves taking the dynamics from Eqs. (8) and (10), removing the leading λ m term, and deriving the loss function that would result in the desired dynamics. One such isotropic 'IsoLoss' function for the Euclidean distance is:

We note that this IsoLoss has the same numerical value as L euc , but the gradient flow is modified by placing the prediction inside the stop-grad and also adding and subtracting z (1) inside and outside of the stop-grad. The associated idealized learning dynamics in our analytic framework are given by:

where the λ m factor (cf. Eq. (9)) disappeared (Table 1). Similarly, for the cosine distance, is one possible IsoLoss, in which W 1 / 2 P = UD 1 / 2 U ⊤ with the square-root applied element-wise to the diagonal matrix D . While this IsoLoss does not preserve numerical equality with the original loss L cos , it achieves the desired effect of removing the leading λ m learning-rate modifier (cf. Table 1).

a

L

euc

noSG

0

1000

Iteration

10

-

1

10

0

λ

L

euc

noPred

0

100000

Iteration

L

euc

0

1000

Iteration

L

euc

iso

0

1000

Iteration

L

cos

noSG

0

1000

Iteration

10

-

1

10

0

λ

L

cos

noPred

0

100000

Iteration

L

cos

0

1000

Iteration

L

cos

iso

0

1000

Iteration

b

c

d

3 Numerical experiments

To validate our theoretical findings (cf. Table 1), we first simulated a small linear Siamese neural network as shown in Fig.1a, for which Theorems 1 and 2 hold exactly. We fed the network with independent standard Gaussian inputs, and generated pairs of augmentations using isotropic Gaussian perturbations of standard deviation σ = 0 . 1 . We then trained the linear encoder with each configuration described above. Training the network with L euc noSG resulted in collapse with exponentially decaying eigenvalues, whereas L cos noSG succumbed to diverging eigenvalues as predicted (Fig. 2a). Training without the predictor caused vanishing updates for L euc noPred and slow learning for L cos noPred , in line with our analysis (Fig. 2b). Optimizing L euc , the representations become increasingly isotropic with all the eigenvalues λ m converging to one (Fig. 2c, top), whereas optimizing L cos also resulted in the eigenvalues converging to the same value, but different from one (Fig. 2c, bottom). The anisotropy in the dynamics of different eigenmodes noted above is particularly striking in the case of the Euclidean distance (Fig. 2c). Training with L euc iso and L cos iso resulted in similar convergence properties as their non-isotropic counterparts, but the eigenmodes converged at more homogeneous rates (Fig. 2d). Finally, we confirmed that these findings were qualitatively similar in the corresponding nonlinear networks with ReLU activation (see Fig. 6 in Appendix A). Thus, our theoretical findings hold up in simple Siamese networks.

a

10

3

10 0

10 3

L

euc

L

euc

iso

L

euc

-

BYOL

25

50

75

Acc.

DP

b

c

0

500

1000

Epoch

10

3

10 0

10 3

0

500

1000

Epoch

0

500

1000

Epoch

L

cos

L

cos

iso

L

cos

-

BYOL

0

500

Epoch

25

50

75

Acc.

DP

d

3.1 Theory qualitatively captures dynamics in nonlinear networks and real-world datasets.

To investigate how well our theoretical analysis holds up in non-toy settings, we performed several self-supervised learning experiments on CIFAR-10, CIFAR-100 [29], STL-10 [30], and TinyImageNet [31]. We based our implementation 1 on the Solo-learn library [32], and used a ResNet-18 backbone [33] as the encoder and the cosine loss, unless mentioned otherwise (see Appendix D for details). As baselines for comparison, we trained the same backbone using BYOL with the nonlinear predictor and DirectPred with the closed-form linear predictor. We recorded the online readout accuracy of a linear classifier trained on frozen features following standard practice, evaluated either on the held-out validation or test set where available.

We found that the eigenvalue dynamics of the representational correlation matrix in the ResNet-18 closely mirrored the analytical predictions for the closed-form predictor. For Euclidean distances (Fig. 3a), the eigenvalues for DirectPred and IsoLoss converged to a small range of values around one. However, the dynamics for BYOL with a learnable nonlinear predictor deviated significantly with the eigenvalues distributed over a larger range. Consistent with our analysis, IsoLoss had faster initial dynamics for the eigenvalues which also resulted in a faster initial improvements in model performance (Fig. 3b). The faster learning with IsoLoss was even more evident for the cosine distance (Fig. 3c). Surprisingly, BYOL, which uses a nonlinear predictor also closely matched the predicted dynamics in the case of the cosine distance. Furthermore, the dynamics showed a stepwise learning phenomenon wherein eigenvalues are progressively recruited one-by-one, consistent with recent findings for other SSL methods [34]. Finally, IsoLoss exhibited faster initial learning (Fig. 3d), in agreement with our theoretical analysis. Thus, our theoretical analysis accurately predicts key properties of the eigenvalue dynamics in nonlinear networks trained on real-world datasets.

3.2 IsoLoss promotes eigenvalue recruitment and works without an EMA target network.

To further investigate the impact of IsoLoss on learning, we first verified that it does not have any adverse effects on downstream classification performance. We found that IsoLoss matched or outperformed DirectPred on all benchmarks (Table 2) when trained with an EMA target network as used in the original studies. Yet, it performed slightly worse than BYOL, which uses a nonlinear predictor and an EMA target network. Because EMA target networks are thought to amplify small eigenvalues [17], we speculated that IsoLoss may work without it. We repeated training for the closed-form predictor losses without EMA to test this idea. We found that L cos iso was indeed robust to EMA removal. However, it caused a slight drop in performance (Table 2) and a notable reduction in the recruitment of small eigenvalues (see Fig. 7 in Appendix A). In contrast, optimizing the standard BYOL/SimSiam loss L cos with the symmetric linear predictor was unstable, as reported previously [17]. Finally, we confirmed the above findings also hold for α = 1 (cf. Eq. (1)) as prescribed by DirectCopy [18] (see Table 3 in Appendix A). Thus, IsoLoss allows training without an EMA target network.

The above result suggests that IsoLoss promotes the recruitment of small eigenvalues in closed-form predictors. Another factor that has been implicated in suppressing recruitment is weight decay [18]. To probe how weight decay and IsoLoss affect small eigenvalue recruitment, we repeated the above simulations with EMA and different amounts of weight decay. Indeed, we observed less eigenvalue recruitment with increasing weight decay for DirectPred (Appendix A, Fig. 8a), but not for IsoLoss (Fig. 8b). However, for IsoLoss larger weight decay resulted in lower magnitudes of all eigenvalues. Hence, IsoLoss reduces the impact of weight decay on eigenvalue recruitment.

4 Discussion

We provided a comprehensive analysis of the SSL representational dynamics in the eigenspace of closed-form linear predictor networks (i.e., DirectPred and DirectCopy) for both the Euclidean loss and the more commonly used cosine similarity. Our analysis revealed how asymmetric losses prevent representational and dimensional collapse through implicit variance regularization along orthogonal eigenmodes, thereby formally linking predictor-based SSL with explicit variance regularization approaches [6, 7, 14]. Our work provides a theory framework which further complements the growing body of work linking contrastive and non-contrastive SSL [24, 35-38].

We empirically validated the key predictions of our analysis in linear and nonlinear network models on several datasets, including CIFAR-10/100, STL-10, and TinyImageNet. Moreover, we found that the eigenvalues of the predictor network act as learning rate multipliers, causing anisotropic learning dynamics. We derived Euclidean and cosine IsoLosses, which counteract this anisotropy and enable closed-form linear predictor methods to work without an EMA target network, thereby further consolidating its presumed role in boosting small eigenvalues [17].

To our knowledge, this is the first work to comprehensively characterize asymmetric SSL learning dynamics for the cosine distance metric widely used in practice. However, our analysis rests on several assumptions. First, the analytic link through the NTK between gradient descent on parameters and the representational changes is an approximation in nonlinear networks. Moreover, we assumed Gaussian i.i.d inputs for proving Theorems 1 and 2. Although these assumptions generally do not

hold in nonlinear networks, our analysis qualitatively captures their overall learning behavior and predicts how networks respond to changes in the stop-grad placement.

In summary, we have provided a simple theoretical explanation of how asymmetric loss configurations prevent representational collapse in SSL and elucidate their inherent dependence on the placement of the stop-grad operation. We further demonstrated how the eigenspace framework allows crafting new loss functions with a distinct impact on the SSL learning dynamics. We provided one specific example of such loss functions, IsoLoss, which equalizes the learning dynamics in the predictor's eigenspace, resulting in faster initial learning and improved stability. In contrast to DirectPred, IsoLoss learns stably without an EMA target network. Our work thus lays out an effective framework for analyzing and developing new SSL loss functions.

Acknowledgments and Disclosure of Funding

This project was supported by the Swiss National Science Foundation [grant numbers PCEFP3_202981 and TMPFP3_210282], by EU's Horizon Europe Research and Innovation Programme (grant agreement number 101070374) funded through SERI (ref 1131-52302), and the Novartis Research Foundation. The authors declare no competing interests.

References

Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems , 33:21271-21284, 2020.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems , 33:9912-9924, 2020.

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision , pages 9650-9660, 2021.

Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15750-15758, 2021.

Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Selfsupervised learning via redundancy reduction. In International Conference on Machine Learning , pages 12310-12320. PMLR, 2021.

Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning. arXiv preprint arXiv:2105.04906 , 2021.

Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Mike Rabbat, and Nicolas Ballas. Masked siamese networks for label-efficient learning. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXXI , pages 456-473. Springer, 2022.

Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning? Advances in neural information processing systems , 33:6827-6839, 2020.

Jane Bromley, Isabelle Guyon, Yann LeCun, Eduard Säckinger, and Roopak Shah. Signature verification using a" siamese" time delay neural network. Advances in neural information processing systems , 6, 1993.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning , pages 9929-9939. PMLR, 2020.

Ting Chen, Calvin Luo, and Lala Li. Intriguing properties of contrastive losses. Advances in Neural Information Processing Systems , 34:11834-11845, 2021.

Manu Srinath Halvagal and Friedemann Zenke. The combination of Hebbian and predictive plasticity learns invariant object representations in deep sensory networks. bioRxiv , accepte, 2022. doi: 10.1101/2022.03.17.484712. URL https://www.biorxiv.org/content/10. 1101/2022.03.17.484712v2 .

Aleksandr Ermolov, Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening for self-supervised representation learning. In International Conference on Machine Learning , pages 3015-3024. PMLR, 2021.

Xi Weng, Lei Huang, Lei Zhao, Rao Anwer, Salman H Khan, and Fahad Shahbaz Khan. An investigation into whitening loss for self-supervised learning. Advances in Neural Information Processing Systems , 35:29748-29760, 2022.

Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning , pages 10268-10278. PMLR, 2021.

Xiang Wang, Xinlei Chen, Simon S Du, and Yuandong Tian. Towards demystifying representation learning with non-contrastive self-supervision. arXiv preprint arXiv:2110.04947 , 2021.

Kang-Jun Liu, Masanori Suganuma, and Takayuki Okatani. Bridging the gap from asymmetry tricks to decorrelation principles in non-contrastive self-supervised learning. Advances in Neural Information Processing Systems , 35:19824-19835, 2022.

Chenxin Tao, Honghui Wang, Xizhou Zhu, Jiahua Dong, Shiji Song, Gao Huang, and Jifeng Dai. Exploring the equivalence of siamese self-supervised learning via a unified gradient framework. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14431-14440, 2022.

Pierre H Richemond, Allison Tam, Yunhao Tang, Florian Strub, Bilal Piot, and Felix Hill. The edge of orthogonality: A simple view of what makes byol tick. arXiv preprint arXiv:2302.04817 , 2023.

Pierre H Richemond, Jean-Bastien Grill, Florent Altché, Corentin Tallec, Florian Strub, Andrew Brock, Samuel Smith, Soham De, Razvan Pascanu, Bilal Piot, et al. Byol works even without batch statistics. arXiv preprint arXiv:2010.10241 , 2020.

Xiao Wang, Haoqi Fan, Yuandong Tian, Daisuke Kihara, and Xinlei Chen. On the importance of asymmetry for siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16570-16579, 2022.

Chaoning Zhang, Kang Zhang, Chenshuang Zhang, Trung X Pham, Chang D Yoo, and In So Kweon. How does simsiam avoid collapse without negative samples? a unified understanding with self-supervised contrastive learning. arXiv preprint arXiv:2203.16262 , 2022.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent. Advances in Neural Information Processing Systems , 32:8572-8583, 2019. URL https://proceedings.neurips.cc/paper/2019/hash/ 0d1a9651497a38d8b1c3871c84528bd4-Abstract.html .

Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv:1312.6120 [cond-mat, q-bio, stat] , December 2013. URL http://arxiv.org/abs/1312.6120 .

Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:2110.09348 , 2021.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.

Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics , pages 215-223. JMLR Workshop and Conference Proceedings, 2011.

Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N , 7(7):3, 2015.

Victor Guilherme Turrisi da Costa, Enrico Fini, Moin Nabi, Nicu Sebe, and Elisa Ricci. sololearn: A library of self-supervised methods for visual representation learning. Journal of Machine Learning Research , 23(56):1-6, 2022. URL http://jmlr.org/papers/v23/21-1155. html .

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

James B Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Albrecht. On the stepwise nature of self-supervised learning. arXiv preprint arXiv:2303.15438 , 2023.

Quentin Garrido, Yubei Chen, Adrien Bardes, Laurent Najman, and Yann Lecun. On the duality between contrastive and non-contrastive self-supervised learning. arXiv preprint arXiv:2206.02574 , 2022.

Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning recover global and local spectral embedding methods. arXiv preprint arXiv:2205.11508 , 2022.

Yann Dubois, Stefano Ermon, Tatsunori B Hashimoto, and Percy S Liang. Improving selfsupervised learning by characterizing idealized representations. Advances in Neural Information Processing Systems , 35:11279-11296, 2022.

Ashwini Pokle, Jinjin Tian, Yuchen Li, and Andrej Risteski. Contrasting the landscape of contrastive and non-contrastive learning. arXiv preprint arXiv:2203.15702 , 2022.

A Supplementary Material

a

0

100

200

Sorted index

1

0

1

Alignment

0

100

200

Sorted index

500

1000

Epoch

L

cos

L

euc

b

a

0

1000

Steps

0

1

λ

m

λ

1

λ

2

L

euc

0

1000

Steps

0

1

λ

m

L

cos

b

a

L

euc

noSG

0

1000

Iteration

10

-

1

10

0

λ

b

L

euc

noPred

0

100000

Iteration

L

euc

0

1000

Iteration

d

L

euc

iso

0

1000

Iteration

e

L

noSG

0

1000

Iteration

10

-

1

10

0

λ

f

L

noPred

0

100000

Iteration

g

L

0

1000

Iteration

h

L

iso

0

1000

Iteration

c

0

500

1000

Epoch

10

2

10 0

10 2

10 4

with EMA

0

500

1000

Epoch

w/o EMA

a

10

3

10 0

10 3

wd

=4×10 4

wd

=10 5

wd

=10 6

0

500

1000

Epoch

10

3

10 0

10 3

0

500

1000

Epoch

0

500

1000

Epoch

b

B Proofs

Lemma 1. (Euclidean and cosine loss in the predictor eigenspace) Let W P be a linear predictor set according to DirectPred with eigenvalues λ m , and ˆ z the representations expressed in the predictor's eigenbasis. Then the asymmetric Euclidean loss L euc and cosine loss L cos can be expressed as:

Proof. Under DirectPred, the predictor is a symmetric matrix with eigendecomposition W P = UDU ⊤ . Since U is an orthogonal matrix, we also have UU ⊤ = I so that we can simplify the losses as follows:

where we used the fact that U is orthogonal and therefore does not change the Euclidean norm. ˆ z = U ⊤ z is the representation rotated into the eigenbasis.

Lemma 2. (Learning dynamics in a rotated basis) Assuming that a given loss L is optimized by gradient descent on the parameters of a neural network with network outputs z , a given orthogonal transformation ˆ z = U ⊤ z and learning rate η , then the rotated representations ˆ z evolve according to the dynamics:

where ˆ Θ t ( X , X ) = ∇ θ ˆ Z∇ θ ˆ Z ⊤ is the empirical NTK expressed in the rotated basis.

Proof. Let θ be the parameters of the neural network. Then we obtain the representational dynamics using the chain rule in the continuous-time gradient-flow setting [2]:

The above is a reiteration of the derivation of Eq. (2) given by Lee et al. [2], with an additional orthogonal transformation on the network outputs.

We proceed by proving the following Lemma which we will use in our proofs of Theorems 1 and 2.

Lemma 3. The NTK for a linear network is invariant under orthogonal transformations of the network output.

Proof. We first note that for a linear network, the parameters θ are just the feedforward weights W . Therefore, for any orthogonal transformation U of the network output:

where ⊗ is the Kronecker product resulting from the fact that every input vector component appears in the update once for each output component.

We now study ˆ Θ t ( X , X ) , the transformed empirical NTK (cf. Lemma 2). The ( M × M ) diagonal blocks in the full ( M | D | × M | D | ) empirical NTK ˆ Θ t ( X , X ) correspond to single samples and the off-diagonal blocks are cross-terms between samples, where | D | denotes the size of the training dataset and M the dimension of the outputs. We can develop a generic expression for each ( M × M ) block ˆ Θ t ( x i , x j ) corresponding to the interactions between samples i and j as:

where we have used the fact that ( A ⊗ B ) ⊤ = A ⊤ ⊗ B ⊤ and ( A ⊗ B )( C ⊗ D ) = AC ⊗ BD . Here, I M is the identity matrix of size M . Noting that Eq. (19) is unchanged when U is just the identity matrix completes the proof.

Theorem 1. (Representational dynamics under L euc ) For a linear network with i.i.d Gaussian inputs learning with L euc , the representational dynamics of each mode m independently follow the gradient of the loss -∇ ˆ z L euc . More specifically, the dynamics uncouple and follow M independent differential equations:

which, after taking the expectation over augmentations yields the dynamics:

Proof. For a linear network with weights W ∈ R M × N , we have from Lemma 3 that the empirical NTK ˆ Θ( X , X ) in the orthogonal eigenbasis is equal to the empirical NTK Θ( X , X ) in the original basis. Furthermore from the proof for the lemma (see Eq. (19) above), each ( M × M ) block of the full ( M | D | × M | D | ) empirical NTK is given by:

where I M ∈ R M × M is the identity. Eq. (20) gives the total effective interaction between the samples i and j from the dataset. For high-dimensional inputs x drawn from an i.i.d standard Gaussian distribution, we have x ⊤ i x j ≈ δ ij by the central limit theorem. Therefore, in the special case of a linear network with Gaussian i.i.d inputs, the representational dynamics (Lemma 2) simplify as follows:

̸

̸

While the assumption of Gaussian i.i.d inputs is quite restrictive, we offer a generalizing interpretation here. Specifically, the above argument also holds when the inputs x are not all mutually orthogonal, but fall into P orthogonal clusters in the input dataset. Then, we would have x ⊤ i x j ≈ δ p i = p j where p i is the 'label' of the cluster corresponding to sample i . If P i is the number of all the samples with

For brevity, we proceed with the simplest case Eq. (21) in which every input is orthogonal. For L euc , the representational gradient ∇ ˆ z i L is then given by:

the same label p i , then Eq. (21) would simply be scaled to give dˆ z (1) i d t = -η P i ∇ ˆ z i L .

Noting that D t is just a diagonal matrix containing the eigenvalues λ m and dropping the sample subscript i for notational ease, we obtain for the m -th component of ∇ ˆ z i L euc :

Substituting this result in Eq. (21) gives us Eq. (8), the expression we were looking for. Finally, introducing ˆ z m ≡ E [ˆ z (1) m ] = E [ˆ z (2) m ] as the expectation over augmentations, we find that each eigenmode evolves independently in expectation value as:

Theorem 2. (Representational dynamics under L cos ) For a linear network with i.i.d Gaussian inputs trained with L cos , the dynamics follow a system of M coupled differential equations:

̸

and reach a regime in which the eigenvalues are comparable in magnitude. In this regime, the expected update over augmentations is well approximated by:

̸

Proof. We can retrace the steps from the proof for Theorem 1 until Eq. (21):

∇ ˆ z i L is a vector of dimension M . Ignoring the sample subscript i for simplicity, and focusing on the m -th component of ∇ ˆ z i L , we get:

̸

̸

proving Eq. (10). Assuming sufficiently small augmentations, ˆ z (1) k and ˆ z (2) k carry the same sign, and the net sign of both terms inside the parenthesis is fully determined by γ m ≡ sign(ˆ z (1) m ) . Hence, we may write:

̸

It is useful to separate out γ m in this manner because every other term in the expression is now non-negative. Then sign( γ m · dˆ z d t ) = sign(ˆ z m · dˆ z m d t ) tells us whether ˆ z m tends to increase or decrease in magnitude, as we have argued in the main text.

Asymptotic analysis. To get a handle on how the different eigenvalues influence each other, we consider two important limiting cases. First, we consider the asymptotic regime dominated by one eigenvalue, and show that it tends towards a more symmetric solution in which the gap between different eigenvalues decreases. Second, we derive asymptotic expressions for the near-uniform regime in which all eigenvalues are comparable in size and show that this solution tends toward the uniform solution (cf. Eq. (11)).

To facilitate our analysis, we define each mode's relative contribution χ m ≡ | ˆ z m | ∥ ˆ z ∥ and evaluate Eq. (10) taking the expectation value over augmentations:

̸

̸

In the second equality we used the fact that the expectation value taken over augmentations is conditioned on the input sample, which makes them conditionally independent.

One dominant eigenvalue. First, we consider the low-rank regime in which one eigenvalue dominates. Without loss of generality, we assume λ 1 ≫ λ k ∀ k = 1 . We then have:

̸

̸

Plugging these values into Eq. (22) gives the following dynamics for the dominant eigenmode:

̸

̸

These updates are always opposite in sign to the representation component, which corresponds to decaying dynamics for the leading eigenmode because γ 1 dˆ z 1 d t < 0 .

For all other modes we have:

̸

so that γ m dˆ z m d t > 0 , i.e, the updates have the same sign as the representation component, which corresponds to growth dynamics. In other words: The dominant eigenvalue 'pulls all the other eigenvalues up,' a form of implicit cooperation between the eigenmodes. We also note that the non-dominant eigenmodes increase at a rate proportional to ϵ , whereas the dominant eigenmode decreases at a slower rate proportional to ϵ 2 . Thus, for sensible initializations with at least one large and many small eigenvalues, the modes will tend toward an equilibrium at some non-zero intermediate value, without a dominant mode. Next we study this other limiting case in which all eigenvalues are of similar size.

̸

̸

Near-uniform regime. To study the dynamics in a near-uniform regime, we note that all χ m are of order O (1) in ˆ z m , whereas the eigenvalues λ m are of order O (ˆ z 2 m ) . In this setting, the effect of the eigenvalue terms λ m on the dynamics is stronger than the χ m terms which are bounded between 0 and 1. With a sufficiently high-dimensional representation, all χ m terms will be centered around 1 / √ M . Based on these observations, we may make the simplifying assumption that the contributions are all approximately equal, i.e, χ i = χ for all i . Substituting this value in Eq (22) gives:

̸

Finally, substituting for χ , which by assumption are all approximately equal:

,

and absorbing back the sign from γ m , we obtain the approximate dynamics in Eq (11):

̸

C Derivation of idealized learning dynamics for different loss variations

C.1 Removing the stop-grad from the Euclidean loss L euc

Omitting the stop-grad operator from L euc gives:

Tracing the steps to prove Theorem 1 and assuming Gaussian i.i.d inputs for a linear network, we write:

which results in decaying representations and thus collapse.

C.2 Removing the stop-grad from the cosine loss L cos

Following the same arguments as above, omitting the stop-grad operator from L cos gives:

̸

so that, when taking the expectation value over augmentations, the dynamics follow:

̸

̸

In the asymptotic regime with dominant eigenvalue λ 1 , we get the dynamics:

̸

̸

Thus, all eigenmodes diverge because γ m dˆ z m d t > 0 .

Similarly, we find divergent dynamics when starting in the near-uniform regime:

̸

selecting the terms with the highest power in the eigenvalues.

Thus, omission of stop-grad precludes successful representation learning for both the Euclidean and the cosine loss, but due to different mechanisms. Euclidean loss yields collapse, whereas the cosine loss succumbs to run-away activity.

C.3 Removing the predictor from the Euclidean loss L euc

To analyze the representational dynamics in the absence of the predictor network, we consider L euc noPred :

̸

The dynamics resulting from this loss function are a special case of the dynamics derived in Theorem 1 with all the eigenvalues equal to one ( λ k = 1 ). In particular Eq. (8) becomes:

which evaluates to 0 under expectation over augmentations. Hence there is no learning without the predictor.

C.4 Removing the predictor from the cosine loss L cos

Similarly, when we remove the predictor from L cos it yields:

so that Eq. (10) becomes:

̸

Here, the near-uniform approximation (Eq. (11)) of ignoring the differences in χ between different eigenmodes is not valid. This is because the λ -terms are no longer present, and the effects of the χ -terms on the dynamics cannot be treated as negligible. In particular, setting W P = I in the dynamics derived in Theorem 2 would yield zero dynamics. However taking the expectation of Eq. (24) over augmentations yields the non-zero dynamics:

̸

which consists of terms of order O ( 1 ˆ z m ) and O (1) in ˆ z m , hinting at slower dynamics compared to Eq. (10). To study these granular effects, we would need to explicitly model the effect of the augmentations, for which we do not have a good statistical model. In lieu of deriving these dynamics analytically, we make an observation which restricts the possible dynamical behavior. Specifically, the sum of the eigenvalues remains constant throughout training. To show this, we begin by writing out the expression for the derivative over time of the sum of all the eigenvalues:

We can derive the term inside the expectation by adding up the dynamics given by Eq. (25) for all the different eigenmodes:

̸

̸

proving that d d t ∑ m λ m = 0 , i.e, the sum of the eigenvalues is conserved. This precludes collapsing dynamics where all eigenvalues go to zero as well as diverging dynamics where at least one eigenvalue diverges.

C.5 Isotropic losses for equalized convergence rates

In Expressions (9) and (11) we see that the overall learning dynamics have a quadratic dependence on the eigenvalues with a root near collapsed solutions, which causes these modes to learn slower. We reasoned that this anisotropy could be detrimental for learning. To address this issue, we sought to derive alternative loss functions that encourage isotropic learning dynamics for all modes.

C.5.1 Euclidean IsoLoss.

We start by deriving an IsoLoss function for the Euclidean case L euc . To avoid the unwanted quadratic dependence, we first note that we would like to arrive at the following expression for the dynamics:

By recalling the Euclidean loss and corresponding dynamics:

we note that the leading λ m term has no influence on the overall sign of the dynamics, and is introduced by the second step in the chain rule:

Based on this realization we see that this second step needs to be modified. To that end, we start with the desired derivative:

and see that several loss functions are possible. The one we have reported in Eq. (15) we derived by applying an appropriate stop-grad while integrating:

to give:

Another alternative loss with the same desired isotropic learning dynamics, but using a different placement of the stop-gradient operators, is given by:

C.5.2 Cosine Similarity IsoLoss.

Since most practical SSL approaches rely on cosine similarity, which suffers from a similar anisotropy of the learning dynamics, we sought to find IsoLosses in this setting. With the same goal as above, we would like to arrive at the dynamics:

starting from the cosine loss and corresponding dynamics:

The IsoLoss in this case can be derived by noting how λ m arises in each of the two terms in Eq. (27), and engineering an alternative loss function corresponding to each term separately.

In the first term, λ m arises from the partial derivative of the numerator λ m ˆ z (1) m SG(ˆ z (2) m ) in the original loss (Eq. (26)). This can be remediated by using ˆ z (1) m SG(ˆ z (2) m ) as the numerator instead.

In the second term in Eq. (27), λ 2 m arises from the partial derivative of ∥ D ˆ z (1) ∥ = √ ∑ k ( λ k ˆ z (1) m ) 2 in the denominator. We can reduce λ 2 m to λ m by instead taking the partial derivative of ∥ D 1 / 2 ˆ z (1) ∥ = √ ∑ k ( λ 1 / 2 k ˆ z (1) m ) 2 . Putting these insights together, we arrive at the desired partial derivative:

and the integrated IsoLoss in eigenspace:

Rotating all terms back to the original space gives the desired IsoLoss for cosine similarity as reported (Eq. (17)):

D Experimental methods

Self-supervised pretraining. We used the CIFAR-10, CIFAR-100 [3], STL-10 [4], and TinyImageNet [5] datasets for self-supervised learning with a ResNet-18 [6] encoder and the SimCLR set of transformations [7]. We also adopted several modifications of ResNet-18 which have been proposed to deal with the low resolution of the images in these datasets [7]. The ResNet modifications comprise using 3 × 3 convolutional kernels instead of 7 × 7 kernels and skipping the first max-pooling operation. The Solo-learn library [8] also provides specialized sets of augmentations that work well for these datasets, which we adopted as well. The configurations we used for each dataset are summarized in Table 4. We used BatchNorm in the backbone and the projector multi-layer perceptron (MLP) in the hidden layer for all methods. For BYOL, we included BatchNorm also in the hidden layer of the predictor MLP.

We used a projection dimension of 256 for the projection MLP using one hidden layer with 4096 units, and the same architecture for the nonlinear predictor for the BYOL baseline. For networks using EMA target networks, we used the LARS optimizer with learning rate 1.0 whereas for networks without the EMA, we used stochastic gradient descent with momentum 0.9 and learning rate 0.1. Furthermore, we used a warmup period of 10 epochs for the learning rate followed by a cosine decay schedule and a batch size of 256. We also used a weight decay 4 × 10 -4 for the closed-form predictor models and 10 -5 for the nonlinear predictor models. For evaluation, we removed the projection MLP and used the embeddings at the pooled output of the ResNet convolutional layers following standard practice. For the EMA, we started with τ base = 0 . 99 and increased τ EMA to 1 with a cosine schedule exactly following the configuration reported in [9]. For DirectPred, we used α = 0 . 5 and τ = 0 . 998 for the moving average estimate of the correlation matrix updated at every step.

Linear evaluation protocol. We reported the held-out classification accuracy on the test sets for CIFAR-10/100 and STL-10, and the validation set for TinyImageNet, after online training of the gradient-isolated linear classifier on each labeled example in the training set during pretraining.

Compute resources All simulations were run on an in-house cluster consisting of 5 nodes with 4 V100 NVIDIA GPUs each, one node with 4 A100 NVIDIA GPUs, and one node with 8 A40 NVIDIA GPUs. Runs on CIFAR-10/100 took about 8 hours each, and the STL-10 and TinyImagenet runs took about 24 hours each.

References

Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning , pages 10268-10278. PMLR, 2021.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent. Advances in Neural Information Processing Systems , 32:8572-8583, 2019. URL https://proceedings.neurips.cc/paper/2019/hash/ 0d1a9651497a38d8b1c3871c84528bd4-Abstract.html .

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.

Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics , pages 215-223. JMLR Workshop and Conference Proceedings, 2011.

Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N , 7(7):3, 2015.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Victor Guilherme Turrisi da Costa, Enrico Fini, Moin Nabi, Nicu Sebe, and Elisa Ricci. sololearn: A library of self-supervised methods for visual representation learning. Journal of Machine Learning Research , 23(56):1-6, 2022. URL http://jmlr.org/papers/v23/21-1155.html .

Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems , 33:21271-21284, 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis article builds on the theory of non-constrastive (nc) self-supervised learning (SSL), with methods such as BYOL or SimSiam. Contrary to the existing litterature, they study directly the cosine similarity loss used rather than a Euclidian loss on the eigenspace, using NTK dynamics. They show that collapse is avoided in this case with implicit variance reduction. With their analysis, they find an undesirable anisotropy which they fix with a new family of losses that outperforms the state-of-the-art, especially without needing an Exponiential Moving Average (EMA) target network.\n\nSTRENGTHS:\nThe theory provides strong results in the linear framework, explaining the eigenvalues dynamics of both euclidian and cosine losses with and without projectors and stop gradients. The use of the NTK dynamics and the analysis of the cosine similarity is novel to me.\n\nThe findings that the eigenvalues affect learning as a learning rate provide a good justification for the new IsoLosses, which gives state-of-the-art results. In particular, the removal of the EMA and the increased dimensionality of the losses ensures their consistency.\n\nWEAKNESSES:\nOne of the main claims of the article is the analysis of the cosine similarity loss. However in that case, if the dynamics are quite different as claimed by the authors, I find the use of the result of Tian et al. that the predictor eigenspace aligns with the one of the correlation matrix of the representations surprising since it was established with a Euclidian loss. I would have hoped atleast a comment on this, if not an analysis of the alignment as done in Tian et al.\n\nI was awaiting an analysis of the eigenvalues similar to Figure 2 for a real network, and not only a comparison with other losses. As such, the analysis of the IsoLoss feels a bit disconnected from the theoretical analysis. Why do the dimensionality increase? Is it because small eigenvalues do not collapse due to their low learning rates ? Do the eigenvalues evolve during training similarly to the linear case, for BYOL/SimSiam and IsoLoss ? (and their euclidian counterparts)\n\nTwo important notions in the theory of NC-SSL are the EMA and the weight decay. Since an important improvement of the method is the removal of the need of the EMA of the target network, I would find natural to mention the effect of the EMA in the dynamic. \nSimilarly, the results of Wang et al. indicate the importance of the role of weight decay to collapse unwanted eigenvalues. I am surprised that the role of weight decay is not mentioned anywhere, especially despite its use in the implementation.\n\nThe formulation of the Isoloss (in particular the cosine similarity one) is very hard to interpret. However, it is logical due to the way it was found.\n\nThis is an interesting theoretical article despite some flaws and questions raised, and I am ready to raise my score if my questions are answered.\n\nQUESTIONS:\nIf would define $\\hat z = U z$ a bit more clearly maybe before Lemma 1.\n\nEquation 3: There is a dimensional issue with the diagonal matrix on the right with the product \"$zD$\" if I am not mistaken. Shouldn't $D$ be on the left?\n\nDefinition 1: I find the different notations often a bit unclear, but most importantly is the sudden usage of the time $t$ index which is not consistent. It is present at line 114 but not at line 111, and appears only on part of the variables. Similarly for lemma 2, where $\\hat z_t$ loses the time index again. The use of bold variables is also often inconsistent.\n\nTable 2: I would have expected the inclusion of DirectCopy and of SimSiam in the comparison.\n\nLIMITATIONS:\nThe authors address adequately the limits of the theory they established, although they do not talk about the effect of EMA and weight decay in their theory or their hypothesis of DirectPred which was established with an Euclidian loss.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe learning dynamics of non-contrastive self-supervised learning is an important problem to understand how these methods avoid collapse without using negative samples. In this paper, the authors provide a rigorous analysis of this problem on a simple linear network with Gaussian inputs, especially the difference between Euclidean and cosine loss functions. It is demonstrated that both losses have implicit regularization effects on the variance of representation, and the role of predictor and stop-grad operations are thoroughly investigated. Based on these insights, the authors propose an isotropic loss to equalize the convergence rate and lead to a better performance in various settings.\n\nSTRENGTHS:\nThe theoretical analysis in this paper is novel and clearly improves the understanding of how loss function, stop grad and projector affect the performance of non-contrastive self-supervise learning. Based on the mathematical understanding, the authors propose a new loss function that could beat the baseline method on different datasets, which illustrates the validity and power of the proposed theory in practice.\n\nWEAKNESSES:\nThe major weakness of this paper is the strong assumptions that require the network to be linear and the input to be iid Gaussian. For real-world datasets and practical network structures, it is not clear if those insights are still valid.\n\nQUESTIONS:\n1. In Table 1, it is predicted that $L_{noPred}$ converges slower than $L$, however, in Figure 2 there is an eigenvalue in $L$ that converge much slower, there seems to be a discrepancy between theory and experiments.\n\n2. In theorem 1 and 2, the dynamics are taken to be the expectation of original dynamics over the distribution of augmentation. It seems that after taking expectations, it is essentially equivalent to inputting the original images without augmentations, which may lose an essential part of contrastive learning.\n\n3. In Theorem 2, it is assumed that the eigenvalues are of comparable magnitude, is this assumption required only on the initial value, or on the whole trajectory? In the latter case, it seems not appropriate to draw conclusions on the behavior of eigenvalues.\n\nLIMITATIONS:\nNo certain limits.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work analyzes the learning dynamics of non-contrastive SSL approaches such as BYOL and Simsiam. Based on the proposed theory, the authors analyze the how the stop-grad and predictor module affect the learning dynamics. Importantly, the authors design a theoretically inspired loss and gain improvement on classification tasks.\n\nSTRENGTHS:\n1. This work provides solid analysis and discussion on the learning dynamics of non-contrastive SSL algorithm, which contributes to our understanding of the algorithm. Especially, they consider the widely used cosine similarity loss, while existing works only discuss Euclidean loss. They also propose an isotropic loss which is inspired by the theory.\n2. The simulation experiments helps understanding the proposed theory.\n\nWEAKNESSES:\n1. Lack of basic introduction to the used techniques, e.g., the neural tangent kernal (NTK).\n\n2. Lack of awareness of dividing the article into paragraphs. The preliminary, theoretical analysis and the simulation experiment are all mixed in chapter 2, which might be confusing for readers.\n\n3. The accuracy for the baseline model (BYOL) is too low. In the well-known SSL repository, solo-learn [1], the top-1 accuracy of BYOL (R-18, 1000ep) is 92.58 on CIFAR-10, and 70.46 on CIFAR-100. However, in this paper the baseline is only 89.4 and 61.1. I am afraid that this experiment cannot justify the effectiveness of the proposed isotropic loss. \n\n4. This paper mainly provides theory to understand how non-contrastive SSL algorithms work (4 pages), and the design of the isotropic loss is inspired by the theory (only half a page). I think the title of this paper is not accurate enough.\n\nQUESTIONS:\n1. What does \"Gaussian i.i.d inputs\" in line 133 refer to specifically? To what extent does this setting fit the real-world scenario?\n\nOthers:\nLine 68: Replace z^{(1/2)} with z^{(1)}, z^{(2)} to avoid ambiguity\nLine 146: typo: \"lammata\"\n$\\ell_{iso}$ first appears in Fig 2. and Tab 1., which is before where it is defined.\n\nIf the author solves my concerns, then I would be willing to raise the rating.\n\n[1] https://github.com/vturrisi/solo-learn\n\nLIMITATIONS:\nSee weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-09"}