T OWARDS D EMYSTIFYING R EPRESENTATION L EARNING WITH N ON - CONTRASTIVE S ELF - SUPERVISION
Anonymous authors
Paper under double-blind review
A BSTRACT
Non-contrastive methods of self-supervised learning (such as BYOL and SimSiam) learn representations by minimizing the distance between two views of the same image.
These approaches have achieved remarkable performance in practice, but it is not well understood how the representation is learned based on the augmentation process.
Tian et al.
(2021) explained why the representation does not collapse to zero and proposed DirectPred that sets the predictor directly.
In our work, we analyze a generalized version of DirectPred, called DirectSet ( α )
We show that in a simple linear network, DirectSet( α ) provably learns a desirable projection matrix and also reduces the sample complexity on downstream tasks.
Our analysis suggests that weight decay acts as an implicit threshold that discard the features with high variance under augmentation and keep the features with low variance.
Inspired by our theory, we simplify DirectPred by removing the expensive eigen-decomposition step.
On CIFAR-10, CIFAR-100, STL-10 and ImageNet, DirectCopy , our simpler and more computationally efﬁcient algorithm, rivals or even outperforms DirectPred.
I NTRODUCTION
Self-supervised learning recently emerges as a promising direction to learn representations without manual labels.
While contrastive learning (Oord et al., 2018; Tian et al., 2019; Bachman et al., 2019;
He et al., 2020; Chen et al., 2020a) minimizes the distance of representation between positive pairs, and maximizes such distances between negative pairs, recently, non-contrastive selfsupervised learning (abbreviated as nc-SSL ) is able to learn nontrivial representation with only positive pairs, using an extra predictor and a stop-gradient operation.
Furthermore, the learned representation shows comparable (or even better) performance for downstream tasks (e.g., image classiﬁcation) (Grill et al., 2020; Chen & He, 2020).
This brings about two fundamental questions: (1) why the learned representation does not collapse to trivial (i.e., constant) solutions, and (2) without negative pairs, what representation nc-SSL learns from the training and how the learned representation reduces the sample complexity in downstream tasks.
While many theoretical results on contrastive SSL (Arora et al., 2019; Lee et al., 2020; Tosh et al., 2020; Wen & Li, 2021) do exist, similar study on nc-SSL has been very rare.
As one of the ﬁrst work towards this direction, Tian et al.
(2021) show that while the global optimum of the noncontrastive loss is indeed a trivial one, following gradient direction in nc-SSL, one can ﬁnd a local optimum that admits a nontrivial representation.
Based on their theoretical ﬁndings on gradientbased methods, they proposed a new approach, DirectPred, that directly sets the predictor using the eigen-decomposition of the correlation matrix of input before the predictor, rather than updating it with gradient methods.
As a method for nc-SSL, DirectPred shows comparable or better performance in multiple datasets, including CIFAR-10 (Krizhevsky et al., 2009), STL-10 (Coates et al., 2011) and ImageNet (Deng et al., 2009), compared to BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2020) that optimize the predictor using gradient descent.
While Tian et al.
(2021) address the ﬁrst question, i.e., why the learned representation does not collapse, they do not address the second question, i.e., what representation is learned in nc-SSL and how the learned representation is related to the data distribution and augmentation process and in turn whether it reduces the sample complexity in downstream tasks.
Main Contributions.
In this paper, we make a ﬁrst attempt towards the second question, by studying a family of algorithms named DirectSet ( α ) , in which the DirectPred algorithm proposed by Tian et al.
(2021) is a special case with α = 1 / 2 .
Our contribution is two-folds:
First, we perform a theoretical analysis on DirectSet ( α ) with linear networks.
Our analysis shows that there exists an implicit threshold , determined by weight decay parameter η , that governs which features are learned and which are discarded.
More speciﬁcally, the threshold is applied to the variance of the feature across different data augmentations (or “views”) of the same instance: nuisance features (features with high variances under augmentation) are discarded, while invariant features (i.e., with low variances) are kept.
We further make a formal statement on the sample complexity of the learning process and performance guarantees of the downstream tasks, in the linear setting similar to Tian et al.
To our knowledge, this is the ﬁrst sample complexity result in nc-SSL.
Second, we show that DirectCopy , a special case of DirectSet ( α ) when α = 1 , performs comparably with (or even outperforms) DirectPred in downstream tasks in CIFAR-10, CIFAR-100, STLand ImageNet.
In DirectCopy, the predictor can be set without the expensive eigen-decomposition operation, which makes DirectCopy much simpler and more efﬁcient than DirectPred.
Related works.
In nc-SSL, different techniques are proposed to avoid collapsing.
BYOL and SimSiam use an extra predictor and stop gradient operation.
Beyond these, BatchNorm (including its variants (Richemond et al., 2020)), de-correlation (Zbontar et al., 2021; Bardes et al., 2021; Hua et al., 2021), whitening (Ermolov et al., 2021), centering (Caron et al., 2021), and online clustering (Caron et al., 2020) are all effective ways to enforce implicit contrastive constraints among samples for collapsing prevention.
We study BYOL and SimSiam as representative nc-SSL methods.
Organization.
The paper is organized as follows.
Section 2-3 introduce DirectSet( α ), prove it learns a projection matrix onto the invariant features, and the learned representation reduces sample complexity in downstream tasks.
Section 4 demonstrates that DirectCopy achieves comparable or even better performance than the original DirectPred algorithm in various datasets, and Sectionshows ablation experiments.
Finally, limitation and future works are discussed in Section 6-7.
P RELIMINARIES 2. N OTATIONS
We use I d to denote the d × d identity matrix and simply write I when the dimension is clear.
For any linear subspace S in R d , we use P S ∈ R d × d to denote the projection matrix on S. More precisely, the projection matrix P S equals UU (cid:62) , where the columns of U constitute a set of orthonormal bases for subspace S .
We use N ( µ, Σ) to denote the Gaussian distribution with mean µ and covariance Σ .
We use (cid:107)·(cid:107) to denote spectral norm for a matrix, or (cid:96) 2 norm for a vector and use (cid:107)·(cid:107) F to denote
Frobenius norm for a matrix.
For a real symmetric matrix A ∈ R d × d whose eigen-decomposition is (cid:80) di
=1 λ
i u i u
(cid:62) i , we use | A | to denote (cid:80) di
=1 |
λ i | u i u (cid:62) i .
If A is also positive semi-deﬁnite, we use A α to denote (cid:80) di =1 λ αi u i u (cid:62) i for any positive α ∈ R .
2. D IRECT S ET ( α ) AND D IRECT C OPY
In nc-SSL, recent methods as BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2020) employ a dual pair of Siamese networks (Bromley et al., 1994): one side is a composition of an online network (including a projector) and a predictor network, the other side is a target network (see Figure 1 for a simple example).
The target network has the same architecture as the online network, but has potentially different weights.
Given an input x , two augmented views x 1 , x 2 are generated, and the network is trained to match the representation of x 1 (through the online network and the predictor network) and the representation of x 2 (through the target network).
More precisely, suppose the online network and the target network are two mappings f θ , f θ a : R d (cid:55)→ R h and the predictor network is a mapping g θ p :
R h (cid:55)→ R h , the network is trained to minimize the following loss:
[Equation]
In BYOL and SimSiam, the online network and the target network are trained by running gradient methods on L .
The target network is not trained by gradient methods; instead, it is directly set with
[Figure]
the weights in the online network (Chen & He, 2020) or an exponential moving average (EMA) of the online network (Grill et al., 2020; He et al., 2020; Chen et al., 2020b).
The DirectSet( α ) algorithm, as shown in Figure 1, directly sets the predictor based on the correlation matrix F of the predictor inputs:
[Equation]
where F = E x 1 f θ ( x 1 )
f θ ( x 1 ) (cid:62) .
In practice, F is estimated by a moving average over batches.
That is, ˆ F = µ ˆ F
+ (1 − µ ) E B [ f θ ( x 1 ) f θ ( x 1 ) (cid:62) ] , where E B is the expectation over one batch.
In the original DirectPred proposed by Tian et al.
(2021), α is ﬁxed at 1 / 2 .
To compute ˆ F 1 / 2 , one needs to ﬁrst compute the eigen-decomposition of ˆ F , and then taking the root of each eigenvalue.
This step of eigen-decomposition can be expensive especially when the representation dimension h is high.
To avoid the eigen-decomposition step, we propose DirectCopy ( α = 1 ), in which the predictor W p is a direct copy of the ˆ F (with normalization and regularization) 1 .
As we shall see,
DirectCopy enjoys both theoretical guarantees and strong empirical performance.
T HEORETICAL A NALYSIS OF D IRECT S ET ( α )
Deep linear networks have been widely used as a tractable theoretical model for studying nonconvex loss landscapes (Kawaguchi, 2016; Du & Hu, 2019; Laurent & Brecht, 2018) and nonlinear learning dynamics
(Saxe et al., 2013; 2019; Lampinen & Ganguli, 2018; Arora et al., 2018).
However, most of them are for supervised learning setting.
Tian et al.
(2021) analyzed nc-SSL on a linear network, but did not analyze their proposed approach DirectPred.
Here, we analyze the representation learning process of DirectSet( α ) on a minimal setting where the online network f θ is a single linear layer.
We also verify DirectSet( α ) works for practical nonlinear deep models and realistic datasets.
3. S ETUP
In this subsection, we deﬁne the network model, data distribution and simplify DirectSet( α ) algorithm for our theoretical analysis.
We consider the following network model (see Figure 1),
Assumption 1 (Linear network model) .
The online, predictor and target network are all single-layer linear network without bias, with weight matrices denoted as W, W p , W a ∈ R d × d respectively.
For the data distribution, we assume the input space is a direct sum of a invariant feature subspace and a nuisance feature subspace.
Speciﬁcally, we assume
Assumption 2 (Data distribution) .
The input x is sampled from N (0 , I d ) , and its augmented view x 1 , x 2 are independently sampled from N ( x, σ 2 P B ) , where B is a ( d − r ) -dimensional subspace.
We denote S as the orthogonal subspace of B in R d .
In this simple data distribution, subspace S corresponds to the features that are invariant to augmentations and its orthogonal subspace B is the nuisance subspace which the augmentation changes.
We will prove that DirectSet( α ) can learn the projection matrix onto S subspace.
Note in the previous work (Tian et al., 2021), they assumed the covariance of the augmentation distribution to be σ 2 I and did not study what representation is learned.
For the convenience of analysis, we consider a simpliﬁed version of DirectSet( α ).
We compute the loss function without normalizing the two representations, so the population loss is
[Equation]
and the empirical loss is
[Equation]
where x ( i ) ’s are independently sampled from N (0 , I ) , and augmented views x ( i ) 1 and x ( i ) 2 are independently sampled from N ( x ( i ) , σ 2 P B ) .
To train our model, ﬁrst, we initialize W as δI with δ a positive real number.
We run gradient ﬂow or gradient descent on online network W with weight decay η , and set the the target network W a = W.
For clarity of presentation, when training on the population loss, we set W p as ( W E x xx (cid:62) W (cid:62) )
α = ( WW (cid:62) )
α instead of ( W E x 1 x 1 x (cid:62) 1 W (cid:62) )
α as in practice; when training on the empirical loss, we set W p as ( W 1 n (cid:80) ni =1 x (
[Equation]
[ x ( i ) ] (cid:62) W (cid:62) )
α .
Here, we set the predictor regularization (cid:15) =
and its inﬂuence will be studied in Section 5.
In the following, DirectSet( α ) is shown to recover the projection matrix P S with polynomial number of samples.
Furthermore, given that the learned matrix is close to P S , the sample complexity on downstream tasks is reduced.
3. G RADIENT F LOW ON P OPULATION L OSS
In this section, we show that DirectSet( α ) running on the population loss with inﬁnitesimal learning rate and η weight decay can learn the projection matrix onto the invariant feature subspace S .
Theorem 1.
Suppose network architecture and data distribution follow Assumption 1 and Assumption 2, respectively.
Suppose we initialize online network W as δI, and run DirectSet ( α ) on population loss (see Eqn. 1) with inﬁnitesimal step size and η weight decay. If we set the weight decay coefﬁcient η ∈ (cid:16) [Equation] [Equation]
and initialization scale
[Equation] ) , then W converges to (cid:16) [Equation] (cid:17) 1 / (2 α ) P S when time goes to inﬁnity.
Theorem 1 shows that when the weight decay is in certain range, and when the initialization is large enough, the online network can converge to the desired projection matrix P S 2 .
In sequel, we explain how the dynamics of W leads to a projection matrix and how the weight decay and initialization scale come into play.
We leave the full proof in Appendix B.1.
We also consider the setting when W p is set as ( W E x 1 x 1 x (cid:62) 1 W (cid:62) )
α in Appendix B.4 and extend the result to deep linear networks in Appendix C.
Due to the identity initialization, we can ensure that W is always a real symmetric matrix and is simultaneously diagonalizable with P B .
We can then analyze the evolution of each eigenvalue in W separately.
Under our assumptions, it turns out that all the eigenvalues whose eigenvectors lie in the
B subspace share the same value λ B , and all the eigenvalues in the S subspace share the value λ S as shown in the following time dynamics:
[Equation]
Next, we show λ B converges to zero and λ S converges to a positive number, which immediately implies that W converges to some scaling of P S .
[Figure]
Similar as the analysis in Tian et al.
(2021), when η > 1 4(1+ σ 2 ) , we know ˙
λ B < 0 for any λ B >and λ B = 0 is a stable stationary point, as illustrated in Figure 2 (top, left).
Therefore, as long as η > 1 4(1+ σ 2 ) , λ B must converge to zero.
When 0 < η < 14 , there are three non-negative solutions to ˙ λ S = 0 , which are 0 , λ − S
= (cid:16) 1 −√ 1 − 4 η(cid:17) 1 / (2 α ) and λ + S = (cid:16) 1+ √ 1 − 4 η(cid:17) 1 / (2 α ) .
As illustrated in
Figure 2 (bottom, left), if initialization δ > λ − S (good basin), λ S converges to a positive value λ + S
if 0 < δ < λ − S
(bad basin), λ S converges to zero.
Thresholding role of weight decay in feature learning:
While Tian et al.
(2021) shows why ncSSL does not collapse, one key question is how nc-SSL learns useful features and how the method determines which feature is learned.
Now it is clear: the weight decay factor η makes a call on what features should be learned.
Nuisance features subject to signiﬁcant change under data augmentation has larger variance σ 2 and 1 4(1+ σ 2 ) <
η , the eigenspace corresponds to this feature goes to zero; on the other hand, invariant features that are robust to data augmentation has much smaller σ 2 and4(1+ σ 2 ) >
η and these features are kept.
In our above analysis, B subspace corresponds to the nuisance features and collapses to zero; S subspace corresponds to the invariant features (whose variance was assumed as zero for simplicity) and is kept after training.
Figure 2 (middle and right) shows the spectrum of F (which is the correlation matrix of the predictor inputs) when the network is trained by DirectSet(1) with and without weight decay η on STL and CIFAR10: when η = 0 , the eigen-spectrum of F in later epochs does not have a sharp drop compared with the case of η = 0 . 0004 .
This means that the nuisance features are not signiﬁcantly suppressed when η = 0 .
Therefore, it is crucially important to choose weight decay appropriately: a too small η may not be sufﬁcient to suppress the nuisance features; a too large η can also collapse the invariant features.
As shown in Section 5, both cases lead to worse downstream performance.
3. G RADIENT D ESCENT ON E MPIRICAL L OSS
In this section, we then proceed to prove that DirectCopy (one special case of DirectSet( α ) with α = 1 ) successfully learns the projection matrix given polynomial number of samples.
Theorem 2.
Suppose network architecture and data distribution are as deﬁned in Assumptionand Assumption 2, respectively.
Suppose we initialize online network as δI, and run DirectCopy on empirical loss (see Eqn. 2) with γ step size and η weight decay.
Suppose the noise scale σis a positive constant, the weight decay coefﬁcient η ∈ (cid:16) 1+ σ 2 / 4 4(1+ σ 2 ) , 1+3 σ 2 / 4 4(1+ σ 2 ) (cid:17) and the initialization scale δ is a constant at least 1 / √ 2 .
Choose the step size γ as a small enough constant.
For any accuracy ˆ (cid:15) > 0 , given n ≥ poly ( d, 1 / ˆ (cid:15) ) number of samples, with probability at least 0 .
99 there exists t
= O (log(1 / ˆ (cid:15) ))
such that (here (cid:102) W t is the online network weights at the t -th step):
[Equation]
The proof proceeds by ﬁrst proving that gradient descent on the population loss converges in linear rate and then couples the gradient descent dynamics on empirical loss and that on population loss.
See the detailed proof in Appendix B.2.
3. S AMPLE C OMPLEXITY ON D OWNSTREAM T ASKS
In this section, we show that the learned representations can indeed reduce the sample complexity on the downstream tasks.
We consider the following data distribution for the down-stream task:
Assumption 3 (Downstream data distribution) .
Each input x ( i ) is sampled from N (0 , I d ) and its label y ( i ) =
(cid:10) x ( i ) , w ∗ (cid:11) + ξ ( i ) , where w ∗ is the ground truth vector with unit (cid:96) 2 norm and ξ ( i ) is independently sampled from N (0 , β 2 ) .
We assume the ground truth w ∗ lies on an r -dimensional subspace S and we denote the projection matrix on subspace S simply as P .
In practice, usually the semantically relevant features ( S subspace here) are invariant to augmentations and the nuisance features (orthogonal subspace of S ) have high variance under augmentations.
Therefore, by previous analysis, we expect DirectSet( α ) to learn the projection matrix
P.
Suppose { ( x ( i ) , y ( i ) ) } ni =1 are n training samples.
Each input x ( i ) is transformed by a matrix ˆ P ∈
R d × d (for example the learned online network W ) to get its representation ˆ Px ( i ) .
The regularized loss is then deﬁned as ˆ L
[Equation]
(cid:107) 2 .
In the below theorem, we show that when (cid:13)(cid:13)(cid:13) P − ˆ P (cid:13)(cid:13)(cid:13) F is small, the above ridge regression can recover the ground truth w ∗ given only O ( r ) number of samples.
Theorem 3.
Suppose the downstream data distribution is as deﬁned in Assumption 3.
Suppose (cid:13)(cid:13)(cid:13) ˆ P − P (cid:13)(cid:13)(cid:13) F ≤ ˆ (cid:15) with ˆ (cid:15) < 1 .
Choose the regularizer coefﬁcient ρ = ˆ (cid:15) 1 / 3 .
For any ζ < 1 / 2 , given n ≥ O ( r +log(1 /ζ )) number of samples, with probability at least 1 − ζ, the training loss minimizer ˆ w satisﬁes
[Equation]
In the above theorem, when n is at least O (cid:16) β 2 ( r +log
[Equation]
1 /ζ )) ˆ
[Equation]
/(cid:17) , we have (cid:13) (cid:13)(cid:13)
ˆ P ˆ w − w ∗
[Equation]
(cid:13)(cid:13) ≤ O (ˆ (cid:15) 1 / 3 ) .
Note that if we directly estimate ˆ w without transforming the inputs by ˆ P , we need Ω( d ) number of samples to ensure that (cid:107) ˆ w − w ∗ (cid:107) ≤ o (1) (Wainwright, 2019).
The proof of Theorem 3 follows from bounding the difference between ˆ P ˆ w and w ∗ by matrix concentration inequalities and matrix perturbation bounds.
The full proof is in Appendix B.3.
E MPIRICAL P ERFORMANCE OF D IRECT C OPY
In the previous analysis, we show DirectSet( α ), and in particular DirectCopy (DirectSet( α ) with α = 1 ), could recover the input feature structure with polynomial samples and make the downstream task more sample efﬁcient in a simple linear setting.
Compared with the original DirectPred (DirectSet( α ) with α = 1 / 2 ), DirectCopy is a simpler and computationally more efﬁcient algorithm since it directly set the predictor as the correlation matrix F , without the eigen-decomposition step.
By our analysis in Theorem 1, DirectCopy also learns the projection matrix P S with larger scale compared with DirectPred, which suggests that the invariant features learned by DirectCopy are stronger and more distinguishable.
Next, we show that DirectCopy is on par with (or even outperforms) the original DirectPred in various datasets, when coupling with deep nonlinear models on real datasets.
4. R ESULTS ON STL-10, CIFAR-10 AND CIFAR-
We use ResNet-18 (He et al., 2016) as the backbone network, a two-layer nonlinear MLP as the projector, and a linear predictor.
Unless speciﬁed otherwise, SGD is used as the optimizer with weight decay η = 0 .
To evaluate the quality of the pre-trained representations, we follow the linear evaluation protocol.
Each setting is repeated 5 times to compute the mean and standard deviation.
The accuracy is reported as “mean ± std”.
Unless explicitly speciﬁed, we use learning rate γ = 0 .
01 , regularization (cid:15) = 0 .
2 on STL-10; γ = 0 .
02 , (cid:15) = 0 .
3 on CIFAR-10 and γ = 0 .
03 , (cid:15) = 0 .on CIFAR-100.
See more detailed experiment settings in Appendix A.
[Table]
[Table]
[Table]
STL-10:
We evaluate the quality of the learned representation after each epoch, and report the best accuracy in the ﬁrst 100/300/500 epochs in Table 1.
DirectCopy achieves substantially better performance than the original DirectPred and SGD baseline, especially when trained with longer epochs.
DirectPred (freq=5) means the predictor is set by DirectPred every 5 batchs [Table] and is trained with gradient updates in other batchs [Table]
which outperforms DirectPred in later epochs
[Table] but is still much worse than DirectCopy. The SGD baseline is obtained by training the linear predictor using SGD.
CIFAR-10/100: For CIFAR-10, DirectCopy is slighly worse than DirectPred at epoch 100, but catches up and gets even better performance in epoch 300 and 500 (Table 1).
For CIFAR-100, at earlier epochs, the performance of DirectCopy is not as good as DirectPred, but the gap gradually diminishes in later epochs.
Both DirectCopy and DirectPred outperfoms the SGD baseline.
DirectPred (freq=5) achieves even better performance, but at the cost of a more complicated algorithm.
4. R ESULTS ON I MAGE N ET
Following BYOL (Grill et al., 2020), we use ResNet-50 as the backbone and a two-layer MLP as the projector.
We use LARS (You et al., 2017) optimizer and trains the model for 100 epochs.
See more detailed experiment settings in Appendix A.
For fairness, we compare DirectCopy to the gradient-based baseline which uses the same-sized linear predictor as ours.
As shown in Table 2, at 100-epoch, this baseline achieves 68.6 top-1 accuracy, which is already signiﬁcantly higher than BYOL with two-layer predictors reported in the literature (e.g., Chen & He (2020) reports 66.5 top-1 under 100-epoch training).
DirectCopy using normalized F accumulated with EMA µ = 0 .
99 on the correlation matrix, regularization parameter (cid:15) = {+
+}
.achieves 68.8 under the same setting, better than this strong baseline.
In contrast, DirectPred (Tian et al., 2021) achieves 68.5, slightly lower than the linear baseline.
O ̇𝜆 " 𝜆 "# 𝜆 "$
GoodBasin BadBasin
O ̇𝜆 "
GoodBasin
O ̇𝜆 " 𝜆 "$ 𝜆 "#
BadBasin
Increase 𝜖
Increase 𝜖 𝜆 " 𝜆 " 𝜆 "
A BLATION S TUDY
In this section, we study the inﬂuence of predictor regularization (cid:15) , normalization method, weight decay and degree α on the performance of DirectCopy.
Predictor regularization:
Table 3 shows that when the predictor regularization (cid:15) increases, the performance of DirectCopy on STL-10 and CIFAR-10 improves at ﬁrst and then deteriorates.
On
STL-10, DirectCopy with (cid:15) = 1 completely fails.
On CIFAR-10, although DirectCopy with (cid:15) =achieved reasonable performance at epoch 300, it’s still much worse than (cid:15) = 0 .
To better understand the role of (cid:15), we analyze the simple linear setting as in Section 3.1 while setting W p =
WW (cid:62) + (cid:15)I. Recall that λ B is the eigenvalue of W in B subspace and λ S is that in S subspace.
When the weight decay is appropriate, λ B still converges to zero.
On the other hand, the dynamics for λ S is as follows:
[Equation]
Increasing (cid:15) shifts the two positive stationary points λ − S , λ + S towards zero.
As illustrated in Figure(left), as (cid:15) increases, when λ + S is still positive, the good attraction basin expands, which means λ S can converge to a positive value from a smaller initialization; when λ + S shifts to zero, λ S converges to zero regardless the initialization size.
See the full analysis in Appendix D.
Intuitively
[Table]
a reasonable (cid:15) can alleviate representation collapse, but a too large (cid:15) also encourages representation collapse.
As shown in Figure 3 (right), when (cid:15) increases from zero, more eigenvalues of F becomes large; but when (cid:15) exceeds 0 . 3 , eigenvalues of F begin to collapse.
Normalization on F :
In our experiments
[Table]
we have been normalizing F by its spectral norm before adding the regularization: W p = F/ (cid:107) F (cid:107) + (cid:15)I.
It turns out that we can also normalize
F by its Frobenius norm or simply skip the normalization step.
In Table 4, we see comparable performance from DirectCopy with Frobenius normalization or no normalization, especially when trained longer
[Table] [Table]
Weight decay:
Table 5 shows that when weight decay η increases, the performance of DirectCopy improves at ﬁrst and then deteriorates.
This ﬁts our analysis on simple linear networks.
Basically, when the weight decay η increases, it can suppress the nuisance features more effectively
[Table]
but a too large weight decay also collapses the useful features
[Table]
[Table]
[Table]
Predictor degree: We compare DirectCopy against DirectSet( α ) with α = 2 , 1 / 2 , 1 / 4 .
Tableshows that DirectCopy outperforms other algorithms on STL-10.
On CIFAR-10, DirectCopy is slightly worse at epoch 100, but catches up in later epochs.
B
As a linear model used to study the behavior of nc-SSL, our model does not capture all of its intriguing empirical phenomena.
For example, we observed that the discarded nuisance features gradually come back after training over longer epochs.
Moreover, whether it comes back or not is related to the downstream task performance.
In Figure 4 on CIFAR-10 dataset, both DirectCopy and BYOL with two-layer nonlinear predictor show this resurgence of nuisance features, as well as strong performance, while BYOL with linear predictor does not seem to learn new features even when trained longer, which might explain its worse performance.
One conjecture is that at the beginning of training, weight decay prioritize the invariant features (i.e., low variance under augmentation) over nuisance ones.
The invariant features then grow, building their own supporting low-level features.
After that, the nuisance feature, which is also useful, are gradually picked up in later stage.
Since the low-level features are already trained through previous steps of back-propagation, the nuisance features are encouraged to use them as the supporting features, rather than creating their own.
In contrast, if we train both the invariant and nuisance features simultaneously, they will compete over the limited pool of low-level supporting features deﬁned by the capacity of the network, leading to worse learned representations.
We believe understanding these phenomena require analysis on the non-linear networks, and we leave it as future work.
C ONCLUSION
In this paper, we have proved DirectSet( α ) can learn the desirable projection matrix in a linear network setting and reduce the sample complexity on down-stream tasks.
Our analysis sheds light on the crucial role of weight decay in nc-SSL, which discards the features that have high variance under augmentations and keep the invariant features.
Inspired by the analysis, we also designed a simpler and more efﬁcient algorithm DirectCopy, which achieves comparable or even better performance than the original DirectPred (Tian et al., 2021) on various datasets.
We view our paper as an initial step towards demystifying the representation learning in nc-SSL.
Many mysteries lie beyond the explanation of the current theory and we leave them for future work.
A D ETAILED E XPERIMENT S ETTING
STL-10, CIFAR-10, CIFAR-100 : We use ResNet-18 (He et al., 2016) as the backbone network, a two-layer nonlinear MLP (with batch normalization, ReLU activation, hidden layer width 512, output width 128) as the projector, and a linear predictor.
Unless speciﬁed otherwise, SGD is used as the optimizer with momentum 0 .
9 , weight decay η = 0 .
0004 and batch size 128.
The EMA parameter for the target network is set as 0.996 and the EMA parameter µ of the correlation matrix ˆ F is set as 0.5.
Our code is adapted from Tian et al.
and we follow the same data augmentation process.
To evaluate the quality of the pre-trained representations, we follow the linear evaluation protocol.
Each setting is repeated 5 times to compute the mean and standard deviation.
The accuracy is reported as “mean ± std”.
Unless explicitly speciﬁed, we use learning rate γ = 0 .
01 , regularization (cid:15) = 0 .
2 on STL-10; γ = 0 .
02 , (cid:15) = 0 .
3 on CIFAR-10 and γ = 0 .
03 , (cid:15) = 0 .
3 on CIFAR-100.
ImageNet : Following BYOL (Grill et al., 2020), we use ResNet-50 as the backbone and a twolayer MLP (with batch normalization, ReLU, hidden layer width 4096, output width 256) as the projector.
We use LARS (You et al., 2017) optimizer and trains the model for 100 epochs, with a batch size 4096.
The learning rate is 7.2, which is linearly scaled from the base learning rate 0.
at batch size 256.
Other setups such as weight decay ( η = 1 e − 6 ), target EMA (scheduled from 0.
to 1), augmentation recipe (color jitters, blur, etc.), and linear evaluation protocol are the same as
BYOL.
B P ROOFS OF S INGLE - LAYER L INEAR N
ETWORKS
B. G RADIENT F LOW ON P OPULATION L OSS
In this section, we give the proof of Theorem 1, which shows that DirectSet( α ) running on the population loss with inﬁnitesimal learning rate and η weight decay can learn the projection matrix onto subspace S .
Theorem 1.
Suppose network architecture and data distribution follow Assumption 1 and Assumption 2, respectively.
Suppose we initialize online network W as δI, and run DirectSet ( α ) on population loss (see Eqn. 1) with inﬁnitesimal step size and η weight decay.
If we set the weight decay coefﬁcient η ∈ (cid:16) [-
 [Equation] 1+-]1 4(1+
[Equation]
2 ) , 14 (cid:17) and initialization scale δ > (cid:16) 1 −√ 1 − 4 η(cid:17) 1 /
[Equation]
[Equation]
) , then W converges to (cid:16) 1+ [Equation]
[Equation]
4 η(cid:17) 1 / (
[Equation]
) P S when time goes to inﬁnity.
As we already mentioned in the main text, Theorem 1 is proved by analyzing each eigenvalue of W separately.
We show that the eigenvalues in the B subspace converge to zero, and the eigenvalues in the S subspace converge to the same positive number, which immediately implies that W converges to a scaling of the projection matrix P S .
Proof of Theorem 1.
We can compute the gradient in terms of W as follows,
[Equation]
Note that the two augmented views x 1 , x 2 are sampled by ﬁrst sampling input x from N (0 , I d ) , and then independently sampling x 1 , x 2 from N ( x, σ 2 P B ) .
Therefore, we know E x 1 x 1 x
[Equation]
=
[Equation]
σ 2 P B and E x 1 ,x 2 x 2 x (cid:62) 1 = I. Recall that we run gradient ﬂow on W with weight decay η, so the dynamics on W is as follows:
[Equation]
where the ﬁrst term comes from the gradient and the second term is due to weight decay.
Since W is initialized as δI, and W a = W, W p = ( WW (cid:62) )
α , so we know initially W, W p , W a , I and
P B are all simultaneously diagonalizable, which then implies ˙
W is simultaneously diagonalizable with W .
This argument can continue to show that at any time point, W, W p , W a , I and P B are all simultaneously diagonalizable.
Since W is always a real symmetric matrix, we have W p
=
( WW (cid:62) )
α = | W | 2 α .
The dynamics on W can then be written as
[Equation]
Let the eigenvalue decomposition of W be (cid:80) di =
λ i u i u
(cid:62) i , with span ( { u d − r +1 , · · · , u d } ) equals to subspace B .
We can separately analyze the dynamics of each λ i .
Furthermore, we know λ 1 , · · · , λ r have the same value λ S and λ d − r +1 , · · · , λ d have the same value λ B .
Next, we separately show that λ B converge to zero and λ S converges to a positive value.
Dynamics for λ B : We can write down the dynamics for λ B as follows:
[Equation]
Similar as the analysis in Tian et al.
(2021), when η >
[Equation]
1+ σ 2 )
, we know ˙
λ B < 0 for any λ B >and λ B = 0 is a critical point.
This means, as long as η > 1 4(1+ σ
[Equation]
, λ B must converge to zero.
Dynamics for λ S : We can write down the dynamics for λ S as follows:
[Equation]
When 0 < η < 14 , we know ˙ λ S > 0 for λ 2 αS ∈ (cid:16) 1 −√ 1 − 4 η
[Equation]
1 − 4 η(cid:17) and ˙ λ S < 0 for λ 2 αS ∈ (cid:16) 1+ √ 1 − 4 η, ∞ (cid:17) .
Furthermore, we know ˙
λ S = 0 when λ 2 αS = 1+ √ 1 − 4 η.
Therefore, as long as 0 < η < 14 and initialization δ 2 α > 1 −√ 1 − 4 η, we know λ 2 αS converges to 1+ √ 1 − 4 η.
Overall, we know when
[Equation]
and δ
[Equation]
(cid:16) 1 −√ 1 −
[Equation]
(cid:17) 1 / (2 α ) , we have λ B converge to zero and λ S converge to (cid:16) 1+ √ 1 − 4 η(cid:17) 1 / (2 α ) .
That is, matrix W converges to (cid:16) 1+ √ 1 − 4 η(cid:17) 1 / (2 α ) P S .
(cid:3)
B. G RADIENT D ESCENT ON E MPIRICAL L OSS
In this section, we prove that DirectCopy successfully learns the projection matrix given polynomial number of samples.
Theorem 2.
Suppose network architecture and data distribution are as deﬁned in Assumptionand Assumption 2, respectively.
Suppose we initialize online network as δI, and run DirectCopy on empirical loss (see Eqn. 2) with γ step size and η weight decay.
Suppose the noise scale σis a positive constant, the weight decay coefﬁcient η ∈ (cid:16) 1+ σ 2 / 4 4(1+ σ 2 ) , 1+3 σ 2 / 4 4(1+ σ 2 ) (cid:17) and the initialization scale δ is a constant at least 1 / √ 2 .
Choose the step size γ as a small enough constant.
For any accuracy ˆ (cid:15) > 0 , given n ≥ poly ( d, 1 / ˆ (cid:15) ) number of samples, with probability at least 0 .
99 there exists t
= O (log(1 / ˆ (cid:15) ))
such that (here (cid:102) W t is the online network weights at the t -th step):
[Equation]
When running gradient descent on the empirical loss, the eigenspace of (cid:102) W t can shift and become no longer simultaneously diagonalizable with P B .
So we cannot independently analyze each eigenvalue of (cid:102) W t as before, which brings signiﬁcant challenge into the analysis.
Instead of directly analyzing the dynamics of (cid:102) W t , we ﬁrst show that the gradient descent iterates W t on the population loss converges to P S in linear rate, and then show that (cid:102) W t stays close to W t within certain iterations.
Lemma 1.
In the setting of Theorem 2, let W t be the gradient descent iterations on the population loss L .
Given any accuracy ˆ (cid:15) > 0 , for any t ≥ C log(1 / ˆ (cid:15) ) , we have
[Equation]
where C is a positive constant.
The proof of Lemma 1 is similar as the gradient ﬂow analysis in Section 3.2.
Next, we show that the gradient descent trajectory on the empirical loss stays close to the gradient descent trajectory on the population loss within O (log(1 / ˆ (cid:15) )) iterations.
Lemma 2.
In the setting of Theorem 2, let W t be the gradient descent iterations on the population loss and let (cid:102) W t be the gradient descent iterations on the empirical loss.
For any accuracy ˆ (cid:15) > 0 , given n ≥ poly ( d, 1 / ˆ (cid:15) ) number of samples, with probability at least 0 . 99 , for any t ≤ C log(1 / ˆ (cid:15) ) , we have
[Equation]
where the constant C comes from Lemma 1.
Then the proof of Theorem 2 directly follows from Lemma 1 and Lemma 2.
Proof of Theorem 2.
According to Lemma 1, we know given any accuracy ˆ (cid:15) (cid:48) , for t =
C log(1 / ˆ (cid:15) ) , we have
[Equation]
where C is a positive constant.
According to Lemma 2, we know given n ≥ poly ( d, 1 / ˆ (cid:15) (cid:48) )
number of samples, with probability at least 0 .
[Equation]
Therefore, we have
[Equation]
Replacing ˆ (cid:15) (cid:48) by ˆ (cid:15)/ 2 ﬁnishes the proof.
[Equation]
In section B.2.1, we give the proof of Lemma 1 and Lemma 2.
Proofs of some technical lemmas are left in Appendix B.5.
B.2. P ROOFS FOR L EMMA 1 AND L EMMA
Proof of Lemma 1.
Similar as in Theorem 1, we can show that at any step t , W t is simultaneously diagonalizable with W a,t , W p,t , I and P B .
The update on W t is as follows,
[Equation]
Let the eigenvalue decomposition of W t be (cid:80) di
=1 λ i,t u i u (cid:62) i , with span ( { u d − r +1 , · · · , u d } ) equals to subspace B .
We can separately analyze the dynamics of each λ i,t .
Furthermore, we know λ 1 ,t , · · · , λ r,t have the same value λ S,t and λ d − r +1 ,t , · · · , λ d,t have the same value λ B,t .
Next, we separately show that λ B,t converge to zero and λ S,t converges to a positive value in linear rate.
Dynamics of λ B,t : We show that
[Equation]
for any step size γ ≤ C 2 , where C 1 , C 2 are two positive constants.
According to the gradient update, we have
[Equation]
We only need to prove that for any λ B,t ∈ [0 , δ ] , we have
[Equation]
This is true since η ∈ (cid:16) 1+ σ 2 / 4 4(1+ σ 2 ) , 1+3 σ 2 / 4 4(1+ σ 2 ) (cid:17) and σ 2 , δ are two positive constants.
Dynamics of λ S : We show that
[Equation]
for any step size γ ≤ C 4 , where C 3 , C 4 are two positive constants.
There are two cases to consider: when the initialization scale δ 2 ∈ [1 / 2 , 1+ √ 1 − 4 η] , we prove
[Equation]
when the initialization scale δ 2 >
[Equation]
, we prove
[Equation]
We focus on the second case; the proof for the ﬁrst case is similar.
According to the gradient update, we have
[Equation]
We only need to show that λ S,t (cid:16) λ 2 S,t − 1 −√
1 − 4 η(cid:17) =
Θ(1) for any λ 2 S,t ∈ [ 1+ √ 1 − 4 η, δ ] .
This is true because η ∈ (cid:16) 1+ σ 2 / 4 4(1+ σ 2 ) , 1+3 σ 2 / 4 4(1+ σ 2 ) (cid:17) and σ 2 , δ are two positive constants.
Overall, we know that there exists constant step size such that after t
=
O (log(1 / ˆ (cid:15) )) steps, we have 0 ≤ λ B,t ≤ ˆ (cid:15) and (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) λ S,t − (cid:114) 1 +
[Equation]
−
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ≤ ˆ (cid:15).
This then implies,
[Equation]
Proof of Lemma 2.
We know the update on (cid:102) W t is
[Equation]
p,t (cid:32) − (cid:102) W p,t (cid:102) W t (cid:32)n n (cid:88) i =
x ( i ) 1 [ x ( i ) 1 ] (cid:62) (cid:33) + (cid:102) W a,t (cid:32)n n (cid:88)
i =
x ( i ) 1 [ x ( i ) 2 ] (cid:62) (cid:33)(cid:33) − γη (cid:102) W t , and the update on W t is
[Equation]
Next, we bound (cid:13) (cid:13)(cid:13)(cid:102) W t +1 − (cid:102) W t − ( W t +1 − W t ) (cid:13) (cid:13)(cid:13) .
According to Lemma 3, we know with probability at least 1 − O
( d 2 ) exp (cid:0) − Ω(ˆ (cid:15) (cid:48)
2 n/d 2 ) (cid:1) , (cid:13)(cid:13)(cid:13)(cid:13) (cid:13) n n (cid:88) i =
x ( i ) 1 [ x ( i ) 1 ] (cid:62) − I − σ 2 P B (cid:13)(cid:13)(cid:13)(cid:13) (cid:13) , (cid:13)(cid:13)(cid:13)(cid:13) (cid:13) n
[Equation]
(cid:88) i =
x (
[Equation]
[Equation] [Equation]
] (cid:62)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13) (cid:13) ,
[Equation]
n n
[Equation]
(cid:15)
[Equation]
[Table]
Combing the above
[Table]
[Table]
[Equation]
where C 1 is a constant larger than 1 .
So for any t ≤ C log(1 / ˆ (cid:15) ) , we have
[Equation]
for some positive constant C 2 .
Choosing ˆ (cid:15) (cid:48)
= ˆ (cid:15) C [-
-]2 +1 , we know as long as n ≥ poly ( d, 1 / ˆ (cid:15) ) , with probability at least 0 .
99 , for any t ≤ C log(1 / ˆ (cid:15) ) , we have
[Equation]
B. S AMPLE C OMPLEXITY ON D OWN - STREAM T ASKS
In this section, we give a proof for Theorem 3, which shows that the learned representations can indeed reduce sample complexity in downstream tasks.
Theorem 3.
Suppose the downstream data distribution is as deﬁned in Assumption 3.
Suppose (cid:13)(cid:13)(cid:13) ˆ P − P (cid:13)(cid:13)(cid:13) F ≤ ˆ (cid:15) with ˆ (cid:15) < 1 .
Choose the regularizer coefﬁcient ρ = ˆ (cid:15) 1 / 3 .
For any ζ < 1 / 2 , given n ≥ O ( r +log(1 /ζ )) number of samples, with probability at least 1 − ζ, the training loss minimizer ˆ w satisﬁes
[Equation]
Suppose { ( x ( i ) , y ( i ) ) } ni =1 are n training samples in the downstream task, let X ∈ R n × d be the data matrix with its i -th row equal to x ( i ) .
Denote y ∈ R n as the label vector with its i -th entry as y ( i ) .
Each input x ( i ) is transformed by a matrix ˆ
P ∈ R d × d to get its representation ˆ Px ( i ) .
The regularized loss can be written as
[Equation]
This is the ridge regression problem on inputs { ( ˆ Px ( i ) , y ( i ) ) } ni =1 , and the unique global minimizer ˆ w has the following close form:
[Equation]
With the above closed form of ˆ w , the proof of Theorem 3 follows by bounding the difference between ˆ P ˆ w and w ∗ by matrix concentration inequalities and matrix perturbation bounds.
Some proofs of technical lemmas are left in Appendix B.5.
Proof of Theorem 3.
Denoting ˆ P as P + ∆ , we know (cid:107) ∆ (cid:107) F ≤ ˆ (cid:15) by assumption.
We can also write y as Xw ∗ + ξ where ξ ∈ R n is the noise vector with its i -th entry equal to ξ ( i ) .
Then, we can divide ˆ w into two terms,
[Equation]
Let’s ﬁrst give an upper bound for the second term that comes from the error term ∆ (cid:62) .
Upper bounding (cid:13)(cid:13)(cid:13)(cid:13)(cid:16) 1 n ˆ P (cid:62) X (cid:62) X ˆ P + ρI (cid:17) − 1 1 n ∆ (cid:62) X (cid:62) ( Xw ∗ + ξ ) (cid:13)(cid:13)(cid:13)(cid:13)
We ﬁrst bound the norm of 1 n ∆ (cid:62) X (cid:62)
Xw ∗ .
According to Lemma 5, we know with probability at least 1 − exp( − Ω( n )) , (cid:13)(cid:13)(cid:13) 1 √ n ∆ (cid:62) X (cid:62) (cid:13)(cid:13)(cid:13) F ≤ O (ˆ (cid:15) ) .
Since Xw ∗ is a standard Gaussian vector with dimension n , according to Lemma 8, with probability at least 1 − exp( − Ω( n )) , (cid:13)(cid:13)(cid:13) 1 √ n Xw ∗ (cid:13)(cid:13)(cid:13)
≤ O (1) .
Therefore, we have (cid:13)(cid:13) [Equation]
(cid:13)(cid:13) ≤ O (ˆ (cid:15) ) .
Then we bound the norm of [Equation]
According to Lemma 8, we know with probability at least 1 − exp( − Ω( n )) , (cid:13)(cid:13)(cid:13) 1 √ n ξ (cid:13)(cid:13)(cid:13) ≤
O ( β ) .
According to Lemma 6, we know with probability at least 1 − ζ/ 3 , (cid:13)(cid:13) ∆ (cid:62) X (cid:62)
¯ ξ (cid:13)(cid:13) ≤ O (cid:16) ˆ (cid:15) (cid:112) log(1 /ζ ) (cid:17) .
Therefore, we have (cid:13)(cid:13)
[Equation]
Since λ min (cid:16) 1 n ˆ P (cid:62) X (cid:62) X ˆ P + ρI (cid:17) ≥ ρ, we have (cid:13)(cid:13)(cid:13) (cid:13) (cid:16) 1 n ˆ P (cid:62) X (cid:62) X ˆ P + ρI (cid:17) − 1 (cid:13)(cid:13)(cid:13) (cid:13) ≤ 1 ρ .
Combining with above bound on (cid:13) (cid:13) 1 n ∆ (cid:62) X (cid:62) ( Xw ∗ + ξ ) (cid:13) (cid:13) , we know with probability at least 1 − exp( − Ω( n ))
− ζ/ 3 ,
[Equation]
Analyzing (cid:16) 1 n ˆ P (cid:62) X (cid:62)
X ˆ P
[Equation]
ρI
[Equation]
1 1 n P (cid:62) X (cid:62) ( Xw ∗ + ξ )
We can write 1 n ˆ P
[Equation]
X ˆ P as 1 n P (cid:62) X (cid:62) XP + E, where
[Equation]
Let’s ﬁrst bound the spectral norm of XP.
Since P is a projection matrix on an r -dimensional subspace S, we can write P as UU (cid:62) , where
U ∈ R
d × r has columns as an orthonormal basis of subspace S.
According to Lemma 4, we know with probability at least 1 − exp( − Ω( n )) ,
[Equation]
Since (cid:107) U (cid:107) ≤ 1 , we have (cid:13)(cid:13)(cid:13) 1 √ n XP (cid:13)(cid:13)(cid:13) =
(cid:13)(cid:13)(cid:13) 1 √ n XUU (cid:62) (cid:13)(cid:13)(cid:13) ≤ O (1) .
According to Lemma 5, we know with probability at least 1 − exp( − Ω( n )) , [Equation] F [Equation]
O (ˆ (cid:15) [Equation]
So overall, we know (cid:107) E (cid:107) ≤ (cid:107) E (cid:107) F ≤ O (ˆ (cid:15) ) .
Then, we can write (cid:18)
[Equation]
According to the perturbation bound for matrix inverse (Lemma 11)
[Equation]
we have (cid:107) F (cid:107) ≤ O ( ˆ (cid:15)ρ 2 ) .
Then, we have
[Equation] [Equation]
We ﬁrst show that the ﬁrst term is close to w ∗ .
Let the eigenvalue decomposition of 1 n P (cid:62) X (cid:62) XP be V Σ V (cid:62) , where V ’s columns are an orthonormal basis for subspace S .
Here Σ ∈ R r × r is the diagonal matrix that contains all the eigenvalues of 1 n P (cid:62) X (cid:62) XP .
According to Lemma 4, we know that with probability at least 1 − exp( − Ω( n )) , all the non-zero eigenvalues of 1 n P (cid:62) X (cid:62) XP are Θ(1) .
Then, it’s not hard to show that
[Equation]
This immediately implies that
[Equation]
Next, we bound the norm of the second term F 1 n P (cid:62) X (cid:62)
Xw ∗ .
Similar as before, we know with probability at least 1 − exp( − Ω( n )) , (cid:13)(cid:13)(cid:13) 1 √ n Xw ∗ (cid:13)(cid:13)(cid:13)
≤ O (1) and (cid:13)(cid:13)(cid:13) 1 √ n P
[Table]
X
[Table]
≤ O (1) .
Therefore, we have
[Table]
[Equation]
[Table]
[Equation]
[Table]
Finally, let’s bound the third term (cid:16)(cid:0) 1 n P (cid:62) X (cid:62) XP + ρI (cid:1) − 1 + F (cid:17)
1 n P (cid:62) X (cid:62) ξ.
We ﬁrst bound the norm of 1 n P (cid:62) X (cid:62) ξ.
with probability at least 1 − exp( − Ω( n )) , we know (cid:107) ξ (cid:107) ≤ 2 β √ n.
Therefore, we know (cid:13)(cid:13) 1 n P (cid:62) X (cid:62) ξ (cid:13)(cid:13) ≤
O ( β/ √ n ) (cid:13)(cid:13) P (cid:62) X (cid:62) ¯ ξ (cid:13)(cid:13) , where ¯ ξ = ξ/ (cid:107) ξ (cid:107) .
According to Lemma 7, with probability at least 1 − ζ/ 3 , we have (cid:13)(cid:13) P (cid:62) X (cid:62) ¯ ξ (cid:13)(cid:13) ≤ √ r + O ( [Equation]
log(1 /ζ )) . Overall, with probability at least 1 − exp
[Equation] [Equation]
n
[Equation]
It’s not hard to verify that for any vector v ∈
R d in the subspace S, we have (cid:13)(cid:13)(cid:13)(cid:16)(cid:0) 1 n P (cid:62) X (cid:62)
[Equation]
(cid:1) − 1 + F (cid:17) v (cid:13)(cid:13)(cid:13)
≤ O ( (cid:107) v (cid:107) ) .
Since 1 n P (cid:62) X (cid:62) ξ lies on subspace S, we have
[Equation]
Combining the above analysis and taking a union bound over all the events, we know with probability at least 1 − exp( − Ω( n )) − 2 ζ/ 3 ,
[Equation]
∗
[Equation]
= O
[Equation]
Suppose n ≥ O (log(1 /ζ )) and setting ρ = ˆ (cid:15) 1 / 3 , we further have with probability at least 1 − ζ,
[Equation]
where the last inequality assumes ˆ (cid:15) < 1 .
We can also bound (cid:13)(cid:13)(cid:13) ˆ P ˆ w − w ∗ (cid:13)(cid:13)(cid:13) as follows,
[Equation]
B. A NALYSIS WITH W p := ( W E x 1 x 1 x (cid:62) 1 W (cid:62) ) α
In this section
[Equation]
we prove that DirectSet( α ) can also learn the projection matrix when we set W p :=
( W E x 1 x 1 x (cid:62) 1 W (cid:62) ) α .
For the network architecture and data distribution, we follow exactly the same setting as in Section 3.2.
Therefore, we know W p :=
( W E x 1 x 1 x (cid:62) 1 W (cid:62) )
α = ( W ( I + σ 2 P B ) W (cid:62) )
α .
Theorem 4.
Suppose network architecture and data distribution are as deﬁned in Assumption 1 and
Assumption 2, respectively.
Suppose we initialize online network W as δI, and run DirectPred ( α ) on population loss (see Eqn. 1) with inﬁnitesimal step size and η weight decay.
Suppose we set W a =
W and W p =
( W E x 1 x 1 x (cid:62) 1 W (cid:62) ) α .
Assuming the weight decay coefﬁcient η ∈ (cid:16) [-1 4(1+ σ 2 ) 
-][Equation]
4(1+ σ 2 ) 1+2 α , 14 (cid:17) and initialization scale δ > (cid:16) 1 −√ 1 − 4 η(cid:17) 1 / (2 α ) , we know W converges to (cid:16) 1+ √ 1 − 4 η(cid:17) 1 / (2 α ) P S when time goes to inﬁnity.
The only difference from Theorem 4 is that now the initialization δ is only required to be larger than4(1+ σ 2 )
+2 α .
The proof is almost the same as in Theorem 1.
Proof of Theorem 4.
Similar as in the proof of Theorem 1, we can write the dynamics on W is as follows:
[Equation]
Dynamics for λ B : We can write down the dynamics for λ B as follows:
[Equation]
When η > 1 4(1+ σ 2 )
+2 α , we know ˙
λ B < 0 for any λ B > 0 and λ B = 0 is a critical point.
This means, as long as η > 1 4(1+ σ 2 )
+2 α , λ B must converge to zero.
Dynamics for λ S :
The dynamics is same as when setting W p
= ( WW (cid:62)
[Equation]
[Equation]
so when 0 < η < 1 4 and initialization δ 2 α > 1 − √ 1 − 4 η, we know λ 2 α S converges to 1+ √ 1 − 4 η.
Overall, we know when
[Equation]
(cid:17) 1 / (2 α ) , we have λ B converge to zero and λ S converge to (cid:16) 1+ √ 1 − 4 η(cid:17) 1 / (2 α ) .
That is, matrix W converges to (cid:16) 1+ √ 1 − 4 η(cid:17) 1 / (2 α ) P S .
(cid:3)
[Table]
Lemma 3.
Suppose { x ( i ) , x ( i ) 1 , x ( i ) 2 } ni =1 are sampled as decribed in Section 3.
Suppose n ≥
O ( d/ ˆ (cid:15) 2 ) , with probability at least 1 − O ( d 2 ) exp (cid:0) − Ω(ˆ (cid:15) 2 n/d 2 ) (cid:1) , we have
[Table]
n
[Table]
Proof of Lemma 3.
For each x ( i ) 1 , we can write it as x ( i ) +
z ( i ) 1 where x ( i ) ∼ N (0 , I ) and z ( i ) 1 ∼
N (0 , σ 2 P B ) .
So we have
[Table]
[Equation]
According to Lemma 9, we know as long as n ≥
O ( d/ ˆ (cid:15) 2 ) , with probability at least 1 − exp( − Ω(ˆ (cid:15) 2 n )) ,
[Equation]
Similarly, with probability at least 1 − exp( − Ω(ˆ (cid:15) 2 n )) ,
[-z ( i ) 
-][Equation] n
[Equation]
Next we bound (cid:13)(cid:13)(cid:13) 1 n (cid:80) ni
=1 x ( i )
[ z ( i ) 1 ] (cid:62) (cid:13)(cid:13)(cid:13) .
We know each entry in matrix 1 n (cid:80) ni
=1 x ( i )
[ z ( i ) 1 ] (cid:62) is the average of n zero-mean O (1)
-subexponential independent random variables.
Therefore, according to the Bernstein’s inequality, for any ﬁxed entry ( k, l ) , with probability at least 1 − exp (cid:0) − ˆ (cid:15) 2 n/d 2 (cid:1) , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:34)n n (cid:88) i = x ( i ) [ z ( i ) 1 ] (cid:62) (cid:35) k,l (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ≤ ˆ (cid:15)/d.
Taking a union bound over all the entries, we know with probability at least 1 − d 2 exp (cid:0) − ˆ (cid:15) 2 n/d 2 (cid:1) ,
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)n n (cid:88) i =
x ( i )
[ z ( i ) 1 ] (cid:62) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 n n (cid:88)
i =
x ( i )
[ z ( i ) 1 ] (cid:62) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
F ≤ ˆ (cid:15).
The same analysis also applies to (cid:13) (cid:13)(cid:13) 1 n (cid:80) n i =
1 z
( i )
[ x ( i ) ] (cid:62) (cid:13) (cid:13)(cid:13) .
Combing all the bounds, we know with probability at least 1 − O
( d 2 ) exp (cid:0) − Ω(ˆ (cid:15) 2 n/d 2 ) (cid:1) , (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)n n (cid:88) i
= x ( i ) 1 [ x ( i ) 1 ] (cid:62) − I − σ 2 P B (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
Similarly, we can prove that with probability at least 1 − O
( d 2 ) exp (cid:0) − Ω(ˆ (cid:15) 2 n/d 2 ) (cid:1) ,
[Equation]
Changing ˆ (cid:15) to ˆ (cid:15) (cid:48) / 4 ﬁnishes the proof.
[Equation]
Lemma 4.
Let X ∈ R n × d be a standard Gaussian matrix, and let U ∈
R
d × r be a matrix with orthonormal columns.
Suppose n ≥ 2 r , with probability at least 1 − exp( − Ω( n )) , we know
[Equation]
Proof of Lemma 4.
Since U has orthonormal columns, we know XU is a n × r matrix with each entry independently sampled from N (0 , 1) .
According to Lemma 9, we know when n ≥ 2 r, with probability at least 1 − exp( − Ω( n )) ,
[Equation]
This immediately implies that
[Equation]
Lemma 5.
Let ∆ be a d × d matrix with Frobenius norm ˆ (cid:15) , and let X be a n × d standard Gaussian matrix.
We know with probability at least 1 − exp( − Ω( n )) ,
[Equation]
Proof of Lemma 5.
Let the singular value decomposition of ∆ be U Σ V (cid:62) , where U, V have orthonormal columns and Σ is a diagonal matrix with diagonals equal to singular values σ i ’s. Since (cid:107) ∆ (cid:107) F = ˆ (cid:15), we know (cid:80) di
=1 σ 2 i = ˆ (cid:15) 2 .
Since U is an orthonormal matrix, we know ˆ X :=
XU is still an n × d standard Gaussian matrix.
Next, we bound the Frobenius norm of (cid:101) X := ˆ X Σ .
It’s not hard to verify that all the entries in (cid:101)
X are independent Gaussian variables and (cid:101)
X ij ∼ N (0 , σ 2 j ) .
According to the Bernstein’s inequality for sum of independent and sub-exponential random variables, we have for every t > 0 ,
[Equation]
Since (cid:80) dj =1 σ 2 j = (cid:107) ∆ (cid:107) 2 F = ˆ (cid:15) 2 , we know max j ∈
[ d ] σ 2 j ≤ ˆ (cid:15) 2 .
We also have (cid:80) j ∈
[ d ] σ 4 j ≤ (cid:16)(cid:80) j ∈ [ d ] σ 2 j (cid:17) 2 = ˆ (cid:15) 4 .
Therefore, we have
[Equation]
Replacing t by n ˆ (cid:15) 2 , we concluded that with probability at least 1 − exp( − Ω( n )) , (cid:13)(cid:13)(cid:13) (cid:101)
X
(cid:13)(cid:13)(cid:13)
[Equation]
≤ 2 n ˆ (cid:15) 2 .
Furthermore, since (cid:13) (cid:13) V (cid:62) (cid:13) (cid:13) =
1 , we have
[Equation]
Lemma 6.
Let ∆ (cid:62) be a d × d matrix with Frebenius norm ˆ (cid:15) and let X (cid:62) be a d × n standard Gaussian matrix.
Let ¯ ξ be a unit vector with dimension n .
We know with probability at least 1 − ζ/ 3 ,
[Equation]
ξ
[Equation]
≤
[Equation]
ˆ
[Equation]
log
[Equation] [Equation]
Proof of Lemma 6.
Let the sigular value decomposition of ∆ (cid:62) be U Σ V (cid:62) .
We know X (cid:62) ¯ ξ is a ddimensional standard Gaussian vector.
Further, we know V (cid:62) X (cid:62) ¯ ξ is also a d -dimensional standard
Gaussian vector.
So Σ V (cid:62) X (cid:62) ¯ ξ has independent Gaussian entries with its i -th entry distributed as N (0 , σ 2 i ) .
According to the Bernstein’s inequality for sum of independent and sub-exponential random variables, we have for every t > 0 ,
[Equation]
Choosing t as O (ˆ (cid:15) 2 log(1 /ζ )) , we know with probability at least 1 − ζ/ 3 , we have
[Equation]
Since (cid:107) U (cid:107) = 1 , we further have
[Equation]
Lemma 7.
Let P ∈ R
d
×
d be a projection matrix on a r -dimensional subspace, and let ¯ ξ be a unit vector in R d .
Let X (cid:62) be a d × n standard Gaussian matrix that is independent with P and ξ .
With probability at least 1 − ζ/ 3 , we have
[Equation]
Proof of Lemma 7.
Since P is a projection matrix on an r -dimensional subspace, we can write P as UU (cid:62) , where
U ∈ R
d × r has orthonormal columns.
We know U (cid:62) X (cid:62) is still a standard Gaussian matrix with dimension
r × n.
Furthermore, U (cid:62) X (cid:62) ¯ ξ is an r -dimensional standard Gaussian vector.
According to Lemma 8, with probability at least 1 − ζ/ 3 , we have
[Equation]
Since (cid:107) U (cid:107) = 1 , we further have
[Equation]
C A NALYSIS OF D EEP L INEAR N ETWORKS
In this section, we extend the analysis in Section 3.2 to deep linear networks.
We consider the same data distribution as deﬁned in Assumption 2.
We consider the following network,
Assumption 4 (Deep linear network) .
The online network is an l -layer linear networks
W l W l − 1 · · · W 1 with each W i ∈ R
d × d .
The target network has the same architecture with weight matrices W a,l W a,l − 1 · · · W a, 1 .
For convenience, we denote W as W l W l − 1 · · · W 1 and denote W a as W a,l W a,l − 1 · · · W a, 1 .
Training procedure:
At the initialization, we initialize each W i as δ 1 /l I d .
Through the training, we ﬁx W p as (cid:0) WW (cid:62) (cid:1) α and ﬁx each W a,i as W i .
We run gradient ﬂow on every W i with weight decay η.
The population loss is
L ( { W i } , W p , { W a,i }
[Equation]
Theorem 5.
Suppose the data distribution and network architecture satisﬁes Assumption 2 and Assumption 4, respectively.
Suppose we train the network as described above.
Assuming the weight decay coefﬁcient
[Equation]
, and initialization scale δ ≥ [Equation] (cid:17) 12 α , we know W converges to cP S as time goes to inﬁnity, where c is a positive number within (cid:18)(cid:16) 2 αl +2 l − 2 4 αl +2 l − 2 (cid:17) 12 α , 1 (cid:19) .
Similar as in the setting of single-layer linear networks, we prove Theorem 5 by analyzing the dynamics of the eigenvalues of W. Note that with constant α , the upper/lower bounds for η and scalar c in the Theorem are always constants no matter how large l is.
Proof of Theorem 5.
For j ≥ i , we use W [ j : i ] to denote W j W j − 1 · · · W i and for j <
i have
W [ j : i ] = I.
We use similar notations for W a, [ j : i ] .
For each W i , we can compute its dynamics as follows:
[Equation]
(cid:1) (cid:0) W [
[Table]
− [Table] ] (cid:1) (cid:62) + (cid:0) W p W a, [ l [Table]
i
[Table]
] (cid:1) (cid:62) W [Table] (cid:0) W a, [ [Table]
] (cid:1)
[Table]
− ηW
[Table]
It’s clear that through the training
all W
i ’s remains the same and they are simultaneously diagonalizable with W p , I and P B .
We also have W a = W and W p =
| W | 2 α .
Since we will ensure that W is always positive semi-deﬁnite so W p = | W | 2 α = W 2 α = W 2 αl i .
So the dynamics for each W i can be simpliﬁed as follows:
[Equation]
Let the eigenvalue decomposition of W i be (cid:80) di
=
ν i u i u (cid:62)
i , with span ( { u d − r +1 , · · · , u d } ) equals to subspace B .
We can separately analyze the dynamics of each ν i .
Furthermore, we know ν 1 , · · · , ν r have the same value ν S and ν d − r +1 , · · · , ν d have the same value ν B .
We can write down the dynamics for ν S and ν B as follows,
[Equation]
Let λ S be the eigenvalue of W corresponding to eigen-directions u 1 , · · · , u r , and let λ B be the eigenvalue of W corresponding to eigen-directions u d − r +1 , · · · , u d .
We know λ S
=
ν lS and λ B = ν l B .
So we can write down the dynamics for λ B as follows,
[Equation]
and similarly for λ S we have
[Equation]
Dynamics for λ B : We can write the dynamics on λ B as follows,
[Equation]
where g ( λ B ) := − λ 4 α +2 −
2 /l B
[Equation]
− η.
We show that when η is large enough, g ( λ B ) is negative for any positive λ B .
We compute the maximum value of g ( λ B ) for λ B > 0 .
We ﬁrst compute the derivative of g as follows:
[Equation]
It’s clear that g (cid:48) ( λ B ) > 0 for λ 2 αB ∈ (0 ,
[Equation]
2 l − 2 (4 αl +2 l − 2)(1+ σ 2 ) ) and g (cid:48) ( λ B ) < 0 for λ 2 αB ∈ ( 2 αl +2 l − 2 (4 αl +2 l − 2)(1+ σ 2 ) , + ∞ ) .
Therefore, the maximum value of g ( λ B ) for positive λ B takes at λ ∗ B = (cid:16) 2 αl +2 l − 2 (4 αl +2 l − 2)(1+ σ 2 ) (cid:17) 12 α and g ( λ ∗ B )
= − (cid:18)
2 αl + 2 l −(4 αl + 2 l − 2)(1 +
[Equation]
2 ) (cid:19) 2+
α − 1 αl (1 + σ 2 ) + (cid:18) 2 αl + 2 l −(4 αl + 2 l − 2)(1 + σ 2 ) (cid:19) 1+ 1 α − 1 αl − η = 2 αl
2 αl + 2 l − 2) 1+ 1 α − 1 αl (4 αl + 2 l − 2) 2+ 1 α − 1 αl (1 + σ 2 ) 1+ 1 α − 1 αl − η.
As long as η >
[Equation]
, we know g ( λ B ) < 0 for any λ B > 0 , which further implies that ˙ λ B < 0 for any λ B > 0 .
So λ B converges to zero.
Dynamics for λ S : We can write down the dynamics on λ S as follows,
[Equation]
S
[Equation]
, where h ( λ S ) = − λ 4 α +2 − 2 /l S + λ 2 α +2 − 2 /l S − η.
We compute the derivative of h as follows:
[Equation]
+
[Equation]
α
[Equation]
2 −
[Equation]
l
[Equation]
(cid:1) .
So h ( λ S ) is increasing in (0 , (cid:16) 2 αl +2 l − 2 4 αl +2 l − 2 (cid:17) 12 α ) and is decreasing in ( (cid:16) 2 αl +2 l − 2 4 αl +2 l − 2 (cid:17) 12 α , ∞ ) .
The maximum value of h for positive λ S takes at λ ∗ S = (cid:16) 2 αl +2 l − 2 4 αl +2 l − 2 (cid:17) 12 α
and we have
[Equation]
−
[Equation]
As long as η < 2 αl (2 αl +2 l − 2)
[Equation]
1 α − 1 αl (
[Equation]
− 2)
+ 1 α − 1 αl , we have h ( λ ∗ S ) > 0 .
Furthermore, since h is increasing in (0 , λ ∗ S ) and is decreasing in ( λ ∗ S , ∞ ) and h (0) , h ( ∞ ) < 0 , we know there exists λ − S ∈
(0 , λ ∗ S ) , λ + S ∈ ( λ ∗ S , ∞ ) such that h ( λ S ) < 0 in (0 , λ − S ) , h ( λ S ) > 0 in ( λ − S , λ + S ) and h ( λ S ) < 0 in ( λ + S , ∞ ) .
Therefore, as long as δ ≥ λ ∗ S > λ − S , we have λ S converges to λ + S .
Since h (1) < 0 , we know λ + S ∈ ( (cid:16) 2 αl +2 l − 2 4 αl +2 l − 2 (cid:17) 12 α , 1) .
Overall as long as η ∈ (cid:18) 2 αl
[Equation]
, 2 αl (2 αl +2 l − 2) 1+ 1 α − 1 αl (4 αl +2 l − 2) 2+ 1 α − 1 αl (cid:19) , we know W converges to cP S , where c is a positive number within ( (cid:16) 2 αl +2 l − 2 4 αl +2 l − 2 (cid:17) 12 α , 1) .
(cid:3)
D A NALYSIS OF P REDICTOR R EGULARIZATION .
In this section, we study the inﬂuence of predictor regularization in a simple linear setting.
In particular, we consider the same setting as in Section 3.2 except that we set W p :=
( WW (cid:62) )
α + (cid:15)I.
Theorem 6.
In the setting of Theorem 1 except that we set W p = ( WW (cid:62) )
α + (cid:15)I.
We have
• when (cid:15) ∈ [0 , 1+ √ 1 − 4 η) , as long as δ > (cid:16) max (cid:16) 1 −√ 1 − 4 η− (cid:15), 0 (cid:17)(cid:17) 12 α , we have W converges to (cid:16) 1+ √ 1 − 4 η− (cid:15) (cid:17) 12 α P S ; • when (cid:15) ≥ 1+ √ 1 − 4 η, W always converges to zero.
Proof of Theorem 6.
We can write the dynamics of W as follows,
[Equation]
Let the eigenvalue decomposition of W be (cid:80) d i =
λ i u i u
(cid:62) i , with span ( { u d − r +1 , · · · , u d } ) equals to subspace B .
We can separately analyze the dynamics of each λ i .
Furthermore, we know λ 1 , · · · , λ r have the same value λ S and λ d − r +1 , · · · , λ d have the same value λ B .
Dynamics for λ B : We can write down the dynamics for λ B as follows:
[Equation]
When η > 1 4(1+ σ 2 ) , we still know ˙
λ B < 0 for any λ B > 0 and λ B = 0 is a critical point.
So λ B converges to zero.
Dynamics for λ S : We can write down the dynamics for λ S as follows:
[Equation]
where the second inequality assumes 0 < η < 1 4 .
We have
• when (cid:15) ∈ [0 , 1+ √ 1 − 4 η) , as long as δ > (cid:16) max (cid:16) 1 −√ 1 − 4 η− (cid:15), 0 (cid:17)(cid:17) 12 α , we have λ S converges to (cid:16) 1+ √ 1 − 4 η− (cid:15) (cid:17) 12 α > 0 ; • when (cid:15) ≥ 1+ √ 1 − 4 η, λ S always converges to zero.
[Equation]
E T ECHNICAL T OOLS E. N ORM OF R ANDOM V ECTORS
The following lemma shows that a standard Gaussian vector with dimension n has (cid:96) 2 norm concentrated at √ n .
Lemma 8 (Theorem 3.1.1 in Vershynin (2018)) .
Let X =
( X 1 , X 2 , · · · , X n ) ∈ R n be a random vector with each entry independently sampled from N (0 , 1) .
Then
[Equation]
where C is an absolute constant.
E. S INGULAR V ALUES OF G AUSSIAN M ATRICES
The following lemma shows a tall random Gaussian matrix is well-conditioned with high probability.
Lemma 9 (Corollary 5.35 in Vershynin (2010)) .
Let A be an N × n matrix whose entries are independent standard normal random variables.
Then for every t ≥ 0 with probability at least 1 − 2 exp( − t 2 / 2) one has
[Equation]
E. P ERTURBATION B OUND FOR M ATRIX P SEUDO - INVERSE
With a lowerbound on σ min ( A ) , we can get bounds for the perturbation of pseudo-inverse.
Lemma 10 (Theorem 3.4 in Stewart (1977)) .
Consider the perturbation of a matrix A ∈ R m × n :
B =
A + E. Assume that rank
( A ) =
rank ( B ) = n, then
[Equation]
The following corollary is particularly useful for us.
Lemma 11 (Lemma G.8 in Ge et al.
Consider the perturbation of a matrix A ∈ R m × n :
B =
A + E where (cid:107) E (cid:107) ≤ σ min ( A ) / 2 .
Assume that rank ( A ) =
rank ( B ) = n, then
[Equation]