U NDERSTANDING S ELF - SUPERVISED L EARNING WITH DUAL DEEP NETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose a novel theoretical framework to understand self-supervised learning methods that employ dual pairs of deep ReLU networks (e.g., SimCLR, BYOL).
First, we prove that in each SGD update of SimCLR, the weights at each layer are updated by a covariance operator that speciﬁcally ampliﬁes initial random selectivities that vary across data samples but survive averages over data augmentations, which we show leads to the emergence of hierarchical features, if the input data are generated from a hierarchical latent tree model.
With the same framework, we also show analytically that BYOL works due to an implicit contrastive term, acting as an approximate covariance operator.
The term is formed by the interplay between the zero-mean operation of BatchNorm and the extra predictor in the online network.
Extensive ablation studies justify our theoretical ﬁndings.
I NTRODUCTION
While self-supervised learning (SSL) has achieved great empirical success across multiple domains, including computer vision (He et al., 2020;
Goyal et al., 2019; Chen et al., 2020a; Grill et al., 2020;
Misra and Maaten, 2020; Caron et al., 2020), natural language processing (Devlin et al., 2018), and speech recognition (Wu et al., 2020; Baevski and Mohamed, 2020; Baevski et al., 2019), its theoretical understanding remains elusive, especially when multi-layer nonlinear deep networks are involved.
Unlike supervised learning (SL) that deals with labeled data, SSL learns meaningful structures from randomly initialized networks without human-provided labels.
In this paper, we propose a systematic theoretical analysis of SSL with deep ReLU networks.
Our analysis imposes no parametric assumptions on the input data distribution and is applicable to stateof-the-art SSL methods that typically involve two parallel (or dual ) deep ReLU networks during training (e.g., SimCLR (Chen et al., 2020a), BYOL (Grill et al., 2020), etc).
We do so by developing an analogy between SSL and a theoretical framework for analyzing supervised learning, namely the student-teacher setting (Tian, 2020; Allen-Zhu and Li, 2020; Lampinen and Ganguli, 2018; Saad and Solla, 1996), which also employs a pair of dual networks.
Our results indicate that SimCLR weight updates at every layer are ampliﬁed by a fundamental positive semi deﬁnite (PSD) covariance operator that only captures feature variability across data points that survive averages over data augmentation procedures designed in practice to scramble semantically unimportant features (e.g. random image crops, blurring or color distortions (Falcon and Cho, 2020; Kolesnikov et al., 2019;
Misra and Maaten, 2020; Purushwalkam and Gupta, 2020)).
This covariance operator provides a principled framework to study how SimCLR ampliﬁes initial random selectivity to obtain distinctive features that vary across samples after surviving averages over data-augmentations.
Based on the covariance operator, we further show that (1) in a two-layer setting, a top-level covariance operator helps accelerate the learning of low-level features, and (2) when the data are generated by a hierarchical latent tree model, training deep ReLU networks leads to an emergence of the latent variables in its intermediate layers.
We also explain why BYOL works without negative pairs , by showing analytically that an interplay between the zero-mean operation in BatchNorm and the extra predictor in the online network creates an implicit contrastive term, thereby explaining empirical observations in the recent blog (Fetterman and Albrecht, 2020).
We also discover that reinitializing the predictor every a few epochs doesn’t hurt the performance in terms of linear evaluation protocol, questioning the hypothesis of “optimal predictor” in the recent version of BYOL (Grill et al., 2020).
To the best of our knowledge, we are the ﬁrst to provide a systematic theoretical analysis of modern
SSL methods with deep ReLU networks that also incorporates data augmentation.
Related Work.
Besides SimCLR and BYOL, we brieﬂy mention other concurrent SSL frameworks for vision.
MoCo (He et al., 2020; Chen et al., 2020b) keeps a large bank of past representa-
[Figure]
tions in a queue as the slow-progressing target to train from.
DeepCluster (Caron et al., 2018) and
SwAV (Caron et al., 2020) learn the representations by iteratively or implicitly clustering on the current representations and improving representations using the cluster label.
(Alwassel et al., 2019) applies similar ideas to multi-modality tasks.
In contrast, the literature on the analysis of SSL with dual deep ReLU networks is sparse.
(Arora et al., 2019) proposes an interesting analysis of how contrastive learning aids downstream classiﬁcation tasks, given assumptions about data generation.
However, no deep models is involved.
O VERALL FRAMEWORK
Notation.
Consider an L -layer ReLU network obeying f l = ψ ( ˜ f l ) and ˜
f l = W l f l − 1 for l = 1 , . . .
L .
Here ˜ f l and f l are n l dimensional pre-activation and activation vectors in layer l , with f 0 = x being the input and f L = ˜
f L the output (no ReLU at the top layer).
W l ∈ R n l × n l −are the weight matrices, and ψ ( u ) :
= max( u , 0) is the element-wise ReLU nonlinearity.
We let
W :=
{ W l } Ll =1 be all network weights.
We also denote the gradient of any loss function with respect to f l by g l ∈ R n l , and the derivative of the output f L with respect to an earlier pre-activation ˜ f l by the Jacobian matrix J l ( x ; W ) ∈ R n L × n l , as both play key roles in backpropagation (Fig. 1(b)).
An analogy between self-supervised and supervised learning: the dual network scenario.
Many recent successful approaches to self-supervised learning (SSL), including SimCLR (Chen et al., 2020a), BYOL (Grill et al., 2020) and MoCo (He et al., 2020), employ a dual “Siameselike” pair (Koch et al., 2015) of such networks (Fig. 1(b)).
Each network has its own set of weights W 1 and W 2 , receives respective inputs x 1 and x 2 and generates outputs f 1 ,L ( x 1 ; W 1 ) and f 2 ,L ( x 2 ; W 2 ) .
The pair of inputs { x 1 , x 2 } can be either positive or negative, depending on how they are sampled.
For a positive pair, a single data point x is drawn from the data distribution p ( · ) , and then two augmented views x 1 and x 2 are drawn from a conditional augmentation distribution p aug ( ·| x ) .
Possible image augmentations include random crops, blurs or color distortions, that ideally preserve semantic content useful for downstream tasks.
In contrast, for a negative pair, two different data points x , x (cid:48) ∼ p ( · ) are sampled, and then each are augmented independently to generate x 1 ∼ p
aug ( ·| x ) and x 2 ∼ p aug ( ·| x (cid:48) ) .
For SimCLR, the dual networks have tied weights with
W 1 = W 2 , and a loss function is chosen to encourage the representation of positive (negative) pairs to become similar (dissimilar).
In BYOL, only positive pairs are used, and the ﬁrst network W 1 , called the online network, is trained to match the output of the second network W 2 (the target), using an additional layer named predictor .
The target network ideally provides training targets that can improve the online network’s representation and does not contribute gradient.
The improved online network is gradually incorporated into the target network, yielding a bootstrapping procedure.
Our fundamental goal is to analyze the mechanisms governing how SSL methods like SimCLR and BYOL lead to the emergence of meaningful intermediate features, starting from random initializations, and how these features depend on the data distribution p ( x ) and augmentation procedure p aug ( ·| x ) .
Interestingly, the analysis of supervised learning (SL) often employs a similar dual network scenario, called teacher-student setting (Tian, 2020; Allen-Zhu and Li, 2020; Lampinen and
Ganguli, 2018; Saad and Solla, 1996), where W 2 are the ground truth weights of a ﬁxed teacher network, which generates outputs in response to random inputs.
These input-output pairs constitute training data for the ﬁrst network, which is a student network.
Only the student network’s weights
W 1 are trained to match the target outputs provided by the teacher.
This yields an interesting mathematical parallel between SL, in which the teacher is ﬁxed and only the student evolves, and SSL, in which both the teacher and student evolve with potentially different dynamics.
This mathematical parallel opens the door to using techniques from SL (e.g., (Tian, 2020)) to analyze SSL.
Gradient of (cid:96) 2 loss for dual deep ReLU networks.
As seen above, the (dis)similarity of representations between a pair of dual networks plays a key role in both SSL and SL.
We thus consider minimizing a simple measure of dissimilarity, the squared (cid:96) 2 distance r := (cid:107) f 1 ,L − f 2 ,L (cid:107) 2 between the ﬁnal outputs f 1 ,L and f 2 ,L of two multi-layer ReLU networks with weights W 1 and W 2 and inputs x 1 and x 2 .
Without loss of generality, we only analyze the gradient w.r.t W 1 .
For each layer l , we ﬁrst deﬁne connection K l ( x ) , a quantity that connects the bottom-up feature vector f l − 1 with the top-down Jacobian J l , which both contribute to the gradient at weight layer l .
Deﬁnition 1 (The connection K l ( x ) ) .
The connection K l ( x ; W ) :
= f l − 1 ( x ; W ) ⊗ J (cid:124) l ( x ; W ) ∈ R n l n l − 1 × n L .
Here ⊗ is the Kronecker product.
Theorem 1 (Squared (cid:96) 2 Gradient for dual deep ReLU networks) .
The gradient g W l of r w.r.t.
W l ∈ R n l × n l − 1 for a single input pair { x 1 , x 2 } is (here K 1 ,l :=
K l ( x 1 ; W 1 ) and K 2 ,l := K l ( x 2 ; W 2 ) ):
[Equation]
We used vectorized notation for the gradient g W l and weights W l to emphasize certain theoretical properties of SSL learning below.
The equivalent matrix form is ∂r/∂W 1 ,l
[Equation]
1 ,l [ J 1 ,l W 1 ,l f 1 ,l − 1 − J 2 ,l W 2 ,l f 2 ,l − 1 ] f T 1 ,l − 1 .
See Appendix for proofs of all theorems in main text.
A NALYSIS OF S IM CLR
As discussed above, SimCLR (Chen et al., 2020a) employs both positive and negative input pairs, and a symmetric network structure with W
= W 2 = W .
Let { x 1 , x + } be a positive input pair, and let { x 1 , x k − } for k = 1 , . .
, H be H negative pairs.
These input pairs induce corresponding squared (cid:96) 2 distances in output space, r + := (cid:107) f 1 ,L − f + ,L (cid:107) 22 , and r k − :=
(cid:107) f 1 ,L − f k − ,L (cid:107) 22 .
The InfoNCE loss with temperature τ then minimizes (maximizes) the positive (negative) pair distances:
[Equation]
When (cid:107) u (cid:107) 2 = (cid:107) v (cid:107) 2 = 1 , we have −(cid:107)
u − v (cid:107) 22 = 2sim( u , v ) − 2 , and Eqn.
2 reduces to what the original SimCLR uses (the term e − 2 /τ cancels out).
We leave analysis of normalization terms that make (cid:107) f L (cid:107)
= 1 as one future work.
One property of InfoNCE is important for our analysis:
Theorem 2 (Property of Contrastive Loss) .
For InfoNCE loss (Eqn. 2), we have ∂L∂r + > 0 and ∂L∂r k − < 0 for k = 1 , . . .
H and ∂L∂r + + (cid:80)
Hk =1 ∂L∂r k − = 0 .
Now, under an approximation in which we neglect variations in ∂L∂r k − across different k in a large minibatch, Theorem 1 and 2 imply the SimCLR gradient is governed by positive semi-deﬁnite (PSD) covariance operator at any layer (some simple contrastive loss (e.g., r + − r − ) satisﬁes it trivially):
Theorem 3 (Covariance Operator from Contrastive Loss) .
Assume that for any { x 1 , x + , x k − } , ∂L∂r k − = − β/H is a constant, then in a large batch limit, the update for W l is
[Equation]
, where vec(∆ W l ( t ))
=
β V x
[ ¯ K l ( x )]vec( W l ( t )) .
(3) where α is the learning rate and ¯ K l ( x ; W ) :
= E x (cid:48) ∼ p aug ( ·| x ) [ K l ( x (cid:48) ; W )] is the expected connection under the augmentation distribution, conditional on the datapoint x .
Above, we use the notation Cov[ X, Y ] :=
E [ XY (cid:124) ] − E
[ X ] E [ Y ] (cid:124) and V
[ X ] :=
Cov[ X, X ] .
The covariance operator V x (cid:2) ¯ K l ( x ) (cid:3) ∈ R n l n l − 1 × n l n l − 1 is a time-varying PSD matrix over the entire training procedure.
Therefore, all its eigenvalues are non-negative and at any time t , W l is most ampliﬁed along its largest eigenmodes.
Intuitively, this covariance operator ignores different views of the same sample x by averaging over the augmentation distribution to compute ¯ K l ( x ) , and then computes the expected covariance of this augmentation averaged connection with respect to the data distribution p ( x ) .
Thus, at all layers, any variability in the connection across different data points, that survives augmentation averages, leads to weight ampliﬁcation.
This ampliﬁcation of weights by the PSD data covariance of an augmentation averaged connection constitutes a fundamental description of SimCLR learning dynamics for arbitrary data and augmentation distributions. [Figure]
𝑧
[Figure]
FC
FC
FC 𝑧 !
𝒙(𝑧 ! ,𝑧 "
[Figure]
SSL
[Figure]
𝑤 (c) 𝑤 ",$ 𝑤 %,$ (d)
[Figure]
data
H OW THE COVARIANCE OPERATOR DRIVES THE EMERGENCE OF FEATURES
To concretely illustrate how the fundamental covariance operator derived in Theorem 3 drives feature emergence in SimCLR we setup the following paradigm for analysis.
We assume the input x = x ( z 0 , z (cid:48) ) is generated by two groups of latent variables, class/sample-speciﬁc latent z 0 and nuisance latent z (cid:48) .
After data augmentation, z 0 remains the same while z (cid:48) changes (Fig. 2(a)).
In this case, the covariance operator is simply V z
[ ¯ K l ( z 0 )] since all nuisance latents z (cid:48) are integrated out in ¯ K l ( z 0 ) .
In this setting, we ﬁrst show that a linear neuron performs PCA within an augmentation preserved subspace.
We then consider how nonlinear neurons with local receptive ﬁelds (RFs) can learn to detect simple objects.
Finally, we extend our analysis to deep ReLU networks exposed to data generated by a hierarchical latent tree model (HLTM), proving that, with sufﬁcient over-parameterization, there exist lucky nodes at initialization whose activation is correlated with latent variables underlying the data, and that SimCLR ampliﬁes these initial lucky patterns during learning.
4. S ELF - SUPERVISED LEARNING AND THE SINGLE NEURON : ILLUSTRATIVE EXAMPLES
A single linear neuron performs PCA in a preserved subspace.
For a single linear neuron ( L = 1 , n L = 1 ), the connection in deﬁnition 1 is simply K 1 ( x ) = x .
Now imagine the input space x can be decomposed into the direct sum of a semantically relevant subspace, and its orthogonal complement, which corresponds to a subspace of nuisance features.
Furthermore, suppose the augmentation distribution p aug ( ·| x ) is obtained by multiplying x by a random Gaussian matrix that acts only in the nuisance subspace, thereby identically preserving the semantic subspace.
Then the augmentation averaged connection ¯ K 1 ( x ) =
Q s x where Q s is a projection operator onto the semantic subspace.
In essence, only the projection of data onto the semantic subspace survives augmentation averaging, as the nuisance subspace is scrambled.
Then the covariance operator in
Theorem 3 is V x
[ ¯ K 1 ( x )]
= Q s V x
[ x ] Q s (cid:124) .
Thus the covariance of the data distribution, projected onto the semantic subspace, governs the growth of the weight vector W 1 , demonstrating SimCLR on a single linear neuron performs PCA within a semantic subspace preserved by data augmentation.
A single linear neuron cannot detect localized objects .
We now consider a generative model in which data vectors can be thought of as images of objects of the form x ( z 0 , z (cid:48) ) where z 0 is an important latent semantic variable denoting object identity, while z (cid:48) is an unimportant latent variable denoting nuisance features, like object pose or location.
The augmentation procedure scrambles pose/position while preserving object identity.
Consider a simple concrete example (Fig. 3(a)):
[Equation]
Here 0 ≤ z (cid:48) ≤ d − 1 denotes d discrete translational object positions on a periodic ring and z 0 ∈ { 1 , 2 } denotes two possible objects 11 and 101 .
The distribution is uniform both over objects and positions: P ( z 0 , z (cid:48) )
= 12 d .
Augmentation shifts the object to a uniformly random position via P aug ( z (cid:48) | z 0 )
= 1 /d .
For a single linear neuron K 1 ( x ) = x , and the augmentation-averaged connection is ¯ K 1 ( z 0 ) =
2 d 1 , and is actually independent of object identity z
(both objects activate two pixels at any location).
Thus V z 0 (cid:2)
¯ K 1 ( z 0 ) (cid:3) = 0 and no learning happens.
A local receptive ﬁeld (RF) does not help.
In the same generative model, now consider a linear neuron with a local RF of width 2 .
Within the RF only four patterns can arise: 00 , 01 , 10 , 11 .
Taking the expectation over z (cid:48) given z 0 (Fig.
3(a2))
yields ¯ K 1 ( z 0 =1) =
1 d [ x 11 + x 01 + x 10 + ( d − 3) x 00 ]
[Figure]
(a1) 𝑧
[Figure] [Figure]
and ¯ K 1 ( z 0 =2)
= 1 d [2 x 01 + 2 x 10 + ( d − 4) x 00 ] .
Here, x 11 ∈ R 2 denotes pattern 11 .
This yields
[Equation]
where u := x 11 + x 00 − x 01 − x 10 .
(5) and V z 0 (cid:2) ¯ K 1 ( z 0 ) (cid:3)
∈ R 2 × 2 since the RF has width 2.
Note that the signed sum of the four pattern vectors in u actually cancel, so that u = 0 , implying V z 0 (cid:2)
¯ K 1 ( z 0 ) (cid:3) = 0 and no learning happens.
Interestingly, although the conditional distribution of the 4 input patterns depends on the object identity z 0 (Fig.
3(a2)), a linear neuron cannot learn to discriminate the objects.
A nonlinear neuron with local RF can learn to detect object selective features .
With a ReLU neuron with weight vector w , from Def. 1, the connection is now K 1 ( x , w ) =
ψ (cid:48) ( w (cid:124) x ) x .
Suppose w ( t ) happens to be selective for a single pattern x p (where p ∈ { 00 , 01 , 10 , 11 } ), i.e., w ( t ) (cid:124)
x p > 0 and w ( t ) (cid:124) x p (cid:48) < 0 for p (cid:48) (cid:54) = p .
The augmentation averaged connection is then ¯ K 1 ( z 0 )
∝ x p where the proportionality constant depends on object identity z 0 and can be read off (Fig.
3(a2)).
Since this averaged connection varies with object identity z 0 for all p , the covariance operator does not vanish and is given by V z 0 (cid:2)
¯ K 1 ( z 0 ) (cid:3) =
c p x p x (cid:124) p where the constant c p > 0 depends on the selective pattern p and can be computed from Fig.
3(a2).
By Theorem 3, the dot product x (cid:124) p w ( t ) grows over time:
[Equation]
Thus the learning dynamics ampliﬁes the initial selectivity to the object selective feature vector x p in a way that cannot be done with a linear neuron.
Note this argument also holds with bias terms and initial selectivity for more than one patterns.
Moreover, with a local RF, the probability of weak initial selectivity to local object sensitive features is high, and we may expect ampliﬁcation of such weak selectivity in real neural network training.
4. T WO - LAYER CASE WITH MULTIPLE HIDDEN NEURONS
Now consider a two-layer network ( L = 2 ) with 1 output ( n L = 1 ) and 1-hidden layer with n
ReLU neurons (Fig. 2(c)).
In this case, the augmentation-averaged connection ¯ K 1 ( z 0 ) at the lowest layer l = 1 can be written as
[Equation]
(7) where u j ( z 0 ) :
= E z (cid:48) | z 0 (cid:2) x
( z 0 , z (cid:48) )
I ( w (cid:124) 1 ,j x ( z 0 , z (cid:48) ) ≥ 0) (cid:3) is the expected activation of the hidden layer.
Note that each u j ( z 0 ) ∈
R d and w 2 ,j is the scalar weight of the second weight matrix that connects node j to the output.
Then we have the following theorem:
Theorem 4.
If Cov z 0 [ u j , u k ]=0 for j (cid:54) =
k , then the time derivative of w 2 ,j and w 1 ,j satisﬁes:
[Equation]
where A j := V z
[ u j ( z 0 )] .
It is easy to see that d | w 2 ,j | 2 / d t = d (cid:107) w 1 ,j (cid:107) 2 / d t and thus w 22 ,j = (cid:107) w 1 ,j (cid:107)
+ c where c is some constant which doesn’t change over time.
Since A j is PSD, w (cid:124) 1 ,j A j w 1 ,j ≥ 0 and w 2 ,j always grows, which in turn accelerates the convergence of w 1 ,j to the largest eigenvector of A j .
This shows top- down modulation : the existence of top-layer weights accelerate the training of the lower layer.
Previous works (Allen-Zhu and Li, 2020) also mention a similar concept in supervised learning, called “backward feature corrections”, for deep polynomial networks with quadratic activations.
Here we provide a rigorous formulation showing similar behaviors happen in SSL with SGD training.
Note that if we consider ReLU neurons, A j changes with w 1 ,j ; while for linear nodes, A j is a constant, since the gating I ( w (cid:124) 1 ,j x > 0) is always 1.
If the data x forms a mixture of two Gaussians: x ∼ 12 I ( z 0 =1) N ( w ∗ 1 , σ 2 ) + 12 I
( z 0 =2) N ( w ∗ 2 , σ 2 ) .
Let ∆ w ∗ := w ∗ 1 − w ∗ 2 , then in the linear case,
A j ∼ ∆ w ∗ ∆ w ∗ (cid:124) and w 1 ,j converges to ± ∆ w ∗ (Fig. 3(b)).
In the nonlinear case with multiple
Gaussians, if one of the Gaussians sits at the origin (e.g., background noise), then dependent on initialization, A j evolves into w ∗ k w ∗ (cid:124) k for some center k , and w 1 ,j → w ∗ k .
Note this dynamics is insensitive to speciﬁc parametric forms of the input data.
[Table]
4. D EEP R E LU SSL TRAINING WITH H IERARCHICAL L ATENT T REE M ODELS ( HLTM )
For multi-layer case, we mainly study a hierarchical latent tree model, where each leaf variable is sampled via a Markov Chain from the root latent variable z 0 .
The model is motivated by the structure of ConvNet: each latent variable at the intermediate layer generates a subset of the observable data x that is sent to the deep ReLU networks, in which the nodes also have local receptive ﬁelds (Fig. 2(d)).
We deﬁne symbols in Tbl.
At layer l , we have categorical latent variables { z µ } , where µ ∈ Z l is the Greek letter index for latent variables.
Each z µ can take discrete values.
The topmost latent variable is z 0 .
Following the tree structure, for µ ∈ Z l and ν 1 , ν 2 ∈ Z l − 1 , conditional independence holds: P ( z ν 1 , z ν 2 | z µ ) =
P ( z ν 1 | z µ ) P ( z ν 2 | z µ ) .
The ﬁnal sample x is the instantiation of the leaf latents (Fig. 2(d)), and thus depends on all latent variables.
Corresponding to the hierarchical tree model, each neural network node j ∈ N l maps to a unique µ = µ ( j ) ∈
Z l .
Let N µ be all nodes that map to µ .
For j ∈ N µ , its activation f j only depends on the value of z µ and its descendant latent variables.
We also deﬁne v j ( z µ ) :
=
E z
[ f j | z µ ] as the expected activation w.r.t z µ .
Given a sample x , the data augmentation is done by resampling all z µ (which are z (cid:48) in Fig. 2), given the root z 0 .
Symmetric Binary HLTM .
Here we consider a symmetric binary case: each z µ ∈ { 0 , 1 } and for µ ∈ Z l , ν ∈
Z l − 1 , P ( z ν = 1 | z µ = 1) =
P ( z ν = 0 | z µ = 0)
= (1 + ρ µν ) / 2 , where the polarity ρ µν ∈ [ − 1 , 1] measures how informative z µ = 1 is.
If ρ µν = ± 1 then there is no stochasticity in the top-down generation process; ρ µν = 0 means no information in the downstream latents and the posterior of z 0 given the observation x can only be uniform.
See Appendix for more general cases.
Now we compute covariance operator V z
[ ¯ K µ ( z 0 )] at different layers, where ¯ K µ ( z 0 ) =
E z (cid:48) (cid:104)
f N ch µ ⊗ J (cid:124) µ | z 0 (cid:105) , we ﬁrst check the term E z (cid:48) (cid:104)
f N ch µ | z 0 (cid:105) and incorporate the Jacobian later.
Theorem 5 (Covariance operator in binary HLTM
[Equation]
Here a N ch
[Equation]
:=
[Equation]
If max αβ | ρ αβ | ≤ γ < 1 , then lim L → + ∞ ρ 0 µ → 0 for µ ∈ N l .
Theorem 5 suggests when ρ 0 µ and (cid:107) a N ch µ (cid:107) is large, the covariance operator is large and training is faster.
For deep HLTM and deep networks, at lower layers, ρ 0 µ → 0 and P 0 µ becomes uniform due to mixing from the progression of the Markov Chain, making V z
[ ¯ K l ( z 0 )] small.
Therefore, training in SSL is faster at the top layers where the covariance operators have large magnitude.
On the other hand, large (cid:107) a N ch µ (cid:107) means that s k :
= ( v k (1) − v k (0)) / 2 is large, or the expected activation v k
( z ν ) is selective among different values of z ν for ν ∈ ch( µ ) .
Interestingly, this can be achieved by over-parameterization (i.e., |N µ | > 1 ):
Theorem 6 (Lucky nodes in deep ReLU networks regarding to binary HLTM ) .
Suppose each element of the weights W l between layer l + 1 and l are initialized with
Uniform (cid:104) − σ w (cid:113) 3 / |N ch µ | , σ w (cid:113) 3 / |N ch µ | (cid:105) .
There exists σ 2 l so that V [ f k | z ν ] ≤ σ 2 l for any k ∈ N l .
For any µ ∈ Z l +1 , if |N µ | =
O (exp( c )) , then with high probability, there exists at least one node j ∈ N µ so that their pre-activation gap | ˜ v j (1) − ˜ v j (0)
| = 2 w (cid:124)
j a N ch µ >
and the activations satisfy:
[Equation]
Intuitively, this means that with large polarity ρ µν (or strong signals sending top-down), overparameterized ReLU networks with random initialization yields selective neurons, if the lower layer also contains selective ones.
For example, when σ l = 0 , c = 9 , if ρ µν ≥ 63 .
3% then there is a gap between expected activations v j (1) and v j (0) , and the gap is larger when the selectivity in the lower layer l is higher.
Note that at the lowest layer, { v k } are themselves observable leaf latent variables and are selective by deﬁnition.
So a bottom-up mathematical induction will happen.
After the gradient update, if we assume the top-down Jacobian is just 1 , then for the “lucky” node j we will have (for brevity, let c := αβρ 20 µ (1 − ρ 20 ) ):
[Equation]
which means that the pre-activation gap ˜ v j (1) − ˜ v j (0)
= 2 w (cid:124)
j a N ch µ grows over time and the latent variable z µ is learned (instantiated as f j ) during training, even if the network is never supervised with its true value.
Similar to Sec.
4.2, once the top layer starts to have large weights, ¯ K l ( z 0 ) for lower layer becomes larger due to large Jacobian and training is accelerated.
We implement the HLTM and conﬁrm, as predicted by our theory, that the intermediate layers of deep ReLU networks do indeed learn the latent variables of the HLTM (see Tbl.
2 below).
A NALYSIS OF INGREDIENTS UNDERLYING THE SUCCESS OF BYOL
In BYOL, the two networks are no-longer identical and, interestingly, only positive pairs are used for training.
The ﬁrst network with weights W 1 =
W is an online network that is trained to predict the output of the second target network with weights W 2 = W (cid:48) , using a learnable predictor W pred to map the online to target outputs (Fig. 1(a) and Fig.
1 in Grill et al.
In contrast, W (cid:48) is an exponential moving average (EMA) of W : W (cid:48)
( t + 1) =
γ ema W (cid:48) ( t ) + (1 − γ ema )( W\W pred )( t ) .
We now theoretically analyze BYOL to explain why the predictor and an additional ingredient
BatchNorm (BN) (Ioffe and Szegedy, 2015) are simultaneously required for BYOL’s success.
Without EMA ( γ ema = 0 ) and the predictor, we obtain W (cid:48)
=
W and the BYOL update without BN for
W is (derivation in Appendix E.1):
[Equation]
This update only promotes variance minimization in the representations of different augmented views of the same data samples and therefore would yield model collapse.
Motivated by a recent blogpost (Fetterman and Albrecht, 2020), we now consider BN in addition to the predictor.
In particular, we analyze a simpliﬁed version of BN in which the mini-batch mean only is subtracted.
When adding predictor, Theorem 1 can still be applied by adding identity layers on top of the target network W (cid:48) so that the online and the target networks have the same depth.
Theorem 5 in (Tian, 2018) demonstrates this version of BN shifts the downward gradients so their mini-batch mean is 0 :
[Equation]
Here g il is the i -th sample in a batch and ¯ g l is the batch average (same for ¯ f l ).
Backpropagating through this BN (vs. just subtracting the mean only in the forward pass 1 ), leads to a correction term:
Theorem 7.
If (1) the network is linear from layer l to the topmost and (2) the downward gradient g l undergoes Eqn.
11, then with large batch limits, the correction of the update is 2 (for brevity, dependency on W is omitted, while dependency on W (cid:48) is made explicit):
[Equation]
and the corrected weight update is (cid:93) ∆ W l := ∆ W l + δW BN l .
Using Eqn.
10, we have:
[Equation]
Theorem 7 makes several predictions.
Without EMA or a predictor ( W (cid:48)
= W ), e.g. SimCLR, then δW BN l = 0 and BN should not matter.
In contrast, as suggested in empirical ﬁndings in a recent blogpost (Fetterman and Albrecht, 2020) as well as our own experiments (Tbl. 3), standard BYOL without BN fails.
Theorem 7 provides a theoretical explanation: when W (cid:48) (cid:54)
=
W and BN is present, δW BN l (cid:54) = 0 gives an implicit contrastive term.
Note that W (cid:48) (cid:54) =
W means there is a predictor or EMA or both.
We discuss the predictor (and leave EMA in Appendix E.3).
Consider the last two terms in Eqn. 13: − V x (cid:2) ¯ K l ( x ) (cid:3) vec( W l ) + Cov x (cid:2) ¯ K l ( x ) , ¯ K l ( x ; W (cid:48) ) (cid:3) vec( W (cid:48) l ) .
Because − V x (cid:2) ¯ K l ( x ; W ) (cid:3) is second order in the Jacobian of W pred while +Cov x (cid:2) ¯ K l ( x ; W ) , ¯ K l ( x ; W (cid:48) ) (cid:3) , is ﬁrst-order, when W pred has small initial weights, the latter term dominates.
In this regime, BYOL with predictor+BN sensibly ampliﬁes variance across data samples through +Cov x (cid:2) ¯ K l ( x ) , ¯ K l ( x ; W (cid:48) ) (cid:3) (last term in Eqn. 13) and minimizes variance across different augmented views of the same data sample through − E x (cid:2)
V x (cid:48) ∼ p aug ( ·| x )
[ K l ( x (cid:48) )] (cid:3) (ﬁrst term in Eqn. 13).
[Table] [Table]
[Table] [Table]
E XPERIMENTS
We test our theoretical ﬁndings through experiments on STL-10 (Coates et al., 2011) and ImageNet (Deng et al., 2009).
We use a simpliﬁed linear evaluation protocol: the linear classiﬁer is trained on frozen representations computed without data augmentation.
This reuses pre-computed representations and accelerates evaluation by 10x.
Hierarchical Latent Tree Model ( HLTM ) .
We implement the HLTM and check whether the intermediate layers of deep ReLU networks learn the corresponding latent variables at the same layer.
The degree of learning is measured by the normalized correlations between the ground truth latent variable z µ and its best corresponding node j ∈ N µ at each layer.
Tbl. 2 indicates this measure increases with over-parameterization and learning, consistent with the analysis in Sec. 4.3.
Factors underlying BYOL performance .
To test our theory, we perform an ablation study of
BYOL on STL-10 by modifying three key components: predictor, EMA and BN.
Tbl.
3 shows that BN and predictor are critical and EMA further improves the performance.
First, without a predictor, neither BN nor EMA give good performance.
A predictor without BN still doesn’t work.
A predictor with BN starts to show good performance ( 63 . 6% ) and further adding EMA leads to the best performance ( 78 . 1% ).
This is consistent with our theoretical ﬁndings in Sec. 5, in which we show that using a predictor with BN yields δW BN l (cid:54) = 0 and leads to an implicit contrastive term.
To further test our understanding of the role played by BN, we fractionate BN into several sub-components: subtract by batch mean ( mean-norm ), divide by batch standard deviation ( std-norm ) and affine , and do ablation studies (Tbl. 4).
Surprisingly, removing affine yields slightly better performance on STL-10 (from 78.1% to 78.7%).
We also ﬁnd that variants of mean-norm performs reasonably, while variants of detached mean-norm has similar poor performance as no normalization, supporting that centralizing backpropagated gradient leads to implicit contrastive terms (Sec. 5).
Note that std-norm also helps, which we leave for future analysis.
We also check whether the online network requires an “optimal predictor” as suggested by recent version (v3) of BYOL.
For this, we reinitialize the predictor ( ReInit ) every T epochs and compare the ﬁnal performance under linear evaluation protocol.
Interestingly, as shown in Tbl. 5, ReInit actually improves the performance a bit, compared to the original BYOL that keeps training the same predictor, which should therefore be closer to optimal.
Moreover, if we shrink the initial weight range of the predictor to make Cov x (cid:2) ¯ K l ( x ) , ¯ K l ( x ; W (cid:48) ) (cid:3) (third term in Eqn. 13) more dominant, and reduce the learning rate, the performance further improves (See Tbl. 10 in Appendix F), thereby corroborating our analysis.
A B ACKGROUND AND B ASIC S ETTING (S ECTION 2) A. T HEOREM
Deﬁnition 2 (reversibility) .
A layer l is reversible if there is a function G l ( x ; W ) so that f l ( x ; W ) =
G l ( x ; W ) f l − 1 ( x ; W ) and g l = G (cid:124) l ( x ; W ) g
l − 1 .
A network is reversible if all layers are.
Note that many different kinds of layers have this reversible property, including linear layers (MLP and Conv) and (leaky) ReLU nonlinearity.
For multi-layer ReLU network, we have:
[Equation]
14) where D l ∈ R n l × n l is a binary diagonal matrix that encodes the gating of each neuron at layer l .
The gating D l ( x ; W ) depends on the current input x and current weight W .
Remark .
Note that the reversibility is not the same as invertible.
Speciﬁcally, reversibility only requires the transfer function of a backpropagation gradient is a transpose of the forward function.
Lemma 1 (Recursive Gradient Update (Extension to Lemma 1 in (Tian, 2020)) .
If the network is reversible (Def. 2), then minimizing the L2 objective
[Equation]
with respect to weight matrix W l at layer l yields the following gradient at layer l :
[Equation]
where V l − 1 ( x ) :
= V l ( x ) G l ( x ) ∈ R n l × n l − 1 and V L ( x ) = I n L × n L .
Proof.
We prove by induction.
Note that our deﬁnition of W l is the transpose of W l deﬁned in (Tian, 2020).
Also our g l ( x ) is the gradient before nonlinearity, while (Tian, 2020) uses the same symbol for the gradient after nonlinearity.
For notation brievity, we let f l ( x 1 ) :
= f l ( x 1 ; W l ) and G l ( x 1 ) :
= G l ( x 1 ; W l ) .
Similar for x 2 and W 2 .
When l = L , by the property of L2-loss, we know that g L =
f L ( x 1 ; W 1 ) − f L ( x 2 ; W 2 ) , by setting V
L ( x 1 ) =
V L ( x 2 )
=
I , the condition holds.
Now suppose for layer l , we have:
[Equation] [Equation]
Then:
[Equation] [Equation]
Note that for multi-layered ReLU network, if we set x 1 = x 2 = x , G l ( x ) =
D l ( x ) W l for each
ReLU+Linear layer and W 1 = W , W 2 = W ∗ (teacher weights), then we go back to the original
Lemma 1 in (Tian, 2020).
Remark.
Note that the same structure holds for blocks of ResNet with ReLU activation.
Now we prove Theorem 1 as follows:
Proof.
We consider more general case where the two towers have different parameters, namely Wand W 2 .
Applying Lemma 1 for the branch with input x 1 , and use ReLU structure G l = D l W l , we have:
[Equation]
where f 1 ,l − 1 :
= f l − 1 ( x 1 ; W 1 ) is the activation of layer l − 1 at tower 1 (similar for other symbols).
The gradient (and the weight update, according to gradient descent) of the weight W l between layer l and layer l − 1 is (here we deﬁne the Jacobian matrix J l ( x ; W ) :
= V l ( x ; W ) D l ( x ; W ) ):
[Equation]
[Equation]
Using vec( AXB )
=
( B (cid:124) ⊗ A )vec( X ) (where ⊗ is the Kronecker product), we have:
[Equation]
Let
[Equation] [Equation]
Note that K l ( x ; W ) is a function of the current weight W , which includes weights at all layers.
By the mixed-product property of Kronecker product ( A ⊗ B )( C ⊗ D ) =
AC ⊗ BD , we have:
[Equation]
where K l ( x 1 ) =
K l ( x 1 ; W 1 ) and K l ( x 2 ) =
K l ( x 2 ; W 2 ) .
In SimCLR case, we have W 1 =
W 2 =
W so
[Equation]
B A NALYSIS OF S IM CLR USING T EACHER -S TUDENT S ETTING (S ECTION 3) B. T HEOREM
Proof.
we have:
[Equation] [Equation]
and obviously we have:
[Equation]
B. T HEOREM
Proof.
First we have:
[Equation]
Then we compute each terms
[Equation]
Using Theorem
[Equation]
we know that:
[Equation]
Since Eqn.
33 holds, K l ( x 1 ) K (cid:124) l ( x 1 ) will be cancelled out and we have:
[Equation]
Then we use the assumption that ∂L∂r 2 k − = − β/n with β > 0 :
[Equation]
Taking large batch limits, we know that E [ K l ( x 1 ) K (cid:124) l ( x + )]
=
E x (cid:2) ¯ K l ( x ) ¯ K (cid:124) l ( x ) (cid:3) since x 1 , x + ∼ p
aug ( ·| x ) are all augmented data points from a common sample x .
On the other hand,
E
[Equation]
( x )
[Equation]
since x 1 and x k − are generated from independent samples and data augmentation.
Therefore,
[Equation]
C T HE DYNAMICS OF TWO - LAYER R E LU NETWORK AND THE INTERPLAYS OF COVARIANCE OPERATORS BETWEEN NEARBY LAYERS (S ECTION 4.2) C. T HEOREM
Proof.
For convenience, we deﬁne the centralized version of u j ( z 0 ) :
ˆ u j ( z 0 ) =
u j ( z 0 ) − E z
[ u j ( z 0 )] and the matrices A jk :=
Cov z
[ u j ( z 0 ) , u k ( z 0 )] =
E z
[ˆ u j ( z 0 )ˆ u (cid:124) k
( z 0 )] .
At layer l = 1 the covariance operator is V z
[ ¯ K 1 ( z 0 )]
= [ w 2 ,j w 2 ,k A jk ] ∈ R n 1 d × n 1 d .
On the other hand, if we check the second layer l = 2 , we could compute ¯ K 2 ( z 0 ) ∈
R n 1 : ¯ K 2 ( z 0 ) =
[Equation]
=
[Equation]
] =
[Equation]
|
[Equation]
[ x (cid:124) ⊗ D 1 ( x ) | z 0 ] vec( W 1 ) (41)
For each component j , we have [ ¯ K 2 ( z 0 )]
j = w (cid:124) 1 ,j u j ( z 0 ) , or: ¯ K 2 ( z 0 ) =
[ w (cid:124) 1 , 1 u 1 ( z 0 ) , w (cid:124) 1 , 2 u 2 ( z 0 ) , . .
. , w (cid:124) 1 ,n 1 u n 1 ( z 0 )] (42)
So at layer l = 2 we can compute the covariance operator V z
[ ¯ K 2 ( z 0 )]
=
[ w (cid:124) 1 ,j A jk w 1 ,k ] ∈
R n 1 × n 1 .
Using the assumption that A jk :=
E
[Equation] [Equation]
u j
[Equation]
z
[Equation]
u
[Equation]
k
[Equation]
=
E
[Equation] [Equation]
u
[Equation]
z
[Equation]
)] E
[Equation]
[ˆ u (cid:124) k
( z 0 )]
= 0 for j (cid:54) =
k .
We arrive at the conclusion that w j is only coupled with v j , yielding the following dynamical system:
[Equation]
D H IERARCHICAL L ATENT T REE M ODELS (S ECTION 4.3)
D.
L EMMAS
Lemma 2 (Variance Squashing) .
Suppose a function φ : R (cid:55)→ R is L -Lipschitz continuous: | φ ( x ) − φ ( y ) | ≤ L | x − y | , then for x ∼ p ( · ) , we have:
[Equation]
[Table]
Proof.
Suppose x, y ∼ p ( · ) are independent samples and µ φ :=
E [ φ ( x )] .
Note that V [ φ ( x )] can be written as the following:
[Equation]
Therefore we have:
[Equation]
Lemma 3 (Sharpened Jensen’s inequality (Liao and Berg, 2018)) .
If function φ is twice differentiable
[Equation]
and x ∼ p ( · )
[Equation]
then we have:
[Equation]
Lemma 4 (Sharpened Jensen’s inequality for ReLU activation) .
For ReLU activation ψ ( x ) := max( x, 0) and x ∼ p ( · ) , we have:
[Equation]
Proof.
Since ψ is a convex function
[Equation]
by Jensen’s inequality we have E [ ψ ( x )] − ψ ( E [ x ]) ≥ 0 .
For the other side
[Equation]
let µ :=
E p [ x ] and we have (note that for ReLU, ψ ( x ) − ψ ( µ ) ≤ | x − µ | ):
[Equation]
where the last inequality is due to Cauchy-Schwarz.
D. M ORE GENERAL ASSUMPTION OF CONDITIONAL DISTRIBUTION IN HLTM
We consider a more general case of HLTM where each z µ is categorical: 0 ≤ z µ < m µ .
For convenience, we deﬁne the following symbols for k ∈ N ch µ (note that |N ch µ | = N ch µ is the number of the children of the node set N µ ):
[Equation]
As an extension of binary symmetric HLTM , we make an assumption for the transitional probability:
Assumption 1.
For µ ∈ Z l and ν ∈ Z l − 1 , the transitional probability matrix P µν :=
[ P ( z ν | z µ )] has decomposition P µν = 1 m ν 1 µ 1 (cid:124)
ν + C µν
where C µν 1 ν = 0 µ and 1 (cid:124) µ C µν =
ν .
Note that C µν 1 = 0 is obvious due to the property of conditional probability.
The real condition is 1 (cid:124) µ C µν =
ν .
If m µ = m ν , then P µν is a square matrix and Assumption 1 is equivalent to P µν is double-stochastic.
Assumption 1 makes computation of P µν easy for any z µ and z ν .
Lemma 5 (Transition Probability) .
If Assumption 1 holds, then for µ ∈ Z l , ν ∈ Z 1 − 1 and α ∈
Z l − 2 , we have:
[Equation]
In general, for any µ ∈ N l 1 and α ∈ N l 2 with l 1 > l 2 , we have:
[Equation]
Proof.
Using Assumption 1, we have
[Equation] [Equation]
since 1 (cid:124) ν 1 ν = m ν , C µν 1 ν =
ν
and 1 (cid:124)
ν C να =
α , the conclusion follows.
Remark .
In the symmetric binary HLTM mentioned in the main text, all C µν can be parameterized as (here q :=
[ − 1 , 1] (cid:124) ):
[Equation]
This is because 1 (cid:124) 2 C µν = 0 2 and C µν 1 2 = 0 2 provides 4 linear constraints (1 redundant), leaving 1 free parameter, which is the polarity ρ µν ∈ [ − 1 , 1] of latent variable z ν given its parent z µ .
Moreover, since q (cid:124) q = 2 , the parameterization is close under multiplication:
[Equation]
D. T HEOREM
Proof.
First note that for each node k ∈ N ch µ :
[Equation]
Note that 1 0 1 (cid:124)
ν v k is a constant regarding to change of z 0 .
So we could remove it when computing covariance operator.
On the other hand, for a categorical distribution ( P ( z
[Equation] [Equation]
With P 0 :=
diag[ P ( z 0 )] , the mean is E z
[ u ] = 1 (cid:124) P 0 u and its covariance can be written as (here 1 = 1 0 ):
[Equation]
[Equation]
Note that each column of V 0 , N ch µ is v 0 k .
Setting u
= v 0 k
and we have:
[Equation]
Note that E z 0 [ v 0 k ] = 1 m ν 1 (cid:124)
ν v k
[Equation]
E z
[ C 0 ν v k ] , since 1 (cid:124) P 0 1 = 1 .
With some computation, we could see Cov z 0 [ v 0 k , v 0 k (cid:48) ]
=
Cov z
[ C 0 ν ( k ) v k , C 0 ν ( k (cid:48) ) v k (cid:48) ] .
The equation above can be applied for any cardinality of latent variables.
In the binary symmetric case, we have (note here we deﬁne ρ 0 :=
P ( z 0 = 1) − P ( z 0 = 0) , and q :=
[ − 1 , 1] (cid:124) ):
[Equation]
Note that in the binary symmetric case, according to remarks in Lemma 5, all C µν = 12 ρ µν qq (cid:124) and we could compute C
ν v k :
[Equation]
where according to Eqn. 61, we have:
[Equation]
69) and the covariance between node k and k (cid:48) can be computed as:
[Equation]
[Equation]
The last equality is due to the fact that due to tree structure, the path from z 0 to all child nodes in
N ch µ must pass z µ .
Therefore we can compute the covariance operator:
[Equation]
When L → + ∞ , we have:
[Equation]
[Equation]
and thus the covariance becomes zero as well.
D. T HEOREM
Proof.
According to our setting, for each node k ∈ N µ , there exists a unique latent variable z ν with ν = ν ( k ) that corresponds to it.
In the following we omit its dependency on k for brevity.
Since we are dealing with binary case, we deﬁne the following for convenience:
[Equation]
We also deﬁne the sensitivity of node k to be λ k := | ( v + k ) 2 − ( v − k ) 2 | .
Intuitively, a large λ k means that the node k is sensitive for changes of latent variable z ν .
If λ k = 0 , then the node k is invariant to latent variable z ν .
We ﬁrst consider pre-activation ˜ f
j := (cid:80) k w jk f k and its expectation with respect to latent variable z :
[Equation]
Note that for each node k ∈ N ch µ we have:
[Equation] [Equation]
Let a N ch µ :=
[ a k ] k ∈N ch µ :=
[ ρ µν s k ] k ∈N ch µ and u + N ch µ :
[Equation]
Then we have ˜ v + j = w (cid:124)
j u + N ch µ and ˜ v − j = w (cid:124)
j u −N ch µ .
Note that w j is a random variable with each entry w jk ∼ Uniform
[Equation]
It is easy to verify that E [ w jk ] = 0 and V [ w jk ] =
σ 2 w / |N ch µ | .
Therefore, for two dimensional vector ˜ v j =
[˜ v + j , ˜ v − j ] (cid:124) , we can compute its ﬁrst and second order moments: E w [˜ v j ] = 0 and
V w [˜ v j ] = σ 2 w |N ch µ | V µ, N ch µ V (cid:124) µ, N ch µ = σ 2 w |N ch µ | [ u + N ch µ , u −N ch µ ] (cid:124)
[ u + N ch µ , u −N ch µ ] .
Deﬁne the positive and negative set (note that a k :=
ρ µν s k ):
[Equation]
Without loss of generality, assume that (cid:80)
k ∈
A + a 2 k ≥ (cid:80) k ∈
A − a 2 k .
In the following, we show there exists j with λ j is greater than some positive threshold.
Otherwise the proof is symmetric and we can show λ j is lower than some negative threshold.
When |N ch µ | is large, by Central Limit Theorem, v can be regarded as zero-mean 2D Gaussian distribution and we have for some c > 0 :
[Equation]
Moreover
[Equation]
if a l (cid:54) =
[Equation]
then the following probability is also not small :
[Equation]
Therefore, when |N µ | = O (exp( c )) , with high probability, there exists w j so that
[Equation]
Since ¯ v N ch µ ≥ 0 (all f k are after ReLU and non-negative), this leads to:
[Equation]
By Jensen’s inequality, we have (note that ψ ( x ) :=
max( x, 0) is the ReLU activation):
[Equation]
On the other hand, we also want to compute v − j :=
E z
[ f j | z µ = 0] using sharpened Jensen’s inequality (Lemma 4).
For this we need to compute the conditional covariance V z
[ ˜ f
j | z µ ] :
[Equation]
Note that 1 (cid:13) is due to conditional independence: f k as the computed activation, only depends on latent variable z ν and its descendants.
Given z µ , all z ν and their respective descendants are independent of each other and so does f k .
(cid:13) is due to the fact that each w jk are sampled from uniform distribution and | w jk | ≤ σ w (cid:113) 3 |N ch µ | .
Here V z ν | z µ [ E z [ f k | z ν ]] =
s 2 k
(1 − ρ 2 µν ) can be computed analytically.
It is the variance of a binomial distribution: with probability 12 (1 + ρ µν ) we get v + k otherwise get v − k .
Therefore, we ﬁnally have:
[Equation] [Equation]
As a side note, using Lemma 2, since ReLU function ψ has Lipschitz constant ≤ 1 (empirically it is smaller), we know that:
[Equation]
Finally using Lemma 4 and ˜ v − j < 0 , we have:
[Equation]
Combining Eqn. 92 and Eqn. 101, we have a bound for λ j :
[Equation]
E T HE A NALYSIS OF BYOLE. D ERIVATION OF BYOL GRADIENT
Note that for BYOL, we have:
[Equation]
under large batchsize, we have (note that we omit W for any term that depends on W , but make dependence of W (cid:48) explicit in the math expression):
[Equation]
For brevity, we write E x [ · ] :=
E x ∼ p ( · ) [ · ] and E x (cid:48) [ · ] :=
E x (cid:48) ∼ p aug ( ·| x ) [ · ] .
Similar for V .
And the equation above can be written as:
[Equation]
when W = W (cid:48) , the last two terms cancelled out, yielding:
[Equation]
Note that the weight update ∆ W l is its negation and the conclusion follows.
E. T HEOREM
Proof.
When BN is present, Eqn. 103 needs to be corrected with an additional term, (cid:103)
∂r∂W l :=
∂r∂W l − δW BN l , where δW BN l is deﬁned as follows:
[Equation]
From the proof of Theorem 1, we know that for each sample i ∈ B :
[Equation]
Since the network is linear from layer l to the topmost layer L , we have D il = ¯ D l .
Since the only input dependent part in J il is the gating function between the current layer l and the topmost layer
L , for linear network the gating is always 1 and thus ¯ J l = J il and is independent of input data.
We now have (note that we omit W for any terms that are dependent on W , but will write W (cid:48) explicitly for terms that are depend on W (cid:48) ):
[Equation]
Therefore we have:
[Equation]
Note that by assumption, since ¯ J l doesn’t depend on the input data, we have
[Equation] [Equation]
Taking large batchsize limits and notice that the batch B could contain any augmented data generated from independent samples from p ( · ) , we have:
[Equation]
An important thing is that the expectation is taking over x ∼ p ( x ) and x (cid:48) ∼ p aug ( ·| x ) .
Intuitively, this is because ¯ f l − 1 and ¯ g l are averages over the entire batch, which has both intra-sample and inter-sample variation.
With augment-mean connection ¯ K l ( x ) we could write:
[Equation]
Plug in δW l,BN into Eqn.
and we have corrected gradient for BYOL:
[Equation]
Note that the weight update (cid:93) ∆ W l is its negation and the conclusion follows.
E. A NALYSIS ON EMA
Consider the following discrete dynamics of a weight vector w ( t ) :
[Equation]
119) where α is the learning rate, λ > 0 represents the excessive magnitude of EMA terms compared to the weight term according to Sec. 5, w ema ( t + 1) = γ ema w ema ( t ) + (1 − γ ema ) w ( t ) is the exponential moving average of w ( t ) . For convenience, we use η := 1 − γ ema . Since it is a recurrence equation, we apply one-sided z-transform (Oppenheim, 1999) on the temporal domain (note that here the symbol z has nothing to do with the latent variables in Sec. 4), where
[Equation]
[ w ( t )]
= (cid:80) + ∞ t =0 w ( t )
z − t .
This leads to:
[Equation]
Note that for w ema ( t ) we have:
[Equation]
If we set w ema (0) = 0 , i.e., the target network is all zero at the beginning, then it gives w ema ( z ) =
η z − 1+ η w ( z ) .
Plugging it back to Eqn.
and we have:
[Equation]
And then we could solve w ( z ) :
[Equation]
[Equation]
Note that the denominator has two roots
z 1 and z 2 :
[Equation]
[Equation]
and w ( z ) can be written as
[Equation]
[Equation]
Without loss of generality, let z 1 < z 2 .
The larger root z 2 > 1 when λ > 0 (which happens according to Sec. 5), so the zero ( z = 1 − η = γ ema < 0 ) in the nominator won’t cancel out the pole at z 2 .
And we have:
[Equation]
where 1 / (1 − z 2 z − 1 ) corresponds to a power series z t 2 in the temporal domain.
Therefore, we could see w ( t ) has exponential growth due to z 2 > 1 and learning happens.
Now let us check how z 2 changes over η , i.e., how the parameter γ ema :=
1 − η of EMA affects the learning process.
We have:
[Equation]
Use the fact that (1 + x ) 1 / 2 ≤ 1 + 12 x for x ≥ 0 , we have:
[Equation]
[Equation]
Compared to no EMA case (i.e., γ ema = 0 or η = 1 ), with a γ ema < 1 but close to 1 (or equivalently, η is close to 0 ), the upper bound of z 2 becomes smaller but still greater than 1, and the exponential growth is less aggressive, which stabilizes the training.
Note that if γ ema = 1 (or η = 0 ), then w ema ( t ) ≡ w ema (0) = 0 and learning also doesn’t happen.
F A DDITIONAL E XPERIMENTS
Based on our theoretical analysis, we try training the predictor in different ways and check whether it still works.
From the analysis in Sec. 5, we know that the reason why BYOL works is due to the dominance of Cov x
[ ¯ K l ( x ) , ¯ K l ( x ; W )] and its resemblance of the covariance operator
V x
[ ¯ K l ( x )] , which is a
PSD matrix.
The dominance should be stronger if the predictor has smaller weights than normally initialized using Xavier/Kaiming initialization.
Also, Cov x [ ¯ K l ( x ) , ¯ K l ( x ; W (cid:48) )] should behave more like a
PSD matrix, if the predictor’s weights are all small positive numbers and no BN is used.
The following table justiﬁes our theoretical ﬁndings.
In particular, Tbl.
10 shows better performance in STL-10 with smaller learning rate and smaller sample range of the predictor weights.
Sample
[Table]
With BN in predictor
[Table]
Without BN in predictor
[Table]
[Table] [Table] [Table]