Learning In-context $n$-grams with Transformers: Sub-$n$-grams Are Near-Stationary Points

We sincerely appreciate the detailed feedback and thoughtful suggestions provided by the reviewer.

References: We 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.

Practicality of Proposition 3.2: Although 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).

Experimental Setting: please see reply to R. 6GFt

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.

$\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}$ .

where $\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,

The 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.

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.

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.

Markov property, other data modalities Although 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.

References:

[1] Guo et al. https://arxiv.org/pdf/2502.13913

[2] Nischani et. al. https://arxiv.org/pdf/2412.06538

We appreciate the reviewer for bringing Chen et al. (2024) to our attention, which we unfortunately overlooked in our discussion of related work.

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:

• 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.

• 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.

We will add this detailed comparison to the paper.

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.

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.

Reviewer xuc4

$k$-gram Construction: To 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..

Lemma 1 For the parameters defined, (a) The attention score $a\_{(i,j)}^{(1,h)}$ between key and query elements $(i,j)$ of head $h$ of layer 1 is $$

$

     a\_{(i,j)}^{(1,h)} = \begin{cases} 1 - O(i e^{-c}) \text{ when } i > h \text{ and } j = i-h, \\\\ 
     1 - O(i e^{-c}) \text{ when } i \leq h \text{ and } j = 1, \\\\
     O(e^{-c}) \text{ o.w.}. 
     \end{cases}.

$

$

a^2_{(t,i)} = \begin{cases} \frac{1}{|M^k_t|} - \frac{O( t e^{-c})}{|M^k_t|} \text{ for } i \in M_t^k, \\ O(e^{-c}) \text{ o.w. } \end{cases}

$

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.

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.

Theorem: For the distangled transformer $p_{\theta}$, the gradients at $\theta^k_*$ are given by

$\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} ) }.$

This 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.

We 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.

k-gram construction. Please see the response to reviewer 6GFt for the formal construction.

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.

In 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.

Other concerns. 10-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.

Other Comments or Suggestions.

1-5,8-9: We have corrected the typos you identified and carefully proofread to make sure we fixed the inconsistencies in the paper.

6: 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.

7: We have fixed the condition to almost sure equality.

10: 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.

14: 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}$.

We 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.