Research map / Architectures and efficiency / Attention and transformers
Infinite Limits of Multi-head Transformer Dynamics
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives DMFT-based infinite-limit descriptions of feature-learning multi-head transformer training as key/query dimension, head count, and depth go to infinity, and checks convergence numerically on vision and causal LM experiments.
Problem
Which parameterizations of transformers admit well-defined infinite width (key/query dimension, heads) and depth limits in which attention layers still learn features during training, and what statistical description do these limits have?
Approach
Applies dynamical mean field theory (MSR path integral) to randomly initialized multi-head transformers trained with SGD, identifying admissible scalings (mu-P attention scaling for N to infinity; head-averaged kernels for H to infinity; L^{-alpha_L} residual scaling for depth) and showing N to infinity collapses heads while H to infinity keeps a distribution over attention heads.
Key results
- Shows that mu-P attention scaling (alpha_A=1) is necessary for a stable N to infinity SGD limit, and that heads then collapse to identical dynamics.
- Shows that at fixed N, the H to infinity limit exists with output logits converging at rate O(H^-1) early in training, with convergence weakening later.
- Finds that with alpha_L=1 the attention weights keep moving as L grows, while alpha_L=1/2 freezes attention layers in the depth limit; scaling depth gives little benefit in C4 experiments.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase the attention head count | Vision Transformer · CIFAR-5M | test loss | Helped |
| Increase the attention head count | Vision Transformer · CIFAR-5M | logit convergence to the infinite-head proxy | Mixed |
| Use μP attention scaling (alpha_A = 1) instead of alpha_A = 1/2 | Multi-head Transformer · CIFAR-5M, 2500 training steps | variance of attention matrices across heads | Hurt |
| Use alpha_A = 1/2 instead of alpha_A = 1 | Multi-head Transformer · CIFAR-5M, 2500 training steps | diversity of attention variables across heads | Helped |
| Increase the key/query dimension N while using alpha_A = 1/2 | Causal Transformer · C4 | final loss | Helped |
| Increase the key/query dimension N while using alpha_A = 1 | Causal Transformer · C4 | final loss | Helped |
| Increase the attention head count | Causal Transformer · C4 | final loss | Helped |
| Use alpha_A = 1/2 instead of alpha_A = 1 | Causal Transformer · C4 | final loss | Helped |
| Increase depth L at fixed N and H | Causal Transformer · C4 | final loss | No clear effect |
| Use residual branch scaling alpha_L = 1 instead of alpha_L = 1/2 | Transformer · CIFAR-5M, large-depth limit | loss | Helped |
| Use residual branch scaling alpha_L = 1 instead of alpha_L = 1/2 | Transformer · Infinite-depth limit | attention-layer feature learning | Helped |
| Use residual branch scaling alpha_L = 1/2 instead of alpha_L = 1 | Transformer · Infinite-depth limit | attention-layer feature learning | Hurt |
Surprising findings
- Under mu-P attention scaling, the attention variance across heads decreases with N, so scaling N yields little benefit over scaling heads, whereas alpha_A=1/2 gives a substantial benefit from larger N.
- Under alpha_L=1/2, the attention pre-activations vanish in the infinite-depth limit, so the MHSA blocks reduce to average pooling.
Limitations
- Exact theory is restricted to SGD; Adam scaling is based only on heuristic arguments.
- Analysis assumes fixed training steps and batch size as widths and depth go to infinity, ignoring compute-optimal regimes.
- The H to infinity limit is computationally intractable, so it is approximated by ensembles of wide models.
Experimental setup
Datasets: CIFAR-5M C4
Benchmarks: CIFAR-5M C4
Metrics: test loss mean square difference of output logits kernel convergence
Baselines: standard 1/sqrt(N) attention scaling mu-P attention (alpha_A=1) alpha_L=1/2 residual scaling
Builds on: mu-P Tensor Programs DMFT for feature learning in infinite width networks
Closest papers by meaning
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 2024 · similarity 0.89)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.88)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.88)
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · similarity 0.88)
- Dynamical Properties of Tokens in Self-Attention and Effects of Positional Encoding (NeurIPS 2025 · similarity 0.88)
- Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers (NeurIPS 2026 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.