Research map / Architectures and efficiency / RNNs and normalization
Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs
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.
Uses Tensor Programs to derive the exact infinite-width limit of a single multi-head attention layer under standard 1/sqrt(n) scaling and finite heads, showing the output is a hierarchical (conditionally Gaussian) non-Gaussian law.
Problem
Infinite-width Gaussian approximations (NNGP, Tensor Programs) fail to describe standard attention layers with 1/sqrt(n) scaling and finite heads; prior limits relied on 1/n scaling or infinite heads.
Approach
Extends the Tensor Programs master theorem with a new class of finite-dimensional dot-product variables that are Gaussian and independent of the Gaussian-limit vectors; the attention output is then Gaussian conditional on these random similarity scores, yielding a non-Gaussian hierarchical limit.
Key results
- Finite-width attention outputs converge to the theoretical non-Gaussian limit: average log-KL decreases with log_4(n) over widths 16 to 1024 across 10 trials.
- Under 1/n scaling, dot-product scores collapse to zero, whereas under 1/sqrt(n) scaling they remain nondegenerate at n=256.
- As the number of heads H increases, the non-Gaussian limit approaches the infinite-head Gaussian limit; theory remains accurate with H=256 at n=256.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase width | Multi-head attention with H=2 · Multi-head attention | log-KL divergence to the infinite-width limit | Helped |
| use 1/sqrt(n) scaling instead of 1/n scaling | Attention dot-product scores · n=256 | score variability | Helped |
| increase the number of heads | Finite-width attention output · n=256, H=1 versus H=256 | Gaussianity | Mixed |
| increase width and head count proportionally with n_H=64 | Low-rank attention · Low-rank attention | convergence to the infinite-width limit | Helped |
| increase spatial dimension from s=4 to s=8 | Infinite-width theory · Multi-head attention | distribution agreement | No clear effect |
| replace clipping activation with ReLU | Infinite-width theory · Multi-head attention | distribution agreement | No clear effect |
Surprising findings
- Under standard 1/sqrt(n) scaling, the single-head attention output is non-Gaussian even at infinite width, due to randomness in the attention similarity scores.
Limitations
- Assumes constant sequence length and batch size, following the Tensor Programs framework.
- Theory for square weight matrices; low-rank setting only validated numerically.
- Non-Gaussianity analysis is limited to a standalone attention layer; deep Transformer extension is future work.
Experimental setup
Metrics: KL divergence between finite-width and limit densities kernel density estimates of output histograms
Baselines: 1/n-scaling attention Gaussian limit infinite-head Gaussian limit (infinite attention)
Builds on: Tensor Programs NNGP infinite attention (Hron et al.) pseudo-Lipschitz master theorem
Closest papers by meaning
- On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.90)
- Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers (ICML 2025 · similarity 0.89)
- An Infinite-Width Analysis on the Jacobian-Regularised Training of a Neural Network (ICML 2024 · similarity 0.89)
- Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.88)
- Scaling Limits of Long-Context Transformers (NeurIPS 2026 · similarity 0.88)
- Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency (NeurIPS 2025 · similarity 0.87)
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.