Skip to content
Sagan

Paper

Uniform Scaling Limits in AdamW-Trained Transformers

Unreadunread

AI 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