Title: What One Cannot, Two Can: Two-Layer Transformers Provably Represent Induction Heads on Any-Order Markov Chains
Venue: NeurIPS 2025
Accepted: yes
Keywords: transformers

interpretability
Abstract: In-context learning (ICL) is a hallmark capability of transformers, through which trained models learn to adapt to new tasks by leveraging information from the input context.  Prior work has shown that ICL emerges in transformers due to the presence of special circuits called induction heads. Given the equivalence between induction heads and conditional $k$-grams, a recent line of work modeling sequential inputs as Markov processes has revealed the fundamental impact of model depth on its ICL capabilities: while a two-layer transformer can efficiently represent a conditional $1$-gram model, its single-layer counterpart cannot solve the task unless it is exponentially large. However, for higher order Markov sources, the best known constructions require at least three layers (each with a single attention head) - leaving open the question: *can a two-layer single-head transformer represent any $k^{\text{th}}$-order Markov process?* In this paper, we precisely address this and theoretically show that a two-layer transformer with one head per layer can indeed represent any conditional $k$-gram. Thus, our result provides the tightest known characterization of the interplay between transformer depth and Markov order for ICL. Building on this, we further analyze the learning dynamics of our two-layer construction, focusing on a simplified variant for first-order Markov chains, illustrating how effective in-context representations emerge during training. Together, these results deepen our current understanding of transformer-based ICL and illustrate how even shallow architectures can surprisingly exhibit strong ICL capabilities on structured sequence modeling tasks.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis is a theory paper on multi-layer Transformer models. The paper studies inductions heads for the conditional k-grams problem: given a sequence of k tokens, can the model predict the next token correctly?  This problem is modeled as a Markov predictions in a k-gram and has connections with k-th order Markov chains and In Context Learning. \nAuthors establish that a 2-layer Transformer model is sufficient for building an induction head for this problem, regardless of the value of k. This is an improvement on previously known result which had established that 3 layers are sufficient. Authors also share a result on the dynamics of gradient descent for converging to the solution.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. the paper is well written and easy to follow\n2. the result is well exposed: reduction of 3 to 2 layers \n3. literature review is complete and nicely exposed\n\nWeaknesses\n\n1. optimization result makes very strong assumptions. I understand that the problem is really hard and unless we make strong assumptions it is not possible to prove anything useful. However we also need to understand the behavior outside the nice regime proposed here. Authors are assuming all layers but one already have the optimal value. \n\n2. lack of lower bound for number of layers\n\n3. optimization is disconnected from realistic situations because very strong assumptions are needed.\n\n4. limited numerical experiments: gap between theory and practice\n\nQUESTIONS:\n1. optimization: strong assumptions are made. What happens if this is not the case that all layers have their optimal values? how far are we from this setup when data is noisy?\n\n2. are 2 layers the fewest needed layers? how hard is it to establish a lower-bound?\n\nLIMITATIONS:\nmore numerical experiments on noisy / real data could help bridge the gap between theory results and practical architecture selection", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper shows that a two-layer transformer with a single attention head per layer can represent the conditional k-gram model for arbitrary Markov order k (Thm 4.1) by using the universal approximation of MLP. It further provides \n- (i) a depth-width trade-off construction with two heads in the first layer (Thm 3.2) \n-  (ii) a gradient-descent analysis showing that a simplified two-layer network learns first-order induction heads (Thm 5.1). \n\nEmpirical evidence is limited to visualizing attention heatmaps on synthetic data (Fig. 2).\n\n\nWhile the authors successfully construct a representation of the  $k$-gram model using a two-layer Transformer with a single attention head per layer, the construction heavily relies on the universal approximation capability of the MLP, which likely introduces a substantial number of parameters. This reliance renders the comparison with alternative architectures—such as deeper models or those with more attention heads but less dependence on the MLP—potentially unfair. Addressing these empirical gaps, along with providing a fair and comprehensive comparison of the total parameter counts across models, would substantially enhance the contribution.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n-  Prior work has demonstrated that representing the conditional $k$-gram model typically requires either a 3-layer Transformer with a single attention head per layer or a 2-layer Transformer with two heads per layer. The proposed new construction only use a 2-layer Transformer with a single attention head per layer. Therefore, they answer the question \"can a 2-layer, 1-head transformer represent the conditional k-gram\" successfully, bridging the gap between $1$-layer and $3$-layer.\n- The paper constructs the representation using the non-non-linearities of MLP. The proof sketch highlights how ReLU + LayerNorm can simulate an additional “virtual head”, a perspective likely useful beyond this setting. \n\n\n**Weakness**\n\n-  The proposed construction necessitates that the MLP in the first layer effectively encodes the entire one-hot vector within its weight matrix, which likely entails a substantial number of parameters. However, the authors appear to account only for the parameter counts in the embedding layer and the query, key, and value matrices, while omitting those in the MLP. As a result, the parameter comparison with alternative architectures may be misleading. For a fair and comprehensive evaluation, the total parameter count should include all components—namely, the embedding layer, attention mechanism, LayerNorm, and MLP—for each architecture under consideration.\n-  The assumptions underlying the gradient descent analysis appear highly artificial—for instance, Assumption (iv) involves a carefully designed initialization coupled with frozen, \"optimal\" MLP blocks. Consequently, it remains unclear whether standard stochastic gradient descent applied to a genuine two-layer, single-head transformer would converge to the proposed representation of the $k$-gram model. To substantiate the theoretical claims, it is recommended to conduct empirical experiments using the actual two-layer single-head architecture and to examine whether the trained MLP replicates the functional behavior of encoding the entire one-hot vector, as prescribed in the hand-crafted construction.\n- The paper lacks empirical comparisons between the two-layer single-head architecture and the three-layer single-head counterpart on realistic datasets such as natural language, music, or other real-world sequences that can be effectively modeled by the $k$-gram framework.\n\nQUESTIONS:\n- Could the authors provide a simplified explanation of the key technical challenges that arise when extending the optimization analysis to the case $k>1$. \n- The reviewer is curious whether the hand-crafted attention pattern proposed in the construction actually emerges in practice. Could the authors present attention scores from real-world tasks to examine whether such patterns are observed empirically?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

comments: SUMMARY:
This paper aims to prove that a transformer with two layers within which contains a single head self-attention is sufficient to realise the k-th order in-context learning for any k. This is so far the tightest (2 layers and 1 head in each layer) theoretical solution against previous works where either more layers or more heads are needed.

STRENGTHS AND WEAKNESSES:
Strengths

- This paper provides a rigorous proof to the requirements of induction heads in the transformer architecture.
- This paper highlights the importance of the MLP module in a transformer layer. With a proper construction of MLP weights, the tightest results are obtained. Whereas previous works which requires more layers just ignored MLPs. This is also an important evidence suggesting that the MLPs in the pratical LLMs can play roles in structural abilities other than merely Key-Value stores.
- Empirical results in a simplified 1st-order markov analysis have accorded with the theorical prediction by a large margin.

Weaknesses

- Despite the simplification of the experiment as stated by the authors, the analysis still relies on synthesized data training other than induction heads in real world problems. The gaps are underexplored.
- In order to convey a concise and clear thought, this paper sacrifices much details. While most of them are acceptable, how the 1st layer MLP constructs $u_n$ is not easy to understand. Since this part is the key difference of this paper's constructional proof, it may worth more spaces (spared from the warm-up section) to describe (or even with more figures).

QUESTIONS:
As suggested above, can you provide a more friendly description on how the MLP weights are constructed?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper shows how a k-th order induction head can be represented by a 2-layer transformer with a single attention head and MLP in each layer. The key idea is that the tokens can be expanded to a sufficiently high dimensional space where they are sparse (one-hot encoded, in this case) and attention weights can be used to combine them in superposition with an exponential moving average -- to represent k-grams. Prior constructions required two different moving averages (of window length k) shifted by one, but this paper shows how one can be used to derive the other since the token embedding at a timestep $x_{n-k}$ can be estimated using a suitable bias & ReLU on the moving average quantity $v_n$, and these can then be linearly combined to estimate $u_n$. In the second layer, the similarity between $u_T$ and $v_n$ is used to trigger an induction-based completion mechanism for $x_{T+1}$ copying $x_{n}$. While that shows that k-th order Markov chains are representable, it leaves open the question of whether they can actually be discovered during training by optimizing the cross-entropy loss on tokens. For this, the paper performs a perturbative analysis of stability of the proposed architecture weights (actually a subset, thereof) to show how optimization (at least within this subspace) converges to the anticipated weight patterns. The paper also points out an apparent \"depth-width tradeoff\" given that a k-gram model could be represented with a 3-layer transformer with a single attention head in each layer or alternatively a 2-layer transformer with two attention heads in the first layer.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n1. The proposal is conceptually elegant, including the minimality in components needed to represent k-th order Markov chains.\n2. The experimental results seem closely matched with the theoretical model.\n3. Explores not just in-principle representability, but also feasibility of discovery by training/optimization.\n4. Such principled understanding will hopefully have a positive impact on architectures and training methods over time.\n\n## Weaknesses\n1. With this framework, the largest allowed \"k\" will be practically limited by the finite-precision of activations. This is even more of a problem when inference uses quantized activations. [addressed in rebuttal]\n2. The observation regarding the depth-width tradeoff seems quite weak / superficial, and of unclear significance. It would be more interesting if the authors could flesh this out further. [partially addressed in rebuttal]\n\nQUESTIONS:\n1. Why should the attention weights have the form $3^i$ instead of any other scale for the exponential? Is this precise number validated by experiment, or only the exponential behavior?\n2. In Figure 2, layer 1 -- might it be interesting to have an additional plot depicting the attention map intensities on a logarithmic scale? Or maybe plot the variation across different diagnoals and inspect whether the weights are exponential as proposed (or whether there might be some other interesting pattern).\n3. In Appendix E (fig 5/9/13), why not compute a pseudo-attention map for the first layer?\n4. In Fig 2 Layer 2 (likewise figs 6-15) -- is it correct to say that the first \"k\" columns should actually be ignored when comparing with the pseudo-attention because the k-th order context is actually ill-defined for the first k tokens? If so, inspecting the \"difference\" between the \"learned attention\" and \"Optimal\"/\"pseudo\" attention might be misleading. The comparison might be less misleading (and also more favorable!) if you consider a sequence of length k+N and showed the attention matrix for a window of length N starting with the (k+1)-th token. Would you agree that's a more interesting plot?\n5. Have you considered the implications for inference with quantized weights/activations? This proposal might point to specific patterns of failure, which could then be compared against LLMs for more circumstantial evidence on the representations within LLMs.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}