# The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit

arXiv: https://arxiv.org/abs/2306.17759 · HTML: https://mlsanity.com/paper/2306.17759/

## Key facts

- Cited by 47 papers in the map (data through 14 Oct 2026).
- Builds on 54 of its 76 references that are in the map.
- First posted on arXiv 30 Jun 2023.

## Builds on (most cited first, 30 of 54)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Deep Neural Networks as Gaussian Processes](https://mlsanity.com/paper/1711.00165/) (2017)
- [Deep Learning Scaling is Predictable, Empirically](https://mlsanity.com/paper/1712.00409/) (2017)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Gaussian Process Behaviour in Wide Deep Neural Networks](https://mlsanity.com/paper/1804.11271/) (2018)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [Exponential expressivity in deep neural networks through transient chaos](https://mlsanity.com/paper/1606.05340/) (2016)
- [Deep learning: a statistical viewpoint](https://mlsanity.com/paper/2103.09177/) (2021)
- [Deep Information Propagation](https://mlsanity.com/paper/1611.01232/) (2016)
- [Feature Learning in Infinite-Width Neural Networks](https://mlsanity.com/paper/2011.14522/) (2020)
- [The Principles of Deep Learning Theory](https://mlsanity.com/paper/2106.10165/) (2021)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [When Do Neural Networks Outperform Kernel Methods?](https://mlsanity.com/paper/2006.13409/) (2020)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [Tensor Programs II: Neural Tangent Kernel for Any Architecture](https://mlsanity.com/paper/2006.14548/) (2020)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks](https://mlsanity.com/paper/2202.08658/) (2022)

## Built on by (most cited first, 30 of 47)

- [LoRA+: Efficient Low Rank Adaptation of Large Models](https://mlsanity.com/paper/2402.12354/) (2024)
- [Don't be lazy: CompleteP enables compute-efficient deep transformers](https://mlsanity.com/paper/2505.01618/) (2025)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)
- [Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit](https://mlsanity.com/paper/2309.16620/) (2023)
- [The Curse of Depth in Large Language Models](https://mlsanity.com/paper/2502.05795/) (2025)
- [The Impact of Initialization on LoRA Finetuning Dynamics](https://mlsanity.com/paper/2406.08447/) (2024)
- [Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains](https://mlsanity.com/paper/2402.04161/) (2024)
- [Infinite Limits of Multi-head Transformer Dynamics](https://mlsanity.com/paper/2405.15712/) (2024)
- [Simplifying Transformer Blocks](https://mlsanity.com/paper/2311.01906/) (2023)
- [Measure-to-measure interpolation using Transformers](https://mlsanity.com/paper/2411.04551/) (2024)
- [There Will Be a Scientific Theory of Deep Learning](https://mlsanity.com/paper/2604.21691/) (2026)
- [Geometric Dynamics of Signal Propagation Predict Trainability of Transformers](https://mlsanity.com/paper/2403.02579/) (2024)
- [Super Consistency of Neural Network Landscapes and Learning Rate Transfer](https://mlsanity.com/paper/2402.17457/) (2024)
- [Two failure modes of deep transformers and how to avoid them: a unified theory of signal propagation at initialisation](https://mlsanity.com/paper/2505.24333/) (2025)
- [Theory of Scaling Laws for In-Context Regression: Depth, Width, Context and Time](https://mlsanity.com/paper/2510.01098/) (2025)
- [Always Skip Attention](https://mlsanity.com/paper/2505.01996/) (2025)
- [Is Random Attention Sufficient for Sequence Modeling? Disentangling Trainable Components in the Transformer](https://mlsanity.com/paper/2506.01115/) (2025)
- [Understanding and Minimising Outlier Features in Neural Network Training](https://mlsanity.com/paper/2405.19279/) (2024)
- [Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers](https://mlsanity.com/paper/2405.15926/) (2024)
- [AltLoRA: Towards Better Gradient Approximation in Low-Rank Adaptation with Alternating Projections](https://mlsanity.com/paper/2505.12455/) (2025)
- [Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning](https://mlsanity.com/paper/2602.06204/) (2026)
- [Towards Training Without Depth Limits: Batch Normalization Without Gradient Explosion](https://mlsanity.com/paper/2310.02012/) (2023)
- [Clustering in Deep Stochastic Transformers](https://mlsanity.com/paper/2601.21942/) (2026)
- [Understanding Catastrophic Forgetting In LoRA via Mean-Field Attention Dynamics](https://mlsanity.com/paper/2402.15415/) (2024)
- [Bayesian Inference with Deep Weakly Nonlinear Networks](https://mlsanity.com/paper/2405.16630/) (2024)
- [Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs](https://mlsanity.com/paper/2506.00846/) (2025)
- [Attention-Only Transformers via Unrolled Subspace Denoising](https://mlsanity.com/paper/2506.03790/) (2025)
- [Differential Equation Scaling Limits of Shaped and Unshaped Neural Networks](https://mlsanity.com/paper/2310.12079/) (2023)
- [Pretraining Graph Transformers with Atom-in-a-Molecule Quantum Properties for Improved ADMET Modeling](https://mlsanity.com/paper/2410.08024/) (2024)
- [ReLU's Revival: On the Entropic Overload in Normalization-Free Large Language Models](https://mlsanity.com/paper/2410.09637/) (2024)

## Abstract (arXiv)

In deep learning theory, the covariance matrix of the representations serves as a proxy to examine the network's trainability. Motivated by the success of Transformers, we study the covariance matrix of a modified Softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. We show that at initialization the limiting distribution can be described by a stochastic differential equation (SDE) indexed by the depth-to-width ratio. To achieve a well-defined stochastic limit, the Transformer's attention mechanism is modified by centering the Softmax output at identity, and scaling the Softmax logits by a width-dependent temperature parameter. We examine the stability of the network through the corresponding SDE, showing how the scale of both the drift and diffusion can be elegantly controlled with the aid of residual connections. The existence of a stable SDE implies that the covariance structure is well-behaved, even for very large depth and width, thus preventing the notorious issues of rank degeneracy in deep attention models. Finally, we show, through simulations, that the SDE provides a surprisingly good description of the corresponding finite-size model. We coin the name shaped Transformer for these architectural modifications.