📊 statistics

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

This paper characterizes feature learning in linear-width two-layer networks by demonstrating that a second step of gradient descent, particularly with reused batches, overcomes the rank-one limitations of single-step updates to learn multiple directions corresponding to target functions with higher information exponents through a sharp spectral analysis of weight updates.

Original authors: Behrad Moniri, Hamed Hassani

Published 2026-05-19

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Behrad Moniri, Hamed Hassani

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching a Robot to See

Imagine you are training a robot (a neural network) to recognize patterns in a massive library of books (data). The robot has a "brain" made of layers of neurons. The goal is for the robot to learn features—meaningful ways of looking at the data, like recognizing that a "cat" has pointy ears, rather than just memorizing the exact pixels of a specific cat photo.

This paper studies what happens when the robot takes two steps to learn, compared to just taking one step. The researchers found that taking that second step changes everything, but only if you do it in a specific way.

The Setting: The "Just-Right" Size

The researchers are looking at a specific type of robot brain:

Not too big, not too small: They call this the "linear-width" regime. Imagine the number of neurons in the robot's brain is roughly the same as the number of books in the library and the number of pages in each book. This is a realistic size for modern AI, unlike the "infinite" brains studied in older theories (which are too perfect) or the "tiny" brains (which are too simple).
The Task: The robot is trying to learn a complex rule (the "target function") that turns book pages into answers. This rule might be simple (like "count the words") or complex (like "detect sarcasm").

The Problem with One Step: The "Flashlight"

Previous research showed that if you let the robot take just one step of learning (one update to its brain weights), it acts like a flashlight.

It shines a single beam of light in one direction.
It can only learn one simple feature at a time.
The Catch: It can only learn features that are "linear" (straightforward). If the rule it's trying to learn is curved or complex (like a wave), the one-step flashlight misses it entirely. It's like trying to find a hidden treasure using a flashlight that only points North; if the treasure is East, you won't find it.

The Discovery: The "Second Step" Superpower

This paper asks: What happens if we let the robot take a second step?

The answer is surprising. The second step acts like a multi-beam searchlight.

Multiple Directions: Instead of just one beam, the robot's brain suddenly develops multiple "outliers" (special, strong directions).
The Magic Number: How many new directions does it learn? It depends on the size of the learning steps (called "step-sizes").
- Think of the step-size as how aggressively the robot adjusts its brain.
- If the robot takes small, careful steps, it learns a few new directions.
- If it takes larger, bolder steps, it learns many new directions.
- The paper provides a mathematical formula to predict exactly how many new "beams" of light will appear based on how big the steps are.

The Crucial Twist: Reusing vs. Fresh Data

The most important finding is about how the robot uses its data during these two steps. The researchers tested two scenarios:

1. The "Reused Batch" (The Same Old Book)

In this scenario, the robot looks at the same set of books for both Step 1 and Step 2.

The Result: The robot becomes a genius at finding complex, curved patterns. Even if the rule it's trying to learn is very difficult (has no simple "linear" part), the second step allows it to figure it out.
The Analogy: Imagine you are trying to solve a puzzle. In Step 1, you look at the puzzle pieces. In Step 2, you look at the exact same pieces again, but this time you use what you learned in Step 1 to see a hidden pattern in them. Because you are looking at the same pieces, the "noise" cancels out, and the complex pattern emerges.

2. The "Fresh Batch" (A New Book)

In this scenario, the robot looks at the first set of books for Step 1, and then grabs a completely new, different set of books for Step 2.

The Result: The robot fails to learn the complex patterns. It gets stuck with only the simple, linear features.
The Analogy: You look at the puzzle pieces, then you throw them away and pick up a brand new box of random pieces. The connection you made in Step 1 is broken. You can't build on your previous insight because the new data doesn't "speak the same language" as the old data. You are back to square one, only able to see simple shapes.

The "Spectral" Secret

The paper uses advanced math (Random Matrix Theory) to describe the robot's brain.

Before learning: The brain's weights look like a smooth, flat ocean (a "bulk" of random values).
After Step 1: A single "spike" or island rises out of the water (one learned direction).
After Step 2 (with reused data): Several new islands rise up! The number of islands depends on the step sizes. These islands represent the new, complex features the robot has learned.

Summary of Claims

One step is limited: It only learns simple, straight-line features.
Two steps are powerful: They allow the robot to learn complex, curved features and multiple directions at once.
Step size matters: The "aggressiveness" of the learning steps determines how many complex features are learned.
Data reuse is key: To learn complex features, the robot must use the same data for both steps. If you switch to new data for the second step, the robot loses its ability to learn complexity and reverts to only learning simple features.

The paper concludes that by understanding these "spectral transitions" (the islands rising from the ocean), we have a better mathematical map for how modern, over-parameterized AI systems actually learn features in their early stages.

Technical Summary: Feature Learning in Linear-Width Two-Layer Networks

Problem Statement

This paper investigates the mechanism of feature learning in two-layer fully connected neural networks within the linear-width regime, where the number of hidden neurons ( $N$ ), sample size ( $n$ ), and input dimension ( $d_X$ ) scale proportionally ( $N \asymp n \asymp d_X \to \infty$ ).

While recent literature has analyzed feature learning via a single step of gradient descent (e.g., [BES+22, MLHD24]), the authors identify fundamental limitations in this approach:

Rank-One Constraint: A single update is approximately rank-one, allowing the model to learn only a single direction of the target function.
Information Exponent Limitation: Single-step updates generally require the target function to have a non-zero linear component (information exponent $\tau = 1$ ). Functions with higher information exponents (e.g., pure cubic or higher-order interactions) are effectively treated as noise and remain unlearned.

The paper addresses the gap in understanding by characterizing the second step of gradient descent. It asks: What features are learned in the second update? How does the spectrum of the weight matrix evolve? Does the second step enable learning beyond a single direction or functions without linear components?

Methodology

The study employs a layer-wise training procedure on a two-layer network $f_{NN}(x) = a^\top \sigma(Wx)$ :

Initialization: Weights $W_0$ and $a_0$ are initialized with Gaussian entries.
Two-Step Update: The first-layer weights $W$ $W$ are updated twice using gradient descent on a correlation loss, while the second layer $a$ $a$ remains fixed at initialization ( $a_0$ $a_{0}$ ).
- $W_1 = W_0 - \eta_1 \nabla_W \mathcal{L}(W_0, a_0)$
- $W_2 = W_1 - \eta_2 \nabla_W \mathcal{L}(W_1, a_0)$
Step-Size Scaling: The step sizes scale with the network width as $\eta_1 \asymp N^{\alpha_1}$ and $\eta_2 \asymp N^{\alpha_2}$ , where $\alpha_1, \alpha_2 \in [0, 0.5)$ .
Batch Settings: The analysis distinguishes between two data usage scenarios:
- Reused Batch: The same dataset $(X, y)$ is used for both gradient computations.
- Fresh Batch: The dataset is partitioned into two disjoint sets $(X_1, y_1)$ and $(X_2, y_2)$ for the first and second steps, respectively.
Theoretical Framework: The authors utilize random matrix theory and Hermite polynomial expansions to derive sharp spectral characterizations of the updated weight matrices. They assume the activation function $\sigma$ and target functions $g_k$ are polynomials (or have rapidly decaying Hermite coefficients).

Key Contributions and Results

1. Spectral Characterization of $W_2$

The paper provides a precise characterization of the updated weight matrix $W_2$ . Unlike the first update ( $W_1$ ), which is a rank-one perturbation of the initialization, $W_2$ behaves as a spiked random matrix with multiple outliers.

The spectrum consists of a bulk (from initialization $W_0$ ) and a set of outliers corresponding to learned directions.
The number of these outliers, denoted by $\Lambda$ , is determined by the step-size scaling parameters:
$\Lambda(\alpha_1, \alpha_2) = \min\left(L-1, \left\lfloor \frac{\alpha_2}{1/2 - \alpha_1} \right\rfloor \right)$
where $L$ is the degree of the polynomial expansion.
This result demonstrates that the second update allows the model to capture multiple distinct directions simultaneously, breaking the rank-one limitation of the first step.

2. Alignment with Target Directions

The authors analyze the alignment between the learned directions $\{\hat{\beta}_{2;k}\}$ and the true target directions $\{\beta^*_{k}\}$ . The results differ significantly based on the batch setting:

Reused Batch Setting:
- The second update can learn directions corresponding to target functions with an information exponent greater than one (i.e., functions with no linear component, such as $g(z) = H_3(z)$ ).
- Specifically, if the step sizes satisfy $\Lambda(\alpha_1, \alpha_2) \geq 2$ , the model aligns with directions where the linear coefficient is zero.
- This confirms that data reuse enables the model to "leap" over the information exponent barrier.
Fresh Batch Setting:
- The second update remains restricted to learning directions with a non-zero linear component (information exponent $\tau = 1$ ).
- Even with a second step, if fresh data is used, the model cannot learn directions where the linear coefficient is zero. The learned directions $\hat{\beta}_{2;k}$ for $k \geq 1$ show no alignment with target directions having $\tau > 1$ .

3. Phase Transitions

The paper identifies phase transitions in the learning dynamics based on the step-size exponents $\alpha_1$ and $\alpha_2$ :

If $\alpha_1 + \alpha_2 < 1/2$ , the second update is effectively equivalent to a single step with step size $\eta_1 + \eta_2$ ; no new directions are learned ( $\Lambda = 0$ ).
As $\alpha_1 \to 0.5^-$ , the number of learnable directions proliferates, allowing the model to capture higher-order features.

Significance and Claims

The paper claims to establish a tractable mathematical framework for studying optimization and feature learning phenomenology in modern, overparameterized networks. Its primary significance lies in:

Beyond One-Step Analysis: It moves beyond the limited "rank-one" analysis of single-step updates to provide a full characterization of the spectral evolution during the early phases of training (specifically the second step).
Validation of Batch Reuse: It rigorously confirms that the benefits of batch reuse (previously observed in finite-width regimes) persist in the high-dimensional linear-width limit. Reusing data allows the network to learn complex, non-linear features (high information exponents) that are inaccessible when using fresh batches.
Scaling Laws: It provides explicit scaling laws for step sizes ( $\alpha_1, \alpha_2$ ) that dictate the capacity of the network to learn multi-index functions versus single-index functions.

The authors position this work as a bridge between the "random features" regime (where features are fixed) and the "mean-field" or "narrow" regimes, offering insights into how structured representations emerge in overparameterized systems through iterative optimization.