Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs

Thank you for the constructive feedback and for suggesting additional literature.

Weakness: On the limitation to a single layer at initialization.

Reply: Thank you for the comment. This is indeed a valid and important point. We intentionally focused on a single attention layer and forward pass at initialization, as even this setting has remained analytically intractable under realistic conditions such as finite heads and standard $1/\sqrt{n}$-scaling. Our goal was to rigorously characterize the behavior of attention without relying on approximations like infinite-head limits or altered scaling, which are common in prior theoretical work.

We view this contribution as a foundational step, analogous to how the development of the Neural Network Gaussian Process (NNGP) framework first established forward-pass analyses for fully connected networks before later extensions to training dynamics and deeper architectures. Similarly, a precise and rigorous understanding of a single layer is essential before tackling the more complex and layered nature of modern Transformers, especially since attention layers exhibit non-Gaussian behavior even at initialization.

Q1: MuP scaling.

Reply: Thank you for this important question regarding training stability. Our work analyzes the $1/\sqrt{n}$-scaling, a choice motivated by the fact that it is the standard used in the vast majority of practical Transformer implementations. This scaling is also a central object of study in much of the relevant theoretical literature (e.g., Hron et al., 2020; Dinan et al., 2023; Cowsik et al., 2024).

While some analyses suggest potential instabilities with this scaling, the precise conditions for stable training are still an active area of research and may depend on factors often abstracted away in simplified theoretical limits. These include the number of tokens, the use of finite low-rank QK projections, and data-specific statistics. Indeed, recent work of Hayou and Liu (2025) has shown that data statistics can significantly alter optimal scaling rules of standard muP.

The theoretical foundation of self-attention is still under active development, and we believe different limiting procedures (e.g., $1/n$, $1/\sqrt{n}$, infinite-head) each offer valuable insights into different phenomena. Our contribution to this ongoing effort is to provide the first rigorous characterization of the initial distribution for the most commonly used scaling. This establishes an essential and previously missing foundational starting point from which to analyze the more complex training dynamics.

Q2: Infinite-head limit.

Reply: Thank you for this insightful question. The fundamental difference is that while the infinite-head limit yields a Gaussian approximation, our work rigorously shows that for any finite number of heads, the limiting distribution is fundamentally non-Gaussian.

To your second point, we agree that for many modern models, the infinite-head limit can be a reasonable approximation. Our claim is not that this limit is always bad, but that our framework offers a more general and precise perspective. In other words, while the infinite-head approximation is only accurate when the number of heads is large, our theory provides a good characterization regardless of whether the number of heads is infinite or finite. As our experiments show, our theory remains highly accurate even as $H$ grows large (Fig. 2(b)) and when $n$ and $H$ grow proportionally (Fig. 3 in the appendix). This demonstrates our theory's ability to accurately capture the behavior of attention mechanisms regardless of the number of heads.

Based on your comments we will emphasize this important suggestion more in the main text.

Q3: Relation with Dinan et al. (2023).

Reply: Thank you for highlighting this important and relevant work. You are correct that Dinan et al. also analyze Transformers under a similar $1/\sqrt{n}$-scaling. The primary difference lies in the scope and granularity of the analysis:

• Dinan et al. (2023) focus on tracking the first two moments (the kernel) of the signal to characterize its propagation properties at initialization and during training.
• Our work, in contrast, rigorously derives the full limiting distribution of the attention layer's output, going beyond just the first two moments.

This distinction is crucial because our approach reveals the fundamentally non-Gaussian nature of the output, a detail that an analysis focused only on the kernel does not capture.

On missed citations.

Reply: Thank you for your suggestions on citations. We will ensure our discussion of the related literature is comprehensive in the revised manuscript.

• Regarding Hron et al. (2020), we would like to clarify that this key work was already cited and discussed in our paper, particularly in the context of the infinite-head Gaussian approximation.
• Regarding Dinan et al. (2023), we agree it is highly relevant. As detailed in our response to your Question 3, our work differs in that we rigorously derive the full limiting distribution of the attention output, whereas their analysis focuses on the kernel.
• We also agree that the work of Bordelon et al. (2024) is an important point of comparison. Their work analyzes the training dynamics under various infinite limits to find stable parameterizations. Our scope is distinct, as we focus on rigorously deriving the exact distribution at initialization, providing a foundational starting point for such dynamical analyses.

Reference

[1] Bordelon et al. (2024). Infinite Limits of Multi-head Transformer Dynamics.

[2] Cowsik et al. (2024). Geometric dynamics of signal propagation predict trainability of transformers.

[3] Dinan et al. (2023). Effective Theory of Transformers at Initialization.

[4] Hayou and Liu (2025). Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size.

[5] Hron et al. (2020). Infinite attention: NNGP and NTK for deep attention networks.

Thank you for the thoughtful review and for highlighting important limitations and additional literature.

Weakness 1 and Q1: Training dynamics.

Reply: Thank you for your insightful question. We agree that understanding the conditions for stable training is a vital research direction.

Our choice of $1/\sqrt{n}$-scaling is deliberate, as it reflects the standard used not only in the vast majority of practical Transformer implementations but also in much of the relevant theoretical literature (e.g., Hron et al., 2020; Dinan et al., 2023; Cowsik et al., 2024). We believe that establishing a rigorous theoretical foundation for this canonical setting is an indispensable first step.

Furthermore, the definitive theory for what ensures stable training is still an evolving area of research (like Bordelon et al., 2024), and may depend on factors not typically considered in standard infinite-width limits. For instance, recent work has shown that data statistics can significantly alter the optimal scaling rules of standard muP (Hayou and Liu, 2025).

The theoretical understanding of self-attention is a collective effort, with different limits (e.g., $1/n$, $1/\sqrt{n}$, infinite-head) each providing valuable insights. Our contribution is to provide the first rigorous characterization of the initial distribution for the most commonly used scaling. We see this as providing the essential, previously missing groundwork for future research on the training stability of these widely-used models.

Weakness 2: Fixed batch size and sequence length.

Reply: Thank you for this important observation. As the reviewer correctly notes, our analysis assumes a constant sequence length and batch size as $n\to\infty$. This is a standard assumption made in the existing Tensor Programs literature (Yang, 2019), which our work builds upon, to keep the theoretical analysis tractable. We acknowledge that modern Transformers often utilize very long sequences, and our setting may not capture all the dynamics present in such cases. As suggested, we will expand our limitations section in the revised manuscript to more explicitly state this assumption and its potential implications. Extending the framework to handle growing sequence lengths is a challenging but important direction for future research.

Weakness 3: Novelty and significance.

Reply: Thank you for this comment and for the opportunity to clarify the novelty and significance of our work in relation to Hron et al. (2020).

As the reviewer notes, Hron et al. (2020) analyzed the large-$n$ and large-$H$ limit. Importantly, they also empirically observed that the finite-head case exhibits non-Gaussian behavior, which poses a major challenge for the standard NNGP framework. This observation motivated their shift to the analytically more tractable infinite-head limit, effectively sidestepping the question of what the exact distribution looks like in realistic finite-head settings.

Our contribution directly addresses this open problem by providing the first rigorous derivation of the limiting distribution in the finite-head regime under standard $1/\sqrt{n}$-scaling. We show that the limiting distribution is a hierarchical Gaussian, which, although more complex than a standard Gaussian, is an exact characterization and captures important aspects of real-world attention mechanisms that previous Gaussian approximations cannot.

Furthermore, we believe this result opens the door to a deeper theoretical understanding of Transformers. In particular, the hierarchical Gaussian structure we uncover at the first layer suggests a pathway to study how non-Gaussianity may compound through depth, possibly leading to heavy-tailed or anisotropic distributions in deep layers. These phenomena could play a critical role in the representational power and feature learning dynamics of Transformers. We see our work as a first rigorous step toward such a refined theory, one that moves beyond infinite-width approximations and closer to practical architectures.

Finally, we would be very happy to engage in further discussion if there are particular directions or questions the reviewer is especially interested in. We welcome any follow-up comments or suggestions.

Q2: The attention scaling considering head counts.

Reply: Thank you for this excellent and highly constructive suggestion. We agree completely that including the $1/\sqrt{H}$-scaling factor provides a more meaningful comparison as the number of heads increases. We will adopt this formulation throughout our revised manuscript and have redone the corresponding experiments.

This change has significantly strengthened our work by enabling a direct comparison with Hron et al. (2020). We found that with this scaling:

• The limiting distribution of the attention scores in our framework perfectly matches the Gaussian distribution in their infinite-head limit.
• The variance of our finite-head non-Gaussian output distribution theoretically matches the variance of their infinite-head Gaussian output distribution.

To illustrate this, we reran the experiment for Figure 2(b) and included the theoretical distribution from Hron et al. (2020). The new plot, which we will include in the revised manuscript, now visually demonstrates that as $H$ increases, our non-Gaussian distribution approaches the Gaussian distribution from Hron et al. (2020). This observation provides strong visual support for our framework as a more general theory that encompasses the infinite-head limit.

Q3: The realism of the theoretical setting.

Reply: This is a very important point. We agree that many state-of-the-art models scale up the number of heads $H$ significantly. Our response is twofold.

First, the main insight of our work is a characterization of the attention mechanism that is broadly applicable, regardless of the number of heads $H$. While our theory is formally derived in the $n\to\infty$ limit, its predictions are not confined to a specific relative scaling speed of $n$ versus $H$. Our experiments confirm this robustness across the practical scenarios. Figure 2(b) shows that our theory remains accurate even as $H$ grows large, and the analysis in our appendix (Figure 3) further validates our framework in the common setting where the dimension-per-head is fixed.

Second, we believe that a rigorous analysis of the finite-$H$ regime remains valuable. It provides a foundational understanding of the attention mechanism's building blocks and is directly applicable to smaller-scale models where the infinite-head approximation may not be suitable.

We will revise our manuscript to emphasize more clearly that our findings provide a robust description for any number of heads, clarifying their broad applicability.

On missed citations.

Reply: We thank you for bringing these key papers to our attention. We will cite and discuss these works in our revised manuscript. The primary differences between their approaches and ours lie in the scope and methodology:

• Cowsik et al. (2024) use mean-field theory to characterize the edge of chaos by analyzing forward/backward signal propagation. Their analysis, however, relies on several simplifying assumptions (e.g., the Gaussianity of the QK product) to make the problem tractable. Our work, in contrast, provides a rigorous derivation of the exact limiting distribution of the forward pass without such assumptions.
• Bordelon et al. (2024) use dynamical mean field theory to provide an important analysis of the training dynamics of Transformers under various infinite limits, including infinite width, heads, and depth. Their primary focus is on identifying parameterizations that ensure stable feature learning over time. While their analysis focuses on the time evolution of second moments (kernels), it does not provide a rigorous derivation of the exact limiting output distribution at initialization. Our work, with its distinct scope, addresses this specific problem by establishing the precise, non-Gaussian initial distribution.
• Noci et al. (2023) aim to describe signal evolution using a stochastic differential equation, which requires modifying the standard softmax function. Our scope is distinct, as we focus on rigorously analyzing the attention layer in its standard, unmodified form.

Reference

[1] Bordelon et al. (2024). Infinite Limits of Multi-head Transformer Dynamics.

[2] Cowsik et al. (2024). Geometric dynamics of signal propagation predict trainability of transformers.

[3] Dinan et al. (2023). Effective Theory of Transformers at Initialization.

[4] Hayou and Liu (2025). Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size.

[5] Hron et al. (2020). Infinite attention: NNGP and NTK for deep attention networks.

[6] Noci et al. (2023). The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit.

[7] Yang (2019). Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation.

Thank you for your detailed review. We acknowledge the concerns regarding clarity, impact, and the scope of our experiments.

Weaknesses: Derivation is somewhat obscure and difficult to follow (at least for non-specialists).

Reply: Thank you for your valuable feedback. As suggested by Reviewer pKBC, we will add a pseudocode example to the appendix to provide a more concrete and intuitive illustration of our setup.

Weaknesses: Impact of analysis is not clear, what can be learned?

Question: What is the impact of your analysis for optimization, initialization or hyper-parameter selection?

Reply: Thank you for this important question. We present specific impacts on current choices as well as future impacts.

As our current work focuses on the forward pass analysis, one concrete impact is that it sheds light on how initialization scale affects the distribution of attention outputs, particularly in realistic settings with finite heads and $1/\sqrt{n}$-scaling. Understanding the output distribution precisely (e.g., its variance and tail behavior) provides a theoretical basis for choosing initialization scales that preserve signal variability and avoid pathological behaviors such as rank collapse. While prior work has offered valuable insights, it often relies on specific regimes, such as infinite heads or tailored scaling, which depart from standard practices. Our contribution is to rigorously derive the infinite-width limit under the commonly used $1/\sqrt{n}$-scaling with finite heads, revealing its non-Gaussian nature.

Looking forward, we believe our framework can serve as a foundation for analyzing backward pass dynamics as well. Extending the analysis to include gradients and weight updates could eventually inform learning rate selection or adaptive scaling strategies in training attention-based models. Much like how Yang (2019) provided the foundation for later studies on training dynamics and hyperparameter selection (Yang and Hu, 2021) for neural networks without attention, we believe our work provides a similar foundation for the attention case. This remains an exciting direction for future work.

We agree that the contributions of our analysis should be made more explicit, and we will revise the paper accordingly to emphasize them.

Weaknesses: Experiments are very rudimentary.

Question: Can you provide stronger experimental evidence to support your analysis?

Reply: Thank you for the suggestion to include more experiments. In response, we have conducted an additional experiment by varying the spatial dimension, $s$ (the number of tokens).

The table below presents the results of the experiment from Section 5.1 with $s=8$. The values represent the KL divergence (averaged over 10 independent trails) between the empirical distribution of the attention output ($y_{1}^{1}$) and our theoretical limit ($Z^{y^{1}}$), corresponding to the analysis in Figure 1(b).

As the table demonstrates, the predictions of our theory remain highly accurate even when the number of tokens is changed. We plan to conduct further experiments with different activation functions in the future. These new results will be incorporated into the appendix of our revised manuscript.

Question: Is there a way you can make the notation more readable?

Reply: We agree that the notation, while consistent with that used in Yang (2020), may pose a burden for readers. To improve clarity, we will include a notation summary table in the appendix and clarify key notational conventions when first introduced in the main text. We hope this will make the paper more accessible and easier to follow, especially for readers new to this line of work.

Furthermore, to improve clarity and make the notation more accessible, we will add a pseudocode example for the multi-head attention layer in the style of the NETSOR programs from Yang (2019) in the appendix of our revised paper.

Reference

[1] Yang (2019). Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are Gaussian processes.

[2] Yang (2020). Tensor programs iii: Neural matrix laws.

[3] Yang and Hu (2021). Tensor programs iv: Feature learning in infinite-width neural networks.

Thank you for the constructive feedback and for highlighting our contribution.

Weakness 2-2: The non-Gaussianity is somewhat expected.

Reply: We agree that the breakdown of Gaussianity in attention mechanisms with finite heads has been heuristically observed in earlier studies, particularly in Hron et al. (2020), which provides compelling empirical evidence and predictions. Our work builds on this intuition by rigorously establishing the exact form of the limiting distribution for a single attention layer under realistic settings (finite heads and standard $1/\sqrt{n}$-scaling). In this sense, while the emergence of non-Gaussianity may be expected qualitatively, we believe that the precise derivation of the hierarchical Gaussian structure and its formal validation is a meaningful contribution to the theory of attention.

Q1: Analysis of deeper models and overcoming the non-Gaussian challenge.

Reply: This is an excellent and challenging question. We predict that the behavior of the attention outputs will likely deviate even further from Gaussianity as layers are stacked. In fact, our analysis indicates that even at the first layer, the output already exhibits heavier tails than a Gaussian. This suggests that as layers accumulate, the output distribution may increasingly diverge from any hierarchical Gaussian form, possibly yielding even heavier-tailed or more structured non-Gaussian behaviors.

While a complete analysis of deep attention networks remains an open and challenging direction, we believe our work offers a necessary first step toward that goal. To our knowledge, a theoretical framework for handling such evolving non-Gaussian distributions does not yet exist. We are currently exploring how the hierarchical structure evolves under layer composition, and we view this as an exciting avenue for future research.

We thank the reviewer again for highlighting this important point. We will incorporate a discussion of this important challenge into the revised version of our paper.

Q2: Implications of non-Gaussianity and differences from a standard Gaussian.

Reply: We agree that a deeper discussion on the implications of our findings is crucial.

The hierarchical Gaussian distribution derived in our work differs from the standard Gaussian in a fundamental way: it introduces randomness in the variance itself. That is, the output distribution is conditionally Gaussian given a set of random similarity scores, but these scores themselves follow a Gaussian distribution, resulting in a scale mixture of Gaussians. It is well-known in probability theory that such hierarchical structures produce heavier tails than standard Gaussians. For instance, when the variance follows a uniform distribution, the resulting distribution is a normal scale mixture; when it follows the square of a Gaussian, the resulting distribution becomes a variance-gamma distribution, both of which are strictly heavier-tailed than the Gaussian.

This inherent non-Gaussianity may significantly affect several aspects of learning in self-attention models. For example:

1. In feature learning, the alignment of activations with specific vector directions is often critical. Heavier-tailed distributions may make such alignment easier by amplifying high-magnitude components, thereby facilitating the emergence of dominant features.
2. In terms of optimization landscapes, the anisotropy induced by non-Gaussianity could lead to more irregular curvature, possibly affecting convergence properties of training dynamics (e.g., through sharper gradients or more prominent saddle regions).
3. Additionally, signal propagation may be influenced by the presence of higher-order moments. For instance, heavier tails might reduce the risk of rank collapse by preserving variability in feature representations across layers.

At the same time, the non-Gaussian nature of attention outputs poses significant challenges for theoretical analysis. Many existing frameworks, such as Tensor Programs, rely heavily on Gaussian assumptions for tractability. Extending these tools (or developing new ones) to rigorously capture non-Gaussian behaviors is an intriguing and important open question, and we believe our work provides a concrete starting point for such future developments.

We appreciate the reviewer’s insight in drawing attention to these broader implications, and we hope our response clarifies the potential impact and theoretical interest of the non-Gaussian structure revealed in our study.

Q3: Readability and pseudocode.

Reply: We appreciate this suggestion. To improve clarity and make the notation more accessible, we will add a pseudocode example for the multi-head attention layer in the style of the NETSOR programs from Yang (2019) in the appendix of our revised paper.

Reference

[1] Hron et al. (2020). Infinite attention: NNGP and NTK for deep attention networks.

[2] Yang (2019). Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are Gaussian processes.