#  What shapes the loss landscape of self-supervised learning?

Liu Ziyin1,2,3 Ekdeep Singh Lubana2,3,4 Masahito Ueda1,5,6 Hidenori Tanaka2,3
1Department of Physics, The University of Tokyo, Tokyo, Japan
2Physics & Informatics Laboratories, NTT Research, Inc., Sunnyvale, CA, USA
3Center for Brain Science, Harvard University, Cambridge, USA
4EECS Department, University of Michigan, Ann Arbor, USA
5Institute for Physics of Intelligence, The University of Tokyo, 7-3-1 Hongo,
Bunkyo-ku, Tokyo 113-0033
6RIKEN Center for Emergent Matter Science (CEMS), Wako, Saitama 351-0198,
Japan

###### 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 provide answers to
these questions by thoroughly analyzing SSL loss landscapes for a linear
model. We derive an analytically tractable theory of SSL landscape and show
that it accurately captures an array of collapse phenomena and identifies
their causes. 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.

##  1 Introduction

Self-supervised learning (SSL) methods have achieved remarkable success in
learning good representations 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; HaoChen et al.,, 2021; Zbontar et al.,, 2021; Caron
et al.,, 2020) or asymmetric dynamics (Caron et al.,, 2021; Grill et al.,,
2020; Chen and He,, 2021) to ensure that the distance between unrelated
samples remains large. In practice, however, SSL training often experiences
the failure mode of _dimensional collapse_ (Jing et al.,, 2021; Tian et al.,,
2021; Pokle et al.,, 2022), where the learned representation spans a low
dimensional subspace of the overall available space. In the extreme case, this
failure mode instantiates as a _complete collapse_ , where the learned
representation becomes zero-rank, and no 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 et al.,, 2021; Bardes et al.,, 2021). However,
recent work by Cosentino et al., (2022) empirically demonstrates otherwise:
learning informative representations can require collapses to be emergent.
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 to the lack of a 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 popular losses used in self-supervised learning, including
InfoNCE (Oord et al.,, 2018), Normalized Temperature Cross-Entropy (NT-xent)
(Chen et al., 2020a, ), Spectral Contrastive Loss (HaoChen et 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. 1.

the interplay between data variation and data augmentation determines the
geometry of the loss;

  2. 2.

the geometry of the loss explains when dimensional collapse can be helpful and
why certain SSL losses are robust against data imbalance, but not the others.

To the best of our knowledge, our work is the first to thoroughly study the
landscape causes of collapse in SSL.

(a) An eigenmode

(b) No collapse

(c) Dimensional collapse

(d) Complete collapse

Figure 1: Landscape in self-supervised learning (SSL). SSL losses generally
depend only on the relative angle between pairs of network outputs (e.g,
$f(x)^{T}f(x^{\prime})$). Thus, the landscapes with a linear network
($f(x)=Wx$) have a global rotational symmetry and are symmetric about the
origin. Our theory finds that the local stability at 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
linear model with a diagonal weight matrix $diag(r_{1},r_{2})$. (a) The $1d$
landscape when fixing one of the parameters. (b-d) The $2d$ landscape. (b) No
collapse: the origin is an unstable local maximum, and surrounding local
minima avoid collapse. The dimensionally collapsed solutions are the saddle
points. (c) Dimensional collapse: the value of $w_{1}$ for all stable fixed
points collapses to zero. (d) Complete collapse: the origin becomes the
isolated local minimum.

##  2 Related Works

SSL and Collapses. Prior literature has often argued collapse as a harmful
phenomenon that can deteriorate downstream task performance (Jing et al.,,
2021; Zbontar et 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 et al.,, 2021; Bardes et al.,, 2021; Ermolov
et al.,, 2021). However, in direct contrast, Cosentino et al., (2022)
empirically showed that dimensional collapses under strong augmentations could
dramatically 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
understanding the success of SSL techniques from different perspectives: e.g.,
learning theory (Arora et al.,, 2019; HaoChen et al.,, 2021; Saunshi et al.,,
2022; Nozawa and Sato,, 2021; Ash et al.,, 2021; Wei et al.,, 2021),
information theory (Tsai et al., 2021a, ; Tsai et al., 2021b, ; Tosh et al.,
2021a, ; Tosh et al., 2021b, ), causality and data-generating processes
(Zimmerman et al.,, 2021; Kugelgen et al.,, 2021; Trivedi et al.,, 2022; Tian
et al.,, 2020; Mitrovic et al.,, 2020; Wang et al.,, 2022), dynamics (Wang and
Isola,, 2020; Tian et al.,, 2021; Tian,, 2022; Wang and Liu,, 2021), and loss
landscapes (Jing et al.,, 2021; Pokle et 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 a competition between the feature signal strength and
augmentation strength can lead to dimensional collapse. In contrast, our focus
is on the landscape and our result implies that this feature-augmentation
competition on its own is insufficient to cause a dimensional collapse. In
fact, we show that there will be no collapse in the setting studied by Jing et
al., (2021).

Interpretability of linear models. Previous works have demonstrated how linear
models are often sufficient to reproduce phenomena observed in non-linear deep
networks (Saxe et al.,, 2013; Kawaguchi,, 2016). If an observed phenomenon
does not occur for linear models, one can conclude that the use of
nonlinearity is a necessary condition for this phenomenon. If a phenomenon
happens for both models, one naturally concludes that non-linearity is not
necessary to cause the phenomena. For example, posterior collapses in Bayesian
deep learning were first thought to be caused by expressivity due to non-
linearity (Alemi et al.,, 2018); however, recent works (Lucas et al.,, 2019;
Wang and Ziyin,, 2022) found that models without non-linearity can still
induce the collapse.

##  3 A Landscape Theory of Self-Supervised-Learning

This section presents the main theoretical results. Let
$\\{\hat{x}_{i}\\}_{i}^{N}$ be a dataset with $N$ data points. For every data
point $\hat{x}$, we augment it with an i.i.d. noise $\epsilon$ such that
$x:=\hat{x}+\epsilon$. To be concrete, we start with considering the standard
contrastive loss, InfoNCE (Oord et al.,, 2018):

$L=\mathbb{E}_{\epsilon}\left[-\sum_{i=1}^{N}\log\frac{\exp(-|f(x_{i})-f(x_{i}^{\prime})|^{2}/2)}{\sum_{j\neq
i}\exp(-|f(x_{i})-f(\chi_{j})|^{2}/2)+\exp(-|f(x_{i})-f(x_{i}^{\prime})|^{2}/2)}\right],$
(1)

where $f(x)\in\mathbb{R}^{d_{1}}$ is the model output; all $x$, $x^{\prime}$
and $\chi$ are augmented data points for some independent additive noise
$\epsilon$ such that
$\mathbb{E}_{\epsilon}[x]=\hat{x}=\mathbb{E}_{\epsilon}[x^{\prime}]\neq\mathbb{E}_{\epsilon}[\chi]=\hat{\chi}$.
We decompose the model output into a general function
$\phi(x)\in\mathbb{R}^{d_{0}}$ and the last-layer weight matrix
$W\in\mathbb{R}^{d_{1}\times d_{0}}$: $f(x)=W\phi(x)$. The covariance of
$\phi(\hat{x})$ is
$A_{0}:=\mathbb{E}_{\hat{x}}[\phi(\hat{x})\phi(\hat{x})^{T}]$, and the
covariance of the data-augmented penultimate layer representation is
$\Sigma:=\mathbb{E}_{x}[\phi({x})\phi({x})^{T}]$. The effect of data
augmentation on the learned representation is captured through a symmetric
matrix $C:=\Sigma-A_{0}$. For a general $\phi$, the eigenvalues of $C$ can be
either positive or negative. When $\phi$ is the identity mapping, $A_{0}$
becomes the empirical data covariance, $C$ becomes PSD and is the covariance
of the noise $\epsilon$, and $\Sigma$ is the covariance of the augmented data.
In some sense, this loss function captures the essence of SSL: the numerator
encourages the representation $f(x)$ to be closer to the representation of
similar data, and the denominator encourages a separation between dissimilar
data.

For a fixed set of noises, we can write the InfoNCE in a cleaner form:

$\displaystyle L_{\epsilon}$
$\displaystyle=\mathbb{E}_{\hat{x}}\left\\{\frac{1}{2}|f(x)-f(x^{\prime})|^{2}+\log\mathbb{E}_{\hat{\chi}}\left[\exp\left(-\frac{1}{2}|f(x)-f(\chi)|^{2}\right)\right]\right\\},$
(2)

where we used $\mathbb{E}_{\hat{x}}$ to denote an averaging over the training
set.

| Hessian | Dim. | Complete
---|---|---|---
InfoNCE | $A_{0}$ | ✗ | ✗
NT-Xent (SimCLR) | $A_{0}-\nicefrac{{C}}{{N}}$ | ✓ | ✓
Spectral Contrastive | $C$ | ✗ | ✗
Barlow Twins | $A_{0}+C$ | ✗ | ✗
\+ Normalization | - | ✓ | ✗
\+ bias | - | ✓ | ✓
\+ Weight Decay | $+\gamma I$ | ✓ | ✓
Table 1: What shapes the SSL landscapes around the origin? For each of the SSL
losses, the combination of data covariance ($A_{0}$), data-augmentation
covariance ($C$), and dataset size ($N$) can affect its stability and thus
determine the presence (✓) and absence (✗) of dimensional/complete collapse
(Here, a ✓means “there exists a hyperparameter setting and data distribution
such that the relevant collapse happens;” see section 3). Beyond collapses,
the theory implies that SCL, whose landscape is formed primarily by data
augmentation, is more robust to data imbalance than InfoNCE, which is affected
primarily by the data (see section 4).

In this notation, we have
$\mathbb{E}_{\epsilon}\mathbb{E}_{\hat{x}}[x]=\mathbb{E}_{x}[x]$ and
$\mathbb{E}_{\epsilon}[L_{\epsilon}]=L$. We first show that the expansion of
the loss function around the origin takes a rather universal form. We then
find analytical solutions to the stationary points of this landscape and study
their relevance to feature learning and collapses. See Table 1 for a summary
of the main results. The proofs are presented in Appendix E. For a
quantitative understanding, we mainly focus on the case when $\phi$ is the
identity function. We discuss the general nonlinear case in Section 4.1.

###  3.1 Landscape of a Linear Model

In this section, we first analyze representative SSL loss functions and
identify properties of the landscapes and show that to leading order in $W$,
the local geometry of SSL losses takes the following form

$L=-{{\rm Tr}[WBW^{T}]}+\frac{1}{8}{\rm Var}[|W(x-\chi)|^{2}].$ (3)

A distinctive feature of Eq. (3) is that its first and third-order terms
vanish. This is because the loss function is invariant to a left rotation of
$W$. We will see that this symmetry in rotation is a 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:

$L=\underbrace{{\rm
Tr}[WCW^{T}]}_{E}+\underbrace{\mathbb{E}_{\epsilon,\hat{x}}\left\\{\log\mathbb{E}_{\hat{\chi}}\left[\exp\left(-\frac{1}{2}|W(x-\chi)|^{2}\right)\right]\right\\}}_{-S}.$
(4)

Expanding the entropy term to the fourth order, we obtain111Throughout, we use
$||\cdot||$ to denote the $L_{2}$ norm for vectors and Frobenius norm for
matrices.

$\displaystyle-S$
$\displaystyle=-\mathbb{E}_{x}\mathbb{E}_{\chi}\left[\frac{1}{2}|W(x-\chi)|^{2}\right]+\frac{1}{8}{\rm
Var}[|W(x-\chi)|^{2}]+O(||W||^{6}).$ (5)

This (perturbative) decomposition of entropy deserves some special attention.
The entropy decomposes into a repulsion term that is second order in $W$, and
a variance term that is fourth order in $W$. The first term encourages a
repulsion between $x$ and its augmentation, which counteracts the effect of
the energy term. The first term can be decomposed into

$\mathbb{E}_{x}\mathbb{E}_{\chi}\left[\frac{1}{2}|W(x-\chi)|^{2}\right]={\rm
Tr}[WCW^{T}]+{\rm Tr}[WA_{0}W^{T}].$ (6)

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 $A_{0}$. It is intriguing to see 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 $A_{0}$ remains. This
means that to leading order, the learned representation has a larger variation
along the directions where the data has a larger variation, which is what one
naively expects. Collecting results, we have obtained the loss landscape in
the neighborhood of the origin as $L=-{\rm Tr}[WA_{0}W^{T}]+\frac{1}{8}{\rm
Var}[|W(x-\chi)|^{2}]+O(||W||^{6})$.

NT-xent (SimCLR). Next, we analyze Normalized Temperature Cross-Entropy (NT-
xent) used in SimCLR (Chen et al., 2020a, ). As in Tian, (2022), we note
InfoNCE can be generalized to encompass NT-xent as follows:

$L=\mathbb{E}_{\epsilon}\left[-\sum_{i=1}^{N}\log\frac{\exp(-|f(x_{i})-f(x_{i}^{\prime})|^{2}/2)}{\sum_{\chi\neq
x}\exp(-|f(x_{i})-f(\chi_{j})|^{2}/2)+\alpha\exp(-|f(x_{i})-f(x_{i}^{\prime})|^{2}/2)}\right].$
(7)

In contrast to InfoNCE, here one of the terms in the denominator is reweighted
by a factor of $\alpha\geq 0$. Two interesting limits are $\alpha=1$, where we
recover the InfoNCE loss, and $\alpha=0$, where we obtain the popular NT-xent
loss used in SimCLR (Chen et al., 2020a, ). For general $\alpha$, 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.

For a perceptron, the same perturbative expansion as Eq. (4)–(6) gives

$L=\frac{1-\alpha}{N}{\rm Tr}[WCW^{T}]-{\rm Tr}[WA_{0}W^{T}]+\frac{1}{8}{\rm
Var}[|W(x-\chi)|^{2}]+O(||W||^{6})+O(||W||^{4}N^{-1}).$ (8)

Now, the Hessian of the origin is no longer guaranteed to be ND. In fact, if
$\frac{1-\alpha}{N}C-A_{0}\geq 0$, $W=0$ becomes an isolated local minimum.

Landscape Analysis. The above discussion shows that the landscape for a class
of situations in self-supervised contrastive learning can be reduced to an
effective form in Eq. (3). The following proposition 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. Then, $L=-{\rm Tr}[WBW^{T}]+{\rm
Tr}[W\Sigma W^{T}W\Sigma W^{T}]$.

When training ends, one expects the model to locate at (at least close to) a
stationary point of the loss. It is thus important to identify all the
stationary points of this loss function.

######  Theorem 1.

Let $d^{*}:=\min(d_{0},d_{1})$. Let the data and noise be Gaussian. All
stationary points $W$ of Eq. (3) satisfy
$W^{T}W=\frac{1}{2}\Sigma^{-1/2}UM\Lambda U^{T}\Sigma^{-1/2}$, where $U\Lambda
U^{T}$ is the eigenvalue decomposition of $\Sigma^{-1/2}B\Sigma^{-1/2}$, and
$M$ is an arbitrary (masking) diagonal matrix containing only zero or one such
that (1) $M_{ii}=0$ if $\Lambda_{ii}<0$ and (2) contain at most $d^{*}$
nonzero terms.

Additionally, if $C$ and $A_{0}$ commute, all stationary points satisfy

$W^{T}W=\frac{1}{2}\Sigma^{-1}B_{M}\Sigma^{-1},$ (9)

where $B_{M}$ denotes the matrix obtained by masking the eigenvalues of $B$
with $M$.

This stationary-point condition implies the direct cause of the dimensional
collapse. Namely, dimensional collapse happens when the eigenvalues of the
matrix $B$ become negative. The eigenvalues of $B$, in turn, depend on the
competition between data augmentation and the data feature. Comparing the
commuting case with the noncommuting case, we see that the main difference is
that when $C$ does not commute with $A_{0}$, the augmentation can also change
the orientation of the learned representation; otherwise, augmentation only
affects the eigenvalues. To focus on the most important terms, we now assume
that the augmentation is well-aligned with the features such that the
augmentation covariance commute with the data covariance.

######  Assumption 1.

From now on, we assume $CA_{0}=A_{0}C$.

For the case of weighted InfoNCE, we have that $B=A_{0}-\frac{1-\alpha}{N}C$.
Let $a_{i}$ denote the $i$-th eigenvalue of the $A$ and $c_{i}$ that of $C$
viewed in a predetermined order; then, the $i$th subspace collapses when
$\frac{1-\alpha}{N}c_{i}\geq a_{i}$, 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 a certain
type of local minima. When weight decay is used, the condition for collapse
becomes $\frac{1-\alpha}{N}c_{i}+\gamma\geq a_{i}$. It becomes easier to cause
a collapse when weight decay is used.

The global minimum of the loss function is quite easy to find. For all
stationary points, the loss function takes a simple form; $L=-\frac{1}{4}{\rm
Tr}[\Sigma^{2}B_{M}B]$. Thus, $L$ becomes more and more negative if the
eigenvalues of $B_{M}$ align with the largest eigenvalues of $B$. Namely, the
global minimum is achieved if $M$ leaves the largest eigenvalues of $B$
intact.

Because the stationary points contain collapsed solutions where the
eigenvalues of $W^{T}W$ are zero, one is naturally interested in how likely it
is to converge to these solutions.

######  Proposition 2.

($W^{T}W$ achieves maximum possible rank) Let $m$ denote the number of
positive eigenvalues $B$. Then, ${\rm rank}(W^{T}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 a maximum
possible rank. In fact, this result implies that the collapses may be 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
$||f(x)||^{2}=c$. When the normalization is applied, only the direction of the
learned representation matters. While this is a 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 as a regularization term:
$R:=(\mathbb{E}_{x}||f(x)||^{2}-c)^{2}$:

$L_{\rm norm}=Eq.~{}\eqref{eq: general linear landscape}+\kappa R.$ (10)

Note that this regularization term achieves two things simultaneously: (1)
$||f(x)||^{2}=c$ for all $x$ is a minimizer of the loss function; (2) the
regularization is invariant to any rotation of the learned representation. For
a linear model, we note that this condition is not entirely the same as a
direct normalization of the representation because it is generally impossible
to achieve $||Wx||^{2}=c$ for all $x$ because a linear model has limited
expressivity. However, it is generally possible to achieve the slightly weaker
condition: the representation has a norm $1$ on 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 $\kappa$ 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 it also encourages the representation to have a constant variance. Our
theory also explains why the VICReg is observed to experience collapses when
$\kappa$ is not large enough. As $\kappa$ tends to infinity, this constraint
will become perfectly satisfied. We thus take the infinite $\kappa$ limit to
study the effect of normalization.

The following proposition gives a condition that all stationary points of Eq.
(10) satisfy.

######  Proposition 3.

Let $\rho(W):={\rm Tr}[W\Sigma W^{T}]$, $B^{\prime}:=B+2\kappa(c-\rho)\Sigma$,
and let $\Lambda_{i}$ be the eigevalues of $B^{\prime}$. Then, every
stationary point of Eq. (10) satisfy
$W^{T}W=\frac{1}{2}\Sigma^{-1}B^{\prime}_{M}\Sigma^{-1}$, where $M$ is an
arbitrary diagonal mask of the eigenvalues of $B^{\prime}$ containing only
zero or one such that (1) $M_{ii}=0$ if $\Lambda_{i}<0$ and (2) contain at
most $d^{*}$ nonzero terms.

Compared with the unnormalized case, the term $2\kappa(1-\rho)\Sigma_{M}$
emerges due to normalization. The effect of normalization is as expected: it
shrinks the norm of the model if $\rho>1$, and it expands the model if
$\rho<1$, and it does not have any effect if we have already achieved
$\rho=1$. Interestingly, this rescaling effect is anisotropic and stronger
along the directions of larger eigenvalues of the covariance of the augmented
data $\Sigma$.

The next theorem gives the explicit form of $\rho$ at the stationary points.

######  Proposition 4.

For any stationary point $W^{*}$, $c-\rho(W^{*})=\frac{c-\frac{1}{2}{\rm
Tr}[\Sigma^{-1}B_{M}]}{1+\kappa d_{M}}$, where $d_{M}$ is the number of non-
zero eigenvalues of $B^{\prime}_{M}$.

For a finite $\kappa$, these results suggest that collapses can still happen.
For VICReg, $B=-A_{0}$, and the complete collapse can happen when
$\kappa\ll||A_{0}||/c||\Sigma||$ – this explains the experimental observation
of collapses for small values of $\kappa$ in VICReg loss (Bardes et al.,,
2021).

Lastly, to understand normalization, we are interested in the case of
$\kappa\to\infty$. Combining Proposition 3 and 4, we have proved the following
theorem, showing that the asymptotic solution converges to a form independent
of $\kappa$.

######  Theorem 2.

Let $W_{\kappa}$ be a stationary point of Eq. (10) at fixed $\kappa$. Then,

$\lim_{\kappa\to\infty}W_{\kappa}^{T}W_{\kappa}=\frac{1}{2}\Sigma^{-1}\left[B_{M}+\frac{2c-{\rm
Tr}[\Sigma^{-1}B_{M}]}{d_{M}}\Sigma_{M}\right]\Sigma^{-1}.$ (11)

The correction term $\frac{2c-{\rm Tr}[\Sigma B_{M}]}{d_{0}}\Sigma_{M}$
emerges as a result of applying normalization. The effect can be easier to
understand if we write the solution as

$W^{T}W=\frac{1}{2}\left[\Sigma^{-1}B_{M}-\frac{{\rm
Tr}[\Sigma^{-1}B_{M}]}{d_{M}}M+\frac{2c}{d_{M}}\right]\Sigma^{-1},$ (12)

where we have used the relation $\Sigma_{M}\Sigma^{-1}=M$. Note the term in
brackets: it subtracts the average eigenvalue of $\Sigma^{-1}B_{M}$ from
$\Sigma^{-1}B_{M}$ and shifts the remaining eigenvalues positively by
$2c/d_{M}$. Because the eigenvalues of $WW^{T}$ must be positive, the
following condition must hold for all solutions:

$\lambda_{i}+{2c}/{d_{M}}>\bar{\lambda},$ (13)

where $\lambda_{i}$ are the eigenvalues of $\Sigma^{-1}B_{M}$ and
$\bar{\lambda}$ is its average. Namely, for the $i-$th dimension not to
collapse, it must be smaller than the average eigenvalues by at most
$2c/d_{M}$. Any smaller eigenvalues 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.1. One finds that the condition for collapse becomes
heavily dependent on the data structure, and there are cases where collapses
become harder, and there are cases where collapses become much easier.
Importantly, it also becomes the case that a 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\to 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 a
translation in the learned representation. However, this effect dramatically
changes if we apply normalization at the same time. This is because
normalization removes the translation symmetry of the effective loss, and the
trivial solution $W=0$, $b=1$ becomes the simplest way to 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 $\mathbb{E}[x]=0$. Then, all stationary points satisfy Eq.
(9), subject to the constraint that $\rho(W)={\rm Tr}[W^{T}\Sigma W]\leq c$.

Namely, the solution reverts to the case where there is no 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
$c<(a_{i}-c_{i})/(a_{i}+c_{i})$ for all $i$, a complete collapse can happen
despite normalization. When $c=1$ and $c_{i}\ll a_{i}$, $\rho\approx d_{M}/2$
and the constraint indicates that $d_{M}\leq 2$: when the augmentation is very
weak, there are at most $2$ nontrivial subspaces. This is too restrictive for
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-similarity loss, and our result explains why
previous works have tried to introduce asymmetry to cosine similarity to avoid
collapses (Grill et al.,, 2020; Chen and He,, 2021).

Relevant Loss Functions. Having developed a 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

$L_{SCL}=-2\mathbb{E}[f(x)^{T}f(x^{\prime})])+\mathbb{E}[(f(x)^{T}f(\chi))^{2}]+const.\quad\quad\text{s.t.
$||f(x)||^{2}=1$}.$ (14)

Let $f(x)=Wx$ be linear, the distributions are zero-mean Gaussian, and ignore
the normalization. This loss function becomes

$\displaystyle L_{SCL}$ $\displaystyle=-2{\rm Tr}[WCW^{T}]+{\rm Tr}[W\Sigma
W^{T}W\Sigma W^{T}].$ (15)

When normalization exists, we can apply the result in Section 3.2. By our
argument, there is no 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 $A_{0}$.

The case of Barlow Twin (BT) (Zbontar et al.,, 2021) is similar. While the
fourth-order term of BT is much more complicated due to the imbalance created
by the $\lambda$ term. The second-order term can be identified easily:
$L_{BT}=-2{\rm Tr}[W\Sigma W^{T}]+O(||W||^{4})$. This also does not collapse.
A difference between the SCL loss and InfoNCE is that the learned
representation has a spread that aligns with the combination of the feature
and the augmentation strength.

(a) 1D landscape

(b) No collapse

(c) Dimensional collapse

(d) Complete collapse

Figure 2: Landscape of Resnet18 on CIFAR10 with SimCLR qualitatively agrees
with our linear theory. (a) Training objective $L$ as a function of a
rescaling of the last layer $W\to aW$. (b-d) $L$ as a function of a $2d$
rescaling of the last 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, a qualitative change that
leads to different types of collapses in qualitative agreement with our linear
theory (cf. Figure 1). Additionally, we also notice the same qualitative
changes in the landscape in simpler nonlinear models (see Appendix A).

##  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 general. In fact, the loss landscape we have studied is
quite close to the most general landscape one can have. Let $L(f(x))$ be a
general SSL loss function for data point $x$. The quality of the learned
representation should be independent of the population-level orientation of
the representation. Therefore, the loss function should satisfy a rotational
invariance. Namely, for any rotation matrix $R$, $L(x)=L(Rf(x))$; this
rotational invariance implies that the loss should expand as
$L(f(x))=af(x)^{T}f(x)+b[f(x)^{T}f(x)]^{2}+O(f(x)^{6})$. Note that all the
odd-order terms of $f(x)$ vanish due to the rotational symmetry. Substituting
$f(x)=W\phi(x)$ in the loss function, we obtain a very general form of
landscape that $W$ obeys:

$L(W,\phi)={\rm Tr}[W^{T}WA]+\sum_{ijklmn}W_{im}W_{jm}W_{kn}W_{ln}Z_{ijki},$
(16)

where $A$ and $Z$ are dependent on $\phi$. 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 $A$
becomes crucial for our understanding. To illustrate, we train a Resnet18 on
CIFAR10 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 matrix of the last
layer $W^{*}_{\rm last}$ by a factor $a$ and compute the loss as a function of
$a$. See Figure 2-a. We then partition the singular values of $W^{*}_{\rm
last}$ into the larger half and the smaller half. We rescale the larger half
by a factor $r_{1}$ and the smaller half by $r_{2}$. We plot the loss as a
$2d$ function of $(r_{1},r_{2})$ in Figure 2. In all cases, the landscape
features qualitative changes comparable to those in Figure 1.

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 $\sqrt{\lambda}$ of $W^{T}W$ are
the order parameters of the system, and the phase transitions happen when
$\lambda$ 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. One can also compare with the recent result on the second-order
phase transition in the context of supervised learning (Ziyin et al.,, 2022;
Ziyin and Ueda,, 2022). What is shared between the collapse in SSL and
supervised learning is the need for an effective second-order regularization
term to appear. What is unique to SSL is that the phase transition can happen
even when there is only one layer, whereas in supervised learning one needs at
least two layers. 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 are effectively
frozen along this direction.

###  4.2 Robustly inducing good collapses

Contrary to previous works, a recent work (Cosentino et al.,, 2022) has
suggested that dimensional collapse can be beneficial and significantly
improve the generalization performance of the model. This observation raises a
question. How can dimensional collapse be beneficial and how can it be
induced? In the following, we first introduce $\beta$-InfoNCE, which can
adjust the degree of dimensional collapse, and analyze the collapse behavior
to elucidate the mechanism of _task-alligned collapse_.

Adjusting the degree of dimensional collapse with $\beta$-InfoNCE. Despite the
potential benefit, existing SSL loss functions cannot robustly induce
dimensional collapse. InfoNCE is insufficient to induce a collapse, and the
collapse induced by SimCLR depends on a vanishingly small parameter $1/N$. One
thus wonders whether there is a loss function that allows us to induce
collapsing behavior in a more predictable 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 competition
for the model Hessian at the origin. For InfoNCE, one way to achieve this is
to weigh the entropy term by a general factor $\beta$:

$\mathbb{E}_{x}\left\\{\frac{1}{2}|f(x)-f(x^{\prime})|^{2}+\beta\log\mathbb{E}_{\chi}\left[\exp\left(-\frac{1}{2}|f(x)-f(\chi)|^{2}\right)\right]\right\\}.$

Due to its similarity with the $\beta$-VAE in Bayesian learning, we call it
the $\beta$-InfoNCE. The leading term in the loss function becomes

$-{\rm Tr}[W(A_{0}-\left(1-\beta\right)C)W^{T}].$

When $1-\beta>0$, the augmentations $C$ pull the representation towards zero.
When the augmentation is as strong as the feature variations, a collapse
happens. One can thus introduce collapse by setting $\beta$ to be sufficiently
small. When $1-\beta<0$, 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 $\beta$; $\beta=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 $\beta$ parameter here directly controls
the level of difficulty of collapse. The next section shows how one can
introduce some controllable collapse to improve the model performance.

Figure 3: Top: Phase diagram of representational collapses. Bottom:
$\beta-$InfoNCE with $\beta=0.5$. The generalization error of a downstream
regression task where the data augmentation (1) is isotropic and
noninformative or (2) aligns with the style. We see that the performance
worsens as collapses happen for the noninformative augmentation and improves
as the collapse happens for the style-targeting augmentation.

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 $d_{c}<d_{0}$ and (2) a task-irrelevant set:
$x=(x_{1},...,x_{d_{c}},...,x_{d_{0}})$. Our result suggests a precise way to
remove the irrelevant features from the learned representation. For the
purpose of causing a robust collapse, we use the $\beta$-InfoNCE with
$\beta=1/2$. For illustration, we consider the simple case $d_{c}=1$ and
$d_{0}=2$. For any input $x=(x_{1},x_{2})$, the label is generated as a linear
function of $x_{1}$: $y=cx_{1}$.

Correspondingly, we consider a structured data augmentation $x=\hat{x}+\sigma
R\xi$, where $R\in\mathbb{R}^{d_{0}\times d_{0}}$ is
$R=diag(\sqrt{1-\theta},\sqrt{\theta})$, where $\theta\in[0,1]$. The parameter
$\sigma$ controls the overall strength of the augmentation, and $\theta$
controls the orientation of the strength. When $\theta=0.5$, we have an
uninformative isotropic noise that has often been used in practice. When
$\theta=1$, the augmentation is only on the task-irrelevant feature, and when
$\theta=0$, the augmentation is only on the content. Since the prediction
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 $z:=f(\hat{x})$ to train a ridge linear regressor that
minimizes
$\min_{G}\mathbb{E}_{\hat{x}}[||Gz-y({\hat{x}})||^{2}]+0.001||G||^{2}$. 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-like behaviors in the
generalization performance and that a data augmentation aligning with the
task-irrelevant dimension improves performance.

Figure 4: Spectral Contrastive loss (SCL) is more robust against data
imbalance than InfoNCE.  We train SimCLR and SCL ResNet-12 models on
imbalanced versions of CIFAR-10. We see that SCL is more robust than SimCLR,
as suggested by our theory. These results are especially pronounced when there
is no projector head.

###  4.3 Robustness to Data Imbalance

Our theory is not only relevant for understanding collapses but can also be
used to understand 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 minority subgroups. As another application of our analysis, we
illustrate the robustness of different techniques is not equal. As we have
seen, the learned model $W^{T}W$ has eigenvalues that, to the leading order,
are proportional to the Hessian $B$, which is different for each loss
function. As previously summarized in Table 1, for InfoNCE and SimCLR, the
learned model aligns with the eigenvalues of the data covariance $A_{0}$,
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 can be less dependent on data and thus more
robust against data imbalance. See Figure 4. More experimental details are
given in Appendix C.

##  5 Conclusion

In this work, we approached the problem of collapses in SSL from a landscape
perspective. We analytically solved an effective landscape that can be
extended to understand the effect of normalization. Our result suggests that
dimensional collapse can be well understood in the minimal setting and is
something neutral to learning on its own. We showed that when task-irrelevant
dimensions are targeted, dimensional collapse can result in dramatically
improved performance, whereas an uninformative noise will (without good luck)
leads to collapses in the dimensions that are relevant to the task. 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. We 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.

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 a convergence towards the saddle points in the landscape
(Ziyin et al.,, 2021), which, as we have shown, are the collapsed solutions.
Investigating the role of dynamics in the collapse is thus a crucial future
problem.

## References

  * Alemi et al., (2018) Alemi, A., Poole, B., Fischer, I., Dillon, J., Saurous, R. A., and Murphy, K. (2018).  Fixing a broken elbo.  In International Conference on Machine Learning, pages 159–168. PMLR.
  * Arora et al., (2019) Arora, S., Khandeparkar, H., Khodak, M., Plevrakis, O., and Saunshi, N. (2019).  A Theoretical Analysis of Contrastive Unsupervised Representation Learning.  In Proc. Int. Conf. on Machine Learning (ICML).
  * Ash et al., (2021) Ash, J. T., Goel, S., Krishnamurthy, A., and Misra, D. (2021).  Investigating the role of negatives in contrastive representation learning.  Int. Conf. on Artificial Intelligence and Statistics.
  * Bardes et al., (2021) Bardes, A., Ponce, J., and LeCun, Y. (2021).  Vicreg: Variance-invariance-covariance regularization for self-supervised learning.  arXiv preprint arXiv:2105.04906.
  * Caron et al., (2020) 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 Information Processing Systems (NeurIPS).
  * Caron et al., (2021) Caron, M., Touvron, H., Misra, I., Jegou, H., Mairal, J., Bojanowski, P., and Joulin, A. (2021).  Emerging Properties in Self-Supervised Vision Transformer.  arXiv, abs/2104.14294.
  * (7) Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. (2020a).  A simple framework for contrastive learning of visual representations.  In International conference on machine learning, pages 1597–1607. PMLR.
  * (8) Chen, T., Kornblith, S., Swersky, K., Norouzi, M., and Hinton, G. E. (2020b).  Big Self-Supervised Models are Strong Semi-Supervised Learners.  Adv. in Neural Information Processing Systems, 33.
  * Chen and He, (2021) Chen, X. and He, K. (2021).  Exploring Simple Siamese Representation Learning.  In Proc. Int. Conf. on Computer Vision and Pattern Recognition (CVPR).
  * Cosentino et al., (2022) 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.
  * Ermolov et al., (2021) Ermolov, A., Siarohin, A., Sangineto, E., and Sebe, N. (2021).  Whitening for self-supervised representation learning.  In International Conference on Machine Learning, pages 3015–3024. PMLR.
  * Grill et al., (2020) 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, k., Munos, R., and Valko, M. (2020).  Bootstrap your own latent: A new approach to self-supervised Learning.  In Proc. Adv. on Neural Information Processing Systems (NeurIPS).
  * HaoChen et al., (2021) HaoChen, J. Z., Wei, C., Gaidon, A., and Ma, T. (2021).  Provable guarantees for self-supervised deep learning with spectral contrastive loss.  Advances in Neural Information Processing Systems, 34:5000–5011.
  * He et al., (2020) He, K., Fan, H., Wu, Y., Xie, S., and Girschick, R. (2020).  Momentum Contrast for Unsupervised Visual Representation Learning.  In Proc. Int. Conf. on Computer Vision and Pattern Recognition (CVPR).
  * Hsu et al., (2019) Hsu, H., Qi, H., and Brown, M. (2019).  Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification.  arXiv, abs/1909.06335.
  * Hua et al., (2021) 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 et al., (2021) Jing, L., Vincent, P., LeCun, Y., and Tian, Y. (2021).  Understanding dimensional collapse in contrastive self-supervised learning.  arXiv preprint arXiv:2110.09348.
  * Kawaguchi, (2016) Kawaguchi, K. (2016).  Deep learning without poor local minima.  Advances in Neural Information Processing Systems, 29:586–594.
  * Kugelgen et al., (2021) 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 and Lifshitz, (2013) Landau, L. D. and Lifshitz, E. M. (2013).  Statistical Physics: Volume 5, volume 5.  Elsevier.
  * Liu et al., (2021) Liu, H., HaoChen, J. Z., Gaidon, A., and Ma, T. (2021).  Self-supervised learning is more robust to dataset imbalance.  International Conference on Learning Representations.
  * Lucas et al., (2019) Lucas, J., Tucker, G., Grosse, R., and Norouzi, M. (2019).  Don’t blame the elbo! a linear vae perspective on posterior collapse.
  * Mitrovic et al., (2020) Mitrovic, J., McWilliams, B., Walker, J., Buesing, L., and Blundell, C. (2020).  Representation learning via invariant causal mechanisms.  arXiv preprint arXiv:2010.07922.
  * Nozawa and Sato, (2021) Nozawa, K. and Sato, I. (2021).  Understanding negative samples in instance discriminative self-supervised representation learning.  Advances in Neural Information Processing Systems, 34:5784–5797.
  * Oord et al., (2018) Oord, A. v. d., Li, Y., and Vinyals, O. (2018).  Representation learning with contrastive predictive coding.  arXiv preprint arXiv:1807.03748.
  * Pokle et al., (2022) 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 et al., (2021) 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 et al., (2022) 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. on Machine Learning (ICML).
  * Saxe et al., (2013) Saxe, A. M., McClelland, J. L., and Ganguli, S. (2013).  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  arXiv preprint arXiv:1312.6120.
  * Tian, (2022) Tian, Y. (2022).  Deep contrastive learning is provably (almost) principal component analysis.  arXiv preprint arXiv:2201.12680.
  * Tian et al., (2021) Tian, Y., Chen, X., and Ganguli, S. (2021).  Understanding self-supervised Learning Dynamics without Contrastive Pairs.  In Proc. Int. Conf. on Machine Learning (ICML).
  * Tian et al., (2020) Tian, Y., Sun, C., Poole, B., Krishnan, D., Schmid, C., and Isola, P. (2020).  What makes for good views for contrastive learning?  Advances in Neural Information Processing Systems, 33:6827–6839.
  * (33) Tosh, C., Krishnamurthy, A., and Hsu, D. (2021a).  Contrastive estimation reveals topic posterior information to linear models.  J. Mach. Learn. Res., 22:281–1.
  * (34) Tosh, C., Krishnamurthy, A., and Hsu, D. (2021b).  Contrastive learning, multi-view redundancy, and linear models.  In Algorithmic Learning Theory, pages 1179–1206. PMLR.
  * Trivedi et al., (2022) Trivedi, P., Lubana, E. S., Heimann, M., Koutra, D., and Thiagarajan, J. J. (2022).  Analyzing data-centric properties for contrastive learning on graphs.  arXiv preprint arXiv:2208.02810.
  * (36) Tsai, Y.-H., Wu, Y., Salakhutdinov, R., and Morency, L.-P. (2021a).  Self-supervised Learning from a Multi-view Perspective.  In Proc. Int. Conf. on Learning Representations (ICLR).
  * (37) Tsai, Y.-H. 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 on Learning Representations.
  * Wang and Liu, (2021) Wang, F. and Liu, H. (2021).  Understanding the behaviour of contrastive loss.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 2495–2504.
  * Wang and Isola, (2020) Wang, T. and Isola, P. (2020).  Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere.  In Proc. Int. Conf. on Machine Learning (ICML).
  * Wang et al., (2022) Wang, Y., Zhang, Q., Wang, Y., Yang, J., and Lin, Z. (2022).  Chaos is a ladder: A new theoretical understanding of contrastive learning via augmentation overlap.  International Conference on Learning Representations.
  * Wang and Ziyin, (2022) Wang, Z. and Ziyin, L. (2022).  Posterior collapse of a linear latent variable model.  arXiv preprint arXiv:2205.04009.
  * Wei et al., (2021) Wei, C., Shen, K., Chen, Y., and Ma, T. (2021).  Theoretical Analysis of Self-Training with Deep Networks on Unlabeled Data.  In Proc. Int. Conf. on Learning Representations (ICLR).
  * Zbontar et al., (2021) Zbontar, J., Jing, L., Misra, I., LeCun, Y., and Deny, S. (2021).  Barlow twins: Self-supervised learning via redundancy reduction.  In International Conference on Machine Learning, pages 12310–12320. PMLR.
  * Zimmerman et al., (2021) Zimmerman, R., Sharma, Y., Schneider, S., Bethge, M., and Brendel, W. (2021).  Contrastive Learning Inverts the Data Generating Process.  In Proc. Int. Conf. on Machine Learning (ICML).
  * Ziyin et al., (2022) Ziyin, L., Li, B., and Meng, X. (2022).  Exact solutions of a deep linear network.  arXiv preprint arXiv:2202.04777.
  * Ziyin et al., (2021) Ziyin, L., Li, B., Simon, J. B., and Ueda, M. (2021).  Sgd with a constant large learning rate can converge to local maxima.  arXiv preprint arXiv:2107.11774.
  * Ziyin and Ueda, (2022) Ziyin, L. and Ueda, M. (2022).  Exact phase transitions in deep learning.  arXiv preprint arXiv:2205.12510.

##  Appendix A Additional Numerical Results

In this section, we validate our theory with numerical results. Unless
specified otherwise, the dimension of the learned representation is set to be
equal to the input dimension: $d_{0}=d_{1}$.

Figure 5: The three smallest singular values of $W^{T}W$ as a function of 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 $\sigma^{4}$ and do not
vanish for any finite value of $\sigma$. Mid: Weight InfoNCE. $\alpha=0.1$,
$\sigma=5$. Collapse happens at the critical dataset size predicted by the
theory. Right: (Sqrt) Eigenvalues of $WW^{T}$ in $\beta$-InfoNCE. The
collapses can be well controlled. Figure 6: A collapse happens easily when the
learned representation is normalized. The smallest eigenvalues of $A_{0}$ are
roughly $0.2$, and the collapse happens much before the noise reaches this
strength.

No Collapse for InfoNCE. 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 $\sigma^{4}$ where
$\sigma^{2}$ is the strength (namely, the variance) of the augmentation. See
the left panel of Fig. 5. We use the vanilla InfoNCE loss defined in (1) with
a linear model. The training set is sampled from $\mathcal{N}(0,I_{32})$. The
training proceeds with Adam with a learning rate of $6e-4$ with full batch
training for $5000$ iterations. We use a simple diagonal Gaussian noise with
variance $\sigma^{2}$ for data augmentation. We see that the singular values
scale as $\sigma^{4}$ and never vanishes, as the theory 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 as the previous
experiment. We set $\alpha=0.1$ and change the size of the training set.
Theory suggests that for a collapse in the $i-$th subspace to happen, the size
of the dataset needs to obey

$N>\frac{a_{i}}{c_{i}(1-\alpha)}:=N_{crit}.$ (17)

See the middle panel of Figure 5. We show the smallest three eigenvalues of
$W^{T}W$ (roughly 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.

Collapses in $\beta$-InfoNCE. With $\beta<1$, one can cause collapses in a
predictable and controllable way. In this experiment, we let $d_{0}=5$ and we
plot all five eigenvalues of $W^{T}W$ as we increase the strength of an
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^{T}W$ when we apply the standard representation
normalization in practice: $f(x)\to f(x)/||f(x)||$. To facilitate comparison,
we also use the same dataset and training procedure as before. See Figure 6.
We see that normalization does cause a collapse in the smallest eigenvalues at
an augmentation strength much smaller than the feature variation.

Figure 7: The Landscape of nonlinear models is very similar to the landscape
of 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 of a nonlinear model with a weak augmentation. Middle: with
intermediate augmentation. Right: with strong augmentation.

##  Appendix 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$: $W_{last}\to
aW_{last}$ and plot the loss function along this direction. See the top panel
of Figure 7 for both the tanh and the ReLU nonlinearity. We then rescale the
two rows of the weight matrix of the model by $r_{1}$ and $r_{2}$
respectively: $W=(w_{1},W_{2})^{T}\to(r_{1}w_{1},r_{2}w_{2})$. We see that the
landscape of the model is qualitatively the same as that of the linear models,
shown in Figure 1.

##  Appendix C Setup for Imbalanced Data Experiments

Creating an Imbalanced Dataset: For our experiments measuring the influence on
imbalanced datasets on SSL training, we use CIFAR-10 by sampling 20000 samples
out of the 50000 training samples. The sampling process is described by a
Dirichlet distribution and is often used to analyze effects of heterogeneity
and data imbalance in Federated Learning problems (Hsu et al.,, 2019).
Specifically, a small value of the distribution parameter yields a highly
imbalanced dataset, while a large value yields a 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, 161, 4633, 1015, 131, 2386]

  * •

No imbalance: [2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000]

Training Setup: We use ResNet-12 models as the backbone for all experiments
due to computational constraints. SimCLR augmentations (Chen et al., 2020a, )
are followed, except for a reduced strength of resized cropping from 0.2 to
0.5. All training involves a standardly used cosine decay learning rate
schedule, starting at 0.03 and decaying to 0.001. When a 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 to 0.01 with a cosine schedule. We note linear evaluation
is used for supervised models as well, following the practice advocated by Liu
et al., (2021).

##  Appendix 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
$\Sigma B_{M}$ are bounded between $-1$ and $1$

$-1\leq\frac{a_{i}-c_{i}}{a_{i}+c_{i}}\leq 1,$ (18)

and $-1$ is achieved when $c_{i}\gg a_{i}$, and $1$ is achieved when $a_{i}\gg
c_{i}$.

When the augmentation is negligibly small, $\Sigma^{-1}B_{M}\approx M$, and
$\lambda_{i}\approx\bar{\lambda}=1$, the condition thus becomes

$\frac{2}{d_{M}}>0,$ (19)

which always holds. Thus, a sufficiently small augmentation will never cause
collapse. Next, when we apply very strong augmentation to the $j$-th subspace
and zero augmentation to the others, the condition for the non-augmented
spaces becomes

$1+\frac{2}{d_{M}}>\frac{d_{M}-2}{d_{M}},$ (20)

meaning that the collapse will not happen. For the $j$-th space, the condition
is

$-1+\frac{2}{d_{M}}>\frac{d_{M}-2}{d_{M}}(\Longleftrightarrow)\frac{4}{d_{M}}>2,$
(21)

which is only possible when $d_{M}=1$, namely, the strongly augmented space is
the only space that does not collapse. This is reasonable when the original
data is 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. In
different word, a strong enough augmentation will cause a collapse in the
corresponding subspace, as is the case without normalization.

It is also interesting to note that having $c_{i}\geq a_{i}$ is no longer
sufficient to cause a collapse. For example, let $c_{1}=0$ and $c_{j}=a_{j}$
for $j\neq 1$. The condition for $j\neq 1$ becomes

$\frac{2}{d_{M}}>\frac{1}{d_{M}},$ (22)

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_{i}<a_{i}$
to cause a collapse. Suppose we add a weak augmentation only to the first
subspace such that $a_{i}-c_{i}=\epsilon>0$, the condition for this dimension
to not to collapse is

$\frac{\epsilon}{a_{i}+c_{i}}+\frac{2}{d_{M}}>\frac{d_{M}-1+\epsilon}{d_{M}},$
(23)

which can be violated whenever
$\epsilon<\frac{(a_{i}+c_{i})(d_{M}-3)}{a_{i}+c_{i}+d_{m}}$. Namely, in some
cases, normalization can in fact facilitate collapse.

##  Appendix E Proofs

###  E.1 Proof of Proposition 1

Proof. The second term in Eq. (3) can be written as

$\displaystyle{\rm Var}[|W(x-\chi)|^{2}]$ $\displaystyle=\mathbb{E}\left[({\rm
Tr}[W(x-\chi)(x-\chi)^{T}W^{T}])^{2}\right]-\mathbb{E}\left[{\rm
Tr}[W(x-\chi)(x-\chi)^{T}W^{T}]\right]^{2}$ (24) $\displaystyle=[first\
term]-4{\rm Tr}[W(A_{0}+C)W^{T}]^{2}$ (25) $\displaystyle=[first\ term]-4{\rm
Tr}[W\Sigma W^{T}]^{2},$ (26)

where we have used the definition $\Sigma=A_{0}+C$. The first term is

$[first\ term]=\mathbb{E}\left[({\rm
Tr}[W(x-\chi)(x-\chi)^{T}W^{T}])^{2}\right]=4{\rm Tr}[W\Sigma W^{T}]^{2}+8{\rm
Tr}[W\Sigma W^{T}W\Sigma W^{T}].$ (27)

Combining the above expressions, we see that Eq. (3) can be written as

$\displaystyle L$ $\displaystyle=-{{\rm Tr}[WBW^{T}]}+\frac{1}{8}{\rm
Var}[|W(x-\chi)|^{2}]$ (28) $\displaystyle=-{\rm Tr}[WBW^{T}]+{\rm Tr}[W\Sigma
W^{T}W\Sigma W^{T}].$ (29)

This finishes the proof. $\square$

###  E.2 Proof of Theorem 1

Proof. All stationary points have a zero gradient:

$-2WB+4W\Sigma W^{T}W\Sigma=0.$ (30)

Multiplying by $W^{T}$ on the left and $B^{-1}$ on the right,

$W^{T}W=2W^{T}W\Sigma W^{T}W\Sigma B^{-1}$ (31)
$(\Longleftrightarrow)\quad\Sigma^{1/2}W^{T}W\Sigma^{1/2}=2\Sigma^{1/2}W^{T}W\Sigma
W^{T}W\Sigma B^{-1}\Sigma^{1/2}$ (32)

Defining $H:=\Sigma^{1/2}W^{T}W\Sigma^{1/2}$, we obtain

$H=2H^{2}\Sigma^{1/2}\Sigma B^{-1}\Sigma^{1/2},$ (33)
$(\Longleftrightarrow)\quad H(I-2H\Sigma^{1/2}B^{-1}\Sigma^{1/2})=0.$ (34)

Because both $H$ and $\Sigma^{1/2}\Sigma B^{-1}\Sigma^{1/2}$ are symmetric,
one can take the transpose of Eq. (33) to find that $H$ and
$\Sigma^{1/2}B^{-1}\Sigma^{1/2}$ commute with each, which implies that $H$ has
the same eigenvectors as $\Sigma^{1/2}B^{-1}\Sigma^{1/2}/2$.

Eq. (34) then implies that the eigenvalues of $H$ is either the inverse of
that of $\Sigma^{1/2}B^{-1}\Sigma^{1/2}$ or zero. This implies that any
stationary point of $H$ can be written in the form

$H=\frac{1}{2}UM\Lambda U^{T},$ (35)

where $U$ is a unitary matrix, $\Lambda$ is diagonal matrix containing the
eigenvalues of $\Sigma^{1/2}B^{-1}\Sigma^{1/2}$, and $M$ is an arbitrary
(masking) diagonal matrix containing only zero or one such that (1) $M_{ii}=0$
if $\Lambda_{ii}<0$ and (2) contain at most $d^{*}$ nonzero terms. This then
implies that the weight matrix $W$ satisfies

$W^{T}W=\frac{1}{2}\Sigma^{-1/2}UM\Lambda U^{T}\Sigma^{-1/2}.$ (36)

Lastly, when $\Sigma$ and $B$ commute, we can compactly write the result as

$W^{T}W=\frac{1}{2}\Sigma^{-1}B_{M}\Sigma^{-1},$ (37)

where $B_{M}$ denotes the matrix obtained by masking the eigenvalues of $B$
with $M$. This finishes the proof. $\square$

###  E.3 Proof of Proposition 2

Proof. For all stationary points, $W^{T}W$ commutes with $B$ and $\Sigma$,
which means that at these stationary points, one can simultaneously
diagonalize all the matrices and the loss function (3) can be written as

$L=-\sum_{i=1}^{d^{*}}\lambda_{i}b_{i}+\lambda_{i}^{2}s_{i}^{2}$ (38)

where $\lambda_{i}$, $b_{i}$, $s_{i}$ are the eigenvalues of $W^{T}W$, $B$,
and $\Sigma$ respectively.

We can thus consider each $i$ separately. When $b_{i}>0$, $\lambda_{i}=0$
cannot be a local minimum because the local Hessian is $-b_{i}<0$. When
$b_{i}\leq 0$, the only stationary point is $\lambda_{i}=0$. This sum covers
at most $d^{*}$ summands, and so, at the local minima, $\lambda_{i}\neq$ if
and only if $b_{i}>0$, and so the number of non-zero eigenvalues of $W^{T}W$
is $\min(m,d^{*})$. $\square$

###  E.4 Proof of Proposition 3

Proof. The regularization can be written as

$\displaystyle R$ $\displaystyle=[(\mathbb{E}_{x}||Wx||^{2}-c)^{2}]$ (39)
$\displaystyle={\rm Tr}[W\Sigma W^{T}]^{2}-2c{\rm Tr}[W\Sigma W^{T}]+c^{2}.$
(40)

By Proposition 1, Eq. (10) reads

$\displaystyle L$ $\displaystyle=-{{\rm Tr}[WBW^{T}]}+{\rm Tr}[W\Sigma
W^{T}W\Sigma W^{T}]+\kappa({\rm Tr}[W\Sigma W^{T}]^{2}-2{\rm Tr}[W\Sigma
W^{T}]+1)$ (41) $\displaystyle=-{{\rm Tr}[W(B+2\kappa c\Sigma)W^{T}]}+{\rm
Tr}[W\Sigma W^{T}W\Sigma W^{T}]+\kappa\rho^{2}.$ (42)

The derivative of $\rho$ is

$\frac{d}{dW}\rho=4\rho W\Sigma.$ (43)

The zero-gradient gradient is thus

$-2W(B+2\kappa c\Sigma-2\kappa\rho\Sigma)+4W\Sigma W^{T}W\Sigma=0.$ (44)

We can define $B^{\prime}:=B+2\kappa c\Sigma-2\kappa\rho\Sigma$ to see that
this condition is the same as Eq. (30) in the proof of Theorem 1. The rest of
the proof thus follows from the arguments. We thus arrive at the theorem
statement:

$W^{T}W=\frac{1}{2}\Sigma^{-1}B^{\prime}_{M}\Sigma^{-1}.$ (45)

We are done. $\square$

###  E.5 Proof of Proposition 4

Proof. Recalling that $\rho={\rm Tr}[W\Sigma W^{T}]$, we multiply $\Sigma$
from the right to both sides of the solution in Proposition 3 and take trace:

$\displaystyle\frac{1}{2}{\rm Tr}[\Sigma^{-1}B^{\prime}_{M}]$
$\displaystyle=\frac{1}{2}{\rm
Tr}[\Sigma^{-1}(B_{M}+2\kappa(c-\rho)\Sigma_{M})]$ (46) $\displaystyle={\rm
Tr}[W^{T}W\Sigma]$ (47) $\displaystyle={\rm Tr}[W\Sigma W^{T}]=\rho.$ (48)

The first line further simplifies to

$\frac{1}{2}{\rm Tr}[\Sigma^{-1}B_{M}]+\kappa(c-\rho){\rm
Tr}[\Sigma^{-1}\Sigma_{M}]=\frac{1}{2}{\rm
Tr}[\Sigma^{-1}B_{M}]+\kappa(c-\rho)d_{M},$ (49)

where $d_{M}:={\rm Tr}[M]$ is the number of nonzero eigenvalues of
$B^{\prime}_{M}$.

This gives an equation of $\rho$ that solves to

$c-\rho=\frac{c-\frac{1}{2}{\rm Tr}[\Sigma^{-1}B_{M}]}{1+\kappa d_{M}}.$ (50)

This proves the proposition. $\square$