Fundamental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds
Authors: Emanuele Troiani
Hugo Cui
Yatin Dandi
Florent Krzakala
Lenka Zdeborova
Venue: ICML 2025 poster (2025)
Keywords: theory
multi-index model
approximate message-passing
replica method
statistical physics
attention mechanism
Primary area: theory->learning_theory
TL;DR: We derive the computational limits of Bayesian estimation for sequence multi-index models in high dimensions, with consequences for deep attention neural networks.

Abstract: In this manuscript, we study the  learning of deep attention neural networks, defined as the composition of multiple self-attention layers, with tied and low-rank weights. We first establish a mapping of such models to sequence multi-index models, a generalization of the widely studied multi-index model to sequential covariates, for which we establish a number of general results.  In the context of Bayes-optimal learning, in the limit of large dimension $D$ and proportionally large number of samples $N$, we derive a sharp asymptotic characterization of the optimal performance as well as the performance of the best-known polynomial-time algorithm for this setting --namely approximate message-passing--, and characterize sharp thresholds on the minimal sample complexity required for better-than-random prediction performance. 
Our analysis uncovers, in particular, how the different layers are learned sequentially.  Finally, we discuss how this sequential learning can also be observed in a realistic setup.

Fundamental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds
Emanuele Troiani
Hugo Cui
Yatin Dandi 1
Florent Krzakala Lenka Zdeborová
|
Abstract
thus constitute function class able represent intricate dependencies between the tokens of sequential covariates. In spite of their ubiquity, the study of such models from a thcoretical viewpoint is still in its infancy.
In this manuscript, study the learning of deep attention neural networks, defined as the composi- tion of multiple self-attention layers, with tied and low-rank weights. We first establish a mapping of such models to sequence multi-index models, a generalization of the widely studied multi-index model to sequential covariates, for which we es- tablish number of general results. In the context of Bayes-optimal learning, in the limit of large dimension D and proportionally large number of samples N, we derive a sharp asymptotic charac- terization the optimal performance as well as the performance of the best-known polynomial- time algorithm for this setting -namely approx- imate peessag-passsing. and characterize sharp thresholds on the minimal sample complexity re- quired for better-than-rane random prediction perfor- mance. Our analysis uncovers, in particular, how the different layers are learned sequentially. Fi- nally, we discuss how this sequential learning can also be observed in a realistic setup.
This situation stands in contrast with the wealth of results on multi-index models, central to many theoretical works e.g. (Arous et al., 2021; Abbe et al., 2022; Veiga et al., 2022; Ba et al., 2022; Arnaboldi et al., 2023; Collins-Woodfin et al., 2024; Damian et al., 2024a; Bietti et al., 2023; Moniri et al., 2024; Berthier et al., 2024). Multi-index models define functions based on dow-dimensional subspaces of covariates and are popular exemplars among theoreticians. However, shallow architectures like multi-index models are structurally simpler than attention models. Specifically, they (a) lack the hierarchical structure from multiple attention layers and (b) act on less structured, non-sequential covari- ates. These distinctions make it unclear whether results from multi--nnex models apply to multilayer attention archi- tectures. Here, we show these models share a deep formal connection, enabling a transfer of insights and analytical approaches.
Namely, our first motivation is to extend the theoretical framework of multi-index models to sequence models. We consider the class of sequence multi-index (SMI) functions introduced by (Cui et al., 2024a; Cui, 2025), and defined over length M sequences of D-dimensional tokens x RDxM of the form
1. Introduction
Recent years have witnessed shift in paradigm in the automated learning from sequential data, such as language (Brown et al., 2020; Kenton & Toutanova, 2019). The backbone of many of these technological advances arguably lies in the use of transformer architectures (Vaswani al., 2017) a parametrization allowing the model to dynamically focus relevant portions of the input, and extricate increasingly complex correlations between tokens through the successive application of atlention layers. They
Wæ a
SMI yW (x)
(1)
where W E RPxx is a learnable projection matrix, and RPXM + RK is (possibly multi-dimensional)) link function. The K-dimensional output could represent to give an example- the predicted class probabilities in sentiment classification task. As we shall see, SMI functions prove very versatile generalizations of multi-index models (M 1) sequence data of length M > that can be studied using related ideas and theoretical tools.
'Statistical Physics Of Computation Laboratory, École Poly- technique Fédérale de Lausanne (EPFL) *Center of Mathemati- cal Seiences and Applications, Harvard University "Information Learning and Physics Laboratory. École Polytechnique Fédérale de Lausanne (EPFL). Correspondence to: Emanuele Troiani <emanuele. somanuclc.troianidee @ch
Our second motivation is to deploy the formalism devel- oped for SMI models to study complex neural architectures. Remarkably. the SMI class encompasses a range of mod- els with a particularly noteworthy example being deep
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
atlention architectures. A deep attention model of depth L, yDA RDxM - XMM, is defined recursively
(B are defined recursively as:
Zi) B ..... (,) (1-1)) [ (Z1,.., 111) 2 Z. (Z),.., Zi 2-1)]]
yDA yDA(#) (x) [aI-1w) WI
-01,1)c
(5)
D
where V/ E L-].
with 0 1M, the identity matrix. The claim (4) then follows from the definitions (5) and (2) by recursion, finally yielding
2 T
1/-1 ".,,/,a.... wiii 1
el M +0 81 = 2i-1
(3)
D
yDA yDA(2) = (Z Z1.)
(6)
with Zo = x E DDX and a RMXM + RMXM non-linear map, taken in this work to be the softmax function. The deep attention function (2) corresponds to a composition of single-head self-attention layers (with tied key and query), specified by the strength E R of its skip connections, and tied key and query matrices wy R P for 1 ,,...,, L. We set the value matrix to identity. We show in Appendix A that the SMI class also describes even more generic multi-layer attention networks.
Eq. (6) can be written more compactly and evocatively
yDA (I) = W- x VD
(7)
introducing the total weights W E RPxD, defined as the vertical concatenation of the weight matrices {wij}= along their first dimension, and denoting / = P1 + +...+PP The link function g in (7) follows from redefinition of B., as can be read from (6), and takes values in RMX z R for K = M One then observes that the equivalent model (7) is an SMI model (1), whose width / is related to the depth of the original deep attention model (2). The recursive structure of the deep network (2) is encoded in the way the highly structured function g acts on the different rows of its argument W** E RPxM Importantly, the P rows generically do not play symmetrie roles, reflecting the structure of the original deep attention function.
Before stating our main contributions, we elucidate the con- nection between the deep attention model (2) and the SMI model (I). While
focus on the single-head tied-weights architecture (2) for clarity, such mapping can also be con- structed for deep architectures with multiple heads, untied weights, and sequence-tor to-eeuut---seeceeeeeeee models with outputs in RDXM, as detailed in Appendix A. The connection between the SMIs and deep attention models constitutes a pathway towards many exciting research avenues in the study of neural architectures beyond the exemplar of multi-layer perceptrons -ucch as classical multi-index models-, while allowing the transfer of ideas from the study of the latter.
Main contributions- - Building on the above mapping, we study the statistical and computational limits of learnability from data generated by the models (1) and (2). Namely, we consider training data composed of N samples of Gaus- sian i.i.d. input sequences, and output labels generated by the SMI model (1) with random i.i.d. Gaussian weights (the same realization of weights for every sample). In the limit of large dimension D and number of samples N with N/ P, M, K = (1), we characterize the minimal pre- diction error achievable information theoretically and also by a class of algorithms conjectured optimal among all polynomial algorithms. This achieved by the following technical contributions:
Mapping deep attention to the SMI model- Starting from a deep attention model (2) with weights 1/1 let us introduce the projections (a.k.a the eponymous indices) {Z,- PR X M' Vi MY= of the input data x RDXM M each weight matrix w RRRxD: Zi w /DD. Layer by layer, the post-activation sequences xı build on the initial sequence by adding increasingly sophisticated correla- tions between tokens. A key observation, however, is that successive post-activations xi retain simple formal expres- sion, and depend non- linearly on the input x only through the indices 71, ..., Zt More precisely, at every layer 1, the post activation क can be written as
information-thhorrttcaall
We show how to rigorously generalize results previously established for multi-inde multi--ndexx models to the SMI models including the notions of phase transition for weak recovery. and optimal message-passingg algorithms (Damian al., 2024b; Troiani et al., 2025).
x = BB Zt),
(4)
where we introduced a function B² RPMXM x ... x
We derive sharp asymptotic expressions for optimal esti- mation emrors achievable information theoretically and with the Approximate Deessage-Passingg (AMP) algorithm, in the limit of large covariate dimension D, and proportionally
RPixM + RMM subsuming the complex inter-token interactions resulting from the propagation of the input through previous layers. The sequence of mixing functions
2

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
large number of samples N, but finite sequence length M, weight rank P and depth L. This characterization builds the analyses of (Donoho et al., 2009; Javanmard & Monta- nari, 2013; Gerbelot & Berthier, 2023) for AMP, which is optimal among first-order methods (Celentano et al., 2020).
implement gradient descent steps, establishing how an opti- mization algorithm can be encoded by such an architecture. Deep attention models have also been studied through the lens of their capacity (Edelman et al., 2022) and expressiv- ity (Hahn, 2020), and scaling limits (Bordelon et al., 2024). The training dynamics of deep attention models were as- certained close to initialization by (Bietti et al., 2023), and fully by (Abbe et al., 2024a), in exchange of the assumption of diagonal weights. Closer to our setting, (Tiberi et al., 2024) study the Bayesian learning of multilayer attention, assuming trainable value matrices but frozen key and query matrices. In this respect, the model is complementary to the one considered in the present work, where in contrast the learning of key and query weights is addressed, while value matrices are fixed identity. Another key difference lies in the fact that the attention matrix at every layer is computed using the bare input æ in (Tiberi et al., 2024), opposed to the output of the previous layer 1 as in (2).
We extend the numerical evaluation and analysis of the resulting asymptotic equations from M- = 1 in (Aubin et al., 2018; Troiani et al.. 2025), and P in (Cui al., 2024a; Cui, 2025) to SMI models with both sequence length M > 1, and number of indices P >>1 1.
Equipped with the above theoretical results, we extract the weak recovery thresholds, namely the sample complexity re- quired for better-than-random estimation, for the weights of every layer in the case of two-layer attention model. We show that different layers are learnable at different sample complexities, implying sequential learning of the different weights. We then briefly discuss how such layer-wise learn- ing also arises in more realistic settings.
Further related works
Learning multi-index models- - Multi-index models cor- respond to target function of the form y(x) = 9 (Wx), specified by a link function g and a finite rank matrix W E RPxD, with P @ and can be viewed as two-layer MLPs with weights W and activation g. They provide natural exemplar for functions that operate on high-dimensional covariates, but only depend on a finite number of directions. Several authors have used multi-in index models as testbeds to explore the ability of neural networks to learn low-dimensional subspaces in high dimensions (Ba et al., 2022; Moniri et al., 2024; Dandi et al., 2023; 2025; Cui al., 2024b). A substantial stream of works has further- more been devoted studying the behavior of (stochastic) gradient descent on such non-convex objectives (Saad & Solla, 1995; 1996; Arous et al., 2021; Bietti et al., 2023;
Analyses of single attention layers- - Many theoretical studies of attention-
models (Vaswani et al., 2017) consider the case of a single attention layer, studying it in isolation (Lopardo et al., 2024; Zhang et al., 2024; Li al., 2023b; Tian et al., 2023; Jelassi et al., 2022). This rich line of works shed light on various aspects of these models, such as their inductive bias (Sahiner et al., 2022; Ataee Tarzanagh et al., 2023; Tarzanagh et al., 2023), training dynamics (Li et al., 2023a), or expressivity (Fu et al., 2024). A fraction this body of works was devoted to characterizing attention models in the limit of large dimensions. (Lu et al., 2024) and (Rende et al., 2024) provide tight asymptotic characterization of the error achieved by linear attention layers, respectively for in-context learning and next-token prediction. (Cui et al., 2024a) uncover phase transition between semantic and positional learning in a model of attention mechanism with low-rank weights. For a related model, (Marion et al., 2024) study its population gradient dynamics, and prove its asymptotic Bayes-optimality for regression task. Because these works focus single attention layers however, they cannot capture the compound effect of multiple successive attention layers in building increasingly complex correlations between tokens present in multilayer models such as (2).
Simsek et al., 2025; Barak et al., 2022; Berthier et al., 2024;
Abbe et al., 2023; Damian et al., 2022; Glasgow, 2023; Amn aboldi et al., 2024b), with the work of (Oko et al., 2024) providing insights for more generie scalings of the number of hidden units P with the dimension D. In closer rela- tion to the present work, (Troiani et al., 2025) ascertain the fundamental computational limits of learning multi-index functions, leveraging an analysis of AMP algorithms. In this manuscript, we establish a deep connection between deep attention models (2) and SMI models (I), a generalization of multi-index models introduced by (Cui et al., 2024a; Cui, 2025), allowing to borrow from the aforementioned wealth of insights and results. The connection to deep attention models was not noticed in (Cui et al., 2024a; Cui, 2025), and (Cui, 2025) studied the empirical risk minimization for the SMI models using the replica method, and has not estab- lished the fundamental statistical and computational limits for the SMI function class.
Analyses of deep attention models- In an effort to super- sede those limitations, recent stream of works has striven to tackle the problem of multiple attention layers, under varying simplifications. (Geshkovski et al., 2024; 2023) study the propagation of a signal through multiple attention layers with frozen weights. (Ahn et al., 2023; Von Os- wald et al., 2023) describe how the successive layers can
3

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
Statistical and computational limits of learning- The optimal prediction errors associated with single-layer neural network functions were first derived in line of founda- tional works (Gardner & Derrida, 1988; Sompolinsky et al., 1 1990; Györgyi, 1990) for perceptrons, and rigorously in (Barbier et al., 2019). The statistical limits of learning in committee machines were explored (Schwarze Hertz, 1992; Schwarze, 1993; Monasson & Zecchina, 1995), and its computational limits in (Aubin et al., 2018; Troiani et al., 2025). The case of large-width networks was covered in (7avatone- Veth et al., 2022) for linear models, and (Cui et al., 2023; Camilli et al., 2023; Maillard et al., 2024) for non-linear models. Closer to our work, (Erba et al., 2024) consider a related model of bilinear sequence regression, characterizing its optimal errors, and proposing an associ- ated AMP algorithm matching these performances. (Abbe et al., 2024b) study the weak learnability of variety of tasks for transformers trained with descent algorithms.
E [wlww [/ [x) over the Bayesian posterior distribution
N 6 (UW-
P(W (WID)xxe+((WW)) II
SMI
(-")
- -VW "w""((")))
SMI
(8)
Asymptotic limit In the limit of large covariate dimension D and large number of samples N, sampling the posterior distribution (8) is generically computationally challenging. It is on the other hand possible, in this limit, to access exact theoretical characterizations of key statistics of the optimal estimators. We detail this analysis in the follow- ing sections. More precisely, we consider the asymptotic limit of D, + oo while = N/D @(). We further assume that the sequence length M, model width P' and output dimension K remain finite : M. P. K @())
Methodology - The following technical results are enabled by the key observation that the SMI model (1) can further be viewed as an ordinary multi-index model, provided the data and weights are properly reshaped. More precisely, to go from an SMI model to an equivalent multi-index model, one can flatten the input E RDxM, viewed length M sequence of D-dimensional tokens, into a vector in RMD. In parallel, the matrix of weights W' E RPxD should be mapped into a block-diagonal matrix of size PM x DM with each of the M blocks on the diagonal being equal to W*, as detailed in Appendix B. In words, an SMI model can be viewed as a multi-index model acting on flattened sequences, with non-separable prior on the parameters.
2. F Fundamentals limits of learning SMI models
The question of determining the statistical and computa- tional limits of learning has traditionally been thoroughly explored for multi-index models (Gardner & Derrida, 1988; Schwarze & Hertz, 1992; Barbier et al., 2019; Aubin al., 2018; Troiani et al., 2025). In this first section, we show how these ideas can be extended to SMI models (1).
2.1. Bayes-optimal learning of SMI models
We are interested in the statistical inference prob- lem of estimating the weights W' of an SMI model SMI = g (W"#/vD) from N observations D = {(x"" wW- (x*)) |= of the function on identically and idependently sampled (i.i.d covariates x E RDDM with Lid standard Gaussian components. We further consider a random instance of the SMI model, where the weights W* have been randomly sampled from prior distribution, which we assume for definiteness to be i.i.d Gaussian over all the components. For this problem, we ask the question of the optimal statistical and computational prediction eтror that can be deduced from the estimation of the weights W from the observations DD what is the information- theoretically lowest reconstruction error? What the lowest error achievable by polynomial-tnn algorithms?
This mapping allows us to borrow from previous results on multi-index models. On the one hand, the computa- tional limits can be determined along the lines of (Troiani et al., 2025), leveraging an analysis of the generalized-A (GAMP) algorithm, which is provably optimal within the class of first-order methods (Celentano et al., 2020). This analysis can be carried out using methods developed in (Ger- belot & Berthier, 2023) on the state evolution of GAMP. The statistical limits of learning may be, on the other hand, ob- tained by leveraging the rigorous results on the free-energy of the posterior measure (8) for multi-index models estab- lished in (Aubin et al., 2018), as sketched in Appendix B.2, or using the non-rigorous replica method from statistical physics, detailed in Appendix D.5.
We answer these questions in the Bayes-optimal setting, in which the statistician has full knowledge of the param- eters of the model except the realization of its weights w The estimator v(-) minimizing the prediction error Ege [[(()) - SMI 12 is then given by the mean y(x)-
2.2. Statistical limits: Bayes-optimal error
To state our main results, we start by defining a finite- dimensional output channel induced by Eq. (1) that will be referred to throughout the subsequent sections:
Y(u, V) = (N + vV2).
(9)
4

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
itd. where 7 E RPxM and Z ~ NN 1). On an intuitive level, w E RPxM and V E Sp (a positive semidefinite matrix od size P) correspond to estimates of the mean and covariance of the imputs to the nonlinearity g-) in Eq. (1).
Algorithm 1 GAMP for SMI model
Inputs: D {-* 21 " Ja=1
Initialize W = -0,, E =1p, 9 = 0pxM
for S mmax do
V O+ E RPxP
With the above definition and assumptions on g specified in Appendix B.2, we are in position to state our first result, characterizing the asymptotic Bayes-optimal error for a general SMI model of the form (1):
D
w/m, = 1
Σ
x
W - V gm E RPXM
vD
m,
i=1
9m., = [ (8"
w VO] E ERP
N.M
Theorem 2.1. Consider the SMI model (1) with non- linearity g. Let & E PRMM denote a matrix with entries id Eij ~ N(0, 1). Ler (Q), denote the conditional en- tropy of the associated oulput channel Y defined by Eq. (9) with w = vQC and V = =1p- - Q for the so-called overlap Q E ST. Suppose that the following sup inf problem:
A' = N
Σ
Cy w// Vt] E RPxP
m.m==
N.M
bf = 1 Σ
x Imm + A'W, ERP
vD p.mm
W = (1p + A')-
(++1 = (1p + AAa
{ (T+(QQ)
end for
sup inf QEs), Qest
=log(1p + +Q)
(10)
+ xHQ
Lemma 2.2 (State evolution (Gerbelot & Berthier, 2023)).
Define gour RMMx + RRXM to be the following "de- noiser":
admits unique global extremizer Q-, Q-, Then the asymp- totic Bayes-optimal prediction error is given by:
E AMI (a)) Ew (MMMI(2) (x))DDi 1| D
Jout WY, = E[]
(12)
(11)
-11)) 9 (V 1 - Q- Q-2-----=2
(QT)t))
where the conditional expectation is w.r.t the output channel in Eq. (9), Let w' denote the iterates of Algorithmm at time t E N. Assuming 9out C2, under the high-diieensionl limit N, D - oo with fixed ratio a= a=N N/D, constunt M and any finite time t, the limiting overlaps converge in probabil- ity to the overlap Q':
and the Bayes-optimal estimation error is [W") Ew [WW T [D] 1|2 + 1 11@-
The proof of Theorem 2.1 is outlined in Appendix B. A complementary non-rigorous derivation using the replica method is provided in Appendix D.5.
1/a w Q. 1/aW'"" W L. (13)
2.3. Computational Limits: Weak-recovery thresholds
with Q' satisfying the state evolution equations
We next move to presenting a sharp asymptotic characteriza- tion of the computational limits of the Bayesian estimation problem (8). Our approach relies on the study of a GAMP algorithm (Donoho et al., 2009), reported in Algorithm 1, building on two of its remarkable properties. First, the Bayes-optimal estima- tion problem is provably optimal among first-order methods (Celentano et al., 2020) and runs in polynomial time. Conse- quently, the sample complexity required for the learnability for GAMP implies a computational lower bound on the class of all first-order algorithms, which includes in partic- ular widely used gradient descent methods. Secondly, the performance of GAMP algorithms admits sharp asymp- totic description in terms of dittesiiinnii variables. This description, known as the slute evolution equations (Bayati& & Montanari, 201 I; Javanmard & Montanari, 2013), allow the analysis of GAMP in the high-dimensional limit D + cc. The state evolution description is formalized through the following Lemma:
Q' = F FF 9oul (Y,
- 1021))
(14)
from an initial condition Qo In (14), we denoted F(() (1p + Q)- and employed the shorthand Y YVQE, /1p-111p- Q), with E RPXM denoting matrix with entries Eig it N(O0 1).
GAMP algorithm associated
In Appendix D.2, we establish that the fixed points of the state evolution equations (14) are in one-to-one correspon- dence with the stationary points of the objective function defined in Equation (10). Consequently, in cases where Equation (10) admits a unique non-zero stationary point Q* E Sp corresponding to the global extremum, the GAMP algorithm, provided it converges to non-zero limit, neces-
1 W IT
converging to Q", It then follows from Equation 11 that GAMP asymptotically achieves Bayes- Sayemmoptimal performance. The proof of Lemma 2.2, which builds upon (Berthier et al., 2020), is outlined in
sarily attains an overlap 1/aW"
5

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
1.0
X
Prediction error
Sample complexity
1.0
Q11
Overlap
10-1
Q2
0.5
0.5
Extor
10
0.0
0.0
0.0
0.5
1.0
0
2
4
Sample complexily = N/D
Skip connection c
Figure 1. (Left) Diagonal elements of the overlap Q (13) (red and bluc) and prediction error (black) achieved by GAMP, as function of the sample complexity a = S/D, for a two-layer attention model P: = P2 = c = M, = 2. Off-diagonal elements of Q are zero, and thus not plotted. Crosses: numerical implementations of GAMP in dimension D 1000, averaged over 16 runs. Convergence was reached with at most 50 iterations. Continuous lines: theoretical prediction from Eqs. 14 and (11). Dashed line: prediction error of GAMP when the first layer weights fixed to zero. (Right) Weak recovery thresholds ct (in red) and oz (in blue), above which the second and first layer weights are respectively retrieved, as function of the skip connection strength c. The red (resp. blue) lines indicate the sample complexities above which the second (resp. first) layers be learned.
Appendix B. We also provide complementary derivation of this result from the Belief Propagation scheme (Pearl, 1988) associated with the inference problem (8) in Appendix D. Lemma 2.2 provides asymptotic description of the perfor- mance of GAMP (Algorithm 1), which can be leveraged to ascertain when reconstruction is computationally achievable, what sample complexity is required to do so, and whether the algorithm achieves the Bayes-optimal performance.
achieve non vanishing overlap in any subspace U without having any significant overlap at all. This thresholds can be characterized in closed-fome in terms of the stability of the fixed-points of the state-evolution (14) of Lemma 2.2. Indeed, the iterates W can escape the neighborhood of the fixed point Q = O in the presence of vanishing side- information À if and only if the fixed-point Q = O becomes unstable. The weak- recovery threshold for initial recovery can thus be determined by performing linear stability anal- ysis. Notice that once this first subspace U1 is recovered, then new threshold (22 emerges: 022 describes the minimal sample complexity such that the fixed point becomes unsta- ble to perturbations orthogonal to U1. This procedure can be then repeated, creating a natural ordering in the learning of row span(W'* as a function of the sample complexity, which we call grand wrand-stairease mechanism analogy with Definition 6 in (Troiani et al., 2025). This discussion, which we prolong in Appendix C, is formalized in the following theorem.
Hereafter, let us consider more particularly the case of even SMI models, namely those invariant under the reflection W <>> - W of their weights. This is the case, for instance, of the tied deep attention model (2). As a consequence of this symmetry, Q = 0 is a trivial fixed point of the state- evolution of GAMP, and the latter fails to reconstruct any part of the row-space of W As is well known for sin- gle and multi-index models (Mondelli & Montanari, 2018; Troiani et al., 2025), the breaking of such a symmetry neces- sitates the introduction of side inferinfoommaion, defined as additional observation AW + ví -i, Xε, with & E RPxD a matrix with Gaussian i.i.d entries, and À E IR. The GAMP algorithm in the presence of side information differs from Algorithm only in the updates of W, Ć, as we discuss in Appendix D. The introduction of side-information allows us define a number of weak-recovery thresholds. Given subspace U RP the weak recovery threshold for U is defined as the smallest a such that the iterates W of the GAMP algorithm when projected along the rows of W* achieve non- vanishnig components along all directions of U. with arbitrarily small side-information coefficient We discuss the formal definitions and existence of such thresh- olds in Appendix C, based on the analysis in (Troiani et al., 2025). A special threshold we can characterize in closed form is the weak-recovery threshold for initial recovery, that is the smallest sample complexity ai such that GAMP can
C2 with G = Ow.oout 001pp) RPXPxMXX denoting the tensor- Jacobian of Jout with respect to W. Define F(X) to be the following linear operator on PxP
Theorem 2.3. Suppose that gout
M
P
F(X);;
= Σ
Σ
Ey (Grkabr (15)
a.b=1k
where Y is distributed as in (9) Then the weak recovery threshold for initial recovery namely for reconstructing any subspace is
1
sup
[F(())
(16)
a1
X
where |[ is the Frobenius norm and S.. is the set of positive semi-definite matrices of size Px P. More generally
6

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
these results seamlessly transfer deep attention models (2), and discuss in particular the consequences of Theorem
C1
2.3 for these architectures.
Cosine similarity
3.1. Layer-wise learning in depth L = 2 attention
We examine for definiteness the case of two-layer atten- tion model (2) (L 2, M 2, P1 P2 1, 1). Fig. (left) illustrates the theoretical predictions of Lemma 2.2 for the prediction error (11) and overlap Q (13) (right) achieved by GAMP, as a function of the sample complexity a. Before discussing the curves, let us highlight two remark- able features of the behavior of GAMP in this model. First, we observe that the values of these metrics convergence remain unchanged if GAMP is initialized with strong side information. This finding is telltale sign of the absence of computational waapoonnnn-t---tttttiiiieet statistical gaps for this particular problem, and strongly suggests that GAMP is in this case Bayes- optimal. Secondly, we notice that the off-diagonal elements Q12, Q21 of the overlap Q (13) vanish at convergence.
100
10
Iterations
Ci
C2
100
102
Epochs
Figure 2. (Top) Evolution of the performance of GAMP with the number of iteration, measured by the cosine similarity C:- ((wf)) ) / between the GAMP estimate wi of the 1-th layer weights and the target weights w", for L 2, P1 P2 = 1.c = M = 2. We display a single run of the algorithm in dimension D= 1000 and sample complexity a= N/D = 1.2 starting from the first iteration. (Bottom) Evolution of the cosine similarity, for the same target function, when training the same model using SGD. We display 8 runs of the algorithm in dimension D = 500 and sample complexity a = 15, with the average indicated in bold. The numerical experiments
A striking observation from Fig. 1 is that the learning pro- ceeds in sequential steps, the sample complexity is in- creased. For a < Q1 R 0.14 (grey regime in Fig. 1), there is insufficient data to learn either set of weights -as signaled by vanishing overlaps Q1 = Qz2 = 0-. As a consequence, the prediction error (black) stays sensibly constant, and does not decrease. In the regime 0i < a aq 0.79 (red regime in Fig. 1), the sample complexity is greater than the sharp threshold for the learnability of the second layer (1 2), and the corresponding overlap becomes non-zero Q22 # O, reflecting a partial reconstruction of w2. This translates into a sizable decrease in the prediction error throughout this phase. Finally, for sample complexities a > a2 (blue regime in Fig. 1), the first layer weights are also learned (Qu + 0), resulting in yet another decrease in prediction error. The different layers are thus learned sequentially, with a clear separation in sample complexity between the consecutive stages. The weak recovery thresholds 01, 02 delineating the different regimes can be computed using the results described in Section 2. Specifically, the initial re- covery threshold 01 is the critical sample complexity above which first subspace -in this case the span of the second layer weights- is recovered, as characterized in Theorem 2.3. on is the critical sample complexity required to recover the subspace spanned by the first layer weights, conditioned on having acquired a non-zero overlap Q22 > 0 with the second layer weights. We describe how to compute az in Appendix E.2. Finally, we chart the boundaries of the three regimes observed in Fig. I (left) in phase diagram (Fig. I (right)), as a function of the skip connection and the sam- ple complexity, evidencing how the phases subsist across a wide range of skip connection strengths c.
performed at A = 1.4 x 10 l = 15, and batch size 200, with each batch used for 3 consecutive iterations.
for a subspace U C RP, we denote by at, the threshold for weak-recovery of all directions in U. defined as:
ay = =inffa:: lim lim TT > 0}, (17)
X-0+ t-000
where Pú denotes the projection onto U and O2 E RPPP are recursively defined by:
2tt1=
=F Fx
a [x Dout (Y, - Q$)
(18)
with FX(2)=
+Q( ((1--))-1
Definition! 17 then specifies ordering of subspaces U1 C U2 C ...Uk C IR such that oV,, < OtU, < <..'U,. We denote au, for j = 1, 3... in the above ordering as ojj, with on matching the description in Equation 9, and in general use the shorthand oi = ave
The existence of the limit defined by Eq. (17) conse- quence of the monotonicity of the asymptotic dynamics of GAMP described by Eq. (18), see Appendix B.
3. Results for deep attention models
The previous section delineates the fundamental limits of learning in the class of SMI models (1). We now show how
7

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
1.0
6 x 10-1
Prediction error
1.0
Error
4 x 10
Similarity
Overlap
Q
3 x 10
0.5
22
2 10
0.5
S
Q33
10-1
S₂
0.0
0.0
0.5
1.0
1.5 2.0
10-1
10'
Sample complexity N/D
Epochs
Figure 3. (Left) Diagonal elements of the overlap Q (13) (red, blue and green) and prediction error (black) achieved by GAMP, as a function the sample complexity a = //D, for a three-layer attention model P1 = P3 = P3 = 1,c = I.M 2. Off-diagonal elements of Q are zero, and thus not plotted. Dashed line: prediction error when the first and second layer weights are fixed to zero. (Right) Similarity S = un(w]] T xf T ( [ II between the weights wn at training step and the last-iterate weights wf function of the training time. A transformer with L = 2 attention layers and fully connected readout is trained on the TREC classification task (Hovy et al., 2001; Li & Roth, 2002). Different colors indicate different attention layers, with the multiple curves representing distinet runs.
3.2. The role of depth
3.3. Dynamics of learning
is natural to wonder whether such layer-wise learning generic feature of deep attention networks or simply phenomenon proper the two-layer case. Fig. 3 (left) shows the diagonal elements of the overlap Q for a depth L = 3 attention model with P1- P22 = P3 = 1, 2, c = 1. As in the L= 2 case (see Fig. 1). the last layer is learned first as the number of samples increases. On the other hand, the shallower layers = 1,2 seem to share the same weak recovery threshold, and are learned sensibly at the same rate. We conjecture this to be general feature of deeper attention models.
The precedent subsection evidenced how the different layers warrant different sample complexities to be reconstructed and are learned sequentially by GAMP more samples be- come available. A similar sequential learning phenomenon can in fact also be observed at a fixed sample complexity, as a function of the number of ilerations. Fig. 2 (top) represents the evolution of the cosine similarity G = w w; /|| be- tween the GAMP estimate un of the layer weights and the corresponding target weights w over the GAMP iterations. The second layer is learned within the first few iterations. In contrast, the first layer is at first not learned, and its re- construction only commences after approximately 10 steps, when the second layer is almost perfectly reconstructed.
Let us make an interesting observation in this direction, which we leave as inspiration for future work. The deep attention model (2) admits an elegant limiting form in the joint limit of large depth - oo and large skip connection BL, with B finite, provided the activation function is simultaneously rescaled (22) = ==//o):
This layer-wise learning phenomenon over training time is not restricted to the GAMP algorithm. In Fig. 2 (bottom), we plot the performance of Stochastic Gradient Descent (SGD) in reconstructing the weights of the same two-layer archi- tecture, from the same dataset D = { (xx) ) by implementing the descent steps
YDA (z) of ((1 E OZ, 21) 2 zu 1 t L. LZ 2i)
の + = Σ V #EB,
u6y wi "i uo (x) - YDd (x)] -xxit
where we remind that 21 w denotes the projection of the input x on the 1-th layer weights w; A detailed derivation of this limiting form is provided in Appendix A.4. This simple expansion shows that, in the large depth limit, the deep attention model can be viewed as a combination of linear attention module Z Zz parametrized by the last layer weights, and "committee" of attention modules cor- responding to the first L - layers. Since these first layers are interchangeable we expect that also the overlaps should share the same invariance, making it possible to simplify the state evolution equations in this limit by assuming that also the overlaps need to have an interchangeable structure, in similar spirit to (Aubin et al., 2018).
with learning rate 1) and weight decay A. B, C D denotes the batch used for the t-th step. Notice how we reuse the same batch for multiple gradient updates (Arnaboldi et al., 2024a). The simulations reveal that, similarly to GAMP, the learning of SGD proceeds in two sequential steps. While SGD starts learning the second layer weights after one epoch, the first layer is only reconstructed later, when the second layer is almost perfectly Iearned, with a clear separation in training time between the two stages. We believe that describing analytically the gradient descent dynamics of SMI models, and in particular of deep atten-
8

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
tions, can be interesting research direction, which can be pursued leveraging the existing literature on the dynamics of multi-index models (Gerbelot et al., 2024; Celentano al., 2021; Dandi et al., 2024; Mignacco et al., 2020). Similarly, a systematic study of layer-wise learning in more realistic transformers could offer valuable insights, especially though the lenses of mechanistic interpretability.
models to sequential covariates. For this larger class of models, in the Bayes-optimal setting, we derive the statistically and computationally optimal estimation errors, and extract from those the weak recovery thresholds. In the case of a depth attention model, our findings uncover how distinct layers are learned sequentially, with increasing number of samples, or over training time, We briefly discuss how a related layer-wise learning phenomenon can also be observed in practical settings.
4. Discussion and conclusions
Among the limitations of our work, the SMI functions do not naturally account for fully connected layers usually interlaid with attention layers. Another limitation is that we consider only learning from Gaussian i..d. inputs; adding input correlation is possible but left for future work.
We have so far focused our investigations on the case of target functions of the form (2), evidencing how the weights of different layers of a two-layer attention architecture are learned sequentially over training time. this last section, we briefly illustrate how similar layer-wise learning phe- nomena can also be observed for more practical transformer architectures, with real datasets.
Acknowledgments
We thank Freya Behrens, Vittorio Erba, Mikhail Terekhov, and Fabrizio Boncoraglio for discussions. We acknowledge funding from the Swiss National Science Foundation grants SNFS SMArtNet (grant number 2 212049), and OperaGOST (grant number 200021 L_003900.. HC acknowledges support from the Center of Mathematical Sciences and Applications (CMSA) of Harvard University.
We consider a classification problem on the Text REtrieval Conference (TREC) dataset (Hovy et al., 2001; Li & Roth, 2002), in which the network is tasked with labeling the object of a question. The data is pre-processed using the uncased base BERT model (Kenton & Toutanova, 2019) to obtain token embeddings that are subsequently fed into simple transformer architecture with two consecutive at- tention layers, and fully-connected readout layer. The network is trained on the cross-entropy loss using the Pytorch (Paszke et al., 2017) implementation of the AdamW (Loshchilov & Hutter, 2019) optimizer, over 40 epochs. Further details can be found Appendix F. Fig. 3 shows the evolution of the similarity metric
Impact statement
The goal of this work is to advance the field of theory of Machine Learning. There are multiple potential societal impacts which can follow, none of which we feel needs to be particularly highlighted.
S = tr((w,] w ) w))))
(19)
/w/l] [u/[|?]] ||
References
between the weights w/ training time t and the last-iterate weights wf, for a given layer I E (1,2), as a function of t. The curves reveal a marked sequential learning process. They are furthermore quite reminiscent of our findings for two-layer attention target functions, gathered in Fig. 2, which we discussed above. Remarkably, however, compared to Fig. 2, the order in which the layers are learned in the TREC task is reversed, echoing the findings of (Chen al., 2023) that shallower layers tend to converge faster over training. We expect the order in which layers are learned to generically depend on the intricate interactions between the architecture, task and algorithm. Ascertaining precisely the factors determining the order in a given setting is challeng- ing task, which warrants careful mechanistic studies. We leave this question an exciting future research direction.
Abbe, E., Adsera, E. B., and Misiakiewicz, T. The merged- staircase property: a necessary and nearly sufficient con- dition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782-4887. PMLR, 2022.
Abbe, E., Adsera, E. B., and Misiakiewicz. T. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory, pp. 2552-2623, PMLR, 2023.
Abbe, E., Bengio, S., Boix-Adsera, E., Littwin, E., and Susskind, J. Transformers learn through gradual rank increase. Advances in Neural Information Processing Systems, 36, 2024a.
To conclude, in the present manuscript, we consider the problem of learning deep attention network, defined as the composition of multiple self-attention layers. We first establish that such architectures pertain to the class of sequence models, which generalize multi-index
Abbe, E.,, Bengio. S., Lotfi, A., Sandon, C., and Saremi, O. How far can transformers reason? the globality barrier and inductive scratchpad. Advances in Neural Informa- tion Processing Systems, 17:27850-27889 2024b.

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn implement preconditioned gradient de- scent for in-context learning. Advances in Neural Infor- mation Processing Systems,
the National Academy of Sciences, 116(12)5551-546)) 2019.
Bayati, M. and Montanari, A. The dynamics of message passing on dense graphs, with applications to compressed sensing. IEEE Transactions on Information Theory. 57 (2)76647755, 2. doi: 10. N0.1109/TIT.200 2094817.
2023.
Arnaooldi, L., Stephan, L.., Krzakala, F., and Loureiro, B. From digh-dimensional & mean-field dynamics to dimen- sionless odes: A unifying approach to sgd in two-layers networks. In The Thirty Sixth Annual Conference on Learning Theory, pp. 11 1227. PMLR, 2023.
Berthier, R., Montanari, A., and Nguyen, P.-M. State evolu- tion for approximate message passing with non-separable functions. Information and Inference: A Journal of the IMA, 9(1:333-79 2020.
Arnaboldi, L.. Dandi, Y., Krzakala, F. Pesce, L., and Stephan, L. Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405. 15459, 2024a.
Berthier, R., Montanari, A., and Zhou, K. Learning time- scales in two-layers neural networks. Foundations of Computational Mathematics, pp. 1-84, 2024.
Arnaboldi, L., Krzakala, F., Loureiro, B., and Stephan, L. Escaping mediocrity: how two-layer networks learn hard generalized linear models with sgd. arXiv preprint arXiv:2305. 10.5.8855 18502, 2024b.
Bietti, A., Bruna, J., and Hillad-Vivien L. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310. 19793, 2023.
Arous. G. B., Gheissari, R., and Jagannath, A. Online stochastic gradient descent on non-convex losses from high-dimensionall inference. Journal of Machine Learn- ing Research, 22(106):1-51, 2021.
Bordelon, B., Chaudhry, H., and Pehlevan, C. Infinite limits of multi-head transformer dynamics. Advances Neural Information Processing Systems, 37:35824-35878, 2024.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners, 2020.
Ataee Tarzanagh, D., Li, Y., Zhang, X., and Oymak, S. Max- margin token selection in attention mechanism. Advances in Neural Information Processing Systems, 36:48314- 48362, 2023.
Camilli, F., Tieplova, D., and Barbier, J. Fundamental limits of overparametrized shallow neural networks for supervised learning. arXiv preprint arXiv:2307.0560 2023.
Aubin, B., Maillard, A., Krzakala, F., Macris, N., Zde- borová, L.,, et al. The committee machine: Computational to statistical gaps in learning two-layers neural network. Advances in Neural Information Processing Systems, 31, 2018.
Celentano, M., Montanari, A., and Wu, Y. The estimation error of general first order methods. In Conference on Learning Theory, pp. 1078-1144 111. PMLR, 2020.
Aubin, B., Loureiro, B., Baker, A., Krzakala, F., and Zde- borová, L. Exact asymptotics for phase retrieval and compressed sensing with random generative priors. In Mathematical and Scientific Machine Learning, pp. 55-
Celentano, M., Cheng, C., and Montanari, A. The high- dimensional asymptoties of first order methods with ran- dom data. arXiv preprint arXiv: 12..0757.. 2021.
73. PMLR, 2020.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., Wu, D., and Yang, G. High-dimensional asymptotics of feature leam- ing: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35; 2022.
Chen, Y., Yuille, A., and Zhou, Z. Which layer is learning faster? systematic exploration of layer-wise conver- gence rate for deep neural networks. The Eleventh International Conference on Learning Representations, 2023.
Barak, B., Edelman, B., Goel, S., Kakade, S., Malach, E., and Zhang, C. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750- 21764, 2022.
Collins- Wooins-Woodfin, E., Paquette, C., Paquette, E., and Seroussi, Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index mod- els. Information and Inference: A Journal of the IMA, 13 (4):iaae028, 2024.
Barbier, J., Krzakala, F., Macris, N., Miolane, L., and Zde- borová, L. Optimal errors and phase transitions in high- dimensional generalized linear models. Proceedings of
Cui, H. High-dimensional learning of narrow neural net- works. Journal of Statistical Mechanics: Theory and Experiment, 2025(2):023022, 2025.
10

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
Cui, H., Krzakala, F., and Zdeborová, L. Bayes-optimal learning of deep random networks of extensive-width. In International Conference Machine Learning, pp. PMLR, 2023.
Edelman, B. L.,, Goel, S., Kakade, S., and Zhang, C. Induc- tive biases and variable creation in self-attention mecha- nisms. In International Conference Machine Learning, pp. 3793-5831. PMLR, 2022.
Erba, V., Troiani, E., Biggio, L., Maillard, A., and Zde- borová, L. Bilinear sequence regression: A model for learning from long sequences of high-dimensional tokens. arXiv preprint arXiv:24100 10. 18858, 2024.
Cui, H., Behrens, F., Krzakala, F., and Zdeborová, L. A phase transition between positional and semantic learning in solvable model of dot-product attention. Advances in Neural Information Processing Systems, 37:36342- 36389, 2024a.
Fu, H., Guo, T., Bai, Y., and Mei, S. What can a single attention layer learn? study through the random features lens. Advances in Neural Information Processing Systems, 36, 2024.
Cui, H., Pesce, L.,, Dandi, Y., Krzakala, F., Lu, Y., Zde- borova, L., and Loureiro, B. Asymptotics of feature learning in two-layer networks after one gradient-step. In International Conference on Machine Learning, pp. 0695-9695..2 PMLR, 2024b.
Gardner, E. and Derrida, B. Optimal storage properties of neural network models. Journal of Physics A: Mathemat- ical and general, 21(1):271, 1988.
Damian, A., Lee, J., and Soltanolkotabi, M. Neural net- works can learn representations with gradient descent. In Loh, P.-L. and Raginsky, M. (eds.), Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Pro- ceedings of Machine Learning Research, pp. 5413-5452. PMLR, 02 02-051M0 Jul 2022.
Gerbelot, C. and Berthier,
Craph-based approximate message passing iterations. Information and Inference: Journal of the IMA, 12(4):2562-2628 2023.
Gerbelot, C., Troiani, E., Mignacco, F., Krzakala, F., and
Zdeborová, L. Rigorous dynamical mean-field theory for stochastic gradient descent methods. SIAM Journal on Mathematics of Data Science, 6(2):400-42) 2024. doi:
Damian, A., Nichani, E., Ge, R., and Lee, D. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2024a.
10. N0.1137/2M15944333
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. A mathematical perspective on transformers. arXiv preprint arXiv:2312. 10794, 2023.
Damian, A., Vivien, L.,, Lee, J., and Bruna, J. Computationallstaiisiie gaps in gaussian single-index models. In The Thirty Seventh Annual Conference on Learning Theory, pp. 1262-1262. PMLR, 2024b.
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Dandi, Y., Krzakala, F., Loureiro, B., Pesce, L., and Stephan, L. How two-layer neural networks learn, one (giant) step at a time. arXiv preprint arXiv:2305. (tt.::335///87000 2023.
Glasgow, M. Sgd finds then tunes features in two-layer neural networks with near-optimal sample complex- ity: A case study in the xor problem. arXiv preprint wrXi::2309.55111 2023.
Dandi, Y., Troiani, E., Arnaboldi, L., Pesce, L.,, Zdeborova, L., and Krzakala, F. The benefits of reusing batches for gradient descent in two-layer networks: breaking the curse of information and leap exponents. In Proceed- ings of the 41st International Conference on Machine Learning, pp. 9991-10016, 2024.
Györgyi, G. First-order transition to perfect generaliza- tion in a neural network with binary synapses. Physical Review A, 41(12):7097, 1990.
Hahn, M. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 171-171, 2020.
Dandi, Y., Pesce, L., Cui, H., Krzakala, F., Lu, Y., and Loureiro, B. A random matrix theory perspective on the spectrum of learned features and asymptotic generaliza- tion capabilities. In The 28th International Conference on Artificial Intelligence and Statistics, 2025.
Hovy, E., Gerber, L., Hermjakob, U., Lin, C.-Y,, and Ravichandran, D. Toward cemantics-based answer pin- pointing, In Proceedings of the first international confer- ence on Human language technology research, 2001.
Donoho, D. L., Maleki, A., and Montanari, A. Message- passing algorithms for compressed sensing. Proceedings of the National Academy of Sciences, 106(45):18914- 18919, 2009.
Javanmard, A. and Montanari, A. State evolution for general approximate message passing algorithms, with applica- tions to spatial coupling. Information and Inference: Journal of the IMA, (22):115-144 2013.

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
Jelassi, S., Sander, M., and Li, Y. Vision transformers provably learn spatial structure. Advances in Neural In-
Mezard, M. and Montanari, A. Information, physics, and computation. Oxford University Press, 2009.
formation Processing Systems, 35:37822-378866 2022.
Mézard, M., Parisi, G., and Virasoro, M. A. Spin glass the- ory and beyond: An Introduction to the Replica Method and Its Applications, volume 9. World Scientific Publish- ing Company, 1987.
Kenton, J.
D. M.-W. C. and Toutanova, L. K. Bert: Pre- training of deep bidirectional transformers for language understanding, 2019.
Lauditi, C., Troiani, E., and Mézard, M. Sparse represen- tations, inference and learning. Journal of Statistical Mechanics: Theory and Experiment, 2024(10):004001, 2024.
Mignacco, F., Krzakala, F., Urbani, P., and Zdeborová, Dynamical mean-field theory for stochastic gradient de- scent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33:9540-95500 -550,, 2020.
Li, H., Wang, M., Liu, S., and Chen, P.-YY A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference Learning Representations, 2023a.
Monasson, R. and Zecchina, R. Weight space structure and internal representations: a direct approach to learning and generalization in multilayer neural networks. Physical review letters, 75(1))24432, 1995.
Li, X. and Roth, D. Learning question classifiers. In COL- ING 2002: The 19th International Conference on Com- putational Linguisties, 2002.
Mondelli, M. and Montanari, A. Fundamental limits of weak recovery with applications to phase retrieval. In Conference On Learning Theory, pp. 1445-1450. PMLR, 2018.
Li, Y., Li, Y., and Risteski, A. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pp. 19689-19729, 1ML... 2023b.
Moniri, B., Lee, D., Hassani, H., and Dobriban, E.
theory of non-linear feature learning with one gradient step in two-layer neural networks. In Proceedings of the 41st International Conference on Machine Learning, pp. 36106 36106-36159, 2024.
Lopardo, G., Precioso, F., and Garreau, D. Attention meets post-hoe interpretability: A mathematical perspec- tive. In Forty-first International Conference Machine Learning, 2024. URL /ttttp::///peeeevvieee www.e/eeeeeeeee////e//mmmmm...// forum? Cerummid-wntee wwm.:iit1nn..1112....
Nishimori, H. Exact results and critical properties of the ising model with competing interactions. Journal of Physics C: Solid State Physics, 13(21):4071, 1980.
Loshchilov, and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2019.
Oko, K., Song, Y., Suzuki, T., and Wu, D. Learning sum of diverse features: computational hardness and efficient gradient-bass based training for ridge combinations. In The Thirty Seventh Annual Conference on Learning Theory, pp.
Lu, Y. M., Letey, M. I., Zavatone-Vert Veth, J. A., Maiti, A., and Pehlevan, C. Asymptotic theory of in-context learning by linear attention. arXiv preprint 1IXiv:2405.1175/ 2024.
PMLR, 2024.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang. E., Veto,, 7.,, Lin, Z., Desmaison, A., Antiga, I.,, and Lerer, A. Automatic differentiation in pytorch. In NIPS-W, 2017.
Maillard, A., Troiani, E., Martin, S., Zdeborová, L., and Krzakala, F. Bayes-optimal learning of extensive- width neural network from quadratically many samples. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 82085-82132. Curran Associates, Inc., 2024. URL opperffilees/ppper/////////// 05J17711190002006668101102111112
Pearl, J. Probabilistic reasoning in intelligent systems: net- works of plausible inference. Morgan Kaufmann, 1988.
Rende, R., Gerace, F., Laio, A., and Goldt, S. Mapping of attention mechanisms generalized potts model. Physical Review Research, 6(2):023057. 2024.
Saad, D. and Solla, S. Dynamics of on-line gradient descent learning for multilayer neural networks. In Touretzky, D., Mozer, M. C., and Hasselmo, M. (eds.), Advances in Neural Information Processing Systems, volume MIT Press, 1996.
Marion, P., Berthier, R., Biau, G., and Boyer, C. Attention layers provably solve single-location regression. arXiv preprint ar arXiv:2410001337. 2024.
12

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
Saad, D. and Solla, S. A. On-line learning in soft committee machines. Physical Review E, 52(4):4225, 1995.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.
Sahiner, A., Ergen, T., Ozturkler, B., Pauly, Mardani, M., and Pilanci, M. Unraveling attention via convex dual- ity: Analysis and interpretations of vision transformers. In International Conference on Machine Learning, pp- 19050-19088. PMLR, 2022.
Wainwright, M. J. and Jordan, M. I. Graphical Models, Exponential Families, and Variational Inference. Foun- dations and Trends in Machine Learning, 2008. doi: 10.1561/220000001.
Schwarze, H. Learning a rule in a multilayer neural network. Journal of Physics A: Mathematical and General, 26(21): 5781, 1993.
Tavatone-Veth, J. A., Tong, W. L.,, and Pehlevan, C. Con- trasting random and learned features in deep bayesian linear regression. Physical Review E, 06411) 2022.
Schwarze, H. and Hertz, J. Generalization in a large com- mittee machine. Europhysics Letters, 20(4):375, 1992.
Zdeborová, L. and Krzakala, F. Statistical physics of infer- ence: Thresholds and algorithms. Advances in Physics, (5(5))453-555 2016.
Simsek, B., Bendjeddou, A., and Hsu, D. Learning gaussian multi-index models with gradient flow: Time complexity and directional convergence. In The 28th International Conference on Artificial Intelligence and Statistics, 2025.
Zhang, R., Frei, S., and Bartlett, L. Trained transform- ers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
Sompolinsky, H., Tishby, N., and Seung, H. S. Learning from examples in large neural networks. Physical Review Letters, 65(13):1683, 1990.
Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. Transformers support vector machines. arXiv preprint arXiv:2308. /6898, 2023.
Tian, Y., Wang, Y., Chen, B., and Du, S. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in Neural Information Processing Systems, 36.71911-71947, 2023.
Tiberi, L., Mignacco, F. Irie, K., and Sompolinsky, H. Dis- secting the interplay of attention paths in a statistical mechanics theory of transformers. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Troiani, E., Dandi, Y., Defilippis, L., Zdeborova, L., Loureiro, B., and Krzakala, F. Fundamental computa- tional limits of weak learnability high-dimensional multi-index models. In The 28th International Confer- ence on Artificial Intelligence and Statistics, 2025.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.. Gomez, A. N., Kaiser, L., and Polosukhin, I. Atten- tion is all you need. Advances in Neural Information Processing Systems, 31, 2017.
Veiga, R., Stephan, I., Loureiro, B., Krzakala, F., and Zde- borová, L. Phase diagram of stochastic gradient descent in sigh-dimensional two-layer neural networks. Advances in Neural Information Processing Systems, 35:23244- 23255, 2022.
13

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
A. Generalization to untied and multi-head attention
The backbone of the analysis of deep attention models (2) reported in Section 3.1 of the main text is the mapping (7) to a SMI model (1). As we extensively discuss in the main text, this formal connection permits the transfer of analytical ideas from the study of multi-index models that of deep attention architectures. Although we considered in the main text attention architectures with tied query and key weights for clarity and ease of exposition, the class of SMI functions (1) in fact encompasses much broader array of deep attention architectures, with untied weights, and generic number of heads, In this Appendix, we detail the mapping of these more complex architectures to SMI functions. A direct consequence of this mapping is that the general results on the statistical (Theorem 2.1) and computational (Lemma 2.2 and Theorem 2.3) limits of learning in SMI models directly apply to such architectures. We discuss sequentially the case of deep attention functions with untied weights, before extending the mapping to multiple heads. Finally, we derive the infinite depth limit I. -> xo discussed the main text. In the following we will use bold characters for all the quantities that with (D) elements, as we did in the main text.
A.1. Untied weights
As such, our main theoretical results consequently hold for untied architectures as well, In this paragraph, we detail the mapping of deep attention network with generically untied weights to SMI. Consider the model
[2T_110)
22)
QT L w xL-1
10) w Q w - (20)
(UDA(x))= =a
AI E [1, 21 = x1 cl M +o
D
D
where wf. w E RR are two sets of independently sampled standard Gaussian weights. As we did in the main text, we can write the untied deep attention model yudA (æ) as an SMI model by introducing two families of indices Kr, Qi for the keys and the queries respectively
w
Q wi
x
Q
(21)
VD
VD
with which we can define a new mixing functions ( Kn, Q1 ..... 2) defined recursively
= 1B &,, (1 & to (B' (TKKBB
(22)
with B = 1,., giving us
(UDA(()) = (K,, KL, Q1...., 21)
(23)
This corresponds to an SMI model (1). To this, denote P Pi ... + Pr, and introduce the concatenated weights WUDA E R2PXD, defined the matrix obtained by vertically stacking w ..... w w/ ,.., the link function JUDA 222x M.M.M
.... w in this order. Then define
(X) = B (XLL)) Xppp)())))))))))
(24)
where X(ss) denotes the (j (j-11 1 1) x submatrix of X comprising the rows indexed between and j. Then, the deep attention architecture YUDA can be rewritten compactly as
VUDA (x) gUDA (UUDA)..
(25)
In other words, the untied deep attention is also a SMI. Consequently, all the technical results reported in the main text directly apply to this architecture.
A.2. Sequence to Sequence multi-layer attention models
The SMI class of models defined by Eq. (1) further subsume sequence-to-sequenee models with output in RDxM given by:
[T__ w
wix!
y =1-1 cl M to
D
instead of the MM output in Eq. (2). Above, at are defined recursively as in Eq. (3) for E [Z].
14

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
The mapping of the above model to SMI follows by writing
x (2. Zi),
(26)
where the functions By are the equivalent for sequence to sequence models of By in (5)
(21;-.. (Z ... Z)) B (Z1, ...... Z1-1) 1) 61,L)c1M + (Z. 71-1 Z Z.BB11 (21 -111))
(27)
and B: 11 a. Intuitively (Z Z2) RRMM denotes the components of the tokens (columns) in along the basis of the input tokens in æ. Since we assume D > P. while itself is high-dimensional, it can be constructed from by specifying the M x M scalar entries of B (2 ...,,. Zi). Hence, the estimation of y given æ is equivalent to the estimation of (2 [Z,,.... Zi) given æ, which reduces to a sequence multi-index (SMI) model with outputs RMXMM.
Note that for single layer L 1 this sequence-tor to to-sequence variant of the model was analyzed in (Cui et al., 2024a).
A.3. Multi-head attention
By the same token, multi-head architectures can also mapped to a SMI, and thus fall within the scope of our theoretical results. Let us take the input data æ E RDDHxM sequence of length M embedded in dimension DH. At each layer the embeddings are split into H vectors of dimension D, which are then projected to dimension Pt. We can model this by first splitting into list of H elements (x* BRPxM14 l==1 and then defining H H sequence indices for each layer both for the keys and the queries
Aw haha 2
Khuha
Qluhe
(28)
VD
VD
Each attention head is defined as
h1 TQ Nhi
T Kwwtt hiz
1-1
x--1
1
Σ
(Qf
T
(29)
HD
I1
h1.hg
h1.h2
The deep multi-head architecture is then defined as
1 "
Co Tccilt ]
VI E [1, LII, of at-1
(1 dr,1)cM to
(30)
H
b2g
with
YMHA (x) = Σ To
(31)
where Un E PRxx are trainable readout weights. As before, we introduce set of functions B (R#*H*RxMy082
x (R#xHxXXX XM + RHXMXM
[ [B_-' IT, (Ohhi) xchhh (B( |ha ) 188--111
B =IB²--11 ((,z)) (1 (1- k to
Σ
(32)
h1.b2
with [B9] 1 Then
UMHA (x) = Σ 1 .... K..,,,, M (12)1 .... 2 1
(33)
defining V/, = =.,, aa.
To expound the connection of this deep, multi-head model to SMI models, first denote P = (PI + + PL) + Pr, and
define the total weights WMHA E Px DH defined as
T MHA = (K W1 ... WW K Wwi.. ... K WW ow].. ... 9w ... 9w] 9w v).
(34)
15

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
with
K w
T
V1
w/th
Q W =
(35)
Qu uoot
Finally, we construct the link function 9MHA RRPxM - RPPXM For any argument X E RPXM we view X as vertical block vector with 2HT. + H blocks {(KX) X/ E @H ((0X] E M ( E RBPxM) JJ, in this order. Each of these blocks is further viewed as block matrix of H vertically stacked blocks, which we index with the subscript W. Then we define
9MHA (X)
Σ VX ( (II((
(I(9X)),,,...
(36)
With those definitions, it follows that the multi-head deep attention model can be rewritten as
UMHA (x) = 9MHA ( WMHA)),
(37)
i.e. this model pertains the class of SMI models. All the technical results established for SMI models in the main text thus directly transfer to multi-head architectures.
This Appendix showcases how fairly intricate deep attention architectures, with generically untied weights and multiple heads, fall into the class of SMI models - which thus proves a truly rich class of models. The study of the phenomenology of these more complex models and the implication of our thcoretical characterization, therefore, are left for future works.
A.4. Large depth limit of tied, single-head attention
In this subsection we study the limit strong skip connection in the model we are considering, showing how the formalism simplifies in this limit. We start by redefining the multi-layer attention by rescaling each layer
[x por]w;wj x1-1 )
VI E [1,L - II, 2 = x1-1 1M + -
(38)
D
With these new definitions, we are interested the large depth and skip connection limit c, L + co, while keeping the two terms proportional by imposing BL. We start by writing the functions B in this new definition
..., (,,,,...,,Z))
Lo 1 ... ....,, Z1-1) (1 _o
Z.
(Z,, 211))]
(39)
with Bº = 11. We would like to expand the functions above for large c, we define the functions A such that
c 1t + (!)
B = 1M +
(40)
For 1 0 we have A = 0. It possible to determine the other values of A by expanding (39)
B¹ B IM + - 1)B Z,B(-1)
1
=(1Ns +
1 _A6-11
IM + Lo
((1ur -a( +
2T22, 2i 1l - - 1
+0 #
(41)
=1M + -
-
1
((IM
/1M+>
1 -
27 2i
1M + 1 -bb-1) +
=1M + Al-1
1
+
1
[ Z₂) 0
1
16

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
from which we can extract
1 -%= +o t
(42)
We can thus provide a closed form expression for A
WE 11,L 11,
A/= = Σ Zr)
(43)
pm1
Finally, we have the result
4-1
1 4-1
(ZT
BL .... Zz) Zz)= =
BL
Z Zz
1
(44)
BL
BL 1=1
1=
This is very special function. In order to highlight its properties, will study generic functions g RMX x RRM RMXM
+
H(J = (1+ 11 J
(1+4)
11
(45)
with the identification
4-1 Σ o (z zi)
(Z, Zp) = ZTT.Z,,
(46)
L
1=1
B. Proof of Theorem 2.1
In this Appendix, we detail the equivalence, discussed in subsection 2.1 of the main text, between SMI function (1) and a classical multi-index model acting on flattened input and with shared weights. This connection instrumental in the application of the results of (Gerbelot & Berthier, 2023; Aubin et al., 2018; Troiani al., 2025) for multi-index models to SMI models. Leveraging this equivalence, we then provide a proof of Theorem 2.1.
B.1. Mapping SMI models block-structured multi-index models
Consider the setup sequence multi-index models defined by Eq. (1) and suppose that the link function g RRXM RK outputs K- K-dimensional vectors yw- (x).
Given sequence E RDxM, denote by De E RDD the corresponding "flattened" vector defined by:
x² = [x]]
x M
(47)
where "1,"" dM denote the columns of
Under the above flattening, the pre-activations Wa a E R. can re-expressed through projections of the flattened inputs xe onto block-structured weight matrix given by:
W
W
0,
W
0
E R
PMXDM
(48)
Equivalently, W can be represented as the Kronecker product W & IM The above representation leads to the following correspondence:
(W*
(W*+e)
(49)
[W**MM)
Therefore, the PM -if-dimensoonal vector (Wix.)) ar.) is obtained by stacking the projections of different columns onto W
17

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
With the above definition, VW* can be equivalently expressed as:
W'
W. aie
VW
SMI
VD
9-
VD
(50)
where ge RRPM - RK is obtained by lifting g through the mapping defined by Eq. (49) to act on inputs in RPM
B.2. Proof of Theorem 2.1
We make the following assumption on derived from (Aubin et al., 2020; Barbier et al., 2019):
Assumption B.1. The non-linearity in SMI (model 1) lies in C² and is entry-wise bounded.
For technical reasons, further assume that Gaussian noise matrix with i.i.d (/ A) entries with arbitrarily small variance A is added to the SMI output (1) and the corresponding channel (9):
Assumption B.2. We consider the limit A - under the modified activation ā(-) g(()) + VDg.
where { (O, Ix) denote noise independent across samples.
The above assumptions required due to the limitations of the adaptive interpolation framework in (Aubin al., 2020; Barbier et al., 2019) and we believe they can be relaxed through more careful control of the error terms.
We establish Theorem 1.1 by combining the above mapping defined by Eq. (50) with the adaptive interpolation scheme in (Aubin et al., 2020; Barbier et al., 2019). The central idea behind the adaptive interpolation scheme to interpolate between the high-dimensional Hamiltonian corresponding to the posterior measure (Eq. 8) and an effective "factorized" Hamiltonian constructed through the oow-dimensional output channel defined by Eq. (9) and an effective input channel on R... Analogous Caupsian-comparison inequalities, one then utilizes Stein's 's Lemma relate the time-derivative of the free-energy over the interpolation path to the overlaps corresponding to the Hamiltonian time t. Using the concentration of overlaps (justified through the Nishimori identity and the inclusion of vanishing side-information), one then obtains family of interpolation paths satisfying systems of ODEs such that the high- nigh-dimensional remainder terms during the interpolation vanish. Selecting particular choices for such interpolation paths then yields matching lower and upper bounds for the free-energy.
The mapping of SMI models to block-structured multi-index models on flattened inputs given in Section B allows us to prove the asymptotic expression for the free-energy through reduction the steps utilized in the proof of Theorem 3.1 in (Aubin et al., 2020) and Theorems 1,2 in (Barbier al., 2019). For brevity, only describe the central differing step. related to the derivation of the time-derivative of the free-entropy described above.
Following (Aubin et al., 2020), we parameterize the interpolation paths through two functions r(f) [0, + Sp and (tt) : [0, 1] + Sp and scalar parameters 61, 62 E R Let Rí (t), (t) denote the solutions to the following ODEs:
Ri(1) = +1 ( (())dss
R2(1) = e + ('((t)a (()dss
f."
(51)
Define Se,u E PPM for u E [M] as the following "interpolating" inputs the effective output channel (Eq. (9)):
Wx VD
==
+ /R2(t) + Vt1 Re(t)
(52)
where Wn. U E RPxM have entries denote unkown parameters.
sid NO0 1), with Wμ playing the role of observed side-information while W, U
We further define the corresponding interpolating variables the input channel for i E [D]:
Y / (t) + Z,
(53)
where Y2, X3 R.P. We thus obtain two sets of "interpolating" parameters, observations:
Ye,, g(S..));
(54)
18

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
for μa E [M] and y for E [D]. These define an "interpolating posterior":
N II
P(W, UY",UIKK (Y Pus Tas
un (ulne(N]) =
1
II e 311%
Rioj. 1
Py (Ye,u)S..))) (55)
(t)
d=1
p=1
(Eq. 17 in (Aubin et al., 2018)) and the associated free-entropy:
1
(t) - Eloo (t)) 71
(56)
which reduces to the original model (Equation 8) at 6. O.
The only difference from the setup in (Aubin et al., 2020) lies in that St are matrices, instead P-dimensional vectors. However, leveraging the flattening defined by Eq. (49), we may express St.p as:
W. Te VD
Sune == -1-t-
+ Wiz()) + ( H.er.
(57)
Under the above flattening, equation 44 (Aubin et al., 2018) (arising from applying Gaussian integration by parts to the interpolating Hamiltonian w.r.t variables æe and ( are modified with the replacement:
WW W.W.
(58)
The mapping is then completed by noting that the flattening operator commutes under multiplication (as one expects from the properties of Kronecker/hensor products). Therefore:
W.W. (W"
(59)
Under the above substitution, the limiting free-entropy is obtaining by following the remainder of the proof of Theorem 3.1 in (Aubin et al., 2018), resulting in the variational problem given by Eq. (10). Subsequently, the resulting prediction error is obtained through steps identical to the proof of Theorem 2 in (Barbier et al., 2019).
C. Proofs of Lemma 2.2 and Theorem 2.3
In this section, we prove the results presented in Section 2.3, namely the tight asymptotic characterization of AMP (Algorithm in terms of its finite-dimensional state evolution equations (Lemma 2.2), and the characterization of weak recovery thresholds (Theorem 2.3). We start by establishing Lemma 2.2 through a mapping to general results for GAMP algorithms with non-separable non-linearities. We remark that Lemma 2.2 and Theorem 2.3 themselves do not establish the optimality of Algorithm or even prescribe its particular form the latter will be derived in Appendix D. However, by comparing the resulting expressions for state-evolution (14) to the extremization criteria in Theorem 2. we observe that GAMP asymptotically achieves the optimal error.
Throughout, we make the following assumptions on the non- non-linearity g:
Assumption C.1. The denoiser Jout defined by 12 lies in C².
is easy to check that the above assumptions holds for any g E C² the presence of vanishing noise.
C.1. Proof of Lemma 2.2
Analogous to the proof Theorem 2.1, the proof of Lemma 2.2 relics on the mapping of the sequence models to multi-index models on flattened inputs, as described in Appendix B. Under this mapping, Algorithm can be expressed the following form:
n' xB') -9'-1 .y)V B++1 g
(60)
(S' y) + fB(B))
(61)
19

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
where S2 E RR-xPM contains rows corresponding the flattened version of w in Algorithm as per Eq. 49 and B RRDMPPM contains columns given by the columns of W. in Algorithm 1, flattened as per Eq. 48. Vi, A, then denote the corresponding "flattened" Onsager terms. Finally, the non-linearities 9r and fe denote the flattened versions of go and (1p + A') respectively:
ge ge( y) Dout (wip Mude PM,
(62)
applicd row-wise and:
f (B") = (1p AA)
(63)
where the notation, e represents the flattening and tensorization defined in Equations 47, 48. We note that the slock-structured multiplication in fi. which maintains the same estimate W for different rows of W is non-separable across the DM rows of B Such non-seperable non-linearitie are however, permitted within the existing analysis in (Gerbelot & Berthier, 2023; Berthier et al., 2020) since it leads to well-defined convergence of empirical averages (Assumptions A5-A7 in Theorem I of (Gerbelot & Berthier, 2023)). Furthermore, using the commutativity of the tensorization and matrix matrix-multipliae multiplicaaion in (59),
the resulting state-evolution equation (Eq. (14)) admits the same form for the regular multi-index models in (Troiani et al., 2025).
Therefore, applying Theorem 1 in (Gerbelot & Berthier, 2023), obtain that for any number of finite time step t N:
W'WW")T D
d=00 , Q'
(64)
where Q' satisfies Eq. 14.
C.2. Definition of the Weak-Recovery Thresholds and the Proof of Theorem 2.3
As explained in the discussion preceding Theorem 2.3, Lemma 2.2 allows us to characterize weak-recovery under vanishing side-information through stability analysis of the state-evolution. We begin by providing the precise definitions for weak-recovery thresholds based on (Troiani et al., 2025):
Definition C.2. Consider Algorithm 1 in the presence of side information of the form XW + 1-X6, where & contains independent entries Eıj ~ (N, 1). We say that ag > 0 is weak-recovery threshold for Algorithm if the following hold with high probability as D + oc:
1. For any a Oe. there exists ō > 0 such that for any (arbitrarily small) À >> 0, 3t Ollg 4) such that 8.
W' 1 D
2. For a do.
((("") [ D
< <ev for some constant e > 0 and any / E N.
Analogously, we define thresholds characterizing the weak-recovery of given a subpsace U C RR:
Definition C.3. For a subspace U C RP, we say that at > 0 is the subspace weak-recovery threshold for U if:
1. For any ag, there exists 8 > 0 such that for any (arbitrarily small) A 0, 3t = O(log })) such that W'W' > 8PU.
2. For a ov, Fe E U with ||w]| such that
Wive vvo" T TW")""""
S e for some constant for any E N.
Under the above definitions and Lemma 2.2, the proof of Theorem 2.3 follows (Troiani et al., 2025). Since the form of state-evolution in Equation 18 exactly matches the form in (Troiani et al., 2025), we only outline the main steps.
The inclusion of side-information XW* + -1- Xg. which amounts to a modification the prior over W, changes the state evolution dynamics in Lemma 2.2 to Equation 18. Under the assumption Jout C2 the following linearization holds for the state-evolution dynamics defined by Equation 18:
Q1+1 ~ xFF(BQ) + VAL + ((al|(2 O(A),
(65)
20

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
The term Vll arising from side-information implies that at / 1, the overlaps satisfy for small enough A:
tr(Q', > A > 0,
(66)
with Q' E Sp. The generalized-Peree Perron Frobenius theorem then implies that F admits leading eigenvector in Sp.
Therefore, whenever a > ac, we obtain recursion of the form:
tr( Q*) 2 (1+ (1+))
(67)
for some K 0, leading the positive part of Theorem 2.3 for a > a through Theorem 4.2 in (Troiani al., 2025). The negative part follows from an analogous reverse inequality as detailed in the proof of Theorem 4.2 in (Troiani et al., 2025).
Similarly, the proof of Proposition 5.1 in (Troiani et al., 2025) based on the monotonocity of the GAMP dynamies, implies the existence and consistency with Definition C.3 of the subspace weak-recovery threshold defined by Equation 17.
D. Derivation of GAMP from Belief Propagation
This Appendix details the derivation of the precise form of the GAMP algorithm (Algorithm 1) associated with the Bayesian inference problem (8), alongside the expressions of the state-evolution equations (14) that asymptotically describe its dynamics. A rigorous proof of the state evolution is presented in Appendix C.
The departure point of this derivation are the Belief Propagation iterations (Pearl, 1988) associated with (8), then simplifying them to offer linear time algorithm that can be practically implemented with ease (Mezard & Montanari, 2009; Wainwright & Jordan, 2008; Lauditi et al., 2024). Contextually, we will derive the State Evolution iteration (14) and show that it describes the performance of GAMP at each iteration. Note that the derivation of GAMP and its state evolution for SMI models was previously reported in (Cui, 2025) for empirical risk minimization problems. We complement this analysis with the case of Bayesian inference considered the present work. While this derivation is heuristic in nature, we feel provides an insightful view of how propose algorithms for Bayesian estimation.
We start by defining factor graph, which is a bipartite graph with two families of nodes: D variable notes, each encoding a weight vector W. E R" and D + N factor nodes that connect them. D of the factor nodes are each linked to one variable node and encode the Gaussian prior, the remaining N are each connected to all the weights and represent the constraint that properly tuned weights should map the input data x. RDxM to the labels yd 'R' based on the examples the dataset through sequence multi-index model linked by g RPxM + RM We give pictorial representation in Fig. 4. We then introduce two families of messages ma Mi-sy which are the marginal probabilities of W. if we remove the edges i - a or μ - i. In the following it will be convenient to call W E Rx the stack of vectors W;
I. 1
8(1" nz
8
W VD
W E R" i = 1, ..., D
(2m)
/1 ...,. N
1, ..... D
mp-ti
Mi+tP
Figure 4. Factor graph representation of a sequence multi-inde model with 4 variable nodes (cireles) and 3 4 factor nodes (squares).
Picture inspired by (Aubin et al., 2018)
21

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
It is possible to define an iterative algorithm the messages, commonly called the Belief Propagation (BP) equations.
IW12
m 3
(Wi)
1
II m (Wi),
-i-t (2x)
(68)
/ Naws j+1
D Σ
W
m W.)
y* - 9
m (W,)
Zp-xi
is impossible iterate the BP equations for this problem in practice, but that's not an issue: estimate the weights w only need the first few moments of the messages. This observation will allow obtain efficient algorithm called Relaxed BP. It is convenient to introduce the sequence indices ZH E PxM using Fourier transform.
Wa
)
8 yd 9
VD
(69)
(2 ( PM
1
JPPXM
dz
dZ"
exp
Σ
(2t)
T
D
W
fa jm
D
ZH
2/
g(Z*))
m=1
This representation can be put in the equation (68) for m p (W,).
taking care of separating the contribution of W;
g(Z²)) exp {
W im
M
m f (V))
dz"d2"0
2H
Σ
Za
PL JRPX
VD
D
M
W jm
x II (Wmm/-ppWW exp
(2m)
i
VD
j
m=1
(70)
is time to define the mean and variance of the messages mi-#
Wi+p [ dw dW,m)-,,
JEP In the limit D + oc the term L can be easily expanded. Then, we can recognize that the lowest terms in the expansion are nothing but simple functions of the messages (71), and obtain a resummed expression at leading order in D.
dw W,, W RP
Wi-n (Wjn))
(71)
Wy
dW,mi-p
DM- (2
VD
exp
M
fm
(H
W
5
M
x
5m 3mm (2th)
RD
2
D
no
m., =1
M Σ
(xe Lejin)
exp
(2th
W. 5m
(2m)
T Cg
2u
2
D
m
m
and finally using the inverse Fourier transform, obtain
mr (W.)) ~
Zp PM
1
[ JapNM az 6(y" - 9((4))
JRPXM
dž
exp
-
M
Zu
T
W
vD
+
M
(2u) T 2H
m==1
m=1
I[
exp
zill jm
ZH
We
M
1
(ar
jmm
T (2
zu
2
D
j
m=1
1
1
I Jpexss dz 8(8" s(2²))
dž
II
M
il
12
2 "
W:
W
(aft
mm 21
(2
C 24
Zyuoi PM
RPXM
m=l
1
1
dZ 8(15" )((²))
II
M
b
7.m
W
whit
(Vini) 1( 1mm
W;
km
Zu-vi / (2xx)PM I dei(V
JRPX
22

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
where we introduced the mean and variance of H
D
Σ
d
jm
W
=
D
Σ
(ad) jm
(72)
371
jyt
As we did, before we will expand the term H im
in the D - oo limit keeping the leading order terms
Hith
(Zm w(m!) T (Vumi)) )
(0 (mm w(m)) (1
1
xdd sm
W
(Vith")
- (Zm -wh't
1 (xx )
im
'W, W W
(V
VD
D
1 (a im )2
W. (V ) (Zm
(V W;
D
As resum the expression above, express it using Jout and Owfout, as defined in (156). For convenience, define gu E RR, G E RPxP
gh [ooue(u' wift Vitit)]
G 1OwHlou (y" w Vi
wilt
(73)
mm
We arrive to the expression
mr
Wi)
=
1
1
M
x
im
W. Im +
M
(ar )2
W. T
Ith (sti") T
W, +
Σ
M
(ar
W, W.T.. W.
Zuss
2D
2D
m=1
m==
m=1
1
1
1+ W.
Br uami
+
W. T B By B (-ii)
T W.
W. T A A/r W;
(74)
2 n-
2
2
det (A
2 (W, - (Apm)) - Blass) Abrne (W, - 1
(W.
V
exp
)2
where we implicitly defined Ap-ii and By-ss:
M B μ-ク
M
(x c )
A/
Grnt
(75)
m==
The approximation in equation (74) can be plugged the first equation of (68)
1 e- 1
(de((A))a) (2T)
exp { 1 (W.
Abmos (W. (A B) }
m (W,) (W.))
II
Z (2x)
N
11 det(A
y W. +W, Σ B Σ (B: T (4)-))
N
1
-
1p-
Σ Ab-1
B
Zi
(2T)P
c
(76)
We have come full circle, as we can now explicitly write the moments of mb-. we first introduced in (71).
-1
N
N
W!+1_ = i-tp
1p + Σ A
Σ
Br
Br
g++1
T-pp
1p +
Σ
A
(77)
What we have obtained is thus a consistent set of equations that iterates on the moments of the BP messages. We give a
23

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
compact description of this algorithm here
D
1
Σ
x
W. jm
VD
jーH
jti
V.l.t. im
Σ (ajj)) itti
)2 j²
[gom (N*t with V] m gt
G
[ww (v" wittt, V.u.t]
x
Be p
the
(78)
(w am )2
G
N
N
W141
1p +
A
Σ By-ss
k-pp
-1
N
6+41 k-p
1p +
Σ
A
D.1. Relaxed BP implies the State Equation
first result we can get from Relaxed BP is the state equation (14). The key ingredient to obtain that noticing how at each step the quantities in (78) either concentrate or converge in distribution Gaussian variables in the limit of large D.
The first step is to specify how y" is generated. We can do by defining the index of the generative model Z" PxM and it is equivalent Z!"² with the weight W, removed
ZH
e
Wj
Z
Wj.
(79)
VD
fm
sm
VD
jm
With this definition the labels are yd ((Z)). Now, in the limit of large D, Z" and wn are going to be jointly Gaussian at each step with the following covariance structure
ZH W
1r
g1)
N
00
Q'
(80)
where Q' the same overlap of the weights estimated at time and the target
W((W;) D
W!(?)
Q'
=
lim D-+00
=
lim D
(81)
D
where the second equality is because of the Nishimori identity. The expectation is taken over the target weights W and the input data x. We can show this by looking at the first few moments of Z and a.t
E
1 0,
E Z (Z
1p.
1
E
0,
E Tiitit (wilit)
=
Wt J
Wi! W
R
Σ
=Q¹
D
D
(82)
1
E
1ith (in TJ
E
W, (W
= Σ
=Q'
D
D
24

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
Next, derive the leading order behavior of Jout and Soutu
Jout (y" with V!!.t) = gout (g( (7 with V!nt)
M
= gout (g( (2: wt V!llt) + VD
Σ m=1
Dmmm 02 [gout (g( wit VV+ut)] "
(83)
Owoout (y" with, Viri) = Iout (g( (Z ). v.rit).
The covariances Vid will concentrate
[
2
E
J-H
D
j=1
D
D 1
(84)
E
dw
dW
(W)
Σ
W
(W
D
RP
j=1
==1
1p
Q'
where in the first expectation on the second line we used Nishimori. We can then decompose the sum of the B as follows
N.M
N.M
B
Jout (g( wft Viui
Σ (1
)2
[0 (o("" w Viuil W
(85)
m
D
S.W;
Here, 5 RR is at each step Gaussian variable with covariance Q²
E(S 0,
(86)
where
Q' = aEzw [oout (g( (), 11p Q') gout (0(Z). 10-,1pp11) -
(87)
and the expectation is taken over Z, w distributed as in (80). It is possible to show with an extremely tedious computation that (*(St)T Q Finally, we have the State Equation
D
Q++l =E
D
w
(w
= (1+
(88)
D
which is exactly the same expression as (14).
D.2. The fixed points of the State Equation are extremisers of the the variational problem (10)
is not immediate by just looking at it, but the State Evolution (14) at its fixed point describes the extremisers of the functional (10). To show this we first rewrite (10) as
sup inf 5(Q, Dest QEst.
(89)
where
S(Q. Q) = -T)(((00)))))
log (or(1r (1)++)+++++)) tap log Zon (g(v1r- Q2 + (QE)), 10 Q)
(90)
where we already introduced for convenience Zont and yout. defined
Jout (v. w, V) = E(anN(w.V)) ((2 (Z--ad V .(2))
Zontly, w, V)- = F ((oWWWw.V)) [8(v ,(2)))
(91)
E ((anN(w.VV) [8(y - o[Z))]
25

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
In the following we will the identities
Oy Zout Zout
Dy Zout 2out (Owgout + gout @ 9out).
(92)
The extremisers are going to satisty the following system.
(4)(QQ.Q)= =0 (0S(Q.Q)
(93)
or equivalently
{ Q 1, log Zout 1MVIp-ZZ VQ
(94)
Q (1p Q) 1p (1p - QQ-
The hard part is computing the derivative with respect to Q of the first equation above. We start by adding dummy variable y. which allows us have the identity
dqEş log Zout (olv1p- - QZ" + +VQE).
(95)
0Q dyy Ey (v. VVV,,, - -Q) log Zout (v. 1r 1r-)
We compute the derivative
00 f dy Zout log Zout.
/ dy Ig (01
Zour Zout log Zout 0q Zous]
-ddy
1 1/2 Zom - Bv Zout [2200 (110 Oy Zout log Zou + 2
-f dy LEg (20-110 (11211/22 2outt (DwSro + Jout @ gout) log Zout
- 2out Zou (Ow + Jout @ gout))
/ dy Le (Ow Gou log Zout) Zout (Ow Dout + @ gout) log Zout + O2 (Zout Zour Eouu((a (Ow + Yout @ Jour))
/ dy Eg (Ba Bou log Bout) Zoue Dout + Jout & Dout) log Zout]
/ dy [g [ott
(96)
where in the 5-th line used Stein's lemma. Removing the dummy variable get
Q at ve (g(V1p Q2" (QEE). @) 82
(97)
Q - (1p - Q)
which is equivalent to the state equation without time indices.
26

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
D.3. From relaxed BP GAMP
The relaxed BP equations (78) describe perfectly usable algorithm with O(D²) time complexity. We can do better and obtain a linear time iterative scheme. Let define the following quantities
witht
z mm
W
VD
j+
D
VRa
(a )²
C²
D
j=1
(98)
W
1p +
Σ B
Bi
1p
A
which are very similar to their equivalent in relaxed BP. The idea is to have set of consistent equations on these new quantities. Since here the moments w, V. W, Ć' have one less index that in relaxed BP, the algorithm will be O(D) instead of O(D²).
For large D we have that at leading order Ö++1 se Ce+1 Similarly, V is changed minimally.
VK
Σ (anj))²(3
(99)
D
The other quantities are bit more delicate. First, notice how w; has non-negligible correction of order O(IvVD)
Wht1
1p
B
im your (v" wiftit V.lio) )
==1
(100)
W!!!
C++1
Jour (3r with V/ m
Similarly, removing one index to and V inside yout produces an extra correction in O(L/VD)
addt iim
Yout (y" wit Viiit (v"" whit Vuv)
1w (v² whit Vuit)] W
(101)
All this results in a correction in the equation for w, typically called Onsager reaction term
D
D
will
W
Σ
(ott (m
[gour whit-1 val-l))
(102)
D
1=1
The moment B p-- similarly receives a correction
M
M
B
B
youu (y whit Vuit]] m D
(x [0 (8" whit vuit]] mm W
(103)
convenient to define two new quantities A and by
N
N.M
4
=
A
Σ
(o
(w whit VHie)
D
μ=1
(104)
N
N.M
6
B
Jout (3g whit VHirj] + A,W?
==1
27

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
from which we finally obtain
W = (1p + A4) b/ G = (1p + At) -1
(105)
At)
This concludes the derivation. At very large D, it is possible to introduce Ó V A as the version without indices of O Vi.aa A7
Õ (1p
A²)
N.M N
A'
[awgout , whit Va.l)
(106)
mmm
V = C
This simplified algorithm is described in pseudo-code in Algorithm in the main text.
D.4. AMP with side information on the target
In this section we generalize GAMP to include some side information on the weights W used to generate the data. Suppose we are provided not only with the dataset = {x" "NN Iu=1. but also with a side information matrix S RPxD
S = 'W + vv-
(107)
where Ç ~ (0, 1) is distributed as standard Gaussian variable and À > 0 is a small constant. This additional information can be encoded in the prior on the weights w
1
IW 1 2(1 - 2)
Wi
(2 - 2)) exp
(108)
We can now go back and go through the derivation again with this new prior. Eq. (76) becomes
II (A)
-w: (1tr10 E 13:0 w, +w; (Sst S
B/-1)
4 E (m) (Ab_a)) 1B
1
m 1- (W;)
Z
(2x(1- - {
(109)
which gives us the new relaxed BP equations leading order in small À
(--1) E B + VAs. Σ
N
N Σ
W (1 (1-AMD++ AAam
(1 (1-1)
Ab-s
Ge+1 = -(1 (1-
A
(110)
and their equivalent in GAMP
WH+ ((1 p + A1) 41)-1(11 x)0, + viS,
Gtt = ((1- ((1-))0p + 19)-1
(111)
The state equation (88) is then slightly modified
(1+ XJQ' + A).
(112)
D.5. Sketch of the derivation of the State Evolution using the replica method
In this section, we provide a sketch of how to characterize the performance of a Bayes Optimal estimation using the replica method, a heuristic procedure inspired by the Statistical Physics of Disordered Systems (Gardner & Derrida, 1988; Mézard al., 1987; Mezard & Montanari, 2009). Despite being non-rigorous, the method very flexible and often provides a result
28

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
that be later proven rigorously (Barbier al., 2019; Aubin al., 2018). We invite the interested reader consult the many reviews available for more detailed explanation (Zdeborová & Krzakala, 2016; Lauditi al., 2024). A closely related characterization of the learning of SMI models using the replica method was first presented in (Cui, 2025) for empirical risk minimization problems.
Let us consider generic sequence multi-index model. We are interested in studying a model with P indices Zp. each describing sequences of length M. By this we mean that we have a dataset = {o*", y² "/> of labeled samples. Each sample x is sequence of M tokens of dimension D taken from standard Gaussian distribution:
xt E R DDM
a s,m
(0,
(113)
The indices are obtained by projecting the input data onto a set P weight vectors wp with standard Gaussian entries
2 E RD,
W ~ NO,,
2 2/= D
(114)
In the following, we will call W E RPxD the weight matrix obtained by stacking the vectors wy. The labels are generated by linking together the indices with thorough a function g such that
g RPRM RM
yp = g 2b).
(115)
Our results will hold in the high dimensional, proportional limit, that when N, D + with finite ratio = N/D. All the other quantities M, M', and P are assumed be much smaller than D. We are interested in characterizing the performance of Bayes Optimal (BO) estimator average over the dataset. The derivation starts with Bayes theorem:
P(WID) P(W) I1 N ==1 P (")
P(W)
N
wjx x
px"
))
II
g
(116)
P(D)
D
vD
where a normalization constant
( wix
Wpx"
Z Ew II
y" g
(117)
VD
On the other hand, what we are interested in the BO estimator on average over the randonness of the dataset, To this purpose,e we are resolved to compute the averaged cumulant generating function (or free entropy):
i lim
1 D
Φ
= lim D-+00
ED Eo log 2|= = lim n-20 On D=
log ED [2")
(118)
D
where the last equality is called the replica trick: we consider a system with n E N copies (or replicas) of the original model, extract the leading behaviour at large D. then analytically prolong the result real valued and take the limit - O. We are thus allowed to simply compute the moments of Z
N.n II _,,a=
W o
))
E [3"] =E .
Uu -9
VD
VD
(119)
To make the expression more symmetric we think of the labels y" as generated by zeroth replica
w x
wpx
8d =
(120)
VD
giving us
N.a II
/w 20
un upat
Ep Zd =Eg Wi
5
9
(121)
vD
VD
VD
29

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
Now notice that all the indices are jointly correlated Gaussian variables:
W
2
R-M
EZ =0,
E Zb. p.m q.1
Qu0 mm
(122)
DD
where we introduced the overlaps Qpq.
D Σ
Qu6
ur 1, w
(123)
Since Qug contains also the covariance of the target weights W we have Quy pg. Our goal will be to show that the overlaps Quh are the only quantity that characterizes the model, and that the partition function can be written supremum over them. With these new definitions, can rewrite the partition function
ED =
f Iprior ( (1Q
(124)
where the quantities Iprior and Chhannel are pieces that depend only on the prior or the channel respectively:
P.
I pprior Ew
II p.q=1,a.b=0
8 (DQm)
(125)
Iehannel
E(g.)((0.Q)
II 8 (g (i..... ..... ZT:) (Z). ..... z3))
(126)
In the second expression we introduced a new variable 2/ NO Q) as the distribution of each of the tokens of Zn, that 2p Zpm. We can massage these two quantities into compact expressions. We start with priio
P.n
prioor =E -w
II
8 (DQ
(127)
1,
Pa
P.
=Ew-
exp
Σ
E
bnp
2 (w))
(128)
p.q= ba
P
Pa
= EE W
exp
Dd Qah
jab 7m
w i
(129)
p.g=1,a.=00
p-g=1,a,b=0
D
P.a
P.n Σ
exp
D pq Pp
Ew-
exp
Snh u 20
(130)
P.g T,a
P.o 100
P.n
cxp
Daah
D log det (2Qpt) pg
(131)
p=
Since are only interested in the large D behaviour, the saddle point approximation
(Qph)
-2
(132)
which gives us the simple expression
{
D log det (Q
Iptior exp
(133)
This is good moment state some properties the overlap Q. First, we consider 1 identical replicas in the model. Qua needs to be symmetric under exchange of replicas ++ commonly called the Nishimori condition (Nishimori, 1980). We then make a "replica symmetric" ansatz for the overlap
Qab = dabōpg + (1 5an)
(134)
30

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
which simplifies the equations and allows us analytically continue the theory. It similarly reasonable to assume that Qug symmetric, positive definite matrices: Qpg Sp. After tedious computation, it is possible to show the following identities
det(Qpp)) det(1p Q)" det(1p nQ), (Q (Q-1)oa (1p (1, (]p +nQ)
(135)
(136)
(Q- Q1 -Q)
(137)
The quantity Aprior can now be simplified even further
pprior exp exp {2 In log (1 Q) + log del + moi} 2 exp
{
(nD [log del Q) + ThO].
(138)
where the approximation is valid for small enough n.
We can now focus on the channel part Jehannel of the partition function. We will show how the distribution over the variables Pp assumes a nice expression in the n + limit:
P.n
dz det Q eap
Σ
(Q
lsab b p
(139)
Pm
Pm
dæ
erp
2 (Q
1,00 aa "
(Q
- 10
b
(140)
det Q
2
2
p.q=1,a=0
p.4=l,aşbb
Pa
Pmn
de
det Q eap
2
pg
(Q (Q-1)
11
pq
011
2
(Q-
110-0
(141)
p.q=1.a=0
p.g=1,a,b=0
dzdg det Q
P.
P.u.
P
erp
Σ
11(Q-1
-1:0 1,00 (Q
01
Sp (Q
101 (Q-1)0)
Sp
(142)
2 p=1
p-y=1,a=0
p.p=1aa00
In the limit - 0 the identities (135) become
del (Qph) 1 pg)
(143)
(Q 10a + 1p
(144)
(Q-1yab= -111--11p1110-- -1)
(145)
As can see in (142), it is possible to write a factorized distribution each replica of 2p, while all of them coupled with an external random variable ç. We will drop the replica index to have:
dadg
1
Pm Σ
1yab bb = dge !1lE1 dz dzeepp exp { 1 pa-9
n+1
det Q exp
(Q
1- 2
- (1, T ()--Q)-22+3 ((1--Q)) Q))
2 b=
(146)
We now do change of variable
> Vlr - QC and massage the result to obtain
dzdg
(
Σ p=
P.
=1(0-1)p0
q dfe -ee--lell 1el1
p Q) exp { (z - VQE) QQ- vQ() }
dz
m+1
det Q eap
2
(147)
The computation is almost done: define Zout (y, w. V) as
Zout (U) w, V) = (en)((n))) 16(y - ((2)))
(148)
31

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
from which we get for small
[EWNrOOn ( (1)l)) (/ -(2))
chhanell
(149)
/Q.. (Q. -Q)
EEw.VOLn Sout
(150)
Zou (9(Z"), VQE, 1p - Q)
+
(151)
expp Ne log Zon (9(2%), VQL. Q })}}
(152)
Bringing everything together we have
1 nD log[=z"==(()) del (1 htgp://p---00tggoteeeeaaeeee//// log Zont (1(2²), 1
1p
(153)
Since we are interested in the large D limit, we can use the saddle point approximation to obtain the free entropy:
Φ = sup S(Q).
(154)
QES}
The supremum @ achieved for 0gS(²)= 0. After tedious computation it is possible to show that the supremum is achieved at the solution of the following fixed point equation, typically called the state equation:
Pout (0(2%) 1p -Q) Dout (0129) Vor.
Q F
a aEEz,
(155)
1111p -1p-QZ1 Qu+ V0s, 2 ~ N(0,1p). ~(0.pp) Lastly, Jout is:
with F[X] = - X0- X 20
[(z~N(w,V) [(2 WVV 8(y x(1))
Jout w. W)
(156)
IF ((onN(w.V) 18(y - -0()))
E. Specialization to deep attention models
The main theoretical results of subsection 2.1 are stated for the broad class of SMI functions (1). These results are formulated, in particular, in terms of the denoiser function yout (12), whose form tributary the specific architecture considered. In this Appendix,
derive the expression of Sout for the special case of rank P. 1 deep attention functions and length M = 2 data thereby specializing the characterizations of Theorem 2.3 and Lemma 2.2 to this specific architecture. This specialization allows on the one hand to plug the specialized Jour in Algorithm 1 and thus implement AMP for deep attention networks, and on the other and to reach the tight theoretical predictions for its performance plotted in Fig. 1.
E.1. Specializing Jout to specific activations for multi-layer attention in the M = 2, Pi = case
We start by rewriting Sout as an integral
(Rem wim Vi2 2m
f d mm
2/
21= (Zym WPm) V.J (Z))
[yout(y. W. V)
(157)
Zim
daz e 1,1 m=1
8 8 BL(
We can now massage the integrals the numerator and denominator. Let us take the denominator as an example. We can
32

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
isolate the integral over ZL
Fl wim MZ mmmm
/ dz
20p
By-
(158)
/ II dZ1
E m=
13 Wiml2 wymml
/ ddzze
1m
Min km waml
1im=l
(159)
8 o (ZI(B) (21, m. 2m-11) B
(To (21, (2 11.-1)21))
(160)
2 1.2
( - wem 12 -mimm
L-1
II dzj
e 1,1
21111
/ dZze
lim=1
13 (a!-11- 2, "1 K( 11 mm
2012
8 8 (Z Z1))
det B ZL-1)
(161)
1-1.2
(Rem wemm Mel wprm)
Σ
L-1
e
1,1' mm!
dZj
To
(162)
dde BC- (Z), Zr.-1)
Intuitively, we mapped an M dimensional integral with a Dirac delta as an argument into a (L- - 1) x M dimensional one, which is much more easily integrated numerically. Lo has be computed analytically. We give two examples with be realized for the numerator.
linear and softmax activation in the following two subsections. A similar manipulation Putting the two procedures together yields, for the case L = 2 discussed in the main text:
Corollary E.1. In the case M = 2, Pi P2 = 1 and softmax activation, the filter (12) becomes
3. unb-
mmmm
1, mmi
((BB-1)-124 - -m
dd
B (71)
Hoou (Y,w, Vm
(163)
X(PE
mmmm
24/1
dd
dz,
det B!(
1
where the sum over Zz is intended over the solutions of the two equations B² = + 12,1), with (xx 8) defined as
1 7(2,8) = Vlog ((x- 1) (y- 1))
log (x- 1)]
(164)
log (y-1 -21-1)
E. 1. COMPUTING I FOR LINEAR ACTIVATION AND M
We state again Ia in a compact form
I b- dz ö(y - 0(2 [122
(165)
where is a generic function of 2. It is convenient to write all the matrices explicitly
Iinear f dec d22 8
[11 U21
V/2
27
2122 =f
F F(z1, 22]
(166)
1/2
21
We can see that y/2 = 1/2, and also that 911, 1/2 >> O. On the other hand, the information on (yı 9/2, y22} is redundant, as one can use either just {911, 312) or {y22, }. We will choose the former, and write the integral
ff dz da2 80yn =2)8(112
Jinear =
22]
(167)
which we can easily solve obtaining
v12 F P[-Vin Vyn,
+ F
1/12 Vyni. /9111
/
linear
(168)
2 1y11
33

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
1.2. COMPUTING l SOFTMAX ACTIVATION AND M 2
Let us state again the problem
Ioftmma -
dzy dz2
8
1/11
V12
softmax
z
21 22
F F
22]
(169)
1/21
1/2
21122
This case is more involved. Let us write out explicitly the softmax
softmax
22
32
(170)
21
22
22
This tell us that y/1 + /12 = 1/2 + 1/11 = 1. A less immediate identity is that y/2 + 3/21 I 1, as we can see from these simple manipulations
912 + 1/2 S 1
6
$
e + =
e +
(171)
221 c
221 S 2 +
The last inequality is always true because the arithmetic and geometric mean inequality. We choose keep only the information in {y12, 3. We will thus get
1goftmax = ff
e
82
dz2g
1/2
e
+
8 U21
e.9
+ 20
Fl21, 22)
(172)
Let us try to solve the following system of equations:
21
V/2
(173)
1/1
which can be rewritten as This is simple system of two second-degree equations in ø1, The solutions are
log
1
2
71 22
812
(174)
log
2.
2/1 2/2
V21
log
+
V
lox
log
P12
v2l
(175)
log
22
Vlog(
log
012
This defines mapping (V12, U21) + (21,22) = #7(112. The Jacobian of the transformation is
det Jnl
(176)
2 2912822 - ()(22 - log [V/2 1)
V121221
Notice that all the quantities above real-valued because y12 + 921 1.
Finally, we obtain
[[ M21) F (r(112, V21)
Isofumax
(177)
2 9128/21 (V12 () 1) log
212321 [v12
10
34

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
0.6
Overlap first layer
0.4
0.2
0.0
0.7
0.8
0.9
1.0
Overlap second layer
Figure Left: Overlap of the first layer Qu in two layers rank one attention with skip connection e and sample complexity 1 as function of the overlap of the second layer Qz2, kept fixed during the iteration. We can see that unless Q22 is almost one, Q not leamned.
E.2. Weak recovery thresholds for two layers of rank one attention
In this section we describe in detail how to obtain the weak recovery thresholds for the model in Section 3.1. We start by computing == fow (Y, 0, p). We can write it compactly as
2.2 (...m=
Zem Zetr
f /dze
(Zim Zum - (w:)0w - g(Z))
Gu'mn
(178)
Zim Zym m 8ly g(Z)
d
Since g is an even function, Giimm' is zero if + This implies that the operator F in (15) diagonal, hence on is simply
1
sup XES
|F max
Σ
Ey
(179)
α1
||||||
m.m!
Implementing the formula above reveals that for these models the maximum is achieved for = Computing 02 is harder, but make the simplifying assumption that for > 02 the second layer is perfectly recovered, so
Q
(180)
On the other hand, the second layer is perfectly recovered, one has to simply compute a linear perturbation on Q along Q11. giving us
Σ VI2
GG )
(181)
We would like stress that having learned the second layer is crucial for the learning of the first, and it's central to the grand staircase phenomenon. We can see this also in the model of attention. Suppose we run AMP and choose not to keep the weights in the second layer at a fixed overlap Q22" with the target: with all other parameters fixed, the value of the overlap of the second layer at convergence depends strongly on Q22", as shown in Figure for a = 1. In this particular case AMP would learn the second layer almost perfeetly and achieve a good overlap in the first one, as we can see in Figure 1 in the main. On the other hand, we prevent the algorithm to learn the second layer, then also the first one is not learned.
E.3. Language and hardware specific details
Here we discuss the coding details of the implementation of GAMP and the state evolution that we use for the figures. The interested reader can look at the code repository at https: / / www.tthoooooooooet /SOSSPCC-OO- SPOC-group/ Sequence IndexModels.
35

---

Fundamental limits of learning in sequence multi-index models and deep attention networks
common ingredient is the denoising function Pout (Y, ww, V) in (163), which is just two-dimensional integral over smooth function. We perform the integral by quadrature using the function db quad in Scipy the set [-3,3] [-3, 3]. Additionally, regularize the integrand by adding 10 3 to the diagonal of V and 10-6 to the square root argument in (163). For GAMP we need also to compute Sout. which we do by centered finite differences computed on two points at distance 10-5, For state evolution we have to perform and expectation over Y, E as in (14). We choose to do so using a Monte Carlo method, that is by taking 1440 samples Y and & and computing Yout for each them. At each iteration we symmetrise the overlap Q.
For each run of GAMP or the state evolution we use 2 Intel Xeon Platinum 8360Y processors and approximately 290 GB of RAM. In Fig. (left) smoothen the prediction error and the overlap theory curves by interpolating with splines above the critical threshold.
F. Details on the experiments on real datasets
In this Appendix, we provide further details on the numerical experiment presented Fig. 3, namely the training of simple transformer model on the TREC (Hovy et al., 2001; Li & Roth, 2002) classification task.
Data The TREC dataset (Hovy al., 2001; Li Roth, 2002) contains 5500 labeled questions of average length 10. divided in 6 classes, classifying the type of the question, e.g. whether it bears concept, human, numerical value. The vocabulary size is 8700. The data is pre-processed using the uncased base BERT model (Kenton & Toutanova, 2019) and padded into sequences of length M = 32.
Architecture We consider small transformer model with two consecutive self-attention layers, with tied key and query matrices, and a fully-connected readout layer. More precisely, the architecture comprises
Embedding layer: embeds the input sequence in dimension D 128
Attention layers: two consecutive self-attention layers
T F ii 24-1
I 2;-1 softmax
with tied weights wy E PPP,, and hidden dimension P1 = Pg = 64.
Pooling: the sequence after the attention layers 22 E PxxM is averaged over its second dimension.
Fully-connected readout: linear projection from dimension D (number of classes).
The architecture is then trained over the cross-entropy loss, using the Pytorch (Paszke et al., 2017) implementation of the Adamii (Loshchilov & Hutter, 2019) optimizer, using learning rate v = 5 x 10-5 and batch size 16. The weights wy at initialization are sampled from a Gaussian distribution with zero mean and standard deviation 10-4
36


Reviews:
[{"review_id": "YuGjjpFOul", "reviewer": "Reviewer_xpYy", "summary": "This paper studied the problem of learning sequential multi-index models. The authors showed a deep connection between the SMI model and deep attention networks, in that the deep attention network can be formulated in the form of a SMI function. The authors studied the limit thresholds of the sample complexity for weak recovery of the hidden directions $W^*$. Numerical experiments are conducted, providing intuitions on the dynamics of the GAMP and SGD algorithms when learning deep attention networks. In particular, they showcased the sample complexity thresholds of which 2 layers of the weights are learned in a hierarchical manner.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive comments.\n\n**Claims And Evidence:** \n\nFig. 3 numerically supports the conjecture, while Sec. 3.2 discusses its infinite-depth limit where inner layers ($\\ell \\leq L - 1$) become interchangeable.\n\n**Relation To ... Literature:**\n\nThe reviewer is right that our work builds on prior techniques. However, as noted, the new structure introduces substantial technical challenges.\n\n**Weakness**\n\n1.  On optimality:\n\n(1) GAMP is a polynomial-time algorithm that converges with few iterations, as predicted by state evolution. All examples in our plots used at most 50 iterations—we will state this explicitly.\n\n(2) First-order optimality: GAMP minimizes prediction error among all first-order algorithms (e.g., Celentano et al. 2020), making our thresholds optimal in that class.\n\n(3) Bayes-optimality: Our Th. 2.1 shows that when state evolution has a unique fixed point, it corresponds to the MMSE estimator.  In such cases, GAMP achieves statistical optimality and our thresholds match the fundamental limits.  This holds in all our examples.\n\n2. Thank you—this helps clarify terminology. The weak recovery threshold for initial recovery is the sample complexity beyond which $\\hat{W}$ escapes the zero fixed point, i.e., gains non-zero row-space. In contrast, the subspace recovery threshold for a subspace $U$ of the target row-space is the sample complexity above which $U$ is included in the row-space of $\\hat{W}$. \nDef. (17), equivalent to Def. 2 in [1], formalizes this:\n$v^\\top Q_t v > 0$ ensures $v \\in U$ is in $\\hat{W}W^\\star$'s row space, where $Q_t = \\hat{W}W^{\\star \\top}/d = \\hat{W} \\hat{W}^\\top/d$.\n\nFor SMI targets, GAMP dynamics are more complex than for single-index: Learning is coupled across subspaces—progress on one can aid others. This manifests as GAMP passing near a series of saddle points. The thresholds in (17) reflect the minimal sample complexities required to escape these saddles, matching those in [1] (Def. 5), found via the stability of (18).\n\nThe threshold in (16) is the minimal complexity to recover any 1D subspace, thus a special case of the more general subspace recovery threshold\n\n3. The weak recovery threshold for initial recovery marks the point where there exists a subspace $U$ that satisfies the condition (17). For example, in sec. 3.1, span$(w^\\star_2)$ is recovered first, so its subspace recovery threshold (17) coincides with the weak recovery threshold for initial recovery (16). In contrast, the threshold for span$(w^\\star_1)$ is higher, and depends on how much of $w^\\star_2$ has been learned. It is not equal to (16): One can define a more precise threshold $\\alpha_{w^\\star_1}(q_2)$—the sample complexity for recovering span$(w^\\star_1)$ given overlap $q_2$ with $w^\\star_2$, assuming minimal $q_2$ in (17). Thus, (17) gives a fine-grained notion of threshold for a specific subspace, while (16) captures the threshold for recovering any subspace.\n\n4. In (18), $Q_t$ denotes the overlap between $\\hat{W}^t$ and $W^\\star$ when the estimator uses side-information of the form $\\lambda W^\\star + \\sqrt{1 - \\lambda} \\times \\text{noise}$ (see l. 254, right column). Eq. (14) corresponds to the case $\\lambda = 0$, with no side-information beyond the dataset ${x^\\mu, y(x^\\mu)}$. Thus, (14) and (18) describe distinct AMP algorithms, depending on the presence of side-information. Introducing side-information is technically required to define the thresholds (17) via the limit $\\lambda \\to 0$. For clarity, we will denote the overlap as $Q^t_\\lambda$ in Theorem 2.3.\n\n5. The thresholds in Fig. 1 indeed instantiate the theory of Sec. 2. \nWe will revise $\\alpha_1$ in (17) to $\\alpha_U$.\nIn subsection 3.1, $\\alpha_1=\\alpha_{{\\rm span} (w^\\star_2)}$, i.e. the sample complexity necessary to retrieve the subspace spanned by the second layer weights, is the weak recovery threshold from (16). The sample complexity $\\alpha_2$ corresponds to the sample complexity required to further retrieve the full row-space $ {\\rm span} (w^\\star_2, w^\\star_1)$.\nThis corresponds to (17) for $U= {\\rm span} (w^\\star_2, w^\\star_1)$.\n\n6. Writing a MIM as an SMIM naturally leads to studying targets where each input is a set of vectors—a common setup in sequence models. SMIMs can be seen as MIMs with a block-structured weight matrix (see Appendix B). While our proofs build on MIM results, the specific block structure we consider is new and well-suited to attention models. SMIMs are more natural for gradient-based training, where gradients are computed for all weights. The equivalent MIMs involve tied/zeroed weights, incompatible with standard gradient descent. This is similar to how CNNs are special cases of FCNs, but still distinct in practice due to gradient computation.\n\n**Questions For Authors**\n\nWe expect GAMP to saturate the SQ bounds, as shown for MIMs in [1] (Lemmas 3.3 and 4.3). \nWeak recovery thresholds relate to generative exponents: 0 if threshold is zero, 1 if finite, $\\geq 2$ if infinite.", "claims_and_evidence": "Almost. There is a conjecture about the learning dynamics of 3 layer attention net, supported by some numerical experiments of a special case.", "methods_and_evaluation_criteria": "Theory paper, N/A.", "theoretical_claims": "I skimmed through the proofs but unfortunately I did not check it line by line.", "experimental_designs_or_analyses": "yes. no significant issues.", "relation_to_broader_scientific_literature": "The key contributions align with a thread of work on the limits of weak recovery of multi-index models using GAMP. The techniques and results are similar to prior works like [1], though I think it needs some more careful analysis with respect to the new structure of the model.\n\n[1] Fundamental computational limits of weak learnability in high-dimensional multi-index models", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Weakness:\n1. I am not sure what conclusions we should draw from this set of results. The authors proposed a bayes-opitmal GAMP method for learning attention networks and SMIs, where it needs calculations of the posterior distribution $E[Z|Y]$, where $Y$ are the labels. The choice of this GAMP method makes sense somehow since this is one of the optimal first-order methods. However, it is not clear it is optimal in what sense. Is this optimal in the sense of sample complexity, or it is optimal in terms of runtime as well? If the algorithm achieves optimal sample complexity but have exponential runtime, then this sample complexity threshold seems not very informative. Or, are the authors trying to provide a kind of lower bound, since the GAMP is 'optimal' first-order algorithm?\n2. I am afraid this paper is not very clear written. It is particularly confusing to me the phrase in Theorem 2.3, what are the differences between 'weak recovery for threshold for initial recovery' and the 'subspace recovery threshold'? What are the differences between these notions and the notion of weak subspace recovery in prior works, like definition 2 in [1]?\n3. Continuing with the second point above, I am very confused about why there should be 2 different thresholds. Is this coming from the theory that there are, say, some local minima or some unstable fixed point that the power iteration cannot overcome? \n4. What's more, there seems to be 2 Q^t's, coming from (14) and (18), which are very confusing to me. Are they measuring the same 'non-vanishing of overlaping' between $W^t$ and $W^*$? Perhaps I am wrong but my understanding about this thread of research it to show that after some threshold $\\alpha$, the fixed point of the AMP process/power iteration Q would be p.d. and is a fixed point. Therefore, when the authors talk about fixed points and weak recovery thresholds (in particular in (16)), which process are the authors referring to? \n5. The authors observed that the second threshold alpha_1 is the sample complexity threshold for learning the first layer weights after learning the second layer. Are there any theoretical justifications?\n6. Finally, I am confused about the difference between SMI and MIM. As the authors mentioned, the SMIs can be essentially viewed as a MIN (in (49)), in particular the new MIM has hidden dimension $PM = O(1)$, then why we need a new set results on SMIs?", "other_comments_or_suggestions": "See weakness.", "questions_for_authors": "When $P = M = 1$, (i.e., considering SIMs), how are the results relates to the sample complexity of efficient SQ algorithms, i.e.,  the generative exponent?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RvaR5m42c6", "reviewer": "Reviewer_KcvH", "summary": "This paper studies the fundamental limits of learning in deep attention networks by establishing a connection to sequence multi-index models. The key contributions include the mapping from deep attention networks to SMI models and theoretical characterization of statistical and computational limits of the models using AMPs.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive comments. The focus of this work is indeed theoretical, and lies in making several theoretical findings on learning in attention models, extending previous related works. We discussed how the uncovered phenomena also hold with Gradient Descent (Figure 2 Bottom) and on some simple real world case (Figure 3 Right), but totally agree that it would be interesting to have a more thorough mechanistic study of layer-wise learning in more involved transformer models. On the other hand, we believe such an analysis is outside the scope of the paper at hand, and hope our work will inspire a more applied set of researchers to conduct such an investigation.", "claims_and_evidence": "The claim that deep attention networks map to SMI models is well verified through formal derivations (Section 2, Appendix A), and the claim AMP achieves Bayes-optimal performance is also demonstrated through both theoretical analysis (Lemma 2.2, Theorem 2.3) and synthetic experiments (Fig. 1–2).", "methods_and_evaluation_criteria": "The theoretical analysis leverages AMP and state evolution, and the synthetic experiments are well-controlled (Gaussian inputs, large dimensions) and verify the performance of AMP. However, the real-world evaluation is limited, since the experiment is based on single dataset and simplified architecture.", "theoretical_claims": "Theorems 2.1 and 2.3 rely on mapping to multi-index models and prior AMP analyses. Proofs in appendices seem rigorous.", "experimental_designs_or_analyses": "The synthetic experiments clearly show phase transitions and layer-wise learning. However, it would be better if there are some real-world evaluations.", "relation_to_broader_scientific_literature": "This work connects to multi-index models and extends single-layer attention analyses, which is highly relavent and important to the area of feature learning theory.", "essential_references_not_discussed": "As far as I know, this work has a comprehensive discussion on literatures. Most of the related literatures that I know are well addressed.", "other_strengths_and_weaknesses": "This work has novel theoretical connection between attention networks and SMIs with rigorous phase transition analysis. However, it needs more real-world evaluation to address its practical implications.", "other_comments_or_suggestions": "Minor typos: Page 3, \"idependently\" → \"independently\"", "questions_for_authors": "I have no other questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aBeZYZv4Xa", "reviewer": "Reviewer_R4QM", "summary": "This paper extends the theoretical framework of multi-index models to sequence models and derives a sharp asymptotic characterization of the optimal performance of the generalized approximate message-passing (GAMP) algorithm. This paper also characterizes sharp thresholds on the minimal sample complexity required for better-than-random prediction performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive comments. \nNumerically, we observe that in the setting of Fig. 2, GAMP and  Gradient Descent  display similar layer-wise learning dynamics. To put these observation on completely rigorous theoretical grounds, one would need to adapt the dynamical-mean-field-theory analysis of (Gerbelot, Troiani, Mignacco, Krzakala, Zdeborová 2024) to SMIMs. While we believe such an extension to be possible, it would warrant significant additional investigation, and we leave this for future work.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No.", "experimental_designs_or_analyses": "I check the soundness of the experimental designs, which makes sense for me.", "relation_to_broader_scientific_literature": "The optimal performance of gradient descent algorithms has been well studied for multi-index models, the key contributions of the paper is   that it extends the theoretical analysis of multi-index models to sequence multi-index (SMI) functions.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n1. This paper is well organized and easy to follow.\n2. The experiments are clear.", "other_comments_or_suggestions": "No.", "questions_for_authors": "Questions:\n1. This paper considers the best-known polynomial-time algorithm (AMP) algorithm. I am wondering if it is possible to get similar theoretical results, if we change the AMP algorithm to gradient descent algorithms.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CFW99vtCzB", "reviewer": "Reviewer_jD8f", "summary": "This paper uses GAMP to quantify the Bayes-optimal performance of deep attention networks. They show that such a deep attention network can be mapped into a sequence multi-index model, which is a variant of the basic multi-index model but applied to a sequence of data of fixed length M.\n\n-- update after rebuttal --\n\nThis paper presents an interesting application of AMP to the attention mechanism, so I maintain my evaluation. The authors have addressed most of my questions, although I look forward to future research on the order of learning (e.g. in Bayes-optimal learning or SGD) that this paper identified as an open problem.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive comments. Their remarks will allow us to improve the clarity of our manuscript greatly.\n\n**Claims And Evidence**\n   \nClarifying the limit: By \"deep,\" we mean constant ($O(1)$) depth. Our results hold in the large $D$, proportionally large $N$ limit.\n\nStationary points vs. minimizers: We agree and will refine the mapping between SE fixed points and the free energy around line 229 (see also response to Reviewer xpYy).\n\n**Methods And Evaluation Criteria**\n\nOrder of learning: The reversed learning order is not specific to the algorithm. On our synthetic task, all algorithms learned the last layer first. In TREC, the reversal likely stems from task, model, and algorithm jointly. Our theory predicts that Bayes-optimal learning proceeds layer-wise from the last layer. While this matches some empirical observations, the learning order depends intricately on task and architecture and deserves further study. We will note this.\n\n \n**Experimental Designs Or Analyses**\n\nThe diagonal structure of $Q$ is task-specific. Iterating SE from generic $Q$ shows it remains diagonal. Intuitively, the target weights are random and approximately orthogonal in high dimensions, implying the same for the GAMP estimates.\n\n*Fig. 1 (right): define weak recovery threshold.*\nDefined in Theorem 2.3: the smallest sample complexity at which GAMP’s estimates have non-zero correlation with the target. We will add a pointer in the caption.\n\n*Fig. 2: nonzero start* The first point shows overlap after one iteration—not initialization. We will clarify.    \n\n**Supplementary Material**\n\nThank you for the note, we will improve that appendix to be better self-contained.\n\n**Other Comments Or Suggestions**\n\n*line 21*  We indeed mean \"proportionally large\", we will reword. \n\n*Suggestions on eqns (2)-(3)*: Thank you, we will follow your suggestion for clarification.     As we discuss in l-157, we consider frozen value matrices fixed to identity. We will stress this point.\n    \n*In (3) $1_M$ denotes*: In (3) $1_M$ refers to the identity matrix.\n\n*barer input and other typos*: We thank the reviewer for their careful reading. These typos will be corrected.\n\n**Questions For Authors**\n1. c is the skip connection strength, introduced in (3).\n \n2. They are indeed PxP matrices. We will stress this point.\n \n3. They are Kroneker deltas: $\\delta_{ab} = 1$ if $a=b$ and zero otherwise.\n \n4. We mean that the prior on the weights $W^*$ is nonseparable. The denoiser $g_{out}$ is defined in (12) and its functional form can be arbitrarily complicated, both column-wise and row-wise non-separable, as  in the case of deep attention models, as detailed in Appendix E.\n \n5. We mean that it would be natural to expect we need less order parameters: in general one would need $PL \\times PL$, but if the first $L$ layers are identical we just need $4P \\times 4P$ by imposing that the overlap $Q$ has a certain structure. More in detail, we would be imposing that the overlap of a layer $l$ and the $L-$th layer is identical for any $l\\le L-1$, that the overlap between any pair of layers $l_1,l_2\\le L-1$ is identical, and finally that the self-overlap of the weights of any layer $l\\le L-1$ is identical. \n \n6. We refer to the main answer above to this same question.\n \n7.  The SE equations do involve the dimension $M$ through the denoiser $g_{\\rm out}: \\mathbb{R}^{M\\times M} \\to \\mathbb{R}^{P\\times M}$, involved in the SE equation (14), and defined in (12). A technical challenge was, in fact, precisely to ascertain the denoiser $g_{\\rm out}$, which we unravel in Appendix E. Another challenge was to actually show the equivalence between SMIMs and MIMs (see appendix A). \n\n \n8. $K$ denotes the output dimension—for sentiment classification, this corresponds to the predicted class probabilities. In our setting, the output is an $M \\times M$ matrix, so $K = M^2$, interpretable as a target attention matrix capturing token interactions. Sentiment analysis, with $K$ being the number of sentiments, is a good example that we will include.\n  \n9. The setup in GB23 and BMN20 (Berthier, Montanari, Nguyen, 2020) for non-symmetric matrices allows for non-separable non-linearities along both the axis i.e across different $N$ samples and across different coordinates of the parameters. See for instance equation $(5)$ in  GB23 and equations $(1),(2)$ in BMN20. In Equations (61)-(62), $B^{t}$ contains the flattened parameters $\\hat{W}^t$ and therefore has dimensions $DM \\times PM$. The number of samples $N$ appears in the dimensions of $\\Omega^t \\in \\mathbb{R}^{N \\times PM}$ which contains the corresponding flattened inputs.\n \n10. Adding fully trainable non-linear layers is a challenging pursuit which would warrant completely new technical ideas, and we defer the study thereof to future work. Non-learnable non-linear layers could be added without a big additional technical challenge.", "claims_and_evidence": "The analysis of this paper follows naturally from the existing analysis for multi-index models.\n\nI think it’s important to clarify early on, maybe around (1), that the scaling regime considered in this paper. By “deep”, we refer to a constant number of layers L, rather than infinite. (there is only one short paragraph on $L\\to\\infty$) The growing quantities are the number of data samples N and the length of the token D.\n\nline 229: my understanding is that AMP (i.e. the fixed point of SE) attains the largest stationary points of the replica potential. The global minimizer of the potential corresponds to the Bayes-optimal performance. Therefore when the potential function has a unique stationary point, then AMP attains the Bayes-optimal performance. Is this correct? If so, I think it would be better to be precise with terminology e.g. stationary points vs global minimiser vs local extremiser. \n\nThe paper is heavy in notation, so I think it's important to carefully define every term and its dimension, especially when you present the transformation from the deep attention network to SMI. See Other Comments Or Suggestions for details.", "methods_and_evaluation_criteria": "AMP is a principled way of characterising the Bayes-optimal performance. \n\nWhat I find confusing though is the order of learning. Do you have any explanation as to why, compared to GAMP or SGD, the order of learning is reversed in your TREC example? ie earlier layers in the network are learned first? I think this is an important behaviour that the paper claimed that AMP should capture correctly, but why is this reversed in practice? Perhaps the method you used for solving the TREC task is far from Bayes-optimal?", "theoretical_claims": "I didn't check the proofs in detail but believe they followed naturally from the existing analysis for multi-index models.", "experimental_designs_or_analyses": "Could you explain why the off-diagonal entries of Q are zeros? I’d imagine the estimate of a given token should be correlated across different layers of the network?\n\nFigure 1 Right: how did you define the weak recovery threshold? please add legend.\n\nFigure 2: why does the similarity curve start at a nonzero value for GAMP?", "relation_to_broader_scientific_literature": "Interesting paper that might help our understanding of the learning dynamics in transformers.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Generally well-written, but notation especially those in the supplementary material can be improved.\n\nI think the title is slightly too long.", "other_comments_or_suggestions": "line 20: Bayes-optimal\n\nline 21: what do you mean by “commensurably large”? It might be more accurate to say “proportional”, large dimension D—> token length? Be more precise. “D-dimensional” is cumbersome, could use “length-D” instead\n\neqns (2)-(3) I’d suggest reverse the order and define x0 first then the recursive update rule for xl, and finally yDA(x). Also, maybe clarify that w_l is the key and query matrix for layer l instead of calling w_l a generic weight matrix. Is the value matrix absorbed into the activation function here? Maybe clarify that.\n\nIn (3) $\\mathbb{1}_M$ denotes the MxM all-one matrix, instead of the identity right? clarify notation here\n\nline 55: check linespace\n\nline 70: “N samples of Gaussian iid input sequences”, clarify that each sample refers to M sequences of len-D tokens.\n\nline 159: barer input not output right?\n\nline 364: Figures 2 top or bottom instead of left or right\n\ntypos in eq (40)\n\nlots of typos around (61)-(62): I thought Bt should be NxPM? instead of the stated DMxPM?", "questions_for_authors": "Clarification questions (addressing these will likely improve the paper and change my evaluation of the paper):\n\n1. line 92: B_c^l what does the subscript c stand for?\n\n2. (13)-(14) can you clarify the dimension of the SE iterates? I think they are PxP matrices?\n\n3. (5): what are the deltas here? I didnt get this formula\n\n4. line 218: “with a non-separable prior on the parameters” can you elaborate what exactly is non-separable? Is your denoiser g_out rowwise non-separable?\n\n5. line 357: “Since these first layers are interchangeable we expect that also the overlaps should share the same invariance, making it possible to simplify the state evolution equations in this limit, in a similar spirit to (Aubin et al., 2018).” Do you mean you can simplify the SE equations treating the first L-1 layers essentially as the same so captured by one set of parameters, while the L-th layer itself can be described by a second set of parameters? So overall your SE will only involve (P1+PL)x(P1+PL) matrices instead of PxP matrices, which will now have infinite dim since L\\to \\infty.\n\n6. Most important question: Do you have any explanation as to why, compared to GAMP or SGD, the order of learning is reversed in your TREC example? \n\n7. The analysis of this paper seems to follow naturally from the existing analysis for multi-index models, in the sense that the SE is dim-free so extending from M=1 to arbitrary M doesnt really change the SE equations. Can you elaborate what are the main technical challenges that you've resolved?\n\n8. $g: \\mathbb{R}^{P\\times M}\\to \\mathbb{R}^K$. What does the K-dimensional output represent? It might worth adding one sentence giving an example. I'm guessing this could represent say the sentiment of the input?\n\n9. I'm confused by your application of the non-separability result from Gerbelot & Berthier 2023. GB23 allows non-separability across the N data samples, whereas it seems that in (61)-(62) the non-separability in your work is within each sample as you treat the N samples as iid. By a \"sample\" I mean M length-D tokens. Can you clarify? Can you also clarify whether Bt should be NxPM in (61)-(62)?\n\n10. Is it possible to incorporate the feedforward layers alongside the attention mechanisms?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper considers learning of attention networks with O(1) depth, with tied and low-rank weights. Authors map such models to sequence multi-index models, and conduct their analysis on under this relationship. Authors study weak recovery and the sample complexity thresholds. They further provide numerical results, on GAMP and SGD algorithms when learning the multi-index model.

This paper was reviewed by four expert reviewers and received the following Scores: 4x Weak Accept. I think paper is studying an interesting topic but authors are not able to convince the reviewers sufficiently well. The following concerns were brought up by the reviewers:

- The paper and its motivation is not clear.

- Results not sufficiently explained in paper.

No reviewers championed the paper and they are not particularly excited about the paper.
I believe majority of the concerns can be addressed. As such, based on the reviewers' suggestion, as well as my own assessment of the paper, I recommend including this paper to the ICML 2025 program.
Decision: Accept (poster)