📊 statistics

Training Infinitely Deep and Wide Transformers

This paper establishes a rigorous mathematical framework for training infinitely deep and wide transformers in the mean-field regime by modeling their dynamics as a neural PDE, proving the well-posedness of forward and backward passes, and demonstrating that gradient flow converges to global minima under specific conditions of Neural Tangent Kernel injectivity.

Original authors: Raphaël Barboni, Maarten V. de Hoop, Takashi Furuya, Gabriel Peyré

Published 2026-05-19

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Raphaël Barboni, Maarten V. de Hoop, Takashi Furuya, Gabriel Peyré

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 Crowd to a Flowing River

Imagine a modern Transformer (the AI brain behind tools like chatbots) as a massive factory assembly line.

The Tokens: The data coming in (like words in a sentence or patches of an image) are the "workers" on the line.
The Layers: The factory has many floors (layers).
The Attention Mechanism: This is the rule that tells every worker how to look at every other worker to decide what to do next.

Usually, we study these factories with a fixed number of floors and a fixed number of workers. But this paper asks a "what if" question: What happens if the factory has infinite floors and infinite workers?

The authors built a mathematical framework to understand how these infinite factories learn. They found that while other deep learning models (like ResNets) behave like a single person walking down a path, Transformers behave like a river flowing through a landscape.

Key Concept 1: The "River" vs. The "Path"

The Old Way (ResNets):
Think of training a standard deep neural network like a hiker walking up a mountain. The hiker takes one step at a time. Mathematically, this is like an Ordinary Differential Equation (ODE). It's a single path where the hiker's position depends only on where they are right now.

The New Way (Transformers):
In a Transformer, every "worker" (token) is constantly looking at every other worker. If one worker changes, it ripples through the whole group.

The Analogy: Imagine a river. You can't just track one drop of water; you have to track the entire flow of the river. The water at one point depends on the water everywhere else upstream and downstream.
The Math: Because the "workers" are constantly influencing each other, the paper shows that training a Transformer is actually controlling a Partial Differential Equation (PDE). It's not just a path; it's a complex, shifting flow of probability distributions.

Key Concept 2: The "Mean-Field" Crowd

To make sense of infinite workers, the authors use a concept called Mean-Field.

The Analogy: Imagine a stadium full of 100,000 people. Instead of tracking every single person's name and location, you look at the crowd as a whole "fluid." You ask: "How dense is the crowd here? How fast is the crowd moving there?"
The Paper's Claim: They treat the "tokens" (data) not as individual items, but as a smooth distribution (a fluid) that evolves as it moves through the infinite layers of the Transformer. They also treat the "attention heads" (the rules the AI uses to pay attention) as a fluid distribution of parameters.

Key Concept 3: The Map and the Compass

The paper proves two very important things about how this "river" moves:

1. The Forward Pass (The Map):
They showed that if you start with a specific distribution of data, there is a unique, well-defined way it will flow through the infinite layers.

Analogy: If you pour a specific amount of blue dye into a river at the source, the math guarantees you can predict exactly how that dye will spread and move as it flows downstream, even if the river is infinitely long.

2. The Backward Pass (The Compass):
To train the AI, you need to know how to adjust the rules (the attention parameters) to make the output better. This is done by "backpropagation" (looking at the error and working backward).

Analogy: Imagine you are at the bottom of the river and you see a rock that shouldn't be there. You need to figure out which upstream current caused that rock to move there. The authors derived a precise "compass" (using something called Adjoint Sensitivity Analysis) that tells the system exactly how to nudge the infinite rules to fix the error.

Key Concept 4: The "No Dead Ends" Guarantee

The biggest worry in training AI is getting stuck in a local minimum.

The Analogy: Imagine you are trying to find the lowest point in a foggy valley. You might get stuck in a small dip (a local minimum) and think you've reached the bottom, when actually there is a much deeper valley nearby.
The Paper's Claim: The authors proved that for Transformers, if you start with a "good enough" initial setup (specifically, if the "Neural Tangent Kernel" is injective, which is a fancy way of saying the AI's internal map is diverse enough), there are no fake dead ends.
The Result: If the AI starts close to the solution, the "gradient flow" (the learning process) is guaranteed to slide all the way down to the global minimum (the absolute best solution) without getting stuck. It's like having a perfectly smooth, bowl-shaped valley where the bottom is the only place you can stop.

Key Concept 5: When Does This Work? (The "Distinctness" Rule)

The paper asks: "When is this 'no dead ends' guarantee true?"
They found a specific condition related to the data (the tokens).

The Analogy: Imagine the data points are different colored marbles. If all the marbles are identical, the AI gets confused and can't learn. But if the marbles are distinct enough (like different colors, shapes, or sizes), the AI can tell them apart.
The Math: They proved that as long as the data distributions are "linearly independent" (mathematically distinct, like distinct Gaussian mixtures or discrete points), the AI's learning map is perfect.
Real-world check: They showed that for almost any random set of data you might actually use (like random words or images), this condition holds true. You don't need to do anything special; nature usually provides data that is distinct enough.

Summary

This paper builds a rigorous mathematical bridge to understand infinite Transformers.

It changes the view of training from a "hiker on a path" to a "river flowing."
It proves that the flow is predictable and well-behaved.
It provides a "compass" to navigate the training process.
Most importantly, it proves that if the data is diverse enough (which it usually is), the AI will not get stuck in bad solutions; it will find the best possible solution.

The authors did this by extending theories used for simpler networks (ResNets) and adapting them to the complex, interconnected nature of the Attention mechanism found in Transformers.

Technical Summary: Training Infinitely Deep and Wide Transformers

Problem Statement
While Transformers have become the dominant architecture in modern machine learning, the theoretical understanding of their training dynamics remains limited, particularly for deep architectures with many layers and attention heads. Existing theoretical frameworks often rely on the Neural Tangent Kernel (NTK) regime, which assumes weights scale as $1/\sqrt{\text{width}}$ , leading to "lazy training" where the network remains close to initialization and feature learning is frozen. Alternatively, mean-field regimes for shallow networks or ResNets (modeled as Ordinary Differential Equations, ODEs) do not fully capture the complexity of Transformers, where the attention mechanism couples the evolution of multiple token distributions. There is a need for a rigorous mathematical framework to analyze gradient-based training in the joint limit of infinite depth (number of layers) and infinite width (number of attention heads), accounting for the simultaneous evolution of token distributions and parameter distributions.

Methodology
The paper develops a mean-field framework for Transformers by extending the theory of conditional optimal transport and Wasserstein gradient flows, previously applied to ResNets, to the Transformer setting.

Dual Mean-Field Representation: The model treats Transformers as infinitely deep (parameterized by a continuous depth variable $s \in [0,1]$ ) and infinitely wide (parameterized by a probability distribution $\rho$ over attention parameters). Crucially, the framework employs two types of measure representations:
- Token Distributions: The evolution of tokens through layers is modeled as a probability measure $\mu(s)$ over the token space.
- Attention Parameters: The parameters at each layer are represented as a measure $\rho(\cdot|s)$ over the parameter space $\Theta$ (query, key, value matrices).
Neural PDE Formulation: Unlike ResNets, which can be viewed as controlling a neural ODE on feature vectors, the training of Transformers in this framework corresponds to controlling a neural Partial Differential Equation (PDE). The token evolution is governed by a continuity equation where the velocity field is determined by the mean-field attention operator, which couples the distributions of all tokens.
$\partial_s \mu(s) + \text{div}(\mu(s) \Phi_{\rho(\cdot|s)}[\mu(s)]) = 0$
The authors reformulate this PDE perspective into a Banach-space-valued ODE acting on flow maps $\Lambda$ , which advect the initial token distribution. This reparameterization avoids the lack of differentiable structure in the space of measures and provides enhanced regularity for the analysis.
Training Dynamics via Conditional Wasserstein Gradient Flow: The training process is modeled as a gradient flow on the space of parameter distributions equipped with the Conditional Optimal Transport (COT) metric. This metric enforces that the marginal distribution of parameters across layers is preserved during training. The authors derive the conditional Wasserstein gradient of the training risk using adjoint sensitivity analysis, obtaining an explicit formula involving adjoint variables governed by backward ODEs.
Neural Tangent Kernel (NTK) Analysis: To establish convergence, the paper investigates the injectivity of the NTK for attention mechanisms. The authors derive necessary and sufficient conditions for the NTK to be positive definite, linking this property to the linear independence of log-sum-exp (cumulant generating) functions of the token distributions modulo affine functions.

Key Contributions and Results

Well-Posedness of Forward and Training Dynamics: The paper establishes the existence and uniqueness of solutions for the forward pass (token evolution) and the training evolution (parameter distribution flow). It proves that the token evolution is characterized by flow maps satisfying ODEs in function spaces and that gradient flow curves exist and are unique in the conditional Wasserstein metric space.
Explicit Gradient Formulation: Using adjoint sensitivity analysis, the authors derive an explicit formula for the conditional Wasserstein gradient of the training risk. This involves adjoint variables that satisfy backward ODEs, providing a rigorous foundation for gradient-based training in the mean-field limit.
NTK Injectivity Conditions: A central technical contribution is the characterization of NTK injectivity for attention mechanisms. The paper proves that NTK injectivity is equivalent to the linear independence of log-sum-exp functions modulo affine functions. This condition is shown to be satisfied by diverse token distributions, including:
- Discrete distributions satisfying a distinctness condition (verified almost surely for i.i.d. samples from absolutely continuous measures).
- Uniform distributions on cubes with distinct radii.
- Symmetric multivariate Laplace distributions with distinct covariance matrices.
- Two-component Gaussian mixtures with distinct separation parameters.
Convergence to Global Minima: Under the assumption that the NTK is positive at initialization (specifically, the $V$ -part of the NTK) and the initial loss is sufficiently small, the paper proves that the gradient flow converges to a global minimizer with a linear convergence rate. This result eliminates spurious local minima from the optimization landscape in the specified regime.

Significance and Claims
The paper claims to provide a rigorous mathematical foundation for the training of infinitely deep and wide Transformers. By extending mean-field theory from ResNets (neural ODEs) to Transformers (neural PDEs), it bridges the gap between the theoretical understanding of deep residual networks and the more complex dynamics of attention mechanisms.

The authors emphasize that their framework captures the coupling of multiple token distributions through the attention mechanism, a feature absent in standard ResNet analyses. The establishment of NTK injectivity conditions provides constructive guidance for initialization, ensuring that gradient-based training converges to global minima without getting trapped in spurious local minima, provided the initialization is sufficiently close to the optimal solution.

The paper modestly notes that its analysis is restricted to encoder-style (non-causal) attention. Extending this framework to decoder architectures with causal masking is identified as a substantial technical challenge due to the breaking of symmetries used in the proofs. Furthermore, the current formulation keeps the head dimension fixed while letting the number of heads grow to infinity; capturing the coupled scaling of embedding dimension and head count remains an open problem.