Title: Bayes optimal learning of attention-indexed models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Bayes optimal analysis

Transformers

Self-attention mechanism

Random matrix theory

Matrix models

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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors have introduced the Attention-Indexed Model (AIM) to analyze learning in deep attention layers. Theoretically, the authors analyze this using statistical mechanics and random matrix theory to get the closed-form predictions for Bayes-optimal generalization error. Also, the authors analyze the case where embedding dimension $d$ goes to infinity and the sample complexity is $\\alpha = n/d^2$.\n\nSTRENGTHS AND WEAKNESSES:\nThe strengths of this paper can be summarized as follows:\n\nFirst, this paper gives a very thorough theoretical analysis. This is the first tractable analysis of attention mechanisms with extensive-width matrices, whereas previous works only focus on narrow attention layers. As far as I can see, all the theoretical analyses are rigorous and correct. \n\nSecond, the authors have also conducted experiments and provided an optimal AMP algorithm. In particular, the result has shown that the practical gradient descent achieves similar performance.\n\nThird, the implementation details and code are included. It makes the experimental results being reproducible. \n\nThe weaknesses of this paper are summarized as follows:\n\nFirst, as the authors have mentioned, they only analyze the single layer case. It would be better if the authors can consider including a more detailed discussion on the generalization to the multi-layer case. For example, what are the potential difficulties one may encounter to generalizing it to multi-layer case?\n\nSecond, the author claims that they analyze the case where $d$ is very large. However, in their experiments, $d = 100$ or $d = 200$ seems pretty small or moderate. Is it possible to try larger $d$’s to make it more aligned with the theoretical results?\n\nQUESTIONS:\nPlease see weaknesses.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes the Attention-Indexed Model, designed to study learning in attention-like architectures, where tokens interact with the model parameters as in the attention mechanism of modern Transformers. For the proposed model, the authors derive the generalisation error in the high-dimensional limit of the number of samples $n$, embedding dimension $d$, and width $r$ (with suitable ratios between these quantities) for the Bayes-optimal estimator. They also propose an AMP algorithm that recovers the overlaps appearing in the asymptotic error, thus matching the Bayes-optimal performance. While the theory covers a general $L$-layer model, the authors also discuss the $L = 1$ case as analytically more tractable.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths \n\n- Novel modeling perspective. By pulling out the bilinear attention indices and studying them directly, the paper isolates what is mathematically distinctive about attention and carries out standard techniques to study the asymptotic behavior. Thus, I see the main contribution as the novelty and relevance of the results more than in their technique.\n\n- I particularly appreciated the discussion on the intractability of the integrals involved which prevent to derive a closed form for the Result 3.1. \n\n- Concrete formulas in the single-layer case. For $L=1$, they can write the AMP’s prior-channel denoiser in closed form via rotationally-invariant eigenvalue shrinkage (RIE) (22) and an explicit state evolution equation (24), leading to analytic phase-transition locations that perfectly match AMP simulations. \n\n## Weaknesses:\n\nI am overall in favor of clear acceptance, as I find the paper technically sound, the results are relevant and novel. These minor weaknesses are to be seen more as suggestions or discussion points\n\n- Background on AMP and Bayes optimality. I think the authors could have given more background on approximate message passing and more explicitly state its connection to Bayes optimal estimation. The paper relies heavily on this equivalence but only mentions it briefly (Appendix C). A short primer in the main text would make the results accessible to a broader audience.\n\n- The paper writes the Bayes free-entropy (or “potential”) of the Attention-Indexed Model as an extremisation problem. Unfortunately, the extrema for $q$ and $\\hat{q}$ in Result 3.1 are in implicit for, and it’s thus a bit hard to interpret. However, the authors extenively discuss why the intrgral in the partition function $I_{in}$ is intractable.  However, the authors provide more “meat” with the discussion of the $L=1$ case, which is tractable because now the $\\Delta_\\ell$ becomes diagonal. I think this compensate for this lack of interpretability.\n\nQUESTIONS:\nThe authors propose a mean-centered attention mechanism with $1/\\sqrt{d}$ scaling. It would be instructive to discuss the alternative $1/d$ scaling (probably, no mean centering needed in this case). Under this \"LLN\" scaling, are the asymptotics different, or would the current proofs carry over?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies Bayes-optimal generalisation error in a theoretical model of a Transformer. The advantage of their proposal is that they can loosen the restriction on the dimension of queries and keys being small. The provide synthetic experiments to validate their theoretical work.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n----------------\nThe paper tackles an interesting problem of studying generalisation in Transformers. The assumptions the authors make are reasonable and seem to push forward the surrounding literature (although I am not an expert in this area).\n\nThe experimental results seem quite convincing and I appreciate the numerical validation of the theory.\n\nWeaknesses\n----------------\nThe paper targets a very specific community and is hard to follow. I have some light background in RMT and found this quite challenging to read. I believe it would be useful to contextualize with some concrete examples the notation (e.g. Equations 4 and 5). For instance, how would a ViT or LLM look like using the notation. I understand it would not match 1-to-1 with actual Transformer architectures, but at least it would help clarify the notation.\n\nI believe the claim \"AIM offers a solvable playground for understanding learning in modern attention architectures\" in the abstract is perhaps over-claiming as I do not believe there is high semblance in the update equations to modern attention architectures (a lot of very non-linear and complicated components are missing).  \n\nOverall, having some light background in RMT, I found the paper very hard to follow and I think that the paper could benefit from a pass to try to make it more accessible more broadly. I do understand that the nature of the work might make it challenging, for this reason I am happy for the paper to be accepted if other reviewers believe it is worthwhile as I am not overly familiar with this area of research and cannot confidently assess the contribution of the submission given what is already out there.\n\nQUESTIONS:\nI believe it's a typo, but should it not be $x_a^T S_l x_b$ to give a scalar in Equation 1? It looks like $x_a \\in \\mathbb{R}^d$ is a column vector and $h_{ab}^{\\ell}$ a scalar. \n\nCould the $L$ \"attention indices\" be interpreted as $L$ individual heads? I am finding the choice of notation quite confusing when compared to the notation used in standard transformer works, for example $L$ is usually the number of layers while $H$ is the number of heads, $h$ could maybe be called $\\alpha$ to signify that it is acting like an attention coefficient.  -- of course this is personal preference but I am finding the current notation hard to parse .   \n\nEquations 4 and 5 are also difficult to parse. How would they resolve in the context of a ViT for example or an LLM in your notation?\n\nLIMITATIONS:\nThe limitations and overall the wording in the paper is often quite fair and honest.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 1, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces attention-indexed models (AIMs) as a theoretical framework to study learning in attention-based models. These models are the analog of single- or multi-index models for MLPs, adapted to the attention setting. The authors formalize the problem of learning an $L$-layered attention-indexed model, which covers the case of stacked attention layers with residual connections. They characterize the Bayes-optimal predictor in the high-dimensional limit where embedding dimension $d$, sample size $n$ (prop. to number of parameters $d^2$), and attention index matrix rank $r$ grow proportionally, while the number of layers $L$ and context length $T$ remain finite. An approximate message passing algorithm is proposed that recovers the Bayes-optimal performance under certain conditions. The Bayes-optimal analysis is worked out in full detail for the single-layer case ($L=1$), with results applied to both hardmax and softmax attention regimes.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**: This work introduces a new class of models inspired by multi-index models in MLPs to study attention layers. I like the idea of studying the equivalent here given the prevalence of the transformer. The paper takes good first steps by characterizing the generalization and estimation error of the Bayes-optimal solution in this class of models. The results also seem very interesting, especially the contrast between the hard and softmax regimes. There are several directions for future work building on this, including extending to multi-layer attention, adding other architectural components, studying ERM with GD, and extending to symmetric matrices, many of which the authors have already mentioned in the paper.\n\n**Weaknesses**:  I want to highlight a few missing points and some related questions here. \n\n1. **Finite vs. Infinite Context Length Regimes**: You study the regime where context length $T$ is finite. Given that real-world transformers operate on increasingly long contexts, have you considered the asymptotics of $T \\to \\infty$? Do you expect qualitatively new phenomena to arise in this regime?\n\n2. **Relevance and Plausibility of Scaling Regime**: In your analysis, $d \\to \\infty$, $n \\propto d^2$ and $r \\propto d$, while $T$, $L$ remain fixed. I understand that taking $d \\to \\infty$ helps to apply tools from random matrix theory, but in practice, $d$ is fixed and moderate, while $nT≫d^2$. How should we interpret the Bayes-optimal results in this context? Do phenomena like phase transitions, sample complexity thresholds, or rank dependencies remain qualitatively valid when these asymptotics don’t strictly hold? Are there empirical observations or architectural patterns in real transformers that support the proportionality assumptions you've adopted? A discussion on this would be helpful. \n\n\n3. **Hardness of $L>1$ Case**: Lines 193–202 say the difficulty arises from cross-layer noise correlations; if the covariance were diagonal the problem would factorize. Could you clarify exactly which analytical steps break down once this correlation is present, and whether any near-diagonal or orthogonal-layer approximations (e.g., treating weak off-diagonal terms perturbatively) might make an $L ⁣> ⁣1$ extension tractable?\n\n4. **Intuition Behind Hardness of Hardmax**: Fig. 2 shows that softmax reaches zero Bayes-optimal error at finite sample complexity, while hardmax never does. Is the barrier for hardmax purely information-theoretic, i.e. the arg-max mapping discards logit magnitudes so the estimator cannot reconstruct them, or is there an additional effect from the activation’s non-smoothness? A brief explanation of which aspect drives the strictly positive asymptotic error would help clarify the contrast with softmax.\n\n\n5. **Sharp Transition in Small Width Limit**: In Corollary 4.2 (Fig. 1 right), you show a weak recovery threshold for hardmax with small width. Could you provide intuition for why performance is flat below this threshold, i.e., the estimator behaves like it has no signal? What changes fundamentally at the threshold?\n\n6. **Rank Sensitivity in Hardmax vs. Softmax**: Figure 2 shows that lowering the rank makes the softmax task noticeably easier, its strong-recovery threshold drops. Whereas Figure 1 suggests the hardmax error curve barely shifts with rank. Why does softmax benefit so strongly from reduced rank while hardmax remains almost insensitive?\n\nQUESTIONS:\nPlease see the section above.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-23"}