PaperlayerView exact revision on arXiv

Source-linked AI summary

High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation

Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, Greg Yang

arXiv:2205.01445v1stat.MLcs.LGmath.ST

TL;DR

The paper asks whether a single gradient step can produce useful feature learning beyond fixed-kernel behavior in high-dimensional neural networks. It analyzes the first-layer update and conjugate-kernel ridge risk in a proportional student-teacher setting, finding that learning-rate scaling separates modest linear-regime gains from stronger improvements beyond the linear regime. The results are established under specific modeling assumptions, including Gaussian inputs and smooth centered activations.

  • Problem

    The paper studies whether early gradient descent can capture feature learning that fixed random features and kernel methods miss, and whether one step improves over the initial kernel in the proportional limit.

  • Method

    The authors analyze one first-layer gradient step and the resulting conjugate-kernel ridge estimator in a proportional student-teacher model, using Gaussian-equivalence and random-matrix methods across two learning-rate scalings.

  • Results

    One step improves on the initial random-features ridge model at η = Θ(1) but remains in the linear regime, whereas sufficiently large learning rates can yield lower risk than the nonlinear-component lower bound for certain teachers and outperform broad kernel classes.

  • Takeaways & Limitations

    A single feature-learning step can substantially improve over random features, and its ability to move beyond linear-kernel behavior depends on learning-rate scaling.

  • Takeaways & Limitations

    The analysis assumes smooth centered activations and Gaussian inputs, and early training may be insufficient when gradient descent undergoes a long initial search before alignment.

Abstract

from arXiv · show

We study the first gradient descent step on the first-layer parameters $\boldsymbol{W}$ in a two-layer neural network: $f(\boldsymbol{x}) = \frac{1}{\sqrt{N}}\boldsymbol{a}^\topσ(\boldsymbol{W}^\top\boldsymbol{x})$, where $\boldsymbol{W}\in\mathbb{R}^{d\times N}, \boldsymbol{a}\in\mathbb{R}^{N}$ are randomly initialized, and the training objective is the empirical MSE loss: $\frac{1}{n}\sum_{i=1}^n (f(\boldsymbol{x}_i)-y_i)^2$. In the proportional asymptotic limit where $n,d,N\to\infty$ at the same rate, and an idealized student-teacher setting, we show that the first gradient update contains a rank-1 "spike", which results in an alignment between the first-layer weights and the linear component of the teacher model $f^*$. To characterize the impact of this alignment, we compute the prediction risk of ridge regression on the conjugate kernel after one gradient step on $\boldsymbol{W}$ with learning rate $η$, when $f^*$ is a single-index model. We consider two scalings of the first step learning rate $η$. For small $η$, we establish a Gaussian equivalence property for the trained feature map, and prove that the learned kernel improves upon the initial random features model, but cannot defeat the best linear model on the input. Whereas for sufficiently large $η$, we prove that for certain $f^*$, the same ridge estimator on trained features can go beyond this "linear regime" and outperform a wide range of random features and rotationally invariant kernels. Our results demonstrate that even one gradient step can lead to a considerable advantage over random features, and highlight the role of learning rate scaling in the initial phase of training.

1 Introduction

The paper studies whether one gradient step can learn useful features beyond fixed random features in the proportional asymptotic regime. It shows that the update creates alignment with the teacher and that learning-rate scaling determines whether trained features merely improve within the linear regime or surpass it.

  • Motivation: Feature learning addresses a limitation of random features, whose fixed representations do not fully explain neural networks’ empirical success.Prior work motivates studying adaptive features because gradient-trained networks can outperform broad classes of linear and kernel estimators.
  • Setup: The analysis studies one first-layer gradient step in a student-teacher regression model as n, d, and N jointly diverge proportionally.It evaluates conjugate-kernel ridge regression before and after updating W from Gaussian initialization.
  • Alignment: The first update is approximately rank-1, producing an information spike plus noise bulk whose isolated singular vector aligns with the teacher’s linear component.The top conjugate-kernel eigenvector also aligns with the training labels.
  • Learning-rate regimes: For η = Θ(1), one step gives nontrivial improvement over the initial random-features ridge model but remains unable to outperform the best linear estimator.The result follows from extending Gaussian equivalence to the updated feature map and characterizing risk with random-matrix tools.
  • Learning-rate regimes: For sufficiently large learning rates, certain teacher functions achieve lower risk after one feature-learning step than the nonlinear-component lower bound, outperforming a wide range of kernel ridge estimators.The comparison includes the neural tangent kernel and is illustrated through prediction-risk curves for erf features.
  • Context: The paper connects these results to proportional-limit kernel asymptotics, where many fixed kernels behave essentially linearly and typically cannot learn nonlinear target components.Training the feature map can overcome this limitation in the n ≍ d scaling for sufficiently large learning rate.

2 Problem Setup and Basic Assumptions

The paper studies first-layer feature learning with gradient descent in the proportional asymptotic limit, then evaluates the resulting representation using ridge regression on fresh data. Its student-teacher analysis assumes Gaussian inputs, sub-Gaussian noise, normalized smooth activations, and a possible initial search stage for some teachers.

  • The network is trained by minimizing empirical squared loss, with the first layer updated while the second-layer coefficients remain fixed at initialization.
  • For certain teacher models, gradient descent may require a long initial search stage before achieving nontrivial alignment.
  • After one feature-learning step, ridge coefficients are estimated on an independent fresh dataset to avoid dependence between updated features and training inputs.
  • Performance is measured by prediction risk R( f̂ ) = E_x( f̂(x) − f*(x) )^2 under the same Gaussian test distribution.
  • The proportional limit sends n, d, and N to infinity with fixed ratios n/d and N/d.
  • The student-teacher setup uses Gaussian inputs, sub-Gaussian label noise, and centered activations with nonzero linear Gaussian component.

4. Gaussian Initialization. √

The paper contrasts random-feature and rotationally invariant kernel baselines with feature learning under mean-field scaling. It uses ridge regression on trained conjugate-kernel features and identifies a lower bound tied to the teacher’s nonlinear component.

  • Mean-field scaling lets neurons move away from initialization and align with the target, unlike the NTK parameterization where parameters remain close to initialization.
  • The experiments include a two-ReLU-neuron teacher and student, with d = 512, ψ1 = ψ2 = 10, and optimization continuing until both training losses fall below 10^-3.
  • The comparison evaluates ridge regression on trained CK features against initial random-feature, NTK, and rotationally invariant kernel estimators.
  • The ridge estimators use the feature maps and kernel matrices associated with CK, NTK, and rotationally invariant kernels.
  • RF and rotationally invariant kernels cannot outperform the best linear estimator unless the teacher is linear, meaning ∥P>1f*∥L2 = 0.

3 How Does One Gradient Step Change the Weights?

One gradient step produces an approximately rank-1 perturbation of the first-layer weights, creating a spike whose singular vector aligns with the teacher’s linear component. The corresponding CK spectrum inherits a spike and label-aligned top component, while the bulk remains largely unchanged.

  • 3 How Does One Gradient Step Change the Weights?: The first-step gradient is approximately rank-1, so updated weights have an information-spike plus noise-bulk structure.
  • 3.1 Almost Rank-1 Property of the Gradient Matrix: The rank-1 perturbation is linked to labels and produces a spike in W1 when the learning rate is sufficiently large.
  • 3.2 Alignment with the Target Function: Increasing the learning rate beyond a threshold enlarges the leading spike, and alignment improves with larger linear components of the activation and teacher.
  • 3.2 Alignment with the Target Function: For η = Θ(1), the leading singular vector aligns with the teacher’s linear component β*, while the remaining singular-value bulk stays asymptotically unchanged.
  • A Spiked Model for CK?: The CK population covariance inherits the same BBP transition, and the empirical CK matrix develops an isolated eigenvalue after one gradient step.
  • A Spiked Model for CK?: The top CK principal component correlates with training labels, as captured by the Gaussian-equivalent model.

4 Do the Learned Features Improve Generalization?

One gradient step improves ridge prediction on trained conjugate-kernel features. With η = Θ(1), the improvement is nontrivial but remains bounded by the best linear model, whereas sufficiently large learning rates can surpass that kernel lower bound for specific teacher–activation settings.

  • Feature alignment: After one gradient step, first-layer weights align with the teacher’s linear component, motivating improved ridge prediction on the trained conjugate kernel.The analysis specializes to single-index teachers and compares prediction risk before and after updating W.
  • Small learning rate: The Gaussian equivalence theorem extends to the trained feature map under η = Θ(1), accurately describing ridge risk at any fixed number of steps.The theorem requires an odd activation and independent data for learning W and estimating the ridge coefficients.
  • Small learning rate: For η = Θ(1), one-step training yields nontrivial improvement over the initial random-features ridge estimator but cannot outperform the best linear estimator.The Gaussian-equivalent model cannot learn the nonlinear component of the target, imposing the stated limitation.
  • Small learning rate: The improvement from one step is always positive under the stated proportional conditions, but becomes negligible as width N/d grows without bound.Theorem 7 reports improvement across positive aspect ratios, while Proposition 9 gives a vanishing improvement in the large-width regime.
  • Large learning rate: With a sufficiently large step, trained features can outperform the kernel lower bound when the nonlinear teacher component is large relative to τ* and n/d is sufficiently large.For erf activations, the resulting risk is O(d/n); for tanh, it is below the nonlinear-component lower bound under the stated conditions.
  • Large learning rate: The large-step separation is specific to activation pairs with sufficiently small τ*, and generally one gradient step may not suffice to learn a good representation.The paper identifies this as a scope boundary for the large-learning-rate result.

5 Discussion and Conclusion

The paper rigorously characterizes how one gradient step creates feature-learning structure in the conjugate kernel and improves prediction risk under different learning-rate scalings. Its main limitations concern data independence, unresolved learning-rate transitions, and an unproved CK-spike hypothesis.

  • One gradient step yields a signal-plus-noise decomposition of the updated weights and improves conjugate-kernel ridge risk under two learning-rate scalings.
  • Limitations: The analysis assumes the trained weights are independent of the data used to compute the conjugate kernel, excluding feature learning and kernel evaluation on the same data.The authors identify pretrained or transfer-learning analysis as a natural setting for this assumption.
  • Limitations: The regimes η = Θ(1) and η = Θ(√N) behave drastically differently, but whether a phase transition connects them remains open.
  • Limitations: The isolated eigenvalue and eigenvector of the trained conjugate-kernel matrix are supported by a Gaussian-equivalence hypothesis whose validity remains an important open question.

A.1 Additional Experiments

The additional experiments examine alignment, prediction risk, and kernel target alignment after feature learning, while also documenting where the Gaussian-equivalence analysis can fail. They further compare mean-field feature learning with kernel-regime behavior and relate the analysis to linear-regression lower bounds and prior kernel methods.

  • Teacher-component learning: The student first aligns with the teacher’s linear component before beginning to learn its nonlinear component.The experiment tracks separate alignment quantities for the linear and nonlinear components of the teacher.
  • Failure cases: Gaussian-equivalence predictions eventually become inaccurate as gradient steps or learning rate increase, and at approximately 30 steps they overestimate prediction risk.The stated failure occurs because the first-layer weights must move sufficiently far from initialization to achieve small training loss.
  • Alignment: Theorem 3 predictions closely match empirical leading-singular-vector alignment, which increases with learning rate η and sample ratio ψ1 = n/d.The alignment is measured between the leading eigenvector of W 1 and the teacher’s linear component β∗.
  • Large learning rate: For SoftPlus, the large-learning-rate CK ridge risk remains non-vanishing, consistent with the predictive value of the upper-bounding quantity τ∗ despite the absence of a precise asymptotic characterization.Here τ∗≈0.03 > 0, so the upper bound in Theorem 11 is non-vanishing.
  • Kernel target alignment: Mean-field trained CK aligns with both training and test labels, whereas the kernel-regime neural network does not exhibit such alignment.Kernel target alignment is tracked during training until the training loss reaches 10^-2.
  • Relation to prior regimes: The analysis focuses on the early phase of mean-field-scaled training, avoiding the nonlinear PDE by considering only one gradient step.The section contrasts this approach with local kernel approximations and mean-field analyses of full parameter evolution.

B.1.1 Norm Bounds for the First Gradient Step

This appendix develops norm and probability controls for the first gradient step under Lipschitz and moment assumptions. The bounds decompose the relevant matrices and control their random-feature, target-function, and noise contributions in the proportional regime.

  • Matrix decomposition: The first-gradient-step analysis begins from the squared-loss gradient matrix and decomposes the key matrix into a rank-one term plus remainder terms.The proof separately analyzes three matrices and later uses G0 − A = B + C.
  • Assumptions and bounds: Under Lipschitz assumptions on f∗ and σ, the appendix establishes norm and probability bounds for the decomposed terms, with constants depending on model moments and ||f∗||L2.The stated assumptions include f∗ ∈ L2(Rd, Γ), and the constants depend on λσ, μ1, σε, and ||f∗||L2(Rd, Γ).
  • Scaling dependence: The analysis tracks dependence on n, d, and N separately rather than simplifying immediately using proportional scaling.A remark states that this separation is needed in some calculations, even though the results are ultimately used in the proportional regime.
  • Concentration arguments: The proof controls random-feature and Gaussian quantities using concentration tools including Bernstein, Hanson–Wright, Lipschitz Gaussian concentration, and sub-Gaussian tail bounds.These controls are applied to operator norms, Frobenius norms, quadratic forms, and sup-norms across the decomposed terms.
  • Extension to fixed steps: For any fixed t ∈ N and learning rate η = Θ(1), Lemma 16 gives a bound on the weight matrix after t gradient steps.This extends the appendix’s control from the first step to any fixed number of steps under constant-order learning rate.

B.2.2 Analysis of Spike in Weight Matrix

The first-layer update is analyzed as a rank-one perturbation of a Gaussian weight matrix. Depending on the signal strength, this perturbation creates an isolated spectral spike whose leading singular vector aligns with the teacher’s linear component.

  • The updated matrix is approximated by a Gaussian matrix plus an independent rank-one perturbation.This reduces the spectral analysis to a spiked random-matrix problem.
  • The non-spike singular values remain at the Marchenko–Pastur bulk edge, while the leading singular value can separate from that bulk.For fixed i > 1, the corresponding singular values converge to 1 + √ψ2.
  • When the signal exceeds the threshold, the isolated eigenvalue converges to λ0 outside the bulk spectrum.When the threshold is not met, the largest eigenvalue instead converges to the bulk edge (1+√ψ2)^2.
  • Above the spike threshold, the leading left singular vector of the updated weights asymptotically aligns with the teacher’s linear component β∗.The alignment transfers from the rank-one approximation to the original updated matrix through Davis–Kahan perturbation control.
  • For η = Θ(1), the Gaussian equivalence property remains valid when the weights stay close to initialization.This property supports the subsequent characterization of prediction risk on the updated features.

C.2 Prediction Risk of the Gaussian Equivalent Model

The prediction risk after one update is analyzed through a Gaussian-equivalent feature model. Its risk separates into bias from learning the teacher and variance from label noise, while the small-learning-rate analysis preserves the variance term.

  • The updated-feature ridge estimator is replaced asymptotically by a Gaussian covariates model using Gaussian equivalence.The resulting risk is denoted RGE(λ).
  • The Gaussian-equivalent prediction risk decomposes into bias from estimating f∗ and variance from the label noise ε̃.This decomposition organizes the comparison between initialized and trained features.
  • For η = Θ(1), the variance term in the risk is unchanged after one gradient step.The analysis therefore focuses on how the update changes the bias terms.
  • The first-step weight update can be reduced to the rank-one form W1 = W0 + ua⊤ without changing the asymptotic bias terms.Here u is proportional to the data-gradient direction and the output-weight vector a supplies the rank-one factor.
  • The teacher function is decomposed into linear and nonlinear components for analyzing the updated kernel risk.This decomposition is used to track the terms contributing to the prediction-risk calculation.

C.3.1 Concentration and Simplification

The concentration analysis uses rotational invariance and concentration inequalities to simplify the rank-one update’s risk terms. The resulting quantities converge to deterministic limits in the proportional asymptotic regime.

  • Each target quantity concentrates around a deterministic limit as n, d, and N grow proportionally.The analysis establishes this concentration for the terms appearing in the risk decomposition.
  • Rotational invariance allows the teacher direction β∗ to be treated as uniformly distributed on the sphere without changing the quantities’ distributions.This simplifies the dependence on the fixed teacher direction.
  • Nonlinear Hanson–Wright concentration controls the quadratic forms involving the random teacher direction.The argument applies concentration after conditioning on bounded-norm events.
  • The expected-kernel contribution is controlled entrywise using the Lipschitz property of the teacher’s activation function.The resulting trace term converges to zero in probability.
  • Rotationally invariant random matrices with bounded operator norm permit further simplification of the quadratic forms.This lemma is used to show that the relevant finite-dimensional terms converge to their deterministic counterparts.

C.3.2 Risk Calculation via Linear Pencils

The risk calculation expresses asymptotic bias and variance terms through self-consistent Stieltjes-transform equations and linear-pencil methods. For η = Θ(1), the resulting risk difference is nonnegative, so one gradient step does not worsen CK ridge risk asymptotically.

  • The asymptotic quantities governing risk are characterized by self-consistent equations involving Stieltjes transforms.These transforms describe the limiting spectra of the initial conjugate kernel and its companion matrix.
  • For each risk-decomposition term Ti, the limiting value is given by the corresponding deterministic quantity τi.Proposition 29 supplies these limits under the stated assumptions and λ > 0.
  • Linear-pencil methods reduce the required trace calculations to large block-matrix limits connected to prior kernel constructions.This approach is used for the difficult terms in the bias–variance expansion.
  • The prediction-risk difference R0(λ) − R1(λ) is expressed through the limiting τi terms and parameters of the proportional asymptotic model.The result applies to fixed η = Θ(1) and positive ridge parameter λ.
  • The asymptotic risk difference is zero when η = 0 or when the teacher’s linear component has zero relevant coefficient.Thus, the improvement depends on both a nonzero update and a nonzero linear teacher signal.
  • The risk-improvement function δ(η, λ, ψ1, ψ2) is nonnegative.It vanishes when at least one of the relevant signal, gradient, or learning-rate factors is zero.

C.4.1 Case I: Large sample limit

In the large-sample limit ψ1 → ∞, the prediction-risk improvement from one feature-learning step converges to δ(η, λ, ∞, ψ2). This limiting improvement is non-negative, vanishes only in specified degenerate cases, and increases with the learning rate.

  • Large-sample limit: The asymptotic analysis reduces the relevant quantities to Marchenko–Pastur Stieltjes transforms and their derivatives.The limits of τi and mi are obtained using self-consistent equations and the Marchenko–Pastur distribution.
  • Large-sample limit: Theorem 31 identifies the limiting risk difference R0(λ) − R1(λ) with δ(η, λ, ∞, ψ2).The limit is taken first over n,d,N and then as ψ1 → ∞.
  • Large-sample limit: δ(η, λ, ∞, ψ2) is non-negative for η, λ, ψ2 ∈ (0, +∞).The proof establishes non-negativity by decomposing the expression into terms whose signs are controlled through Stieltjes-transform identities.
  • Large-sample limit: δ = 0 if and only if one of µ1, µ∗1, or η is zero.Thus the limiting improvement is strictly positive outside these stated degenerate cases.
  • Large-sample limit: δ(η, λ, ∞, ψ2) increases with the learning rate η ≥ 0.The argument differentiates δ with respect to the η-dependent parameter θ2 and verifies positivity.

D Proof for Large Learning Rate (η = Θ(N))

For η = Θ(N), Gaussian equivalence is unavailable, so the proof constructs an oracle second layer and bounds the trained-feature kernel ridge risk through approximation and concentration arguments. The constructed predictor is tied to neurons whose updated weights approximately match the teacher direction.

  • Proof strategy: For η = Θ(N), the proof replaces Gaussian equivalence with an upper bound on the conjugate-kernel ridge prediction risk.The argument constructs an oracle second layer and then transfers its risk bound to ridge regression.
  • Oracle construction: An oracle second layer is constructed by averaging neurons whose first-layer weights approximately match the target direction β∗.The selected neurons form a subset Ar, and their coefficients are averaged to construct the predictor.
  • Approximation: The trained subnetwork is approximated by an expected single-index model with an O(d/n)-type substitution error.The expected model shifts Gaussian weights by a multiple of β∗, where the multiple depends on the learning rate and nonlinearities.
  • Risk reduction: The oracle risk reduces to a one-dimensional expectation τ∗ after choosing an appropriate neuron subset.This reduction yields a high-probability upper bound for the constructed predictor up to Θ(d/n)-error.
  • Ridge-risk transfer: When n ≫ d, conjugate-kernel ridge regression on the trained weights has prediction risk bounded above by τ∗.The transfer uses classical kernel ridge regression analysis after constructing a predictor with the required risk.

D.2 Constructing the “Oracle” Estimator

The oracle estimator averages a large subset of neurons selected through their initialized second-layer coefficients and aligned first-layer updates. Its risk is controlled by a one-dimensional Gaussian objective, with explicit small values for several activation pairs.

  • Estimator construction: For any constant α, many initialized coefficients ai are close to α with high probability, enabling construction of the subset Ar.The subset size scales polynomially with N for suitable r.
  • Estimator construction: The oracle predictor averages neurons from a subset Ar whose updated first-layer weights are aligned with β∗.The construction replaces the full gradient matrix by its rank-1 component and controls the residual terms separately.
  • Approximation: The expected predictor differs from the constructed subnetwork by a vanishing approximation error as the selected subset grows.The bound scales as 1/√Nr under the stated concentration argument.
  • Risk characterization: The oracle risk τ∗ is obtained by optimizing a finite one-dimensional Gaussian expectation over the coefficient parameter α.The equivalent parameterization uses κ, which depends on α, the learning-rate constant, and model nonlinearities.
  • Examples: τ∗ ≈ 3 × 10^-4 for σ = σ∗ = tanh, with κ∗ ≈ 1.6.These values come from numerical integration.
  • Examples: τ∗ ≈ 0.03 for σ = σ∗ = SoftPlus, with κ∗ ≈ 0.96.These values come from numerical integration.

D.3 Prediction Risk of Ridge Regression

The ridge-regression analysis decomposes prediction risk into approximation, bias, and variance components and controls the empirical feature covariance by concentration. With sufficiently strong regularization and a proportional asymptotic regime, the estimator’s risk is bounded by the oracle risk τ∗.

  • Covariance concentration: The analysis controls the trained-feature conjugate-kernel ridge estimator through a concentration event for the empirical feature covariance.The concentration argument uses sub-Gaussian feature vectors and high-probability covariance bounds.
  • RKHS comparison: The ridge estimator is compared with an RKHS predictor obtained by projecting the target onto the trained-feature function space.The target is decomposed as f∗ = f̌ + f⊥, where f⊥ is orthogonal to the RKHS.
  • Risk decomposition: The prediction risk decomposes into B1, B2, V1, and V2, separating approximation and variance-related contributions.The decomposition is used to establish separate bounds before combining them.
  • Regularization: For λ = Ω(n^(ε−1)) with small ε > 0, the variance contribution V1 + V2 is o_d,P(1).Under the same scaling, the remaining B1 and B2 terms are finite.
  • Final bound: When n, d, and N grow proportionally and n/d exceeds a threshold ψ∗1, the trained-feature ridge risk is bounded by τ∗.The bound follows from the constructed oracle predictor and the regularized risk decomposition.
Loading 2205.01445v1…