揭示AdamW训练的Transformer在深度极限下的统一动态规律
Uniform Scaling Limits in AdamW-Trained Transformers
- 将Transformer隐藏状态建模为通过注意力耦合的粒子系统
- 证明隐藏态与反向传播变量以速率O(L⁻¹+L⁻¹⁄³H⁻¹⁄²)收敛到常微分方程组
- 首次实现对初始条件紧凑集上误差的全局一致控制，不依赖序列长度
我们研究了使用AdamW训练的Transformer在大深度极限下的行为，将隐藏状态动力学建模为通过注意力机制耦合的相互作用粒子系统（IPS）。在适当的注意力头缩放下，我们证明隐藏状态与反向传播变量的联合动态在$L^2$意义下，以速率$\ ext{O}(L^{-1}+L^{-1/3}H^{-1/2})$，在初始条件上一致收敛到前向-后向常微分方程组的解。其中，$L$和$H$分别表示Transformer的深度和注意力头数。当注意力头不包含因果掩码时，该极限系统可识别为麦凯恩-弗拉索夫常微分方程（MVODE）。通过利用该MVODE对应的流映射并结合测度集中技术，我们在紧凑初始条件集上得到了离散与连续模型之间的误差界，且无需覆盖论证，因此常数不依赖于词元数量。此外，在对AdamW进行适当调整后，这些界也独立于词元嵌入维度。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。