summaries: The paper provides a new method for self-supervised learning (SSL), called DirectCopy, based on a previous work DirectPred. An important contribution of this work is theoretical analyses on DirectSet($\alpha$), a theoretical model on linear neural networks that work with arbitrary $\alpha$. The authors prove that the weight decay coefficient $\eta$ has the ability to filter out unnecessary features and hence helps learn useful representations. DirectCopy is a special case of DirectSet($\alpha$) when $\alpha=1$. The new algorithm does not require to compute the burdensome eigen-decomposition while enjoying similar empirical performance as DirectPred.  In this paper, the authors make theoretical progress on understanding non-contrastive self-supervised learning (ncSSL).  ncSSL has previously demonstrated strong empirical performance, even outperforming contrastive learning, but the theory behind it is still unclear.  In this work, the authors build off of prior analysis by [1], and showcase the role weight decay has in learning a desirable representation; it acts as a threshold that discards noisy features with high variance introduced by the data augmentation, and keeps stable features with low variance.    [1] Tian et al., Understanding Self-Supervised Learning Dynamics without Contrastive Pairs, 2019. This paper attempts to investigate the reasons behind why non-contrastive SSL methods such as BYOL and Sim-Siam do not collapse to trivial solutions, how they learn representations that are related to the data distribution and augmentation process, and how these reduce the sample-complexity of downstream tasks. It heavily draws upon DirectPred (Tian et al) and generalizes the method for directly setting the weights of the predictor network via a parameter that seems to be tied to the strength and distinguishability of the learnt features.   Tian et al: https://arxiv.org/abs/2102.06810 The paper "Towards demystifying representation learning with non-contrastive self-supervision" presents and analyzes a family of algorithms, DirectSet(\alpha), for non-conctrastive self-supervised learning with only positive pairs. The theoretical analysis assumes linear layers, and focuses on a special data distribution assumption, where the input space is separated in two linear subspaces, one invariant under the 'data augmentations', the other being the complement. It is then shown that the proposed algorithm converges to the projection matrix onto the invariant sub-space. Further it is shown theoretically that this has the down-stream advantage reduced sample complexity for learning on this representation. Empirically, it is shown that the method performs on par, or sometimes slightly better, than the previously proposed closely related method DirectPred (Tian et al. 2021). 
abstractText: 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 efficient algorithm, rivals or even outperforms DirectPred.
authors.name: ING WITH
authors.name: NON-CONTRASTIVE SELF-SUPERVISION
references.authors: Sanjeev Arora Nadav Cohen Elad Hazan
references.title: On the optimization of deep networks: Implicit acceleration by overparameterization
references.venue: In ICML. PMLR,
references.authors: Sanjeev Arora Nadav Cohen Noah Golowich Wei Hu
references.title: A convergence analysis of gradient descent for deep linear neural networks
references.venue: In ICLR,
references.authors: Philip Bachman R Devon Hjelm William Buchwalter
references.title: Learning representations by maximizing mutual information across views
references.venue: arXiv preprint arXiv:1906.00910,
references.authors: Adrien Bardes Jean Ponce Yann LeCun
references.title: Vicreg: Variance-invariance-covariance regularization for self-supervised learning
references.venue: arXiv preprint arXiv:2105.04906,
references.authors: Jane Bromley Isabelle Guyon Yann LeCun Eduard Säckinger Roopak Shah
references.title: Signature verification using a“ siamese” time delay neural network
references.venue: NeurIPS,
references.authors: Mathilde Caron Ishan Misra Julien Mairal Priya Goyal Piotr Bojanowski Armand Joulin
references.title: Unsupervised learning of visual features by contrasting cluster assignments
references.venue: NeurIPS,
references.authors: Mathilde Caron Hugo Touvron Ishan Misra Hervé Jégou Julien Mairal Piotr Bojanowski Armand Joulin
references.title: Emerging properties in self-supervised vision transformers
references.venue: arXiv preprint arXiv:2104.14294,
references.authors: Ting Chen Simon Kornblith Mohammad Norouzi Geoffrey Hinton
references.title: A simple framework for contrastive learning of visual representations
references.venue: arXiv preprint arXiv:2002.05709,
references.authors: Xinlei Chen Kaiming He
references.title: Exploring simple siamese representation learning
references.venue: arXiv preprint arXiv:2011.10566,
references.authors: Xinlei Chen Haoqi Fan Ross Girshick Kaiming He
references.title: Improved baselines with momentum contrastive learning
references.venue: arXiv preprint arXiv:2003.04297,
references.authors: Adam Coates Andrew Ng Honglak Lee
references.title: An analysis of single-layer networks in unsupervised feature learning
references.venue: In International conference on artificial intelligence and statistics,
references.authors: J. Deng W. Dong R. Socher L.-J. Li K. Li L. Fei-Fei
references.title: ImageNet: A Large-Scale Hierarchical Image Database
references.venue: In CVPR,
references.authors: Simon Du Wei Hu
references.title: Width provably matters in optimization for deep linear neural networks
references.venue: In ICML,
references.authors: Aleksandr Ermolov Aliaksandr Siarohin Enver Sangineto Nicu Sebe
references.title: Whitening for selfsupervised representation learning
references.venue: In International Conference on Machine Learning,
references.authors: Rong Ge Qingqing Huang Sham M Kakade
references.title: Learning mixtures of gaussians in high dimensions
references.venue: In Proceedings of the forty-seventh annual ACM symposium on Theory of computing,
references.authors: Jean-Bastien Grill Florian Strub Florent Altché Corentin Tallec Pierre H Richemond Elena Buchatskaya Carl Doersch Bernardo Avila Pires Zhaohan Daniel Guo Mohammad Gheshlaghi Azar
references.title: Bootstrap your own latent: A new approach to self-supervised learning
references.venue: arXiv preprint arXiv:2006.07733,
references.authors: Kaiming He Xiangyu Zhang Shaoqing Ren Jian Sun
references.title: Deep residual learning for image recognition
references.authors: Kaiming He Haoqi Fan Yuxin Wu Saining Xie Ross Girshick
references.title: Momentum contrast for unsupervised visual representation learning
references.venue: In CVPR,
references.authors: Tianyu Hua Wenxiao Wang Zihui Xue Yue Wang Sucheng Ren Hang Zhao
references.title: On feature decorrelation in self-supervised learning
references.authors: Kenji Kawaguchi
references.title: Deep learning without poor local minima
references.venue: NeurIPS,
references.authors: Alex Krizhevsky Geoffrey Hinton
references.title: Learning multiple layers of features from tiny images
references.authors: Andrew K Lampinen Surya Ganguli
references.title: An analytic theory of generalization dynamics and transfer learning in deep linear networks
references.venue: In ICLR,
references.authors: Thomas Laurent James Brecht
references.title: Deep linear networks with arbitrary loss: All local minima are global
references.venue: In ICML,
references.authors: Jason D Lee Qi Lei Nikunj Saunshi Jiacheng Zhuo
references.title: Predicting what you already know helps: Provable self-supervised learning
references.venue: arXiv preprint arXiv:2008.01064,
references.authors: Aaron van den Oord Yazhe Li Oriol Vinyals
references.title: Representation learning with contrastive predictive coding
references.venue: arXiv preprint arXiv:1807.03748,
references.authors: Andrew M Saxe James L McClelland Surya Ganguli
references.title: Exact solutions to the nonlinear dynamics of learning in deep linear neural networks
references.venue: arXiv preprint arXiv:1312.6120,
references.authors: Andrew M Saxe James L McClelland Surya Ganguli
references.title: A mathematical theory of semantic development in deep neural networks
references.venue: Proc. Natl. Acad. Sci. U. S. A.,
references.authors: Gilbert W Stewart
references.title: On the perturbation of pseudo-inverses, projections and linear least squares problems
references.venue: SIAM review,
references.authors: Yonglong Tian Dilip Krishnan Phillip Isola
references.title: Contrastive multiview coding
references.venue: arXiv preprint arXiv:1906.05849,
references.authors: Yuandong Tian Xinlei Chen Surya Ganguli
references.title: Understanding self-supervised learning dynamics without contrastive pairs
references.venue: arXiv preprint arXiv:2102.06810,
references.authors: Christopher Tosh Akshay Krishnamurthy Daniel Hsu
references.title: Contrastive learning, multi-view redundancy, and linear models
references.venue: arXiv preprint arXiv:2008.10150,
references.authors: Roman Vershynin
references.title: Introduction to the non-asymptotic analysis of random matrices
references.venue: arXiv preprint arXiv:1011.3027,
references.authors: Roman Vershynin
references.title: High-dimensional probability: An introduction with applications in data science, volume 47
references.authors: Martin J Wainwright
references.title: High-dimensional statistics: A non-asymptotic viewpoint, volume 48
references.authors: Zixin Wen Yuanzhi Li
references.title: Toward understanding the feature learning process of self-supervised contrastive learning
references.venue: arXiv preprint arXiv:2105.15134,
references.authors: Yang You Igor Gitman Boris Ginsburg
references.title: Large batch training of convolutional networks
references.authors: Jure Zbontar Li Jing Ishan Misra Yann LeCun Stéphane
references.title: Deny. Barlow twins: Self-supervised learning via redundancy reduction
references.authors: Ge
references.title: The following corollary is particularly useful for us
sections.text: 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 efficient algorithm, rivals or even outperforms DirectPred.
sections.heading: 1 INTRODUCTION
sections.text: 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 classification) (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 first 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 find a local optimum that admits a nontrivial representation. Based on their theoretical findings 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 first 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 first 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 specifically, 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. (2021). To our knowledge, this is the first 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, STL-10 and ImageNet. In DirectCopy, the predictor can be set without the expensive eigen-decomposition operation, which makes DirectCopy much simpler and more efficient 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 Section 5 shows ablation experiments. Finally, limitation and future works are discussed in Section 6-7.
sections.heading: 2 PRELIMINARIES
sections.heading: 2.1 NOTATIONS
sections.text: We use Id to denote the d×d identity matrix and simply write I when the dimension is clear. For any linear subspace S in Rd, we use PS ∈ Rd×d to denote the projection matrix on S. More precisely, the projection matrix PS equals UU>, where the columns of U constitute a set of orthonormal bases for subspace S. We useN (µ,Σ) to denote the Gaussian distribution with mean µ and covariance Σ. We use ‖·‖ to denote spectral norm for a matrix, or `2 norm for a vector and use ‖·‖F to denote Frobenius norm for a matrix. For a real symmetric matrix A ∈ Rd×d whose eigen-decomposition is∑d i=1 λiuiu > i , we use |A| to denote ∑d i=1 |λi|uiu>i . If A is also positive semi-definite, we use Aα
  to denote ∑d i=1 λ α i uiu > i for any positive α ∈ R.
  2.2 DIRECTSET(α) AND DIRECTCOPY
  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 x1, x2 are generated, and the network is trained to match the representation of x1 (through the online network and the predictor network) and the representation of x2 (through the target network). More precisely, suppose the online network and the target network are two mappings fθ, fθa : Rd 7→ Rh and the predictor network is a mapping gθp : Rh 7→ Rh, the network is trained to minimize the following loss:
  L(θ, θp, θa) := 1
  2 Ex1,x2 ∥∥∥∥∥ gθp (fθ(x1))∥∥gθp (fθ(x1))∥∥ − StopGrad ( fθa(x2) ‖fθa(x2)‖ )∥∥∥∥∥ 2 .
  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
  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:
  Wp = Fα
  ‖Fα‖ + I,
  where F = Ex1fθ(x1)fθ(x1)>. In practice, F is estimated by a moving average over batches. That is, F̂ = µF̂ + (1− µ)EB [fθ(x1)fθ(x1)>], where EB is the expectation over one batch.
  In the original DirectPred proposed by Tian et al. (2021), α is fixed at 1/2. To compute F̂ 1/2, one needs to first 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 Wp 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.
  3 THEORETICAL ANALYSIS OF DIRECTSET(α) 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.
sections.heading: 3.1 SETUP
sections.text: In this subsection, we define 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,Wp,Wa ∈ Rd×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. Specifically, we assume Assumption 2 (Data distribution). The input x is sampled from N (0, Id), and its augmented view x1, x2 are independently sampled from N (x, σ2PB), where B is a (d − r)-dimensional subspace. We denote S as the orthogonal subspace of B in Rd.
  1Computing the spectral norm of F̂ is much faster than computing the eigen-decomposition of F̂ , because the former only needs the top eigen-vector of F̂ . Table 4 shows that the spectral norm can also be replaced by Frobenius norm or no normalization, and similar performance can be achieved.
  In this simple data distribution, subspace S corresponds to the features that are invariant to augmentations and its orthogonal subspaceB 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 σ2I and did not study what representation is learned.
  For the convenience of analysis, we consider a simplified version of DirectSet(α). We compute the loss function without normalizing the two representations, so the population loss is
  L(W,Wa,Wp) := 1
  2 Ex1,x2 ‖WpWx1 − StopGrad (Wax2)‖ 2 , (1)
  and the empirical loss is
  L̂(W,Wp,Wa) := 1
  2n n∑ i=1 ∥∥∥WpWx(i)1 − StopGrad(Wax(i)2 )∥∥∥2 , (2) 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), σ2PB). To train our model, first, we initialize W as δI with δ a positive real number. We run gradient flow or gradient descent on online network W with weight decay η, and set the the target network Wa = W. For clarity of presentation, when training on the population loss, we set Wp as (WExxx>W>)α = (WW>)α instead of (WEx1x1x>1 W>)α as in practice; when training on the empirical loss, we set Wp as (W 1n ∑n i=1 x
  (i)[x(i)]>W>)α. Here, we set the predictor regularization = 0 and its influence will be studied in Section 5.
  In the following, DirectSet(α) is shown to recover the projection matrix PS with polynomial number of samples. Furthermore, given that the learned matrix is close to PS , the sample complexity on downstream tasks is reduced.
sections.heading: 3.2 GRADIENT FLOW ON POPULATION LOSS
sections.text: In this section, we show that DirectSet(α) running on the population loss with infinitesimal 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 infinitesimal step size and η weight decay. If we set the weight decay
  coefficient η ∈ (
  1 4(1+σ2) , 1 4
  ) and initialization scale δ > ( 1− √ 1−4η 2 )1/(2α) , then W converges to(
  1+ √ 1−4η 2 )1/(2α) PS when time goes to infinity.
  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 PS 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 Wp is set as (WEx1x1x>1 W>)α 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 PB . 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:
  λ̇B = λB [ −(1 + σ2) |λB |4α + |λB |2α − η ] , λ̇S = λS [ − |λS |4α + |λS |2α − η ] . (3)
  Next, we show λB converges to zero and λS converges to a positive number, which immediately implies that W converges to some scaling of PS .
  2Note that Theorem 1 also holds with negative initialization δ < − ( 1− √ 1−4η 2 )1/(2α) , in which case W
  converges to − ( 1+ √ 1−4η 2 )1/(2α) PS . Our other results can be extended to negative δ in a similar way.
  Similar as the analysis in Tian et al. (2021), when η > 14(1+σ2) , we know λ̇B < 0 for any λB > 0 and λB = 0 is a stable stationary point, as illustrated in Figure 2 (top, left). Therefore, as long as η > 14(1+σ2) , λB must converge to zero. When 0 < η < 1 4 , there are three non-negative solutions
  to λ̇S = 0, which are 0, λ−S = ( 1− √ 1−4η 2 )1/(2α) and λ+S = ( 1+ √ 1−4η 2 )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 significant change under data augmentation has larger variance σ2 and 14(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 and
  1 4(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 STL10 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 significantly suppressed when η = 0.
  Therefore, it is crucially important to choose weight decay appropriately: a too small η may not be sufficient 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.
sections.heading: 3.3 GRADIENT DESCENT ON EMPIRICAL LOSS
sections.text: 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 defined in Assumption 1 and 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 σ2
  is a positive constant, the weight decay coefficient η ∈ (
  1+σ2/4 4(1+σ2) , 1+3σ2/4 4(1+σ2)
  ) and the initialization
  scale δ is a constant at least 1/ √
  2. Choose the step size γ as a small enough constant. For any accuracy ̂ > 0, given n ≥ poly(d, 1/̂) number of samples, with probability at least 0.99 there exists t = O(log(1/̂)) such that (here W̃t is the online network weights at the t-th step):∥∥∥∥∥W̃t − √ 1 + √ 1− 4η 2 PS
  ∥∥∥∥∥ ≤ ̂. The proof proceeds by first 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.
sections.heading: 3.4 SAMPLE COMPLEXITY ON DOWNSTREAM TASKS
sections.text: 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, Id) and its label y(i) = 〈 x(i), w∗ 〉 + ξ(i), where w∗ is the ground truth vector with unit `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̂ ∈ Rd×d (for example the learned online network W ) to get its representation P̂ x(i). The regularized loss is then defined as L̂(w) := 12n ∑n i=1 ∥∥∥〈P̂ x(i), w〉− y(i)∥∥∥2 + ρ2 ‖w‖2 . In the below theorem, we show that when ∥∥∥P − P̂∥∥∥ 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 defined in Assumption 3. Suppose∥∥∥P̂ − P∥∥∥ F ≤ ̂ with ̂ < 1. Choose the regularizer coefficient ρ = ̂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̂ satisfies ∥∥∥P̂ ŵ − w∗∥∥∥ ≤ O(̂1/3 + β√r +√log(1/ζ)√
  n
  ) .
  In the above theorem, when n is at least O ( β2(r+log(1/ζ))
  ̂2/3
  ) , we have ∥∥∥P̂ ŵ − w∗∥∥∥ ≤ O(̂1/3). Note that if we directly estimate ŵ without transforming the inputs by P̂ , we need Ω(d) number of samples to ensure that ‖ŵ − w∗‖ ≤ 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.
sections.heading: 4 EMPIRICAL PERFORMANCE OF DIRECTCOPY
sections.text: 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 efficient in a simple linear setting. Compared with the original DirectPred (DirectSet(α) with α = 1/2), DirectCopy is a simpler and computationally more efficient 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 PS with larger scale 3
  3Recall that in Theorem 1 under DirectSet(α), online matrix W converges to ( 1+ √ 1−4η 2 )1/(2α) PS . So
  with a larger α, the scalar in front of PS becomes larger.
  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.
sections.heading: 4.1 RESULTS ON STL-10, CIFAR-10 AND CIFAR-100
sections.text: 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 specified otherwise, SGD is used as the optimizer with weight decay η = 0.0004. 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 specified, we use learning rate γ = 0.01, regularization = 0.2 on STL-10; γ = 0.02, = 0.3 on CIFAR-10 and γ = 0.03, = 0.3 on CIFAR-100. See more detailed experiment settings in Appendix A.
  Num of epochs 100 300 500
  STL-10 DirectCopy 77.83±0.56 82.01±0.28 82.95±0.29 DirectPred 77.86±0.16 78.77±0.97 78.86±1.15
  DirectPred (freq=5) 77.54±0.11 79.90±0.66 80.28±0.62 SGD baseline 75.06±0.52 75.25±0.74 75.25±0.74
  CIFAR-10 DirectCopy 84.02±0.37 89.17±0.12 89.62±0.10 DirectPred 85.21±0.23 88.88±0.15 89.52±0.04
  DirectPred (freq=5) 84.93±0.29 88.83±0.10 89.56±0.13 SGD baseline 84.49±0.20 88.57±0.15 89.33±0.27
  CIFAR-100 DirectCopy 55.40±0.19 61.06±0.14 62.23±0.06 DirectPred 56.60±0.27 61.65±0.18 62.68±0.35
  DirectPred (freq=5) 56.43±0.21 62.01±0.22 63.15±0.27 SGD baseline 54.94±0.50 60.88±0.59 61.42±0.89
  Table 1: STL-10/CIFAR-10/CIFAR-100 Top-1 accuracy of DirectCopy. The numbers for DirectPred, DirectPred (freq=5) and SGD baseline on STL-10/CIFAR-10 are obtained from Tian et al. (2021).
  epochs 100
  ImageNet DirectCopy 68.8 DirectPred 68.5 SGD Baseline 68.6
  Table 2: ImageNet Top-1 accuracy of DirectCopy, DirectPred and BYOL baseline.
  STL-10: We evaluate the quality of the learned representation after each epoch, and report the best accuracy in the first 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, and is trained with gradient updates in other batchs, which outperforms DirectPred in later epochs, 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.
sections.heading: 4.2 RESULTS ON IMAGENET
sections.text: 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 significantly 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 = 0.01 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.
sections.heading: 5 ABLATION STUDY
sections.text: In this section, we study the influence of predictor regularization , normalization method, weight decay and degree α on the performance of DirectCopy.
  Predictor regularization: Table 3 shows that when the predictor regularization increases, the performance of DirectCopy on STL-10 and CIFAR-10 improves at first and then deteriorates. On STL-10, DirectCopy with = 1 completely fails. On CIFAR-10, although DirectCopy with = 1 achieved reasonable performance at epoch 300, it’s still much worse than = 0.3.
  To better understand the role of , we analyze the simple linear setting as in Section 3.1 while setting Wp = WW> + 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:
  λ̇S = −λS ( λ2S + − 1− √
  1− 4η 2
  )( λ2S + − 1 + √
  1− 4η 2
  ) .
  Increasing shifts the two positive stationary points λ−S , λ + S towards zero. As illustrated in Figure 3 (left), as 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, a reasonable can alleviate representation collapse, but a too large also encourages representation collapse. As shown in Figure 3 (right), when increases from zero, more eigenvalues of F becomes large; but when exceeds 0.3, eigenvalues of F begin to collapse.
  Normalization on F : In our experiments, we have been normalizing F by its spectral norm before adding the regularization: Wp = F/ ‖F‖ + 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.
  Weight decay: Table 5 shows that when weight decay η increases, the performance of DirectCopy improves at first and then deteriorates. This fits our analysis on simple linear networks. Basically, when the weight decay η increases, it can suppress the nuisance features more effectively, but a too large weight decay also collapses the useful features.
  Predictor degree: We compare DirectCopy against DirectSet(α) with α = 2, 1/2, 1/4. Table 6 shows that DirectCopy outperforms other algorithms on STL-10. On CIFAR-10, DirectCopy is slightly worse at epoch 100, but catches up in later epochs.
sections.heading: 6 BEYOND LINEAR MODELS: LIMITATIONS AND DISCUSSION
sections.text: 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 defined 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.
sections.heading: 7 CONCLUSION
sections.text: 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 efficient 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.
sections.heading: A DETAILED EXPERIMENT SETTING
sections.text: 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 specified 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. (2021) 4, 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 specified, we use learning rate γ = 0.01, regularization = 0.2 on STL-10; γ = 0.02, = 0.3 on CIFAR-10 and γ = 0.03, = 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.45 at batch size 256. Other setups such as weight decay (η = 1e−6), target EMA (scheduled from 0.99 to 1), augmentation recipe (color jitters, blur, etc.), and linear evaluation protocol are the same as BYOL.
sections.heading: B PROOFS OF SINGLE-LAYER LINEAR NETWORKS
sections.heading: B.1 GRADIENT FLOW ON POPULATION LOSS
sections.text: In this section, we give the proof of Theorem 1, which shows that DirectSet(α) running on the population loss with infinitesimal 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 infinitesimal step size and η weight decay. If we set the weight decay
  coefficient η ∈ (
  1 4(1+σ2) , 1 4
  ) and initialization scale δ > ( 1− √ 1−4η 2 )1/(2α) , then W converges to(
  1+ √ 1−4η 2 )1/(2α) PS when time goes to infinity.
  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 PS .
  Proof of Theorem 1. We can compute the gradient in terms of W as follows,
  ∇L(W ) =Ex1,x2W>p (WpWx1 −Wax2)x>1 =W>p ( WpWEx1x1x>1 −WaEx1,x2x2x>1 ) .
  Note that the two augmented views x1, x2 are sampled by first sampling input x fromN (0, Id), and then independently sampling x1, x2 fromN (x, σ2PB). Therefore, we know Ex1x1x>1 = I +σ2PB and Ex1,x2x2x>1 = I. Recall that we run gradient flow on W with weight decay η, so the dynamics on W is as follows:
  Ẇ =W>p (−WpW (I + σ2PB) +Wa)− ηW, where the first term comes from the gradient and the second term is due to weight decay.
  SinceW is initialized as δI, andWa = W,Wp = (WW>)α, so we know initiallyW,Wp,Wa, I and PB are all simultaneously diagonalizable, which then implies Ẇ is simultaneously diagonalizable
  4Their open source code is at https://github.com/facebookresearch/luckmatters/tree/main/ssl
  with W . This argument can continue to show that at any time point, W,Wp,Wa, I and PB are all simultaneously diagonalizable. Since W is always a real symmetric matrix, we have Wp = (WW>)α = |W |2α . The dynamics on W can then be written as
  Ẇ = |W |2α (− |W |2αW (I + σ2PB) +W )− ηW =W ( −(I + σ2PB) |W |4α + |W |2α − η ) .
  Let the eigenvalue decomposition of W be ∑d i=1 λiuiu > i , with span({ud−r+1, · · · , ud}) 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:
  λ̇B = λB [ −(1 + σ2) |λB |4α + |λB |2α − η ] Similar as the analysis in Tian et al. (2021), when η > 14(1+σ2) , we know λ̇B < 0 for any λB > 0 and λB = 0 is a critical point. This means, as long as η > 14(1+σ2) , λB must converge to zero.
  Dynamics for λS: We can write down the dynamics for λS as follows:
  λ̇S = λS [ − |λS |4α + |λS |2α − η ] .
  When 0 < η < 14 , we know λ̇S > 0 for λ 2α S ∈ ( 1− √ 1−4η 2 , 1+ √ 1−4η 2 ) and λ̇S < 0 for λ2αS ∈(
  1+ √ 1−4η 2 ,∞ ) . Furthermore, we know λ̇S = 0 when λ2αS = 1+ √ 1−4η 2 . Therefore, as long as
  0 < η < 14 and initialization δ 2α > 1− √ 1−4η 2 , we know λ 2α S converges to 1+ √ 1−4η 2 .
  Overall, we know when 14(1+σ2) < η < 1 4 and δ > ( 1− √ 1−4η 2 )1/(2α) , we have λB converge to
  zero and λS converge to ( 1+ √ 1−4η 2 )1/(2α) . That is, matrix W converges to ( 1+ √ 1−4η 2 )1/(2α) PS .
sections.heading: B.2 GRADIENT DESCENT ON EMPIRICAL LOSS
sections.text: 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 defined in Assumption 1 and 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 σ2
  is a positive constant, the weight decay coefficient η ∈ (
  1+σ2/4 4(1+σ2) , 1+3σ2/4 4(1+σ2)
  ) and the initialization
  scale δ is a constant at least 1/ √
  2. Choose the step size γ as a small enough constant. For any accuracy ̂ > 0, given n ≥ poly(d, 1/̂) number of samples, with probability at least 0.99 there exists t = O(log(1/̂)) such that (here W̃t is the online network weights at the t-th step):∥∥∥∥∥W̃t − √ 1 + √ 1− 4η 2 PS
  ∥∥∥∥∥ ≤ ̂. When running gradient descent on the empirical loss, the eigenspace of W̃t can shift and become no longer simultaneously diagonalizable with PB . So we cannot independently analyze each eigenvalue of W̃t as before, which brings significant challenge into the analysis. Instead of directly analyzing the dynamics of W̃t, we first show that the gradient descent iterates Wt on the population loss converges to PS in linear rate, and then show that W̃t stays close to Wt within certain iterations.
  Lemma 1. In the setting of Theorem 2, let Wt be the gradient descent iterations on the population loss L. Given any accuracy ̂ > 0, for any t ≥ C log(1/̂), we have∥∥∥∥∥Wt − √ 1 + √ 1− 4η 2 PS ∥∥∥∥∥ ≤ ̂, where C is a positive constant.
  The proof of Lemma 1 is similar as the gradient flow 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/̂)) iterations.
  Lemma 2. In the setting of Theorem 2, let Wt be the gradient descent iterations on the population loss and let W̃t be the gradient descent iterations on the empirical loss. For any accuracy ̂ > 0, given n ≥ poly(d, 1/̂) number of samples, with probability at least 0.99, for any t ≤ C log(1/̂), we have ∥∥∥W̃t −Wt∥∥∥ ≤ ̂, 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 ̂′, for t = C log(1/̂), we have ∥∥∥∥∥Wt − √ 1 + √ 1− 4η 2 PS ∥∥∥∥∥ ≤ ̂′, where C is a positive constant.
  According to Lemma 2, we know given n ≥ poly(d, 1/̂′) number of samples, with probability at least 0.99, ∥∥∥W̃t −Wt∥∥∥ ≤ ̂′. Therefore, we have∥∥∥∥∥W̃t − √ 1 + √ 1− 4η 2 PS ∥∥∥∥∥ ≤ ∥∥∥∥∥Wt − √ 1 + √ 1− 4η 2 PS
  ∥∥∥∥∥+ ∥∥∥W̃t −Wt∥∥∥ ≤ 2̂′. Replacing ̂′ by ̂/2 finishes the proof.
  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.
sections.heading: B.2.1 PROOFS FOR LEMMA 1 AND LEMMA 2
sections.text: Proof of Lemma 1. Similar as in Theorem 1, we can show that at any step t, Wt is simultaneously diagonalizable with Wa,t,Wp,t, I and PB . The update on Wt is as follows,
  Wt+1 = Wt + γWt ( −(I + σ2PB)W 4t +W 2t − η ) .
  Let the eigenvalue decomposition of Wt be ∑d i=1 λi,tuiu > i , with span({ud−r+1, · · · , ud}) 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
  0 ≤ λB,t ≤ (1− γC1)tδ
  for any step size γ ≤ C2, where C1, C2 are two positive constants.
  According to the gradient update, we have λB,t+1 = λB,t + γλB,t [ −(1 + σ2)λ4B,t + λ2B,t − η ] .
  We only need to prove that for any λB,t ∈ [0, δ], we have
  −(1 + σ2)λ4B,t + λ2B,t − η = −Θ(1). This is true since η ∈ (
  1+σ2/4 4(1+σ2) , 1+3σ2/4 4(1+σ2)
  ) and σ2, δ are two positive constants.
  Dynamics of λS: We show that 0 ≤ ∣∣∣∣λ2S,t − 1 +√1− 4η2 ∣∣∣∣ ≤ (1− γC3)t ∣∣∣∣δ2 − 1 +√1− 4η2 ∣∣∣∣
  for any step size γ ≤ C4, where C3, C4 are two positive constants. There are two cases to consider: when the initialization scale δ2 ∈ [1/2, 1+ √ 1−4η 2 ], we prove
  0 ≤ 1 + √
  1− 4η 2
  − λ2B,t ≤ (1− γC3)t ( 1 + √
  1− 4η 2
  − δ2 ) ;
  when the initialization scale δ2 > 1+ √ 1−4η 2 , we prove
  0 ≤ λ2B,t − 1 + √
  1− 4η 2
  ≤ (1− γC3)t ( δ2 − 1 + √ 1− 4η 2 ) .
  We focus on the second case; the proof for the first case is similar.
  According to the gradient update, we have λS,t+1 =λS,t + γλS,t [ −λ4S,t + λ2S,t − η ] =λS,t − γλS,t ( λ2S,t − 1− √
  1− 4η 2
  )( λ2S,t − 1 + √
  1− 4η 2 ) We only need to show that λS,t ( λ2S,t − 1− √ 1−4η 2 ) = Θ(1) for any λ2S,t ∈ [ 1+ √ 1−4η 2 , δ]. This is
  true because η ∈ (
  1+σ2/4 4(1+σ2) , 1+3σ2/4 4(1+σ2)
  ) and σ2, δ are two positive constants.
  Overall, we know that there exists constant step size such that after t = O(log(1/̂)) steps, we have
  0 ≤ λB,t ≤ ̂ and ∣∣∣∣∣λS,t − √ 1 + √ 1− 4η 2 ∣∣∣∣∣ ≤ ̂. This then implies, ∥∥∥∥∥Wt − √ 1 + √ 1− 4η 2 PS ∥∥∥∥∥ ≤ ̂.
  Proof of Lemma 2. We know the update on W̃t is
  W̃t+1 − W̃t = γW̃>p,t ( −W̃p,tW̃t ( 1
  n n∑ i=1 x (i) 1 [x (i) 1 ] >
  ) + W̃a,t ( 1
  n n∑ i=1 x (i) 1 [x (i) 2 ] >
  )) − γηW̃t,
  and the update on Wt is Wt+1 −Wt = γW>p,t ( −Wp,tWt ( I + σ2PB ) +Wa,t ) − γηWt.
  Next, we bound ∥∥∥W̃t+1 − W̃t − (Wt+1 −Wt)∥∥∥ .According to Lemma 3, we know with probability
  at least 1−O(d2) exp ( −Ω(̂′2n/d2) ) ,∥∥∥∥∥ 1n n∑ i=1 x (i) 1 [x (i) 1 ] > − I − σ2PB ∥∥∥∥∥ , ∥∥∥∥∥ 1n n∑ i=1 x (i) 1 [x (i) 2 ] > − I ∥∥∥∥∥ , ∥∥∥∥∥ 1n n∑ i=1 x(i)[x(i)]> − I ∥∥∥∥∥ ≤ ̂′.
  Recall that we set W̃a,t = W̃t and set Wa,t as Wt, so we have ∥∥∥W̃a,t −Wa,t∥∥∥ = ∥∥∥W̃t −Wt∥∥∥ .
  Also since we set W̃p,t = W̃t ( 1 n ∑n i=1 x (i)[x(i)]> ) W̃>t and set Wp,t = WtW
  > t , we have∥∥∥W̃p,t −Wp,t∥∥∥ = O (∥∥∥W̃t −Wt∥∥∥+ ̂′) since ‖Wt‖ = O(1).
  Combing the above bounds and recall γ is a constant, we have∥∥∥W̃t+1 − W̃t − (Wt+1 −Wt)∥∥∥ = O (∥∥∥W̃t −Wt∥∥∥+ ̂′) . Therefore, ∥∥∥W̃t −Wt∥∥∥ ≤ Ct1̂′, where C1 is a constant larger than 1. So for any t ≤ C log(1/̂), we have∥∥∥W̃t −Wt∥∥∥ ≤ CC log(1/̂)1 ̂′ ≤ (1/̂)C2 ̂′, for some positive constant C2. Choosing ̂′ = ̂C2+1, we know as long as n ≥ poly(d, 1/̂), with probability at least 0.99, for any t ≤ C log(1/̂), we have∥∥∥W̃t −Wt∥∥∥ ≤ ̂.
sections.heading: B.3 SAMPLE COMPLEXITY ON DOWN-STREAM TASKS
sections.text: 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 defined in Assumption 3. Suppose∥∥∥P̂ − P∥∥∥
  F ≤ ̂ with ̂ < 1. Choose the regularizer coefficient ρ = ̂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̂ satisfies ∥∥∥P̂ ŵ − w∗∥∥∥ ≤ O(̂1/3 + β√r +√log(1/ζ)√
  n
  ) .
  Suppose {(x(i), y(i))}ni=1 are n training samples in the downstream task, let X ∈ Rn×d be the data matrix with its i-th row equal to x(i). Denote y ∈ Rn as the label vector with its i-th entry as y(i). Each input x(i) is transformed by a matrix P̂ ∈ Rd×d to get its representation P̂ x(i). The regularized loss can be written as
  L(w) := 1
  2n ∥∥∥XP̂w − y∥∥∥2 + ρ 2 ‖w‖2 .
  This is the ridge regression problem on inputs {(P̂ x(i), y(i))}ni=1, and the unique global minimizer ŵ has the following close form:
  ŵ =
  ( 1
  n P̂>X>XP̂ + ρI
  )−1 1
  n P̂>X>y (4)
  With the above closed form of ŵ, the proof of Theorem 3 follows by bounding the difference between P̂ ŵ andw∗ 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 ‖∆‖F ≤ ̂ by assumption. We can also write y as Xw∗+ ξ where ξ ∈ Rn is the noise vector with its i-th entry equal to ξ(i). Then, we can divide ŵ into two terms,
  ŵ =
  ( 1
  n P̂>X>XP̂ + ρI
  )−1 1
  n P̂>X>y
  =
  ( 1
  n P̂>X>XP̂ + ρI
  )−1 1
  n P>X> (Xw∗ + ξ) +
  ( 1
  n P̂>X>XP̂ + ρI
  )−1 1
  n ∆>X> (Xw∗ + ξ)
  Let’s first give an upper bound for the second term that comes from the error term ∆>.
sections.heading: Upper bounding
sections.text: ∥∥∥∥( 1n P̂>X>XP̂ + ρI)−1 1n∆>X> (Xw∗ + ξ)∥∥∥∥ We first bound the norm of 1 n∆ >X>Xw∗. According to Lemma 5, we know with probability at least 1 − exp(−Ω(n)),∥∥∥ 1√n∆>X>∥∥∥F ≤ O(̂). Since Xw∗ is a standard Gaussian vector with dimension n, according
  to Lemma 8, with probability at least 1 − exp(−Ω(n)), ∥∥∥ 1√nXw∗∥∥∥ ≤ O(1). Therefore, we have∥∥ 1
  n∆ >X>Xw∗ ∥∥ ≤ O(̂). Then we bound the norm of 1n∆ >X>ξ. According to Lemma 8, we know
  with probability at least 1 − exp(−Ω(n)), ∥∥∥ 1√nξ∥∥∥ ≤ O(β). According to Lemma 6, we
  know with probability at least 1 − ζ/3, ∥∥∆>X>ξ̄∥∥ ≤ O (̂√log(1/ζ)) . Therefore, we have∥∥ 1
  n∆ >X>ξ ∥∥ ≤ O(β̂√log(1/ζ)√ n ) .
  Since λmin ( 1 n P̂ >X>XP̂ + ρI ) ≥ ρ, we have ∥∥∥∥( 1n P̂>X>XP̂ + ρI)−1∥∥∥∥ ≤ 1ρ . Combining with above bound on ∥∥ 1 n∆ >X> (Xw∗ + ξ)
  ∥∥, we know with probability at least 1−exp(−Ω(n))−ζ/3,∥∥∥∥∥ ( 1 n P̂>X>XP̂ + ρI )−1 1 n ∆>X> (Xw∗ + ξ) ∥∥∥∥∥ ≤ O ( ̂ ρ + β̂ √ log(1/ζ) ρ √ n ) .
  Analyzing (
  1 n P̂ >X>XP̂ + ρI )−1 1 nP >X> (Xw∗ + ξ) We can write 1n P̂
  >X>XP̂ as 1 nP >X>XP + E, where
  E = 1
  n ∆>X>XP +
  1 n P>X>X∆ + 1 n ∆>X>X∆.
  Let’s first bound the spectral norm of XP. Since P is a projection matrix on an r-dimensional subspace S, we can write P as UU>, where U ∈ Rd×r has columns as an orthonormal basis of subspace S. According to Lemma 4, we know with probability at least 1− exp(−Ω(n)),
  Ω(1) ≤ σmin (
  1√ n XU
  ) ≤ σmax ( 1√ n XU ) ≤ O(1).
  Since ‖U‖ ≤ 1, we have ∥∥∥ 1√nXP∥∥∥ = ∥∥∥ 1√nXUU>∥∥∥ ≤ O(1).
  According to Lemma 5, we know with probability at least 1− exp(−Ω(n)),∥∥∥∥ 1√nX∆ ∥∥∥∥ F ≤ O(̂).
  So overall, we know ‖E‖ ≤ ‖E‖F ≤ O(̂). Then, we can write(
  1 n P̂>X>XP̂ + ρI
  )−1 = ( 1
  n P>X>XP + ρI
  )−1 + F.
  According to the perturbation bound for matrix inverse (Lemma 11), we have ‖F‖ ≤ O( ̂ρ2 ). Then, we have(
  1 n P̂>X>XP̂ + ρI
  )−1 1
  n P>X> (Xw∗ + ξ) =
  ( 1
  n P>X>XP + ρI
  )−1 1
  n P>X>Xw∗
  + F 1
  n P>X>Xw∗
  +
  (( 1
  n P>X>XP + ρI
  )−1 + F ) 1
  n P>X>ξ
  We first show that the first term is close to w∗. Let the eigenvalue decomposition of 1nP >X>XP be V ΣV >, where V ’s columns are an orthonormal basis for subspace S. Here Σ ∈ Rr×r is the diagonal matrix that contains all the eigenvalues of 1nP
  >X>XP . According to Lemma 4, we know that with probability at least 1 − exp(−Ω(n)), all the non-zero eigenvalues of 1nP
  >X>XP are Θ(1).
  Then, it’s not hard to show that∥∥∥∥∥ ( 1 n P>X>XP + ρI )−1 1 n P>X>XP − P ∥∥∥∥∥ ≤ O(ρ). This immediately implies that∥∥∥∥∥ ( 1 n P>X>XP + ρI )−1 1 n P>X>Xw∗ − w∗
  ∥∥∥∥∥ ≤ O(ρ) Next, we bound the norm of the second term F 1nP >X>Xw∗. Similar as before, we know
  with probability at least 1 − exp(−Ω(n)), ∥∥∥ 1√nXw∗∥∥∥ ≤ O(1) and ∥∥∥ 1√nP>X>∥∥∥ ≤ O(1). Therefore, we have ∥∥∥∥F 1nP>X>Xw∗ ∥∥∥∥ ≤ ‖F‖ ∥∥∥∥ 1√nP>X> ∥∥∥∥∥∥∥∥ 1√nXw∗ ∥∥∥∥ ≤ O( ̂ρ2 ) . Finally, let’s bound the third term ((
  1 nP >X>XP + ρI
  )−1 + F ) 1 nP >X>ξ. We first bound the
  norm of 1nP >X>ξ. with probability at least 1 − exp(−Ω(n)), we know ‖ξ‖ ≤ 2β
  √ n. Therefore, we know ∥∥ 1 nP >X>ξ ∥∥ ≤ O(β/√n)∥∥P>X>ξ̄∥∥ , where ξ̄ = ξ/ ‖ξ‖ . According to Lemma 7, with probability at least 1 − ζ/3, we have
  ∥∥P>X>ξ̄∥∥ ≤ √r + O(√log(1/ζ)). Overall, with probability at least 1− exp(−Ω(n))− ζ/3,∥∥∥∥ 1nP>X>ξ ∥∥∥∥ ≤ O (√ rβ + √ log(1/ζ)β√ n ) .
  It’s not hard to verify that for any vector v ∈ Rd in the subspace S, we have∥∥∥(( 1nP>X>XP + ρI)−1 + F) v∥∥∥ ≤ O(‖v‖). Since 1nP>X>ξ lies on subspace S, we have∥∥∥∥∥ (( 1 n P>X>XP + ρI )−1 + F ) 1 n P>X>ξ ∥∥∥∥∥ ≤ O (√ rβ + √ log(1/ζ)β√ n ) .
  Combining the above analysis and taking a union bound over all the events, we know with probability at least 1− exp(−Ω(n))− 2ζ/3,
  ‖ŵ − w∗‖ = O ( ρ+ ̂
  ρ +
  ̂ ρ2 + β̂ √ log(1/ζ) ρ √ n + √ rβ + √ log(1/ζ)β√ n )
  Suppose n ≥ O(log(1/ζ)) and setting ρ = ̂1/3, we further have with probability at least 1− ζ,
  ‖ŵ − w∗‖ =O ( ̂1/3 + β̂2/3 √
  log(1/ζ)√ n +
  √ rβ + √ log(1/ζ)β√ n )
  ≤O ( ̂1/3 + β √ r + √ log(1/ζ)√ n ) ,
  where the last inequality assumes ̂ < 1.
  We can also bound ∥∥∥P̂ ŵ − w∗∥∥∥ as follows,∥∥∥P̂ ŵ − w∗∥∥∥ =∥∥∥P̂ ŵ − Pŵ + Pŵ − Pw∗∥∥∥ ≤ ∥∥∥P̂ ŵ − Pŵ∥∥∥+ ‖Pŵ − Pw∗‖ ≤ ∥∥∥P̂ − P∥∥∥ ‖ŵ‖+ ‖P‖ ‖ŵ − w∗‖ ≤̂O ( 1 + ̂1/3 + β √ r + √ log(1/ζ)√ n ) +O ( ̂1/3 + β √ r + √ log(1/ζ)√ n )
  ≤O ( ̂1/3 + β √ r + √ log(1/ζ)√ n )
  B.4 ANALYSIS WITH Wp := (WEx1x1x>1 W>)α
  In this section, we prove that DirectSet(α) can also learn the projection matrix when we set Wp := (WEx1x1x>1 W>)α. For the network architecture and data distribution, we follow exactly the same setting as in Section 3.2. Therefore, we know Wp := (WEx1x1x>1 W>)α = (W (I+σ2PB)W>)α. Theorem 4. Suppose network architecture and data distribution are as defined 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 infinitesimal step size and η weight decay. Suppose we setWa = W and Wp = (WEx1x1x>1 W>)α. Assuming the weight decay coefficient η ∈ ( 1 4(1+σ2)1+2α , 1 4
  ) and initialization scale δ > ( 1− √ 1−4η 2 )1/(2α) , we know W converges to ( 1+ √ 1−4η 2 )1/(2α) PS when time goes to infinity.
  The only difference from Theorem 4 is that now the initialization δ is only required to be larger than 1 4(1+σ2)1+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:
  Ẇ =W>p (−WpW (I + σ2PB) +Wa)− ηW = ∣∣W 2(I + σ2PB)∣∣α (− ∣∣W 2(I + σ2PB)∣∣αW (I + σ2PB) +W )− ηW
  =W ( −(I + σ2PB)1+2α |W |4α + |W |2α − η ) .
  Dynamics for λB: We can write down the dynamics for λB as follows:
  λ̇B = λB [ −(1 + σ2)1+2α |λB |4α + |λB |2α − η ] When η > 14(1+σ2)1+2α , we know λ̇B < 0 for any λB > 0 and λB = 0 is a critical point. This means, as long as η > 14(1+σ2)1+2α , λB must converge to zero.
  Dynamics for λS: The dynamics is same as when setting Wp = (WW>)α,
  λ̇S = λS [ − |λS |4α + |λS |2α − η ] .
  so when 0 < η < 14 and initialization δ 2α > 1− √ 1−4η 2 , we know λ 2α S converges to 1+ √ 1−4η 2 .
  Overall, we know when 14(1+σ2)1+2α < η < 1 4 and δ > ( 1− √ 1−4η 2 )1/(2α) , we have λB converge to
  zero and λS converge to ( 1+ √ 1−4η 2 )1/(2α) . That is, matrix W converges to ( 1+ √ 1−4η 2 )1/(2α) PS .
sections.heading: B.5 TECHNICAL LEMMAS
sections.text: Lemma 3. Suppose {x(i), x(i)1 , x (i) 2 }ni=1 are sampled as decribed in Section 3. Suppose n ≥ O(d/̂2), with probability at least 1−O(d2) exp ( −Ω(̂2n/d2) ) , we have∥∥∥∥∥ 1n n∑ i=1 x (i) 1 [x (i) 1 ] > − I − σ2PB ∥∥∥∥∥ , ∥∥∥∥∥ 1n n∑ i=1 x (i) 1 [x (i) 2 ] > − I ∥∥∥∥∥ , ∥∥∥∥∥ 1n n∑ i=1 x(i)[x(i)]> − I
  ∥∥∥∥∥ ≤ ̂. 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, σ2PB). So we have
  1
  n n∑ i=1 x (i) 1 [x (i) 1 ] > = 1 n n∑ i=1 ( x(i)[x(i)]> + z (i) 1 [z (i) 1 ] > + x(i)[z (i) 1 ] > + z (i) 1 [x (i)]> ) .
  According to Lemma 9, we know as long as n ≥ O(d/̂2), with probability at least 1 − exp(−Ω(̂2n)), ∥∥∥∥∥ 1n n∑ i=1 x(i)[x(i)]> − I ∥∥∥∥∥ ≤ ̂. Similarly, with probability at least 1− exp(−Ω(̂2n)),∥∥∥∥∥ 1n n∑ i=1 z (i) 1 [z (i) 1 ] > − σ2PB
  ∥∥∥∥∥ ≤ ̂. Next we bound
  ∥∥∥ 1n∑ni=1 x(i)[z(i)1 ]>∥∥∥ . We know each entry in matrix 1n∑ni=1 x(i)[z(i)1 ]> is the average of n zero-mean O(1)-subexponential independent random variables. Therefore, according to the Bernstein’s inequality, for any fixed entry (k, l), with probability at least 1− exp ( −̂2n/d2
  ) ,∣∣∣∣∣∣ [ 1 n n∑ i=1 x(i)[z (i) 1 ] > ] k,l
  ∣∣∣∣∣∣ ≤ ̂/d. Taking a union bound over all the entries, we know with probability at least 1− d2 exp ( −̂2n/d2
  ) ,∥∥∥∥∥ 1n n∑ i=1 x(i)[z (i) 1 ] > ∥∥∥∥∥ ≤ ∥∥∥∥∥ 1n n∑ i=1 x(i)[z (i) 1 ] > ∥∥∥∥∥ F ≤ ̂.
  The same analysis also applies to ∥∥∥ 1n∑ni=1 z(i)1 [x(i)]>∥∥∥ . Combing all the bounds, we know with
  probability at least 1−O(d2) exp ( −Ω(̂2n/d2) ) ,∥∥∥∥∥ 1n n∑ i=1 x (i) 1 [x (i) 1 ] > − I − σ2PB
  ∥∥∥∥∥ ≤ 4̂. Similarly, we can prove that with probability at least 1−O(d2) exp ( −Ω(̂2n/d2)
  ) ,∥∥∥∥∥ 1n n∑ i=1 x (i) 1 [x (i) 2 ] > − I
  ∥∥∥∥∥ ≤ 4̂. Changing ̂ to ̂′/4 finishes the proof. Lemma 4. Let X ∈ Rn×d be a standard Gaussian matrix, and let U ∈ Rd×r be a matrix with orthonormal columns. Suppose n ≥ 2r, with probability at least 1− exp(−Ω(n)), we know
  Ω(1) ≤ λmin ( 1
  n U>X>XU
  ) ≤ λmax ( 1
  n U>X>XU
  ) ≤ O(1).
  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 ≥ 2r, with probability at least 1− exp(−Ω(n)),
  Ω(1) ≤ σmin (
  1√ n XU
  ) ≤ σmax ( 1√ n XU ) ≤ O(1).
  This immediately implies that Ω(1) ≤ λmin ( 1
  n U>X>XU
  ) ≤ λmax ( 1
  n U>X>XU
  ) ≤ O(1).
  Lemma 5. Let ∆ be a d×d matrix with Frobenius norm ̂, and let X be a n×d standard Gaussian matrix. We know with probability at least 1− exp(−Ω(n)),∥∥∥∥ 1√nX∆ ∥∥∥∥ F ≤ O(̂).
  Proof of Lemma 5. Let the singular value decomposition of ∆ be UΣV >, where U, V have orthonormal columns and Σ is a diagonal matrix with diagonals equal to singular values σi’s. Since ‖∆‖F = ̂, we know ∑d i=1 σ 2 i = ̂ 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 X̃ := X̂Σ. It’s not hard to verify that all the entries in X̃ are independent Gaussian variables and X̃ij ∼ N (0, σ2j ). According to the Bernstein’s inequality for sum of independent and sub-exponential random variables, we have for every t > 0,
  Pr ∣∣∣∣∣∣ ∑
  i∈[n],j∈[d]
  X̃2ij − n̂2 ∣∣∣∣∣∣ ≥ t  ≤ 2 exp[−cmin( t2∑ i∈[n],j∈[d] σ 4 j , t maxj∈[d] σ 2 j )] .
  Since ∑d j=1 σ 2 j = ‖∆‖ 2 F = ̂ 2, we know maxj∈[d] σ2j ≤ ̂2. We also have ∑ j∈[d] σ
  4 j ≤(∑
  j∈[d] σ 2 j
  )2 = ̂4. Therefore, we have
  Pr ∣∣∣∣∣∣ ∑
  i∈[n],j∈[d]
  X̃2ij − n̂2 ∣∣∣∣∣∣ ≥ t  ≤ 2 exp [−cmin( t2 n̂4 , t ̂2 )] .
  Replacing t by n̂2, we concluded that with probability at least 1 − exp(−Ω(n)), ∥∥∥X̃∥∥∥2
  F ≤ 2n̂2. Furthermore, since ∥∥V >∥∥ = 1, we have∥∥∥∥ 1√nX∆ ∥∥∥∥ F = ∥∥∥∥ 1√nX̃V > ∥∥∥∥ F ≤ ∥∥∥∥ 1√nX̃ ∥∥∥∥ F ‖V ‖ ≤ O(̂).
  Lemma 6. Let ∆> be a d×dmatrix with Frebenius norm ̂ and letX> be a d×n standard Gaussian matrix. Let ξ̄ be a unit vector with dimension n. We know with probability at least 1− ζ/3,∥∥∆>X>ξ̄∥∥ ≤ O(̂√log(1/ζ)). Proof of Lemma 6. Let the sigular value decomposition of ∆> be UΣV >. We know X>ξ̄ is a ddimensional standard Gaussian vector. Further, we know V >X>ξ̄ is also a d-dimensional standard Gaussian vector. So ΣV >X>ξ̄ has independent Gaussian entries with its i-th entry distributed as N (0, σ2i ). According to the Bernstein’s inequality for sum of independent and sub-exponential random variables, we have for every t > 0,
  Pr [∣∣∣∥∥ΣV >X>ξ̄∥∥2 − ̂2∣∣∣ ≥ t] ≤ 2 exp [−cmin( t2
  ̂4 , t ̂2
  )] .
  Choosing t as O(̂2 log(1/ζ)), we know with probability at least 1− ζ/3, we have∥∥ΣV >X>ξ̄∥∥2 ≤ O (̂2 log(1/ζ)) . Since ‖U‖ = 1, we further have∥∥∆>X>ξ̄∥∥ = ∥∥UΣV >X>ξ̄∥∥ ≤ ‖U‖ ∥∥ΣV >X>ξ̄∥∥ ≤ O (̂√log(1/ζ))
  Lemma 7. Let P ∈ Rd×d be a projection matrix on a r-dimensional subspace, and let ξ̄ be a unit vector in Rd. Let X> be a d× n standard Gaussian matrix that is independent with P and ξ. With probability at least 1− ζ/3, we have∥∥P>X>ξ̄∥∥ ≤ √r +O(√log(1/ζ)). Proof of Lemma 7. Since P is a projection matrix on an r-dimensional subspace, we can write P as UU>, where U ∈ Rd×r has orthonormal columns. We know U>X> is still a standard Gaussian matrix with dimension r × n. Furthermore, U>X>ξ̄ is an r-dimensional standard Gaussian vector. According to Lemma 8, with probability at least 1− ζ/3, we have∥∥U>X>ξ̄∥∥ ≤ √r +O(√log(1/ζ)). Since ‖U‖ = 1, we further have∥∥P>X>ξ̄∥∥ = ∥∥UU>X>ξ̄∥∥ ≤ ‖U‖ ∥∥U>X>ξ̄∥∥ ≤ √r +O(√log(1/ζ)).
sections.heading: C ANALYSIS OF DEEP LINEAR NETWORKS
sections.text: In this section, we extend the analysis in Section 3.2 to deep linear networks. We consider the same data distribution as defined in Assumption 2. We consider the following network, Assumption 4 (Deep linear network). The online network is an l-layer linear networks WlWl−1 · · ·W1 with each Wi ∈ Rd×d. The target network has the same architecture with weight matrices Wa,lWa,l−1 · · ·Wa,1. For convenience, we denote W as WlWl−1 · · ·W1 and denote Wa as Wa,lWa,l−1 · · ·Wa,1.
  Training procedure: At the initialization, we initialize each Wi as δ1/lId. Through the training, we fix Wp as ( WW> )α and fix each Wa,i as Wi. We run gradient flow on every Wi with weight decay η. The population loss is
  L({Wi},Wp, {Wa,i}) := 1
  2 Ex1,x2 ‖WpWlWl−1 · · ·W1x1 − StopGrad(Wa,lWa,l−1 · · ·Wa,1x2)‖ 2 .
  Theorem 5. Suppose the data distribution and network architecture satisfies Assumption 2 and Assumption 4, respectively. Suppose we train the network as described above. Assuming the weight decay coefficient
  η ∈ (
  2αl(2αl+2l−2)1+ 1 α − 1 αl
  (4αl+2l−2)2+ 1 α − 1 αl (1+σ2)1+ 1 α − 1 αl , 2αl(2αl+2l−2)
  1+ 1 α − 1 αl
  (4αl+2l−2)2+ 1 α − 1 αl ) , and initialization scale δ ≥(
  2αl+2l−2 4αl+2l−2
  ) 1 2α
  , we know W converges to cPS as time goes to infinity, where c is a positive number within ((
  2αl+2l−2 4αl+2l−2
  ) 1 2α
  , 1 ) .
  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 WjWj−1 · · ·Wi and for j < i have W[j:i] = I. We use similar notations for Wa,[j:i]. For each Wi, we can compute its dynamics as follows:
  Ẇi =− ( WpW[l:i+1] )> ( WpW (I + σ 2PB) ) ( W[i−1:1] )> + ( WpWa,[l:i+1] )> Wa ( Wa,[i−1:1] )> − ηWi.
  It’s clear that through the training all Wi’s remains the same and they are simultaneously diagonalizable with Wp, I and PB . We also have Wa = W and Wp = |W |2α . Since we will ensure that W is always positive semi-definite so Wp = |W |2α = W 2α = W 2αli . So the dynamics for each Wi can be simplified as follows:
  Ẇi = −W 4αl+2l−1i (I + σ 2PB) +W 2αl+2l−1 i − ηWi.
  Let the eigenvalue decomposition of Wi be ∑d i=1 νiuiu > i , with span({ud−r+1, · · · , ud}) 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,
  ν̇S =− ν4αl+2l−1S + ν 2αl+2l−1 S − ηνS , ν̇B =− ν4αl+2l−1B (1 + σ 2) + ν2αl+2l−1B − ηνB .
  Let λS be the eigenvalue of W corresponding to eigen-directions u1, · · · , ur, and let λB be the eigenvalue of W corresponding to eigen-directions ud−r+1, · · · , ud. We know λS = νlS and λB = νlB . So we can write down the dynamics for λB as follows,
  λ̇B = lν l−1 B ν̇B =− lν 4αl+3l−2 B (1 + σ 2) + lν2αl+3l−2B − lην l B
  =− lλ4α+3−2/lB (1 + σ 2) + lλ 2α+3−2/l B − lηλB ,
  and similarly for λS we have
  λ̇S = −lλ4α+3−2/lS + lλ 2α+3−2/l S − lηλS .
  Dynamics for λB: We can write the dynamics on λB as follows,
  λ̇B = lλBg(λB),
  where g(λB) := −λ4α+2−2/lB (1 + σ2) + λ 2α+2−2/l B − η. 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 first compute the derivative of g as follows:
  g′(λB) =− (4α+ 2− 2/l)(1 + σ2)λ4α+1−2/lB + (2α+ 2− 2/l)λ 2α+1−2/l B
  =λ 2α+1−2/l B ( −(4α+ 2− 2/l)(1 + σ2)λ2αB + (2α+ 2− 2/l) ) .
  It’s clear that g′(λB) > 0 for λ2αB ∈ (0, 2αl+2l−2(4αl+2l−2)(1+σ2) ) and g ′(λB) < 0 for λ2αB ∈ ( 2αl+2l−2(4αl+2l−2)(1+σ2) ,+∞). Therefore, the maximum value of g(λB) for positive λB takes at λ ∗ B =(
  2αl+2l−2 (4αl+2l−2)(1+σ2)
  ) 1 2α
  and
  g(λ∗B) =− (
  2αl + 2l − 2 (4αl + 2l − 2)(1 + σ2)
  )2+ 1α− 1αl (1 + σ2) + ( 2αl + 2l − 2
  (4αl + 2l − 2)(1 + σ2)
  )1+ 1α− 1αl − η
  = 2αl(2αl + 2l − 2)1+ 1α− 1αl
  (4αl + 2l − 2)2+ 1α− 1αl (1 + σ2)1+ 1α− 1αl − η.
  As long as η > 2αl(2αl+2l−2) 1+ 1 α − 1 αl
  (4αl+2l−2)2+ 1 α − 1 αl (1+σ2)1+ 1 α − 1 αl , 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,
  λ̇S = lλSh(λS),
  where h(λS) = −λ4α+2−2/lS + λ 2α+2−2/l S − η. We compute the derivative of h as follows:
  h′(λS) = λ 2α+1−2/l S ( −(4α+ 2− 2/l)λ2αS + (2α+ 2− 2/l) ) .
  So h(λS) is increasing in (0, ( 2αl+2l−2 4αl+2l−2 ) 1 2α ) and is decreasing in ( ( 2αl+2l−2 4αl+2l−2 ) 1 2α ,∞). The maxi-
  mum value of h for positive λS takes at λ∗S = ( 2αl+2l−2 4αl+2l−2 ) 1 2α and we have
  h(λ∗S) = 2αl(2αl + 2l − 2)1+ 1α− 1αl
  (4αl + 2l − 2)2+ 1α− 1αl − η.
  As long as η < 2αl(2αl+2l−2) 1+ 1 α − 1 αl
  (4αl+2l−2)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 ∈ ( ( 2αl+2l−2 4αl+2l−2 ) 1 2α , 1).
  Overall as long as η ∈ (
  2αl(2αl+2l−2)1+ 1 α − 1 αl
  (4αl+2l−2)2+ 1 α − 1 αl (1+σ2)1+ 1 α − 1 αl , 2αl(2αl+2l−2)
  1+ 1 α − 1 αl
  (4αl+2l−2)2+ 1 α − 1 αl
  ) , we know W
  converges to cPS , where c is a positive number within ( ( 2αl+2l−2 4αl+2l−2 ) 1 2α , 1).
sections.heading: D ANALYSIS OF PREDICTOR REGULARIZATION.
sections.text: In this section, we study the influence of predictor regularization in a simple linear setting. In particular, we consider the same setting as in Section 3.2 except that we set Wp := (WW>)α + I. Theorem 6. In the setting of Theorem 1 except that we set Wp = (WW>)α + I. We have
  • when ∈ [0, 1+ √ 1−4η 2 ), as long as δ > ( max ( 1− √ 1−4η 2 − , 0 )) 1 2α , we have W con-
  verges to ( 1+ √ 1−4η 2 − ) 1 2α PS;
  • when ≥ 1+ √ 1−4η 2 , W always converges to zero.
  Proof of Theorem 6. We can write the dynamics of W as follows, Ẇ =W>p (−WpW (I + σ2PB) +Wa)− ηW
  =W ( −(I + σ2PB) ( |W |2α + I )2 + ( |W |2α + I ) − η ) .
  Let the eigenvalue decomposition of W be ∑d i=1 λiuiu > i , with span({ud−r+1, · · · , ud}) 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:
  λ̇B = λB
  [ −(1 + σ2) ( |λB |2α + )2 + ( |λB |2α + ) − η ]
  When η > 14(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:
  λ̇S =λS [ − ( |λS |2α + )2 + ( |λS |2α + ) − η ]
  =− λS ( |λS |2α + − 1− √
  1− 4η 2
  )( |λS |2α + − 1 + √
  1− 4η 2
  ) ,
  where the second inequality assumes 0 < η < 14 . We have
  • when ∈ [0, 1+ √ 1−4η 2 ), as long as δ > ( max ( 1− √ 1−4η 2 − , 0 )) 1 2α , we have λS con-
  verges to ( 1+ √ 1−4η 2 − ) 1 2α > 0;
  • when ≥ 1+ √ 1−4η 2 , λS always converges to zero.
sections.heading: E TECHNICAL TOOLS
sections.heading: E.1 NORM OF RANDOM VECTORS
sections.text: The following lemma shows that a standard Gaussian vector with dimension n has `2 norm concentrated at √ n.
  Lemma 8 (Theorem 3.1.1 in Vershynin (2018)). Let X = (X1, X2, · · · , Xn) ∈ Rn be a random vector with each entry independently sampled from N (0, 1). Then
  Pr[ ∣∣‖x‖ − √n∣∣ ≥ t] ≤ 2 exp(−t2/C2),
  where C is an absolute constant.
sections.heading: E.2 SINGULAR VALUES OF GAUSSIAN MATRICES
sections.text: 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(−t2/2) one has
  √ N − √ n− t ≤ smin(A) ≤ smax(A) ≤ √ N + √ n+ t
sections.heading: E.3 PERTURBATION BOUND FOR MATRIX PSEUDO-INVERSE
sections.text: 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 ∈ Rm×n : B = A+ E. Assume that rank(A) = rank(B) = n, then∥∥B† −A†∥∥ ≤ √2 ∥∥A†∥∥∥∥B†∥∥ ‖E‖ . The following corollary is particularly useful for us. Lemma 11 (Lemma G.8 in Ge et al. (2015)). Consider the perturbation of a matrix A ∈ Rm×n : B = A+ E where ‖E‖ ≤ σmin(A)/2. Assume that rank(A) = rank(B) = n, then∥∥B† −A†∥∥ ≤ 2√2 ‖E‖ /σmin(A)2.
title: TOWARDS DEMYSTIFYING REPRESENTATION LEARN-