Towards Understanding the Universality of Transformers for Next-Token Prediction
Authors: Michael Eli Sander
Gabriel Peyré
Venue: ICLR 2025 Poster (2025)
Keywords: Transformers
In-Context Learning
Deep Learning Theory
Primary area: learning theory
TL;DR: We investigate how causal Transformers can predict the next token in a sequence by constructing models that use a causal kernel descent method to learn context-dependent functions.

Abstract: Causal Transformers are trained to predict the next token for a given context. While it is widely accepted that self-attention is crucial for encoding the causal structure of sequences, the precise underlying mechanism behind this in-context autoregressive learning ability remains unclear. In this paper, we take a step towards understanding this phenomenon by studying the approximation ability of Transformers for next-token prediction. Specifically, we explore the capacity of causal Transformers to predict the next token $x_{t+1}$ given an autoregressive sequence $(x_1, \dots, x_t)$ as a prompt, where $ x_{t+1} = f(x_t) $, and $ f $ is a context-dependent function that varies with each sequence.
On the theoretical side, we focus on specific instances, namely when $ f $ is linear or when $ (x_t)$ is periodic. We explicitly construct a Transformer (with linear, exponential, or softmax attention) that learns the mapping $f$ in-context through a causal kernel descent method. The causal kernel descent method we propose provably estimates $x_{t+1} $ based solely on past and current observations $ (x_1, \dots, x_t) $, with connections to the Kaczmarz algorithm in Hilbert spaces. We present experimental results that validate our theoretical findings and suggest their applicability to more general mappings $f$.

Published as a conference paper at ICLR 2025
TOWARDS UNDERSTANDING THE UNIVERSALITY OF TRANSFORMERS FOR NEXT-TOKEN PREDICTION
Michaël E. Sander & Gabriel Peyré
Ecole Normale Supérieure, CNRS
Paris, France mi chael wwwderr@polyteehnil.uuuuuu org, gabriel
peyreeens. ..briel.peyreeen............ fr
ABSTRACT
Causal Transformers are trained to predict the next token for a given context. While it widely accepted that self-attention is crucial for encoding the causal structure of sequences, the precise underlying mechanism behind this in-context autoregressive learning ability remains unclear. In this paper, we take a step towards understanding this phenomenon by studying the approximation ability of Transformers for next- token prediction. Specifically, we explore the capacity of causal Transformers to prediet the next token Z++-1 given an autoregressive sequence (æ1,.... 24)
prompt, where It+11 = f(xx), and f is context-dependent function that varies with each sequence. On the theoretical side, we focus on specific instances, namely when f is linear or when (#))(21 is periodic. We explicitly construct a Transformer (with linear, exponential, or softmax attention) that learns the mapping in-context through a causal kernel descent method. The causal kernel descent method we propose provably estimates #t+-1 based solely on past and current observations (zi,...: ), with connections to the Kaczmarz algorithm in Hilbert spaces. We present experimental results that validate our theoretical findings and suggest their applicability to more general mappings
INTRODUCTION
The transformative impact of deep learning on artificial intelligence has led to increasingly powerful architectures, with Transformers (Vaswani al., 2017) at the forefront. These models have in particular revolutionized natural language processing (NLP) (Devlin et al., 2018), and now serve as the foundation for large-scale language models, such as GPT (Radford et al., 2018; Brown et al., 2020), significantly advancing artificial intelligence's 's capabilities in both understanding and generating human language, setting new benchmarks across various tasks.
Most recent large language models (Hoffmann al., 2022; Team al., 2023; Jiang et al., 2023; Dubey et al., 2024) are causal Transformers pretrained to predict the most likely next token x++1 from finite vocabulary given a context F1: = (I1, .x2). These models excel such tasks and beyond, demonstrating what is known as in-context learning: after training, they show remarkable few-shot learning capabilities, inferring patterns from just few examples within the context (Brown et al., 2020). Recent studies suggest that in-context learning capabilities emerge from the Transformer performing optimization on an inner objective during its forward pass, where the attention matrix plays crucial role (Von Oswald et al., 2023a; Mahankali et al., 2023; Ahn et al., 2023; Zhang et al., 2023; Kim & Suzuki, 2024). In particular, Cheng et (2024) demonstrate that trained Transformers can perform in-context learning through kernel ridge regression. However, the question of why causal Transformers excel at general autoregressive prediction remains open.
In this paper, we propose kernel interpretation for the autoregressive setting. We introduce framework to rigorously analyze the expressivity of deep Transformers next-token prediction. Specifically, consider sequences generated according to Z f(xx), with context-dependent function a vector-valued Reproducing Kernel Hilbert Space (RKHS) associated with positive semi-definite kernel k. This generalizes the works of Von Oswald et al. (2023b); Sander et al. (2024). Within this framework, we explore how successive attention layers solve a causal kernel least square regression problem to predict the next token accurately, as described in Figure 1.

---

Published as a conference paper at ICLR 2025
1ist
To
cu
T
c
= Pet ~ de+1
Figure 1: Illustration of the method proposed in this paper. Given sequence 21:t, a first layer To computes augmented tokens ei Next, stack of n identical Transformer layers 7 with residual connections iteratively update the tokens ei... following the causal kemnl descent method introduced Section 4. For autoregressive sequences presented in Assumption I, and under specific instances outlined in Assumption 2, projecting of with projector P yields estimate ut of #++1 as and approach +oo, stated in Theorem 1.
More precisely, we make the following contributions:
In Section 3, we formalize the sequence generation model that serves as the foundation for our theoretical results. We then present our main result in Theorem 1, demonstrating that there exists Transformer model, with an explicit construction, that, given the first tokens 21:6, accurately predicts the next token 2++1 as & tends to +00, for specific instances, namely when f linear or when the sequence (#4))22 is periodic.
In Section 4, present the methodology underlying the proof of Theorem 1. We introduce family of causal kernel descent methods that build an estimate u, of It+1, based only on past and current observations F1:1- This approach modifies least squares gradient descent method to account for causality while preserving the parallelization benefits the Transformer architecture. In Theorems 2, 3, and 4, prove that for the specifie cases considered, uy 2 converges to 0 as +00, drawing connections to the Kaczmarz algorithm (Kaczmarz, 1937) in Hilbert spaces. Finally, Proposition 5 shows that the causal kernel descent methods be implemented with Transformer model, as illustrated in Figure 1.
In Section 5, we first present experimental results that validate our theoretical findings and extend them more general class of mappings f beyond those studied in Sections 3 and 4. We then empirically show that the Transformer models constructed in Theorem can be successfully fine-tuned obtain faster convergence of the estimate with the sequence length
BACKGROUND AND RELATED WORK
Transformers. Transformers (Vaswani et al., 2017) process sequences tokens (#1, Fr) or arbitrarily length T. After embedding the sequence into new sequence (eI. er), the Transformer consists of a series of blocks with residual connections (He al., 2016). Each block comprises two primary components: multi-head self-attention mechanism and feedforward multi- layer perceptron, with the latter operating independently on each token. We will almost always disregard the feedforward layer in this study (except when building augmented tokens), theoretical analyses of in-context learning (Mahankali et al., 2023; Ahn et al., 2023; Zhang et al., 2023). In contrast, the multi-head self-atte relf-attention mechanism involves pairwise interaction between the tokens. This module consists in applying multiple self-attention operations in parallel, parametrized by of weight matrices (W'. Whh W/) SH, where // denotes the number of attention heads (Vaswani et al., 2017; Michel et al., 2019). The output of the multi-head self-attention mechanism is given by:
common
e, 5 e2 + T(e.,) with T(e ΣΣ Ahe
(1)
where Ah, the attention matrix, determines the attention weights between tokens and is typically defined as: Ab, = =N Wkee)), with (:.) representing a dot product and / being normal- ization function. The standard choice is to consider N = softmax, i.e.
A = WWg., We
One can also consider the unormalized attention when = exp. Another approach is to con- sider N id, which corresponds to what is known (despite being non-linear) as linear attention
2

---

Published as a conference paper at ICLR 2025
(Katharopoulos et al., 2020), enabling faster inference. One significant advantage of equation is that the updates on the e 's can be computed in parallel during training, leveraging modern hardware for faster computations.
Expressivity. The universal approximation properties of encoder-only Transformers are well estab- lished. Yun et al. (2019); Nath et al. (2024); Furuya et al. (2024) demonstrate that Transformers can approximate cnnnuttatiooequivariainn functions. A more constructive approach, though applicable to narrower class of functions, is proposed by Wang models, the expressivity of Transformers for next-token prediction is not yet fully understood, though a popular recent line of works studies the in-context learning ability of Transformers.
(2024). When it comes to decoder-only
In-context learning. A major property of Transformers is that they adapt their computations given the context. In particular, given a context (I1, g(z1),... " a trained large Transformer can infer the next output g(x..) without parameter updates. Many recent studies have contributed to understanding this phenomenon. The seminal work of Von Oswald et al. (2023a) considers functioon g of the form g(x) w x for some w and construct a linear Transformer for which the forward pass is equivalent to a single step of gradient descent on mean squared error loss. Theoretical guarantees are provided by Mahankali al. (2023); Ahn et al. (2023); Zhang al. (2023), showing trained one-layer linear Transformer implements one step of (preconditioned) gradient descent. Other works study the softmax attention without accounting for training dynamics (Garg et al., 2022; Akyürek et al., 2022; Li et al., 2023). Of particular interest to us is the recent work of Cheng et al. (2024) which shows that there exists a simple parameter configuration of non-linear Transformers such that they implement gradient descent in the function space with respect to the RKHS metric induced by the attention kernel. In this work, we take a step further by considering a more general next-token prediction task and propose causal kernel descent method that can be implemented by Transformer to solve it. For this, we extend the autoregressive in-context learning setting introduced by Sander et al. (2024), where tokens are generated according an autoregressive process of order 1: Iz-1 - (xx). where [(x) = Wx for a context-dependant parameter W, varying with each sequence. The autoregressive in-context learning ability is described as the model's capacity to decompose its prediction into two steps: first, estimating W through an in-context mapping, and then applying a straightforward prediction function, which is either equal to or closely related to x + Wr. However, Sander et al. (2024) focus solely on linear Transformers, whereas our study encapsulates linear, exponential, and softmax Transformers. Our key contribution is showing how such Transformers can implement optimization algorithm-termed causal kernel weshent-within their forward pass. This can be interpreted as a mntaioopii mechanism (Von Oswald et al., 2023b).
Neural Ordinary Differential Equations. Neural ODEs (Weinan, 2017; Chen et al., 2018) are a class of implicit models where a neural network F parameterizes a vector field in an ordinary differential equation (ODE), as follows:
de dr (+)=((()))))
where r denotes the continuous depth of the network. For a given input e00). neural ODE outputs deep representation e*, which is the solution (when it exists) of equation with initial condition e(0), at some finite or infinite time horizon. Neural ODEs can be viewed as continuous-depth analog of deep Residual Networks (He et al., 2016), where the latter corresponds to an Euler discretization for solving equation (Marion et al., 2023). As such, neural ODEs are widely employed to better understand the theoretical properties of Residual Networks. In the context of Transformers, the neural ODE framework has emerged as a valuable tool for studying attention-basee models, such as the impact of the choice of normalization function N model behavior (Sander et al., 2022) and the emergence of clusters (Geshkovski et al., 2024). However, these studies do not address the autoregressive setting, which tackle in this work.
3 TRANSFORMERS FOR AUTOREGRESSIVE IN-CONTEXT LEARNING
Notations.
Throughout the paper, we denote d the dimension of 2t ||-I1 the & norm, O(d) the orthogonal manifold, 1d the unit sphere in IRd and the adjoint.
In this section, we describe the types of autoregressive sequences we consider and present our main results, which show that we can explicitly construct Transformer models that approximate the next token in the sequence.
3

---

Published as a conference paper at ICLR 2025
3.1 PROPOSED FRAMEWORK
Reproducing Kernel Hilbert Space. For k RR x IR + I a positive definite kernel, we define H as the vector-valued Reproducing Kernel Hilbert Space (RKHS) associated with the feature map p F + H. where (zz) (x.-).. Therefore, for any function f E H. there exists a linear map W Rd such that f() Wy(()), where, denoting W (w1, wa) as vectors of H.
[| == Ei= || llx is finite.
We consider autoregressive sequences of order 1, which we formalize in the following assumption. Assumption! (Autoregressive sequences.). We consider sequences of order 1, defined as follows:
Initial State: The sequence starts with some #1 E 5d-1
Hidden Variable: We suppose there is a hidden variable f E such that the subsequent states are generated autoregressively as 1t+11 = [(=+)](e)))) for 1.
Note that although we focus first-order recursions here, higher-order recursions can be considered by embedding tokens in a
space. Indeed, for recursions of the form 2[+1 = g( -T-r) )) for some context mapping defining y/ == (xt, ... ,It-r) E R(++1)d one has Ut++1 (It-1,--. --+ = (g It-r),.. xz++-7) which only depends on yt. We thus have V-++1 = f(v) for some mapping f. Therefore as long as the recursion memory is finite, our approach can be generalized. The formulation in Assumption 2 differs from the classical in-context learning setup, where sequences consist of input-output pairs (m Vr). that are often supposed to be iid. We argue that modeling sequences with It+1 = ((x+) better reflects the nature of real-world sequences on which causal Transformers, such large language models (LLMs), are trained. Indeed, our proposed model incorporates autoregressive relationships, aligning more closely with the data LLMs encounter during training. However, similarly to in-context learning, in order to accurately predict the next token T++1 given the previous states Iiat as inputs, a Transformer would have to implicitly estimate the hidden map f. In this paper, propose general method to provide an estimate of r++1 given Z14- However, to prove that such an estimate can be built with a standard Transformer, and in order prove our universality results in Theorem 1, consider specific choices for the kernel and the function f, summarized in the following assumption.
Assumption (Specific considerations.). We define the following instances.
(I) k(x, y) ka, v) = (x, v). fxx Wx for some W E (O)).
(2) k(x,, u) k2, = c( and [(x = Qx for some Ω Odd).
(3) l: = keep. the sequence (zz)121 is periodic and xe E 5tt- for all t.
Note that for these specific kernel choices, because ||æe|| = 1, one has k( xx) 1(x11 a1) for all t. This follows because k( 2x) depends only on ||=r,|| = ||
crucial step in our construction is the building of augmented tokens for a sequence (I1,".. .rT). Indeed, in Von Oswald et al. (2023b) and Sander et al. (2024), augmented tokens are used, explicitly encoding the relative positions of tokens I₁. This approach is similar to in-context learning, where sequences often consist of pairs (Ii, yyi) encode both input and output information. However, such construction with attention based models is non-trivial. Sander et al. (2024) propose method for constructing augmented tokens using general positional encoding. In this work, we provide a more detailed result, showing that augmented tokens can be computed with a one-layer Transformer that employs dot-product absolute positional encoding. We introduce coanning-offssqqeenee token == 022 and consider the extended sequence To.T = ( Z1, .T).
Augmented tokens.
In this work, we are going to consider the augmented tokens e/ E IR 4d42+ defined as
c == ( 0, xt, 1, 0z) for > 1, and et = (0a, 1, xr, 1, Od, 0a)-
The intuition behind augmented tokens 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;b), directly worked on augmented tokens, one of our contribution lies in showing how they can be computed using positional encoding-based attention. The augmented tokens incorporate standard
4

---

Published as a conference paper at ICLR 2025
beginning-ov boaimiing-of-squueneeee of-seueence tokens, while the I's correspond to concatenated positional encodings, which are technically required in our proofs. More precisely, the following proposition shows that the of can be approximated with Transformer layer.
Proposition 1. There exists sequence of one-layer and 2-heads causal Transformer To: with N = softmax followed by a feedforward layer, such that To( = -imm,++0o To" (Fo:t) = ee.
proof of this result, along with the explicit construction of the corresponding model, is provided in reginning-ff-sqquee token. The 0's and
Appendix A.I. Here, ef should be interpreted I's in El correspond to specific choices positional encodings, which are crucial for handling the softmax normalization in main Theorem 1. From this point onward, will consider the new augmented tokens c.
3.2
MAIN RESULT
We now present our main theorem. Using the augmented tokens et, we demonstrate that stack of Transformer layers can approximate the next token Iz41 based solely eq (and thus solely on 2ir as well) as increases.
Approximation with Transformers. We consider a model MM composed of n identical Trans- former layers T with residual connections, followed by a projection. The model iterates from the begimning-of-sequence token eg as:
M" (x = Pej where e- T(ee)) 0 0sk k Ann-ni
with gdd+2 + d the projector selecting the last d coordinates, which can be integrated in the form of a token-wise feedforward layer.
We have the following main theorem, which shows that for the specific instances we consider, there exist Transformer layers T such that lim lim MM1I1 I++1 = 0.
Theorem 1 (On the expressivity of Transformers for Next Token Prediction). For each instance in Assumption there exists an attention-only, one-layer, two-head causal Transformer T with attention normalization N such that, for any autoregressive sequence (æ: (2))121 generated according Assumption MM (x11:)) converges exponentially fast as goes to infinity. Furthermore, denoting M(11a) = lim (MIr))) one has lim;- (M(za))
x++1) 0. More specifically:
For instance (1), N = id, and the convergence of (M(11a) x++1) ) to 0 is exponentially fast in t for almost all 2i and W.
For instance (2), N = exp or N softmax.
For instance (3), N' = exp or N = softmax, and the convergence of (M(21a) 1+) to is exponentially fast in t.
Finally, for any of the instances above, when N' = id or N exp, we have (² M(x1a) as long as 2 t, so that lim (2M x(+1)= = 0.
We emphasize that the models correspond to explicit constructions. Specifically, we have:
W/. Wk E Rdx(10+2) W W (1d4))) W W ((d+1)((1++2)
More precisely, one has
WC2)
O0tixdt1 Jasi. O241xza),
( w
(2)
Mann Ozt1x Ozerxza)
0441xd+1
Oct1xa+1
Oa-rx
w
04+1x4411
Oge1xe+1
Og-
Ogetxo
02+1xd+1
0a-1xd=1
Ozeixe
0aeixd
04xd+1
Osxa-1
W
02s1xa11
0g-1xd:1
0ssixd
04xi+1
Oaxd-1
Oaxe
Ia
Oaxd-1
0axdt1
0
Opxe
0axd-1
02v441
Iu
Oaxd
The proof of Theorem 1, which in Appendix A.9, relies on demonstrating that the model M implements causal kernel descent method, which describe and analyze in Section 4.
5

---

Published as a conference paper at ICLR 2025
Remark 1. In Theorem we have || (z1a) r++1ll I Ei(n,t) + Ey(t), with (1(n.t)
(M1a) MMxx1)) and E2(t) = [[M(a - Z The proof of Theorem 1 reveals that to E1 (n. hence the last statement of the theorem. When N softmax, we conjecture that E1 (n. - 0 as x++1) 0. We provide
and a (t) = O. When N id N exp, E (n,,t) 0 if t,
exp.
n,t. - +00 and t/n 0, which implies pt and andt/n (MM evidence for this conjecture in the last paragraph of Appendix A.9.
The model M introduced in Theorem also has a continuous-timm interpretation, which we now formulate.
Neural ODE Interpretation. of Transformer layers goes to infinity can be interpreted as continuous time neural ODE (Chen al., 2018). Specifically, M satisfies
The model M defined in Theorem the limit when n n-the number
de: (T)= T((nr(()) dr
M(x1a) lim Peş(T), where T--00
with ec(0) To(
(4)
We now present the theory behind our constructions to guarantee consistent approximation of the next token 1 as the sequence size increases.
Comparison with RNNs. Even though consider autoregressive sequences, it is not straightfor- ward that recurrent neural networks (RNNs) can effectively capture these models. This is because estimating W in-context requires computations with long-range dependencies. Determining the optimal W indeed requires inverting the data covariance matrix. Attention mechanisms inherently handle such procedure, which is proven in this work. We believe RNNS would require to "propagate" such information. While each RNN layer is less computationally expensive than an attention layer, the overall cost might be similar. Investigating this is complex and beyond the scope of this article, which demonstrates how current Transformer-basee architectures are particularly well-suited for in-context learning due to their global attention mechanism.
layers
4 CAUSAL KERNEL DESCENT
In this section, introduce causal method instantiate the iterations in equation and prove Theorem 1. Specifically, we propose a causal kernel descent method that incorporates causality into standard gradient descent for least squares minimization.
4.1 CAUSAL DESCENT
Non-Causal Descent. We consider sequence F1.7 (x1, .TT) Rxxd The goal is to
solve the least squares minimization problem of minimizing Σx=1 [x((x) x++ with respect to
E H. Recall that for any f E H. there exists linear map W H - Rdd such that f(x) W (1r)
for all x E d Thus, consider the least squares optimization problem:
T-1
min E(W) = Σ W
Wop(s.) - 113.
(5)
We solve equation 5 using gradient descent with step size 2. starting from initial W and iterating for 0 < t: < n 1:
T-1
Wh+1 -W* Σ
p(x.) - x=+1)9((2.'.
==1
We define prediction variable u W (xx.) ). By right-multplyingg the gradient descent equation by (xx)), we obtain:
u u 1 Σ k(z,, I.+1),
A+1
(6)
which corresponds to a least squares descent on the predictions. However, this descent is non-causal because the update of u, depends u1r-1 and Fi.r, making it unsuitable for implementation in causal Transformer. We now propose a causal formulation for equation 6.

---

Published as a conference paper at ICLR 2025
Causal Descent. Inspired by the descent in equation 6, propose modified least squares descent that introduces causality, ensuring that each estimate of 2(+1 is based solely on past and current observations Z1:t- To achieve this, define the following causal kernel descent, which incorporates both an unnormalized and row-wise normalized framework. Starting from any initial uf. the descent iterates for 0 v k I n as follows:
(((e; 2o); k( 2x) E. k(
k = hid
k++1 u 7) Σ u1
(u, lsct#o+1) with At,s
k(a),z,)
if k: Kexp
(7)
We denote A the corresponding lower triangular matrix. Note that this descent is causal in the sense that th." depends only on Z1:t- Note also that each iteration in equation 7 can be parallelized.
For well-chosen step sizes rj, the method in equation 7 converges. Specifically, we have the following proposition:
Proposition 2. For each causal kernel descent in equation 7, there exists nf" and u, such that, for all 17 < 17,, ut - Un exponentially fast as n goes to infinity. Specifically:
When At,s k.),,, then 17 =(x1,z1) Moreover, when 1j k((1,01) up uz if
k(u...)) When At,s E!-1 k(z,,,)) then 17" = 2.
Our proof, in Appendix A.2, relies on the lower triangular property of the matrix A.
Our objective is to show that in "close" to #1+1 when is sufficiently big, demonstrating that the causal kernel descent accurately tracks the future states provided it has seen sufficiently long context. We have the following result.
Proposition 3. Let (1e)22 be the unique sequence of vectors satisfying
Σ Eukk(Eo, = I1+1. VE 2 1.
(8)
Then for all t 1, F++1 it 11)14
For proof, see Appendix A.3. From Proposition 3, see that lom (ut x++1) 0 is
equivalent to having lim fe = 0. We provide a dual interpretation for μ in Appendix B.
4.2 CONVERGENCE IN THE SEQUENCE LENGTH
In this section, we prove that under the instances defined in Assumption 2, lim, (už 2141) = O. For any μ satisfying equation 8, we define the following map from H to Rd:
..... P2
1.0
W₁ = Σ (usp(a))
0.5
By construction, we have for all t: Wep(x;)) Iz-1. To prove that oom Pr = 0 (and therefore tom (ut- F1+1) 0), we will actually prove a stronger result, showing that for the instances in Assump- tion 2 and in a specific limit sense, Wr W. For this, we define
0.5
1.0
1.0
0.5
-055
0.0
0.5
0.5
1.0
1.0
((xx)
((()) (x))()))
For some random vectors (green) and V1,---Mo in S², we display P.u (grey), Ps Pav (orange), ... and P Pav.
Figure
==
and P. = (T v.y?)
(10)
Therefore, P. is the orthogonal projection onto the subspace orthogonal to p(x.). With these notations, have the following proposition:
7

---

Published as a conference paper at ICLR 2025
Proposition 4 (Estimate Update Recursion). For all 2 1, we have
We W P
For a proof, see Appendix A.4. Intuitively, the recursive relationship in Proposition 4 shows how the difference between W; and W is progressively reduced by successive orthogonal projections. Such a process is illustrated in Figure 2, where each point corresponds to an iterate 1/. defined as Vt Pi P.v. Note that without further assumptions, there is no guarantee that (W, - W) w))p(ar) converges to for arbitrary æ. However, for the specific instances considered in Assumption 2, we are able establish convergence for vectors æ in a certain space.
We first focus on instance (1), where we can derive convergence speed results.
Linear recursions and dot-product kernel. In this section, we consider instance (1) from As- sumption 2. Under these assumptions, we prove the following theorem:
Theorem 2 (k = kid, linear recursions). Under instance (1) in Assumption 2, one has that lim;-++00 (ut -I++1) 0. In addition, for almost all W and Iı, W, + W exponentially fast,
Proof sketch. We first show that W - W = ( - x o almost all W and 21, p(W xxxe 1, where p denotes the spectral radius.
We then establish that for
See Appendix A.5 for full proof. We now turn to the case where k = kiexp
Linear recursions and exponential kernel. In this paragraph, we consider instance (2) in Assump- tion 2. We present the following theorem, with complete proof provided in Appendix A.6.
Theorem 3 (k = hexp- linear recursions). Under instance (2) in Assumption 2, for any x E RRd. (W: W*) xr - O. In particular, lim (u: xx+1) O.
Proof sketch. We consider the subset of H comprising functions x + a.K., for E R and 7 2 This subset a pre-Hilbert space under the inner product inherited from H. By completing this space with respect to the induced norm from H. we obtain new Hilbert space Ht. We show that P.Pt-1 ... P1 + strongly in H' as t + +oc. For this, we observe that the convergence of P.P.- ... P1 is equivalent to the convergence of the Kaczmarz algorithm (Kaczmarz, 1937), A sequence 221 of unit vector for which P. P.- ... P1 - O strongly as + +oo referred to as effective. The specificities of our case twofold: first, V/ H'' which is potentially of infinite dimension, and second, the vectors V, follow an autoregressive relation. We show that the sequence (va)<21 is effective in H'. Note that because Ω E O(d), one has for any positive integers r that (Ustr: P(4r) = (u,, M1) Such sequence is called stationary. Bochner's theorem states that there exists measure a on the unit circle -'alllled spectral measure-such that, for all 2 1,
a (1)+1, = f 2 dz(
We then use the following characterization from Kwapien & Myciclski (2001); Rainis Haller (2005).
Theorem (Effectiveness stationary sequences (Kwapień & Mycielski, 2001)), A stationary sequence of unit vectors which is linearly dense in Hilbert space is effective if and only if its spectral measure either coincides with the normalized Lebesgue measure or is singular with respect to the Lebesgue measure.
This characterization combined with Fourier analysis results shows that the sequence (v. )x2) is effective. While this does not necessarily imply the strong convergence of W. W WP, P to 0, this enough to prove that lim (u, x(+1) 0.
Periodic recursions. In this paragraph, we turn to instance (3) in Assumption 2, where i: = keep and the sequence (3r))21 is assumed to be periodic. The motivation for considering periodic sequences is better align with real-world next-token prediction tasks, where the vocabulary has a finite size and is frequently repeated in eyclical patterns. We have the following theorem (proof in Appendix A.7).
Theorem 4 (l: = kesp- periodic recursions.). Under instance (3) in Assumption 2, one has that ti (uf x1+1) O exponentially fast.
8

---

Published as a conference paper at ICLR 2025
4.3 IMPLEMENTATION WITH TRANSFORMERS
We can now use the previous results stating that the output of the causal kernel method uz approaches Tt-1 increases construct the models presented in Section 3.
Expressing equation with a Transformer. We have the following proposition, which, combined with Theorems and 4, allow us prove Theorem
Proposition 5. For any > 0, for each configuration of equation 7, there exists an attention-only, one- layer, twe-head causal Transformer T with attention normalization. N such that, for any autoregres- sive sequence (Iz)121 generated according to Assumption 1, defining ex == (It-1, O, It. 1, 2, ut) for >> 1 and ef == (0a, 1, I2, 1, Od, uf), ei solves equation 3 if and only if u solves equation 7. More specifically.
When At,s k( vz)) and l kid, have N id.
When Ar,s k( and l keag. have N exp.
A(au.z)) When At,s Disi k(x,,2,) and k kiezy, we have N softmax.
Here again, we stress that the transformer layers T correspond to an explicit construction. See Appendix A.8 for constructive proof. With Proposition 5, we prove Theorem in Appendix A.9.
Keep
kd
10-2
1 U,
10-1
10
5000
10000
15000
2500
5000
7500
10000
0 2500
5000 7500
10000
t
Figure 3: Evolution of the squared error ||u] - = with for different scenarios. The curves averaged over five sequences T1:4. Left: instances (1) and (2) (with random W and f2). illustrating Theorems and (d 15). Center: instance (3), illustrating Theorem (d 15, the period tp randomly sampled between 20 and 40, and random sequence repeated tp times), Right: instance (4) described Section (d = 4).
EXPERIMENTS
In this section, we present experimental results. Our code will be open-sourced.
Illustration of the Theorems. We first illustrate the theoretical results of Theorems 2, 3, and 4. For the three instances in Assumption 2, compute u, defined in Proposition 2. We then examine the evolution | I2+1ll2 with t. The corresponding curves are shown in Figure 3 (left and center), illustrating the convergence of ut Ft+11 zero for the considered instances. As predicted by our theory, the convergence is exponential for instances (1) and (3).
More Complex Iterations. We also explore the potential convergence of ut - to zero for a new instance (4), described as follows. The key idea is apply "non-linear rotations" to 2-dimensional chunks of the input. Each of these rotations is parameterized by an angle and a scalar q, and defined as 2 + Raz", where Rg is dimensional rotation angle 6. We now formalize this idea. We consider sequences {z,} in 2R generated according to the process in Assumption 1, where the mapping / 2p 2p defined as follows. We first convert the real vector J, into complex (2)
vector 2, E CP, where 2 (2j-1) + i x (2)) for j = 1, 2, p.. We then apply unitary matrix U E CPXP to 2,, defining 2 Uz, Next, we modulate the magnitude and phase of 2/ with scalar q E R and bias vector & E IR as follows:
x exp((i)) (1='l @ exp (ig arg (=()))

---

Published as a conference paper at ICLR 2025
where @ denotes element- wise multiplication, | is the element- wise magnitude, and arg(z{) is the element-wise phase of z/- We then compute 24+1 = U+2 2f', where U is the Hermitian transpose U. Finally, convert 2 back to a real vector 1++1 E R The sequence starts from an initial vector Ii E R normalized to unit length. The unitary matrix U and bias vector are randomly initialized. The parameter q controls the non-linearity applied to the phase of the transformed vector. In Figure (right), we plot l|u z++1 Il against for d = 4, q = 2, and k = Kexp- We observe convergence to zero, suggesting that our causal kernel descent method may generalize to more complex settings.
Training M In Theorem 1, the number of layers, is taken infin- ity to obtain the estimate if of I++1- However, we experimentally demon- strate that fine-tuning the model M equation leads trained model that competes with the infinitely deep model M.. For this, we take N softmax to build M", and fine-tune the corresponding weights. Note that, for simplicity, consider square trices for the parameters, by complet- ing the parameters with zeros. To em- phasize the parameter dependency, we denote by Mg the corresponding ini- tialization (which therefore satisfies
à [[ - M
ā alle+1 1 (x:0011
3 blx++1 - - (x
10-2
10-3
0
20
40
60
80
100
t
Figure 4: Errors 1 (1(())))) 24+1 2 against for G MM Results are averaged over the whole test set. Theorem 1). We take d 15, 6, and consider instance (2) with randomly generated Lt's and a1 is for dataset with 2 elements, that we split into train, validation, and test sets with respective sizes of 60%, 20%. and 20% of the original dataset. We train the model using Adam (Kingma Ba, 2014) on the Mean Squared Error (MSE) loss for next-token prediction on sequences of length = 100, i.e., we minimize:
1 T-1 ( = d 1=1
Σ
(M - x241l1?
(
the parameters of M", starting from the initialization fg. We train for 5000 epochs with early stopping. We denote by MM the corresponding model. We then examine how the error a (I1z) 22 behaves with i. We find that not only does Mÿ, significantly outperform Ma but it also outperforms the infinite depth model M when t > 20, as shown in Figure 5.
CONCLUSION
In this paper, we took a step towards understanding the universality of Transformers for next-token prediction by considering sequences of the form 2++1 = f(xz) for some hidden variable f. We demonstrated in Theorem that an explicitly constructed Transformer can accurately predict the next token Ø[+1 as + +00, in specific cases where f is linear or the sequence (zx)121 is periodic. Our construction corresponds the Transformer implementing causal kernel descent methods, which provably provide consistent estimates of the next token 2++1 under the specific cases considered in this paper. Experimental results validated our theoretical findings and indicated that these methods can be extended to more general mappings f. paving the way for future investigations to generalize our theoretical results.
One current limitation of our approach is the deterministic aspect of the sequences considered in the paper. A possible extension of our work would be considering noisy dynamics the form _t-1 = f(x) E, where the E 's are i..d random variables. Another open problem concerns the ability of other sequential architectures such as RNNs or state-space models to learn in-context such autoregressive processes. These questions are left for future research.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The work M. Sander and G. Peyré was supported by the French government under the management of Agence Nationale de la Recherche as part of the "Investissements d' T'avenir" program, reference ANR-19-P31A-000 (PRAIRIE 31A Institute). The work of G. Peyré is supported by the European Research Council (ERC project WOLF). M. Sander thanks Francisco Andrade and Pierre Marion for fruitful discussions.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306. 002293 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.1566) 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1900 2020.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. In Forty-first International Conference Machine Learning, 2024.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04805, 2018.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv preprint 240i: 2 2024.
Takashi Furuya, Maarten V de Hoop, and Gabriel Peyré. Transformers are universal in-context learners. arXiv preprint arXiv:2408.01367 2024.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Jordan Hoffimann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint 1x.v:2310.00825, 2023.
MS Kaczmarz. Approximate solution of systems of linear equations, 1937.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns:: Fast autoregressive transformers with linear attention. International conference on machine learning, pp. 5156-5165. PMLR, 2020.
11

---

Published as a conference paper at ICLR 2025
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean- mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258, 2024.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Stanistaw Kwapieñ and Jan Mycielski. On the kaczmarz algorithm of approximation in infinite- dimensional spaces. Studia Mathematica, 1(148):75-86, 2001.
Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference Machine Learning, pp. 19563-19594, PMLR, 2023.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint wrXiv:2307.03576. 2023.
Pierre Marion, Yu-Han Wu, Michael E Sander, and Gérard Biau. Implicit regularization of deep residual networks towards neural odes. arXiv preprint XiXx:2309.002233. 2023.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in neural information processing systems, 32, 2019.
Swaroop Nath, Harshad Khadilkar, and Pushpak Bhattacharyya. Transformers are expressive, but are they expressive enough for regression? arXiv preprint 244.:::200.,54778, 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre- pre-training. 2018.
Ryszard Szware Rainis Haller. Kaczmarz algorithm in hilbert space. Studia Mathematica, 169(2): 123-132, 2005. URL http: / /eddml......mmm// , org doc//84
Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pp. 3515-3530. PMLR, 2022.
Michael E Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How do trans- formers perform in-context autoregressive learning? arXiv preprint arXiv:2402.05787, 2024.
Gemini Team, Rohan Anil, Sehastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk. Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.1 1805, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023a.
Johannes Von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agilera y Arcas, Max Vladymyrov, Razvan Pascanu, and João Sacramento. Uncovering mesa-optimizationn algorithms in transformers, 2023b.
Mingze Wang and Weinan E. Understanding the expressive power and mechanisms of transformer for sequence modeling, 2024. URL http:://aaaavvo/aa org/ wwwp/rgg//bbb///ee//../ abs/240/222222 00522.
Weinan. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 2017.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequencc functions? arXiv preprint arXiv: 2019, 2019,
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprint M:Xiv:2306009927 2023.
12

---

Published as a conference paper at ICLR 2025
A PROOFS
A.1 PROOF OF PROPOSITION
Proof. Let n 0 be an integer.
We define positional encodings Pt E R as
Pr (-1)'nd
We concatenate the input embeddings and positional encodings and define
x (c) E TR
We define the weight matrices (h(h) Q
((()) w for each head h = 1, as follows:
- =10xx 1], W 01xds -1], w = -O2dxd
Oax1
02dxd
02dx1
Odxd
0ax1
W( = =0 1], W [ixds 1]. w
Ia
0ax1
0axd
0ax1
With such constructions, one has
(wil) Te
wh xgt) -PaP and W x² 0,0).
(2) xg)) (W(?) 18, w
P.P: and W xF (0, 20,
The corresponding attention scores are
E e(-1)" Et (--1)))pp-(-1)ppppp
When h 1, one has that (-1)"PPPt (-1) st is maximal and strictly positive when 8 = t 1. Similarly, when 2, one has that (-1)* p.Pt = (-1)* 2 maximal and strictly positive when t. Therefore, Al.. 8 and A 8a
400.
We therefore consider the forward rule defined as, for 2 1:
(*OOT) ALe 0,0) + Σ Al. (0,
s=1
As - +00, one has F (xoor) + (It-1, 0) = F(ror)).
We then apply feedforward map on each token (2:-1, 0) defined as g(a,b,c) (1mm0s) ()) ())
La e) to obtain the desired augmented tokens ex. The feedforward map can be approximated with sigmoid based perceptron G" using the fact that T la=0, as n, + +00.
Denoting To the Transformer composed of the attention module F and the feedforward module g conludes the proof.
13

---

Published as a conference paper at ICLR 2025
A.2 PROOF OF PROPOSITION 2
Proof. Recall that the matrix A is defined as:
M1),, Ar,a K(Tt, or Eie k(ze)))) lsst if keep
if k kia
We have in matrix notations:
" = (11=((+n))")A-- diaggda))aaaa (2))
Importantly, when Ar.a k(x;, vz), for all 0 < M(ja)) one has (4 (4))" 7 Oxxt exponentially fast. In this case, we therefore have 17" k(ai.a))" Even more, when k(Zi,z,) the matrix (Ie 7) is nilpotent, and (I+ nA)" Otxt for m 2
k())..)) When Ar,s Ei., k(x,,a=)' for all 0 < 1) one has (Le nA)" + Otxt exponentially fast. In this case, we therefore have 17" = 2.
In both cases, when 1) 1f",, we have lim, u1s uia AA diag(A))
A.3 PROOF OF PROPOSITION 3
Proof. When 0 17",, convergence, we have:
wia = T2441 - A-1 diag( A2221111
Writing = A-1 diag A) and Ma == M(21,21]' we obtain the recursions Jc:
Σ (akk( = k(s,, wi))
k(x1, T1+1 = Wy(at), 11 2 1,
(11)
where the second equality comes from the fact that || = ||01||
We have it I++1 - k( Z1 P1)...
A.4 PROOF OF PROPOSITION 4
Proof. We have
W+1 W. + #411(1)41))
Right multiplying by (0((+1)) gives
2411) = x142 Wip(11+1)
Therefore,
W41 W,(I p(041)(((041))
((1441:2441) p((141)).
x442
Since
x142p(2411* = W
we obtain
WA1 = (W, - W) ((x(+1))(()+)) + W
k
Since W, = Wryff, follows that W, W = -PPPP ... P,, where P. = (I-v,1)) is the orthogonal projection onto the subspace orthogonal v..
14

---

Published as a conference paper at ICLR 2025
A.5 PROOF OF THEOREM 2
Proof. We first prove the second statement of the theorem, that is that W - W + 0 exponentially fast for almost all ø1 and W.
Let's first fix W and 21. In what follows we denote x == I1 to ease the notations.
Let A, W. W. We have
A41=A,W' xx W
Unrolling, we obtain
D A1(W(1 -xx ===)))===-((- xx T
because A W(I xx
For fixed W and æ, let y be an eigenvector of norm 1 of W(I - xw T. There exists such that
W(1 - xx Jy = Xy.
If [X] 1, then because W preserves the norm, it follows that
[ xx T >vll = || = 1.
Developing, we get
|, w/112 O.
But since
|| v>112 2 (x, v)²
we must have (x. 3) = O, and therefore Wy = Ay. so that is also an eigenvalue of W.
Therefore, to show that for almost all and W. p(W xrT)) T) < 1, it suffices to show that for almost all y V(W):
(x. い) + 0,
where V(W) denotes the eigen vectors of unit norms of W.
Let (O(d) be the subset of (d) of orthogonal matrices with distinct eigenvalues and let W DO(d).
Let '1,... Aa be the eigenvalues of W and bi...., ba be d corresponding eigenvectors. Let E V(W), and be the corresponding eigenvalue.
Writing y Ein yibi, one has on the one hand that
Wy = Ay
and on the other hand
Wy yr
1=1
Identifying the coefficients gives À = A₁ whenever Yi + O. Since the X; 's are all distinct, necessarily exactly one Yi non-zero, and y Cb..
Therefore, if (x. v) = 0, then x E + However, for almost all æ E Sd-1, (Cb.)) Therefore, for almost all W E DO(d) and r E gd-1
(x.
Since for almost all W E (d), W E DO(d), we conclude that p(W(T .x: )) < 1 for almost all Sd-1 and W Odd). As consequence, for almost all Sdd-1 and W O(d), W. W + 0 as t + +o0. It already implics that, for almost all æ and W, paa - 0. Therefore, Tr(140] x.44 > O. But since x, x = 1 and Tr(4y44) [|] it gives ia 0, and therefore ui -141 + for almost all Σ1 and W.
However, fully prove the first statement of the theorem, that is that, for all and W. ut Fi+1 0, we follow the exact same demonstration as for Theorem 3 (corresponding to the next proof in Appendix A.6), simply by replacing exp by id in equation 12.
15

---

Published as a conference paper at ICLR 2025
A.6 PROOF OF THEOREM 3
Proof.
As mentioned in the proof sketch, we consider the subset of H of linear combinations of the
v;:
ag for a, R T
This subset is a pre-Hilbert space under the inner product inherited from H. By completing this space with respect to the induced norm from H. we obtain a new Hilbert space H!'
We first show that P.P.-1.--. -..------ Pi + strongly in H', using a characterization from Kwapień & Mycielski (2001).
A sequence (Ve)a21 of unit vector for which P.Pt-1-- ... P1 + 0 strongly is referred to as effective.
The specificities of our case are twofold: first, Va H', which is potentially of infinite dimension, and second, the vectors v, follow an autoregressive relation. In our case, we can still show that the sequence (v5)221 is effective in H'.
Note that because Ω E O(d), one has for any positive integers t, s, that (Va+r; Metr) = (V v/). Such sequence is called stationary.
Bochner's theorem states that there is a measure on the unit circle S called spectral measure-such that, for all 1 1,
az = (1)+1, = / z' dz)..
We are going to use the following characterization from Kwapień & Mycielski (2001); Rainis Haller (2005).
Theorem (Effectiveness of stationary sequences (Kwapień & Mycielski, 2001)). A stationary sequence of unit vectors which is linearly dense in a Hilbert space is effective if and only if its spectral measure either coincides with the normalized Lebesgue measure is singular with respect to the Lebesgue measure.
Now, let Ω be orthogonal matrix. One has
1
at (Vt+1, v1) k( k( ((E+41,
exp ((x1, 1'xx)))
k(x1,
Since Ω an orthogonal matrix, it can be diagonalized via rotations with angles 0,, 02, We therefore can write (x1, 25x1) as
.., Op-
1 (di + cos(O,1) +
m
Fi
x
Tie
1=1
22++
i=m+1
for some p and m. Therefore, one can write
aj = g(t( ... ,0)))
where is defined as
glyi.... Up) = exp -1
( 12)--) + 24(2)) cos(y)) + Σ
Σ
2
Σ
Ti
(12)
1=1
1=p+1
i=
The function g is periodic each variable ys with period 2n. Since g is also C²,, this allows us to represent g Fourier series:
g(v) Σ (ef(u.))
KEZP
Therefore,
it0,))
0t
kEEP
16

---

Published as a conference paper at ICLR 2025
Therefore, the spectral measure satisfies
(0,))
REBP
Note that this necessarily implies that Cg 2 by Bochner's theorem. We now going prove that is singular with respect to the Lebesgue measure.
We recall that two positive measures and 3 defined on measurable space C' are singular if there exist two disjoint measurable sets A, B such that A U B = C, such that is zero on all measurable subsets of B, while zero on all measurable subsets of A. Here, a a, and / is the Lebesgue measure on C = S1
We have
UkEz (0)
is measurable and the Lebesgue measure zero on A, while by definition, a is zero on all measurable subsets of B = 5 A.
Therefore, a is singular with respect to the Lebesgue measure, which, using Theorem from Kwapieñ & Mycielski (2001) in the Hilbert space H' shows that the sequence (Va)221 is effective in H'.
Now, we have, by taking the adjoint of Wt W:
W W P..".
We recall that we identify W with Wwe so that W Ht' Rdd and W Rd + H.. One has for all E Rdd
(Wr W' Σ (( W Pe PWW* +
because T x E H'.
Therefore, for all E Rd. ((xx) 2) - 0 in H.. Because [1x(xx) k(x1, > O, we
necessarily have (2),x) + for aorall E Rdd This is equivalent to 0, and therefore ut T++1
A.7 PROOF OF THEOREM 4
Proof. Let tp denotes the period of (2z) We define
II = ......py.
One has supvewt [|]]== |||b||| < (see Rainis Haller (2005), page 2). Therefore, 11Th - 0 as m - +oc.
As such, for any such that (p(x) H'' Wep(2) + Wip(z) exponentially fast.
In particular, for x = I1, we get that
Euok(os,
converges exponentially fast. Since k(x,) > exp(-1), it implies that + 0 exponentially fast.
A.8 PROOF OF PROPOSITION 5
Proof. Building positional encodings and parameters.
Similarly to the previous proof, we consider positional encodings pi and pi E R as
Pt 6==1, Dr 1.
17

---

Published as a conference paper at ICLR 2025
We concatenate the input embeddings and positional encodings and define
Ct.p (2 P. 2z, pi: Ft, ut) (z1-1. 21, 1, F2, ue) R44+2
if 1
and
1 = (0a, pl 21, P1. Oa, u1) = (Oa, 1, 1, 1, Oa. u1). (h) (h)
We define the weight matrices (h( K with for each head / 1, as follows:
Q
w(
10axd+1: La 0ax2441], and
0+++xd+1
0441xd+1
Oa+1xd
Ostixd
w
0d+1xd+1
0441xd+1
Oatixd
0at1xd
Odxd+1
Odxd+1
Odxd
Odxd
Oaxd+1
0axd+1
Oaxa
La
w:(2)
(0++1x4
Ta+1,
02+1x2a],
[T 02+1xd+1, 04+1x2a), W
0041xd+1
0441xd+1
Oa+1xd
Odtixd
0s41xd41
0441xd+1
Oa+1xd
Oa+1xd
9
04xd41
Oaxd41
Odxd
Oaxa
0axd41
Oaxd+1
Ia
Oaxd
With such constructions, one has
(We em
Wh (s.)) (1)) and W Cs.p -1y(0+1,, 0a+1; Od; Uy).
1 (w(() px K eap)
(x1, + p&pi = (Ee, + 8am1 and w es.p
n(0011, O441, Og, I.) if > 1 and W e1 = (0a+1, O2+1, O2, Oz).
Implementing equation 7 when Ar,, 2(
Let T one-layer, two-head Transformer with the parameters above (with N = id of N = =exp). Its forward rule is defined as, starting from ei Ct.p
k+I ctp
Σ Al,400411 Oc41, Oa. ( +
Ah, Od41, O2,
#=1
s=1
Therefore, the (3d 2) first coordinates of et,p are not modified and the d last coordinates uf are updated as
k+I = u 17 Σ k( k( x I.) +
Σ
k(Ie, 13-1)20
(13)
x=2
The second term in equation 13 reads
Σ
2,),..
2=2
Therefore, equation 13 reads
k+1 u
(( loceF.41),
u
17
which corresponds exactly to equation 7.
Implementing equation 7 when A Ex_, k(
Let T be a N = softmax-based, one-layer, two-head Transformer with the parameters above. Its forward rule is defined
A+1 t.p
Σ (0 Oc41, Oa, ut) + Σ (0a+ O4+1, 0a, x.).
E.p
18

---

Published as a conference paper at ICLR 2025
Therefore, the (3d + 2) first coordinates of Ctp not modified and the d last coordinates un updated as
e (
+ ( Io.
k+1 4 ut
(14)
elasat,)
E+=1
elzvee 11t8,-3 11
==2
Because ||,||| = for all 2 and
0, one has
Σ clanat))
that the second term in equation 14 reads
(
t-1 eleeRte)
+6 Is
2 ((te)at)))
1-= e(ee,))) Ts41.
Therefore, equation 14 reads
==1 2 eloine)) (u, Loc1#o+1);
which corresponds exactly to equation 7.
A.9 PROOF OF THEOREM 1
Proof. We first use Proposition 5. For any 0 < < n7,, where ry" of defined in Proposition and for each configuration of equation 7, there exists an attention-only, one-layer, two-head causal Transformer T such that, for any autoregressive sequence (2x)(21 generated according to Assumption defining
cx = (II-1, 0,, Ft, ut
for > 1 and
ej = (0a, 1, v, 1, Oa, ut),
we have that, for any n, e solves equation if and only utir solves equation Thanks to Proposition 2, we know that u²" too ut, and define M(x1=) == ug- Then:
When Ar,, = (1r, v.).. taking V id in the construction T in the proof of Proposition 5, and under instance (1), we have M(F1.)) It+1 O, with exponential speed for almost all 21 and W, thanks to Theorem 2.
When A,, = e( (resp. Ar, Σ e(r)))))) ) taking V = exp (resp. / = softmax) in the construction of in the proof of Proposition 5, and under instance (2), we have M(r..t) - F+41 + 0 thanks to Theorem 3.
mg)
When Ar,. ( (resp. Ar,. Σtzi e(R)))))) taking V exp (resp. A softmax) in the construction of T in the proof of Proposition 5, and under instance (3), we have M(x P141 + 0, with exponential speed, thanks to Theorem
The last statement of Theorem follows from Proposition 2. When V' - id N exp, by choosing k(z1.=))> we have that M = u² 2j = M((11:) whenever 2 t.
When N = softmax however, choosing for instance
1 in the proof of Proposition A.2, we get -i - ui = (4: u Since p(4, A) < - 4. is reasonable to conjecture that
|lup uill is order (1 - 1)". In this case, is sufficient to impose 1 - 0 as t, + +o0 to have
E (n, = |!u6" uill - 0 as n, 400.
19

---

Published as a conference paper at ICLR 2025
B ADDITIONAL RESULTS
Dual interpretation. We provide a dual interpretation for μ defined in Proposition 3. We conve- niently write equation minw ( - Il where for μ in RdxT-1, (W) = (W ((x))) and + (4) = 2t (e)( T Gradient flow on E reads:
W W--VE(W) - (Φ(W) (
At optimality, has W (u) for some R ("kerne trick"), so that consider the energy F(u) = E (1)). Gradient flow on reads: j = - #2:7+1| where := Q T
[,, Ni) (ss, Since K is positive, instead consider the equivalent flow
[K 12.4+1)
Making this flow causal by replacing K with its masked counterpart A recovers to the results of the previous paragraph, since, at convergence, l satisfies equation
20


Reviews:
[{"review_id": "8hkyBnibHF", "reviewer": "Reviewer_gDtN", "summary": "This paper extends a line of work showing what transformers are capable of learning. In particular, it shows that causal transformers are capable of learning kernelized linear functions, in context, which means that each token is the same linear function of the previous token. Previous work had focused on linear attention whereas this work extends to softmax attention. It also handles period sequences.", "questions": "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)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper furthers our understanding of why transformers, like the amazing LLMs, are able to learn in context. It extends to learning linear functions with softmax attention. It provides nontrivial, rigorous guarantees as well as experimental evidence corroborating the theory.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "QNlMOgrFlQ", "reviewer": "Reviewer_Se7F", "summary": "This paper considers the capacity of causal Transformers to the prediction of next token $x_{t+1}$ given autoregressive sequence $(x_1,\\ldots,x_{t})$ as a prompt where $x_{t+1}=f(x_t)$. They explicitly construct a Transformer that learns the mapping $f$ in-context through a causal kernel descent method (which connects to the Kaczmarz algorithm) and prove long scope guarantees.", "questions": "The questions and suggestions are given by the weakness section. The reviewer believes that this work requires significant improvement to reach his/her expectation and doesn't seem to be possibly addressed within the time frame of ICLR 2025 rebuttal phase.\n\nHowever, if the author is able to resolve at least (1) or (2) in the question section, the reviewer might be able to improve his/her opinion.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The strength comes mainly from the problem setup of studying next-token prediction, which seems to be quite an important task. The proof looks sound. The presentation of the work is clear. The authors perform empirical evaluation for their theoretical results. The discussion of theoretical results are provided in the work.", "weaknesses": "The major concern over this work is listed as follows.\n\n(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. 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.\n\n(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.\n\n(3) The limitations of this work is not properly discussed at the conclusion part.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "EKWAxMmgmA", "reviewer": "Reviewer_LWyw", "summary": "The paper studies the approximation abilities of causal transformers when predicting the next token in-context for autoregressive sequences. In particular, it focuses on specific instances where the context-dependent map $f$ determining the next token $x_{t+1} = f(x_t)$ in the sequence belongs to the RKHS of a given kernel $k$ and is either linear or the sequence is periodic.  The authors introduce the causal kernel descent method and theoretically analyze its convergence properties. They provide a construction of the transformers layers with linear, exponential and softamx attention that can implement this method.", "questions": "1. Could the authors clarify how their results using particular kernel choices and corresponding attention mechanisms contribute to understanding the universality of transformers in learning sequence-to-sequence functions and what they mean by universality in this context?\n3. In Definition 2 could the authors clarify how is the periodicity enforced in instance (3). \n4. In Proposition 1, it’s unclear what $n$ refers to, as elsewhere it denotes the number of layers, but here it relates only to the first layer that computes augmented tokens. In Appendix A.1, $n$ seems tied to the specific positional encoding. Could the authors clarify this?\n5. Is it correct that in the presented model, the first layer computing the augmented tokens always uses softmax attention and does not have a skip connection, whereas the identical layers can employ different normalizations (softmax, linear, exponential) and include skip connections?\n6.  In the paragraph: Augmented tokens,  could the authors clarify why $e_1^0 := (0_d,1,x_t,1,0_d,0_d)$ rather than $(0_d,1,0_d,1,x_t,0_d)$ based on the feed-forward map in appendix A.1.\n7. (Minor) In Remark 1, there may be a typo in the definition of $\\epsilon_1$.\n8. In Equation 8 could the authors provide more insights regarding the interpretation of $\\mu$ and its relation with the data-generating process?\n9. For instance (4) described in section 5, did the authors try to train a model, and if so, what were the results?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written. \n- The problem is conceptually well-motivated.\n- The proposed method and analysis are clearly explained with all necessary details.\n- The manuscript includes both rigorous proofs and empirical experiments that validate the findings.\n- The proof technique and its connection to the Kaczmarz algorithm in Theorem 3 are insightful.", "weaknesses": "1. The choice of normalization in the attention appears to be specifically tailored to match the kernel used in the data-generating process (see Definition 2 and Theorem 1). This design decision raises questions about the generalizability of the findings, making it unclear how these results might extend to broader settings.\n2. Since the analysis is confined to the class of functions determined by the chosen kernel, it is not clear how the results illustrated in the paper contribute towards understanding the universality of transformer models.\n4. The theoretical results presented in the paper rely on taking the limits as the context length $t$ and the number of layers $n$ approach infinity. However,  the empirical experiments reported in Figure 4, demonstrate that Transformers with a finite number of layers, trained using the Adam optimizer, can achieve better performance than the proposed infinite-layer model. This discrepancy suggests that finite-sized trained Transformers may implement a different algorithm than the causal kernel descent method.\n5. Exponential convergence for the case of exponential kernel and exponential or softamax attention requires periodic sequences.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ADBC8frGfT", "reviewer": "Reviewer_6GYT", "summary": "This paper theoretically studies the universality of Transformers for next-token prediction. The authors specifically consider the sequences of the form $x_{t+1} = f(x_t)$. In particular, they theoretically analyze two types of instances:  for linear $f$ and periodic sequences (exponential kernel). They construct augmented tokens and show that an explicitly constructed Transformer can learn in context to accurately predict the next token asymptotically, through a causal kernel descent method.", "questions": "1. 250 typo $\\epsilon_1(n,t) = \\mathcal{M}^n(x_{1:t}) - \\mathcal{M}(x:1,t)$?\n\n2. In the experiment section, how did you compare fine-tuning $\\mathcal{M}^n$ with infinitely deep model? Did you just take $\\mathcal{M}$ to be the ground truth? (Figure 5 is missing)", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Very mathematically sound and interesting. The universality of this next-token predictability is a strong result. \n\n2. Causal kernel descent is a very intriguing framework to analyze sequence models in context prediction, taking the important aspect of masking out the future into account.", "weaknesses": "1. 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 cases seem 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.\n\n2. 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. \n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "h5qNoMnn5R", "reviewer": "Reviewer_mAs7", "summary": "The paper aims to advance our understanding of causal Transformers for next-token prediction in autoregressive settings by showing how to construct networks that can solve the next-token prediction in-context learning task.", "questions": "- L178: why is your setup better reflective of how LLMs are trained? Please elaborate.\n- L194: Why does the unit-norm of $x$ imply that the kernel values for $k(x_{i},x_{i})$ are same for all $i$?\n- Eq (3): why do you need to model the projection (in this exact way of picking the last $d$ coordinates)?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "**Despite not being very familiar with kernel methods and learning theory (but somewhat familiar with Transformers and in-context learning), I believe that the authors' contributions are relevant and interesting, though I cannot assess their (learning theoretical) details.**\n\nParticular, I find the paper is strong in the following aspects:\n- crisp and clear setup in the abstract\n- intuitive Fig 1 (for the setting, not for the results)\n- very well summarized contributions list\n- the results seem to be strong (I am uncertain about how realistic the setting is)\n\nNonetheless, as I will detail below, **I think the presentation needs to be significantly improved**.", "weaknesses": "### Major points\n- the paper is very dense, even though the authors could have used an additional 10th page. I'd suggest more explanations for the result, and less formulas in the main text\n- the intuition seems to be missing (for me) as to why you need the augmented tokens. What do they mean?\n- 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).\n- **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 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\n- What is the rationale behind the construction of $f$ in the _\"More Complex Iterations\"_ paragraph?\n\n### Minor points\n- please use equation numbers for the main equations \n- please define all quantities before you use them\n- Definition 1: what is $S^{d-1}$?\n- Definition 2: what is $O(\\cdot)$?\n- Definitions 1 and 2 would be better suited as assumptions\n- L206: please explain what $e$ is, what the indices stand for, and why you add multiple beginning-of-sequence tokens and tokens of \"1\".\n- Proposition 1: what is $\\mathcal{N}?$\n- L246: what makes these constructions explicit? As far as I understand, you specify the dimensions, but that leaves many degrees of freedom (ie, all the elements can be of any value)?\n\t- From Prop. 5 and checking A.8, I get a sense what these matrices would be. Thus, as you have a contruction, please refer to this fact to the reader, otherwise, calling a matrix explicit without saying that you can calculate all the values can be misleading.\n- L286: do you mean the update of $u_{t}$ depends on $x_{1:t}?$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)