The Geometry of Statistical Feature Learning in Mean-Field Langevin Dynamics
Abstract
We introduce a geometric formulation of statistical feature learning for supervised regression. Feature learning is defined through a base--fiber decomposition: the base is the feature-side geometry produced by training, and the fiber is the learned feature space where estimation is performed. We prove this property for spherical mean-field Langevin dynamics, viewed as the Wasserstein gradient flow of a negative entropy-regularized empirical risk. In Gaussian multi-index models, the low-temperature stationary distribution concentrates near the hidden indices, forms a multi-spike structure, and yields parameter recovery with high probability, even though negative entropy regularization penalizes concentration. This concentration has a sharp transition at temperature λasymp 1. In Gaussian single-index models, the stationary measure satisfies a concentration property, with parity determining whether it lives on S2d−1 or ℝℙd−1. The induced learned feature space aligns the regression signal and yields rates d/N and Md/N, up to logarithmic factors.
Related
- Statistical and Privacy Guarantees for Deep M-Regression via Stochastic Gradient Langevin Dynamics
- Langevin dynamics for high-dimensional optimization: the case of multi-spiked tensor PCA
- Uniform-in-N Log-Sobolev Inequality for the Mean-Field Langevin Dynamics with Convex Energy
- Reversible Stratonovich perturbations for accelerating convergence to equilibrium of Langevin dynamics
- Initialization Stability and Convergence of Projected Langevin Dynamics
BibTeX
@misc{shang2026geometry,
title = {{The Geometry of Statistical Feature Learning in Mean-Field Langevin Dynamics}},
author = {Zong Shang and Tomoya Wakayama and Guillaume Lecué and Taiji Suzuki},
howpublished = {arXiv},
year = {2026},
month = jul,
doi = {10.48550/arxiv.2606.31429},
eprint = {2606.31429},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2606.31429},
}
RIS
TY - UNPB
TI - The Geometry of Statistical Feature Learning in Mean-Field Langevin Dynamics
AU - Zong Shang
AU - Tomoya Wakayama
AU - Guillaume Lecué
AU - Taiji Suzuki
PY - 2026
DA - 2026/07/30
PB - arXiv
DO - 10.48550/arxiv.2606.31429
UR - https://arxiv.org/abs/2606.31429
AB - We introduce a geometric formulation of statistical feature learning for supervised regression. Feature learning is defined through a base--fiber decomposition: the base is the feature-side geometry produced by training, and the fiber is the learned feature space where estimation is performed. We prove this property for spherical mean-field Langevin dynamics, viewed as the Wasserstein gradient flow of a negative entropy-regularized empirical risk. In Gaussian multi-index models, the low-temperature stationary distribution concentrates near the hidden indices, forms a multi-spike structure, and yields parameter recovery with high probability, even though negative entropy regularization penalizes concentration. This concentration has a sharp transition at temperature λasymp 1. In Gaussian single-index models, the stationary measure satisfies a concentration property, with parity determining whether it lives on S_2^d−1 or ℝℙ^d−1. The induced learned feature space aligns the regression signal and yields rates d/N and Md/N, up to logarithmic factors.
ER -