Mean-Field Theory

Follow this topic Ask about this topic

Momentum

1 paper in the last four weeks, with none the four weeks before. 0.0% of all new papers.

Jul 13Week of Sep 28

Latest papers 29

Oct 7, 2026cs.LG

Fluctuations of Nonlinear Observables in Mean Field Neural Network Training

Mean field limits describe the training dynamics of wide neural networks through the evolution of the empirical distribution of their parameters. Although functional central limit theorems characterize the asymptotic fluctuations of this distribution, quantities of practical interest are typically nonlinear observables of the parameter distribution rather than the distribution itself. In this work, we show how these mean field fluctuations propagate to finite dimensional nonlinear observables for shallow neural networks trained by stochastic gradient descent. Working in the weighted Sobolev space in which the limiting fluctuation process is constructed, we apply a functional Delta method under ordinary Fr{é}chet differentiability, without requiring Lions derivatives with respect to the measure variable. We obtain a central limit theorem for the observables and, under a suitable representation of their differentials, an explicit covariance formula inherited from the underlying mean field fluctuation theory. We also study whether prescribed quantities of interest can be recovered from the selected observations. Under a constant rank assumption, we prove that a quantity of interest factors locally through the observation functional if and only if, throughout a neighborhood, the kernel of the differential of the observation is contained in that of the quantity of interest. Thus, a differential condition expressed directly in the ambient Sobolev space yields an exact nonlinear local factorization. These results provide a framework both for quantifying finite-width uncertainty on observable, statistically or physically meaningful quantities and for assessing whether the chosen observations contain the information required to identify them.

Arnaud Descours, Geoffrey Lacour

UCBL · MaIAGE

Neural Network Training Dynamics Uncertainty Quantification Mean-Field Theory Shallow Neural Networks

Sep 14, 2026cs.LG

Attention Mean Fields Predict Average Representation Dynamics and Reveal Context-Specific Computation

A language model's representation geometry is not predetermined; it evolves as the model runs. A faithful account of that geometry must capture that dynamic process, and so cannot be based solely on model-independent statistics such as co-occurrence. Here we introduce a mean-field analysis of attention. The average attention from one token to another defines a kernel that carries representations layer to layer and can be iterated through the network to model how the geometry is transformed. We condition this average two ways. Conditioned on a whole corpus, the kernel predicts the average-case evolution of representation geometry. Conditioned instead on a single context, it predicts the expected geometry for that context. A head's departure from that prediction, its \emph{mean-field deviation}, isolates the context-specific computation that the mean field misses. Under the corpus-conditional reading, the kernel yields an open-loop model: from the input embeddings and the frozen weights alone, we can iterate the kernel and the model's own MLPs over token representations, never consulting a measured deviation at any layer. The resulting prediction is highly accurate. In early training the model and its corpus mean field are indistinguishable. Replace every attention head with its mean field, and the substitution leaves the loss on real text unchanged. Around the onset of induction, the two diverge, and the gap widens as representations become contextualized. Under the context-conditional reading, deviation from the mean field is a task-agnostic measure of context-specific computation. The residual decomposes additively into unusual attention routing and contextualization of the transported values. Across controlled induction and few-shot settings, greater deviation tracks greater reliance on in-context information.

Micah Adler, John W. Byers, Mark Crovella

Department of Computer Science, Boston University · Faculty of Computing & Data Sciences, Boston University

Attention Head Analysis In-Context Learning Mean-Field Theory Representation Geometry in Language Models

Aug 8, 2026math.PR

A cylindrical neural approximation theorem for conditional laws of McKean-Vlasov equations with common noise

We introduce conditional cylindrical neural networks for approximating functionals of conditional laws in McKean-Vlasov equations with common noise. Fourier moments of the initial law and truncated signatures of the time augmented common noise are mapped by a mixture density network to a Gaussian mixture approximation of the conditional law. A cylindrical neural network then evaluates the target functional through analytic integrals against this predicted measure. Rough path well posedness and stability provide a conditional law map that is continuous in the initial distribution and the rough driver and agrees almost surely with the classical conditional law at the Itô Brownian lift. Combining this continuity with Fourier separation, signature uniqueness, Wasserstein density of Gaussian mixtures, and neural universal approximation, we prove an

L^2

universal approximation theorem for continuous square integrable functionals. The numerical study implements the resulting two stage procedure on six examples, including non Gaussian initial laws, nonlinear drift, multiplicative common noise, and a two dimensional state. Independent particle references are used when no closed form law is available. The learned conditional law and functional approximations consistently improve on the empirical particle plug in, and additional experiments examine feature sensitivity, training from one terminal observation per common noise scenario, and Itô--Stratonovich consistency.

Nacira Agram, Reda Hmioui, Jan Rems

Department of Mathematics, KTH Royal Institute of Technology, 100 44 Stockholm, Sweden. · Digital Futures, Sweden. · Grenoble INP – Ensimag, Université Grenoble Alpes, Grenoble, France. +1

Mean-Field Theory Neural Network Approximation Theory

Jun 12, 2026math.OC

Operator Calculus for Population-Based Optimization: Modular Convergence and Finite-Population Guarantees

Population-based optimizers combine update rules such as mutation, selection, and recombination. When one rule changes, it is often unclear which convergence guarantees survive or how the new combination should be assessed. We develop an operator calculus: an operator is a population-update rule, and the calculus specifies how separately checked effects can be combined. Under explicit regularity and small-step conditions, the leading changes caused by the updates add, yielding reusable building blocks for convergence analysis. The framework distinguishes finding and retaining a good solution, reducing the population's mean objective, and concentrating candidates near an optimizer, and identifies the extra approximation conditions needed for finite evaluation-budget guarantees. Applications include distribution adaptation, recombinative evolution, and consensus dynamics, with verified nonconvex cases. Controlled experiments on a common nonconvex problem collection show how component effects change with population geometry and the performance measure: a rule can worsen the mean objective yet produce better candidates.

Pekka Malo, Lauri Viitasaari, Patrik Nummi +3

Aalto University, Department of Information and Service Management, Finland. · Aalto University, Department of Information and Communications Engineering, Finland. · Indian Institute of Management Ahmedabad, Operations and Decision Sciences, India. +1

Evolutionary Optimization Optimization Convergence Analysis Metaheuristic Optimization Nonconvex Optimization Mean-Field Theory

Jun 7, 2026cond-mat.soft

Discovering and decoding latent mean-field structure with variational autoencoders

Generative models are increasingly used to capture correlations in many-body systems, but the representations they learn remain largely opaque to physical interpretation. Here, we establish an intuitive criterion that quantifies the capacity of a variational autoencoder (VAE) to faithfully reconstruct the joint probability distribution of a many body system. In a nutshell, a bound on the VAE capacity is obtained by comparing the rate of the latent channel to the bipartite mutual information of the data. Using this bound, we show that the conditionally independent decoder of any successful VAE is structurally identical to a finite-size mean-field factorization. Hence, a successful reconstruction is direct evidence for a latent mean-field theory and the microscopic parameters of that theory can be read off the trained decoder. We validate these conclusions on a hierarchy of solvable models with scalar (Curie-Weiss), vector (Hopfield) and tensor (Maier-Saupe) order parameters, recovering the full Hopfield pattern matrix from equilibrium samples alone. We find that, when applied to Salamander retinal recordings, a two-latent VAE reproduces the population statistics with only two effective collective variables allowing us to recover the `stored patterns' of the neural population and write a generalized Hopfield model which correctly models the experimental data.

Marco Biroli, Max Welling, Vincenzo Vitelli

Department of Physics and the James Franck Institute, University of Chicago, Chicago, IL, USA · CuspAI, Cambridge, England · AMLab, University of Amsterdam, Amsterdam, Netherlands +1

Variational Autoencoders Mean-Field Theory

Jun 5, 2026cs.LG

Beyond Linear and Overcomplete Regimes: A Mean-Field Analysis of Bottleneck Autoencoders

Autoencoders (AEs) learn low-dimensional representations by mapping data into a latent space while minimizing reconstruction error. Despite their empirical success, theoretical understanding remains limited and largely restricted to linear models or settings without a bottleneck. In this work, we study nonlinear AEs with a fixed finite-dimensional bottleneck in the mean-field (MF) regime. We derive explicit MF learning dynamics for both encoder and decoder, providing a tractable characterization of training in the nonlinear setting. We show that, over finite time horizons, the empirical risk of finite-width networks trained with stochastic gradient descent closely tracks the MF risk trajectory with high probability. At optimality, we further establish that the finite-width risk converges to the MF optimum, demonstrating that finite networks are sufficiently expressive to approximate the infinite-width solution.

Santanu Das, Ramyak Bilas, Pascal Esser +1

STCS department Tata Institute Of Fundamental Research Mumbai, India-400005 · Department of Mathematics Indiana University Bloomington Indiana, USA-47405 · Department of Mathematics Ludwig-Maximilians-Universität München Germany, München-80799 +1

Autoencoders Mean-Field Theory Neural Network Approximation Theory Neural Network Training Dynamics

Jun 2, 2026math.OC

Nonlocal Mean Field Schrödinger Bridge with Learned Interactions

The Schrödinger Bridge Problem connects an initial distribution to a terminal one along a minimum-energy stochastic process. Its mean-field extension, the Mean-Field Schrödinger Bridge, governs interacting populations whose dynamics and costs depend on the collective distribution. When these interactions are nonlocal, their direct evaluation scales quadratically with the population size, making large ensembles intractable within FBSDE-based solvers. We replace these terms with neural surrogates in state and time, trained on empirical interaction values along sampled trajectories and embedded in a four-stage alternating scheme that updates the forward and backward potentials and the surrogates in turn, while preserving forward--backward consistency and the prescribed endpoint marginals. We derive Grönwall-type stability bounds quantifying how surrogate errors propagate to the generated trajectories under a small-gain condition. On crowd-navigation and high-dimensional opinion-dynamics benchmarks, the surrogates reproduce the trajectories obtained with exact evaluation at reduced training cost. The advantage is most significant when the interaction is a nonlinear functional of the measure, such as the normalized bounded-confidence drift, for which random-batch subsampling is biased and unstable whereas the learned surrogate remains accurate.

Daisuke Inoue, Dante Kalise, Mathieu Laurière

Department of Mathematics, Imperial College London, London SW7 2AZ, United Kingdom. · Shanghai Frontiers Science Center of Artificial Intelligence and Deep Learning, NYU-ECNU Institute of Mathematical Sciences, NYU Shanghai, Shanghai 200126, People’s Republic of China.

Neural Surrogate Modeling Schrödinger Bridges Mean-Field Theory

Jun 1, 2026cs.IR

The Crowded Embedding Space: A Mean-Field Mechanism for Emergent Marginalization in Retrieval-Augmented Agents

Retrieval-augmented generative agents rely on retrieval for grounding, yet are typically evaluated on a query-by-query basis. This isolates interactions that are geometrically coupled in a shared embedding space. For example, we show that the high document density required to serve majority interests (e.g., generic "Crime" movies) can geometrically overcrowd the retrieval neighborhood of a semantically similar minority (e.g., "Film Noir"), effectively expelling minority content from top-

k

results. We introduce a formal framework to analyze how such goal collisions in dense retrieval induce fundamental performance limits and emergent fairness issues inherent to spatial crowding. In our static analysis, we demonstrate that for a fixed embedding space, a phase transition occurs where minority user goals suffer a catastrophic collapse in performance as the density of majority goals increases. We then extend this to a dynamic model and derive a non-linear Fokker-Planck equation that governs the evolution of document embeddings as the agent updates them to maximize retrieval accuracy. Our analysis reveals that this local relevance objective triggers an emergent global mechanism that systematically marginalizes minority interests. We prove that such objectives drive the system to self-organize into a state that exclusively serves majority interests. These results provide a theoretical foundation for understanding a critical grounding failure mode in retrieval-augmented agents.

Shwan Ashrafi, Dan Roth

Oracle AI

Algorithmic Fairness Mean-Field Theory Agentic RAG Dense Retrieval

May 28, 2026cs.LG

Mean-Field Diffuser: Scaling Offline MARL to Thousands of Agents

Diffusion-based planning has achieved strong results in single-agent offline reinforcement learning, yet scaling to many-agent systems remains intractable due to the curse of dimensionality in the joint trajectory space. We introduce MF-Diffuser, a framework that lifts trajectory planning to the Wasserstein space of trajectory distributions, where the propagation of chaos ensures a small representative subset of agents captures the full population dynamics. Our approach features a value-weighted chaotic entropy objective that reconciles generative fidelity with return maximization, and a hierarchical coarse-to-fine strategy that progressively grows the agent population during denoising. We establish end-to-end suboptimality bounds with four interpretable terms, revealing that mean-field approximation error scales as

O(H^2/\sqrt{N})

while offline distribution shift provably does not grow with population size

N

, and prove the generated policy is an approximate mean-field Nash equilibrium with explicit convergence guarantees. Experiments on three mean-field RL benchmarks -- spanning stage games, sequential dynamics, and adversarial team competition -- show MF-Diffuser achieves the best return in the majority of settings, with the largest gains on suboptimal offline data and at extreme scales (

N \geq 10^3

).

Wenhao Li, Xiangfeng Wang, Bo Jin

Tongji University · East China Normal University

Multi-Agent System Optimization Mean Field Games Diffusion-Based Planning Mean-Field Theory Multi-Agent Reinforcement Learning Offline MARL

May 28, 2026math.OC

Kernel-based potential mean-field games with unbiased random Fourier $U$ -statistics

We study the subclass of potential mean-field games in which the running interaction cost and the terminal target cost are both expressed through reproducing-kernel maximum mean discrepancy (MMD) penalties, and develop a computational framework that exploits this kernel structure. Both costs are estimated from finite-sample empirical distributions using a random Fourier U-statistic representation that is unbiased and has linear cost in the batch size. The drift of the controlled diffusion is parametrized by a neural network and trained via stochastic gradient descent. For population near-minimizers we prove convergence to the terminal-constrained problem as the penalty diverges, and show that the same limit is recovered almost surely when the learned controls are evaluated on independent finite samples under explicit coupling conditions on the penalty, random-feature count and sample size. The framework includes the kernel-MMD-penalty Schr{ö}dinger bridge problem as the special case of a vanishing interaction cost. Numerical experiments illustrate the method on the Schr{ö}dinger bridge problem in dimensions up to one hundred, and on an electric vehicle charging coordination problem with per-vehicle physical heterogeneity, where an aggregate-demand congestion cost represents price-feedback competition at the population level and the terminal MMD penalty shapes the state-of-charge distribution at the deadline.

Yumiharu Nakano

Department of Mathematical and Computing Science, Institute of Science Tokyo

Multi-Agent Coordination Mean Field Games Schrödinger Bridges Mean-Field Theory

May 27, 2026cs.MA

The incremental voter model: mean-field analysis and convergence to equilibrium

We introduce the incremental voter model (IVM), a discrete-opinion multi-agent system where agents undergo step-wise transitions biased by the opinion of a randomly selected persuader. Our incremental voter model comprises a large population of interacting agents, each holding an opinion represented by an element of the discrete set

\{-k,\ldots,0,\ldots,k\}, k \in \mathbb{N}_{+}

. At each update step as time progresses, a pair of distinct agents are selected independently and uniformly at random from the population, and the first agent (viewed as the listener'') updates its opinion based on that of the second (viewed as the persuader''), adopting a new opinion that differs from its current one by at most one unit. By deriving the mean-field system of nonlinear ordinary differential equations (ODEs) that governs the large-population limit of the agent-based model, we develop a rigorous mathematical framework to study the asymptotic behavior of the opinion distribution in the mean-field limit. These results contribute to a deeper understanding of social influence processes in complex systems, particularly in modeling opinion polarization, and may guide the formulation of more advanced models in future research.

Fei Cao, Xiaoqian Gong

Amherst College - Department of Mathematics, Amherst, MA 01002, USA

Opinion Dynamics Computational Social Science Mean-Field Theory Agent-Based Modeling

May 27, 2026cs.LG

Thinned Mean Field Langevin Dynamics

Several important learning tasks can be formulated as minimizing an entropy-regularized objective over an appropriate space of probability distributions. Mean-field Langevin dynamics (MFLD) facilitate computation in this general context, casting the minimizer as the invariant distribution of a McKean--Vlasov process, which can be numerically discretized using

N

particles and thus simulated. However, simulating this interacting particle system has computational complexity of order

N^2

. Motivated by recent research into \emph{kernel thinning}, we propose \texttt{KT-MFLD}, in which each particle interacts only with a thinned particle coreset of size

\mathcal{O}(N^{\frac{1}{2}})

. \texttt{KT-MFLD} thus reduces the computational complexity to order

N^{\frac{3}{2}}

while, under mild regularity conditions, achieving the same convergence guarantees (up to logarithmic factors) as MFLD. Our theoretical analysis is empirically confirmed on tasks including the training of student-teacher neural networks, quantization with maximum mean discrepancy, and computation of predictively-oriented posteriors in a post-Bayesian framework.

Zonghao Chen, Heishiro Kanagawa, François-Xavier Briol +2

University College London · 2Newcastle University · 3The Alan Turing Institute +1

Langevin Dynamics Mean-Field Theory

May 23, 2026cs.LG

Feature Learning in Wide Neural Networks under $μ$ P: Identifiability and Sparse-Dictionary Decomposition of the Mean-Field Limit

We establish four structural results for feature learning in wide two-layer neural networks under the Maximal Update Parametrization (

μ

P). First, we prove global existence and uniqueness of the mean-field limit of noisy gradient descent under

μ

P, identifying the maximal admissible weight

w^*

on the moment sequence of the initialization as the reciprocal parameter-moment-growth boundary, and hence the largest weighted moment class propagated by the flow. The finite-particle approximation has uniform-in-time squared-Wasserstein rate

O(N^{-1})

. Second, we characterize identifiability of the mean-field limit: two admissible parameter measures induce the same network function in

L^2

exactly when their active components agree modulo the finite-rank realization symmetry of the architecture. The orbit depth

D^*_{\mathrm{orb}}

is separated from the moment-variety depth

D^*_{\mathrm{var}}

. Third, under the Barron-Hermite target condition the active support of the long-time limit measure admits a sparse-dictionary decomposition: it is supported on at most

S^*

atoms modulo finite-rank realization symmetry, with

S^*

bounded by an explicit coefficient-threshold number. Fourth, we derive the total feature-learning-error decomposition into statistical, optimization, propagation-of-chaos, and sparse-residual components, with a target-dependent Hermite/Barron tail replacing any initialization-only residual. The four results are tied together by an architectural identity: the triple

(w^*, D^*_{\mathrm{orb}}, S^*)

-- the maximal admissible weight, the orbit identifiability depth, and the sparse-dictionary depth at which the target is realizable -- is the natural learning cell of the architecture-data pair

(σ, ρ)

. The proofs are self-contained except for standard results from

μ

P and mean-field Langevin theory.

Akmal Xodarev

Independent Researcher · Tashkent, Uzbekistan

Representation Learning Parameter Identifiability Neural Network Optimization Mean-Field Theory

May 22, 2026stat.ML

Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer

We develop a gradient flow on the space of probability measures defined on matrix-valued parameters induced by regularized Muon, an analytically smoothed version of the idealized Muon optimizer. The key observation is that the regularized orthogonalization map is the gradient of a smooth Fenchel-dual smoothing of the nuclear norm. This identifies the (regularized) Muon update as a mirror/prox step in the update variable, with momentum acting as the dual coordinate. We use this structure to lift Muon from a single matrix parameter to finite-particle probability objectives of the form

J(ρ)=R\left(\int F d ρ\right)

, a setting motivated by mean-field descriptions of neural-network training, and derive the inertial continuous-time limit. Using this structure, we derive the finite-particle continuous-time limit under the inertial scaling of step size and momentum, and then pass to a phase-space mean-field equation over probability laws on parameter-momentum pairs. The resulting flow can be shown to be a damped Hamiltonian probability dynamics whose kinetic energy is induced by the regularized Muon mirror potential. We prove an exact Hamiltonian dissipation identity, showing that the Hamiltonian energy decreases monotonically. While the target objective itself need not be monotone along the inertial Muon dynamics, under additional gradient-dominance, bounded-momentum, and curvature/alignment assumptions, we obtain continuous and discrete-time exponential convergence rates for the objective gap. We also study the well-posedness of the mean-field limit equation and establish propagation of chaos guarantees for the interacting particle system. Finally, we extend the formulation to Hilbert-valued feature maps on product matrix spaces, yielding a blockwise Muon probability flow applicable to smooth transformer mixture-of-experts models.

Aratrika Mustafi, Soumya Mukherjee, Bharath K. Sriperumbudur

Department of Statistics, Pennsylvania State University

Momentum Methods Muon Optimizer Mean-Field Theory Neural Network Training Dynamics

May 21, 2026stat.ML

Uniform-in-Time Weak Propagation-of-Chaos in Shallow Neural Networks

We consider one-hidden layer neural networks trained in the feature-learning regime using gradient descent, and relate the output of the finite-width network

f_{\hatρ_t^m}

to its infinite-width counterpart

f_{ρ_t^{MF}}

, which evolves in the mean-field dynamics. While constant-time horizon bounds for

\|f_{ρ_t^{MF}} - f_{\hatρ_t^m}\|

may be obtained via standard Grönwall estimates, the long-time behavior of the fluctuation is a more delicate matter. Uniform-in-time bounds often rely on (local) strong convexity in the landscape or Logarithmic Sobolev inequalities present in noisy gradient dynamics. In this work, we establish non-asymptotic weak propagation-of-chaos that holds uniformly in time, obtained by exploiting instead the convergence rate of the mean-field deterministic Wasserstein-gradient-flow dynamics. Specifically, denoting by

L_t

the mean-field excess MSE loss at time

t

and

m

the number of neurons, under standard regularity assumptions and the condition

\int_0^\infty L_t^{1/2} dt =O(\log d)

, we obtain the uniform in time bound

\|f_{ρ_t^{MF}}- f_{\hatρ_t^m}\|^2 \lesssim \text{poly}(d) m^{-\min(1,c/6)}

whenever

L_t \lesssim t^{-c}

. Our result holds in a noiseless setting and does not make any assumptions on the geometry of the landscape near the optimum, and extends seamlessly to other forms of discretization, including finite number of samples and time discretization. A key takeaway of our result is that whenever the convergence rate of the mean-field, population-loss dynamics is faster than

t^{-2}

, we can attain a loss of

ε

with only

\text{poly}(d/ε)

neurons, training samples, and GD steps.

Margalit Glasgow, Joan Bruna

Massachusetts Institute of Technology · Courant Institute School of Mathematics, Computing and Data Science, New York University

Shallow Neural Networks Mean-Field Theory Neural Network Training Dynamics Feedforward Neural Networks Gradient Descent Dynamics

May 20, 2026cs.LG

Dropout Universality: Scaling Laws and Optimal Scheduling at the Edge-of-Chaos

We ask whether the standard treatment of dropout as a static hyperparameter is optimal, or whether its utility can be improved by letting it vary over depth. We answer this by developing a mean-field theory of dropout near the edge of chaos, identifying distinct universality classes for smooth and kinked activations, together with their scaling exponents. The resulting propagation theory, constrained by maximizing the regularization delivered by dropout, motivates concentrating dropout near the input. Experiments on vision, speech and financial time series show gains most consistently in MLPs, with smaller gains in Transformers.

Lucas Fernandez Sarmiento

Carnegie Mellon University (CMU), Pittsburgh, Pennsylvania, USA.

Mean-Field Theory

May 11, 2026math.AP

Quantifying Concentration Phenomena of Mean-Field Transformers in the Low-Temperature Regime

Transformers with self-attention modules as their core components have become an integral architecture in modern large language and foundation models. In this paper, we study the evolution of tokens in deep encoder-only transformers at inference time which is described in the large-token limit by a mean-field continuity equation. Leveraging ideas from the convergence analysis of interacting multi-particle systems, with particles corresponding to tokens, we prove that the token distribution rapidly concentrates onto the push-forward of the initial distribution under a projection map induced by the key, query, and value matrices, and remains metastable for moderate times. Specifically, we show that the Wasserstein distance of the two distributions scales like

\sqrt{{\log(β+1)}/β}\exp(Ct)+\exp(-ct)

in terms of the temperature parameter

β^{-1}\to 0

and inference time

t\geq 0

. For the proof, we establish Lyapunov-type estimates for the zero-temperature equation, identify its limit as

t\to\infty

, and employ a stability estimate in Wasserstein space together with a quantitative Laplace principle to couple the two equations. Our result implies that for time scales of order

\logβ

the token distribution concentrates at the identified limiting distribution. Numerical experiments confirm this and, beyond that, complement our theory by showing that for finite

β

and large

t

the dynamics enter a different terminal phase, dominated by the spectrum of the value matrix.

Albert Alcalde, Leon Bungert, Konstantin Riedl +1

Department of Mathematics Friedrich-Alexander-University Erlangen-Nürnberg · Institute of Mathematics CAIDAS University of Würzburg · Mathematical Institute University of Oxford Reuben College +1

Transformer Inference Transformer Transformer Attention Mean-Field Theory

May 10, 2026stat.ML

Quantitative Local Convergence of Mean-Field Stein Variational Gradient Flow

Stein Variational Gradient Descent (SVGD) is a deterministic interacting-particle method for sampling from a target probability measure given access to its score function. In the mean-field and continuous-time limit, it is known that the flow converges weakly toward the target, but no quantitative rate is known for the last iterate. In this paper, we establish quantitative local convergence in strong norms for this dynamics, when the interaction kernel is of Riesz type on the

d

-dimensional torus. Specifically, assuming that the initial density and the target are smooth and close in

L^2

-norm, we obtain explicit polynomial convergence rates in

L^2

-norm that depend on the dimension and on the regularity parameters of the kernel, the initialization and the target. We further show that these rates are sharp in certain regimes, and support the theory with numerical experiments. In the edge case of kernels with a Coulomb singularity, we recover the global exponential convergence result established in prior work. Our analysis is inspired by recent results on Wasserstein gradient flows of kernel mean discrepancies.

Lénaïc Chizat, Maria Colombo, Roberto Colombo +1

Dynamical Systems Optimization Convergence Analysis Mean-Field Theory Variational Inference Stein Variational Gradient Descent

May 9, 2026math.AP

Kinetic theory for Transformers and the lost-in-the-middle phenomenon

We study causal self-attention dynamics -- a toy model for decoder Transformers -- which we interpret as a non-exchangeable interacting particle system. Adapting cumulant expansions to the triangular causal dependency structure of the model, and appealing to non-hierarchical methods to estimate correlations using Glauber calculus, we prove a quantitative mean-field limit result and a next-order characterization of correlations. For iid uniformly distributed tokens, the limiting correlation equation can be solved in closed form and we obtain a rigorous explanation of the empirically observed \emph{lost-in-the-middle} phenomenon: the token retrieval profile, as a function of the source position in the prompt, is

\mathsf{U}

-shaped, with primacy, recency, and a unique interior minimum under an explicit smallness condition.

Mitia Duerinckx, Borjan Geshkovski, Stefano Rossi

Self-Attention Transformer Attention Mean-Field Theory

May 8, 2026stat.ML

Learnability and Competition in High-Dimensional Multi-Component ICA

Independent Component Analysis (ICA) is a foundational tool for unsupervised representation learning, yet its high-dimensional theory remains largely limited to single-component recovery. We develop an asymptotically exact mean-field theory for multi-component online ICA, capturing the coupling induced by simultaneous learning and orthogonalization. In the high-dimensional limit, the joint empirical distribution of learned estimates and ground-truth components converges to a deterministic process, yielding a closed ODE system for the overlap matrix between learned directions and true components. This characterization reveals a genuinely multi-component, initialization-driven phase structure: a decoupled regime, where estimates align with distinct components and evolve nearly independently, and a competition regime, where overlapping initializations induce orthogonality-driven conflicts, slow reorientation, and delayed convergence. Our steady-state analysis gives explicit learnability boundaries and competition conditions linking step size, data moments, and initialization. These conditions show that larger higher-order moments and competition shrink the stable learning-rate window, increase convergence times, and predict a staircase phenomenon in which the number of recoverable components changes discretely with the learning rate. Experiments on synthetic data and hyperspectral remote sensing data validate the predicted trajectories and phase behavior.

Eser Ilke Genc, Samet Demir, Zafer Dogan

MLIP Research Group, KUIS AI Center, Koç University · Department of EEE, Koç University

Independent Component Analysis Unsupervised Representation Learning Mean-Field Theory Statistical Learning Theory Online Learning

May 8, 2026cs.LG

Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers

Transformers perform inference by iteratively transforming token representations across layers. This layerwise computation has been studied empirically, and recent mean-field theories of Transformer dynamics explain how attention can drive token distributions toward clustering. However, existing mean-field analyses largely treat model parameters as prescribed, leaving open how training reshapes this clustering picture. We study this question in a noisy mean-field Transformer in which only a parameter-linear FFN is trained under

L^2

regularization. We find and analyze a training-induced phase in the dynamics: after initially following attention-driven clustering, the token distribution can leave the clustered regime near the final layers. Our mathematical analysis is based on an entropy-regularized interaction energy that captures the clustering bias of attention. More broadly, our results point toward a training-aware mean-field theory of Transformer dynamics, in which training and inference dynamics are treated together.

Noboru Isobe, Daisuke Inoue, Masaaki Imaizumi

1RIKEN Center for Advanced Intelligence Project · 2Imperial College London · 3The University of Tokyo +1

Transformer Clustering Transformer Attention Mean-Field Theory Neural Network Training Dynamics

May 5, 2026cs.LG

Population-Aware Imitation Learning in Mean-field Games with Common Noise

Mean Field Games (MFGs) provide a powerful framework for modeling the collective behavior of large populations of interacting agents. In this paper, we address the problem of Imitation Learning (IL) in MFGs subject to common noise, where the population distribution evolves stochastically. This stochasticity compels agents to adopt population-aware policies to respond to aggregate shocks. We formulate two distinct learning objectives: recovering a Nash equilibrium and maximizing performance against an expert population. We investigate two imitation proxies: Behavioral Cloning (BC) and Adversarial (ADV) divergence. We then establish finite-sample error bounds showing that minimizing these proxies effectively controls both the policy's exploitability and its performance gap relative to the expert. Furthermore, we propose a numerical framework using generalized Fictitious Play and Deep Learning to compute expert population-aware policies. Through experiments on three environments we demonstrate that standard population-unaware policies fail to capture the equilibrium dynamics. Our results highlight that learning population-aware policies is crucial to avoid being misled by the randomness inherent in common noise.

Grégoire Lambrecht, Mathieu Laurière

Center for Data Science, New York University and NYU Shanghai, New York, United States of America · Shanghai Center for Data Science; NYU-ECNU Institute of Mathematical Sciences at NYU Shanghai; NYU Shanghai, Shanghai, People’s Republic of China.

Mean Field Games Adversarial Imitation Learning Imitation Learning Mean-Field Theory Behavior Cloning

Apr 30, 2026math.OC

Continuous-time q-learning for mean-field control with common noise, part-I: Theoretical foundations

This paper investigates the continuous-time counterpart of the Q-function for entropy-regularized mean-field control (MFC) with controlled common noise, coined as q-function by Jia and Zhou (2023) in the single agent's model. We first show that, under discretely sampled actions, the value function in the exploratory formulation converges to the one in the relaxed control formulation as the time grid refines. Leveraging the relaxed control formulation, we derive the exploratory Hamilton-Jacobi-Bellman (HJB) equation, in which the controlled common noise gives rise to an additional nonlinear functional of policy, rendering the policy iteration intricate. Under certain concavity condition, we establish the existence and uniqueness of the optimal one-step policy iteration via a first-order condition using the partial linear functional derivative with respect to policy. The policy improvement at each iteration is verified by relating to an entropy-regularized optimization problem over the space of policies. In the mean-field setting, we introduce the integrated q-function (Iq-function) defined on the state distribution and the policy, and it is shown that an optimal policy is identified as a two-layer fixed point to the argmax operator of the Iq-function. Finally, we provide the explicit characterization of an optimal policy as a Gaussian distribution in the general linear-quadratic (LQ) setting.

Zhenjie Ren, Xiaoli Wei, Xiang Yu +1

LaMME, Universit´e ´Evry Paris-Saclay, ´Evry, France · Department of Applied Mathematics, The Hong Kong Polytechnic University, Kowloon, Hong Kong · Department of Industrial Engineering and Operations Research, Columbia University, New York, USA

Stochastic Optimal Control Q-Learning Maximum Entropy RL Mean-Field Theory Optimal Control

Feb 4, 2026cs.LG

Gradient Flow Through Diagram Expansions: Learning Regimes and Explicit Solutions

We develop a general mathematical framework to analyze scaling regimes and derive explicit analytic solutions for gradient flow (GF) in large learning problems. Our key innovation is a formal power series expansion of the loss evolution, with coefficients encoded by diagrams akin to Feynman diagrams. We show that this expansion has a well-defined large-size limit that can be used to reveal different learning phases and, in some cases, to obtain explicit solutions of the nonlinear GF. We focus on learning Canonical Polyadic (CP) decompositions of high-order tensors, and show that this model has several distinct extreme lazy and rich GF regimes such as free evolution, NTK and under- and over-parameterized mean-field. We show that these regimes depend on the parameter scaling, tensor order, and symmetry of the model in a specific and subtle way. Moreover, we propose a general approach to summing the formal loss expansion by reducing it to a PDE; in a wide range of scenarios, it turns out to be first-order and solvable by the method of characteristics. We observe a very good agreement of our theoretical predictions with experimental results.

Dmitry Yarotsky, Eugene Golikov, Yaroslav Gusev

Applied AI Institute, Moscow, Russia · Steklov Mathematical Institute of Russian Academy of Sciences, Moscow, Russia.

Tensor Decomposition Mean-Field Theory Neural Tangent Kernel Gradient Descent Dynamics

Feb 2, 2026cs.LG

Local exponential stability of mean-field Langevin descent-ascent and associated particle system

We study the mean-field Langevin descent-ascent (MFL-DA), a coupled optimization dynamics on the space of probability measures for entropically regularized two-player zero-sum games, together with its associated interacting particle system. For general nonconvex-nonconcave payoffs, Wang and Chizat (COLT 2024) asked whether the original single-timescale MFL-DA converges to the mixed Nash equilibrium and, if so, at what rate. We prove a local affirmative answer in Wasserstein space: if the initial datum is sufficiently close to the mixed Nash equilibrium, then the mean-field dynamics converges to it exponentially fast at a quantitative rate. We further show that the finite-

N

particle system inherits this stability up to times exponential in

N

, with an

N

-independent exponential rate modulo a finite-particle error floor. Combined with the recent counterexample of Mourrat and Pillaud-Vivien for MFL-DA, which shows that global convergence cannot hold in general, our theorem completes the positive local counterpart of the Wang-Chizat question: the mixed Nash equilibrium has a robust basin of attraction, stable under both the mean-field flow and its finite-particle approximation.

Geuntaek Seo, Minseop Shin, Pierre Monmarché +1

Nash Equilibrium Lyapunov Stability Game Theory Wasserstein Gradient Flows Zero-Sum Games Mean-Field Theory+1

Nov 26, 2025cs.LG

Mean-Field Model for Two-Layer Neural Networks Trained with Consensus-Based Optimization

We study Consensus-Based Optimization (CBO) for two-layer neural network training. We compare the performance of CBO against Adam on two test cases and demonstrate how a hybrid approach, combining CBO with Adam, provides faster convergence than CBO. Additionally, in the context of multi-task learning, we recast CBO into a formulation that offers less memory overhead. The CBO method allows for a mean-field model formulation, which we couple with the mean-field model of the neural network. To this end, we first reformulate CBO within the optimal transport framework. As the number of particles tends to infinity, we lift the corresponding dynamics to the Wasserstein-over-Wasserstein space and show that the variance decreases monotonically. We confirm numerically that both mean-field models converge.

William De Deyn, Michael Herty, Giovanni Samaey

Institut für Geometrie und Praktische Mathematik, RWTH Aachen University, Germany · Department of Computer Science, KU Leuven, Belgium · Extraordinary Professor, Department of Mathematics and Applied Mathematics, University of Pretoria, South Africa

Neural Network Optimization Mean-Field Theory Neural Network Training Dynamics

Aug 29, 2024cs.LG

DeepSPoC: A Deep Learning Based Sequential Propagation of Chaos

Classical particle methods based on propagation of chaos (PoC) have been developed for solving mean-field stochastic differential equations and their associated nonlinear Fokker--Planck equations. However, direct PoC implementations are difficult to apply to high-dimensional problems because they require simulating and storing large numbers of interacting particles, often with high particle-particle interaction costs. Motivated by these limitations, we build on the recently proposed sequential propagation of chaos (SPoC) framework, which replaces the fully interacting particle system in PoC with a sequential interaction mechanism. Based on this structure, we present DeepSPoC, a neural particle method that embeds a neural density representation into the sequential particle dynamics. DeepSPoC simulates particles batch by batch, while the neural network represents the evolving empirical law and is substituted into the coefficients of the mean-field SDE, thereby replacing direct particle-particle interactions with particle-network interactions. In DeepSPoC, a recently developed normalizing flow model called KRnet is used to approximate the empirical measure of particles. Compared with direct particle implementations, DeepSPoC substantially reduces memory consumption and evaluates interaction terms more efficiently, thereby improving scalability for high-dimensional problems. We apply DeepSPoC to a wide range of mean-field equations and verify its effectiveness and computational advantages.

Kai Du, Yongle Xie, Tao Zhou +1

Shanghai Center for Mathematical Sciences, Fudan University, Shanghai, 200438, China. · Shanghai Center for Mathematical Sciences, Fudan University, Shanghai, 200433, China. · LSEC, Institute of Computational Mathematics and Scientific/Engineering Computing, AMSS, Chinese Academy of Sciences, Beijing, China.

Stochastic Differential Equations Neural PDE Solvers Mean-Field Theory

Oct 28, 2022cs.LG

A Functional-Space Mean-Field Theory of Partially-Trained Three-Layer Neural Networks

To understand the training dynamics of neural networks, prior studies have considered the mean-field limit of two-layer neural networks as the width tends to infinity, establishing theoretical guarantees for its convergence under gradient flow training as well as approximation and generalization capabilities. In this work, we study the infinite-width limit of a type of three-layer neural network where the first-layer weights are randomly sampled and untrained. To rigorously define the limiting model, we extend the mean-field theory by lifting the representation of neurons from Euclidean to functional spaces. This allows us to establish the mean-field training dynamics as a functional gradient flow with a time-varying kernel that remains positive-definite under suitable assumptions, thus proving a linear-rate convergence of its training loss. Furthermore, we define novel function spaces that contain the solutions obtained through the mean-field training dynamics and prove Rademacher complexity bounds for these spaces. Notably, our analysis applies to a range of scaling choices of the model, resulting in two distinct regimes of the mean-field limit that both exhibit feature learning through training.

Zhengdao Chen, Eric Vanden-Eijnden, Joan Bruna

Google Research Mountain View, CA 94043 · Courant Institute, New York University New York, NY 10003

Neural Network Generalization Neural Network Optimization Mean-Field Theory Neural Network Training Dynamics

Date pendingmath.ST

On Universality of Non-Separable Approximate Message Passing Algorithms

Mean-field characterizations of first-order iterative algorithms -- including Approximate Message Passing (AMP), stochastic and proximal gradient descent, and Langevin diffusions -- have enabled a precise understanding of learning dynamics in many statistical applications. For algorithms whose non-linearities have a coordinate-separable form, it is known that such characterizations enjoy a degree of universality with respect to the underlying data distribution. However, mean-field characterizations of non-separable algorithm dynamics have largely remained restricted to i.i.d. Gaussian or rotationally-invariant data. In this work, we initiate a study of universality for non-separable AMP algorithms. We identify a general condition for AMP with polynomial non-linearities, in terms of a Bounded Composition Property (BCP) for their representing tensors, to admit a state evolution that holds universally for matrices with non-Gaussian entries. We then formalize a condition of BCP-approximability for Lipschitz AMP algorithms to enjoy a similar universal guarantee. We demonstrate that many common classes of non-separable non-linearities are BCP-approximable, including local denoisers, spectral denoisers for generic signals, and compositions of separable functions with generic linear maps, implying the universality of state evolution for AMP algorithms employing these non-linearities.

Max Lovig, Tianhao Wang, Zhou Fan

Department of Statistics and Data Science, Yale University · Halicioglu Data Science Institute, University of California, San Diego

Mean-Field Theory