# Ultra-fast feature learning for the training of two-layer neural networks in the two-timescale regime

arXiv: https://arxiv.org/abs/2504.18208 · HTML: https://mlsanity.com/paper/2504.18208/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 73 references that are in the map.
- First posted on arXiv 25 Apr 2025.

## Builds on (most cited first, 29 of 29)

- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Kernel Mean Embedding of Distributions: A Review and Beyond](https://mlsanity.com/paper/1605.09522/) (2016)
- [Breaking the Curse of Dimensionality with Convex Neural Networks](https://mlsanity.com/paper/1412.8690/) (2014)
- [Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit](https://mlsanity.com/paper/1902.06015/) (2019)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [On the linearity of large non-linear models: when and why the tangent kernel is constant](https://mlsanity.com/paper/2010.01092/) (2020)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence](https://mlsanity.com/paper/2006.02509/) (2020)
- [On Learning Gaussian Multi-index Models with Gradient Flow](https://mlsanity.com/paper/2310.19793/) (2023)
- [Mean-Field Langevin Dynamics: Exponential Convergence and Annealing](https://mlsanity.com/paper/2202.01009/) (2022)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [Mean-Field Langevin Dynamics and Energy Landscape of Neural Networks](https://mlsanity.com/paper/1905.07769/) (2019)
- [KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support](https://mlsanity.com/paper/2106.08929/) (2021)
- [Generative Sliced MMD Flows with Riesz Kernels](https://mlsanity.com/paper/2305.11463/) (2023)
- [Global convergence of neuron birth-death dynamics](https://mlsanity.com/paper/1902.01843/) (2019)
- [Leveraging the two timescale regime to demonstrate convergence of neural networks](https://mlsanity.com/paper/2304.09576/) (2023)
- [Gradient Descent on Infinitely Wide Neural Networks: Global Convergence and Generalization](https://mlsanity.com/paper/2110.08084/) (2021)
- [Wasserstein Steepest Descent Flows of Discrepancies with Riesz Kernels](https://mlsanity.com/paper/2211.01804/) (2022)
- [(De)-regularized Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/2409.14980/) (2024)
- [Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces](https://mlsanity.com/paper/2402.04613/) (2024)
- [Mean-field Analysis on Two-layer Neural Networks from a Kernel Perspective](https://mlsanity.com/paper/2403.14917/) (2024)
- [Wasserstein Gradient Flows of the Discrepancy with Distance Kernel on the Line](https://mlsanity.com/paper/2301.04441/) (2023)
- [Mean-Field Langevin Dynamics for Signed Measures via a Bilevel Approach](https://mlsanity.com/paper/2406.17054/) (2024)
- [A note on convergence of solutions of total variation regularized linear inverse problems](https://mlsanity.com/paper/1711.06495/) (2017)

## Built on by (most cited first, 6 of 6)

- [Closed-Form Last Layer Optimization](https://mlsanity.com/paper/2510.04606/) (2025)
- [Towards a Unified Analysis of Neural Networks in Nonparametric Instrumental Variable Regression: Optimization and Generalization](https://mlsanity.com/paper/2511.14710/) (2025)
- [From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD](https://mlsanity.com/paper/2510.21020/) (2025)
- [Rethinking Neural Network Learning Rates: A Stackelberg Perspective](https://mlsanity.com/paper/2605.15530/) (2026)
- [How are linear representations learned? Exact solutions to the dynamics of abstraction](https://mlsanity.com/paper/2607.08843/) (2026)
- [Two-Timescale Fine-tuning Provably Learns New Features for Two-Layer ReLU Networks](https://mlsanity.com/paper/2609.34667/) (2026)

## Abstract (arXiv)

We study the convergence of gradient methods for the training of mean-field single-hidden-layer neural networks with square loss. For this high-dimensional and non-convex optimization problem, most known convergence results are either qualitative or rely on a neural tangent kernel analysis where nonlinear representations of the data are fixed. Using that this problem belongs to the class of separable nonlinear least squares problems, we consider here a Variable Projection (VarPro) or two-timescale learning algorithm, thereby eliminating the linear variables and reducing the learning problem to the training of nonlinear features. In a teacher-student scenario, we show such a strategy enables provable convergence rates for the sampling of a teacher feature distribution. Precisely, in the limit where the regularization strength vanishes, we show that the dynamic of the feature distribution corresponds to a weighted ultra-fast diffusion equation. Recent results on the asymptotic behavior of such PDEs then give quantitative guarantees for the convergence of the learned feature distribution.