Uniform Scaling Limits in AdamW-Trained Transformers

TL;DR

This paper analyzes the large-depth limit of transformers trained with AdamW, modeling hidden-state dynamics as an interacting particle system and deriving convergence results to a system of ODEs.

Contribution

It introduces a novel approach to study the asymptotic behavior of AdamW-trained transformers using interacting particle systems and McKean--Vlasov ODEs.

Findings

Convergence of hidden states to a system of ODEs at rate O(L^{-1}+L^{-1/3}H^{-1/2})

Uniform bounds on the difference between discrete and continuous models independent of tokens

Bounds become independent of token embedding dimension with suitable AdamW adaptation

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.