Provable Benefits of Unsupervised Pre-training and Transfer Learning via Single-Index Models

Thank you for the thoughtful comments and suggestions.

• Use of parentheses around citations---Thank you for pointing out this issue. We will be sure to correct this in the final version.

• 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.

  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. 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)). 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. 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.

Thank you for the thoughtful comments and suggestions.

• 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.

  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.

• 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.

• 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.

Thank you for the thoughtful comments and suggestions.

• 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.

• $\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.

• 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.

• Information exponent clarification---Thanks for pointing this out. We will include your interpretation in the final version. Please see the response to reviewer DGqD.

• 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.

• 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.

Thank you for the thoughtful comments and suggestions.

• 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.

• 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.

• '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.

• 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. 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.

• 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.

• 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.

• 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.

• 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.

• 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.