Uniform Scaling Limits in AdamW-Trained Transformers

Authors: William Gibson, Christoph Reisinger

arXiv · PDF

Summary

The authors prove that as transformer depth L grows, the layer-by-layer hidden-state dynamics and backpropagation variables (under AdamW training) converge uniformly to a system of ordinary differential equations—essentially a continuous-depth limit. The convergence rate is O(1/L + 1/(L^(1/3) H^(1/2))), where H is the number of attention heads. When attention heads don't use causal masking, the limiting ODE has a McKean–Vlasov (mean-field) form; concentration-of-measure techniques yield bounds uniform over compact sets of initial conditions, with constants independent of the number of tokens (avoiding a covering argument).

Main takeaways:

• Models transformer hidden states as an interacting particle system and proves L²-convergence to a forward–backward ODE system as depth L → ∞
• Convergence holds uniformly over initial conditions and improves with more attention heads H
• Without causal masking, the limit is a McKean–Vlasov ODE (mean-field interaction); with masking, a more general forward–backward system
• Bounds are independent of token count, thanks to concentration of measure (no covering argument needed)
• Under a suitable AdamW adaptation, bounds become independent of token embedding dimension as well

Relevance

Tangential to my midtraining and persona-installation work. Relevant only if I ever need to reason about deep-layer accumulation of persona signals or continuous-depth asymptotics of fine-tuning dynamics.

Abstract

arXiv:2605.11059v1 Announce Type: new Abstract: We study the large-depth limit of transformers trained with AdamW, by modelling the hidden-state dynamics as an interacting particle system (IPS) coupled through the attention mechanism. Under appropriate scaling of the attention heads, we prove that the joint dynamics of the hidden states and backpropagated variables converge in $L^2$, uniformly over the initial condition, to the solution of a forward--backward system of ODEs at rate $\mathcal O(L^{-1}+L^{-1/3}H^{-1/2})$. Here, $L$ and $H$ denote the depth and number of heads of the transformer, respectively. The limiting system of ODEs can be identified with a McKean--Vlasov ODE (MVODE) when the attention heads do not incorporate causal masking. By using the flow maps associated with this MVODE and applying concentration of measure techniques, we obtain bounds on the difference between the discrete and continuous models that are uniform over compact sets of initial conditions. As this is achieved without resorting to a covering argument, the constants in our bounds are independent of the number of tokens. Furthermore, under a suitable adaptation to AdamW, the bounds become independent of the token embedding dimension.