We thank the reviewer for the helpful feedback. We address the comments and questions below.

Impact of multiple trigger tokens

Extending the analysis to the case of multiple triggers is an interesting and natural direction, and we are actively working on this extension.
The key idea behind our current analysis is that the relative strength of two peaks in the attention key-query matrix—namely, the positional memory and the induction head components corresponding to terms (a) and (b) in Equation (3.1)—determines which of the two algorithms is implemented by the transformer. Our results suggest that when the pretraining dataset (i.e., the positions of the triggers) is diverse, the peak associated with the positional memory becomes dispersed across multiple trigger positions, thereby making the induction head dominant.

In contrast, when there are multiple triggers—drawn from a set of $K$ tokens, as you suggested—the induction head signal, which depends on the semantics of the triggers, also becomes dispersed. This yields two implications:

• Transition Point. Since multiple triggers make the induction head more difficult to implement, we expect that a higher level of diversity in the pretraining data is required to induce a phase transition from positional memorization to the induction head regime.
• Sample Complexity. As the induction head signal is attenuated due to the semantic dispersion across multiple triggers, the number of pretraining samples required to prevent the signal from being buried in noise—in other words, the sample complexity required for implementing the induction head—is expected to increase and dependent on $K$.

We believe these implications have the potential to make the transition point more complex but to provide a rich characterization of the competition between positional and semantic signals.

Formalizing the Max-Sum Ratio's Role

As you pointed out, the effective magnitude of the positional shortcut and that of the induction head component in $s_t (= x_{t}^\top W_{KQ} x_{T})$ are proportional to the numerator and the denominator, respectively. In short, the max-sum ratio governs the relative strength of the positional shortcut and the induction head in the attention logits. In light of your comment, we will discuss this intuition in the main text.

On the experiment

In practice, the precise location of the transition point can depend on architectural and embedding-specific details, and thus cannot be fully predicted in a quantitative manner based solely on the max-sum ratio.

Being said, some qualitative insights can still be obtained. For instance, if we consider a uniform distribution over an interval of width $K$, namely $[\ell_0, \ell_0 + 1, \ldots, \ell_0 + K - 1]$, and we shift this distribution to the right (i.e., toward longer contexts) while keeping the total width $K$ fixed, the max-sum ratio decreases. As suggested by Theorem 7, a smaller max-sum ratio is favorable for OOD generalization. Therefore, shifting the distribution to the right reduces the required width $K$ for achieving OOD generalization.　This is consistent with the experimental result.

Simplified setup

We provide the following remark concerning our simplified setup:

• Simplified Embedding. Our simplified embedding scheme is intended to model the behavior in deeper architectures, where information from previous tokens flows into the residual stream [1]. While the analysis may become more involved, a simplified multi-layer architecture equipped with residual connections and relative positional embeddings [2] would reproduce our theoretical insights.
• Single-step analysis. We remark that our learning setup — based on one gradient descent step with a large learning rate — is intended to model the behavior of parameters escaping from saddle points during the initial phase of gradient-based algorithms. This setup has been extensively studied in the neural network theory literature, and in many feature learning settings it is known that alternative gradient-based training algorithms such as one-pass SGD yield similar statistical efficiency [3,4].

[1] Bietti et al. Birth of a Transformer: A Memory Viewpoint. [2] Wang and Sato. Rethinking Associative Memory Mechanism in Induction Head. [3] Abbe, Boix-Adsera, and Misiakiewicz: SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. [4] Barak et al: Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit.

We would be happy to answer any follow-up questions in the discussion period.

We thank the reviewer for the helpful feedback. We address the comments and questions below.

OOD length generalization

I do not quite grasp the authors initial statement and assumption on the transformers' inability to…

Thank you for raising this point. While LLMs tend to work well even with very long contexts, they sometimes rely on heuristic or spurious features for solving certain tasks, which can manifest itself in a failure to generalize to certain OOD scenarios. These failures to generalize have been observed, e.g., in linguistic tasks (e.g. [1]), or in simple algorithmic tasks requiring length generalization (e.g. [2,3]). Generally, these tasks can be much more complicated than simple few-shot learning, and often involve the composition of small algorithmic primitives that need to be robust to changes in their inputs.

[1] McCoy, Pavlick and Linzen. Right for the Wrong Reasons: Diagnosing Syntactic Heuristics in Natural Language Inference.

[2] Anil et al. Exploring Length Generalization in Large Language Models.

[3] Zhou et al. What Algorithms can Transformers Learn? A Study in Length Generalization.

Our setup describes possibly the simplest algorithmic scenario in a transformer that is prone to such failures, namely copying a token based on heuristic positional information instead of learning a general copy mechanism that works regardless of position. The simplicity of our task allows us to obtain a precise characterization of the phenomena at play, such as how the choice of data distribution affects gradient dynamics, and how much diversity is needed in order to properly avoid the heuristic “shortcut” solutions and obtain reliable models. Our study paves the way towards understanding more complex scenarios which may involve compositions of such small algorithmic components.

Value of theoretical works on transformers

Although I concur with the authors on that developing theoretical frameworks…

This is a good question. While it is clear at this point that the community has realized the importance of pre-training data and data efficiency, many important questions remain open in this area, and the potential to further improve data efficiency could unlock significant improvements in model performance, as well as cost savings in pre-training. This is even more the case when you start considering synthetic / simulated data, instead of filtering existing available data, as in some of the referenced works. An important long-term outcome that we’re aiming towards is a way to find the smallest possible dataset that would allow us to train a model that perfectly solves a range of complex algorithmic reasoning tasks. Even just formalizing this question is very difficult. Our work takes a first step by considering possibly the simplest algorithmic task (copying with induction heads), and trying to understand how the data distribution affects the learning dynamics, and to characterize how diversity leads to a transition from using spurious positional heuristics to generalizable copy mechanisms. While the task in itself is quite basic, you could imagine extensions of our work that compose many such small algorithmic components in a larger transformer in a reliable way to achieve more complex algorithmic tasks. We hope this clarifies our motivation, and the significance of this initial step towards a much grander goal.

On the numerical experiment

Thank you for the feedback. We will elaborate on the experimental setting in the revision.

• Data Model. The data used in the experiment of Figure 1 is synthetic data following the data-generating process in Definition 1. In particular, We set the vocabulary size to $N = 36$ and maximum sequence length to $L = 36$. The vocabulary itself consists of natural numbers $\{1, 2, \dots, N=36\}$; each sample sequence is of length up to $L = 36$, such as $[10, 8, 1, 4, 5, 8, 1, \dots]$, and the probability distribution over data sequences is defined in Definition 1. We do not treat token 1 as a special token; rather, we embed it in the same manner as any other token, as described below (This is because it is required to construct the induction head without explicitly signaling that token 1 carries a different role from the other $N-1$ tokens.)
• Embeddings. Each token $z_i \in [N]$ (where $i \in [L]$) is represented by the concatenation of a zero vector of length $L$ and a one-hot vector of length $N$ whose $z_i$-th component is one. For positional embeddings, we use the concatenation of a one-hot vector of length $L$ (with one at the $i$-th position) and a zero vector of length $N$. See also Appendix E.1 for further details.
• Architecture and Training. Based on the above, we train a two-layer architecture as described in Section 4. We perform optimization using the Adam optimizer with a learning rate of 0.0001, with early stopping every 250 steps. We plan to integrate the explanations of architectural and training details, which are currently split between Section 4 and Appendix E.1.

We would be happy to answer any follow-up questions in the discussion period.

We thank the reviewer for the helpful feedback. We address the comments and questions below.

Implications of the max-sum ratio for realistic settings

While the max-sum ratio is derived under our simplified theoretical setup, we believe it provides a high-level guideline for designing datasets that promote better generalization.

Primarily, the max-sum ratio encourages diversity in the task distribution, as discussed in the main text. Further intuition behind the max-sum ratio is that the task distribution, when reweighted by the inverse of the context length (i.e., $q_\ell$ weighted by $\ell^{-1}$), should be approximately uniform in order for the max-sum ratio to attain its minimum. In other words, $q_\ell$ associated with smaller $\ell$—which corresponds to a task where positional information can be memorized easily—should take relatively smaller values. Accordingly, the max-sum ratio suggests the following simple high-level principles for dataset design: (i) selecting a more diverse range of tasks, and (ii) placing greater emphasis on tasks in which positional signals are harder to obtain.

Architectural and Embedding Considerations

Thank you for the suggestion. Indeed, it is challenging to straightforwardly extend our theoretical framework to multiplicative embeddings like RoPE or to SSMs that depart from the transformer architecture with absolute embedding.
That said, the high-level insights regarding dataset design discussed in the previous section—such as promoting some notion of diversity in the pretraining data distirbution and emphasizing tasks where positional signal is harder to obtain—are not specific to any particular architecture or embedding scheme. As such, we believe it is worthwhile to explore whether these dataset design principles remain effective in these more advanced settings.

A step toward the RoPE setting is to incorporate an additional layer. While the analysis may become more involved, the use of a residual stream potentially allows us to extend our one-layer architecture to a multi-layer architecture equipped with relative positional embeddings [1,2]. Although further extensions would be required to fully capture the multiplicative structure of RoPE, we view this as a meaningful future step in that direction.

[1] Bietti et al. Birth of a Transformer: A Memory Viewpoint. [2] Wang and Sato. Rethinking Associative Memory Mechanism in Induction Head.

We would be happy to answer any follow-up questions in the discussion period.

We thank the reviewer for the helpful feedback. We address the comments and questions below.

Extension to the full training dynamics and simultaneous training on the all layers

We remark that our learning setup — based on one gradient descent step with a large learning rate — is intended to model the behavior of parameters escaping from saddle points during the initial phase of gradient-based algorithms. This setup has been extensively studied in the neural network theory literature, and in many feature learning settings it is known that alternative gradient-based training algorithms such as one-pass SGD yield similar statistical efficiency [1,2].

Moreover, the experiments described in Section 4 were conducted using a multi-step Adam optimizer, where all layers were trained jointly. This suggests that the transition from memorization to out-of-distribution generalization is not merely an artifact of the one-step gradient descent setting assumed in our theory. We will include a more thorough empirical analysis of how the phenomena we identified unfold in multi-step optimization, including an in-depth examination of inner mechanisms such as attention patterns.

[1] Abbe, Boix-Adsera, and Misiakiewicz: SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics.

[2] Barak et al: Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit.

Insights on the realistic NLP settings

Thanks for the question. Our goal is to provide a first step towards understanding how to design the most efficient data distributions for learning a range of algorithmic tasks in LLMs. While our focus is on a simple synthetic task involving copying, you could imagine scaling this up to much more involved algorithmic tasks with synthetic data, and possibly turning them into natural language reasoning tasks with the help of LLMs, something which has been recently successful (e.g., [3]). Indeed, while the copying operation is just a tiny step towards this grander goal, copying is plausibly an important building block in many more complex algorithms, and failing to have a reliable copy mechanism would easily cause the larger tasks to also fail. We hope that our precise characterization of how data affects training dynamics, and how diversity helps the model transition from relying on spurious (positional) heuristics to generalizable solutions, will help to study data efficiency of more complex tasks consisting of multiple such components.

[3] Gunasekar et al. Textbooks Are All You Need.

Merged key and query matrices

In our theoretical analysis, we apply a simplification that ties $W^K$ and $W^Q$ in our model, yet we still observe the emergence of induction heads. This suggests that a factorized form is not strictly necessary for induction heads to appear (although they may appear at different speeds in the two models due to different dynamics and implicit biases, as studied in the referenced paper “How Transformers Get Rich”). On the other hand, the experiments described in Section 4 use a setting where keys and queries are separated. Taken together, these findings indicate that whether keys and queries are tied or separated is not essential for reproducing our theoretical insight.

We would be happy to answer any follow-up questions in the discussion period.

I appreciate the authors' detailed explanations. However, my main concerns remain unresolved:

• Theoretical Analysis: The analysis remains incomplete. As acknowledged in the rebuttal, the paper examines only the dynamics of a single gradient step, without addressing the full training trajectory. While some theoretical works focus on the first step under simplified assumptions, they typically also establish global convergence—something this paper does not.

• Relevance to realistic settings: Although this is a theoretical paper, its data-centric viewpoint is interesting. Nonetheless, it is important to clarify its relevance to realistic NLP scenarios. At least, the paper should discuss the connection between the studied task and standard NLP tasks, and examine the extent to which the findings may generalize.

Given these concerns, I will maintain my current score.

Thank you for the followup comment. We make the following remarks.

the paper examines only the dynamics of a single gradient step, without addressing the full training trajectory

• In our shallow transformer setup, the optimization landscape is non-convex and highly complex, and hence the single-gradient-step simplification is common in the literature – see for example [1,2]. As mentioned in our previous response, insights in these one-step analyses can extend to different gradient-based training algorithms. To our knowledge, existing theoretical works on similar n-gram settings that go beyond the first gradient step all worked with a particular layer-wise training paradigm [3,4,5], and often assumed a known value matrix. Moreover, the novelty of our work among these prior studies lies in characterizing the transformer's algorithm selection from a data-centric perspective — this turns out to be technically nontrivial even in the one GD step setting. We therefore believe that our theoretical analyses should not simply be disregarded as incomplete.
• To connect our optimization analysis to a more “global” characterization, note that Proposition 5 and Theorem 7 imply that, under a small max-sum ratio, the learned transformer “generalizes OOD”, in the sense of Definition 2: namely, for all OOD input sequences of length up to $L$, the label that receives the highest predicted probability coincides with the correct target token. Hence if the temperature (or equivalently, the scale of $W^V$) is taken to be sufficiently large, the cross-entropy loss becomes arbitrarily close to zero. Consequently, the only remaining gap from (approximate) global optimality lies in the scaling of $W^V$; hence one may interpret our result as suggesting after one GD step, the model has already learned the “global” optimal directions.
• Last but not least, in Section 4 we have empirically demonstrated a similar transition from failure to success in OOD generalization under standard training with the Adam optimizer. We will include additional experiments that examine the full training dynamics in greater detail (e.g., quantitative assessment of the transition and the learned mechanism).

the paper should discuss the connection between the studied task and standard NLP tasks

The induction head mechanism has been demonstrated to play a critical role in enabling the in-context learning capabilities of LLMs; many prior theoretical works have studied the gradient-based learning of this mechanism in synthetic copying tasks similar to ours, and as mentioned in our previous response, the understanding of such tasks allows us to potentially infer the model performance on larger and more realistic tasks that involves induction head. To our knowledge, our work is the first to characterize how the diversity of data distribution dictates a transition from a spurious positional solution to the “correct” induction head mechanism; while we respect the reviewer's practice-oriented perspective, we believe that our theoretical framing of how the data quality affects the learned algorithm is a nontrivial contribution in itself.

We hope these clarifications are helpful and appreciate your time in engaging with our work.

[1] Bietti et al. Birth of a Transformer: A Memory Viewpoint, 2023.

[2] Wang and Sato. Rethinking Associative Memory Mechanism in Induction Head, 2025.

[3] Chen et al. Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers, 2024.

[4] Wang et al. How Transformers Get Rich: Approximation and Dynamics Analysis, 2025.

[5] Nichani et al. How Transformers Learn Causal Structure with Gradient Descent, 2024.

We thank the reviewer for the helpful feedback. We will correct the typos in the manuscript. We address the comments and questions below.

Intuition behind the max-sum ratio

As you pointed out, the max-sum ratio is not shift-invariant, since the inverse of the context length — measured by $\ell^{-1}$ — varies under a global shift of the distribution. For instance, if we consider a uniform distribution over an interval of width $K$, namely $[\ell_0, \ell_0 + 1, \ldots, \ell_0 + K - 1]$, and we shift this distribution to the right (i.e., toward longer contexts) while keeping the total width $K$ fixed, the max-sum ratio decreases. As suggested by Theorem 7, a smaller max-sum ratio is favorable for OOD generalization. Therefore, shifting the distribution to the right reduces the required width $K$ for achieving OOD generalization. The intuition is that the task distribution, when reweighted by the inverse of the context length (i.e., $q_\ell$ weighted by $\ell^{-1}$), should be approximately uniform in order for the max-sum ratio to attain its minimum. In other words, $q_\ell$ associated with smaller $\ell$—which corresponds to a task where positional information can be memorized easily—should take relatively smaller values.

A more precise interpretation of the max-sum ratio, then, is that it primarily encourages diversity in the task distribution, while assigning greater weight to longer contexts where positional signals are harder to obtain. In light of your feedback, we plan to incorporate this precise implication into the paper.

Regarding experiments in Fig 1

We address each of your points in the following:

a) As mentioned above, the wider the rightward shift of the distribution over the period length $\ell$ during pretraining, the narrower the required support of the pretraining distribution becomes for implementing an induction head. Thus, this point is consistent with the prediction made by the max-sum ratio.

b) Our learning setup — based on one gradient descent step with a large learning rate — is intended to model the behavior of parameters escaping from saddle points during the initial phase of gradient-based algorithms. When using other algorithms such as one-pass SGD, we expect that the length of pretraining influences the implemented algorithm. This is analogous to the compatibility observed in some feature learning settings between the sample complexity of one-step GD and the computational complexity of one-pass SGD [1,2].

[1] Abbe, Boix-Adsera, and Misiakiewicz: SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics.

[2] Barak et al: Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit.

c) One method to internally verify the mechanism by which each algorithm is implemented in the transformer is to examine the heatmap of the key-query matrix that governs the attention scores. Theoretically, this is formalized in Equation (3.1) (more precisely, Lemmas 10 and 11 in Appendix B), which show that the key-query matrix exhibits peaks corresponding to each algorithm (positional memorization and induction head). Furthermore, using the same data model and algorithm as in our theory, we conducted experiments demonstrating that both peaks indeed appear in the key-query matrix, and that increasing the width of the pretraining task leads to the selective attenuation of the positional memory peak. We will include these experimental results in the revision.

We would be happy to answer any follow-up questions in the discussion period.