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

arXiv:2605.07772v2 Announce Type: replace Abstract: 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 analyses 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 identify a training-induced phase of the inference dynamics across layers: after clustering in the early layers, the token distribution can leave the clustered regime near the final layers. A single gradient-descent update already produces this escape, and minimizers of the training objective satisfy an exponential turnpike bound under which clustering persists through the interior layers. Our analysis is based on an entropy-regularized interaction energy that captures the clustering bias of attention, for which we establish a double-well landscape and a local Polyak-{\L}ojasiewicz inequality. More broadly, our results point toward a training-aware mean-field theory in which training and inference dynamics of Transformers are treated together.

Abstract: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 analyses 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 identify a training-induced phase of the inference dynamics across layers: after clustering in the early layers, the token distribution can leave the clustered regime near the final layers. A single gradient-descent update already produces this escape, and minimizers of the training objective satisfy an exponential turnpike bound under which clustering persists through the interior layers. Our analysis is based on an entropy-regularized interaction energy that captures the clustering bias of attention, for which we establish a double-well landscape and a local Polyak-Łojasiewicz inequality. More broadly, our results point toward a training-aware mean-field theory in which training and inference dynamics of Transformers are treated together.

54 pages, 6 figures. Accepted at NeurIPS 2026

Machine Learning (cs.LG); Analysis of PDEs (math.AP); Dynamical Systems (math.DS); Optimization and Control (math.OC)

arXiv:2605.07772 [cs.LG]

(or arXiv:2605.07772v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2605.07772 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.

aiqbench analysis

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

Related coverage