A Mean-Field Theory of Transformers: Well-Posedness of the Coupled Data--Parameter Dynamics and Global Convergence of Training

What the math gives to ML

The paper formulates transformer computation and training as two coupled mean-field evolutions: a token law transported through depth and a probability law over attention-head parameters evolved by Wasserstein gradient flow. The transferable asset is replacing a fixed finite collection of heads with a trainable distribution whose particle approximation naturally supports head diversity, variable-width inference, and principled entropy-induced exploration. A practical first test is a particle attention layer in which each head parameter is updated by the empirical risk gradient plus calibrated Langevin noise, while the layer output averages the particle heads. The available extraction does not expose the paper's detailed attention kernel or convergence constants, so this proposal is an MVP of the stated mean-field mechanism rather than a complete reproduction of every theorem.