Attention with Markov: A Curious Case of Single-layer Transformers

We thank the reviewer for the encouraging feedback and insightful comments. We address the individual questions below.

Assumptions: We would like to highlight that Section 3 in the paper already delineates the assumptions behind our current theoretical results (Theorems 1-3). In particular, they are:

• 1. Data: We assume that the input sequence is a binary Markov chain with state transition probabilities $p$ and $q$, as illustrated in Figure 2b of the paper.

• 2. Model: We consider a single-layer transformer with single-head attention, without layer norm. For the input embedding, without loss of generality, we consider the (Uni-embedding) structure in Line 213 in the paper. The rest of the architecture is the same as the one illustrated in Figure 2a and Lines 124-131 in the paper. For this architecture, we study both the weight-tied and non weight-tied settings in Sections 4.1 and 4.2 respectively.

Vocabulary size: Please note that while our current theoretical results are for the binary setting, they readily generalize to the multi-state scenario as discussed in Section 4.3 in the paper. This demonstrates the generality of our results and observations.

Comparison to [1]: The suggested work [1] on ICL for single-layer transformers differs fundamentally from our setting. While [1] examines gradient flow dynamics in a single-layer transformer, their focus is on a linear regression task, which is markedly distinct from our sequential learning task with Markovian inputs. Therefore, a direct comparison between our results and theirs is not applicable.

References

[1] Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality. COLT, 2024.

Dear Reviewer FHWC,

We sincerely appreciate the time you have taken to provide valuable feedback for our work. As we are getting closer to the end of the discussion period, could you let us know if our responses above have adequately addressed your concerns? We remain at your disposal for any further questions.

Sincerely, The Authors

We thank the reviewer for the helpful feedback and insightful comments. We address the individual concerns below.

Characterizing the learning gap: This is an excellent question. The primary motivation behind our work is to demystify the curious empirical behavior of single-layer transformers when trained on Markov chains. This is what we precisely characterize in our paper through the loss landscape results in Theorems 1-3, corroborating the empirical observations in Figure 3. This can be viewed as a static explanation of these phenomena through the loss landscape. With regards to the learning dynamics and the presence of local/global minima for $p+q \lessgtr 1$, the following are two key points:

• (1) Existence results and what we empirically observe: Although our theorems provide existence arguments, they align closely with empirical solutions for both $p+q > 1$ and $p+q < 1$. For $p+q > 1$, as discussed in Lines 370–405, the transformer model empirically converges to local minima that nearly replicate the theoretical construction in Theorem 2. Similarly, for $p+q < 1$, the global minima described in Theorem 1 perfectly match the empirically observed minima, as shown in Lines 264–295 and Appendix D. Together, these results confirm that our theoretical constructions rigorously characterize the solutions learned by transformers.

• (2) Learning dynamics and why we observe above phenomena: Indeed, as noted in Line 407, understanding why the model converges to bad minima for $p+q > 1$ but not for $p+q < 1$ requires analyzing its learning dynamics. A follow-up work [1] addresses this directly and identifies the key reason behind this phenomenon. Specifically, as shown in Figure 1 of [1], the standard Gaussian initialization around zero falls in the basin of attraction for local minima when $p+q>1$, causing the model to get stuck at these minima. Conversely, for $p+q < 1$, the same initialization falls within the basin of attraction for global minima, enabling convergence to global minima.

In other words, our results can be viewed as a static landscape picture of the curious phenomena exhibited by single-layer transformers, whereas [1] complements it through learning dynamics.

Line 363, proof sketch of Theorem 2: Indeed, when $\boldsymbol{e} = \boldsymbol{a} = 0$, the model ignores input statistics entirely, yielding a constant prediction probability, as shown empirically in Figure 3c. While this holds regardless of weight-tying, Theorem 2 demonstrates that such solutions correspond to bad local minima with weight-tying, whereas they become saddle points without weight-tying (Theorem 3). This distinction arises from the impact of weight-tying on the Hessian at these critical points, as detailed in the proof of Theorem 2 (Appendix B.4).

References

[1] Local to Global: Learning Dynamics and Effect of Initialization for Transformers, https://arxiv.org/abs/2406.03072

We sincerely thank the reviewer for their thoughtful evaluation and for the strong appraisal of our results and contributions. With many interesting research directions emanating along this theme, your positive feedback is a definite encouragement to continue this line of research.

We thank the reviewer for their strong appraisal and encouraging comments about our work. We address the individual questions below.

Role of attention mechanism: Indeed, surprisingly, it turns out that a single-layer transformer can effectively model a first-order Markov chain without relying heavily on the attention mechanism, leveraging the skip connection instead. Note that this behavior arises naturally from the inherent structure of the transformer architecture itself and the input being first-order, not from any explicit design on our part. While the diminished reliance on attention might initially seem uninteresting, what is striking is that the model automatically learns to disregard attention during gradient-based training. Our experimental results validate this theoretical insight, as evidenced by the approximation error between the attention output $\boldsymbol{y}_n$ and skip-connection $\boldsymbol{x}_n$, and the underlying attention patterns visualized here: Figures 1 and 2 of https://anonymous.4open.science/r/Markov-6B43/Rebuttal/iclr-rebuttal.pdf

On the other hand, if we increase the depth of the transformer, it starts exhibiting in-context learning (ICL) capabilities when trained on Markovian inputs. Here, the attention layer plays a key role in realizing the induction-head mechanism, a cornerstone component behind ICL. The following works study this in detail: https://arxiv.org/abs/2407.17686v1, https://arxiv.org/abs/2402.14735

Pure MLP layers: Thank you for your question! If you're asking about training the transformer with the attention layer removed but keeping everything else, including the MLP, intact, our results suggest that the model would indeed succeed in learning the Markovian kernel. This setup simplifies the task compared to the current scenario. In the present setting, even though the attention layer is included, the model eventually learns to disregard it. Removing the attention layer outright would make it easier for the model to directly learn the first-order kernel. Figure 3 here indeed confirms this observation: https://anonymous.4open.science/r/Markov-6B43/Rebuttal/iclr-rebuttal.pdf