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

AI summary

Overview

• Research area: Deep learning theory; specifically the statistical and computational complexity of gradient-based feature learning for Gaussian single-index models.
• Technical level: Advanced. The paper is a pure theory contribution built on Hermite expansions, high-dimensional asymptotics, and statistical query lower bounds. It reports one experiment (Figure 1) but the core content is theorems and corollaries.
• Scope (one sentence): The paper characterizes how the choice of learning rate hyperparameters determines whether online gradient-based algorithms for learning Gaussian single-index models operate in an "information exponent regime" or a "generative exponent regime," and proves that the transition between them is a phase transition.

What This Paper Is About

Learning a Gaussian single-index model — where the label is a nonlinear link function applied to a one-dimensional projection of a high-dimensional Gaussian input — has become a standard testbed for understanding feature learning. Prior theory shows that ordinary online SGD needs roughly $d^{(p-1) \lor 1}$ samples, where $d$ is the input dimension and $p$ is the target's information exponent, while algorithms that take several gradient steps per sample can be governed instead by the smaller generative exponent $p_*$. The puzzle this paper addresses is that this improvement is only observed when the learning rates used in those extra steps are large enough — and no prior work explained exactly when, or why. The authors build a single framework that captures both behaviors and pin down the learning-rate thresholds at which the complexity changes.

Key Contributions

1. A learning-rate-dependent framework for online gradient algorithms. Section 3 introduces a generic update rule with a global learning rate $\gamma$ and a second learning rate $\eta$ that scales the non-correlational part of the update. The paper's main result (Theorem 3.3) gives sample-complexity bounds in terms of Hermite coefficients $\mu_i(\eta)$ that depend explicitly on the learning rates.

2. A unified treatment of online SGD and batch reuse SGD. The framework recovers vanilla online SGD (Section 4.1) and batch reuse SGD (Section 4.2). For the latter, the analysis interpolates between the batch-reuse complexity $n = \Theta(d^{(p_*-1) \lor 1})$ and the online SGD complexity $n = \Theta(d^{(p-1) \lor 1})$ as the learning rates decrease, with proven phase transitions as a function of the learning rate for the first update on each batch.

3. A new layer-wise algorithm, "alternating SGD." Section 4.3 introduces a two-timescales algorithm that uses different learning-rate scalings for the first and second layers. With squared error and no sample reuse, it goes beyond correlational queries. When the second-layer learning rate is sufficiently large, it achieves almost linear sample complexity provided the square of the link function has information exponent 1 or 2.

4. Extension to deeper sparsely-connected networks. The layer-wise analysis is extended to a sparsely-connected network with $D$ layers, with the same conclusion holding (under further assumptions) when the $D$-th power of the link function has information exponent 1 or 2.

Main Findings

• The learning rate is as important as the algorithm design. The abstract states directly that the choice of learning rate is as important as the design of the algorithm in achieving statistical and computational efficiency. This is the paper's central claim.

• Main theorem (Theorem 3.3). Under Assumptions 3.1 and 3.2, starting from $\langle \bm{\theta}*, \bm{w}^{(0)} \rangle \asymp d^{-1/2}$, if $\gamma$ is at most $C\delta \max{1\le i \le r} \mu_i d^{-(\frac{i}{2} \lor 1)}$, then $T(\eta) = \min_{1 \le i \le r, \mu_i > 0} \tilde{\Theta}\big(\gamma^{-1} (\mu_i(\eta))^{-1} d^{\frac{i-2}{2} \lor 0}\big)$ iterations are necessary and sufficient for weak recovery, with probability at least $1-\delta$.

• Two learning rates play different roles. $\gamma$ is a "global" learning rate and $\eta$ controls the scale of any non-correlational terms in the update oracle $\psi_\eta$. Both affect sample complexity, but only $\eta$ induces the phase transition of interest. The largest feasible $\gamma$ is constrained by $\eta$, and otherwise the algorithm can diverge.

• Optimal learning-rate choice (Remark 3.4). Taking $\gamma \asymp \max_{1\le i \le r} \mu_i(\eta) d^{-(\frac{i}{2} \lor 1)}$ gives $T(\eta) = \min_{1 \le i \le r, \mu_i>0} \tilde{\Theta}\big((\mu_i(\eta))^{-2} d^{(i-1) \lor 1}\big)$. When all $\mu_i$ are non-decreasing in $\eta$, the best sample complexity is obtained by taking $\eta$ as large as possible.

• Why the transition is a genuine phase transition. The $\min$ in the complexity expression is nonsmooth, so complexity transitions sharply when the minimizing index changes. The paper identifies transitions by finding $\eta$ where $\mu_i^{-1}(\eta) d^{\frac{i-2}{2}} = \mu_j^{-1}(\eta) d^{\frac{j-2}{2}}$ for $i \neq j$.

• Online SGD baseline (Corollary 4.1). With $\psi_\eta(y,z) = y\sigma'(z)$ there is only one correlational term, no dependence on $\eta$, and no phase transition. The complexity is $\tilde{\Theta}(d^{(p-1)\lor 1})$ when $\gamma \asymp d^{-(\frac{p}{2} \lor 1)}$, matching the constraint and bound in BAGJ (21).

• Batch reuse SGD transitions (Corollary 4.2). With Algorithm 1, $T(\eta) = \min_{1 \le i \le r} \tilde{\Theta}\big((\eta d)^{-2(i-1)} d^{(p_i-1)\lor 1}\big)$. The threshold between indices $i$ and $j$ is $\eta \le d^{\frac{[(p_j-1)\lor 1] - [(p_i-1)\lor 1]}{2(j-i)} - 1}$. Small $\eta \lesssim d^{-\frac{p+1}{2}}$ recovers $T = \Theta(d^{(p-1)\lor 1})$; large $\eta \gtrsim d^{-1}$ gives $T = \tilde{\Theta}(d)$, matching LOSW (24).

• Alternating SGD transitions (Corollary 4.3). With Algorithm 2, $T(\eta) = \tilde{\Theta}(d^{(p-1)\lor 1}) \land \tilde{\Theta}(\eta^{-2} d^{(p_2-1)\lor 1})$, where $p_2 := \mathrm{IE}(\sigma_*^2)$. The transition occurs at $\eta \asymp d^{-\frac{1}{2}[(p-p_2) \lor (p-2)]}$. So alternating SGD beats online SGD when squaring the target reduces its information exponent and the second-layer learning rate is large enough.

• Experiment (Figure 1). A grid over combinations of learning rate $\eta$ and sample size $n$ shows which pairs achieve alignment $\langle \bm{w}, \bm{\theta}* \rangle \geq 0.5$ for a network with $N = 1$ trained by alternating SGD, in the setting $\sigma* = \sigma = \mathsf{He}_3$ with $d = 50$. Appendix D describes a separate batch reuse SGD experiment exhibiting the phase transition.

• Two exponents, one ordering. The information exponent is $\mathrm{IE}(g) := \min{k > 0 : u_k(g) \neq 0}$ over the Hermite expansion of $g$; the generative exponent is the smallest information exponent over all $L^2$ transformations of $g$. The paper notes $\mathrm{GE}(g) \le \mathrm{IE}(g)$ for all $g$. Lemma 2.3 (from LOSW (24)) guarantees an integer $I$ with $\mathrm{IE}(\sigma_^I) = p_$, and $I \le C_q$ if $\sigma_*$ is polynomial of degree at most $q$.

• The motivating puzzle. Full-batch gradient flow reuses the whole dataset at every iteration, yet the best known upper bounds for it on squared loss still depend on the information exponent (BBSS 22; MHWSE 23). The authors identify the neglected learning rate as the missing ingredient explaining this apparent contradiction.

Methodology in Plain English

The authors study a single abstract update rule rather than a specific algorithm. At each iteration the first-layer weight is nudged in a direction determined by a function of the current label $y$ and the current pre-activation $\langle \bm{x}, \bm{w} \rangle$, then renormalized onto the unit sphere. This "general gradient oracle" can represent one gradient step, two gradient steps on the same sample, or a layer-wise alternation — the differences are entirely captured by one function and by the two learning-rate parameters.

To analyze learning, the authors project the weight dynamics onto the direction of the true signal $\bm{\theta}_*$ and track how that alignment grows over iterations. Expanding the update function in Hermite polynomials turns the problem into tracking a small number of scalar coefficients, which the paper calls $\mu_i(\eta)$. Each coefficient measures how strongly the update correlates with the $i$-th Hermite component of the target. The sample complexity is then governed by whichever $i$ gives the smallest required number of iterations, producing a minimum over $i$. Because each $\mu_i(\eta)$ grows with $\eta$ at a different rate (different powers of $\eta d$ appear for different indices, as in the batch reuse Taylor expansion), increasing $\eta$ changes which index wins the minimum — and that switch is the phase transition. The authors then instantiate this recipe for three concrete algorithms and check that it reproduces known bounds, and they run a small experiment to visualize the predicted threshold.

Why This Matters

Impact on research. The paper reconciles two disconnected literatures: work on information-exponent-limited online SGD and work on generative-exponent-limited batch-reuse or alternative-loss algorithms. It shows these are not different algorithms so much as different learning-rate regimes of the same family, and it supplies a general theorem that prior analyses of one-pass SGD and batch reuse can be recovered from as special cases. It also introduces a mechanism — two-timescale layer-wise updates — that escapes correlational-query limits without reusing samples or changing the loss away from squared error, which JMS (24) had identified as an open direction.

Real-world applications (the paper is theoretical and does not report these applications; they are plausible implications of the theory):

• Choosing learning rates and schedules in fine-tuning pipelines, where the theory suggests a step size that is too small can silently push training into a sample-hungrier regime.
• Two-stage or layer-wise training recipes already common in practice, which the alternating SGD analysis gives a theoretical rationale for.
• Sample-efficiency planning for high-dimensional regression problems where data collection is expensive and the input dimension is large.
• Diagnosing training runs in which loss curves plateau despite abundant data — the theory predicts a threshold below which progress stalls until enough samples are accumulated.

Industry relevance. Practitioners rarely set learning rates optimally; the paper explicitly frames its motivation as understanding the effect of non-optimal choices in practice. A precise account of when a too-small step size changes the required sample size is directly relevant to compute budgeting and to deciding whether to buy more data or retune hyperparameters.

Future Directions

• What other training components matter? The paper notes that JMS (24) showed alternative loss functions can also break the information-exponent curse, and explicitly leaves open the study of other components of the training algorithm.

• Closing the gap to SQ-optimality. The authors observe that the weight perturbation and averaging of CWL+ (25) is complementary to their work and can push the algorithms they study toward SQ-optimality when both $\gamma$ and $\eta$ are chosen as large as possible — an avenue they do not pursue.

• From proof of concept to practice. The authors describe alternating SGD as useful from a theoretical perspective and a proof of concept, implying that the question of whether such two-timescale mechanisms are at play in real neural network training remains open.

• Beyond single-index and Gaussian settings. The paper's related work points to general multi-index models, which remain difficult to analyze, and to recent work on spherically symmetric input distributions — both natural extensions of the framework. The $D$-layer extension in Appendix C.4 also holds only under further assumptions on $\sigma_*$ and the activations $\sigma_j$.

Target Audience

Theoretical machine learning researchers working on feature learning, single-index and multi-index models, statistical query lower bounds, and the optimization dynamics of gradient descent. It will also be useful to mathematically literate practitioners who want a principled account of how learning rate affects sample complexity, and to readers already familiar with the information exponent and generative exponent literature who want a unified framework connecting them. Readers without a background in Hermite analysis and high-dimensional asymptotics will find the technical core demanding, though the framing question — how much does the learning rate matter? — is broadly accessible.

Authors’ abstract

To understand feature learning dynamics in neural networks, recent theoretical works have focused on gradient-based learning of Gaussian single-index models, where the label is a nonlinear function of a latent one-dimensional projection of the input. While the sample complexity of online SGD is determined by the information exponent of the link function, recent works improved this by performing multiple gradient steps on the same sample with different learning rates -- yielding a non-correlational update rule -- and instead are limited by the (potentially much smaller) generative exponent. However, this picture is only valid when these learning rates are sufficiently large. In this paper, we characterize the relationship between learning rate(s) and sample complexity for a broad class of gradient-based algorithms that encapsulates both correlational and non-correlational updates. We demonstrate that, in certain cases, there is a phase transition from an "information exponent regime" with small learning rate to a "generative exponent regime" with large learning rate. Our framework covers prior analyses of one-pass SGD and SGD with batch reuse, while also introducing a new layer-wise training algorithm that leverages a two-timescales approach (via different learning rates for each layer) to go beyond correlational queries without reusing samples or modifying the loss from squared error. Our theoretical study demonstrates that the choice of learning rate is as important as the design of the algorithm in achieving statistical and computational efficiency.