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 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.

• The paper presents a base–fiber decomposition to formally separate feature distribution and function estimation in mean-field neural networks.
• It rigorously analyzes long-time dynamics of the nonlinear Fokker–Planck equation, revealing multi-spike concentration and self-regularization phenomena.
• It establishes minimax-optimal estimation rates in single- and multi-index Gaussian models and connects geometry with adaptive feature alignment.

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

Overview

This paper presents a rigorous geometric-statistical framework for understanding feature learning in supervised regression, focusing on mean-field shallow neural networks trained by Mean-Field Langevin Dynamics (MFLD). The authors conceptualize feature learning via a base–fiber decomposition of learned representations, provide a precise geometric formulation, and analyze the statistical and dynamical mechanisms that enable adaptive feature learning within this setup. Key technical contributions include sharp analysis of the long-time dynamics of the nonlinear Fokker–Planck equation in the mean-field regime, detailed study of feature alignment and multi-spike concentration phenomena, and the establishment of minimax-optimal estimation rates in Gaussian single- and multi-index models.

Geometric Formulation of Feature Learning

The central conceptual advance is the decomposition of statistical feature learning into base (feature distribution) and fiber (function estimation given features) components. In the mean-field model class, predictors are parameterized by probability measures over network weights. The "base" corresponds to the distribution over hidden-layer weights that evolves during training; the "fiber" is the induced function space (a data-dependent RKHS) in which regression is performed using the learned features.

A precise feature-learning property is formalized: a method possesses this property if the learned feature space is both statistically aligned with the signal and effectively exploited by the estimator, in a problem-dependent, data-driven fashion. This is differentiated from fixed-feature methods (e.g., random features, kernel methods), where the representation is pre-specified and not adaptive to the underlying structure of the regression function.

Mean-Field Langevin Dynamics and Variational Characterization

MFLD is modeled as the Wasserstein gradient flow of a negative entropy-regularized empirical risk functional over the space of probability measures on parameter space. The long-time limit of these dynamics is characterized as a negative-entropy regularized empirical risk minimizer (RERM), yielding a Gibbs-type self-consistency condition for the stationary measure. The analysis focuses on the compact parameter space setting, for which convergence and rates to equilibrium are established.

A key finding is that self-regularization emerges: the output layer estimator after feature learning is itself the solution to a strongly convex RERM in the learned feature space with a random, data-dependent regularizer.

Feature Learning in Gaussian Index Models

Multi-Spike Concentration at Low Temperature

The main probabilistic result is that, in the low-temperature regime ($\lambda = o(1)$), the stationary hidden-layer marginal (i.e., the distribution over hidden weights) develops a multi-spike structure in Gaussian multi-index problems. Explicitly, this measure concentrates probability mass near the latent directions (indices) of the underlying regression function. This phenomenon is proven to exhibit a sharp phase transition in the temperature parameter, and occurs despite negative entropy regularization favoring spread-out measures—i.e., feature learning aligns the model with the problem structure even against regularization pressure.

For single-index models, the stationary measure exhibits a Lévy–Milman concentration phenomenon: for odd information exponents, concentration occurs on the sphere; for even exponents, mass splits evenly between the antipodal directions, reflecting the inherent symmetry.

Dimension Reduction and Feature Alignment

The geometric analysis rigorously connects the low-temperature concentration to dimension reduction: although the model class is infinite-dimensional, the learned solution—by virtue of multi-spike concentration—effectively restricts prediction to a low-dimensional space determined by the latent signal structure. Quantitatively, this enables the construction of a "top-$k$" signal approximation: the regression function is well-approximated in the learned feature space by a small number of leading directions aligned with the true signal.

The alignment property is proven: the estimator exploits these directions to yield excess risk bounds controlled by the tail energy outside the leading directions, as encoded in the learned feature space covariance.

Sharp Estimation Rates

Under Gaussian single- and multi-index regression with well-specified link functions, the analysis yields the following strong statistical guarantees, all up to logarithmic factors:

• Single-Index Models: Estimation error rate $\mathcal{O}(d/N)$.
• Multi-Index Models: Estimation error rate $\mathcal{O}(Md/N)$, where $M$ is the number of indices.

These match the minimax optimal rates for the respective (unknown) low-dimensional structure, despite the estimator and its feature space being entirely data-driven. Parameter recovery—the convergence of the empirical moments of the learned distribution to those of the ground truth—is also sharply quantified.

Furthermore, these results are robust under structural variants, including restricted isometry properties for the index set and different regularization scales. The convergence analysis shows that achieving such concentration can require very long run times (exponential in $N$), reflecting the rough landscape induced by nonconvexity and metastability in Wasserstein space.

Theoretical and Practical Implications

Theoretical Implications

• The presented base–fiber geometric framework allows systematic separation of feature learning from regression, making possible the precise study of alignment and adaptivity.
• The analysis demonstrates that Langevin-type statistical dynamics can provide implicit regularization that shapes the energy landscape to favor problem-adaptive representations, even in the absence of explicit sparsity constraints or prior knowledge.
• The theory connects feature learning to classical statistical concepts (such as minimax theory and moment methods) and to the dynamic formation of data-dependent RKHSs, potentially influencing model selection and regularization schemes in high-dimensional regimes.

Practical Implications and Future Directions

• The results suggest that mean-field neural networks trained with noise (as in MFLD or continuous-time Langevin dynamics) can adaptively exploit unknown low-dimensional or multi-index structure, without manual feature engineering.
• Understanding the transition, run-time, and metastability behavior of mean-field training remains an open avenue, particularly outside the well-specified Gaussian regime.
• Extending these insights to more general (e.g., deep, non-shallow) architectures, non-Gaussian inputs, and practical two-stage or multi-layered training procedures is a promising future direction.
• The geometric-statistical approach may also inform diagnostics and validation for learned features in deep learning pipelines, contributing to interpretability and model compression strategies.

Conclusion

This work provides a mathematically rigorous and conceptually clean foundation for understanding statistical feature learning in mean-field neural networks, rooted in a geometric base–fiber decomposition. The approach precisely characterizes how, why, and under what conditions data-driven, statistically optimal feature learning can occur—even in highly overparameterized, nonconvex, and high-dimensional regimes. The technical contributions unify landscape geometry, convergence dynamics, and statistical risk analysis, and open new directions for both theoretical study and practical algorithm design in adaptive representation learning.