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

In modern theoretical analyses of neural networks, which often rely on Gaussian approximations in the infinite-width limit, this work rigorously identifies that attention layers depart fundamentally from Gaussianity under realistic architectural dimensionality and standard scaling.

The gist

The limiting distribution of outputs of a single attention layer converges to a hierarchical Gaussian distribution, being Gaussian conditional on the random similarity score, and this score variable itself converges to a Gaussian.

Non-Gaussian Limiting Distribution

The study investigates the infinite-width limit distribution of outputs of a single attention layer under the common scaling and number of heads, specifically using standard 1/√n-scaling with n dimensionality. The key finding is that the output distribution converges to a hierarchical Gaussian distribution, which is a type of non-Gaussian. Specifically, the limiting distribution is a Gaussian conditional on the random similarity score, and this score variable itself converges to a Gaussian. This result departs from existing studies that have limited themselves to special regimes like infinite heads or 1/n-scaling.

Novel Proof Technique with Dot-Products

The paper develops a novel proof technique focused on analyzing the similarity score variables by the dot-products of an attention layer. The analysis involves showing a convergence of the score variables to their limiting distributions, and then prove a conditional weak convergence of the outputs of an attention layer. This analysis characterizes the output distribution by incorporating the intrinsic randomness from the dot-product. This approach is crucial because it explicitly defines scalar dot-products as a Gaussian vector independent of other Netsor program variables, facilitating the characterization of non-Gaussianity.

Consistency with Numerical Experiments

The theoretical predictions are validated through numerical experiments, confirming that our theory accurately captures the non-Gaussian behaviors exhibited by attention mechanisms. The results show that even when the width is finite, our theory proves sufficiently accurate, provided that the dimension is large enough. Furthermore, simulations demonstrate convergence: the density of y1 converges rapidly to that of Z1 as n increases, and the KL divergence between finite-width and infinite-width limits shows a consistent decay with growing width.

Robustness Across Architectural Variants

The framework demonstrates robustness across different practical settings. The study investigates two scaling regimes: the 1/√n-scaling, which is motivated by its prevalence in practice, and the 1/n-scaling. While the 1/n-scaling regime makes all similarity scores converge to zero in the infinite-width limit, it is the 1/√n regime that preserves variability away from zero at large n. Additionally, the theory is robust to changes in hyperparameters: experiments with different spatial dimensions (s=4 vs. s=8) and activation functions (clipping vs. ReLU) confirm that our theory remains accurate under these variations, suggesting its applicability to a broader range of practical model architectures.

Sub-Gaussianity of Attention Outputs

The paper provides a detailed analysis of the sub-Gaussianity of the limiting distribution. For any output variable, such as the single random variable Z1 for a specific output element, it is shown that Z1p˚(1)1:s ∼ N (0, a1:s (p˚(1)1:s)⊤Var(Z v˜1:s)a1:s (p˚(1)1:s)) when considering the single-head setting. This implies that the outputs are sub-Gaussian, meaning their tail probability decays at a Gaussian rate, even though they cannot be approximated by a single Gaussian distribution with mean and variance.

Implications for Deep Architectures

The findings suggest profound implications for learning dynamics in deep Transformer architectures. The non-Gaussianity of attention layers implies that the presence of higher-order moments associated with such non-Gaussianity could influence signal propagation by preserving feature variability across layers, thereby reducing the risk of rank collapse. This suggests that a new framework distinct from existing Gaussian-based approaches is essential for analyzing deep architectures to understand how non-Gaussianity affects optimization landscapes and training dynamics.

Conclusion

The work establishes a rigorous foundation for analyzing attention mechanisms in the infinite-width regime, proving that their outputs exhibit non-Gaussian behavior under standard scaling, and provides a concrete starting point for developing unified theories of deep Transformer architectures. The framework is shown to be effective even when incorporating low-rank attention settings where the number of heads increases proportionally with network width. This research lays the groundwork for future extensions to deep Transformer architectures, predicting that MLP layers will also converge to non-Gaussian distributions in this limit.

References

[AAA+23] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.

This is a rigorous analysis of how the theoretical framework presented in Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs (Sakai et al., arXiv:2506.00846v3) can be applied to improve AI systems.

The core innovation lies in moving beyond Gaussian approximations to capture the non-Gaussian hierarchical structure inherent in attention mechanisms under realistic scaling and finite head counts.

Here are the specific improvements and capabilities this research enables for AI systems:

) Specific Improvements & Capabilities Enabled by the Research:

2. Non-Gaussian Distribution Modeling for Attention Layers:

3. The system can be theoretically designed to explicitly model the non-Gaussian output distribution of a single attention layer, which is crucial because current standard methods rely on Gaussian approximations that fail under realistic finite-head scenarios (where the number of heads, H, is finite).

5. Robust Attention Mechanism Design:

6. The framework allows for the design of attention mechanisms where the output distribution is characterized as a hierarchical Gaussian distribution conditional on a random similarity score. This explicit structure can guide architectural choices to ensure desired statistical properties in downstream tasks, moving beyond simple second-moment matching (which is what infinite-head limits often achieve).

8. Accurate Scaling Rule Identification:

9. The research provides a rigorous comparison between different scaling regimes (e.g., the standard 1/√n vs. the 1/n regime) and head counts, allowing researchers to empirically determine which scaling strategy yields the most accurate approximation for their specific model size and task requirements, preventing reliance on potentially misleading infinite-head or 1/n approximations.

11. Improved Training Dynamics Analysis:

12. The mathematical proof technique developed (Theorem 3.1) allows for the analysis of how the randomness in dot-product scores propagates through the network structure, providing deeper insights into feature variability and potential risks like rank collapse by characterizing higher-order moments that Gaussian models miss.

14. Sub-Gaussianity Guarantees for Stability:

15. The research provides a proof of sub-Gaussianity for the limiting distribution of attention outputs (Corollary 3.2). This is vital for understanding the tail behavior and stability of the model, allowing engineers to better predict and mitigate extreme output values during training or inference in large models.

) What Improved AI Systems Can Do:

2. Enhanced Transformer Architectures:

3. The theoretical foundation allows for the construction of next-generation Transformers where attention layers are designed not just for performance, but for explicit statistical control over their output distribution, leading to models with more predictable and stable behavior during complex training regimes (e.g., training on massive datasets).

5. Fine-Tuning and Hyperparameter Optimization:

6. By understanding the dependence of the limiting distribution on the random similarity score (the conditioning variable), developers can create more sophisticated hyperparameter tuning strategies that account for this intrinsic randomness, leading to better generalization across different data distributions than standard methods allow.

8. Task-Specific Model Adaptation:

9. The framework's ability to handle realistic finite-width and finite-head settings means that specialized attention mechanisms can be developed tailored precisely to the computational constraints of specific hardware while maintaining high theoretical fidelity, leading to more efficient and accurate models for specialized tasks (e.g., low-rank attention settings where H scales with n).

11. Deeper Theoretical Understanding:

12. The system gains a unified theory that bridges the gap between simplified Gaussian approximations and the complex reality of Transformer dynamics, providing a necessary theoretical tool for analyzing deep, wide architectures beyond what existing NNGP or NTK methods can offer in this specific context.