The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning
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 ﬁnd 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 diversiﬁed 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 ﬁrst end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with a trainable prediction head and normalization.
Introduction
Self-supervised learning is about learning representations of real-world vision or language data without human supervision, and contrastive learning [62, 43, 41, 24, 20, 34] is one of 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.e., pairs of different data points.
However, the surprising ﬁnding of the Bootstrap Your Own Latent ( BYOL ) method by Grill et al.
[37] initiated the research of non-contrastive self-supervised learning , which refers to contrastive learning methods without using the negative pairs.
BYOL achieved state-of-the-art results in various computer vision benchmarks and there are plenty of follow-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 φ ( x (cid:48) ) , where x and x (cid:48) 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 φ ( · ) is a constant vector whose variance is zero.
Another trivial solutions called dimensional collapse by [44] is when all the coordinates φ i ( · ) are exactly aligned.
Nevertheless, adding a trainable prediction head on top of (one branch of) φ ( 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.
Since the proposition of BYOL , there have been lots of empirical studies trying to understand noncontrastive 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 yet .
To explain the non-contrastive learning, it is inevitable to study how the solutions are chosen during the optimization.
Therefore, our research questions are:
[Equation] of prediction [Equation]
Why do most non-contrastive self-supervised 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 diversiﬁed representations in non-contrastive self-supervised learning?
Due to the existence of trivial collapsed optimal solutions of the 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 a high level, the results in this paper are summarized as follows:
Our empirical contributions.
In non-contrastive self-supervised 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 veriﬁed that even when the prediction head is an identity-initialized matrix (with ﬁxed diagonal entries), its off-diagonal entries display a rise-and-fall pattern, and it does not always converge to a symmetric matrix.
See Figure 3.
Our theoretical contributions.
We based our theory on a 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 a 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 to 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 diversiﬁed 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 effect .
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
[Figure]
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 diversiﬁed 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 out that all the claims below are derived only in our theoretical setting and are partially veriﬁed 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 a theoretical statement that (symmetric) linear prediction head will converge to a matrix that commutes with the covariance matrix of linear representations at 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., the eigenspace alignment result in [78]) is not sufﬁcient to understand the trainable prediction head.
Can the symmetric prediction head explain the trainable prediction head?
In the paper [78], experiments over the STL-10 dataset showed that the linear prediction head converges to a 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).
Speciﬁcally, under their linear network setting, where W is the weight matrix of the base encoder, they manually set the prediction head W p at iteration t to be W ( t ) p [Equation] [Equation] 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
[Figure] Figure 1 and Figure 3b show that even if we initialize the prediction head using a symmetric matrix (identity), the trainable prediction head can be very asymmetric at the early training stage when the encoder network learn most of its features .
Actually, in the presence of feature imbalance (e.g., E x 1 x 1 x (cid:62) 1 has huge eigen-gap), the symmetric prediction head is also likely to learn a rank-one matrix where W focus on learning the largest eigenvector of E x 1 x 1 x (cid:62) 1 .
The role of stop-gradient and output-normalization.
It is 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-gradient and output-normalization together can turn the features substituted via the prediction head into a factor in 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 ﬁrst to explain the effects of output-normalization in optimizing nonlinear neural networks in self-supervised learning.
Preliminaries on Non-contrastive Learning
In this section, we formally deﬁne what is non-contrastive self-supervised learning.
To do this, we ﬁrst introduce contrastive learning following [24, 85] as background.
We use [ N ] as a shorthand for the index set { 1 , . . .
, N } .
Background on contrastive learning.
Letting φ W ( · ) be the neural networks, contrastive learning aims to learn good representations φ W via contrasting representations of similar data samples to those of dissimilar ones.
Usually we are given a batch of data points { X i } i ∈ [ N ] , and we construct for each i ∈ [ N ] a positive pair ( X (1) i , X (2) i ) by applying random data augmentations to X i , and collect negative pairs ( X (1) i , X (2) j ) for i (cid:54) =
j ∈ [ N ] .
Now given z i =
φ W ( X (1) i ) , z (cid:48)
i =
φ W ( X (2) i ) , i ∈
[ N ] , we train the network φ W to minimize the contrastive loss:
[Equation]
where sim ( · , · ) is the similarity metric, often deﬁned as the cosine similarity, and τ is the so-called temperature hyper-parameter.
Intuitively, minimizing the contrastive loss can be roughly viewed as trying to classify the representation z
i as z (cid:48)
i instead of z (cid:48)
j , j (cid:54) =
i .
It is a common belief that in order for the network φ W to be able to “distinguish” data points X i from X j , j (cid:54)
=
i , merely minimizing the positive term of contrastive loss is not sufﬁcient.
Non-contrastive self-supervised learning.
We choose the SimSiam method [26] as our primary framework, whose differerence with BYOL is a EMA component that is proven inessential in [26].
Following the same notations as above, except that z (cid:48) i =
StopGrad [ φ W ( X (2) i )] is detached from [Figure] gradient computation, the loss objective become: (the symmetric network version)
[Equation]
which is just the positive term in contrastive loss (2.1) (not divided by τ ).
Clearly there exist plenty trivial global optimal solutions for this objective.
For example, the complete collapse refers to when φ W ( · ) learns some constant vector.
Another solution called dimensional collapse [44] is when all the coordinates [ φ W ( · )]
i has correlation ± 1 .
The dimensional collapsed solution can minimize the objective (2.2) even when φ W ( · ) is BN-normalized to avoid learning a constant vector [44, 88].
However, by adding a trainable prediction head on top of z i , the training miraculously succeeds and outputs a state-of-the-art feature extractor.
Let g ( · ) be a shallow feed-forward network (often one or two-layer, or even simply linear), we train g and φ W simultaneously on the following objective:
[Equation]
) (2.3) where z (cid:48) i is still detached from gradient computation.
The g ( z i ) =
g ◦ φ W ( X (1) i ) and the detached part z (cid:48)
i =
StopGrad [ φ W ( X (2) i )] are often called the online network and the target network respectively following [39], known as two branches of non-contrastive learning.
Note that the trainable prediction head can represent identity function, so the objective (2.3) still has the collapsed optima.
Problem Setup
In this section, we present the setting of our theoretical results.
We ﬁrst deﬁne the data distribution.
Notations.
We use O, Ω , Θ notations to hide universal constants with respect to d and (cid:101) O, (cid:101) Ω , (cid:101) Θ notations to hide polynomial factors of log d .
We denote a = o (1) if a → 0 when d → ∞ .
We use the notations poly ( d ) , polylog ( d ) to represent large constant degree polynomials of d or log d .
We use N ( µ, Σ) to denote standard normal distribution in with mean µ and covariance matrix Σ .
We use the bracket (cid:104)· , ·(cid:105) to denote the inner product and (cid:107) · (cid:107) 2 the (cid:96) 2 -norm in Euclidean space.
And for a subspace V ⊂ R d , we denote V ⊥ as its orthogonal complement.
We use 1 B to denote the indicator function of event B .
We use I m to denote the m × m identity matrix.
Following the standard structure of image datasets, we consider data divided into patches, where each patch can contain either features or noises.
Deﬁnition 3.1 (data distribution and features) .
Let X ∼ D be X
=
( X 1 , . . .
, X P ) ∈ R d × P where each X i ∈
R
d is a patch.
We assume that there are two feature vectors v 1 , v 2 such that (cid:107) v (cid:96) (cid:107) 2 = 1 , (cid:96) = 1 , 2 and are orthogonal to each other.
To generate a sample X , we uniformly sampled (cid:96) ∈ [2] and generate for each p ∈
[ P ] :
[Equation]
where z p ( X ) is the latent vector of X , ξ p is the noise vector of patch p ∈
[ P ] whose assumption will be given in Assumption 3.2.
We denote S ( X ) =
{ p : z p ( X ) (cid:54) = 0 } ⊆
[ P ] as the set of feature patches, where z p ( X ) =
z p (cid:48)
( X ) ∈ { 0 , ±
α (cid:96) } , ∀ p, p (cid:48) ∈ [ P ] , where α (cid:96) will be picked afterwards.
We assume P = polylog ( d ) , S ( X ) ≡ P 0 = Θ(log d ) for every X .
A ﬁgurative illustration is given in Figure 4.
Strong and weak features.
We pick α 1 = 2 polyloglog ( d ) and α
2 = α 1 / polylog ( d ) .
Hence v 1 is the strong feature and v 2 is the weak feature , and we want the learner network to learn both v 1 , v 2 (but by different neurons) as their learning goal.
This is a simpliﬁcation of the real scenario.
Intuitively, we can think of the strong features in a dataset are the ones that show up more frequently or with larger magnitude, and weak features as those that show up rarely or with smaller magnitude.
Assumption 3.2 (noise) .
Denoting V = span( v 1 , v 2 ) , we assume
ξ p ∈ V ⊥ is independent for each p ∈
[ P ] \ S ( X ) , where X =
( X p ) p ∈
[ P ] ∼ D , and:
(a) For any unit vector u ∈ V ⊥ , E [ (cid:104) ξ p , u (cid:105) ] = 0 , and E [ (cid:104) ξ p , u (cid:105) 6 ]
= σ 6 for some σ = Θ(1) ; (b) It holds for some (cid:37) ∈ [0 , 1 d Ω(1) ] it holds | E [ (cid:104) u 1 , ξ p (cid:105) 3 (cid:104) u 2 , ξ p (cid:105) 3 ] | ≤ (cid:37) and | E [ (cid:104) u 1 , ξ p (cid:105) 5 (cid:104) u 2 , ξ p (cid:105) ]
| ≤ (cid:37) for any two vectors u 1 , u 2 ∈ V ⊥ that are orthogonal to each other
Remark 3.3 .
A simple example of our noise ξ p is the spherical Gaussian noise in V ⊥ .
Assumption 3.2b ensures that the prediction head cannot be used to cancel the noise correlation between different neurons.
We point out that the features in our data can be learned via clustering, but we emphasize that we do not intend to compare our algorithm with any clustering method in this setting since our goal is to study how the prediction head helps in learning the features.
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 = ( w 1 , . . . , w m
) ∈ R d × m be the weight matrix, where w i ∈ R d , the encoder network f is deﬁned by
[Equation]
Here we use the cubic activation function σ ( z ) =
z 3 , as polynomial activations are 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 deﬁned as a matrix E =
[ E i,j ] ( i,j ) ∈ [ m ] 2 with
E i,i ≡ 1 , i ∈ [ m ] , where only the the off-diagonals E i,j , i (cid:54) =
j are trainable parameters.
The online network (cid:101) F is deﬁned by: given j ∈ [ m ] , we let F j ( X ) :=
f j ( X ) + (cid:80) r (cid:54) =
j E j,r f r ( X ) , and
[Equation]
where the batch normalization BN 1 here is deﬁned as follows:
Given a batch of inputs { z i } i ∈ [ N ] ,
[Equation]
And the we deﬁne the target network G as (cid:101) G j ( X ) :
= BN ( G j ( X ))
=
BN ( f j ( X )) , ∀ j ∈ [ m ] .
3.2 Training Algorithm
Data augmentation.
We use a very simple data augmentation: for each data X =
( X p ) p ∈
[ P ] , we randomly and uniformly sample half of the patches P ⊆
[ P ] to generate the positive pair :
[Equation]
Our data augmentation is similar to the common random cropping used in contrastive learning [22, 76].
It is also analogous to the data augmentations studied in theoretical literatures [85, 47, 58].
Non-contrastive loss function.
Now we deﬁne the loss function as follows: we sample N data points { X i } i ∈
[ N ] , X i i .
i .
d .
∼
D and apply our data augmentation (3.2) to obtain S
=
{ X ( i, 1) , X ( i, 2) } i ∈
[ N ] .
Now we deﬁne
[Equation]
where the StopGrad operator detach gradient computation of the target network (cid:101) G ( · )
This form of objective (3.3) is ﬁrst deﬁned in [37] and is equivalent to (2.3) in Chen and He [26] when (cid:101) F and (cid:101)
G share the same encoder network f ( · ) and their outputs are normalized.
Intuitions of the data augmentation and collapse.
In Deﬁnition 3.1, the features v 1 , v 2 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 ) = (cid:80) p σ ( (cid:104) v 1 , X p (cid:105) ) or φ ( X ) =
(cid:80) p σ ( (cid:104) v 2 , X p (cid:105) ) is sufﬁcient.
We consider learning the same feature v i in all the neurons f j in the encoder network f as the dimensional collapsed solution.
Initialization and hyper-parameters.
At t = 0 , we initialize W and E as W (0) i,j ∼ N (0 , 1 d ) and
E (0)
=
I m
and we only train the off-diagonal entries of E ( t ) .
For the simplicity of analysis, we let m = 2 , which sufﬁces to illustrate our main message.
For the learning rates, we let η ∈ (0 , 1 poly ( d ) ] be sufﬁciently small and η E ∈ [ η/α O (1), η/ polylog ( d )] , which is smaller than η 2 .
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 t = 0 , . .
, T − 1 , we sample a new batch of augmented data S t =
{ X ( t,i, 1) , X ( t,i, 2) } i ∈
[ N ] and update
[Equation]
If we do not train the prediction head, we just simply keep E ( t ) ≡ I m .
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 )) :
= E X
∼D [( ψ ( X ) − E [ ψ ( X )]) 2 ] be the variance of any function ψ of X ∼ D , we denote the correlation Corr ( ψ ( X ) , ψ (cid:48)
( X )) of any two function ψ, ψ (cid:48) over D as
[Equation]
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 > 2 , let
N ≥ poly ( d ) , η ∈ (0 , 1 poly ( d ) ] be sufﬁciently small, and η E ∈ [ η α O (1), η polylog ( d ) ] .
Then with probability 1 − o (1) , after training for T = poly ( d ) /η
many iterations, we shall have for some (cid:96) ∈ [2] : w ( T )=
β 1 v (cid:96) + ε 1 , w ( T [Equation] with | [Equation]
1 | [Equation] | [Equation]
Furthermore, the objective converges: E S∼D N [ [Equation] + [Equation]
Where OPT stands for the global optimum 3 .
Theorem 4.1 clearly shows the network learn all the desired features, even under huge imbalance between v 1 and v 2 .
This leads to the following corollary.
Corollary 4.2.
Under the same hyper-parameter in Theorem 4.1, with probability 1 − o (1) , after training for T = poly ( d ) /η many iterations, then the encoder f avoids dimensional collapse:
[Equation]
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, we 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 N ≥ poly ( d ) ,
η
= o (1) and the number of neurons m ≤ O ( α 1 /α 2 ) .
Suppose we freeze E ( t ) ≡ I m , then with probability 1 − o (1) , after training for T = poly ( d ) /η many iterations, we shall have:
[Equation]
β
[Equation]
ε
[Equation]
with | β j | = Θ(1) , (cid:107) ε j (cid:107) 2 ≤ (cid:101) O ( 1 √ d ) for all j ∈ [ m ]
Furthermore, the objective converges: E S∼D N [ L S ( W ( T ) , E ( T ) )] ≤ OPT + 1 poly ( d ) ≤ O ( 1 log d ) .
This means the collapsed solution also reaches the global 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 a certain variance in the outputs.
Immediately, we have a corollary.
Corollary 4.4.
Under the same hyper-parameter in Theorem 4.3, with probability 1 − o (1) , after training with E ( t ) ≡ I m for T = poly ( d ) /η
many iterations, we shall have dimensional collapse :
[Equation]
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.
[Equation]
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 ﬁrst three phases explain how the prediction head can help learn the base encoder, and the last phase explains why the off-diagonal entries often shrink later in training.
5.1 Phase I: Learning the Stronger Feature
At the beginning of training, the stronger feature v 1 enjoys a much larger gradient as opposed to the weaker feature v 2 , so naturally, v 1 will be learned ﬁrst.
Without loss of generality, let us assume at initialization, the neuron f 1 has larger v 1 between f j , j ∈ [2] , then we can show:
Lemma 5.1 (learning the stronger feature, see Lemma C.13) .
After some t ≥ T 1 = d 2+ o (1) /η , the feature v 1 in neuron f 1 will be learn to (cid:104) w ( t ) 1 , v 1 (cid:105) = Ω(1) , while all other features (cid:104) w ( t ) j , v (cid:96) (cid:105) =
o (1) for ( j, (cid:96) ) (cid:54) =
And the prediction head (cid:107) E ( t ) − I 2 (cid:107) 2 ≤ d − Ω(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 v 1 in neuron f 1 is learned to a certain degree, and neuron f 2 remains largely unlearned.
5.2 Phase II: The Substitution Effect
To illustrate the substitution effect, let us keep assuming that neuron f 1 has already learned some signiﬁcant amount of the strong feature v 1 , say w 1 = β 1 v 1 + residual with | β 1 | =
Ω( (cid:107) residual (cid:107) ) , then we have: (recall f j ( · ) , j ∈ [2] are the neurons of the encoder network)
Lemma 5.2 (substitution effect, formal statement see Lemma D.8) .
After |(cid:104) w ( t ) 1 , v 1 (cid:105)| = Ω(1) , we shall have | E ( t ) 2 , 1 | increasing until | E ( t ) 2 , 1 f 1 ( X (1) )
| (cid:29) | f 2 ( X (1) ) | when X is equipped with the strong feature v 1 {+, for T 2 − T 
 = o ( T 1 ) iterations.+}
Intuition of the substitution effect.
After the stronger feature is learned in neuron f 1 , the optimal way to align two positive representations F 2 ( X (1) ) and G 2 ( X (2) ) is not learning features in weight w 2 , but use the prediction head to “substitute” the features in f 1 into F 2 .
This is how the substitution effect happens when trained with a prediction head.
5.3 Phase III: The Acceleration Effect
After the substitution of v 1 in F 2 , our concern is, w ( t ) 2 will learn v 2 and only v 2 eventually, according to the acceleration effect in the following lemma.
Lemma 5.3 (acceleration effect, formal statement see Lemma E.8) .
After E ( t ) 2 , 1 is learned in Lemma 5.2, learning v 2 in w ( t ) 2 will be much faster than v 1 , until (cid:107) w ( t ) 2 − β 2 v 2 (cid:107) =
o (1) for some β 2 = Θ(1) .
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 −∇ w 2 F 2 ( X (1) ) ·
StopGrad [ G 2 ( X (2) )] to learn f 2 , this negative gradient will only try to maximize f 2 ( X (1) ) · f 2 ( X (2) ) , rather than to maximize F 2 ( X (1) ) · f 2 ( X (2) ) .
This is because the stop-gradient is on G not on F : while F 2 has a large component of v 1 borrowed from f 1 using E ,
G 2 does not have this component.
So the gradient of F 2 is to align with the features in G 2 that does not contain many v 1 , while the gradient of G 2 is to aligned with F 2 that contains a lot of v 1 .
What is the role of the output normalization?
Again due to the StopGrad operation, the gradient of (cid:101) F 2 is taken with respect to the ratio f 2 ( X (1) ) /
(cid:112) Var
[ F 2 ( X (1) )] .
As gradient descent tries to maximize this ratio, a direct computation gives ∇ w
[Equation]
which borrow the substituted feature v 1 from f 1 ( · ) to adjust the gradient of v 2 in f 2 ( · ) , via the prediction head E ( t ) 2 , 1 .
Without the output normalization, the learning of v 1 will dominate that of veven 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 t ≥ T = poly ( d ) /η iterations, we shall have (cid:107) E ( t ) − I 2 (cid:107) F ≤ 1 poly ( d ) .
While we admit that only some of our real-world experiments in Figure 3 show the convergence to zero for the off-diagonal 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.
Additional Related Work
Self-supervised learning Self-supervised learning has created huge success in natural language processing [30, 86, 18] and established the pretrain-ﬁnetune paradigm for deep learning.
In vision, contrastive learning [75, 41, 24, 20, 27, 28, 34, 64, 33] became dominant in many downstream tasks recently.
Another approach is the generative learning [65, 16, 42], which also gives promising results.
Applications such as [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 a feature learning view, and inspired our analysis in the non-contrastive setting.
For generative learning, [51, 74] provides downstream performance guarantees.
[70, 84] studied the natural language tasks.
[58] gave a 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 on 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.
Conclusion
In this paper, we showed how the prediction head can ensure the neural network learns all the features in non-contrastive learning through theoretical investigation.
Our key contribution is 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 a 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.