Title: Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs
Venue: NeurIPS 2025
Accepted: yes
Keywords: attention mechanism

infinite-width

tensor programs

limit distribution
Abstract: In modern theoretical analyses of neural networks, the infinite-width limit is often invoked to justify Gaussian approximations of neuron preactivations (e.g., via neural network Gaussian processes or Tensor Programs). However, these Gaussian-based asymptotic theories have so far been unable to capture the behavior of attention layers, except under special regimes such as infinitely many heads or tailored scaling schemes. In this paper, leveraging the Tensor Programs framework, we rigorously identify the infinite-width limit distribution of variables within a single attention layer under realistic architectural dimensionality and standard $1/\sqrt{n}$-scaling with $n$ dimensionality. We derive the exact form of this limit law without resorting to infinite-head approximations or tailored scalings, demonstrating that it departs fundamentally from Gaussianity. This limiting distribution exhibits non-Gaussianity from a hierarchical structure, being Gaussian conditional on the random similarity scores. Numerical experiments validate our theoretical predictions, confirming the effectiveness of our theory at finite width and accurate description of finite-head attentions. Beyond characterizing a standalone attention layer, our findings lay the groundwork for developing a unified theory of deep Transformer architectures in the infinite-width regime.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the distribution of attention layer outputs at random initialization. Unlike previous settings (such as standard MLPs or CNNs) where the limiting distribution of preactivations becomes (at infinite width) Gaussian with variance controlled by a moment of the previous layer's distribution, the authors prove that the distribution for an attention layer at large $n$ is only Gaussian conditional on the dot product scores $p = \\frac{1}{\\sqrt{n}} k \\cdot q$ where $k$ is key and $q$ is query. As a consequence the full distribution of the attention outputs can be non-Gaussian but with a simple hierarchical distribution that the authors characterize. They prove that the $p$ variables converge to Gaussian variables and\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The result is very intuitive and the authors provide rigorous derivations which establish this limiting distribution. \n2. The authors also provide a few experiments that show the validity of their results. \n\nWeaknesses:\n1. The authors focus on properties of the transformer at random initialization. The authors do not argue or show that this scaling rule yields stable training in the $n\\to\\infty$ limit (indeed prior works showed that the $1/n$ scaling is necessary if $k, q$ variables move by $\\Theta(1)$ since their updates are correlated https://arxiv.org/abs/2203.03466, https://arxiv.org/abs/2405.15712). \n2. The analysis also requires batchsize and sequence length (total tokens) to be constant with respect to $n$ as $n \\to \\infty$. While this is commonly employed in signal propagation papers, it could be disregarding important effects when describing modern transformers. \n3. One weakness is that the paper may be significance / novelty in that large $n$ and large $H$ limit were analyzed in Hron et al. This paper removes the need for a large $H$ limit, but at the cost of a more complicated hierarchical distribution for the output variables. However, I am willing to improve my assessment of novelty / significance during conversations with the authors.\n\nQUESTIONS:\n1. Do the authors have thoughts on training dynamics at large $n$? The $\\mu$P scaling of Yang et al suggests that $1/n$ attention layer scaling would admit stable training if you desire constant scale changes to keys and queries. While this doesn't have an interesting large $n$ limit at initialization, it is becoming more commonly employed in practical LLM training. \n2. I recommend the authors to define \n$$y = \\frac{1}{\\sqrt H} \\sum_{a=1}^H  v^{ a } \\text{Softmax}( p^{a} )$$\n, since when you include the $\\frac{1}{\\sqrt H}$ it keeps the variance stable as heads $H$ is increased and would make Figure 2b's comparison more reasonable. \n3. The authors claim that heads are not infinite and that this makes large head limits unrealistic. To be fair, the $n$ is not infinite in practice either. Wouldn't it be more reasonable to study which dimensions are **scaled up** as practitioners vary model size? In many papers (such as the LLAMA paper  https://arxiv.org/abs/2302.13971 Table 2 , or the GPT paper https://arxiv.org/abs/2005.14165 Table 2.1), the $n$ (dimension per head) is either not scaled up or is scaled up much more slowly than the number of heads. \n\nThis work misses some key papers that have studied large $n$ limits of transformers, both at initialization and during training\n[1] https://arxiv.org/abs/2403.02579 which study the signal propagation of deep transformers with the $1/\\sqrt{n}$ scaling at initialization.\n\n[2] https://arxiv.org/abs/2405.15712, which studies the training dynamics of infinite $n$, infinite $H$, and infinite depth $L$ limits of transformers under different scalings. They argue that the large $n$ limit with $1/\\sqrt{n}$ has an interesting initial distribution, it does not admit stable training. If you adopt a $\\frac{1}{n}$ scaling, then it allows stable training but taking $n \\to \\infty$ causes all heads to collapse to the same values. Instead the large head limit at fixed $N$ maintains diversity across attention heads.\n\n[3] https://arxiv.org/abs/2306.17759, where the authors study the signal propagation of *deep shaped transformers* where the depth and width diverge proportionally and the softmax attention is rescaled. In this paper the authors derive a stochastic differential equation for the evolution of the cross-token correlation along the residual stream of the network. \n\nIt would be helpful to the reader to understand the similarities / differences between their approach and the findings of these other works. \n\nIn general, I am very open to increasing my score upon some of these issues being resolved.\n\nLIMITATIONS:\nThe authors could expand do acknowledge that this work focuses on the initialization regime. They could consider adding information that they currently require samples & sequence length to be fixed as $n \\to \\infty$.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the infinite-width limit of a single attention layer by extending the Tensor Programs framework to handle the standard $1 / \\sqrt{n}$ attention scaling used in transformers. The authors find that in comparison to previously studied limits which use a $1 / n$ scaling, the output distribution is non-Gaussian due to a hierarchical structure where the output is Gaussian when conditioned on the jointly Gaussian random attention scores. The authors then provide some numerical experiments demonstrating the accuracy of the limiting distribution for finite widths.\n\nSTRENGTHS AND WEAKNESSES:\nThe results is novel and interesting because it characterizes a different infinite-width limit for a practically relevant instantiation of attention. However, the result is limited to only a single attention layer and only describes the forward pass at initialization.\n\nQUESTIONS:\nI am a bit confused about the proposed scaling regime and its relations to other works. \n\n1. muP scaling [1]: My understanding is that the $1/n$ scaling (as opposed to $1/\\sqrt{n}$ scaling) is *necessary* to prevent blow-up after one step. I believe that [2] also supports this claim. Will this be an issue in the scaling of this paper? \n\n2. Both [2] and [3] also analyze the infinite head limit. But [2] seems to make the head dimension finite, while [3] (like in this paper) also grows the head dimension. Is there a fundamental difference here to be aware of? It also seems that in practice the number of heads is very large, so is the infinite-head limit actually a bad approximation?\n\n3. What is the relation of all these works to the scaling in [4]? It seems there is also a $1/ \\sqrt{n}$ scaling studied here?\n\nAs a comment, I think [3], [4] should be cited and discussed.\n\n1. Yang et al. (2021), Tuning large neural networks via zero-shot hyperparameter transfer.\n2. Bordelon et al. (2024), Infinite Limits of Multi-head Transformer Dynamics\n3. Hron et al. (2020), Infinite attention: NNGP and NTK for deep attention networks\n4. Dinan et al. (2023), Effective Theory of Transformers at Initialization.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents an analysis of the infinite width limit for transformers with $1/\\sqrt{n}$ scaling of random weights, improving prior analyses that had to resort to the less interesting $1/n$ scaling. Moreover, the  paper aims to avoid a limit where the number of attention heads is taken to infinity. The arguments are viable in my view as otherwise the limit case diverges from what is observed in large - but implementable - networks. \n\nThe paper makes use of (and extends) the existing tensor program framework, which is well established in the literature. Compared to previous work, this yields meaningful results for the case of finite width and a finite number of heads. The key result is theorem 3.1, which handles the limiting distributions of pre-activations (Gaussians) as per previous results, but extends these by deriving limiting distributions for the (pre-softmax) attention inner products.  As far as I understand the key idea is to calculate these directly instead of going through the limiting distribution of the Gaussian pre-activations. This is - in principle - a viable extension of the so-called Nestor master theorem. A sketch of the theorem is given, but the proof is rather technical. An example for the multi-head case is added to illustrate the use of the novel analysis.\n\nSTRENGTHS AND WEAKNESSES:\n+ Extends our understanding of infinite-limit analysis to attention-based models\n+ Innovative way of handling inner products of random variables (attention pre-activation)\n+ Focuses on key limitations of existing analyses \n+ Provides a non-trivial technical result\n- Derivation is somewhat obscure and difficult to follow (at least for non-specialists)\n- Impact of analysis is not clear, what can be learned?\n- Experiments are very rudimentary\n\nQUESTIONS:\nCan you provide stronger experimental evidence to support your analysis?\nWhat is the impact of your analysis for optimization, initialization or hyper-parameter selection?\nIs there a way you can make the notation more readable? Currently it requires quite some effort.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the infinite-width limit of a single multi-head attention layer under a finite number of heads and the standard $1/\\sqrt{n}$ scaling for attention. The authors rigorously prove that under these conditions, the limiting distribution of the layer's output is non-Gaussian. They characterize this distribution as a hierarchical Gaussian, where the output is Gaussian conditional on the random attention scores, which themselves converge to a Gaussian distribution.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n- The paper's main strength is its focus on a realistic setting for attention ($1/\\sqrt{n}$ scaling, finite heads), which has been a gap in prior infinite-width analyses. This work moves beyond the simplifying assumptions of infinite heads or $1/n$ scaling to provide a more faithful theoretical model.\n- The mathematical derivation, leveraging tensor programs, is rigorous and technically sound, with numerical experiments supporting the theoretical results.\n\n**Weaknesses:**\n-  The paper proves the non-Gaussianity but does not discuss what this means for the behavior of attention layers or deep Transformers. A deeper discussion of these implications is missing, which limits the paper's overall contribution to our understanding of these architectures (see question 2).\n- The paper demonstrates that the Gaussianity breaks down at the very first attention layer but offers no discussion or potential solution for how to proceed with the analysis of deeper models. While rigorously proven, this result is somewhat expected by the community ([HBSDN20]).\n- While the paper is generally well-organized, the mathematical notation is dense and can be difficult to follow.\n\nQUESTIONS:\n- The core result is that the limiting distribution is non-Gaussian, which breaks the standard analysis for deep models. Do the authors have any insights, even speculative, on how this challenge might be overcome? For example, are there frameworks or modifications to tensor programs that could handle such hierarchical Gaussian distributions as inputs for subsequent layers?\n- What are the main differences between the hierarchical Gaussian distribution and the standard Gaussian? How might this inherent non-Gaussianity affect signal propagation (e.g. rank collapse), feature learning, or the optimization landscape in self-attention models?\n- To improve readability, the authors may consider adding a pseudocode example for a single attention layer, in the style of the Netsor programs from [Yan19b]. This would be helpful for readers.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}