Provable Benefits of Unsupervised Pre-training and Transfer Learning via Single-Index Models
Authors: Taj Jones-McCormick
Aukosh Jagannath
Subhabrata Sen
Venue: ICML 2025 spotlightposter (2025)
Keywords: Single-Index Model
Pre-training
Transfer-learning
sample complexity
Stochastic Gradient Descent
Primary area: theory->learning_theory
TL;DR: We prove that unsupervised pre-training can dramatically reduce sample complexity in single-index models

Abstract: Unsupervised pre-training and transfer learning are commonly used techniques to initialize training algorithms for neural networks, particularly in settings with limited labeled data. In this paper, we study the effects of unsupervised pre-training and transfer learning on the sample complexity of high-dimensional supervised learning. Specifically, we consider the problem of training a single-layer neural network via online stochastic gradient descent. We establish that pre-training and transfer learning (under concept shift) reduce sample complexity by polynomial factors (in the dimension) under very general assumptions. We also uncover some surprising settings where pre-training grants exponential improvement over random initialization in terms of sample complexity.

Provable Benefits of Unsupervised Pre-training and Transfer Learning via Single-Index Models
Taj Jones-McCormick
1
Aukosh Jagannath * 12 Subhabrata Sen" *3
Abstract
initialization. However, access to additional data opens the door to other avenues for initializing model weights.
Unsupervised pre-training and transfer learning are commonly used techniques to initialize train- ing algorithms for neural networks, particularly in settings with limited labeled data. In this paper, we study the effects of unsupervised pre-training and transfer learning on the sample complexity of high-dimensional supervised learning. Specifi- cally, we consider the problem of training a single- layer neural network via online stochastic gradient descent. We establish that pre-training and trans- fer learning (under concept shift) reduce sample complexity by polynomial factors (in the dimen- sion) under very general assumptions. We also un- cover some surprising settings where pre-training grants exponential improvement over random ini- tialization in terms of sample complexity.
A prominent setting with additional data is semi-supervised learning, where one might have an abundance of unlabeled data. Unsupervised pre-training has emerged as a popular strategy in this context (Devlin et al., 2019; Brown, 2020). The core idea behind pre-training is to train model (on the unlabeled data) on task that is related or, even bet- ter, necessary precursor to the supervised task of interest. Initializing from a pre-trained model, the hope is that the model will have already learned useful features from the unlabeled data and hence solve the supervised task with reduced sample complexity. The model trained in this way can then be used to initialize model for the supervised task in various ways e.g., by removing the final layer of the network and replacing it with output head relevant to the labeled task, such as a classification head.
Another prominent setup with additional data is transfer learning (Pan & Yang, 2009), where one has access to sam- ples from related supervised tasks. In this case, natural idea is to initialize the model weights on the target dataset using the model trained on the upstream task.
1. Introduction
The canonical pipeline of modern supervised learning as follows: given supervised data, (i) choose appropriate model/estimator (usually specified by deep neural net- work), (ii) choose a loss function and set up a suitable em- pirical risk minimization problem, and (iii) minimize this (possibly non-convex) empirical risk using stochastic gra- dient descent (SGD). Extensions of this "basic" approach have been successfully deployed to train -tate-of-the-art models in diverse domains. As deep learning models be- come larger and more complex, one has wrestle with the issue of model weight initialization during training. With- out additional information, one usually resorts to random
The precise scheme for unsupervised pre-training or trans- fer learning can vary significantly across applications. For example, BERT models employ self-supervised representa- tion learning by predicting "masked" tokens based on the observed ones (Devlin al., 2019). Similar self-supervised pre-training algorithms form core components in the train- ing of modern language models such as GPT (Radford et al., 2019; Brown, 2020; Achiam et al., 2023) and have attracted widespread attention recently. Many other forms of both pre-training and transfer leaning are applied in countless works spanning many different fields within machine leamn ing (Wang al., 2016; He et al., 2017; Devlin al., 2019; Schneider et al., 2019). Their popularity underscores the importance of understanding the effects of these different methods of weight initialization on supervised tasks.
"Equal contribution Department of Statisties and Ac- tuarial Science, University of Waterloo, Canada "Cheriton School of Computer Science, University of Waterloo, Canada 'Department of Statistics, Harvard University, United States of America. Correspondence to: Taj Jones- McCormick <Iejonessmecormickke @ @uwaterloo.ca. Aukosh Jagannath <a.jagannath @ uwaterloo.ca>. Subhabrata Sen <subhabratasene fas. harvar.edu...
The main goal this work is to build towards a theoretical understanding of the benefits of pre-training, particularly in terms of the effect on sample complexity for solving supervised learning tasks in high dimension that involve
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Benefits of pre-training via single-index models
optimizing non-convex losses. In general, characterizing the performance of neural networks on supervised learn- ing tasks is challenging problem. Recent works focus specific classes of problems such as single-index models learned with single-layer networks (Ben Arous et al., 2021) or two-layer networks (Lee et al., 2024), and characterize the number of samples required for recovery of the latent signal. We study the effects of distinct initializations on the sample complexity for closely related class of problems. Specifically, we show provable benefits of pre-training and transfer learning in terms of reducing sample complexity for single-layer networks. We also highlight surprising com- plexities and powerful benefits of pre-training- we discover simple scenarios under which one cannot hope to succeed with random initialization, but the problem can be solved easily with suitable pre-training. Finally, we demonstrate our findings empirically in finite dimensional settings with simulations.
Long et al. (2017); Wu et al. (2019); Sagawa et al. (2019); Gerace al. (2022) and the references therein for in- complete list). From a learning theoretic perspective, recent works study the generalization performance as a function of the discrepancy between the source and the target domains (Albuquerque et al., 2019; Ben-David et al., 2010; David et al., 2010; Hanneke & Kpotufe, 2019; Tachet des Combes et al., 2020; Zhao al., 2019).
In (Damian et al., 2022), the authors study the benefits of transfer learning in the setting of single/multi-index models, They keep the representation fixed across the source and target, and vary the link function across the two tasks. In contrast, we keep the link function constant (and assume that the link is known), and study settings with distinet (but correlated) representations in the source and target tasks.
2.2. Understanding Sample Complexity for single-index models
2. Related Work
Single-index models have emerged as popular toy-models for understanding the sample complexity of training of neu- ral networks. This is due to the fact that they are both nigh-dimensional and non-convex. From a statistical per- spective there has been work on the fundamental thresholds of inference in these problems (Barbier et al., 2019; Maillard et al., 2020b) and its landscape geometry (Sun et al., 201 8;; Maillard et al., 2020a; Dudeja & Hsu, 2018). From the per- spective of sample complexity, a substantial amount of deep work in this direction focused on the sample complexity of spectral methods or related algorithms, particularly in relation to the Phase Retrieval problem, (Candes et al., 2015; Barbier et al., 2019; Lu & Li, 2020).
We summarize some related works in this section, and com- pare these prior works with the contributions in this article.
2.1. Pre-training and Transfer Learning Theory
There has been significant recent progress in understanding the benefits of distinct unsupervised pre-training methods. In Lee et al. (2021), the authors provide rigorous evidence of the benefits of self- supervied pre-training (SSL). They explain the benefits of SSL via specific conditional inddpee- dence relations between two sets of observed features, given the response. In a related direction, Arora et al. (2019); Tosh et al. (2021a:b) examine the benefits of contrastive pretraining, while Zhai et al. (2023) examines the effects of augmentation-basen self-supervised representation learning- In Wei et al. (2021), the authors explore the benefits pre- trained language models, while Zhang & Hashimoto (2021), explores the inductive bias of masked language modeling by connecting to the statistical literature on learning graphical models. Finally, we highlight the work (Azar & Nadler, 2025), which exhibits provable computational benefits of semi-supervised learning under the low-degree likelihood hardness conjecture.
More recently there have been tight analyses of the sample complexity for online stochastic gradient descent from ran- dom initialization. In particular, it was shown in (Ben Arous et al., 2021) that the sample complexity in the online set- ting is characterized by the Information Exponent. Since then there has been a tremendous body of work around complexity exponents, such as the Information Exponent, Leap Exponent (Abbe et al., 2023), or Generative Exponent (Damian et al., 2024b). In particular, these exponents have enabled studies which contrast the performance of various learning paradigms such as Correlational Statistical Query (CSQ) versus Statistical Query (SQ) bounds (Damian et al., 2024b), feature learning versus kernel methods (Ba et al., 2024), better choices of loss function (Damian et al., 2024a), and the importance of data reuse (Dandi et al., 2024; Lee et al., 2024; Arnaboldi al., 2024). We note here that there has been quite lot of recent important work the case of multi-indes models which we do not explore here, see, e.g.. (Abbe et al., 2023; Bietti et al., 2023; Ren & Lee, 2024) for small selection of this rich literature.
The paucity of high-quality labeled data has directly moti- vated inquiries into the properties of transfer learning across diverse application domains. The recent literature focuses on several distinct notions of transfer learning (e.g., covari- ate shift (Heckman, 1979; Huang et al., 2006; Shimodaira, 2000), model shift (Wang& & Schneider, 2015; Wang al., 2014), target shift (Maity et al., 2022), conditional shift Muiionero-Candell et al., 2022; Storkey, 2008) etc) and de- velops distinct rigorous methods to ensure successful knowl- edge transfer in these settings (see Shimodaira (2000); ?);
To our knowledge, most of this work has focused largely
2

---

Benefits of pre-training via single-index models
on the setting of isotropic Gaussian features (though note Zweig et al. (2024); Pesce et al. (2023) for work on univer- sality), However, given that pre-training only has access to the features, one requires that the features have some correlation with the underlying spike. Inspired by the re- cent works of (Mousavi-Hosse 2024), we model this via a spiked covariance model. There are some additional recent works citedandi2024randomm cui2024asymptotics, ba2022high which analyze the high dimensional asymptoties of neural network training after single gradient step and also consider spiked and isotropic index models.
with initialization Xo. where V denotes the spherical gradi- ent with respect the parameters X.
We want to understand the benefit of pre-training and so we consider two methods of initializationss random and with pre-training. For random initializations we let Xo ~ Uniform (Sa-1). To model pre-training, we use Prin- cipal Component Analysis (PCA) on the unlabeled data (a)) obtain an estimate of the spike direction v. We then use this to initialize SGD for our supervised task, that is we let Xo ë. PCA is arguably the ideal starting point for a rigorous investigation into unsupervised pre- training. From a statistical perspective, PCA is the sim- plest dimension-reductiin algorithm; further, it's properties are well-understood high-dimensions (Bai & Silverstein, 2010). More importantly, has been shown that more advanced representation learning algorithms, such as re- construction autoencoders, also implement PCA in certain regimes (Bourlard & Kamp, 1988; Baldi & Hornik, 1989; Nguyen, 2021). In this light, we will restrict ourselves to PCA-based pre-training in this paper. We note that this approach is distinet from the spectral methods introduced for single-index models and, in particular, phase retrieval. There, the methods are supervised in that they use both knowledge of the label and features, whereas in unsuper- vised pre-training one only has access to the features.
et al., 2023; Ba et al.,
3. Pee-Preeiinn Fre-Training
3.1. Problem Set Up and Notation
We consider a single layer supervised network with specified activation function f- We consider Gaussian features with spiked covariance, where the spike is correlated with the parameter vector of interest.
Let the labeled data be (ye, at). with each (yi: a.) inde- pendent and identically distributed. We have the following relationship between a and y: y, = f(a vg) + 6,, for some e; independent of ag with mean 0 and finite fifth moment. The parameter vector wish to estimate is uo and f is known activation function. Throughout we assume that f is twice differentiable almost everywhere with f, f", f of at most polynomial growth. We would like the model we consider to capture the essence of pre- pre-training. To perform pre-training, one should have access to additional unlabeled data. We thus assume access to some unlabeled (a)) with af 2 az- In order for pre- pre-training to be useful, there is an implicit assumption that the unlabeled data contains some information in its structure that is related to the su- pervised task. We thus let 0g NO, la + \ou"), with V.00 = 7/2 E [0, 1 > 0 (71, À are dimension indepen- dent) and u, ng E gd-1 (the unit sphere Rd)). Thus, the features are Gaussian with spiked covariance, where the spike vector has some correlation with the unknown param- eter vector of interest Vo. In this way, our model captures the significance of pre-training by allowing the unlabeled feature data to contain information hidden in its covariance structure that is directly correlated with the solution of the supervised learning problem. The value of 17, measures the strength of this correlation. We define 1/2 /1 n/.
With these two methods of initialization, our goal is to contrast their respective sample complexity requirements for solving the supervised learning task (recovering the un- known vector vp). We are in particular interested in the high dimensional regime. We consider spherical SGD (hence forth referred to as simply SGD) with the total number of steps (and samples of (v:, a:)) given by N = agd. Thus the number of samples we observe is a function of the di- mension, and we are then interested in analyzing the high dimensional limit d - oc.
3.2. Main Results
In this section we state our main results. Firstly we state a few definitions and assumptions. We defer the proofs of all results to Appendix A. Throughout we will often refer to the 'population loss':
((X)= = =E[S(X- a) -y]²
= =E[[(x-0 a) a)-f(vo aa)2 BR
We note that there are two important directions of interest in this problem, namely uo and the residual direction of the spike vector, after subtracting off the projection onto on, that is [ mov.). Without loss of generality, we let the first two basis vectors be written as = vo and 22 = 1 (v 7)),) so that v = 1 + 7/22- We can rewrite the population loss which is a function of X, solely through the correlation of X with each of these directions.
Our goal is to estimate the unknown vector Uh with parame- ter vector X E 1 by using SGD on the following loss function: C(X,y) y) = [f(x a) = We use spherical gradient descent with siep-size 5/d to optimize parameters X, given by the following stochastic updates:
m/2
X.- 1 VV(()))))
X = [X V
3

---

Benefits of pre-training via single-index models
Let m)() X- eq 11 and m2 (X) X e2 = 2/2-
SGD on the given loss with N ad steps where a = (1). a52 o(1), we have that the sequence of initializations x = 0g the PCA estimators of Va obtained with N' =
then
((X) [B]((a X) f(a- ei e1)]² 1Ee
o' d unlabeled samples where o' = (1), are Effective.
[f (a1 + 0232 + v =t - -2zg)= - f(a1))² +E²
The theorem above states sufficient conditions on & and the correlation between v and by such that with pre-training, we are able to recover the true parameter vector with high probability in large enough dimensions. Further, we see that a (1), and hence our recovery with N = ad steps is
= 22)
where [a1, a2, 9] is jointy Gaussian with mean 0 and + a1, 02 and Eg² = 1. The covariance of a1, d2 is given by:
just beyond linearly many steps in the dimension.
1 iny
For our next two results we will work with activations f which satisfy
Xn1
1 + x1B)
82
Throughout we will use the term 'population flow' which is simply the discretized gradient flow on the population loss l.
()) - = E 0ga f(0) > 0
(1)
for ~ (O, Recalling the notion of information ex- ponent from Ben Arous et al. (2021) which is restated in the following section, (1) is equivalent to requiring the in- formation exponent 3 or greater for f and the information exponent less than or equal to 2 for f². We now consider second main result which considers recovery from random initializations.
daz1 Definition 3.1. A sequence of initializations (X.( )a21 is Effective for SGD with N agd steps of stepsize //din if mi (x + in probability as d -> 20.
A sequence of initializations is considered Effective if SGD with some number of steps and stepsize, initialized from the given sequence, recovers the true solution in the high dimensional limit. To that end, we say that a sequence of initializations is Ineffective if it not effective. We have defined Effective initializations for SGD based on the con- vergence of the SGD) process, and we can also consider these definitions for initializations of population flow, defined by the convergence of the population flow process in place of SGD.
Theorem 3.4. Suppose that I satisfies (1). Then for spher- ical SGD with N = ad steps where a < d, a5² = O1), for the sequence of initializations x² Un we have that m. (Xx) - O in probaability, as d - co.
Our second main result states that under appropriate moment conditions on f. we have that in order to recover the un- known parameter vector Up. we require at least. N = (dd)) samples. We emphasize that this result does not inform us of when we can recover the true parameter vector, only suf- ficient conditions for showing that we cannot recover with less than quadratic samples in the dimension.
Assumption 3.2. We say that Assumption 3.2 holds with point m² = (mmj, m²) E (2, 1) if there exists a point m², such that:
V .1 61 > 0 sgn(m2( (X)) 62 <0
We now state our third main theorem which is the most surprising result and brings to light the complexity of the single layer supervised network with Gaussian features and spiked covariance. Let H {X E g : |m (X) < r). Theorem 3.5. Suppose that satisfies (1) g ~ ~N(01)) 1). When 7 = 1. for spherical SGD with N = ad steps where a = =1), a52 < dd// we have that there exists some dimension independent value > 0 such that for all se- quences of initializations x(( E 11
for all X such that (m1(X), (x)) E (2 1) :
mi (X) 2 mi, [x)( < mmz. Where (2 1) is the
2-d unit ball.
We note that the ability of to meet Assumption 3.2 de- pends entirely on the choice of activation function It is clear that population flow, initialized within the rectangle defined by the points m and (1, 0) will recover the correct solution (see Figure for an example), hence (emphasizing that m² is independent of the dimension) this assumption provides a simple way to verify that a sequence of initial- izations is Effective for population flow. Now we state our main result regarding initializing with pre-training.
(d) then we have that: mi (Xx) + 0 in probability d - co. Further we have that for all e 0:
P(sup |m (X.) > >r
in probability as d + xx.
Theorem 3.3. Suppose that Assumption 3.2 holds with point m². Further in 2 mī and |122 S mž. Then for spherical
This result demonstrates the surprising fact that in the sim- ple scenario where the spike perfectly aligned with the
4

---

Benefits of pre-training via single-index models
unknown parameter vector, recovery is not possible from random initializations with any amount of data, given appro- priate stepsize. Even more surprising, not only is recovery not possible from random initialization, but even initializiing with some fixed correlation, can result in not only failure to recover but further a loss of the initial correlation. We also have that the maximum correlation attained over the course of SGD is contained in a ball around with radius slightly larger than the ball containing the initializations. Taking into account Theorem 3.3, we see that there exists problems such that pre-training can allow us to solve the problem in linear time, whereas the problem is unsolvable from random initialization in the given scaling regime, regardless of the amount of labeled data.
information exponent can be written in terms of the Her- mite coefficients of the activation function f. which can be similarly expressed as moment conditions on f. In light of the results in the isotropic Gaussian case, Theorems 3.3 and 3.4 may not seem that surprising. We emphasize here that the introduction of the spike to the covariance, makes the problem much more complicated. This is made clear by Theorem 3.5, where in contrast to the isotropic feature case where initializing with some fixed correlation puts us immediately into the descent phase allowing for recovery with near linear sample complexity (Ben Arous et al., 2021), with the introduction of the spike with any positive magni- tude, a local minima appears around m1 (X)= = 0 and hence with random initialization or initializing with some fixed correlation that is within the attractor region of the local optima, SGD tends the local optima, in effect learning nothing and perhaps destroying the initial information.
It is important here to note that a similar negative result was observed by HlousaviHHsseeint et al. (2023) for population gradient flow when fitting two-layer network with ReLU activation. There the authors propose to correct for this via preconditioning the gradient. By contrast, here we use this to illustrate the power of pre-training.
We quickly point out that Theorem 3.5 holds even with exponential data, stating that with the given stepsize, SGD does not recover the unknown parameter vector. We also note that the step-size specified in Theorem 3.5 is in fact more general than in 3.3 and 3.4. Hence, this is reasonable range of stepsizes for which one would expect to solve the problem with sufficiently many samples.
3.3. Discussion
Together, the first two theorems above tell us that for cer- tain activation functions f such that both the assumptions of Theorem 3.3 and the assumptions of Theorem 3.4 are met, we establish a significant separation of the required samples for recovering the unknown parameter vector. With pre-training, we can achieve convergence with N ad whenever = w(1). That is, we can recover with less than log-linear samples in the dimension. For random initializa- tions, we require N = (dd2) at minimum. This gives us a separation of dc for all Ç < 1. Theorem 3.5 not only high- lights the complex scenarios that can arise by introduction of spike vector but also serves as a demonstration of the pow- erful effects of pre-training. In the scenario where the spike vector is equal to the parameter vector, pre-training alone is sufficient for solving the problem (of course this information would not be available to any practitioner) and without pre- training. no amount of data is enough to solve the problem from random initialization under the given regime. How- ever, provided enough correlation between the spike vector and unknown parameter vector, the problem is solvable via pre-training with just over linearly many samples.
3.4. Meeting Assumptions
We now take moment to consider the assumptions in our theorems. Assumption 3.2 requires the existence of some point m² such that within the rectangle defined by this point and the global optima (X²()), m2())) = (1, the population dynamics are well behaved, tending to the global optima at (m, (x)) = (1,0) in linear time. When it comes to the moment conditions on f required to apply Theorems 3.4 and 3.5, one can easily check (see Lemma A.10 in the supplementary material) that the Her- mite polynomials with degree 2 3 satisfy them. While this claim does not extend to all linear combinations of Hermite Polynomials, it can be extended to linear combinations of Hermite Polynomials of degrees 3 or greater, with the added constraint that any two coefficients in the Hermite expansion that are exactly 2 degrees apart, must have the same sign, i.e. ( (g) (g)hh-2))))) 2 0. This provides class of functions which demonstrate our Theorems and thus the effects of pre-training.
We point out that Theorems 3.3 and 3.4 have established a lower bound on the benefit of pre-training. As we see from Theorem 3.5, there are scenarios which deviate significantly from the lower bound provided here. Past works (Ben Arous et al., 2021) have shown that when the features are isotropic Gaussian, the sample complexity is governed by a quantity called the information exponent, which is essentially the order of the first non zerz term in the Taylor expansion of the population loss. In the case of single-index models the
4. Transfer learning
We also consider related problem which we find lends itself better to the notion of transfer learning- We con- sider a related scenario under which we once again have labeled data (Js, a:) according to a single layer network with known activation function f and unknown parame- ter vector Up. We assume that f is differentiable almost
5

---

Benefits of pre-training via single-index models
everywhere with f; of at most polynomial growth. How- ever, we now consider the case of isotropic Gaussian fea- tures: az (0, Ia). Without the spike in the covariance, there is only one correlation variable of interest, namely mi (X) as previously defined. It can be shown that for f differentiable almost everywhere and f' of at most polyno- mial growth, the population loss & (X)) can be expressed as 6(m) (X)) E C², with (x(x) < O, Vz E (0, 1) and further, the sample complexity for solving this problem with SGD is well understood (Ben Arous et al., 2021).
> dec(k-2) (log d)a. Fur- from a >> d 2 d to a
ther in the case that Ç = i.e., initialize with fixed correlation independent of the dimension, we see that only a = w(1) is required, and hence we can recover in nearly linear sample complexity regardless of the information ex- ponent. This offers a substantial polynomial reduction in the event that the information exponent is large. We do note however that even in the case the of the information exponent 2, we still benefit from a complexity reduction of factor of (log d)². In the case of information exponent I, we do not benefit from transfer learning, however, in this case we already recover with nearly linear sample complexity, hence we have no need to perform transfer learning. We also note here that the reduction in sample complexity is the same as observed in Theorem 7 of Soousvvi-Hosseeee et al. (2023).
To introduce the notion of transfer learning, we consider the scenario where we have access to some sequence of vectors v((d) with v(d) 220
(44 7/-- We may consider these correlated
vectors v(d) to be a sequence of estimates of some vector correlated to up which were obtained via SGD on some related task. For the sake of analysis we are not concerned with how these correlated vectors are obtained, only the benefit provided by having access to them for initializing SGD. We are again interested in the sample complexity function of dimension and how this complexity is affected by initializing with transfer learning in contrast to uniform random initializations.
5. Examples
In this section we show some examples of our theorems and provide simulations to empirically verify our claims in large finite dimensions.
Recall the information exponent from (Ben Arous et al., 2021).
5.1. The Third Hermite Polynomial
We will now show how one can apply our theorems from the past section. As noted in the past section, for all Hermite polynomials with degree 3 or greater, all assumptions are met for applying Theorems 3.4 and 3.3, provided large enough correlation between and vo, ie for large enough mh Hence in order to apply our theorems to specific problem set up with some Hermite polynomial in place of ʃ and some values of X, 772 we must simply verify whether or not 7)1 is large enough. To do this, we must identify the region of Effective initializations for population flow and determine whether (n, n2) falls inside.
Definition 4.1. We say that population loss & has infor- mation exponent k if de e ((-1( 1)) and there exist C, 0 such that:
d'd (0) 0 do 2 (0) S -c -c<0 0, d' dm (m)
for 1 S <<ck ,
SC.
We now state result concerning transfer learning with isotropic Gaussian features as described above. The details on sample complexity for this model are well understood from the work of Ben Arous et al. (2021), which allows us to easily analyze the effects of transfer learning.
Note that taking f to be any polynomial function, we solve for (XX) explicitly, by expanding and computing the mo- ments of Gaussian random variables. After computing the explicit loss, one can compute the spherical gradients with respect to Fi and x2 and analyze their signs in order to iden- tify a value of m to apply Assumption 3.2. Below we plot the phase diagram for population flow when f(x)) = x²3 -3r, the third Hermite polynomial. We identify point m² to Apply assumption 3.2.
Theorem 4.2. Let k 2 be the information exponent of
. -to = nd = 5(d-5), with Ç E [0, 1/2). Then for spherical SGD with N = ad steps with a > (p((k-2) (log d)2/[c>0) and -1 & < a 1/2 X, = v,
6. Let v(d)..
we have that: mx(X) + in probability as d > co. Here
1 l5 0) the indicator function, taking value 0 if < 0 and I otherwise.
5.2. Simulations
The proof of this theorem follows almost exactly from (Ben Arous et al., 2021),, noting that their arguments still hold under slightly different initializations. Contrasting this theorem with the results of (Ben Arous et al., 2021), we notice that when p(d)_ -/6 = Na = O(d-5) for some < E (0, 1/2) and the information exponent is 3 or greater, we benefit from a polynomial sample reduction
We conduct a few simulations to empirically demonstrate our claims in finite dimensions. In the first simulation, we consider letting f'xx)) = x² 3x and setting À = 1, = 0.45. We then conduct SGD from both random initializa- tions and from estimates of v obtained via PCA. We use dimension d 1000 and let SGD run for 3 dd 1, 500, 000
6

---

Benefits of pre-training via single-index models
1 steps of size 1oa (10, 000, 000)-1 We select the param- eters such that we would expect to be able recover the true parameter vector from random initialization had we been in the case À = O. We determine this scaling based on the results of Ben Arous et al. (2021) and some experimenting. See Figure 1.
06(21,32) 221/ee1 + 22/11/p222 + ((1(()) xr)
Om
02)) 221pex2 + 22m1n2cr + ((1(21) 22)
Org
We also perform simulations under the setting 0.5 and ?1 = in Theorem 3.5. We then perform SGD with the dimension, step size and number steps given above, only we consider initializing uniform randomly, conditional on fixing the correlation mi (Xo)= = 0.1. See Figure 1.
for some positive constant c. We remind the reader here of the definition (x), (m), m²(X)). Analyzing the linearized system, we see that the first-order terms are orthogonal to (and point towards) the line L = {(21, x2) 21 112 x2), and are equal to 0 on L. Letting T/ (x1, 222) m
measure the distance of (22) to L, we have that the first terms exceed higher order terms when outside the set C = { X22 Ty (I1, > e1 (01,,)))1 for some constant C1. This set provides a cusp surrounding the line L. To prove our result, we carefully construct stopping times in order to observe the process over specific regions of the space, such as the Cartesian quadrants and positioning rel- ative to the set C. We show that in quadrants 1 and 3, the process tends to quadrant 2 or quadrant 4. Once in these quadrants we bound the distance of the process to the line L, ultimately ensuring the process gets close to L. Once this happens we show that the given sample complexity is not sufficient to leave some fixed ball around the origin. The proof of this theorem is the most involved. While this idea can be understood at the population level, the extension to SGD requires controlling the variance of the martingale along various important directions with martingale inequali- ties.
6. Proof Ideas
We make use of the 'bounding flows' approach from Ben Arous et al. (2020b), Ben Arous al. (2020a), and Ben Arous et al. (2021). This approach was applied to Single-Index Models in (Ben Arous et al., 2021). The key difference here is that the populations dynamics here cannot be reduced to a 1-dimensional correlation variable but a 2-dimensional vector of correlations. As such more delicate analysis of the phase portrait and the martingale fluctuations are involved.
6.1. Theorem 3.3
Our first main theorem, provides sufficient condition to check when pre-training for initializing SGD can recover the unknown parameter vector in almost linear time. The proof of has two main components. First note that under Assumption 3.2, there exists a rectangle that, when initial- ized within, the population dynamics will find the global optima. Next Lemma A.3 shows that when initializing with some fixed initial correlations, SGD behaves like the pop- ulation dynamics, which under Assumption 3.2 converge to the global optima when initialized in the rectangle. The second part of the proof of Theorem 3.3 is simply applying a few well-known facts regarding PCA in high dimensions (see Appendix B). These facts along with the assumption on the strength of the correlation of the spike to the un- known parameter vector, tell us that when initializing with pre-training. we find ourselves in the region of effective initializations as given by Assumption 3.2 and hence we can recover the unknown parameter vector in approximately linear time.
6.3. Theorem 3.5
When ? = 1 (and hence 7/2 0) as in Theorem 3.5, the system no longer depends on 12 (Recall the definition of o in the start of section 3.2). We then have the following 1-d system:
00(11, 221 + O(I((11, x2) Oz
From this system we there is local optima m, (X) O, whose attractor region is fixed and not dependent on the dimension. Hence the remainder of the proof of Theorem 3.5 is showing that in the high dimensional limit, the ran- domness of SGD is insufficient to escape the attractor region in the given scaling regime.
6.2. Theorem 3.4
We consider the Taylor expansion of the population loss in the two correlation variables of interest around O. Recall that
7. Conclusion
In this paper we consider natural statistical models for which one can analyze the effects of pre-training and transfer learn- ing. Namely single-index models with Gaussian features with spiked covariance and isotropic covariance. We ana- lyze the ability to recover the unknown parameter vector
random uniform initializations yield initial correlations on
the order of m1 (Xu), (Xo) = 0 ((22) (Vershynin, 2018). Under the Assumptions on of Theorems 3.4 and 3.5 we have the following system for the population dynamics:
7

---

Benefits of pre-training via single-index models
1.00
0.50
0.4
0.25
Mg00
Mx(X)
0.25
-0.4
-0.25
0.00 M2(X))
0.25
1.00
-1.00
0.50
M2(X))
(a)
(b)
Legend Random initia ization
Legend Random Initializatiinn
M1(Xg
Mi[X)]
Pre-trein
Pre-trair
Fixed Initie izzatio 1
Inttialization
Fixed Initia izattion
Fixed Initialization
0.0
-0.2
0.0
0.2
0.4
0.5
3.0
1.0
12
1.4
0.2
0.4
0.5
0.0
1.0
12
1.4
106
106
(c)
(d)
Figure The figures on the left correspond //I 0.45, and right correspond f/i 1, A = 0.5. Each figure displays SGD runs. One random initialization, initialization via pre-training with PCA and fixed initializations. The left figures feature fixed initializations (mi (Xo), m2()))) = (0.25, / 0.252) and (m1 (Xo), m2 (X))) = (0.25, -0.75). We observe that pre-training results in finding the global optima (m (Xo), m2((Xo) (1,0) and random initialization makes little progress. Fixing the value of mi (.o). we also observe very different behaviors simply by varying the value of m²(XNu, highlighting the complexity that arises when introducing second dimension to the population dynamies. In the case of = 1 the fixed initializations mi (Xo) 0.25. In addition to noticing the behavior of random initialization versus pre-training, we observe very different behavior for fixed initializations. For small enough my (Xo) SGD tends towards the local optima (X) suggested Theorem 3.5. For larger enough mi (Xa), SGD tends towards the global optima.
(Xo) 0.1 and
in these models using stochastic gradient descent and the required sample complexity in the high dimensional regime, from both random initializations and with pre-training transfer learning. In both scenarios we prove polynomial separation in the sample complexity as function of the dimension required solve, for class of functions which contains the Hermite polynomials of degree 3 greater. We also highlight the complexity of amalyzing recovery der single-index model with Gaussian features and spiked covariance, by highlighting a simple case (the case where the spike vector is equal to the parameter vector) in which a local optima arises and traps random initializations. This paper contributes to the growing body of work attempting add theoretical justification for common practices pre-
training and transfer learning.
A Acknowledgments
A.J. acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada Research Chairs programme, and the Ontario Re- search Fund. Cette recherche a été enterprise grâce, partie, au soutien financier du Conseil de Recherches Sci- ences Naturelles et en Génie du Canada (CRSNG), [RGPIN- 2020-04597, DGECCR-0000-00199) 90999999-00 er du Programme des chaires de recherche du Canada. S.S. gratefully acknowl- edges support from NSF (DMS CAREER 2239234), ONR (N00014-23-1-2489) and AFOSR (FA9950-23..) 23-1-0429. T.J
8

---

Benefits of pre-training via single-index models
acknowledges the support of the Natural Sciences and Engi- neering Research Council of Canada (NSERC).
/ /ppppple.math.eth...tttot Bandei www..nnneessreeeeee-------aaaaaa
Barbier, J., Krzakala, F,, Macris, N., Miolane, L., and Zde- borová, L. Optimal errors and phase transitions in high- dimensional generalized linear models, Proceedings of the National Academy of Sciences, 116(12)5551155460 2019.
Impact Statement
This paper focuses on advancing the theoretical foundations of Machine Learning, and we not anticipate significant socictal consequences from our work.
Ben Arous, G., Gheissari, R., and Jagannath, A. Algorithmic thresholds for tensor pea. The Annals of Probability, 48 (4):2052-2087) 2020a.
References
Abbe, E., Adsera, E. B., and Misiakiewiez, T. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Anmual Conference Learning Theory, pp. 2552-2623. PMLR, 2023.
Ben Arous, G., Gheissari, R., and Jagannath, A. Bounding flows for spherical spin glass dynamics. Commmiccations in Mathematical Physics, 773:1011-1445 2020b.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. I.., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXivv23033 087744 2023.
Ben Arous, G., Gheissari, R., and Jagannath, A. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learn- ing Research, 22(106)11-51, 2021.
Albuquerque, I., Monteiro, J., Darvishi, M., Falk, T. H., and Mitliagkas, I. Generalizing to unseen domains via distribution matching. arXiv preprint arXiv:/91 19911.0000444 2019.
Ben-David, S., Blitzer, J., Crammer, K., Kulesza, A., Pereira, F., and Vaughan, J. W. A theory of learning from different domains. Machine learming, 79:151-175, 2010.
Arnaboddi, L., Dandi, Y., Krzakala, F. Pesce, L., and Stephan, L. Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405.15453 2024.
Bietti, A., Bruna, J., and Pillaud- Pillaud-Vivien, L. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:23100 10. 19.79333, 2023.
Arora, S., Khandeparkar, H., Khodak, M., Plevrakis, O., and Saunshi, N. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:/902.09222, 2019.
Bourlard, H. and Kamp, Y. Auto-association by multilayer perceptrons and singular value decomposition. Biological cybernetics, 59(4):291-294, 1988.
Brown, T. B. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
Azar, E. and Nadler, B. Semi- supervupe sparse gaussian classification: Provable benefits of unlabeled data. Ad- vances in Neural Information Processing Systems, 2025,
Candes, E. J., Li, X., and Soltanolkotabi, M. Phase re- trieval via wirtinger flow: Theory and algorithms. IEEE Transactions Information Theory, 61(4):1985-2007, 2015.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, 7.., and Wu, D. Learning in the presence of low-dimensional structure: spiked random matrix perspective. Advances in Neurul Information Processing Systems, 36, 2024.
Damian, A., Lee,, J., and Soltanolkotabi, M. Neural net- works can learn representations with gradient descent. In Conference on Learning Theory. pp. 5452. PMLR, 2022.
Bai, Z. and Silverstein, J. W. Spectral analysis large di- mensional random matrices, volume 20. Springer, 2010.
Baik, J., Ben Arous, G., and Péché, S. Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. 2005.
Damian, A., Nichani, E., Ge, R., and Lee, J. D. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2024a.
Baldi, and Hornik, K. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58, 1989.
Damian, A., Pillaud-Vivii L., Lee, J., and Bruna, J. Computational-statee stattsssica gaps in gaussian single-index models. In The Thirty Seventh Anmual Conference on Learning Theory, pp. 1262-1262. PMIR, 2024b.
Bandeira, A. S., Singer, A., and Strohmer, T. Math- ematics of Data Science. 2020. URL https:

---

Benefits of pre-training via single-index models
Dandi, Y., Troiani, E., Arnaooldi, L.., Pesce, L.., Zdeborová, L.,, and Krzakala, F. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. International Conference on Machine Learning, 2024.
Long, M., Zhu, H., Wang, J., and Jordan, M. Deep transfer learning with joint adaptation networks. In International conference machine learning, pp. 2208-2217. PMLR, 2017.
Lu, Y. M. and Li, G. Phase transitions of spectral initial- ization for high-dimensional non-convex estimation. In- formation and Inference: A Journal of the IMA, 9(3): 507-541, 2020.
David, S. B.. Lu, T., Luu, T., and Pál, D. Impossibility theorems for domain adaptation. In Proceedings of the Thirteenth International Conference on Artificial Intelli- gence and Statistics, pp. 129-136. JMLR Workshop and Conference Proceedings, 2010.
Maillard, A., Ben Arous, G., and Biroli, G. Landscape complexity for the empirical risk of generalized lin- ear models. In Lu, and Ward, R. (eds.), Proceed- ings of The First Mathematical and Scientific Machine Learning Conference, volume 107 of Proceedings of Machine Learning Research, pp. 287-327. PMLR, 20- 24 Jul 2020a. URL https: / http::::/pccceedeeeeeeeo .mlr. press/ /v1 07/ / ma i 11a.d20a..hmmm.
Devlin, J., Chang, M.-W.. Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, volume 1, pp. 4171-4186, Minneapolis, Minnesota, 2019. Association for Computational Linguistics.
Maillard, A., Loureiro, B., Krzakala, F., and Zdeborová, L. Phase retrieval in high dimensions: Statistical and computational phase transitions. Advances Neural In- formation Processing Systems, 33:1 13:11071-1108, 2020b.
Dudeja, R. and Hsu, D. Learning single-index models in gaussian space. In Conference On Learning Theory. pp. 1887-1930. PMLR, 2018.
Gerace, F,, Saglietti, L., Mannelli, S. S., Saxe, A., and Zdeborová, Probing transfer learning with a model of synthetic correlated datasets. Machine Learning: Science and Technology, 3(1:0015330, 2022.
Maity, S., Sun, Y., and Banerjee, M. Minimax optimal approaches to the label shift problem in non non-parametric settings. Journal of Machine Learning Research, 23(366): 1-45, 2022.
Gross, I. Gronwall's ' inequality and its applications. The Journal of Mathematical Analysis and Applications, 20 (3):359-333) -00, 1967.
Mousavi-Hosseini, A., Wu, D., Suzuki, T., and Erdogdu, M. A. Gradient- based feature learning under structured data. Advances in Neural Information Processing Sys- lems, 36:71 (4771449-71485, 2023.
Hanneke, S. and Kpotufe, S. On the value of target data in transfer learning. Advances in Neural Information Processing Systems, 32, 2019.
Nguyen, P.-M. Analysis of feature learning in weight-tied autoencoders via the mean field lens. arXiv preprint arXiv:2102.08373: 2021.
He, K., Gkioxari, G., Dollár, P., and Girshick, R. Mask r- cnn. In Proceedings of the IEEE international conference on computer vision, pp. 2961-2969, 2017.
Pan, S. J. and Yang. Q. A survey on transfer learning. IEEE Transactions on knowledge and data engineering, 22(10): 1345 1359, 2009.
Heckman, J. Sample selection bias as a specification error. Econometrica: Journal of the econometric society, pp- 153 161, 1979.
Pesce, I.,, Krzakala, F., Loureiro, B., and Stephan, L. Are gaussian data all you need? the extents and limits of uni- versality in high-dimensional generalized linear estima- tion. In International Conference on Machine Learning, pp. 27680-27708. PMLR, 2023.
Huang. J., Gretton, A., Borgwardt, K., Schölkopf, B., and Smola, A. Correcting sample selection bias by unlabeled data. Advances in neural information processing systems, 19, 2006.
Cathonero-Candllae J., Sugiyama, M., Schwaighofer, A., and Lawrence, N. D. Dataset shift in machine learning. Mit Press, 2022.
Lee, J. D., Lei, Q., Saunshi, N., and Zhuo, J. Predicting what you already know helps: Provable self-supervised learning. Advances in Neural Information Processing Systems, 34:309-323, 2021.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Lee, J. D., Oko, K., Suzuki, T., and Wu, D. Neural network learns dimesinnlli polynomials with sgd near the cettmmatton-thhorrete limit. arXiv preprint arXiv:2406.01 01 581 2024.
Ren, Y. and Lee, D. Learning orthogonal multi-index models: A finegnrained information exponent analysis. arXiv preprint arXiv:24 10. 10.000709 2024.
10

---

Benefits of pre-training via single-index models
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case gen- eralization. In International Conference on Learning Representations (ICIR), 2019.
Williams, D. Probability with Martingales. Cambridge University Press, 1991.
Wu, Y., Winston, E., Kaushik, D., and Lipton, Z. Domain adaptation with asymmetricalll- distribution align- ment. In International conference on machine learning, pp. 6872-6881. PMLR, 2019.
Schneider, S., Baevski, A., Collobert, R., and Auli, M. wav2vec: Unsupervised pre-training for speech recog- nition. arXiv preprint 1rXivv1904.05866 2019.
Zhai, R., Liu, B., Risteski, A., Kolter, Z., and Ravikumar, P. Understanding augmentation saaeeetttttoobadaedddd self-supervised repre- sentation learning via rkhs approximation and regression. International Conference on Machine Learning, 2023.
Shimodaira, H. Improving predictive inference under covari-
ate shift by weighting the log-likelihood function. Jour- nal of statistical planning and inference, 90(2):227-244, 2000.
Zhang, T. and Hashimoto, T.
On the inductive bias of masked language modeling: From statistical to syntactic dependencies. arXiv preprint arXiv:2104. 1445544004,05644, 05699, 2021.
Storkey, A. When training and test sets are different: char- acterizing learning transfer. 2008
Zhao, H., Des Combes, R. T., Zhang, K., and Gordon, G. On learning invariant representations for domain adapta- tion. In International conference on machine learning, pp. 7523-7532. PMLR, 2019.
Sun, J., Qu, Q.. and Wright, J. A geometric analysis of phase retrieval. Foundations of Computational Mathematics, 18:1131-1198, 2018.
Tachet des Combes, R., Zhao, H., Wang. Y.-X., and Gordon, G. J. Domain adaptation with conditional distribution matching and generalized label shift. Advances in Neural Information Processing Systems, 33: 13:19976-19288 2020.
Zweig, A., Pillaud- Vivien, L., and Bruna, J. On single- index models beyond gaussian data. Advances in Neural Information Processing Systems, 36, 2024.
Tosh, C., Krishnamurthy, A., and Hsu, D. Contrastive esti- mation reveals topic posterior information to linear mod- els. Journal of Machine Learning Research, 22(281):
1-31, 2021a.
Tosh, C., Krishnamurthy, A., and Hsu, D. Contrastive learn- ing, multi-view redundancy, and linear models. In Algo- rithmic Learning Theory, pp. 1179 1179-1206. PMLR, 2021b.
Vershynin, R. Digne-Dimesiooall Probability: An Introduc- tion with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018
Wang, J. X., Kurth-Nelson, Z., Tirumala, D., Soyer, H., Leibo, J. Z., Munos, R., Blundell, C., Kumaran, D., and Botvinick, M. Learning to reinforcement learn. arXiv preprint arXiv: .57663, 2016.
Wang, X. and Schneider. J. G. Generalization bounds for transfer learning under model shift. In UAI, pp. 922-931, 2015.
Wang, X., Huang, T.-K., and Schneider, Active transfer learning under model shift. In International Conference on Machine Learning, pp. 1305-1313. PMLR, 2014.
Wei, C., Xie, M., and Ma, T. Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning. Advances Neural Information Processing Systems, 34:16158-16 16177, 2021.

---

Benefits of pre-training via single-index models
A. Proofs of Main Results
For convenience we define the sample wise error as (X. - (Xx,)) throughout our proofs.
- - (X).. We first prove lemma which used
Lemma A.1. There exist constants C1 C2 > such that the following moment bounds hold uniformly in d:
E (z y) 0a) S C₁
sup -,0
y) 2 C₁ (1((2 V
sup #E84-1
Proof. See the proof of proposition B.1 (Ben Arous al., 2021) and not that similar arguments apply when the features have spiked covariance.
A.1. Proof of Theorem 3.3.
We first prove Theorem 3.3.
Theorem 3.3. Suppose that Assumption 3.2 holds with point m². Further 9h 2 my and [12al S m.z Then for spherical SGD on the given loss with N a d steps where a = w(1), a6² = o(1), we have that the sequence of initializations x(d) = ugg the PCA estimators of obtained with N' = a' unlabeled samples where a' = w(1), are Effective.
Proof. The proof of Theorem 3.3 an immediate consequence of the sequence of PCA initializations being Effective for population flow as given by Proposition A.4 and the convergence between SGD and population flow as given by Lemma A.3 (below).
Lemma A.3. Suppose that Assumption 3.2 holds with point m Fix any point (m-nit, -m3ri) with mi bnit > m ma < m For sequence of initializations (m (Xx)), m2(X5 m2(X6))a21 converging to (mi .m²ir), spherical SGD on the given loss with the given initializations and N add steps where a =(1), ad² = o(1) yields the following:
- (m m2) (xxi))l22 +0
sup K
in probability as + co. Here X, is population flow with the same initialization.
We will show that the spherical projections are negligible (arbitrarily small with probability o(1)). We can then consider the linearized paths of (X,) m2( (X,)) and (m1 (X.), mx()))). Bounding their difference with Gronwall's inequality (Gross, 1967) and Doob's Inequality (Williams, 1991) control the martingale term.
Proof. Let V z denote the usual Euclidean gradient. Then consider the spherical gradient:
V = Vz - (V) X)X
That is the euclidean gradient projected onto the orthogonal space of X. We know that the population loss can be written as function of (X), m²(X)) (x1, 22) (recall that (X)(X), my = (I1, due to the without loss of generality
assumption that ti e and (122 - %)) 62). Hence:
06(1), dx
00(21,02)
V X
sup
sup -
v
X|12 S A
2
Orz
where A is some constant independent of d, using that E C² and X E gd-1 (compact). Additionally with the above, one can show that there exists K (independent of d) such that for X, Y Sdd1
( V s K|
Y|
(3)
12

---

Benefits of pre-training via single-index models
We will this fact later. Now let
r++ = | d
(V - VH(XI, 36))l1
8 ((IV(()))1B + (V ut)l
$ V
(4)
21+ 82
d )
(5)
In (4) we use that X, and ( (H(,,, 9r) are orthogonal as the gradient spherical. In (5) we use that for
> O, vl+uuuu 1 1+0 and the spherical gradient has bounded norm shown in (2). L, || (X,, /dl. noting that while Le is random, it's expectation is bounded by a constant independent of d by Lemma A.1. Note that the quantity we have defined rs+1 is simply the radius of X, after the gradient update, but before projecting back onto the sphere. We have that |r++1 - 1| S 8² 4 + (z:). This bounds the distance between X₂ with itself, had not been projected onto the sphere after the last gradient update:
||, - (Xt-1 3 2 (0 (. 1I-x))))
I I
4
| 112 S 82
(
7
7
By iterating this bound, we have the following:
sup
|X
(Xo-
8 Σ VH E (x.)) +
11) S
N-1
Σ
82
A 1
+
)
E
By Markov's inequality, the probability that the right hand side greater than some is:
N
L > >
P
82
sup EL,)
d²
d
ISN
1=0
Which is o1) given at o(1) and supigN EL, co by Lemma A.1. It thus suffices to consider the linearization of (X)) which for the two correlation variables of interest, we denote:
Y, = m(oo
Σ (X)( - VH(,.. 1=0
where m(x) = (m, mz(x).. We note that this linearization (Y:) is not the same linear SGD. This process is equivalent to performing spherical SGD, but adding back all of the projection vectors each stage, that were used to map X, to the sphere after each gradient update. Which is also not equivalent to doing regular gradient descent with spherical gradients. Redoing the above computations with respect to X, one would see that deterministically, for large enough d we have that X₁ also within e of Yú, the linearization of gradient flow, given by:
Y = mx
()()))
Hence to prove our result it enough to show the convergence probability between Y; and Y,. To do so, let us consider the martingale term given by A Et Di (H(,,, y1). Applying Doob's inequality with p 2 that
2Ca82
P(sup N
M (X, vl)l > 2)
2d = o(1/d)
0=0
13

---

Benefits of pre-training via single-index models
for some constant C, independent of the dimension. To show that sup,SN 1|2
Yell2
first consider for some fixed
T, the quantity; supigr5-1d || Yell2. Then on the set { - Y: || - Yell2 < e}. for all t S T
8
|| Yell2 -
(mxx)( v+(X)))]
m V (X,, v1)ll
8 1-1
K
KIImiXx xx))12 + 7
+7
(V,
1=0
==0
1-1
+ 8
<TTK<+ +
Σ
- 11;1l
Σ
d
k=0
s (2TK + 1)6
Σ
K| - Vill2
==0
with probability 1 o(1), by applying Doob's inequality and noting the use of (3). Thus applying the discrete Gronwall inequality (Gross, 1967), obtain:
sup 15T6-id
||Y - Y212 S (2TK + 1)
(6)
For any 7 > 0 the above can be made less than 7/5 by choice of <(>; T). We now let T be such that
[I (Y)), m2(())) 0)l2 < </5
sup T8- T8-1 dSIEN
This T exists constant (which is independent of d) result of Assumption 3.2 and the fact that (m,(Yd)) m2 (22)) + (minit, minit). To better understand this, recall that with constant initialization, the 2-d population dynamics of mi and my are otherwise unaffected (other than through stepsize and number of steps) by the dimension. By Assumption 3.2, we have that the population dynamics converge to the intended solution at (mi, m2)) = (1, as d + co. Now consider:
8
6 Σ m m(VH(X,,))lll m(V((X.)) (1,0)12 + 112
sup T8-1dSIEN
|| (1,0)]l2 |
Σ
d
i=T8-i
i=
8
1-1
$ 25/5 + 11
Σ
m
(X))l1
i=78-1d
The above bound comes from applying triangle inequality separate out the martingale term and then noting that at time T8-1 'd we have Y, is within 7/5 of Y, which is within 7/5 of (1, 0) and the gradient with respect to the population loss only moves Y, closer to (1,
sup T8 dSISN
| V:112
< sup
(||Y (1,0)l2 + 1 - -11,0)l22
T8 T8-IN
8
1-1
52n
11
Σ m (VH(X,,))ll 1=T8-id
Sy
with probability 1 o()). This follows from another application of Doob's inequality onto the projection of the high dimensional martingale onto the two fixed directions of interest e 62 keeping in mind Lemma A.1. We have thus shown that for any 7 > 0, with probability 1 - o(1), that sup ( (X,,,X,), (x.))) - (m1 (X.), m2(()))))))))))))))))))))))))m ll2 < 7 which concludes the proof.
14

---

Benefits of pre-training via single-index models
Proposition A.4. Suppose Assumption 3.2 holds with point m². Further 1/1 2 mi and 1/2 mzz The the sequence of PCA estimators (6z)a21 of va each obtained with Na
d samples,a w(1), is Effective for population flow.
Proof. We will use well-known facts about PCA in high dimensions, which we recall for the reader's ' convenance in section of the appendix. If we consider using fixed linear portion of our samples for conducting PCA, we can choose the fraction 7 = d/N be any constant we desire. Choosing y sufficiently small, we can ensure both that y X² and ma
m1(0) (0) mi and
(ū) < m2- To this consider the following: Let A1 |m] 7//. A2 Imm 72/ c min miin(Ap. A2)2. We have:
1- - 7/2 || - 2112 | + ||01 - 20 < 2(1
1
+ 7/X2 + < 4 minl mmmnn1.....
for small enough 7. The second equality follows for large d using a well known result regarding the limiting correlation of u and ê as can seen Bandeira et al. (2020).
Thus by triangle inequality, we have that (+) 2 mi and [m² S m² for sufficiently large Letting the choice of y tend to 0, it is clear the limit of (m1(Xo))) m mx( exists and is equal to (m1(v), m²(e)) = (77, m2).
A.2. Proof of Theorem 3.4
Theorem 3.4. Suppose that f satisfies the following:
0² ag² f(g)
(g) EE) 0,
for g (N, 1). Then for spherical SGD with N ad steps where < d, axd2 (1), for the sequence of initializations
X ~ Uniform($4 we have that: |mm (XN) - 0 in probability, as + 00.
The overall strategy of the proof is show that for for any y 0, if X, (m1 (X2), m2().)) is in the d-1/6-ball, i.e. {z E R | < d d/1/6,, we have that X₁ enters the /d-1/66ball "times out' at f N, before it leaves the y-ball with probability 1 - o(1)
We want to analyze the population dynamics and we will consider doing so via a second order Taylor expansion of the population loss
0(21, x2) (af + 0222 + V 1 +2)) l(,1)] EE2
in (m)(X), (X)) = (x1. x2), around the origin. Noting that, under our assumptions on f. we can differentiate under the expectation, we perform the following computations:
0o(0.0)
dz
=
0o(0,0)
(2) -(a){()
Oug
(0,00)) =
(0, 0)
Duz 8 (0,00)
-
0r,02z
Evaluating the expectations we have:
a
((0,0) (2(
- f(a))
2(1 + xn))
dx
15

---

Benefits of pre-training via single-index models
Using that 01 1 g, Baj O, and Gaussian integration by parts: [f((z)a; var(a))BB((aa) for dj Gaussian
600, 0) = f(a) (
Using that (f (I by parts and given the following decomposition of 02:
2 Ef(a))a1 Which follows by another application of Gaussian integration A+
2m
(2
a2 = + An
a1 +
1 + 2n
1+ Xn/
where of 1 and var(at) 1.
8 0u7
(0, 0) 000000000000100000-0000000000000
(8)
= =2[(1 [0 =2[(1+ 2001::/0333330000000000000000000 -
(9)
(10)
From (8) to (9) above, we used the fact that Ef (g)f(())) = Egf(g) which simply follows from Gaussian
integration by parts. From (9) to (10) we used the fact that ( = {5(a)(g² 1) which follows from two applications of Gaussian integration by parts.
0 (0,0) 0) 1000000--0-0000---0000000--11--1
0m3
82
(00,0) 0) = =t
dr, Oxz
Using that (Bf(a))a10) =
Am BanEEE(o (a af by tower property and decomposition of 0z. We use the above calculations to 1+
compute the second order Taylor expansion of & around (0,
a
x2)1,22) = $00 0)
(0. 0).
(0,000
Oxz
1
1 02
0
0²
+ 2 dm
(0,0)=²
(0,000 +
(00, (0,00)212))
R(x1, x22)
Or3
Or 0z
where R(I1,02) = ((11(11,22)133 Which allows us to compute the derivatives of (() as follows:
00(11,
2 (0, 0)æ +
8²
(0000)= +
(0, 0)2z + 0(1((1,,22113
Oz
0m
Ox,Or,
=-2 Xn
+
1)]=1 + (/0 + -I12 ((I((1,, x2)1
16

---

Benefits of pre-training via single-index models
06(
8
8²
0() +
00,0)a2 +
(), + ((1
0x2
0m2
0x,002
+ { 1)]22
- + Xnf) "|a, +((l()),
Using our assumptions on [f) we have the simplified system:
36()))
( + 22 +(]()),)))
(11)
dz1
a (2) ( (0 12)
(12)
0z2
We recall that the spherical gradient correction term is higher order:
00(21. x22
06(21, x2) x2)
dx x1
Ox²
Analyzing the linearized system, given by (11) and (12), see that the first order terms are orthogonal to the line L = [tuy :t E RR [211 x22 I2 21 x1), where 2.22 (72 -m) On the line L, the first order terms are both O, and hence the magnitude of the first order effects tends as (æ1, x²2 tends the line The magnitude of the first order terms, exceeds the magnitude of all higher order terms when the distance of (x),))) the line L, exceeds cI) x2) for some constant
Before proceeding we define some notation. Let us introduce the four quadrants Q1 {xx,0) : 0, 0), Q2 x < 0, y > 0} Q3 = {(x, x < 0, < 0}, Q+ = ={ y) x > 0, < 0}. We will consider the variable ((x, = m = (X,,, mx,))) = (01, 22).-
We define two operators TLI and TL2 as follows: (X8) u/ PL where Pz is the orthogonal projection operator onto L. Tra is defined such that for X in Q4, (X) We define Tra -Tr, so that (x,2(X) > for X E Q2- m x and
We also define the 'left and right half-spaces' with respect to L follows: Hi {(11,
I2
m2
H, ((11, F2
- 21
We now define two additional operators T,t and T follows. Let U2 = (m, 12) and then let T+* 02 (I (X) T measures the signed distance X to the line L, with (X)) > when X E Q1. Note that T1,, T,- Ti are all linear operators.
With notation above now define the following set: C ((11) : (x1,
[2]. cl(( (z) For X $ C, we have that the first order terms in the gradient of the population loss, exceed the higher order terms.
Let Lj be the line {I11 x2) I2 = - / 22/g11 } if m > 72 or {(11. I2 = - 2m /n/r1) if < T2- In essence this line is the line half way between and whichever quadrant boundary is closer to L.. We define it's counter part Lī as the line with the reciprocal slope, i.e. the slope of Lý is over the slope of Li. Without loss of generality we assume 7n < 7/2- Similar to the 'left and right half-spaces' with respect to L defined above, we can define the left and right half-spaces of L and Li. H11 Hr1, H12 Hr2. Then we consider the set Q* (H1r n U O H11). We now notice that if we consider the set C intersect the 7-ball, that for sufficiently small 7, the intersection a subset of Q". For better understanding of the definitions given above see Figure 2. We will make use of the set C in Lemma A.7 below.
We now state three lemmas which we will use to complete our proof of Theorem 3.4. We defer the proofs of these lemmas until after the proof of Theorem 3.4. Lemma A.8 will tell us that when X, is in Q or Q3. it must leave, entering Q2 or Q+ Lemma A.6 will tell us that X cannot leave the b-bal through Q2 or Q3 without first exiting the quadrant, provided arrived small enough. Lemma 3.4 provides a bound on how far X can move away from the
before it's 's norm grows by 1/3 1/3
17

---

Benefits of pre-training via single-index models
line L before entering the set C defined above. Together these lemmas allow control the magnitude of X, before it enters the set C, and then show that once X, is in C it cannot escape the y--bll before re-entering the 1/2d-1/6-bbll ball..
0.6
0.4
0.2
0.0
-.2
-0.4
y-ball
{X1
x2}
L1, L2
-0.6
ac
Q"
-0.6
-0.4
-0.2
0.0
0.2
0.4
0.6
Figure visual to help guide understanding the definitions of L1, La and set clear that regardless m, by decreasing the vale of 7, we find that C intersect the b-ll subset of Q"
18

---

Benefits of pre-training via single-index models
-
d-1/6
2 d/
d-166 1/6
L
{x1
n22
n1
X23
Bound on distance L
d- 1/6
d-1/6
2d
2/61/6
Figure 3. A visual to help guide in the understanding of the proof of Theorem 3.4 and the use of the three lemmas. When initialized in Qa the d' d-ll--ball. X must leave Qa before leaving the 3/2d 1/6 ball. Suppose enters Q4. red line provides boundary X, cannot cross before entering the set (see Figure 2).
Lemma A.6. There exists some sequence Ka growing to infinity such that under the assumptions of Theorem 3.4 and restricting the set E, (supegn l Ei= (X)))(((XX))) < K 1/101 then
we have that if X,- E Q4, and |Xt- S
1|&ill2 n, Vt [t" min((Tn N)]
(13)
where TQz is the stopping time for the next time X, leaves Q4. On that event, the same statement is true with Q2 in place of Q4 as well. Further, the P(F.)) o(1).
19

---

Benefits of pre-training via single-index models
This lemma says that under the E1, for X₂ with norm less than d-11/10 and in Q1 X2 must leave the quadrant Q1 before it leaves the y-ball. Similarly, for X in Q2 with norm less than d d-1/10, X exits Q2 before it leaves the y-ball.
Lemma
X- T77(X) under the assumptions of Theorem 3.4 under the event: E2 (supign 11 Ei=1 VI, (M(H( < 1/3,, when Xt. E Hi, || S 2d 1/6 we have that:
A.7. Letting
|| S 11 d ////vvv Vt [t", min(re, T1/20-1/0; N)]
where TO is the next time X. enters the set C and 11/2d- 1/8 is the next time X, enters the 1/2d-1/6-bbn Further, the same statement holds when replacing T with T and H3 with Hy.
This lemma says that if X, is in the 2d 1/6 ball, then it's distance from the line L can only increase by up to d' 1/3 before it enters the set C or re-enters the 1 1/2d-1/6-bal -bll.
Lemma A.8. Under the assumptions of Theorem 3.4 and on the sets Ej 1 d 2/5), j = 2, when Xx- E Q3. we have that:
{supIen 1a E VH(X,) ezl <
10
|Xolla S 11 d 1/1 Vi 11", T05
Where Taz is the stopping time for the next time X, leaves the quadrant Q3. Further the above statement holds replacing Q3 with Q1.
This lemma says that the maximum amount the norm of X, can increase while in Q3 before leaving is d 1/5
Proof of Theorem 3.4. We intend to show that for any > 0, P(supigN (m,,, me(X.))) > 7) 0 as d - cc.
The proof of this result will make use of the three lemmas above to show that if X, is in the d d-1/l-ball, it will re- enter the d1/2-11 ball before it leaves the y ball. To make of the three lemmas above, we note that the sets considered in these lemmas: E1, Ez, Ei Eg are all probability 1 o(1) by Doob's inequality and hence is their intersection E = E, n Eg n E n Ez We now remind the reader that random initializations yield correlations on the order of m1 (Xo), (Xo) O (2) < 1 We now argue that under the set E for X, in the d' d - ball, we have
deterninistically that X; re-enters the 1 1/2d-1/6 ball or reaches = N before it leaves the 7-ball, hence by the Markov property, this will conclude the proof.
We will prove this in a case by case manner, firstly considering the case that X, is in the d' b/ll and in Q3-
By Lemma A.8, X, must leave Qs before can leave the 3/2d-1/6-ball.. Hence for X, to leave the 7-ball, it must first exit the quadrant. So suppose X, enters Q4 (the case where X, enters Q2 follows by similar argument). Since it entered small enough (by small enough y-ball before exiting the quadrant. Additionally, by Lemma A.7, it's distance from L cannot increase by more than d/1/3 before entering the set C. Now also notice that for X, to enter Qa \ 2d-1/6-ball, it would require it's distance from L to increase by a quantity that is order 1/6 > d 1/3 which cannot happen before X, enters C. Thus X, cannot exit the 7-ball through Q3 until entering C. Xx may re-enter Q3 or even Q2. but any entry to Q4 or Q2 requires || 1l < d and hence by by lemma A.6, X, cannot exit the b-all through quadrants Q2 or Q4. Hence X, cannot exit the y-ball through any quadrant without re-entering the 1 1/2d-1/6__ -ball, timing out first entering the set C. So suppose X, enters C in Q4. X, still cannot leave the 7 ball until first leaving the quadrant. But now recall that C is contained in Q" and hence to leave the quadrant without first re-entering the 1/2d-1/6_ 1/lllall,, requires X, to exit Q", and further move a distance on the order of d' 1/6 away from the line, which by lemma A.7 cannot happen. This completes the proof for the case where X, is in Q3. Note that the case where X, is in Q1 follows by symmetric arguments. Now finally notice that the cases where X, is in either Q2 or Q4 follow by the same arguments as above, noting that the Q1 and Q3 cases reduce to the Q2 and Q4 cases.
mean ||X2lla < 1/10, allowing to invoke lemma A.6), by Lemma A.6, it cannot leave the
Proof of Lemma A.6. We will prove the case where X E Q1 using Tıi and simply note that the case of X E Q2 with T2e follows by the same arguments.
20

---

Benefits of pre-training via single-index models
Note that for any / and any value 7 > 0, Je, such that for Q1, 11112 7= (2((2) ey. We will now show that the event Pa contains the following event:
(TLI e, Vt [t", min(TQ, NN]
Observe that for such E [t", (in(TT;;NN)
- (O - 8 -
TT -
<TnI(X.)) Tril d Σ m(vv(X,))
m (VH(()))
Note that this follows from the fact that Tia X, > 0 for as long as X, E Q4. and that the radius of X₁ plus a gradient step, is deterministically greater than one as it is spherical gradient. Now given X, E Q1 and recalling that the first order terms in the gradient are orthogonal the line L, we have that
(X)))) most second order in Tra (X,). Hence:
1-1 Σ
1-1
(r)(X)) & (X +
2 - Tinl Σ
m m(VH(X))))
For some constant Note that the event(supien 1 (supsenssssen 12 E (TL(m(((XX))] < Ka/Vd). Thus we have that:
Ei (TL(M(((XX))
KaaVVd contains the event
T +
Σ m=t
T Ka/Va,vt It", min(To, NN]
follows from this inequality that Tr. TriX, Ey 4l E [t", (i",min(T,;" N)). To see this, observe that by the discrete Bihari-LaSalle inequality C that for some constant cj:
Ka Va (1 cj0Kad -
(X)(X) S
so long as (1,() S Kg (which is true so long as d to s K2). For sufficiently large d, the right hand side above is va
smaller than 6y for all t < where:
i i=
for some > 0 sufficiently small, provided Ka is growing slower than d2/2-5 for some < > 0.
Thus if we choose Ka be diverging appropriately slowly, we have that i N (recalling that va o(Vd) and Va 0(8-1)) We thus choose Ka such that Tr. (Xx.) S d to S Vd and Ka S dz/2-5 (for some Ç > 0), for example
K₂ = 11/2-1/111 then Ka/ = 1/11 > 1/10 and further P(sup 1 Ei= (TTm((((()))) < Ka/Vd) ( (o)))//K))00 = 1 - o(1).
Proof of Lemma A.7. We prove the result in the case that Xx- E H2 where X₂+ (Xx-) noting that this implies X, H2 VI E [t", min(Te, T1/2a-1/e, The case where Xx- H, and replacing T with T, follows by an identical argument. Observe that
x < 8
- -
6
- Tm 5 (
Once again the first inequality comes from the fact that the spherical gradient always greater than and X > for X, in Hi. The second inequality follows from two observations. The first observation is that removing the spherical projections provides an upper bound given X- > O in the time interval considered here. The second observation is that removing
21

---

Benefits of pre-training via single-index models
the gradient of the population loss term, also provides an upper bound. This follows from our analysis of the population dynamics (provided in the discussion before the proof), which told that the first order terms in the population dynamics, point orthogonally towards the line L. Hence when the first order terms exceed the higher order terms (i.e. when X; C), the gradient of the population loss term, is negative under the operator Tj when X Hi. Expanding the above we have:
|| | 122 S 1 + IT
Σ
= 11x +
1/3 +d
under E2.
Proof of Lemma A.8. We prove the case of Q1 and make note that the case of Q3 follows by symmetric argument. We note that in Q1 we have that mi (Xt), m2(X,) > 0 and (Xx, -e) ee > 0, (X) e2 > 0 (this follows from the discussion on the population dynamics, given prior to the proof of 3.4). We now consider mi (X₂) X2 e1:
e1 I 1 61
vo((x-1) e1 -
(X 61
I Xx+
( e1 d
VH(X)) -e 61
imt
S Xp -e1
(H( 1el
11
i=t'
- X1- é + 2d22/5
The inequalities above follow by similar arguments to the previous two lemmas noting that the sign of X, e and the gradients are always the same while in Q₁. The bound on the martingale term follows under the set E7. A similar argument follows for X, 62. We conclude the proof noting that:
1 5
1& S V (X 5 22 + (X,- 62 +
2/33/
$ V (Xr- 01)2 + (Xe- -2)² + d 2/5
S V || 11 + d + 2d /5 X,
V (1 ||2 |
= 18-112 + 115
A.3. Proof of Theorem 3.5
Theorem 3.5. Suppose that f satisfies the following:
E 8²
E)) EE) = 0,
f(g) > 0
0g²
for g NO,, 1). When nn for spherical SGD with N ad steps where a = w(1), as2 < d/13, we have that there
exists some such that for all sequences of initializations x((() m v d, then we have that:
(xx)) 0
22

---

Benefits of pre-training via single-index models
in probability as d + co. Further we have that for all > 0:
P(sup |mm (Xx) > r e) +0
KN
in probability as d + co.
Proof. From the Taylor expansion used in the proof of Theorem 3.4, specifically equations (11) and (12), we see that for 7h = 1 + 7/2 = O, the first order term in the population dynamics sends 21 to O. This is to say that Er > 0 VX s.t. mi (X) E (B0, 2r) (the I-dimensional open ball of radius 2r) we have that (X) . vosgn(X . vo) > 0. This tells us that the population gradient step sends 21 towards 0, for all 21 in the specified ball. We note that the value of is independent of the dimension.
We now start by restricting to the set E = {supigN 1/ E VH(X;) vol < <d-1/s3.
We note that by Doob's inequality
we have that:
VH( - eol > d= < C a08
P(sup N
d1/3 =(1)
1=1
(recalling the assumption that ad² < d/3 Then note that the set E implies the set {suposiSIEN 19 VH(X). - vol 2d 1/33.
We now claim that under the set E, we have that deterministically mi mI(xx) - d - co. Consider the sequence of stopping times, (Tx)x21, corresponding to where mi (X,) crosses zero, i.e., let ты- inf(t > Tk-1 sgn(m (X+)) of sgn(mi (X+-1)) and To 0. Note the maximum single step that mi (X,) can take is bounded and tending to 0 in the dimension (recall the stepsize is & and the gradient of the loss is bounded, as seen in Lemma A.I and the proof of Theorem 3.3 Thus we have that m1 (X,.) - 0 as d + co whenever ты < N. This is because the distance of mi (X,,) to 0 is bounded by the maximum single step.
We will proceed by breaking up the interval [0, _] into the union of U220 [Te, min(74+1, N + 1)). We separately consider the first interval [0, ( N + 1)) and each other interval [T,, min(Te++1, N + 1)). & 2 0, starting with the latter.
Fix k >>0 0 and suppose that E [7k, min(Ts + N,N + 1)). Suppose, without loss of generality, that m. (X+.) 2 O. Now for t [Tx, min(mn+1,n N + 1)), this process remains positive and we have that:
mi (X(+)+) = Xe+1 00
+ (X, 1 (VX(X))+ VX)))) -o
XX, (X, + VX(,)) l22
< X, d (V + VH(()))) 000
8
$
(Xr, d
(V + VHH(X))) co
1=
$ (Xn, d Σ v - vo 1/3 1/3
8
I mi (XTz)
+ 1/3 1/3
< 4d 1/3
The first inequality follows from the fact that the spherical gradient step always results in a point with norm greater than equal to 1. The second inequality follows from the fact that removing the spherical projections for each mi (X.) provides an upper bound as the process is positive over the time interval considered. The third inequality simply applies the restricting set E. The fourth inequality comes from the fact that the V&(X,)) up > 0 whenever < m (X,) < 2r. The final incquality comes from the maximum one step change of m1 (X,) (which is o < d 1/3) which provides an upper bound on mi (Xr,) for k >0 0.
23

---

Benefits of pre-training via single-index models
Now that m1 (X,) 2r, AI E [Tk, min(Tx+1, N + +)), suppose for the sake of contradiction, that my (X,) > 2r for some time tx E [T,, min (Tk+1; N 1)) and further 1 is the first time mi (X,) exceeds 2r. Repeating the inequalities above (noting that mi < 2r, Vt E [Tx, (*)) we have that
mi(xi-) I 4d 4d 1/3 5 2r
a contradiction. We thus have that for any interval [Tk, min(T N + +1), the value m2 (X₂) is upper bounded by 4d 4d'1/3, Note that we only showed this for the case mi (X=,)= ) > 0, but the case mi (Xxx) < 0 follows by similar argument.
The first statement of the proof, i.c. that mi (XXw) + 0 in probability as d oo would then follow so long as TI < 1. Which is to say that if the process m1(X) (Xe) crosses 0 least once, it will remain within d of 0 and hence be there at time N, proving the first statement of the Theorem. The second statement of the proof would also follow so long as (N mm (Xe) < r + e for any fixed e > 0, deterministically under the restricted set E, which has probability 1 ((1)
We now proceed to finish the proof by considering the time interval [0, min(ri, N + 1)). We once again assume without loss of generality that mi (Xo) > i.c that the process is positive over the time interval considered. Then we again have
for in this interval:
m (x(X11))= Xe+1
(X, 9 (V)(X)) + VH()))) 20
(1 X (V+(X,) VH()))) -vol2
8 (V0(X,) VH(X))) -00
$ Xx d
S (X, v.(())) 40 + 8 Σ
VH(X)) v0
i=0
1=0
$ (Xo
(
10 + 2d3
d i=0
Once again, so long m (X,) S 2r, we have that (X(),)) Ugg > 0. By the same contradiction based argument as before,
we can show that mi (X.)= s 2r for all E [0, min(mi, N 1)) (recalling that m. (Xo) r). Hence we have the bound:
m mi(X++1) S 2d 1/3
Which completes the proof of the second statement of the Theorem, i.e. that for all > 0:
P(sup |m (X:) > r - -) 0
in probability as d - oo. We now complete the proof of the first statement of the theorem by showing that m. (Xx) - 0 in probability as + oo in the case that T > N. Suppose for the sake of contradiction that there exists some c > 0 such that:
lim sup IP( inf |m (X,) c)
(14)
d
ISN
Before proceeding, consider the converse of this assumption which is that for every e > O:
lim sup inn KN
inf [m (X,)) > c) =
d
P(
| (x)) =c) N
Then fixing any > 0 and restricting to the - (1) probability set (infigx [m (X+) S c} n F, let Te = in
(X,)) < c}. Assume without loss of generality that m, (X,) > 0, t E [0, To]. Once again repeating the same sequence of inequalities used previously and the same contradiction argument that allows us to invoke those inequalities, we have that:
(x(x)) S mx(X..) + 1/3 1/3 < c + 2d <
24

---

Benefits of pre-training via single-index models
Which is to say that for any e O, mi (Xx)= S 2c with probability 1 o(1), which would complete the proof. Hence, return to the assumption given by (14). Now recall that we have already shown that restricting En ( N + 13 we have mi mi(x, (X+) < 4d/3//3 deterministically, we thus restrict to the set EM(T) > N Now may assume that 0 S m (X;) s 2r for all & E [). The lower bound follows under { < N + 1t. the upper bound m. (X;) < 2r holds generally with probability - o(1) due to the second statement of the theorem with e = r, which has already been proven at this point. Now by the previous analysis of the population loss via Taylor expansion, we have that there exists some constant dj such that (X.) vo 2 (X,)(X)) for E [0, N) and further on the subsequence d&; k 2 such that:
lim P(inf [m (X,)) c) IEN
de
we restrict to the positive probability set that (infigN Img (X)) > c} again assuming without loss of generality
mi (X.) > 2AO E NJ. we have (X,) u0 2 eic for all [0, N). We then see:
mx(,,) S (Xo-
V
to 2d
N Σ 1=0
< (Xo
0o + 2d, 1//3 c1c)
< m, (Xo) cc + 2d 1/3
Which is diverging to -00 as d2 + co. This is contradiction this places negative upper bound on my (X;) which is strictly positive on the positive probability set considered. This completes the proof.
A.4. Proof of Lemma A.10 and Theorem 4.2
Lemma A.10. For all Hermite Polynomials with degree 3 or greater, h2 (z), 2 3:
02
(g) = ( = 0, E
f(g)² > 0
0g²
for g ~ 10)
Proof. For all Hermite polynomials with degree or greater, Ehh- 0. Further it is well known the Hermite polynomials satisfy the following:
n!
(x))
h (x)
(n - k)!
k!
-BA (R)
(g)
El 00 Vk <n
k)!
This gives us that BN"(g)= (g) O. Now consider E 2
82 1())²
Ehi, (h)(g) Now we have that:
(N)())))(()) = k(k 11-1) (()))
by orthogonality. This completes the claim.
Proof of Theorem 4.2. The proof of this theorem follows from (Ben Arous et al., 2021). Specifically in section 2.1 of their paper they show that the single-inder index model covered here meets the assumptions required for there main results which apply more generally. Then noting that by replacing random uniform initialization of order d-1//2 with fixed initialization of order S all their arguments for Theorem 1.3 of their paper still hold, with the new sample complexity provided above (depending on (). In the case of ( = O, the result simply follows by applying Theorem 3.2 of their paper, noting the arguments used to prove this theorem apply just as well when considering a sequence of initializations which itself is not constant but is bounded above and below by constants.
25

---

Benefits of pre-training via single-index models
B. PCA in High Dimensions
consider applying PCA high dimensions in Theorem 3.3. There are a number of well known results about applying PCA the high dimensional limit including the BBP transition(Baik al., 2005). We informally state few of these results here for completeness, however, we refer the reader to Bandeira al. (2020) for more details.
Consider a d-dimensional isotropic Gaussian vector Z (O,, Ia). Letting X E Rdx= denote the data matrix (n rows of observations of 2). If we let both and d grow to infinity, keeping their ratio fixed d/n 7 the distribution of the cigenvalues of the matrix J x X (the sample covariance) will in the limit, follow the Marcenko- Pastur distribution given by:
(x) = 2TTy V (7+ - Y --
with = + 71² and 7- (1 - 7)2. Thus in high dimensions, we can expect to observe top eigenvalues of 1 _ size (1 2 even when there no covariance structure X all. In order detect the spike require À vd/n. n. The limiting squared correlation between the top eigenvector of the sample covariance matrix and the true spike can be shown to be |0- 1o-0)²= = 1-//x- when À 2 / vd/n and 0 otherwise.
X up
1+5/X2
C. Discrete Bihari-LaSalle Inequality
The discrete Bihari-LaSalle Inequality claims that for sequence m, satisfying the following for some k 2, >> 0:
my S a
bm k-1
t=0
then we have that me 2
(1-bo =イ)===
For proof see Appendix C (Ben Arous al., 2021).
D. Effect of 17 on Learning
To investigate the impact of the strength in correlation between the spike and parameter vectors, 7h we run the following experiment. We fix function f which we again take as the third Hermite polynomial, and we vary the value of π and perform training with PCA initialization (and all other settings the same as described in section 5.2. We observe as expected that for large enough values of 7/ SGD with PCA initialization is able to recover the unknown parameter vector and there a point at which 7/ becomes too small and SGD is unable to recover even with PCA initialization. See the figure below.
26

---

Benefits of pre-training via single-index models
1.0
0.8
0.6
n 0.5
M1(Xt)
n = 0.4
0.4
A 0.3
n = 0.2
0.2
f = 0.1
n = 0.01
0.0
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75 2.00
1e6
Figure The figure above displays the correlation values of (XX)) the course of training via SGD with PCA initializations. As expected we see that there is some threshold for 7jI below which the spike vector provides insufficient information about the final parameter initialization for SGD realize the parameter on reasonable time scale.
27


Reviews:
[{"review_id": "fLbtmSfUGs", "reviewer": "Reviewer_DGqD", "summary": "This paper investigates the benefits of unsupervised pre-training and transfer learning in the context of high-dimensional supervised learning, specifically in single-layer neural networks trained via online stochastic gradient descent. The authors establish that pre-training and transfer learning can significantly reduce sample complexity and, in some cases, even lead to exponential improvements over random initialization. The analysis relies on single-index models with spiked covariance and characterizes the effects of different initialization schemes. A key result is that pre-training enables efficient learning in scenarios where random initialization fails due to the existence of attractor basins. The paper also extends its analysis to transfer learning and shows how pre-trained feature representations improve learning efficiency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful comments and suggestions. \n\n+ Controlling the information exponent---We indeed control the information exponent of $f$. However, this assumption is satisfied by all *Hermite polynomials of degree at least three*. We can also check this condition for many polynomials constructed as a combination of these monomials. Consequently, our assumption captures a reasonably generic class of functions, and is not an edge case. \n\n+ Taylor expansion in Theorem 3.5 \\& minimum correlation threshold--- Please note that this is a Taylor expansion of the population loss and is necessarily fundamental to the model. The sign of the coefficient depends on properties of the data distribution, namely $f$ and $\\lambda$ and cannot be changed at will. The desired sign on the coefficient holds for a broad class of functions (see above). There does exist an $r$ (depending on $f$) such that one-pass SGD fails to learn the latent signal for initial correlations less than this value. This correlation value $r$ is fundamental to the model, and is obtained by analyzing the population loss landscape around the origin ($m(X_0)=0$) for functions $f$ satisfying the assumptions stated above. Note that we do not aim to exactly characterize the class of functions $f$ which lead to this phenomenon.  Instead our goal is to illustrate the benefits of pretraining on a class of commonly studied models. \n    \n+  'Phase transitions' between not-learning and learning--- We do not use the terms 'phase transitions' and 'sharp thresholds' in our paper. However, as a consequence of Theorem~ 3.5 one can show a sharp phase transition for the sample complexity for a fixed $f$ as $\\lambda$ varies: When $\\lambda>0$ (the case studied in the paper) the information exponent is infinite so the sample complexity is super-polynomial. If, however $\\lambda\\leq 0$, the sample complexity is polynomial. \n\n+  Other pretraining methods---We study PCA as a natural pretraining algorithm for this data model for the reasons described in Section 3.1. Please see the references therein for various settings where it is argued that more sophisticated pre-training methods effectively implement PCA. It would be interesting to extend the analyses to more realistic pre-training algorithms. \n    That said, we note that our arguments would naturally extend to other pretraining algorithms if one is able to characterize their `overlap' with the planted signal. This, however, is beyond the current paper. \n   \n\n+  Theoretical novelty--- The main technical novelty compared to Ben Arous et. al. arises from the analysis of the population gradient flow. In that work, the population gradient flow reduces to a one dimensional system whereas here we have a two-dimensional dynamical system in our study, which exhibits subtle properties due to the presence of local traps. The analysis of this dynamical system is our main technical contribution. \n\n+  Missing reference and comparison with prior work--- Thank you for catching the missing Arnaboldi et. al. reference. We will certainly add this reference to the camera-ready version. We will also add a comparison with and further contextualize the recent works on the spiked random feature model.   \n\n+  Multi-pass SGD---We agree that it would be interesting to study multi-pass SGD in this setup, adapting the ideas in Lee et. al. and Arnaboldi et. al. \n\n+  Empirical validation beyond stated framework---Pretraining and transfer learning are routinely employed in modern practice. Our results are an attempt to explain the empirical success of these methods. In this sense, our results provide theoretical insight into empirical phenomena  already observed in previous works. \n\n+  Empirical study of effect of correlation at initialization---We are happy to add a simulation study noting the effects of initial correlation on the output of one-pass SGD.", "claims_and_evidence": "The paper makes several claims regarding the benefits of pre-training and transfer learning. While the analysis generally follows well-established approaches in the study of high-dimensional inference, some claims appear somewhat engineered. For instance, the role of the Information Exponent is controlled to illustrate benefits, which raises concerns about whether the observed advantages of pre-training arise naturally or are artificially emphasized. The negative result in Theorem 3.5—establishing a threshold below which learning does not occur—is an interesting contribution, but it might be an artifact of the specific Taylor expansion used rather than a fundamental property of the model. The claim that pre-training yields exponential benefits in certain cases is intriguing, but more discussion is needed on how realistic those cases are in practical settings.", "methods_and_evaluation_criteria": "The methodology follows a standard approach for analyzing the sample complexity of single-index models. The authors leverage stochastic gradient descent in a spiked covariance model and analyze the evolution of correlations over training. However, the connection to prior work, particularly the analysis of Ben Arous et al., is very strong, making the contribution appear incremental. The paper could benefit from a clearer comparison to these prior works to highlight what is truly novel in the methodology.\n\nAdditionally, the use of PCA as a pre-training strategy is reasonable but somewhat simplistic, as it does not capture more sophisticated pre-training methods commonly used in deep learning.  Howver, this might eb expected for analysis purpose, and i agre with the justification of the author (who citie the statistical properties of PCA)", "theoretical_claims": "The theoretical results certainly correct, but they rely on assumptions that may limit their applicability. The proof techniques align with prior work on high-dimensional learning dynamics and information exponents. The claim that pre-training reduces sample complexity is well supported, though it is not surprising given prior literature. Theorem 3.5 presents a negative result that suggests the existence of a minimum correlation threshold for learning, which is an interesting insight. However, the ability to set the leading term in the Taylor expansion to a negative value raises concerns about whether this result is intrinsic to the problem or simply an artifact of the chosen expansion.", "experimental_designs_or_analyses": "The empirical results support the theoretical claims but do not provide strong additional insights. The numerical simulations primarily verify the sample complexity predictions rather than exploring the broader practical implications of pre-training. One limitation is that the experiments focus on synthetic settings with controlled spiked covariance structures, which may not fully capture real-world scenarios. Moreover, the figures suggest that the transition between learning and non-learning regimes is relatively smooth rather than sharply defined, raising questions about whether the term “phase transition” is appropriate in this context.", "relation_to_broader_scientific_literature": "The paper is well-situated within the existing literature on pre-training and transfer learning in high-dimensional learning models.\n\nOn a side note, the spiked random feature model has been investigated in prior works such as, with a different movitations, sich as:\n- Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., Wu, D., & Yang, G. (2022). High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35, 37932-37946.\n- Cui, H., Pesce, L., Dandi, Y., Krzakala, F., Lu, Y. M., Zdeborová, L., & Loureiro, B. (2024). Asymptotics of feature learning in two-layer networks after one gradient-step. arXiv preprint arXiv:2402.04980.\n- Dandi, Y., Pesce, L., Cui, H., Krzakala, F., Lu, Y. M., & Loureiro, B. (2024). A random matrix theory perspective on the spectrum of learned features and asymptotic generalization capabilities. arXiv preprint arXiv:2410.18938.", "essential_references_not_discussed": "When the paper cite Lee et al 2014, they should cite the near identicial but prior results of Arnaboldi et al  arXiv:2405.15459. In fact, i would be curious about the effect of training discussed in Lee et al et Arnadoli et al in the present paper.", "other_strengths_and_weaknesses": "Strengths:\n\nProvides a theoretical framework for understanding pre-training and transfer learning in high-dimensional learning problems.\nDerives sample complexity results that confirm the benefits of pre-training in structured feature models.\nHighlights the role of initial feature correlation in determining the success of training, which is an important insight.\nWeaknesses:\n\nThe theoretical results, while correct, are largely incremental relative to prior work.\nThe pre-training strategy (PCA) is somewhat simplistic compared to modern deep learning methods.\nThe empirical validation is limited to synthetic experiments and does not explore broader practical implications.\nSome results, such as the failure of learning below a certain correlation threshold, may be artifacts of the mathematical formulation rather than fundamental properties of the problem.", "other_comments_or_suggestions": "The authors should more explicitly compare their results to prior work, particularly in relation to the spiked random feature model and information exponents.\n\nThe numerical experiments could be extended to investigate how different levels of correlation affect learning outcomes.\n\nThe discussion of Theorem 3.5 should clarify whether the threshold phenomenon is a fundamental property of the model or an artifact of the analysis.", "questions_for_authors": "Theorem 3.5 suggests that below a certain correlation threshold, learning does not occur. Could this result be an artifact of the chosen expansion, rather than a fundamental limitation?\nHave you considered alternative pre-training methods beyond PCA, and how would your results extend to those cases?\nThe numerical experiments show smooth transitions rather than sharp phase transitions. How do you justify using the term “phase transition” in this context?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ry5WLDtWOo", "reviewer": "Reviewer_Z477", "summary": "This paper studies the dynamics of online SGD for learning single-index models for spiked Gaussian data, where the spike in the Gaussian distribution is meaningfully correlated with the ground truth direction. The authors model pretraining and fine-tuning in practical machine learning settings as learning the spike direction and learning the ground truth direction. In that sense, the authors prove that under a certain assumption on the target function and sufficient correlation between the spike and the ground truth direction, SGD initialized from PCA pretraining is successful at recovering the ground truth direction with only a linear number of samples. In contrast, the authors prove negative results that show $\\Omega(d^2)$ samples are necessary from random initialization, and further when the spike and ground truth directions are identical, even constant initial correlation is not sufficient, regardless of the amount of data used.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful comments and suggestions. \n\n+  Notation and dimension dependence of parameters---We assume throughout that $\\lambda>0$ is dimension independent. $B_2(0,1)$ is indeed the two-dimensional ball of radius 1 centered at the origin, and $m^*$ is also dimension independent in our analysis. We will clarify these dependencies in the final version. \n\n+  $\\lambda$ and $\\eta_1$ scaling---We assume that $\\lambda$ and $\\eta_1$ are dimension independent in our work. We expect that the analysis we provide here can be carried over to the case (following your notation) $r_1=0$ and $r_2>0$. We agree that the case $r_1>0$ is interesting. However it will require substantially more work: One would have to carry out a similar “bounding flows'' style argument, but around the point $(0,1)$ instead of $(0,0)$. In light of the tight turn around for this conference, we defer this to future work.\n    \n+  Comparison with Mousavi-Hosseini et. al. (2023, 2025)---Thank you so much for pointing us to these results. As remarked above, if $\\lambda$ and $\\eta_1$ scale with $d$, our regime is related to the prior work. Thank you also for pointing us to the sample complexity improvement in Mousavi-Hosseini et. al. (2023). We will add a discussion of these connections in our final version.  \n\n+  Information exponent clarification---Thanks for pointing this out. We will include your interpretation in the final version. Please see the response to reviewer DGqD.\n\n+  Assumption 3.2, dimension dependence of $m^*$ and dimension dependence of $r$ in Theorem 3.5---Thanks for these comments. Indeed, the parameter $r$ in Theorem 3.5 is dimension independent. Assumption 3.2 is equivalent to the condition that the global minimizer of the population risk (i.e. $(1,0)$) is locally attractive. Thus, the population dynamics initialized in a neighborhood of the global minimum will converge to this minimum. Finally, note that $m^*$ is dimension independent---to see this, note that $m^*$ can be specified in terms of the function $\\phi$, which is dimension independent. \n\n\n+  Step-size for SGD---If the step size is larger, the SGD does not follow the population gradient flow. This invalidates the main analysis strategy introduced in Ben Arous et. al. We adapt these ideas in our analysis and are thus also unable to handle SGD with large step size.", "claims_and_evidence": "The claims are clearly supported by proofs and some experiments.", "methods_and_evaluation_criteria": "Not applicable.", "theoretical_claims": "I did not explicitly check the correctness of the proofs, but the statements seem plausible and in agreement with prior results.", "experimental_designs_or_analyses": "Not applicable as this is mostly a theory paper.", "relation_to_broader_scientific_literature": "The paper provides a valuable contribution to the high-dimensional SGD literature, especially since learning single-index models has seen a lot of recent attention as a model of recovering low-dimensional structure in high-dimensional spaces and feature learning with neural networks. While the results may not be very surprising, they offer a nice perspective on the principles of pretraining and fine-tuning, even for simple models such as learning single-index models.", "essential_references_not_discussed": "Essential references seem to be sufficiently discussed. I provide additional discussion suggestions below.", "other_strengths_and_weaknesses": "**Strengths**: \n\nThe paper presents a nice story with intuitive arguments that are easy to follow, and the negative result of Theorem 3.4 is particularly strong as it covers many initializations and any number of samples.\n\n**Weaknesses**: \n1. The role of certain parameters in the results is not clear. For example, how would $\\lambda$ or $\\eta_1$ that scale with dimension affect Theorem 3.3? Also, Assumption 3.2 is a bit opaque. The authors provide some arguments to support this assumption in Section 5, but I still don't have a good intuition of how $m^*$ behaves with respect to dimension in typical settings. I would be happy to raise my score if there is more discussion and examples on these points.\n\n2. A limitation of the lower bounds is that they only apply to sufficiently small step size, $\\alpha\\delta^2 = O(1)$ and $\\alpha\\delta^2 = O(d^{1/3})$. Intuitively, SGD updates should become unstable beyond this step size, but this is not theoretically shown in this work or prior works such as Ben Arous et al., 2021.", "other_comments_or_suggestions": "Please see below.", "questions_for_authors": "1. Theorem 3.3 does not mention the assumption on $\\lambda$. I believe it is assumed that $\\lambda > 0$ is dimension-independent, in which case it would be helpful to explicitly mention this. Also, what happens if we have $\\lambda$ and $\\eta_1$ that scale with dimension?  Mousavi-Hosseini et al., 2023 consider $\\lambda \\asymp d^{r_2}$ and $\\eta_1 \\asymp d^{-r_1}$, and show that the sample complexity of full-batch gradient flow with PCA initialization is $d^{1 + \\max((k-1)(2r_1 - r_2), 0)}$ for information exponent $k$, while Mousavi-Hosseini et al., 2025 consider gradient descent with exponential compute and show the sample complexity is $d^{1 - \\max (r_2 - 2r_1,0)}$ independent of information exponent. It would be helpful if the sample complexity from Theorem 3.3 could be compared with these results. \n\n2. Moreover, it can be helpful to note the sample complexity improvement of Theorem 4.2 is identical to Theorem 7 of Mousavi-Hosseini et al., 2023. The authors there turn gradient flow on spiked covariance into gradient flow on isotropic data with non-random initialization, and achieve a sample complexity of $d^{1 + 2\\zeta(k-1)}$. Theorem 4.2 of this paper provides an improved $d^{2\\zeta}$ factor. This is similar to the improvement of online SGD (Ben Arous et al. 2021), over gradient descent/flow (Bietti et al., 2022) on isotropic Gaussian data .\n\n3. I believe Equation (1) asks for f to have information exponent >= 3 and f^2 to have an information exponent <= 2. It would be helpful to discuss the relationship between information exponent and Equation (1).\n\n4. My interpretation was that $r$ in Theorem 3.5 is not dimension-dependent. If that is the case, it would be helpful to explicitly mention it in the theorem statement.\n\n5. I think $B_2(0,1)$ is not defined in the text, does it mean the 2-dimensional ball of radius 1 centered at 0? Also, is there any way to make Assumption 3.2 more interpretable, e.g. for single Hermite polynomials? In which case, can we explicitly know the dimension-dependency of $m^*$?\n\nReferences:\n\nA. Mousavi-Hosseini et al. \"Gradient-Based Feature Learning under Structured Data.\" NeurIPS 2023.\n\nA. Mousavi-Hosseini et al. \"Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics.\" ICLR 2025.\n\nG. Ben Arous et al. \"Online stochastic gradient descent on non-convex losses from high-dimensional inference.\" JMLR 2021.\n\nA. Bietti et al. \"Learning Single-Index Models with Shallow Neural Networks.\" NeurIPS 2022.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HCDew3ph4r", "reviewer": "Reviewer_bDFS", "summary": "The paper provides a theoretical analysis of how model initialization learned via unsupervised pre-training benefits supervised learning tasks when training a single-layer neural network using online stochastic gradient descent.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful comments and suggestions. \n\n+ Weakness of Assumption 3.2---Please note that Assumption 3.2 pertains to gradient flow on the *population* loss (population gradient flow) and not to online SGD. Indeed, one of our main contributions is to establish rigorously that if the step sizes are reasonably small, the evolution of the correlations of the SGD trajectory with the latent directions is well approximated by that of population gradient flow.\n\n\n  We believe that Assumption 3.2 is quite mild. Note that $(1,0)$ is the global minimizer of the population loss. Assumption 3.2 is essentially equivalent to the assumption that the global optimizer is a local attractor for population gradient flow. This ensures that if population gradient flow is initialized in a small neighborhood of $(1,0)$, it will converge to the optimum in short time. We agree that the precise form of Assumption 3.2 could be potentially weakened---we stated it in a form which was easy to verify. In our paper, we verify this condition formally when $f$ is a Hermite polynomial. \n\n+ Assumption 1---There has been significant recent progress in understanding the properties of gradient descent algorithms on single index models with isotropic gaussian features. Prior work establishes that the sample complexity of one pass SGD is governed by the \"information exponent'' of the problem. In this setting, the information exponent is the index of the first non-zero term of the expansion of the function $f$ in the Hermite basis. Assumption 1 is in the same spirit---we assume that $f$ has information exponent at least three, while $f^2$ has information exponent at most two. We note in the paper that this assumption is satisfied by all Hermite polynomials of degree $\\geq 3$, and suitable linear combinations---thus this assumption is satisfied by a broad class of non-linearities $f$. We note that the RELU function does not satisfy this condition; however, the RELU activation function has information exponent one, and thus can be learned with approximately linear sample complexity. The condition focuses on more challenging non-linearities which have polynomial sample complexity for the isotropic single index model. \n\n+ Data sample vs. gradient updates---We focus exclusively on single pass SGD in this work. We agree that it would be interesting to study other variants e.g. batched SGD or multi-pass SGD in future work. We believe that our analysis can be extended to settings where each epoch uses fresh batches of data, but the batch sizes are relatively small. This variant will not change the sample complexity of the algorithm.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I did not check the proofs.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The setting is interesting, but it requires some strong assumptions.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strength:\nThe paper provides a clear and well-motivated setup, and is strong in terms of readability and structure. The problem formulation and theoretical contributions are presented clearly and logically; in particular, the discussion sections, as well as the intuitive explanations of theorems and assumptions, are quite helpful.\n\nWeakness:\nAssumption 3.2 significantly simplifies the analysis and may be overly restrictive. It requires that the gradient with respect to all data samples let the update towards the direction $e_1$ and be far from $e_2$. However, in realistic scenarios, gradient descent updates could have more complex trajectories while still converging optimally. This assumption might limit the general applicability of the theoretical insights.", "other_comments_or_suggestions": "See questions.", "questions_for_authors": "1. Could the authors elaborate on the assumption (1) of $f$ (Line 168)? I appreciate the discussion in Section 3.4. However, is the commonly used ReLU function not satisfied? Could we weaken the third requirement to let the ‘>0’ to be ‘’$\\geq 0$’?\n2. Since in your setting the data sample and the gradient update steps are the same, it leads to a lower bound of sample complexity in Theorem 3.4. However, in reality, these two terms are not the same. If we let them to be two separate terms, what can you get?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XKVUoSscFX", "reviewer": "Reviewer_Fd1t", "summary": "This paper investigates the benefits of unsupervised pre-training in supervised learning tasks, focusing particularly on GLMs with parameter vectors correlated to a spike in the data covariance. The authors show that a \"PCA initialization\" derived from an independent sample of unlabeled data can drastically reduce the sample complexity threshold for parameter vector recovery through SGD, compared to random initialization. They highlight a setting in which random initialization fails regardless of the amount of data, while PCA initialization does not. Additionally, the authors discuss the benefits of pre-training in transfer learning, which can lead to a polynomial reduction in sample complexity, provided that the correlation between the source and target is strong enough. The findings are validated through finite-dimensional simulations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful comments and suggestions. \n\n\n+ Use of parentheses around citations---Thank you for pointing out this issue. We will be sure to correct this in the final version. \n\n + Limitations of the data distribution---We are happy to include additional comments on the motivations for our data distribution, and its limitations. Let us recap here: Our goal is to understand the benefit of pre-training. For this, one must have a solid baseline where the sample complexity without pretraining is exactly understood. We work with single index models/GLMs as they are, to our knowledge, the only non-trivial models of supervised learning with neural networks where the sample complexity of online SGD is exactly understood.\n    \n    For pre-training to be successful, the (unlabeled) features must have some non-trivial information about the downstream prediction task. For single-index models, this correlation must occur either as a non-trivial mean for or in the covariance of the features. The question is which more naturally models what is seen in practice. \n    In the context of modern language models, a popular pre-training strategy is to construct artificial prediction tasks by 'masking' tokens in the unlabeled text data. Intuitively, one hopes to construct 'clozed masks', i.e., masks that are helpful to predict the downstream label. This masking strategy has been deployed in well-known models such as BERT (Devlin et. al. (2018)). \n    More generally, it is common to perform pre-training via denoising autoencoders. In light of this, and the references in Section 3.1 regarding autoencoders, it seems more natural to relate the features to the downstream prediction tasks via the feature covariance.\n    Given this, the spike covariance model is a natural, albeit simple starting point. Let us end by emphasizing that we are not proposing PCA as a general purpose pretraining algorithm. Instead, we use it as a natural model of pretraining in our set up.  We hope that this simple analysis will motivate further investigations into richer models.", "claims_and_evidence": "All claims are supported by convincing evidence.", "methods_and_evaluation_criteria": "The methods are well suited to this problem.", "theoretical_claims": "I checked the proofs for the theorems in section 3, I have no issues to discuss.", "experimental_designs_or_analyses": "I have no issues to discuss.", "relation_to_broader_scientific_literature": "This paper builds on existing literature in supervised learning with GLM target functions, extending methods and results from the Gaussian iid setting with random initialization to the spiked covariance problem. The authors introduce a pre-training-based initialization procedure, which improves learnability in the considered setting, offering a significant advancement over traditional random initialization approaches in high-dimensions.", "essential_references_not_discussed": "I am not aware of any essential references that have been omitted.", "other_strengths_and_weaknesses": "The paper is clear in its exposition and addresses the important issue of the theoretical understanding of the benefits of pre-training, deriving interesting results that support the argument. One weakness is the artificial nature of the data distribution, which may seem tailored to make the proposed pre-training method work. In my opinion, the authors do not discuss this artificiality sufficiently, nor do they highlight how this toy model could resemble more realistic settings and how the results may translate to them. Nonetheless, this setting provides clear theoretical insights and original results, and could serve as a baseline for future work in more complex or realistic contexts.", "other_comments_or_suggestions": "A minor issue is the inconsistency in the use of parentheses around citations (for example, on lines 327-328, second column). I suggest revising the paper to correct this.", "questions_for_authors": "Following the discussion in the strengths and weaknesses section, I would ask the authors to provide further comments on their choice of setting and whether they can offer a stronger justification for it. I suggest expanding this discussion in the paper as well.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This is a technically solid paper that studies the role of unsupervised pre-training and of transfer learning in the context of the single index model, for which many properties of the learning procedure and hardness can be understood in great detail. This paper obtains numerous interesting results in this direction and constitutes an interesting contribution to ICML. 

I would note that similar models were considered in some additional existing literature that is worth including in the related work: 
E.g. a model similar to the one used for pre-training in https://proceedings.mlr.press/v202/pesce23a/pesce23a.pdf , and another model related to the one used for transfer learning in https://iopscience.iop.org/article/10.1088/2632-2153/ac4f3f
Decision: Accept (spotlight poster)