The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning
Authors: Zixin Wen
Yuanzhi Li
Venue: NeurIPS 2022 Accept (2022)
Keywords: Deep Learning Theory
Self-supervised Learning
Non-convex Optimization

Abstract: The surprising discovery of the BYOL method shows the negative samples can be replaced by adding the prediction head to the network.  It is mysterious why even when there exist trivial collapsed global optimal solutions, neural networks trained by (stochastic) gradient descent can still learn competitive representations. In this work, we present our empirical and theoretical discoveries on non-contrastive self-supervised learning. Empirically, we find that when the prediction head is initialized as an identity matrix with only its off-diagonal entries being trainable, the network can learn competitive representations even though the trivial optima still exist in the training objective. Theoretically, we characterized the substitution effect and acceleration effect of the trainable, but identity-initialized prediction head. The substitution effect happens when learning the stronger features in some neurons can substitute for learning these features in other neurons through updating the prediction head. And the acceleration effect happens when the substituted features can accelerate the learning of other weaker features to prevent them from being ignored. These two effects enable the neural networks to learn diversified features rather than focus only on learning the strongest features, which is likely the cause of the dimensional collapse phenomenon. To the best of our knowledge, this is also the first end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with a trainable prediction head and normalization. 

The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning
Zixin Wen
Yuanzhi Li
Machine Learning Department Carnegie Mellon University zixinw@andrew. cmu edu
Machine Learning Department Carnegie Mellon University yuanzhi Candrew cmu. edu
Abstract
The surprising discovery of the BYOL method shows the negative samples can be replaced by adding prediction head to the neural network. is mysterious why even when there exist trivial collapsed global optimal solutions, neural networks trained by (stochastic) gradient descent can still learn competitive representations. In this work, we present our empirical and theoretical discoveries on non-contrastive superviedd learning. Empirically, we find that when the prediction head is initialized as an identity matrix with only its off-diagonal entries being trainable, the network can learn competitive representations even though the trivial optima still exist in the training objective. Theoretically, we characterized the substitution effect and acceleration effect the trainable, but identity-initializz prediction head. The substitution effect happens when learning the stronger features in some neurons can substitute for learning these features in other neurons through updating the prediction head, And the acceleration effect happens when the substituted features can accelerate the learning of other weaker features to prevent them from being ignored. These two effects enable the neural networks to learn diversified features rather than focus only on learning the strongest features, which is likely the cause of the dimensional collapse phenomenon. To the best of our knowledge, this is also the first end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with trainable prediction head and normalization.
Introduction
Self-supervised learning is about learning representations of real-world vision language data without human supervision, and contrastive learning [62, 43, 41, 24, 20, 34] is one the most successful self-supervised learning approaches. It has been known that the behavior of contrastive learning depends critically on the minimization of the negative term, which corresponds to contrasting the representations of negative pairs, i.c., pairs of different data points. However, the surprising finding of the Bootstrap Your Own Latent (BYOL) method by Grill et al. [37] initiated the research of non-contrastive self-supervised learning, which refers contrastive learning methods without using the negative pairs. BYOL achieved ttate-of-the-an results in various computer vision benchmarks and there plenty of follow-up up works [39, 26, 21, 17, 33, 87, 44, 61] in this direction.
On a high level, in non-contrastive self-supervised learning, one wishes to learn a network ø such that ((x) aligns in direction with d(x'), where and x' are called the positive pair, generated by random augmentations from the same sample. Without the negative samples, collapsed global optima exist in the training objectives. The complete collapse is when o(-) is constant vector whose variance is zero. Another trivial solutions called dimensional collapse by [44] is when all the coordinates (.).) exactly aligned. Nevertheless, adding a trainable prediction head on top of (one branch of) d(x) magically avoids learning such solutions, even though the prediction head can possibly learn the identity mapping and render itself useless. A more formal introduction will be given in Section 2.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Since the proposition of BYOL, there have been lots of empirical studies trying to understand non- contrastive learning. The SimSiam method by Chen and He [26] shows the exponential moving average (EMA) is not necessary for avoiding collapsed solutions while stop-gradient is necessary. [68] empirically disproved using batch normalization (BN) is the reason why BYOL can avoid collapse. [21, 88) further explored other similar approaches. If one wishes to work without both asymmetry and the negative pairs, one must add extra diversity-enforcing structures as in Barlow Twins [87] or [33, 44, 17]. Although these previous papers provided some empirical insights, in theory, the question of how the prediction head helps in learning those diverse features is still unanswered.
Despite the great empirical progress, there is very little theoretical progress towards explaining them. Most of existing theories focus on contrastive learning, especially from the statistical learning perspective [79, 81, 14, 80, 40, 82, 13, 15, 47, 45, 59]. However, due to the existence of trivial collapsed global optimal solutions (even with the prediction head) of the non-contrastive methods, to the best of our knowledge, there is no well-established statistical framework for those methods yer. To explain the non-contrastive learning, it is inevitable to study how the solutions are chosen during the optimization. Therefore, our research questions are:
Our theoretical questions: the role of prediction head
Why do most non- con-contra self supe-vvuee methods learn collapsed solutions when the so-called prediction head is absent in the network architecture? How does the trainable prediction head help optimizing the neural network to learn more diversified representations in non-contrastive self-supervised learning?
Due to the existence of trivial collapsed optimal solutions of the non- non-contrastive learning objective, we need to understand the implicit bias in optimization posed by the prediction head. However, to the best of our knowledge, all of the previous implicit biases theories focus only on the supervised learning tasks, and thus cannot be applied to our question. On high level, the results in this paper are summarized as follows:
Our empirical contributions. In non-contrastive self- superviued learning, we obtain the following experimental results:
We discover empirically that even when the prediction head is linear and initialized as an identity matrix with only off-diagonal entries being trainable, the performance of learned representation is comparable to using the usual non-linear two-layer MLP or randomly initialized (trainable) linear prediction head. See Figure 1.
We empirically verified that even when the prediction head is an identity-initialized matrix (with fixed diagonal entries), its off-diagonal entries display a rise-and-fall pattern, and it does not always converge to symmetric matrix. See Figure 3.
Our theoretical contributions. We based our theory on very simple setting, where the data consist of two features: the strong feature and the weak feature. Intuitively, the strong features in a dataset are the ones that show up more frequently or with large magnitude, and weak features as those that show up rarely or with small magnitude. We consider learning with two-layer non-linear neural network with output normalization using (stochastic) gradient descent. Under this setting:
We prove that without a prediction head, even with BN on the output to avoid complete collapse, the networks will still converge dimensional collapsed solutions, which provides a theoretical explanation to the dimensional collapse phenomenon observed in [44].
We prove that the trainable prediction head, combined with suitable output normalization and stop-gradient operation, can learn diversified features to avoid the dimensional collapse problem. We characterize two effects leveraged by the prediction head: the substitution effect and the acceleration effect, as intuitively described below:
The effects of the trainable prediction head
In our setting, we prove that the trainable prediction head can help to learn diversed features by leveraging two effects: the substitution effect and the acceleration effectt The substitution effect happens when by learning the prediction head, the learned stronger features in some neurons can substitute for learning the same features in other neurons. The acceleration effect happens when the substituted features from the prediction head further accelerate learning the weaker features in those substituted neurons.
Besides the above effects, we also explain in our setting, how the two common components in non-contrastive learning: stop-gradient operation and output normalization, can assist the prediction
2

---

DIA6.10
S7L-20
05-
city dil Biew pooditent
sirtie an frou coll ES,
head,
lrok,
candemwin:: mear Lineme
biy val wor and book, and
82
aborddy
12907
Eouch
(a) CIFAR-IO Accurayy
(b) CIFAR- 100 Accuracy
(c) STL-10 Accuracy
57-30
D-lyer PINd THIS LOFOTY EMA
2rageer WAR LIFaY, EMA haod
2-laye EMA
1.0
Seriiiy Braar bacd
(ie Biy VUE WALL LICRTS
Linam,
Sontiiiy hoad Cheam,,
AMA
15
Lowers LMS
1.0
0.5
00
33
25
30
15
140
125
150
175
303
25
50
139
125
150
375
203
02007
(d) CIFAR-10 10 Loss
(c) CIFAR-100 Loss
(1) STL-10 Loss
Figure Performances using different prediction heads. Here CIFAR-10, CIFAR- 100 and STL-10, identity-initialized linear prediction head can achieve good accuracies comparable to commonly used two-layer non-lincar MLP andooly-initializeed linear head. All the prediction hcads are trainable, while for identity-init prediction head only the off-diagonals trainable. Here BN L2norm represents the output normalization, and EMA represents using exponential moving average to update the target network as in BYOL [39].
head in creating those effects during training, which will be further discussed in Section 5.3. There are already some theoretical papers [78, 83, 63) that try to address similar questions. Our results provide a completely different perspective compared to them: We explain why training the prediction head can encourage the network to learn diversified features and avoid dimensional collapses, even when the trivial collapsed optima still exist in the training objective, which is not covered by the prior works, as shall be discussed below.
1.1 Comparison to Similar Studies
In this section, we will clarify the differences between our results and some similar studies. We point that all the claims below are derived only in our theoretical setting and are partially verified in experiments over datasets such as CIFAR- 10, CIFAR-100, and STL-10.
Can eigenspace alignment explain the effects of training the prediction head? The paper [78] presented theoretical statement that (symmetric) linear prediction head will converge to matrix that commutes with the covariance matrix of linear representations the end of training, and they provided experiments to support their theory. However, our theory suggests that the intermediate stage of training the prediction head matters more to the feature learning of the encoder network than the convergence. Indeed, as shown in Figure 3, in many cases, the trainable projection head will converge back to identity after training, which commutes with any covariance matrix. Moreover, the experiments in Figure 3a shows the training trajectory of the prediction head displays a clear two-stage separation, which demonstrates that the convergence result (e.g., g., the eigenspace alignment result in (78)) is not sufficient to understand the trainable prediction head.
Can the symmetric prediction head explain the trainable prediction head? the paper [78], experiments over the STL-10 dataset showed that the linear prediction head converges to symmetric matrix during training. And the follow-up paper [83] established a theory under the symmetric prediction head (which is not trained but manually set at each iteration). Specifically, under their linear network setting, where W is the weight matrix of the base encoder, they manually set the prediction head W, at iteration / to be w <- [( (W(0)T and the outputs of both online and target network are not normalized. Under this manual update rule of the prediction head, they proved a subspace learning result over spherical gaussian data. Nevertheless, our experiments in
3

---

0.39
OPAAIE CHAA 183
OPARID CINA LLG
074-16
0.0
0P48-10 12
0.2
0.0
0.00
000
0.0
200
10050
(a) Average off-diag entries
(b) F-norm of off-diag matrix
(c) Maximum of off-diag entries
Figure 2: Trajectories of the identity-iiitialzzed prediction head with a (min, max) confidence band, average over 3 runs. In all three datasets, we observe consistent rise and fall trajectory pattern.
Figure and Figure 3b show that even if we initialize the prediction head using a symmetric matrix (identity), the trainable prediction head can very asymmetric the early training stage when the encoder network learn most of its features. Actually, in the presence of feature imbalance (e.g.. Ex, F has huge eigen-gap), the symmetric prediction head is also likely to learn rank-one matrix where W focus learning the largest eigenvector of Ezy 21
The role of stop-gradient and output- normalitoooiiiiiionnn. It discussed in the theory of Tian et al. [78] that without the stop-gradient, the linear network will learn the zero (constant) solution. [83] also incorporated the stop-gradient into their theory, but did not explain why it is necessary for their setting. As a comparison, we proved in our setting that the stop-graud noo-gradient and output-normalization together can turn the features substituted via the prediction head into factor the gradient of the slower learning neurons, thereby creating the acceleration effect. In contrast, analyses in [78, 83] did not incorporate the output normalization, even though their experiments have used certain forms of normalizations. To the best of our knowledge, our paper is the first to explain the effects of output-normalizaiion in optimizing nonlinear neural networks in seff-super learning.
Preliminaries on Non-contrastive Learning
In this section, we formally define what is non-contrastive self-supervised learning. To do this, first introduce contrastive learning following [24, 85] as background. We use [N] as shorthand for the index set NN.
Background on contrastive learning. Letting (w(-) be the neural networks, contrastive learning aims to learn good representations dw via contrasting representations of similar data samples to those of dissimilar ones. Usually we are given a batch of data points (X; [X))E[N]. and we construct for
each = [N] a positive pair (xx2) x x(2)) by applying random data augmentations to X,, and collect xx²)) for $ j [N]. Now given 2, ow (x("),, 2f ow (x ( i [N].
negative pairs (xx"",, xx("",x(²)
we train the network ow minimize the contrastive loss:
Loontrastive (dw)
Σ KE[M]
sim (2i, 20/T log
Σ JE[N
exp (sim 2f)/T)
(2.1)
positive term
negutive term
where sim(...) is the similarity metric, often defined as the cosine similarity, and 7 is the so-called temperature flyee-paraaettr Intuitively, minimizing the contrastive loss be roughly viewed as trying to classify the representation as z instead of z/..i i * i. It is common belief that in order for the network Φw be able to "distinguish" data points X, from X,, + merely minimizing the positive term of contrastive loss is not sufficient.
Non-contrastive self-supervised learning. We choose the SimSian method [26] primary framework, whose differerence with BYOL is EMA component that is proven inessential [26]. Following the same notations as above, except that 2f StopGrad|ow (x.(2) (xx(2)) is detached from

---

CIFAR 10, off diog
CIF4R STL.10
CIFA8-10, off pillc--oooddaa
1.0
off.dieg
Robonnus
CIFAR-100, off diag Slag-off_dag off_dag"T
Rabor no
STL-10, off_ding
08
off dogg
0.5
25
30
75
100
125
150
175
200
Spoth
(a) ||off-diag( E [] and || ( 1|
(b) | (E / (10
Figure 3: Trajectories of the identity-initialle initialized prediction head. off off-diag EE) obtained by setting the diagonal of E' be zero. (a), we discover
Frobenius norm of our identit--nittilizeed initialized prediction head's off-diagonal matrix clearly display a two stage separation. more precisely, rise and fall pattern: In (b), The off-diagonal matrix of the prediction head is not symmetric in CIFAR- 10 and CIFAR-100.
gradient computation, the loss objective become: (the symmetric network version)
Lisinia (ow) = N EiE(N) sim( 2()
(2.2)
which is just the positive term in contrastive loss (2.1) (not divided by r). Clearly there exist plenty trivial global optimal solutions for this objective. For example, the complete collapse refers to when (-(-) learns some constant vector. Another solution called dimensional collapse [44] is when all the coordinates (ow (-)]; has correlation +1. The dimensional collapsed solution can minimize the objective (2.2) even when (-(-) is BN-normalized avoid learning constant vector [44, 88).
However, by adding a trainable prediction head on top of z;, the training miraculously succeeds and outputs state-of-the-art feature extractor. Let g(-) be a shallow ferd-forward network (often one or two-layer, or even simply linear), we train 9 and øw simultaneously on the following objective:
Lsinsian Low g) ΣiE(N) - -sim(((2))))))) (g(2),
(2.3)
where z/ is still detached from gradient computation. The g(z:)( = (X!")) and the detached part 2 = StopGrad (x(²)) are often called the online network and the target network respectively following [39], known as two branches of contrasste learning. Note that the trainable prediction
head can represent identity function, so the objective (2.3) still has the collapsed optima.
Problem Sctup
In this section, we present the setting of our theoretical results. We first define the data distribution.
Notations. We use 02 θ notations to hide universal constants with respect to d and Ō, Π2 θ notations hide polynomial factors of log d. We denote o(1) if the notations poly(d), polylog(d) (d to represent large constant degree polynomials of d or log d. We use (H, Σ) to denote standard normal distribution with mean μ and covariance matrix Σ. We use the bracket (-,-) to denote the inner product and ||- ||2 the 4--mr in Euclidean space. And for subspace V C RRd, we denote V- as its orthogonal complement. We use I to denote the indicator function of event B. We Tm denote the m m identity matrix.
when d + co. We use
Following the standard structure of image datasets, we consider data divided into patches, where each patch can contain either features noises.
Definition 3.1 (data distribution and features). Let X D be X (XI...,, Xp) E dd where each X; E Rd is patch. We assume that there are two feature vectors such that [uell2 = l 1,2 and orthogonal to each other. To generate sample we uniformly sampled i E [] and generate for each p E [P]:
Xp 2X))er + 1 1p
Exwole 00 Vp e [P]
5

---

X
X()
x(2
f
11
22
Random Data Augmentations
a2 124
s
a, 24
fs
a, 12
tu
6o
a;
a
Feature Patch
6p
Noise Patch
Masked Patch
Figure 4: Illustration of the data distribution and data augmentations. Each data is equipped with feature, either in or U2. and contains a lot of noise patches. After the data augmentations, the positive pair (X X (x)) is constructed by randomly masking out half of non-owt non-overlapping patches for each positive sample. The reason for constructing positive pair with non-overlapping patches because of the strong noise assumption made Assumption 3.2 and the feature decoupling principle in [85].
where 2p (X) is the latent vector of X, Ep is the noise vector of patch p [P] whose assumption will be given in Assumption 3.2. We denote S(X) {p (X) # 0} C [P] as the set of feature patches, where (x) Znd (X) (0, #az), Vp, E [P], where a will be picked afterwards. We assume P = polylog(1)) (d), S(X) = Po Olog d) for every X. A figurative illustration is given in Figure 4.
Strong and weak features. We pick 01 2polyloglog(d) and a₂ = a1 ppolylog(d). Hence V1 the strong feature and U2 is the weak feature, and want the learner network to learn both u1, 02 (but by different neurons) as their learning goal. This is a simplification think of the strong features in dataset the ones that show up more frequently with larger magnitude, and weak features as those that show up rarely or with smaller magnitude.
the real scenario. Intuitively,
Assumption 3.2 (noise). Denoting V = span(us U1, v2), we assume Ep E V independent for each p
[ \ (()), where X (Xp)re(P) ~ D, and:
(a) For any unit vector u E V Bl(sp u}] = O, and (S u) σ for some o = =(1);
(b) It holds for some E [0, zumr it holds [E[(u,, &p) (u),,p))) and [B[( 5p) (uz; Ep)i
@ for any two vectors u1, Ug E V that are orthogonal each other.
Remark 3.3. A simple example of our noise &p is the spherical Gaussian noise in V Assumption 3.2h ensures that the prediction head cannot be used to cancel the noise correlation between different neurons. We point out that the features in data that we do not intend to compare our algorithm with any clustering method this setting since our goal is to study how the prediction head helps in learning the features.
be learned via clustering, but we emphasize
3.1 Learner Network
Following the SimSiam framework, the online and target network share the same encoder network in our setting, as explained in Section 2. We consider the base encoder network f as a simple convolutional neural network: Let W (wi..... Wm) E Rdxm be the weight matrix, where w; E Rdd the encoder network f is defined by
(x) EpEly o(ww3, X,)),
Vj E [m]
Here use the cubic activation function (2) 23, as polynomial activations standard in literatures of deep learning theory [9, 35, 50, 2, 52, 23] and also has comparable performance in practice [2]. The (identity initialized) prediction head is defined as matrix É [E ( with E,, i E [m], where only the the off-diagonals Ex of are trainable parameters. The online

---

network F defined by: given j E [m]. we let FX(X) = fx(X) + Enzs Ex, and
Fxx)== = BN (F)()) = BN E (o((w,,X,)) + Mn Ej o (wwr, (wrr)
where the batch normalization BN here is defined follows: Given batch of inputs (Zi)ve(N),
N >iEENN 2F
BN (z))
(3.1)
E EIENN o - (N EIELNI 21
And the define the target network G as T,(X) == BN (G, BN (fx(X)) vij E [m].
3.2
Training Algorithm
Data augmentation. We use a very simple data augmentation: for each data X = (Xp)pE(Pl. we randomly and uniformly sample half of the patches P - [P] to generate the positive pair:
(
X (Xp 11 pPP)
(3.2)
Our data augmentation is similar to the common random cropping used in contrastive learning [22, 76]. is also analogous to the data augmentations studied in theoretical literatures [85, 47, 58).
Non-contrastive loss function. Now we define the loss function as follows: we sample N
data points (XI) X lid. DD and apply data augmentation (3.2) to obtain S
{X(11)) X (()))) HEIN). Now we define
Ls(W E) =
Σ KE[N]
F(x StopGrad 2
(3.3)
N
where the StopGrad operator detach gradient computation of the target network G(-). This form of objective (3.3) is first defined in [37] and is equivalent to (2.3) in Chen and He [26] when and G share the same encoder network f(() and their outputs are normalized.
Intuitions of the data augmentation and collapse. In Definition 3.1, the features vI, 02 appear in multiple patches, but the noises are independent across different patches (see Figure 4). As our data augmentation produces positive pairs with non-overlapping patches, learning to emphasize noises cannot align the representations of the positive pair, but learning either one of the features <(X) = Ep x()) or o(X) =xx==E, E, ((12, Xp)) is sufficient. We consider learning the same feature U; in all the neurons f in the encoder network f as the dimensional collapsed solution.
Initialization and hyper- paremeparaeters. At 0, we initialize W and E as w! ~ (O, 1) and
e Im and we only train the off-diagonal entries of E(9), For the simplicity of analysis, we let m = 2, which suffices to illustrate our main message. For the learning rates, we let 17 E (0, polytay)
be sufficiently small and n [n/a] O(1) / polylog(d)). which is smaller than 172
Optimization algorithm Given the data augmentation and the loss function, we perform (stochastic) gradient descent on the training objective (3.3) as follows: at each iteration f = 0, ..., T - we sample a new batch of augmented data S, {.X (1,i2)) he[[ and update
(((+1) (tt) Vs E(!),
(1 Ls, E(U), Vi EE! N VE, Ls, 4s,
we do not train the prediction head, we just simply keep (x 1m-
We use batch normalization output- nommaliatoniiiiioonnn method, rather than for the supposed implicit negative term effects as disproved in Richemond et al. [68].
We conjecture that by modifying certain assumptions for the noise (especially by allowing the noise to span the feature subspace V). prove similar result for the case Te 1).
7

---

Statements of Main Results
In this section, we shall present our main theoretical results on the mechanism of learning the prediction head in non-contrastive learning. To measure the correlation between neurons, we introduce the following notion: letting Var(*(X)) == Ex-p [(v(() - Ele(())) be the variance any function v of X D. we denote the correlation Corr( (X), d (X)) of any two function 1,, v over D as
Corr (xrr(((x),((x)) (X)) = [((((X) - -(
V VVar()(X)) (X)))
Now we present the main theorem of training with a prediction head, and set m = 2.
Theorem 4.1 (learning with prediction head and BN, see Theorem F.2). For every d ler N 2 poly(d), E (0, polytay be sufficiently small, and Te E ( polylog(ar) Then with probability
1 - (1), after training for T poly(d))/ many iterations, we shall have for some E E [2]:
w
(T) BLUE E11 W2
(T)
82
with 18.1. @()), 1 S Ō
Furthermore, the objective converges: ES-DN (Ls(W(), < OPT + poly(d) S 0 1ng
Where OPT stands for the global optimum'.
Theorem 4.1 clearly shows the network learn all the desired features, even under huge imbalance between 11 and U2. This leads to the following corollary.
Corollary 4.2. Under the same hyper-paramee parameter Theorem 4.1, with probability 1- o11,), after training for T = poly(d)/n many iterations, then the encoder f avoids dimensional collapse:
[Corr([(X)), <
Vd
In contrast, learning without the prediction head will create strong correlations between any two neurons. To emphasize that this problem cannot be alleviated by having more neurons, let the number of neurons m be any positive integer in the following theorem.
Theorem 4.3 (learning without prediction head but with BN, see Theorem G.1), Let 2 poly(d), o(1) and the number of neurons m S ol21/a2). Suppose we freeze E(t) = Ims then with probability o(1), after training for T poly(d)/r many iterations, we shall have:
( -8,11+ E with |8,| =0]), ||5y|l2 S Ō Vd w
for all j [m]]
Furthermore, the objective converges: ES-DN [Ls (W [T)) < OPT polytd] S O Joga This means the collapsed solution also reaches the globul minimum of the objective. Again OPT stands for the global optimum.
Note that since we have used BN as our output normalization, the learner is immune to complete collapse and must have certain variance in the outputs. Immediately, have corollary.
Corollary 4.4. Under the same training with E()) for T poly(d)/1 many iterations, shall have dimensional collapse:
in Theorem 4.3, with probubility 1 - -o(1), after
(Corr()((X), 2 00 Va
for all i, E [m]]
In the following section, we shall give some intuitions by digging through the training process and separately discuss the four phases of the training process.
2 Under our setting described Scction 2, the global minimum of our objective (3.3) in population is -PT=2- El(((x) PI (S(X)
\ PI] e log d
R| PP

---

5 The Four Phases of the Learning Process
We divide the complete training process into four phases: phase I for learning the stronger feature, phase II for the substitution effect, phase III for the acceleration effect, and the end phase for convergence. The first three phases explain how the prediction head can help learn the base encoder, and the last phase explains why the off off-diagonal entries often shrink later in training.
5.1 Phase I: Learning the Stronger Feature
At the beginning of training, the stronger feature Vi enjoys a much larger gradient as opposed to the weaker feature 02, naturally, ii will be learned first. Without loss of generality, let us assume at initialization, the neuron fi has larger Vi between f E [2]. then we can show:
Lemma 5.1 (learning the stronger feature, see Lemma C.13). After some 2 = ( 17, the feature vn in neuron will be learn to (w(") v1)
2(11, while all other features (w]) ( ,v) = (1)
for (j,() + (1,1). And the prediction head [| - T2 v d-1((1) is still close to the initialization. In this phase, the prediction head has not come into play. The substitution effect can only happen after the feature 21 in neuron f1 is learned to certain degree, and neuron f2 remains largely unlearned.
5.2 Phase II: The Substitution Effect
To illustrate the substitution effect, let us keep assuming that neuron fi has already learned some significant amount of the strong feature U1, say u1 31 01 + residual with |3g [x( residual)) )) then we have: (recall f. [2] are the neurons of the encoder network)
Lemma 5.2 (substitution effect, formal statement Lemma D.8). After ((0)0 v1)i (2)1),
shall have | increasing until [E(4) (x(()) >> (/2(X())) when X is equipped with the strong feature vi, for T22 Ti = (Ti) iterations.
Intuition of the substitution effect. After the stronger feature is learned in neuron f1, the optimal way to align two positive representations (x) and (X(()) is not learning features weight W2, but use the prediction head to "substitute" the features in fi into Fz. This is how the substitution effect happens when trained with a prediction head.
5.3 Phase III: The Acceleration Effect
After the substitution of 0j in F2, our concern is, w22 (1) will learn 02 and only U2 eventually, according to the acceleration effect in the following lemma.
Lemma 5.3 (acceleration effect, formal statement see Lemma E.8). After Egy is learned in Lemma
5.2, learning 12 in w/ will be much faster than U1, until |! - B202l = o1) for some 33 (())
The acceleration effect is caused by the interactions between the prediction head, the stop gradient operation, and the normalization method (which in this case is the batch normalization).
What is the role of the stop-gradient? Thanks to the stop-gradient operation, when we compute the gradient- -V F2(((1)) StopGrad[G2(X112) to learn fz. this negative gradient will only try to maximize f2(X(()) f ())()),,, rather than to maximize F2 (x(()) (x(())). This is because the stop-gradicn is on G not on F: while F2 has large component of v1 borrowed from fi using E, G2 does not have this component. So the gradient of F2 is to align with the features in G2 that does not contain many U1, while the gradient of G2 is to aligned with F2 that contains lot of en.
What is the role of the output normalization? Again due to the StopGrad operation, the gradient of F2 is taken with respect to the ratio f2(x(")) ht/l///ar///a/ (x As gradient descent tries to maximize this ratio, a direct computation gives
(x(
Var (P)((1))))) ExE(2) (LE5) (wf)) Us 03- 3 Varl/2(X()1) (V,, f2(((()))) veLUe
which borrow the substituted feature ii from fi(-) to adjust the gradient of 02 in fa(-), via the prediction head E.)) Without the output normalization, the learning of 01 will dominate that of 02 even when we train the prediction head.

---

5.4 The End Phase: Convergence
As the weak features are learned, we have already obtained a good encoder network f(-) as shown in Theorem 4.1. The rest of Theorem F.2 also contains the following result: Proposition 5.4 (convergence of the prediction head, see Theorem F.2c). After some poly(d)/n iterations, we shall have E I2]|F poly(d))
While we admit that only some of our real-world experiments in Figure 3 show the convergence to zero for the off afi-oiagonal entries of the prediction head, most of the experiments do display a rise and fall trajectory pattern of off-diagonal entries consistently, which supports our theory to some degree.
6 Additional Related Work
Self-supervised learning Self- supervueee learning has created huge success in natural language processing [30, 86, 18] and established the pretrain- finetune paradigm for deep learning. In vision, contrastive learning [75, 41, 24, 20, 27, 28, 34, 64, 33J became dominant in many downstream tasks recently. Another approach is the generative learning [65, 16, 42), which also gives promising results. Applications such [64, 66) also illustrate the power of contrastive learning in multiple domains.
Theory of self-supervised learning The theoretical side of self-supervised learning developed quickly due to the success of contrastive learning. Since [12], plenty of papers have studied the contrastive learning. [25, 69] discussed many interesting phenomena associated with the negative term. Saunshi et al. [71] provided evidence that function class agnostic analyses is vacuous. [85] took feature learning view, and inspired our analysis in the non-contrastive setting. For generative learning. [51, 74J provides downstream performance guarantees. [70, 84] studied the natural language tasks. [58] gave recovery guarantee for tensors under hidden Markov models. [4] provided an optimization guarantee for GANs trained by stochastic gradient descent ascent.
Feature learning theory of deep learning Our theoretical results are also inspired by the recent progress of the feature learning theory of neural networks [55, 56, 5, 3, 49, 90, 46). [55] initiate the study of the speed difference in learning different types of features. [1] developed theory for learning two-layer neural networks beyond the neural tangent kernel (NTK) [7, 8, 6, 32, 11]. [5, 3, 2] further studied how features are learned in different deep learning tasks. Before this recent progress, [77, 89, 19, 72, 31, 53, 54] also studied how shallow neural networks can learn certain simple data distributions, but all of them focus on the supervised learning. There are also plenty of studies [73, 38, 10, 60, 48, 67, 29] on the implicit bias of optimization in deep learning, but none of their techniques are designed for analyzing self-supervised learning.
7 Conclusion
In this paper, we showed how the prediction head can ensure the neural network learns all the features con-contrastive learning through theoretical investigation. Our key contribution that we proved the prediction head can leverage two effects called substitution effect and acceleration effect during the training process. We also gave an explanation for the dimensional collapse phenomenon. We believe our theory, although based on very simple setup, can provide some insights into the inner workings of non-contrastive self-supervised learning.
Acknowledgments and Disclosure of Funding
Funding in direct support of this work includes NSF Award 2145703.
References
[1] Zeyuan -hu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems, 32, 2019.
[2] Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413, 2020.
[3] Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint 20Xx::2012.09816, 2020.
10

---

[4] Zeyuan Allen Zhu and Yuanzhi Li. Forward super-resolution: How can gans learn hierarchical generative models for real-world distributions. arXiv preprint arXiv:2106.02611 2021.
[5] Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 62nd IEEE Annual Symposium on Foundations of Computer Science, FOCS 2021, Denver, CO, USA, February 7-10, 2022. pages 977-988. IEEE, 2021. URL
attps://dol: org/10..0000...00 S119/FFCS55997 20219997..02......
[6] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. In NeurlPS 2019 : Thirty-third Conference on Neural Information Processing Systems, pages 6158 (188-616,,40 2019.
[7] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via pver-parameterization. In ICML 2019 Thirty-sixth International Conference on Machine Learning, pages 242-252, 2019.
[8] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. In Advances in Neural Information Processing Systems, pages 6676-6688, 2019.
[9] Alexandr Andoni, Rina Panigrahy, Gregory Valiant, and Li Zhang. Learning polynomials with neural networks. In International conference on machine learning, pages 1908-1916, PMLR, 2014.
[10] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo, Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[11] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine- fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference Machine Learning, pages 322-332. PMLR, 2019.
[12] Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09222 2019.
[13] Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Dipendra Misra, Investigating the role of negatives in contrastive representation learning. arXiv preprint arXiv:2106.. 09943,99 2021.
[14] Yamini Bansal, Gal Kaplun, and Boaz Barak. For self-supervised learning, rationality implies generalization, provably. arXiv preprint arXiv:20 2020.
[15] Han Bao, Yoshihiro Nagano, and Kento Nozawa. Sharp learning bounds for contrastive unsupervised representation learning. arXiv preprint 1r00022100.02501, 2021.
[16] Hangbo Bao, Li Dong, and Furu Wei. Beit: Bert pre-training of image transformers. arXiv preprint arXiv:2106.008: 082544 2021.
[17] Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variane-invarincc--ovrria regulal ization for self-supervised learning. arXiv preprint arXiv:2105.04996, 2021.
[18] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[19] Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for convnet with gaussian inputs. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 6654 JMLR. org, 2017.
[20] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems, 33:9912-40 12-9924, 2020.
[21] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9650-9660, 2021.
11

---

[22] Shuxiao Chen, Edgar Dobriban, and Jane H. Lee. A group-theoretic framework for data augmentation. Journal of Machine Learning Research, 21(245):1-71, 2020.
[23] Sitan Chen, Jerry Li, Yuanzhi Li, and Anru R Zhang, Learning polynomial transformations. arXiv preprini urXiv:2204.042000 2022.
[24] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ICMI. 2020: 37th International Conference on Machine Learning, volume 1, pages 1597-1607, 2020.
[25] Ting Chen, Calvin Luo, and Lala Li. Intriguing properties of contrastive losses. Advances in Neural Information Processing Systems, 34, 2021.
[26] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVE Conference on Computer Vision and Patlern Recognition, pages 15750-15758, 2021.
[27] Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint XXix:2033304297 2020.
[28] Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9640-9649, 2021.
[29] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference Learning Theory, pages 1305-1338. PMLR, 2020.
[30] Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-LLT, pages 4171 4186, 2019.
[31] Simon S Du, Jason D Lee, and Yuandong Tian. When is convolutional filter easy to learn? In International Conference Learning Representations, 2018.
[32] Simon S. Du, Jason D. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In 36th International Conference Machine Learning. ICML 2019, pages 1675-1685, 2019.
[33] Aleksandr Ermolov, Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening for self-supervised representation learning. In International Conference on Machine Learning, pages 3015-3024. PMLR, 2021.
[34] Tianyu Gao, Xingcheng Yao, and Danqi Chen. Simese: Simple contrastive learning of sentence embeddings. arXiv preprint arXiv:2104.088 08820, 2021.
[35] Antoine Gautier, Quynh N Nguyen, and Matthias Hein. Globally optimal training of generalized polynomial neural networks with nonlinear spectral methods. Advances Neural Information Processing Systems, 29, 2016.
[36] Sven Gowal, Po-Sen Huang, Aaron van den Oord, Timothy Mann, and Pushmeet Kohli. Self- supervised adversarial robustmess for the low-label, high-data regime. In International Confer- ence on Learning Representations, 2020.
[37] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self- learning. Advances in Neural Information Processing Systems, 33:21271-21288 2020.
[38] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems, 31, 2018.
12

---

[39] Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, -aan-Bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap latent- predictive representations for multitask reinforcement learning. In International Conference on Machine Learning, pages 3875-3886. PMLR, 2020.
[40] Jeff Z. HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self- supervised deep learning with spectral contrastive loss. arXiv preprint arXiv:2106.04155 04156, 2021.
[41] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. 2020 JEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 9729-9738, 2020.
[42] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. arXiv preprint arXiv:2111.06:3 2021.
[43] Devon Hjelm, Alex Fedorov, Samuel Marchi-nn Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In International Conference on Learning Representations, 2019.
[44] Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao. On feature decorrelation in self-supervised learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9598-9608, 2021.
[45] Weiran Huang, Mingyang Yi, and Xuyang Zhao. Towards the generalization of contrastive supervupe learning. arXiv preprint arXiv:2111.00743, 2021.
[46] Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization deep learning, 2022. URL https:
[47] Wenlong Ji, Zhun Deng, Ryumei Nakada, James Zou, and Linjun Zhang. The power of contrast for feature learning: A theoretical analysis. arXiv preprint arXiv:21 2021.
[48] Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference Learning Theory, pages 1772-1798. PMLR, 2019.
[49] Stefani Karp, Ezra Winston, Yuanzhi Li, and Aarti Singh. Local signal adaptivity: Provable feature learning in neural networks beyond kernels. Advances in Neural Information Processing Systems, 34, 2021.
[50] Joe Kileel, Matthew Trager, and Joan Bruna. On the expressive power of deep polynomial neural networks. Advances neural information processing systems, 32, 2019.
[51] Jason D Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo Predicting what you already know helps: Provable self-supervised learning. Advances in Neural Information Processing Systems, 34, 2021.
[52] Yuanzhi Li and Zehao Dou. Making method of moments great again? ?-ow can gans learn distributions. arXiv preprint arXiv:2003.04033 2020.
[53] Yuanzhi Li and Yang Yuan. Convergence analysis of two-layer neural networks with relu activation. In Advances in neural information processing systems, pages 597-607, 2017.
[54] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularizaionnn in over- wer-parameterrize matrix sensing and neural networks with quadratic activations. In COLT 2018: 31 Annual Conference on Learning Theory, pages 2-47, 2018.
[55] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In NewrlPS 2019: Thirty-third Conference on Neural Information Processing Systems, pages 11674-11685, 2019.
[56) Yuanzhi Li, Tengyu Ma, and Hongyang R. Zhang. Learning over wverpparametrized two-layer relu neural networks beyond ntk. In COLT, pages 2613-2682, 2020.
13

---

[57] Bingbin Liu, Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. Analyzing and improv- ing the optimization landscape of noise-contrastive estimation. arXiv preprint arXiv:21 2021.
[58] Bingbin Liu, Daniel Hsu, Pradeep Ravikumar, and Andrej Risteski. Masked prediction tasks: parameter identifiability view. arXiv preprini urXiv:2202.09305. 2022.
[59] Zeping Luo, Cindy Weng, Shiyou Wu, Mo Zhou, and Rong Ge. One objective for all models- self-supervised learning for topic models. arXiv preprint arXiv:2203.03533 2022.
[60] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint
2019.
[61] Daisuke Nüzumi, Daiki Takeuchi, Yasunori Ohishi, Noboru Harada, and Kunio Kashino. Byol for audio: Self-supervised learning for general-purpose audio representation. In 2021 International Joint Conference on Neural Networks (LJCNN), pages 1-8. IEEE, 2021.
[62] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1807. 80.7077480, 2018.
[63] Ashwini Pokle, Jinjin Tian, Yuchen Li, and Andrej Risteski. Contrasting the landscape of contrastive and non-contrastive learning. arXiv preprint arXiv:2203. 15702, 2022.
[64] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et Learning transferable visual models from natural language supervision. In International Conference on Machine Learning. pages 8748 8763. PMLR, 2021.
[65] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8831. PMLR, 2021.
[66] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. OpenAl blog, 2022. URL https: / cdn.
openai com httpa//ppppes//daaa/ da.1-e-2.pd......
[67] Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. Advances in neural information processing systems, 33:21174-21187, 2020.
[68] Pierre H. Richemond, -ran-enstien Grill, Florent Altché, Corentin Tallec, Florian Strub, Andrew Brock, Samuel Smith, Soham De, Razvan Pascanu, Bilal Piot, and Michal Valko. BYOL works even without batch statistics. arXiv:2010. 10241 [es, stat], October 2020.
[69] Joshua Robinson, Li Sun, Ke Yu, Kayhan Batmanghelich, Stefanie Jegelka, and Suvrit Sra. Can contrastive learning avoid shortcut solutions? Advances in Neural Information Processing Systems, 34, 2021.
[70] Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora. A mathematical exploration of why language models help solve downstream tasks. arXiv preprint wrXiv:2010.03644 2020.
[71] Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. arXiv preprint 22 14037, 2022.
[72] Mahdi Soltanolkotabi. Learning relus via gradient descent. In Advances Neural Information Processing Systems, volume 30, pages 2007-2017, 2017.
[73] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasckar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1):2222-2877 2018
[74] Jiaye Teng, Weiran Huang, and Haowei He. Can pretext-based self-supervised learning be boosted by downstream data? theoretical analysis. arXiv preprint arXiv:2103.03566 2021.
14

---

[75] Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In European conference on computer vision, pages 776-794. Springer. 2020.
[76] Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning. In Advances in Neural Information Processing Systems, volume 33, 2020.
[77] Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. In Proceedings of the 34th International Conference on Machine Coariing-Volume 70, pages 3404-3413. JMLR. org. 2017.
[78] Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning, pages 10268-10278. PMLR, 2021.
[79] Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive estimation reveals topic posterior information to linear models. arXiv preprint arXiv:2003.02234, 2020.
[80] Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive learning, multi-viww redundancy, and linear models. In Algorithmic Learning Theory, pages 1179-1206. PMLR, 2021.
[81] Yao-Hung Hubert Tsai, Yue Wu, Ruslan Salakhutdinov, and Louis-Philippe Morency. Self- supervised learning from a multi-view perspective. arXiv preprini arXiv:2006.05557 2020.
[82] Julius Von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. Advances in Neural Information Processing Systems, 34, 2021.
[83] Xiang Wang, Xinlei Chen, Simon Du, and Yuandong Tian. Towards demystifying repre- sentation learning with non-contrastive self-supervision. arXiv preprint arXiv:21 10.04947, 2021.
[84] Colin Wei, Sang Michael Xie, and Tengyu Ma. Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning. Advances Neural Information Processing Systems, 34, 2021.
[85] Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning, pages 11112-11122. PMLR, 2021.
[86] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in neural information processing systems, pages 5754-5764, 2019.
[87] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self- supervised learning via redundancy reduction. In International Conference on Machine Learn- ing. pages 12310-12333 12320. PMLR, 2021.
[88] Chaoning Zhang, Kang Zhang, Chenshuang Zhang, Trung Pham, Chang D Yoo, and In So Kweon. How does simsiam avoid collapse without negative samples? unified understanding with self-supervised contrastive learning. arXiv preprint arXiv:2203. 16262, 2022.
[89] Kai Zhong, Zhao Song, Prateek Jain, Peter L Bartlett, and Inderjit S Dhillon. Recovery guarantees for one-hidden-layer neural networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 4140 JMLR. org, 2017.
[90] Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of adam in learning neural networks with proper regularization. arXiv preprint 11X. 2021.
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits. hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [No]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [No]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [No]
(b) Did you mention the license of the assets? [No]
(c) Did you include any new assets either the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "sXql4LcFh2", "reviewer": "Reviewer_hrov", "summary": "This paper aims to theoretically analyze why non-contrastive self-supervised learning methods (e.g., BYOL and SimSiam) can avoid trivial solutions even without negative pairs. To this end, the authors assume a simple setting: (i) data contains only two feature vectors, (ii) the backbone encoder architecture is a single-layer CNN with a non-linear activation, and (iii) the prediction head architecture is a linear layer. Under these assumptions, this paper proves that the trainable prediction head can learn CNN to avoid trivial solutions. Furthermore, this paper provides experimental results that support the theoretical results under more standard (i.e., real-world) settings.\n", "questions": "**Q1.** Why $|E_{2,1}|$ is increasing in Phase II?\n\nAt the ends of Phase I and II (Lemma C.4 and D.8, resp.), $B=(B_{ij})=(w_i^\\top v_j)$ does not change with $B_{1,1}=\\Theta(1)$ and $B_{ij}=\\Theta(1/\\sqrt{d})$ for $(i,j)\\neq(1,1)$, but $E_{2,1}$ increases from $\\tilde{O}(\\eta_E / \\eta d)$ to $\\Theta(\\sqrt{\\eta_E / \\eta})$. It means that $F_2$ chagnes but $G_2$ does not in Phase II. This does not make sense because our objective (3.3) minimizes the distance between $F$ and $G$.\n\nAlso, the intuition of the substitution effect (L276-L279) is also hard to understand. Why does borrowing the features in $f_1$ minimize the distance between $F_2$ and $G_2$?\n\n---\n\n**Q2.** After Phase II, what is $w_2$?\n\nThis question is closely related to Q1. If Lemma 5.2 holds, then $F_2\\approx E_{2,1} f_1$. Due to the objective (3.3), $G_2=f_2 \\approx F_2$ anytime. It means that $w_2$ should be aligned with $w_1$ which is also aligned with $v_1$ by Lemma 5.1. Namely, $w_2\\approx v_1$ at the end of Phase II. Is it right? If so, why does $w_2$ change from $v_1$ to $v_2$ in Phase III? I think $w_2$ will remain as $v_1$ after Phase II.\n\n---\n\n**Q3.** Why use Batch Normalization (BN) instead of l2 normalization?\n\nBN is an unnecessary component in non-contrastive SSL methods, as shown in [68]. Also, BYOL and SimSiam apply l2 normalization without BN at their output layers. However, $\\tilde{F}$ and $\\tilde{G}$ defined in L198-200 use BN, and the objective (3.3) does not use l2 normalization. This gap may undermine this paper's motivation.\n", "limitations": "This paper does not address the limitations and potential negative societal impact.\n\n**Limitations.** This paper uses strong assumptions on data distribution (i.e., only two feature vectors) and network architectures (i.e., 1-layer CNNs). Also, this paper does not focus on l2 normalization, but it is widely utilized in various SSL methods.\n\n**Negative societal impact.** This paper is a theory paper. I think there is no negative impact.\n", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I think the main **strength** of this paper is the theoretical analysis of non-contrastive SSL methods. I agree that why the methods are working is less studied, so such a theoretical study is important in the field of SSL. Furthermore, the empirical observations with linear prediction heads support the validity of the theoretical assumptions well.\n\nMy main concern with this paper is the presentation quality. In general, it is hard to understand statements and their intuitions. First of all, Definition 3.1 is hard to understand:\n- There is no Figure 3 about Definition 3.1 (e.g., see L173 and L213).\n- Confused notations: $\\mathcal{S}(X)$ and $S(X)$ (e.g., see L170, L172, L180).\n- Does all $X$ in $\\mathcal{D}$ share the same feature vectors $\\{v_\\ell\\}_{\\ell\\in[2]}$? If so, Definition 3.1 seems to assume that the data distribution $\\mathcal{D}$ has only two data points (with some noises). Is this reasonable?\n- I think the constraint in L171, i.e., \"all feature patches have the same direction of $v_\\ell$\", is also too strong.\n- In L171, it would be better to use $z_p(X)=z_{p'}(X)\\in\\{\\pm\\alpha_\\ell\\}$ rather than $\\{0,\\pm\\alpha_\\ell\\}$ because for any $p\\in\\mathcal{S}$, $z_p\\neq 0$ by definition.\n- You should describe what is $\\xi_p$ first. Until p5, it is hard to know whether $\\xi_p$ is a random variable, a constant vector, or something else.\n\nAnd, what are the intuitions of theorems in Sections 4 and 5? I can understand what the theorems state, but it is hard to know how they can be derived. For example, how does borrowing the features in $f_1$ by increasing $E_{2,1}$ minimize the distance between $F_2$ and $G_2$? I feel that the authors should allow more space for explanations about the substitution and acceleration effects since they are key components of this paper. I also feel it would be better to provide empirical results based on the theoretical setup described in Section 3 to understand the changes in Phase I ~ IV.\n\nTo sum up, I think this paper's contributions are meaningful, but the presentation is unsatisfactory, and some results are still hard to understand. Hence, I vote for borderline rejection now, but I'm willing to increase my score after a strong rebuttal.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "o2AVbdNzK9A", "reviewer": "Reviewer_X3Gx", "summary": "This paper provides original theorems to explain the convergence of self-supervised learning methods based on positive views only (i.e., BYOL-like). Under reasonable assumptions, the authors demonstrate that the predictor module prevents the appearance of dimensional collapse. Interestingly, the authors could mathematically and empirically highlight different training regimes. They also introduce the concept of weak vs. strong features to intuit the learning dynamics. \n\nDISCLAIMER: I have not checked the Appendix yet, for they are 40 pages long. I *may* have additional comments during the discussion as the Appendix requires quite some extra time to parse. Therefore, this review did not check the mathematical details. Besides, I cannot put high positive/negative scores for intellectual ethic reasons.\n", "questions": "I have many questions listed below:\n\nA) It is a recurrent complaint about the description of BYOL-like methods: BYOL does not optimize a loss nor have a global minimum. It is a dynamic system (or a game) that has equilibria; as mentioned in the original paper, similar to GAN, BYOL has no optimum too (because of the stop gradient). As far as I could tell from the proof, the authors work at the gradient level, i.e., training updates, and not loss minimization, so the proof still holds. I would thus recommend that the authors replace the mention of loss with gradient updates or equilibria (depending on the context); especially l49 and 248.\n\nB) As mentioned by the authors in Section 1.1., this paper does not study the convergence from an eigenvalue perspective. However, the concept of weak/strong features could be related to high/low eigenvalues. Do you have any intuition about it? From an empirical perspective, I would try to plot the eigenvalues (linear layer only), and see whether they may relate to the training phase you mention.\n\nC) I may be missing something... but I do not understand how dotted lines may match the magnitude of non-dotted lines in Figure 2. Indeed, I would expect $E^t - (E^t)^T$ to be around zero (or around 1). I do understand that the off-diagonal is non-symmetric, but it here means that there exist a few leading values where the transpose is null (which sounds weird). Can you explain to me why the difference is so tight?\n\nD)  In corollary 4.2, the authors demonstrate that dimensional collapses are avoided. However, do the authors have a similar results that show that without the predictor, the dimensional collapse will happen with high probability? Here, I am afraid of a potential confirmation bias.\n\nE) In the footnote, 1 page 7, the authors mentioned that the same proof might be feasible with $\\eta_E = \\eta$. On a personal note, I do not think that it is possible in the BYOL setting. Indeed, BYOL requires the predictor to be close to optimal. Therefore, $\\eta_E$ must be smaller than $\\eta$, cf Appendix I.1 of the original BYOL paper. Note that this remark may be wrong in the SimSiam setting because of the inner optimization loop. Would it be possible to discuss this point and potentially update the paper to link it to the optimal predictor observation? \n\nF) In the paper, the authors explored different normalizations, showing that a non-trainable BN may be enough to avoid the collapse (even if it only returns mediocre results). It is indeed quite interesting and pushes further the limit of non-negative SSL methods. Differently in Appendix F.6 - Figure 7 of BYOL, the authors also noticed that the predictor works without normalization if and only if it is combined with weight decay. In such a case, BYOL diverged, but it is then stabilized because of the WD. In the absence of WD, the representation just diverged to infinity. More generally, the encoder WD is one of the most crucial parameters in BYOL. From your analysis, can you explain/extend this empirical (and reproducible) phenomenon?\n\nG) It is quite an open remark. One of the strengths of a hypothesis is to be able to be predictive and thus... be contradicted. If you had infinite resources, what would be the experiment you would do (and the observation that you would get) that could refute your paper? Please, note that I am not trying to trap you here but to understand the limit cases better.\n\nRemarks:\n  - When mentioning BN, I would recommend mentioning that you are using non-trainable BN, and only the centering/normalizing components\n - The size/depth of the predictor may have a bigger/smaller impact given the task difficulty. Therefore, I am not surprised that you empirically observed little difference between method accuracy in Figure 1. Differently, did you gather some intuition regarding the predictor of inner/outer dimensions during your work? (i.e., some quantities that may depend on $|z|$ or generalization requirement)?  \n - While the mathematical notation is very well explained, it may take time to apprehend them for some people. The authors mentioned that they give an illustration in Figure 3 (Appendix). Yet, the figure barely ground the notation, and it is simply the classic non-negative SSL pipeline. I would strongly recommend that the authors create a new large figure to illustrate all the mathematical variables. In other words, the paper impact could be even more significant if they use sketches to explain the results.\n - Please add all your training hyperparameters in the Appendix for your experiments.\n\n\nNit:\n  - non-contrastive methods may be too much of an approximation for BYOL-like methods. Indeed, pretext task methods would fall under this umbrella. Yet, I do not have a better naming, and BYOL-like is awkward too.\n - Put the OPT definition in the main paper instead of the footnote. I looked for it for a few minutes!\n - sigma is used both for activations and std; please use another notation for activation\n - In equation 2.1, I would add j != i in the sum of the negative \n\nIn summary, I am quite positive about the paper. Even though I still have many questions, most of them only deal with opening directions. ", "limitations": "n/a", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper is impressive for multiple reasons. \nFirst, the theoretical contribution is partially resolving a recent open problem in self-supervised learning: why non-contrastive methods such as BYOL or SimSiam do not collapse while they have trivial poor minima. Furthermore, the authors even describe learning dynamics, which is quite remarkable and even unique to my knowledge. \nSecond, the authors did a fantastic job explaining their theoretical results, making it extraordinarily accessible for most self-supervised learning practitioners. The structure with a formal definition, the lemme, and the intuition combine rigor and pedagogy. \nThird, the authors complete their theoretical insights with experiments. While less impressive and may still be completed, they perfectly complete the papers.\nOverall, this paper and [1] now provide the basics to understand the impressive and surprising results of non-contrastive methods \n\nThe main paper's weakness also results from its strength: the level of granularity of the learning dynamics may be too precise and, therefore, more open to debate. For instance, the time complexity may overlap between the lemma (if I am not mistaken, they rely on different parameters). Therefore, the authors may have wanted to make the story a bit too much compelling: while they describe it as an iterative process, the training stage phenomenon may occur jointly. \n\n[1] Tian, Yuandong, Xinlei Chen, and Surya Ganguli. \"Understanding self-supervised learning dynamics without contrastive pairs.\" International Conference on Machine Learning. PMLR, 2021.  ", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "W2hXPXYfudB", "reviewer": "Reviewer_a9Hp", "summary": "The authors explore why BYOL-style self-supervised learners do not collapse to trivial solutions from a theoretical perspective, focusing on the prediction head component of the architecture. They use a simplified setting with a strong feature and a weak feature that compete. They define various phases of training and describe what their theoretical results suggest should happen in those phases. Some empirical results are presented to support this. ", "questions": "What limitations does using the strong/weak feature setup introduce?\nCould you elaborate more on how batch normalization interacts with the stop-gradient to contribute to the acceleration effect?\nAs someone outside of deep learning theory, I would like to know more about how the assumptions posed here differ from those in other works, as well as some justification of these difference.", "limitations": "See questions. I am mostly interested in the limitations posed by the assumptions utilized in this work. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- The authors consider an important problem in a way that is simplified but feels justified. \n- Simple steps are provided to qualitatively understand what the theory reveals about the training process. \n- The previously-introduced dimensional collapse problem and reasons behind it are made conceptually clear in this work. \n- I like the organization of the paper and the intuition provided at different steps. \n\nWeaknesses: \n- I found the acceleration effect section relatively harder to follow. \n- The writing could use some revision. A review with an eye for grammar in particular would be helpful. \n- The strong/weak feature assumption could be better motivated.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ByHk6X37PEp", "reviewer": "Reviewer_6azf", "summary": "This paper studies how the prediction head in non-contrastive learning avoids trivial collapsed solutions (full collapse or dimension collapse). The empirical observations made in the paper are: (1) initializing the prediction head to identity matrix and training only the off-diagonal entries can lead to competitive performance, (2) the off-diagonal entries display a rise and fall behavior. The paper then provides theory for a simple setting to explain this phenomenon. It analyzes the dynamics of a two-layer non-linear convolutional style  neural network for a setting where data contains a strong (higher weight) and a weak feature (lower weight). Here the following can be shown: (1) without a prediction head the network will demonstrate dimensional collapse, even with batch norm (BN), (2) with a trainable prediction head + BN.+ stop-gradient can avoid such a collapse. The analysis identifies two effects that help with this, (1) substitution effect where once the stronger feature is learned in a neuron, the prediction head helps decrease the learning speed of the strong feature in other neurons and (2) acceleration effect where the strong features, via the prediction head, can further accelerate the learning of weaker features in the substituted neurons. These two steps also benefit from BN and stop-gradient.", "questions": "- Is there a way to check for substitution + acceleration in real networks? For instance, noticing that some features are trained faster and others.\n- What is the distribution of $z_p$ in Definition 3.1?\n- L173: Figure 3 is missing or wrongly referenced? Figure 3 in the appendix is just a simple framework.\n- L205: [84] is cited twice\n- OPT is undefined\n- L241: will creates -> will create\n- L51: “no well-established statistical framework” [50] connects non-contrastive methods like SimSiam with a identity covariance constraint to non-linear CCA. Although not the case where there are collapsed solutions", "limitations": "Some limitations discussed, The contributions being theoretical, there is no immediate negative societal impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- Studies a relevant problem of understanding why non-contrastive methods do not collapse and makes concrete progress towards it. To my knowledge this is the first paper that studies this for a non-linear neural network. I believe this paper would help subsequent work on non-contrastive methods.\n- The use of linear predictor + identity init + trainable off-diagonal entries is new and interesting. The empirical findings about rising and falling off-diagonals and theoretical results capturing this behavior seem non-trivial.\n- The paper is well written and mostly easy to follow\n\n**Weaknesses**\n\nThe main “weakness” in my opinion is lack of discussions about certain aspects:\n\n- Role of non-linearity: Analyzing non-linear models is key to opening up the deep learning black-box. That said, the paper currently lacks a discussion about whether and why non-linearity is crucial for the example being studies. In particular, what would happen with a linear model in this case? Will there be dimension collapse? Will substitution and acceleration kind of phenomena not happen? Including a discussion about this could be insightful either way.\n- Proof sketch: The discussion about the various phases in Section 5 is much appreciated, as it also helps understand the role of various components like trainable predictor, BN and stop-gradient. However I felt that the paper could benefit from from some more details in the proof sketch, especially to describe the substitution and acceleration effects. In particular, it might help to spell out (in terms of gradient computations say) what happens in phase III when there is no normalization and/or stop-grad.\n\nGiven these (relatively minor) concerns I am currently assigning a score of weak accept. I was unable to read the proofs in the appendix and so the above discussion will also help with increasing confidence in the results. I would be happy to raise the score if the above points are addressed in the author response, because I believe that the overall contribution of the paper is positive.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept