What shapes the loss landscape of self supervised learning?
Authors: Liu Ziyin
Ekdeep Singh Lubana
Masahito Ueda
Hidenori Tanaka
Venue: ICLR 2023 poster (2023)
Keywords: loss landscape
self-supervised learning
collapse

Abstract: Prevention of complete and dimensional collapse of representations has recently become a design principle for self-supervised learning (SSL). However, questions remain in our theoretical understanding: When do those collapses occur? What are the mechanisms and causes? We answer these questions by deriving and thoroughly analyzing an analytically tractable theory of SSL loss landscapes. In this theory, we identify the causes of the dimensional collapse and study the effect of normalization and bias. Finally, we leverage the interpretability afforded by the analytical theory to understand how dimensional collapse can be beneficial and what affects the robustness of SSL against data imbalance.

Published as a conference paper at ICLR 2023
WHAT SHAPES THE LOSS LANDSCAPE OF SELF SUPER- VISED LEARNING?
Liu Ziyin 12.33 Ekdeep Singh Lubana?. Masahito Ueda 15.6, Hidenori Tanaka 2:3
Department of Physics, The University of Tokyo, Tokyo, Japan
2 Physics Informatics Laboratoiies NTT Research, Inc., Sannyvale, CA, USA Center for Brain Science, Harvard University, Cambridge, USA EECS Department, University of Michigan, Ann Arbor, USA
Institute for Physics of Intelligence, The University of Tokyo, 7-3-1 Hongo, Bunkyo-ku, Tokyo RIKEN Center for Emergent Matter Science (CEMS), Wako, Saitama, Japan
ABSTRACT
Prevention of complete and dimensional collapse of representations has recently become design principle for self-supervised learning (SSL). However, questions remain in our theoretical understanding: When do those collapses occur? What are the mechanisms and causes? We answer these questions by deriving and thor- oughly analyzing an analytically tractable theory of SSL loss landscapes. In this theory. we identify the causes of the dimensional collapse and study the effect of normalization and bias. Finally, we leverage the interpretability afforded by the analytical theory to understand how dimensional collapse can be beneficial and what affects the robustness of SSL against data imbalance.
INTRODUCTION
Self-supervised learning (SSL) methods have achieved remarkable success in learning good repre- sentations without labeled data (Chen et al., 2020b). Loss functions used in such SSL techniques promote representational similarity between pairs of related samples while using explicit penalties (Chen et al., 2020a; He et al.. 2020; Zbontar et al., 2021; Caron et al., 2020) or asymmetric dynam- ies (Caron et al., 2021; Grill al.. 2020: Chen and He. 2021) to ensure that the distance between unrelated samples remains large. In practice, however, SSL training often experiences the phe- nomenon of dimensional collapse (Jing et al., 2021; Tian al., 2021; Pokle et al., 2022), where the learned representation spans low dimensional subspace of the overall available space. In the extreme case, this failure mode instantiates as complete collapse, where the learned representation becomes zero-rank, and informative features can be extracted.
Prior work has primarily positioned such collapses in SSL as enemies of learning, arguing that they can negatively impact downstream task performance (Zbontar et al., 2021; Jing al., 2021: Bardes et al., 2021). However, recent work by Cosentino et al. (2022) empirically demonstrates otherwise: quality of representations can be improved when there is degree of collapse. These conflicting results indicate that despite extensive empirical explorations, a gap remains in our understanding of the collapse phenomenon in SSL training. We argue that this gap is due the lack of theoretical framework to analyze the mechanisms promoting collapsed representations. We aim to close this gap by carefully studying the loss landscapes of SSL.
In this work, we analytically solve the effective landscapes of linear models trained on several pop- ular losses used in self-supervised learning, including InfoNCE (Oord et al., 2018), Normalized (NT-xent) (Chen et al., 2020a), Spectral Contrastive Loss (HaoChen
Temperature Cross-Entrope al., 2021), and Barlow Twins VICReg (Zbontar et al., 2021; Bardes et al., 2021). The main thesis of this work is: the local geometry of the SSL landscapes around the origin crucially decides the learning behavior of SSL models. Technically, we show that
1. the interplay between data variation and data augmentation determines the geometry of the loss; the geometry the loss explains when dimensional collapse can helpful and why certain SSI. losses are robust against data imbalance, but not the others.
To the best of our knowledge, our work is the first to study the landscape causes of collapse in SSL thoroughly.
Work done during an internship Physics & Informatics Laboratories, NTT Research.

---

Published as a conference paper at ICLR 2023
Stronger augmentation
L
Larger
data variation
12
T2
T2
(a) An eigenmode
(b) No collapse
(c) Dimensional collapse (d) Complete collapse
Landscape in self-super learning (SSL). SSL losses generally depend only on the relative angle between pairs of network outputs (e.g, ((x)"{(=')) Thus, the landscapes with linear network (f()) Wæ) have global rotational symmetry and symmetric about the origin. Our theory finds that the local stability the origin decides the collapse, and larger data variation (green) prevents collapse, while strong data augmentation (red) can promote collapse. We plot the loss for a toy lincar model with a diagonal weight matrix diag(ri, r2). (a) The 1d landscape when fixing one of the parameter. (b-d) The 2d landscape. (b) No collapse: the origin is unstable local maximum, and surrounding local minima avoid collapse. The dimensionally collapsed solutions the saddle points. (c) Dimensional collapse: the value of wi for all stable fixed points collapses to zero. (d) Complete collapse: origin becomes the isolated local minimum.
Figure 1:
RELATED WORKS
SSL and Collapses. On the one hand, prior literature has often argued collapse as a harmful phe- nomenon that can deteriorate downstream task performance (Jing al., 2021; Zbontar al., 2021). Preventing such collapsed representations is a frequently discussed topic in literature (Hua et al., 2021; Jing et al., 2021; Pokle et al., 2022; Tian et al., 2021) and has motivated the design of several SSL techniques (Zbontar al., 2021; Bardes et al., 2021; Ermolov et al., 2021). On the other hand, Cosentino al. (2022) empirically showed that dimensional collapses under strong augmentations could significantly improve generalization performance. Our work demystifies these conflicting results by finding analytic solutions to loss landscapes of several standard SSL techniques.
Theoretical Advances in SSL. Recently, several advances have been made towards understand- ing the success of SSL techniques from different perspectives: e.g. learning theory (Arora et al., 2019; Saunshi al., 2022; Nozawa and Sato, 2021; Wei et al., 2021), information theory (Tsai et al., la;;;;;; Tosh et al., 2021), causality and data-generating processes (Zimmerman et al., 2021; Kugelgen al., 2021; Trivedi al., 2022: Tian al., 2020; Mitrovic et al., 2020; Wang al., 2022), dynamics (Wang and Isola, 2020; Tian et al., 2021; Tian, 2022; Wang and Liu, 2021; Simon et al., 2023), and loss landscapes (Pokle al., 2022). These advances have unveiled practically useful properties of SSL, such as robustness to dataset imbalance (Liu et al., 2021) and principled solutions to avoid spurious correlations (Robinson et al., 2021).
The work by Jing et al. (2021) is the closest to ours in problem setting. In that paper, the authors focused on studying the linearized learning dynamics and suggested that competition between the feature signal strength and augmentation strength can lead to dimensional collapse. In contrast, our focus on the landscape and our result implies that this antuure-ugmmentation competition on its own is insufficient to cause dimensional collapse. In fact, show that there will be collapse in the setting studied by Jing et al. (2021).
A LANDSCAPE THEORY OF SELF-SUPERVISED-LEARREE
This section presents the main theoretical results. Let { N be dataset with N data points. For every data point â. we augment it with an i.i.d. noise e such that x := + To be concrete, we start with considering the standard contrastive loss, InfoNCE (Oord et al., 2018):
exp(-ls(a)=f(a)))P
L-
E,
log (xxxex((-1(a)=--)(()))) +xxp-))))))((( ((:))1))))
(1)
where f(x) dRd1 is the model output; all x, x' and x augmented data points for some indepen- dent additive noise such that E. = à = =,[='] [x = X. We decompose the model output into a general function (xx) e d and the last-layer weight matrix W Rdixdo; f(x)) W @xx).. The covariance of (()) is A {z()())((2)")) and the covariance of the aata-augmented penul- timate layer representation is Σ [=(((=)((1)) The effect of data augmentation the learned
2

---

Published as a conference paper at ICLR 2023
representation is captured through symmetric matrix C-=E-An For general ø, the eigenvalues of C' can be either positive or negative. When & is the identity mapping, Ao becomes the empirical data covariance. C' becomes positive semi denindefinite and is the covariance of the noise c, and Σ is the covariance of the augmented data. In some sense, this loss function captures the essence of SSL: the numerator encourages the representation (() to be closer to the representation of similar data, and the denominator encourages separation between dissimilar data.
For fixed set of noises, we can write the InfoNCE in cleaner form:
L Ex [1)(()) - -(x')) log Rx exp (-110(00--
where we used E; to denote an averaging over the training set.
In this notation, we have E, Eg [x] Eg [x] and E. [.]] = L. We first show that the expansion of the loss function around the origin takes a rather universal form. We then find analytical so- lutions to the stationary points of this landscape and study their relevance to feature learning and collapses. See Table for summary of the main results. The proofs are presented in Appendix E. For a quantitative understanding, we mainly fo- the when is the identity function. We discuss the general nonlinear case in Sec- tion 4.1.
Hessian
Dim.
Complete
InfoNCE
Ao Au Ao-cl CC/N
X
X
NT-Xent (SimCLR)
V
Spectral Contrastive
C
X
Barlow Twins + Normalization
Ao + C
X
X
X
+ bias
Weight Decay
+77
V
Table What shapes the SSL landscapes around the origin? For cach of the SSL losses, the combi- nation of data covariance (Ao), data angtenttgmnttoionn variance (C), and dataset size (N) can affect its sta- bility and thus determine presence (V) and ab- sence (X) of dimensionnllcompllte collapse (Here, a / /caas "there exists a hyperparameter setting and data distribution such that the relevant collapse hap- pens:" see section 3). Beyond collapses. the theory implies that SCI., whose landscape is formed primar- ily by data augmentation, is more robust to data im- balance than InfoNCF, which is affected primarily by the data (see section 4).
3.1
LA LANDSCAPE OF LINEAR MODEL.
We first analyze representative SSL loss func- tions and show that to leading order in W, the local geometry of SSL losses takes the following form
L= - [W _ Var[[W (x- x) (3)
distinctive feature of Eq. (3) is that its first and third-order terms vanish. This is because the loss function is invariant left rotation of W. We will that this symmetry in rotation is crucial and general feature of the SSL loss functions that allow us to treat them in a universal way. We discuss how rotation symmetry can cause collapses in nonlinear settings in Section 4.
InfoNCE. The loss function simplifies to:
= Tr[WCW" E.,0 log B&g exp (-3 (a- (=x)= }}
E
-S
Expanding the entropy term the fourth order, we obtain'
-S--E,E, W
[ []((=-))]++ I ())
(5)
This (perturbative) decomposition of entropy deserves some special attention. The entropy decom- poses into a repulsion term that second order in W, and a variance term that is fourth order in W. The first term encourages a repulsion between and its augmentation, which counteracts the effect of the energy term. The repulsion term can be decomposed into
E,Ey W - (x)====))] Tr[WCWT] TH[WA,W "WW
The first term encourages an expansion of W along the direction of the augmentation C, while the second term encourages an expansion along the directions of feature Ag. It is intriguing to see
Throughout, we || denote the T.2 norm for vectors and Frobenius for matrices.
3

---

Published as a conference paper at ICLR 2023
that the repulsion term dominates the attraction of the energy term: the motion along the direction of C' completely cancels out, and only the expansion along Ag remains. This means that to leading order, the learned representation has larger variation along the directions where the data has larger variation, which is what one naively expects. Collecting results, we have obtained the loss landscape in the neighborhood of the origin as I. -Tr[W A,W")+ 'I Var[IW(< =x)=]+ O(I]W]]).
NT-xent (SimCLR). As an additional example, we analyze Normalized Temperature Prtss-Entropy loss (NT-xent) used in SimCLR (Chen al., 2020a). Tian (2022) shows that InfoNCE can be generalized to encompass NT-xent as follows:
exp( (000((())))
L
Ee
Σ log 111
(7)
Exer exp(-IX((;)) ((x)]))) -(ae))))-()-+-pp(((((s)a)))))))) -(xp
In contrast to InfoNCE, here one of the terms in the denominator reweighted by factor of 2 O. Two interesting limits are 1, where we recover the InfoNCE loss, and a 0, where we obtain NT-xxent For general a, we refer to this loss as the weighted InfoNCE. We will see in section 3 that this weighted InfoNCE can have a mild dimensional collapse problem.
The same perturbative expansion as Eq. (4)-(6) gives
== 1 N - Tr[WA, Varll x) O(IWI1)) O(IWW))) (8)
Now, the Hessian of the origin is no longer guaranteed be negative definite. In fact, 1 C- Ao 2
0, W = 0 becomes an isolated local minimum.
Landscape Analysis. The above discussion shows that the commone loss landscapes in self- supervised contrastive learning can be reduced to an effective form in Eq. (3). The following propo- sition shows that the variance term of the loss takes a specific form when the data is Gaussian.
Proposition 1. Let the data and noise be Gaussian. Tr(WEWT WEWT]
Then, L -Tr[W -Tr[WBW"]
When the training ends, one expects the model to locate at (at least close to) a stationary point the loss. is thus important to identify all the stationary points of this loss function.
Theorem 1. Let d' := min(do, d1). Let the data and noise be Gaussian. All stationary points W of Eq. (3) satisfy W W - MPUMAU where UAU is the eigenvalue decomposition of Σ- /1-1/22-22 and M is arbitrary (masking) diagonal matrix containing only zero such that (1) Ma 0 if A& < 0 and (2) contain at most d' nonzero terms.
Additionally, if C and Ag commute, all stationary points satisfy
WTW = -W=----BB--'
(9)
where By denotes the matrix obtained by masking the eigenvalues of with M.
This stationary- point condition implies the direct cause of the dimensional collapse. Namely, dimen- sional collapse happens when the cigenvalues of the matrix B become negative. The eigenvalues of B, in turn, depend the competition between data augmentation and the data feature. Com- paring the commuting case with the noncommuting case, we see that the main difference is that when C does not commute with Ag, the augmentation representation; otherwise, augmentation only affects the eigenvalues. To focus on the most impor- tant terms, we now assume that the augmentation is well-aligned with the features such that the augmentation covariance commute with the data covariance.
also change the orientation of the learned
Assumption 1. From now on, we assume CA AOC.
For the case of weighted InfoNCE, we have that B = Au - 1-a C. Let denote the i-th eigenvalue of the and e; that of C' viewed in a predetermined order; then, the ith subspace collapses when c 2 a, namely, when the variation introduced by the noise dominates that of the original data. Importantly, this collapse is a property shared by all stationary points of the landscape, and one cannot hope to fix the problem by, say, biasing the gradient descent towards certain type of local minima. When weight decay is used, the condition for collapse becomes 1-a c + + 2 ag; it becomes easier to cause collapse when weight decay is used.
4

---

Published as a conference paper at ICLR 2023
The global minimum of the loss function is also easy to find. For all stationary points, the loss Thus, becomes more and more negative if the
function takes a simple form: -= T eigenvalues of By align with the largest eigenvalues of B. Namely, the global minimum is achieved M leaves the largest eigenvalues of B intact.
Because the stationary points contain collapsed solutions where the eigenvalues of W W are zero, one is naturally interested in how likely it is to converge to these solutions.
Proposition 2. W achieves maximum possible rank) Let m denote the number of positive eigenvalues B. Then, rank(W"W)) min(m, d²) for any local minimum.
This proposition implies that the loss landscape of contrastive SSL (with a linear model) is rather benign because all local minima must achieve maximum possible rank. In fact, this result implies that the collapses may he well controllable by carefully controlling and tuning the eigenvalues of the matrix B, which directly depends on the nature of the data augmentation we use.
3.2 LANDSCAPE WITH NORMALIZATION
It is common in practice to normalize the learned representation such that [|[())|} = When normalization applied, only the direction of the learned representation matters. While this is simple trick in practice, its implication on the landscape is poorly understood. In this section, we extend our theory to analyze the effect of normalization.
We model the effect of normalization
regularization term: R (ELl]](=)]
Luorm Eq. (3) + R.R
(10)
Note that this regularization term achieves two things simultaneously: (1) [[f())] e for all is minimizer of the loss function; (2) the regularization is invariant to any rotation of the learned representation. For linear model, we note that this condition is not entirely the same as direct normalization of the representation because it is generally impossible to achieve || = for all æ because linear model has limited expressivity. However, is generally possible to achieve the slightly weaker condition: the representation has norm 1 average. This loss function can also be seen as a mathematical model of the VICReg loss (Bardes et al., 2021), where R effectively models the variance regularization term of VICReg loss and k is its strength. This modeling is necessary because the variance term of the original VICReg is not differentiable and thus cannot be expanded. The proposed term R captures the essence of the variance term because also encourages the representation to have a constant variance. Our theory also explains why the VICReg is observed experience collapses when x not large enough. As tends infinity, this constraint will become perfectly satisfied. We thus take the infinite k limit to study the effect of normalization.
The following proposition gives a condition that all stationary points of Eq. (10) satisfy. Proposition 3. Let p(W) == TTr:==TrrWS Tr[WEWTT] B' := 2x(c2((()) E)Σ, and let A be the eigevalues of B'. Then, every stationary point of Eq. (10) satisfy WTW 1 diagonal mask of the eigenvalues of B' containing only zero or one such that (1) Ma = if A, < and (2) contain at most d' nonzero terms.
E' E-1 where M is an arbitrary
Compared with the unnormalized case, the term 2x(1 - Ea emerges due to normalization. The effect of normalization is as expected: it shrinks the norm of the model if p > 1, and it expands the model if 1, and it does not have any effect if we have already achieved م 1. Interestingly, this rescaling effect is anisotropic and stronger along the directions of larger eigenvalues of the covariance of the augmented data Σ.
The next theorem gives the explicit forn of p at the stationary points.
Proposition 4. For any stationary point W e- p(W*)= c-f-n// Barl of non-zero eigenvalues of B'y
where das is the number
1-xda
For finite K, these results suggest that collapses can still happen. For VICReg, == -Ao, and the complete collapse can happen when < ||Aa||/c|E||| this explains the experimental observation of collapses for small values of x in VICReg loss (Bardes et al., 2021).
Lastly, understand normalization, are interested in the case of к Combining Proposi-
tion 3 and 4, we have proved the following theorem, showing that the asymptotic solution converges to form independent of
5

---

Published as a conference paper at ICLR 2023
Theorem 2. Let Wx be stationary point of Eq. (10) at fixed K. Then,
1 Σ-1 Bas 2c Tr(E-1 BB
lim W W.
IM E
(11)
dm
2e Te EM emerges result of applying normalization. The effect can da
The correction term be easier to understand if we write the solution as
w W=
1 E
BM
TH
M
2c
Σ
(12)
2
dM
dM
where we have used the relation MM -ー1 M. Note the term in brackets: it subtracts the average eigenvalue of E-1 Bar from 5- Bs and shifts the remaining eigenvalues positively by 2e/das. Be- cause the eigenvalues of WWT must be positive, the following condition must hold for all solutions:
X + 2c/dM X,
(13)
where A; are the eigenvalues of Σ-1 Bas and À is its average. Namely, for the i-th dimension not to collapse, it must be smaller than the average eigenvalues by most 2c/dM. Any smaller eigen- values must collapse. Compared to the case without normalization, normalization makes collapses dependent on the relative strength of each feature and augmentation. In the following discussion, we let c 1 to simplify the discussion. We present a detailed analysis of this condition in Section D. One finds that the condition for collapse becomes heavily dependent the data structure, and there are cases where collapses become harder, and there are cases where collapses become much easier. Importantly, also becomes the case that sufficiently strong augmentation can always cause a collapse in the corresponding subspace.
Effect of Bias. Lastly, we study the effect of explicitly having a bias term: Wx Wx + b. First of all, when there is no normalization, the bias term does not affect the solution because the loss landscape is invariant to translation the learned representation. However, this effect dramatically changes if we apply normalization the same time. This is because normalization removes the translation symmetry of the effective loss, and the trivial solution W = 0, b = becomes the simplest way achieve the norm-1 constraint. Our result shows that the addition of bias dramatically affects the stationary points.
Theorem 3. Let f(x) Wx + b and [[x] = 0. Then, all stationary points W satisfy Eq. (9), subject the constraint that Tr[W [w]sc sc
Namely, the solution reverts to the case where there is normalization at all, except that the norm of the solution can no longer be larger than c. This upper bound can make collapses much easier to happen. For example, if < (a, c)/ c:) for all complete collapse can happen despite normalization. When c = 1 and c; < a; P s d 2 and the constraint indicates that d S 2: when the augmentation is very weak, there learning a meaningful representation, which helps us understand why dimensional collapse can harm learning in practice. The fact that simple normalization cannot prevent collapse has been noticed for a while for the simplest case of a cosine-simiin similarity loss, and our result explains why previous works have tried introduce asymmetry to cosine similarity to avoid collapses (Grill et al., 2020; Chen and He, 2021).
most 2 nontrivial subspaces. This is too restrictive for
Relevant Loss Functions. Having developed framework for understanding normalization, we show that other common loss functions in SSL can also be written in the form given in Eq, (3). The spectral contrastive loss (SCL) (HaoChen et al., 2021) reads
LSCL 22 f(x')] E((/()))) (()](()((())))))))))))))) + const.
s.t. 1| =
(14)
Let f(x) Ww be linear, the distributions are zero-mean Gaussian, and ignore the normalization.
This loss function becomes
LSCL = -
(15)
When normalization exists. collapse in this loss function. The difference with InfoNCE loss is that the learned feature spreads along the directions of the augmentation C, not along the directions of the feature Ap-
apply the result in Section 3.2. By argument, there is no

---

Published as a conference paper at ICLR 2023
1,5
5.44
1,5
sug
3.7
0
5.0
12
12
5.6
Weal dets
390
0.0
1.5
372
0.0
1.5
0.0
1.5
5.54
1
(a) Landscape of ResNet
(b) No collapse
(c) Dimensional collapse
(d) Complete collapse
1,5
534
1.4
548
1.5
5.89
Gtann dats aug
5.8
5778
3.7
12
2
Week augs
4.67
53
1,5
5.39
0.0
1.5
5.1
(e) Landscape of VIT
(I) No collapse
(g) Dimensional collapse
(h) Complete collapse
Figure 2: Landscape of Resnet18 (upper) and vision transformers (lower) on CIFARIO with SimCLR qualitatively agrees with our linear theory. Training objective L function of rescaling of the layer W aW. (b-d) as function
2d rescaling of the layer where the data augmentation strength is (b) small, (c) intermediate, and (d) strong. Red indicates areas of high loss, blue indicates areas of low loss, and stars locate local minima. The use of data augmentation changes the stability of the origin, qualitative change that leads different types of collapses in qualitative agreement with our linear theory (cf. Figure 1). Additionally, also notice the qualitative changes of landscape in simpler nonlinear models (see Appendix A). (e-h) are the same setting but for VIT.
The case of Barlow Twin (BT) (Zbontar al., 2021) is similar. While the fourth-order term of BT is much more complicated due to the imbalance created by the À term. The second-order term can This also does not collapse. difference
be identified easily: .BT -TTr -2T:WWW"") + O(I]W]]*. between the SCI. loss and InfoNCE is that the learned representation has spread that aligns with the combination the feature and the augmentation strength.
4 IMPLICATIONS
In this section, we explore some theoretical and practical implications of our results. In Appendix Section A, we also present numerical simulations that directly validate the predictions of the theory.
4.1 RELEVANCE TO NONLINEAR MODELS
An important question is how much of the analysis is relevant for deep nonlinear models in gen- eral. In fact, the loss landscape we have studied is quite close the most general landscape one can have. Let L(f(x)) be general SSL loss function for data point x. The quality of the learned representation should be independent of the population- level orientation of the representa- tion. Therefore, the loss function should satisfy rotational invariance. Namely, for any rotation matrix R, L(x)) = (x)) this rotational invariance implies that the loss should expand as L(f())) af af
((f f(=]²² O(f()))). Note that all the odd-order terms of f(x) vanish due to the rotational symmetry. Substituting f(z) = W Wo(z) ) in the loss function, we obtain
very general form of landscape that W obeys:
L(W,6) Tr[WTTTWW WA]+ Win om Wjm Win Win Zijkit
(16)
where and Z are dependent on é. Note how all the examples we have studied take this form. For W. its collapse entirely depends on the stability of the matrix A. Thus the study of the stability of the matrix becomes crucial for our understanding. To illustrate, we train a Resnet18 on CIFARI0 with the SimCLR loss with normalization and with weight decay strength 10-3 until convergence to obtain the converged weights W*. The representation has a dimension 128. We rescale the weight function of a. See Figure 2-a.
matrix of the last layer W by factor and compute the loss We then partition the singular values of Whitt into the larger half and the smaller half. We rescale the larger half by factor 71 and the smaller half by 72. We plot the loss 2d function of (r1, r2) in Figure 2. We also perform experiments for vision transformers (ViT) in the lower row (Dosovitskiy al., 2020). In all cases, the landscape features qualitative changes comparable to those in Figure 1.

---

Published as a conference paper at ICLR 2023
A connection to Landau theory in physics. Those familiar with statistical physics should note that the proposed theory is analogous to the Landau theory of second-order phase transitions. When treating the loss function as the free energy, the square root of the eigenvalues VÀ of W are the order parameters of the system, and the phase transitions happen when À turns from 0 to positive. These transitions (collapses) happen because of symmetry breaking (Landau and Lifshitz, 2013); the loss function (2) is symmetric in the sign of W. Yet, for any nontrivial learning, W must be nonzero; thus, a symmetry breaking of the sign of W needs to happen for learning. The recent work by Ziyin and Ueda (2022) suggested how symmetry breaking around the origin and Landau theory could explain various types of collapses in deep learning. Therefore, the dimensional collapse could be related to neural collapses in supervised learning (Papyan et al., 2020; Ziyin et al., 2022a) and posterior collapse in Bayesian deep learning (Wang and Ziyin, 2022). Because second-order phase transitions should come with the divergence of the correlation function, one might also wonder what is "divergent" in the SSL problem. Here, the learning time scale for the collapsing dimension is divergent at the critical point because the second-order term vanishes in this direction, and so the dynamics effectively frozen along this direction.
4.2 ROBUSTLY INDUCING GOOD COLLAPSES
Task-ireleeann
1.0
Contrary to previous works, recent work (Cosentino et al., 2022) has suggested that dimensional collapse can be benefi- cial and significantly improve the generalization performance of the model. This observation raises a question. How can di- mensional collapse be beneficial and how can it be induced? In
0: Direction of deta aug.
Good dim-collapse
0.65
No collapse
Complete c
0.5
the following, we first introduce 3-InfoNCE, which can adjust the degree of dimensional collapse, and analyze the collapse behavior to elucidate the mechanism of task-alligned collapse.
Bad dim-collapse
0.0
Task-relevant 1.0
2.0
3.0
a²: Strength of data aug.
Adjusting the degree of dimensional collapse with 3- InfoNCE. Despite the potential benefit, existing SSL loss functions cannot robustly induce dimensional collapse. In- foNCE is insufficient to induce a collapse, and the collapse induced by SimCLR depends on a vanishingly small param- eter 1/N. One thus wonders whether there is a loss function that allows us to induce collapsing behavior in more pre- dictable matter so that one might controllably extract some benefits from collapse. Our result suggests that one way to directly control collapses is through the strength of the compe- tition for the model Hessian at the origin. For InfoNCE, one way to achieve this is to weigh the entropy term by a general factor 3:
T: Theory 0.5
2
Test loss
0.65
10-3-
1.0
10
1.0
2.0
3.0
a²: Strengtt of data aug.
Figure 3: Top: Phase diagram of representational collapses. Bottom: I-InOCCCE with / 0.5. The gen- eralization error of a downstream re- gression task where the data augmen- tation (1) is isotropic and noninforma- tive or (2) aligns with the style. We that the performance worsens as col- lapses happen for the noninformative augmentation and improves as the col- lapse happens for the style-targeting augmentation.
E. G (1(())-((()))
Blog E, exp -1100)--(0)))
Due to its similarity with the 3- in Bayesian learning, we call it the B-InfoNCE. The leading term in the loss function becomes
(1-
B> > O, the augmentations C pull the representation towards zero When the augmentation is as strong as the fea- ture variations, collapse happens. One can thus introduce collapse by setting 3 to be sufficiently small. When 1 - 3 < O, the augmentations push the weights away from the origin along its direction, resulting in no collapse at all: When one really wants to avoid collapse, one can use a rather large 3: 8 = 1 is thus at the boundary of this bifurcating behavior. We note that existing loss functions often do not have a parameter that is directly controlling the collapse behavior (see Table 1). The / parameter here directly controls the level of difficulty of collapse.
When 1-
Achieving invariance with dimensional collapse. Here, we closely study an illustrative minimal example to demonstrate how collapses can be beneficial. Consider the following structured data generating process where the input features can be separated into two sets: (1) a task-relevant set with dimension de do and (2) a iask-irreeeaant set: (21, Fdd Tdd.) Our result suggests
8

---

Published as a conference paper at ICLR 2023
precise way to remove the irrelevant features from the learned representation. For the purpose of causing a robust collapse, we use the 3-InfoNCE with / = 1/2. For illustration, we consider the simple case d, 1 and dg 2. For any imput r (2), the label is generated as a linear function of 11 y cD1-
Correspondingly, we consider structured data augmentation E a RS, where R e RRdoxdo is R diag( V,),,, where 0 E [0, 1.. The parameter controls the overall strength of the aug- mentation, and controls the orientation of the strength. When = 0.5, we have an uninformative isotropic noise that has often been used in practice. When 0 1, the augmentation only on the lask-irrelevant feature, and when 0 = 0, the augmentation is only on the content. Since the predic- tion target only depends on the content, we want to learn a representation invariant to the style. For the downstream regression task, we use the learned representations == f()) to train a ridge linear regressor that minimizes ming [E - y [ 0.001||G]]? See Figure 3. The top panel shows the phase diagram of this problem with different combinations of the augmentation strengths and orientations. The bottom panel shows that collapses introduce phase-transition-iiee behaviors in the generalization performance and that a data augmentation aligning with the task-irrelevan dimension improves performance.
Without Projector
With Projector
Spectral
Spectral
4.3 ROBUSTNESS TO DATA IMBALANCE
Contrastive
Contrastive
Our theory not only relevant for understand- ing collapses but can also be used to under- stand how an SSL model encodes the feature. Liu et al. (2021) recently showed that compared with supervised learning, SSL techniques are relatively more robust to imbalanced datasets that have disproportionately represented minor- ity subgroups. As another application of our analysis, we illustrate the robustness of differ- ent techniques is not equal. As we have seen, the learned model W W has eigenvalues that, the leading order, are proportional to the Hes- sian B, which is different for each loss function. As previously summarizze in Table 1, for In- foNCE and SimCLR, the learned model aligns with the eigenvalues of the data covariance Ag. which varies hugely as different classes of a dataset become more and more imbalanced. In comparison, the model trained with SCL aligns purely with the augmentation covariance C, which is independent of the data imbalance. This suggests that the SCL landscape be less dependent on data and thus more robust against data imbalance. See Figure 4. More experimental details are given in Appendix C.
Test accuracy
55
65
50
InfoNCE
InfoNCE
45
60
None
Moderate Large
None
Moderate
Large
Data Imbalance
Data Imbalance
Figure 4: robust against data imbalance than InfoNCE. We train SimCLR and SCL ResNet-12 models imbal- anced versions of CIFAR-10. We see that SCL more robust than SimCLR, as suggested by our theory. These results are especially pronounced when there is no pro- jector head.
Spectral Contrastive loss (SCL) is more
CONCLUSION
In this work, approached the problem of collapses in SSL from loss landscape perspective. We analytically solved an effective landscape that can be extended to understand the effect of normaliza- tion. Our result suggests that dimensional collapse can be well understood in the minimal setting and is something neutral to learning on its own. With the help from the theory, we also showed that when task-irrelevant dimensions are targeted, dimensional collapse can result in improved performance, whereas an uninformative noise will (without good luck) leads to collapses in the dimensions that thus important for practitioners to devise targeted data augmentation
are relevant to the task. mechanisms that incorporate the correct domain knowledge. Also, we advocated the thesis that the local geometry of the loss landscape around the origin an essential component for understanding collapses, and this should invite more future work to understand the landscape around the origin.
The limitation of our work is clear; our result only identifies the causes of the collapse that can be directly attributed to the low-rank structure of the local minima of the landscape. One possible alternative cause of the collapse is dynamics. For example, having a large learning rate and small batch can sometimes cause convergence towards the saddle points in the landscape (Ziyin et al., 2022b), which, as we have shown, are the collapsed solutions. Investigating the role of dynamics in the collapse is thus crucial future problem.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
This work was supported by a KAKENHI Grant No. JP18H0114, 145 from the Japan Society for the Promotion of Science. Ziyin has been financially supported by the JSPS fellowship and thanks Zihan for the generous help during the writing of this paper. ESL was partially supported via NSF under the award 1NS-2008151.
REFERENCES
Arora, S., Khandeparkar, H., Khodak, M., Plevrakis, O., and Saunshi, N. (2019). Theoretical Analysis of Contrastive Unsupervised Representation Learning. In Proc. Int. Conf. on Machine Learning (ICML).
Bardes, A., Ponce, J., and LeCun, Y. (2021). Vicreg: Variaace-invaraane-ccvaaannn regularization for sulf-supervised learning. arXiv preprint 21ti::2005.040066
Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. (2020). Unsupervised Learning of Visual Features by Contrasting Cluster Assignments. In Proc. Adv. on Neural Infor- mation Processing Systems (NeurIPS).
Caron, M., Touvron, H., Misra, I., Jegou, H., Mairal, J., Bojanowski, P., and Joulin, A. (2021). Emerging Properties in Supervise Vision Transformer. arXiv, abs/2104. 142994
Chen, T., Korbblith, S., Norouzi, M., and Hinton, G. (2020a). A simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1597-
1607. PMLR.
Chen, T., Kornblith, S., Swersky, K., Norouzi, M., and Hinton, G. E. (2020b). Big Self-Supervised Models Strong Semi-Supervised Learners. Adv. in Neural Information Processing Systems, 33.
Chen, X. and He, K. (2021). Exploring Simple Siamese Representation Learning. In Proc. Int. Conf. on Computer Vision and Pattern Recognition (CVPR).
Cosentino, R., Sengupta, A., Avestimehr, S., Soltanolkotabi, M., Ortega, A., Willke, T., and Tepper, M. (2022). Toward a geometrical understanding of self-supervised contrastive learning. arXiv preprint arXiv:2205. 06926.
Dayan, P. and Abbott, L. F. (2005). Theoretical neuroscience: computational and mathematical modeling of neural systems. MIT press.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2020). An image is worth 16x16 words: Trans- formers for image recognition scale. arXiv preprint arXiv:2010.11929,
Ermolov, A., Siarohin, A., Sangineto, E., and Sebe, N. (2021). Whitening for self-supervised repre- sentation learning. In International Conference on Machine Learning, pages 3015-3024. PMLR.
Grill, J.-B., Strub, F. Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, Z., Gheshlaghi Azar, M., Piot, B., kavukcuoglu, Munos, R., and Valko, M. (2020). Bootstrap your own latent: A new approach to self-superviier Learning. In Proc. Adv. on Neural Information Processing Systems (NeurIPS).
HaoChen, Z., Wei, C., Gaidon, A., and Ma, (2021). Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in Neural Information Processing Systems, 34:5000-5011.
He, K., Fan, H.. Wu, Y., Xic, S., and Girschick, R. (2020). Momentum Contrast for Unsupervised Visual Representation Learning. In Proc. Int. Conf. Computer Vision and Pattern Recognition (CVPR).
Hsu, H., Qi, H., and Brown, M. (2019). Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification. arXiv, abs/1909.00333
10

---

Published as a conference paper at ICLR 2023
Hua, T., Wang, W., Xue, Z., Ren, S., Wang, Y., and Zhao, H. (2021). On feature decorrelation in self-supervised learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9598-9608.
Jing, L.,, Vincent, P., LeCun, Y., and Tian, Y. (2021). Understanding dimensional collapse in con- trastive self-supervised learning. arXiv preprint arXiv:21
Kugelgen, J., Sharma, Y., Gresle, L., Brendel, W., Scholkopf, B., Besserve, M., and Locatello, F. (2021). Self-Supervised Learning with Data Augmentations Provably Isolates Content from Style. arXiv, abs/2106.04619,
Landau, L. D. and Lifshitz, M. (2013). Statistical Physies: Volume 5, volume Elsevier.
Liu, H., HaoChen, J. 7 Gaidon, A., and Ma, T. (2021). Self-supervised learning is more robust to dataset imbalance. International Conference on Learning Representations.
Mitrovic, J., McWilliams, B., Walker, J., Buesing, I.., and Blundell, C. (2020). Representation learning via invariant causal mechanisms. arXiv preprint arXiv:2010.07922.
Nozawa, K. and Sato, I. (2021). Understanding negative samples in instance discriminative sulf-supe representation learning. Advances in Neural Information Processing Systems, 34:5784-5797.
Oord, A. v. d., Li, Y., and Vinyals, O. (2018). Representation learning with contrastive predictive coding. arXiv preprint arXiv: /807.03748
Papyan, V., Han, X., and Donoho, D. L.. (2020). Prevalence of neural collapse during the ter- minal phase of deep learning training. Proceedings of the National Academy of Sciences, 117440))2452-24663.
Pokle, A., Tian, J., Li, Y., and Risteski, A. (2022). Contrasting the landscape of contrastive and non-contrastive learning. arXiv preprint arXiv:2203.15702
Robinson, J., Sun, L., Yu, K., Batmanghelich, K., Jegelka, S., and Sra, S. (2021). Can contrastive learning avoid shortcut solutions? Advances in neural information processing systems, 34:4974- 4986.
Saunshi, N., Ash, J., Goel, S., Misra, D., Zhang, C., Arora, S., Kakade, S., and Krishnamurthy, A. (2022). Understanding Contrastive Learning Requires Incorporating Inductive Biases. In Proc. Int. Conf. Machine Learning (ICML).
Simon, B., Knutins, M., Ziyin, L., Geisz, D., Fetterman, A. J., and Albrecht, J. (2023). On the stepwise nature of self-supervised learning. arXiv preprint vrXiv:2303.15433
Tian, Y. (2022). Deep contrastive learning is provably (almost) principal component analysis. arXiv preprint arXiv:2201.12683
Tian, Y., Chen, X., and Ganguli, S. (2021). Understanding self-supervised Learning Dynamics without Contrastive Pairs. Proc. Int. Conf. on Machine Learning (ICML).
Tian, Y., Sun, C., Poole, B., Krishnan, D., Schmid, C., and Isola, P. (2020). What makes for good views for contrastive learning? Advances Neural Information Processing Systems, 33:6827- 6839.
Tosh, C., Krishnamurthy, A., and Hsu, (2021). Contrastive estimation reveals topic posterior information to linear models. J. Mach. Learn. Res., 22:281-1.
Trivedi, P., Lubana, E. S., Heimann, M., Koutra, D., and Thiagarajan. J. (2022). Analyzing data-
centric properties for contrastive learning graphs. arXiv preprint arXiv:2208.02814
Tsai, Y.H., Wu, Y., Salakhutdinov, R., and Morency, L.-P. (2021a). Self-supervised Learning from Multi-view Perspective. In Proc. Inf. Conf. Learning Representations (ICIR).

---

Published as a conference paper at ICLR 2023
Tsai, Y.HH. H., Ma, M. Q.. Yang, M., Zhao, H., Morency, L.-P., and Salakhutdinov, R. (2021b). Self-supervised representation learning with relative predictive coding. International Conference Learning Representations.
Wang, and Liu, (2021). Understanding the behaviour of contrastive loss. In Proceedings the IEEE/CVE conference on computer vision and pattern recognition, pages 2495-2504.
Wang, T. and Isola, P. (2020). Understanding Contrastive Representation Learning through Align- ment and Uniformity on the Hypersphere. In Proc. Int. Conf. on Machine Learning (ICML).
Wang, Y. Zhang, Q. Wang. Y. Yang. J., and Lin, Z. (2022). Chaos is ladder: new theoretical understanding of contrastive learning via augmentation overlap. International Conference Learning Representations.
Wang, Z. and Ziyin, L. (2022). Posterior collapse of a linear latent variable model. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K., editors, Advances in Neural Information Processing Systems.
Wei, C., Shen, K., Chen, Y., and Ma, T. (2021). Theoretical Analysis of Self-Training with Deep Networks on Unlabeled Data. Proc. Int. Conf. on Learning Representations (ICLR).
Zbontar, J., Jing, L., Misra, I., LeCun, Y., and Deny, S. (2021). Barlow twins: Self-supervised learning via redundancy reduction. In International Conference Machine Learning, pages 12310-12320. PMLR.
Zimmerman, R., Sharma, Y., Schneider, S., Bethge, M., and Brendel, W. (2021). Contrastive Learn- ing Inverts the Data Generating Process, In Proc. Int. Conf. on Machine Learning (ICML).
Ziyin, L., Li, B., and Meng, X. (2022a). Exact solutions deep linear network. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K., editors, Advances in Neural Information Processing Systems.
Ziyin, L.,, Li, B., Simon, J. B., and Ueda, M. (20226). SGD can converge to local maxima. In International Conference Learning Representations.
Ziyin, L. and Ueda, M. (2022). Exact phase transitions in deep learning. arXiv:200..12510.
arXiv preprint
12

---

Published as a conference paper at ICLR 2023
A ADDITIONAL NUMERICAL RESULTS
In this section, we validate our theory with numerical results. Unless specified otherwise, the di- mension of the learned representation is to be equal to the input dimension: do d1.
We showed that there is no collapse at all for the vanilla InfoNCE, no matter how strong the augmentation is. Our result implies that the smallest singular of the model W scales as o² where o² is the strength (namely, the variance) of the augmentation. See the left panel of Fig. 5. We set is sampled from (0, 132). The training proceeds with Adam with learning rate of 6e - 4 with full batch training for 5000 iterations. We use simple diagonal Gaussian noise with variance o² for data augmentation. We see that the singular values scale as o and never vanishes, as the theory
No Collapse for InfoNCE.
the vanilla InfoNCE loss defined in (1) with a linear model. The training
predicts.
Nonrobust Collapses of Weighted InfoNCE. We now demonstrate that, as the theory predicts, collapses of weighted InfoNCE depend strongly on the dataset size. We use the same dataset and training procedure the previous experiment. We set 0.1 and change the size of the training set. Theory suggests that for collapse in the i-th subspace to happen, the size the dataset needs to obey
a, c: (1-aa)
== Nerst.
(17)
See the middle panel of Figure 5. having similar magnitudes), and the critical dataset size for the smallest eigenvalue. We see that the theoretical threshold of collapse agrees well with where the collapse actually happens.
We show the smallest three eigenvalues of WTW W (roughly
Collapses in N-InfoNCE. With 1, one can cause collapses in a predictable and controllable W as we increase the
way. In this experiment, we let do 5 and we plot all five eigenvalues of W strength of isotropic augmentation. As the numerical results show, collapses happen at the points predicted by the theory.
Normalization Causes Dimensional Collapse. We also plot the three smallest eigenvalues of W when apply the standard representation normalization in practice: f(x) f()) To facilitate comparison, we also use the same dataset and training procedure as before. See Fig- ure 6. We that normalization does cause collapse in the smallest eigenvalues augmentation strength much smaller than the feature variation.
B LANDSCAPE OF A NONLINEAR MODEL
In this section, we plot the landscape of the layer of nonlinear models on the same synthetic dataset we outlined in the previous section. We train a three-layer nonlinear network with output dimension 2 with SGD until convergence. We then rescale the optimized weight of the last by a factor a: Wiast a aliast and plot the loss function along this direction. See the top panel of Figure 7 for
0.3
0.004
Nott
Theory
10
0.003
10-2
15
0.002
10-3
10"4
10a 0.000
81 1
0.0
100
60
80
100
120
0.75
1.00
1.25
1.50
g²
Figure 5: The three smallest singular values of W W as function the augmentation strength. We see that our effective landscape theory around the origin accurately captures collapses in learning. Left: Vanilla InfoNCE. As the theory suggests, the singular values scale as a and not vanish for any finite value of Mid: Weight InfoNCE. 0.1, a 5. Collapse happens at the critical dataset size predicted by the thcory. Right: (Sqrt) Eigenvalues of W WWT in I-IffoNCE. The collapses can be well controlled.
13

---

Published as a conference paper at ICLR 2023
0.004
0.003
0.002
0.001
0.025 0.050
0.100
d²
Figure 6: A collapse happens easily when the learned representation is normalized. The smallest eigenvalues of Au are roughly 0.2, and the collapse happens much before the noise reaches this strength.
6.8
50
6.2
5.-
0.50
702
8.00
6.34
7.76
0.00
0.00
6.78
0.00
7.44
-0
0.25
0.25
-050
-0.50
642
-0.50
-5
6.4
-0.7
-0.75
6.48
-100
6.00
-1.00
6.18 -150
6.16
Figure 7: The Landscape of nonlinear models is very similar to the landscape linear models (cf. Figure 1). Top: 1d projection of the landscape of a two-layer tanh and ReLU network. Bottom Left: the landscape of a 2D projection of the last layer nonlinear model with a weak augmentation. Middle: with intermediate augmentation. Right: with strong augmentation.
both the tanh and the ReLU nonlinearity. We then rescale the two of the weight matrix of the model by Fı and r2 respectively: W = (w) W2) + (r1w,,2w2)) We see that the landscape of the model is qualitatively the same that of the linear models, shown in Figure 1.
C SETUP FOR IMBALANCED DATA EXPERIMENTS
Creating Imbalanced Dutaset: For our experiments measuring the influence imbalanced datasets on SSL training, we use CIFAR-10 10 by sampling 20000 samples out of the 50000 train- ing samples. The sampling process described by Dirichlet distribution and is often used to analyze effects of heterogeneity and data imbalance in Federated Learning problems (Hsu et al., 2019). Specifically, small value of the distribution parameter yields a highly imbalanced dataset, while large value yields perfectly balanced dataset. We evaluate our models in three scenarios, for which we report below the number of samples per class:
High imbalance: [4890, 87, 5000, 0, 74, 0, 0, 212, 4788, 4947]
Medium imbalance: [4268, 4296, 1741, 420, 945, 4633, 1015, 131, 2386]
No imbalance: [2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000]
Training Setup: We ResNet- 12 models as the backbone for all experiments due to computational constraints. SimCLR augmentations (Chen et al., 2020a) are followed, except for reduced strength of resized cropping from 0.2 to 0.5. All training involves standardly used cosine decay learning
14

---

Published as a conference paper at ICLR 2023
rate schedule, starting at 0.03 and decaying to 0.001. When projector module is used, it involves a two-layer MLP with hidden dimension of 512 and BatchNorm layer in between. We use SGD for optimization and perform the standardly used linear evaluation protocol for measuring the quality of the final representation. For training the linear layer, we use an initial learning rate of 10 and decay it 0.01 with cosine schedule. We note linear evaluation is used for supervised models well, following the practice advocated by Liu et al. (2021).
D ADDITIONAL THEORETICAL CONCERNS
D.1 COLLAPSE CONDITION FOR NORMALIZATION
The important condition for collapse in Eq. (13) can be better understood by considering the extreme cases. First of all, note that the eigenvalues of EB are bounded between -1 and
0g 11
1
(18)
0;
and -1 is achieved when o; >> ag and is achieved when 0 >> 4--
When the augmentation is negligibly small, B M, and X Ā 1, the condition thus becomes
dM 0.
(19)
which always holds. Thus, sufficiently small augmentation will never cause collapse. Next, when apply very strong augmentation to the j-th subspace and zero augmentation to the others, the condition for the -on-augmented spaces becomes
dM
(20)
das
das
meaning that the collapse will not happen. For the j-th space, the condition is
2 dM -2 dm
4
+
2,
(21)
dM
dM
which is only possible when dM = 1, namely, the strongly augmented space is the only space that does not collapse. This is reasonable when the original data rank-1 because the normalization will ensure that this space does not collapse, but when the original data is not rank-1, this stationary point will be a saddle and will not be preferred by gradient descent. different word, strong enough augmentation will cause collapse in the corresponding subspace, as is the without normalization.
It is also interesting example, let c 0 and Cj ag for j # 1. The condition for j # becomes
that having 6 a; is longer sufficient to cause collapse. For
(22)
dM
dM
which always holds. At the same time, it does not mean that collapsing has become harder in general. For example, it is also possible for c; < a; cause a collapse. Suppose we add a weak augmentation only to the first subspace such that 0; - 44 = e O, the condition for this dimension to not to collapse is
2 dm 1+e
+
(23)
04 +
dM
dm
which be violated whenever fact facilitate collapse.
(dur3) a
Namely, in some cases, normalization in
15

---

Published as a conference paper at ICLR 2023
E PROOFS
E.1
PROOF OF PROPOSITION
Proof. The second term in Eq. (3) can be written as
(24)
= [first term] (A)
(25)
= [first term] - <T>[WW>"]]
(26)
where we have used the definition Ao C. The first term
[first term] E [[T:[ x) - (
"
(27)
Combining the above expressions, we see that Eq. (3) can be written
L [T[[WBBT]]]
VarllW x]
(28)
8
Tr[WBWT]] T-WEW""EW""
(29)
This finishes the proof. D
E.2 PROOF OF THEOREM
Proof. All stationary points have gradient:
BWB WEW
(30)
Multiplying by W on the left and B on the right,
wwW- = WSW
(31)
(32)
Defining #= WTTTWT obtain
2
(33)
(=) H(1 - 1
(34)
Because both И and 1 are symmetric, one can take the transpose of Eq. (33) to find that H and 1/2 11/2B-'y'0 1/2 commute with each, which implies that has the same eigenvectors as 1/ 1/ /2.
Eq. (34) then implies that the eigenvalues of is either the inverse of that of 11/2 or zero. This implies that any stationary point of H can be written in the form
1 H U=-UMAU!
(35)
where U is unitary matrix, A is diagonal matrix containing the eigenvalues of 2//22-1E//2 and M is an arbitrary (masking) diagonal matrix containing only zero or one such that (1) Ma = 0
Aa < 0 and (2) contain most d² nonzero terms. This then implies that the weight matrix W satisfies
wTw 2
"PMMUT
(36)
Lastly, when E and B commute, we can compactly write the result as
y'W=-y--BBN3
(37)
where Ba denotes the matrix obtained by masking the eigenvalues of with M. This finishes the proof. 0
16

---

Published as a conference paper at ICLR 2023
E.3 PROOF OF PROPOSITION
Proof. For all stationary points, WTW commutes with B and Σ, which means that at these sta- tionary points, one can simultaneously diagonalize all the matrices and the loss function (3) can be written as
-=-Σ X;b; + X i=1
(38)
where X ba, 8 are the eigenvalues of WTW, B, and Σ respectively.
We can thus consider each separately. When b, > 0, A = 0 cannot be a local minimum because the local Hessian is -b; < O. When b; s 0, the only stationary point is X; 0. This sum covers at most d' summands, and so, at the local minima, X + and only if b; > 0, and so the number of non-zero eigenvalues of W W min(m,
E.4 PROOF PROPOSITION
Proof. The regularization can be written as
= [[Ba[[W=]]==e)
(39)
TT[WEW"]3 -
(40)
By Proposition 1, Eq. (10) reads
L T=-TTWBW"]]- "WW"SW <WW-:<<""""""""-----""""""""""" +1
(41)
=-Tr[W(EB [
kp?
(42)
The derivative of is
d dw P 4pWE.
(43)
The graddera gradient thus
2W 2xcE - 2xpE) 4WEW
(44)
We can define 1² B+ 2xe 2xp2 see that this condition the same as Eq. (30) the proof
of Theorem 1. The rest of the proof thus follows from the arguments. We thus arrive at the theorem statement:
WWW
E B 2
(45)
We are done.
E.5 PROOF OF PROPOSITION
Proof. Recalling that p= Tr Tr[WEW"] multiply E from the right to both sides of the solution in Proposition 3 and take trace:
TH
Byy = THE (Bur + p)
(46)
2
= [T(W"WE]]
(47)
TTWWE"
(48)
The first line further simplifies to
1 "T['-' +
((CC-P))-E=>'>,)) - - BM] ] x(c- pd
(49)
2
where dm == Tr[M] is the number of nonzero eigenvalues of B'y
This gives an equation that solves to
C- TTr(2-1 BM]
p 1 + kd
(50)
This proves the proposition.
17

---

Published as a conference paper at ICLR 2023
F ADDITIONAL THEORETICAL CONCERNS
F.1 CASE OF DATA-INDEPENDENT NON-GAUSSIAN AUGMENTATION
In the main text, mainly considered the case when the noise is Gaussian. In this section, we consider a case where the noise is data-dependent and non-Gaussian We show that the results we discussed in the main text still hold qualitatively. The general form of the loss function in Eq. (3) still applies:
L Tr ([WBBW)+ Var[IW(<=-X x)
(51)
8 We consider a global rescaling augmentation for each datum æ:
(52)
where 8 ~ czp(b) obeys an exponential distribution with mean b and variance D² Note that even Gaussian, the augmented data is no longer Gaussian. In particular, the augmentation now becomes data-dependent. This augmentation can also be seen structured, biologically plausible data augmentation that encourages the model to be scale-invariant, which what Wien's 's law for biological perception demands (Dayan and Abbott, 2005): matter whether image is dark or bright, the content of the image is the same.
Under this augmentation, the noise covariance is dependent and no longer Gaussian:
E a 2bb Ao-
(53)
We also obtain that
C=E[ s) xx =bb2-n -b²Ap-
(54)
The second term in Eq. (3) can be written as
Var[IW x E
x(x x) - ""'1"
(55)
= [first term]
(A C
(56)
= [first term] - 4TH[WEW"]],
(57)
where we have used the definition Σ An C. The first term is
[first term] = [E[[I:[W (x(
(58)
However, for fixed rescaling factor 82 and 8x. each W(x - x) obeys a multivariate Gaussian distri- bution with variance 2(x² " Ag, and so we have
[first term] E EE,,,, [s +x (4TT AoW
W.
(59)
where E,, .x [63 ) 56b Combining terms, we obtain that
Var[WW(x x) 4802 x 4Tr(W [W + 566* x 8Tr[W W WAW'")
(60)
The loss function thus: L= TT-TT[BBWT1)
565 566 Ao J.
(61)
Note that this loss function is a special case of the loss function in Eq. (10) where 0 and k 24h (and with a rescaled fourth-order term). As in the main text, B is different according to different choices of loss functions. Because B commute with Ag by construction, one expects collapses to happen at locations predicted by Proposition 3 and under suitable choices of parameters. Also note that the odd terms vanish as discussed, and so the local stability of the origin should decide the collapsing behavior of this situation.
This shows that collapse also happen when the data augmentation is structured. We comment that the analysis in this section is minimal, and one important future direction is to provide more precise and insightful conditions of collapse under structured data augmentation.
18


Reviews:
[{"review_id": "wEEBNCbTLYT", "reviewer": "Reviewer_ffca", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Self-Supervised Learning (SSL) algorithms like BarlowTwins and BYOL (both contrastive, non-contrastive) are known to suffer from $\\textit{dimension collapse}$, where the effective rank of the representations is much smaller than the dimensionality of the representation. Empirical analysis of such algorithms presents conflicting evidence on the role of $\\textit{dimensionality collapse}$ in downstream generalization. The authors give a different perspective, where they study the loss-landscape of SSL with simplified linear models. In particular, this analysis highlights the role of data augmentation & data-variation in inducing dimension collapse in the representations that could be beneficial for generalization by tuning out task-irrelevant dimensions. Building on these insights, the authors perform experiments w/ Resnet-18 on CIFAR datasets to demonstrate qualitative agreement with the theory.", "summary_of_the_review": "In summary, the authors study the local geometry of the loss landscape with self-supervised learning objectives to understand the dimension-collapse in feature space. Using simplified analytic and linear models, the authors demonstrate that dimension collapse emerges by controlling the strength of data augmentation and data variation at train time and characterize conditions that lead to generalization-friendly dimension collapse. By extending experiments to non-linear ResNet-18 models on CIFAR datasets, the authors show qualitative similarities in loss landscape under different strengths of data augmentation. ", "strength_and_weaknesses": "Some strengths of the paper are outlined below:\n+ The article is well-written and presents a coherent argument for disentangling the role of data augmentation in self-supervised learning via the lens of loss-landscape analysis. \n+ The authors provide empirical evidence that suggests some qualitative alignment in the behavior of the SSL learning objectives on realistic models/datasets.\n\nSome weaknesses of the current paper include the following:\n- The analysis focuses almost entirely on linear models, with little discussion on the validity of the nonlinear regime. \n- Though not essential, it would be helpful to have similar qualitative/quantitive experiments for more architectures (in particular transformers) to ablate the role of architecture design from a choice of the loss function.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "nbwcKnmP_d", "reviewer": "Reviewer_Q5CS", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper investigates the landscape causes of collapse in self-supervised learning.\n", "summary_of_the_review": "In general, Although the paper has some strong assumption, it still gives some insight about how the collapse happens in SSL. ", "strength_and_weaknesses": "Strength\n\n1. Existng work have conflict opinion about the collapse in SSL. This work try to fill this gap by studying the geometry of the SSL.\n\n2. This work shows that the interplay between data variation and data augmentation determines the geometry of the loss.\n\n3. They found the geometry of the loss explains when dimensional collapse can be helpful and why certain SSL\nlosses are robust against data imbalance, but not the other.\n\nSuch findings may help design a better loss in SSL.\n\n\nWeakness:\n1. The work may rely on the assumption that the augementation is x = x +e, and all the data and noise must be Gaussian also, the nerual is linear. This assumption may be strong, since they are many non-noise based augementation. It is not clear if other data augumentation still has similar results as in this work.\n\n2. If the data is imbalance as shown in 4.3 How the entire analysis appiled to this case, like if data is not Gaussian?\n\n3. The code is not available, the reproducibility is unclear.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "OrzxDmPFFyr", "reviewer": "Reviewer_m5m1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this work, the authors analyze the problem of collapses in SSL from a landscape perspective. They solved a landscape that can be extended to understand the effect of normalization. Their result suggests that dimensional collapse can be explained in the minimal setting and is something neutral to learning on its own. They showed that when task-irrelevant dimensions are targeted, dimensional collapse can result in dramatically improved performance, whereas an uninformative noise will lead to collapses in the dimensions that are relevant to the task. The authors believe that It is thus important for practitioners to devise targeted data augmentation mechanisms that incorporate the correct domain knowledge. The proposed theory can serve as a theoretical foundation and baseline of any advanced theory of collapses because a correct theory should agree with our results when restricting to the case of a linear model. The authors advocated the thesis that the local geometry of the loss landscape around the origin is an essential component for understanding collapses, and this should invite more future work to understand the landscape around the origin.\n\n", "summary_of_the_review": "Please refer to the above sections to find the points that can be revised to improve the paper.", "strength_and_weaknesses": "Pros: The authors conduct a solid and comprehensive theoretical analysis of the SSL landscape. For example, in Section 3.1, the authors show that the landscape for a class of situations in self-supervised contrastive learning can be reduced to an effective form in Eq. (3). In Proposition 1, the authors show that the variance term of the loss takes a specific form when the data is Gaussian. Moreover, in Section 4 and Appendix A, the authors illustrate some theoretical and practical implications of their analytic results.\n\nCons: Although focused on analyzing the reasons for the collapse in SSL, this work cannot account for all the SSL collapses but only identifies the results that can be directly attributed to the low-rank structure of the local minima of the landscape. The paper's numeric analysis is limited. The language and organization of the draft have space to be improved. The paper is not easy to understand due to the lack of intuitive explanation.\n\nThe authors are suggested to improve the draft by improving the points mentioned in the Cons to improve the paper.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster