stat.MLMay 18, 2026

Feature Learning in Linear-Width Two-Layer Networks: Two vs. One Step of Gradient Descent

Authors: Behrad Moniri, Hamed Hassani

Organizations: University of Pennsylvania, Philadelphia, PA, USA

Abstract

We study feature learning in two-layer neural networks within the linear-width regime, where the number of hidden neurons, sample size, and input dimension scale proportionally. While recent work has analyzed feature learning via a single step of gradient descent on the first layer weights in this regime, such one-step update schemes are fundamentally limited: the update to the weights is approximately rank-one, captures only a single direction, and requires the target function to have an information exponent of one. In this paper, we go beyond one-step updates to provide a full characterization of the features learned during the \textit{second step} of gradient descent with step-sizes η1≍Nα1η_1\asymp N^{α_1} and η2≍Nα2η_2 \asymp N^{α_2} for α1,α2∈[0,0.5)α_1, α_2 \in [0,0.5), where NN is the number of hidden neurons. We derive a spectral characterization of the updated weights, demonstrating they behave as a spiked random matrix with multiple outliers, each corresponding to a learned direction. We show that the number of the outliers is determined by the parameters α1,α2α_1, α_2 through ⌊α21/2−α1⌋\lfloor \frac{α_2}{1/2 - α_1} \rfloor. Furthermore, by analyzing the alignment between the learned directions and the target function, we identify a gap between training with independent versus reused batches. While independent batches restrict learning to directions with an information exponent of one, batch reuse enables the second update to capture directions even when the information exponent exceeds one, provided that α1,α2α_1, α_2 are chosen properly. This shows that the benefits of batch reuse, previously observed in narrow-width regimes, persist in the linear-width limit as well. By characterizing these early-phase evolutions, our work proposes a tractable framework for studying optimization and feature learning phenomenology in modern overparameterized networks.

Explore similar work

CardsList
May 23, 2026cs.LG

Feature Learning in Wide Neural Networks under μP: Identifiability and Sparse-Dictionary Decomposition of the Mean-Field Limit

We establish four structural results for feature learning in wide two-layer neural networks under the Maximal Update Parametrization (μμP). First, we prove global existence and uniqueness of the mean-field limit of noisy gradient descent under μμP, identifying the maximal admissible weight w∗w^* on the moment sequence of the initialization as the reciprocal parameter-moment-growth boundary, and hence the largest weighted moment class propagated by the flow. The finite-particle approximation has uniform-in-time squared-Wasserstein rate O(N−1)O(N^{-1}). Second, we characterize identifiability of the mean-field limit: two admissible parameter measures induce the same network function in L2L^2 exactly when their active components agree modulo the finite-rank realization symmetry of the architecture. The orbit depth Dorb∗D^*_{\mathrm{orb}} is separated from the moment-variety depth Dvar∗D^*_{\mathrm{var}}. Third, under the Barron-Hermite target condition the active support of the long-time limit measure admits a sparse-dictionary decomposition: it is supported on at most S∗S^* atoms modulo finite-rank realization symmetry, with S∗S^* bounded by an explicit coefficient-threshold number. Fourth, we derive the total feature-learning-error decomposition into statistical, optimization, propagation-of-chaos, and sparse-residual components, with a target-dependent Hermite/Barron tail replacing any initialization-only residual. The four results are tied together by an architectural identity: the triple (w∗,Dorb∗,S∗)(w^*, D^*_{\mathrm{orb}}, S^*) -- the maximal admissible weight, the orbit identifiability depth, and the sparse-dictionary depth at which the target is realizable -- is the natural learning cell of the architecture-data pair (σ,ρ)(σ, ρ). The proofs are self-contained except for standard results from μμP and mean-field Langevin theory.
Akmal Xodarev
Two-Layer Neural NetworksFeature LearningMean-Field LimitNeural NetworkIdentifiability
AskPDF
May 29, 2026cs.LG

Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks

We study optimal learning-rate selection in two-layer and three-layer linear neural networks trained to learn linear target functions. In particular, we derive the exact closed-form expressions for the gradients and test loss after one and two steps of gradient descent, enabling a precise characterization of early training dynamics. We characterize how learning rates should scale under the gradient approximation in the first two steps, and prove that performing updates with this approximation yields a tractable surrogate loss with a tight, small approximation error. This formulation enables the theoretical analysis of layer-wise learning rates and reveals a distinct early-training regime: test loss can be minimized by unequal learning rates at the initial step, while equal learning rates become optimal in subsequent steps. Our numerical experiments validate the theory and demonstrate the importance of balancing layer-wise learning rates early during training. The code is available at: https://github.com/TDCSZ327/Layer-Balancing.
Tianyu Pang, Vignesh Kothapalli, Shenyang Deng +3
Learning RatesTwo-Layer Neural NetworksGradient DescentNeural Scaling Laws
AskPDF
May 8, 2026cond-mat.dis-nn

Spectral Dynamics in Deep Networks: Feature Learning, Outlier Escape, and Learning Rate Transfer

We study the evolution of hidden-weight spectra in wide neural networks trained by (stochastic) gradient descent. We develop a two-level dynamical mean-field theory (DMFT) that jointly tracks bulk and outlier spectral dynamics for spiked ensembles whose spike directions remain statistically dependent on the random bulk. We apply this framework to two settings: (1) infinite-width nonlinear networks in mean-field/μμP scaling and (2) deep linear networks in the proportional high-dimensional limit, where width, input dimension, and sample size diverge with fixed ratios. Our theory predicts how outliers evolve with training time, width, output scale, and initialization variance. In deep linear networks, μμP yields width-consistent outlier dynamics and hyperparameter transfer, including width-stable growth of the leading NTK mode toward the edge of stability (EoS). In contrast, NTK parameterization exhibits strongly width-dependent outlier dynamics, despite converging to a stable large-width limit. We show that this bulk+outlier picture is descriptive of simple tasks with small output channels, but that tasks involving large numbers of outputs (ImageNet classification or GPT language modeling) are better described by a restructuring of the spectral bulk. We develop a toy model with extensive output channels that recapitulates this phenomenon and show that edge of the spectrum still converges for sufficiently wide networks.
Clarissa Lauditi, Cengiz Pehlevan, Blake Bordelon
Deep NetworkNeural NetworkFeature LearningMean-Field LimitTransfer LearningGrowth+1
AskPDF