Towards Understanding the Universality of Transformers for Next-Token Prediction

Thank you for your explanations! Now I better understand the contributions of the paper, and will consider raising my score, but only if there will be a revision of the manuscript during the rebuttal period. You have agreed to carry out many changes, and I would only be confident if I could read the new version, and ask questions in case further clarifications are required.

Thank you for reviewing our paper and for the insightful remarks. Below, we address each of your points in detail, hoping to convince the reviewer to increase their score.

The paper is very dense, even though the authors could have used an additional 10th page. I'd suggest more explanations for the results and fewer formulas in the main text.

Thank you for pointing this out. While we adhered to the recommended 9-page limit, we will make use of the 10th page in the final version to improve the readability of the paper. Specifically, we will add more explanations to clarify key results and streamline the presentation of the formulas.

The intuition seems to be missing (for me) as to why you need the augmented tokens. What do they mean?

The intuition behind augmented tokens is that they enable the model to directly access two successive elements in the sequence. Importantly, while prior works, such as von Oswald et al. (2023a, 2023b), directly worked on augmented tokens, our contribution lies in showing how they can be computed using positional encoding-based attention.

Additionally, von Oswald et al. (2023b) provide experimental evidence that such tokens are formed numerically in real-world models, validating their practical relevance.

Figure 2 needs to be explained; the one-sentence reference in L342 does not explain (for me, at least) what all the quantities in the figure are and how I should think about them.

We will provide a more detailed explanation for Figure 2 in the revised version. Briefly, each point in the figure corresponds to an iterate $\nu_t$, defined as $\nu_t = P_1 \cdots P_t \nu$. The figure illustrates the effect of iterating projections on a vector $\nu$ in finite dimensions, demonstrating how the vector evolves under these transformations.

It is unclear to me how a Transformer can implement Eq. (7) if one problem is with the non-causal descent. You still need to modify the training method, right? Please elaborate on how this works/correct me if I am wrong (I see how you can construct the Transformer, but I do not see how the Transformer itself can account for the change in the descent method). If this point is about the "meta-optimization" during solving the in-context learning task, then please say so explicitly.

There is no training method in our paper. Instead, we construct a model that directly implements the causal kernel descent method in its forward pass. As the reviewer notes, this is related to meta-optimization: our key contribution is showing how the Transformer architecture can implement an optimization algorithm.

To clarify, we explicitly construct a Transformer that performs causal kernel descent and prove its convergence for specific choices of sequences. We will add a remark in the final version to explicitly relate this to meta-optimization, even though the term may not fully apply in our causal setting.

What is the rationale behind the construction in the "More Complex Iterations" paragraph?

We will expand on this explanation in the revised version. The key idea is to apply "non-linear rotations" to 2-dimensional chunks of the input. Each of these rotations is parameterized by an angle $\theta$ and a scalar $q$, and defined as $F_{\theta, q}(z) := R_{\theta} z^q$, where $R_{\theta}$ is a 2D rotation of angle $\theta$. Our construction generalizes this by randomly rotating the input before applying these transformations.

Thank you very much for your review and for the nice comments. Please see our responses to your concerns below, hoping they can convince the reviewer to increase their score.

The theoretical setting is limited to when $f$ is linear and when the sequence is periodic, with each point in the context being on the unit sphere. Neither case seems too complicated in the first place, and so the "universality" does not seem well justified in the context of Theorem 1. Perhaps the authors can justify the reduction to these sub-classes of functions more explicitly.

We would like to emphasize that supposing linear dependency in the data is common in the in-context learning literature; see the related background section, l. 119. The difficulty in our work lies in the fact that the parameter $W$ is different for each sequence, which makes the problem non-trivial.

We would like to stress that, to the best of our knowledge, the problem we study cannot be addressed using existing results. The wording "universal" is to be understood as universal within a class of functions, which here are autoregressive processes. We will change the title of the theorem, as we agree it is overclaiming the result. We apologize for this confusion.

Both the process of generating the augmented token $e_t^0$ from $x_{0:t}$, as well as the causal kernel descent iteration, theoretically involves an infinitely deep Transformer as $n \to \infty$. It might be more informative to discuss the theoretical implications of finite depth.

Even though $e_t^0$ is obtained from $x_{0:t}$ by considering a limit, it is not an infinite depth limit but rather a scaling limit. The reviewer is absolutely correct that, to perfectly approximate the processes, we would need infinitely many layers. However, this is also true for standard gradient descent on any non-trivial optimization problem.

We carefully discuss the theoretical implications of finite depth in Remark 1, as well as in l. 242 of Theorem 1.

The universality of next token predictability result is asymptotic in the context length $t$. Although the convergence is fast enough (exponential), most empirical observations are in the non-asymptotic regime in terms of context length. It might be good to explicitly provide analysis on this.

We are happy to provide larger values of $t$ in the curve, even though we already consider $T = 15k$. Please note that it is computationally expensive to compute attention matrices for large sequences and that modern implementations use clever hardware-aware tricks that we did not use in our experiment. The goal of the experiment is to illustrate the theoretical results, which guarantee exponential convergence.

Questions

Typo.

Absolutely, thank you. Fixed!

In the experiment section...

Yes, absolutely. Please note that $u_t^\star$ can be computed in closed form using Proposition 3, by solving a linear system.

Thank you very much for your positive review and your suggestions. Please find our responses to your concerns below.

Weaknesses

It is true that while linear attention is particularly adapted to linear autoregressive processes, considering linear processes when using softmax attention is still restrictive. However, it remains a non-trivial problem because the matrices $\Omega$ in this case are different for each sequence. In our opinion, it is interesting to see that softmax attention can learn such processes.

We would also like to emphasize that handling the row-wise normalization of the attention matrix is tricky, and we manage to address this in the article. An extension of our results would be to consider more general mappings $W \phi(x_t)$ directly.

For the exponential kernel, the associated RKHS is universal, meaning that any function could be represented as $W \phi(x)$. However, we focus only on a subset of these functions that are linear in $x$, thereby proving universality within a smaller subset class of functions. This is reflected in the wording "towards understanding" in the title.

We agree, however, that the title of our main Theorem 1 may be misleading, and we will revise it accordingly.

Absolutely, the finite-depth model can outperform the infinitely deep model, even though the difference is small compared to the difference with the untrained finite-depth model. This is not surprising, as the trained model's parameters can be adjusted to model much faster optimization methods.

We do not state in the paper that a trained Transformer necessarily implements the causal kernel descent. Instead, we propose this method as a tool to understand the expressivity of Transformers.

Questions

We prove that Transformers are universal within a certain class of functions. By universal, we mean that any function in this class can be approximated. Universality is always to be understood as universal within a specific function class.

Periodicity is simply enforced by replicating sequences of a certain length.

Thank you for spotting this inconsistency. In the context of the first layer, $n$ refers to a scaling factor. We will update the notation to avoid any confusion.

Absolutely. Since skip connections are useful primarily when stacking many layers, we do not believe it is an issue to omit them in the first layer.

Thank you for spotting this typo in the definition of $e^0_1$, which should indeed match your suggestion and is consistent with the expressions used in the proof. We will correct this typo on line 206.

Yes, thank you. Fixed.

Thank you for the interesting question. Actually, we provide an interpretation for $\mu$ using duality in Appendix B. We hope this clarifies the matter.

No, we have not tried it. However, this is an excellent suggestion.

Thank you very much for reviewing our work and for the relevant remarks on the limitations of our approach. We will carefully highlight these limitations in the revised version of the submission.

To specifically address your questions:

(1)

The problem setup is not realistic enough in the sequential relationship. In particular, the reviewer does not believe either the language or the time series admit simple relationship of $x_{t+1} = f(x_t)$. Hence, the reviewer cannot understand how this work contributes to the understanding of Transformers on these modern ML tasks, which is believed to be very crucial to the ML community.

We would like to emphasize that we specify in line 177 that higher-order recursions of the form $x_{t+1} = f(x_t, \cdots, x_{t-\tau})$ could be considered by embedding the tokens in a higher-dimensional space. More precisely, in this case, defining $y_t = (x_t, \cdots, x_{t-\tau})$, one has $y_{t+1}= (x_{t+1}, \cdots, x_{t+1-\tau}) = (f(x_t, \cdots, x_{t-\tau}), \cdots, x_{t+1-\tau})$ which only depends on $y_t$ and can therefore be written as $F(y_t)$ for some function $F$. Therefore, by embedding the tokens in dimension $(\tau +1 )d$, our approach shows universality for autoregressive sequences of memory $\tau +1$. As such, as long as the recursion memory is finite, our approach can be generalized. We commit to add a precise remark about this fact in the paper. We hope we convince the reviewer on this point.

The other crucial point we would like to emphasize is that, compared to previous theoretical works, we consider a strong dependency between the tokens in the sequence. This contrasts with all previous works on in-context learning, which assume i.i.d. data. Therefore, even though we acknowledge that our token encoding choice does not fully capture the structure of language, it is still better adapted than previous approaches.

Moreover, the reviewer believes that RNN might be able to do the same task as the results claimed in this work for Transformers. If this is true, the results might universally hold for many sequential models. Then, the unique advantage of Transformer is not highlighted.

Even though we consider autoregressive sequences, it is not at all clear that RNNs can effectively capture these models. This is because estimating $W$ in context requires computations with long-range dependencies. We kindly remind the reviewer that the in-context function $f$ varies with each sequence and that determining the optimal $W$ requires inverting the data covariance matrix. This does not rigorously prove that RNNs cannot do it, but attention mechanisms inherently handle this, which is what we prove in our work.

We believe RNNs would require more layers to "propagate" the information. While each RNN layer is less computationally expensive, the overall cost might be similar. Investigating this is complex and beyond the scope of this article. The key point is that we demonstrate how current Transformer-based architectures are particularly well-suited for in-context learning due to their global attention mechanism.

(2)

The problem setup is not realistic enough in the deterministic assumptions. The reviewer believes that either Markov chains or some random autoregressive models are necessary for the purpose of studying the universality of next token prediction of Transformers. At the current stage, this fully deterministic dynamic system seems rather naive.

We believe this is a very relevant critique. A possible extension could involve considering noisy dynamics of the form $x_{t+1} = f(x_t) + \varepsilon_t$. We will add a remark in this direction in the final version.

We would like to stress that considering stochastic models would make the analysis even more challenging. Nevertheless, our results are not trivial and shed light on how attention mechanisms leverage pairwise interactions to estimate in-context quantities.

We strongly hope that the reviewer will consider increasing their score as suggested in their review and would be happy to answer any further questions.

Thank you for your review and for taking the time to evaluate our work.

This is not a weakness of the work as much as a fit for the conference. This work seems important but also could arguably be more appropriate for a specialized theoretical audience like various theory conferences (COLT/ALT/STOC/FOCS). Previous works have covered linear functions and the extension to kernels and softmax attention is undoubtedly important for our understanding, but the question is how many of the conference attendees will appreciate these contributions. Of course, theory is a topic on the CFP and transformers are a key interest for ICLR so it's not a bad fit.

We would like to emphasize that many theoretical works are published at ICLR or similar conferences such as NeurIPS and ICML, and we believe our work would be of interest to the audience. For instance, consider the following examples:

• A theoretical understanding of shallow vision Transformers, ICLR 2024: link

• Generalization in diffusion models arises from geometry-adaptive harmonic representations, a purely theoretical paper that received an outstanding award last year at ICLR 2024: link

• Emergence of clustering in attention models using PDE theory, NeurIPS 2023: link

• Implicit bias in vision transformer models showing they can learn the spatial structure of images, NeurIPS 2022: link

• Attention layers provably solve single-location regression, submitted this year to ICLR with good reviews: link

We would be very grateful if the reviewer could reconsider this remark in light of the examples provided. We are convinced that some quick investigation from the reviewer would confirm that our work fits the conference topics.

Just to make sure I understand, is the main technical piece here handling the softmax activation? If the goal was to extend linear activation to the kernelized setting, would that just be adding the "kernel trick" to Sander et al (2024)?

We want to emphasize that our work differs significantly from Sander et al. (2024):

• First, the reviewer is correct to observe that, in contrast to Sander et al. (2024), we consider softmax attention models. This makes the analysis not only more difficult but also much more aligned with practical applications.
• Second, and importantly, we consider more general sequences with no commutativity assumptions on the context matrices $W$ or $\Omega$. This leads to a significantly more challenging problem.
• Third, we consider deep non-linear Transformers, in contrast to Sander et al. (2024) who considered shallow linear Transformers. Having many layers is necessary in our case due to the non commutativity of the matrices $W$.

Sander et al. (2024) focus on linear attention with one layer because the problem they study is simpler. Additionally, their work characterizes the minima of the loss landscape in that case, whereas we are interested in universality results. We stress that our work does not merely involve adding a "kernel trick"; we prove non-trivial results on the convergence of new causal kernel flows, which we believe can have a significant impact on the theoretical understanding of autoregressive attention-based models, such as LLMs.

For these reasons, we would be grateful if the reviewer could consider increasing their score. We are happy to provide further clarifications should the reviewer have additional questions.