Investigating the Early Signatures of Feature Learning in Deep Linear Networks

Shaunak Bhandarkar; Boris Hanin

stat.MLOct 27, 2026

Investigating the Early Signatures of Feature Learning in Deep Linear Networks

Authors: Shaunak Bhandarkar, Boris Hanin

Organizations: Princeton University

Abstract

Deep neural networks learn data-dependent features, a capacity widely credited for their empirical success. However, explicitly characterizing what features emerge -- and when they improve generalization -- is challenging, since training is nonlinear and non-convex and the data is often anisotropic. Here, we give a closed-form account of feature learning for a deep linear network in a mean-field parameterization, trained by one gradient step from initialization on high-dimensional data carrying a low-dimensional (spiked) signal. In the infinite-width mean-field limit, one step updates each layer's feature kernel by a target-aligned rank-one spike, amplified by both the learning rate and the layer depth. We next assess how these features impact prediction through kernel ridge regression (KRR), finding that KRR on the learned feature kernel yields scalar inflation of the KRR estimator at initialization -- an "anti-shrinkage" mechanism that reinflates task-relevant directions suppressed by ridge, while also amplifying noise. To study when this inflation improves generalization, in the overparameterized proportional limit we derive closed-form expressions for the bias, variance, and generalization errors of these estimators in terms of key macroscopic parameters -- notably the spike strength and spike-target alignment. Crucially, we find that a one-step generalization advantage vanishes for isotropic inputs and appears when both the spike strength and an effective signal-to-noise ratio are sufficiently large, demonstrating that a successful first gradient step both requires and benefits from low-dimensional task structure.

Figures & tables

Figure 1: A phase diagram of the predicted minimal asymptotic generalization error (over ridge strengths $\lambda\geq 0$ ) arising from learned one-step layer- $\ell$ features (with effective update strength $k_{\ell}:=\frac{\eta^{2}\ell(\ell+1)(2\ell+1)}{6D}\to 10$ ) relative to baseline features at initialization. Here, $n/D\to\psi=0.2$ and the noise level is $\sigma^{2}=1$ . Blue denotes a feature learning advantage, and red denotes the opposite. The shaded region marks an “SNR window” in which a global feature learning advantage provably holds: certain one-step ridge estimators ( $k_{\ell}>0$ ) can outperform any baseline ridge estimator (see Sec. 4.3 ). To the right, generalization error curves over $\lambda$ are shown for one-step ( $k_{\ell}\to 10$ ; blue) and baseline (red) estimators for selected values of $\gamma$ and the SNR ( $\rho^{2}/\sigma^{2}$ ). A. $\gamma=0.25$ , $\text{SNR}=1.0$ . B. $\gamma=20$ , $\text{SNR}=0.0049$ . C. $\gamma=10$ , $\text{SNR}=0.25$ . Dots mark empirical simulations (with SEM; $D=5000$ , $n=1000$ ) and lines indicate theory. Stars show location of the minimum generalization error attainable by either model. See Appendix A for more details.

Figure 2: Empirical validation that one GD step in a deep linear network asymptotically updates

\Phi^{\ell}

by the rank-one spike

\frac{\eta^{2}d_{\ell}}{n^{2}}K_{x}yy^{T}K_{x}

. Empirical simulations are shown for various widths

N

(with SEM) and dashed lines show theory. Here,

L=10

D=30

n=10

, and

\eta=0.5

. A. The feature kernel update at each layer

\ell

structurally matches the rank-one matrix

K_{x}yy^{T}K_{x}

at large

N

. B. Empirical estimation of the depth coefficient

d_{\ell}

matches theory at large

N

Figure 3: Asymptotics of $c_{\lambda}$ across ridge strengths ( $\lambda$ ) and feature learning strengths ( $k_{\ell}$ ); here, $\psi=\sigma=0.5$ . A. $c_{\lambda}$ across $\lambda$ for different choices of $k_{\ell}$ , both for isotropic data (left) and spiked data (right). For the choice of spiked data (right), even small $k_{\ell}$ saturates $c_{\lambda}$ . Dots represent empirical simulations (with SEM; $D=1000$ , $n=500$ ) and lines indicate theory ( $\bar{c}_{\lambda}$ ). B. Heatmaps of theoretically predicted $\bar{c}_{\lambda}$ values, shown for isotropic (left) and spiked (right) data.

Figure 4: Comparison between the asymptotic generalization curves over ridge strengths (

\lambda

) arising from KRR with a one-step kernel with

k_{\ell}=10

(blue) versus that of KRR with the baseline kernel

K_{x}

(red), shown across a range of

(\gamma,\rho)

pairs. Here,

\psi=\sigma=0.5

. Markers represent empirical simulations (with SEM;

D=1000

n=500

), and solid lines indicate theory. Stars indicate the minimum generalization error attained by either kernel.

Figure 5: Decomposition of generalization error into bias and variance, for baseline and one-step estimators, as a function of ridge (

\lambda

). Here,

\psi=\sigma=0.5

, and

k_{\ell}=10

for the feature learning estimators. Left: isotropic inputs (

\gamma=0

). Right: spiked, target-aligned inputs (

\gamma=10

\rho=0.6

). Empirical simulations shown as markers (

D=1000

n=500

) and theory shown as lines. Stars mark the minimum generalization error attained by either type of estimator.

Figure 6: A diagram of

\left.\frac{\partial^{2}\tilde{\mathcal{G}}}{\partial k_{\ell}\partial\lambda}\right|_{(0,0)}

as a function of the SNR. For various nonzero spike strengths, this mixed partial becomes negative for sufficiently large SNRs, implying that

\left.\frac{\partial}{\partial\lambda}\right|_{\lambda=0}\tilde{\mathcal{G}}\left(\hat{f}_{\lambda;k_{\ell}}\right)<\left.\frac{\partial}{\partial\lambda}\right|_{\lambda=0}\tilde{\mathcal{G}}\left(\hat{f}_{\lambda;0}\right)

for

k_{\ell}>0

(see Eq. 4.13 ). As

\gamma

grows large, the critical SNR at which a local advantage appears descends to

\text{SNR}_{\text{crit}}:=\frac{\psi^{2}}{(1-\psi)^{2}}

. Here,

\psi=0.5

and

\sigma=0.4

; as a result,

\max_{\rho}\rho^{2}/\sigma^{2}=6.25

and

\min_{\gamma>0}\frac{(1+\gamma\psi)^{3}}{\gamma(1+\gamma)(1-\psi)^{3}}=6\sqrt{3}\approx 10.39

, meaning that Eq. 4.14 automatically holds. Thus, in this setting, the sign of

\left.\frac{\partial^{2}\tilde{\mathcal{G}}}{\partial k_{\ell}\partial\lambda}\right|_{(0,0)}

indicates whether a local FLA holds. See Appendix A for more details regarding this mixed partial.

Figure 7: Comparison of asymptotic generalization error curves over

\lambda

for one-step estimators (

k_{\ell}=10

) and baseline estimators, shown for a setting where the inputs are isotropic (left) and a setting where the spike and target directions are orthogonal (right). Here,

\psi=\sigma=0.5

. Markers represent empirical simulations (with SEM;

D=3000

n=1500

) and solid lines indicate theory. Stars indicate the minimum generalization error attained by either kernel.

Deep Linear Networks

Deep linear networks study how compositions of trainable linear maps learn representations, optimize, and acquire depth-dependent biases.

See all 33 papers

May 18, 2026

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

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 second step of gradient descent with step-sizes

η_1\asymp N^{α_1}

and

η_2 \asymp N^{α_2}

for

α_1, α_2 \in [0,0.5)

, where

N

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

through

\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

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.

Behrad Moniri, Hamed Hassani

University of Pennsylvania, Philadelphia, PA, USA

Representation Learning Neural Network Training Dynamics Gradient Descent Dynamics

Ask PDF

May 18, 2026

How does feature learning reshape the function space?

Feature learning is widely regarded as the key mechanism distinguishing neural networks from fixed-kernel methods, yet its impact on the induced function space remains poorly understood. In this work, we precisely characterize how the function space spanned by the features of a two-layer neural network evolves during gradient descent training. We prove that, in the high-dimensional proportional regime, after a large gradient step the post-update feature distribution is well approximated by a target-dependent spiked Gaussian covariance. This induces a data-adaptive kernel that reshapes the function space and modifies its spectral structure. Our analysis reveals that feature learning can be interpreted as a distributional transformation in either parameter space or input space, equivalently as the introduction of a target-dependent kernel. In particular, it selectively amplifies eigenvalues aligned with the target direction and mixes leading eigenfunctions, coupling the top radial mode with a target-aligned quadratic harmonic. Overall, our results provide a precise function-space perspective on early-stage feature learning: rather than just rescaling a fixed kernel, gradient descent induces a data-adaptive deformation that preferentially enhances directions aligned with the signal in the data.

João Lobo, Bruno Loureiro, Long Tran-Than +1

Department of Computer Science, University of Warwick, United Kingdom · Departement d’Informatique, École Normale Supérieure, PSL & CNRS · School of Mathematical Sciences, Institute of Natural Sciences and MOE-LSC, Shanghai Jiao Tong University, China. Part of work was done at Department of Computer Science, and Centre for Discrete Mathematics and its Applications (DIMAP), University of Warwick, United Kingdom

Representation Learning Two-Layer Neural Networks Kernel Methods Neural Network Training Dynamics

Ask PDF

Sep 22, 2026

A Spectral Theory of Grokking: Weight Decay induces Feature Learning

In grokking an early fit to the training data separates from a much later improvement in generalization. During this delay, training can move from a fixed neural tangent kernel (NTK) regime to one in which task-relevant kernel eigendirections continue to evolve. We provide a quantitative theory for how this transition from lazy to rich learning can produce delayed generalization. For homogeneous networks trained with squared loss and

L_2

weight decay, we show that a finite residual remains after memorization, with larger residual fractions in target components associated with smaller NTK eigenvalues. These residuals feed back into the dynamics of the NTK itself, and projecting the resulting dynamics onto task-relevant spectral directions yields a reduced system in which residual-driven kernel growth competes with weight decay. This system predicts that the grokking timescale is controlled by the product of learning rate and weight decay, that feature learning slows logarithmically near a critical decay above which task-aligned NTK structure can no longer support generalization, and that stronger decay can prevent fitting altogether. We test these predictions in modular addition. In a homogeneous MLP, task-aligned Fourier structure continues to emerge in the NTK after training accuracy has saturated, and an 84

\times

90-grid of trained networks across varying learning rate and weight decay recovers the predicted phase geometry and inverse-product scaling of the generalization time with learning rate and weight decay. A one-block Transformer shows similar macroscopic phase structure in a 42

\times

45-grid, as well as the same transition-time scaling despite violating exact homogeneity. Together, these results provide a mechanistic derivation connecting post-fit feature learning to both the onset of generalization and its phase structure in the learning rate and weight decay plane.

Lenz Pracher, Pascal de Jong, Oskar Lieshaus +2

Arnold-Sommerfeld-Center for Theoretical Physics, Ludwig-Maximilians-Universität München, Munich, Germany · Department of Applied Physics, Standford University, Standford, CA, USA · Department of Mathematics, University of Cambridge, Cambridge, UK

Neural Network Generalization Weight Decay Modular Addition Neural Tangent Kernel Grokking

Ask PDF

Abstract

Deep Linear Networks

Explore similar work

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

How does feature learning reshape the function space?

A Spectral Theory of Grokking: Weight Decay induces Feature Learning