How to think step-by-step: A mechanistic understanding of chain-of-thought reasoning

1 Introduction

• Mechanistic analysis of CoT faces scale, adaptive backup-component, and factual-knowledge challenges, while earlier approaches offered only indirect views of the neural algorithm.The paper positions its investigation as an effort to address these challenges.Source: §1 · p. 2
• The study uses Llama-2 7B and fictional-ontology questions decomposed into decision-making, copying, and inductive-reasoning subtasks.The setting is intended to reduce MLP effects and focus on reasoning from provided context.Source: §1 · p. 2
• Task-specific attention-head sets significantly overlap, with induction-like mechanisms and early-layer information transfer between ontologically related entities.The initial layers also show a gradual transition from pretraining-biased to in-context representations.Source: §1 · p. 2
• Multiple attention heads write answers in parallel, drawing on generated CoT, question context, and few-shot examples.Answer-writing heads appear at or after the 16th decoder block.Source: §1 · p. 2
• Generated-CoT use varies across subtasks, while parallel pathways also collect answers directly from question context.The authors present this as empirical evidence about whether LLMs use generated CoT to answer questions.Source: §1 · p. 2

2 Related work

• The section reviews CoT reasoning research alongside mechanistic interpretability literature on language models.It covers theoretical and empirical investigations into CoT characteristics as well as different forms of CoT reasoning.Source: §2 · p. 3
• Nye et al. (2021) introduced scratchpads as intermediate computation, while Wei et al. (2022b) showed that natural-language CoT can support versatile reasoning.Nye et al. (2021) argued that intermediate answers can bypass a constant-depth attention bottleneck.Source: §2 · p. 3
• Theoretical and empirical studies address CoT’s computational properties, robustness, scale, causal use, and links to training data structures.Feng et al. (2023) analyzed circuit complexity; Tan (2023) found that intermediate-step information can contribute causally to answers.Source: §2 · p. 3
• Mechanistic interpretability has progressed from toy Transformers to GPT-2 and Alpaca-7B, but scaling circuit analysis is complicated by adaptive components and backup circuits.McGrath et al. (2023) call compensatory behavior after ablation the hydra effect.Source: §2 · p. 3

3 Background

• Transformer attention and MLP blocks update token representations through residual streams, and the unembedding projection maps the final stream to token logits.The model autoregressively predicts the next token from the input sequence.Source: §3 · p. 4
• PrOntoQA frames reasoning as stepwise question answering over fictional-entity ontologies, enabling subtask-level analysis while reducing interference from pretrained factual knowledge.False ontologies instead require models to override real-world knowledge.Source: §3 · p. 4
• A circuit is a model-computation subgraph responsible for specific tasks, with components as nodes and their interactions as edges.The paper follows Wang et al. (2023) notation.Source: §3 · p. 4
• Activation patching compares clean and corrupted forward passes to test whether replacing an activation restores the original output token.The example corrupts the name Mary to Anne and checks whether patching restores Mary.Source: §3 · pp. 4–5
• Knockout prunes circuit nodes; the study uses mean activations from false-ontology inputs for mean-ablation rather than destructive all-zero ablation.The notation also supports knocking out a set of heads and specifies when the final residual stream is considered.Source: §3 · p. 5

4 Task composition

• Decision subtasks select a reasoning path, copying subtasks reproduce context statements, and induction subtasks infer relations such as A is C from A is B and B is C.The ten-step CoT setting includes five decision, three copying, and two induction subtasks.Source: §4 · p. 5, Figure 1 · p. 6

4.1 Task-specific head identification

• Heads identified for induction perform across all subtask types, and the decision, copying, and induction head sets share many essential heads.The authors suggest induction heads combine representation mixing with contextual pattern matching, which can also support selecting and copying information.Source: §4.1 · p. 7, Figure 2 · p. 7
• Head importance scores alone do not determine pruning effects: accuracy changes with deviation from the mean importance score, δ(hj,k).The authors attribute strong contributions to heads at both low and high importance scores, depending on the associated KL-divergences.Source: §4.1 · p. 7
• For Llama-2 7B, pruning from the central bin of heads with the smallest δ(hj,k) values reduces its 1,024 heads to about 400 while retaining >90% of original accuracy.Source: §4.1 · p. 7
• Subtasks 0 and 5 are most sensitive to pruning; subtask 0 has less guidance for choosing a fact, while subtask 5 performs induction without a prior.Subtask 9 has a prior because its answer is framed as a question in the input.Source: §4.1 · pp. 7–8, Figure 3 · p. 8

5 Token mixing over fictional ontology

• A four-layer ReLU probe distinguishes positive, negative, and unrelated ontology pairs, whereas the linear classifier fails to distinguish them.Layer-specific classifiers are used to account for variation in residual-stream representations across layers.Source: §5 · p. 9
• Pair distinguishability improves across depth after early attention mixing, then diminishes after a certain depth, possibly as other task information accumulates.The authors point to successive attention heads continuing to move information between related entity tokens.Source: §5 · p. 9
• More in-context examples do not significantly raise peak classification performance, while classification performance becomes unstable as examples increase.The authors suggest task-specific information may accumulate earlier with more examples.Source: §5 · p. 9

6 Circuitry of step-by-step generation

• A task-specific subset of attention heads retains aggregate accuracy of 0.9 on inputs the full model answers correctly; unembedding maps residual-stream outputs to token space.The unembedding projection is interpreted as identifying the token most likely to follow an intermediate representation.Source: §6 · pp. 9–10

6.2 Subtask-wise answer generation

• Multiple heads write subtask answers to the last residual stream, with some high-probability answer-writing heads reused across subtasks.Examples of reused heads span subtasks 1 and 3, 2, 4, and 8, and 5 and 9.Source: §6.2 · p. 11, Figure 6 · p. 12

6.1 In-context prior vs pretraining prior

• Around layer 16, ontology-related token distinguishability peaks in layers 10–15, context abidance rises in deeper layers, and answer-writing heads appear after the transition.The depth-dependent context pattern is reported as task-agnostic; the authors associate the transition with a shift from pretraining to contextual priors.Source: §6.1 · pp. 10–11, Figure 4 · p. 10, §6.2 · pp. 11–12

6.3 Parallel pathways of answer processing

• The analysis traces answer information backward from answer-writing heads through attended residual streams, identifying earlier heads that write the attended content.Tracing stops at the first decoder block or a residual stream corresponding to the first input token.Source: §6.3 · p. 12
• Answer-writing pathways collect answer tokens from question, generated, and few-shot contexts, and their primary sources differ across subtasks.Generated-context sources occur for subtasks 4, 5, 8, and 9; subtasks 0–3 lack answer tokens in generated context, while 6 and 7 require the question index.Source: Figure 7 · p. 13, §6.3 · pp. 13–14
• The pathways implement different algorithms: few-shot answer-collection paths can use contextually mismatched entities, and their presence decreases as generation proceeds.Generated-context collection distinguishes inductive subtasks 5 and 9 from decision-making and copying; subtasks 4 and 8 can draw from the preceding step.Source: §6.3 · p. 14

• Despite differing reasoning requirements across CoT stages, functional components remain almost the same, with different algorithms composed of induction circuit-like mechanisms.Source: §7 · p. 14
• Attention heads move information between ontologically related or negatively related tokens, making their representations distinguishable from early layers through the middle.This occurs zero-shot, while in-context examples exert pressure to mix other task-specific information more quickly.Source: §7 · p. 14
• Multiple attention-head pathways write each CoT subtask’s answer token to the last residual stream, with differing probabilistic certainty.Source: §7 · p. 14
• Answer pathways collect tokens from generated CoT, question, and few-shot contexts, providing empirical evidence that the model uses generated context while answering.Source: §7 · p. 14
• At layer 16, residual-stream content and head functions shift: pretraining associations dominate earlier, while the in-context prior takes over around and after the rift.Answer-writing heads appear after the rift; heads that collect answers from few-shot examples are bounded by the earlier half.Source: §7 · p. 15
• Causal interventions that alter reasoning pathways should account for parallel answer-generation routes; future work could test whether the rift also involves MLP-based memorization.Source: §7 · p. 15
• The study analyzes fictional-ontology reasoning with a restricted few-shot CoT template and does not address MLP roles, so free-form and zero-shot dynamics remain outside its scope.Source: §7 · p. 15

A PrOntoQA question example

• The examples include both fictional and false ontologies.The fictional example derives that Max is bright; the false-ontology example derives that Rex is spicy.Source: A PrOntoQA question example · p. 18
• The examples trace a queried entity through ontology statements to a final true-or-false answer.One example also includes an irrelevant statement whose entities are outside the ontology relevant to the question.Source: A PrOntoQA question example · p. 18

B Knockout

• For knockout analysis, false-ontology activations are stored for each subtask and mean activation is knocked out.The accompanying example shows prompts progressing through successive reasoning prefixes.Source: B Knockout · pp. 18–19

C.1 Decision-making heads

• Decision-making heads are identified by patching individual heads after corrupting the final input token and recording answer-token probabilities.The input sequence and correct answer token are denoted SDecision and sans, respectively.Source: §C.1 · p. 19

C.2 Copy heads

• Copy heads are identified by measuring attention to the context token that supplies the tail entity after a statement’s head entity is chosen.For example, after choosing Rompus, the model must copy grimpus from “Rompus is grimpus.”Source: §C.2 · p. 19

C.3 Inductive reasoning heads

• Inductive-reasoning heads are identified by patching activations at three token positions in the pattern “[A] is [B]. [B] is [C]. [A] is” and measuring KL divergence.The answer token in this example is [C].Source: §C.3 · p. 20

C.4 Performance of inductive reasoning heads for each subtask

• For each subtask from 0 to 9, attention heads are grouped by their inductive-reasoning importance score.The reported statistics concern head removal according to importance in inductive reasoning.Source: §C.4 · p. 20
• Heads in selected importance ranges are knocked out, and the resulting accuracy is recorded.The authors report using this procedure to examine context-abidance and answer-writing pathways.Source: §C.4 · p. 20
• The selected ablations aim to retain relative accuracy near 0.9 of the full model’s correct predictions.Relative accuracy is defined as the ablated model’s correct predictions divided by the full model’s correct predictions.Source: §C.4 · p. 20

D Probing for token mixing

• Positive pairs are immediate ontology neighbors, negative pairs include incompatible entities, and unrelated pairs come from disjoint ontologies.Distractor entities seed the unrelated pairs.Source: D Probing for token mixing · pp. 20–21
• The probe uses 28,392 training pairs and 9,204 test pairs, with positive, negative, and unrelated pairs equally represented.The entities do not overlap between training and testing.Source: D Probing for token mixing · p. 21

E Prompts used and model performance

• Generation uses six-shot CoT examples, with input statements shuffled and two distractors added to reduce positional patterns.The reported performance evaluation modifies statements by removing articles and quantifiers and replacing negative entities with incorrect entities.Source: E Prompts used and model performance · pp. 21–22
• Accuracy is 52.5 on false ontology and 35.9 on fictional ontology for Llama-2-7B.Accuracy requires a complete match of all intermediate reasoning subtasks and the final true-or-false statement.Source: E Prompts used and model performance · p. 22

F Attention probability

• Attention to task-specific tokens increases in absolute value from the first to the sixth few-shot example.The reported pattern is observed across subtasks.Source: F Attention probability · p. 22
• For subtask-0, attention peaks at tokens aligned with subtask-0 across the first six examples, including Max in the fifth and Wren in the sixth.The figures show slight or small peaks in the first through sixth examples.Source: F Attention probability · p. 22, Figure 10 · p. 23, Figure 11 · p. 24
• The plotted scores are averages over all heads’ attention probabilities at the last residual stream, shown over input-context tokens and few-shot examples.Figures 10–12 cover the examples in pairs; Figure 9 covers input-context tokens.Source: F Attention probability · p. 22

G Information flow

• To trace information flow, the authors start at answer-writing heads and follow the heads that wrote their source information backward until reaching the first decoder block.They identify source and target residual streams and the information written by each head.Source: G Information flow · p. 26
• In subtask 1, seven heads write the answer token, while eight heads collect “imp” from the input question context.The example predicts “imp” after “Stella is.”Source: Figure 13 · p. 29
• In subtask 5, three heads write the answer token, collecting it through four heads each from the input question and generated context.The example predicts “j” after the generated response “Stella is impus. Impus is jompus. Stella is.”Source: Figure 14 · p. 30