From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD

We thank the reviewer for their feedback. Please find our response below.

Notation for $\mu(\eta)$: We agree with the reviewer's suggestion to improve readability. In the final version, we will define $\mu_i(\eta) := \mathbb{E}[\psi_{\eta}(\sigma_*(a),b) He_i(a) He_{i-1}(b)]$, where $a,b$ are independent standard normal random variables.

Clarification on Contributions: Our core technical contribution is a careful analysis of different SGD variants to learn single-index models, while keeping track of dependencies on the learning rate, achieved through a unified set of tools by looking at a general update rule. Through this analysis, we are able to characterize how the query class and complexity of SGD-type algorithms change depending on the learning rate, which is our main conceptual contribution. These contributions are different from those of [3] on the technical side by a precise analysis that keeps track of the learning rate and on the conceptual side as we cover regimes in between CSQ and SQ.

As the reviewer mentioned, in the specific case of batch reuse SGD, our analysis reduces to that of [1], with the exception of explicit tracking of learning rate. However, to our knowledge, the fact that alternating SGD can also modify the query class of SGD is entirely novel.

Non-Polynomial Link Functions: We can in fact admit any link function that has at most polynomial growth (as in [2]) without affecting our main result. We will modify Assumption 3.1 in the final version to indicate this. This is a much larger function class that includes ReLU and its variants, splines, as well as any bounded link function such as the logistic and probit. Assuming at most polynomial growth for $\sigma_*$ ensures that $\langle \boldsymbol{\theta}_*, \boldsymbol{g} \rangle$ is sub-Weibull, allowing us to make concentration arguments.

Realistic Simulations of Alternating SGD: We would like to highlight that we only study alternating SGD as a simplified proof of concept, as it provides a novel mechanism for modifying the statistical query type of SGD, one that comes from updates on the second layer rather than the first. In practice when both layers are trained simultaneously for a number of steps on the same batch, this effect will be combined with the effect of updates on the first layer alone, which is captured by Algorithm 2. Therefore alternating SGD here should be seen as a theoretical construction (similar to the batch reuse algorithm of prior works) to allow us to understand more practical settings, rather than a new optimization algorithm. We will better emphasize this fact in the paper.

Varying $d$ in Experiments: We repeated the experiments in Figures 1 and 4 with $d = 25$ and $d = 75$. Taking $d$ larger is not feasible on our GPU due to the large number of samples (and thus running time) required. The theory predicts that the sample complexity will be "flat'' until a phase transition at $\eta \asymp d^{-1/2}$, after which it smoothly decreases until $\eta \asymp 1$. There will not be a sharp vertical transition as described by the reviewer. Our experimental results match this prediction. In particular, the threshold for the phase transition decreases as $d$ increases.

In our experiments, taking larger values of $\eta$ than those appearing in Figure 1 leads to instability and poor generalization for both alternating and batch reuse SGD. We have a similar observation for all three algorithms when $\gamma$ is taken too large.

We will add these additional experimental results to the final version.

Insights on Learning Rate Tuning: We would like to stress that our main goal is not to find an optimal learning rate for practitioners, but to understand how learning rate can affect the oracle access/query model and the subsequent sample complexity of SGD. Nevertheless, we agree with the reviewer on the importance of being able to provide practical recommendations.

The suggestion from our results is that learning rates $\gamma$, $\eta$ should be taken as large as possible (while still ensuring convergence) so as to avoid incurring unnecessary additional training cost. In particular, in the cases of alternating SGD and batch reuse SGD, the coefficients $\mu_i(\eta)$ are non-decreasing in $\eta$ (as long as Assumption 3.2 holds so that the relative Hermite coefficients are positive). In fact, for any algorithm where $\eta$ has the interpretation of a learning rate and thus appears in $\psi_{\eta}$ only as a scaling factor, this property will hold. It is therefore immediate that the practitioner should take $\eta$ as large as possible. From our condition $\psi_{\eta} = O(1)$ — which ensures control over the errors arising from normalization, SGD noise, and label noise — this means taking taking $\eta \asymp d^{-1}$ for batch reuse SGD (as in [1]) and $\eta \asymp 1$ for alternating SGD.

The Upper Bound on $\gamma$: The reviewer's intuition is exactly correct. Taking $\gamma$ sufficiently small is necessary to ensure that the “informative” terms in the gradient update dominate the errors arising from both the normalization and the noise. The bound we provide on $\gamma$ is tight in the sense that it is the largest possible value for which these errors remain controlled. Moreover, as alluded to in an earlier answer, we observe in experiments that taking $\gamma$ larger than the theoretical threshold leads to instability.

Deep Networks: We can generalize our example with $\sigma(z) = z^2$ to any depth. This is because, up to constants, $\psi_{\eta}(y,z)$ will always take the form $\sum_{i=0}^{D-1} \eta y^{i+1} z^{k_i}$, where all $k_i$ are odd. The odd Hermite coefficients of odd monomials are always positive, while their even Hermite coefficients are zero. Hence, we can recover in $\tilde{\Theta}(d)$ time if $\eta \asymp 1$ and at least one of $u_2(\sigma_*^k) > 0$ for $k \in [D]$. We will add a detailed discussion to illustrate this point in Appendix C.3.

SQ-Optimality through Smoothing: Yes, it is possible to achieve the SQ lower bound through a technique inspired by the smoothing approach in [2]. The recent paper [3] shows that this can be done by aggregating the update $\psi(y,\langle \boldsymbol{x},\tilde{\boldsymbol{w}}\rangle)$ for a large number of random perturbations $\tilde{\boldsymbol{w}}$ of the current iterate $\boldsymbol{w}$ (see Algorithm 1 and Theorem 4.2 in [3] for details). This achieves the SQ lower bound whenever $\mu_{p_*}(\eta)$ is positive and of constant order (see Assumption 4.1(b) in [3], where $\mathbb{E}[\zeta_{s*}(y) \psi_{s*-1}(y)]$ is exactly $\mu_{p*}$ in our notation). This will be the case for batch reuse SGD whenever $\eta$ is sufficiently large (on the order $d^{-1}$), whereas for alternating SGD it will hold whenever $\mathrm{IE}(\sigma_*^2) = p_*$ and $\eta \asymp 1$.

Emphasis on Learning Rate: We highly appreciate the reviewer's suggestion to better frame our contributions. In the final version, we will establish a clearer distinction between the hyperparameters $\gamma$ and $\eta$, emphasizing that only the latter yields phase transitions in sample complexity whereas the former affects the sample complexity in a smooth way (due to the $\gamma^{-1}$ factor appearing the rate in Theorem 3.3). We will also highlight that our framework allows for the study of other hyperparameter and design choices, such as the role of depth in the alternating SGD algorithm.

We are happy to answer any further questions.

[1] Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit. NeurIPS 2024.

[2] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D Lee. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index models. NeurIPS 2023.

[3] Siyu Chen, Beining Wu, Miao Lu, Zhuoran Yang, and Tianhao Wang. Can neural networks achieve optimal computational-statistical tradeoff? an analysis on single-index model. ICLR 2025.

We thank the reviewer for their feedback. Please find our response below.

Clarifying the $\mu(\eta)$ Notation: The formalism of the generic update rule $\psi(y,\langle \boldsymbol{x}, \boldsymbol{w} \rangle)$ was introduced in [1]; our only modification is to track the explicit dependence on a hyperparameter $\eta$, which appears in SGD variants that employ label transformation, for example. In fact, the expression denoted $\mathbb{E}[\zeta_{s*}(y) \psi_{s*-1}(y)]$ in Assumption 4.1(b) of [1] is exactly $\mu_{p_*}$ in our notation. While the authors of [1] assume that the latter quantity is of constant order (which in our examples amounts to taking $\eta$ as large as possible), we consider the sample complexity for a full range of $\eta$ values. We will add a discussion on this effect in Section 3 where we introduce the formalism and the $\mu(\eta)$ notation. We will also preview how this formalism captures batch reuse SGD (as per the reviewer's suggestion) for easier readability.

Incremental Progress Towards Matching SQ Lower Bounds: We would like to highlight that our goal is not to find the optimal sample complexity under some oracle assumption, but rather to provide a detailed characterization of the failure modes of existing SGD variants that break out of CSQ. Namely, we show that incorrect choice of hyperparameters governing the size of non-correlational updates can lead to sample complexity as poor as online SGD. Moreover, any algorithm that fits our framework (including our novel alternating SGD) and satisfies $\mu_{p_*} = \tilde{\Theta}(1)$ can be made SQ-optimal with the smoothing-inspired approach in [1].

1) Smoothing Hyperparameters: We believe that it is possible to adapt our framework to capture landscape smoothing and treat the associated hyperparameter (called $\gamma$ in [1] and $\lambda$ in [2]) as our $\eta$. This boils down to casting the algorithm's update as $\psi(y,\langle \boldsymbol{x},\boldsymbol{w} \rangle)$ and computing the coefficients $\mu_i$. This is more intricate than the three examples we consider in our work due to the random perturbations of the weight vector $\boldsymbol{w}$. In particular, for the algorithm in [2], the gradient update takes the form $\frac{1}{\sqrt{1+\lambda^2}} y \mathbb{E} [\sigma’(\frac{\langle \boldsymbol{w}, \boldsymbol{x} \rangle + \lambda \langle \boldsymbol{\xi}, \boldsymbol{x} \rangle}{1 + \lambda^2})]\boldsymbol{x}$, where $\boldsymbol{\xi}$ is drawn from the uniform distribution on $\mathbb{S}^{d-1}$ conditioned on being orthogonal to $\boldsymbol{w}$. Note that the distribution of $\langle \boldsymbol{\xi}, \boldsymbol{x} \rangle$ depends on $||\boldsymbol{x}||$ in addition to $\langle \boldsymbol{x},\boldsymbol{w} \rangle$. We believe that the proof techniques from the appendices of [2] to handle this can be of use here.

On this point, it is also worth mentioning that that Theorem 1 of [2] establishes that that landscape smoothing SGD has sample complexity $\tilde{O}(d^{p-1} \lambda^{-2(p-2)})$ for $\lambda \in [1,d^{1/4}]$. Hence, there is a smooth (i.e., no phase transitions) interpolation between the online SGD complexity $\tilde{O}(d^{p-1})$ and the CSQ lower bound as a function of $\lambda$.

The situation is more complicated for the SQ-optimal algorithm of [1], where the analysis assumes that the polarization parameter $\gamma$ is fixed at its optimal value. This choice directly influences the number of samples required per step to ensure concentration of the empirical gradient (Proposition E.4 in [1]). Moreover, the expected alignment between the gradient and $\theta_*$ depends on $\gamma^{p_*-1}$ (Proposition E.3 in [1]). We leave it to future work to understand the behaviour for suboptimal $\gamma$ and attempt to cast the algorithm into our framework.

Please note: We think the references [1,2] are missing in the review.

2) More General Link Functions: We can in fact handle any link function that has at most polynomial growth (as in [3]). We will modify Assumption 3.1 in the final version to indicate this. This is a much larger function class that includes ReLU and its variants, splines, as well as any bounded link function such as the logistic and probit. Assuming at most polynomial growth for $\sigma_*$ ensures that $\langle \boldsymbol{\theta}_*, \boldsymbol{g} \rangle$ is sub-Weibull. Indeed, the class of sub-Weibull random variables includes sub-Gaussian and sub-exponential random variables and is closed under transformations with at most polynomial growth (up to changing the tail parameter). This condition on the noise ensures that we can make a concentration of measure argument to control the errors arising from label noise and the stochastic gradient.

3) The Choice of $\gamma$: The hyperparameter $\gamma$ must on the order of its largest possible value as per Theorem 3.3 so that optimal sample complexity can be achieved. This constraint is necessary to control the errors in the update to $\langle \boldsymbol{\theta}_*,\boldsymbol{w} \rangle$ that arise due to normalization, stochastic gradient noise, and label noise (see proof sketch in Section 5). It is analogous to the constraint placed on the hyperparameters in batch reuse SGD (see Theorem 2 in [3], Theorem 1 in [4]).

4) The Implications of the Phase Transition: We had two main motivations to study the effect of learning rate on generalization:

a) Our first motivation was to understand when gradient-based algorithms on squared loss can break out of CSQ. While prior works [3,4] focus on the role of batch reuse, we believe learning rate is another important aspect that determines the oracle class of gradient-based algorithms, as demonstrated by our examples.

b) In practice, it has been observed that larger learning rate leads to better generalization [5]. There is no consensus on why this is the case, and works on edge of stability (EoS) are more concerned with the optimization dynamics than generalization behaviour. In that regard, we believe studies that link sample complexity and learning rate can be valuable. Specifically, our work shows that significant improvements in the generalization behaviour can occur even when one is not operating in the EoS regime. Further studies to better characterize practical settings where this is the case are interesting, but outside the scope of our current paper.

We are happy to answer any further questions.

[1] Siyu Chen, Beining Wu, Miao Lu, Zhuoran Yang, and Tianhao Wang. Can neural networks achieve optimal computational-statistical tradeoff? an analysis on single-index model. ICLR 2025.

[2] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D Lee. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index models. NeurIPS 2023.

[3] Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows SGD to learn high-dimensional multi-index functions. arXiv 2024.

[4] Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit. NeurIPS 2024.

[5] Yuanzhi Li, Colin Wei, Tengyu Ma. Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks. NeurIPS 2019.

We thank the reviewer for their feedback. Please find our response below.

The Class of Models under Consideration: We would like to point out that despite their simplicity compared to practice, the class of Gaussian single-index models provides a tractable formalism for the theoretical study of feature learning [1,2,3]. The assumptions inherent to this setting are crucial to establish rigorous sample complexity bounds for SGD. In the paper, we assume that the link function $\sigma_*$ is polynomial, which is rather restrictive. We can in fact extend this to any link function with at most polynomial growth and will modify Assumption 3.1 in the final version to reflect this. This broadens our theory to capture much more general link functions in statistics and machine learning, such as ReLU (and its variants), the logistic, and the probit.

Additional Experiments: As per the reviewer's suggestion, we have run additional experiments for online and batch reuse SGD. To start, we focused on our existing example of $\sigma_* = \operatorname{He}\_3$ since $\operatorname{IE}(\sigma_*) = 3$ and $\operatorname{IE}(\sigma_*^2) = 2$. For online SGD, there is no phase transition with respect to the learning rate $\gamma$ and the test error decreases at approximately a $\Theta(\gamma^{-1})$ rate, as expected from theory. On the other hand, we performed a larger-scale version of the experiment in Figure 4 for batch reuse SGD and confirmed that this algorithm does exhibit a phase transition due to the parameter $\eta$. Since we are not allowed to share supplementary figures during the rebuttal, we can only include these experiments in the final version.

We would also like to highlight the fact that experiments in practical settings have already observed that learning rate can affect generalization (and thus sample complexity), see e.g. [6]. Our study can be one way to approach a theoretical justification of this observation.

A General Discussion of the Phase Transition in Section 3: We thank the reviewer for this suggestion, which will indeed help the readers better parse the main result. We provide a brief discussion below, and will it to the final version of the paper in Section 3.

For simplicity, assume that all $\mu_i$ are positive. Consider the key equation $T(\eta) = \min_{1 \leq i \leq r} \tilde{\Theta}(\gamma^{-1} \mu_i(\eta)^{-1} d^{\frac{i-2}{2} \vee 0})$. The sample complexity is determined by the smallest term, whose index we can identify by $\alpha(\eta) = \operatorname{argmin}_{1 \leq i \leq r} \tilde{\Theta}(\gamma^{-1} \mu_i(\eta)^{-1} d^{\frac{i-2}{2} \vee 0})$. We say that the sample complexity exhibits phase transitions when $\alpha$ changes with $\eta$. We can identify potential phase transitions by comparing $\mu_i^{-1} d^{\frac{i-2}{2}} \leq \mu_j^{-1} d^{\frac{j-2}{2}}$. In particular, at fixed $d$, if $\mu_i$ is growing faster with $\eta$ than $\mu_j$, then the $i$th term will be more important in determining sample complexity for larger $\eta$. This is the phenomenon we illustrate with alternating SGD and batch reuse SGD in Section 4.

Please let us know if we can clarify this further.

Referencing Issues: We thank the reviewer for pointing these issues out. We have addressed these for the final version.

Limitations of the Paper: We agree with the reviewer that our discussion of limitations should be more explicit, and we will rectify this in the final version. Currently we are limited to the setting of Gaussian single-index models under the specific assumptions in Section 3 in order for the analysis to be theoretically tractable. In the conclusion, we list possible extensions that would broaden the scope of our theory (e.g., multi-index models [4], more general input distributions [5]).

We hope our responses have addressed the reviewer's concerns, and are happy to answer any further questions.

[1] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. JMLR 2021.

[2] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. NeurIPS 2022.

[3] Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit. NeurIPS 2024.

[4] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. SGD learning on neural networks: Leap complexity and saddle-to-saddle dynamics. COLT 2023.

[5] Nirmit Joshi, Hugo Koubbi, Theodor Misiakiewicz, and Nathan Srebro. Learning single-index models via harmonic decomposition. arXiv 2025.

[6] Yuanzhi Li, Colin Wei, Tengyu Ma. Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks. NeurIPS 2019.

We thank the reviewer for their feedback. Please find our response below.

The Optimal Learning Rate: Indeed, the optimal $\eta$ depends on the form of the $\mu_i(\eta)$ in general. However, for the cases alternating SGD and batch reuse SGD, these coefficients are non-decreasing in $\eta$ (as long as Assumption 3.2 holds so that the relative Hermite coefficients are positive). In fact, for any algorithm where $\eta$ has the interpretation of a learning rate and thus appears in $\psi_{\eta}$ only as a scaling factor, this property will hold. It is therefore immediate that the practitioner should take $\eta$ as large as possible while still ensuring convergence. From our condition $\psi_{\eta} = O(1)$ — which ensures control over the errors arising from normalization, SGD noise, and label noise — this means taking $\eta \asymp d^{-1}$ for batch reuse SGD (as in [1]) and $\eta \asymp 1$ for alternating SGD.

The Link Functions under Consideration: We can in fact handle any link function that has at most polynomial growth (as in [2]). We will modify Assumption 3.1 in the final version to indicate this. This is a much larger function class that includes ReLU and its variants, splines, as well as any bounded link function such as the logistic and probit. Assuming at most polynomial growth for $\sigma_*$ ensures that $\langle \boldsymbol{\theta}_*, \boldsymbol{g} \rangle$ is sub-Weibull. Indeed, the class of sub-Weibull random variables includes sub-Gaussian and sub-exponential random variables and is closed under transformations with at most polynomial growth (up to changing the tail parameter). This condition on the noise ensures that we can make a concentration of measure argument to control the errors arising from label noise and the stochastic gradient.

The Artificial Nature of Alternating SGD: We thank the reviewer for raising this concern, which allows us to clarify the point of studying Algorithm 1. Although this algorithm is somewhat contrived when compared with practice, we believe that its study is useful from a theoretical perspective. We see it as proof of concept to demonstrate what mechanisms might be at play in neural network training so that they achieve near-optimal sample complexity. Specifically, it reveals another mechanism for obtaining benefits of data repetition — one that comes from updates on the second layer rather than first. In practice when both layers are trained simultaneously for a number of steps on the same batch, this effect will be combined with the effect of updates on the first layer alone, which is captured by Algorithm 2.

We also highlight the connection between Algorithm 1 and the literature on studying the two-timescales optimization dynamics of neural nets. By studying this algorithm, we are able to link the discrepancy in the learning rates to sample complexity, which is a novel perspective in this literature. Once again, the ideal algorithm would take multiple steps on the same batch and observe the combined effect of Algorithms 1 and 2, but studying this idealized setting is out of reach for our current techniques.

We will better highlight these points in the final version to motivate the study of Algorithm 1.

We are happy to answer any further questions.

[1] Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit. NeurIPS 2024.

[2] Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows SGD to learn high-dimensional multi-index functions. arXiv 2024.