📊 statistics

Uniform Scaling Limits in AdamW-Trained Transformers

This paper establishes that the joint dynamics of hidden states and backpropagated variables in AdamW-trained transformers converge uniformly to a forward-backward system of ODEs (specifically a McKean-Vlasov ODE in the absence of causal masking) as depth and the number of attention heads increase, providing dimension-independent error bounds without relying on covering arguments.

Original authors: William Gibson, Christoph Reisinger

Published 2026-05-13

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: William Gibson, Christoph Reisinger

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: From a Chaotic Crowd to a Smooth River

Imagine a Transformer (the type of AI behind modern chatbots) as a massive, multi-story building.

The Floors: The building has $L$ floors (layers).
The Workers: On each floor, there are $H$ teams of workers (attention heads) who look at the information coming up from below.
The Data: The information is a stream of "tokens" (words or image patches) moving up the building.
The Training: The building is being "trained" using an optimizer called AdamW. Think of AdamW as a very strict, smart foreman who adjusts the workers' tools to minimize mistakes, while also gently shrinking their tool sizes to keep them from getting too big and unstable (this is called weight decay).

The Problem:
When this building is huge (thousands of floors and millions of workers), it becomes impossible to track every single worker's movement. It's like trying to predict the exact path of every grain of sand in a massive sandstorm. Usually, mathematicians say, "If we want to understand the whole storm, we have to count every grain," which makes the math depend on how many grains (tokens) there are.

The Breakthrough:
This paper proves that you don't need to count every grain. Even if you have a billion tokens, the behavior of the whole system converges to a smooth, predictable flow that looks the same regardless of how many tokens you have.

The Core Analogy: The "Interacting Particle System"

The authors model the hidden states (the data moving through the network) as an Interacting Particle System (IPS).

The Old Way: Imagine a room full of people (tokens) trying to talk to each other. If you want to know what the group thinks, you have to listen to every single conversation. If the room gets bigger, the math gets harder.
The New Way: The authors show that if you have enough people, you can stop listening to individuals and just listen to the "average vibe" of the room.
- Instead of tracking Person A talking to Person B, you track how the "average person" interacts with the "average vibe."
- This turns a chaotic, discrete mess of individual updates into a smooth, continuous River of Data.

The "Uniform" Magic: Why It Doesn't Matter How Many Tokens You Have

The most surprising part of the paper is the word "Uniform."

In math, when you try to prove something works for all possible inputs, you usually have to worry about the "worst-case scenario." If you have 100 tokens, the math is one thing. If you have 1,000,000 tokens, the math usually gets much messier, and the error bounds (how far off your prediction might be) get worse.

The Paper's Claim: The authors prove that the error between the real, messy Transformer and their smooth, continuous "River" model does not get worse as you add more tokens.
The Metaphor: Imagine trying to predict the weather. Usually, if you add more weather stations (tokens), your prediction model gets more complex and harder to solve. This paper says, "Nope. Once you have enough stations, the weather pattern becomes a smooth, predictable curve that is just as easy to calculate whether you have 10 stations or 10 million."

They achieved this by avoiding a mathematical trick called a "covering argument" (which is like trying to map every single street in a city to understand traffic). Instead, they used a technique called concentration of measure, which is like realizing that in a huge crowd, the average behavior is so stable that the outliers don't matter.

The Role of AdamW: The "Decoupled" Foreman

The paper specifically looks at AdamW, the standard optimizer for training these models.

The Issue: In many math models, the "tools" (parameters) the workers use can grow infinitely large and wild, making the math explode.
The Fix: AdamW has a special feature called decoupled weight decay. It's like a foreman who says, "You can adjust your tools to fix mistakes, but I will also gently shrink your tools back to a safe size every day."
The Result: This keeps all the workers' tools within a fixed, safe "box" (a compact set). Because the tools never get too wild, the math stays stable, and the authors can prove that the "River" model works perfectly, even over long training sessions.

The "Flow Map" and the "Backward River"

The paper doesn't just look at data moving forward (Input $\to$ Output). It also looks at backpropagation (how the model learns from its mistakes).

Forward River: Data flows up the building.
Backward River: Gradients (the "lessons learned") flow down the building.
The System: The authors show that both the forward data and the backward lessons converge to a system of Ordinary Differential Equations (ODEs).
- Think of this as a pair of synchronized rivers flowing in opposite directions.
- They proved that the discrete steps of the real computer (jumping from floor to floor) are almost identical to the smooth flow of these mathematical rivers.

The Main Result (The "Theorem")

The paper provides a specific formula for how close the real Transformer is to their smooth mathematical model. The error depends on:

$L$ (Depth): How tall the building is.
$H$ (Heads): How many worker teams there are.

The error shrinks as the building gets taller and has more teams. Crucially, the number of tokens ( $N$ ) does not appear in the error formula.

In plain English:
If you build a Transformer with infinite depth and infinite width, trained with AdamW, its behavior becomes perfectly predictable and smooth. You can describe the entire system with a simple set of equations, and it doesn't matter if you are processing 10 words or 10 billion words; the rules of the game remain the same.

Summary of Contributions

Smoothness: They turned the chaotic, step-by-step training of a Transformer into a smooth, continuous flow (ODEs).
Token Independence: They proved this smoothness holds true no matter how many tokens you feed the model. This is a rare and powerful result because it removes the "curse of dimensionality" regarding the number of tokens.
AdamW Stability: They showed that the specific way AdamW shrinks weights keeps the system stable, allowing them to prove these results without the math blowing up.
Dimension Independence (Bonus): If they use a specific version of AdamW (Blockwise), the math also stops caring about the size of the word embeddings (the "vocabulary size" of the math), making the model even more scalable.

The Takeaway:
This paper gives us a mathematical "lens" that lets us see the forest instead of the trees. It tells us that as these AI models get bigger and bigger, they don't just get more complex; they actually become simpler and more predictable, governed by smooth laws that are independent of the sheer volume of data they process.

Technical Summary: Uniform Scaling Limits in AdamW-Trained Transformers

Problem Statement
The paper addresses the theoretical challenge of understanding the training dynamics of Large Language Models (LLMs), specifically Transformers, as they scale in depth ( $L$ ) and width (number of attention heads, $H$ ). While empirical scaling laws drive the industry to increase model size, a rigorous theoretical understanding of how design choices affect generalizability and robustness in the infinite-depth and infinite-width limit remains incomplete.

Existing literature on the convergence of gradient descent to global optimality often relies on limiting models (e.g., Neural ODEs or Mean-Field limits). However, transferring results from these limits to practical models requires convergence bounds that hold uniformly over the model's input. A significant gap exists in the literature regarding Transformers trained with the AdamW optimizer. Previous works often assume gradient descent or require parameters to be shared or dependent across layers, which does not reflect the standard practice of independent and identically distributed (i.i.d.) initialization. Furthermore, the cubic growth of gradients in softmax self-attention (due to $Q, K, O, V$ matrix products) complicates theoretical guarantees, often leading to bounds that grow exponentially with training steps. Finally, standard convergence bounds often depend on the number of tokens ( $N$ ) via covering arguments (e.g., Dudley's Entropy Integral Theorem), which limits their applicability to the uniform convergence over arbitrary input distributions.

Methodology
The authors model the hidden-state dynamics of an unmasked, attention-only Transformer trained with AdamW as an interacting particle system (IPS) coupled through the attention mechanism. The analysis proceeds through the following steps:

Mean-Field Formulation: The discrete dynamics of the $N$ tokens and $H$ heads are approximated by a continuous-time limit. The hidden states and backpropagated (adjoint) variables are shown to converge to the solution of a forward-backward system of Ordinary Differential Equations (ODEs).
McKean–Vlasov ODEs (MVODEs): When causal masking is absent, the limiting system is identified as a McKean–Vlasov ODE. The velocity field of this ODE is defined by the expectation over the parameter distribution.
Handling AdamW and Non-Linearity: A critical component of the methodology is the treatment of the AdamW optimizer. Unlike standard SGD, AdamW decouples weight decay. The authors leverage this property to prove that the support of the parameter measures remains within a compact set uniformly across all training steps. This allows for global Lipschitz constants for the attention mechanism and its derivatives, preventing the exponential growth of bounds typically associated with the cubic gradient growth in self-attention.
Uniform Convergence without Covering: To achieve bounds independent of the number of tokens ( $N$ $N$ ), the authors avoid traditional covering arguments. Instead, they utilize:
- Concentration of Measure: Applying techniques to bound the $L^1$ norm of martingale differences associated with the parameter initialization.
- Blockwise Taylor Approximation: Splitting the layers into blocks and approximating the flow maps within each block using Taylor expansions.
- Tensor Structure: Exploiting the specific algebraic structure of the attention mechanism (expressed as ratios of linear functions of measures) to show that the supremum over initial conditions can be reduced to a finite-dimensional space, thereby eliminating dependence on $N$ .
Blockwise AdamW: The authors introduce a variant of AdamW (Blockwise AdamW) that regularizes blocks of parameters separately. This adaptation allows the derived bounds to be independent of the token embedding dimension ( $d$ ) as well.

Key Contributions

Uniform Convergence Theorem: The main result (Theorem 7) establishes that the unmasked attention-only Transformer, with parameters initialized i.i.d. from a compactly supported distribution, converges in $L^2$ uniformly over initial conditions to an interacting particle system of mean ODEs. The convergence rate is $O(L^{-1} + L^{-1/3}H^{-1/2})$ .
Independence from Token Count: Unlike previous bounds that scale with $\sqrt{\log(N)}$ or depend on $N$ via covering arguments, the derived bounds are independent of the number of tokens $N$ . This implies the results hold for initial conditions not present in the training data.
AdamW Adaptation: The work extends large-depth convergence results to Transformers trained with AdamW. It demonstrates that the decoupled weight decay of AdamW ensures the parameters remain in a compact set, providing global-in-training Lipschitz continuity for the attention mechanism.
Dimension Independence: By utilizing Blockwise AdamW, the bounds are shown to be independent of the token embedding dimension $d$ .
Explicit Derivatives: The paper explicitly computes Lions' derivatives for the attention mechanism on the measure component, establishing local Lipschitz continuity with respect to the $p$ -Wasserstein distance for all $p \in [1, \infty]$ .

Results
The paper proves that under appropriate scaling of attention heads and with parameters initialized i.i.d., the joint dynamics of hidden states and backpropagated variables converge to the solution of a forward-backward system of ODEs.

Convergence Rate: The error between the discrete Transformer model and the continuous mean-field limit is bounded by $C(L^{-1} + L^{-1/3}H^{-1/2})$ , where $C$ is a constant independent of $N$ and, under Blockwise AdamW, independent of $d$ .
Uniformity: The convergence is uniform over compact sets of initial conditions.
Parameter Bounds: The authors show that AdamW updates confine the parameters to a compact set $S(R_\theta)$ uniformly over all training steps, a property crucial for establishing global Lipschitz constants.

Significance and Claims
The paper claims to provide a rigorous theoretical foundation for understanding the scaling behavior of Transformers trained with the dominant optimizer, AdamW. The significance lies in:

Bridging the Gap: It connects the discrete training dynamics of practical, large-scale Transformers to continuous mean-field limits, offering a tractable setting for analysis.
Robustness to Input Size: By establishing bounds independent of the number of tokens, the results suggest that the theoretical behavior of the model is intrinsic to its architecture and initialization, rather than artifacts of the specific dataset size. This is particularly relevant for generalization to unseen data.
Optimization Stability: The analysis highlights the stabilizing role of AdamW's decoupled weight decay in controlling parameter growth, which is essential for the theoretical guarantees of Lipschitz continuity in the presence of non-linear attention mechanisms.

The authors acknowledge a trade-off: the convergence rate $O(L^{-1} + L^{-1/3}H^{-1/2})$ is slightly slower than the $O(L^{-1} + (LH)^{-1/2})$ rate seen in some ResNet studies, a sacrifice made to achieve uniform bounds independent of $N$ . They note that extending these results to causal masking (standard in language modeling) and SDE scaling regimes remains an open area for future research.