Attention layers provably solve single-location regression

Answer to weaknesses: Even in the setting we consider, which is indeed simpler than practice, the analysis already involves non-convex optimization dynamics, which requires advanced mathematical tools, mainly from dynamical systems theory. We refer to the proof sketch in Appendix A, which details the main technical steps and difficulties of the proof. Regarding softmax, we refer the reviewer to the general rebuttal for a discussion.

We thank the reviewer for the references. The article [2] is already cited in our paper, but we were not aware of [1] and [3], and we will cite them in the next version. While we agree with the reviewer that from a high-level perspective our approach is similar to [1]-[3], we believe that there are significant technical differences between our setting and theirs. In particular, this makes it far from obvious that their proof techniques for handling softmax could be directly borrowed in our case. We will discuss these papers in the next version, and provide some elements at the end of this rebuttal in case the reviewer is interested.

Q1: Our proof technique relies on the use of the Center-Stable Manifold Theorem, a tool from dynamical systems theory. This tool does not provide quantitative rates of convergence. Obtaining a rate is a tricky matter because it requires quantifying the distance to the saddle points of the risk (since the dynamics is slower near saddle points), which in turn requires other tools of analysis and potentially additional assumptions. We will include this discussion in the next version. More generally, as mentioned above, Appendix A provides a 2-page proof sketch (that unfortunately does not fit within the page limit of the main paper).

Q2: Informally, by the Central Limit Theorem (CLT), summing $L$ independent and zero-mean random variables with standard deviation $\lambda$ gives a sum of the order of $\lambda \sqrt{L}$. As noted in the paper, the paragraph you refer to gives an informal intuition, while the proof provides a rigorous justification. We will mention in line 334 that the magnitude comes from the CLT.

Q3: A possible direction to handle a general initialization is discussed in lines 500-505 of the article. In summary, the idea is to show that the manifold is stable in the sense of dynamical systems (i.e., if the dynamics are close to the manifold at a given time, they remain close to the manifold). Our numerical experiments suggest that such a generalization should hold, and that the temperature parameter $\lambda$ should play a key role in this analysis.

Comparison with [1]-[3]:

• Solving the task in [1] requires learning a fixed latent causal graph over the positions of the tokens, making positional encodings a critical element of the analysis. In contrast, our task is invariant under permutations of the tokens. Moreover, in [1], the output is expressed as a function of the last token, with the previous tokens providing the necessary context for this computation. In our setup, however, the output depends on a token whose position varies and must be identified within the context.
• In [2], the argument of softmax (i.e., a matrix $A \in \mathbb{R}^{L \times L}$) is directly a parameter of the model, instead of being a product of the data and of some parameters. This is a radically different structure from the usual attention, and from our setup, where the data appear in the nonlinearity $\sigma(X_\ell^\top k)$.
• The closest work to our setup is the recent paper [3], which also incorporates a notion of token-wise sparsity: the output is computed as the average of a small subset of tokens, where the subset is identified by comparing the positional encodings of each token with that of the last token. A key difference is that, in our setting, the tokens also encodes an output projection direction ($v^\star$) on top of the information on the position of the informative token ($k^\star$). In other words, our task involves learning a linear regression in addition to identifying the relevant token, which is not the case in [3].

W1: This is a good point, and we agree with the reviewer that there is no direct connection. However, we felt it was important to mention these models because they have some similarities to our approach.

1. In De Veaux (1989), the conditional distribution of output $Y$ given input $X$ follows a mixture model, expressed as $Y|X \sim \lambda*N(\alpha_1+ \beta_1 X, \sigma^2) + (1-\lambda)*N(\alpha_2+\beta_2 X, \sigma^2)$. The connection to our model lies in the marginal distribution of each token, which in our case follows a Gaussian mixture $N(0, I)$ and $N(\sqrt{\frac{d}{2}} k^\star, \gamma^2 I)$.

2. In McCullagh & Nelder (1983), single-index models correspond to generalized linear models, represented as $Y = g(\beta^\top X) + \mathrm{noise}$, where $g$ is an unknown regression function learned with a nonparametric model (e.g., a neural network). A slight generalization is the multi-index model, where $\beta$ is a matrix rather than a vector. In our case, the input is not a vector but a sequence of vectors. Nevertheless, the output $Y$ is computed as a nonlinear function of the projection of $X$ onto the directions $k^\star$ and $v^\star$. Thus, our model can be expressed as a specific instance of a multi-index model, where all tokens are stacked into a single large vector $X \in \mathbb{R}^{dL}$, $\beta = ((k^\star, \dots, k^\star), (v^\star, \dots, v^\star))$, and with a specific link function $g$.

W2 and W3: Thank you for pointing this out. We agree that the role of softmax remains an open question and will reformulate our discussion to avoid overstating its importance. We also acknowledge the critical role of normalization (even in the presence of softmax) and will emphasize this aspect more thoroughly, especially when referring to Wortsman et al. (2023). The importance of normalization is also evident in our experiments when the initialization is outside the manifold: the normalization parameter $\lambda$ plays a central role in learning the parameters $k^\star$ and $v^\star$.

Q1: The entire structure of the proof would break if softmax were used; see the common rebuttal for a specific example. The reviewer is correct in pointing out that the choice of erf over another nonlinear, bounded, increasing, equal to 0 at 0, and differentiable function is primarily for technical reasons. Specifically, erf is used because closed-form formulas exist for the expectation of erf and its derivatives applied to Gaussian random variables (see, e.g., Lemma 18). In principle, however, the behavior should remain similar for any such nonlinearity. We will include a discussion of this point in the next version.

Q2: You raise an interesting point. Indeed, one could imagine an MLP designed specifically for this task, where the weights have a diagonal structure with respect to the sequence index. In such a setup, the first layer could learn the projections along $k^\star$ and $v^\star$, while subsequent layers could learn the link function $g$ (see above). However, this architecture is far from resembling those used in practice (and arguably halfway between attention and a standard MLP). If we do not assume a diagonal structure and instead use traditional MLPs, the number of parameters must scale at least linearly with the sequence length, which is highly suboptimal and may lead to very slow training. This highlights the efficiency of attention layers, which perform single-location regression with a fixed number of learnable parameters, independent of the input length. We will add a discussion on this matter.

Q3: The tokens are indeed independent conditionally on $J_0$; thank you for pointing this out. This independence is mainly used to simplify the cross-token interaction terms in the risk evaluation, especially in the proof of Lemma 6 (see e.g. line 994). We will include a mention of this in the next version of the paper.

Q4 and Q5: Thank you for the minor remarks. We will include the link to the code in the de-anonymized version of the paper.

We agree with the reviewer that our model is simplified with respect to practical settings, but our goal in this paper is precisely to present and analyze a simplified model in order to understand real-world phenomena in a tractable setting.

W1. The interesting but more complex case where the information is shared across multiple tokens instead of just one is related to multi-head attention. We have added an additional experiment and discussion on this topic (see the common answer for more details).

W2. and W3. Simplifying the architecture is a common approach in theoretical studies of Transformers, often involving adjustments such as linear attention, omitting skip connections, or removing normalization, among others. That being said, we also performed an additional experiment demonstrating that our simplifications do not alter the phenomena under investigation (see common answer for more details).

W4. Our proof technique relies on the use of the Center-Stable Manifold Theorem, a tool from dynamical systems theory. Unfortunately, this tool does not provide quantitative rates of convergence. Obtaining a rate is a challenging task as it would require quantifying the distance of the iterates to the saddle points of the risk (the dynamics is indeed slower near saddle points), which in turn requires other tools of analysis and potentially additional assumptions. We will include this discussion in the next version of the paper.

W5. We indeed restrict the analysis of the dynamics when initializing on the invariant manifold, as discussed in the paper (lines 468-475). Our numerical experiments suggest that similar convergence behaviors should hold for a more general initialization. A possible direction to extend the proof is discussed in lines 500-505, requiring in particular to show that the invariant manifold is stable in the sense of dynamical systems (i.e., if the dynamics are close to the manifold at a given time, they remain close to the manifold).

W6. Our approach is to start from the experimental evidence accumulated in the literature and provide a simplified model with similar patterns and tractable analysis. As a result, the reviewer is correct in stating that our contribution is mostly theoretical in nature, and does not consist of adding additional real-world empirical evidence to already well-established phenomena such as sparsity or internal linear representations. Besides, our experiment on linear probing on a pretrained BERT architecture, which is intended to verify the role of the [CLS] token as well as how BERT handles information contained in a single token, is closer to practice, although we acknowledge that the NLP task is synthetic.

W1. Our approach is to start from the experimental evidence accumulated in the literature and provide a simplified model with similar patterns and tractable analysis. As a result, the reviewer is correct in stating that our contribution is mostly theoretical in nature, and does not consist of adding real-world empirical evidence to already well-established phenomena.

For example, with respect to sparsity, multiple papers (some of which we cite in lines 141-142) show that a wide range of NLP tasks exhibit token-wise sparsity, meaning that the relevant information is concentrated in $s \ll L$ tokens. It is true that in our paper we consider the limiting case $s=1$ and provide an explanation of how attention layers deal with sparsity in this case. Studying the case $s > 1$ is a very interesting next step, and we provide additional preliminary experiments in the revised version of the paper (see the common rebuttal for more details).

In addition, our experiment on linear probing on a pretrained BERT architecture, to verify the role of the [CLS] token as well as how BERT handles information contained in a single token, is closer to practice, although we acknowledge that the NLP task is synthetic.

Q1. Regarding Figure 1, the synthetic data is described in detail in Appendix E.1. To summarize, the label $Y = \pm 1$ is indeed related to the input sentence, since we set the label to $1$ when a positive sentiment adjective (such as nice, cute, etc.) appears in the sentence, and to $-1$ in the case of a negative sentiment adjective. The color shading in Figure 1(a) was intended solely for visual illustration, but we acknowledge that the current figure and caption may be confusing (especially because we use both "output" and "label" in the caption to refer to the same quantity $Y = \pm 1$). This will be clarified in the next version.