# Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics

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

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 73 references that are in the map.
- First posted on arXiv 14 Aug 2024.

## Builds on (most cited first, 30 of 45)

- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Breaking the Curse of Dimensionality with Convex Neural Networks](https://mlsanity.com/paper/1412.8690/) (2014)
- [Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss](https://mlsanity.com/paper/2002.04486/) (2020)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [When Do Neural Networks Outperform Kernel Methods?](https://mlsanity.com/paper/2006.13409/) (2020)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks](https://mlsanity.com/paper/2202.08658/) (2022)
- [Regularization Matters: Generalization and Optimization of Neural Nets v.s. their Induced Kernel](https://mlsanity.com/paper/1810.05369/) (2018)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [Optimal Errors and Phase Transitions in High-Dimensional Generalized Linear Models](https://mlsanity.com/paper/1708.03395/) (2017)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [Limitations of Lazy Training of Two-layers Neural Networks](https://mlsanity.com/paper/1906.08899/) (2019)
- [Neural Networks Efficiently Learn Low-Dimensional Representations with SGD](https://mlsanity.com/paper/2209.14863/) (2022)
- [Mean-Field Langevin Dynamics: Exponential Convergence and Annealing](https://mlsanity.com/paper/2202.01009/) (2022)
- [Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models](https://mlsanity.com/paper/2305.10633/) (2023)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [On Learning Gaussian Multi-index Models with Gradient Flow](https://mlsanity.com/paper/2310.19793/) (2023)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [Stochastic Particle Gradient Descent for Infinite Ensembles](https://mlsanity.com/paper/1712.05438/) (2017)
- [The Benefits of Reusing Batches for Gradient Descent in Two-Layer Networks: Breaking the Curse of Information and Leap Exponents](https://mlsanity.com/paper/2402.03220/) (2024)
- [Mean-Field Langevin Dynamics and Energy Landscape of Neural Networks](https://mlsanity.com/paper/1905.07769/) (2019)
- [Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit](https://mlsanity.com/paper/2406.01581/) (2024)
- [Fundamental Limits of Weak Recovery with Applications to Phase Retrieval](https://mlsanity.com/paper/1708.05932/) (2017)
- [Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions](https://mlsanity.com/paper/2405.15459/) (2024)
- [Learning Polynomials of Few Relevant Dimensions](https://mlsanity.com/paper/2004.13748/) (2020)
- [Gradient-Based Feature Learning under Structured Data](https://mlsanity.com/paper/2309.03843/) (2023)

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

- [Survey on Algorithms for multi-index models](https://mlsanity.com/paper/2504.05426/) (2025)
- [Optimal Spectral Transitions in High-Dimensional Multi-Index Models](https://mlsanity.com/paper/2502.02545/) (2025)
- [Robust Feature Learning for Multi-Index Models in High Dimensions](https://mlsanity.com/paper/2410.16449/) (2024)
- [Propagation of Chaos for Mean-Field Langevin Dynamics and its Application to Model Ensemble](https://mlsanity.com/paper/2502.05784/) (2025)
- [Learning a Single Index Model from Anisotropic Data with vanilla Stochastic Gradient Descent](https://mlsanity.com/paper/2503.23642/) (2025)
- [When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective](https://mlsanity.com/paper/2503.11272/) (2025)
- [Mirror Mean-Field Langevin Dynamics](https://mlsanity.com/paper/2505.02621/) (2025)
- [Dynamic Bundling with Large Language Models for Zero-Shot Inference on Text-Attributed Graphs](https://mlsanity.com/paper/2505.17599/) (2025)
- [Uniform-in-Time Weak Propagation-of-Chaos in Shallow Neural Networks](https://mlsanity.com/paper/2605.22010/) (2026)
- [From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD](https://mlsanity.com/paper/2510.21020/) (2025)
- [Learning Hierarchical Polynomials of Multiple Nonlinear Features with Three-Layer Networks](https://mlsanity.com/paper/2411.17201/) (2024)
- [Feature learning is decoupled from generalization in high capacity neural networks](https://mlsanity.com/paper/2507.19680/) (2025)
- [Provably Extracting the Features from a General Superposition](https://mlsanity.com/paper/2512.15987/) (2025)
- [Born Discrete, Made Smooth: Variational Formulation of Shallow Neural Networks](https://mlsanity.com/paper/2607.02003/) (2026)
- [Learning Orthogonal Multi-Index Models Beyond Small Initialization: Incremental Learning, Competitive Dynamics and Symmetry](https://mlsanity.com/paper/2609.10879/) (2026)

## Abstract (arXiv)

We study the problem of learning multi-index models in high-dimensions using a two-layer neural network trained with the mean-field Langevin algorithm. Under mild distributional assumptions on the data, we characterize the effective dimension deff that controls both sample and computational complexity by utilizing the adaptivity of neural networks to latent low-dimensional structures. When the data exhibit such a structure, deff can be significantly smaller than the ambient dimension. We prove that the sample complexity grows almost linearly with deff, bypassing the limitations of the information and generative exponents that appeared in recent analyses of gradient-based feature learning. On the other hand, the computational complexity may inevitably grow exponentially with deff in the worst-case scenario. Motivated by improving computational complexity, we take the first steps towards polynomial time convergence of the mean-field Langevin algorithm by investigating a setting where the weights are constrained to be on a compact manifold with positive Ricci curvature, such as the hypersphere. There, we study assumptions under which polynomial time convergence is achievable, whereas similar assumptions in the Euclidean setting lead to exponential time complexity.