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

TL;DR

This paper analyzes how neural networks can efficiently learn multi-index models in high dimensions by leveraging mean-field Langevin dynamics, highlighting conditions for polynomial-time convergence and the role of low-dimensional structures.

Contribution

It introduces a framework for understanding the sample and computational complexities of neural networks learning multi-index models, and explores conditions for polynomial-time convergence of the mean-field Langevin algorithm.

Findings

Sample complexity scales almost linearly with effective dimension.

Neural networks adapt to latent low-dimensional structures.

Polynomial time convergence is possible on manifolds with positive Ricci curvature.

Abstract

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 $d_{\mathrm{eff}}$ 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, $d_{\mathrm{eff}}$ can be significantly smaller than the ambient dimension. We prove that the sample complexity grows almost linearly with $d_{\mathrm{eff}}$, 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 $d_{\mathrm{eff}}$ 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.