Research map / Architectures and efficiency / RNNs and normalization
Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces Alternating Gradient Flows (AGF), a two-step utility-maximization/cost-minimization approximation of small-initialization gradient flow in two-layer networks that predicts saddle-to-saddle loss drops across linear, transformer, and modular-addition settings.
Problem
Feature learning in two-layer networks trained from small initialization exhibits staircase-like saddle-to-saddle loss curves; there is no unified theory predicting the order, timing, and magnitude of the feature-learning drops.
Approach
Approximate small-init gradient flow by alternating between maximizing a utility over dormant neurons (projected gradient on the unit sphere, tracking accumulated utility to find the jump time) and minimizing the loss over active neurons; each iteration activates one neuron and yields a loss drop. The framework is proven to match gradient flow in diagonal linear networks and recovers known greedy
Key results
- In diagonal linear networks, AGF's predicted loss curves converge to gradient flow's as the initialization alpha goes to 0, with a proof that AGF reproduces the sequence of Algorithm 1 of Pesme and Flammarion.
- For modular addition with p=20, the network learns dominant Fourier frequencies in decreasing order of coefficient magnitude, and the numerical AGF output closely tracks gradient descent.
- For linear regression with an attention-only linear transformer, AGF reproduces the sequential principal component regression with predicted loss levels and jump-time lower bounds.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use AGF to predict gradient-flow dynamics as initialization scale decreases | Diagonal linear networks · Diagonal linear network training | trajectory agreement | Helped |
| Use AGF to predict feature-learning dynamics | Two-layer quadratic network · Modular addition, p = 20 | dynamics agreement | Helped |
| Train from small initialization with numerical AGF predictions | Two-layer ReLU network · Subset of CIFAR-10 | loss-curve agreement | Mixed |
| Increase hidden width (ablation) | Two-layer ReLU network with AGF · Subset of CIFAR-10 | loss-curve agreement | Hurt |
| Learn Fourier components sequentially, dominant frequencies first | Two-layer quadratic network · Modular addition, p = 20 | frequency acquisition order | Helped |
| Apply AGF to fully connected linear networks | Two-layer fully connected linear network · Gaussian inputs with power-law covariance | singular-value dynamics agreement | Helped |
| Apply AGF to an attention-only linear transformer | Attention-only linear transformer · Linear regression in context | dynamics agreement | Helped |
Surprising findings
- On natural data (a subset of CIFAR-10 with a ReLU network), loss curves are often not visibly stepwise even at tiny initialization, and AGF's correspondence weakens as width grows and cost-minimization phases overlap.
Limitations
- Convergence of AGF to gradient flow is proven only for diagonal linear networks; a general proof remains open and the conjecture may fail in general settings.
- Focuses on two-layer networks; extension to deeper and more realistic architectures is open.
- AGF's correspondence weakens when many neurons activate in close succession, as in wider ReLU networks on CIFAR-10.
Experimental setup
Datasets: synthetic linear regression data modular addition (p=20) with cosine-wave encoding subset of CIFAR-10
Benchmarks: modular addition (p=20) CIFAR-10 subset
Metrics: training loss singular value dynamics power spectrum over time jump times
Baselines: gradient flow gradient descent Greedy Low-Rank Learning Algorithm 1 of Pesme and Flammarion (path-following)
Builds on: greedy low-rank learning saddle-to-saddle dynamics of diagonal linear networks Mean-field / maximal update parameterization
Closest papers by meaning
- Asymptotics of feature learning in two-layer networks after one gradient-step (ICML 2024 · similarity 0.89)
- How does Gradient Descent Learn Features --- A Local Analysis for Regularized Two-Layer Neural Networks (NeurIPS 2024 · similarity 0.89)
- Three Mechanisms of Feature Learning in a Linear Network (ICLR 2025 · similarity 0.89)
- A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks (ICML 2024 · similarity 0.88)
- From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning (ICML 2025 · similarity 0.88)
- The Benefits of Reusing Batches for Gradient Descent in Two-Layer Networks: Breaking the Curse of Information and Leap Exponents (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.