Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs
Authors: Mana Sakai
Ryo Karakida
Masaaki Imaizumi
Venue: NeurIPS 2025 poster ()
Keywords: attention mechanism
infinite-width
tensor programs
limit distribution
Primary area: deep_learning
TL;DR: We rigorously identify the infinite–width limit distribution of neurons within a single attention layer under realistic architectural dimensionality

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.

Infinite-Width Limit of Single Attention Layer: Analysis via Tensor Programs
Mana Sakai
Ryo Karakida
Masaaki Imaizumil
The University of Tokyo
2 National Institute of Advanced Industrial Science and Technology 'RIKEN Center for Advanced Intelligence Project
mana. sakai. 777.ma70mmm....m com, serrakida.ro@aist...o imai zumi wcc..ecc...... w-cc.u.tokyo..o ac. jp
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 wnfinite-width limit distribution of variables within a single attention layer under realistic architectural dimensionality and standard 1/ Vitssalngg with n dimensionality. We derive the exact form of this limit law without resorting infinite- head approximations or tailored scalings, demonstrating that it departs fundamentally from Gaussianity. This limiting distribution exhibits non- -auusasssyy from hierarchical structure, being Gaussian conditional on the random similarity scores. Numerical experiments validate our theoretical predictions, confirming the effectiveness of our theory finite width and accurate description of finite-head attentions. Beyond characterizing standalone attention layer, our findings lay the groundwork for developing a unified theory of deep Transformer architectures in the infinite-width regime.
Introduction
useful approach to understanding the complex probabilistic behavior neural networks through the study of parameter distributions in the infinite-wii width limit. Notable examples include the neural network Gaussian process (NNGP) [LBN 18 HBSDN20), which approximates the limit of stochastic parameter distributions with Gaussian processes; the neural tangent kernel (NTK) (JGH18, which represents the model near the initial value with kernel functions; mean field theory, which describes the update of parameter distributions [MMN18;; and the Tensor Program, developed by [Yan19]] Yan20a) YL21 | 1 YHB'21 21 YYZH24) which is general probabilistic analytical framework that unifies the representation of the infinite-width limit for wide range of neural architectures and multiple layers. These methods provide probabilistic models that closely approximate the complex phenomena of neural networks.
One challenge in the studies to properly represent the attention mechanisms used in Trans- formers [VSP*17] which frequently appear in recent large-scale architectures (DCLT19) AAA+23] ZZL Unlike ordinary multi-laeer perceptrons, an attention layer has interactions between query and key variables by multiplication, which makes the infinite-width limit distribution considerably
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

more complex as shown in [HBSDN20], for example. To avoid this difficulty, the existing studies have mainly limited themselves to two special settings compute the parameter distribution of the limit: (i) Infinite-heat regime: [HBSDN20] considers the NNGP for the attention mechanism with infinite heads, resulting in a Gaussian approximation. (ii) 1/n-scaling regime: the Tensor Programs [Yan approximate the multiplication by changing the scale of the input variables for attention layers from 1/vn to I/n, where n is a dimension of variables. However, these simplifications compromise the expressiveness and structure of actual attention mechanisms (and consequently, of Transformers). Specifically, the limit distribution at the infinite-head not an effective approximation of the actual attention mechanisms because it differs significantly from the finite-head case. Also, the 1/n-scaling regime makes all similarity scores converge to zero in the infinite-width limit, making the model equivalent to not measuring similarity between key and query vectors. Therefore, the limit parameter distribution of the attention mechanism is still under development.
In this study, we investigate the infinite-width limit distribution of outputs of a single attention layer under the common scaling and number of heads. To achieve this, we apply the Tensor Programs framework and analyze new class of variables defined by the multiplication of intermediate variables, and derive a corresponding limit distribution. This new class of variables allows us to represent the dot-product score by the multiplication of keys and queries in the attention mechanism.
We summarize our contributions as follows:
We study a distribution of outputs of a single attention layer with the 1/ V/ and finite heads, and demonstrate that in the infinite wiith limit, the output distribution converges to a hierarchical Gaussian distribution, which is a type of non-Gaussian. Specifically, the limiting distribution is Gaussian conditional on the random similarity score, and this score variable itself converges to a Gaussian.
Non-Gaussian limiting distribution:
Consistency with numerical experiments: Our experiments justify that our theoretical limit distribution accurately captures the non-Gaussian behaviors exhibited by attention mechanisms. Specifically, even when the width finite, our theory proves sufficiently accurate, provided that the dimension is large enough.
Novel proof technique with dot- hot-products: We develop novel proof technique focused on analyzing the similarity score variables by the dot-products of attention layer. More concretely, we first show a convergence of the score variables to their limiting distributions, and then prove a conditional weak convergence of the outputs of attention layer. This analysis characterizes the output distribution of an attention layer by incorporating the intrinsic randomness from the dot-product.
1.1 Related Works
The concentration of measures in neural networks plays fundamental role in both theoretical analysis and practical applications of machine learning. Initially, studies in this area aimed to characterize the feedforward signal propagation in wide neural networks with random weights, inspired by statistical mechanics [Ama77] [ (SCS88). The outputs of such random neural networks converge to Gaussian process, and the computation of signal propagation reduces to the composition of kernel functions of Gaussian processes utilized in machine learning NNa96 Wi]96] DFS16 LBN* 18 dGMHR 18). They are often referred to as the NNGP. Since the NNGP kernel captures intrinsic inductive biases of architectures, its prediction performance correlates well with trained networks across different architectures [LBN* Moreover, one can interpret NNGP as a network with random initialization of optimization, which naturally leads quantitative insight into desirable weight scales to avoid exploding/vanishing signal and gradient problems [PLR+ 16 SGGSDI7).
Depending the network architecture, various types of kernels can be obtained, not only for fully-connected neural networks, but also convolutional neural networks (CNNs) IXL*119. skip connections [YS17]. naive or gated recurrent neural networks (RNNs) [CPSI8] Yan 19b) and more [YPR* 19 GHL.233 While some classical works assumed random Gaussian weights generated in an i..dd manner, recent research has shown that the same NNGP be derived even in networks with non- -on-Gausiian [GY22) or weakly correlated weights [ST224] The NNGP can also be obtained for network with narrow bottleneck layer sandwiched between wide layers [AP120)
Two pioneering studies have investigated the NNGP of self-attention layers (Yan HBSDN20)
Initially, [Yan pointed out that an unconventional scaling factor of 1/n in the softmax function

---

enables straightforward NNGP evaluation. Greg Yang has introduced the theoretical framework of Yan20a Yan20b], systematically composing kernel functions, including NNGPs, for modern neural networks. Transformers with the 1/n-scaling fall within the scope of the applicability of Tensor Programs. However, the 1/n-scaling is rarely used theoretical or practical contexts, leaving the more realistic 1 //iiscaliig unresolved. To attack this problem, [HBSDN20) analyzed self-attention by varying the number of heads. They numerically demonstrated non-Gaussian behavior emerging in the single-head case due to stochasticity and correlations within the attention matrix, even at infinite embedding dimensions. They further showed that if we take the infinite limit of the number of heads, Gaussian behavior emerges in the self-attention output and defines an NNGP kernel termed infinite attention. Although this infinite attention empirically improved performance certain NNGP regression benchmarks, its suitability a theoretical foundation for realistic self-attention remains uncertain. This is because practical attention implementations typically use only 1-128 heads EXW 24), far fewer than the embedding dimension.
Tensor Programs [Yan 9a
Beyond the classical NNGP and Tensor Programs analyses, several recent works have examined Transformers under the standard 1/ vn scaling and related asymptotic or dynamical regimes. [DYZ23 analyzed Transformers by tracking the first two moments (the kernel) of the signal to characterize propagation at initialization and during training. [CNQG24] applied mean-field theory to characterize the edge of chaos via forward and backward signal propagation, assuming Gaussianity of the QK product. [BCP24] employed dynamical mean field theory to study training dynamics under various infinite limits, including infinite width, heads, and depth, identifying parameterizations that ensure stable feature learning over time. [NLL*23) modeled signal evolution with a stochastic differential equation, which requires modifying the softmax function.
2 Preliminary
2.1 Notation and Setup of Neural Networks
We define the notation for a standard neural network and its usage. In what follows, we denote by n the dimensionality corresponding to the network's width. Although we can vary n across different layers or architectures, for simplicity we here treat every layer as having the same width n. A comprehensive summary of the notations used throughout this paper is provided Appendix A
Neural network We define notation for standard neural networks, excluding the attention mechanism. In particular, we adopt notations inspired by the framework of Tensor Programs [Yan19a Yan 19b), which allow us to describe a broad class of neural network architectures.
We describe an architecture of vectors h le.....,"' which is inductively generated by the following rule. We fix a nonempty subset " Vo c .il.....i ( called the set of initial vectors (input layer). For each index k with h² & Vo. the vector ht is generated either by matrix multiplication (MatMul) or by a coordinatewise nonlinearity (Nonlin). In the MatMul rule, given a weight matrix W R and some j # k, one sets h* Wh. In the Nonlin rule, given & & Jmm c [J] and function & Rm - R. one sets
neural networks as a finite set of Ra-uul random
h o(h .......""),.. h = (hhi...., I("") (a [n]).
This type of Tensor Program is called NETSOR, and it covers broad range of neural network architectures, including a perceptron layer, a convolutional layer, a recurrent layer, and many others
Yan 9a Y9n]
We present the perceptron layer as a specific neural network represented by NETSOR. Suppose x E R" is the input of the (-th layer. It generates pre-activation variable z E R" and the input of the next layer xf E <R"M by
== W= x ((2))
where /' R is weight matrix and & : R (potentially nonlinear) coordinatewise activation function. The vectors zf and x' are generated by MatMul and Nonlin, respectively.
Meati-head attention We define the attention layer. Let and H denote the spatial dimension and the number of heads, respectively. Suppose wQ.a W..aa W.v. wo.a e R weight
3

---

matrices for head a E [H]. With an input sequence of random vectors x .'.'.'''' R", we define Re as matrix whose i-th row x', i.e. X [x' r'). For each head a e [H]. define Qla) (WW..a) k(a) = X(Wk.a)T, via) = x(WV.a) so that Qla), K(a), vla) are Roxn matrices. The scaled dot-product score matrix is given by
(a) _ (x (Pij)) ].jjE(s] E R
Pijj = Vn -(WR.ax)) ((Wkka)))) (t).
(I)
Here, the common scaling 1/vn in the definition of Pi.j is key issue in this study row-wise softmax function to the score matrix, define
(a)
Applying the
A (a) (o RR*s,
A.e SoftMax (0) ................................ ................................ Pi.. p(
We then weight the values to form each head's output Head = (() E R***, Finally, we sum across all H heads to recover the full attention output MultiHead H Σ Head (WO.a)T (WO.a) Rsxa, where each row MultiHead,, is given by
H
(MultiHead,))
woaaw SoftMax,
y 1.... p( ER" (ie
(2)
VH
Weight initialization We set the weight matrices such as w', wQ.a wK.a, wVa, and w and initial vectors h Vo as follows:
(i) Each weight matrix W is independent. Each (a, element of W sampled i..d. from Was ~ (0,, /n), where ow > O is a constant that may depend W.
(ii) Let % (Z h E Vo} |/Vl be a multivariate normal distribution. For each a E [n]. the collection of a-th components of all initial vectors in Vo. denoted by {he h Vo), is sampled i.i.d. from ZoV,
This setup of weight matrices is common in practice (see GBIO for summary) and also identical to
that in the existing Tensor Programs Yan 19a Yan 19b Yan20b
2.2 Distribution of Attention Mechanism in Previous Setup
We discuss the existing challenges on characterizing the output distribution the attention layer, which exhibits -an-aassiinn behavior. To derive this distribution, regimes typically considered: the choice of scaling in the similarity computation, and the number of heads.
1/n-scaling regime Yan]9b] studies attention layer whose scaling term 1/vi in the scaled dot product of Eq. T is replaced by 1/n. In this regime, the Tensor Programs framework can show that the outputs of the attention layer converge to Gaussian distribution: as it holds that
MultiHead; N(O, x²) (ar [n]),
with some x > 0 (see Appendix and Theorem E.8 in Yan for details). Intuitively, in the 1 //n-scaling regime, the dot-product score p converges to for every pair (i. [s] and head
a E [H], which simplifies the non-Gaussian behavior of the attention outputs.
HBSDD20 considers the regime in which many dimensionalies diverge to infinite, then shows the convergence of the attention outputs to Gaussian. Specifically, it is shown that n, H co, it holds that
Infinite-heads regime
MultiHead: - N(O, (K))), (ar E [n]).
'Our analysis focuses on the 1/ 1/Vi-ical choice motivated by its prevalence in both practical Transformer implementations and the theoretical literature (e.g.. [HBSDN20] DYZ3 CNQG24). While some analyses have raised questions about its stability 2Y [CP24]. the precise conditions stable training remain active area of research and may depend on factors often abstracted away in simplified limits, such as the number of tokens or data-speca speciaic statistics. For instance, recent work has shown that data statistics can significantly alter optimal scaling rules (HL25).

---

with some x' > 0 (see Theorem 3 in [HBSDN20). The variance (x')² is described by a covariance (a)
of the nonlinearly transformed dot-product score Pi.j In the result, by letting the number of heads H grow infinite, the complex non-Gaussian effects from the dot-product score smoothed out, yielding convenient Gaussian distribution.
While these approximations are analytically appealing, they have notable limitations. In practice, Transformers typically use 1/vi-scaling (see Equation (1) in [VSP*17) and finite number of heads, resulting in the frequently observed non-Gaussian behaviors of attention layers. Although the above regimes simplify the behavior into a Gaussian form for tractability, a more precise theory is needed to capture the true behavior of attention layers in practice.
3 Main Theorem
3.1 Limiting Distribution
We introduce the limiting distribution for two types of quantities: variables generated within the NETSOR program, and the scalar dot-producce that arise in attention mechanisms. The convergence of these variables to their limiting distributions is formally established in our main result, Theorem 3.1
Definition 3.1 (Limiting Distribution). (A) Limiting distribution for vectors in the NETSOR program: For each vector h R" in the NETSOR program, there exist corresponding random variable Z as follows:
(i) If h is an initial vector from Vo, then Z" follows the distribution specified in Section 2.1 (ii) If g'.....g' .6. are generated by MatMul, then (Zx',. ..... Z=*) ) is zero-mean Gaussian vector. Specifically, for g' w' and g W their covariance given by
if W and W are different matrices), of B zdd ( w and W are the same matrix).
Cov(Zs', 22'' ,28')=
(iii) If h generated by Nonlin, h s(h² h²), then zh is defined as Zh 6( &(Z"' Z
(B) Limiting distribution for sealar dot-products: For each sealar p; (g' g² / where g' and g are outputs of MatMul within the NETSOR program, we define p, as follows:
(iv) (éi..... (,) a zero-mean Gaussian vector. This vector is statistically independent of all z" variables defined in (A). The covariance of (pI.....ph .r) ) is given by
Cov(pi. Px) B z z z
the formulations for the limiting random variables 2 associated with initial vectors, MatMul outputs, and Nonlin outputs are developed by the existing Tensor Programs framework [9b] Specifically, the limiting distribution of initial vectors follows that given in Section 2. variables generated by MatMul converge to Gaussian distribution; and variables generated by Nonlin nonlinear transformations of the corresponding limiting distributions.
In Definition 3.1
The distinct aspect of our analysis lies in the treatment of the limiting distribution for the scalar dot- products (p1...., py), which represent pre-softmax attention scores. While the marginal Gaussian distribution of these scores is consistent with that in the infinite-head limit (see Theorem in HBSDD20)). ). framework explicitly defines them Gaussian vector that is independent of the limiting distributions z within the NETSOR program. This independence facilitates the computation of the limiting distribution of the attention outputs, it allows to treat the attention scores separately from the other variables in the NETSOR program.
3.2 Statement
In this section, we show the convergence of the distribution of variables in the presence of an attention layer as the main result. Here, we refer to a function as pseudo-Lippshittz it is neudo-Lipschitz of order d for some d [2, oo). See Definition C.I for the definition of pseudo-Lipschin functions.
5

---

As preparation, we first introduce the results of the convergence of the NETSOR program without an attention layer. The following theorem is slight simplification of Theorem 5.4 Yan 19b ?
Fact 3.1 (NETSOR Master Theorem Wan19h)) Consider NETSOR program. Suppose all initial vectors and weight matrices are sampled as explained in Section 7 and all nonlinearities used in
Nonlin are aseudo-Lipschile For any positive integer k, let h ....,."." h be any vectors in the NETSOR program. Then, for ant pseudo- Lipoo-ippppchhiuneceece
v B - R. we have
Σ n
1
11a1 0.3 B[
B
z"
as Here, Z²' tha defined in Definition 3.1
(i) (i)-
Building on the NETSOR master theorem, we now present our main result. Before that, we introduce some assumptions specific the attention layer.
Assumption 3.1. Let and be positive integers that satisfy 2r. Suppose g'............ .'...'""" n R" vectors in the NETSOR program generated by MatMul. We assume that a subset (gh): : [r]. j [2]] C {g .....g" """) can be expressed as, without loss of generality
g = w
x = o''
(i [r). (2)),
where each ai is bounded and gveudo-Lipschitzfuneee The weight matrices Widd e R=xx satisfy two conditions:
(a) (WiJ :e [r], j [2]} is specific the generation of (g :ie [r], [2]) and not used for any E {g .s...... gm }
(g/- : [r], [r]. [2]]
(b) It is permissible for wi- to be the same matrix as W unless ='. # j' (e.g., WI.1 could be the same matrix as 1
Finally, we define the sealar dot dot-products {p; : [r]] by
(g (r VI
Theorem 3.1. Consider NETSOR program, and suppose all nonlinearities used in Nonlin are Eseuud-Lipschits. We adopt the settings and notations from Assumption which defines vectors g'.....g and scalar dot-products Pl.....Ph Further, suppose all initial vectors and weight matrices are sampled as explained in Section 2.1 Now, let h ,...h h²t R." be vectors whose elements given by
N/. = (gl ...
8 Pl....,pr)
pr) (ar [m], je [k]).
where each σ! is v : R - B. we have
Then, for any bounded and psenuo-LLipchhittffuctto
v (h (h........ ................................ nttt) d Blu(Z") ...., ...,,Z") PI.... ....... .p.l
as - Here, z is given by Zn 9 (Z = .I......p pr) for j [k], and (Zs', 7 .p..... Pr) ) is defined in Defimition 3.1
It may be noted that in the statement of Theorem 3.1 the boundedness of Hi as in Assumption 3.1 not essential, and we can drop this condition. However, we keep it for simplicity of the proof. Also, note that in this statement, PI...., Py are random variables with positive variance, and the convergence distribution " described in the statement refers to convergence the distribution of them. Unlike the conventional master theorem derived from Tensor Programs Yan Yan 19b),
19h]9hh considers controlled functions instead of pseudo-Lipschitz functions, the former of which is generalization of pseudo-Lipschitz functions.
*The Tensor Programs framework generates vectors inductively, precluding circular dependencies (see Section Consequently, for each n (g' !........"" .")... any arguments from {g fg'.....g" "......") that effectively contribute to the output xd (i.c., those upon which the function gf-1 actually depends) must be defined prior to g'> W in this inductive sequence.

---

this theorem simultaneously characterizes the randomness of the ordinary variables like g' e!......"" 8"" along with the effects of dineessieniii random variables PI.....Pr Pr.
Theorem 3.1 implies that, due to the randomness of (f1....., -r); the overall outputs resemble a hierarchical distribution, where the variance of one limiting distribution may depend realization of another limiting distribution (øI...., py). This result differs from existing results by the Tensor Programs where the variance depends on moments of the other distributions. This result leads to non-Gaussian limiting distribution in the attention case, whose details will be provided in Example 3.1
The following corollary is analogue of Theorem A.5 in Yan20b)
Corollary 3.2 (Coordinatewise Convergence). Assume the same premise in Theorem 3.1 Then, for all a 2 I, we have
(H'....,
1t - (zh)
We present specific multi-head attention results as an example of the application of Corollary 3.2 This example implies that the output of the attention mechanisms described by non-Gaussian distribution. Specifically, when y' is an output of a multi-head attention layer as MultiHead; defined in Eq. 22 (Z' ...... Z follows a hierarchical Gaussian distribution whose conditioning (fI..... P.) = (D.)) i. E [s], a e [H]} itself is a random variable, causing (Z" ..... ....."" 2") to follow a non-Gaussian distribution overall. We provide the details as follows:
Example 3.1 (Multi-Headl Attention). We consider the multi-head attention in Eq. that is the spatial dimension and H is the number of heads. We sample new weight matrices wQ.a, W..a, WV.a, W0.a E Rax for a [H]. Letx .et.x.......' E R" be vectors within the NETSOR program. We assume these input vectors are generated by Nonlin, where the nonlinearity is bounded and pseudo-Lipschite For j E [s]. a [H], we define the value vectors ya. = WV.xx' E RR (j [s]. E [#[] and its further transform pa.j = W.. aa e R.. Finally, with the score Pi.j in Eq. Q.. we rewrite an element of the multi-head head attention Eq. as
Recall
S ..... p/. va
o (va) je a [M]), (pLs) :ij E [s]. a (Mi).
by introducing functions &f e [s]), each of which is a pseudo-Lipschitz function Then, by Corollary 3.2 we have
(a.....ya d (2' .....,"') (n an--00). oo),
where z is defined by Zz 11 ΣN Σj SoftMax, (D p! y By Definition 3.1
( j [s(s), e [H]] and (p10) :i.j E [s], a e [H]] satisfy the following.
(zpmi j [s]. a [M] is independent of (p(9) :ij E [s], a [H]]. (zpas j E [s], e [H] is jointly Gaussian with zero mean. For a, d' E [M] and 1. [s], the covariance is given by
(a#a)),
02 @yyon
ww.. Blza' (a=
is jointly Gaussian with zero mean. For a, a' E [M]] and i,'',,j'' e [s], the covariance is given by
(P iij E [s]. a E [M]
(a # a'),
Cov(p!o)
(a) la')
)
02 oh
B("
Zr'
(a = d').
4See Proposition C.3 and Lemma
C.4 Note that SoftMax is Lipschitz so that it is also Secudo-Lipschit.
(Fact C.3
7

---

Note that 12* f E [s]} can be computed by Definition 3.1
The pseudocode for this example is presented in Algorithm in Appendix B.1
Remark 3.1. The non-Gaussianity the limiting distribution of the attention layer highlights significant challenge for theoretical analysis, since many existing frameworks, such as NNGP and Tensor Programs, are developed under regimes where the limiting distributions are Gaussian (or can be expressed Gaussian components with correction terms). Extending these tools to rigorously handle such non non- Gausiaa behaviors remains an important open problem, and we believe our work provides a concrete starting point for such future developments.
Remark 3.2. Our result in Example 3.1 has close relationship with the result of [HBSDN20] First, the limiting distribution of the attention scores in our framework perfectly matches the Gaussian distribution in their infinite-head limit. Second, the variance of our finite-head non- Gaussian output distribution theoretically matches the variance of their infinite-head Gaussian output distribution. This alignment suggests that while the infinite-head limit correctly captures the second-order statistics of the output, our finite-head analysis provides a more complete picture by characterizing the non-Gaussian aspects of the distribution that are not captured by second-order statistics alone.
Remark 3.3. We show that the limiting distribution z is sub-Gaussian provided that bounded nonlinearities are used for the attention scores (e.g., the softmax function as in 1xamplee..1p. In other words, although Z cannot be approximated by a single Gaussian distribution with mean (B(²") and variance Var Var(Z*') ), its tail probability decays Gaussian rate, possibly with larger variance parameter. A detailed discussion provided in Appendix
Proof Sketch of Theorem 3.1
We sketch the proof of Theorem 3.1 First, observe that it is equivalent to the following statement:
Theorem 4.1. Assume the same premise as Theorem 3.7 Then, for any bounded and pseudo- Supschitzfunction v,, we have
as n + 0o, where (Zs' ........ 7 .I..... pr) is defined in Definition 3.1
Clearly, Theorem 3.1 implies Theorem Conversely, applying Proposition
C.3 establishes the
opposite direction.
In what follows,
we outline the proof of Theorem
Hereafter, æ1.k denotes the vector ....)..,,)......
and likewise, xllk denotes the vector (x (.......
Fix bounded and pseudo-Lipschitz function v/ By Lemma C.2 it suffices to show that, for any bounded and Lipschitz function
x.))
Bf
Σ
v (ge." lon PIr)) Bf (Elu(Z .D.)) 1 (1 +
holds. Note that by Fact 3.1 have
(igh")) (79
for any pseudo-Lipschitz function C. Using this result, we first show that Pi: converges in distribution to Pir, which is Gaussian vector defined Definition 3.1 (Appendix D.I.Ip. Then, consider
Bf
v (gght 1am p(r)) -B/ Bf [B(u P pi.l)
s S + S2.
F
where Si and S2 are given by
-Bʃ (B [ (Z9' -Pir) pir ||
Si = B
-
1
Σ
v (ggu I.m P1rr)
a=l
S2 Bf [B
W(Za' .Pr)) per )) ( (BI PPr)) p(r))
We separately show Si - 0 (Appendix D.I.2 and So - 0 (Appendix D.1.3)

---

n-16
D
D-256
n=1024
log(KL)
-
o0 width
-6
3.5 4
4.5
0
log_ 4(m)
Distribution of y and Z
(b) KL divergence with bars.
Figure Comparison of the distribution of the attention output yt and its infinite-width limit 2> in Example 3.1 (a) Kernel density estimates of the empirical distribution (via Monte Carlo sampling) of y for widths E (16, 64, 256, 1024) (dashed lines) alongside that of z (solid line), showing the convergence of the finite-width distribution to its limit. (b) Average of the log-KI. divergence log (L(Dist(y))) ()Ilss(((" over 10 independent trials, plotted against (og,()) with error bars indicating one standard deviation, confirming a decreasing trend.
5 Simulation and Discussion
We perform simulations to validate the infinite-widd width limit distributions derived in Example 3.1 Details on the simulation setting can be found in Appendix B All simulation codes are available at https://gi http:/ggthub..oooo..... hatb.//ooommmaaaaaaaiiitiiii te- width-attentee
5.1 Effect of Finite Width
To assess how well the infinite-width theory of Theorem 3.1
aligns with finite-width multi-head attention behavior and to verify that discrepancies diminish width grows, we perform 10 independent experiments for each width E (16. 64, 256, 1024). In each trial, we draw samples of the multi- head output y with H
Example 3.1 estimate its density via kernel density estimation, and compute
the KI. divergence to its theoretical limit z the density of which is also approximated via Monte Carlo sampling. Figure M ()) plots the estimated densities our first trial, showing that the density of y converges rapidly to that of z as increases. Figure bb) quantifies this convergence by plotting the average log-KL KL divergence against (n), with error bars showing one standard deviation across the 10 trials, demonstrating a consistent decay with growing width.
5.2 Scalings of the Dot-product Score and Finiteness Heads
We next investigate facets of finite-width behavior in Example 3.11 which are the effect of different scaling rules on the variability of the dot-product score and the impact of finite number of heads on the attention output. For comparison, [an]99]] assumes 1/n-scaling, and the infinite-heat result by
HBSDN20 applies only in the limit n, H c0.
(1) Figure aa shows the histogram of the empirical distribution of the dot-product score Pi,1 at width 256 under the two scaling schemes 1/vn and 1/n, together with the infinite-width limit distribution ø1. Even at this moderate width, the 1 1/n-scaled scores are tightly concentrated around zero, whereas the 1/ scalled scores remain spread out in a nondegenerate fashion, confirming that only the latter preserves variability away from zero at large
Figure b) reports histograms of the attention output y for 256 with head counts H and H 256, overlaid with the infinite-widre densities of zx which are approximated via Monte Carlo sampling. The close agreement between finite-width histograms and their infinite-width limit curves, even when H grows on the same order as demonstrates the robustness of our infinite-width approximation in the growing number heads. Furthermore, is observed that as H increases, our non-Gaussian distribution approaches the Gaussian distribution from [HBSDN20)

---

n=256, 1/va- scalsciin
n-256, H=1
0.6
n=256, 1/0 scaling
256 H-256
width, 1/vh scaling
0.5
se width, H=1
width, 1/0- //a-scalin
0.4
width, H=256
se widthihhhh we-bead
0.3
0.2
0.1
0-
-2
2
Dot-product score under different scalings.
(b) Varying head count of the attention output.
Figure 2: Visualization the dot-product score Pi.1 and attention output as defined Example 3.1 comparing finite-width behavior to their infinite-width limits. (a) Histogram of the empirical distribution of Pi,1 for = 256 alongside the plot of its infinite-width limit distribution i under two scaling schemes; 1/va and 1/n. The 1/n-scaled score collapses to zero the infinite- width limit, while the 1/ s/cieee score retains nondegenerate distribution. (b) Kernel density estimates of the empirical distribution of y for π = 256 (dashed lines) alongside the plot of its infinite- width
limit distribution z (solid lines), varying head counts H (1, 256). The black solid line represents the density of the infinite-heat limit distribution from [HBSDN20). This demonstrates that our theoretical prediction remains accurate even when H grows, and approaches the infinite-head limit.
Beyond the choice of scaling, our work provides a general perspective on the number of attention heads: our theory provides accurate characterization for both finite and large numbers of heads, effectively subsuming the large-H regime within a single framework. This robustness is confirmed by our experiments, which show theoretical predictions remain highly accurate H grows large (Figure[bb), as well as in low-rank attention setting where the number of heads H increases proportionally with the network width n (Appendix B.2).
Conclusion
In this paper, we rigorously analyzed the infinite-widh limit distribution of outputs from a single attention layer using the Tensor Programs framework. Specifically, we theoretically showed and empirically confirmed that the attention outputs converge a non-Gaussian distribution under realistic conditions with finite heads and standard 1 1/Vi-scaling.
Looking forward, believe framework serve foundation for future extensions deep Transformer architectures. We predict that when layers are stacked as in Transformers, not only attention layers but also MLP layers converge non-Gaussian distributions in the infinite-width limit. The nan-Gausyyy of attention layers (and possibly MLP layers) is not merely a theoretical curiosity, but it has profound implications for the learning dynamics of iitention-based models. For instance, the presence of higher-order moments associated with such non- non-Gaussianity could influence signal propagation by preserving feature variability across layers, thereby reducing the risk of rank collapse. Furthermore, the anisotropy induced by non-Gaussianity may lead to more irregular curvature in the optimization landscape, possibly affecting the convergence properties of training dynamics (e.g., through sharper gradients more prominent saddle regions). Consequently, we argue that a new framework distinct from existing Gaussian- based approaches is essential for analyzing deep architectures. Our work provides an important first step in this direction.
As limitation of this study, our analysis assumes constant sequence length and batch size, following the Tensor Programs framework. Extending the framework to handle growing sequence lengths is an important direction for future research.
Acknowledgements
We are grateful the reviewers for their insightful comments. We also thank Yuta Matano for helpful discussions and for pointing out the sub- aub-Gaussianity of the limiting distribution of the attention
10

---

output. Mana Sakai was supported by RIKEN Junior Research Associate Program. Ryo Karakida was supported by JST FOREST (Grant No. JPMJFR226Q) and JSPS KAKENHI (Grant No. 22H05116). Masaaki Imaizumi was supported by JSPS KAKENHI (Grant No. 24K02904), JST CREST (Grant No. JPMICR21D2), and JST FOREST (Grant No. JPMJFR2161).
References
[AAA*23] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Flo- rencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.087 2023.
[Ama77] S.-I. Amari. Neural theory of association and concept-formation. Biological Cyber- netics, 2653):175-185, 1977.
[APH20] Devanshu Agrawal, Theodore Papamarkou, and Jacob Hinkle. Wide neural networks with bottlenecks are deep Gaussian processes. Journal of Machine Learning Research, 21(175): 2020.
[BCP24] Blake Bordelon, Hamza Chaudhry, and Cengiz Pehlevan. Infinite limits of multi- head transformer dynamics. In Advances in Neural Information Processing Systems, volume 37, pages 35824-35878. Curran Associates, Inc., 2024.
[BCRT58] J. R. Blum, H. Chernoff, M. Rosenblatt, and H. Teicher. Central limit theorems for interchangeable processes. Canadian Journal of Mathematics, 10:222-229, 1958.
[BM11] Mohsen Bayati and Andrea Montanari. The dynamics of message passing dense graphs, with applications to compressed sensing. IEEE Transactions on Information Theory, 57(2):764-775, 2011.
[CNQG24] Aditya Cowsik, Tamra Nebabu, Xiao-Liang Qi, and Surya Ganguli. Geometric dynamics of signal propagation predict trainability of transformers. arXiv preprint arXiv:2403.02579, 2024.
(CPS18) Minmin Chen, Jeffrey Pennington, and Samuel S. Schoenholz. Dynamical isometry and mean field theory of RNNs: Gating enables signal propagation in recurrent neural networks. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 873-882, PMLR, 2018.
[DCLTI9] Jacob Devlin, Ming-We Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4186 Minneapolis, Minnesota, 2019. Association for Computational Linguistics.
[DFS16] Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and dual view on expressivity. In Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
[dGMHR*18 18] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard F. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
[DYZ23] Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of Transformers at initialization. arXiv preprint ar04v: 02034, 2023.
EX] Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A. Alemi, Roman Novak, Peter J. Liu, Izzeddin Gur, Jascha Sohl- Sohl-Dickstein, Leslie Pack Kaelbling. Jaehoon Lee, and Jeffrey Pennington. Scaling exponents across parameterizations and optimizers. In Proceedings of the 4/st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 2266612700 PMLR, 2024.
11

---

[GB10]
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Yee Whye Teh and Mike Titterington, editors. Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume of Proceedings of Machine Learning Research, pages 249-256, Chia Laguna Resort, Sardinia, Italy, 2010. PMLR.
[GHLG23]
Tianxiang Gao, Xiaokai Huo, Hailiang Liu, and Hongyang Gao. Wide neural networks as Gaussian processes: Lessons from deep equilibrium models. In Advances in Neural Information Processing Systems, volume 36, pages 1-94555... Curran Associates, Inc., 2023.
[GY22] Eugene Golikov and Greg Yang. Non-Gaussian tensor programs. In Advances in Neural Information Processing Systems, volume 35, pages 21521-21533. Curran Associates, Inc., 2022.
[HBSDN20] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: NNGP and NTK for deep attention networks. In Proceedings of the 37th International Conference on Machine Learning, volume 1 199 of Proceedings of Machine Learning Research, pages 4376-4386. PMLR, 2020.
[HL25] Soufiane Hayou and Liyuan Liu. Optimal embedding learning rate in LLMs: The effect of vocabulary size. arXiv preprint arXiv:2506.1502 2025.
[JGHI8] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Con- vergence and generalization in neural networks. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[LBN*18] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S. Schoenholz, Jeffrey Penning- ton, and Jascha Sohl-Dickstein. Deep neural networks as Gaussian processes. In International Conference on Learning Representations, 2018.
[MMN18] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape two-layer neural networks. Proceedings of the National Academy of Sciences, 1 115(33) 2018.
[Nea96] Radford M. Neal. Bayesian Learning for Neural Networks, volume 118 of Lecture Notes in Statistics. Springer Science & Business Media, 1996.
[N]]*23] Lorenzo Noci, Chuning Li, Mufan (Bill) Li. Bobby He. Thomas Hofmann, Chris J Maddison, and Daniel M. Roy. The shaped Transformer: Attention models in the infinite depth-and-winn limit. In Advances in Neural Information Processing Systems, volume 36, pages 19250-54288 Curran Associates, Inc., 2023.
[9] Roman Novak. Lechao Xiao, Jachoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Sohl-Dicksteni Bayesian deep convolutional networks with many channels are Gaussian processes. In International Conference on Learning Representations, 2019.
[PLR+16] 16] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In Advances Neural Information Processing Systems, volume 29, Curran Associates. Inc., 2016.
[SCS88] H. Sompolinsky, A. Crisanti, and H. J. Sommers. Chaos in random neural networks. Physical Review Letters, 61:259-262, 1988.
[SGGSD17] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.
[SNT24] Thiziri Nait Saada, Alireza Naderi, and Jared Tanner. Beyond IID weights: sparse and low-rank deep neural networks are also Gaussian processes. In The Twelfth International Conference on Learning Representations, 2024.
12

---

[Vaa98] A. W. van der Vaart. Asymptotic Statistics. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 1998.
[VSP*17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[Wi196]
Christopher Williams. Computing with infinite networks. Advances Neural Information Processing Systems, volume 9. MIT Press, 1996.
[Yan19a] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04766 2019.
[Yan19b] Greg Yang. Tensor programs i: Wide feedforward recurrent neural networks of any architecture are Gaussian processes. Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[Yan20a] Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
[Yan20b] Greg Yang. Tensor programs ii:: Neural matrix laws. arXiv preprint arXiv:2009. 2020.
[YH21] Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite- width neural networks. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11727-11737. PMLR, 2021.
[YHB*21] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In Advances in Neural Information Processing Systems, volume 34, pages 17084 17094-17097. Curran Associates, Inc., 2021.
[YL.21] Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11762-11772. PMLR, 2021.
[YL23] Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit. arXiv preprint arXiv:2308.018/1 2023.
[9PR*19] Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Gohl-Dickstein, and Samuel S. Schoenholz. A mean field theory of batch normalization. International Conference Learning Representations, 2019.
[YS17) Greg Yang and Samuel S. Schoenholz. Mean field residual networks: On the edge of chaos. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[YYZH24] Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning infinite depth neural networks. In The Twelfth International Conference on Learning Representations, 2024.
[Z] Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yinggian Min. Beichen Zhang, Junjie Zhang, Zican Dong, et al. A survey of large language models. arXiv preprint 203.v:2303.18223 2023.
13

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We clearly stated the main claim in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize to other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We stated the limitation in the last section for discussion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should diseuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

We have described the full assumption, statements, and proof in the main body and the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We have described the detailed experimental details in the corresponding section in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but should possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code. with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We will open the source code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurIPS code and data submission guidelines (https://nips. hhtpss//ijs..c/publc/uuddde getp///cgguiee///eeeeeee
for more details. The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
SondeSSumissionoo.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We have clearly described the experimental details in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We have repeated the experimental multiple times and report the standard deviation comes from the replication.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
16

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [No]
Justification: Our experiments are small-scale and implementable by small laptop. Also, we do not pursue the computational cost in this study, so the computational resource out of our focus.
Guidelines:
The answer NA that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurIPS Code of Ethics https:
Answer: [Yes]
Justification:
We have checked the code.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consideration due laws or regulations in their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [No]
Justification; A main focus of this study is fundamental, so there almost no effect on social impacts.
Guidelines:
The answer NA means that there is societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address socictal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification:
Since this paper is theoretical, the outcome does not have a high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
We have not used any existing assets.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswithcode com/ datatsts has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
18

---

For existing datasets that are re- ee-pakagedd both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out the asset's creators.
13. New assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer:
[No]
Justification:
We have not created a new asset throughout this study.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/codemmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [No]
Justification:
We have not performed the crowdsourcing experiments and others.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [No]
Justification: Since this study is fundamental, there no potential risk on this point.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
19

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [No]
Justification:
We have used LLM only for the formatting purposes.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to LLM policy for what should or should not be described.
20

---

A Notation Summary
This section summarizes the key notations. We adopt two main conventions for simplicity. First, for finttewidth vectors and matrices, superscripts serve as indices to distinguish variables (e.g.. x' E R", W R=x"), while subscripts denote their components (e.g., x/a W E R). Second, the aB
dependence on the network width n for all finite-width variables kept implicit (e.g., we write xf instead of x' (n)). For other notations, please refer to Table
Table 1: General Notations
Symbol
Description
The dimensionality of vector spaces, corresponding the network's width.
The number of heads in the multi-head attention mechanism.
The spatial dimension of input sequences (i.e., sequence length).
W
A generic weight matrix in R"x0 with elements Wap N(O, w//)..
2
A random variable for the infinite-width limit of vector h.
A random variable for the infinite-width limit scalar dot-product p.
61:k
The vector (X), . )
x l:k
The vector xx
B
Simulation Details and Supplementary Analysis
B.1 General Experimental Setup
Unless otherwise noted, the simulations presented in this paper set the spatial dimension to s = 4. The core experimental setup follows that described in Example 3.1 which is outlined Algorithm
Algorithm 1 Multi- Head Attention (Example 3.1
Input: (x) hels)
R" input vectors for sequence of length & Rx weight matrices for H heads
Input: (W2.a, W..a, wV.a, wo.a )o.a)ae(H) ]ue[H]
for a E [H] do
for [s] do
aa < w
MatMul: Query vectors
aa W
MatMul: Key vectors
pa + W..y
MatMul: Value vectors
gad wo.upaii
MatMul: Output value vectors
end for
end for
for a E [H] do
for i E [s] do
for j [s] do
Pi.j -
(a)
Scaled dot- dot-product scores
end for
end for
end for
for i E [s] do
(a) y' - - IH Σソ=1 (p/ ..... . pa.j
Attention output vectors
end for
Output: (5 (r")etts]
> R² output vectors
21

---

Specifically, let let ,..,x"...,x" be outputs of Nonlin. These are obtained by applying clipping activation function & to preceding vectors h', ..... ...,, ht R.:
x = = -C -C + Ha s C) c1+ > (a [m], i [s]), (4)
where 1(-} denotes the indicator function and a positive constant. The vectors h!...., / are outputs of MatMul, defined by
W w' [
Each element of the initial vector & R" sampled independently from standard normal distribution.
For all weight matrices involved in the attention mechanism we.a, w.k., wV.a, w0.a and the matrices W generating x', set
wo.. o WK.a = o o o -o... = - W =.
The elements of these weight matrices are independently sampled from N(0, ww/)), where oW is the respective variance (here, l).
Under this setup, the input vectors x' to the attention layer are designed such that their infinite-width limits zx/ uncorrelated for # i.e.,
B[Z zM = (i.1' [s]. j j').
Furthermore, the indinite-width limit zd is standard normal random variable, leading
2 (1- (C)((C 24((C-
where C' is the constant used in the clipping activation, and & and & are the cumulative distribution function and probability density function the standard normal distribution, respectively. As C - co, this value converges to 1. In our experiments, we set C-= = 100, and we approximate B[(2*' ) by 1. Consequently, the covariances of the limiting variables for the vectors ga.J and dot-product scores Pi.j simplify as described Example 3.1
(a))
Co 50.1 bu [B( (a=d',
(otherwise),
and
Cov(A - =1 aa= j j'),,
(otherwise).
To estimate the empirical distributions of finite-width attention outputs and their corresponding widthh limits, we employ Monte Carlo sampling. For each such estimation, 50, 000 samples are drawn, unless otherwise noted. Kernel density estimation (KDE) is used to visualize these empirical distributions.
B.2
Analysis of Low-Rank Attention
B.2.1 Specific Setup for Low-Rank Attention
In practice, large-scale Transformers typically assume a specific embedding dimensionality for multi-head self-attention layers. For head counts H, the embedding dimension set linearly as Нии, where denotes the head dimension and determines the sizes of weight matrices as w0.a W.Kaa V.. e *H Thus, the QK product becomes low-rank relative the embedding dimension n, and the scaling factor is given by 1/ 1//HH follows:
( k. (i.j E [s], a [H]).
Pi.
For example, the original Transformer architecture sets H = 8 and aи = 64. Large-scale models often increase the number of heads be on the order of the hidden embedding dimension (EXW+24) seen in GPT-3 (175B parameters), which sets H 96 and nu 128.
22

---

n-64, H=1
0.6
252 H=4
-5.5
0.5
n-20 H-16
-6
log(KL)
width, H=1
de width, H=4
-6.5
o width, H=16
-7
0.2
-7.5
0.1
0
1.5
-3
log 4(n)
Distribution of 5 and Z
(b) KL divergence with bars.
Figure 3: Comparison of the distribution of the attention output y and its infinite-width limit z under the low-rank setting. (a) Kernel density estimates of the empirical distribution (via Monte Carlo sampling) of y for various widths n and head counts H (with Hy n H 64 is fixed, dashed lines) alongside that of Z (solid lines). (b) Average of the log- KL divergence log KL(Dist(y:) Diiss(()))))))""))))) over independent trials, plotted against log, (n) with error bars indicating one standard deviation.
Additionally, since the output from each head is also of dimension ny through the value matrix, an output weight a0,a E n is applied to map it back to the n-dimensional input for the subsequent layer:
SoftMax,
(a) ..... Pi,s wo.w (i [s]).
Note that, to ensure the dot-proe dot-product scores and the attention outputs are of order 1, the weight matrices are randomly initialized with the following scales:
aB w...a W.. N(0, ow/n). aB ~ -N(0,o (a, B (n)).
For simplicity, ow 1.
B.2.2 Results and Discussion
Finally, we investigate the behavior of the attention output y' in the low-rank regime described above. where the number of heads H increases proportionally with n. Fixing the head dimension n/H = 64, we perform 10 independent experiments for each (n, H) {(64, 1). (256, 4), (1024, 16)). Figure 3 shows the estimated densities of our first trial, and Figure plots the average log-KI. divergence between the distribution of y and the corresponding infinite-width limit distribution, with error bars showing standard deviation across the 10 trials. These figures show the convergence to the infinite-width limit as n and H increase proportionally.
Notably, even in these practically relevant settings employing low-rank attention (where head- specific projections are лt n or X ny rather than the X matrices primarily assumed in our formal derivations in Theorem 3.I.. our infinite-width framework continues to provide an excellent approximation. This agreement suggests that the core principles of convergence captured by our theory extend robustly to common architectural variants like low-rank attention (with appropriate scaling considerations). underscoring the practical utility of our theoretical predictions under these structural assumptions common in modern attention models.
B.3 Additional Experiments on Robustness
To further validate the robustness of our theoretical predictions, conducted additional experiments by varying key hyperparameters. These experiments investigate the impact of the spatial dimension (i.e., the number of tokens) and the choice of activation function.
B.3.1 Varying the Spatial Dimension
were conducted with a spatial dimension of Here, we present the results for increased spatial dimension of s 8. All other experimental settings remain
Our main experiments in Section 5
23

---

n-256
0.6
log(KL)
n=1024
width
-6
0.2
3.5
4.5
-3
limit 2 with log, 4(m)
Distribution of 5 and Z
(b) KL divergence with bars.
Figure 4: Comparison of the distribution of the attention output y and its infinite-width limit z when 8. (a) Kernel density estimates of the empirical distribution (via Monte Carlo sampling) of y for widths n E (16, 64, 256, 1024] (dashed lines) alongside that of Z (solid line). (b) Average of the log-KL divergence log KL.(Dist(y)) (y ([(Dist(()DDis(()))))))))) |[Dist(Z>")) over 10 independent trials, plotted against (n) with error bars indicating one standard deviation.
B=16
n=64
n=256
0.6
n=1024
log(K))
-6
widtth
-8
2
2.3
3.5
4
4.5
-3
0
2
3
log 4(m)
(a) Distribution of y and 2²
(b) KL divergence with error bars.
Figure 5: Comparison of the distribution of the attention output y and its
ReLU activation function. (a) Kernel density estimates of the empirical distribution (via Monte Carlo sampling) of y for various widths и and head counts H (with пи n H = 64 is fixed, dashed lines) alongside that of Z (solid lines). (b) Average the log-KI divergence log (L(Dis((y))) ()iss(Zz')) over 10 independent trials, plotted against log,(n) (n) with error bars indicating one standard deviation.
identical to those described in Appendix B.1 The results, shown in Figure demonstrates that our theory remains accurate even when the number of tokens is changed. This suggests that the convergence to the theoretical limit robust changes in the sequence length.
B.3.2 Varying the Activation Function
The experiments the main text utilize a clipping activation function. To ensure our findings are not specific to this choice, we also perform experiments with the ReI.U activation function, i.e., & ((h))) W2, v in Eq. 4 which is widely used in modern neural networks 5 The covariances the limiting variables for the vectors pa.j and the dot- dot-product scores p!.) are
Cov 2
1a = d'}
+
1( = 1r)
2x
2x
and
Cov(pl))ppo))) 1(a aa}
-
1(1
1(i=/))
1
10-1
1 2n
10
2T
2x
2x
2
*We note that the ReLU function is not bounded and consequently, our theory does not directly apply. However, mentioned in the main text, the boundedness assumption not essential theory.

---

For each experiment, we draw 100,000 samples for both the finite-width attention output and its corresponding infinite-width limit, an increase from 50,000, The experimental setup otherwise identical to that described in Appendix B.1
Figure 5 confirms strong agreement between the empirical distributions and theoretical predictions. This indicates that our theory is robust to this change in activation function, strengthening its applicability broader range of practical model architectures.
C Mathematical Tools
C.1
Basics
Lemma C.1. For S m and a)...... ........e ag B. we have
lail
aai
Proof. The first inequality is an application of the triangle inequality. The second inequality follows from Jensen's inequality. Since m the function + x on [0, oo) is convex, and thus Jensen's ss inequality implies
Σ
ail
= km
ail
S k
a th m
aa m
desired.
Lemma C.2 (Portmanteau lemma (Lemma 2.2 in Vaa98 )) The following conditions are equivalent.
(i) Xn 4 X.
(ii) B[f(X,)) B[f(X)) for all bounded and continuous function f.
(iii) B[S(X,)] - B[(())) for all bounded and sipschitz,fuctticn f-
(iv) P(X, E B) P(X e B) for all Borel sets with P(X 6B) = O, where oB denotes the boundary of B.
Fact C.1. Suppose (KnIeeN is sequence of integrable random variables that converges in probability
X. Then the following statements are equivalent.
(i) The sequence (XnhneN is uniformly integrable.
(ii) B((X,)) + B(Xx)) < co.
Remark C.1. If X, converges to probability, then by Fact C.I we have
B(Ix,) = o(1) - (XuhneN is uniformly integrable.
Moreover, since |B(X,)) & B(Ix,)) D, follows that B(X,) = o(1).
Fact C.2. Suppose there exists 6 > 1 such that sup,, B((x,1) co. Then the sequence (Xn Enen is uniformly integrable.
C.2 Pseudo-Lipschitz Functions
Definition C.1 (Pseudo-Lipschie functions [BM1)). Let > 1 be constant. A function ( R² - R pseudo-Lipschitz of order if there exists a constant C > such that, for x, R
lf(x) f(y)l C|xx yl + ||x||_-I + -1))
holds.
Fact C.3. The following statements hold.
(i) A Lipschitz funcciot is pseudo-Lipschitz of order d for all d > I.
25

---

(ii) A pseudo- Lipschita,uuctii (of any given order) is continuous.
In this paper, we refer function pseudo-Lipschitz if it is pseudo-Lipschitz of order for some E o2).
Proposition C.3. Suppose f: B R. and 8i RR - [k]) are Lipschtt. Then the function h R R defined by h(x)
[x)x)) is also pseulo-Lipschitz.
Proof. Suppose pseudo-Lipschitz. of order do + and each gi ppeudo-Lipschitzz of order di + Define function R - RR and constant d 2 1 by
gx) (xx))).........)
max .d...
Applying the pseudo-Lipschin bounds for f and each gi gives
- ((x))-((x') llg(x) g(x'Il [1 (lx(x)) (1))))
and
- -x(x))_:(xx)) & ||x - x'1l (1 (+ |||||| + 1 s ||x - x'll ||||| + -1x))1))
The last inequality implies
1/2
(x 8(x')1 = (8((x) -((())-s((()))))
||x - x''xx11 (1 ||]|| + I1x'11))
On the other hand, since
k
llg())11 Σ 1xx(x) =
Σ
[[:(())
holds in general, Lemma C.I implies
& s
du
Σ
(x)
$
The pseudo-Lipschitz property of each gi yields
[gr(x)) s (0)) (x) (0) $ ||x]|(1 |xx]|*) s + ||]|| $ ||x||
Hence, by Lemma C.I we have
[| $ +1+ d s 1 [|x]]
Combining these elements gives
[/(x) - h(x'l s 1|x x'll 1 ||*|| + -1xx1(0) (1 (1+ ||x||| dbld + ||xx (166dd)l))
We expand the product as
(1 + ||x||d + (Ix)) + d(d+t) + ||x (1x((1d(()))))
= =1 ||x||" ||x''| (od+)) + ||x' ([ + (lx)))) ||x | (ddd
+ ||' + ||x [1
Observe that each of the first seven terms are bounded by 1 ||x|| + ||x'|1 where a is given by d do(d + 1).
For the remaining two terms, we apply the weighted AM-GM inequality to obtain
|| ||x [=ooddo()+)) = (1Ix119) (11x' dutdel)
do(d 1) 11'11" s ||x|| + 1
I 5 ||x|||
The same bound applies to ||x || [x][||bldd))
Therefore the entire product satisfies
(1 + ||=||_ || ) + d[=])d)(++1)) + [[x''ldhid+t)) s + ||x||| + ||x 11a
which implies that h is Lscudo-Lipschitz of order
26

---

Lemma C.4. Define f - f(x,y) xy. Then f is pseudo-Lipschitz of order d for every (0)
Proof. For (x', y²) E RR? we have
lf(x,) y) - f(x - xxlly
(x', yx(x))] ly').
Observe that for any d 2 2, we have
|x| & (x y) = 1 + (l ||I, and consequently, we have
by |<'1,>> s 1+ II*x,
xx| ly'l $ 1+ 11x, |1(x',
This gives us
- (('.)) & (1(x,)) y2 [1xx, 11xx, + [1(x', )(11-)))
which shows that is Epeudo-Liisschit of order d.
D Remaining proofss
D.1 Proof of Theorem
In this section provide detailed proofs for the results sketched in Section 4 thereby completing the proof of Theorem 4.1
Throughout, B[- X] denotes the conditional expectation with respect to the o-algebra (X). Since conditional expectations are only defined up to almost-sure equality, we omit "a.s." when writing .a.5 in this context.
D.1.1 Weak Convergence of the Dot Products
In this section we prove the following proposition.
Proposition D.1. Under the assumptions of Theorem 4.1 the vector (Pi....,Py pr) converges in distribution to (p1..... pr), which is the Gaussian vector defined in Definition 3.7
The proof of Proposition D.I relies the next lemma, which is application of Theorem 2 in BCRT58).
Lemma D.2. For each n E N. let (Xa)ae(n) be an exchangeable sequence of random variables satisfying
B[Xa] =0,
B|X-] = o - (n oo).
Ser S = Σ Xa/ Vn. Assume the following conditions:
(a) X2) o(1/)).
(b) lim B(X?x3) a..
(c) 3(X ol Vn).
Then, S converges in distribution to where the random variable Z satisfies
25 if =020
z N=N00,03)
otherwise.
We introduce the following notation. For any two matrices Wi. and w". appearing in the program, define d ((')))) JJ) by
(1))
(if Wis and wi. are the same matrices),
((-)
(i.)) =10 (otherwise).
27

---

is important to note that (('('); is always deterministic value that is independent of nt and is fixed by the program architecture. According to the sampling rule explained in Section 2.1 the matrices WI. and w'' are sampled independently whenever d (f'''' is zero. In particular, Assumption 3.1 gives
(i.f)
(2.) (1.1)
0 (i [r]).
(5)
Let I, be arbitrary constants. We define
S
Σ
IPI
W wha di 12
1 Vn a=1
Σ Xa.
1=1
E=1
@=1 y) y==
where Xa is given by
Xa =
Σ i=1 71-72=1
Will was 2..1 1
ay Wayy+11 - "2
show that the sequence (Xa) [e[]] satisfies the conditions of Lemma D.2 The proof of Proposition D.1 is completed by applying Lemma to S - X /vi and then invoking the Cramér- Wold device.
Lemmas D.3 D.8
Lemma D.3. The sequence (Xataefm) is exchangeable.
Proof. By the sampling rule, an element w of the random matrix Wi- independently and identically aB
follows N(0, of wws/n) for a. E [n]. Hence, conditional on {xtd : e [r]. E [2]). the random variables XI...., Xn are i.i.d. Hence, by de Finetti's 's theorem, [X is exchangeable.
Lemma D.4. B(Xa) = 0 holds for every a [n].
Proof. We compute B(X,)) Eim Ey 1, (w()), B (w12.) B (1h1 O.
Lemma D.5. We have lim B(X²) = ?f? with
o, = Σ
78 7 22 7 73
li B
Σ
Σ
wr, 1z zebu
E 2 zulr
2 (bw) ( (H,s") (1,1) ((.1)((1)) (11.2)
11.2==1 (f.F)e)
where we defined J ((1,2), (2. 1)).
Proof. For any a E [n], we have
B(X²)
=E
Σ
in wi.l 2hi2 12
Σ
Wh.t 2 12.2 22.2 L
ay ay 71 2m
ar ax y3 2.4
71-7=1
73.741
Σ
E Who whi 2
E ( 2 12.1 (222))
crys /3364
Σ
Σ
1
(w B (w
B
(x
xxi xx2 x
((oo)) ((aa') (1)22)
= Σ
Σ
Wi. 12 BB
1 xty x
Σ
4.hel (JJ))e)
Σ
Σ
Whn o2
B Izzi zaz2
B
1 2
الاn "(ir')) (1.1) (h.2)
28

---

where the convergence follows from Lemma
D.9
Finally, Eq. ( and Definition 3.1 imply that (l22)) (l27") (6,1) (6,2)
oyr. o w4, 1 z
28 28" 28 28
B
holds for any 11.12 E [r] and (j.j'
Lemma D.6. XxX 0 holds for every a.ß [n]. * B.
Proof. For a # B. we have
B(X XB) B
Σ
Σ
Tn Why
x22
Σ
Σ
Fo Wht
Bys 2 yys
ay
|1 y1.72=
h= y3.y41
Σ
Σ
In
W
Bya
B
bh.2
74=1
0
as desired.
Lemma D.7. lim,- B(X o holds for every a. [n], a # B.
Proof. By calculation similar to Lemma D.5 we have
B(X,
Σ
Σ
Wh.1 1 Wis. wis (anl)
22
o (lani) (1255) (latt) (inn)))
((r.) (ir2) (a.1) (1.2)
Bwni 22 X22
ta Lant 1yg Ings
Σ x
ia.h 2
Xys
n1
n=
x4=1
n=0000
Σ
Σ
Iii 11 Wij, o
((z.))) ((b-5)) glant) (lil (ii1) (i.2)
Wh.,
W
4=
BB 2
B
where the convergence follows from Lemma D.9
Observe that this limit equivalent to o.
Lemma D.8. 3([Xa13) ol vn) holds as n - oo for every a [n].
Proof. By the Lyapunov inequality, we have
((xtt)
1
sup B(X
Vn
Thus, it suffices show that sup,, B(X holds. We can express B B(X as
B(X²) = Σ Σ
whah ay dry ays aya ays who why wis
E
N=11
B
Ays A x x/x 2
h..
Applying the Cauchy-Schwar inequality yields
E
4. 1 fo.1 in in 2.2 in.2 2i.22 Xyy
72
Xys
x
(1)):118)
[ E [a(t)t)8]
1 A=1
(o
[
E
E
A=1
k=l
B (xf.1)8)
<
sup
/e[r]. JE[2]
[E[2]
29

---

By the boundedness of xf1 (i E [r], j E [2]) the last term is bounded uniformly in n, and consequently, holds that
sup
B
n A22 4
di.1 12.1 fa.1 21
474 26
24.2 1 t98
< 00.
Hence, we compute
B(Xx) s
Σ
Σ
B
Wind wiy ways ay ay2
w
w why? wi.2 ays
ary 0%
41
Σ
Whl
Wils Wilh Wi?s wiig wh? w
Yay
Σ
B
ay ays ary ara ays aye ary aass
wii w
11 byii M
Σ
whe2 ary ary crys whyt arys aya ary whya]
3
B
Σ
What wint ay whyl way 2 aa wayz w
6
B
1= fy (e(())))) Ys.....y=
n
Σ
way May Wind Wind wayt wind wi 2 w way an
B
i d2m P
all distinct
=: A 4A 3A3+ + 6A4 + As.
Define or by o= =
:E [r], j E [2]]. Then, we compute A as
A
E Will wisl will wi.l ar. ays aayl E [w12 w was wa? ars aas
1= M
B ((wi) ay
)'] E [(W??) )²]
1= >1.221
+3
Σ
(wat)) y+
(win) )²|
(Wiy))
(JJ')eJ 71,22= y2
9
Σ Σ Σ
Σ
B
[(wi))
an )
[wii]] arys )
B [w ays 2
B [ any )
1=1 71.72=1 yi 74
904 04
1) wi 0 3(n
(n 1)20
Wr.1 04 n²
Wr,1 WL3
WiS
wia
3
+9
n
n²
(.FP)J
Wi2 $ 8
Wi.1
Likewise, we have
A2 Σ Σ
Σ
B
Way wind wht/ B
w wh, 2 arys
(1,1) (i.2)
1 bith (J. 78=
9 Σ Σ
Σ
(1.11) (1.2) s 8.8
(i.1)
GA (J-')e
30

---

and
A3
Σ
Whit wiit why Wayl E
w wh whyl
ars am ay Way d (inddi (1.2)
ay
ays
i (J-P)ed M
= Σ Σ 1= Ge U-f"ee
04 wi Wii Wh an (lon glnr1) 0..
(1,1) (112)
Applying a similar argument, we can also show that A4 s 0 and As s R holds. Therefore, we conclude that B(X+) $ o holds, which completes the proof.
Lemma D9. Take ki.....k Kg E :iE [r], j E [2]} arbitrarily, Then, the following statements hold as co.
(i) For (x*1
x*), we have
Σ
72
y2
yi
11
y==
92m1
2
B
2 %
B
(ii) For (x (xb.....,.*)))
xt), we have
B
Σ
Z E
z 42
B
7--l
Proof. Define the residual term Re by
1
12 k2
E 2112*00 x
Rτ
Xye
Xyc
for each / E [4]. Define constants and R by
C max [c]4]
B
Z
max B (R [c|4]
Note that is bounded by the boundedness of x for all i [r] and j E [2] (see Definition 3.1 Then, write
k2y
kit Nye
22 -1
2/
Nye
n
y==
1 $24 2 124
-l 536
E
Rs
+ B
[=1
R
R²
4CR* Rd.
where the last inequality follows from the Likewise, have
inequality and the Lyapunov inequality.
2
2
1/ -1
Xye
Xyy
Z
Z
y==1
2
2
12p-1 25 124
E
Re + B
2
[=1
2CR R2
31

---

Thus, remains only to prove that R converges to 0 as π - co. This can be achieved by showing that B(RR) converges to 0 for all & E [4]. By Fact and Lemma 1.m for all & [4], we know R converges almost surely to O. The continuous mapping theorem then yields
To upgrade this to convergence in expectation, Facts C.1 and C.2 imply it suffices to show the existence
of constant 1 that satisfies
sup B ( < co.
But since each x and its infinite-wind width limit nnd are bounded, such a exists. Therefore, we conclude that B(Rz) converges to 0 for all e e [4], and consequently, R does well.
D.1.2 Si Converges to 0
We study the convergence of the term Si defined in Section 4 Specifically, show
Si Bf Σ v (gg.M pir))
Bf
B v (Z
PPrr)
Pir )) O.
Fix a small O. Since Pır is a Gaussian vector (see Definition compact set K c R that satisfies
it is tight. Hence there exists
P(Pr: K) 1-E
Set
1 y
Σ a=l
v ( LM Pir), ==B v ( .Pir)) |Pr
Then, we have
S = BB(
f()(f()) K}]1+ [B((()) f()) K}]l
- ((((()))((())1pp)) K}]1 + +[]f(y)) - f())l(pi>> $ K)].
(6)
By the boundedness of f. the second term of Eq. 6 is bounded as
{}} S CP(piz $ K)
with some constant Furthermore, noting that K Borel set, applying Lemma C.2 Proposition D.I giveives
and
gives
lim P(P E K) P(Pr e K) <E n-000
Therefore, for large enough
we have
B(lf(y) [()) E K}] < Ce.
For the first term of Eq. (6) we have
[[(((y)
(Pz)
E K]]1
a
v ( l.M )
-B - M (Z9
(Z9 PM .air)
sup E
f
a
72
sup Φ,
a
where P R defined by
v (g. (1w)
-B -B(5 (B 1BB4(22 ( .ai>)
Φr =
Define Vay. M - R by
J Jayn (uI:M) ((uuM.,.))))
32

---

Observe that Ja., is bounded by the boundedness of v. We later show that Van. is also pseudo- Lipschitz (Lemma D.10). Therefore, by Eq. 3 and Lemma C.2 we have
lim (arr)) 0 (a R').
Moreover, noting that f and are bounded and continuous (see Fact C.3)
apply the
(uncoditional and conditional) bounded convergence theorem to show that + is continuous at every point. Therefore, can apply Lemma D.11 to show that sup EK Φn ((rr)) converges 0 as
n 00.
Lemma D.10. Define var, RRM R. Eq. Then, is pseudo-Lipschitz
Proof. Suppose us is pseudo-Lipschin of order d + with d 1. Then, we have (ur.m) Bai, (M(MM) I ( a1r) - // (u air)) an.) ((LLM,
x (u( (u( + (I
( a1r) 114).
By Lemma C.I
we can bound ((01::..a=))))
[I(um: a1:) s (I((um))l [|a1.)1)* [|(rr) 2d-1 (I(())))]] + 11a1:111)..
Thus, have
1 + (1r))l [l( a1 z 2d-1 (I(()))))) (K +ICu) 2dd ||rrll
$ [(uu)))) [1(u'r)))"
where the last inequality holds because air is fixed.
Lemma D.11. Let K c R be compact set. Suppose for each n E N, fn R - R is continuous
function that satisfies
lim (a (air)== =0
(8)
n
for any constant E K. Then, we have
lim sup [f((ar)) =0. n an eK
Proof. Fix é > O. Let B(arr= e) denote the ball of radius centered at air E R". By the continuity of fn. for each dii K. there exists 60ir such that
(f(()))) fa(b e/2 (bir B(a1:r66a.))
holds. Note that K is covered by the union of balls as
Kc U B(air. Sarr).
QKiEK
Since is compact, we can cover K with finitely many balls, say
Ü B, B(a dd..)
where B, is given by B; B(a ). By Eq. 8 for each i E [/], there exists N; e N such that
Is (a ( </22 (n 2 Nj)
holds.
Let N denote the maximum of (Nr : [4]]. Then, for any N, have
sup [f(()))) z max sup [f((ar=)] N max sup (Ifa(ar))) fa (a(a),))) (a.,)) iE[/] ei
air K
iE[/] cB
max iE[/]
which implies the convergence of supai,ek ek (f(:r) I to zero goes to infinity.
33

---

D.1.3 S2 Converges to 0
We study the convergence of the term S2 also defined in Section 4 Specifically, we prove
$ = Bf E (Z
.P:r) ptr) Bf [[o LM .Pr)) piz) | - O.
Define function V: R - R by
P(arr) = (( .nr)).
then it bounded since / is bounded. By the bounded convergence theorem, it is also continuous. Therefore, by the continuous mapping theorem, have
P(P..) - P(P..).
Since Zo MM
is independent of pin, have
E (Z9 TM .P:.)) p1:
= :((p:))
Therefore, have
(20 TM .Pi:)) ptr
d v((::))
Note that P(pr)) can be expressed as
v(pi:)) = (ZO .D::)) p1:
where za! is independent of P By Lemma
C.2 this implies S2 - 0.
D.2 Proof of Corollary 3.2
Let dh R' R be bounded and Lipschitz function that satisfies |d]| I C. Note that by Fact C.3 it also pseudo-L.issciizz. Define bounded and continuous function f - R by
f(x) = = < C) + xxfx e [-C, C(x C).
Then, Lemma C.2 and Theorem 3.1 imply that
Bl
==)==B E
Σ
v (h a. ........,,.. 11y) = f Σ wcn A'd
aml
B (BIu(z")..... z"') pp.....,pphy
p. =B u(
holds Since the above convergence holds for all bounded and Lipschitz function d. by Lemma this implies the desired convergence in distribution.
E Sub-Gaussianity of the Attention Outputs
In this section, discuss the sub-Gaussianity of the limiting distribution of the attention outputs from Example 3.1 In this appendix, we provide a proof of the sub-Gaussianity of the single random variable 2 for any i e [s], and consider the single-head attention setting (i.e. H = 1) for simplicity. Specifically, we consider
SoftMax,( ".........)""" 2z
The same proof strategy also applies to prove the Sabb-aussiantty of the vector (Z)), ...., .....ZY") with H being any positive integer.
For each j [s), define a (P(1) by
(1)
(1) a (p (13) SoftMax (1)
(p
...... .P..)
34

---

It follows that
z 1'Ip 1,1s
NO. (1s(P,)) ( z goliin
Ja (1(1)))
a
Since each a (f(P lies in the interval [0, I], we have
|| (1assPP/111 P Σ ( =1
and therefore
a1ss (Psss) (1) ) Var( z alula Jais (Pi,
|a (p. (.l [ (ar 2 lll ||Var Z 61.1
lll..
where || lll is the operator norm. As a result, for any 2 0, it holds that
P(122 2 r) =
P(IZ"1 2 p ((1)))
1
$ 2B exp
201 (1)
Varr sh.s
)a1s (p(1)
Pi.ss
$ 2 exxp
Var Z
bllito)) 112
by the conditional sub-Gause Sabb-aussianity of 2 given pl..s-
In this argument, we can also show the sub- Gaussianity the limiting distribution attention outputs of the form y' = 1 Σh Σj=1 O5 (P ..... Pi,s p! ) where the softmax functions are replaced by Vи bounded and pscudo-Lipschitz nonlinearities $,.
35


Reviews:
[{"review_id": "vT0Ni4AT6R", "reviewer": "Reviewer_pKBC", "summary": "The 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.", "questions": "- 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.", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The paper provides a rigorous analysis of the infinite-width limit of a realistic attention setting (finite heads, standard scaling), addressing a gap in prior work. While the breakdown of Gaussianity is somewhat expected, the precise characterization of the hierarchical Gaussian distribution is a meaningful contribution. The main limitation is that the implications, both theoretical and practical, are discussed only speculatively, without concrete analysis or validation. Nevertheless, the authors acknowledge this and propose future directions. Overall, I think the paper offers valuable theoretical insights that warrant a borderline accept.", "rebuttal": "Thank you for the constructive feedback and for highlighting our contribution.\n\n**Weakness 2-2: The non-Gaussianity is somewhat expected.**\n\n**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.\n\n**Q1: Analysis of deeper models and overcoming the non-Gaussian challenge.**\n\n**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.\n\nWhile 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. \n\nWe 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.\n\n**Q2: Implications of non-Gaussianity and differences from a standard Gaussian.**\n\n**Reply:** We agree that a deeper discussion on the implications of our findings is crucial. \n\nThe 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.\n\nThis inherent non-Gaussianity may significantly affect several aspects of learning in self-attention models. For example:\n1. 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.\n2. 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).\n3. 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.\n\nAt 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.\n\nWe 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.\n\n**Q3: Readability and pseudocode.**\n\n**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.\n\n**Reference**\n\n[1] Hron et al. (2020). Infinite attention: NNGP and NTK for deep attention networks.\n\n[2] Yang (2019). Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are Gaussian processes."}, {"review_id": "UWpoSdPrRR", "reviewer": "Reviewer_Tp28", "summary": "The 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.", "questions": "Can 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.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "+ 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", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We acknowledge the concerns regarding clarity, impact, and the scope of our experiments.\n\n**Weaknesses: Derivation is somewhat obscure and difficult to follow (at least for non-specialists).**\n\n**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.\n\n**Weaknesses: Impact of analysis is not clear, what can be learned?**\n\n**Question: What is the impact of your analysis for optimization, initialization or hyper-parameter selection?**\n\n**Reply:** Thank you for this important question. We present specific impacts on current choices as well as future impacts.\n\nAs 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.\n\nLooking 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.\n\nWe agree that the contributions of our analysis should be made more explicit, and we will revise the paper accordingly to emphasize them. \n\n**Weaknesses: Experiments are very rudimentary.**\n\n**Question: Can you provide stronger experimental evidence to support your analysis?**\n\n**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).\n\nThe 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).\n\n| $n$ | 16 | 64 | 256 | 1024 |\n|---|---|---|---|---|\n| log-mean | -2.78 | -5.02 | -6.68 | -7.31 |\n| log-mean - 1SD | -2.93 | -5.22 | -6.89 | -7.58 |\n| log-mean + 1SD | -2.63 | -4.82 | -6.48 | -7.03 |\n\nAs 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.\n\n**Question: Is there a way you can make the notation more readable?**\n\n**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.\n\nFurthermore, 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.\n\n**Reference**\n\n[1] Yang (2019). Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are Gaussian processes.\n\n[2] Yang (2020). Tensor programs iii: Neural matrix laws.\n\n[3] Yang and Hu (2021). Tensor programs iv: Feature learning in infinite-width neural networks."}, {"review_id": "HTYH6JXzBS", "reviewer": "Reviewer_NP2J", "summary": "The 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.", "questions": "I 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.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "The 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.", "was_revised": true, "final_justification": "I believe the setting is interesting due to the foundation nature and the results are solid. However, I am still not completely convinced that this setting is very important due to the fact that it is limited to one layer, has theoretical instabilities after training, may be well approximated by the infinite head limit.", "rebuttal": "Thank you for the constructive feedback and for suggesting additional literature.\n\n**Weakness: On the limitation to a single layer at initialization.**\n\n**Reply:** Thank you for the comment. This is indeed a valid and important point.\nWe 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.\n\nWe 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.\n\n**Q1: MuP scaling.**\n\n**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).\n\nWhile 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.\n\nThe 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.\n\n**Q2: Infinite-head limit.**\n\n**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.\n\nTo 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.\n\nBased on your comments we will emphasize this important suggestion more in the main text.\n\n**Q3: Relation with Dinan et al. (2023).**\n\n**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:\n\n- 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.\n- Our work, in contrast, rigorously derives the full limiting distribution of the attention layer's output, going beyond just the first two moments.\n\nThis 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.\n\n**On missed citations.**\n\n**Reply:** Thank you for your suggestions on citations. We will ensure our discussion of the related literature is comprehensive in the revised manuscript.\n\n- 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.\n- 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.\n- 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.\n\n**Reference**\n\n[1] Bordelon et al. (2024). Infinite Limits of Multi-head Transformer Dynamics.\n\n[2] Cowsik et al. (2024). Geometric dynamics of signal propagation predict trainability of transformers.\n\n[3] Dinan et al. (2023). Effective Theory of Transformers at Initialization.\n\n[4] Hayou and Liu (2025). Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size.\n\n[5] Hron et al. (2020). Infinite attention: NNGP and NTK for deep attention networks."}, {"review_id": "yvLm1uEGOQ", "reviewer": "Reviewer_9XUB", "summary": "This 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", "questions": "1. 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.", "limitations": "The 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$.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "The authors analyze an important problem (the feature distribution of randomly initialized transformers at various head counts and head dimensions). With their promised revisions and clarifications of the limitations of their approach, I am in support of acceptance of this work.", "rebuttal": "Thank you for the thoughtful review and for highlighting important limitations and additional literature.\n\n**Weakness 1 and Q1: Training dynamics.**\n\n**Reply:** Thank you for your insightful question. We agree that understanding the conditions for stable training is a vital research direction.\n\nOur 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.\n\nFurthermore, 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).\n\nThe 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.\n\n**Weakness 2: Fixed batch size and sequence length.**\n\n**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.\n\n**Weakness 3: Novelty and significance.**\n\n**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).\n\nAs 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.\n\nOur 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.\n\nFurthermore, 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.\n\nFinally, 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.\n\n**Q2: The attention scaling considering head counts.**\n\n**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.\n\nThis change has significantly strengthened our work by enabling a direct comparison with Hron et al. (2020). We found that with this scaling:\n\n- The limiting distribution of the attention scores in our framework perfectly matches the Gaussian distribution in their infinite-head limit.\n- The variance of our finite-head non-Gaussian output distribution theoretically matches the variance of their infinite-head Gaussian output distribution.\n\nTo 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.\n\n**Q3: The realism of the theoretical setting.**\n\n**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.\n\nFirst, 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.\n\nSecond, 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.\n\nWe will revise our manuscript to emphasize more clearly that our findings provide a robust description for any number of heads, clarifying their broad applicability.\n\n**On missed citations.**\n\n**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:\n\n- 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.\n- 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.\n- 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.\n\n**Reference**\n\n[1] Bordelon et al. (2024). Infinite Limits of Multi-head Transformer Dynamics.\n\n[2] Cowsik et al. (2024). Geometric dynamics of signal propagation predict trainability of transformers.\n\n[3] Dinan et al. (2023). Effective Theory of Transformers at Initialization.\n\n[4] Hayou and Liu (2025). Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size.\n\n[5] Hron et al. (2020). Infinite attention: NNGP and NTK for deep attention networks.\n\n[6] Noci et al. (2023). The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit.\n\n[7] Yang (2019). Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation."}]

Author rebuttal:

Meta-review: This paper gives a rigorous characterization of the infinite-width limit of a single attention layer under the standard scaling with a finite number of heads. Using Tensor Programs, the authors prove that the layer’s outputs converge to a hierarchical Gaussian law. Finite-width simulations (including variations in token count and low-rank settings) closely match theory. Together with comparisons to infinite-head regimes, the result cleanly fills a gap between prior analyses relying on infinite heads or nonstandard scalings and the setting used in real Transformers.

Reviewers broadly agreed the problem and setting are important and the mathematics sound. Strengths highlighted include: (i) rigorous limit law; (ii) technically careful derivations within Tensor Programs; and (iii) empirical checks that substantiate the limit in realistic ranges. Main weaknesses are scope and presentation: the result is for one layer and the forward pass at initialization only; implications for deep models and training dynamics remain largely speculative; notation is heavy; prior relevant work is not thoroughly discussed; and experiments are modest. During rebuttal the authors (a) clarified relations to infinite-head limits and adopted an explicit 1/√h factor to make head-count scaling comparisons fair, showing convergence toward the Gaussian infinite-head law; (b) added discussion of consequences of non-Gaussianity (heavier tails/anisotropy) and how this might propagate with depth; (c) provided additional experiments varying the number of tokens; and (d) committed to clearer presentation and to broadening related-work coverage. All reviewers voted for acceptance. I recommend acceptance as well.
Decision: Accept (poster)