Learning In-context $n$-grams with Transformers: Sub-$n$-grams Are Near-Stationary Points
Authors: Aditya Varre
Gizem Yüce
Nicolas Flammarion
Venue: ICML 2025 poster (2025)
Keywords: in context learning
markov chains
transformers
n-grams
Primary area: deep_learning->theory
TL;DR: For the next token population cross entropy loss over in-context $n$-grams,  $k$-grams ( for $k \leq n$ ) are near stationary points.

Abstract: Motivated by empirical observations of prolonged plateaus and stage-wise progression during training, we investigate the loss landscape of transformer models trained on in-context next-token prediction tasks. In particular, we focus on learning in-context $n$-gram language models under cross-entropy loss, and establish a sufficient condition for parameter configurations to be stationary points. We then construct a set of parameter configurations for a simplified transformer model that represent $k$-gram estimators (for $k \leq n$), and show that the gradient of the population loss at these solutions vanishes in the limit of infinite sequence length and parameter norm. This reveals a key property of the loss landscape: {sub-$n$-grams are near-stationary points of the population cross-entropy loss}, offering theoretical insight into widely observed phenomena such as stage-wise learning dynamics and emergent phase transitions. These insights are further supported by numerical experiments that illustrate the learning dynamics of $n$-grams, characterized by discrete transitions between near-stationary solutions.

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Aditya Varre" 1
Gizem Yüce
Nicolas Flammarion
Abstract
behaviors with plateaus such grokking (Power et al., 2022) and stage-wise transitions (Olsson et al., 2022) where the training process frequently lingers in regions of slow progress before going through a sudden phase transition after which ICL abilities are acquired. Mechanistic inter- pretability studies suggest that after these plateaus, specific circuits, such induction heads (Olsson et 2022), syn- tactic structures (Chen et al., 2024a), are gradually learned. This raises a natural question:
Motivated by empirical observations of prolonged plateaus and stage-wise progression during train-
ing, we investigate the loss landscape of trans- former models trained on in-context next-token prediction tasks. In particular, we focus on learn- ing in-context n-gram language models under cross-entropy loss, and establish sufficient con- dition for parameter configurations to be station- ary points. We then construct a set of parameter configurations for a simplified transformer model that represent k- k-gram estimators (for le & n), and show that the gradient of the population loss these solutions vanishes in the limit of infinite sequence length and parameter norm. This re- veals a key property of the loss landscape: sub-n- grams are near-stationary points of the population cross-entropy loss, offering theoretical insight into widely observed phenomena such as stage- wise learning dynamics and emergent phase tran- sitions. These insights are further supported by numerical experiments that illustrate the learning dynamics of n-grams, characterized by discrete transitions between near-stationary solutions.
Why does training linger at plateaus before developing such abilities?
The setting of Edelman et al. (2024) provides an avenue to take closer look at this question with the specialized task of learning in-context n- o-grams (Shannon, 1948; Chom- sky, 1956; Brown et al., 1992). Here, the training unfolds in a hierarchical, phase- manner. It begins with the transformer making uniform predictions, progresses through unigram and bigram predictions, and potentially general- izes to higher-order n-grams. Figure/ 1 illustrates how these training phases overlap with the losses of the sub-n-gram estimators.
Building on this observation, we provide a theoretical foun- dation for why training lingers long plateaus -the loss landscape of transformers trained on in-context sequential data exhibit stationary points aligned with sub-hierarchical or sub-syntactic solutions. These points act as intermedi- ate solutions where gradients vanish, causing training to stagnate before transitioning to the next hierarchical level.
1. Introduction
Transformers (Vaswani et al., 2017) have become central to modern machine learning, due to their capabilities such as in-context learning (ICL) (Brown, 2020) the ability of models to perform new tasks by leveraging a few examples provided within the context, without the need for parameter updates or retraining.
Concretely, this paper explores the loss landscape of trans- former models trained on in-context n- n-gram language mod- els predict the next token. We show a sufficient stationar- ity condition for the cross-entropy loss that the sub-n-gram constructions satisfy, shedding light on the incremental and phase-wise learning phenomena observed during training for in-context n-grams. Our main contributions can be sum- marized as follows:
Recent empirical studies have revealed that the dynamics that result in in-context learning abilities often deviate from simple monotonic decrease in loss, exhibiting complex
Equal contribution 'Theory of Machine Learning Lab, EPFL, Switzerland. Correspondence to: Aditya Varre <aditya. cptl.hh>,, Gizem Yüce <gizem.yuce@eppt.l.
In Section 3, we provide sufficient condition for the solutions to be the stationary points of the cross-entropy loss in the next-token prediction task when model's score function depends solely on a sub-sequence of the input history. This characterization provides a powerful tool for
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
2.0
the landscape of transformers on data from binary first- order Markov chains (bigrams) but does not consider the in-context setting. Bietti et al. (2024) investigates the forma- tion of induction heads needed to learn bigrams with specific trigger tokens. Additionally, Nichani et al. (2024) studies the formation of induction heads through gradient descent in learning causal structures. A closely related work to ours is Edelman et al. (2024). They report the stage-wise dy- namics of transformers while learning in-context n-grams, and identify that these stages correspond to sub-n-grams. However, their theoretical analysis is limited to one step of gradient descent on binary bigrams, while we characterize the sub-n-grams as near-stationary points for unrestricted n. and vocabulary size, explaining the long plateaus once the training reaches a state expressing these solutions. Chen et al. (2024b) study the same in-context n-gram prediction task, but with an architecture that involves a feed-forwart network layer that does the token selection. However, they have a threee-stage training procedure and an initializa- tion scheme that ensures different heads attend to different tokens from the start, climinating the stage-wise learning dynamics we study.
Unigram
Bigram
0.9
Trigram
Test Loss
4-gram
0.7
0.6
100
200
300
.400
500
600
700
800
Step
Figure 1: The stage-wise behavior of the test loss during training for the 4-gram language model. The dashed lines represent the cross entropy loss of the k-gram estimators for k == 1, 2, 3, 4. The plateaus of the test loss overlap with the losses of k-gram estimators.
analyzing the structure of the loss landscape for various tasks, including but not limited to n-grams.
In Section 4, we demonstrate that a set of solutions repre- senting sub-n-gram estimators near-stationary points, by verifying that they satisfy the conditions sufficient for the gradient of the population loss to converge to zero in the infinite weight and sequence length limit.
Sequential learning. Fukumizu & Amari (2000) analyzes the plateaus in the loss curve and their relationship to crit- ical points for supervised learning with neural networks. The characterization of dynamics, including jumps between these stationary points, has also been studied for simpler models such as matrix and tensor factorization (Razin et al., 2021; Jiang et al., 2022), matrix sensing (Arora et al., 2019; Li et al., 2021; Jin al., 2023), diagonal networks (Gissin et al., 2020; Berthier, 2022; Pesme & Flammarion, 2023), linear networks (Saxe et al., 2019; Gidel et al., 2019; Jacot et al., 2021; Varre et al., 2023), ReLU networks (Boursier et al., 2022; Abbe et al., 2023) and transformers with diago- nal weight matrices (Boix-Adsera et al., 2023).
In Section 5, we present empirical evidence that illus- trates the structural evolution of these models during training and how transitions between training phases are in alignment with the predictions of our theory.
1.1. Related Work
In-context Learning. To understand in in-context learning (ICL) (Brown, 2020), previous works have explored various approaches. One approach is mechanistic interpretability, which has revealed the emergence of circuits called induc- tion heads during ICL (Olsson et al., 2022). Another ap- proach involves studying ICL on specific hypothesis classes to understand how transformers solve these tasks in context. A common feature of the training dynamics in these studies is the presence of plateaus (Chen et al., 2024a; Kim et al., 2024), after which the models acquire certain capabilities. Examples include studying ICL on regression tasks (Garg et al., 2022; Von Oswald al., 2023; Ahn al., 2024), boolean functions (Bhattamishra et al., 2023), regular lan- guages (Akyürek et al., 2024), and n-grams.
Large Language Models(LLMs) as n-grams Nguyen (2024) investigates whether LLM predictions can be approx- imated by simpler, interpretable statistical rules based on n-gram frequencies, and shows that LLMs exhibit curricu- lum learning during maaining-starting with simpler n n-gram patterns and progressively capturing more complex ones. In related vein, Zekri et al. (2025) leverages the equivalence between auto-regressive models and n-gram models with long contexts to derive generalization bounds.
(In-context) Learning n-gram language models. n-gram language models, or higher-order Markov chains, are ef- fective mathematical models for generating sequential data and capturing certain aspects of natural language (Shannon, 1948; Jelinek, 1998; Jurafsky & Martin, 2009). As a re- sult, many studies have focused on analyzing transformers through the lens of Markov sequences. Svete & Cotterell (2024) examines the representational limits of transform- ers on n-grams, while Rajaraman et al. (2024) focuses on the in-context counterpart. Makkuva et al. (2024) explores
2. Problem Setting
In this section, we define the in-context next token predic- tion loss, the landscape of which we aim to understand for transformers. Next, we introduce the n- n-gram language mod- els and formally describe the disentangled attention-only transformer architecture.
2

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Notation. Let [S] (1.2.... be finite alphabet for S E N. The Kleene star [s]* denotes the set of all finite- length sequences whose elements are in [S]. For / & k, of = f 1, E [S]k-/-| denotes subse- quence, and when not specified, I = 1. For an element x; in sequence, its (- /histoory refers to -;_ and its (-f)-token refers to _._ A k denotes the probability simplex over R A language model is distribution over [S]=. We use ef denote the hh elementary basis vector of Rdd simpli- fied to e when d = S. For matrix E I r[[] E Rdd denotes the vector representation of it's 1th row.
these assumptions restrict their expressive power, they re- tain certain key characteristics of natural language, such as causal dependence the prior context. Under these assumptions, the probability of the sequence x" can be expressed using the chain rule as
Pr (+" = II
Pr (2) /x'
))=
II Pr (11)))))))))) LE[7]
LE[T]
Estimators. For the next token prediction task with the in-context n-gram language model, the k-gram estimator of particular interest.
2.1. In-context Next Token Predictions
Consider language model p and parametric model S-1, A which predicts the probability of the next token. Given sequence x" E [S]", the perfor- mance of the model in predicting the next token is evaluated using the cross-entropy loss (CE):
Definition 2.1 (k-gram estimator). Given a sequence x" and i [S], the k-gram estimator is defined as
p()..): [S]] +
Σ Mat = - =
e (0,xx") = Σ Ep(E -00t,|=s|0"" log (p(0.>")))) (s))
Dk (i)===
#ES
1a k11 k
In the in-context setting, we assume that the ground truth language model is a mixture of multiple language models. Specifically, for every sequence, a language model pr is sampled from prior distribution P. Then, the sequence x" sampled from? pr.
Intuitively, this estimator checks whether the (k-1): histories match and counts the tokens that follow. For k = 1 (unigram), computes the empirical frequency of each to- ken in the sequence. For k 2 (bigram), the estimator computes the empirical frequency of tokens that follow those matching the 7 token in the sequence. The n-gram estimator is the "in-context" maximum likelihood estimator (MLE) for the n-gram language model. Moreover, Han et al. (2021) shows that the smoothed version of this estimator achieves the minimax optimal rate for the next-token prob- ability estimation. We include all the k-gram estimators where k < n, in the definition of sub-n-grams.
The in-context population loss is defined as [())
= JE
EE L (0,xx") x² ~pp
P-~P
We focus on the task of learning in contentex n-grams where Pr 's are modeled as n-gram language models.
2.2. In-context n-grams task
We start with the definition of n-gram language model and discuss some estimators for this setting.
We note that our choice of the n-gram language model and k-gram estimator is primarily for ease of presentation. The results this paper naturally extend to more general time- homo--oomgeneuuss causal dependencies in the sequence heyond n-grams in Nichani et al. (2024). For exam- ple, our results also apply to causal graphs where token 2, depends only on specific parent tokens, such as 2;-2 and øı-1. Similarly, while the k-gram estimator matches a contiguous (k- liiitoy by definition, a similar estimator can match non-contiguous histories and count what follows. For instance, let the ()) of token z, refer to (Ix-1, xt-3); an estimator could leverage this pattern to pre- dict subsequent tokens in the same way the 3-gram estimator does for the contiguous 2-history. We use the term sub-n- grams also to include such cases, which we discuss further in Appendix F.2. Next, we introduce the attention-only transformer architecture used for this in-context learning task.
The n-gram Language Model. A language model pr an n-gram language model if it satisfies the two key assumptions:
(a) Markov property: The conditional probability of token depends only on the (n-1)-history rather than the entire history, i.e.
Pr (x, P- (x.)= for l >
(b) Time Homogeneity: The transition probabilities are inde- pendent of the position in the sequence. Formally, for all Ł, I E [T] and sequences s" E [S]".
Pr (z.=s,|aa - = Pr (z,=s,12,))
The assumptions above distinguish n-grams from general language models by limiting the dependency range and as- suming uniform transition dynamics over time. Although
3

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
2.3. Disentangled Transformer Architecture
Embeddings. The token embedding S [S] + Rd is orthogonal, which means that the set (St)aE[s) forms an orthogonal family in R Here, 8; denotes the embedding of the token i. For such embedding exist, d t S. For any sequence x², the input is encoded as,
Given an input sequence r'=[S]', the transformer first mmps each token
d-dimensional embedding using token-wise semantic and positional embeddings, defined as E: [S] + Rd and P: [T] - IE The core of the trans- former is the self-attention mechanism, which allows the model to weigh different parts of the input sequence based on learned similarity scores. Given a sequence embedding r E Rxxxd self attentinn computes a weighted sum of token embeddings as
To = [[
Sigg
E Txd
No explicit positional embeddings are used'.
First Attention Layer. The first attention layer contains m attention heads, and each attention matrix is parame- terized by a single learned matrix A (h) RTxT and is independent of the input embeddings. The output of each head is therefore given by
SA (r) =
where (Q, K, V), are the query, key and value matrices of the attention head, and the masked soft-max a (-) is defined as
r r(h)) a (A!")) ro (V!*)
R2xd
(2)
explate) Eeg. euplaie)
jEi
(z) [bb
This construction is equivalent to standard attention head where token embeddings are concatenated with one-hot positional embeddings and the attention only relies on the latter. Thus, the output of the first layer is given by
0,
olherwise
Traditionally, the outputs of self-attention layers added to the residual stream. For ease of interpretation and analy- sis, we instead consider disentangled architecture in which the outputs of individual heads and layers are concatenated rather than summed. This framework was proposed by Fried- man et al. (2023) and formalized by Nichani et al. (2024). The disentangled architecture retains the same representa- tional power as the standard formulation (see Theorem 3 of Nichani et al. (2024)) and is formally defined below.
(0)
(I)
(m)
Tx (m+
r1 = ri ri
where (0) To is the skip connection.
Second Attention Layer. The second attention layer con- tains a single attention head, where the value matrix is fixed to V [Ta: 0z: Oal :.ll:.:))((nn)))) (m that reads the first
Definition 2.2 (Disentangled Attention-only Trans- former). Let T, be the depth, {he)eE(z) be the number of heads per layer, d be the embedding dimension, de be the dimension of layer f, da be the hidden dimension of the model parameters, and dows be the output dimen- sion. For the tt head in the 9t layer, let Q K. E
block. Consequently r, v! = ri = =r and the output of the second layer is given by
T2 = o (r, Q K2r,)) ro E Txd
(4)
We note that there no concatenation to the residual stream in the second layer, which corresponds to not using a resid- ual connection in the standard transformer.
(h)
Rdda,, vi()) IRde be the query, key, and value ma-
trices of the respective head and layer, let SA (h) (-) = SA, "o"" wilhy (-) and let U Rdo. be the un- embedding matrix. Given an input sequence x", the dis- entangled transformer outputs the logits TF(0) for 0 {19, (h) ((A))
Finally, the unembedding matrix is given by
S
U
Ejs
(5)
v (h) (re(t))M((Ma) U given by,
where (e,)'s are canonical basis of R Note that if the token embedding S is the one-hot encoding, then U is the identity matrix. Finally, given sequence x², ,, the probability of the next token estimated by the model, denoted by pe po(x") is
To = [E(x") ) P(=*)] E Xx
r. = [r-.., SA (1)
SA (b) (r.-.)
TF(0) = T.
Po Pa(=²) U [T).
To simplify the presentation of our theoretical results, we consider two-layer simplified disentangled transformer with specific design choices. These modifications, detailed below, include orthogonal token embeddings and a fixed value matrix in the second layer.
0 (A (b)
(1) Jam (K2, Q2) denotes the set of parameters of our simplified disentangled model.
Positional information is implicitly used through the attention matrix the first layer.

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
3. A Sufficient Stationary Condition for Population CE on Sequences
We leverage this proposition formally show that the pop- ulation loss gradient vanishes when the model correctly learns the true k-gram conditional probabilities. However, the proposition is not specific to the setting of k-grams and holds generally for cross-entropy loss in next-token predic- tion tasks. A generalization of this result, which applies to models that depend on any arbitrary subset of tokens (not just a contiguous suffix), discussed in Appendix C.2.
In this section, we analyze the gradient of the population next-token cross-entropy loss for sequential data. This analy- sis provides fundamental insights into the training dynamics. We begin by presenting a lemma that provides a closed-form expression for the gradient.
Lemma 3.1. Consider any parametric model po(-) (S A that maps a sequence of states to a proba- bility vector on the states. The derivative of the population cross-entropy loss [(0) with respect to parameter 0; I
4. Theoretical Insights into Stage-wise Dynamics Through the Loss Landscape
is
In this section, we first provide the representations of the sub-n-grams with the simplified disentangled transformer architecture discussed in Subsection 2.3. We then prove that these sub-n-grae --rammmmmmmmmmm constructions are near-stationary points of the loss.
0g, [())==
E EE P. Pa~Px" upr
(po(x")= (112"), , de, log logpo(a")).
where 0e is the partial derivative with respect to 0g and log(.) denotes component-wise logarithm.
4.1. Representing Sub-n-grame with Simplified Transformer
The lemma 3. presented for fixed length T for notational simplicity, generalizes to sequences of arbitrary length. This result extends prior analyses such as Bietti et al. (2024, Lemma 1) and Makkuva et al. (2024).
We now construct a disentangled transformer to represent the n-gram estimator in Definition 2.1. We then extend this construction to represent any k-gram estimator for [n]] by deactivating specific heads. Parameters are often expressed as sums of outer products of orthogonal vectors, emphasizing their role associative memories in the spirit of Bietti et al. (2024).
The gradient expression in Lemma 3.1 is the expectation of an inner product between two key terms: (a) the prediction residual, po(x") - (-I x²"" representing the error between the model's s estimate and the true next-token distribution; and (b) the score function, which is the gradient of the model's logits. In the following proposition, we demonstrate that when the score function's ' structure depends only on a specific sub-sequence of the input, the population loss gradient simplifies significantly.
Transformer representing n--gram. Consider the simpli- fied transformer model presented in Section 2.3 with (n-1) attention heads. 2 For the 11th head in the first layer, we set the parameters as follows:
E. such that Oo=o. log p=²"")== g (?r, ie., the score is solely function of the context Py and the last T -t+11 elements of the sequence a², the gradient of the population loss L can be written as
Proposition 3.2. For any 0_
T-1 Σ ef
h-1 Σ e
A
(h)
(ET
+c
(e
L=h
1=0
V (
Σ
858}
VC(0))=
E De Perp xTupT mpp
E
(x-) (x (.1:) (Pr.
The second layer query and key matrices are assigned such that
Furthermore, if for such 0. E RR, the model estimates the conditional probability of the next token Pr (-) x²)) i.e. Po. (==) (-] x²)) almost surely for Pr ~ P, then O, a stationary point.
S 8-
(Q2)
K2 =
(T))T
(6)
j=1
Proposition 3.2 provides a sufficient condition for stationar- ity. It reveals that if the score function g depends only on the context and suffix 2!,, then the gradient is no longer driven by the error against the true probability conditioned on the full sequence, Pr (-|æ"), but rather by the error against the probability conditioned on the shorter suffix, Pr (-|x?"). Therefore, the population loss gradient vanishes when the model correctly learns to estimate the latter.
where
> 0 is a constant scaling factor and
att E Rnd is block vector defined as (6h)T 0 0
S
0_
à times
times
Intuitively, the 1hth head in the first layer is constructed to attend the token at relative position -h. The value matrix
TThe heads can more than n-
5

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
2nd Head
1 Head
2
Skip
2
0
2
2mm
C2-
D
0
pry--1
C
2
2
2
20y--2
2
First attention layer the (k-io--istory the (k-1)- (k-1))-histry and aggregates the following token
Second layer matches
First attention layer -iiitoyy
Second layer matches (k-1)) 1-------tttory and aggregates the following token
(a)
(b)
Figure 2: Transformer representing r-gram and k-gram estimator. (a) The task we consider is learning in-context tri-grams (n=3). Here, we illustrate how the transformer given by 0" constructed in section 4.1 operates the sequence (..., z,. 2, 0, 2, =z=,,.... 0, =2) and computes tri-gram(k=3) estimator. Skip. 1*e-Head, 2 denotes the outputs of the skip connection, 1s and the 2nd head. The first layer creates the (2)-history and the second layer compares it with the (2)-history of token 7+1, i.e, Fr = The second layer compares these histories and attends to tokens f, and averages the tokens at that position. (b) * represents that 2nd head is deactivated. Hence, the first layer creates the (1)-history and the second layer matches with the (iitorytory of token T+,,, i.e, Ct. -2 and attends and averages the tokens at 2-
acts as an identity map. copying the embedding of this (-h)- token into the th block of the first layer's output, r((h) As illustrated in Fig. 4, for any position f, the first layer's output aggregates the embeddings of the previous n-1 tokens (in the limit c - oo):
T) exclusively selects past positions whose (n - 1) --histtory matches the history at position T'. The model's output is then the average of the tokens following these matches, thereby computing the n-gram MLE estimator (see Appendix B. for a formal proof).
Transformer representing sub-n-gram. The n-gram con- struction can be directly adapted to represent k-gram for any & < n. The key insight is that the hthh head is respon- sible for retrieving the (-h)-token. Hence, to compute a k- A-gram counting estimator, suffices to deactivate the heads responsible for history elements beyond h 1, as shown in Figure 4(b). The heads in the first layer are thus divided into two groups:
[
T
E nd
(7)
The second attention layer then compares these histories using the specific structure of the query and key matrices defined in Equation (6). The pre-softmax attention score between the first-layer embeddings the 7th and jihh tokens r. [i] and r.] is computed as:
a) Activated Heads. These heads compute the (-h)-token for h [K [k1].
(Krr, (Q)) = Σ 10
b) Deactivated Heads. These heads output zero vector, effectively ignoring history beyond k - tokens.
In the limit as - 0o, where the softmax converges to a hardmax, this construction ensures that the attention for predicting the token + (using query from position
The second layer structure is similar to n a-grams case but with components related the deactivated heads set to zero.
6

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
This is implemented by modifying the parameters as follows:
formers: points in parameter space that implement simpler solutions, k-gram estimators for k < n, are near-stationary points. This provides a formal basis for understanding stage- wise learning as a process of moving between these near- stationary points.
T-1
h-1 Σ 1=
A (k)_
c Σ e.. (eI_)²
tc
ee
(e5)
for he[k-1].
arbitrary
otherwise.
Note that our result holds for any sequence length T. therefore also applies to the commonly used loss function, which averages the loss over sequences of varying lengths, (8b) i.e. (0) Σi= (e) where, with a slight abuse of notation, Lt refers to L in Equation (1) for sequences of length t. In the averaged loss, the gradients from the longer sequences suffer from the vanishing gradient problem as a result of the above theorem.
(8a)
Σ 858 for h E [k-1],
0 otherwise,
s h-1
(Q2) K2
(T))
(8c)
This construction ensures that the second layer only com- pares the (k: -iitisttriess The model attends uniformly to all past positions that share the relevant k-history, thereby computing the k-gram MLE estimator.
4.3. Proof Sketch
Our proof hinges on applying the stationarity result from Proposition 3.2. This requires establishing two key con- ditions for the parameter configuration: (a) demonstrating that model's prediction Pog (z") converges the true con- ditional probability P-(- =44))2) and (b) proving that the score function depends solely on the (k-1) -histtory.
We denote the point given by Equations (8a), (8b), (8c) by 0 = ((A(() (b) v((h j== K2 U22) defined precisely in App. (10). In Lemma B. we provide formal results showing that 0 implements the k-gram estimator in Defini- tion 2.1 in the limit c + 00.
It is easy to see that the first condition holds asymptotically. As c -> co, the model pø:: (x") implements the k-gram esti- mator (see Lemma B. and as T - xo, the g-gram estima- tor converges to the conditional probability (see Lemma H.6 in the Appendix for the formal argument). Finally, ob- tain non-asymptotic bounds on the gradient, we use these asymptotic computations while carefully controlling for the perturbation caused by non-asymptotic conditions.
The specific construction described above is not unique. For instance, multiple heads could be assigned compute the same (-I)-token for I < k. and the result would hold as long head outputs the history for 1 > k. See Appendix F for details.
Other possible constructions.
4.2. Sub-n-grams Are Near-Stationary Points
The main technical challenge is proving the second condi- tion: that the score function has a restricted dependency structure. Our argument proceeds by analyzing the structure of the model's output derivatives. The model's prediction is weighted sum of one-hot vectors:
Having constructed the points 0% that implement k-gram estimators, we now show that they are first-order station- ary points of the population cross-entropy loss in the large context and large norm asymptotic. This result follows from the general characterization of stationary points for the cross-entropy loss, provided in Section 3.
po (x
Σ aar Ezee
Theorem 4.1. For the simplified disentangled transformer pe. the following on holds on the norm of the gradient at 0
where (2) are the attention scores in the second layer for key i and query T (see the Appendix B for details). The derivative with respect to any parameter (11) (22 in the first or second layer involves the derivatives of these attention scores:
|10g_0: L(O)I
Vc E E O(Ipp= (1/07-40)) - Pss 1 O(T PT-
Opo(z") @po(x") 00 00
T
da 0r, Dr, [1] 00
da 30
This theorem quantifies the gradient norm at the 0% con- struction. The first term is driven by the statistical error of the k-gram MLE estimator and decays as e eT( (Penev, 1991). Hence, in the limit of @(T) + o., the gradi- ent vanishes. As the stationarity holds asymptotically as T, c + xo, we term these points near-stationary.
Cay
eag
A erucial observation is that the softmax function has self- toff-bouniig property: the magnitude of the derivative of a softmax output is bounded by (and scales with) the output value itself. Consequently, when an attention score is close to zero, its gradient will also be close to zero. This property will play a key role our subsequent analysis.
This result reveals a remarkable feature of the landscape of learning in-context n-gram language models with trans-
7

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
At the k-gram estimator parameterized by 0%. the at- tention scores are highly sparse. Let M C [T], be the set of past positions whose k-history matches that of the (T + 1)th token. Formally defindd as M {i E (T) 1 (ar- k = .12)) . By construction, the at-
incorporated in the architecture. However, in this case, the transformer's output would no longer be in A requiring
normalization via softmax at the end. This restriction on the value matrix can also be alleviated by using an MLP layer to approximate the logarithm (see Appendix F.I).
tention scores for 0 are given by:
On the Emergence of Syntactic Structure. Consider the behavior of gradient-based method at a sub-n-gram, say O². As is a stationary point, training remains this point for prolonged period, leading plateau in the training curve. However, as training progresses, the model eventu- ally escapes due to landscape curvature or stochastic noise, allowing it to learn a new syntactic structure- attending to (-k)-token before reaching the next stationary point. This phenomenon is general and has been empirically reported- emergence of syntactic structure (Chen et al., 2024a; Wei et al., 2022), phase transitions (Olsson et al., 2022; Edel- man et al., 2024). By carefully analyzing the loss landscape of relevant yet simple in-context task, we demonstrate that the stationary points correspond to underlying syntactic structures. Consequently, our work provides insights into why these syntactic structures emerge following extended plateaus.
for i E M
a
(2) (4/77)
M
O.W..
Due to the self- seff-bounding property, the summation in the gradient expression collapses, receiving significant contri- butions only from tokens within the matching set M
(xp(x") (0)2)
(x")
RM
(1)))
le=
Daliiny
Cg.
(0(2) (2)
Ör,
KEMA
REM
A key result we establish next is that the derivatives of da [i] depend exclusively on the (k-1)-history of to- ken i. Several reinforcing factors contribute to this result. Due to the specific structure of the key and value matrices in the second layer at 0 the gradients with respect to the parameters of the deactivated heads vanish. The deactiva- tion, in turn, ensures that the embeddings after the first layer only contain the embeddings of the (K-1))-istoop Finally, for i MA. its (k-1)-history is identical to that of token T+1. Together, these factors imply that the derivatives of the second layer are solely a function of (k-1)-history of token T+1 The complete proof is in Appendix B.2.
Limitations. Our theoretical results hold in the asymptotic limit of infinite sequence length (T - oo) and parame- ter scale (c + oo). The dependence on c is mild, as the gradient decays exponentially for finite c. However, there is an inherent difficulty in moving beyond the assumption of infinite sequence length. Existing works on theoretical analysis of learning in-context n- n-gram like sequences with transformers, such as Rajaraman et al. (2024) and Nichani et al. (2024), also rely on this assumption. It is even un- clear which estimator transformers learn at finite sequence lengths, even at the end of training. This makes it partic- ularly challenging to determine what transformers exactly compute during intermediate stages of training.
4.4. Extensions and Perspectives
Beyond contiguous history. Our main analysis focuses on contiguous k-gram histories (suffixes). However, as discussed previously in section 2.2, the dependency on the last & tokens are not more important than the other n-k for the n-gram data generating process under a uniform prior on the transition matrices. Our framework can be extended to show that estimators matching arbitrary subsets of the (n- -hiisory (e.g., matching only the (-i)-token for some i > 1 as counterpart of the bigram estimator) also correspond to stationary points under certain assumptions. See Appendix F.2 for details.
5. Experimental Evaluation
In this section, we perform experiments on the disentangled transformer introduced in the previous section to examine the stage-wise wise learning behavior and analyze the different so- lutions the transformer learns during different stages of train- ing. The code is available at https: //git com/ tml epfl/ / subb n- grams are stationary.
Experimental Setup. We train our model on data generated from trigram (n = 3) language model over vocabulary of size S = 5. Each in-context sequence has a length of T = 32, and the transition probabilities are drawn from a uniform Dirichlet prior, Dir(a1), with a = 0.5. The model is the two-layer simplified Transformer analyzed in our theory, with two heads in the first layer and embedding dimension of d = 5. We use one-hot token embeddings
Towards general transformer architecture. Our parame- ter construction, techniques and methodology extend natu- rally to general transformer architecture. First, positional encoding, which was not explicitly used in the simplified model, can be incorporated using one-hot positional en- coding. The concatenation of the attention head output and residual connections can be replaced with the additioo. Additionally, the value matrix in the second layer can be
8

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
h h=1
Attention Scores or (A (100)
115
Lo5
(- 5000
h h= 2
16000
t 115
f = 5000
16000
Step
Training Step
(a)
(b)
The evolution of the attention heads in the first layer during training. (a) Progression of the test loss during training. The highlighted points are the iterations the plateaus for which demonstrate the attention matrices. (b) The evolution of attention scores of the heads of the simplified transformer architecture during training representing the tokens it is attending. First, both of the attention heads attend to all the previous tokens uniformly. At the second plateau, they both attend to the previous token. Finally, as the model escapes this plateau, the second attention head learns to attend to (-2) (-2)-token at the end training.
Figure 3:
and train for 24 iterations using the Adam optimizer with constant learning rate of 0.01 and a batch size of 128. No weight decay is used. The test loss evaluated on a separate set of 216 sequences.
within simplified Transformer architecture that provably implement sub-n-gram estimators. Our main theoretical result then establishes that these configurations correspond to near-stationary points of population cross-entropy loss. This finding creates a mechanistic link between the geometry of the loss landscape and the empirical observa- tion of training plateaus corresponding to sub-n-grams: the model's learning trajectory naturally pauses at these regions of vanishingly small gradient, which represent mastery of simpler, sub-syntactic skills, before eventually escaping to learn more complex dependencies.
Discussion. To accurately predict the next token, the model needs attend to the previous 1 2 tokens in the sequence. Figure 3a shows that learning occurs in distinct phases, where the model remains in a plateau for an ex- tended period before quickly jumping to the next one. In Figure 3b, we illustrate the evolution of the attention maps from both heads in the first layer at various plateaus during training, providing fine-grained view of the structure of the model at these stages. Initially, the attention maps uniform, and the model does not consider token history in its predictions. Upon reaching the first plateau, both atten- tion heads focus on the previous token (the (-1)-token), and the model behaves like a bigram estimator at this stationary point. In the later stages of training, the second attention head learns shift its focus the (-2)-token, transforming the model into trigram estimator. The same phenomenon holds for the general attention-only transformers, see Fig 5.
Our analysis provides a foundational step, and several ex- citing avenues for future research remain. The theoretical results are derived under idealized conditions, namely for the population loss, in the limit of infinite sequence length, and without regularization. Extending these findings to the sante-ssmmle regime and incorporating practical consider- ations like finite context windows and the effect of weight decay is important direction. Furthermore, our work focuses the static properties of the loss landscape. A complementary and crucial line of inquiry involves analyz- ing the training dynamics themselves. Characterizing how stochastic gradient- methods navigate this structured landscape to escape plateaus and transition between solu- tions, as explored by Nichani et al. (2024), presents a key challenge for future research.
6. Conclusion
In this work, we investigate the problem learning in- context n-grams with transformers, specifically focusing on simplified yet insightful setting to gain deeper under- standing of the dynamics of complex large language mod- els. We first constructed specific parameter configurations

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Impact Statement
Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https:
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
/ /oppnneevew..ee www.uoiiooriiio... iUxQ1I.
Brown, P., Dellapietra, V., Souza, P., Lai, J., and Mercer, Class-based n n-gram models of natural language. Compu- lational Linguistics, 18:467-479, 01 1992.
Acknowledgements
Brown, T. B. Language models are few-shot learners. arXiv
preprint arXiv:2005. /4/65, 2020.
This work was supported by the Swiss National Science Foundation (Grant No. 212111) and an unrestricted gift from Google. A.V. V. acknowledges funding from Swiss Data Science Center Fellowship. The authors extend their gratitude to Adway Girish for his valuable feedback the manuscript and to anonymous reviewers for their insightful comments, which significantly improved the final version.
Chen, A., Shwartz-Ziv, R., Cho, K., Leavitt, M. I.., and Saphra, N. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in MLMs. In The Twelfth International Conference on Learning Represen- tations, 2024a. URL https: openreview net./ forum derummid=Monn weriidddMOPiHHBBB
References
Chen, S., Sheen, H., Wang, T., and Yang, Z. Unveiling induction heads: Provable training dynamics and feature learning in transformers. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b.
Abbe, E., Adsera, E. B., and Misiakiewicz, T. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory, pp. 2552-2623. PMLR, 2023.
Chomsky, N. Three models for the description of language. Transactions on Information Theory, 2(3):113-224, 1956. doi: 10.1109/TIT.1956. 1056813.
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-contex. learning. Advances in Neural Infor- mation Processing Systems, 36, 2024.
Edelman, B. L.. Edelman, E., Goel, S., Malach, E., and Tsilivis, N. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint arXhv:2402.11004 2024.
Akyürek, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: Arhitectures and algorithms. arXiv preprint arXiv:2401.12991 2024.
Friedman. D., Wettig, A., and Chen, D. Learning trans- former programs. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Arora, S., Cohen, N., Hu, W., and Luo, Y. Implicit regular- ization deep matrix factorization, Advances in Neural Information Processing Systems, 32, 2019.
Fukumizu, K. and Amari, S. Local minima and plateaus in
Berthier, R. Incremental learning in diagonal linear net- works. arXiv preprint arXiv:2208.14673 2022.
hierarchical structures of multilayer perceptrons. Neural Networks, 13(3):317-327, 2000. ISSN 0893-6080. doi: URL siieene/artaaaee///////s90000000
Bhattamishra, S., Patel, A., Blunsom, P., and Kanade, V. Understanding in-context learning in transformers and IIms by learning to learn discrete functions. arXiv preprint arXiv:2310.03016, 2023.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of sim- ple function classes. Advances Neural Information Processing Systems, 555:00833305998 2022.
Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, L. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Gidel, G., Bach, F., and Lacoste-Julien, S. Implicit regu- larization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Boix-Adsera, E., Littwin, E., Abbe, E., Bengio, S., and Susskind, J. Transformers leam through gradual rank increase. arXiv preprint arXiv:2306.077 070442 2023.
Gissin, D., Shalev-Shward S., and Daniely, A. The im- plicit bias of depth: How incremental learning drives generalization. International Conference on Learning Representations, 2020.
Boursier, E., Pillaud-VVe Vivien, L., and Flammarion, N. Gradi- ent flow dynamics of shallow reLU networks for square loss and orthogonal inputs. In Oh, A. H., Agarwal, A.,
10

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Han, Y., Jana, S., and Wu, Y. Optimal prediction of markov chains with and without spectral gap. Advances in Neural Information Processing Systems, 34:1 1233-1124666 2021.
Penev, S. Efficient estimation of the stationary distribution for exponentially ergodic markov chains. Journal of Statistical Planning and Inference, 27(1):105-123, 1991.
Jacot, A., Ged, F., Şimşek, B., Hongler, C., and Gabriel, F. baddle-to-sadde dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106. 2021.
Pesme, S. and Flammarion, N. Saddle-to-saddle dynamies in diagonal linear networks. Advances in Neural Informa- tion Processing Systems, 36:7475-777 7505, 2023.
Power, A., Burda, Y. Edwards, H., Babuschkin, I., and Misra, Grokking: Generalization beyond overfit- ting small algorithmic datasets. arXiv preprint arXiv:2201.022 12227707777 2022.
Jelinek, F. Statistical methods for speech recognition. MIT Press, Cambridge, MA, USA, 1998. ISBN 0262100665.
Jiang, L., Chen, Y., and Ding, L. Algorithmic regulariza- tion in model-free overparametrized asymmetric matrix factorization. arXiv preprint wrXiv:2203.0288.00 2022.
Rajaraman, N., Bondaschi, M., Ramchandran, K., Gast- par, M., and Makkuva, A. V. Transformers on markov data: Constant depth suffices. arXiv preprint arXiv:2407. 17686, 2024.
Jin, J., Li, Z., Lyu, K., Du, S., and Lee, J. D. Under- standing incremental learning of gradient descent: A fine-grate fine-grained analysis of matrix sensing. arXiv preprint arXiv:2301. 1500, 2023.
Razin, N., Maman, A., and Cohen, N. Implicit regulariza- tion in tensor factorization. International Conference on Machine Learning, pp. 8913-8924. PMLR, 2021.
Jurafsky, D. and Martin, J. H. Speech and Language Pro- cessing: An Introduction to Natural Language Process- ing, Computational Linguistics, and Speech Recognition. Prentice Hall, Upper Saddle River, NJ, 2nd edition, 2009.
Saxe, A. M., McClelland, L.., and Ganguli, S. A mathe- matical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sci- ences, 1 116(23):11577-1154) 2019.
Kim, J., Kwon, S., Choi, J. Y., Park, J., Cho, J., Lee, J. D., and Ryu, E. K. Task diversity shortens the icl plateau, 2024. URL https: //apii org/ abs/ 2410. 05448.
Shannon, C. E. A mathematical theory of communication.
The Bell system technical journal. ((13):379-2333 1948.
Svete, A. and Cotterell, R. Transformers can represent n-- gram language models. arXiv preprint arXiv:2404.1 14994, 2024.
Li, Z., Luo, Y., and Lyu, K. Towards resolving the im- plicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2021.
Varre, A. V., Vladarean, M.-L., Pillaud-Vivien. L., and Flammarion, N. On the spectral bias of two-layer lin- ear networks. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: // //operreview net/ FoormmiiddFFdrreaas
Makkuva, A. V., Bondaschi, M., Girish, A., Nagle, A., Jaggi, M.. Kim, H., and Gastpar, M. Attention with markov: A framework for principled analysis of transformers via markov chains, arXiv preprint arXiv:2402.0416). 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., and Gomez, A. N. L. u. kaiser, and polo- sukhin, "attentio is all you need.". Advances in neural information processing systems, 30:5998-6008 2017.
Nguyen, T. Understanding transformers via n-gram statis- tics, 2024. URL 12034.
Von Oswald, J., Niklasson, F.,, Randazzo, F., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pp.
Nichani, E., Damian, A., and Lee, D. How transformers learn causal structure with gradient descent, 2024. URL https: / arxiv. org/ abs///b////// / 2202..22222 14735.
PMLR, 2023.
Odonnat, A., Bouaziz, W., and Cabannes, V. Cluster- ing head: A visual case study of the training dynamics in transformers, 2025. URL https: /ttppss////xvvvoo/o// 2oo/2210.2440
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., et al. Emergent abilities of large language models. arXiv preprint wwxiv:2206.07622 2022.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., et al. In-context learning and induction heads. arXiv preprint ArXiv:2209 //895, 2022.
Zekri, O., Odonnat, A., Benechehab, A., Bleistein, L., Boullé, N., and Redko, I. Large language models as markov chains, 2025. URL https: / /////aavvv.////////////a///.aaav axxiv.org/ abs/2410.022

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
A. Organization of the Supplementary Material
A.1. Links to Materials Referenced The Main Text.
First, we present an index of the supporting material referenced in the main text.
The proofs of Lemma 3.1 and Proposition 3.2 are provided in Section and the discussion on the stationarity condition for subsequences beyond suffixes is provided at Remark C.2.
The proof of the construction of k-grams is given in Subsection B.1 and the proof of Theorem 4.1 is provided in Subsection B.2.
The discussion on the possible alterntte representation for the stationary distribution is provided in Section and the extension for general transformer architecture is given in Subsection F.I. Stationary points conditioned on subsequences that are not suffixes for fixed 7 are further discussed in the Subsection F.2
A.2. Outline of the Supplementary Material.
Section B, provides the construction of the k-grams and the proof of the stationarity of the construction (in Subsec- tion B.I and Subsection B.2 respectively).
Section C provides the proofs of the results on the gradient of the cross-entropy and sufficient stationary conditions, i.e., of Lemma 3.1 and Proposition 3.2.
Section provides a technical lemma related computing the gradient for two-layer simplified transformer.
In Section E, a technical lemma related the k-gram representations is provided.
In Section F, discuss the extensions of the results general transformer architecture in Subsection F.J and the stationary points conditioned subsequences that are not suffixes in Subsection F.2.
In Section G, present the derivatives of a single layer celf-attention map.
A.3. Notation and Definitions
Notations. We use & denote the Kronecker product. We use operator for flattening the matrix to vector. We use eº to denote the :th elementary basis vector of IR. We drop the superscript when d = S and denotes the 7th elementary basis vector of s
Definition A.1 (Jacobian of function). Let f Rm RRP be Ci-function defined variable X. af denotes the Jacobian which is function from R=xn + Rxxmn
Definition A.2. Define 0
A
(b) 1Am1 hE(n-1)
(h)
(22 Q2) as the of parameters given by the following expressions
S 1-1
(22) T Ve
(5/)T
(10a)
j=1 #=1
k-1
K2 = Vc ΣD
(o))),
(10b)
j=1
E er T for he(k- 1).
Σ e (ei_.)
AM)
(10c)
arbitrary
otherwise,
((h)
Σ 8 for h E [k-1]]
(10d)
1=1
0 O.W.
12

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
2nd Head
2T-22
1s Head
Skip
O
O
20
St-
*
2
20'11
4
0
20'1
<
2
T 2
0
0 2
0
2
2
2
2
T
20-22
*
27--2
35y 2
2P
2T
8t
First attention layer creates the (k- historyy
Second layer matches (k-1) (A=-1)-histor 1)-sssttory and aggregates following token
Second layer matches creates the (k- hi-httoo the (k-1)) 1i-historr and aggregates the following
First attention layer
(h)
Figure 4: Transformer representing n-gram and k--ram estimator. (a) The task we consider task is learning in-context tri-grams (n=3). Here, we illustrate how the transformer given by 0 constructed in section 4.1 operates on the sequence (..., =0. . 2, 0, 2, x.=1,... .. ... =2) and computes tri-gram(k=3) estimator. Skip. 1st 1ee_Head, 2nd -Head denotes the outputs of the skip connection, 1s and the 2nd head. The first layer creates the (2)-history and the second layer compares it with the h2)-history of token T+1, i.e, Fr The second layer compares these histories and attends tokens l, and averages the tokens at that position. (b) * represents that 2nd head is deactivated. Hence, the first layer creates the (1) histtrr and the second layer matches with the -histoyy of token T ie, 27 =2 and attends and averages the tokens at t, 2,
B. Construction and the Stationarity of the k-gram solutions
In this section, we present the proofs technical lemmas and theorems presented in the main text. First, we present the parameter configurations representing k- d-grams and then prove the stationarity of these configurations.
B.1. Representing k-grams with Simplified Transformer
Forward pass of the transformer. Before presenting the proofs, we give an alternate form of the forward pass of the simplified transformer. First, we express the embeddings after the first layer as follows:
r = w r.
())
(.4)
((4)
(34)
h=
j=1
13

---

Learning In-context nigrams with Transformers: Sub-n-grams Are Near-stationary Points
where d denotes the attention scores for key and query i of the attention head & in the 1st layer. The matrices W are used for concatenation. Formally, the attention scores and the matrices are given as follows:
eA
(w) [1,
(LA))
(..9)
Ei expfA 12 [1,1)
Σ $J
==1
For the 1th token, the output embedding of the second layer writes
Ta [t]
Ga Tolo,
where (n denotes the attention scores in the second layer for key i and query t, and are given by
exp (K2r [i]. Q2ri [t]
(2)
(11)
(..))
exp (K (b], Q2r,
The final output probabilities of the token after the unembedding are given by
po(x') = Ur.[t]
a (2)
Uro
Σ (2) Ex.
0 represents a k-gram. Moving forward we provide the proof for the k-gram MLE constructions. To support the proofs,
we define a subset tokens, M, S [t], which consists tokens whose k-history matches the k-history of the (t + 1) token. Formally, defined as follows:
M {i: i m12))-
We begin with a formal lemma to show that 0% represents a k-gram in the asymptotic limit where - co. To make the lemma easier to read on its own, we have restated the standard notations.
Lemma B.1. Let (1,b) denote the attention score of the key and query element (i,j) in the 1 head of the forst layer and let
aic.s denote the attention score between elements i, t in the second layer. Let M = E t]: (w/ = x; =(a))2) } be the set of tokens which match the k-history of the (t+1)th token.
For the parameters 02 defined in Eq. (8), in the limit oc, the attention scores of the activated heads in the first layer, i.e., for heads h where h & k 1, are given by,
[1
when i >> h and j i - h, when i & h and j = 1, otherwise
{
(1,b) (33)
=
(12)
and the attention scores in the second layer are given by
for E M, otherwise
[M
(13)
(1.1)
O
In order to prove the above lemma, we provide two lemmas one for each layer of the transformer, which detail the attention scores the respective layers compute. The lemmas below are non-asymptotic in and provides the bounds on the pertubation in the case of finite c. Intuitively, the attention scores of the 1hth head in the first layer are computed such that the each token attends to its corresponding (-h)-th token, The attention scores of the second layer ensures that tt token attends tokens whose historr matches the h--history of the (1+1)" token. The notation hides the terms polynomial in c.
14

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Lemma B.1. [First-Layer] With the set of parameter 0% given in Def. A.2,
(1.h) (a) The attention score of head h in layer I of key and query j denoted by d(j44) is
- 99 when i > h and j = - h,
O(ie when i < h and j = 0,
(14)
9) o.w.
(b) The first layer outpus the embeddings:
+ O(ie 1 for k 1,
Ti [i]
(15)
20 + ( 1 fori < k k-
Lemma B.2. [Second-Layer] With the set of parameter 0% given in Def. A.2, the attention scores of the second layer are
Olle
for i M
a (..
[MF O(E [MP]
IMP
(16)
For the proof of the above lemmas, see Section E.
Proof of the construction for k-gram MLE the limit e < 00. We give the proof of construction any and when we get the n-gram estimator. The final output probabilities writes
po
(x*)
Czy-
In the limit - co, the attention scores from the second layer from Lemma B.2 is given by,
for M
MMT
6.0
O.W.
Using this the output probabilities are only supported by the tokens in M. and is given by,
po (o(x')
Σ Exp
I
iEM
Now
Po (z =
M;) s}. s).
M
which exactly matches the k--gam MIE estimator in Definition 2.1.
B.2. Proof of Stationary Points with Simplified Transformer
In this subsection, we will show that the set of parameters given by 0g are indeed near-stationary by providing a bound on the norm of the gradient.
Theorem 4.1. For the simplified disentangled transformer pe. the following on holds on the norm of the gradient at e
|10 L())Il
E O([pp (11=)_..)) Dr. 1 O(T Vc P 2
Proof The proof is an application of Lemma D.I for the transformer parameters that compute the k-gram MLE estimator 0% given in Def. A.2.
15

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Parameters of the second layer.
From Lemma D.I, the derivatives of pe with respect to the second layer
Opa
(ex,) @ vec (Q2n )) (1.)
0K2
Ope
Pides (ex;) vee 7(",
aQ₂
Before computing these quantities at O.,, we gather the attention scores and weighted embeddings from Lemma B.I, B.2.
k-1
+
sh di-h
+
for i e 1,
ri
h=
(17)
k-1
+
Σ
$ + O(ie - for < k - 1. + Σ h=i+1
h=1
Oite=)) for E M
((.,)
((e
|
[M
(18)
O.W.
MM
The average embedding,
2 (1.)
ri
Σ
+
Σ
(2) (s,2) T1+
EEM;
M
The deviation due finite weights can be controlled as the following.
O(te IM
Σ
Σ
G.0 sup
||]]
EE(T]
igM
igM
as | < for and (a.o from Eq. (18) for & M Now, we consider the summation Eieny
Σ
Σ
1
9
r
Σ
r
O(te |MP1
1.
M
|
M
KEM
M
M
Recall that for E M r, from Eq. gives
A-1 + Olie O
r
[
Sz,
Σ
h
9) 1.
Now we use the definition of the set M to simplify the above expressions of r.[d] and F()) Note that d:-h = -1+11-1 for E [k - 1] E M.. Using this, we have,
T
[
+ O(ie-)) 1,
h=
h-1 Sutia
Σ
r
82
Olle 1.
IM
M
EEM
iEM
Combining them, we get
k- 1
Σ
Σ
0
= Ole
M
8x:
REME
EM
1=1
1
Σ
+ +(te
M
TEM!
16

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
i M
EiEmp so
r - P()) - =62
+ 9tte 1,
M
(n - (( = O(tvoe 1,
Q2(r, vo Σ
S + +((VEE-9)- 11
hm1
For all
(e - [ &
K2(n. [i] ((")((i))) < Vc.
Recalling the gradients,
0pg
Σ
(cx;) @ vec -(.))) =(")T)²
OK2
(10))
Ope
die (ex,) vec -+"
aQ2
We split the gradients supported on M and its complement,
Ope
Σ (i (ex,) vee -x (T))))))))
+ Σ
( (ex.) @ vec (K2(r.[i] (t,(t))
iQ2
iEM?
iEM?
= - 1 Olt/ce -1
|M
The final gradient,
Opo OltVOE - 1+ Oltice 9). 1 8Q₂
(19)
|M
On the similar lines,
Opo Vc OK2
IM$I
1
Σ TEMs
Ca, & vec
Σ
Bx41 1-h
0
Enemp 52
IM/
1))
OttVoe-9) 11
Denote
R- 1
0(1
612)
h=1
M
MEMA
Note that the first term in the above expression is only function of k-history xh Using this,
Opo
VC Vo
Σ e., vec (d())_)))) 10.. -5555)) $[]] + -
OK2-
M
|
iEMI
= c/12 |M
1
Σ
#laE M,} (ea) @ vec (1( - + + 1.
aES
17

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Note that the second term og is function of the k-gram MLE Pk. precisely,
#1a E M
Pk [a]
Using this se - P where 5 is the Bunxnd matrix where the first block is the embedding matrix and O's everywhere else. Using this
Ope 2K₂
Vc
Evtlale Palale & vee (o(a) x12) ( s' 1) O(tvce
(20)
aES
Parameters of the first layer.
The derivatives with respect to the first layer parameters are given by,
Ope (V( (h)
dio ((ea. - -, -(: (F: [n,1t] Q2 K2W!" (r, T Q2
Öpe
(333) (1.e) (ez, 1
0A (h)
[i,]]
r
Q2 K.W! (h) V() (rolil - -(o + (r. (!!) K2 Q2 (rols) - -m(aa))
where the averages of embeddings weighted with attention scores are given by,
(1,) Cz,:
((0.A)
Ga)) ro roli].
(1,)
j=1
For the heads that not activated, the derivatives with A (h) averaged embedding is given by,
[ are since the v(( is 0. For the activated heads, the
ralil - h] + Oieep((-c) 1.
Now, consider two cases for the derivatives.
For j = i - h, roli] (00,)) = O(exp(-et), hence the derivative is (icexp((-c)
For j * i- h. due to the property of the softmax function, the attention scores are O(exp(-ct), hence the derivative is again (cceppf-c)).
For the derivative with respect to (kk, we again have two cases,
For the non-activated heads f k, K Σ E ( -j=1 1/ 0. Similarly for the other
term, (h) O. Hence the derivative is O.
For the activated heads h & le 1, using the previous computations have, K2(r.0 - -(") O(tvce
1-1
Q2(ri(1)) = Ve
+ O(LVEe 1,
2111
-=1
k-1
(y = Vc
Σ
1=1
18

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
The product is solely function of k-history ot 612 and does not depend on i. Computing the multiplicative factor in the front,
dio ((ez, - -. (F! = Σ
(ie ((ex, - 2 (F! + ((exp(-c))
1=1
iEME
Note that _ SpI- + Oxii and for i MA, we have, Sx1-h S1141-h' Using this,
((ex, - - [ (+: = Σ
(1.)) (lex, - (5xe=1e1) T) + ((epp--c))
(1,))
MMM
Σ
(ex,)
Σ
(8xe11-)) + O(exp(-c)
LiMM
iEM
The term in the square brackets is O(exp(-c) (for hard attention it is zero). Hence the derivative with respect to v((h) is (exxp(-c)).
0pe O(expt-c)) av(h)
O(vVoE-9)- 1
Vc
(21)
(b)
Final Gradients
Bringing things together, from Equations (19), (20), (21),
Öpe OlLVee 9. 1 OlLVEE
aQ2
|M
Ope
h-1
= O(exp(-cf) Ve
h M141-A + -
@v (K)
Ope
= Ve Σ
Th [a] ea 0 vec (D(ri
[00- S0t 99.
OK,
aES
Note that the only non-vanishing gradient is the gradient with respect to K2. Now, the gradient can be written as
Öpe
Vc
Σ
Pa Pa(aaee ea @ vec (d(x)_)=)) 1s - sp O(t
OK2
aES
We p. to denote the vector pr (.) (1127 the gradient using the difference as,
Opo
Vc Σ pr(a] ea & vec (5 [ --pp1 + O(Ipp- Ph 113) 09
OK₂
( + O(Ip, - px 112) (t 1,
where vi is appropriately defined. The first term here only depends on the (k)-history and we can use Proposition 3.2, to show that does not contribute to the final gradient. To give the final computation,
OC
(po(x") (oox"")pp )-p (.1x"). I log (ogpo((")
E
E mp-
PawP
OK2
Do -p- Pr (.12²), Ooog (x()))
E
E
Ok,
i log
2 llgg (x
E Pr~P
E
Pr (.127 Pr x
IF
E. l - (.) (.12T_Ea) |11|
Ok, 11
OK2
P+~P TTmpp
E
E
pr
(-1=.] 11)) -P (-1x7), { kiz))
E PeNp
E O(1pp (-1=if +12) Do 112) OlVice 9) 1,
PP~P
19

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
The first term vanishes due to Proposition 3.2 and this finishes the proof.
C. Supporting Lemmas for The Derivatives of Cross-entropy Loss
In this section, we provide the proofs the lemmas presented in Section 3 of the main text. the end, we remark about the extension of Proposition 3.2.
Lemma 3.1. Consider any parametric model po(.): 18|T + A that maps sequence of states to probability vector on the states. The derivative of the population cross-entropy loss [(0) with respect to a parameter 0g E R is
(po(=")) (27x(11 (.1="), dg. log logpo(x")):
= E
E up-
P.~P
where dg, is the partial derivative with respect to 0; and log(-) denotes component-wise logarithm.
Proof Recalling the definition of population cross-entropy loss from Eq. (1), we have,
[(0)=
E x
E {00,x²),
E(0,x²)
Σ
(1r+) log (po (0, (s)),
So, De [s] Po (x [s]
do,E (0,xx)))))
Σ
pr (#0+1==(x²))
Using the fact that po(x") E A5-1
$ po(a") ps()")[ss
Taking the derivative of the above expression, we have,
Po. [s] 0.
Adding this to the derivative of the loss gives,
s
Pr (= Po.Po(>")(s po (x
s Jo (x² [s))
0o.6 (0,xx))) x²)) = Σ
m=1
s=1
= (pp((=")pp (112"), Oo log logpo(x"))
Remark C.1. In general, a parametric model computes function fe [S] + RS after which normalizing function like soft-max used to project onto the simplex A-1 The partial derivative w.r.t to any parameter 6; in this case simplifies to
C(0)
E E p-mp at mp,
(po(=")) P+ (-1x²), 00, fo.
Proposition 3.2. For any 0, E Rp such that Oo=o, logpo(x²)= (pr, x²'), i.e. the score is solely function of the context Pr and the last 1-++1 elements of the sequence I², the gradient of the population loss L can be written as
vL(0.)= =
E Tupp
(pp.(x")- (.1x7) .9 (pr x()).
Furthermore, if for such 0, RP, the model estimates the conditional probability of the next token pr (.1x7), iee Po. (xT)=p=(( (.) xx)' almost surely for Pr P, then 0_ is a stationary point.
20

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Proof From the above lemma, the partial derivative of the population loss is given by,
0s (()= = P of upp
(po(x")-p (11=") de log po()))
Using our assumption on the score from the proposition, the above expression is rewritten as follows:
E Ppr (-1 (11="). , de, log po poo))("))))))
E (pr (.1x²) .9 (pr. 22.)
xTupr
Now the sequence x" is split into two random variables (a² ,
aT>p, Pr (-1=²) 0o, log po (x²)
Pr (11(='),,,))) (prs
(22)
(z pup,
We use the following property on factorization of the expectation, for any two random variables, A, B. Let (AA,, (B) be the support of the respective random variable, we have the following fact,
E Pr (A) A.B
((B) = Σ
Σ
Pr(A a, B = b))r)pr((p
a, () (B))
BET(B) (A)
Σ
Pr(B b)p(B)
Σ
Pr(A= = alB b) (-( a, B b))
BE
(A)
Σ BEA(B)
Pr(B b)o(B)Pr (-B=b).
E BEe(B)
Pr(B b)b(B) (.
(.B=b) Σ
Pr(A alB= = b),
aEσ(A)
Σ
Σ Pr(A B- (-BB=b)
BE (EO(A)
( (:BB))(B)
E Pr
A.B
Using the above property, we can rewrite the expression Eq. (22) using A = x² and B = x.. the per task loss can be written as,
(pr (.|x"), 0o log logpo(z")) = E aT~pr
(p- (1 (pr,
E up--
This gives us the desired result. The remaining argument is direct: the residue vanishes when the probability estimate matches the conditional probability.
Remark C.2. The proof of the proposition hinges on dividing the sequence x" into prefix x'-1 and a suffix x;', as illustrated in Eq. (22). This particular split is chosen for clarity in presentation and discussion. However, the proof remains valid for any partition of the sequence into two disjoint subsequences, which need not be contiguous or follow a prefix-suffix suffix structure. Consequently, the result extends to conditioning on arbitrary non-contiguous subsequences that are not necessarily suffixes.
C.1. Conditional Probabilities: Definition and Proper Asymptotics
In the main paper, we have used a somewhat informal treatment of conditional probabilities. Here, we provide formal definitions to supplement our discussion. First, we define them for any general sequences of length T for any generic probability distribution on [S]*
Definition C.3. Given the propability distribution py (x Fr, (1) ) of sequence of random variables Vi, ..,... Tr taking their values in (S). We define the following conditional probabilities of the next token given (T t) history is:
Ex Pr (xT11 = 11)
Pr
Fr41
x..
Ee Pr (x = (7)
21

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
where the marginal distribution for sequence length / is given by
Pr (x
Σ PrCss)
pr pr #tt)
Under the assumption that pr is 7-gram language model, it remains establish that the conditional probabilities defined in time-home homogenoggs manner and that the k-gram estimators converge to them asymptotically. These aspects are addressed in detail in Section H.
D. Derivatives of The Simplified Transformer
In this section, we focus on the derivatives of the simplified two-layer transformer. Using the derivative of the masked self-attention map Lemma G.I, the derivatives for the two layers computed here.
Lemma D.I. Using Lemma G.I, we define the following quantities, which can be interpreted as attention-weighted embeddings for both layers.
7()
(Go) [i]]
(23)
(
a (2)
(c,e) Czp
(24)
plO.A)
a (1,0) roli).
(25)
1=1
With the above notations the derivatives of the ousput probabilities after tokens with respect to the parameters of the model are given by,
Öpe
a (i
(ez.) @ vee - -
(26)
0K,
Öpo
(a.) (ex,) & vec ((=
(27)
DQ₂
Ope
(.o ((ex, - (+:0)) s (r.11 T o k2W!') (r.(1] ((o)) TK, Q2m!!!)
(28)
0V(3) [M]
Ope
G. (.0) (Ezz - 2 ) x
0A (h) [7,1]
[n][ Q K (h) (r.lil - F(O.)) + (r.10)=
- - J.J Q2 (r.li) -F!o.))
(29)
Proof
The final output probabilities given by the model are
po(x') = Ur2[t] =
Σ
Ur.li = Σ Urof
Gi Exs.
where the attention scores write
exp K2r, [i], Q2i; [1]
(2) (6.
Σ
exp (Knr [i], Q2r [t])
1=1
22

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Using the Lemma G.I, the derivatives of the output probabilities after / tokens with respect to the parameters in the second layer and the embedding of the first layer are given by,
Ope
(ex,) @ vec (Q2n/ft(rr T)
2K₂
aQ2 (1t)) (ex:) vec (K2 - (ri)tt))) T
Öpe
@ra[t]
(in) (ex, - - 2 (K2 Q24i [t]) ;
For *t, t,
[r
0r,[t]
[r, (3.e) (ez, " @ K Q21i [t])
+ Σ
(2) (ex, - - @ (Q
(r - -!
The output embeddings of the first attention layer.
r.ld = w ro +
W V (h)
robi,
exp A (1)
[
S-1
where
a (1.) (3.9)
W = Σ
ST
Li=1 exp A (h) [K,
j=1
The Jacobian of the first layer embeddings r, w.r.t to attention and value matrix of the head h,
òr. ( T @ W!!) ov (1)
@r,
(h) (ro[i) =(O."))
0A 0A [1
Combining them, we get,
Öpe (b)
Öpe 0r
Ör, [i] (v(h) [b)
(ez, - (() & Q2r (2r()r1111) 7] (r:O)T @ wi(
(.0)
1=1
(ez: - ह( ② (Q K2 - f!")) ( W!!h)
(1t)
i=1
Σ
(2) (itt)
(lex, E(") (F!"))) 2 (n Q2 k² +r, イ(प) T K2 Qu!!!i)
Ope
Opo
Ir, [
0A
[i.j]
Or, [i] a [i
(2) (1,6)
P (ea,
2 8 K. - + - -((") T K, W (rabl 1001-Fo0 -F!o.)))
(n.6 Q
In the above simplifications, we have used the following property of Kronecker product,
(A () @ D) (AC) 8 (BD),
for matrices A, B, C, D of appropriate dimensions.
23

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
E. Proofs of Representation with Simplified Transformers
Lemma B.I. [First-Layer] With the set of parameter Og given in Def. A.2,
(a) The attention score of head layer 1 of key and query j denoted by d(j.)) is
(ii when i 3 h and j = h,
(1,10 (3.4)
1 - 99 when i < h and j 0,
(14)
Ole 9) o.w.
(b) The first layer outpus the embeddings:
1
22
+
hh
Bi-h
+ 0 1
for 3 h 16-1,
r
=
h=1
k-1
(15)
+
+
Σ Amm+1
sh + O(ie .1 for < k - 1.
Bi-h
Am1
Proof Recall that the attention scores of head the first layer for key i and query are given by
expA1 (h) [
(3
Ei= exp A [
For the activated head h, the A (h) is given by Eq. (8a),
T-1
A (b)
> c.eo
I=0
Using the above two equations for i > h, the attention scores can be simplified as,
T-1
A (h) [. [ =T
Σ 1=h
GET + c
> eee
ej ccj ii),),
1=0
exp A hh) [. (i-1) +e,
1
For j * i h
((ep(-cc),
1) cr
6
and, for j = i - h, Coming to the embeddings after the first layer,
(
=1 O(iexp(-c)).
- 1) + en
r w! [0] + Σ Σgx.
with ((())
(h)
ro[i).
Gig
h=1j=1
where,
W(h) = > St.J..
j=1
Note that for > & the term v(h) = 0, and for h < k, the term v((h) given by Eq, (8b). Using the above two equations and rg] [i]] Sz.. we can write the embeddings after the first layer as,
r [1] = w! Sz, +
W v(A) 8zy.
h=
j=1
24

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Note that V Am; = 82, for h < ls and W 82, 2/.. Furthermore, for E [S],
||أا|ا5 & 1, and [ sylloo 1.
Using this, the above equation can be written as,
[ 50
R + O(ie "9e 1.
similar computation for < h gives the required result.
Lemma B.2. [Second-Layer] With the set of parameter 0% given in Def. A.2, the attention scores of the second layer are
Ole
for i M.:
a
MMT
IMP
(16)
O(e
MMT
O.W.
Proof The product of the query and key in the second layer from Eq. 8c is given by,
k-1
QQ K2
ΣΣ $ (s)).
Using this the attention scores in the second layer for any key t and query are given by,
exp (Kr Sgr [1]. Q2r [t])
(2)
6..
exp (Krr Q2r/(t)
To compute the inner product,
(K2r [1]. Q2r, (r.14. Q K2r, [i])
We know from Lemma B.I that the embeddings after the first layer for > are given by,
+ - 11
Now for >
k-1 s
k-1
(Q2) Ken [1] =c
T
R.L + 9
h=1j=1
1==
+ 0 9).
Zink
k-1
k- 1
(r.(t)
'Kr =c
h-1
+ + ))
Fi-h
h=0
h=l
h-1
h-1 P4ーh
+ +0
#++1-k
h=1
25

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Now the fact that the embeddings s/ .8 are orthogonal for + j hı + h2. Using this, we write the above expression as,
(r. (Q2) (i) (i))
8 h-1
S + O(e
=
B241-h
1 - [ O(te 9). Σ h=1
If the k-history of t + and i match, i.e., i M,. then the summation is maximum at (k 1)e² otherwise it will be (k- - 22c there is atleast one mismatch. Using this, write the attention scores as,
Σ exp (Knr. Q21. < |M/] 1)c} + (t [M,]) exp( - -)cc..
Hence, the attention scores are given by neglecting the terms with double exponentiation (i.c., exp(exp(-c)))
exp (A: 1)c²
exp (K 1)c²
For i M
My exp 1)e} |M, exp 22)c²
Σ exp (Krr, Q2r.())
1
(t- Mel) exp (As 2)c}
t M
expl
M
(..0)
[M.] [M exp 1)0} + (1 [M,] expl 2)c}]
M,
exp (k 2)c O(exp(-c)) o.w..
For i & M
(s.1)
[M.] exp (k- 1)e
Hence, the attention scores are given by,
10(exp(-c))
For i M
(
M
|M
For i & M
(4) (6.
= O(exp(
o.W..
This completes the proof of the lemma.
F. Possible Extensions of The Results
In this section, we discuss the possible extensions of the results presented the main text. First we discuss how the results can be extended to a general transformer architecture. Then we discuss the possible constructions for stationary points and beyond suffixes.
Other stationary points. Before we begin, we note that there are other possible constructions of the stationary points. In the main text, we focus on the set parameter configurations where for each h in [1, 1] a single activated head computes the (-h)-token, while the other heads are turned off. However, there can be multiple heads dedicated to computing the (-h) (-h)-tokens for f & k 1. Using symmetry, the current approach can be extended to show that the gradients parameters of these heads vanish, similar the proof of Theorem 4.1.
For example, consider the bigram MLE estimator; in the construction given in Eq. (10) a head is activated to compute the (-1)-token and the other heads are turned off. Alternatively, there exist parameter configurations where all heads are activated and compute the (-1)-token similar to Figure 3. Our proof of Theorem 4. works in this case too, with the arguments of symmetry, the gradients of these heads will vanish. This can be shown using the same arguments as in the proof of Theorem 4.1. We give an explicit construction for this simple bigram case.
26

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
n-1
(Q2)
ve
(30a)
J=1
K2 Vc
(
(30b)
T-1
A
(M)
Σ
e (e())
for h E [n-1],
(30c)
1=1
v Vc Σ 8ys, for [n-1].
(30d)
Here, all heads are switched on and the token with all the previous heads (notice in blue in the superscript for all in Q2). This can generalized to any k-gram, where the (-hi)ttken for h & is computed by different heads with varying multiplicities.
is computed by all the heads. The query is used to compare the It
F.1 General Transformer Architecture
The results presented the main text are given for a simplified transformer architecture. In this section, we discuss how results can extended to a general transformer architecture. provide brief proof sketches for two extensions: adding a value matrix the second layer and replacing the concatenation of head embeddings and the first-layer skip connection.
Token and position embeddings. First define some block embeddings in dimension Rd and later lift them into the dimension of the transformers, i... nd and sequences upto length T.
Token embeddings S0, 51, ....... 88-1 E IR
Position embeddings - Po. Pl PI...., pr- Rd
The embeddings mutually orthogonal, i.e.,
8; 1 83. for all i + E [N].
P: . Pj. for all i i j E [T].
1 8j, for all E [7], j E [N].
For h [Ac + 1], we denote the st p. Rand the lift of the block embeddings which are defined the following,
0a Oa
Oa
Oad
Pt
blocks
The embedding layer maps to go E RTxxnd xnd where each row is the embedding of 7th element in the sequence along with its positional embedding, i.e., 4o [i] s6, + p. nnd
The attention layers. The first layer has h attention heads and also has a skip connection. For head h E [k), let (((h) denote the query, key and value matrices. The forward pass on go writes
K (1)))
P1 = To
Σo o
ro
(1(h)
Q
(1) T :.) no (v(()))
Here (.)) is a row-wise softmax operator with a causal masking. Note that here we are adding in comparison to the simplified architecture where we are concatenating.
27

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
The second layer has just one head and also skip connection (the skip connection can be included and the value matrix should be scaled appropriately so that it is ignored after normalization). Let K2, Q2, V/2 be the key, query and value matrices of the second layer. The forward pass for the second layer writes
T2 = (r, (Q2) K2r,) (V2)
Now that r, is not row wise normalized and the result does not give an output vector on the simplex, we have normalize them using softmax operator, i.e., (r²) get the final output. In this case, the counting algorithm (MLE) implemented by the attention-only transformer previously discussed cannot compute logits of the probability Pr. so that softmax normalization over log pr gives the true probability distribution py. Hence, an MLP layer is used to compute the logits.
An MLP layer. As discussed above, there is an MLP layer that computes the logarithm of its input. Let M be the set of parameters of the MLP layer and m be the function implemented, m(-) RS R.S does component wise logarithm m(x)= log(1z) + e) for some small epsilon.
Writing it together, the forward pass of the transformer writes,
Σ rro (Q ) T r. no (v((H))²
21 To +
h=1
r2 (r. (Q2)
K2 r.. r. (V2)
rf m(Urg),
pa (xx) = (r.t)
This model generalizes the simplified version by adding a value matrix in the second layer, moving beyond simple concatenation for head embeddings and the first-layer skip connection, and explicitly using positional encoding. While the simplified model's representation results extend to this generalized version, proving that the gradients vanish requires additional work, particularly for the value matrices.
For the 1h token, the output after the second layer of the transformer be written as,
Ya [4]
(2) V2r [t].
where the attention scores in the second layer
(2) for key i and query are given by,
(.20) exp (Ker [1]. Q2n. [t]
exp (Keri [i], Q2n, [1]
The final output probabilities are given by
po(x))=Ui = Σ
20 a
Uro = Σ (G0
(1,)) Ex.
Here, we show that the gradient with respect to V, vanishes. The others follow the same pattern as the simplified model and will be not be precisely computed here. The partial derivative with respect to gradient gives us,
0r 0V2
Σ ai Is @ &r.1 Is @ F!!
28

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Define 0 by
T-1
) TKT""""" (1()))
(h)
Ch
Σ
Pi-NPi
where for [k - 1]
(32)
1=h
0
O.W.
8
(((h)
Σ
(69)
for h e E
(33)
==1
0
o.w.
(Q2) _ K2 = e Σ Σ
(3))T,
(34)
1=1 KEN
1/2 Σ (5))))
(35)
Intuitively, this is just an expansion on simplified transformer with lifting the first layer and multiplication with value made explicit. Recalling si
M Diem; s2 using this,
0 P() = cs
R1i11 + O(exp(-c)).
1=1
At the limit T + oo, 5? depends only the context and (k 1)-history of the token t and the other summation only depends on the (k= 1)-history of the token + 1. Hence, the gradient vanishes as -> co. This gives glimpse to how the computation can be extended to a general transformer architecture.
F.2. Other Constructions for Stationary points and Beyond Contigous History
Note that we only consider estimators conditioned on the contiguous suffices. Consider again the simple example of bigram, which is conditioning on a single element in the ni-history. There can be many such estimators, for example, conditioning on the first element the n-history p(r, = -)), the second element in the n o-history p(z,, = 22 and so on. is natural to wonder if such an estimators are stationary and we can extend our results to all these possible estimators.
To generalize this, we define --MLE estimators indexed by the / S [n- 1]. For any A, define subsequence o = (Wtth)NEN N --history and the M-MLE estimator is given by,
Σ (1ly = 1 (m = i]
Σ M( = x3²}
which compares the A -hiistory of the token + 1 with the --history of all the tokens the sequence.
There is small technical challenge for transformer to implement this. To compare the V-history, must compare the (-h)-token of token i with the (-h)-token of the token (t+1) (not f) for N. This essentially requires comparing the (hoo-tnken of with the (-(h-1))-tokee of l. As a result, an induction head must attend to the (-(h-1))-token in the first layer. When N forms suffix, this step isn't explicitly necessary since h-1 will also be in N. Otherwise, it must be computed explicitly using additional heads or changing the attention matrix specifically depending on the last token.
When the inputs are of fixed length the later approach could used. We specify the construction first and then give the intuition on why would work.
29

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
(Q2) T Ve
(6)))
(36a)
1=1 WEN
K₂ Vc
(
(36b)
5=1 KEN
T-2 Σ
A
ef_, (e.)
+ - (e%_1)) for h E N,,
(36c)
l=h
S 8,3j for h EN,
VC Σ
(h)
(36d)
O.W.
The main modification is in the attention matrix of the first layer, i... now head / computes (-h) token for i o T but (-(h-1)) token for = T. The last token acts as a special token (Nichani et al., 2024) and the attention matrix is modified accordingly. The gradients of the heads will vanish similar to the proof of Theorem 4.1. However, the major drawback is that it only works for fixed length inputs and does not even work for arbitrary sequence lengths. This not a concern for the previous constructions corresponding to contiguous suffix.
G. Derivatives of the self attention map.
We calculate the derivatives of the masked self-attention map with respect to key, value, query and the input embeddings. Let E Txd be the input embeddings, Q, K, V E Rdd be the query, key and value matrices. Let denote the 7th row of For any t, the output embeddings of the th -token of the transformer layer can be written as
qt
P.V Vi..
P. exp (Kq, Qu:)
Σ exp (Kaj. Qu.)
Lemma G.1. Consider self attention map defined by
cxp (Ky, Q4i)
4
pt V.v.
where
p:
exp (Kaj, Q9:)
Define q = Ei Pege. The partial derivatives are given by,
Oyt
Σ
Ps (q @ Ia) q_ & la.
av
==
Dqt
Σ Pg (Vq;) @ vec (Qu. T) T
OK
0at
pj (Vqz) & vec (K(g) - 9)9
0Q
1=1
0a
For i # t.
= p. V + ps (V - 7)) @ (KJ Qq))
0q.
0q
P. V + pp (V(a ))) 2 (K Qa:) T
Σ
p3 (V @ (0_TQQ
Klg -
T.
0q.
j=1
30

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Proof First taking the derivative with respect the value V gives us,
dqp QV
Σ ==1
Pi (1 & Ia).
The derivative wrt to di gives us
Out
0pj
p4 V + Σ (Va)) 80
0q:
0q:
dqt
Σ
(Vai) 8
Öpj
0K
OK'
/=1
0g.
Öp:
Σ
(Va,)
iQ
aQ
To compute the derivative of p wrt q, Q, K, we begin with definition of intermediate functions,
R² + R' where g(x) exp 24
Σ expxr
1=1
h: BRddd x paxd x IRtxd + R,, where h(K, ((Kg,, Qu:))
Using the above definitions, p can be written as
= g(((k(((K00) Q))).
The partial derivative of g writes
0gj
9,1 (k- = 5) Oh9j.
Ork
ah
1sk = i) Qqi) for ftritt,
Dq:
(QQ Kqx)T 1(k =I==]] (K Q9:)
for =
1lk- = i) (K Qui) +1(i (((2Kqg)
ah
vec (Q 9:9 _T) )
0K
Oh
vec (K 9bg
aQ
31

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Using the chain rule take the derivative gives us,
apj
dgj Thg
0q:
Ohu Og:
= Σ Ipy (k= = j) PEPSl [[k = i3 (x)) T + +1(i= "
k=1
Ps [1 = i) (KTQQ)) T + +1fi = =t](QT K
=
PiPs (K T =Q))² T - -iii= = =tpps Σ pa (Q Kg)) 1=
0pj
Öpj 0hk Thk
Σ
1p;1 (k- = -) -PoPy] vec (Q 4:4%)
aK
k=1
= vyeee (Q Pj
Σ pavec (Q9:%)) "
Define q Σ Ps
k=1
0pj
pyve (Qqe(a)- --
aK
0pj
Op
Ōhk
Σ
11 (A: =
phPp) vec (K qez)))
aQ
Thg
eQ
k=1
k=1
vec (K q;9,)) T
Σ
vec (2 (Q
Pj
k=
= vec (K(a) - -)99)) T.
Substituting the following derivates,
0pj pjvec (Qqe(1) 9 7) T
0K
0p;
P (K 979 )T
QQ
dqt
0K Op,
Σ (Vq,) 8
Σ
Pj (Vq,) @ voe (Q4:(1; - 7)
OK
1=1
1=1
dat
Op,
Σ (V
8 aQ
Σ j=1
pj (V 2 vec (K(g) - ) T
eQ
1=1
Coming to the derivatives wrt g., we have,
Opj
p, [19] =i) =i) (KJ (eeTQqq) T 1+i = t)(QT K
0q.
PPpp (K T Qq)) T -1(i- =1p, Σ pu p Kqx)
32

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Writing the above expression by case basis, get,
Pr(1 Pr) (KTQui Pj (Q ((gj - 7))
when i - j== =
P5(1 - Pÿ) (K Q4i) T when i + t.
dpj
0q:
PiPj (K' Qq:)
when + j and i i t,
PPps (q:)
+ Ps (Q ((; - 3) when i =t and i i J.
First computing the derivative wrt qi for i we get,
04
pp
= V + ==1
(Va,) 3
0q.
0q:
= P + Pi (Vo) 39 (K Q4)) - pi Σ Pj (Vo)) 2 (KT (uTQu)"
j=1
V + (Vg) @ (KTQu) Q4)) T Di (VQ) @ (kQQ))
+ 3)) @ (K (QQ)))
0qtt
P V Pe (Vo)) 8 K (qqQqq) Pt Σ
ps (Va)) [K (ái) T + Σ ps ( (QTK(g 7))
0q:
1=1
p V + Pa (V 7)) 3 (K Qa)
+
Σ Ps (Va,) 8 (Q_ 1) 7)
j=1
This proves the lemma
H. Higher Order Markov Chain
This section first recalls key properties of higher-order Markov chains and their stationary distributions. We then leverage these properties to establish that the k-gram estimator has a limit as the sequence length, T, approaches infinity. Definition H.1 (Higher Order Markov Process). A Markov process pr of order k generates sequence of random variables E2.... [S] such that the conditional distribution of Ie+1 given x'' depends only on 2 i.e., for any t k
Pr
Fetl
x²)
Pr Zef1
Now, we define the transition tensor with is given by Pr. Definition H.2 (Transition Tensor). The transition tensor of a Markov process of order is a k-dimensional tensor P such that
Pu+t
it+1 x² 1¹
Now, we lift the process into the higher dimension and write simple Markov process of order 1. We create a Markov chain on the space [S]* order Markov chain with transition matrix T. Index the states by j E [S]* Now the transition probabilities for this are defined as
+ 0
i = (it, int)), for some int E [
T
yx #
otherwise
The transition probabilities are explicitly given by
T
*t+ 14
Patt.
33

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Stationary distribution. Let π E RiS!* denote the stationary distribution of the chain E, i.e.,
T π
In the summation form, for the coordinate indexed by 2 we get that,
Ty+ T Σ 5
2/ j
Note that T
1++1 j
-0 whenever j # i.. Using this, we get,
sh+1 5²) (1 (I 12 15 1 = 1 {;))
Σ
Tt T
ik+1
j = Σ Tith T
Σ Tigh -T
ittz jr { 12 = {
1
Σ
(
1
(,,
Σ
T
jet
(+)
Using this expansion, we get the following stationarity condition,
T Σ Tee
T
ikl
(37)
H.1. Generating Sequences with The Markov Chain E
In this subsection, we show the subsequences of length follow the stationary distribution when generated from the Markov Chain E. To generate sequences from the Markov Chain E, we do the following generation,
a) Generate sequence of length l sampling from the stationary distribution π (say i*)
b) Sample the next token from the distribution Prt,.)) say i4t1
After generating sequence if of length k, the next token is generated by the sampling from the distribution
P (.))
Lemma H.3 (Stationarity of k-tuples), Let x' be the random variable representing the sequences generated E, the distribution of subsequence 2i... of length lo is given by n.
Proof We will prove this using induction. For the base case, observe that for 1 k, the value sampled from æ by construction. Use the induction hypothesis, any & length subsequence x obeys the law given by n. remains show that 4:44 distributed according π.
Pr (xit1
=1+ 1 = Σ
Pr (x- = i(+1),
--b+1
Σ k11
pr (x' = (x'=
Fies 144
x' =1)
Using the Markov property of the sequence generation, we have,
Pr I141
x' =1') = p-
1i+s in Il )
34

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Substuting this expression, have,
Σ Py (a' = = 2141 = itt zlens (
Pr (2/+1 i/41 ()(212) =
Σ
Σ
Pr (x' = =1' pr Iits hns Fi-e4s time))
11-++1 p1-6
= Σ
Σ Lit-k
pr (x' = =it) pr (Si = his af &/ ((
ii-++1
Using the fact that -x Pt (I' = = 1¹) = pr (3/ = 82 ()))))
Using the induction hypothesis, we have that
Pr (1) h11 o x11) k ==41"
pr (zi 1112 (/-b4a)
Σ
Tef -
Pr
Tets hess al-ets = )
4--11
Σ
Tsf lokk1
T
l-k+2
yy (
11-4+1
Using the Eq. (37),
Σ Tel T
Pr (T
i1-441
This proves the hypothesis for length / 1. Hence, by induction, the hypothesis holds for any length.
Lemma H.4 (Stationarity of sub-k-tuples), For any > 0 and & k, we have the following shift invariant property for the marginals,
Pr (x' 1') Pr (x! =²))
Proof Using Lemma H.3, we have that the distribution of the subsequence xfft-11 given by i for any > 0, i.e.,
(x* = =²² Pr (I.th-1 =i²) Tta.
Summing the above expression over all possible values i//4, we get,
Σ Pr (x" = =²²)= = Σ 1 Σ Tpp.
$11
1211
We know that
Σ pr (xx = =²) = Pr (x' = 1¹),
1+1
Σ Pr (2.th-1 1²) Pr (x/- i')
This proves the lemma.
35

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
H.2. Lower-order Conditional Probabilities
The following lemma show that the lower order conditional probabilities are time homogenous and are given by the stationary distribution Lemma H.5. For the sequences generated by the markov chain for & ki, we have,
Σ T
Pr xr+1 =, 211 = =8) = Pr I141 = x' =
Proof For / = k., this holds due the Markov property of order k and
Ikll (3N11 =. x² =t 0+)
Pr TT =.
2- 1²) pr
Now, for 1 < k; we have,
Pr (I: = =2112)
Pr TT+ h441 x- = 8
Pt (xT = =8)
Using the shift-invariant property of the marginals, we have that,
Pr (1 T-1 in1 =x+tt) = pr (xlt+=10t1) = Σ
=
Tst
1+2
Pr (xz 111 =)) pr (x' =1)) = Σ
Tgh.
1411
This proves the lemma.
Conditional Probability for Any Subset of Sequence. Previously, we defined it for contiguous history; now, we can extend the definition to any subset of sequence as well. Consider any set K C [0, S
1]. [S] - is an ordered set defined as
{k- i K}. We define the sequence r' K as the ordered set {x, : i E [S] K}. We define the conditional probability
of the next token given the sequence 2:(5-
PIS)-K (119)- ((+11) = Pr (xx-1 let1 |0 = |=10)-K_1101-8- ((
Pr (Dezs i r(5)-K ((0)--x -x)
Pr (x18)- = [1
Σ Pr (xrt12 = jut1
Ky
Σ Pr (xx = i²)
[[8)-K K)
Consistency of the sub-k counting estimators. Here, we use the ergodic properties of the chain E to show that the lower order estimators convergence as T - oc.
Lemma H.6. Consider the 1+1-gram estimator for 1 & k - 1, then as T + oo, pi(x") + Pr (-Ix").
Proof Using the ergodicity of the chain E (Penev, 1991), the 1--mmmm estimator converges the stationary distribution as it equivalent to counting the empirical frequency on the chain of higher order. Now we write the I-gram estimator in the following way:
ET Ha/ti = i
Σ/= 113-011 )
36

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
Now can write the following expression
fzftin - =
1 1 = 11
j=1
Σ
Σ
1
[x |+/ =1
F j=l
Now, using the ergodicity of the lifted markov chain, we can say that
Ej=1 1a/tthr
=
T
Σ == (2f
= 12
20000
T
12
Using the similar argument, the ratio can now be written as
E 1aftii, = 1 Eg Tgtt
pr (2141
Σ/=1 1ax 111 it 1
En teat
The last equality flows from the lemma H.5.
I. More Experiments
I.1. Experiments with Attention Only Transformer
In this section, we repeat the main experiment in the paper with the common non-disentangled transformer architecture, demonstrating the generality of our results. Figure 5, shows the evolution of attention matrices during training with both one-hot and learned embeddings for the general architecture.
For the transformer experiments, we use vocabulary size =. 3. The length of the input sequences is T 32, and the sequences are sampled from an in-context tri- tri-gram language model, i.e., n The transition matrices for fixed context sampled from a Dirichlet prior with a 1. The embedding dimension set d S + T, to be consistent in both one-hot learned embeddings for both positional and semantic embeddings in different experiments. The transformer is trained with Adam with a weight decay of 0.0001 for 4096 iterations, with a constant learning rate of 0.005 and a batch size of 128. The test loss is evaluated over 2 test sequences.
setups, and
1.2. The Contiguous Solutions are Preferred during Training
In Figure 6, we repeat the procedure in Figure 3 for different seeds and plot the attention heads in the second plateau. Through all the experiments transformer more often attends to the (-1)-token first, rather than the (-2)-token. Recall that the underlying mechanism behind the sub-n-gram estimators is checking if the history of the token æ741 and æy matches and adding žj if they do. We conjecture that, since the token xT is always provided through the skip connection, is casier learn to match Fr with 1 than, for example, Tr-1 and 2--2-
L.3. Plot of norms of the gradients along the trajectory
In Figures 7, 8, we plot the norm of the gradient during training along the similar lines as (Odonnat et al., 2025). The plots demonstrate how the norm gradients stays low during the plateau stages and spikes during the jump between the plateaus.
37

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
1.60
0.09
Loss
Gradient
0.08
1.55
1.50
0.06
Gradient Norm
Test Loss
1.45
0.04
1.40
1.35
0.02
1.30
0.01
0.00
0
2000
4000
6000
8000
10000
12000
14000
16000
Step
Figure 7: Norms of gradients for simplified transformer. A two layer simplified transformer for S = 5 and sequence length of L = 128 from a 3-gram language model. The plots show that the norm of gradients stays low during the plateau stages and spikes during the jump between the plateaus.
0.8
1.62
Loss
Gradient
0.7
1.60
0.6
Gradient Norm
0.5
Test Loss
1.58
0.4
0.3
1.56
0.2
1.54
0.1
0.0
2000
4000
6000
8000
10000
12000
Step
Figure 8: Norms of gradients for transformers with MLP layers. A two layer simplified transformer for S 5 and sequence length of L 64 from 3-gram language model. The plots show that the norm of gradients stays low during the plateau stages and spikes during the jump between the plateaus. However, there is only single intermediate plateau, unlike attention-only transformers, as a single head attends to both the (-1, too-ns,, and it emerges after the platcau. The plateau corresponds to the unigram.
38

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
1.0
1/28
h
Attention Scores
2.44
300
1.00
0.58
800
200
400
650
800
115
t= 5000
16000
Step (t)
Training Step (1)
(a)
(b)
3.08
50
Attention Scores
3.04
200
3.00
h 2
0.08
800
0.98
200
400
620
800
== 115
== 5000
f 16000
Step (t)
Training Step (1)
(c)
(d)
Figure 5: The evolution of the attention heads the first layer during training of an attention-only transformer with one-hot (a-b) and learned embeddings(c-d), (a-c) Progression of the test loss during training. The highlighted points are the iterations on the plateaus for which demonstrate the attention matrices. (b-d) The evolution of attention scores of the heads of the transformer during training representing the tokens it is attending. First, both of the attention heads attend to all the previous tokens uniformly, i.c., the induction heads are not formed. At the second plateau, they both attend to the previous token, one head is not formed yet while the other attends to the previous token. Finally, the model escapes this plateau, the second attention head learns to attend to (-2)-token at the end of training.
39

---

Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points
h=1
Attention Scores or (A
h=2
2
Random Seed (s)
Figure 6: Attention maps of the two heads in the second plateau for different random seeds denoted by s. It shows how the transformers attends to (-1)-token first and never attends the (-2)-token before attending (-1) (-1)-token across random seeds.
40


Reviews:
[{"review_id": "m9I2TpnAb2", "reviewer": "Reviewer_xuc4", "summary": "This paper contributes the following.\n\n1. A sufficient condition for the population cross entropy loss to vanish in\n   the setting of $n^{\\text{th}}$-order Markovian sequence modelling.\n   The condition is based on expressing the derivative of the model's logits\n   independently of  $n-1$ tokens in the input sequence, plus some other\n   conditions.\n\n2. A construction of a transformer architecture and parameter that, in the hard\n   attention limit (i.e. infinite weight norm limit), implements an in-context\n   $k$-gram estimator.\n   The architecture is a two-layer attention-only transformer with a particular\n   distribution of heads and some frozen weights.\n\n3. (Main theorem) Putting these together, a proof that the constructed transformer parameters\n   from (1) meet the sufficient condition from (2) in the infinite weight norm\n   limit and the infinite sequence length limit.\n   They establish that:\n\n   1. such infinite (size and norm) parameters are first-order stationary\n      points of the population cross entropy; and, moreover,\n   \n   2. in the finite regime, the sufficient condition is approximately\n      satisfied, such that the constructed parameters are 'near-stationary'\n      (gradient close to zero).\n\n4. A brief discussion in the main text about how the construction and proof\n   could potentially be generalised to a slightly more standard transformer\n   architecture (with one-hot positional encodings and without frozen\n   parameters). There is a longer discussion in appendix F.\n\n5. A brief discussion in the main text about how the sufficient condition\n   generalises to non-contiguous sub-$n$-grams\n   and how the main theorem could potentially be generalised to these cases.\n   There is a longer discussion in the appendices.\n\n6. A brief numerical experiment shows that training a transformer with the\n   constructed architecture exhibits several 'plateaus' separated by 'phase\n   transitions' in the loss over training.\n   \n   The loss of the plateaus roughly lines up with the loss of idealised\n   $k$-gram estimators for increasing $k$, and during the plateaus the\n   first-layer attention patterns as predicted in the construction for these\n   values of $k$.\n   Therefore, the experiments show that the transformer's training trajectory\n   proceeds through these 'near-stationary' points.\n\nI wrote a pretty long review. Therefore, I include a summary of my review as\nfollows.\n\n* I think this paper has the potential to make a substantial contribution to\n  the field because it represents a clear analysis of a simplified example of a\n  phenomenon of current interest.\n\n* I didn't check all of the technical details of the proofs, but the results\n  seem plausible to me based on my limited study (so far).\n\n* However, I identified two major issues with the paper on the basis of which I\n  have initially recommended 'weak reject:'\n\n  1. The paper is insufficiently clear about the distinction between\n     contributions 3.1 and 3.2 in my list above.\n  \n  2. The presentation of construction (contribution 3 above) contains numerous\n     technical errors.\n\n  I think these issues limit the accessibility of the contributions of the\n  paper and make it not yet ready for publication.\n\n* I detail some other minor issues and recommendations in the remainder of my\n  review. If the authors show they can address the above clarity concerns, I'd\n  be willing to raise my score.\n\n**Edit to add:** Rebuttal addressed my two major concerns. Raising score to 3 for now.\n\n**Edit to add:** After further discussion with authors and wvkz, wvkz raised score to 3. I see no remaining reasons to reject the paper (assuming the revisions proposed by the authors are correct and included in the final version). Therefore, I am raising my score to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and careful review of our paper. We sincerely appreciate the time and effort you have put into reviewing our work and providing constructive feedback. Your comments have helped us identify areas where we can improve the clarity and precision of our presentation. Below, we address the main concerns you have raised and outline the changes we have made in response.\n\n**Distinction between the Finite and Limit Claim**  We give an updated version of Theorem 4.1 that captures the non-asympotatics and make the theorem inline with the title and referenced contributions of the paper.\n\n**Theorem**: For the distangled transformer $p_{\\theta}$, the gradients at $\\theta^k_*$ are given by \n\n$$ \\partial\\_{\\theta = \\theta_*^k}  L(\\theta) = \\sqrt{c}   {  \\mathbb{E}\\_{p\\_{\\tau} \\sim \\mathcal{P}}   \\mathbb{E}\\_{x^T \\sim p\\_{\\tau}}  \\color{green}{   \\| \\hat{p}\\_k - p\\_{\\tau}(. \\big| x^{T}\\_{T-k+2})  \\|^2} } + { \\color{blue} O( \\sqrt{c} T e^{-c} )  }.$$ \n \nThis error contribution for finite $c,T$ is given by the term above. The blue term represents the error due to finite $c$, while the green term accounts for the finite length and non-stationarity of the chain and vanishes as $e^{-\\Theta(T)}$ ( Penev 1991). The limit claim now holds when $c = \\Theta(T) \\to \\infty$.  We hope this clarifies the distinction between the finite and limit cases. \n\nWe apologize for being lenient with the distinction between near-stationary and stationary in the text. We agree with the reviewer's concern and will revise the paper to clarify this nuance, explicitly indicating wherever referenced that **stationarity holds only in the limit**. Similarly, we make sure to refer to the points as \"near-stationary\" instead of \"near stationary\" throughout the paper. \n\n**k-gram construction.** Please see the response to reviewer 6GFt for the formal construction. \n\n**A tool to analyse landscape.**  Analyzing the stationary points of cross-entropy loss is inherently challenging, particularly compared to regression settings, where the loss landscape is often analytically tractable. Prior work has largely focused on regression due to the tractability, leaving a gap in understanding for sequence-based tasks that rely on cross-entropy loss.\n\nIn this proposition, we leverage the factorization of sequence probabilities (e.g., via chain rule or autoregressive decomposition) to derive stationary points. Specifically, we demonstrate how models can converge to stationary distributions by estimating only one factor of the joint probability at a time. This provides a tractable pathway to analyze optimization dynamics in settings where closed-form solutions are otherwise infeasible, see examples given in \"practicality of proposition 3.2\" section for Reviewer wvkz. \n\n**Other concerns.** \n10-11.  The construction for general attention-only transformer is not given as it mostly follows on the similar lines as the simplified transformer, hence the neccesary changes and proof sketches were informally given in the Appendix F. We will give a precise formal statement for this case with a proof in the appendix. \n\n**Other Comments or Suggestions.**\n\n1-5,8-9: We have corrected the typos you identified and carefully proofread to make sure we fixed the inconsistencies in the paper.\n \n6: We used the term context with a slight abuse of terminology because ${p_{\\tau}}$ is the underlying generative process governing the task instance, and determines the distribution of the context $x_1^T$. For clarity we refer to it now as context/task distribution and clarify it in the previous section. \n\n7: We have fixed the condition to almost sure equality.\n\n10: We have changed the main question to \"Why does training linger at plateaus before developing such abilities?\" and editted the introduction in alignment with this point. \n\n14: Recall that the underlying mechanism behind the sub-$n$-gram estimators is checking if the histroy of the token $x_{T+1}$ and $x_j$ matches and adding $x_j$ if they do. Our conjecture is that, since the token $x_T$ is always provided throught the skip connection, it is easier to learn to match $x_T$ with $x_{j-1}$ than for example $x_{T-1}$ and $x_{j-2}$. \n\nWe believe that the revised version effectively addresses your concerns and presents our contributions more clearly. Given these improvements, we hope that you will consider raising your evaluation of our work.", "claims_and_evidence": "**Sufficient condition.**\n\nI checked carefully the proofs for Lemma 3.1 and Proposition 3.2. They appear\nto be sound.\n\n1. In several places, the paper refers to this condition as a \"characterization\n   of the stationary points of the cross-entropy loss\" or similar. This\n   includes the contribution list in the introduction. It seems to me that this\n   terminology risks creating confusion, since a \"characterization\" in\n   mathematics to me suggests a necessary *and* sufficient condition, whereas\n   the paper gives only a sufficient condition. Terms such as \"sufficient\n   condition\" or \"criterion\" might be appropriate, to avoid giving the false\n   impression that a necessary condition has been obtained.\n\n2. Line 045R: \"a powerful tool for analyzing the structure of the loss\n   landscape for various tasks including but not limited to $n$-grams.\" It's\n   not clear to me to what this refers. Is it just non-contiguous\n   sub-$n$-grams? Or is something more general meant?\n\n**$k$-gram simple transformer construction.**\n\nI have not fully verified that the construction implements the $k$-gram\nestimators, but it appears to be sound based on what I have checked. However, I\nthink the clarity of the construction requires attention.\n\n3. I attempted to follow the description construction in the main text, and\n   uncovered numerous technical errors. Each seems relatively minor and I\n   believe a correct construction could be recovered. I list specific issues I\n   noticed here.\n\n   1. Equation (7) defining\n        $A_1^{(h)} = c \\sum_{l=h}^T e_{l-h} e_l^\\top$:\n      This definition yields an upper-triangular matrix, but I think you want a\n      lower-triangular matrix. Also, $A_1^{(h)} \\in \\mathbb{R}^{T \\times T}$,\n      suggesting we want $e_i \\in \\mathbb{R}^{T}$, in which case $e_T$ is not\n      defined (since we start with $e_0$.\n\n      Best I can tell, you probably meant\n        $A_1^{(h)} = c \\sum_{l=h}^{T-1} e_{l} e_{l-h}^\\top$.\n\n      This same issue appears to arise in other constructions too.\n\n   2. Equation (7) defining\n        $V_1^{(h)} = \\sqrt{c} \\sum_{j=0}^{\\mathcal{S}-1} s_j s_j^\\top$:\n      Since $s_j$ was defined for\n        $j \\in [\\mathcal{S}] = \\{1,\\ldots,\\mathcal{S}\\}$,\n      I think this is meant to be a sum from $j = 1$ to $\\mathcal{S}$.\n\n      This same issue appears to arise in almost all other sums over $j$,\n      except in equation (15), where it's $\\sum_{j=1}^{\\mathcal{S}-1}$.\n\n   3. Equation (9) neglects the role of the softmax. Perhaps some discussion is\n      needed to explain that this equation is meant to hold in the hardmax\n      limit. Though, see the next two issues.\n\n   4. Equation (9) appears to be out by a factor of $\\sqrt{c}$ (the coefficient\n      of the value matrix).\n\n   5. Equation (9) involves terms that are undefined, namely\n        $s^\\top_{x_{t-n+1}}$ for $t < n$,\n      or\n        $s^\\top_{x_{t-1}}$ for $t=1$.\n      In place of these terms, depending on the corrected definition of\n      $A_1^{(h)}$, there should be something like a uniform mixture of the\n      initial few tokens, or simply $s^\\top_1$.\n\n4. In my opinion, the construction should be defined in full, correct detail in\n   the main text, and its correctness should be stated formally as a numbered\n   theorem/proposition. This would support easier verification of the\n   construction's correctness.\n\nAppendix B.1 presents a more detailed description of the construction, with a\npartial formalisation of the correctness claim (lemmas B.1 and B.2). This\nversion seems to be free from some of the technical gaps in the main text\nversion, but not all of them. I haven't verified the correctness in full detail\n(but, once again, it appears to be sound).\n\n\n**$k$-gram simple transformer satisfies sufficient conditions.**\n\nI have not attempted to verify the soundness of theorem 4.1. I didn't notice\nany obvious problems in the proof sketch in the main text. I haven't read\nappendices B.2, D, or E containing the full proof.\n   \nIn my opinion, the implications of this theorem are not clearly articulated in\nthe paper. In my summary, I distinguish between two claims made by the paper:\n\n* **Limit claim:** Infinite (size and norm) parameters are first-order\n  stationary points of the population cross entropy.\n\n* **Finite claim:** In the finite regime, the sufficient condition is\n  approximately satisfied, such that the constructed parameters are\n  'near-stationary' (gradient close to zero).\n\nThis distinction is essential to understanding the contribution of this paper.\nTheorem 4.1, as stated, establishes the limit claim, but the titular\ncontribution of the paper is the finite claim. However, the distinction is not\nmade clear enough in the paper in my judgement.\n\n5. The finite claim is only ever made explicit informally (in the paragraph\n   below theorem 4.1---line 301R---\"The result holds asymptotically as\n   $T,c\\to\\infty$,\" where I take 'asymptotically' here to mean that the\n   gradient of constructed parameters approaches zero as the parameters\n   approach the infinite norm and infinite length limit). As I understand, the\n   proof is implicit in the proof of theorem 4.1. I believe this additional\n   conclusion and its proof should be made explicit in formal detail.\n\n6. The limit claim and the finite claim are conflated in the rest of the paper.\n   The passage quoted in the previous point continues, \"leading us to term the\n   $k$-gram estimator as *near*-stationary.\" However, up to and after this\n   point, the paper normally just describes the constructed parameters as\n   \"stationary\" in a technical context without appropriate qualification that\n   they are exactly stationary only in the limit:\n\n   1. Section 4 opening (line 181R): \"... showing that the sub-$n$-grams are\n      stationary.\" Probably this should say they are \"near-stationary.\"\n   \n   2. Section 4.3 opening (line 277R): \"We show that the $k$-grams constructed\n      ... are first-order stationary points ...\" It should probably say that\n      are \"near-stationary points.\"\n   \n   3. Theorem 4.1 statement (line 289R): \"i.e., it is solely a function of the\n      $k$-history and context as $T\\to\\infty$.\" That should be \"almost solely\",\n      if I understand correctly.\n\n   4. Line 365L: \"the derivative is supported only on tokens in $M_t$.\" If I\n      understand correctly, it would be more correct to say the derivative is\n      \"approximately\" only supported on such tokens, or \"most of its support\"\n      is on these tokens, etc.\n\n   5. Line 362R: \"As it is a stationary point,\" should probably read\n      \"near-stationary.\"\n\n   6. Conclusion (line 415R): \"... the stationary points encountered along the\n      training trajectory.\" It should probably read \"near-stationary.\"\n\n   There are possibly more instances in the main text or appendices.\n\n[CONTINUED IN NEXT SECTION \"METHODS AND EVALUATION CRITERIA\"]", "methods_and_evaluation_criteria": "[CONTINUED FROM PREVIOUS SECTION \"CLAIMS AND EVIDENCE\"]\n\n7. Generally when the role of limits are discussed early in the paper, the\n   discussion could be more precise. A non-exhaustive list of specific\n   examples:\n\n   1. The claim in the paper's title, \"Sub-$n$-grams Are Near Stationary\n      Points,\" is not an accurate description of what is proven. Due to the infinite weight norm required, the\n      constructed transformers are actually infinitely far from the identifies\n      stationary points (even in the infinite length limit). \"Near-Stationary\"\n      or \"Nearly Stationary\" would be more appropriate than \"Near Stationary\".\n   \n   2. There is no mention of the infinite norm limit in the abstract (arguably\n      acceptable if the discussion were more precise elsewhere).\n   \n   3. Contribution list item 2 (line 050R): \"solutions representing\n      sub-$n$-gram estimators are stationary points,\" needs qualification.\n   \n   4. Later that same paragraph, \"... satisfy the conditions sufficient for the\n      gradient of the population loss to converge to zero presented in section\n      3.\" In my understanding, the condition in section 3 does not discuss\n      convergence to zero gradient, but characterises the gradient and shows\n      conditions under which it equals zero.\n\n   In my opinion, it would be better to spell out the limit claim and the\n   finite claim on page 1 (perhaps this would be a good opportunity to make it\n   clear what \"near-stationary\" means).\n\nNote that my concern about this contribution is not that the finite parameters\nare not exact stationary points per se. Approximate stationarity is sufficient\nfor the discussion concerning plateaus. It's just important that the paper\navoids giving the false impression that the finite parameters are exactly\nstationary, for example.\n\n**Numerical experiments.**\n\n8. The experiments clearly show that the loss of the plateaus roughly lines up\n   with the loss of idealised $k$-gram estimators for increasing $k$, and\n   during the plateaus the first-layer attention patterns as predicted in the\n   construction for these values of $k$. I think this is certainly suggestive\n   that the transformer's training trajectory proceeds through these\n   'near-stationary' points.\n\n9. I also followed the claims and figures in appendix I and found them\n   reasonable.\n\n---\n\nI also discuss some claims that were of some interest to me but were mainly\ndiscussed in the appendices, as follows.\n\n**Discussion of general transformers.**\nIn general, as architectures change (e.g. weights are unfrozen) the gradient\nmay cease to vanish because of the new dimensions. Section 4.5 and appendix F.1\npertain to how the construction and proof of near-stationarity can be extended\nto a more realistic transformer architecture.\n\n10. The paper claims the construction itself can be easily extended. This claim\n    seems unproblematic to me.\n\n11. It is unclear to me whether the paper claims that theorem 4.1 would\n    continue to hold in the new construction, without diving into the technical\n    details. Appendix F.1 says \"we discuss how the results can be extended,\"\n    but there is no precise statement of the extension that has been achieved\n    (i.e., numbered theorem that makes the claim explicit, with proof that\n    could be verified).\n\n\n**Discussion of general non-contiguous sub-histories.**\n\n12. If I understand correctly, the appendices claim that the sufficient\n    condition generalised to non-contiguous sub-histories (remark C.2, though\n    there is no precise statement); that the construction can be generalised\n    (appendix F.2, again no precise statement); and that some elements of the\n    proof of theorem 4.1 can be ported to the new construction, but a full\n    generalisation to this context is left as an open question. I haven't\n    verified the mathematical details, but this seems plausible to me.", "theoretical_claims": "See above.", "experimental_designs_or_analyses": "See above.", "relation_to_broader_scientific_literature": "A current priority for this field of \"science of deep learning\" is to\nunderstand the dynamics of emergent structure in transformers. Before we can\nmake progress understanding large-scale models with hundreds of billions of\nparameters, we can develop hypotheses in smaller-scale systems. This paper\noutlines a pretty neat setting for principled empirical and theoretical work in\nthis direction.\n\n1. It helps a lot to isolate clean, small-scale examples of phenomena of\n   interest, such as stage-wise learning. This paper shows a particularly crisp\n   example of stage-wise learning in higher-order Markovian sequence\n   prediction. Such was already established by Edelman et al. (2024), however\n   this work further simplifies the setting through their experiments with the\n   disentangled transformer architecture, which is even simpler than the\n   transformers studied by Edelman et al. (2024), but still exhibits stage-wise\n   learning.\n\n2. It also helps to have a detailed, mechanistic understanding of how\n   particular parameters implement particular in-context learning algorithms.\n   The paper contributes a pretty simple construction for a range of\n   interesting solutions (sub-$n$-grams). If I understand correctly, this\n   construction is original, though I am not familiar with the prior work on\n   the disentangled transformer.\n\n3. The main result itself, that these sub-$n$-gram transformers have\n   approximately zero gradient, is a pretty interesting observation that fills\n   an important piece of the puzzle of the emergence of structure in this\n   setting, especially if it does turn out to be true for more general\n   transformers.\n\nIn summary, I think this paper's contributions are valuable and of interest to\nthe science of deep learning literature, and could be a step towards a deeper\nunderstanding of the phenomenon of emergent structure in transformers.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths.**\n\nAs I said above, I think this paper has the potential to be quite a valuable\ncontribution to the field of \"science of deep learning.\"\n\n**Weaknesses.**\n\nAs I said, the paper's contributions are somewhat marred by technical issues\nand, in places, a lack of precision and clarity regarding the exact\ncontributions.", "other_comments_or_suggestions": "I encountered several minor grammar issues and punctuation/notational/technical\ninconsistencies while reading the main text. I think the manuscript would\nbenefit from careful proof reading to suit a more polished conference paper. A\nlist of the issues I noticed is as follows.\n   \n1. Line 089R: I suppose $R^k$ should be $\\mathbb{R}^k$.\n\n2. Throughout there are several terms involving hyphens (such as \"$n$-gram,\"\n   \"$(-l)$-history,\" \"$(-l)$-token\") for which the hyphens seem to be, in some\n   cases, formatted as en-dashes or minus signs. The authors can of course\n   choose their preferred punctuation but I invite them to choose one\n   preference and use it consistently.\n\n   Relatedly, in rare cases, \"$n$-grams\" etc. are improperly formatted as text.\n\n3. Line 110R: I suppose \"the $x_T^{th}$ token\" should be either \"$x_T$\" or \"the\n   $T^{\\text{th}}$ token.\"\n\n4. Line 110R: I suppose \"$n$-gram estimator\" should be \"The $n$-gram\n   estimator.\"\n\n5. Line 145R: $Q$ and $K$ here are transposed in comparison to definition 2.2.\n   The same is true for $Q_2$ and $K_2$ defined implicitly through equations\n   (5).\n\n6. Proposition 3.2 statement, part (a): In what sense is $p_\\tau$ a \"context\"?\n   I wasn't able to make sense of this term in this setting.\n\n7. Proposition 3.2 statement, part (b): Since $p_\\tau$ is a random variable, I\n   believe the condition should specify that\n     $p_{\\theta_\\ast}(x^T) = p_\\tau(\\cdot | x_t^T)$\n   holds almost surely for $p_\\tau \\sim \\mathcal{P}$.\n\n8. Line 364R: Missing period.\n\n9. Remark C.2, line 1018: \"conditioning on the any subsequences\" typo.\n\nI also take minor issue with the phrasing at several points in the paper's\ndescription of the motivation and results, and I invite the authors to\nreconsider these parts or add further justification.\n\n10. Line 013R: the motivating question is \"why must training go through plateaus\n    to develop such abilities?\" It seems to me that this question does not have\n    that much relation to the contributions of the paper. The paper establishes\n    that approximate stationary points exist, but their existence does not imply\n    that training \"must\" spend any time near such stationary points. A more\n    pertinent question, in my opinion, would be \"why might training linger at\n    plateaus on the way to its eventual solution\" (answer: there are partial\n    solutions that are stationary points); or perhaps \"what parameters comprise\n    these plateaus?\" (answer: this paper suggests they are sub-$n$-gram\n    estimators).\n\n    Similarly, line 026R: \"a theoretical foundation for why plateaus\n    emerge\"---is this really what the paper contributes? It seems to be that the\n    existence of plateaus and the fact that, once in a plateau, training\n    stagnates, is separate from why training ends up attracted to these\n    solutions on the way to the improved solution in the first place.\n\n11. Lines 010R and 081L: References to \"circuits, such as induction heads.\" In\n    my understanding, the induction head is only one part of the induction\n    circuit studied by Olsson et al. (2022). The \"circuit\" also involves\n    previous-token heads and their composition with induction heads, and is not\n    synonymous with the induction heads themselves.\n \n12. The abstract refers to these numerical experiments as \"comprehensive\" which\n    appears to oversell the magnitude of the experimental sections.\n\nFinally I collect some miscellaneous comments.\n\n13. I found the comment at the end of section 4.3 regarding the average loss\n    over sequences of varying length interesting. However, I think it\n    highlights how results in the infinite $T$ limit are limited in their\n    ability to offer insight into training in this setting. For any sequence,\n    early in training, any $n$-gram or sub-$n$-gram estimator won't have\n    converged to the stationary distribution, and so, predictably, the\n    conditional distribution won't align with the estimator's predictions and\n    the sufficient condition for zero gradient will not be satisfied.\n\n14. I found the comments in section I.3 interesting. Do you have any\n    theoretical insight into why this might be the case? If so, it could be\n    valuable to add to the appendix.\n\n15. I invite you to reconsider the inclusion of definition 2.2 in the main\n    text. Since the focus is on the simplified architecture, it might make more\n    sense to describe the simplified architecture in more formal detail in\n    place of this.\n\n16. For what it's worth, I don't feel like there is adequate simplicity payoff\n    in leaving some details to the appendix such as positional encoding,\n    summation vs. concatenation, and value matrices. The construction and proof\n    in the main text is only marginally simpler than the full construction.\n\n    If you want to prioritise simplicity, it might make sense to devise an even\n    simpler architecture with infinite sequence length and hard max attention\n    built into the architecture, and present the construction and proof in full\n    detail in the main text for this architecture, then state the construction\n    and the asymptotic version of theorem 4.1 for the full architecture with\n    softmax, value matrix, MLP, etc. in the main text and provide the detailed\n    proof in the appendix.\n\nOf course, all of these suggestions are offered to the authors based on my\nimpressions, could be based on misunderstandings, and are only meant for their\nconsideration.", "questions_for_authors": "As I said, I think this paper makes a valuable contribution, yet I have\nrecommended weak reject. The reason is because I think there are a number of\nissues that get in the way of the paper being accessible. The most serious ones\nare the apparent technical errors in presentation of the construction in the\nmain text (claims, section, item number 3 above) and the distinction between\nthe \"limit case\" and the \"finite case\" (claims section, items number 5 to 7\nabove). These issues are at the crux of my recommendation. I invite the authors\nto focus their response on these issues. If they can address my concerns, one\nway or another, I would be pleased to increase my rating.", "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": "02dTC6qUiL", "reviewer": "Reviewer_6GFt", "summary": "This paper explores the loss landscape of next-token prediction in a synthetic setup, where the model is trained to learn the transition probabilities of a Markov chain of order $n$ in-context. The authors use a simplified two-layer transformer (disentangled transformer) and theoretically analyze the population cross-entropy loss, showing that 1) there exists a parameter configuration of the model that is equivalent to a sub $k$-gram model ($k<n$), 2) any sub-$k$-gram is a stationary point of the loss in the asymptotic regime of context window $ T \\to \\infty $. They empirically verify these findings, showing that the plateaus in the test loss progression reported in Eldeman et al. (2024) correspond to sub-$k$-grams based on the attention maps.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for bringing Chen et al. (2024) to our attention, which we unfortunately overlooked in our discussion of related work.\n\n**Comparison with Chen et. al. 2024.** The paper examines the same task using a disentangled transformer, with the primary differences being a **three-stage training procedure** and **the inclusion of a specific FFN layer**, whereas we focus on an attention-only transformer. The key differences are as follows: \n\n- **Selection of History:** In their approach, the selection of previous $n-1$ tokens for $n$-grams is performed by an FFN layer. In contrast, in our model, this selection is handled by the attention heads in the first layer. \n\n- **Stage-Wise Dynamics:** A crucial difference is the absence of stage-wise dynamics in the analysis of Chen et al. This can be attributed to their three-stage training process and, in particular, their initialization scheme (Assumption 3.3), which ensures that different heads attend to different tokens from the start. Such an initialization eliminates stage-wise learning dynamics. In contrast, as seen in our experiments, for general initializations attention heads progressively learn to attend to new tokens at each stage. However, the initialization used by Chen et al. almost pre-configures the heads to attend to the required tokens from the outset, effectively removing the stage-wise learning process. However, Chen et. al. studies the training dynamics where our paper just studies the properties of the landscape. \n\nWe will add this detailed comparison to the paper. \n\n**Comparison with Rajaraman et. al. 2024.** The method for constructing $n$-grams aligns with Rajaraman et al. (2024). However, the construction of $k$-grams (sub-$n$-grams) introduces a subtle but important modification to their  $n$-gram formulation. This difference stems from the **activation state of attention heads**—specifically, whether heads are active or inactive.  This nuance is critical for interpreting the **stage-wise dynamics** observed in experiments: unlike $n$-grams (which rely on all heads), $k$-gram formation depends on *which heads are active at intermediate stages*, as heads progressively activate during training. \n\n**Additional Experiments and General Comments on the Setting.**  We thank the reviewers for this suggestion and we provide additional experiments with the gradients for simplified and regular architectures at https://anonymous.4open.science/r/icml_rebuttal_2025-C11E/main.pdf. The reason for the choice of sequence length $T=32$ for the plots is due to the impracticality of showing larger attention maps on paper. Please see the attached figures, where we provide the loss curve and the gradient norms not only for the simplified model with longer sequence length and larger vocabulary size but also for a 2-layer transformer architecture **involving mlps**. \n\n>Reviewer xuc4\n\n**$k$-gram Construction:**\nTo provide an intuitive understanding, we initially presented it in an informal manner. However, we recognize that this approach introduced several minor technical inaccuracies. To ensure rigor, we complement with the additional lemmas for $k$-gram construction which are partially stated in the App.. \n\n$$ \n\\begin{align*}\n    A_{1}^{h} &= c \\sum_{l=h+1}^{T} e_{l} e_{l-h}^{\\top} + c \\sum_{l = 1}^{h} e_l e_1^{\\top}, \n    \\quad V_{1}^{h} = \\sum_{j=1}^{S} s_{j} s_{j}^{\\top}, \\text{ for } h \\leq k-1, \\text{ o.w. } 0  \\\\\\\\\n    K_{2} &= \\sqrt{c} \\sum_{j=1}^{S}  \\sum_{h = 1}^{k-1} s_{j}^{h} \\left(s_{j}^{h}\\right)^{\\top}, \\quad \n    Q_{2} = \\sqrt{c} \\sum_{j=1}^{S}   \\sum_{h = 1}^{k-1}  s_{j}^{h} \\left(s_{j}^{h-1}\\right)^{\\top}.  \n\\end{align*}\n$$\n\n**Lemma 1** For the parameters defined, \n(a) The attention score $a\\_{(i,j)}^{(1,h)}$ between key and query elements $(i,j)$ of head $h$ of layer 1 is \n    $$\n    \\begin{align}\n         a\\_{(i,j)}^{(1,h)} = \\begin{cases} 1 - O(i e^{-c}) \\text{ when } i > h \\text{ and } j = i-h, \\\\\\\\ \n         1 - O(i e^{-c}) \\text{ when } i \\leq h \\text{ and } j = 1, \\\\\\\\\n         O(e^{-c}) \\text{ o.w.}. \n         \\end{cases}.\n    \\end{align} $$\n\n(b) Considering $M_k^t$ defined in the paper, the attention score $a^2\\_{(t,i)}$ of layer 2 between elements $i,t$ is\n$$\\begin{align*}\na^2\\_{(t,i)} = \\begin{cases}\n    \\frac{1}{|M^k\\_t|} - \\frac{O( t e^{-c})}{|M^k\\_t|} \\text{ for } i \\in M\\_t^k, \\\\\\\\\n    O(e^{-c}) \\text{ o.w. }\n\\end{cases}\n\\end{align*}\n$$\n\nThis fixes the inaccuracies identified, i.e., \n  - The attention matrix and the attention scores after softmax are lower triangular which fixed our previous oversight where we defined it as the transpose of the current matrices. We will ensure consistency across the paper.\n   - We choose ${e\\_1,\\ldots, e\\_T}$ as the elementary basis and fixing the notation in the summation. \n   - We also fixed the bordercase, now the token out of the range will just attend the first token $(x\\_1)$ that acts as a sink.\n   - Eq (9) now holds correctly in the limit $c \\to \\infty$ and for tokens $t>n$.", "claims_and_evidence": "The theoretical results are based on a simplified attention-only transformer where the residual connection is replaced with concatenation — a simplification used in previous works to make the analysis tractable. The experiments verify the results on both the simplified and standard attention-only transformer, suggesting that the simplification does not distort the optimization landscape much from this perspective.\n\nTo verify the claims empirically, they show that the loss plateaus during training correspond to sub-$k$-grams by confirming that the attention maps follow the construction discussed in the theory. However, verifying that the gradient of the loss around these points is near zero would further strengthen the empirical validation.", "methods_and_evaluation_criteria": "See above", "theoretical_claims": "I followed the proof sketch in the main body and selectively reviewed parts of the supplementary material without checking all the details.", "experimental_designs_or_analyses": "see above", "relation_to_broader_scientific_literature": "Several recent works have investigated the optimization landscape of transformers by considering Markov chains as data models. The empirical observation of multiple loss plateaus during training on in-context Markov chains was first reported in Edelman et al. (2024), as noted in the paper. Edelman et al. (2024) also showed empirically that each plateau corresponds to a model that is close (in terms of KL divergence) to a sub-$k$-gram, but their theoretical analysis was limited.\n\nIn this sense, showing concretely that sub-$k$-grams are stationary points is a valuable follow-up to this result.\n\nThere are two related works that could benefit from further discussion on results/differnces in the paper:\n- The idea of the construction of the $k$-gram solutions is very close to the construction in Rajaraman et al. (2024). \n- [1] studies the gradient flow dynamics in the same Markov chain setup. \n\n[1] Chen S, Sheen H, Wang T, Yang Z. Unveiling induction heads: Provable training dynamics and feature learning in transformers.", "essential_references_not_discussed": "See above", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "I do find the result about the stationarity of $k$-grams a useful addition to the body of work on Markov chains for in-context learning. However, as mentioned earlier, a more detailed discussion on how your results compare and fit into the broader context of related works would help clarify the scope of the contribution. Can you comment on the results in  Rajaraman et al. (2024) and [1]?\n\n\n[1] Chen S, Sheen H, Wang T, Yang Z. Unveiling induction heads: Provable training dynamics and feature learning in transformers.", "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": "uUrGC03oxa", "reviewer": "Reviewer_wvkz", "summary": "In this work, the authors study the loss landscape of transformers in the task of next-token prediction. To that end, they consider a simplified two-layer transformer and focus on learning n-gram language models in context with the cross-entropy loss. They rely on a constructive approach of transformer that mimics k-grams with $k \\leq n$ such that the gradients of the population loss tend to zero when the sequence length goes to infinity. The authors deduce from this result that k-grams are stationary points of the population cross-entropy loss, which provides insights into the stage-wise learning dynamics and phase transition in transformers. Finally, they conduct empirical validation of their theoretical insights that illustrate the dynamics of learning n-grams that lead to jumps between stationary points.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the detailed feedback and thoughtful suggestions provided by the reviewer. \n\n**References:**\nWe thank the reviewer for bringing concurrent works Zekri et al. (2024) and Nguyen (2024) to our attention. We acknowledge their relevance to our work, especially in motivating studying learning $n$-grams with transformers and we will cite them in the final version. \n\n**Practicality of Proposition 3.2:**\nAlthough we formally prove that the assumptions of Proposition 3.2 hold for sub-$n$-gram constructions and in the context of learning $n$-gram language models, we conjecture that the conditions are satisfied more generally and this is the reason why there are long plateaus corresponding to **marginal estimates** of the conditional target distribution in many in-context learning setups with cross-entropy loss. One such pattern can be observed in two-hop reasoning tasks (see [1]), where there is a long plateau corresponding to an estimate of $p(end)$ (marginal estimate), where the goal is the predict $p(end|start, bridge)$ (true conditioned estimate). Another example is factuall-recall setting of [2], where the plateau is $p(.|\\text{subject})$ (conditioned only on one token), where the goal is to predict  $p(.|\\text{subject,relation})$ (true estimated conditioned on two tokens). \n\n**Experimental Setting:** please see reply to R. 6GFt \n\n**Scaling Factor $c$:** The scaling factor $c\\in\\mathbb{R}^+$, determines the norm of the parameters in the construction. Please see the response to Reviewer xuc4 that provides a formal lemma for the construction. \n\n**$\\epsilon$-stationary points and error control**.  To precisely give the error bounds and delineate the dependence of the gradient on $c,T$, we give the following computation. To illustrate, we give the computation of gradient w.r.t the  key matrix of the second layer $K_2$ in $\\mathbb{R}^{nd \\times nd}$ . \n\n$$ \\begin{align*}\n\\frac{\\partial p_{\\theta}}{\\partial K_2} = \\mathbb{E}\\_{\\tau} \\mathbb{E}\\_{x^T}   \\sqrt{c} \\sum\\_{a \\in S} ( \\hat{p}\\_k[a] - p_{\\tau} (x_{T+1} = a \\big| x^T_{ \\color{red} T - n + 1}) )  \\phi(x\\_{\\color{green} T-k+2}^T) (s^0_{a} - S^{0} \\hat{p}\\_k )^{\\top}  + { \\color{blue} O( \\sqrt{c} T e^{-c} )  },\n\\end{align*} $$\n\nwhere $\\hat{p}_k$ is the $k-$gram estimator defined in Eq.(2) of the main paper, $\\phi = \\sum\\_{h=1}^{k-1} s^h\\_{x\\_{T+1-h}}$ and $S^0$ is a block matrix in $\\mathbb{R}^{nd \\times nd}$ where the first block is given by the embedding matrix S and $0$'s everywhere else. Now we can split the gradient further, \n\n$$ \\begin{align*}\n\\frac{\\partial p_{\\theta}}{\\partial K_2} = \\mathbb{E}\\_{\\tau} \\mathbb{E}\\_{x^T}   \\sqrt{c} \\sum\\_{a \\in S} ( p_{\\tau}(x_{T+1}{=}a \\big| x^T_{ \\color{red} T - k + 2}) - p_{\\tau} (x_{T+1}{=}a \\big| x^T_{ \\color{red} T - n + 1}) ) ~  \\phi(x\\_{\\color{red} T-k+2}^T) (s^0_{a} - S^{0} p_{\\tau}(x\\_{T+1}{=} \\cdot \\big| x^T_{ \\color{red} T - k + 2})  )^{\\top}  + \\sqrt{c}  {\\color{green} \\mathbb{E}\\_{\\tau} \\mathbb{E}\\_{x^T}   \\\\| \\hat{p}\\_k - p\\_{\\tau}(. \\big| x^{T}\\_{T-k+2})  \\\\|^2 } + { \\color{blue} O( \\sqrt{c} T e^{-c} )  },\n\\end{align*} $$\n\nThe term in blue represents the impact of finite $c$ and the term in green the impact of finite $T$.  The first term goes to zero as a result of Proposition 3.2.\n\n**Feed forward layers:** With feedforward networks, the representation question itself remains unclear—specifically, what solutions transformers learn.  One key difference observed in our experiments is that feedforward networks can solve the task with fewer heads, and each head can attend to multiple previous positions. In contrast, in attention-only transformers, each head attends to a specific previous token. Consequently, the number of possible intermediate stages is uncertain if feed-forward layers are involved. For example, in our experiments with an order-2 chain, we observe only one intermediate stage.  \n\n**Generalization of our result to broader architectures:** Our proof sketch provides intuition for determining stationarity. Roughly, if the first layer computes a function $\\psi(x^i_{i-k+1})$ of the $k$-history and the second layer aligns with it, the core argument—that gradients are supported by $M_t^{k}$ and that the gradient depends only on the $k$-history—can be carefully extended. In this way, we believe we can extend the analysis to general transformers. However, we tone down the claim in l.343 to say general *attention only transformers*.\n\n**Markov property,  other data modalities**\nAlthough we formally prove that the assumptions of Proposition 3.2 hold for sub-$n$-gram constructions and in the context of learning $n$-gram language models, the proposition itself does not require any Markov assumption and holds for all the next token prediction tasks using cross-entropy loss, e.g. tokenized code, time series or other in-context learning settings discussed above. \n\nReferences:\n\n[1] Guo et al. https://arxiv.org/pdf/2502.13913\n\n[2] Nischani et. al. https://arxiv.org/pdf/2412.06538", "claims_and_evidence": "The theoretical claims are well supported by clear and detailed proofs, and the authors also provide experimental validation of their theory in a controlled setting.", "methods_and_evaluation_criteria": "The authors provide theoretical results to better understand the sequential learning dynamics of transformers. They conduct toy experiments to validate their findings. Hence, the methods make sense for the problem at hand.", "theoretical_claims": "The theoretical findings are supported by detailed and clear proofs.", "experimental_designs_or_analyses": "I checked the soundess of the experiments. I believe that, while most of the contributions are theoretical, the experiments are too lightweight with a focus on small-scale, synthetic data and simplified models. I think the submission would benefit from experiments in more practical settings and/or discussions on how the current contributions extend to such scenarios.", "relation_to_broader_scientific_literature": "I find that related work and prior works are well introduced and compared. The submission's contributions are interesting and are part of a growing interest in the literature on the theoretical understanding of transformers. As such, the authors borrow a lot from prior works (data framework [1, 2], disantangled transformer model [2, 3] sub-gram representation [2, 4]), and some theoretical results are, as the authors acknowledge, extensions of existing results (e.g., Lemma 3.1 generalizing [5, 6]).\n\n*References*\n\n[1] Edelman et al. The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains, NeurIPS 2024 \n\n[2] Nichani et al. How transformers learn causal structure with gradient descent, ICML 2024.\n\n[3] Friedman et al. Learning transformer programs, NeurIPS 2023.\n\n[4] Rajaraman et al. Transformers on Markov data: Constant depth suffices, NeurIPS 2024\n\n[5] Bietti et al. Birth of a Transformer: A Memory Viewpoint, NeurIPS 2023\n\n[6] Makkuva et al. Attention with markov: A framework for principled analysis of transformers via Markov chains. ICLR 2025", "essential_references_not_discussed": "To the best of my knowledge, there were no essential references not discussed in the current submission. I mention below some works that I think are relevant to the current submissions and might be discussed.\n\n- [1] demonstrated that any transformer-based LLM has an equivalent Markov chain formulation, which enabled them to study LLMs' generalization capabilities in pre-training and in-context learning on Markov chains. \n\n- [2] studies how well n-gram statistics approximate the next-token transformer prediction, which enables the author to understand better learning dynamics and overfitting. \n\nI acknowledge that both [1] and [2] could be considered as concurrent work, hence, I do not require the authors to discuss them. But I think that they can provide valuable context on the current submission's contributions and help argue the benefits of studying n-grams and Markov chains data to better understand transformers and LLMs.\n\n*References*\n\n[1] Zekri et al. Large Language Models as Markov Chains, arXiv 2024 (appeared on arXiv on 3 October 2024, according to https://arxiv.org/abs/2410.02724)\n\n[2] Timothy Nguyen. Understanding Transformers via N-gram Statistics, NeurIPS 2024 (appeared on arXiv on 30 June 2024, according to https://arxiv.org/abs/2407.12034)", "other_strengths_and_weaknesses": "**Strengths**\n\n- The paper is well written with detailed prior work, clear notations, and a technical background.\n- The problem tackled is interesting\n- The theoretical findings provide a better understanding of the sequential learning of transformers in in-context learning of n-grams\n\n**Weakness**\n\nI list below what I think are weaknesses, but I would be happy to be corrected if I misunderstood some important aspects of the authors' contributions.\n\n- I have trouble seeing how the insights provided by the authors could help in more practical scenarios since strong assumptions are made on both the task and the model. This work rather seems to provide a better understanding of the sequential learning observed in [1], which already considers a specific case of in-context learning of n-grams. \n\n- I believe the asymptotic nature of the results on the stationary points is not quite satisfactory since, in practice, $T$ is finite and $c$ should be too (see questions section). As such, it would have been more interesting, in my opinion, to control the gradient norms to know how far from $0$ they are (see questions section)\n\n- The experiments seem limited to a simplified model with synthetic data and vocabulary size $S=3$ and sequence length $T=32$ (far from the asymptotic regime). It would have been interesting to see how well the provided theory can explain more practical scenarios. An idea to improve the submission (inspired by [2]) could be to consider a small but realistic model (e.g., 2-layer transformer with feed-forward layers) and observe learning curves in parallel with gradient norms (easily recovered with automatic differentiation), hoping to see stage-wise learning and near-stationary points (low gradient norms) appear in tandem.\n\nOverall, I find the paper interesting and the theory well derived. I appreciate controlled settings to conduct rigorous analysis, but I believe that the setting is too oversimplified, which impacts the submission's contributions. This is the reason for my current score but I remain open to modifying my score, provided the authors clarify the points mentioned in the weaknesses section.\n\n**Update after rebuttal**: I increased my score from 2 to 3 based on the authors' answers to my rebuttal comments (I thank them for the additional clarifications and experiments; despite the simple setting, the authors motivate well how it can be generalized to more practical ones).\n\n*References*\n\n[1] Edelman et al. The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains, NeurIPS 2024 \n\n[2] Odonnat et al. Clusering Head: A Visual Case Study of Training Dynamics in Transformers. arxiv 2024", "other_comments_or_suggestions": "None", "questions_for_authors": "1) Could the authors describe practical scenarios in which the assumption of Proposition 3.2 holds (when the derivative of the model’s logits depends solely on a subsequence of the input history)? \n\n2) The authors introduce $c$ as a constant scaling factor in Eq. (7) and (8), but I don't understand where it comes from or how it is defined. Moreover, the authors make it vary up to $\\infty$ in the following theoretical results; however, in my understanding, a constant scaling factor should be fixed or depend on the problem's parameters. Could the authors elaborate on that? \n\n3) The authors state l. 343, second column, that \"the same techniques and methodology extend naturally to a general transformer architecture\". I believe this is somewhat of a bold statement as it would lead to significantly more involved computations, and because of that, it is not clear, at least to me, how much of the current insights can be extended to this more general setting. Could the author elaborate on that?\n\n4) What the authors call \"near-stationary points\" can be seen as to approximate stationary points, also called $\\epsilon$-stationary points [1], i.e., their gradient is bounded by some $\\epsilon$. In this submission, $\\epsilon$ depends on how large the sequence length is since the convergence to zero is asymptotic in $T$ (and $c$). Could the author elaborate on how to control this $\\epsilon$ value?\n\n5) A naive bound on Theorem 4.1 (b) could, for instance involve Cauchy-Swartz and Jensen (since the absolute value is convex) and make three terms appear: the average difference in probabilities, the average norm of the derivative at $\\theta^*$, and the term in $e^{-c}$. How do those terms depend on $T, c$ and the data? Answering this question could help to control $\\epsilon$. If the authors have other ideas to bound the gradient to obtain some control over $\\epsilon$, I would be happy to hear them.\n\n6) Could the authors elaborate on how much considering the feed-forward layers would impact the current analysis? \n\n7) In practice, language does not verify the Markov property, nor do other data modalities (e.g., tokenized code, images, and time series). How well do the insights provided by the current work translate to real-world data commonly processed with transformers?\n\n*References*\n\n[1] Hollender et al. The computational complexity of finding stationary points in non-convex optimization, In Mathematical Programming  2024", "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:

Meta-review: This paper aims to study the loss landscape of next-token prediction for transformers where cross-entropy defines the underlying loss function. More precisely, assuming that the underlying data follows an n-gram language model, the paper analyzes the loss landscape of a simplified attention-only transformer architecture, namely disentangled transformers. The paper shows that k-gram estimators for k <= n form asymptotic stationary points or non-asymptotic near-stationary points. The paper connects these stationary points (for k < n) to the stage-wise learning phenomenon where the training goes through plateaus. The paper also presents synthetic experiments to corroborate the theoretical findings.

The reviewer raised multiple valid questions/concerns about the initial submission. The reviewers pointed out various issues with the presentation of the results and the claimed extension in the initial submission. The reviewers asked for further empirical results to strengthen the main takeaways of the initial submission. The reviewer also asked for further justification about various assumptions in the paper. The authors were able to adequately address most of these concerns during the author-reviewer discussion phase. At the end of the discussion phase, all the reviewers agreed on the importance of the findings in the paper and recommended acceptance provided the final paper takes all the key points discussed during the discussion phase into account.

The authors are strongly recommended to make all the promised changes. In addition, it would be useful to bring up disentangled transformers much earlier (e.g., in the introduction) as the results in the paper are mostly restricted to those.
Decision: Accept (poster)