Bayes optimal learning of attention-indexed models

1. Generalization to the multi-layer case:

We point to the prior-channel equations and the related denoising problem with correlated noise (for $L>1$ layers with $L=O(1)$) in Eq. (20) and the paragraph around it. This is actually an interesting random matrix theory problem that may be tractable, but would constitute a result of its own interest in matrix denoising problems. We will emphasize this interesting technical open problem even more in the revised version.

2. Large dimension in numerical experiments:

As the reviewer correctly pointed out, the analytical results obtained are valid in the limit of large embedding dimension $d\to \infty$. Analogously, we also send to infinite the number of samples $n$ and the width of each weight matrix $r_\ell$, with the particular scalings $\alpha = n/d^2 =O(1)$ and $\rho_\ell = r_\ell /d =O(1)$. Regarding the numerical experiments both for AMP, for gradient descent and its averaged version: we use moderate sizes for the embedding dimension, namely $d=100$, and we obtain curves that are strikingly compatible with the Bayes-Optimal theoretical ones (obtained for $d\to \infty$). We point out the fact that reproducing the large d results already at such moderate sizes is exactly one of the main points that go in favor of the high-dimensional theoretical approach.

In the revised version, we will add a figure showing that as $d$ increases the numerical experiments will get closer to the theoretical prediction at $d\to \infty$. We expect to be able for explore the range of $d$ from $30$ up to $1000$ for some of the small values of $\rho\sim0.2$, which is not far from the embedding dimensions used in practice.

1. Background on AMP and Bayes optimality:

Thank you for this suggestion. Indeed, this is a key point and we will explain it better in the text. The analogy works on the two following lines. On one side we express the posterior distribution of the model with respect to S, which can be manipulated in the considered limit to end up with an optimization problem of a free-entropy in Eq. (15). On the other side, a specific AMP algorithm provably reaches the BO error. The specificity is in the choice of two denoising functions that are taken in such a way that the AMP fixed points are extremizers of the same free-entropy, namely its iterates follows the state evolution equations in Eq. (19). We will make such connections clearer in a revised version of this work.

2. Interpretation of the order parameters:

Regarding the interpretation of the order parameters $q$ and $\hat{q}$: $q$ plays the role of an overlap and, at the fixed points of the free-entropy, is ${\rm Tr}(\hat{S}S^\star)$. On the other hand, $\hat{q}$ is its conjugate parameter and plays the role of a signal-to-noise ratio of a high-dimensional denoising problem in the prior channel.

The fixed point equations for $q$ and $\hat{q}$ are indeed complex for $L>1$ (discussed around Eq. (20)). For $L=1$ they are simpler and numerically easy to evaluate, and hence we focus the discussion of the results in Sec. 4 on this already rich and interesting case.

3. Scaling used:

We are interested in studying the leading order behavior of the attention index. Given that we take the input data to have mean zero and variance one, we expect the dot-product attention to also be a collection of $T\times T$ random variables with mean zero and variance one. In the Bayes-Optimal setting and for our model, dividing the dot product by $d$ instead of $\sqrt{d}$ is just a matter of conventions, as the statistician knows about it and can map between the two by rescaling the inputs.

If the data has a mean of order one, the $1/d$ scaling makes fluctuations in the dot-product attention subleading and the output would concentrate. This would significantly simplify the problem, and the relevant sample complexity scaling will be much lower than $\mathcal{O}(d^2)$. We thus have not considered this case of a target function.

There may be other combinations of scalings that make the problem non-trivial. We are happy to consider them if the referee is curious about any particular one and let us know.

1. Weakness 1. and Question 3. : Actual Transformer, and eqs. (4) and (5).

We thank the referee for the pointers regarding notation. We introduce an abstract attention-indexed model (AIM) in eqs. (1) and (2). Then, in eq. (3) we state a simplified multi-layer self-attention with a single head in every layer and no MLP layers. An actual Transformer architecture, such as used in ViT or LLM, would be very close to eq. (3) with multiple heads and attention layers interlayed with token-wise MLP layers. ViT and LLM are usually trained by next-token prediction while the task we consider is supervised. Equations (4) and (5) are then a result detailed in appendix B of how to map eq. (3) on (1-2), i.e. how to map the multi-layer self-attention on the AIM. These equations indeed do not look very transparent, but the existence of such mapping is our key contribution as it enables results from the AIM to apply to multi-layer attention. We will clarify this in the revised version.

Architectures with multiple heads could also be mapped to the AIM form and then the number of indices $L$ in the AIM would be the total number of heads in all the layers. The only true limitation of the mapping of the transformer architecture to the AIM are the MLP layers. We did not find an analogue of Eqs (4) and (5) that would map the transformer including the MLP layer to the AIM.

2. Weakness 2. : AIM claim:

We agree that "AIM offers a solvable playground for understanding learning in modern attention architectures" is a misleading choice of wording. We will rather say: "AIM offers a solvable playground for understanding learning in self-attention layers, that are key components of modern architectures"

3. Question 1. : Row-vectors notation

The notation we used for the quantity $x_a S_\ell x_b^\top$ is correct, as we consider each token $x_a$ to be a row-vector, namely the rows of the whole sequence matrix $X\in \mathbb{R}^{T\times d}$. We understand however that this is in contrast with standard notation of considering column vectors if not better specified. We will change the transposes and make the notation conform to standard practices of vectors being columns in a revised version of this work.

4. Question 2. : Interpretation of the attention indices L:

In this work we focus on multi-layer self-attention (3) where $L$ indicates the number of layers, and thus $\ell$ is a layer-index with $\ell=1,\dots,L$. However, a collection of attention heads can also map to the AIM. We will add an appendix with such a case and refer to it from the main text to illustrate the generality of the AIM. The observation of the reviewer is indeed spot-on and contained in our analysis.

We sincerely thank the reviewers for their questions. We will reflect our answers in the revised manuscript to improve its clarity.

1. Finite vs. Infinite Context Length Regimes:

Our theoretical results indeed require $T<<d$. However, we can take the limit $T\to \infty$ after the high-dimensional limit. In this case, we do not expect a qualitative change. E.g. the location of the phase transition just requires a rescaling of the sample complexity, see Fig. (2) right panel. Indeed, notice that in Result (4.3) we specify Eq. (19) for $\hat{q}$ in the case of a softmax channel, obtaining a simple expression which do not depend on the number of tokens anymore once we do the simple rescaling $\bar{\alpha}=\alpha (T^2 +T-2)$. A qualitative change would likely arise with more complex structures of the data; proposing a suitable model along that line is an interesting avenue of future exploration.

2. Relevance and Plausibility of Scaling Regime:

In general, one would want to explore as rich range of scalings as possible. Previous work [22-24] treated the case $d \to \infty$, and $r, T, L$ finite. In this respect, our paper explores a more general case that allows $r \propto d$ as often chosen in practice. The choice of the sample complexity is given by the fact that it is in this scaling that the performance changes from trivially bad to perfect. Our requirement of $T, L$ remaining finite is indeed a limitation that future work will hopefully be able to remove. One point to make is that while our theory results are derived in the $d \to \infty$ limit, the results do apply for very moderate values of $d$, see Fig. 1 where experiments with AMP are performed with $d=100$. Analogously, GD and its averaged version also show the predicted recovery thresholds (found theoretically for infinite d) with only $d=100$. The condition asked by the reviewer $nT>>d^2$ corresponds as $\bar{\alpha}T>>1$, so technically this scaling is already contained in our analysis by taking $\bar{\alpha}$ very large. We will clarify the discussion about the reasoning behind the considered scaling limit in the revised paper.

3. Hardness of $L>1$ Case:

For $L>1$ the off-diagonal terms of the overlap matrix are of the form ${\rm Tr}(S_1 S_2)$, which do not have mean zero but scale with $\rho^2$ and is thus not small for moderate $\rho$. Diagonal uncoupling could be present in the low $\rho$ regime, where a perturbative approach could be possible, but we were mainly interested in varying this width-ratio parameter and investigate its impact on the properties of the model. Extending to the $L>1$ case would imply solving a high-dimensional denoising problem with correlated noise for the prior channel, which we presently do not know how to do. We left this for future work and progress in random matrix theory.

4. Intuition Behind Hardness of Hardmax:

Generating labels using a hardmax output channel corresponds to learning binary associations over the token sequence $X$, with discrete labels $y \in \{0,1\}$. This task is intrinsically harder than the softmax case. Conceptually, this mirrors the difference between binary classification with a single-layer perceptron—where labels are generated by applying a sign function to a random linear projection—and linear regression on a linear target. In the classification case, perfect generalization typically requires more than a number of samples linear in the input dimension. By contrast, in the regression setting, zero test error can be achieved once the number of samples exceeds the input dimension. In this analogy, the softmax behaves like the linear regression case, while the hardmax mimics the more difficult binary classification case. We will reflect this discussion in the revised paper.

5. Sharp Transition in Small Width Limit:

The mechanism behind the weak-recovery threshold is that of the BBP transition. When the width of the model decreases and becomes of order $r=O(1)$ in the dimension d, the problem is related to that of a rank-1 spike denoising. This can be noticed mathematically in the computations in Appendix D.2. The spike is contained in the bulk before some $\bar{\alpha}_c$ value in this small width regime. The spike instead emerges from the bulk above such value, namely the MMSE starts decreasing and some information starts being recovered even at finite width $r=O(1)$.

6. Rank Sensitivity in Hardmax vs. Softmax::

In case the referee was commenting on Fig. (1) left: both for the hard-max and soft-max the task is easier as lower rank (Fig. 2 left and Fig. 1 left). We point out that the scale of Fig. (2) is logarithmic and that the difference in MMSE among the various curves is more pronounced when visualizing the curves in linear scale as in Fig. (1). In case the referee was commenting on Fig. (1) right: Note the x-axes in Fig. (1) right is rescaled by r to illustrate better the low-r limit. But Fig. (1) right cannot be interpreted as the task getting harder as the rank is lowered.

There are only two days remaining in the author-reviewer discussion phase. The authors have made a concerted effort to address your concerns, and it's important that you respond to let them know whether their rebuttal addresses the issues that you raised.