Bayes optimal learning of attention-indexed models
Authors: Fabrizio Boncoraglio
Emanuele Troiani
Vittorio Erba
Lenka Zdeborova
Venue: NeurIPS 2025 poster ()
Keywords: Bayes optimal analysis
Transformers
Self-attention mechanism
Random matrix theory
Matrix models
Extensive-rank analysis
Primary area: theory
TL;DR: We introduce and analyze the Attention-Indexed Model (AIM), a theoretical framework for analyzing learning in deep attention layers.

Abstract: We introduce the attention-indexed model (AIM), a theoretical framework for analyzing learning in deep attention layers. Inspired by multi-index models, AIM captures how token-level outputs emerge from layered bilinear interactions over high-dimensional embeddings. Unlike prior tractable attention models, AIM allows full-width key and query matrices, aligning more closely with practical transformers. Using tools from statistical mechanics and random matrix theory, we derive closed-form predictions for Bayes-optimal generalization error and identify sharp phase transitions as a function of sample complexity, model width, and sequence length. We propose a matching approximate message passing algorithm and show that gradient descent can reach optimal performance. AIM offers a solvable playground for understanding learning in self-attention layers, that are key components of modern architectures.

Bayes optimal learning of attention-indexed models
Fabrizio Boncoraglio
Emanuele Troiani
Statistical Physics of Computation Laboratory EPFL, Switzerland
Statistical Physics of Computation Laboratory EPFL, Switzerland
Vittorio Erba
Lenka Zdeborová
Statistical Physics of Computation Laboratory EPFL, Switzerland
Statistical Physics of Computation Laboratory EPFL. Switzerland
Abstract
We introduce the attention attention-indexet model (AIM), theoretical framework for analyzing learning in deep attention layers. Inspired by multi-index models, AIM captures how token-level outputs emerge from layered bilinear interactions over high-dimensional embeddings. Unlike prior tractable attention models, AIM allows full-width key and query matrices, aligning more closely with practical transform- ers. Using tools from statistical mechanics and random matrix theory, we derive closed-form predictions for Bayes-optimal generalization error and identify sharp phase transitions as a function of sample complexity, model width, and sequence length. We propose matching approximate message passing algorithm and show that gradient descent can reach optimal performance. AIM offers a solvable play- ground for understanding learning in self-attention layers, that key components modern architectures.
Introduction
The transformer architecture [1] has transformed machine learning, achieving state-of- sire----tthetttt results natural language processing [2, 3], computer vision [4), and beyond. Its core innovation-th self- models to capture long-range dependencies between tokens. Despite their empirical success, transformers remain poorly understood theoretically, especially regarding how data structure, attention bias, and training dynamies interact in finite-sample regimes. While mechanistic interpretability has shed light on trained models, the learning process statistically and computationally learnable from limited remains unexplained. A common strategy toward progress is to study simplified models in high-dimensional regimes, where the blessing of dimensionality [5] can yield tractable characterizations of learning. key ingredient in this approach is synthetic data model that captures salient aspects of real-world structure.
attention mechanism-enate
is
A common feature of natural data is that structure arises from pairwise interactions between elements.
In language, for example, sentence meaning depends on both the semantics of individual words and their syntactic roles. This can be modeled by assigning each word semantic embedding and contextual embedding z. For instance, in "Max eats chocolate," one level of understanding identifies what each word means (c.g.. Max is person), while another captures grammatical roles (Max the subject). These two embeddings are linked: contextual meaning emerges by comparing semantic content across words. Similar relational patterns occur in images, molecules, and graphs. The empirical success of transformers suggests that multi-layer self tetfraateettowihiieee pairwise comparisons in a hierarchical way-is natural mechanism for modeling such interactions. To analyze this, we construct synthetic data model that reflects pairwise structure between tokens.
performs weighted
Theoretical understanding of fully connected neural networks has advanced significantly through the analysis of Gaussian single-index and multi-index models in the high-dimensional limit [6, 7, 8, 9,
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

10, 11, 12, 13, 14, 15]. In statistical physics, similar models appear as teacher- perceptrons [16, 17, 18, 19] or committee machines [20, 21]. These setups typically assume i.i.d. Gaussian inputs, with targets depending on a small number of random vrojections-"inddeess "indices" of the input. They provide a rich theoretical playground for jointly analyzing learning dynamics, generalization, and architectural biases.
Recent work has extended this framework model key aspects of transformers, introducing the sequence multi-index (SMI) model [22, 23, 24j. While insightful, existing SMI models require the width of the key and query matrices to be much smaller than the token embedding dimension regime where only narrow attention layers can be analyzed. In contrast, practical transformers typically use key and query widths comparable to the embedding dimension. This motivates our contribution: high-ding -igh--imensionall yet analyzable model where learnable matrices have extensive rank. We call this the attention iiteettion-indeeed model.
The aatention-indexet model (AIM). We introduce a new class of high-dimenniooal functions designed to model pairwise relationships between tokens. Analogous to classical multi-index models, the altention-indexed model defines outputs y as nonlinear functions of high-dimensional token embeddings x. Rd for L,..., T, arranged rows of the input sequence of tokens Xo -X e
IRxx xd for a = l...., T. We define L attention indices ()()) E RTxT with components (1()) The labels for each input X E Txxxd are generated via a general output function RLXTXT + RTXT
x St 2bb - d0b Sr
y= g (1h10)
(I)
Here each St Rdxd a learnable matrix. The diagonal mean is subtracted to avoid divergence as 00, ensuring the fluctuations of (()) remain (1).
While our theory applies to general rotationally invariant S,, a motivating example is when S. = QeK? BRdxd in self-attention [1], with key and query matrices K2, Qi E Rdxry, We refer re as the width of the 4th layer; it typically controls the rank of Sr. though we also consider ra > d. For analytical simplicity, we assume tied key and query, Qe
Ke We, that
S. Voed WoW, E Rdxd
We E Rdxrs
Multi-layer attention. A key instance of the output function g arises when contextual embeddings are built via stacked attention layers. Each layer updates the representation by mixing semantic embeddings through weighted pairwise interactions. Formally, we define the multi-layer attention mechanism
y (x(Hi(X1-1)))
-[IIr+ X = [cIr + (o X
cIr
(3)
1 He(Xe-1) = Vd X S - -x
where X, <R is the input sequence, e is the residual strength, and L is the number of layers. The function RTxT - R is typically the softmax with inverse temperature 83 increasing B sharpens the focus on dominant pairwise interactions. Finally, the expectation in H₂(X) intended over the input data Xo (empirical average).
particularly interesting case is the limit of infinite 3, where softmax becomes the hardmax function that selects the maximum entry in each row of the attention logits. This can be interpreted as enforcing sparse, winner-take take-all selection among candidate interactions, capturing the case where for each token only single other token carries the relevant signal. Such mechanisms are especially relevant in symbolic, compositional, or relational reasoning, where outputs depend on identifying a unique best match per query.
This form fits into the attention -ndeexed framework by choosing the output function g in (1) as:
1( s(119)]))) = =ag (B. -1 ...., 11)1 ((11) BB.-1(16) o..),,,,..... (..0,,...(-1))) ()11))))
(5)

---

where the recursion is defined by:
Bº = Tr, B = cUr + or B 100 1T- Be-1 t= I=1,..., L.
=
(6)
In partieular, with this formulation the map in (4) can rewritten as:
(ee(ee1) 1 X S.
Sr Vd
BL-1"
Be- 1000 B
(7)
This formulation also accommodates sequence-level outputs of the form y yX corresponding to the output map g B Details of the corresponding mappings given in Appendix B.
The existence of such mapping, from deep attention to a low-dimensional collection of attention indices, is key contribution: it allows information-theoredd and algorithmic results derived for AIMS to transfer to multi-layer attention. In App. B we extend such mapping to multi-head and seq2seq variants. Full Transformer blocks (e.g.. ViT/LLMs) interleave multi-head self-attention with token-wise MLPs, and are commonly trained an auto-regressive fashion; our supervised setup isolates the attention learning problem, to which results apply directly.
Our contributions. We initiate a theoretical study of learning from data generated by the attention- indexed model, without restrictions the width of the attention matrices S,, in the high-dimensional limit. Specifically, we consider large embedding dimension d, with attention widths ry (or ranks of Se) scaling proportionally with d, while keeping sequence length 7 and depth I. finite. The number of samples scales quadratically with dimension to span the regime where the optimal test error changes from as bad as a random guess to zero:
d d000
n = ((),
Pe = T()( L --
(8)
d
In this limit, for Gaussian i..d. inputs and suitable priors over Se, we analyze the Bayes-optimal esti- mator thatt is, the posterior mean predictor. Concentration of measure allows key observables such as the test error-to to become deterministic, enabling a description in terms of low-dimesiional fixed- point equations. We derive those equations using tools from statistical mechanics and random matrix theory. Beyond analytical tractability, this asymptotic setting captures relevant scalings, where model and data dimensions grow together.
Our analysis uncovers phase transitions in Icarnability as a function of sample complexity sequence length T. and attention width p. We show that attention mechanisms with extensive width can efficiently recover latent structure, particularly when the target model exhibits sparsity. We further study the role of the inverse temperature parameter 3, contrasting the smooth regime < B too with the hard assignment case 8 = +oc, where the task reduces to discrete token association. We derive an approximate message passing (AMP) algorithm that matches Bayes-optimal performance in the studied setting. We also show empirically that an averaged version of gradient descent natural loss achieves similar optimality, aligning theory with practical training dynamics.
Further related work. The attention-indexed model (AIM) is motivated by generative perspective, capturing how structured token-level outputs arise from layered bilinear interactions between high- dimensional embeddings mirroring attention computations in transformers. The idea of modeling learning through such structured synthetic data dates back to the pracher-studen setting in the perceptron literature [16, 17, 20], and more recently to single-index and multi-index models [6, 7, 8, 9, 10, 11, 12, 13, 14, 15].
While many theoretical studies explore simplified transformer variants, most do not rely on or benefit from the high-dimensional limit. These include works that analyze one-layer attention under finite embedding dimension [25, 26, 27, 28, 29], or study training dynamics in the linear, kernel, or random feature regimes [30, 31, 32]. Others use infinite-width approximations without access to generalization error [33, 34]. By contrast, theoretical analysis of nonlinear attention layers with trainable key and query matrices in the limit of high embedding dimension together with sharp control of generalization-
aware, only a few works address this regime [23, 22, 35, 36], and they all assume attention matrices of finite width.
less explored. As far
Methodologically, our approach builds on techniques from high-dimensional multi-index models, particularly those developed [21, 37], and their recent generalizations to sequence learning with multiple low-width self-attention layers [35]. The main technical challenge addressed in this paper is
3

---

extending these tools to the case where the width of the attention matrices scales proportionally with the embedding dimension-ies i.e. the extensive-width gegime-going beyond the key limitations of prior analyses.
To tackle this, we leverage recent results on the ellipsoid fitting problem [38, 39] and its connection to two-layer neural networks with quadratic activations and extensive width [40, 41J. Remarkably, the linear AIM model with T' L 1 is mathematically equivalent such quadratic networks, allowing us to adopt these methods. We generalize this connection to arbitrary T,, L. This is enabled by central conceptual tool, the AIM index, which disentangles the complexity of deep attention models. It allows us to split the problem into two subproblems: (i) how structure propagates across layers and tokens, and (ii) how attention matrices are learned from those structures. This separation is crucial in extending the theory to multiple layers and tokens.
Finally, we note that focus here the tied case Q K for clarity. The untied setting Q + K is amenable to similar analysis following [42], and we leave its treatment for future work.
2 Setting
We consider dataset D = {y² X } of n samples indexed by μa, where X fart Pa=1 E R'xd has rows x T_ Each sample consists of the embeddings of I' tokens x/ Rd taken as standard Gaussian 2: ~ (0, IT)) and of T T matching output matrices yd encoding pair-wise information on the original tokens. We stress that the Gaussian assumption for the data can be relaxed in the same spirit as [41, Assumption 2.2].
We generate y" using an attention indexxd model as given in (1) with matrices {S/]t=1....4 that are symmetric and extracted independently from a rotationally invariant ensemble Pss(S) = Ps SO) for any d d rotation matrix O. We fix the normalizations such that E [] Eps [Tr S²] = K2d and with K1, R2 = (1). We assume that the empirical spectral distribution of P3 converges directions, allowing for different prior distributions P for different layers, as well as considering non-symmetric matrices [43].
Kjd and
distribution ja for d + +oo. This setting can be relaxed in several
We consider the Bayes-optimal (BO) learning setting: the statistician knows the generative process of the dataset, i.c. the non-linearity in (1) and the prior distribution P's, and observes a dataset D but not the specific set of weights {St It=1....L used to generate said dataset. The task is then to optimally estimate either the weights S (estimation task), i.e. find the estimator S(D) that minimizes
1
Eest (S) = ED.s
-
(9)
1=1
or the label associated to a new input sample Xnew (generalization task), i.c. find the estimator ((v D) that minimizes
Egen ED,s- -.o.s-lue Eynne E Nuow (Xeww, D) ynew If+
(10)
where (Unew, Neww) is new sabelesample pair generated with the weights S We will call the error achieved by the optimal estimators, respectively, the BO estimation error and BO generalization error.
Both BO estimators can be computed from the knowledge of the posterior distribution, i.e. the probability that a given set of weights S was used generate the observed dataset
(ISI...,,SIP1 ((D)
1
L
Ps
II (v" -9 (11, ==1
n
(11)
where the attention indices 12( RTXT were defined in (1) and ((D) is a normalization factor. The BO estimator with respect to the estimation error is the mean of the posterior distribution, while BÔ estimator with respect to the generalization error is the mean of the predicted label under the posterior distribution, i.c.
Š = (s)
((x) = ES-P(SID) gg (S). x) (1
(12)
Under the figh-dimensional limit (8), we will show that the estimatiion/generaliaaaioo error achieved by the BO estimators are characterized through the overlaps g, Q defined as
qe d (S-P())) [Tr S.SA
Quk J Es- Ps [Tr ST
(13)

---

3 Statistical and computational limits for ALMs
In this Section we provide results for the cntemmttionthhoretiicn and computational limits on attention- indexed models in very general setting, which we then specify to the attention models in Section 4. In order state our results let us define, for two symmetric positive-definite matrix q, ST.
/
L I dse (S())
L
3m ((((())/2
exp
d
lek Tr( +
Σ 2,k=1
Ve Tr((
f,k=1
/ DYL... ... DYL Z (Y....., YL: log nn (Y,,.... YL:9)
(4) = lim
d=v00 de
(14)
where DY stands for integration over GOE(d) (Wigner) matrix Y, and
Zout (y, V)
/
II dt has N (hasi Wabi Vab) 8 (v- / 8ab )abl(=1)
II dyub / Dm ... DmL Zout w V) log Zout (y. w, V)
T
Lout (g) =
(15)
a,b=1
with:
w
24
(R) Nab
(((h) = 2(QEx (ekk)
where Dnj stands for integration over a L x T x T tensor symmetric in the token indices and with independent entries (O, 1). Moreover, let us define
Jout (y, w, V) 0., In Zout (y.w, w, V),
9mn ( -gyray In (g)
(16)
Our first result provides a description of the error in the high dimensional - oc limit, with finite sample complexity n/d?, number of tokens T and number of layers L.
Result 3.1 (Performance of information-thhorrtiinn optimal estimators). Consider the extremization problem
Tr 4
inf sup GEBT gEst.
Im)()) alou(o))
(17)
where In (g) and Ion (q) are defined in (14) and (15), Call (q" (²)") the global extremizer of (17). Then, in the high dimensional limit (8), the BO estimation error is given by
lim E (Spo) 11Q 4"1li..
(18)
while the BO generalization error is given by
(8)
T.L
-9 (( V(ck)) ab
T.L
Egen
Eyp.
(19)
2
dab
aSb,t=1
/2 Jab
aat,8=1
F
where& & L T T tensor symmetric the token indices and with independent entries N(0, 1) and w. V.n are defined in (15). Finally h(w((), v(ek)) = (() + EE-iVVa ((s)
These expressions provide prediction for the information-theerrtieei generalization errors, and as such they constitute a sharp information theoretical bound on the performance of any algorithm. We derive Result 3.1 I in Appendix C using the heuristic replica method, but believe that a rigorous treatment is possible by adapting [41] to the case of multiple attention indices L and multiple tokens T 1.
optimal estimation and
In our next result we provide an efficient polynomial-tine approximate message passing (AMP) algorithm that in the high-dimensional limit saturates this bound under the condition that there is unique local extremizer of (17) (if multiple extermizers are present, computational-to-SSe -o-statistical gaps may arise [44]).
5

---

Algorithm 1: AMP
Result:
The estimators S₁
Input:
Observations y 1x' and "sensing matrices"
Laj = (a + x x i,b /bb)) 2d((1 + dub) E R;
Initialize S1 Ps and Cu=0= 2(R22 - ki)_L;:
while not converging do
Estimation of the variance and mean
V
-2 and a/y,ab ab = Tr (t 5*1 (1 - oor gout (v" w/// vi- V's'' E eR';
Variance and mean of St estimated from the "output" observations;
n.P.I
db
40
Σ
Jout (yr* why V') Jour (yr" why ())())
and
n
Rij- / + (4¹) 1 n.T.T
Doutly" (y" why V'Jas Job E R p.agt
Σ
gout
Update of the estimation of S with the "input" information;
1
$t Il (R))
and
C
VR 9ta (R¹, (²) d2
t t=t+1 1;
end
Result 3.2 (State evolution of AMP). Call St the time-t iterate of the AMP algorithm 1. In the high dimensional limit (8) of large d and for finite number of iterations, we have that
TH(S + dée
Tr[S
de
(20)
where
V2 ((
Jout (s hiww",, Viuu,
}) w V
(0)
+++
4alEg, ,
8ab
o5b
ab
x Yout g ((hhiuu) V2)
1) we Ve
(K)
(21)
V2 Oab
ab
lim Es,y Tr 19tul (S,
Jegin (S, )
Aek
d
with y = 2(Q )¹) and wi and with ne a standard Gaussian in every component. h(w" V*), and & are defined as Result 3.1. Finally we have (S,, = S₁ + Et. VA. all E,,. are GOE(d) and all S, ~ Ps-
EEin( yek The where (a, b) are token and (E, k) layer indices
Result 3.2 (which we derive in Appendix C) is a classic statement in the theory of AMP, that here we adapt take into account multiple attention indices and multiple tokens. We remark that though the denoiser fin in (16) is a complicated non-separable function, state evolution still holds [45, 46]. Additionally, show in Appendix C that the fixed point of (21) are the extremisers of (17). The AMP algorithm is thus both a tool that gives us theoretical guarantees though (21) and a practical algorithm that be efficiently implemented and run machine. We discuss the implementation in Appendix C.
The prior-related elements In and Om of (14) and (21) require discussion. Both of them involve integrals dimension (LLd) whose large d asymptotics is highly non-trivial.
This is due to the prior Ps on the weights being non-separable, and crucially to the coupling between weights of different layers caused by the non-diagonal matrix q in (14) and (21). We remark that for generic number of layers L > 1, the evaluation of these two equations is equivalent to the evaluation

---

in the high-ding -igensimnnional limit of the free entropy and of the Bayes-optimal estimator for the following L-wise matrix denoising problem
E
Yt = Se+ +
Σ m=
(22)
where from heavily correlated set of L noisy observations Y Rdxd, L,.,,...,4. L, one needs to estimate back the L independent matrices Sm ~ Ps. Here A E S7 is an L-dimensional symmetric positive-defiiite matrix acting noise-to-signal ratio, and E are independent GOE(d) matrices acting as noise. To the best our knowledge, for L > and non-diagonal A this is still a challenging open problem (in case of diagonal A, the problem factorizes in L independent matrix denoising problems), In practice, one can approximate otationally-invarian matrix denoisers by low-degree spectral polynomials, following [47]; this provides practical route AMP for multi-layer ( when the exact HCIZ-based denoiser is unavailable [48].
For L = 1, the problem (22) reduces Bayes-optimal denoising of rotationally invariant matrix [40, 43]. We start by remarking that for = 1 the quantities Q, q, 4 governing the BO performance are scalar.
For Result (3.1) at leading order in d oo we have
Q q log q (((1)/))- 1
Im (q)
where E(u) Ex log |r yl.
(23)
4
and for the AMP algorithm Result (3.2) we have
91a,L=1 = OT 0 (RIE(A, where fRIE(A, = =A,- 24 -A /
dua(t)
(24)
A,
where Q =2, Mu is the asymptotic spectral density of matrix
K2, 7 ~ GOF(dd,, and X O AO is the degoposccomooiiiion of X (see [40] for details).
VAZ with Ps and
Similarly, at leading order in d we have
4=²A²
f ddps()p
Iin. (X,
(25)
Finally, the state evolution equation for in (21) reads
4x² 9 3q2
/
q Q
dt [1])
(26)
4 Results for single-layer attention
We now apply general results to the single-layer (L 1) tied-attention model
x S xb
dab Tr
Vrd L
TWW 2b
Sab
Tn
WWT)
Yab 0g
0g
vrd
(27)
Vd
Vd
where we parametrized the weight matrix S as a tied-attention with extensive-width r = pd and W dxxr has independent entries Waj (O, 1). For the activation, we consider the case of Hardmax haard and Softmax Jsoft, both applied row-wise in (27):
O1 ........2.) 8(i arg max 2z), and (20 ... 2T); (2
(28)
Σ)= e
We stress that both these tasks are well-defined only for T 2 2, the T 1 the output of both activations equals regardless of the input. As discussed in the introduction, the model with hardmax provides interesting token-association task.
Hardmax target. The BO treatment of the hardmax activation for generic number of tokens 7 is challenging due to the complex form of the state equation à (21). We provide an explicit solution in the I = 2 case.
7

---

Result 4.1 (Bayes-optimal errors for hardmax tied-attention, T = 2). Consider the model (27)
with hardmax activation. In the high--diensional limit (8), the asymptotic BO estimation and generalization errors are given by (18) and (19), where (q:9) is the solution of (26) with Q = 1+11
and
6(k), k2 V6(Q q) D ((k,, c)
(- (-(81+))
/281 + 82)
-(11 + 82)
Jour (u, aaa
(29)
V282
ab
where @(k1, 122, c) is the p.d.f. of bi-variate Gaussian with zero mean, variances 1/(1 22) and covariance c/(1- )), and (kk,, k2, c) is its c.d.f (see Appendix A). Moreover, Sa 2yaa 1, ka = (v w12)/ /6QQ - q). S182/3 33 and Wab = /2q Nab-
We detail the derivation of Result 4.1 in App. C. We plot the estimation error given in Result (4.1) in Figure left, for several values of the attention width ratio p. comparing with runs of the associated AMP Algorithm at size 100.
We observe that for all finite the estimation error is strictly positive, and that it approaches zero as a grows with rate compatible with O(1/a). Moreover, as soon as > O. we observe that the estimation eтror is smaller than 1, i.e. the value achieved in the absence of data. Intuitively, the hardmax output function enforces discrete token association per row. This is akin to multiclass classification target, which explains the lack of a finite strong-recovery threshold and the observed power-law decay at large a (Fig. 1, left).
In the limit of small width Result 4.1 simplifies. Notice that in this limit the correct sample scale given by a/p = n/ (drrr), as the matrix to infer not extensive- anymore. In this limit there appears a calcel weak recovery threshold, value of sample complexity below which the estimator reaches the same performance as there were no data. We characterize it as follows.
Corollary 4.2 (Small width limit for hardmax activation). Consider the model (27) with hardmax activation and T = 2. In the high- high-diimesiooal limit (8), the equation for q of Result 3.1 simplifies to q = [max(1 t, 0)}² under the rescaling a = ap and q = tp. in particular, the BO error the same as that of the data-less estimator for all < harddaax where weak
1
hardmax O weak
0.563.
(30)
4Ey,
Etso Jou (M(ws, V), w, V) -=0
Corollary 4.2 follows by combining Result 4.2 with the small-width analysis of [40]. We remark that (30) holds for all activation (using the appropriate gout function) and all values of 2 2. We plot the analytical prediction for the BO estimation error given in Corollary 4.2 in Figure I right.
Softmax target. We now discuss the target function that uses a softmax non-linearity (28). This choice of activation allows for an analytic treatment for any number of tokens T 2 2, and any finite value of the softmax inverse temperature 3 R.
Result 4.3 (Bayes-optimal errors for softmax iie-aateetionn T 2). Consider the model (27) with softmax activation, T 2 and inverse temperature 3 E R24. In the high-dimensional limit (8), the asymptotic BO estimation and generalization errors are given by (18) and (19), where (g. Q) E IR3 is the solution of (26), Q + p and aT - -
4)-
We plot the BO estimation error given in Result 4.3 in Figure 2 left, and observe that contrary to the hardmax case, the BO estimation error vanishes a finite value of a (the so-called strong recovery threshold). Interestingly, the BO errors given in Result (4.3) is independent of the value of the inverse temperature 3 and reduces to the case of single-token model with linear activation [40], modulo a rescaling of the sample ratio a to 2 + T - 2) (notice that the rescaling is not just given by the total number unordered couples of tokens T(T + 1)/2, as it would be in the case of multi-token case with bijective activation, see App. C). The softmax activation is almost invertible, meaning that given the output, the input is fully determined apart for a common additive shift (acting as a noise correlated with the data), and is additionally constrained by the symmetry of the attention matrix. Result (4.3) precisely quantifies the amount of samples required to estimate this undetermined shift. More precisely, fix a given estimation error. Then, achieving this error with the BO estimator in the softmax case with T 2 2 requires a factor + (T + 1)) more samples than the case of a fully bijective activation.
8

---

109
1.00
0.60
Extimation
0.5
0.40
10-1
2.0
0.20
5.0
0
10
Sample complexity w/d!
Sample complexity n/dr
Figure 1: (Leff) The Bayes optimal-error for the single-layer attention-indexed model with T = 2 tokens and hardmax activation for and several values of the width ratio p (Result (4.1)). The log-log scale highlights a large power-law decay of the BO estimation error, strikingly different from the softmax behaviour (see Figure 2). We also plot the corresponding errors achieved by the AMP algorithm (dots) at d 100, averaged over 16 realizations of the data and teacher weights, Error bars are computed with respect to the mean. We find good agreement even for such a moderate size. (Right) Focus the small width Bayes-optimal error case (Corollary (4.2)) of the same model. We rescale the sample complexity to â = a/ p, and we highlight the theoretical prediction of the weak recovery thresholds (gray vertical line).
On the other hand, we remark that the AMP algorithm for the softmax activation at T 2 2 is not a simple rescaling the AMP for the single-token linear-activation case given in [40]. The AMP output function gout is given
obut, =
Tab T
Σ
T
Ted - TodWodl
+
T T-1
Σ o.d
2 Oed + Fabora wab + (1 Pab (31)
V
c
where Tab / 5ab. Φ(y)ab = 8 Log(Vab/ /arr and w, V defined in (15). Thus, AMP processes the data in nriviaivial, optimal way to perform this effective inversion of the softmax activation. We plot experiments for AMP at d 100 in the T = 2,3 case in Figure 2 right (purple and blue dots, be compared with the prediction agreement. We also remark that while the BO performance is independent of the inverse temperature B, as long as is finite, again AMP output function not.
Result (4.3) given by the black line), and observe nice
Thanks to the mentioned reduction, one can transfer directly several results from [40] to the case of softmax tied-attention, including an explicit prediction for the strong recovery threshold (the value of after which the BO error is zero), the slope of the error at strong recovery, and the small-width and large-width limits (see App. D). In particular, the strong recovery threshold satisfies
2
p²/2 if 0<p<L 1/2
softmax yecooveryy
(32)
T +
if p21
We remark again that this threshold does not coincide with the naive counting argument, which would give factor T(T+1) at denominator instead. In particular, the 2/(T*+T-2) T(T+1) factor reflects the near-invertibliity of the row-wise softmax under global row-shift, modulo the symmetry constraint; cf. App. D. Finally, contrarily to hardmax (which implements a discrete winner-takes-all assignment and exhibits power-law error decay), softmax behaves like smooth regression target with finite strong recovery threshold (Eq. (32)), see Fig. vs Fig. 2.
We finally remark that analysis keeps T, L finite while d, n predictions remain accurate as long as T grows much slower than d; the curves obtained after the T-dependent rescaling of a (Sec. 4) remain valid for very large T provided T <<<<<>> d.
with n/d
(1). The
Finally, we consider the performance gradient descent minimizing the loss
xr
WW
dab Tr WWT
((W)
yit 0g
(33)
Vr
6=

---

Estimation error
1.0
P
#=
0.1
0.2
T
AMP
x
T-2
P 0.3
T
T
X
T
P 0.5
T
0.5
P
T
X
T=0
0.0
0.0
0.1
0.2
0.2 0.4
0.6
0.8
1.0
Sample complexity n/d
Sample complexity à = m/d + - 2)
Figure 2: (Left) Illustration of the Bayes-optima estimation error for the softmax feed-attention
model (Result 4.3), eq. (27) for any 0 < 3 < +oo and T tokens, and for several values of the attention width ratio r/d. The model reaches zero BO error at finite a depending on (eq. (32)). (Right) We show, in black dashed lines the theoretical prediction of the BO estimation error computed for the sample complexity rescaled by the number of tokens a (7² + TT' the performance of the corresponding AMP algorithm for 7 2, tokens, correctly achieving the BO error. We also compare the BO performance with those of Adam GD and its averaged version AGD with d = 100. We average each numerical experiment (GD.AGD.AMP over 16 realizations of the data and teacher weights. Error bars are the standard deviation on the mean.
- 2) and p = 0.5. We show
with training generated by Eq. (27) (we optimize using the ADAM optimizer [49]). In line with previous work [40, 42], we also consider the Averaged GD (AGD) estimator given by
M
W-final (W firfinaly T m
SGD.av
m
(34)
M
Vrd
m=
where we average over M initial matrices ( and Wfiinal is the corresponding set of weights at convergence. We plot the results of our numerical experiments at d 200 for both GD and AGD in Figure right. As already observed in [40,42], AGD reaches performances compatible with the BO estimation error, while GD has worse error. We remark that both variants seem achieve perfect recovery at the BO threshold (32). This phenomenon, at this point well documented within this class of models, is still understood.
Limitations
The creentionindexed model introduced in this work and its high oigh-dimensional analysis provide promising stepping stone to the analysis of learning multi-layer attention. We, however, so only analyzed the Bayes-optimal performance and the associated AMP algorithm in the single layer case, gradient descent was only explored numerically and clearly its theoretical analysis in this class of models is interesting topic for future work. While the concept of attention-indexed model captures multi-layer attention networks, the matrix denoiser (22) needed to study the optimal performance is challenging open problem for random matrix theory and thus the analysis of the model for more than one layer remains open. While out results rely on exact statistical physics tools, mathematically rigorous proof of the obtained results for T > 2 is open problem. Future work should also explore the inclusion of multiple heads and MLP layers to mimic yet closer practicalities of the current transformer architectures. The structure of the input data considered in this paper is very simple, future work should also include studied of how the performance depends on the input structure.
Acknowledgments and Disclosure of of Funding
We thank Antoine Maillard, Florent Krzakala, Hugo Cui and Yizhou Xu for insightful discussions related to this work. We acknowledge funding from the Swiss National Science Foundation grants SNSF SMArtNet (grant number 212049).
10

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[2] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. 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 11166, 2019.
[3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-19011 2020.
[4] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
[5] David I. Donoho et al. High-dimensional data analysis: The curses and blessings of dimension- ality. AMS math challenges lecture, 1(200)):32, 2000.
[6] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from nigh-dimenssonal inference. Journal of Machine Learning Research, 22(106):1-51, 2021.
[7] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged -trreed-ttirrase property: necessary and nearly sufficient condition for sgd learning of sparse functions two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
[8] Rodrigo Veiga, Ludovic Stephan, Bruno Lourciro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. Advances in Neural Information Processing Systems, 35:23244-2322 2022.
[9] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptoties of feature learning: How gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946 2022.
[10] Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high- dimensional & mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. In The Thirty Sixth Annual Conference on Learning Theory, pages 1199-1227. PMLR, 2023.
[11] Elizabeth Dollins-Woodfing Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models. Information and Inference: A Journal of the IMA, 1344)iaae028, 2024.
[12] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2024.
[13] Alberto Bietti, Joan Bruna, and Loucas Pillaud- Vivien. On learning gaussian multi- indixi-iiexxxxx models with gradient flow part i: General properties and two-timescale learning. Commmmications on Pure and Applied Mathematics, 78(12)2334-24435 2025.
[14] Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. A theory of non-linear feature learning with one gradient step in two-layer neural networks. Proceedings of the 41st International Conference on Machine Learning, pages 36106-36159, 2024.
[15] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. Foundations of Computational Mathematics, pages 1-84, 2024.
11

---

[16] Elizabeth Gardner and Bernard Derrida. Three unfinished works on the optimal storage capacity of networks. Journal of Physics A: Mathematical and General, 22(12):1983, 1989.
[17] Haim Sompolinsky, Naftali Tishby, and H. Sebastian Seung. Learning from examples in large neural networks. Physical Review Letters, 65(13):1683, 1990.
[18] T.L.H. Watkin, A. Rau, and M. Biehl. The statistical mechanics of learning rule. Reviews of Modern Physics, 65(2):499-556, 1993.
[19] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks the teacher-student setup. Advances in neural information processing systems, 32, 2019.
[20] Andreas Engel. Statistical mechanics of learning. Cambridge University Press, 2001.
[21] Benjamin Aubin, Antoine Maillard, Florent Krzakala, Nicolas Macris, Lenka Zdeborová, et al. The committee machine: Computational to statistical gaps in learning a two-layers neural network.
Advances in Neural Information Processing Systems, 31, 2018.
[22] Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition between positional and semantic learning in a solvable model of dot-product attention. Advances in Neural Information Processing Systems, 37:36342-36389, 2024.
[23] Hugo Cui. High-dimensional learning of narrow neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2025(2):02340,, 2025.
[24] Hugo Cui, Florent Krzakala, and Lenka Zdeborová. Bayes-optimal learning of deep ran- dom networks of extensive-width. Journal of Statistical Mechanics: Theory and Experiment, 2025(1):014001, January 2025. Publisher: IOP Publishing.
[25] Bingqing Song, Boran Han, Shuai Zhang, Jie Ding, and Mingyi Hong. Unraveling the gradi- ent descent dynamics of transformers. Advances in Neural Information Processing Systems, 17:92317-92335 2024.
[26] Yuandong Tian, Yiping Wang. Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in I-layer transformer. Advances in Neural Information Processing Systems, 36:71911-7199 2023.
[27] Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Hyeji Kim, Michael Gastpar, and Chanakya Ekbote. Local to global: Learning dynamics and effect of initialization for transformers. Advances in Neural Information Processing Systems, 37:86243 066624366688,,, 2024.
[28] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In Proceedings of the 41st International Conference on Machine Learning, pages 38018-38070, 2024.
[29] Hongru Yang, Bhavya Kailkhura, Zhangyang Wang, and Yingbin Liang. Training dynamics of transformers recognize word co-occurrence via gradient flow analysis. The Thirty-eight Annual Conference on Neural Information Processing Systems, 2024.
[30] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning, pages 4376-4386. PMLR, 2020.
[31] Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What single attention layer learn? study through the random features lens. Advances in Neural Information Processing Systems, 36:119122 11951, 2023.
[32] Yue M Lu, Mary I Letey, Jacob A Zavatone- Veth, Anindita Maiti, and Cengiz Pehlevan. Asymptotic theory of in-context learning by linear attention. CoRR, 2024.
[33] Blake Bordelon, Hamza Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head trans- former dynamics. Advances in Neural Information Processing Systems, 37:35824-35878. 2024.
12

---

[34] Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and-widd limit. Advances in Neural Information Processing Systems, 36:54250-54222 54281, 2023.
[35] Emanuele Troiani, Hugo Cui, Yatin Dandi, Florent Krzakala, and Lenka Zdeborova. Fun- damental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds. Forty-second International Conference on Machine Learning, 2025.
[36] Pierre Marion, Raphaël Berthier, Gérard Biau, and Claire Boyer. Attention layers provably solve single-location regression. arXiv preprint arXiv:2410.01537, 2024.
[37] Emanuele Troiani, Yatin Dandi, Leonardo Defilippis, Lenka Zdeborova, Bruno Loureiro, and Florent Krzakala. Fundamental computational limits of weak learnability in high-dimensional multi-index models. International Conference on Artificial Intelligence and Statistics, pages 2467-2475. PMLR, 2025.
[38] Antoine Maillard and Afonso S Bandeira. Exact threshold for approximate ellipsoid fitting of random points. arXiv preprint arXiv:2310.05787, 2023.
[39] Antoine Maillard and Dmitriy Kunisky. Fitting an ellipsoid to random points: predictions using the replica method. /EEE Transactions on Information Theory, 2024.
[40] Antoine Maillard, Emanuele Troiani, Simon Martin, Lenka Zdeborová, and Florent Krzakala. Bayes-optimal learning of an extensive-widit neural network from quadratically many samples. Advances Neural Information Processing Systems, 37:82085-822 -22132, December 2024.
[41] Yizhou Xu, Antoine Maillard, Lenka Zdeborová, and Florent Krzakala. Fundamental limits of matrix sensing: Exact asymptoties, universality, and applications. In Nika Haghtalab and Ankur Moitra, editors, Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 1 of Proceedings of Machine Learning Research, pages 5757-5823. PMLR, 30 Jun-04 Jul 2025.
[42] Vittorio Erba, Emanuele Troiani, Luca Biggio, Antoine Maillard, and Lenka Zdeborová. Bilinear sequence regression: A model for learning from long sequences of high-dimensional tokens. Physical Review X, 15(2)0020092, 2025.
[43] Emanuele Troiani, Vittorio Erba, Florent Krzakala, Antoine Maillard, and Lenka Zdeborová. Optimal denoising of rotationally invariant rectangular matrices. In Bin Dong, Qianxiao Li, Lei Wang, and Zhi-Qin John Xu, editors, Proceedings of Matematical and Scientific Machine Learning, volume 190 of Proceedings of Machine Learning Research, pages 97-112. PMLR, 15-17 Aug 2022.
[44] Lenka Zdeborová and Florent Krzakala. Statistical physics of inference: Thresholds and algorithms. Advances Physics, 65(5):453-552, 2016.
[45] Raphael Berthier, Andrea Montanari, and Phan-Minh Nguyen. State evolution for approximate message passing with non-separable functions. Information and Inference: A Journal of the IMA, (1):333-7, 2020.
[46] Cédric Gerbelot and Raphaël Berthier. Graph-based approximate message passing iterations. Information and Inference: A Journal of the IMA, 12(4):2562-2622 2628, 2023.
[47] Guilhem Semerjian. Matrix denoising: Bayes-optimal estimators via low-degree polynomials. Journal of Statistical Physics, 191(10): 2024.
[48] Joël Bun, Romain Allez, Jean-Philippe Bouchaud, and Marc Potters. Rotational invariant estimator for general noisy matrices. IEEE Transactions on Information Theory, 62(12):7475- 7490, 2016.
[49] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
13

---

[50] Hidetoshi Nishimori. Exact results and critical properties of the ising model with competing interactions. Journal of Physics C: Solid State Physics, 13(21):4071, 1980.
[51] Antoine Maillard, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. Perturbative con- struction of mean-field equations in extensive-rank matrix factorization and denoising. Journal of Statistical Mechanics: Theory and Experiment, 2022(8), 2022.
14

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The paper aims at studying Attention Index Models and its connection with self attention mechanism. We do so by providing analytical and numerical results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We state our limitations in the main.
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 a 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 or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not 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 discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
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]
15

---

Justification: We presented in this paper exact claims and conjectures based on well-known analytical methods. Specifically, we developed an application of the so-called replica method in this new context, so we can think of our results at the level of rigor of theoretical physies. A fully rigorous derivation of our results is a very technical and lengthy avenue, and is left for more suitable mathematical venue.
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 referenced in 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 a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: We release the code and discuss the details of the implementation in the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to 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 to 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, it 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 accomplish this, but reproducibility 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, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) 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 (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
16

---

reproducible, they 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We release the code and the data for reproducing the figures.
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
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is 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 to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips.co https.cccpubllii public/ htplllc/guidee guwdee/iieeeee /CodeSummd ss asionPolicyy 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.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments 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 as 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 test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Yes, in the supplementary material
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper 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, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide error bars in each numerical experiment and explain how they were constructed.
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, confi- dence intervals, or statistical significance tests, 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 with given experimental conditions).
17

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one 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 figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: In the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, 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 as 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 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: hwwpar/ps.ccppppaattt/ttciiitooo
Answer: [Yes]
Justification: Our work is theoretical in nature, and we respected the code of ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is theoretical on syntetic data.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
18

---

The conference expects that many papers will be foundational research and not tied 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 an 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 as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work is theoretical in nature and on synthetic data.
Guidelines:
The answer NA means that the paper poses 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 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. Lieenses for existing assets
Question: 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?
Answer: [NA]
Justification:
We use synthetic data.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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 the package should be provided. For popular datasets, paperswi wwwode..hooe...c.. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
19

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We synthetic data that generate ourself.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
We use synthetic data that we generate ourself.
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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We use synthetic data that we generate ourself.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
20

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https:/ /atee:///uuuririi /tw./ruir.cc/////e//// /aaaarppp...coooneeeeee////////2 for what should or should be described.
21

---

A Notations and model description
In this appendix we first remind all the notations and settings of the attention-indexed models. We then remind mathematical concepts and definitions that are present in the main text.
I,.... L as the layer index where L is the total number of layer matrices used, while a, = ...... T are the token index and is the total number tokens. Throughout this work, we use L, k:
Then i 1,... .,....d d are the indices for the dimensions and d is the embedding dimension of each token, and μ = 1... n is the sample index and n is the total number of samples. We will also use u, 0... ... as the replica indices from 0 to m.
We list below the specifics of our model:
X = Xo E Rxxd The matrix of T tokens (rows), each token of embedding dimension d. S₁ E Rdxd symmetric matrices for & 1.... and extracted independently from rotationally invariant ensemble (S(S)= Ps Ps SO) for any rotation matrix O. We fix the normalizations such that Epg [Tr S] = d and Eps [Tr S = R2d and with K1, K2 = (()). Contextually, assume that the empirical spectral distribution of S will converge to well defined measure μ.s. For the purpose framework symmetric matrices of the form S. = W W/vrrd where W Rdxr with entries Wi ~ NO, 1). We refer to the finite quantities pt > as the width ratios of each layer.
the analysis, we will specify our general
We define the AIM as the following model:
y= 9 (1 12)
(35)
with the generic map g B LxTXT
RR TxT which depends on the quadratic preactivations
x St as
dab Tr St
Nab
(36)
Vd
In the following appendix, we will show the tight link between the generic definition of the AIM with deep attention networks.
In the rest of this appendix, we recall the definition of the semicircle and Marchenko-Pastue laws in the context of random matrix theory. In particular
V4
x 1|
2)
Osc, A
2
xx
(37)
(1 p))(x) +
if p 1
(x)
2
p p>> 1
2TT
Finally, recall the following following definitions.
Standard normal pdf and edf
((z) e 222
= f @(t) dt =
((2)
(38)
Bivariate normal density and edf with correlation e
cxp 2² 2-
2(1-c²)
22(u,u;;) /" S di2
Φ2(u, U.C
di (39)
2π
We also remark that we formally define the Dirac delta function 8(z) (i the limit to zero variance of a centered Gaussian.
We finally define the row-wise softmax function with inverse temperature 3 acting on the matrix exp(Shrat)
h TxT matrix:
(hat)) Softmax (Shaa) E, (xp(3hab)
(40)
22

---

B From deep self-attention to the attention-indexed models
In this appendix we highlight the connection between the AIM models defined in Eq. (1) with those of two crucial architectures employed in the analysis of Large Language Models (LLMs), namely deep attention networks and their sequence-to--sqqueec (seq2seq) version. In particular, we show that both the deep atf-ntttntionn encoder and its sequence to to-seuuence (seq2seq) variant can be rewritten exactly as an attention indexxd model of the form (1).
We keep the notation of the main text and the previous appendix: tokens are indexed by a, b E [T]. embeddings by T E ER and every layer E [L] carries a tied key-query weight matrix S E Rdxd with extensive width r7 pid and rolationally-imvariiat prior Ps.
Deep encoder. Let X, E T'd be the matrix whose rows are the token embeddings, (Xo)a: = x
x
deep self-attention network with residual (skip) connection and readout strength e 2 0 given by the recursive formula:
X = [Tr + o Xe-1,
L f= 1, L,
(41)
where a : RTXT + R TxT is the row-wise softmax with inverse temperature 3 implicitly contained in the symbol (.). The function H is given by:
1 Vd
Hx(Xe11)
X Se x
Vd BulXXe 1
(42)
where the expectation is intended over the input data Xo. Define the preactivations
h a = Vd Fa
SeDD
Sr dab,
1= L, o, b [7),
(43)
Vd
and the sequence of token space operators
B = ll,, One verifies inductively that
Bº := [Il + o0 100-'110-B-100
BC == I=....,,4
(44)
X xB=
(= (=
(45)
so that every hidden representation depends on the data only through the collection (h In this way, we can write:
k00).
1
1
He(X,-1) = Ne-1 Ne-18
(x
(46)
Vd
Vd
1 Ne-1S.XT_
1
1
Vd Xe-1
Ftr BB--1o, St x, B
(47)
Vd
1
X S
TrS;
B
B.
(48)
Vd
Vd
Which is exactly the formula shown in (4). Furthermore, this map does only depend the preactiva- tions (5) in the following way:
TrS, Bé-1B
He(Xe-1) Vd (B- St (Bb-1xo)
(49)
Vd
TrS. B ( Vd XoSeXo T BL-IT_
BÊ-1 1---BB-1 -
(50)
vd
BC- (M(0) (b-1(M(0)+ +
TrS Iy BC-1T- TrSg 1
BC-IT
(51)
Vd
Vd
= BC-1 (2) B(-1"
(52)
For simplicity we restrict to the single head, tied setting: extending to multi head merely introduces an additional block index.
23

---

where we used the definition in (1). This completes our mapping in (7). In particular the output given by:
03 (HL(XL-1)) - Idoopl (1(1)
E RT'X
(53)
with gdoop (h ..... and has the intention-indexee model structure (1): the whole deep network collapses to deterministic multivariate function Sdeep of the L bilinear indices n ~ { S
(B (1(:--1)) (1) 1 Equation (53) is exact
Seq2seq variant. If the last layer keeps the token embeddings instead of collapsing them, i.c.
X 03 (H(((1)))) XI.-1,
(54)
with exactly the same algebra
Gseq (1) h....., Xo.
( = oz BB BL- (B B
(55)
Thus the seq2seq readout also an attention-indexed model: a (matrix-valued) function of the same quadratic statistics, followed by a fixed linear map Xo-
Note that in the particular case of just L = 1 layer the seq2seq map simplifies into:
X1 = Ovey (1.0) Tash)
Xo.
(h = a ( (1) B Be = (l) = Ideep (h()) (56)
hence conclude that, shown in equations (53) and (55), any L-layer tied self-attention network with extensive-width weights is information-theoretialll equivalent to an iilention-indexed model with L indices. Consequently all the Bayes-optimal analysis carried out in Secs. 2-4 applies verbatim to deep self-attention and to its seq2seq counterpart: learning the matrices { under the deep architecture is statistically equivalent to learning them under the iiiention-indexee model (1).
From this paragraph
Multi-head self-attention. Let heads m = 1,... M with (tied) weights s(m) and per-head Etr v7 Xe- ((()) XT1] standard multi-head layer
logits H (X(-1) S (m)) X
computes head-wise weights 0g ( and then aggregates (by concatenation or averaging) before a token-wise linear map; for our purposes, averaging:
1 X = C + M
Σ og a (m( (X--1)) Xe-1-
(57)
Applying the transport identity headwise and averaging gives
1 M M
M
H!!T (Xe-)) = B
1
Σ m=
BE-IT.
(58)
M
m==1
If we define the quantity (e())
M Em= ((())))) hence, the closed recursion is
1 M
Bº
cIrr + 00 (()) BC B(-1T) B
Bº = Lr
(59)
with output
y 0g 1 (2(2) Bd Bb-
(60)
and the seq2seq variant
X of 1 B! Bb-1 Xo.
(61)
Finally, also this multi-head variant of the model is again an AIM. Here, the collection of indices is simply enlarged (h (f.m)) Jem I.. m= with each ((..m) defined as in (1). The extensive-rank (m) (m)
regime corresponds to ranks rt =pp d.
"The explicit form of goeep obtained by inserting (45) with / = L 1.
24

---

C Bayes optimal analysis of attention-indexed models (AIM)
We study model described by the general setting:
( St Z
dab Tr S ab==l,.,
Da ~ Pout
(62)
vd
E=
with a. rows of X E R2'xd Sy E Raxd symmetric and y E RTXT Indices range from = ....
samples, with d, n >> 1. Instead the number of tokens and layers T, << d: we interpret Eq. (62) as Uy outputs generated by model of attention from data X that are processed in bilinear way through:
xli St 26 - bab Tr St L
8/ 9
vd
li=1
(63)
or following Eq. (6):
2/l T y..., (1)) BB¹¹ (1
Se 2
dab TI St
Ur Ideep (h
labb=l...)
18-1...
Rxx
(64)
vd
and
80b Tr Se a.b=I....TT Se x x/f
60v B:: x/
Se a
8ub Tr Se)
Pout
(65)
Vd
E=l....
Vd
In setting, the matrices St are symmetrical for each layer l and consider multiple layers indices 1,....,.. L. x is the a-th row of X for a, b = 1,..., T. Each row x E R has i.idd Gaussian entries, z ~ NO 1).
We define the preactivations
h P
x
Se x
sub Tr St
(66)
Vd
Since the matrices St are symmetric, are the preactivations of the model. Finally, for convenience, we rewrite the preactivations of the model in terms of the symmetrized sensing matrices
Bij ( j.b " 20rz8ar)/ 22d(1 8ah) RR
(67)
The preactivations of the model can thus be expressed as:
2 8ab Rab
(E) # Tr(S: ( (8)
(68)
In the rest of the analysis, we will refer to this equivalent representation of the model by considering symmetrized data ( that will just recall (8) μ while incorporating the factor /2 dab in the output function part.
C.1 Replica analysis of AIM and their state evolution
Starting from the posterior distribution of the model:
1 II Ps P(S1 ......................... SL(D) Z(D) (Sa)
(SI, ,..., N,. (SL, (69)
E=1
the replicated partition function of the model in Eq. (62) is:
(Z(D) = Eg.x / Emo1 umll
/IIII asy Po(S?)
IIII Poul Eoa("" asb
1 M.
(M Tr(S)2*)) (70)
V2 dab
where PS() the rotational invariant prior distribution each S,, and 1 (e) preactivations in terms of the symmetrized data as explained in (68). is the replica index, we work Bayes optimal setting. Above, / E (1,..., n. enumerates data samples, E (1...., ,, L.} indexes the distinct layers, {O...., m) indexes the replicas, and b (1...., T) are the token indices.
are the replicated
25

---

We compute the expectation with respect to the data exploiting the Gaussian Gaussian-equivalle equivalence principle:
Exó
h ah
p0 Tr(S) Z
Pa (.), ab
(71)
where Ph is joint Gaussian distribution with the means and covariances:
E [ (0... TT(S)Z") 00 Cov ( n (2) osd
d (2 dacobal Tr (Sp
(72)
We introduce the order parameters measuring the Sy overlaps:
QHe d Tr S (t). for l, 1, ..... u, u .....,
(73)
We enforce the definitions of the overlaps by inserting d-functions:
II II
of
Qix Tr St
(74)
t,k=1 aSe=0
and introduce the corresponding conjugate fields Que We insert
(ffQ dTr[Sp stl) = dyær expf Que (dO d Tr(S; S:l)
(75)
Hence the replicated partition function can be schematically written:
(Z(D) dsy dsy Po(S;)
/ III
f(I
II
dQwr dQux)
a.t x expl
uSe.ek
2
Qu (d Q
id x exp
Q Tr(S
uSa.ek
U.u.fR
p.)
ah M0 (hi ab
II
Pour (3 19
Rab }ab)
(76)
V2 Bub
In (RS) scenario, we let
Quk:
(u v),
Qui
(77)
gekr
(u + v).
and:
Qexr
if u =
iQu
(78)
Gtk:
if u 전
Hence, e.g. the exponent Eek, Qua d2 Qua becomes
d m + 1) Qu Qu
m(m +
Tek 4ek
(79)
4
Likewise, - - u, Que TTS can be reorganized in a form that leads in the limit m + 0 to
typical terms Q = or similar. Moreover Quy
dce.
2
So finally the replicated partition function, hence, takes the following form:
(Z(D) = II
dQue d'ue exp
d2
Σ
QUE
Inntout
(80)
2
U.S.B.k
with:
d / II
id Σ
dSy (o7 exp
Que Tr (S
(81)
u.t
uSo.lk
(e)
lo)
II
dh
v p
(e) 4 hab
Il Pout
v/{
ob
(82)
22
5ab
26

---

The free entropy per degree of freedom of the problem is defined as
1 lim lim lim de n-+00 m-+
1
(n()
(83)
After introducing n ad data samples, the free entropy decomposes into prior contribution and an output contribution:
== extr
qr
+ In + aToun00))
(84)
Thus obtaining the state equations:
10gTnn (4) i = da 8g (g)
(85)
(86)
C.2 Prior Term Computation
First we compute under the RS ansatz:
id
Σ QU Tr (E)
i d
Σ
m
Qui Tr (SP5K)
(-qe)) Tr (S:54)
(87)
2
2
MS==0 0
Quc d
m
qek
Tr (S
+
Tr (5
(88)
2
2
Qtx 2
d
Gek 2
Tr (K
ges d
m Σ
Tr (K)
(89)
We remind that each Se is a rank- ped rotationally invariant matrix of order O(d d). The prior factor that emerges from the partition function, after decoupling the replica indices by applying Strbtooo-tr transformation. reads:
In
dsy Po(S))
exp
i
(w.,)) Tr(S
(90)
2
(=1u=0
t,k=1 uSP=D
L
| dSP exp
Qen +
gek
Tr (S
+ 0,==0
gek
Σ
Tr (S)
(91)
2
2
4
t,k
1==0
= -fasP exp
Ge.k d
gnL
Tr(S
+
(92)
E
Ek
0,2
= |dSP((3)((Y) exp
qek d
Tr(S
L
Vaek d
Tr(S)))
(93)
4
2
F.k
f,k
= =/D ds dSP((S) exp
L
m+1
Σ
qea (r(S +
Vach
Tr(SaKr)
(94)
2
f.k
E.k:
where D(Y)) are GOE(d) measures V [7.] and Yt Raxd and also S RRdxd] In (78) used the identity:
Gy~ _ TYSY|]
= =
Finally, taking the zero replica m - O limit, we can write the prior contribution to the free entropy of the model as:
1
/ DY DY.... DYL (Y,, Y2: log Zin (Y,...... Y2:
Ii (n) = lim
d= d²
L
Zin ((Y)) () =
II dse (St)
(95)
d 1.
Σ
âe Tr(S) + Σ
Van Tr(Y:Sw).
x exp
4
t,k=1
E,k=1
27

---

The matrices Y E Rdxd are the auxiliary fields introduced by the Hubbard Strrattoottttt transfor- mation. Notably, they can be interpreted as "noisy measurements" of the S, matrices with coupled indices. In particular, the denoising problem which is solved by the free-entropy contribution of the prior is:
Y = Σ Vâe S + Z Vi.j.e
(96)
with Ze GOE(d) matrices and Se E Rdxd rotationally invariant matrices, leading exponential term of the form of - De Tr((2, v Sk y,) Such equivalence between the matrix denoising problem in (96) and (95) is analogous to those of [42, 40].
C.3 Output Channel Computation
Starting from the replicated partition function in Eq. (80), we can see that the output channel contribution to the free entropy of the model, factorized with respect to the data, is given by:
/ dy
(V
(P)
]
Tout( (1)
II u,e,aSh
dh
(5) P (1a
"aab)
I u,
Pout
ab
(97)
dab
where we consider only the upper triangular token indices S b. Pr. (1aa (0)m))
bab) is multivariate Gaussian distribution with means and covariance:
Cova. (M Megd 7 (2 Sacdna] Tr (S) Sg) dad Qik
(k)
ab 0,
(98)
Under the RS ansatz and in the limit m - O. we can decouple the replicas through another Hubbard- Stratonovich transformation. The exponent involving (E) becomes:
Σ E
Σ
(5)
(2% ) ab,ed
we.t
(k)
(99)
w,o=Dasbesdd E,k
Substituting back, the output term becomes:
Ig) (4)
11
2h exp
-
-E 10 (Ex Inshes
Pous
(100)
For a fixed channel and for each token pair (a, b) with S b, the covariance in the replica space is given by
(In) (En)aste
uo.ck = Sacdbal Qa [2 Sacdral [(et gek) duv + g0x
(101)
Because of the Kronecker structure dac the Gaussian law over all (h (f),
Ja factorizes over
token pairs (a, b). Hence it is sufficient treat fixed pair (a, b) and then take the product
a & For notational clarity in the next steps, we temporarily fix pair (a, b) and write (f)u h (f) For this pair, the covariance across replicas and layers reads
E (1 100 [00] = 2 (Qex - (ex) dwv + 4ek
M1, v. = O,.., m, 1.=11..... Lt.
(102)
We now construct explicitly family of Gaussian random variables with covariance (102). Introduce a shared Gaussian vector w = (1) (4)) with mean zero and covariance
E (k)
2
(103)
For each replica u = 0g...,my m. an independent Gaussian vector Ed = (g())),,.. "el)) with mean zero and covariance
E B (+)) 2(Qux qek) dav-
(104)
is independent of all (E" 1m
u Define for each replica and layer e: h (0) == (() (()
(105)
28

---

Then, for all e, k and
B (000
E (wll) (w))) + (())))
=E (()
B e(0ug(h)e)
2((0xx gex + qee)
2qe: + 2(Qxk - qex jmu .
(106)
which is exactly (102). Therefore the law of (h(5)) t,u is exactly the RS Gaussian law.
)e, From (105), conditional on w the replicas are independent and
(h(0)) 18=1 (N(W! v).
V(Ek) = =2(ee -9ek)-
(107)
Equivalently, for each u,
o( (n (n(00 ww)
1
ΣΣ t=1k=1
(n())u -w IV- (Ni w(th)
exp
det(22TV)
(108)
Moreover i N(),,24, i.e.
1 L
p(w) det 2g)
exp
w()) [(2g) -le wo(k)
(109)
2
8m1l km
By marginalization over wi the joint density of {h" " is therefore:
p (()00)e))) d' p(w) I
(1000)e
(110)
u=0
We now express (0, 2q) through standard normals. (1)), ...,,.. (()) NN 1) and define
Introduce iidd auxiliary variables
24)er (r),
(1II)
where (v2q) is any matrix square root such that, for all e, ki
Σ
((2g)err 24))kr (24)ek-
(112)
r=1
Then w has covariance (103) since
(ww = [ 24) (v 2g) (1)0) = Σ
(2g)ex.
(113)
r.s
Consequently, for any function F(w),
Σ I(V2))onn (2 )" 2
f N': F(w)
dr 1100093 F II
(114)
2
Restoring the token indices, for each (a, b) we introduce independent (r) N 1) and define
Mab
24 nhb
((B)) 2(QQe gek).
(115)
Using (107)-(114), the Gaussian exponent for each replica and each token pair (a, b) is therefore exactly
1 1.
hab
(6)
(() |V 1V= (1)
-((a))
Mab )a.
(116)
2
2
29

---

which is the desired form with w and V identified as in (115).
Define the auxiliary measure
(0)
(4) (7ab
Dn= =
exp
(117)
2T
2
aSbi=1
For fixed v, the replicas independent and identically distributed through the conditional Gaussian N (haati Wab, V) on layers l. Hence, introducing the single-replica output partition function
hab N (hhaiiabb V) Pout (1) y / II ag
(1) (1,2222- V2 dab asd, Lass, )
Zou (y,
(118)
we can write
Lour (g) = [[ dy / Dn (Zont(s, m-17
(119)
where w and V are the functions of (g. Q) defined in (115). Expanding Eq. (119) for small number of replicas m - and using Am-1 A A in (1 + mnn A + o(m)). the contribution the free entropy is governed by
Tout(q) = f dy
Dij Zont (v, w. V) In Zont (y. w, V);
(120)
We aim to compute 0, out (q) to finally reach the state equation in Eq. (86), namely:
Te 4a f dy Dn [1 In Zou (v,w, w. V)]
i Zout (y. V)
(121)
aye
It is convenient to rewrite the 77- n-integral as integral over itself. For each token pair (a, b) have Wab ~ NO, 24) with independent draws across a < b. Thus we may equivalently write
Tout (1) = dy [[I d N ((was 0, 24) Zout (y, w, V) In Zont (v, w. V), o5b
(122)
where w (w Shshe Now define for each (a, b) and layer &
(gout w, V = a In Zont (y. w. V).
(123)
In (122), q appears in: (i) in the Gaussian measure Wab ~ 24) and (ii) in V = =(Q 9)) inside Zout- Let
F(w, q) = I |dy dy Zomt w, V) In Zont (3) w, V).
f
Then
Tout (4) = / d' ((wuii 0, 24) F(w.g).
(124)
asb
Now fix single pair (a, b), under Wab any smooth G(was):
NO Σ) with E 29. we get the Gaussian identity, for
8
@GG(wab)
B Buss -NV (G) = Buss
(125)
(8) aw
Since 2y, we have 0/0qm 20/00200 hence
0
(GG(war)
Eu.s (1. [G(wat)] = E.i..
(126)
aw ow (R)
ah
Applying (126) to (124) (and summing over independent pairs) yields
Ø Tont (g)
82 g"
+ OF ((
|II
d wab N (Wasi 0, 24)
Σ a<o
(127)
Oqek
(f) (R)
Oger
(x)
30

---

where 0F/0qlw is the explicit derivative through V = 2(Q q) at fixed w.
Since (rr)) = 2(Qrs - (rs), we have (v(rs) /0qux = -2 hence
@F(w,g) OF(w. 9)
(128)
Eqca
0V(ck)
Plugging (128) into (127) gives
oout (g) dqck
/ |II
wa (Wab: 0, 2q)
Σ
F2F
OF
(129)
Ow
(2) Ow
(k)
(v
ash
Now fix w, and define for brevity
Z(y) = Zone (y, w. V),
0. (v) = (oout(( V)
(130)
Then
((w))=// dy Z(y) In Z(y).
(131)
We compute the two derivatives in (129) exactly, Since Z is an integral of a Gaussian density N (habi Wab, V), then by definition:
(Z(y) (2) = Z(y) (v).
8 Z(y)
dg (v)
(k) = (1v) (v) 9ab (v) + ((y)
(132)
aw aw
aw
Therefore,
OF
/ dy (1 + In (2))
(1(y)
/ dy (1 + In ((1)) (vv) (k)
Fab (v),
(133)
(k)
Owyb (k)
F
dy (Z(y) 9ab (v) deb
(6 (v) + (1 Z(w)) Z(v) 2 (w) a% (v) Z ب (v))
(k)
(134)
Ow 10 Ow
For Gaussian integral, the derivative of in Z with respect to V satisfies the identity
In ((v)
dgab (v)
(E)
Dab 9ab (v)
(135)
@V(())
2
Ow (k)
Using Syrcs 2 dyyess and (135), we obtain
OF
= / dy In (()))
(Z(y) (VIR)
2 dy In ((v)) Z(v)
(9.) (v) 9ah (.22
(v)
0. a. (w)) (136)
0V((R)
Subtracting 20F/Ov(th) from F2F/ (Ow ow)()) using (134) and (136), all terms proportional to (1 + In ().....) cancel exactly, yielding
OF aviek)
S dy Z(v) gab (v) (8)
2
gab (v)-
(137)
Ow
(C) dw
(k)
Plugging (137) into (129) and summing over S gives
0l.n((g)
d' d' N 0, 20)
Σ as*
dy (y.w,, V) (Bout (y. w, Vi (Nows (v. VDD
(138)
0qp
Therefore the output-channel state equation is
Σ (gout (v. w, V) (gout (o, w, (A)
gea 4a Ewryy
(139)
which is equivalent to
qex- = 4a En,y
Σ
(goue(y. w, V)) (youe(y, w, V))
(140)
asb
for
l t= L,.,.., a asb= I...,T
31

---

The expectation ((na)) is taken over the joint measure
II Di
(141)
t=1
and the output is drawn from the channel density
Poun y (1)
nat ~ Mw!b ).l.m)
(142)
with:
Pontly (1:01 12) 8((uab a((h(a) (5) Ive)ab
(143)
or particularly, for the deep attention case:
Pounty) y
(1 8((yab Bº (t Jab
(144)
C.3.1 Recap of the state equations
In this section we summarize the findings of the previous appendices. We performed the Bayes optimal analysis of the attention-indexed models (A (AIM) defined in Eq. (1): we found that the problem can be split in two components, the former involving the (extensive width) rotationally invariant prior channel and the latter involving the output channel part of the model. Through replica analysis, we found that the prior channel is described by the following function:
/
I
L
dse Ps(Si)
(145)
Σ t,k=1
Tek Tr(S.St)+ + Σ
Va Tr(YoSa)).
exp
t,k=1
The denoising function associated the prior channel assumes the form:
9in (Y = Oaray (Y,
(146)
The free entropy contribution coming from the prior channel is:
1 lim ] d=o00 d
DY....
Ii (g)
DV, ... DY, Zn (Y, ..... log 2m (Vn. (V....., Y2:
(147)
where DY stands for integration over GOE(d) (Wigner) matrix Y.
Notably, this problem equivalently mapped to same posterior distribution the following matrix denoising problem:
L
Y(S, A), = St
Σ
Vm Em GOE(d) St Ser Ps
(148)
m=1
On the other hand, the output channel of the model is described by the function:
/
II d hab N (habs Wabi Vab) 8 (v- -(t / / //2 Sorll
Zout (y,, w, V) v)=
(149)
with:
Wab
24
Vab
(k)
((k)
2(Qxx - qFk)
The denoising function associated to the output channel takes the form:
Jout (,,w, V) 0 In Zout (y, w, V)
(150)
The free entropy contribution coming from the output channel is:
II dyab / DDn... ... Dnn Zomt (v, V) log Zoot (y, w, V)
Lour (g)
(151)
0,5=1
32

---

where Dn stands for integration over a L x T X T tensor symmetric in the token indices and with independent entries 10)
To conclude this section, in the multi-layer setting described by the AIM framework in Eq. (1), we found the following state equations:
qek
4alBg,
Σ Dout
(9 ((
V2 6ab }) ,w,,
Vab
(6)
ash
9 ({(( h(w, V)ab
(1))
(152)
x gout
/2 - dab :4;
lim EES Tr [9ral (Y (), Degm (Y ()) d0++0 d
Qtk
where E intended as the average over L T symmetric in the token indices and Gaussians with zero mean and unit variance. Moreover, the average Es,y is with respect to respect to Y as given Eq. (148) and Ps. Finally:
[h(w, V)at w/(0)) + Σ
VV
(EK) (k)
Sab
(153)
C.4 The fixed point of AMP is described by the state equations
We start by defining new variable w ab such that Uμ = g( ((( ab/ //2/2/221 dab) Jab)ag)), where we can assume that [[( at (w), ab)& 2Q Our first step is to define the quantities m' and q' on the iterates of AMP
m = Tr$$t
de - Tr(
(154)
We now claim, in analogy with [44, 35, 40, 42] that for every sample μ and every couple of tokens S. 6 the variables wit ab at each time converge to independent centered Gaussian variables with the following covariances
[[[ ab)e (w), at)k] = 24/x=
[[(W% ab)e (w/, at)s
2m'ex
(155)
By Nishimori's ' identities [50] assume m² q² The first equation of (21) is now immediately recovered (modulo the substitution - 2(Q - d²) which will come after)
T
Pe 4adE w
Jout (y, wt, v2y(0) yout (y, w Vt)(ss Jab
(156)
a
where
(f
w
T
Wh w ab
~ N (0,
29'
2m")
)
9
8ab
2m
2Q
(157)
asb
Again in [44, 35, 40, 42] will have that in distribution
R S + (1')
(158)
We are ready to close the circle: going back to the definition of 9² we write
d/s ER Tr [9m (R', (2,.4") 9in RR' (e)."')) /dd
(159)
which is exactly the second equation (21). Notice how the expectation is taken over the random variable R2 in (158). The last step is to notice that
C = Tr(( ST)(St S = Q
(160)
such that V²= =2(QQ -2QQ-9) g¹).
33

---

D The case of L = 1 layer
In this Appendix we restrict the theoretical results derived for an arbitrary number of layers to the particular case of L layer. In this particular case, the order parameters g and ① become scalar quantities. Moreover, in the following analysis we specialize to the extensive-rank choice:
S
WWT R Vrd
We WeRd gdxr (W);
(161)
with rank ratio p p= p=r/d= = O(). Thus, the spectral distribution of the symmetric matrix is that of the Marcenko-Pastunn law for Wishart matrices described in App. (A).
D.1 Prior channel state equation
Starting from Eq. (95) for = layer, we get::
/ DY ZYn (Y: (Y:
[i](q)
lim
d=000 d2
(162)
m); /
ds [ds Po(5) exp (-2 (
!
Tr Tr(ST)) - Vÿd
Tr (YTS)).
2
Again, at the O-replica order Q = 0 and integrating over Y:
/ DY QY;
(163)
= =f DY ( ds PoS))e) (- qd Tr(ST 5) + Vēd
(-4
Tr (VTs) S)
1 Tr (yTy)
2
Po(S)exp Po())(() = -f ffss
_ Tr (STs) exp (- Tr (ST (ST))
qd
[ as (o())
(164)
Now, note that the exponent in 2n (Y: q) can be rearranged as:
vdd TrsS")) 4 -
4d Tr(S 5)+
(165)
2
Observe
T( = s) 2 TTTT) + IYY"
(166)
Hence
- + v
I()Ts) - tT(rTyy)
(167)
2
Therefore:
10()= = exp Ir(Y x fds P6(5) exp -
(168)
Ignoring the factor exp( Tr( that is independent of S, we see that
Po(S) as as which plays the role of posterior density for S given Y V + In the large-d limit, let us parametrize S by its eigenvalues:
f as
exp - Tr (T(VVs->)*
(169)
with Z a GOE(d) noise.
S = AU
(170)
34

---

where A = diag(A: ,....., Then
ds
dd, A(1))) du with A(1a))
II
(171)
Igisjdd
Then the exponent
TT(-A(VS--y)
(172)
becomes
- Tr - (VQUAU=-X)
(173)
We can factor out the integral over U E O(d) and for large:
Soco exp qd Tr|AA DU
exp de LHCIZ (q: Pa. px)]
(174)
where /HCIZ is an explicit functional in the limit d + xo of dimension 2/d times the log of that integral, and μa is the limiting spectral distribution of A/Vd. The prior contribution of the free entropy is given by
4 prior (q)
lim d+o0 d²
E [l))(
(175)
or more explicitly:
prior (4)
lim d=+00
E
[in
/ Po(S)e - T as].
(176)
d²
This term can be explicitly computed and mapped matrix estimation problem. i.c. a denoising problem as follows:
prior (g) lim By Iol) -Q + IncIz (q: no, po (sc,1/vā + const, (177)
1Q
d=ano
where Q 1=1+ p.Then, one has the relation from [51]:
1 [(ma)= +
3 + In +
4q Eng [x
2 [HCIz (i: Mu- (4)
49 Eno [x = 0,
(178)
2
8
where we have defined f4g = nr 2/ and E(s) is the noncommutative entropy:
Σ(m) = / ((de))(dd)) In ja y/-
In our normalization (with Q = 1 + p), rearranging yields
} IncIzzIs no, p4) 1 [(m)) 1 2Qq +
1 In q.
(179)
2
Plugging back into the free entropy, we obtain
① prior (4) = q Q +
1 + const.
(2Q4 1))
In
(180)
8 4
Taking the derivative with respect to yields the "prior state" equation. In fact, differentiating we obtain
OΦ
Q
1
+
(pa)) = 0.
(181)
Oq
4
49
2
0q
Using the derivative:
8
2nT
/ Mg(x)² dx.
(11b)
(182)
0a
39²
this condition becomes
9
+
Q
+
7²
[x( do =0
(183)
4
4
49
30
35

---

which is exactly our desired state equation.
To sum up, in the problem
Y V Z, ~~ GOE(d))
(184)
the law of Y is asymptotically Mss
+ sc, 1/V7 we finally get:
/ /[ur(o)] dx, 3q²
4T²
1 1= Q
(185)
with My = jas 1A se, 1/v
For the computation of py, we recall that if ju = 145 ss we can write
Ray (z) (2) + Ro (2).
(186)
For the semicircle of radius we have Rosc, a (2) a² 2. For ps (Marchenko-Paste distribution with parameter p), we have
R Rumro (2)
P
(187)
vP
In our case 1/Vq. then
Ruv (2) Vp-2 + a² a²
(188)
From Ru (z) = ( 1 one obtains an equation for guv (2), with:
9uv guv (z) = / (y((dx)
(189)
So, using the identity = + Ray (I), where x = guv (2). So we get
x +
P
+
a²I..
(190)
vp- x x
Hence the final polynomial in is:
1 1 - VP 02)) + + Vp
1p
x Juy (2).
(191)
VP
We look for the solution of this equation with largest imaginary part. Moreover, we compute the discriminant of this third order equation in order to correctly quantify the edges of the spectral density we want to numerically compute.
Recalling a² 1/9. the imaginary part yields by (Stieltjes- Perron inversion), i.e.
1 Im guy (zo ie).
(x() lim c->0-
(192)
D.2 Small/Large width limit of the prior channel
We recall that the state equations of the form
Q Q- 1 4x²
de #1)0(x33
37
(193)
= 2Q q. 9)
where M1/4 is the spectral distribution of S. +
1 and Q = d Tr(S?). In our examples, S" is vp
vp times a standard Wishart, with Q + p.
36

---

D.2.1
Small width limit
We follow [40, Section E.1.1]. Call = p/q. and ā = =//p. Call v the distribution of vp(S. Vq 7) vpS. vIZ, i.e.
-
v(y) - /1// 1/2
(194)
Notice that this is precisely the v defined in [40, Eq. 56]. Then we have
1
4T²
/ de dx
Q q
3q
t 2x
dy 10 1/2 12/01111000000 1/3 112))=
122
P
3p2
t 2
dyys
(195)
P
3p
t
4x²t
/ dy m(y)
((2
t)
if t
if t>1
where used [40, Eq. 57 and following] take the limit of small leading order. Thus, the equations can be recast to
t(2 t)
if t<1 t>l
Q q
(196)
1
2QF - 9)
In particular, we have a weak recovery threshold. Indeed, as long as
ā
(197)
2F(1)
we have that Q q 1, i.e. the same error as the average from the prior (BO estimator with no data).
D.2.2 Large width limit
Recall that 1 p and q E 10, 1 + p), so that - q E [0, 1] even in the p limit, Then have
4m²
Q
dx m1/((x)²
3q
4
/ dy dy[vPHH/A(VPP)))
1
3qp
4
/ dy (1/1p(()))
(198)
3pp
+ q
1 1
where we used [40, Section E.2].
D.3 Output channel state equation
For L = layers we obtain the state equation for the output channel contribution:
q = 4a (na)
(gou (v,w, V))a
(199)
a
37

---

Where we remind Nat ~ (0, 1) with I = 1,.... T and wab v29 Mab. V = =2(Q - 9).
Q 1 + p. The denoising function is given by:
(gout (y. w. a) = Owns In Zone (y, w; V)
(200)
and:
Zout (y, w. V) II Caaaaa((haaa Wab V) V)g(y f(h)) agb
/
(201)
where fhh depends on the precise choice of the model. particular, in the following sections consider the three cases dealt in the main text. In the following, we first consider a linear output channel for generic number of tokens T. This simple case serves as baseline for the more interesting case of the softmax channel, namely the self-attention layer for an arbitrary number of tokens. We also consider the hardmax variant of the model treated in the main text for T 2 tokens.
D.4 Linear output channel for generic number of tokens
We consider Paun (yab haat) 8(yab
Nab V2-626
Then
Zout W, V)= / |II N MhabiWab. v)]
hab
II 8 yob 8
dh.
(202)
V2 dab
asb
a5
Enforcing had = /2- dabyab this gives directly:
Zont (y, w, V)
exp
1V2-barMaa - wab)
(203)
1T
2V
Hence
(v2 - dutilab wab)
In Zont (y, w, V)
Σ
- 1n(2 V)
(204)
2V
asb
We recall that (a) depends linearly 1abb e.g.
Wab /2q Nab, Vab = 9(Q- Q = 1+
(205)
Then
9 (V2 dobyab wabl (1))²
= (v2 dabMes wab)
(Domt (2,w,, 42 Mas
In Zont (v, V)
Owlab
Dwab
2V
(206)
Thus
Σ
(goue -))a==
[(v2
Sablab war(1)) 2(Q 9))
(207)
ash
We can compute the expectation:
E
Σ
Jout (v, w, V))ab
azd
(208)
/II a<b
daab
dyoo
Zout (y.w, V) Σ
(gout (U, u. V))a0-
V2T
asb
We can simply use:
f /dyas
1 /2T
(V2
dariyab wab)2
exp
[V2
Sabllab Wan]? V.
(209)
2
Therefore,
[II
[Tdya) Zone (ys asb
(y, wi V)
Σ
(oout (v, w, v))a.
Σ
2(Q ))
1 )2 vl
(210)
asb
asb
38

---

Lincar chauael with T ==
-005
0.8
Tormnrion anr
0.6-
0.1
0.2
0.00
0.10
0.15
0:20
Sumple comupeesity
n//2
Figure 3: Illustration of the Bayes-optimal error for the linear output channel baseline in Eq. (202), for T tokens and several values the width ratio p p= r/d. The model reaches zero BO error at finite a. The recovery threshold matches perfectly the one find by the simple counting argument in (215), plotted in short vertical lines.
Thus each term becomes
1
1
2(Q - 9)
(211)
2(Q 9)
2(Q 4))
Hence the entire sum
Σ asb
1
T(T + 1)
1
(212)
2(Q 7)
4
Notice that this result does not depend on Consequently, the outer integral over becomes I. L.Hence we arrive the final form of the linear output channel state equation:
T(T
q = 4 cz (m.w))
Σ (goue (y, we Vass
4a: Σ
Pon (y, V)) 4a
(213)
4(Q ④)
a5b
which finally simplifies into the output channel state equation:
T(
(214)
Q Q-q
As example, in Fig. 3 we show the fixed point solution for the state evolution equations for the linear output channel. The prior equation (26) remains unchanged, while we use Eq. (213) for simulating the linear output channel results. We also show in vertical dashed lines the recovery threshold found by the simple counting problem:
T(T + 1)
p² 2
Oroount
(215)
2
The linear output channel matches perfectly the counting recovery threshold, unlike in the softmax case shown in Eq. (32).
D.5 Softmax output channel for generic number of tokens
We compute the quantity:
Zont (y. w, V) = / II dhab
Ohat w
2Vab
II asb
8(yab Softmax (
har))) (216)
12 Vab
/2 Jab
asb
where we remind the factor v2 dab is present due the symmetrization of the problem (i.e.
multiply and divide by /2 dah). allowing much simpler treatment of the BO analysis in change of this slight modification of the output channel.
39

---

From now we define the quantity Tab = V2-
dub-
We thus aim to compute the quantity:
/ II
hah
Zout (y, V)
dh o(
II
(habs wab. Vab)
(217)
Tab
agd
asb
We introduce the variable Zah
Nab/Tai Tah and exploit dNW (h, u1 o) dzN (2/T we get:
(y V) = / II
dzab 8 o())
II
Hab
Vgge
Jab
agb
asb
(218)
II agbeT
dtab (tass
Matk
Sa
Vuu
Il ds N
Wart Tar
Vgg
where in the last equality we introduced the inverse mapping of the row-wise softmax function, defined in Eq. (40). In particular, we introduce:
Beas
B2ab
Mto
Yab
va bb<<r
(219)
Eb=1 BBaas
_b==1 @xaa
VaT
which leads to:
Yab
tab
log(
Dab(y) Va <b<T
(220)
yoT
while for = T:
Bara
BENT
Vra
Man (TTT) - 8a = STT + Φra(u) Va <T
(221)
VTT
BeTr
@0oT
having introduced the change of variables:
Zab
tab
Zab
2aT
Zab tab Sa Dub ΦTa STT
Va Sasb< <
(222)
and
TTT 9. TT Z 8TT ΦT. Va < T
(223)
Having this mapping clear and introducing the short-hand notation W w/T and V V/T²,
STT = I, we can see that we can reduce the computation Eq. (218) to that one simple sealar integral in the variable x namely:
Zout (y, wi V) = / drN VTT) II
N(x + (rray), Var)
x II (Vhalu)) + Φra (v) + xr, Wab: Vat)
(224)
asb<T
= =f dr exp
(rr + ΦTa 2)
(dab + ΦTO + wab) 2
Var
Vab
aSb<T
We thus obtain a simple gaussian integral whose exponential is of the form:
1
Øra
Wal
Σ
Dab + Φra Wal
r
+ 2x
2
Vor
Val
ash
a=1
aSbeT
(225)
(hTo
woT)?
(Oab + ΦTa
Wab)) p)]
}
VoT
Vab
Having computed this simple gaussian integral, we can hence compute the quantity of interest:
log 2 =
1
Ta
War
Dab + ΦTa Wor 2
2V
Var
Vab
asbeT
(226)
(oTa -2ar)2
(dab - ÔTa wab))
cost
Var
Vas
aster
40

---

with V = Eess Vabl and again (ab(v) = log Sakb. Wab V2-0ab
Vab
2- we 22quab
Vab
yar'
Vas = V = 2(Q q). h ~ (W V), y = (hh). The constant term contains those terms independent from w, as we are finally interested the denoising functiin, which the derivative:
(V, V)ab = Owab log Zont (V, V)
(227)
We thus compute the denoising function deriving with quantity log Zont (V, w, V) with respect to W thus computing Tz Tuy log Zone (y, w. V) for j T and for j = T. We also consider that Eggs ab (2 dab) +.
2(Q g) q).
Finally, we obtain the final form of the denoising function of the softmax output channel in Eq. (27) for arbitrary number of tokens, substituting back the original V and
T
T-1
(ggout =
Tij T
TabWab
Σ
2 bab
+ + wiy +ō(j < T) (228)
a
which is exactly the same form appeared in the main text in Eq. (31).
We now complete this appendix by computing the quantity Byy Eoso (fout) To do so, we exploit the following relations:
Dab = hab har h NN V) - Tabhat = 22 Mos + vV Eab asbst
(229)
with Jab, Eab ~ (O( and
ΦT hTo hrr haT hTT
(230)
We thus substitute these relationships inside Eq. (228) and finally compute En Easo (gout) ) After long but simple algebraic calculation, it is possible to show that the denoiser function reduces to simply:
T
(gout))s = == ==VVVVETT
Tab
Vab +
Ty V
T
T/T
a
Tij
VV + -(j < TWVVE 8(j < <) Túj
VVETT
TTT
TrT
T2
(231)
Tij
Tab VV + VV T) + +06 TVVVE
a
2
Tj
Tab VVab + VVEij
=
2
ab
which finally gives:
11
2
E V
Fontt 135
Tij Tah Eib Eab
2
T2
15
$5
osb
1
2
Ted BSa
T
ogch osd
(232)
T(T 11)
T
T
T
T
Tab
2
T
7
i5j
asb
2
2
2
Hence, we can finally conclude that the output channel state equation we obtain for self-attention layer with an arbitrary number of tokens is:
4a
+ - 2)
aTT TT 2)
g = 4a Σ (oout) )7
(233)
2V
Q-9 q
41

---

1.0
Softmax T=4
Estimation error
Softmax T=5
AMP T=4
AMP T=5
0.5
0.0
0.02
0.04
0.06
Sample complexity a = n//
Figure 4: Comparison between the fixed points solutions of the state equations for softmax output channel in Eq. (26) and Eq. (233) for T = 4, tokens. We compare the theoretical solution with their corresponding AMP algorithm run over 16 different realizations and with 120. The error bars in the AMP dots are computed with respect to the mean value.
Softmax
1.0
p 0.00
0.8
01
Esiamtion arrur
0.6
0.4
0.2
0.0
0.1
0.2
0.3
0.1
0.5
Sample complexity
Figure 5: Low width limit of the self-attention model for L layer and 7 2 tokens in Eq. (27). We rescale the sample ratio â n/ dr and plot several values of the width ratio p r/d. We correctly predict the weak recovery threshold in Eq. (197).
which the result presented the main text in Eq. (32). We highlight this final result holds for any value of the softmax inverse temperature 0 < 3 < +00. In Fig. 4 we show for completeness the state equations for the softmax output channel described by (216) for T 4, tokens and its corresponding AMP run for 16 different realizations and with d = 120. The error bars in the AMP dots are computed with respect to the mean value. We find a good agreement also in this case. In Fig. 5 we plot the low width behavior of the self-attention model for L layer and T' 2 tokens Eq. (27), for which we recover the simple output channel state equation in Eq. (233), thus giving:
- T
F(Q ()
(234)
2(Q -) g)
Regarding the large width result in Eq. (198) in this softmax output channel case, we get the equation:
q a(T2
+ = aa(T2 T 2) + +)
(235)
Q-9
so we get
a T +
q = 1 - aT? - T + 2)
(236)
which gives the large p result:
MMSE
= 1 - a(12 -
(237)
1 q
42

---

D.6 Hardmax output channel for tokens
We now discuss the hardmax output channel case, in the special case of T 2 tokens. Following Eq. (28) in the main text, we need to compute the quantity:
(y wa V)
/ II
dhab ag
1
(hab mab ) 2a
II (yab hhard(
1 hob Jaab))). (238)
2 Vab
v2 dab
with:
Thard (ZI...ZT)) = 8(i arg max Tg)
(239)
j
In this setting, in particular when T = 2, the output label of e.g. y1 becomes /11 0hh1 - h12)
(240)
and similarly for the other labels. Here (u)) is the Heaviside function:
u > 0,
( 0, u <.
(241)
For the computation of the quantity in Eq. (238), it is convenient to make change of variables by introducing the differences:
h12
h12 V2
== -1111
h22
(242)
v2
We thus rewrite in the T = tokens:
Loue v) = / dd dhhniw(hh Via) (EERIOO)
h12 h12 du N(u+
wi. Kin
(243)
x {1 (VE((w)) du N( +
h12
w02, V22)
V2
where the integration ranges defined by the ttreshhhehhodd:
R(yn) = (u 0}, if yı = 1, (u 0}, if yn 0, Now, by shifting the Gaussian factors we have
{0 0},
yf 1,
R(y22)
{1r 0),
if y22 0.
hi2 V11),
N(u + h12 Vin) = Ai wn1
(244)
V2
and similarly for the o-integral. Thus, the expression becomes
Zou (w, V) = f dhi2 dh12 ((122:1122 Viz) Fin
h12 w) Fez h12
w)
(245)
/2"
/2"
with
Au
Fu1 h12 w) w))
JuER(UN) du Nau; W11
/
h12
Vnn)
$ =((111
W11
/2
(246)
2
VVn
v
Aya
1122 F22
h12 2
W22
w)
du N(; w22
V22)
+ 922
v2
(247)
V2"
JuER(
VV22
where (2)) is the standard Gaussian CDF and
+1, 811 = 1, 1, U/1 0,
+1,
U/2 = 1,
811 2y - -1
$22 2y22
1,
Y/2 0.
Thus, in the thrd-threshold limit the output channel integral is given by:
too
dhiz ((112:012, Vi2) + 811 ((11
W11
W22
h12
Zont(y: w, V)
822
(248)
VVI
V22
43

---

We can further manipulate this expression.
Writing 112 w12 vV12 Z with Z ~N( 1); then, using independence,
Zout (U w. V) = BB 2 ((u1 X1 2 - >12)]
(249)
where
v 12011 W12
V2W32 - W12
= S11
U2 = 822
(250)
V2V
V2V22
V12 A = /
V12
8221 = =221
(251)
2V
2V
A classical identity for jointly Gaussian variables gives
(x + bZ) ( + dZ)] 2
VI+b
VI+42
hd ((+++)(++4))
(252)
Where 22 is the edf of the bivariate normal density defined in Appendix (A). Applying this relation to our model yields:
Zont (3), we V) = D2(( R1, N25 c)
(253)
with the compact parameters
12w11 W12
V2u22 W12
K1 = 811
K2 = 522
(254)
12 + V/2
V2V22 + +122"
V12
1
$21
$11822
(V11 1/2 Vi2).
(255)
(2V11 + V22) + V12)
3
We can hence compute denoising function:
(gout (y, V))
In Zout Zout V).
(256)
Owab
Because Vas is cu-independent, the chain rule gives
i (i2; = 0ki
dkL owi1
V2511
K22
Φ2(111, Ke: P12),
(257)
Owi1
Own
22V11
V12
The four independent derivatives therefore
/2811
(2 P12)
Pout (y.w, w, V)11
(258)
V2V11
+ V12
12 K2: P12)
V2822
d (n2, K1: P12)
gout (y,w, w; V)22
(259)
V2V2 + V12
22 K2i P12)
Jout w,, w, V)12 This expression can be compactly rewritten as:
((22
811
+
822
2 (Ni, K2; P12)
(260)
22V1 + V12
22 + V12
22 K22 p12)
6lk,,22,0) k2, Jourly, w. V)ab v 6[Q 9) 4 (k), k2,
221 ((11 + 82)
- (81 + 82) V282
(261)
ab
where o(k1, A2, c) is the p.d.f. of bi-variate Gaussian with zero mean, variances 1/(1 c²) and covariance c/(1 c²), and & k2, c) is its c.d.f (see Appendix A). Moreover, Sa = 2yao 1, k, = s,( 2ww2 - w12)/ /6(Q- - q). c = 8182/3 and Wab = /2g Nab- This is precisely the result shown in the main text in Eq. (29).
We finally compute the state equation corresponding to the output channel, namely:
q 4aE n.u
Jour (v, w, V)2
(262)
a
where Tab ~N(0,, 1) for a 5 b b= 1,.... T and y Zout (v, V).
44

---

D.7 Generalization error and sequennce-fo-sequenee version of the model
In this section we draw some consideration on the generalization error presented in Eq. (10), in the setting of a self-attention layer as in (3) and its sequence-to serenceet--seqeenee version as in (54).
In the main text, we showed the expression of the Bayes-optimal estimation enror. In the case of one layer of self-attention this reads:
East 1 d -1s- -
(263)
Regarding the generalization error, instead aim to compute and plot the different quantity shown in Eq. (10), namely:
Egen (
EsD,s S. Yuyo Nroow 11 (Fneww D) Vnew 11?
(264)
with:
BO VD
(Xesst) = E
Blest
Xiest
D)
E2 [fs (Xes P( D)ds
Recalling the fact that, for one layer of self-attention, we simply have the relation y as(h)
ae(fhat/ /22 - dab lab), we can introduce the change of variables hab = 2. Sas- dab S and get the Va
expression:
Egen
Σ a.b
dhuadhaalloo
hob Jab
hob /2 dab
2. Sw,
sab Tr S 10thaa
Bas T x. Sx
vd
(265)
We now exploit the fact that, as we know, the preactivations concentrate to:
Suy der TrS
18(hus
5. Su,
60s Tr
= .N(
(((aa))
N1o0 (oos)
(0)(( q Q- 2-22)
P(haa Rrat)
Vd
(266)
Then, the overall generalization emror is given by
Egen
(e [o(1
/2 - dab Jab) olf
hab
/2 8ab las)]
Nab
(267)
0,6=1
which exactly matches the result presented the main text in Eq. (19), where the extension to the multi-layer setting is trivial.
Now we slightly modify our model of self-attention layer by considering its sqquence-t--seqeeee (seq2seq) version y a o
hut /2- 8,0 Jat)ar Ig''d In particular we aim to compute and plot the
generalization error of Eq. (10) in this new setting.
To do so, we define y = Ar and ý As with A (h) and Å = o.) where we leave the factor /2 - Rub implicit. We exploit the concentration of our input data, order to compute the Frobenius norm of y - ÿ = (A A)x. Recalling the fact that the input data are with xt ~ (0, 1), we use the fact that
Tegaing ss de
(268)
with high probability when d is large. Hence:
| A)x1 = >>> Aue - A
ΣΣ
(Ar - AuMAu - Ame)er (269)
1,
2
but using the concentration property of x finally get:
[I(4 - 1) - (Aw - Aw (Au" Ae
E
(270)
We hence have shown that in the case of layer, the sequence-to--sqqeecee version of the model shows the same identical state evolution with respect single self-attention layer.
45

---

D.8 Details on the numerical implementation
The code used to produce all the figures and the experiments is available https: thtbb............ thub. com/ SPOC- Exxtensi EextensiveaankAtee MeeankAttentii Our gradient descent experiments are done in Py PYTorh 1.12.1 by minimizing the following loss using Adam
12
WWW
2t Sub TI WWT
((W)
Σ 4=1 m1,=1
Vab - 00g
(271)
Vrd d
In our implementation we sample both the input data and the weights of the target as standard Gaussian. Notice that we appropriately adjusted the loss be consistent with the main test. We choose learning rate 0.1 and keep the other hyperparameters at their default parameters and initializing the weights as a standard Gaussian.
When running the averaged version of the algorithm we run the optimization procedure 32 times for a fixed experiment, and average the matrix S WWT //rd at the end of training.
Regarding the state equations in the two L 1 cases of softmax and hardmax output channel: in the former case, we simply find the fixed points iterations of the state equations in Eq. (26) and Eq. (233). In the latter case, finally, compute the expectation in Eq. (262) with Monte-Carke Carlo over Nsamples = 20000 samples. In particular, to allow for more stable results, we iterate the state equations for T 150 iterations and compute the mean overlap over the last 30 iterations of the state equations.
46


Reviews:
[{"review_id": "HDkDZioq54", "reviewer": "Reviewer_5j3j", "summary": "This work introduces attention-indexed models (AIMs) as a theoretical framework to study learning in attention-based models. These models are the analog of single- or multi-index models for MLPs, adapted to the attention setting. The authors formalize the problem of learning an $L$-layered attention-indexed model, which covers the case of stacked attention layers with residual connections. They characterize the Bayes-optimal predictor in the high-dimensional limit where embedding dimension $d$, sample size $n$ (prop. to number of parameters $d^2$), and attention index matrix rank $r$ grow proportionally, while the number of layers $L$ and context length $T$ remain finite. An approximate message passing algorithm is proposed that recovers the Bayes-optimal performance under certain conditions. The Bayes-optimal analysis is worked out in full detail for the single-layer case ($L=1$), with results applied to both hardmax and softmax attention regimes.", "questions": "Please see the section above.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**: This work introduces a new class of models inspired by multi-index models in MLPs to study attention layers. I like the idea of studying the equivalent here given the prevalence of the transformer. The paper takes good first steps by characterizing the generalization and estimation error of the Bayes-optimal solution in this class of models. The results also seem very interesting, especially the contrast between the hard and softmax regimes. There are several directions for future work building on this, including extending to multi-layer attention, adding other architectural components, studying ERM with GD, and extending to symmetric matrices, many of which the authors have already mentioned in the paper.\n\n**Weaknesses**:  I want to highlight a few missing points and some related questions here. \n\n1. **Finite vs. Infinite Context Length Regimes**: You study the regime where context length $T$ is finite. Given that real-world transformers operate on increasingly long contexts, have you considered the asymptotics of $T \\to \\infty$? Do you expect qualitatively new phenomena to arise in this regime?\n\n2. **Relevance and Plausibility of Scaling Regime**: In your analysis, $d \\to \\infty$, $n \\propto d^2$ and $r \\propto d$, while $T$, $L$ remain fixed. I understand that taking $d \\to \\infty$ helps to apply tools from random matrix theory, but in practice, $d$ is fixed and moderate, while $nT≫d^2$. How should we interpret the Bayes-optimal results in this context? Do phenomena like phase transitions, sample complexity thresholds, or rank dependencies remain qualitatively valid when these asymptotics don’t strictly hold? Are there empirical observations or architectural patterns in real transformers that support the proportionality assumptions you've adopted? A discussion on this would be helpful. \n\n\n3. **Hardness of $L>1$ Case**: Lines 193–202 say the difficulty arises from cross-layer noise correlations; if the covariance were diagonal the problem would factorize. Could you clarify exactly which analytical steps break down once this correlation is present, and whether any near-diagonal or orthogonal-layer approximations (e.g., treating weak off-diagonal terms perturbatively) might make an $L ⁣> ⁣1$ extension tractable?\n\n4. **Intuition Behind Hardness of Hardmax**: Fig. 2 shows that softmax reaches zero Bayes-optimal error at finite sample complexity, while hardmax never does. Is the barrier for hardmax purely information-theoretic, i.e. the arg-max mapping discards logit magnitudes so the estimator cannot reconstruct them, or is there an additional effect from the activation’s non-smoothness? A brief explanation of which aspect drives the strictly positive asymptotic error would help clarify the contrast with softmax.\n\n\n5. **Sharp Transition in Small Width Limit**: In Corollary 4.2 (Fig. 1 right), you show a weak recovery threshold for hardmax with small width. Could you provide intuition for why performance is flat below this threshold, i.e., the estimator behaves like it has no signal? What changes fundamentally at the threshold?\n\n6. **Rank Sensitivity in Hardmax vs. Softmax**: Figure 2 shows that lowering the rank makes the softmax task noticeably easier, its strong-recovery threshold drops. Whereas Figure 1 suggests the hardmax error curve barely shifts with rank. Why does softmax benefit so strongly from reduced rank while hardmax remains almost insensitive?", "was_revised": true, "final_justification": "I believe this is a good work that puts forward a nice analogous model to single/multi-index models for attention, that can be used by future works to study different aspects (optimization, more complex transformer architectures—with MLPs, multi-head attention, etc.). In my initial review my, main questions were about interpreting results in the $T \\to \\infty$ regime—i)whether authors expect a quantitative change ii) relevance of the finite $T$ regime, iii) intuition on results in the hardmax vs softmax case. The authors improved my understanding on i) and iii), and agreed that ii) is a limitation that future works could look into, which is fine in my opinion, as this one of the first works studying properties of such planted models. As a result, I bumped my score from a 4 to a 5.", "rebuttal": "We sincerely thank the reviewers for their questions. We will reflect our answers in the revised manuscript to improve its clarity.  \n\n1. *Finite vs. Infinite Context Length Regimes*: \n\nOur theoretical results indeed require $T<<d$. However, we can take the limit $T\\to \\infty$ after the high-dimensional limit. In this case, we do not expect a qualitative change. E.g. the location of the phase transition just requires a rescaling of the sample complexity, see Fig. (2) right panel. Indeed, notice that in Result (4.3) we specify Eq. (19) for $\\hat{q}$ in the case of a softmax channel, obtaining a simple expression which do not depend on the number of tokens anymore once we do the simple rescaling $\\bar{\\alpha}=\\alpha (T^2 +T-2)$.  A qualitative change would likely arise with more complex structures of the data; proposing a suitable model along that line is an interesting avenue of future exploration.\n\n2. *Relevance and Plausibility of Scaling Regime*: \n\nIn general, one would want to explore as rich range of scalings as possible. Previous work [22-24] treated the case $d \\to \\infty$, and $r, T, L$ finite. In this respect, our paper explores a more general case that allows $r \\propto d$ as often chosen in practice. The choice of the sample complexity is given by the fact that it is in this scaling that the performance changes from trivially bad to perfect. Our requirement of $T, L$ remaining finite is indeed a limitation that future work will hopefully be able to remove. \nOne point to make is that while our theory results are derived in the $d \\to \\infty$ limit, the results do apply for very moderate values of $d$, see Fig. 1 where experiments with AMP are performed with $d=100$. Analogously, GD and its averaged version also show the predicted recovery thresholds (found theoretically for infinite d) with only $d=100$. The condition asked by the reviewer $nT>>d^2$ corresponds as $\\bar{\\alpha}T>>1$, so technically this scaling is already contained in our analysis by taking $\\bar{\\alpha}$ very large.  We will clarify the discussion about the reasoning behind the considered scaling limit in the revised paper. \n\n3. *Hardness of $L>1$ Case*: \n\nFor $L>1$ the off-diagonal terms of the overlap matrix are of the form ${\\rm Tr}(S_1 S_2)$, which do not have mean zero but scale with $\\rho^2$ and is thus not small for moderate $\\rho$. Diagonal uncoupling could be present in the low $\\rho$ regime, where a perturbative approach could be possible, but we were mainly interested in varying this width-ratio parameter and investigate its impact on the properties of the model. Extending to the $L>1$ case would imply solving a high-dimensional denoising problem with correlated noise for the prior channel, which we presently do not know how to do. We left this for future work and progress in random matrix theory.\n\n4. *Intuition Behind Hardness of Hardmax*:\n\nGenerating labels using a hardmax output channel corresponds to learning binary associations over the token sequence $X$, with discrete labels $y \\in \\{0,1\\}$. This task is intrinsically harder than the softmax case. Conceptually, this mirrors the difference between binary classification with a single-layer perceptron—where labels are generated by applying a sign function to a random linear projection—and linear regression on a linear target. In the classification case, perfect generalization typically requires more than a number of samples linear in the input dimension. By contrast, in the regression setting, zero test error can be achieved once the number of samples exceeds the input dimension. In this analogy, the softmax behaves like the linear regression case, while the hardmax mimics the more difficult binary classification case. We will reflect this discussion in the revised paper. \n\n5. *Sharp Transition in Small Width Limit*: \n\nThe mechanism behind the weak-recovery threshold is that of the BBP transition. When the width of the model decreases and becomes of order $r=O(1)$ in the dimension d, the problem is related to that of a rank-1 spike denoising. This can be noticed mathematically in the computations in Appendix D.2. The spike is contained in the bulk before some $\\bar{\\alpha}_c$ value in this small width regime. The spike instead emerges from the bulk above such value, namely the MMSE starts decreasing and some information starts being recovered even at finite width $r=O(1)$.\n\n6. *Rank Sensitivity in Hardmax vs. Softmax:*: \n\nIn case the referee was commenting on Fig. (1) left: both for the hard-max and soft-max the task is easier as lower rank (Fig. 2 left and Fig. 1 left). We point out that the scale of Fig. (2) is logarithmic and that the difference in MMSE among the various curves is more pronounced when visualizing the curves in linear scale as in Fig. (1).  In case the referee was commenting on Fig. (1) right: Note the x-axes in Fig. (1) right is rescaled by r to illustrate better the low-r limit. But Fig. (1) right cannot be interpreted as the task getting harder as the rank is lowered."}, {"review_id": "qyOpB22zNX", "reviewer": "Reviewer_kAGS", "summary": "The paper studies Bayes-optimal generalisation error in a theoretical model of a Transformer. The advantage of their proposal is that they can loosen the restriction on the dimension of queries and keys being small. The provide synthetic experiments to validate their theoretical work.", "questions": "I believe it's a typo, but should it not be $x_a^T S_l x_b$ to give a scalar in Equation 1? It looks like $x_a \\in \\mathbb{R}^d$ is a column vector and $h_{ab}^{\\ell}$ a scalar. \n\nCould the $L$ \"attention indices\" be interpreted as $L$ individual heads? I am finding the choice of notation quite confusing when compared to the notation used in standard transformer works, for example $L$ is usually the number of layers while $H$ is the number of heads, $h$ could maybe be called $\\alpha$ to signify that it is acting like an attention coefficient.  -- of course this is personal preference but I am finding the current notation hard to parse .   \n\nEquations 4 and 5 are also difficult to parse. How would they resolve in the context of a ViT for example or an LLM in your notation?", "limitations": "The limitations and overall the wording in the paper is often quite fair and honest.", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n----------------\nThe paper tackles an interesting problem of studying generalisation in Transformers. The assumptions the authors make are reasonable and seem to push forward the surrounding literature (although I am not an expert in this area).\n\nThe experimental results seem quite convincing and I appreciate the numerical validation of the theory.\n\nWeaknesses\n----------------\nThe paper targets a very specific community and is hard to follow. I have some light background in RMT and found this quite challenging to read. I believe it would be useful to contextualize with some concrete examples the notation (e.g. Equations 4 and 5). For instance, how would a ViT or LLM look like using the notation. I understand it would not match 1-to-1 with actual Transformer architectures, but at least it would help clarify the notation.\n\nI believe the claim \"AIM offers a solvable playground for understanding learning in modern attention architectures\" in the abstract is perhaps over-claiming as I do not believe there is high semblance in the update equations to modern attention architectures (a lot of very non-linear and complicated components are missing).  \n\nOverall, having some light background in RMT, I found the paper very hard to follow and I think that the paper could benefit from a pass to try to make it more accessible more broadly. I do understand that the nature of the work might make it challenging, for this reason I am happy for the paper to be accepted if other reviewers believe it is worthwhile as I am not overly familiar with this area of research and cannot confidently assess the contribution of the submission given what is already out there.", "was_revised": true, "final_justification": "Unfortunately my confidence on this paper is quite low so please consider reviewers with higher confidence.", "rebuttal": "1. *Weakness 1. and Question 3. : Actual Transformer, and eqs. (4) and (5).*\n\nWe thank the referee for the pointers regarding notation. We introduce an abstract attention-indexed model (AIM) in eqs. (1) and (2). \nThen, in eq. (3) we state a simplified multi-layer self-attention with a single head in every layer and no MLP layers. An actual Transformer architecture, such as used in ViT or LLM, would be very close to eq. (3) with multiple heads and attention layers interlayed with token-wise MLP layers. ViT and LLM are usually trained by next-token prediction while the task we consider is supervised. \nEquations (4) and (5) are then a result detailed in appendix B of how to map eq. (3) on (1-2), i.e. how to map the multi-layer self-attention on the AIM. These equations indeed do not look very transparent, but the existence of such mapping is our key contribution as it enables results from the AIM to apply to multi-layer attention. We will clarify this in the revised version.  \n\nArchitectures with multiple heads could also be mapped to the AIM form and then the number of indices $L$ in the AIM would be the total number of heads in all the layers. \nThe only true limitation of the mapping of the transformer architecture to the AIM are the MLP layers. We did not find an analogue of Eqs (4) and (5) that would map the transformer including the MLP layer to the AIM.\n\n2. *Weakness 2. : AIM claim*:\n\nWe agree that \"AIM offers a solvable playground for understanding learning in modern attention architectures\" is a misleading choice of wording. We will rather say: \"AIM offers a solvable playground for understanding learning in self-attention layers, that are key components of modern architectures\"\n\n\n3. *Question 1. : Row-vectors notation* \n\nThe notation we used for the quantity $x_a S_\\ell x_b^\\top$ is correct, as we consider each token $x_a$ to be a row-vector, namely the rows of the whole sequence matrix $X\\in \\mathbb{R}^{T\\times d}$. We understand however that this is in contrast with standard notation of considering column vectors if not better specified. We will change the transposes and make the notation conform to standard practices of vectors being columns in a revised version of this work.  \n\n\n4. *Question 2. : Interpretation of the attention indices L*:\n\nIn this work we focus on multi-layer self-attention (3) where $L$ indicates the number of layers, and thus $\\ell$ is a layer-index with $\\ell=1,\\dots,L$. \nHowever, a collection of attention heads can also map to the AIM. We will add an appendix with such a case and refer to it from the main text to illustrate the generality of the AIM. The observation of the reviewer is indeed spot-on and contained in our analysis."}, {"review_id": "gVAIAnIADo", "reviewer": "Reviewer_MHZz", "summary": "The paper proposes the Attention-Indexed Model, designed to study learning in attention-like architectures, where tokens interact with the model parameters as in the attention mechanism of modern Transformers. For the proposed model, the authors derive the generalisation error in the high-dimensional limit of the number of samples $n$, embedding dimension $d$, and width $r$ (with suitable ratios between these quantities) for the Bayes-optimal estimator. They also propose an AMP algorithm that recovers the overlaps appearing in the asymptotic error, thus matching the Bayes-optimal performance. While the theory covers a general $L$-layer model, the authors also discuss the $L = 1$ case as analytically more tractable.", "questions": "The authors propose a mean-centered attention mechanism with $1/\\sqrt{d}$ scaling. It would be instructive to discuss the alternative $1/d$ scaling (probably, no mean centering needed in this case). Under this \"LLN\" scaling, are the asymptotics different, or would the current proofs carry over?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## Strengths \n\n- Novel modeling perspective. By pulling out the bilinear attention indices and studying them directly, the paper isolates what is mathematically distinctive about attention and carries out standard techniques to study the asymptotic behavior. Thus, I see the main contribution as the novelty and relevance of the results more than in their technique.\n\n- I particularly appreciated the discussion on the intractability of the integrals involved which prevent to derive a closed form for the Result 3.1. \n\n- Concrete formulas in the single-layer case. For $L=1$, they can write the AMP’s prior-channel denoiser in closed form via rotationally-invariant eigenvalue shrinkage (RIE) (22) and an explicit state evolution equation (24), leading to analytic phase-transition locations that perfectly match AMP simulations. \n\n## Weaknesses:\n\nI am overall in favor of clear acceptance, as I find the paper technically sound, the results are relevant and novel. These minor weaknesses are to be seen more as suggestions or discussion points\n\n- Background on AMP and Bayes optimality. I think the authors could have given more background on approximate message passing and more explicitly state its connection to Bayes optimal estimation. The paper relies heavily on this equivalence but only mentions it briefly (Appendix C). A short primer in the main text would make the results accessible to a broader audience.\n\n- The paper writes the Bayes free-entropy (or “potential”) of the Attention-Indexed Model as an extremisation problem. Unfortunately, the extrema for $q$ and $\\hat{q}$ in Result 3.1 are in implicit for, and it’s thus a bit hard to interpret. However, the authors extenively discuss why the intrgral in the partition function $I_{in}$ is intractable.  However, the authors provide more “meat” with the discussion of the $L=1$ case, which is tractable because now the $\\Delta_\\ell$ becomes diagonal. I think this compensate for this lack of interpretability.", "was_revised": true, "final_justification": "The paper provides a novel treatment of the generalisation error in the high-dimensional limit for attention-like models. Novelty of the results and relevance in the modern theory of deep learning are the main reason I am in favor of acceptance.", "rebuttal": "1. *Background on AMP and Bayes optimality*:\n\nThank you for this suggestion. Indeed, this is a key point and we will explain it better in the text. The analogy works on the two following lines. On one side we express the posterior distribution of the model with respect to S, which can be manipulated in the considered limit to end up with an optimization problem of a free-entropy in Eq. (15). On the other side, a specific AMP algorithm provably reaches the BO error. The specificity is in the choice of two denoising functions that are taken in such a way that the AMP fixed points are extremizers of the same free-entropy, namely its iterates follows the state evolution equations in Eq. (19). We will make such connections clearer in a revised version of this work.\n\n2. *Interpretation of the order parameters*:\n\nRegarding the interpretation of the order parameters $q$ and $\\hat{q}$: $q$ plays the role of an overlap and, at the fixed points of the free-entropy, is ${\\rm Tr}(\\hat{S}S^\\star)$. On the other hand, $\\hat{q}$ is its conjugate parameter and plays the role of a signal-to-noise ratio of a high-dimensional denoising problem in the prior channel.\n\nThe fixed point equations for $q$ and $\\hat{q}$ are indeed complex for $L>1$ (discussed around Eq. (20)). For $L=1$ they are simpler and numerically easy to evaluate, and hence we focus the discussion of the results in Sec. 4 on this already rich and interesting case.  \n\n3. *Scaling used*:\n\nWe are interested in studying the leading order behavior of the attention index. Given that we take the input data to have mean zero and variance one, we expect the dot-product attention to also be a collection of $T\\times T$ random variables with mean zero and variance one. In the Bayes-Optimal setting and for our model, dividing the dot product by $d$ instead of $\\sqrt{d}$ is just a matter of conventions, as the statistician knows about it and can map between the two by rescaling the inputs.\n\nIf the data has a mean of order one, the $1/d$ scaling makes fluctuations in the dot-product attention subleading and the output would concentrate. This would significantly simplify the problem, and the relevant sample complexity scaling will be much lower than $\\mathcal{O}(d^2)$. We thus have not considered this case of a target function. \n\nThere may be other combinations of scalings that make the problem non-trivial. We are happy to consider them if the referee is curious about any particular one and let us know."}, {"review_id": "4cG2ilDKwq", "reviewer": "Reviewer_VRk2", "summary": "In this paper, the authors have introduced the Attention-Indexed Model (AIM) to analyze learning in deep attention layers. Theoretically, the authors analyze this using statistical mechanics and random matrix theory to get the closed-form predictions for Bayes-optimal generalization error. Also, the authors analyze the case where embedding dimension $d$ goes to infinity and the sample complexity is $\\alpha = n/d^2$.", "questions": "Please see weaknesses.", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The strengths of this paper can be summarized as follows:\n\nFirst, this paper gives a very thorough theoretical analysis. This is the first tractable analysis of attention mechanisms with extensive-width matrices, whereas previous works only focus on narrow attention layers. As far as I can see, all the theoretical analyses are rigorous and correct. \n\nSecond, the authors have also conducted experiments and provided an optimal AMP algorithm. In particular, the result has shown that the practical gradient descent achieves similar performance.\n\nThird, the implementation details and code are included. It makes the experimental results being reproducible. \n\nThe weaknesses of this paper are summarized as follows:\n\nFirst, as the authors have mentioned, they only analyze the single layer case. It would be better if the authors can consider including a more detailed discussion on the generalization to the multi-layer case. For example, what are the potential difficulties one may encounter to generalizing it to multi-layer case?\n\nSecond, the author claims that they analyze the case where $d$ is very large. However, in their experiments, $d = 100$ or $d = 200$ seems pretty small or moderate. Is it possible to try larger $d$’s to make it more aligned with the theoretical results?", "was_revised": true, "final_justification": "In general, I think this is a good paper, and the authors’ responses are convincing. I have no further questions and have chosen to maintain my original score, as I believe it is already sufficiently high.", "rebuttal": "1. *Generalization to the multi-layer case*:\n\nWe point to the prior-channel equations and the related denoising problem with correlated noise (for $L>1$ layers with $L=O(1)$) in Eq. (20) and the paragraph around it. This is actually an interesting random matrix theory problem that may be tractable, but would constitute a result of its own interest in matrix denoising problems. We will emphasize this interesting technical open problem even more in the revised version.  \n\n2. *Large dimension in numerical experiments*:\n\nAs the reviewer correctly pointed out, the analytical results obtained are valid in the limit of large embedding dimension $d\\to \\infty$. Analogously, we also send to infinite the number of samples $n$ and the width of each weight matrix $r_\\ell$, with the particular scalings $\\alpha = n/d^2 =O(1)$ and $\\rho_\\ell = r_\\ell /d =O(1)$. Regarding the numerical experiments both for AMP, for gradient descent and its averaged version: we use moderate sizes for the embedding dimension, namely $d=100$, and we obtain curves that are strikingly compatible with the Bayes-Optimal theoretical ones (obtained for $d\\to \\infty$). \nWe point out the fact that reproducing the large d results already at such moderate sizes is exactly one of the main points that go in favor of the high-dimensional theoretical approach. \n\nIn the revised version, we will add a figure showing that as $d$ increases the numerical experiments will get closer to the theoretical prediction at $d\\to \\infty$. We expect to be able for explore the range of $d$ from $30$ up to $1000$ for some of the small values of $\\rho\\sim0.2$, which is not far from the embedding dimensions used in practice."}]

Author rebuttal:

Meta-review: This paper proposes the Attention-Indexed Model (AIM), a theoretical framework for studying learning in deep attention layers. AIM generalizes prior tractable models by allowing full-rank key and query matrices, making it more faithful to modern Transformer architectures. The paper derives closed-form predictions of Bayes-optimal generalization error and characterizes phase transitions in learning as functions of embedding dimension, sample complexity, width, and sequence length. The paper further proposes an Approximate Message Passing (AMP) algorithm that achieves Bayes-optimal performance under certain conditions, and demonstrates that gradient descent can converge to achieve optimal performance. Synthetic experiments validate the theoretical predictions.

The paper was reviewed by four referees who agree on the paper's primary strengths and weaknesses. The AC and reviewers find that the proposed AIM framework provides a novel, technically sound means of analyzing generalization error and phase transitions in the high-dimensional limit for attention-based models. All four reviewers emphasize the significance of the results, which are very thorough, interesting, and convincing. The framework and the theoretical results contribute to our collective understanding of learning in architectures with deep attention layers and provide several interesting directions for future work. Some reviewers raised concerns with the original submission, noting that the presentation would benefit from improved accessibility. In particular, they suggested drawing explicit connections to architectures such as LLMs and ViTs to better contextualize the work, and including background material on Bayes optimality and AMP in the main paper. The authors addressed the reviewers' primary questions and concerns during the rebuttal period, as reflected in the reviewers' Final Justification.
Decision: Accept (poster)