Iteration Head: A Mechanistic Study of Chain-of-Thought
Authors: Vivien Cabannes
Charles Arnal
Wassim Bouaziz
Xingyu Alice Yang
Francois Charton
Julia Kempe
Venue: NeurIPS 2024 poster (2024)
Keywords: chain-of-thought
mechanistic interpretability
controlled experiments
Primary area: interpretability_and_explainability
TL;DR: We elicit an abstract head that enables to implement iterative algorithms through chain-of-thought for auto-regressive LLMs.

Abstract: Chain-of-Thought (CoT) reasoning is known to improve Large Language Models both empirically and in terms of theoretical approximation power.
However, our understanding of the inner workings and conditions of apparition of CoT capabilities remains limited.
This paper helps fill this gap by demonstrating how CoT reasoning emerges in transformers in a controlled and interpretable setting.
In particular, we observe the appearance of a specialized attention mechanism dedicated to iterative reasoning, which we coined "iteration heads".
We track both the emergence and the precise working of these iteration heads down to the attention level, and measure the transferability of the CoT skills to which they give rise between tasks.

Iteration Head:
A Mechanistic Study of Chain-of-Thoughe
Vivien Cabannes FAIR, Meta Al
Charles Arnal Datashape, INRIA
Wassim Bouaziz FAIR, Meta Al
Alice Yang FAIR, Meta AI
Francois Charton FAIR, Meta AI
Julia Kempe
Courant University and Center for Data Science, NYU & FAIR, Meta Al
Abstract
Shain-of-Thoughee (CoT) reasoning is known to improve Large Language Models both empirically and in terms of theoretical approximation power. However, our derstanding of the inner workings and conditions of apparition of CoT capabilities remains limited. This paper helps fill this gap by demonstrating how CoT reasoning emerges in transformers in a controlled and interpretable setting. In particular, we observe the appearance of specialized attention mechanism dedicated iterative reasoning, which we coined "iteration heads". We track both the emergence and the precise working of these iteration heads down the attention level, and measure the transferability of the CoT skills to which they give rise between tasks.
Introduction
In the rapidly evolving field of artificial intelligence, Large Language Models (LLMs) have emerged as pivotal component [45]. Their ability to understand, generate, and manipulate human language has opened up new avenues towards advanced machine intelligence. Interestingly, despite being primarily trained on next-token prediction tasks, L.MMs are able to produce much more sophisticated answers when asked to generate steps of reasoning 30 58 This phenomenon, often referred as Chain-of-Thoughh (CoT) reasoning, and illustrated on Table I appears paradoxical: on the one hand, LLMs not explicitly programmed to reason; the other hand, they are capable of following logical chains of thoughts to produce relatively complex answers.
Table 1: Chain-of- Chain-of-Thought consists in eliciting reasoning steps before answering (A) question (Q). [A] 210.
[Q] What x 8 3? [Q] What is 8 x 8 x 3? Take it step by step.
[A] 8 x 8 = 64, 64 3 = 192. It is 192.
Recent studies have shown that the class of problems transformer can solve with single-token prediction, i.c. by outputting single token meant to be the correct answer, is rather limited 24
15]. contrast, when transformers are allowed to freely generate tokens before providing a final answer, they can use those generated tokens as tape to emulate a form of Turing machine [48] This enables them to solve larger class of problems [38 (39 35) However, our understanding of why and how transformers gain CoT abilities when trained with next-token predictions remains limited. We aim to provide insights on the matter.
Summary of Contributions. We adopt a "mechanistic interpretability" approach [see 17 we work with simple, controlled problems and architectures that capture the key aspects of the problem and allow us observe and analyze, down the network's weights and attention, the emergence of CoT in our models. In practice:
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

We describe the simple yet rich setting of iterative tasks and iterative algorithms, including three simple examples: copying. polynomial iteration, and the parity problems.
We explain why such problems are hard to solve for transformers with single-token pre- diction. Conversely, we describe how a certain distribution of weights within the first two attention layers of transformer, which we call an "iteration head", enables a transformer to solve iterative tasks with CoT reasoning with relative ease.
We hypothesize that iteration heads naturally appear in transformers trained on (hard enough) iterative tasks, and verify this hypothesis in small-scale experiments.
Ablation studies demonstrate the impact of the training set and choice of hyperparameters' in their emergence. We also observe the good transferability of the iterative reasoning skills granted by the attention heads from one iterative task to another, from which we deduce the usefulness of data curation.
Our controlled yet illustrative experimental setup sheds light on the emergence of CoT capabilities in larger L.LMs, whose attention patterns are much harder to interpret. In particular, our experiments suggest that transformers are likely to develop "inner circuits" specially dedicated to multistep reasoning, which can then be applied, in combination with more specialized skills, to a variety of tasks that share the same underlying logical structure. This gives credible explanation of the strong CoT reasoning capabilities of current state-of-th offte-of-thh-rr the-art LLMs, as their training corpora nauman-writter texts, computer code) include many examples of complex multistep reasoning.
Related Work. This work is set in the realm of mechanistic interpretability [e.g., 43 A top- down line of work is trying to explicit algorithms implemented by transformers in the wild [e.g. 57 20 25]. although some findings might be fallacious (9] A bottom-up line of work, to which
belong, consists in building understandings from small models that are relevant for bigger models, in particular regarding in-context learning [see, e.g. 61 19
g 23 34
49
16. 60). In-context learning
relates to the reproduction of reasoning patterns that appear in a prompt or context [TO]. In contrast, our study of CoT relates to reproducing reasoning patterns that appear in the training set.
2 Controlled Setup: Learning Iterative Algorithms
Human language and human reasoning are often organized a multistep, cumulative fashion, with each new thought group of sentences building upon the ones that precede to work towards some final conclusion. LLMs naturally benefit from learning such reasoning patterns: not only are they prevalent through much of their training data, but they also represent an efficient way divide the total processing effort required into easier intermediate steps. In what follows, we choose to focus on iterative algorithms and iterative tasks as a controlled proxy for more general forms of CoT reasoning. Indeed, though conceptually simple, iterative algorithms exhibit key property: they are simultaneously hard learn for transformers using next-token predictions, and comparatively easy to learn using CoT reasoning. As such, iterative tasks are ideally suited to illustrate the usefulness of CoT reasoning, and to study its emergence.
We define iterative algorithuss, or iterative schemes, as follows: iterative algorithm is the combination of an input sequence, denoted as Sequence, and made of L elements (zi=)E[[L) (with [T. = [1, L)), and inter- nal state, denoted as 8, initialized to some default value so Init, and updated according to some rule s, = F(s xx) ) for some func- tion F. Pseudo-code illustrating the concept is provided by Algorithm I see also Figure I By extension, we informally call iterative task task which is naturally solved by outputting the end product of some iterative algorithm applied some input sequence. As an example, consider the parity problem, i.e. the problem of computing the parity of the sum of a sequence of Os and 1s: it can be easily framed as an iterative task. Using the notations of Algorithm|] [] Jr and otherwise. Then the final 8,, gives the parity of the sum. Although this task could also be solved in non-iterative fashion, the iterative solution can be seen as simpler and more parsimonious.
Algorithm Iterative Schemes
Init for in Sequence do 5 F(, end for return 8
the sequence is processed
let the initial state Init be equal to 0, and let F(s. x) be equal to 0 if is equal to

---

(Internal States)
8t-1
31
(Tokens)
-
2
-
Figure I: Arguably, reasoning involves updating an internal state (red) as new information is processed (green). The diagram above. where each element represents a piece of information. is an abstract depiction of this idea. This observation motivates our of iterative tasks proxy for more general reasoning processes. At first glance, a limitation of transformers their lack of an internal state, which makes challenging to implement this diagram
(Additional Layer)
(Internal Layer)
8y
Ft
(Tokens)
Tt-3
Ft-2
Tt-1
Ft
Figure 2: A single transformer layer cannot implement the diagram from [ outputs. This limitation be bypassed by stacking transformer layers, illustrated here. The red arrow indicates residual connection. This naive method requires as many layers as there are reasoning hops.
it cannor access its previous
(Intemal Layer)
(Tokens)
Tt-3
It 2
It's
I
Figure Alternatively, transformer could compute each s; from scratch. This implementation does not require additional layers, not parsimonious, which could lead computational inefficiencies. This explains the difficulty for a transformer to output the final answer of a chain of reasoning within token (i.e., with text-token prediction, and without chain- of of-uhhhught)
(Output Tokens)
8t-1
(Input Tokens)
Tt-3
Tt-2
Tt-1
8c-1
Figure 4: Chain-of-thouuht addresses this issue by explicitly representing the reasoning process in token space. The auto-regressive nature of LLMs (blue arrow) allows for the implementation of iterative algorithms, as long as the states are encoded in token space. A concrete implementation of such mechanism, which we call ieertiin head, is described in Section In practical applications of LLMs, could imagine earlier layers summarizing the 1-th input sentence information of varying token length) into It, as well as summarizing the generated CoT sentences into some 8t-1, with later layers translating the state s, into readable text.
other coherent semantic
3

---

Given certain number of transformer blocks, a transformer can only apply corresponding number of cross- operations to predict the next token. This limits its ability to learn even relatively simple iterative tasks, see Figures 2 and 3 E.g., consider the task where the input sequence is (I Iz)). possibly restricted to æ; E a. bj for some a < b, and the desired output is the product I[ x;. The product is multilinear in the entries of the input sequence (æ1..... Z1.)- If we model the output of an attention layer as sums of monomials of degree at most three in its input variables (due to key-query-value interaction), this makes learning the task quite hard for transformer, and bounds Can Transformer Learn Iterative Algorithms? Briefly summarized a transformer is composed of series of transformer blocks and operates on the space of sequences. A transformer block performs cross-operations that combine elements of a sequence through the use of attention heads to generate new sequences, and parallel operations applied to each element the use of feedforward layers (or MLP, i.e., multi-layer perceptrons). Auto-regressive transformers in particular are trained to perform next-token prediction; in other words, from a training corpus that contains sequences (zx) of tokens, the transformer is trained to output Æt++1 from the truncated sequence S, = (z+) (E(4)
sequence separately through
the maximum length of the sequences that a transformer with a given number of blocks can correctly process [see[5] and related literature for formal discussions on the matter that capture this log-depth dependency]-
However, when transformers are allowed to generate many tokens before providing an answer, which implicitly lifts the constraint on the number of operations performed by the transformer (see Figure D [see also [4, 21]. In particular, Figure 5 explained in the next subsection, illustrates how a two-layer transformer can implement what we named an "iteration head". This potentially enables it to learn any iterative algorithm, assuming that its second layer MLP is big enough to implement any successor function F (Ze, St- + St.
the picture changes [48 39 35
Synthetic Data. To study the emergence of CoT in controlled settings, we introduce two simple iterative problems. The first problem straightforward instance of Algorithm where the tokens and the states are elements of the finite field F, = Z/p2 (for some prime number p). i.e., integers modulo p. and the iterative step is the evaluation of polynomial function P F2(X, YJ in those two variables:
I I EE... Ep. Init 0,
2(,, 2(s.2)
(Polynomial Iteration)
Letting P(s, x) s + x and p = 2, the problem reduces the so-called parity problem:
x x 10, Init = 0,
F(s, = 8 + x.
(Parity Problem)
For ease of study, we also consider an even simpler problem: the copying problem, where the goal is simply to output an exact copy of the input sequence.
x E 10,
2(s,z)=2
(Binary Copy)
Note that there small abuse of notation here, since we interested in the unrolled sequence of states produced iteratively by Algorithm rather than the last token only. While copying may seem like overly simplistic task, it should be put in parallel with the seminal work of Olsson et al. 44 that advocates studying a copying mechanism to better understand in-context learning.
For each of our problems, encode the data, i.c. the sequences (z)). in the following form:
[Problem] [x1] [x2]]
[æz] [EoI] [s1] [$2]
[SL] [Eos].
first token indicates the problem generating the sequence (e.g., ''copy", or "parity"), after which
L input tokens F, are provided. The end of the input is specified by an end-of-input token (Eol).
Subsequent tokens encode the states of Algorithm iat at each iteration, until termination, which is
indicated by end-of-sequence token (EoS).
3 One Head to Rule Them All
We have discussed how transformers are limited in the iterative tasks that they can efficiently solve using only next-token prediction. By contrast, we describe in this section a certain distribution of weights which, if correctly learnt, would allow a two-layer transformer to efficiently implement iterative algorithms by using cfain-of-thought reasoning. After that, we perform various experiments to identify the conditions under which this theoretical circuit does appear.
We assume that the reader is familiar with the transformer architecture [see, e.g., 56 36 for details).

---

xt
Eol
8t-1
Sequence
[..]]
[...]
[...]
1xt Attn layer: Pt
PL+1
PL+t
k: "I am Eol."
qi "Are you Eol?"
PL+1
Pt
PL++
2nd Attn layer:
k: "I am pt.
q: "Are you
8t-1
Figure 5: Implementation of an iteration head with a two-layer transformer. Contiguous box: superposition in high-dimensional space. Blue: information brought to working space thanks to residual connections. Red: information brought thanks attention. Green: next-token prediction. The first layer MLP implements subtraction = (I. f) - (I. + 1) + for the second attention to he able query pa from (PL+11 (Pt441))t...)) The second layer MLP implements F to be able to predict 8; from (st-1, with the "end-of-input" mark assimilated to the initial state of Algorithm
3.1 Theoretical Circuit
This subsection describes natural way implement iterative algorithm with transformer. Let us consider prompt (xtx (e)))L.) to which we append special "end-of-input" (Eol) token that marks the end of the input. The completion sequence will be generated with the (-th element encoding for the state variable after / steps of Algorithm [ The 1-th new element (i.e. the 4. 1 + 1-th token of the full sequence) produced as follows. The first attention head is tasked with retrieving the position of the end of the initial prompt, i.c. the position of the Eol token. As illustrated in Figure 5 does so using query-key combination which informally encodes the question "Are you Eol?" and the answer "I am Eol.". Thus it extracts the positional encoding PL+1 (which is the value associated to the L + 1-th token) regardless of the sequence length L N, and brings PL-t into its working space as well through the residual connection (we formalize such statements in the next paragraph). As shown further below in Figure 3 the next attention head then generates query "Are you p.?" from P and PL++t: which is answered positively by key "I am p; associated to the /-th position. Hence the head retrieves the value associated to this position, which di. also obtains 8;-1 (or rather the approximation of that was produced at the previous step) through the residual stream. The MLP can finally compute the new state s; = Fst 1., from 8;-1 and Zt. This can always be done by large-enough MLP assuming that the second attention layer outputs all the relevant information regarding 1t-1 and vr, as a result of universal approximation [27. Note that the operations performed by the two attention layers the choice of F; their only goal to retrieve Ia and 8--1- We call the pattern of weights that realize these operations, well as the underlying algorithm, "iteration head".
totally independent from the precise iterative task considered, i.e. from
Information Superposition in Working Spaces. In our description of an iteration head, we have rather informally said that some variable x is "extracted" "obtained".. Formally, transformer transforms a sequence (xz)te(L) into series sequences (εi,)) where an index specifying layers, and ét,1 E Rdd, with Rd being referred to as the "working space". The input tokens and their positions are brought into working spaces using embeddings that are typically learned, then added together. Assuming that the working spaces are nigh-dimensional enough, and because those embeddings are learned, a transformer can use different parts of gd to simultaneously store token and positional information, as if those embeddings were actually concatenated rather than added. Likewise, transformer layers output variables are learned functions of their input; if needed, and assuming that d is large enough, 6t,41 can superpose some e and for different s, S t, in which case one may consider 6 as somewhat equivalent to the concatenation of ex,1 and e,,1- This is why our exposition above focuses on "information pathways", i.e. which variable is generated using which variable, and sentences such as "x, and 8,-1 are brought to the working space" should be understood as "some vector encoding the relevant information of both and st-1 is produced"'.
5

---

First attention map
Second atlention map
Train accuracy
Test accuracy
1.0
Query token
Accuracy
60
0.0
50
50
1000
0
1000
Key tokens
Key tokens
Epochs
Epochs
Figure 6: Left: attention maps learned for the parity problem when processing sequence of length L = 29.
Yellow indicates high attention score. The yellow line on the left plot shows that all the queries after the Fol token at position = 30 point to the Eol token. other terms, the first attention implements "Are you Eol?" query of Figure [ while the second implements the "Are you p.?" query. Right: accuracy dynamics for different sequence lengths when learning the parity problem. We observe fast learning of short sequences (we used the tabl0 color scheme of Matplodlib 28 with L E {8, 11, 14, 17, 32)) and characteristic staircase behaviors.
Approximate Iteration Heads. Iteration heads are an efficient, flexible and parsimonious way to implement iterative algorithms; as such, we expect them naturally emerge during training. Nonetheless, transformers have flexible architectures that can perform similar operations in different ways. Hence we also expect to see some variations with respect to the schematic architecture described above (see Figures 8 and 13). in particular when the embedding dimension becomes too small for the information superposition from the previous paragraph to be correctly implemented.
3.2 Learning an Iteration Head
In this subsection, examine the circuit that transformer actually learns when trained iterative task with chain-of of-thought.. We observe that the theoretical circuit described in the previous subsection does appear in practice.
Experimental Design. Unless otherwise stated, our experimental setup is as follows. Data was generated for the binary-copy, parity, and polynomial iteration problem with P(X,Y) XY + in Fn. For each length L from Lmin = 1 to Lmax 32, we generated n = 1024 input sequences of length L (corresponding to a total sequence length of 2L + 3) uniformly at random for both training and testing sets, creating datasets of N 16, 384 = 16 x 1024 sequences in total We utilized auto-regressive transformers [10] with two layers and one attention head per layer. The embedding dimension was set to d = 128, with learned absolute positional encoding added to the learned token embedding. The weights were optimized over 1000 epochs with Adam
29 a batch
size of 256, and a fixed learning rate set to 7 = 3 10-4, with default PyTorch parameters otherwise [46.. Our source code is available at at/htpps: : /ttl//gitttt http:////ccccbbooeeeeee...cc experiments consumed 12k V V100-hours.
Our
Attention Heads. In our initial experiment, we trained a transformer to solve either the parity task or the copying task only. The "iteration head" pattern of weights, described in the previous sub-section, can be seen in the attention maps of the first and second attention layers: an example is reported in Figure 6 Namely,
observe that when the model produces the (L + t 1)-th token (meant to be st), the following happens. The attention of the first transformer block is fully focused on the position of the Eol token, corresponding the informal query "Is this token equal to Eol?", creating a yellow line on the left of Figure This allows the first attention layer to retrieve the positional encoding PL+1 of the Eol token, in addition the positional encoding PL++ of the last token of the current sequence (the state st-1) coming from the residual stream. Using this information, the second attention layer is able to generate the informal query "Is this token in position t?", to extract some encoding of the token T,. Consequently, the attention of the second layer is fully focused on the position of the t-th entry, creating the yellow off- diagonal line on the second plot of
NNte that, even when I, E F2. the expected percentage of collision between the training and testing set decreases exponentially fast with L. ensuring minimal contamination between training and testing.

---

Test accuracy
Without CoT
With one layer
60
1.0
Emb. dim.
0.5
20
0.0
10
20
30
10
30
10
20
30
Max seq. len.
Max seq. len.
Max seq.
Figure 7: Test accuracy (where red indicates better performance) after learning the polynomial iteration task with P(X, Y) XY + in Fu for 1000 epochs. The accuracy is reported as a function of the embedding dimension (on the y-axis), and the maximum sequence length Lmax (on the z-axis). The learning conducted with two-layer transformer with CoT (left), without CoT (middle), with one-layer transformer with CoT (right). This illustrates usefulness CoT and two-layer architectures.
Figure 6
Using the information of and 8t, the following MLP can then compute = F( st). For given sequence length L, the learned attention maps were found to invariant to the input token (EL)) the standard deviation of attention patterns computed over all the data was negligible.
Successor Function. We empirically verified that after training a two-layer transformer with one attention head per layer the copying dataset, fine-tuning only the second layer MLP parity data enabled us to achieve 100% accuracy on the parity problem. In the context described previously, this transfer accomplished in fewer than 20 epochs of fine-tuning on the parity dataset. This confirms that the feed-forward layer of the second transformer block is computing the successor function F. In more general contexts, we found the successor function to be implemented jointly by the second layer MLP, the second attention values and output matrices, as well as the un-embedding matrix.
Position Subtraction. The accuracy the model decreases with the embedding dimension as shown on the left of Figure 7 Figure suggests that when dl is small, the first attention layer remains capable to accurately locate the "Fol" token, the second attention layer struggles retrieve dt- iteration head, the first layer's
This be explained as follows: in model that implements MLP, in conjunction with the second attention key and query matrices, is expected to generate the query-key pair "Are you pr" and "I am p."' by transforming Pt the one hand, and superposition of PL+t and PL+1 on the other hand, so that the end results are aligned. More abstractly, this encodes In high dimensions, is relatively easy to find
the positional subtraction L + t (L + 1) + set of weights to align a large number of vectors (viz., the ones encoding for (PL+1. PL+++) and for Pr). However, in lower dimensions, this only be achieved when the vectors form certain special geometrical patterns [see e.g. 40 633. which the model struggles to learn in our setting, at least with the optimization choices made.
Evaluation Dynamics. With a sufficiently small model, we might expect to understand the training dynamics quite well, which could in turn provide insights on design choices for larger models to minimize training costs. While detailed study of the training dynamics of our two-layer transformers beyond the scope this paper [see e.g., 41, we note several interesting facts that align with recent findings in the literature, such as the staircase profile of accuracy plots in Figure 6
as well as the usefulness of small batch sizes and large learning rates reported in Figure in Appendix
1 3 13
[I]]
despite the risk of loss spikes [12 59).
3.3 Ablation Studies
In addition to visualizing the attention map, we validated the learning of iteration heads through attention patching, i.e., intervening "patch" certain attention maps. Specifically, patching the ideal attention maps (i.c., zeroing out other routes) does not disrupt perfect accuracy. In contrast, zeroing out the focus the Eol by the first attention head, on p: by the second, reduced performance to near random.
observed that
Next-token Prediction; One or Two Layers. As an initial ablation study, we considered the polynomial iteration problem with P(X,Y) XY + 1 in F and compared the performance of
7

---

First attention map
Second attention map
1 Attn. peak.
2nd Attn. peak.
60
1.0
Query tnken
Emb. dim.
20
40
0.5
20
0.0
60
10
20
30
10
20
30
50
Max seq. len.
seq. len.
Key tokens
Key tokens
Figure 8: Left: attention peakiness score after 1000 epochs of learning with the polynomial iteration task parameterized by P) = XY + in F11 as a function of the embedding dimension d and the maximum sequence length Lmax. Right: example of attention maps of sub-sampled iteration heads.
CoT reasoning with next-token prediction (i.e., without Co'T), as well as CoT with single layer transformer. Two parameters come into play: the length of the sequence, which can be seen difficulty parameter regarding the data; and the embedding dimension, which can be seen as a model capacity parameter 53 53). The results, unequivocal favor of CoT and two-layer transformers,
[5
are reported in Figure 7
Alternative Circuits. Next, we explored other circuits that a two-layer transformer can learn to perform the same tasks as iteration head. We proceed by assigning a score to measure how closely the attention maps follow the patterns Figure 6 For the first attention map, we would like a measure of the concentration of the attention at the "Are you Eol? I am" query-key pairs, which correspond to the vertical yellow line from the left of Figure 6 For the second attention map, we would like measure of the concentration of the attention at the "Are you pr? I am" query-key pairs, which corresponds to the yellow off-diagonal from Figure 6 To avoid scaling issues, we define an attention a; "peaky" if it is greater than 50% after softmax averaging. We then measure the average number of peaky scores (within one sequence, and over sequences). i.e., we compute E 1.>.5 5 instead of Σ 04- This provides clear measure of the degree which transformer is implementing the attention mechanisms described in the previous section.
On the left of Figure 8 we report the average peakiness found for the first and second attention layers when training a transformer for different maximum sequence lengths Lmax and embedding dimensions d. We only one experiment per pair (Lmax, d) with fixed random seed. The texture of the figure indicates a certain randomness between runs for similar pairs (Lmax. d). The first attention layer almost always learns the "Are you Eol?" query-key combination, except when the maximum sequence lengths are very small. these cases, the transformer might find different circuits to solve for different sequence lengths.
The second attention layer tends to vary more. particular, for small embedding dimensions, the position subtraction might be challenging for the transformer to perform, leading to find alternative mechanisms. For instance, the first layer attention might perform a previous token copy when processing the input tokens, superposing the current token x; and the previous one xz-1 in the current working space. This allows the second layer to solely point at every other position, e.g., only attend positions t 2 N, either recovering the current token 2t, the previous 21-1- Implementing position subtraction towards even positions only reduces the learning capacity needed by the first layer MLP [T for related scaling laws]. Such sub-sampling mechanism is notably observed on the right of Figure g
Skill Transfer
Some skills might be easier to acquire when trained on certain data rather than others, highlighting the importance of data curation when training LLMs. For example, datasets of code or math [e.g., 26] might exhibit formal reasoning structures that compel LLMs to learn multistep reasoning patterns when trained on them, leading to improved reasoning abilities of the final model, even in plain English [see, e.g., (37. Our synthetic problems are ideally suited to highlight the mechanisms at play in these observations. This section illustrates how strategic data curation can facilitate learning to solve the parity problem. The crux find a dataset that helps the creation of iteration heads, which, once present, significantly eases the learning of the parity problem by transformer.

---

Attention Peakiness
Test Accuracy
1.0
1.0
polynomial
0.5
0.5
parity
transfer
0.0
0.0
0
500
1000
0
500
1000
Epochs
Epuchs
Figure 9: Left: Test accuracy as function of the numher of epochs, averaged over 100 runs, when learning the polynomial iteration task with P(X, Y) = =XY + Right: The second attention peakiness score indicates whether the network is learning the iteration head described Figure The green curve corresponds the accuracy on the parity problem when learning the polynomial iteration for the first 200 epochs before switching the dataset to learn the parity problem.
in F11 (blue) and the parity problem (orange).
4.1 Inducing Induction
We now address a simple question: can we "pretrain" model on a task A, and then "finetune" it on task in order to learn to solve the task B with a smaller total number of flops than if we were to leamn the task B from scratch? We will see that the answer is positive.
Figure 9 compares three learning scenarios. The learning of the polynomial iteration task corre- sponding to P(X, Y) XY + in F1 is reported in blue. The learning of the parity problem is reported in orange. Finally, the green curve represents training on the polynomial iteration task for 200 epochs (these epochs are not reported in the graph, hence the curve offset), before switching tasks and continuing the training the parity problem. When switching from the polynomial iteration dataset to the parity dataset, we chose to reset the Adam buffers to zero. Moreover, our default experimental parameters were changed to Lmax 16 and n 512, generating training and testing sets of N 8, 192 16 x 512 sequences. The left side of Figure(Jrepon reports testing accuracy averaged over 100 runs, along with its standard deviation. The polynomial iteration task is learned relatively quickly, while the parity problem takes longer. The right side of Figure Ireports the second attention peakiness score, capturing whether or not the second attention implementing the "Are you p. ?" query. After 200 epochs of training with the polynomial iteration task, the iteration head is formed, and fine-tuning the network on the parity problem for less than 30 epochs enables the reuse of this circuit on the parity data (green curve, right plot), thus solving the parity task (green curve, left plot). Overall, the data curation represented by the green plot enables the computation of parities in less than 300 epochs, compared to 1000 epochs when learning solely with parity data.
This example provides a controlled setup understand the usefulness of data curation when training larger models. It biases the model toward the implementation of specific circuits. In particular, adding code or math datasets the training of LLMs might induce the learning of more circuits that implement various forms of reasoning patterns. These could be viewed as atomic skills that could be reused to solve more generic problems [see, e.g.. for further discussions on skill factorization].
4.2 The Role of Inductive Biases
To illustrate the usefulness of data curation and skill transfer, we needed to find a problem that is hard to learn from scratch. The parity problem was well-suited to play this role in our synthetic setting. On the other hand, the polynomial iteration task with P(X, Y) XY + 1 in Fn was the easier task. One might wonder why learning with P(X, Y) = XY + in F11 turned out to be a simpler task than learning parities. which corresponds to P(X, Y) = X + Y in F2. Our intuition is that the parity problem can be solved in many different ways, which leads to competing signals in the gradient for updating the weights, reminiscent of the theoretical study by Shalev-Shwarz et al. [52] [see also 50 63). For example, we see on the right of Figure that the standard deviation of the attention peakiness score is quite high when learning with parity data. This can also be observed from the texture in Figure IT in the Appendix. This creates a challenging optimization landscape. In contrast, the polynomial P(X, Y) = XY was chosen to make the final state dependent on token order. Removing permutation invariance is useful to reduce the variety of circuits that can solve the polynomial iteration task, and seems to speed up the training dynamies. Finally, starting from pretrained model that already implements iteration head creates strong inductive bias toward the

---

Test Accuracy (CoT)
Without CoT
With one layer
1.0
Emb. dim.
100
0.5
50
10
20
30
10
20
30
10
20
30
Max seq. len.
Max seq. len.
Max seq. len.
Figure 10: Same as Figure except that we considered the parity problem, and 5000 training epochs.
iteration head circuit to solve the parity problem, allowing the parity problem to be learned within very small number of epochs.
To deepen our understanding of the parity problem, we conducted the same scaling study depicted in Figure 7 using the parity dataset. The results presented in Figure TO Training over 5000 epochs. The data generation process was slightly modified to generate all sequences of length less than L = log, (1048), and these were evenly split between training and testing (instead of generating redundant random sequences). In some sense, the parity problem can be considered relatively easy solve in a non-iterative fashion: simply add all the elements of sequences, and reduce the sum modulo two. In theory, single-layer transformer can use uniform attention to bring all the input tokens into superposition as input for the two-layer MLP layer, which is a universal approximator [27] As result, the parity with fixed sequence length can be solved with such an architecture [see e.g. e.g.E5l [5). Indeed, the bottom right of the left and middle plots in Figure [[]indicate that next-token prediction performs better than sfaiin-of-thought for sequence lengths up to Lmax = 32 with an embedding dimension of d 32. This is due to the difficulty of performing position subtraction necessary for the CoT circuit, compared to the relative ease of performing addition of up to bits with our two-layer architecture. Similarly, we found that a one-layer transformer was able to learn to produce correct CoT sequences for this task, demonstrating the existence of circuits fundamentally different from our iteration head to solve it. Anecdotally, the top right of the left and middle plots of Figure| 10 indicate that as the model capacity increases, next-token prediction tends to overfit the training data, while CoT induces the transformer toward understanding the underlying structure that generated the data.
conducted
Conclusion
In this paper, have explored the emergence of Chain-of-Thought (CoT) reasoning in Large Language Models (LLMs) through the lens of iterative algorithms. We have shown that, despite being trained next-token prediction tasks, transformers can learn to solve iterative tasks efficiently using CoT reasoning. In particular, we have demonstrated that a two-layer transformer can implement what named "iteration head",, enabling it learn any iterative algorithm, assuming that has enough feedforward layers following its two transformer blocks.
We have also shown that data curation can play significant role in guiding the model towards the implementation specific circuits. While our study has focused on simple, controlled problems and architectures, we hope that our findings shed light on the emergence of CoT capabilities in larger LLMs, whose attention patterns are much harder to interpret. particular, they suggest that transformers likely to develop "inner circuits" dedicated to multistep reasoning, which can then be applied to a variety of tasks that share the same underlying logical structure.
Interestingly, our work also highlights a limitation of the transformer architecture: they are stateless models. Indeed, our CoT implementation of Algorithm requires the generated states (st) to have rrequiress
token representation. This allows us to recover the state of the iterative algorithm at the root (i.e., the input) of the transformer. For complex iterative algorithms, or generic language modeling, would be more logical to maintain state internal to the model in the embedding space. The fact that GPT architectures do not allow this is arguably shorteoming of the current transformer architecture [see
32 6
22
47 62 for interesting discussions], [see also 42 33
10

---

Acknowledgments. useful discussions.
The author thanks Alberto Bietti, Carles Somingo-Ennichh and Denny Wu for
References
[1] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. The merged-staircase prop- erty: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks, 2022.
[2] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma. and Denny Zhou. What learning algorithm is in-context learning? Investigations with linear models, 2023.
[3] Maksym Andriushchenko, Aditya Varre, Loucas Pillaud- Vivien, and Nicolas Flammarion. SGD with large step sizes learns sparse features, 2023.
[4] Sanjeev Arora and Anirudh Goyal. A theory for emergence of complex skills in language models, 2023.
[5] Boaz Barak, Benjamin I.. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit, 2023.
[6] Adrien Bardes, Quentin Garrido, Jean Ponce, Xinlei Chen, Michael Rabbat, Yann LeCun, Mahmoud Assran, and Nicolas Ballas. Revisiting feature prediction for learning visual repre- sentations from video, 2024.
[7] David Bau, Jun-Yan Zhu, Hendrik Strobelt, Agata Lapedriza, Bolei Zhou, and Antonio Torralba. Understanding the role of individual units in deep neural network. Proceedings of the National Academy of Sciences, pnas.1907377 1907375177
(1148)))00071300777) September 2020. ISSN 1091-6490. doi: 10.1073/
/ dx. doi. wwwp/o......... 1073 pnas 1907375117
[8] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint, 2023.
[9] Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fernanda Viégas, and Martin Wattenberg. An interpretability illusion for BERT, 2021.
[10] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan. Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
[11] Vivien Cabannes, Elvis Dohmatob, and Alberto Bietti. Scaling laws for associative memories, 2024.
[12] Vivien Cabannes, Berfin Simsek, and Alberto Bietti, Learning associative memories with gradient descent, 2024.
[13] Angelica Chen, Ravid Shwartz-Ziv, Kyunghyun Cho, Matthew L. Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in MLMs, 2024.
[14] Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers, 2024.
[15] Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Mareus Hutter, Shane Legg, Joel Veness, and Pedro A. Ortega. Neural networks and the Chomsky hierarchy, 2023.
[16] Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning Markov chains, 2024.
[17] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
11

---

[18]
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: A theoretical perspective, 2023.
[19] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes, 2023.
[20] Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models, 2023.
[21] Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens, 2024.
[22] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
[23]
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? A case study on learning with representations, 2023.
[24] Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 2020.
[25] Michael Hanna, Ollie Liu, and Alexandre Variengien. How does GPT-2 compute greater-than? Interpreting mathematical abilities in pre-trained language model, 2023.
[26] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset, 2021.
[27] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural Networks, 1989.
[28] J. D. Hunter. Matplotlib: A 2D graphics environment. Computing in Science & Engineering, 9 (3):90-95, 2007. doi: 10. 109/MCSF.2007.55.
[29] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
[30] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large
language models are zero-shot reasoners, 2023.
[31] Andrey Kolmogorov and Vladimir Tikhomirov. E-entropy and e-capacity of sets in functional spaces. Uspekhi Matematicheskikh Nauk, 14(2):3-86, 1959.
[32]
Yann LeCun. A path towards autonomous machine intelligence, 2022.
[33] Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Kütteer, Mike Lewis, Wen tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Retrievallauggmented generation for knowledge-intensive NLP tasks, 2021.
[34] Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning, 2023.
[35] Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems, 2024.
[36] Tianyang Lin, Yuxin Wang. Xiangyang Liu, and Xipeng Qiu. A survey of transformers, 2021.
[37]
Yingwei Ma, Yue Liu, Yue Yu, Yuanliang Zhang, Yu Jiang, Changjian Wang, and Shanshan Li. At which training stage does code data help LLMs reasoning?, 2023.
[38] Eran Malach. Auto-regressive next-token predictors are universal learners, 2023.
[39] William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought, 2024.
[40] Neel Nanda, Lawrence Chan, Tom Licberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. 2023.
[41] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent, 2024.
[42] Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models, 2021.
12

---

[43] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 223911//istilll000244444 http://diil..pm/22222000/ee///oo
[44] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, 7ac Hatfield-Dodde Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022.
[45]
OpenAl et al. GPT-4 technical report, 2024.
[46]
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach Veti,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library, 2019.
[47]
Bo Peng. Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao. Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: Reinventing RNNs for the transformer era, 2023.
[48] Jorge Pérez, Javier Marinković, and Pablo Barceló. On the Turing completeness of modern neural network architectures, 2019.
[49] Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task, 2023.
[50] Elan Rosenfeld and Andrej Risteski. Outliers with opposing signals have an outsized effect on neural network optimization, 2023.
[51] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth, 2024.
[52] Shai Chalev-Shwartz Ohad Shamir, and Shaked Shammah. Failures of gradient- based deep learning, 2017.
[53] Steve Smale and Ding-Xuan Zhou. Learning theory estimates via integral operators and their approximations. Constructive Approximation, 26(2):153-172, 2007.
[54] Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. Transformers as recognizers of formal languages: A survey on expressivity, 2023.
[55] Vladimir Vapnik. The Nature of Statistical Learning Theory. Springer, 1995.
[56) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023.
[57] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small, 2022.
[58] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Fd Chi, Quoe Ie,, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models, 2023.
[59] Jingfeng Wu, Peter I.. Bartlett, Matus Telgarsky, and Bin Yu. Large stepsize gradient descent for logistic loss: Non-monotonicity of the loss improves optimization efficiency, 2024.
[60] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression?, 2024.
[61] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference, 2022.
[62] Jianyu Zhang, Niklas Nolte, Ranajoy Sadhukhan, Beidi Chen, and Léon Bottou. Memory mosaics, 2024.
13

---

[63] Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks, 2023.
Societal Impact. Mechanistic interpretability focused on understanding the key mechanisms at play in deep learning systems by tracing them down to the weights. It is often associated with AI safety, hoping that a deeper understanding of these systems can help us steer them to more "aligned" with "human values", and prevent AI dystopia scenarios. On the other hand, it could also prove useful in training more powerful models, which associated with significant societal issues linked to the rise of advanced AI systems. These issues too broad to be discussed in this paragraph.
A Additional Figures & Findings
Figure 11 studies the probability of finding the iteration head when learning with the parity data. As mentioned in the main text, showcases the bigger probability of learning other eircuits when learning the parity problem only.
st Atln. peak.
2nd Atin. peak.
60
Emb. dim.
40
0.5
0.0
10
20
30
10
20
30
Max seq. len.
Max seq. len.
Figure 11: Same figure Figure yet when learning with the parity dataset. The whitening of the row around d = 20 is duc to GPU failure, and should not be considered when parsing this figure.
Figure T2 showcases the usefulness of large learning rates and small batch size when using SGD, and the correction brought by Adam.
SGD
Adam
1.0
Band
108
0.0
10
10
Zaaminn 10
10
10
10
10
10
Learning
Figure 12: Test accuracy for SGD and Adam after 100 epochs.
Finally, Figure 13 plots some attention maps when learning with three layers transformer with two attention heads per layer. Knowing the iteration head cireuit, we are able observe similar circuit, yet with work shared across heads and layers.
Position embeddings. When we started this project, we were expecting to find some grokking structure emerged from the need to perform position subtraction. In particular, as mentionned in the main text, we were expecting this mechanism to appear as the position embedding dimension was small. When learning on the parity problem, we found that the network was not implementing the iteration head when the position embedding dimension was really small. This can notably seen on and
Figures and [4] We notably observed that freezing the position embedding does not change much the picture, which can be seen as a result of overparameterization. Similar type of observation were observed when learning with the polynomial iteration problem, reported on Figure 16
14

---

2dd layer, 1st head
3dd layer, 2dd head
Query token
20
30
20
20
Key tokens
Key tokens
Figure 13: Recovering the "who is ?" key key-query association, yet shared across layers and heads when training three layers transformer with two attention heads per layer.
Acc. learned pos, emb.
Acc. frozen pos. emb.
30
1.0
Emb. dim.
0.5
10
0.0
10
20
30
10
20
30
Max seq. len.
Max seq. len.
Acc. learned pos, emb.
Acc. frozen pos. emb.
1.0
dim.
20
0.5
10
0.0
20
10
20
Max
len.
Max seq. len.
Figure 14: The effect of small embeddings when learning the parity problem. The top corresponds to what has been learned after 1000 epochs. The bottom one corresponds to 5000 epochs.
Learned: 1" Altn, peak.
Learned: ndd peik.
30
Eaah. din
Max Jen.
Max len.
Atta, peak.
2dd peak.
30
Biik din.
20
0.5
10
10
20
30
Jel.
Figure 15: Attention learned when studying frozen learned positional embedding.
15

---

Learned PE
Frozen PF
Pos. cmb. dim.
30
1.0
20
0.5
10
0.0
5000
5000
10000
Epochs
Epochs
Learned, 1st Attn. peak.
Learned, 2nd Attn. peak.
Pos. emb dim.
30
1.0
0.5
10
0.0
5000
10000
5000
10000
Epochs
Epochs
Frozen, 1 Attn. peak.
Frozen, 2nd Attn. peak.
emb. dim.
30
20
0.5
0.0
0
5000
10000
5000
10000
Epochs
Epochs
Figure 16: Attention learned when studying frozen vs learned positional embedding. The setting is slightly different, we fixed the token embedding dimension to 32, and added the position embedding only on the first p dimension, where p was varying from to 32.
16

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the
paper's contributions and scope?
Answer: [Yes]
Justification: We composed the abstract and introduction after establishing the main results. Guidelines:
The answer NA means that the abstract and introduction do not include the claims
made in the paper.
The abstract and/or introduction should clearly state the claims made, including the
contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors? Answer: [Yes]
Justification:
The limitations of our synthetic data study should be clearly apparent to the
reader.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms
and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faireess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: The results were stated after being demonstrated. Some results that were obvious to us were not demonstrated (assuming one does not have to go back to Peano axiomatic to write additions).
Guidelines:
The answer NA means that the paper does not include theoretical results.
17

---

All the theorems, formulas, and proofs in the paper should be numbered and cross
referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes)
Justification: The codebase contains all the necessary details, and random seeds set to ensure full reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibilityy the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The code is freely available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurIPS code and data submission guidelines (https://ni cc/
public wttpdddde//ooedddeeeee ss onn ccy) for more details.
18

---

reproducible, they While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips. cc/
CodeSSumiise. ss oooooiicii for more details. The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments
should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The codebase is available with all the training and testing details. Some details (e.g., pre-norm, embedding dropout rate) were omitted from the main paper to prevent overwhelming the reader with technicalities that do not form the core of the paper. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Please note that in some 2D plots, "error bars" can be inferred from the texture of the contour lines, which provide clear sense of the randomness the process.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2- 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
19

---

For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiment??
Question:
Answer: [Yes]
Justification: We briefly comment on the resources used. The code is open source and the experiments are relatively light. An interested reader could easily calculate a more precise estimate of the energy cost reproduce the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpt//tw///aaaiooooo/////////aa
Answer:
[Yes]
Justification: Adherence guidelines was of significant importance us.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: This work is focused enhancing our understanding of machine learning systems. While this endeavor comes with clear issues, they are too broad for us to provide any meaningful comments within this context. Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
20

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper was deemed too theoretical to necessitate significant consideration regarding safety concerns.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We didn't cite everything (for instance, we didn't 't mention Python or CUDA), but we did cite PyTorch, the transformer paper, as well as GPT-3.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, paper thcode com/ / dat Bataseta t.s has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We have documented our synthetic datasets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/codeemoodll as part their
submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots. if applicable, as well as details about compensation (if any)? Answer: [NA]
21

---

Justification:
We have not conducted any experiments involving human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code Ethics, workers involved data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: See above. Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
22


Reviews:
[{"review_id": "K1XwlYsWxC", "reviewer": "Reviewer_YdTm", "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.", "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).\n\n[1] Hahn and Goyal, 2023, A Theory of Emergent In-Context Learning as Implicit Structure Induction, https://arxiv.org/abs/2303.07971", "limitations": "The paper should be more upfront about some limitations, as described under \"Weaknesses\".", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper addresses a timely question, with little prior work on the same question.\n* The paper links the outcome of probing with an intuitive explicit construction.\n* 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.\n* 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.\n* 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.\n\n[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]\n\n[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]\n\n[3] Anil et al, 2022, Exploring Length Generalization in Large Language Models, https://arxiv.org/abs/2207.04901", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and valuable feedback on our paper. We appreciate your concerns and will address them below.\n\nWe acknowledge that there is some ambiguity between the concepts of \"scratchpad\" and \"chain-of-thought\" and will review our manuscript to make it clearer. In our view, scratchpads reasoning is a generic method where a LLM edits a scratchpad before giving an answer. In contrast, CoT consists in asking the model to output intermediate steps of reasoning, as per [1], before providing a final answer. The original CoT paper is a bit confusing as it considers a few-shot learning setting (which relates to in-context learning). Thank you for bringing this to our attention.\n\nStudying the training dynamics of iteration heads seems feasible but not trivial, as the dynamics depend on many details that modify the dynamics (Adam, full batch, architecture, data...) and could be hard to quantify precisely. This is an exciting avenue for future work, especially if it sheds light on the importance of data curation, the possibility of \"sharded\" circuits, or the difficulty to choose between circuits for the parity dataset.\n\nRegarding length generalization, in our setting, the first layer MLP has to align superposition of embeddings to perform a subtraction. For a position $t$ never seen before, $p_t$ would be a random embedding, and there is no way to generalize. However, if the lengths are sampled with a Zipf law, a transformer may learn to perform well for long sentences, even if these were sparse in the dataset (to generalize, it only has to learn position subtraction, which can be done in a sample-efficient manner).\n\nThank you very much for the pointer to [2]. This paper is quite interesting, taking an information-theory and minimum description length perspective to shed light on ICL, CoT, composition in LLMs. In contrast, we focus on iteration heads as one \"reasoning template\" that could be learned by LLMs, focusing on its emergence and transferability when varying the training data. It allows us to perform more targeted ablation studies and give more concrete insights, for example, we also observe \"abrupt growth of training and testing losses\", but no \"grokking\" of the position subtraction (that stays in a \"memorization\" regime).\n\n[1] Kojima et al., 2023. Large Language Models are Zero-Shot Reasoners, [https://arxiv.org/abs/2205.11916](https://arxiv.org/abs/2205.11916)  \n[2] Hahn and Goyal, 2023, A Theory of Emergent In-Context Learning as Implicit Structure Induction, [https://arxiv.org/abs/2303.07971](https://arxiv.org/abs/2303.07971)\n\nWe hope this clarifies our position and addresses your concerns. Thank you again for your valuable feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C6d1CDBJPU", "reviewer": "Reviewer_uh9V", "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.\n\nThe 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.", "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.\n- How does the iteration head relate to other copying mechanisms in language models such as induction heads?", "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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper proposes an original copying mechanism that will be of relevance to other researchers studying information flow and circuits in transformers.\n- The theoretical circuit for copying information to the correct token positions is well explained.  Figure 1 clearly illustrates how the mechanism works.\n- 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.\n- 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.\n- 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.\n- 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].\n\n\n[1] Olsson, et al., \"In-context Learning and Induction Heads\", Transformer Circuits Thread, 2022.\n\n[2] Kaiyue Wen, Yuchen Li, Bingbin Liu, Andrej Risteski. \"Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammars.\" (2023).\n\n[3] Guillaume Alain, and Yoshua Bengio. \"Understanding intermediate layers using linear classifier probes.\" (2018).\n\n[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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and valuable feedback on our paper. We appreciate your concerns and will address them below.\n\nRegarding the relationship between iteration heads and other copying mechanisms in language models, such as induction heads, we understand that induction heads were introduced as an example of how to move information across \"registers\" (the $(e_{t,l})$) with copying mechanisms. The induction head circuit is based on a copy of the previous token computed thanks to position information, together with a copying mechanism based on semantic information. The iteration head relies on another copying mechanism to move information around, in particular, the iteration copy mechanism is purely positional. As for the $Q$-composition you mentioned, some operations in an iteration head could be performed jointly by several matrices (i.e., by \"composing\" them). For example, the position subtraction could be performed jointly by the first attention value and output matrices, the first layer MLP and the second attention query and key matrices.\n\nAs mentioned in the general rebuttal, our goal was not to elicit the appearance or non-appearance of iteration head in production-size LLMs, but come up with controlled experiments to shed lights on LLMs. As a side note, one may generalize iteration circuit to natural language examples beyond pointer arithmetic. For example, one may design a first layer that recognizes coherent semantic information of varying token length, and only associates to the end of these coherent blocks some position information encoded as a vector $p_t$. From there, the circuit would only iterate over these blocks of varying length, allowing to generalize the iteration circuit to natural language examples.\n\nRegarding the use of attention patterns as evidence for mechanistic theories, we remark that the attention operation is the sole one that moves information between registers across time (the $(e_{t,l})$ for varying $t$). As a consequence, the attention maps capture information regarding movement across positions. Moreover, in our setting, linear probing of the first keys will necessarily recover the EoI if the attention is focused on EoI, since the scalar product of the keys with a query could be seen as a linear probe (with weights defined by the query) looking for EoI among the different keys.\n\nThank you very much for your questions, we will integrate these answers to enhance our manuscript readability and relate it more explicitly with induction heads, and other probing techniques.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nnbwuHJQVG", "reviewer": "Reviewer_sSSN", "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\n- 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. \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. \n- They convincingly explain how this can be achieved theoretically\n- 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\n- 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.", "questions": "# Major Comments\n- *The following are things that, if addressed, might increase my score*\n- 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. \n- 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\n- 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\n- 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)\n- 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.\n\n# Minor Comments\n- *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*\n- 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)\n- 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. \n- 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\n- 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.\n- 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.\n- 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.\n- 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)\n- 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.\n- Line 64: This was confusing as CoT *is* \"next-token prediction\". Maybe \"Single step next-token prediction\" would be clearer?\n- 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\n- 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", "limitations": "See Questions", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Mechanistic interpretability of chain of thought is an important problem, and this work makes some progress\n- 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\n- The result that pretraining on polynomial iteration is a fast way to learn parity was surprising to me.\n- 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.\n- No evidence is provided that anything like iteration heads are learned in LLMs, the paper purely studies toy models\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and thoughtful review of our paper. We appreciate your feedback and your many valuable suggestions. To answer your major concerns, we detail our perspective below.\n\nRegarding the relevance of our toy model to real-world LLMs, we agree that our setup is simple and may not capture the full complexity of chain-of-thought reasoning in LLMs. One may complexify iteration heads by designing a first layer that segments tokens into different coherent blocks of varying lengths and outputs positional information $p_t$ at the end of each block, allowing then to iterate over blocks of varying length rather than individual tokens. Regarding real-world experiments, while we appreciate your experimental suggestions, we believe that a serious study requires a different mindset and goals than our paper, which aimed to gain intuition in a controlled setup. With this perspective in mind, we believe that our toy model offers many valuable insights.\n\nOur study offers a plausible story for the usefulness of training on structured data like code: it increases the likelihood of developing clean reasoning circuits, akin to the iteration head, that could be reused beyond code. To this end, we remark that if we blend both the parity and polynomial data, we learn a unique circuit used for both tasks, the `problem` token being recovered at the second layer to condition the MLP implementing the successor function. We did not emphasize these results in our original draft and will revise it. Moreover, the percentage of polynomial data in the blend will modify the training time needed to achieve 100% accuracy on both tasks, hence our remark on \"the usefulness of data curation\".\n\nAlthough we have not studied how CoT emergences from a training dynamics perspective, we have a good intuition of the conditions under which the iteration head appears. For example, the parity problem is sometimes solved by recomputing partial sums before spitting out each new output token. This is particularly the case when the sequence lengths are short compared to the capacity of the transformer (which relates to the embedding dimension). On the contrary, when sequences are long, we learn the iteration head, which is a more parsimonious solution. When the sequence length is really long, and the successor function is really simple, the first layer MLP has to learn many position subtractions, while the second MLP layer has not much to learn. This creates excessive pressure on the weights of the first MLP layer, which tends to be alleviated by doing partial pre-computation earlier one, and only learning to perform subtraction for even key positions (Figure 4, right).\n\nWe hope that our answer will help you better understand and appreciate our perspective. We also thank you for your minor concerns that we will integrate when reworking our draft.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3bdVufpDnO", "reviewer": "Reviewer_r3gz", "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.", "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.\n\n- 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?\n\n[2] Elhage, et al. Elhage, et al. A Mathematical Framework for Transformer Circuits. 2021 (https://transformer-circuits.pub/2021/framework/index.html)\n\n[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)", "limitations": "Yes. The authors are aware of the scope of their work and its limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper addresses a highly relevant topic: Understanding the mechanism underlying CoT capabilities of transformers.\n- The paper is well-written, and studies a simple, controlled setting where results can be easily replicated.\n- 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.\n\n- 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.\n\n[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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive review.\n\nWe will integrate a longer discussion to compare our circuit with induction heads. We understand that induction heads were introduced as an example of how to move information across \"registers\" (the $(e_{t,l})$) with copying mechanisms. The induction head circuit is based on a copy of the previous token computed thanks to position information, together with a copying mechanism based on semantic information. The iteration head relies on another copying mechanism to move information around, in particular, the iteration copy mechanism is purely positional.\n\nWe should have been clearer that the goal of our paper is not to study iterative circuits in production-size LLMs, but to introduce a controlled setting that enables rigorous study of many phenomena that are relevant for LLMs. In particular, the Llama and Mistral models use rotary embedding, while we consider absolute position embeddings (APE).\n\nWe test in practice that a rotary positional embedding (RoPE) can not easily implement iteration heads with only two layers. Indeed, rotary embedding easily implements an offset attention focused on the position $t-L$ when processing the position $t$ for some fixed $L$ (no need for MLP here). Yet, iteration heads need to set $L$ to be equal to the random sequence length. To recover the length $L$, a first layer could count the number of tokens before EoI, and a second layer could be focused on EoI to recover the length $L$, before using this information to recover a head that implements the right attention \"$t-L$\" in a third layer. The attention maps learned in practice are reported in the rebuttal pdf.\n\nWe appreciate your curiosity regarding RoPEs, and hope that our rebuttal will help you better understand and appreciate our perspective on our work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We appreciate the time and effort the reviewers spent reviewing our paper. The overall reception has been positive. The reviewers appreciate the "iteration head" definition and experiments regarding its transferability across tasks.

Reviewers expressed concerns that will help us greatly improve our work. In particular, we should have been clearer that the goal of our paper is not to study iterative circuits in production-size LLMs (indeed, ChatGPT was not able to solve our parity problem at writing time), but to introduce a controlled setting that enables the rigorous study of phenomena that are relevant for LLMs.
Following the reviewers suggestions, we will add more references to related works, and compare iteration heads with induction heads more thoroughly.

In order to help reviewers better appreciate our perspective, let us emphasize some results we appreciate from our study:
* Iteration heads illustrate how transformers learn generic reasoning patterns that could be used at inference time. Moreover, it showcases how they may learn to generate many tokens before providing a final answer through teacher-forcing. We expect similar (yet more complex) phenomena to explain the emergence of CoT in larger scale models. 
* We showcase how the apparition of these circuits depends on the training data mix (see Figures 3, 5, and 6).
* Iteration heads appear frequently in our practical settings. However, in extreme situations, we observe the emergence of alternative circuits: for example, Figure 9 shows that when using three layers, the parsimonious circuit can be "sharded" across layers, making the transformer more difficult to comprehend if we were not aware of iteration heads.


An avenue we are excited to explore in future work is studying how to "un-shard" circuits (i.e. going from the attention maps of Figure 4 or 9 to the ones of Figure 2). We believe that our controlled setup provides a good testbed toward architectural and optimizer changes in order to ensure that semantically coherent operations (e.g. the position subtraction, or the successor function) are performed in single layers.

While we hope that our generic answer will help you better understand and appreciate our perspective, we will stay at the reviewer's disposal during the discussion period to clear any remaining concerns.

Meta-review: This paper studies how small transformers are trained to perform simple algorithmic tasks (polynomial iteration, parity, binary copy) with scratch pads. The authors demonstrate that transformers exhibit iteration heads, which iterate through the input and perform index arithmetic to determine which position of the original input to process. The authors also conduct some experiments to verify their findings and show empirically that simple two-layer transformers can learn to implement the iteration-head algorithm when trained on iterative tasks.

The reviewers found the paper to be technically sound and impactful, with a recommendation for acceptance after addressing some weaknesses and concerns. This work will contribute to the understanding of the mechanism behind transformers performing algorithmic tasks.
Decision: Accept (poster)