Iteration Head: A Mechanistic Study of Chain-of-Thought

Summary

This paper describes an algorithm called the "iteration-head" by which autoregressive transformers can implement simple iterative functions, which may be related to the chain of thought (CoT) reasoning capabilities of language models. The paper shows empirically that simple two-layer transformers can learn to implement the iteration-head algorithm when trained on iterative tasks.

Strengths

- The paper addresses a highly relevant topic: Understanding the mechanism underlying CoT capabilities of transformers. - The paper is well-written, and studies a simple, controlled setting where results can be easily replicated. - The empirical findings support the theoretical intuition provided. Using transfer fine-tuning to verify the role of the MLP as $F$ in the iteration head algorithm is a neat contribution.

Weaknesses

- While the paper presents empirical evidence showing iteration heads can arise in small models, CoT prompting has been shown to have the strongest effects in larger models [1]. Because of this, the paper would benefit by including an analysis of whether the iteration-head algorithm is replicated in larger models, which benefit the most from CoT prompting on more complicated reasoning tasks. - Even in settings where the tasks studied are relatively simple (polynomial iteration, binary copying and parity problem), the "iteration head" algorithm is shown to be implemented in multiple ways [Lines 251-259, Figure 4;right, and Figure 9]. The fact that the implementation can be distributed across multiple heads/layers suggests that the algorithm description provided is helpful, but incomplete. Relating how a potentially distributed implementation of the iteration head algorithm might work would strengthen the argument that this sort of computation could arise generally in pretrained transformers. [1] Wei, et al. Chain-of-thought prompting elicits reasoning in large language models, 2023. (https://proceedings.neurips.cc/paper_files/paper/2022/hash/9d5609613524ecf4f15af0f7b31abca4-Abstract-Conference.html)

Questions

- The "iteration head" algorithm introduced and described here is very similar to the “induction head” algorithm from Elhage, et al, and Olsson et al. [2,3], with a few key differences. (1) The iteration head matches query-key vectors by position instead of matching by token/content. (2) The iteration head algorithm also allows for more complex operations beyond literal copying by using a general function $F$ instead. A discussion contrasting the iteration head algorithm with the induction head algorithm would provide some additional clarity on why this algorithm (the iteration head) is needed, and perhaps why induction heads would not be sufficient for enabling CoT reasoning in the cases you've presented. - Just curious -- if relative/rotary positional embeddings are used instead of absolute positional embeddings would you still need an MLP to help compute the positional difference (i.e. $p_{(L+t) - (L+1)+1}$), or could it be done solely via attention? [2] Elhage, et al. Elhage, et al. A Mathematical Framework for Transformer Circuits. 2021 (https://transformer-circuits.pub/2021/framework/index.html) [3] Olsson, et al. In-context learning and induction heads. Transformer Circuits Thread, 2022. (https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html)

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes. The authors are aware of the scope of their work and its limitations.

Summary

- The authors try to understand Chain of Thought in LLMs by studying tiny toy autoregressive transformers trained on algorithmic tasks (like parity, or polynomial iteration) that are much easier to solve iteratively, and which try to be a proxy for problems LLMs solve with CoT - These toy tasks have the form "take as input a sequence of n tokens, and output a sequence of n tokens, where the kth output token is a function of the kth input token and the k-1th output token", where the same model can handle a range of values of n. - A crucial subtask is thus to attend from the k-1th output token to the kth input token. The authors propose a circuit to solve this consisting of an "attend to EoI" head and an "iteration head" composing in two different layers. - They convincingly explain how this can be achieved theoretically - They provide some evidence that it is learned in practice by showing that the predicted attention patterns are learned, but do not mechanistically show the function predicted is implemented - The authors show that a model learns to do the parity task much faster if it is first "pretrained" on polynomial iteration, then finetuned on parity. Their interpretation is that the iteration head is a general algorithm for iterative tasks and can be repurposed for parity, but was easier to learn via polynomial iteration.

Strengths

- Mechanistic interpretability of chain of thought is an important problem, and this work makes some progress - The iteration head is an elegant algorithm for a non-trivial problem, and my best guess is that it is, in fact, learned in their setting - The result that pretraining on polynomial iteration is a fast way to learn parity was surprising to me. - They did many careful and thorough hyper-parameter sweeps, which resulted in interesting heatmaps

Weaknesses

- The iteration head task is significantly easier than real chain of thought, making the applicability of this work to LLMs unclear. In particular, real chain of thought has varying numbers of token per step, meaning the algorithm must be much more complex than a head attending with a fixed offset per input. - No evidence is provided that anything like iteration heads are learned in LLMs, the paper purely studies toy models - The evidence that iteration heads are, in fact, learned in the toy model is only circumstantial - it is showed that there's some of the predicted dependence on hyper-parameters, and that the attention patterns predicted form, but alternate hypotheses for the same attention patterns are not discussed or ruled out. I would be more compelled by weights-based analysis, as in A Mathematical Framework for Transformer Circuits, or Progress Measures for Grokking via Mechanistic Interpretability.

Questions

# Major Comments - *The following are things that, if addressed, might increase my score* - While toy models *can* be extremely enlightening if well-chosen, I think that most toy models work falls short due to failures in the model - if the toy model is a poor proxy for the task of interest, then the quality of the analysis is irrelevant. I broadly believe the author's story about iteration heads in the setting of this toy task, but their relevance to CoT in LLMs is unclear, and that is the main thing that could make this paper interesting. - The toy tasks here are much simpler than eg a MATH problem, as each step implements exactly the same algorithm (across *all* prompts) and each step has the same offset across an input (though not across prompts) and each step only needs to output one token (rather than generating a series of tokens, before figuring out the next step). It wouldn't surprise me if real LLMs had something like a "soft" iteration head that attends somewhere in the corresponding next sentence in a MATH problem, but I don't think the paper provides much evidence for this beyond my existing priors - I would be excited if the authors could either convince me that their toy task is a better proxy than I think it is, or show evidence that iteration head like things arise in LLMs - Example experiment: Find some CoT questions with structured input and output, eg multi-hop factual recall or simple maths questions, and an LLM that can do them (maybe with few shot learning so it keeps to the right structure). Look for iteration-y heads which attend from the k-1th line in the output (or final token of the line) to the kth (or k-1th) line in the input (taking the sum over all tokens in that line). Try ablating each head one at a time, and look at damage to CoT performance. Show that this is anomalously high when ablating iteration-y heads. (This may have issues with hydra effects where ablating a head does little - in this case, try ablating all iteration-y heads at once) - But I still think the paper is a weak accept rather than a reject, as it is technically solid work, iteration heads are an interesting idea and contribution that are plausibly relevant, and I found the finetuning results striking. # Minor Comments - *The following are unlikely to change my score, but are comments and suggestions that I hope will improve the paper, and I leave it up to the authors whether to implement them* - I found Figure 1 very clarifying, and would recommend moving it to page 1 or 2, where it would have helped me substantially (eg making clear that you're working with a two layer transformer) - The abstract and intro do not clearly state that you are working with a toy model on a toy task, which I was confused by, and think is somewhat misleading. I recommend editing to clarify, since this is a crucial detail. - I disagree with your use of superposition to describe the residual stream containing multiple pieces of information. Superposition specifically refers to the model compressing in more meaningful directions than it has dimensions, resulting in non-orthogonal vectors. But in toy models, there could easily be fewer directions than dimensions in an overparametrised model (see eg Progress Measures for Grokking via Mechanistic Interpretability), and so these could be orthogonal. I would just call it a linear representation - The "information superposition in working spaces" section is, in my opinion, a fairly standard idea in mechanistic interpretability, as discussed as "the residual stream as an information bottleneck" in A Mathematical Framework for Transformer Circuits, which could be helpful to cite. - The abstract implied that the paper may study the emergence of iteration heads over training ("how CoT reasoning emerges"), which doesn't really happen, and so felt a bit misleading. - In line 24 the work is motivated as our understanding of why transformers gain CoT abilities is limited. But isn't this clearly because they are trained on examples of people's thought process, eg worked answers to maths problems? More broadly, I think the authors somewhat overstate the degree to which CoT is surprising - isn't it obvious that if you let an LLM break a problem down into small steps, and do each in a separate forward pass, it performs better? The interesting question is how and when this is done. - Line 40: I found "from which we deduce the usefulness of data curation" confusing on first read, and wasn't sure what was meant (even after reading the whole paper I'm a big confused) - I thought the Related Work section was fairly sparse (though I understand the pain of conference page limits!). I might flesh it out by discussing how other toy models work explore applicability to LLMs, or comparing to other work that tries interpreting chain of thought. - Line 64: This was confusing as CoT *is* "next-token prediction". Maybe "Single step next-token prediction" would be clearer? - Line 202: It is shown that fine-tuning just the final MLP layer enables high accuracy on parity, and claimed that this shows it is computing the successor function. This seems too strong to me - it shows that the final MLP layer is *sufficient*, but not that other parts of the model aren't doing successor too - Line 341: The comment on "inner circuits" dedicated to multistep reasoning seems a bit of an overclaim to me. This is a plausible claim in general, but doesn't clearly follow from your results, as you only study models trained on single tasks (or trained on one then finetuned on another), while having a single circuit simultaneously work for many tasks is harder

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

See Questions

Summary

The authors propose a simplified formatting of the input, that represents the structure of chain-of-thought reasoning. They then propose a theoretical circuit that can efficiently copy information from different indices in the input to the appropriate position in the chain-of-thought to perform an additional iteration step. The authors create a new metric - attention peakiness - that measures how similar an attention pattern is to one that would have been generated by the theoretical circuit. They show that several transformer models trained on various iterative tasks in the specified format learn *iteration heads* as measured by the peakiness scores. The authors follow up with an experiment to show how transfer learning involves repurposing existing circuits for a new task. The models take longer to learn the parity task than the other tasks in the paper. However, when you first train a model on a simpler task that involves learning the iteration head, and then finetune it on the parity task, the model can learn parity with roughly an order of magnitude less total compute.

Strengths

- This paper proposes an original copying mechanism that will be of relevance to other researchers studying information flow and circuits in transformers. - The theoretical circuit for copying information to the correct token positions is well explained. Figure 1 clearly illustrates how the mechanism works. - The results from the inductive biases section provide evidence for the generality of iteration heads, showing how they are useful for several different tasks. This is a novel methodology that could be incorporated into other papers in the field of mechanistic interpretability. - The authors test for iteration heads in models trained on a variety of tasks (Polynomial Iteration, Parity), on a variety of embedding sizes, and sequence lengths.

Weaknesses

- There wasn't any discussion in the paper regarding how the iteration head relates to previously found copying circuits such as the induction head [1]. Induction heads can also apply "pointer arithmetic on position embeddings" to find token positions to copy information from (see Q-composition in induction heads). In general, this paper would benefit from more detailed discussion about how this newly discovered mechanism relates to other known internal mechanisms from the mechanistic interpretability literature. - Chain-of-thought in natural language lacks the structure present in the tasks from the paper. The model cannot just use "pointer arithmetic" to locate the relevant information in the context, as there is no guarantee on how the input will be structured. As a result, it is unclear whether actual language models would develop iteration heads to use in chain-of-thought reasoning. - The primary evidence for trained models learning iteration heads comes from observing attention patterns. Other works in the field of mechanistic interpretability have moved beyond using attention patterns as definitive evidence for mechanistic theories as they are insufficient to pin down the exact functionality of the model [2]. The paper lacks any additional evidence for the mechanisms beyond analyzing attention scores, such as linear probing [3], or activation patching [4]. [1] Olsson, et al., "In-context Learning and Induction Heads", Transformer Circuits Thread, 2022. [2] Kaiyue Wen, Yuchen Li, Bingbin Liu, Andrej Risteski. "Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammars." (2023). [3] Guillaume Alain, and Yoshua Bengio. "Understanding intermediate layers using linear classifier probes." (2018). [4] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, Jacob Steinhardt. "Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small." (2022).

Questions

- Including an additional metric for detecting iteration heads in addition to the peakiness score would greatly benefit the paper, such as a linear probe after the first layer trained to predict the index of the EoI token. - How does the iteration head relate to other copying mechanisms in language models such as induction heads?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors mention the limits of the toy models and controlled settings in the conclusion. The authors also mention some of the potential broader impacts of mechanistic interpretability research on society.

Summary

This paper provides a mechanistic study of small transformers trained to perform simple algorithmic tasks (polynomial iteration, parity, binary copy) with scratchpads. Findings show that trained transformers exhibit attention patterns dubbed "iteration heads": heads that iterate through the input, performing index arithmetic to determine which position from the original input to process in the next scratchpad step. The paper also provides a theoretical circuit construction performing this. A variety of experiments are provided to examine the behavior.

Strengths

* The paper addresses a timely question, with little prior work on the same question. * The paper links the outcome of probing with an intuitive explicit construction. * The paper includes convincing analysis of the learned circuit, with ablation and evidence of skill transfer.

Weaknesses

* I believe the biggest weakness is an overclaim issue: While the title and the abstract frame the paper in terms of "chain-of-thought", the setup considered is in fact that of *supervised* training with scratchpads [2], not that of classical chain-of-thought prompting where the ability to do CoT emerges from general pretraining [1]. The paper provides no results from CoT in language models. While the paper claims implications for larger LLMs (line 42-48), there is very little support for these claims. I believe the term "scratchpad" would be more appropriate than "chain-of-thought" for the subject matter of the paper. * There is no theoretical understanding about why this particular circuit would come out of SGD training. Learning theory is admittedly hard, but ths should be acknowledged clearly as a limitation. * Empirically, an important benefit of scratchpads lies in improved length generalization [3], whereas the setup here tests only on in-domain lengths. It would be interesting to know what happens when evaluating at longer lengths. [1] Wei et al, 2022, Chain-of-Thought Prompting Elicits Reasoning in Large Language Models, https://arxiv.org/abs/2201.11903 [cited in the paper] [2] Nye et al, 2021, Show Your Work: Scratchpads for Intermediate Computation with Language Models, https://arxiv.org/abs/2112.00114 [cited in the paper] [3] Anil et al, 2022, Exploring Length Generalization in Large Language Models, https://arxiv.org/abs/2207.04901

Questions

Related findings about mechanistic analysis of CoT-like computation (only with 2 steps, but in a setting where it actually emerges from general next-token prediction) is provided in [1] (particularly their Figure 14). The paper under review here might benefit from clearly pointing out how it provides novelty over that result (e.g., in that it considers longer chains and more extensive analysis). [1] Hahn and Goyal, 2023, A Theory of Emergent In-Context Learning as Implicit Structure Induction, https://arxiv.org/abs/2303.07971

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The paper should be more upfront about some limitations, as described under "Weaknesses".

Thank you for the reply

Thank you for the additional information and results regarding Transformers trained with RoPEs, I found it helpful. I also look forward to the mentioned future work idea related to “un-sharding” circuits where semantic operations may be split across layers. All of my other questions/concerns have been addressed and I think the authors’ proposed clarifications will strengthen the paper, though it is already a solid submission. As such, I have raised my score to a 7.

Reply to Rebuttal

Thank you for the reply. I believe that adding additional discussion about induction heads will greatly contribute to the paper. > As a consequence, the attention maps capture information regarding movement across positions. My point was that even low attention scores indicate some information movement. Even if the attention scores were completely uniform, there is still a chance that the model might be doing an "iteration head"-type operation. This makes it problematic to rely purely on attention patterns as an interpretability tool. I will keep my current score.

Extended thanks

As the discussion period comes to a close, we would like to extend our sincere appreciation to all of the reviewers for their time and effort. We were thoroughly impressed by the high quality of the reviews and the valuable suggestions provided to enhance our work. It is not uncommon to hear complaints about the quality of reviews and the low level of commitment to serving the community, but we want to emphasize that this does not at all reflect our experience. The dedication and constructive feedback provided by the reviewers were truly exceptional. Thank you very much to each and every one of you.