Thank you for your constructive comments. Please see our response below.

Gradient flow analysis: We focus on gradient flow to simplify the theoretical analysis, as it is a common approach for understanding the training dynamics of complex models. We believe that the results can be extended to gradient descent with a sufficiently small step size. Given that our model includes two layers of multi-head softmax attention and feedforward neural networks, analyzing the training dynamics is challenging even with gradient flow. Furthermore, we use gradient descent in our experiments, and the results align with our theoretical findings. This demonstrates that our gradient flow analysis is consistent with the gradient descent implementation.

Comparison to Nichani et al. (2024): The major distinctions between our work and Nichani et al. (2024) lie in the model structure and data structure. We will include the following table for a clear comparison to Nichani et al. (2024).

Relative Positional Encoding: Our findings highlight the benefits of using Relative Positional Encoding (RPE). Previous works have shown that transformers with RPE outperform those with standard positional encoding. For instance, Allen-Zhu et al. (2023) demonstrate that transformers with RPE are more effective at learning hierarchical language structures.

In our work, as we consider an $n$-gram data model, RPE is a suitable choice for positional encoding. Specifically, RPE has been utilized to effectively capture potential parent sets, thereby enhancing the length generalization ability of the model. As long as the model accurately identifies parent sets with RPE, it will also perform well when tested on input sequences of varying lengths $L$.

While RPE may not be the best choice for different datasets, a generalization study of RPE is beyond the scope of our current work.

On Generalization ability:

• Generalization to test data: The generalization capabilities are tied to the model's ability to infer the true parent sets from the data. Our results demonstrate that the parameters are updated to maximize the modified chi-square mutual information. We find that the modified chi-square mutual information is an effective metric for identifying the parent sets as described in Section 3.3. For instance, when the number of parents is known a priori, the modified chi-square mutual information is maximized by the true parent set.

We have conducted additional experiments to evaluate the model's generalization performance on sequences having lengths different from the training data. Please see the attached pdf for the experimental results under the general rebuttal. These additional experiment results indicate that the trained model successfully captures the true parent sets, implying the strong performance on test data. Even if the test data have a different distribution or length $L$, the model will generalize well as long as they maintain the same parent structure as the training data.

• Generalization to other types of data/model: The investigation of data structures beyond the $n$-gram model and more general transformer model is beyond the scope of our current study. We believe it is an important direction for future work to study the training dynamics of transformers in more general settings.

On Additional experiments: The goal of our paper is to examine the induction head mechanism of transformers within the $n$-gram data model framework. The $n$-gram model is widely used to evaluate transformers' capabilities in language modeling from both experimental and theoretical perspectives. Our experiments show that the empirical results align with our theoretical findings that the model implements an induction head. We leave it as future work to explore the model's performance on more complex tasks and datasets.

Autoregressive Transformer in analysis: We want to clarify that our model is autoregressive in that it can predict the $(L+1)$-th token based on the previous $L$ tokens by outputing a distribution over the possible tokens. Importantly, $L$ can be arbitrarily large, provided it is sufficiently large to capture the context needed for accurate predictions.

Limitation: We will include discussion on the limitations of our work in the revision.

Typo: Thank you for pointing out the typo. We will correct it in the revision.

Reference Allen-Zhu, Zeyuan, and Yuanzhi Li. "Physics of language models: Part 1, context-free grammar." arXiv preprint arXiv:2305.13673 (2023).

Thank you for your constructive comments. Please see our response below.

First layer word embeddings: As stated in our general rebuttal, we study a simplified model for a better understanding of the functionality of each component in detail. Although the first attention layer does not contain trainable word embeddings, it can be considered as

where we fix the word embeddings $W \in \mathbb{R}^{d \times d}$ as a random matrix with Gaussian entries with mean $0$ and variance $1/d$. This ensures that $x_l^\top Wx_{l'} \approx 0$ for all $l$ and $l'$ when the dimension $d$ is large. It can be interpreted as fixing $W$ at the initialization of the training, and it is common practice to set it as a random Gaussian matrix. Similar arguments can also be found in Bietti et al. (2023).

We remark that theoretical analysis of training dynamics becomes challenging when parameters for word embeddings are included. Previous works, such as Nichani et al. (2024) and Jelassi et al. (2022), often consider only positional embedding in the attention layer to focus on the positional information. Also, it has been empirically observed that the first attention layer's QK matrix has almost zero values for the token embedding parts during training (Nichani et al., 2024). This suggests that the word embeddings do not contribute to the induction head mechanism.

For our problem, when we incorporate trainable word embeddings, theoretical analysis of the training dynamics in Stage 2 becomes much more complicated. Since the first attention layer depends on both positional and word embedding weights, the dynamics of $w^{(h)}$ become intertwined with the dynamics of the word embedding weights. Then, simultaneous tracking of both dynamics is necessary, resulting in no closed-form solutions for these stochastic differential equations. A comprehensive analysis of the dynamics involving word embedding weights may be addressed in future work.

References

• Bietti, Alberto, et al. "Birth of a transformer: A memory viewpoint." Advances in Neural Information Processing Systems 36 (2024).,
• Nichani, Eshaan, Alex Damian, and Jason D. Lee. "How transformers learn causal structure with gradient descent." arXiv preprint arXiv:2402.14735 (2024).,
• Jelassi, Samy, Michael Sander, and Yuanzhi Li. "Vision transformers provably learn spatial structure." Advances in Neural Information Processing Systems 35 (2022): 37822-37836.

Thank you for your constructive comments. Please see our response below.

On $n$-gram data, simplified model & split training phases: We remark that the $n$-gram Markov chain data assumption is much more expressive and challenging than the bi-gram data assumption in previous works (Bietti et al, (2023), Nichani et al, (2024)). Rigorously analyzing the training dynamics for in-context learning of $n$-gram Markov chain is a significant open problem identified by Nichani et al, (2024), let alone with a more sophisticated transformer model that also include RPE, FFN, layer normalization, and residual link.

In this work, we address these challenges and reveal that the model learns a generalized induction-head (GIH) mechanism for $n$-gram data. We split the training phases to focus on the main hurdles. Experiments on simultaneous training of all parameters match our theoretical findings (see Figure 6 in Appendix C.1). Please refer to the general rebuttal for more details on model simplification and training phase splits.

On the modified $\chi^2$-MI: In the generalized induction-head mechanism, the modified $\chi^2$-MI criterion arises naturally from gradient calculations, not by design. For intuition, see Eqn (3.6) in Section 3.3, which relates the modified $\chi^2$-MI to the standard $\chi^2$-MI under certain symmetry conditions:

This criterion enables the model to balance maximizing mutual information and minimizing model complexity by selecting parents with high mutual information with the token to be predicted. For instance, selecting a larger parent set in the GIH mechanism results in less biased prediction, while also incurring larger estimation variance. In addition, we discuss in line 354-358 that the modified $\chi^2$-MI will become equivalent to the standard $\chi^2$-MI if the data is bi-gram, which reproduces the previous results on bi-gram data.

On the comparison to Bietti et al, 2023: The major differences between our work and Bietti et al. (2023) are as follows:

• (Data structure) Bietti et al. (2023) only consider 2-gram Markov data, while we consider general $n$-gram data. As discussed in the general rebuttals, going from bi-gram to $n$-gram data is highly nontrivial.
• (Model structure) Bietti et al. (2023) do not consider an MLP layer or layer normalization and only use single-head attention for their construction of the induction head mechanism. In stark contrast, we keep all these key components and show they are necessary for learning $n$-gram data and exhibit a rich interplay. As our model is more complex, the original analysis in Bietti et al. (2023) does not apply.
• (Theoretical understanding) Bietti et al. (2023) construct a hypothetical model with empirical evidence also on a synthetic dataset and provide an analysis for three gradient steps. Their analysis is essentially just around the initialization, which has no guarantee on the training trajectory or the limiting model. We theoretically understand the full training dynamics that lead to the emergence of the generalized induction head mechanism.

Can we "scale up" the analysis from bi-gram to $n$-gram? Simply scaling up the bi-gram analysis in Bietti et al. (2023) does not work for $n$-gram data, as the original mechanism handles only bi-occurred tokens. Let us follow the construction in Bietti et al. (2023) and show a counterexample. Consider the simplest $3$-gram data, where each token depends on the previous two tokens as parents:

• Let $x_l = w_E(l) + p_l$, where $w_E(l)$ is the word embedding for token $l$ and $p_l$ is the positional embedding for position $l$. Assume positional embeddings are orthogonal to each other, and $w_E(l)$ is orthogonal to all positional embeddings.
• The positional information is packed into the $W_{QK}$ matrix for the first layer, i.e., $W_{QK} = \lambda \sum_{l=2}^L \sum_{s\in \{1, 2\}} p_l p_{l-s}^\top$ with $\lambda$ being a large constant.
• After the first layer's attention, the model will output

Even if this model can identify the parent sets, it fails to capture order information, i.e., distinguishing which parent is first and which is second. Therefore, without changing the model structure, the original construction in Bietti et al. (2023) fails for $n$-gram data.

Major challenges for $n$-gram data: The counterexample shows why single-head attention without additional nonlinearity fails on $n$-gram data. Our work leverages multi-head attention, FFN, and layer normalization to learn history-dependent features. The main challenge lies in analyzing the dynamics of a more sophisticated model with cross-head interference and nonlinearity, which Bietti et al. (2023) did not address. We will include these discussions in the revision.

On the proof and "$\approx$" symbol: We want to clarify that Appendix E contains a proof sketch instead of a formal proof, and we use the "$\approx$" symbol to facilitate understanding of the proof logic. The complete and rigorous proofs are provided in Appendix F.

References:

• Bietti, Alberto, et al. Birth of a transformer: A memory viewpoint. 2023.
• Nichani, Eshaan, et al. How transformers learn causal structure with gradient descent. 2024.

Thank you for your constructive comments. Please see our response below.

On the model simplification:

• (Residual link) We want to clarify that the residual link is indeed contained in our simplified model via the use of disentangled transformer. As can beseen in Eqn (2.5), the second layer’s attention uses the original input $X$, which is copied from the first layer via a residual link. We apologize for the confusion and will make this point more clear in the revision.
• (Model Simplification) Simplifications in the model are commonly employed by previous work. For example, Bietti et al. (2023) also excludes the token similarity in the first layer of their simplified model, and Nichani et al. (2024) assumes symmetry conditions on the data which makes the second layer's QK embedding matrix as $W_{QK} = a I + b \bf{1} \bf{1}^\top$, where the all one part $b \bf{1} \bf{1}^\top$ assigns the same score among all tokens and the model is essentially the same as only using a single scalar $a$.

Please see also our response in the general rebuttal on the data/model simplification and the split of training phases. We remark that our model still contains the core components of the original transformer, including multi-head attention with relative positional embeddings, FFN, layer normalization, and also the residual link, which distinguishes our setting from the previous ones. Even with split of training stages, the analysis is highly non-trivial due to the cross-head interference and the nonlinearity in the model.

On the experiments: Our work is theory-oriented, thus we mainly conduct experiments on the synthetic data and the two-layer transformer model to validate our theoretical findings. Moreover, as we have pointed out in the general rebuttal, the $n$-gram Markov structure is already more expressive than the bi-gram data studied in previous works (Bietti et al. (2023)).

Questions on attention classifier: If our understanding is correct, the reviewer is asking if we can replace the second attention layer by a linear layer for classification. The answer is negative. The second attention layer plays the role of a classifier that first compares the query's feature (which contains its parents' information) with previous tokens' features, and then aggregates the tokens with the same parents. A simple linear layer cannot do this task as it requires the ability of processing sequential information. See Olsson et al. (2022) for more details on the induction-head mechanism. Please let us know if our understanding is correct.

On the typo: Thank you for pointing out the typo. Yes, $\psi_{\mathcal{S}^\star}(i)$ in line 247 should be $\psi_{\mathcal{S}^\star}(l)$ instead. We will correct this in the revision.

On the definition of the feature $\psi_{\mathcal{S}}(l)$: Thank you for pointing out the display issue in Eqn (D.3). In the definition of $\psi_{\mathcal{S}}(l)$ in Eqn (D.3), which we copy in the following line,

each $i_s$ refers to an index of elements in vector $x_{l-s}$ for $s\in\mathcal{S}$. Here, we have $|\mathcal{S}|$ vectors in $\mathcal{S}$ and each $i_s$ has $d$ choices. The feature $\psi_{\mathcal{S}}(l)$ is thus of dimension $d^{|\mathcal{S}|}$ as it is defined element-wise by iterating over all the indices $i_s$ for all $s\in\mathcal{S}$. We will correct this in the revision.

References:

• Bietti, Alberto, et al. Birth of a transformer: A memory viewpoint. 2023.
• Nichani, Eshaan, et al. How transformers learn causal structure with gradient descent. 2024.
• Olsson, Catherine, et al. In-context learning and induction heads. 2022.