mlsanityData to 14 Oct 2026

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

arXiv 2504.18208 ↗First posted 25 Apr 2025cs.LG math.OCcited by 5

Not on the map. Why?

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.

Key facts. Cited by 5 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

29 of 73 references.

References · 29

  1. 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,621
  2. 2016Optimization Methods for Large-Scale Machine Learningcited by 1,150
  3. 2018A Convergence Theory for Deep Learning via Over-Parameterizationcited by 787
  4. 2019Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descentcited by 738
  5. 2018On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transportcited by 426
  6. 2016Kernel Mean Embedding of Distributions: A Review and Beyondcited by 364
  7. 2014Breaking the Curse of Dimensionality with Convex Neural Networkscited by 331
  8. 2019Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limitcited by 187
  9. 2015Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measurescited by 139
  10. 2020On the linearity of large non-linear models: when and why the tangent kernel is constantcited by 129
  11. 2019Maximum Mean Discrepancy Gradient Flowcited by 127
  12. 2020SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergencecited by 64
17 more
  1. 2022Mean-Field Langevin Dynamics: Exponential Convergence and Annealingcited by 61
  2. 2023Learning time-scales in two-layers neural networkscited by 57
  3. 2023On Learning Gaussian Multi-index Models with Gradient Flowcited by 57
  4. 2022Convex Analysis of the Mean Field Langevin Dynamicscited by 55
  5. 2019Mean-Field Langevin Dynamics and Energy Landscape of Neural Networkscited by 50
  6. 2021KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Supportcited by 37
  7. 2023Generative Sliced MMD Flows with Riesz Kernelscited by 37
  8. 2019Global convergence of neuron birth-death dynamicscited by 22
  9. 2023Leveraging the two timescale regime to demonstrate convergence of neural networkscited by 19
  10. 2021Gradient Descent on Infinitely Wide Neural Networks: Global Convergence and Generalizationcited by 16
  11. 2022Wasserstein Steepest Descent Flows of Discrepancies with Riesz Kernelscited by 16
  12. 2024(De)-regularized Maximum Mean Discrepancy Gradient Flowcited by 16
  13. 2024Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spacescited by 15
  14. 2024Mean-field Analysis on Two-layer Neural Networks from a Kernel Perspectivecited by 13
  15. 2023Wasserstein Gradient Flows of the Discrepancy with Distance Kernel on the Linecited by 6
  16. 2024Mean-Field Langevin Dynamics for Signed Measures via a Bilevel Approachcited by 6
  17. 2017A note on convergence of solutions of total variation regularized linear inverse problemscited by 1

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 2; 2026: 3.

5 papers.

  1. 2025Closed-Form Last Layer Optimizationcited by 5
  2. 2025From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGDcited by 2
  3. 2026Rethinking Neural Network Learning Rates: A Stackelberg Perspectivecited by 1
  4. 2026How are linear representations learned? Exact solutions to the dynamics of abstractioncited by 0
  5. 2026Two-Timescale Fine-tuning Provably Learns New Features for Two-Layer ReLU Networkscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026Dichotomy of Feature Learning and Unlearning: Fast-Slow Analysis on Neural Networks with Stochastic Gradient Descent
  2. 2021On Learnability via Gradient Method for Two-Layer ReLU Neural Networks in Teacher-Student Setting
  3. 2026Sharp description of local minima in the loss landscape of high-dimensional two-layer ReLU neural networks
  4. 2025Dynamical Decoupling of Generalization and Overfitting in Large Two-Layer Networks
  5. 2026High-Dimensional Analysis of Gradient Flow for Extensive-Width Quadratic Neural Networks
  6. 2024Nonlinear Stochastic Gradient Descent and Heavy-tailed Noise: A Unified Framework and High-probability Guarantees
  7. 2024A Survey on Statistical Theory of Deep Learning: Approximation, Training Dynamics, and Generative Models
  8. 2023On the Impact of Overparameterization on the Training of a Shallow Neural Network in High Dimensions