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

TL;DR

This paper analyzes feature learning in two-layer neural networks during the second gradient descent step, revealing how multiple learned directions emerge depending on learning rates and batch reuse, extending beyond one-step analysis.

Contribution

It provides a spectral characterization of features learned after two gradient steps in high-dimensional linear-width networks, highlighting the impact of batch reuse on learning directions.

Findings

Multiple learned directions correspond to outliers in the spectral distribution.

Batch reuse enables learning of directions with higher information exponent.

Number of learned directions depends on the scaling parameters of the gradient steps.

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, such updates are fundamentally limited: they are approximately rank-one, capturing only a single direction, and require 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 $\eta_1 \asymp N^{\alpha_1}$ and $\eta_2 \asymp N^{\alpha_2}$ for $\alpha_1, \alpha_2 \in [0,0.5)$. We derive a sharp 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 these outliers is determined by the scaling parameters $\alpha_1$ and $\alpha_2$ through $\lfloor \frac{\alpha_2}{1/2 - \alpha_1} \rfloor$. Furthermore, by analyzing the alignment between these learned directions and the target function, we identify a qualitative 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, under the condition that $\alpha_1, \alpha_2$ are chosen properly. This confirms that the benefits of batch reuse, previously observed in finite-width regimes, persist in the high-dimensional linear-width limit. By characterizing these early-phase spectral transitions, our work establishes a tractable mathematical framework for studying optimization and feature learning phenomenology in modern overparameterized networks.