The Geometry of Statistical Feature Learning in Mean-Field Langevin Dynamics

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 $\lambda\asymp 1$.

In Gaussian single-index models, the stationary measure satisfies a Lévy--Milman concentration property, with parity determining whether it lives on $S_2^{d-1}$ or $\mathbb{RP}^{d-1}$.

The induced learned feature space aligns the regression signal and yields rates $d/N$ and $Md/N$, up to logarithmic factors.