Selective Induction Heads: How Transformers Select Causal Structures in Context

We thank the Reviewer for their feedback and questions. We provide our responses below:

Weaknesses:

The scope of this study is a bit narrow. The authors only discussed the detection of interleaved Markov chains

We clarify that our paper focuses on next-token prediction in interleaved Markov chains, not detection. This framework serves as a tool to investigate and demonstrate how transformers dynamically select causal structures in context. The task enables us to derive an interpretable solution, providing insights into transformers' capabilities beyond this specific setup. This effort is part of the general line of research on mechanistic interpretability, which investigates how transformers learn to solve next-token prediction tasks in simplified setups such as Markov Chains, linear regression, and other dynamical systems.

It is unclear whether an LLM can detect interleaved Markov chains. Maybe the authors can do some inferences on open-source models to verify this.

We clarify that this study does not aim to assess LLMs’ performance on interleaved Markov chains. Instead, it uses this controlled setting to investigate the underlying capabilities of transformers in sequence modelling. While applying this framework to open-source LLMs is an interesting direction for future research, it is beyond the scope of our work.

It remains uncertain whether an LLM can retain its ability to detect interleaved Markov chains when trained on data that is unrelated or only loosely related to such structures. Given the artificial nature of this task, it's more effective to assess detection ability in an out-of-distribution setting rather than through direct training on the same task.

Our preliminary experiments show that transformers do not generalize to causal structures unseen during training, indicating limited generalization to out-of-distribution data. However, a full evaluation of this interesting behaviour in LLMs is beyond the scope of our work, as mentioned above.

Questions:

Do you have any recommendations for an LLM to improve its understanding and learning of causal structures?

In our experiments, we observed that setting the softmax temperature in the first attention layer to $1$—instead of using the standard scaling factor $\sqrt{d_{QK}}$—appears to accelerate the learning of causal structures in the data. This finding suggests that exploring similar adjustments could potentially enhance large language models’ ability to learn and represent causal structures.

I appreciate the authors' detailed response during the rebuttal phase. Upon reviewing it, I believe the authors would benefit from adopting a broader perspective to more effectively convey why this problem/setting is significant. While there is a distinction between large language models (LLMs) and transformers based on model size and behavior, the attempt to simplify the analysis using a few transformer blocks does not necessarily provide conclusions that generalize to larger models.

For example, in response to the question, "Do you have any recommendations for an LLM to improve its understanding and learning of causal structures?" the authors suggested, "setting the softmax temperature in the first attention layer to 1 instead of $\sqrt{d_K}$ using the standard scaling factor appears to accelerate the learning of causal structures in the data." While interesting, this ad-hoc adjustment is unlikely to translate effectively to larger models trained on extensive datasets and optimized with fine-tuned hyperparameters.

Given my continued skepticism regarding the practical implications for real-world large model scenarios, my score remains unchanged.

Thank you for your positive evaluation of our work. Regarding the noted weakness, we agree that the scope of the synthetic task is limited compared to real-world complexities. However, as outlined in our general response, our primary goal is to use this controlled framework to isolate and examine the mechanisms transformers employ for causal structure selection. We consider this a valuable step towards addressing more complex tasks in future research.

We thank the reviewer for the useful comments, which will help improve the quality of our manuscript. Below, we provide responses to their questions:

The synthetic task of modelling interleaved Markov Chains is interesting and definitely a step in the right direction. But do you think it is yet a good enough proxy to model the complex dynamics of token interactions? In language modelling, often these interactions are informed by the semantics of the tokens and the context. Whereas I don't think here the token meanings are important in any way.

Our synthetic task is indeed a simplification of natural language, but it captures a critical aspect: context-dependent causal relationships.

Modelling such relationships, where causality depends on both context and token semantics, is inherently complex. Our interleaved Markov chain setup is a simplification where the causal structure is fixed within a sequence but is allowed to change across sequences. Token semantics remain important as they directly determine transition probabilities and, thus, the interactions between tokens. Transformers rely on the semantics of the tokens in a sequence to learn both the transition matrix P and to identify the correct causal relationship between tokens (by computing normalized likehoods based on the learned P).

To further approximate real-world language, where causal relationships can vary within a single context, a natural extension for future works would involve dynamically sampling the lag for each context, such that $p(k_t|x_{t-1},\dots,x_{t-c})$ for some fixed context length $c$.

In my opinion, one of the things that made the induction head work interesting is that they found heads behaving in similar ways in actual large language models which were trained on real textual data. I am not really sure if the same can be said about the selective induction head circuit you found here. If this is something that can only be observed in a synthetic task, do we really care about it? Or, do you have any hypotheses about how selective induction heads might manifest in models trained on natural language texts? And, how do you plan to investigate this hypothesis (maybe in future work)?

The reviewer raised an important point about the relevance of selective induction heads for real-world language tasks. The induction head mechanism, while insightful, is relatively simple: it detects repeated patterns and can copy tokens by focusing on previous occurrences of the same tokens. This mechanism is well-suited for specific types of positional dependencies but does not extend well to more nuanced, context-dependent relationships.

In contrast, our selective induction head mechanism is more complex. It involves not only identifying a relevant position in the past by simply comparing the semantics of single tokens but extracting and accumulating some statistics (likelihood of single transitions in our case) from the entire context and using it to select the position from which to copy. This allows the model to make decisions about which prior tokens carry the most relevant information based on the entire history, rather than relying solely on pattern repetition.

We believe that a similar mechanism may occur in real-world language models, where different layers specialize in distinct roles: some layers focus on extracting and encoding relevant information therefore having attention concentrating on fewer tokens, while others accumulate and integrate this information to determine which prior tokens are most significant therefore having more uniform attention scores. Similar findings have been shown in [1] in the context of attention sinks.

However, given the higher complexity than in standard induction heads, we anticipate that detecting selective induction heads in LLMs might require designing ad-hoc NLP tasks: if the causal structures in such tasks are sufficiently clear, we will expect to see similar attention maps to those observed in our synthetic task. We leave these investigations to future work.

We appreciate the Reviewer’s comments and questions. We have considered each point and outlined our responses below:

Weaknesses:

More discussion in the empirical validation would be useful, to explain the findings and walk through the plots in Figure 4. For instance, it would be beneficial to explain the difference between the 1 head (1H) and 3 heads (3H) cases, as the proposed construction scales the heads with the number of lags, and detail the implications of these differences.

We refer the reviewer to Section 4.4 where the paragraph Single-head Transformers discusses the implications of using fewer heads than the number of lags. Furthermore, we conducted new experiments on this topic, reported in Appendix C. Further details are provided in the general response Additional experiments for scalability to more heads and layers.

The scalability of the method remains unexplored; that is, what are the empirically-realized sets of lags that appear in real-world tasks, and how do they vary across tasks? Many natural language reasoning tasks have a non-linear and more complex causal structure, and it would be good to see how / whether this method extends to such tasks, beyond synthetic ones.

We acknowledge that the causal structures in our synthetic setting are simplified compared to the complexities of real-world tasks, where the number and nature of lags can vary significantly depending on the data. However, evaluating the applicability of lag selection to real-world tasks is beyond the scope of this work. Our primary goal is to provide a clear and interpretable analysis of the mechanisms transformers use to dynamically select causal structures in controlled environments.

While Proposition 1 is defended by construction (and the empirical findings), it is unclear how Claim 1 is supported empirically (although it is acknowledged that this claim has yet to be formally proven) through the results in Section 5.

We performed additional experiments validating our Claim 1 and added them in Appendix E. More details are also included in the general response under Additional experiments validating Claim 1.

Questions:

Continuing from the note on scalability in the “weaknesses” section: is it necessarily true that there is such an idea as a fixed “correct lag” that ever appears in real data? This is partially addressed in the many lags and non-contiguous lags setting, but appears to require further generalization for such a transformer model to be truly useful (for example, in natural language tasks).

We do not believe there is a single "correct lag" in real data. Instead, each sentence has its own causal structure. Predicting the next token requires inferring the correct causal structure from the context. Our interleaved Markov chain framework allows us to model this phenomenon, where each source generates sentences according to a lag, and the transformer needs to learn from the context which lag to use to predict the next token in the sequence.

Are there any core (architectural) considerations that might need to be made for scaling to larger data, from an efficiency standpoint? How many heads would be needed, and can this be reduced? This is addressed to some extent in the many lags setting in Appendices G and H.

As illustrated by our construction, the second layer stores single transitions for each lag within the embedding of the tokens. In the worst case, this would require the embedding dimension to scale with the sequence length. Nevertheless, in practice, we observe that transformers can find more efficient ways to store this information. For example, when training standard transformers, we fixed the embedding dimension to $64$ while having sequences of length $128$. Regarding the number of heads, we conducted additional experiments with varying the number of heads and reported them in Appendix C, Figure 5 (left). The results show that for lags $1,2,3,4,5$, transformers with less than 5 heads can't achieve the optimal performance.

We thank the Reviewer for their detailed comments and questions, which helped improve the quality of our manuscript. We addressed each point and provided our responses in the following:

The paper is well written, and great care is taken in the presentation of the technical details to the theoretical solution. However, the task here is pretty idiosyncratic and a specialized transformer architecture is used to make the analysis easier to perform. The connection to actual in-context learning scenarios seems like it might be a bit fraught, I’m not sure how much closer a lagging markov chain is to natural language than a normal markov chain.

We would like to clarify that we do not claim interleaved Markov chains are inherently closer to natural language. Instead, they are a specific type of order-$\hat{k}$ Markov chains that provide a useful framework for analyzing how transformers learn to select causal structures.

These Markov chains have the unique feature of being learnable through two distinct methods, both implementable by transformers:

1. In-context learning of the $\hat{k}$-gram (studied in [1]): This approach involves learning the Markov chain of order $\hat{k}$, corresponding to the largest lag in the sequence. However, it is sample-inefficient, with sample complexity growing exponentially with $\hat{k}$
2. In-context selection of the correct lag: In this more efficient method, the transformer selects the correct lag from the context to predict the next token, achieving optimal performance.

Our theoretical construction shows that a 3-layer attention-only transformer can achieve optimal sample complexity using the second method. Empirically, we observe that trained transformers (both standard and disentangled) naturally learn this strategy, which confirms that they implement mechanisms for lag selection rather than relying on inefficient high-order Markov chain learning.

By having transformers focus on in-context selection rather than in-context learning, our task offers an ideal framework for studying how transformers infer causal structures.

I also get the feeling that moving to a standard transformer architecture would make things a lot less pretty...

To show that our findings are not a byproduct of the specialized architecture of the disentangled transformers, we added some additional experiments to validate the compatibility of our construction and solution with standard transformers. See general response Disentangled vs standard transformers, additional experiments: for more details.

Generally, I am amenable to technical papers that deal with toy settings, so I think the core content of this paper is solid. However, I don’t think that the title and framing of the paper is accurate. The process of selecting causal structures in context is incredibly general and complex, and selective induction heads are just a step in the direction of understanding how in context learning occurs. I think the title and framing needs more transparently communicate that this work is largely theoretical. How you do this is totally up to you and I’m willing to raise my score if you address this concern.

We agree with the reviewer that selective induction heads are just the first step towards understanding the process of selecting causal structures in context. In the revision, we clarified in both the abstract and introduction that the focus is on a synthetic task, that we analyze 3-layer attention-only transformer, and the theoretical nature of the work (see modifications in red in the revised manuscript).

We are open to modifying the title if other reviewers agree. We propose to add Markov chain as "Selective induction heads: transformers and in-context causal structure selection in Markov Chains" to highlight the theoretical nature of the work.

[1] Nichani, Eshaan, Alex Damian, and Jason D. Lee. "How transformers learn causal structure with gradient descent." arXiv preprint arXiv:2402.14735 (2024).

4. Additional experiments for scalability to more heads and layers (Appendix C): As suggested by Reviewers acym and 7pgv we included a discussion of the cases with more layers or more heads in the second layer and provided additional experiments in Appendix C. We train standard transformers with learned positional and semantic embeddings in the same setup as for the experiments in the main text.
   • Varying number of heads: Figure 5 (right) includes scenarios with fewer, equal to, or more than $K$ heads. As predicted by our construction increasing the number of heads leads to performances that get closer to the maximum likelihood up to having the number of heads equal to the number of lags in the set $K$. Beyond this point adding more heads does not change the performance, as expected since ML is optimal.
   • Figures 6, 7, and 8 illustrate the attention maps for a 3-layer transformer with only 1, 2, and 3 heads respectively in the second layer, despite the task having 5 lags. Remarkably, even with fewer than $5$ heads, the attention maps remain consistent with our theoretical construction, displaying analogous patterns.
   • Varying number of layers: In Figure 5 (left) for lags in $\mathcal{K} = \{1,2,3\}$, we show the behaviour of the model with 2 layers and different combinations of heads $[1,1],[3,1],[1,3],[3,3]$. The results clearly show that transformers with $2$ layers can't solve the task. Moreover, increasing the number of layers beyond $3$ does not change the performance.