How to think step-by-step: A mechanistic understanding of chain-of-thought reasoning
Authors: Subhabrata Dutta
Joykirat Singh
Soumen Chakrabarti
Tanmoy Chakraborty
Venue: Accepted by TMLR ()

Abstract: Despite superior reasoning prowess demonstrated by Large Language Models (LLMs) with Chain-of-Thought (CoT) prompting, a lack of understanding prevails around the internal mechanisms of the models that facilitate CoT generation. This work investigates the neural sub-structures within LLMs that manifest CoT reasoning from a mechanistic point of view. From an analysis of Llama-2 7B applied to multistep reasoning over fictional ontologies, we demonstrate that LLMs deploy multiple parallel pathways of answer generation for step-by-step reasoning. These parallel pathways provide sequential answers from the input question context as well as the generated CoT. We observe a functional rift in the middle layers of the LLM. Token representations in the initial half remain strongly biased towards the pretraining prior, with the in-context prior taking over in the later half. This internal phase shift manifests in different functional components: attention heads that write the answer token appear in the later half, attention heads that move information along ontological relationships appear in the initial half, and so on. To the best of our knowledge, this is the first attempt towards mechanistic investigation of CoT reasoning in LLMs.

Published in Transactions on Machine Learning Research (07/2024)
How to think step-by-step: A mechanistic understanding of chain-of-thought reasoning
Subhabrata Dutta* ITT Delhi, India
subha0009@@gaall.com com
Joykirat Singh' Independent
j
Soumen Chakrabarti IIT Bombay, India
soumen@cse..itb.omm
Tanmoy Chakraborty ITT Delhi, India
tanchak@üttd.o oc.
Reviewed on OpenReview: https: //
// openreview. net/ forum? id- wtLDkQVtyC
Abstract
Despite superior reasoning prowess demonstrated by Large Language Models (LLMs) with Chain-of-Thought (CoT) prompting, a lack of understanding prevails around the internal mochanisms the models that facilitate CoT generation. This work investigates the neural sub-structures within LLMs that manifest CoT reasoning from a mechanistic point of view. From an analysis of Llama-2 7B applied to multistep reasoning over fictional ontologies, we demonstrate that LLMs deploy multiple parallel pathways of answer generation for step-by- step reasoning. These parallel pathways provide sequential answers from the input quest.ion context as well as the generated CoT. We observe functional rift in the middle layers of the LLM. Token representations in the initial half remain strongly biased towards the pretraining prior, with the in-context prior taking over the later half. This internal phase shift mani- fests in different functional components: attention heads that write the answer token appear in the later half, attention heads that move information along ontological relationships ap- pear in the initial half, and so on. To the best of our knowledge, this the first atlempl towards mechanistic investigation of CoT reasoning in LLMs. The source code and data are made available at https: h/ httts::/gg..ub..cm/////traaaa ttttubb..co//ooy//aattt///o To-Think- Step-by-Stey
1 Introduction
Recent advancements with Large Language Models (LLMs) demonstrate their remarkable reasoning prowess with natural language across diverse set of problems (OpenAT 2024 Kojima al., 2022: Chen et al., 2023). Yet, existing literature fails to describe the neural mechanism within the model that implements those abilities how they emerge from the training dynamics and why they are often brittle against even unrelated changes. One of these capabilities of LLMs that has boosted their potential in complex reasoning is Thouugt (CoT)) prompting (Wei et al. 2022b; Kojima et al. 2022). Instead of providing a direct answer to the question, in CoT prompting, we expect the model to generate a verbose response, adopting a step-by-step reasoning process to reach the answer. Despite the success of eliciting intermediate computation and their diverse, structured demonstration, the exact mechanism of CoT prompting remains mystorious. Prior attempts have been made to restrict the problem within structured, synthetic reasoning to observe
"Equal contribution.
Work done Research Assistant IIT Bombay #Corresponding Author

---

Published in Transactions on Machine Learning Research (07/2024)
the LLM's CoT generation behavior (Saparov & He 2023 ; context perturbation towards causal modeling
has also been used (Tan 2023). A few recent approaches seek to associate the emergence of CoT reasoning ability in LLMs with localized structures in the pertaining data (Prystawski et al. 2024 Wang & Wang 2023). However, these endcavors produce only an indirect observation of the mechanism; the underlying neural algorithu implemented by the LLM remains in the dark.
Recent developments in the mechanistic interpretation of Transformer-baset models provide hope for uncov- ering the neural 'algorithms' at work inside LLMs (Elhage et al., 2021; Nanda, 2022). Typically, mechanistic methods seek to build a causal description of the model; starting from the output, they localize important components of the model (o.g., an attention head or multilayer perceptron (MLP) block) by activation patching (Wang et al. engineering CoT-prompting in foundational models. The capability of generating CoT is largely dependent on model scale (Wei et al. 2022al Saparov & He 2023). On the other hand, reverse-engincering a large model becomes a wild goose chase due to the hydra effect (McGrath et al., 2023 - neural algorithmic components within LLMs are adaptive: once we 'switch off' one functional component, others will pitch in to supply the missing functionality. Furthermore, in most real-world problems involving multi-step reasoning, there are implicit knowledge requirements. LLMs tend to memorize factual associations from protraining as key-value caches using the MLP blocks Geva et al. (2021) Meng et al. (2022). Due to the large number of parameters in MLP blocks and their implicit polysemanticity, interpretation becomes extremely challenging.
2023 Zhang & Nanda 2023). However, there are some implicit challenges in reverse-
In this paper, we seek to address these challenges, thus shedding light on the internal mechanism of Transformer- based LLMs while they perform CoT-based reasoning. Going beyond the typical toy model regime of mechanistic interpretation, we work with Llama-2 7B (Touvron et al. 2023). To minimize the effects of MLP blocks and focus primarily on reasoning from the provided context, make use of the PrOntoQA dataset (Saparov & He 2023) that employs ontology-bar based question answering using fictional entities (see Figure [ for an example). Specifically, we dissect Co reasoning on fictional reasoning as a composition of a fixed mumber of subtasks that require decision-making, copying, and inductive reasoning
Section
We draw on three prominent techniques for investigating neural algorithms, namely, activation patch- ing (Nanda 2 probing classifiers (Belinkov) 2022), and logit lens (nostalgebraist, 2020) to disentangle different aspects of the neural algorithm implemented by Llama-2 7B. We find that despite the difference in reasoning requirement of different subtasks, the sets of attention heads that implement their respective algorithms enjoy significant intersection Section 4.1b. Moreover, they point towards the existence of mech- anisms similar to induction heads (Olsson et al., 2022 working together. In the initial layers of the model (typically, from first to 16th decoder blocks Tlama-2 7B), attention heads conduct information transfer between ontologically related fictional entities; e.g., for the input numpuses are rompuses, this mechanism copies information from numpus to rompus so that any pattern involving the former can be induced with the latter Section 5 Interestingly, in this same segment of the model, we find gradual transition in the residual stream representation from pretraining before in-context prior, i.c., the contextual information replaces the bigram associations gathered via pretraining Section 6.1
Following this, we seek to identify the pathways of information that are responsible for processing the answer and write it to the output residual stream for each subtask. To deal with the abundance of backup circuits (Wang et al. ) we look for such parallel pathways simultaneously without switching off any of them. We find that multiple attention heads simultaneously write the answer token into the output, though all of them appear at or after the 16th decoder block (*r Section (.2). These answer tokens are collected from multiple sources as well (i.c., from the few-shot context, input question, and the generated CoT context), pointing towards the coexistence of multiple neural algorithms working in parallel.
Our findings supply empirical answers to a pertinent open question about whether LLMs actually rely on CoT to answer questions [Tan] 2023 Lampinen et al. 2022): the usage of generated CoT varies across subtasks, and there exists parallel Ieel pathways of answer collection from CoT as well directly from the question context
Scction 6.3) Here again, we observe the peculiar functional rift within the middle of the model: answer tokens, that are present in the few-shot examples but contextually different from the same token in the question, are used as sources by attention heads, primarily before the 16th decoder block.
2

---

Published in Transactions on Machine Learning Research (07/2024)
To the best of our knowledge, this is the first-ever in-depth analysis of CoT-mediated reasoning in LLMs in terms of the neural functional components.
2 Related work
In this section, we provide an overview of literature related to this work, primarily around different forms of CoT reasoning with Transformer-basod LMs, theoretical and empirical invostigations into characteristics of CoT, and mechanistic interpretability literature around language models.
Nye et al. (2021) first showed that, instead of asking to answer directly, letting an autoregressive Transformer generate intermediate computation, which they called scratchpads, elicits superior reasoning performance. They intuitively explained such behavior based on the observation that, in a constant depth-width model with ) complexity attention, it is not possible to emulate an algorithm that requires super-polyun super-pooyynoiil computation; by writing the intermediate answer, this bottleneck is bypassed. Wei et al. (2022b) demon- strated that LLMs enabled with unstructured natural language expressions of intermediate computations, aka CoT, can demonstrate versatile reasoning capabilities.
Multiple recent attempts have been made toward a deeper understanding of CoT, both empirically and theoretically, Feng et al. (2023) employed circuit complexity theory to prove that - (i) Transformers cannot solve arithmetic problems with direct answers unless the model depth grows super-polynomially with respect to problem input size, and (ii) Constant depth Transformers with CoT generation are able to overcome the
said challenge. Liu et al. (2022) showed that Transformers can learn shorlcul solutions to simulate automata with standard training (without CoT); however, they are statistically brittle, and recency- biased scratchpad training similar to that of Nye ct al.) (2021) can help with better generalizability and robustness. [Saparov & He (2023) dissected the behavior of LLMs across scale on synthetically generated multistep reasoning tasks true, false, and fictional ontologies. In somewhat counterintuitive finding, Lampinen et al. (2022) observed that even if an LLM is prompted to generate an answer followed by explanation, there is a significant improvement over generating an answer without any explanation. While this might point toward a non-causal dependence between the explanation and the answer, Tan (2023) showed that L.Mss do utilize, at least partially, the information generated within the intermediate steps such that the answer causally dependent on the CoT steps. Prystawski et al. (2024) thcoretically associated the emergence of CoT with localized statistical structures within the pertaining data. Wang & Waan (2023) investigated reasoning over knowledge graph to characterize the emergence of COT. They identified that reasoning capabilities are primarily acquired in the pertaining stage and not downstream fine-tuning while empirically validating Prystawski et al. (2024)'s theoretical framework in a more grounded manner. Bi et al. (2024) found that while reasoning via program generation, LLMs tend to favor an optimal complexity of the programs for best performance.
The above-mentioned studies either 1) dissect the model behavior under controllable perturbations in the input problem or 2) construct theoretical frameworks under different assumptions to prove the existence of certain abilttiesddiaabiii ddsabiiiiisss. Mechanistic interpretability techniques go one step deeper and seck to uncover the neural algorithm deployed by the model to perform a certain task. [Elhage et al. (2021) studied one-, two- and three-layer deep, attention-only autoregressive Transformers and their training dynamics. A crucial finding of their investigation was the existence of induction heads a composition of two attention heads at different layers that can perform pattern copying from context. Olsson et al. (2022) empirically observed the simultancous emergence of induction heads and in-context learning the training dynamics of similar toy Transformer models. Similar analyses on phenomena like polysemanticity, superposition, and memorization have been performed in the toy model regime (Elhage et al. 2022 Henighan et al. 2023). Beyond these toy models, Wang et al. (2023) analyzed circuits in GPT-2 responsible for Indirect Object Identification. Wu ct al. 2023) proposed causal abstraction model to explain prico-tagging in the Alpaca-7B model: given a range of prices and a candidate price, the task is to classify if the candidate price falls within the given range. Extrapolating mechanistic observations from the regime of toy models to 'production' LLMs is particularly challenging; with billions of parameters and deep stacking of attention heads, identifying head compositions that instantiate an algorithm extremely difficult. Furthermore, as McGrath et al. (2023) suggest, different components of the LLM are loosely coupled and adaptive in nature, for which they coined the term hydra
3

---

Published in Transactions on Machine Learning Research (07/2024)
effect ablating an attention layer may get functionally compensated by another. Wang al. (0233 also identified similar information pathways called backup circuils that can take over once the primary circuits are corrupted.
3 Background
In this section, for completeness, we briefly introduce the concepts and assumptions necessary for delving into the interpretation of CoT reasoning.
The Transformer architecture (Vaswani et al. 2017 consists of multiple, alternated atlention and MLP blocks, preceded and followed by an embedding projection and a logit projection, respectively, Given the vocabulary as V, we denote the embedded representation of the i-th token 9; V in an input token sequence S of length N (i.e., the sum of embedding projections and position embeddings in case of additive position encoding) as z) E Rd, where d is the model dimension. We denote the content of the residual stream corresponding to the i-th token, input to the j-th decoder block, as .)-1, which becomes æ)-1 after the attention layer reads and writes on the residual stream. The initial content of the residual stream is the same as the token embedding xh. Assuming H number heads for each attention layer, the operation of k-th attention head at j-th decoder block on the i-th token's residual stream denoted as yý.k = h,,x(1))1 (I)) vi.k E Rd. Then #) = + Ex vi is the content of the residual stream immediately after the attention operations. Each attention head hy,k is parameterized using four projection matrices: query, key and value projection matrices, denoted by w.'k, w w.s.. kk I #, respectively, and output projection matrix wibh R xd, Note that in the case of Llama models, the position encoding is incorporated via rotation of query and key projections before computing dot-product (Su et al. 2023); we omit this step for brevity, Similarly, the action the MLP block can be expressed as 23 = (m)), with æ = # + 23 denoting the content of the residual stream after decoder block j. A After processing the information through L number of decoder blocks, a feedforward transformation U E R2 VI, commonly known as the unembedding projection, maps the content of the rosidual stream into the logit space (i.c., distribution over the token space). Given sequence of input tokens S {si, ........ , the autoregressive Transformer model predicts output token 8;+1:
8141 = arg max LM (afoogiil
(1)
Flgg
Fictional ontology-basee reasoning, as proposed by [Saparov & He (2023) as PrOntoQA, provides a tractable approach to dissect Coll generation. The reasoning problem is framed as question-answering on a tree-based ontology of fictional entities (sce Appendix[A for examples of ontologics and reasoning prob- lems framed). This eases Iwo major challenges in our case: (i) Mechanistic interpretation requires input or activation perturbation and recording the results of such perturbations. With CoT, one then needs to re- peat such perturbation process for all the constituent subtasks. Unlike free-form CoT reasoning, PrOntoQA provides clearly demarcated sequence of successive steps that can be analyzed independently. (ii) The solution to most real-world reasoning problems heavily requires factual knowledge, so much so that a sound reasoning process might get misled by incorrect fact retrieval. A fictional ontology ensures zero interference between the entity relationships presented in the question and the knowlede acquired by the model in the pertaining stage. This further minimizes the involvement of the MI.P layers in the neural mechanisms implemented by the model as they typically serves as the parametric fact memory of an LM (Geva et al., 2021 Meng al.) 2022p. Additionally, PrOntoQA provides reasoning formulation over false ontologies (see example in Appendix A). False ontology grounds the reasoning over statements that are false in the real world. An important demarcation between fictional and false ontological reasoning is that while the former minimizes the effects of factual associations memorized as pretraining prior, the latter requires the LM to actively eclipse such memorized knowledge to solve the problem successfully.
Circuits, in mechanistic interpretability rescarch, provide the abstractions of interpretable algorithms im- plemented by the model within itself. Typically, circuit is a subgraph of the complete computational graph of the model, responsible for specific set of tasks. We primarily follow the notation adopted by Wang et al.
4

---

Published in Transactions on Machine Learning Research (07/2024)
2023), with nodes defined by model components like attention heads and projections and edges defined by interactions between such components in terms of attention, residual streams, etc.
Activation patching is a common method in interpretability research. Activation patching begins with two forward passes of the model, one with the actual input and another with a selectively corrupted one. The choice of input corruption depends on the task and the type of interpretation required. For example, consider the following Indirect Object Identification (IOI) task (Wang et al. 2023): given an input John and Mary went to the park. John passed the bottle to, the model should predict Mary. Further, corrupting the imput by replacing Mary with Anne would result in the output changing to Anne. Let x, represent the original and corrupted residual streams at decoder block j, depending on whether Mary or Anne was injected at the input. Now, in a corrupted forward pass, for given j. if the replacement of ar by 2 Malr results in the restoration of the output token from Anne to Mary, then one can conclude that attention mechanism at decoder block j is responsible for moving the name information. This is an example of patching porrupted-to-clenn activation, often called cousal tracing (Meng et al., 2022). Activation patching, in general, refers to both clean-to-come corcorupted as well as corrupted-to-clentt patching (Zhang & Nanda, 2023
Mary and a
Knockout is a method to prune nodes in the full computational graph of the model to identify tpsk-specific circuits. Complete ablation of a node is equivalent to replacing the node output with an all-zero vector. Our initial experiments suggest that such an ablation destructively interferes with the model's computation. Instead, we follow Wang et al. (2023) for mean-ablation to perform knockouts. Specifically, we construct inputs from the false ontologies provided in the PrOntoQA dataset and compute the mean activations for each layer across different inputs:
b Knock = Mean; (4a/ ls; E S E )
(2)
where Dratse denotes the false outology dataset. Then, the language model function with head hia knocked out for the residual stream corresponding to the I-th token, can be represented as,
Knock
arg max LM lngie
j,k (s v¹, = 22 Kukckck)
(3)
More often than not in this work, we will need to knock out set of heads H;; we will denote the corresponding language model as LM Also, if we perform knockout the last residual stream (which, when projected to the token space, gives us the output token at the current step of gencration), we will drop the superscript I.
4 Task composition
We seek to discover the circuits responsible in Llama-2 7B for few-shot CoT reasoning on the examples from PrOntoQA fictional ontology problems. Consider the example presented in Figure [ we ask the model Tumpuses are bright. Lempuses are tumpuses. Max is a lempus. True or False: Max is bright" (for brevity, we omit the few-shot examples that precede and the CoT prompt Let's think step by step that follows the question). Llama-2 7B generates a verbose CoT reasoning sequence as follows: "Max is a lempus. Lempuses are tumpuses. Mar is tumpus. Tumpuses are bright. Mar is bright." " We can observe that such a reasoning process constitutes three critical kinds of subtasks:
The model decides on the path of reasoning to follow. In Figure D given the three possible entities to start with, namely Tumpus, Lempus, and Mar, Llama-2 7B starts with the last one (i..,, Mar is lempus). One would require multiple such decision decision-making steps within the complete reasoning process.
Decision-making:
2. Copying: The LM needs to copy key information given in the input to the output. Typically, decision- making precedes copying, as the model needs to decide which information to copy.
Induction: The LM uses set of statements to infer new relations. Again, a decision-making step precedes as the model must decide on the relation to infer. In the majority of this work, we focus on 2-hop inductions where given statements of the form is B and B is C, the model should infer A is C.
5

---

Published in Transactions on Machine Learning Research (07/2024)
Tumpuses are bright. Lempuses are tumpuses,
Max
lempus,
True False: Max is bright.
Subtask 0
Max
Subtask 1
is lempus
Subtask 2
Lempuses
Subtask 3
+ are tumpuses
Subtask 4
Max
is tumpus
Subtask 5
Subtask 6
Tumpuses
Subtask 7
are bright
Subtask 8
Max
Decision
Copy
Induction
is bright
Subtask
An working example of task decomposition of CoT generation on fictional ontology. Decision-type subtasks (0, 2, 4, 6, and 8) choose which reasoning path to follow; such decisions may be about the order of the information to be copied from the question to the generated output or whether inductive reasoning will be performed. Copy-type subtasks (1, 3, and 7) follow Decision subtasks and copy statements from the context to output. Induction-type subtasks (5 and 9) perform the reasoning sleps of the form "if A is B and B C, then A is C?.
Figure 1:
In all our analyses henceforth, we follow a structure where the model needs to generate CoT responses that solve the overall task using ten steps subtasks, each belonging in Figure [ (See Appendix D for details of few-shot examples used along with the overall performance of Llama-2 7B.) A natural hypothesis would be that distinct, well-defined circuits corresponding to each category of subtasks (i.e., decision decision-making, copying, and inductive reasoning) exist that work in conjunction to lead the CoT generation. However, as our findings in the next section suggest, this is not the case in reality.
of the three categories, as presented
4.1 Task Task- specific head identification
As Wang et al. 2023) suggested, the very first step toward circuit discovery is to identify the components. Since we presume the attention heads as the nodes circuit, the goal then becomes to identify those heads in the language model that are most important for a given task. We define the importance of the k-th head at the j-th decoder block, hy,e for particular task as follows.
Let the task be defined predicting token 8,41 given the input S = {si, (2,1,82.-1. $2. s.). For example, in the demonstration provided in Figure [ predicting the token Mar given the context Tumpuses bright. Lempuses are tumpuses. Max is lempus. True or False: Mar is bright. Response: Let us think step by slep. Tumpuses are bright. Lempuses are lumpuses. We assign a score (1) to each head hj.a proportional to their importance in performing the given task. Following the intuitive subtask demarcation presented carlier, we start with scoring the attention heads for cach different subtask. We provide the detailed procedure of calculating HTaak k for each category of subtasks decision-making, copying, and induction, in Appendix C

---

Published in Transactions on Machine Learning Research (07/2024)
(8) Decision heads
(b) Copy heads
(c) Inductive reasoning heads
(d) Task-only heads
Fraction of of involved (ted)
Accuracy (blue))
0.8
0.8
0.15
0.8
0.6
0.6
0.4
0.4
0.4
0.05
02
0.2
0.2
0
2
3
4 Sub-Task Task
6
7
8
0
1
5 Sub-Tese Task
4
6
8
9
0
0
3
5 Sub- Tackk
4
7
8
9
0
3
5
6
8
Sub-Tack Task
Figure 2: Performance of attention heads identified for each different subtask type across dif- ferent subtasks. We show the performance of (a) decision-making. (b) copy, and (c) inductive reasoning heads for each subtask O to 9 (blue bars show accuracy when the rest of the heads are knocked out; red bars denote the fraction of heads involved). Bold green indices denote the subtasks belonging to the task category used for attention head removal. (d) Task-only heads are only those that are not shared with other tasks, e.g., only those copying heads for subtask 4 that are not decision-making or inductive reasoning heads. Inductive reasoning heads are consistently functional across all the subtasks with the least number of heads involved.
Figure 2 demonstrates how each of these three categories of attention heads performs on its respective subtasks and the rest of the subtasks (note that we perform head ablation here for each subtask index independently). We also show the fraction of the total number of heads involved in performing each subtask. A higher accuracy with lower number of head involvement would suggest that we have found the minimal set of heads constitute the responsible circuit. As we can observe, there are no obvious mappings between the tasks used for head identification and the sublasks that constitule the COT. Quile surprisingly, heads that we deem responsible for inductive reasoning can perform well across all the subtasks (Figure 2 (c)) and requires the least number of active heads across all tasks. Furthermore, the three sets of heads share a significant number of heads that are essential for the subtasks. In Figure [d), for each subtask, we use the heads that are not shared by subtasks of other categories. For example, in subtask 2, since we assumed it to be decision-makes making subtask, we took only those decision-making heads that were not present in the copying inductive rcasoning heads. It is evident that these tasks are not structurally well differentiated in the language model. A good majority of heads share the importance of all three subtasks. The existence of backup circuits justifies such phenomena; however, with large models, the "self-repairing" tendency (McGrath et al.) 2023) is much higher. We hypothesize that there exists even further granularity of functional components that collectively express all three types of surfaco-level subtasks. We seek a solution in the earlier claims that pattern matching via induction heads serves the progenitor of in-context learning (Olsson et al. 2022 We observe that the inductive reasoning subtasks of the form if [A] is [B] and [B] is [C] then [A] is [C] essentially requires circuits that can perform 1) representation mixing from [A] to [B] to [C], and 2) pattern matching over past context by deciding which token to copy and then copying them to the
output residual stream. Thorefore, heads responsible for inductive reasoning can perform decision-making and copying by this argument.
For empirical validation of this argument, we analyze the effects of knocking out heads on the accuracy of each of the subtask indices. For each subtask, construct histogram of head counts over ((h),,)) It can be seen that the corresponding density distribution is Gaussian in nature. Furthermore, actual effect of pruning a head on the accuracy is not dependent on the actual value of μ(h),,) but varies with 6(hg.x) = (u(hz,x) - Meanja (p(hz,)))² From the definition of p(hy,)) provided in Appendix C we can explain the behavior as follows: a very low value of p((y,,)) is possible when the KL-divergence between the noisy logits and original logits is almost same as the KL--diveggence between patched and original logits. Conversely, u(h),,) can be very high when the KL-divergence between the patched logit and the original logit. is very small. Both of these suggest a strong contribution of the particular head.
observe that the
Figure shows the accuracy of the model at each different subtask indices with heads pruned at different 6(hg,) (we show the histogram over p(hg,)) and start pruning heads from the central bin corresponding to the smallest values of 0(hz,))) Among the 1, 0224 heads of Llama-2 7B, we can see that pruning the model
7

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Subtask
Subtask
Subtask
Subtask 3
300
300
Elimirate)
294
250
200
200
200
192
150
150
103
1.0
100
50
94
0.300
0.305
0.310
5.315
0.370
0.000 2.300 0.305 0.310 0.50 6.325
0,200 5,200 5.300 5,305 0.310.33333 0.320 2.325
M(Aya)
((A3A)
Subtask 4
Subtask 5
Subtask 6
100
300
300
Eliminated
250
230
250
200-
200
150
0.8
102
0.6
0.790 0.705 c. 0.22a 8:125
2.:00.295.0 000 365-0 3160 c.320 0.325
0,310
0.315
0.320
Accuracy
p(),2)
Subtask 7
Subtask
Subtask
200
300
Bimiinaed
250
250
350
200
0.2
130
152-
130-
100
0.0
p.10 0305 2(1ya)
na10 adis 0:355 m(A.)))
p(h.a)
Figure 3: Importance of inductive reasoning heads across subtasks. We plot the counts of attention heads binned over relative head importance (m(hz,x)) computed on inductive reasoning task across all subtask indices (0-9, left to right and top to bottom). For a given subtask and a given bin of attention heads (corresponding a range of p(h5,)) knocked out, we show the accuracy of the pruned model relative the original model (color-coded in blue). A wider spread of dark-colored bins for subtask signifies that a higher number of attention heads can be knocked out without hurting the model performance; for example, in subtask 7, the removal of attention heads with p(hj,)) [0.295, 0.325] retains similar performance compared to heads removed in subtask 5 with m(h;a) E [0.305, 0.315), signifying that more attention heads are required to remain active in subtask 5 compared to 7.
as low as 400 heads retains >90% of the original accuracy (see Table Appendix for subtask-wise statistics).
Further patterns in allocating attention heads for different subtasks be observed in Figure 3 Typically, subtask indices 0 and 5 are most sensitive to pruning attention heads; the accuracy of the model quickly decreases as we knock out attention heads with ((h;,.) deviated from the mean value. Subtask-0 is a decision- making subtask where the model decides on the reasoning path. One can further demarcate between subtasks vs 2, 6, and 8, while we initially assumed all of these four to be decision- decision-making stages. Referring to the example provided in Figure [] in subtask-0, the model has very little information about which ontological statement to choose next, compared to subtask indices 2, and where the model needs search for the fact. related to the second entity in the last generated fact. Subtask 5 is the actual inductive reasoning step.
8

---

Published in Transactions on Machine Learning Research (07/2024)
Compared to subtask 9, which is also an inductive reasoning step, subtask 5 does not have any prior (the answer generated in subtask 9 is framed as a question in the input).
5 Token mixing over fictional ontology
In the earlier section, we observed that the set of heads that we identify as responsible for inductive reasoning performs consistently across all subtasks with the least number of heads remaining active, retaining more than 90% of the original performance with as low as 40% of the total number of heads. We continue to investigate the inductive reasoning functionality with this reduced set of heads. Recall that we first need to explore if there is indoed a mixing of information happening across tokens of the form [A] is [B]. Specifically, we pose the following question given the three pairs of tokens ( ([A,], [B,]), ( ([A2], [B2]) and
([Aa], [Ba]), such that there exists positive relation [A,] is [B,], negative relation [A2] is not [Bg],
and no relations exist between [Ag] and [Bg], is it possible to distinguish between [ar) x "1 [an =xf"))
and (> =of"",, where |. denotes the concatenation operator, for a given decoder layer j? Given that these tokens denote entitics from a fictional ontology, the model could not possibly have memorized similar representations with the embedding/ MLP blocks for the tokens. Instead, it needs deploy the attention heads to move information from the residual stream of one token.
We translate the posed question into lcarning probing classifier (Belinkov 2022 C: X + V. where for each |x," B.] E X, the corresponding label y = [-1,0, should be -1, -1 0, or +1, if A and BN are negatively related, unrelated, or positively related, respectively. There can be multiple occurrences of a token in the input context (e.g., there are two occurrences of [B] in the above example). As a result, there will be multiple residual streams emanating from the same token. To decide which residual streams to pair with (in positive/negatien
pairing), we rely on immediate occurrence. For example, given the input [A] is [B]. [A] is not [c], we take the first occurrence of [A] to be positively related to [B], while the second occurrence of [A] as negalively related to [C]. We experiment with linear and a non-linear implementation of C as feedforward networks. Specifically, in the linear setting, we use a single linear transformation of the form y = (Wx x + B) followed by a softmax; in the nonlinear setup, we utilize stacked linear transformations with ReLU as intermodiate non-linearity. We provide the full implementation details of this experiment in Appendix D Given that the baseline distribution of the residual stream representation might vary across layers, we employ layer-specific classifiers.
Figure 4 shows the w-w classification performance in terms of accuracy across different layers using 4-layer ReLU networks. The following conclusions can be drawn from the observations:
Non-linear information movement between residual streams. Our experiments with the linear classifier fail to distinguish between ontologically related, unrelated, and negatively related entities.
Distinguishability of the residual stream pairs improves across depth. While the very first layer of attention provides substantial degree of required token mixing, we observe that the classification performance goes better gradually, pointing towards the existence of multiple successive attention heads that continue moving information from Aj to By if they are related. However, after a certain depth, the distinguishability starts diminishing, most likely due to accumulation of other information related to the tasks.
Token mixing is not boosted from contextual prior. Across different numbers of in-context exam- ples provided, we do not observe any significant difference in peak classification performance compared to zero-shot regimes. However, as the number of examples increases, we observe unstable classification performance. It likely that with more in-context examples, task-specific information is accumulated carlier than the zero-shot regime.
6 Circuitry of step-by-step generation
In Section 4.1
we observed that the head importance identified via inductive reasoning task could serve as a proxy to identify the important heads across all the subtasks. This provides us with an opportunity to reduce
9

---

Published in Transactions on Machine Learning Research (07 /024)
(a) Zero-shot
(b) 1-shot)
(c) 2-shot
0.70-
0.700
0.68
0.675
0.66
Accuracy
0.66
Accuracy
0.650
0.625
0.62
0.62
0.600
0.60
0.60
0.575
0.58
0.58
0.550
5
10
20
25
30
Decoder block
Decoder bloock
Decoder block
(d) 3-shot
(e) 4-shot
0.700
0.700
0.675
0.675
Accuracy
Acturacy
0.650
0.650
0.625
0.625
0.600
0.600
0.575
0.575
0.550
0.550
0.525
20
25
30
Decoder
Decoder block
Figure 4: How does the LLM mix information among tokens according to ontology? We plot the performance (in terms of accuracy) of classifying whether two tokens ontologically related, unrelated, negatively related using their residual stream representations at different layers. We demonstrate the depth- wise accuracy profile using different numbers of few-shot examples (0 to 4). Typically, mixing information between tokens according to their relation does not require in-context learning. Information mixing between related (or negatively related) token pairs results in better distinguishability: gradually increasing from the starting decoder blocks and achieving a peak between decoder blocks 10-15.
the total number of attention heads to analyze when investigating the step-by-step generation procedure. For each subtask, we use a threshold range for Ph.k to select a subset of attention heads that are kept intact while the rest of the heads are knocked out (sce Appendix C for the subtask-wise threshold range used and the corresponding subtask accuracy). This chosen subset of the model gives an aggregate accuracy of 0.9 over the inputs for which the full model generates correctly.
Next, we proceed to look into the exact information that is being read and written by the attention heads from and into the residual streams. Given head hy.6 that writes v to the residual stream xÿ-1, we apply the unembedding projection U on vr,, and select the token with the highest probability, gi.k The uncmbedding projection provides the opportunity the residual stream and its subspaces (nostalgebraist, (020)) However, should be noted that applying unembedding projection typically corrosponds to Bigram modcling (Elhage et al. 2021). Therefore, when map any intermediate representation (attention output, residual stream, etc.), essentially retrieve the token that is most likely to follow.
look into the token-space representation directly associated with
6.1 In-context prior pretraining prior
We start by exploring the depth which the model starts following the context provided as input. Specifi- cally, we check for a given token s; in the sequence S, if aiik the token projected by the attention head hy is such that (8; gidy is bigram present. in S. We compute a contert-abidance score, Gj.k for the head h5 as the fraction of tokens for which the above condition holds true.
In Figure we plot the context abidance distribution for each head across different subtask indices. We observe visible correlation between the depth of the head and how much context abiding is
10

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Subtask 0
Subtask 1
Subtask
Subtask 3
10
10
10
Layer
Layer
15
15
15
15
20
20
25
25
25
25
30
30
30
Head Subtask 4
Head Subtask 5
Head Subtask 6
Head
1
10
10
10
Layer
115
15
20
20
20
Context abidance score
25
25
25
30
30
30
Head Subtask7
Head Subtask
Head Subtask 9
0.5
10
10
Layer
Laver
155
15
15
20
20
20
25
25
25
30
30
20
20
Head
Head
Head
Figure 5: When does the LLM start following the context? Distribution of context-abidance score Gj. of cach head (darker blue shade signifies higher (j&) with layer (j) on y-axis and head index (k) on axis. Atlention heads in red denote those with zero contexl-abidance. We show the distributions for all the subtask indices 0-9 (left to right, top to bottom). Typically, context abidance is task-agnostic and emerges after the 16th decoder block.
the LM starts focusing on the conteatual information at deeper layers. Given that our experiment setup is based on fictional ontology, Cj,k provides strong demarcation between pretraining prior (i.e., language model statistics memorized from pretraining) and contextual prior (i.e., language model statistics inferred from context) since there are negligible chances of the LM to memorize any language modeling statistics containing the fictional entities from the pretraining data. Therefore, it predicts correct bigram only when it is able to follow the context properly. However, there is no visible demarcation across different subtasks; we can claim that this depth-dependence implicit to the model and does not depend on the prediction task.
6.2 Subtask wise answer generation
We investigate the information propagation pathway through the attention heads that constitute the step- by-step answering of the subtasks. We start with the heads that directly write the answer to the particular subtask into the last residual stream by mapping their output to the token space using U. Figure(pplots these answer-writing heads along with the probability of the answer in the attention head output across different subtasks. The existence of multiple answer ansver-writing heads suggests that LMs exploit multiple pathways generate the same answer, with each such pathway reinforcing the generated output. Looking at the heads
11

---

Published in Transactions on Machine Learning Research (07/2024)
Sublask 0
Sublask
Subtask 2
Subtask 3
0.6
0.4
10
0.2
0.2
0.2
0.2
Layer
0
15
15
0
15
20
20
20
20
25
10
30
20
30
20
Head
Head
Subtask 4
Subtask
Subtask 6
Subtask 7
0.4
0.4
0.3
0.3
0.2
0.2
0.2
10
10
10
0.1
10
Layer
15
1.5
15
15
20
25
25
25
30
30
30
30
10
20
10
30
10
20
30
30
20
30
Head
Head
Head
Head
Subtask
Subtask
0.6
0.4
0.2
10
0.1
Layer
15
20
20
25
30
10
20
30
Head
Figure 6: Which heads are writing the answer? For each subtask (0-9), we show the attention heads that write the answer to the corresponding task into the last residual stream, along with the probability of the answer token in the attention head output.
with the highest answer probability, we can observe that some top heads are utilized across different subtasks, e.g., subtasks and 3, subtasks 2, 4, and 8, subtasks 5 and 9, etc.
Sharp change in wisth-wise functionality. Interestingly, the 16-th layer appears region of functional transition in the model with all the answer-writing heads appearing after this particular layer. Our earlier experiments also associate this layer with some phase shifts; in the token-mixing analysis (Section 5], the distinguishability of ontologically related entities is observed to achieve peaks near this layer and starts falling after that; in Section [.1 as well, we see certain overall risc in context-abidance in the heads after the 16-th layer (see Figure 5)- All three findings suggest. a funclional rift in LLMs that happens to lie almost at the halfway point from embedding to unembedding, The embedding layer associates tokens with information available from pretraining. The initial half assists information movement between residual streams and aligns the representations to the contextual prior. The latter half of the model employs multiple pathways to write the answer to the last residual stream.
6.3 Parallel pathways of answer processing
Once we observe the existence of multiple answer-writers within the model, the natural step forward is to wonder if they all process the answer from the input using the same mechanism. We employ a recursive strategy to identify the flow information through the attention heads (see Appendix G for a detailed description of the procodure). Specifically, we start from the answer-writing heads, follow which residual streams are being attended by these heads, identify the content of these residual streams via unembedding projection, and identify the heads in the previous layers that are writing that content into those residual
12

---

Published in Transactions on Machine Learning Research (07 (07/202) /024)
Subtask 0
Subtask
Subtask 2
Subtask
0
5
5
10
10
10
10
Layere
Layers
15
15
15
15
20
20
20
20
25
25
25
25
30
30
5
10
15
20
25
30
5 10
15
20
25
30
5
10
15
20
25
30
5 10
15
20
25
30
Heads
Heads
Heads
Heads
Subtask 4
Subtask 5
Subtask 6
Subtask 3
0
0
5
5
10
10
Layers
15
Jayes
15
Layers
15
15
20
20
20
20
25
25
25
30
30
0
10
15
20
25
30
10
15 Heads
20
25
30
5
10
15
20
25 30
5
10
15 Heads
20
25
30
Heads
Heads
Subtask
Subtask 9
Generated context
5
Question context
Layers
aver
15
Few-shot context
20
20
0
10
15 20 Heads
25
30
0
5
10
15 20 Heads
25
30
Figure 7: Where do the answer-writing heads collect their answers from? For each answer-writing head and each subtask, we show those attention heads that attend to residual streams corresponding to answer tokens. We demarcate heads that collect the answer tokens from the generated context (green), question context (blue), and few-shot context (red).
streams. We continue till one of the two conditions met: (i) we reach head in the first decoder block, (ii) we reach residual stream corresponding the first token in the input token sequence. With such procedure, construct trees of attention heads rooted at the answer writing heads.
Now, let us consider the following example context plus generated token sequence: Lempuses are tumpuses. Tumpuses are rompuses. Rompuses are blue. Mar is lempus. True or false: Mar is blue. Response: Let us think step by step. Max is lempus. Lempuses are tumpuses. Max is with the desired prediction being tumpus (note that this input follows the few shot examples that are not shown here for brevity). We segregate the input context into three parts: few-shot context, question context, and generated context. We can see that the LLM can collect the answer tumpus from either generated or question context. Also, if the same token exists in the few-shot context, it can decide to collect information from there well (however, that should not be ideal since the contextual rolc of that token would be different). We proceed to identify from the information-flow trees all those heads that altend to the streams corresponding to the answer Loken (in this example, tumpus). Note that there can be multiple other heads that might be attending to these same tokens; but only those heads that are elements of the tree rooted at answer writing heads are actually contributing the flow of information from the answer tokens in the context of the output Figure demonstrates these attention heads that attend to the residual stramms corresponding to the answer tokens for different. subtasks,
13

---

Published in Transactions on Machine Learning Research (07/2024)
present in the generated context (green), question context (blue) and few-shot context (red). The following observations can be drawn from the ligures:
1. Coexitent pathways of answer generation. Smaller models like GPT-2 on simpler tasks like 101 implements a unique neural algorithm (Wang et al., 2023). On the contrary, a larger model like Llama-2 exploits different parallel pathways answer propagation while performing reasoning There are different heads that are directly connected the answer writers, and they collect answer information from different places in the context.
2. Different primary sources of answer for different subtasks. In the initial stages of generation (i.c., subtasks 0, 1, 2, 3), the answer tokens are not present in the generated context. For subtasks 6 and 7, again, the answers need to be collected from the question index. For subtasks 4, 5, 8, and 9, we can observe the presence of heads that collect answer tokens from the context generated via earlier subtasks (in fact, subtask 5 has more such heads).
Given the fact that (i) there are multiple answer writing heads, and (ii) there are different heads that are directly connected to the answer writers, and they collect answer information from different places in the context, we can conclude that the different pathways implement different algorithms as well. In the few- shot examples, there are fictional entities that have boen used in the question context as well, though in different contextual roles. The fact that there are attention heads collecting those entities as answers suggests that there are pathways prone to collecting similar information from the few-shot context. Although these pathways collect the same token as the answer, they actually deviate from the correct reasoning algorithm. However, the presence of such pathways decreases the generation progresses from subtask 0 to 9. Note that this is different from few-shot examples providing the necessary pattern via in-context learning. Such patterns are more position-specific, i.c., while gencrating the answer for a certain subtask in the question, there is an overall increased atlention provided to the same subtask tokens in the few-shot examples (see Figures 9 11, and 12 in Appendix). Finally, the pattern of answer collection from question context and generated context across different subtasks clearly demarcates the inductive reasoning subtasks (5 and 9) from the decision-making and copying. Tasks in the former category require pathways that collect answers from the earlier generated context, but the latter does not. Moreover, the decision-making step right before inductive reasoning
extract answers from the preceding step. This explains an albeit small number of heads that collect the answer from the generated context in subtasks 4 and 8. See Figures examples of such information flow towards subtasks and 5, respectively.
13 and for
7 Conclusion
Our findings. This work sought to disentangle the functional fabric of CoT reasoning in LLMs. Specifically, we explored Llama-2 7B a few-shot CoT regime for solving multi-step reasoning problems on fictional ontologies of the PrOntoQA dataset. We observed that:
Despite different reasoning requirements across different stages of CoT generation, the functional compo- nents of the model remain almost the same. Different, neural algorithms are implemented compositions of induction circuit-like mechanisms.
Attention heads perform information movement between ontologically related (or negatively related) tokens. This information movement results in distinctly identifiable representations for such token pairs. Typically, this distinctive information movement starts from the very first layer and continues till the middle. While this phenomenon happens zero-shot, in-context examples exert pressure to quickly mix other task-specific information among tokens.
Multiple different neural pathways are deployed to compute the answer, that too in parallel. Different attention heads, albeit with different probabilistic certainty, write the answer token (for each CoT subtask) to the last residual stream.
14

---

Published in Transactions on Machine Learning Research (07/2024)
4. These parallel answer generation pathways collect answers from different segments of the imput. We found that while generating CoT, the model gathers answer tokens from the generated context, the question context, as well as the few-shot context. This provides a strong empirical answer to the open problem of whether LLMs actually use the context generated via CoT while answering questions.
5. We observe a functional rift at the very middle of the LLM (16th decoder block in case of Llama-2 7B), which marks a phase shift in the content of residual streams and the functionality of the attention heads. Prior to this rift, the model primarily assigns bigram associations memorized via pretraining; it drastically starts following the in-context prior to and after the rift. It is likely that this is directly related to the token-mixing along ontological relatedness that happens only prior to the rift. Similarly, answer answer-writing heads appear only after the rift. Attention heads that (wrongly) collect the answer token from the few-shot examples are also bounded by the prior half of the model.
The design of this study imposes certain limitations in its scope. First and foremost, we analyze a very specific type of reasoning problem. While the fictional ontology with a restricted CoT template provides ease of analysis, free-form reasoning can introduce further complex dynamics not captured in this study. We also could not. address the role of MLPs in reasoning. While existing literature points to their role as factual memory, it should be noted that the model also memorizes a diverse set of token-token associations pertaining to the structure language and not just factual associations. Such associations are likely to play a decisive role in a grounded reasoning setup that our analysis ignores. Finally, we use a few-shot prompting regime to ensure that the model follows specific structure of rcasoning steps. With zero-shot CoT, more complex mechanisms are bound to emerge.
Limitations.
Implications for future research. These findings bear important ramifications towards the ongoing research around language modcling and interprotability. A natural extension to this work would be to incorporate pretraining memorization in terms of MLP blocks - precisely, whether the functional rift across layers bears a similar mechanism when the LLM starts mixing factual associations that have been stored within the MLP neurons. The existence of parallel auswer-generation process is extremely important for causal interventions model behavior (Li ct al. 2024): changing how model should reason via up- down- scaling certain neural pathways should take all the parallel pathways into account. Although real-world reasoning problems are far more dynamic than our strictly controlled sctup of ontology-based reasoning, many of the results hold in the case of CoT in the wild. The existence of multiple neural pathways of answer generation is intrinsic to the model (although in weal-world reasoning, such pathways will contain MLP blocks as a source of the answer, instcad of the context only). The token-mixing mechanism that we uncovered remains true for any ontological relationship, not just fictional. Additionally, the more out-of-distribution (with respect to pretraining data) the reasoning problem becomes, the stronger the rift in terms of context abidance we will observe.
Acknowledgments
The authors would like to acknowledge the financial support of DYSL-AL
References
Yonatan Belinkov. Probing Classifiers: Promises, Shortcomings, and Advances. Computational Linguistics, 48(1):207-219 04 2022. ISSN 0891-2017. doi: 10.1162/coli 00422. URL https://aoo.orr//00
coli a 0_00222.
Zhen Bi, Ningyu Zhang, Yinuo Jiang, Shumin Deng, Guozhon Zheng, and Huajun Chen. When do program- of-thought works for reasoning?, 2024.
Wenhu Chen, Xucguang Ma, Xinyi Wang, and William W. Cohen. Program of thoughts prompting: Dis- entangling computation from reasoning for numerical reasoning tasks. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL httppp///gaaaaeeeee.eeer/rrrrrrs
15

---

Published in Transactions on Machine Learning Research (07/2024)
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuniao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Trans- former Circuits Thread, 1, 2021.
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hathiel-Doddss Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposition. arXiv preprinl ar Xiv:2209. 10652, 2022.
Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. arXiv preprint arXiv:2305. 15408, 2023.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memorics. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp.
Association for Computational Linguistics, 2021. doi: 10.18653/v1 /20./1...mmmmmmmiinn...
URL https://aclanthologg. org/2021.... emw.pemmmnppmai.n. 446
Tom Henighan, Shan Carter, Tristan Hume, Nelson Elhage, Robert Lasenby, Stanislav Fort, Nicholas Schiefer, and Christopher Olah. Superposition, memorization, and double descent. Transformer Circuils Thread, 2023.
Takeshi Kojima, Shixiang (Shane) Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large lan- guage models zero-shot reasoners. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 22199 22213. Cur- ran Associates, Inc., 2022. URL https://prooeeddnng: www.:eeddddnnn.nooois....... wwwouris..c/paaaapp httppppprrffllpppppprreeaee / (alpo///aaaaacfffff1110001111680 uttpatpss//pppoccdggg.oooooossso
pdf
Andrew Lampinen, Ishita. Dasgupta, Stephanie Chan, Kory Mathewson, Mh Tessler, Antonia Creswell, James McClelland, Jane Wang, and Felix Hill. Can language models learn from explanations in context? In Findings of the Association for Compulational Linguistics: EMNLP 2022, pp. 537-563. Association for Computational Linguistics, 2022. doi: 10.18653/v1/2000 http://////22..nnnnneeeeeeee indings-emmlp.33.
Maximilian Li, Xander Davies, and Max Nadeau, Circuit breaking: Removing model behaviors with targeted ablation, 2024.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts automata. arXie preprint arXiv:2210.10749 2022.
Thomas McGrath, Malthew Rahtz, Janos Kramar, Vladimir Mikulik, and Shane Legg. The hydra effect: Emergent self-repair in language model computations. arXie preprint orXiv:2307.15711, 2023.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual as- sociations in gpt. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 17359 17372. Curran Associates, Inc., 2022. URL http:::/pprooeedng...eerreeeeeee http://rraop.ppppppaoaaaa/a////p filees/paprr/00///ieel
Necl Nanda. A comprehensive mechanistic interprotability explainer & glossary, 2022.
nostalgebraist. interpreting GPT: the logit lens Loss Lross - col/oots/ERBBaaaaddddtrrneeeeeep
lesswrong, com. https://wwwz wttp://wwww.ssw.o......o
2020. [Accessed 09-02-2024].
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXim:2112.0011 2021.
Catherine Olsson, Nelson Elhage, Ncol Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXin:2209.11800 2022.
16

---

Published in Transactions on Machine Learning Research (07/ (07/2024)
OpenAL. Gpt-4 technical report, 2024.
Ben Prystawski, Michael Li, and Noah Goodman. Why think step by step? reasoning emerges from the locality of experience. Aduances in Newral Information Processing Systems, 36, 2024.
Abulhair Saparov and He He. Langnage models are greedy reasoners: A systematic formal analysis of chain-of-thought. In The Eleventh International Conference on Learning Representations, 2023. URL
https://openreviers seet//forun21==FFWB22222
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
Juanhe TJ Tan. Causal abstraction for chain-of-thought reasoning in arithmetic word problems. In Pro- ceedings of the 6th BlackborNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, pp. 155-168, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bash- lykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Fer- rer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Tgor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subra- manian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama Open foundation and fine-tuned chat models, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Newral Information Processing Systems, volume 30, 2017. URL -aapceaaaaaaaaaaaa..a.aa
wwwpeeeddiinnnn.nns..e... pttpeeeererpp..pppprrrffii filess/paper/2211///he
Kevin Ro Wang, Alexandre Variengien, Arthur Commy, Buck Shlegeris, and Jacob Steinhardt. Interpretabil- ity in the wild: a circuit for indirect object identificatinn in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023. URL http:::://ppprreveeeeeeee//uuuup httt:///oorruuuddddaassSN
Xinyi Wang and William Yang Wang. Reasoning ability emerges in large language models as aggregation of reasoning paths: A case study with knowledge graphs. In Workshop on Efficient Systems for Foundation @odels@ ICML2023, 2023.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models. Transactions on Ma- chine Learning Research, 2022a. ISSN 2835-8856. URL //tpppe//operreviex net peet/forumiidyyyasssszz Survey Certification.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed Chi, Quoe V Te, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 24824-24837. Curran As- sociates, Inc., 2022b. URL
nwwriiipnnnn.ni.s.oooe...rripiep nwwrinerriicoom newppp.ccccppppprrrr files//paper/222///1ll Confereneeorreee..dcc. Soreraa0000.7ttaaaaaaaoooororrer
Zhengxuan Wu, Atticus Geiger, Christopher Potts, and Noah D Goodman. Interpretability at scale: Iden- tifying causal mechanisms in alpaca. arXiv preprint arXiv:2305. 08809, 2023.
Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metries and methods. Xiu preprint ar Xri::29009 16040, 2023.
17

---

Rex is spicy. Published in Transactions on Machine Learning Research (07 (022)
zumpus
wumpus
shumpus
impus
brimpus
not red
hot
happy
not bitter
Jompus
lorpus
no small
Figure 8: An example of fictional ontology in the PrOntoQA dataset.
A PrOntoQA question example
Following two QA examples from fictional and false ontologies, respectively:
Fictional ontology example:
Contezt: Tumpus is bright. Lempus is tumpus. Max is lempus.
Query: True or false: Max is bright.
Answer: Max is lempus. Lempus tumpus. Max is tumpus. Tumpus bright. Max is bright. True
False onotology example:
Context: Carnivore is spicy, Vertebrates are carnivores. Rex is vertebrate.
Query: True or false: Rex is spicy.
Answer: Rex is vertebrate. Vertebrates are carnivores. Rex is carnivore. Carnivore is spicy. Rex is spicy, True
Figure demonstrates an example ontology tree along with distractors. Hore, zumpus + shumpus translates to zumpuses are shumpuses. From this example, one can construct the following multistep reasoning problem: Wumpuses are shumpuses. Shumpuses are impuses. Mar is a wumpus. Wumpuses are not red. Impuses are happy. Jompuses are lorpuses. True or False: Mar is happy. In this example, the statement Jompuses are lorpuses serves as distractor neither of the entities part of the ontology relevant to the question.
B Knockout
To perform knockout average activation from false ontology the PrOntoQA dataset is used. For each subtask, false ontology activation is stored and mean activation knockout is performed. We provide step-by-step examples of subtask-wise generation over false ontology as follows:
Context: Carnivore is spicy. Vertebrates are carnivores. Rex is vertebrate.
Query: True or false: Rex spicy.
Answer: 'Rex is vertebrate. Vertebrates are carnivores. Rex is carnivore. Carnivore is spicy. True',
'prompt 0': "Carnivore is spicy, Vertebrates are carnivores. Rex is vertebrate. True false: Rex is spicy, Let us think step by step.',
'prompt. 1:: 'Carnivore is spicy. Vertebrates are carnivores. Rex is vertebrate. True or false: Rex is spicy. Let us think step by step. Rex is", 'prompt__" 2": "Carnivore is spicy. Vertebrates are carnivores. Rex vertebrate. Truc or false:
Rex is spicy.
Let us think step by step. Rex is vertebrate.',
"prompt. 3": "Carnivore is spicy. Vertebrates are carnivores. Rex is vertebrate. True false: Rex is spicy.
18

---

Published in Transactions on Machine Learning Research (07 (274))
Let us think step by step. Rex is vertebrate. Vertebrates are", 'prompt. 4": "Carnivore is spicy. Vertebrates are carnivores. Rex is vertebrate. True or false: Rex is spicy. Let us think step by step. Rex is vertebrate. Vertebrates are carnivores.', 'prompt_55 _": "Carnivore is spicy. Vertebrates are carnivores. Rex vertebrate. True or false: Rex is spicy. Let us think step by step. Rex is vertebrate. Vertebrates are carnivores. Rex is', "prompt__0 _ "Carnivore is spicy. Vertebrates are carnivores. Rex vertebrate. True or false: Rex is spicy. Let us think step by step. Rex is vertebrate. Vertebrates carnivores. Rex is carnivore.', 'prompt_ 7": 'Carnivore is spicy, Vertebrates are carnivores. Rex is vertebrate. True false: Rex is spicy, Let us think step by step. Rex is vertebrate. Vertebrates are carnivores. Rex is carnivore. Carnivore is', 'prompt_S 8"; "Carnivore is spicy, Vertebrates are carnivores. Rex is vertebrate. True or false: Rex is spicy, Let us think step by step. Rex is vertebrate. Vertebrates are carnivores. Rex is carnivore. Carnivore is spicy.",
"prompt_9r" "Carnivore is spicy. Vertebrates are carnivores. Rex vertebrate. True or false: Rex is spicy. Let think step by step. Rex vertebrate. Vertebrates are carnivores. Rex carnivore. Carnivore is spicy. Rex is spicy.',
C Task specific head identification
C.1 Decision-making heads
To identify the heads that actively participate in decision- making subtasks, incorporate activation patch- ing on individual heads over decision-making subtasks.
Let DDecision denote the input loken sequence for a particular decision-making subtask with Sans being the token corresponding to the correct answer. Also, let Porg denote the output token probability distribution of the full model corresponding Decision as input, i.c.,
Porg SoftMax (LM(a (SDorision))
Here we omit the token position superscript over Dlogit for brevity since we are taking the output corre- sponding to the last input token. We store the activations Yj.k corresponding each head hj,k-
Next, corrupt the input corresponding the last token in Decision: which is equivalent to knocking off the set of all heads, Hial.. We record the corresponding output token probability as,
Peorrupt SoftMax (LM (Mloii Deccston))
Finally, for each head hi.s, restore its corresponding original output yju and record the output token probability distribution as
pi.k patched SoftMax (LM (bea)) (Flagit (Dcissoo))
Then, we compute the importance score of head hj,k as
Porg (x Sans) P Peorrupt (x: Sans)
PDecision (h2,k) Porg (x = Sans)
P Ppatched(2 (2: = Sans)
C.2 Copy heads
Copying subtasks follows decision-making subtasks immediately. Once the head entity of a statement is decided, copying requires moving the tail entity corresponding to that statement in the input context to the output (c.g., if there is a statement Rompus is grimpus the input context and decision-making heads have decided to output Rompus, then copy subtask requires moving grimpus to the output). For this, we simply look into the attention probability assigned by each head hy the token to be copied source (i.e., key),
19

---

Published in Transactions on Machine Learning Research (07 2024)
denoted
exp (w 2 ( & x5 Σ, exp (W ا T ( W mr))
PCCopy (1gaa)
where an ar3tm denote the residual streams corresponding to the last token and the answer token, respec- lively, al j-th decoder block.
C.3 Inductive reasoning heads
Let the input token sequence for the inductive reasoning subtask be denoted as Smd, which is of the form
[A] is [B] [B] is [c] [A] is, and let Sans denote the answer token (which is [c] in this example).
We reprosent the first and sccond occurrences of cach token [A] and [B] as A1, and B1, B2, respectively. For given Sinducionn can perform activation patching on three different residual streams, æj 2
B: B2 and
of to identify the responsible heads. Let the original token probability distribution be
Porg SoftMax (ZM(z)
/ Induccion))
Next, for each I E (Bi, B2, C),, we compute the corrupted forward pass as follows:
Peorrupt, = SoftMax (LM (Flogit Sinduction)
followed by patching original activation to the head hj.4 as plk patched,1 = SoftMax (LM M Mhraa (Zlogit SInduction)
Then, the inductive reasoning head importance is computed as:
Pladuction (11ga)) = Mean; D'est D
pattched
D
where D is the KL divergence between Porg and P Poorupt,. and D Paatched is the KL divergence between Porg and P.orupt,.
C.4 Performance of inductive reasoning heads for each subtask
For each subtask 0-9, group the attention heads based on their respective Ainduction (hj,x). We knockout heads within certain threshold head importance range (Mmin: Mmmx) by knockout and record the accuracy. In Table [ we report the attention head statistics that we used in the analysis for context abidance and answer writtng pathways. Note that we sought to keep the relative accuracy (i.e., the fraction of correct prediction by the ablated model over the correct predictions of the full model) close to 0.9.
Probing for token mixing
Towards probing ontological relatedness residual streams (sce Section 5 first, we need extract the positive, negative and unrelated entities from the sentence. PrOntoQA provides with tree structure for each sentence, and with the help of the structure, we can extract the pairs. Consider the following example.
Sentence: Each shumpus is zumpus. Shumpuses are wumpuses. Every shampus is hot. Impuses are shumpuses. Each impus is brimpus. Each impus is happy. Wumpuses are not red. Brimpuses are not biller. Lorpuses are jompuses. Each yumpus is not hot. Lorpuses are not small. Mar is an impus. Max is a torpus.
Here all immodiate entities will be termed posive pairs, such as shumpus <> impus, zumpus <> shumpus, impus<>> happy, etc. Similarly, following are few examples of negatively related entities: wumpus
20

---

Published in Transactions on Machine Learning Research (07/2024)
Subtask index
Accuracy
Heads removed
Threshold range
475
0.30487806 0.31463414
0.93
554
0.30243303,0.31 0.33777117
0.96
617
0.3 - 0.3195122
3
1
617
0.3 - 0.3195122
0.99
554
0.30243903 0.31707317
0.93
475
0.30487806 0.31463414
0.94
475
0.30487806 0.31463414
0.88
617
0.3 - 0.3195122
0.95
617
0.3 - 0.3195122
0.91
663
0.297561 0.3219512
Table 1: Statistics of subtask-wisc attention head removal according to their importance in inductive rea- soning.
red, brimpus <> bitter, etc. The distractor entities (jompus, lorpus and small) serves as seed for unrelated entities; any two entities that belong two disjoint ontologies are taken as unrelated.
During training and testing, we also make sure that there are no overlap between entities. During train- ing, we use entities such as "wumpus', 'yumpus', "zumpus"; "dumpus' and during testing 'zonkify', 'quiblitz', "flimjam,, "zizzleuump" 'snickerblat'. Following are the implementation details for the probing classifier:
Total training pairs: 28392; Total testing pairs: 9204. All three types of pairs (positively and negatively related and unrelated) are present in equal proportion in the training and testing data.
4-layer MLP model, 4096 * -> 128 -> 64 -> 32 -> 3. With ReLU in between each Linear layer.
Learning rate: 0.00005
Number of epochs: 120.
E
Prompts used and model performance
We use 6-shot examples of CoT for generation in all the experiments.
### Input:
Gorpus twimpus. Alex is rompus. Rompus is gorpus. Gorpus small. Rompus is mean. True or false: Alex is small. Let think step by step.
### Response:
Alex is rompus. Rompus is gorpus. Alex is gorpus. Gorpus is small. Alex is small. True
### Input:
Gorpuses are discordant. Max is zumpus. Zumpus is shampor. Zumpus is gorpus. Gorpus is earthy. True or false: Max is small. Let us think step by step.
### Response:
Max is zumpus. Zumpus is gorpus. Max is gorpus. Gorpuses are discordant. Max discordant. False
### Input:
Borpin are wumpus. Wumpuses are angry. Wumpus is jempor. Sally is lempus. Lempus is wumpus. True or false: Sally is floral. Let us think step by step.
### Response:
Sally is lempus. Lempus is wumpus. Sally wumpus. Wumpuses are angry. Sally is angry. False
# # Input:
Gorpus is jelgit. Yumpuses loud. Gorpus is yumpus. Yumpus is orange. Rex is gorpus. True false:
21

---

Published in Transactions on Machine Learning Research (07 (07/202) /024)
0.06
0.05
Probability
0.04
0.03
0.02
0.01
0
duu
are
happy 5888
.589
Ster 590
pus, 591
is 1,592
us. 594
1.555
Sam
is 1597
ster 5998
600
True 601
or 602
false, 603
1.604
Sam 605
is 60666
happy, 607
1.08
Let 60999
think 611
siep 612
by 613
Siep 614
615
616
Mai 617
# 618
Response 619
620
621
610
See
988
593
596
589
Token
Figure 9: Average attention probability distribution (across all heads) over question context in an example
subtask 0. Here, the model needs to predict the token Sam.
Rex is loud. Let think step by step.
### Response:
Rex is gorpus. Gorpus is yumpus. Rex is yumpus. Yumpuses are loud. Rex is loud. True
# ## Input:
Lempus is tumpus. Max lempus. Tumpus fruity. Tumpus bright. Lempus is dropant. True or false: Max is bright. Let think step by step.
# ### Response:
Max is lempus. Lempus is tumpus. Max is tumpus. Tumpus is bright. Max is bright. True
### Imput:
Sterpuses are dull. Impus medium. Impuses are sterpuses. Wren impus. Sterpus is daumpin. True or false: Wren is melodic. Let us think step by step.
### Response:
Wren impus. Impuses are sterpuses. Wren is sterpus. Sterpuses are dull. Wren dull. False
To ensure that the model does not learn any positional pattern from the few-shot example, we jumble the statements in the Input context and add two extra distractors in the statement.
Following are the performance of Llama-2-7B on two subdomains of PrOntoQA ontology: 52.5, (ii) accuracy on fictional ontology: 35.9.
accuracy on false
Accuracy calculated by a complete match of all the intermediate reasoning subtasks and the final True True/False statement. We make the following modification to the PrOntoQA dataset: (i) removal of articles like a and from statements, (ii) removal of every and each from statements, (iii) replacing negative entities with incorrect entities (e.g., Max is not cold - Max is dog).
False
F Attention probability
We take the average of attention probability at the last residual stream of all the heads. Figure shows the atlention probability over the input context lokens, and Figures TO 11 and I2 show attention probability over the 6-shot examples used. This example is for subtask-0, and can be seen that there are spikes in attention probability score for tokens the same position as subtask O, i.c., token just after "Response:". Similar patterns are observed for all the subtasks. Furthermore, the absolute value of attention the fask-specific tokens in the fhw-shot context increases we move from the 1st to 6th example.
22

---

Published in Transactions on Machine Learning Research (07/202) 2024)
180
Figure 10: Attention probability distribu- tion over few-shot amples. In this plot, we show the average atten- tion scores over tokens corresponding to the first (left) and second (right) examples (for the rest of the examples, see Fig- ure II and 12 among the 6-shot context cor- responding to the ques- tion in Figure g Here, the model is at subtask- O and needs to predict Sam. A slight peak in attention probability
181
Falso 180
.88
.13
small 808
art. 178
1s_79
ond 177
dix 176
_777
is. 175
small 76
Mia 1M4
173
15
art 177
pus. 74
ond 171
Gor 73
dix 173
.22
are 160
pus.
or_ 70
167
2.
Gor 166
165
164
Alex
163
.66
65
161
Mia
3.63
159
62
158
61
156
157
15,155
.99
154
58
25152
15_56
151
Alex 55
pus,155 158
54
is 148
149
Mia
the token corresponding
.445
Reagonse_ss
144
to subtask-0 in the sec- ond example (Mar in this case) is observable.
Ad.
Response 140
143
All
siep.
.41
by
1_18
step stop_45
138
think, 44
by_
us_43
dink, 135
_41
small 40
133
19
small
132
15_130
false, 36
128
fabe.
725
_33
mean_32
earth_122
pus_ 30
15_121
Rom 29
.28
Gor_115 119
118
26
pus, 25
or_116
114
313
pus, 22
21
X_111
2_20
is_ 19
1.9
pus_ 18
Rom_ 17
pus.
rom, 14
1
13
.11
Us
19
TW
6
and
dax
G
G.50
Input
#dd 0
0.1
0.08
0.06
0.02
0.1
0.08
0.06
0.04
0.02
Probability
Probability
23

---

Published in Transactions on Machine Learning Research (07 (07/2022 2024)
Figure 11: Attention probability distribu- tion over few-shot ex- amples (continued af- ter Figure [O]. In this plot, we show the average altention scores over tokens correspond- ing to the third (left) and fourth (right) exam- ples among the few-shot context corresponding to the question in Figure 9
F81
engyy_279
engyy
.315
US
art
uses
ump
7,365
364
252
US.
US
ump
Y_3E1
5_360
Here, the model at subtask-O and noods predict Sam. Small
US,
emi
.35
2
u5
_254
umc
05, 253
y_354
NO 252
1_353
peaks in attention prob-
Ie_251
pus 352
Ger_351
ability at tokens corre-
aly
335
pus 349
sponding to subtask-0 in the third (S from Sally)
247
or 14338
_244
8.347
Resconse_245
346
345
and fourth (R from Rea)
FE
.222
343
example are obscrvable. For the rest of the exam-
241
342
Response
step.
543
ples, Figure 12
by
#6 339
stpp 233
338
think 236
337
uS
stp 336
Let 734
by 335
235
Stg, 394
o'al 252
think 338
231
US. 332
230
Lec 331
ally 229
.30
3,228
leud. 329
227
K_388
libe 726
317
or 225
R_326
True, 224
325
223
Take 324
22_222
or. _23
ung. 221
Trux, 322
W_220
221
21219
pus 120
U1_218
0r_315
217
8_316
L216
K,317
.215
316
21,214
31515
mg_213 213
.114
orange 313
211
312
ally 210
211
209
ump. 310
228
1.309
empor. 207
316
206
337
205
ump 305
US 284
3.5
203
IS 304
W 202
PaS 308
231
Gor 302
200
3G1
ore. 199
lood 300
193
299
197
Uses 296
W. 196
Uno
T55
Y_256
US, 194
285
193
294
192
2L289
a9e_191
$_222
PRL
Or. 189
290
1.888
187
288
186
287
InpuL_185
286
1.484
2.225
A8.183
284
0.1
0.08
0.06
0.04
0.02
0.1
0.08
0.06
0.04
0.02
Probability
Probability
24

---

Published in Transactions on Machine Learning Research (07 (07/202) 2024)
.573
Figure 12: Attention probability distribu- tion over few-shot ex-
572
8.571
-500
amples (cont after Figure 10 and LT]. In this plot, we show the average altention scores over tokens correspond- uL559
568
Fase 567
.72
471
True 470
d.364
469
bright 468
ren
1s.
W
Max 465
550
A55
bright 464
d.158
ing to the fifth (left) and sixth (right) exam- ples among the few-shot context corresponding to the question in Figure g Here, the model at subtask-O and noods predict Sam. Small peaks in attention prob- ability at tokens corre- sponding to subtask-0 in the fifth (Mar) and sixth fen,
Is, .63
are 557
US, 462
uses 516
LTO. 461
p_555
465
Ster
459
-53
UR, 455
pus, 552
LITO
sier
T.465
ku550
Is
ren, 549
Max 454
W_548
AG1
-_477
UR, 452
uses 546
Ump. 451
p_545
1.450
ster
is
are
un, 448
uses 542
smp
AAE
.50
445
uz 444
Ino.
mp.
It.
is 441
W V15
( ( from Wren) example are obscrvable.
Max 440
439
A38
Response_532
Respomen_48
A 2.36
435
454
_508
A33
Stea
siep, 432
by_526
4y1
Sies, 4.0
thikk
think
Ut. 478
LeL
521
A25
i. 520
bright A25
01
Is
meL
Max 423
lu,117
422
ren
fake
W
420
Tue. 419
lase_5133
.418
OL_ W_S12
511
srop. drap_416
1R,415
it
Uz 41.
umo 508
emp 413
6A_507
Is
..11
pUS
br1ght411
us. 403
UL_
UTO. 407
T,435
.005
ity 404
W.
FU 403
F.412
Uses 456
p_
Uz, 403
Ster
UTOO 309
ane. 485
1,335
Uses, 452
..33
ur_ .96
mpp_305
medium
lt
Is. 1.8488
is 153
US,
Max 392
UI
UITO 389
13388
is 187
USES 4B1
p_ 460
amp, 185
lec 475
2344
383
.33
CX76
Imput
ITS 475
2.000
474
#R_579
#4_473
0.1
0.08
0.06
0.04
0.02
0.1
0.08
0.06
0.04
0.02
Probability
Probability
25

---

Published in Transactions on Machine Learning Research (07/2024)
G Information flow
Assume given head h2g that writes some information to the residual stream corresponding to the token sq via output vj,x- Furthermore, let the maximum attention weight allocated by hy, be placed on the residual stream corresponding to token st. Then, we represent the head with the triplct (st, $1, Sinto = arg max U Uy),)) denoting the source residual stream (source of information written by the head), larget residual stream (target of information written by the head), and the information that is being written.
We start with the answer writing heads, ", For (sL, Si. Sinfo) associated with cach h " we identify the information content present within the source residual stream as arg max Uz! (i.e., the information present within the source residual stream right before the answer writer attended to it). Then, we look for all the heads in decoder blocks < j > responsible for writing that information into that residual stream. This continues till we reach a head that is in the starting decoder block. Figures 13 depict information flow extracted for a single example question for subtask 1. The tokenized input is as follows:

_0 
_1 
_ 
_3 
_< 4 
_5 ###_ Input_7 7 : :_8 <0x0A>_ _9 G_10 or
or_11 11 pus_12 is_ tw w__14 imp_15 15 us us_16 ._17 Alex_ 18 1s_19 rom_20 20 pus 21 ._22 Rom 23 pus_24 is 25 is_25 g_26 or or_27
pus_ 28 ._29 Gor
Gor_30 30 pus_ ___3 is_30 _2 small_33 -_34 Rom_ 35 pus__p 36 is_37 37 mean _8 ._39 True_40
or_4: 41 false_40 42 :_43 Alex_ 44 is_ 15 small_48 46 ._47 Let_48 us_ _s_49 think 50 step_51 51 by_52
step_53 53 ._54 <0x0A>_55 55 ##_56 #_57 Response_55 58 :_59 <0x0A>_ __0 Alex_ 61 is 62 rom_ 63 pus_6e 64
._65 Rom 66 pus_ 67 is_68 68 8_69 o0 pus_T1 71 ._72 Alex_73 73 is_74 g_75 or_76 pus_77 77 ._78
Gor_78 _7_ pus_80 80 is_81 81 small_82 __83 Alex_ _4 is_85 small_ 86 -_77 True_ 88 <0x0A> 89 <0x0A>_00 90
## 91 #_92 Input_98 :_94 <0x0A>_ 95 G_96 orp_97 uses_ 98 are_99 disc_100 ord_101 ant ant_102
_00 fax_104 is_105 zum_ 106 pus_107 -_108 z_109 ump_110 us _s_111 is 1s_112 sh_113 amp_
or or_115_ ._116 z_117 ump_118 us
_1_119 119 is_120 g_121 or or_122 pus_123 _as___23____ __12 Gor Gor_125 pus_ 126
is 127 earth_ 128 y_ y_129 -_ 130 True 131 or 132 false 133 :_134 Max 135 is _s_136
136 small _17
138 Let 1e9 us_ ss_ think 141 step_ 142 by_ 143 step_144 144 __145 _146 <0x0A>_ _117__141111_ ## ##148
#_149 Response_ _50 :_151 <0x0A> 152 Max_ 153 is_154 154 zum_ zum_155 _us_156 1_157 Z_158 ump_159 159
us us0160 160 is 161 g_162 or or_163 pus pus_164 _64
_s_161__
__65 Max 166 is 167 is_167 g_
8_168 or 1__ pus_ __171
Gor Gor_172 P_173 uses 174 are_175 disc 176 ord_17 ant 178 ._179
180 1s_181 disc_18p
ord_ 183 ant 184 ant_184
__185 False_188 <0x0A> _1000>117000000000 <0x0A> 10x00>_188 ## ##_189 # 190 input_19t : :_192
<0x0A> 11900 193 191 or or_195 pin_196 are are_197 197 w_198 ump_19t us_ __201 M_202 ump_ __2 uses_20 _04
are _05 angry_ 206 __207 W_ v_208 ump_ 209 us_ _s_210 is
is_211 211
1_212 empor_ 213 __214 s 215 ally_216
is 217 1e_ ___218 mp_219 us _s_220 __221 L_ L_222 emp_ 223 us _s_224 is_225 w w_226 ump_22 227 us_ _2_ _29
True_ 230 or 231 false_ 232 : :_233 s_ 234 ally_235 is 236 fl 21_237 oral_ __8 __239 Let_ 240
us 1s_2211 think_ 242 step_200 243 by b4_244 step_ 245 __246 _247 <0x0A>_244 ## ___ # #_250 Response_ 2511
:_252 c0x0A>_253 s 254 ally_
256 2__25 mp_258 us _s___9 __0 L_261 amp_262 us us_263
is_264 w_ v_265 ump_ 266 us us_267_ __2 $_9 ally _20 is_271 w_ 272 ump __3 us us_274_ __275 W W_276
ump_ 277 uses 278 are are_279 angry_ 280 _281 $_282 ally 283 is 284 angry_288
__286 False _277 287
<0x0A> 288 <0x0A> 289 ## 290 # #_291 Input 292 : :_293 <0x0^> 294 G_295 or or_296 pus_297
is 298 jel_29 299 git_300 301 Y 302 ump_
m03 303 uses_30d 304 are_ 3re_305 loud_ 106 __307 Gor 308
pus_ pus_309 is a10 y_311 ump 312 us 313 _314 Y Y_315 ump_ ump_316 316 us 317 is 318 orange_31 319 __320
R_321 ex_ s2_322 is 323 g_ 324 or 325 sus_326
is_323
.27 True ___ or 3_9 aalse_330 : :_331 R_332
ex_ ex_333 is 334 loud_ _03 _33 Let_ 337 us us_33t _38 think 339 step_ 340 by_ 341 step_ step_342 ._343
<0x0A> 344 ## =45 # 146 Response 347 :_348 <0x0A> _349 R_ R_350 ex 351 is_352 g_ $_353 or or_354
pus_ 355_355 ._356 Gor
Gor_357 357 pus_ 358 is $59 y_3 ump_361 361 us 362 __363 R R_364 ex ex_365 is ss_666 y_367
ump_3e 368 us _s_369 -_370 Y_33 Y_371 ump_37 372 uses 373 are 374 loud 375 __376 R_377 ex_ ex_378 is is_379
loud_ 380 3_38 True_ True_382 <0x0A>_a0 383 <0x0A>_ _384 ##_385 # 386 Input_380 _87 :_388 <0x0A> 389
L_390 emp_391 391 us_ us_392 $s_393 t t_394 ump_395 395 us
_s_396 396 ._397 Max_398 is _s_399 1e 400 mp_401 401
us_ us_402 ._ 4.4 T_404 ump_40 405 us_ us_406 is 107 f_408 ru 409 ity_ -10 __411 T_412 ump_ _43 us_40 us_414
is_415 bright_4f1 _16 .47 L_418 emp_ 419 us 420 is_421 drop_ 422 ant _23 ..4 Tr5e_425 425
or or_426 false_ 227 :_428 pax_429 430 bright_43 431 __2 Let_ us us_434 think think_435 st6p_436
by_437 step_490 438 . 4_4 <0x0A>_44>> ## _411 442 Response 443 __444 <0x0A> 445 Max_
is is_447 248 mp_449 us us_450 _51 L_452 emp_453 453 us_454
is_455 t_456 ump_ us_45t__ .59
Max_ 460 is 461 t 462 ump 463
464 465 T 466
ump_ 467 us _s_ 468 is 469 bright_49 470 471
26

---

Published in Transactions on Machine Learning Research (07/2024)
Max_ 472 is_473 bright_ 474 ._475 True_476 (0x0A>_47) <0x0A> _48 ##_479
#_480
Input_48 _81___41________
: :_482 <0x0A>_ 483 s_ 1_484 ter_ _85 P- 4.6 uses_48) are_ 488 d Q_489 ul1_490 __491 Imp_ 29_ us us_493
29_494 medium_ 495 496 Imp_ Imp_497 uses_ 498 are_ are_499 ster_50 p_501 uses_50p 5022_
_53 W_ W_504
ren_ 20n_505 is is_506 imp_507 507 us 508 _509 Ster_ Ster_510 pus pus_51i 511 is_512 da_51s 2a_513 ump_514 in_515 ._516
True_517 517 or 518 false_ 519 : __520 W_ w_521 ren_ 522 is 2_3 mel 524 od 525 ic 526 __527 Let_ Let_528 __8
us us_529 think_530 530 step_533 531 by_532 532 step_55, 533 _5334-- <0x0A>_553 ## ##_ #3666 #_537 Response_55 __38
: :_539 <0x0A>_544 540 W_541 ren 542 is 2_5443 imp_544 544 us us_544 545 __546 Imp_ 547 uses 548 are_ 549
ster_ 550 p_551 uses_552 .5553 W_554 ren 555 is 556 ster_ 557 pus_ 558 ._559 Ster_ 560 p_561
uses_ 562 are are_563 d_ d_564
_566 W_567 ren_ ren_568 is_569 d 570 ul1_571 __572 False_573
<0x0A>_574 <0x0A>_577 575 ##_576 # #_577 Input_575 _78 : :_579 <0x04>_500 580 J_581 om_ mo_ pus _as__33 i8_584
wooden_589 _5886 Imp_ 587 us
us_588 588 is is_589 j_590 om om_591 pus_
592
__593 St_ 594 alla_595 is_
imp_599 597 us_598 599 T00e_600 or or_601 false_602 602 : 603 St_604 ella_
is_606 wooden_607
_68 Let_609 us us0 chink_611 611 step_612 612 by_613 step_614 _615 mxx0A>_616 ## ##_617 # #_618
Response_613 619 __0 <0x0A>_62 _21 St 622 ella_ 623 _2_
Figure 14 depict information flow extracted for single example question for subtask 5, The tokenized input is as follows:

 0 
_1 
_2 <B>_3 
 4 
_5 
_6 ###_ Input 8 :_9 <0x0A> 10 G_ 11 or _r_1 pus_ __
is is_14 14 tw_ imp_16 16 us _s_17 17 -_18 Alex_19 19 is_20 rom_ 21 pus 22 ._23 Rom_ 24 pus 25 is_26 g_27
_us_22
pus_25
or 28 pus_ pus_29 ___29__ ._30 Gor_ 31 pus_ 32 is_33 small_ __11__3_____ ._35
Rom 36 pus_ ___ is_38 mean_30 39 ._40
True_4t or_42 false_4 _3 : __44 Alex_45 is_46 small_ 4.7 .488 Let Let_49 49 us 50 think 51 step_5s _2
by_53 step_5e 54, ._55 <0x0A>_ _56 ## ##_57 #_58 Response 59
:_60 <0x0A>_60 61 Alex 62 is 63 rom 64
pus_68 65 ._ 66 Rom_6f pus_ 68 is 19_69 g_70 or or_71 p2 ._73 Alex_74 74 is_ 8_76 or_ 77 pus_78
._79 Gor_ pus_ 81 is_82
82 small_83 _88 Alex_ 85 is_ is_86 small_81 1 __88 True_ 89 <0x0A>_90
<0x0A>_91 ##_92 #_93 Input_ _4 :_95 <0x0A>_96 G_97 orp_ _8 uses_99 99 are_100 disc_ 101 ord_
ant_103 ._104 Max_100 1s_106 zum_107 pus_108 ._109 z_110 ump_ 111 us_ 1s_113 sh_114
amp_115 or_116
__117 2_118 ump_119 us 120 is_12: g_ 121 or_128 pus 124
125
Gor_ 1126
pus pus_127 is_128 earth_129 129 y y_130 131 True_130 132 or or_133 false false_13411 134 135
bax_136
is_ 137
small_130 138 1_9 Let Let_140 140 us_ us_141 think 142 step_ 143 by_144 step 145 _164
_17 <0x0A> 148
## ##_149 # #_150 Response 151 __152 <0x0A>_ 153 Max 154 is_155 155 zum 156
1100>>1133000
pus_ 157 _15 Z_159
ump_ 160 us se_161 is 162 g_3 or 164 pus_166 165 __166 Max 167 is 168 g 169 or _r_170 pus_171
._12 Gor_ 173 _or_173 p_ 174 uses_ 175 are_ 176 disc_ 177 ord_ 178 ant_
180 Max 181 _s_182
disc_188 183 ord_184 184 ant_185 __886 False_18 187 <0x0A> _18 <0x0A> __9 #90 #_191 Input_ 192
: :_193 <0x0A>_194
B_195 or 196 pin_ 197 are_ are_198 w_ w_199 ump_ 200 us us_201 ._202 W_20 __203 ump_ 204
or_196
uses_ _0s are are_206 angry_ _27 __208 W_ N_209 ump_ 220 us 211 _s_212 j_213 empor __1122241_11 __215 S $_216
ally_21 217 is_21 s__218 le_ 1e_219 mp_22 mp_220 us 221 __222 L_223 emp 224 us us_ 225 is 226 w_ 2_227 ump ump_228 228
is_226
us_ 2s_229 __230 True_ 231 or_ 2r_222 false 233 : :_234 $32 ally_ _36 is 2s_ fl 238 oral 239
false_233_
f1_238
__240 Let let_241 us_ 2s_ think_ 243 step_ _44 by_245 step_ 246 ._247 _248 <0x0A>_ __9 ## ##_250
#_251 Response 252 : :_253 <0x0A>_ _54 $_ 255 ally_25 256 is 257 1e_258 mp_259 us us_260 ._261
L_262 emp_ 263 us ua_264 is_265 w v_266 ump 267 us
emp_263
268 us_268
__269 s 270 ally_21 271 is $s_272 v_273
ump_274 us us_275 __276 W_277 ump_278 uses __9 are_280 angry_28
_22 S_283 ally_28d is is_285
angry_
_266 286 ._287 False_288 288 <0x0A> 289 <0x0A>_29 _20 ## _9_911 # 292 Input_ 293 : :_294 <0x0A> 295
G_296 or or_997 pus_ pus_298 is 29_ jel_300 300 git_30
__302 Y_303 ump 304 uses 305 are are_306 loud_ 307
308 Gor _09 pus_310 is is_311 y_312 ump_31r 313 us _s_314 ._315 y_316 ump_3m 317 us _s_318 is_319
orange_300 320 __321 R_322 ex_ ex_323 is is_324 g_ g_325 or or_322 p_326 pus 327 __328 True_ __9 or_ rr0 False_331
:_332 R_ 333 ex_ ax_334 is 1s_335 loud_ 336 __337 Let _et_338 us 339 think 340 step_34t
by_343
step_ 343 __344 <0x0A>_344 ## 346 # 347 Response 348 __349 <0x0A> 350 R R_351 ex xx_352 is is_353
g_354 or rr355 pus_ sas_6666_ ._357 Gor _3____8 pus_ pus_359 is ss_60 y_361 ump_362 _62 us
_s_363 363 __364 R_365 ex 366
is_367 y_ y_368 ump_ 369 us us_370 370 3_311 y_372 ump_ 373 uses
uses_374 374 are
are_375 375 loud _ ._377 R_378
ex_379 is _s_380 loud_ 381
__382 True_ 383 <0x0A> _384 <0x0A> 385 ##_386 # #_387 Input_ __88
: :_389 <0x0A>_390 390 L_391 emp_
__92 us _s3 is 1s_394 t t_395 ump_3st 396 us _97 __398 Max_399 is 400
1e_401 m0_402 us_403 v03 404 T_405 ump_40ppp 406 us_ 407 ia_408 f_409 ru, ru_410 410 ity_411 __412 T T_413
ump_4t ump_414 us_ 21_ is is_416 416 bright_411
__418 L_419 emp_ -20 us_ us_421 is 422 drop_423
423 ant_424 __425
True_ 426 or 427 false 428 : __429 Wax_430 is_431 bright_430 432 433 Let 434
False_428_
us 435 think 436
27

---

Published in Transactions on Machine Learning Research (07 (22/))
step_ 437 by_438 4.8 step_43;
439 ._440 <0x0A>_ 441 ##_42 # 443 Response 444 : :_445 <0x0A> 4446
Max_ 447 is 448 le_441 449 mp_ 450 us 15_4511 -_452 L. 453 emp 454 us 655 is_ 456 t -_457 ump_
us_459 -_460 Max_460 is__i__860 462 t_ f_463 ump_464 464 us $8_465 1.666 T T_467 ump_ _4_ us us_ 469 is_470 470
bright_44 _41 ._472 Max_477 is 4744 bright_477
_77
True_477 <0x0A> <ex0>>_478 _478 <0x0A>_ _479 ##_480
#_481 Input 482 : :_483 <0x0A> 484 s $_485 ter 486 P_487 uses 488 are are_489 489 d 490 ull_ _911
._492 Imp_493 us_ 29_494 is_495 medium 496 _497 Imp_498 498 uses 499 are 500 ster ster_501 P_502
_ses_ 503 __504 *_505 ren_ 506 is 507 imp_ 508 us us_ 509 510 Ster 511 pus_ 512 is is_513 513 da da_
is_507
da_514 514
ump_ uP__15 in 516 _57 True_ 518 519 false 520 __521 522 ren 523 is 12_ 524 mel_525 od od_526
ic ic_527 528 Let_529 us 530 think 531 step_ 532 by_ 533 step 534 _555 <0x0A>_553 ## 537
# #_538 Response_5a0 539 __500 <0x0A>_54h 541 w_542 ren 543 is 144 imp 1mp_545 us_ 246 __54
1mp_548
uses_ 549 are 550 ster_5tt
p_552 uses 553
554 W 5555 ren 556 is 557 ster
558 pus _59
_560 Ster_55e 561 P_ P_562 uses 563 are_564 d_565
wll_ 566 567 W_ 568 ren_ 569 is_570 d_571
u11_572 ._573 False _74 <0x0A> 10x0A>_575 575 <0x0A> _566
574
## 577 #_878 Input_577
:_580 <0x0A>_50
J_582 om om_583 pus_584 584 is ss_85 wooden_558 586 __587 Imp_588 588 us 589 is $__900 j_591 om
am__s______000 592 pus pus_593
.9594 St_50 St_595 ella_599 596 is 597 imp_ us _99
600 True_600 601 602 false _03 6_4 St_
ella_ 606 is_607 wooden_608 608 6,9 Let_610 us 611 think 612 step_ 613 by_614 614 step_615 615 __616
<0x0A>_ _677 ##_618
# #_619 Response_ _200 620 621 <0x0A>_ 622 St St_623 ella___ll_00le 624 is_628 625 imp_626
us us_627 -62 Imp_ 629 us 630 is 631
j. 632 om 633 pus 634 635 St 636
ella 631 is_638
28

---

Published in Transactions on Machine Learning Research (07 (07/202) /204)
220 Toaerpp
<05 26wiiwe-ammpp
<6 19-im->-777~ujp>
d1 50
<22 20001=--00--t
+3 3roapee
<22 ---0-=-77--->>>
< dipeap-t--depp
42:25mus
<I-1 4-10-ttmp-dd>>-tt
4t 2io-dspo Iondmpo-
.00 30 Siippps
22 2wwiiisssissiinn
45 20-226--mmp--00hta...s 1o--tpp--997-aae
d 1ot
Figure 13: Infomation flow through attention heads towards answer-writing heads in Llama- 2 7B for subtask 1. Each node in this DAG (pruned for brevity) is labeled <layer index-head <source residual content><source residual position><head output content>. The sub- task here is to predict impus given the generated response Stella is (see Appendix G for the complete tokenized input). In this example, there are seven answer writing heads (in green) that are writing <imp> to the last residual stream corresponding to In this particular example, the answer token imp is collected from token no. 597 in the input question context by 8 different (in bluc).
29

---

Published in Transactions on Machine Learning Research (07 (07/202) /224)
<20 30mpoe
01:30xpp4622aam>>>>
+13 22-00
of 22200-22-----itpp
-12 12-20-----2---
ot 40 300=om--522--nn
+ $3-0mp-d2--pp
<1 <1.3aepxs00::aurpp
<t 1xxctht----dlliieeeee
<15 45-12wamo--B--ging 4 Siben
< Si
22 22-20=p-652-amm
24 24280-p-622-mm
2 20- St
-13 43-1--aa0--522ppp
<6 1820mann--S22apss
47. 47:2220
+ 1=gp-99--amps
<17. 11:11==tppee1nt
<0-0 60:0:wetaa--wwsippp
S 3.t0::tthe==//dddppp
<)4 24-190-ogx040oiiepp
<13 20-boog---eeeppp
oi 2--00so-t---ppp-
43-72-o020-0osieppp
Figure 14: Infomation flow through attention heads towards answer- heads in Llama-2 7B for subtask 5. The representations of each node are similar to Figure I3 The subtask here is to predict jompus given the generated response Stella is impus. Impus is jompus. Stella is (sec Appendix G for the complete tokenized input). In this example, there are three answer writing heads (in green) that are writing <j> the last. residual stream corresponding to is. In this particular example, the answer token imp is collected from position, token no. 591 in the input question context heads (in blue), and token no. 632 in the generated context by heads (in red).
30


Reviews:
[{"review_id": "0U0e548lmj", "reviewer": "Reviewer_KCnw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The research topic of understanding CoT is very important and has attracted increasing interest in the community. This work presents a comprehensive analysis of the mechanism of CoT reasoning in LLMs. \n\n2. The empirical results are comprehensive. Some results well align with the observations in related CoT studies.\n\nWeaknesses:\n\n1. There are possible over-claims.  Various prior studies have tried to understand CoT reasoning either theoretically [1] or empirically [2][3]. Though the analysis approach may vary, it is risky to claim that \"To the best of our knowledge, this is the first attempt towards mechanistic investigation of CoT reasoning in LLMs.”\n\n2. There are some findings that may not well generalize to other LLMs. This work ambitiously aims for a mechanistic understanding of\nchain-of-thought reasoning in LLMs. However, the analysis is purely based on the Llama-2 7B model (w/ few-shot prompting), which is not widely adopted or optimized for real CoT-based tasks. Besides, there are some \"magic\" observations, e.g., \"a functional rift at the very middle of the LLM (16th decoder block in case of Llama-2 7B)\" It is not clear whether those findings could be different in other LLMs. I understand the potential computation cost when using a larger model. However, as the vanilla  Llama-2 7B is not a strong baseline for CoT reasoning, it would be more convincing if CoT-tuned Llama (e.g., Alpaca or Vicuna series)  could be adopted to support the findings.\n\n3. The task definition is not clear. This paper mentions subtasks 0-9 many times. Are all the tasks the same as those in Figure 1, say, Decision-type subtasks (0, 2, 4, 6, and 8), Copy-type subtasks (1, 3, and 7), and Induction-type subtasks (5 and 9)? A clear categorization of those tasks would be useful.\n\nReferences:\n\n[1] Ben Prystawski and Noah D Goodman. Why think step by step? Reasoning emerges from the locality of experience. Advances in Neural Information Processing Systems, 2023.\n\n[2] Xinyi Wang and William Yang Wang. Reasoning ability emerges in large language models as aggregation of reasoning\npaths: A case study with knowledge graphs. In Workshop on Efficient Systems for Foundation Models@ ICML2023,\n2023.\n\n[3]Zhen Bi, Ningyu Zhang, Yinuo Jiang, Shumin Deng, Guozhou Zheng, Huajun Chen. When Do Program-of-Thought Works for Reasoning?. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 38, No. 16, pp. 17691-17699), 2024.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work seeks for a mechanistic understanding of chain-of-thought (CoT) reasoning in large language models (LLMs). To this end, this work leverages Llama-2 7B in few-shot prompting for multi-step reasoning over fictional ontologies. Various findings have been reported after the analysis. Empirical studies show that  LLMs deploy multiple parallel pathways of answer generation for step-by-step reasoning. These parallel pathways provide sequential answers from the input question context as well as the generated CoT. Besides, there is a functional rift at the very middle of the LLM (16th decoder block in case of Llama-2 7B), which marks a phase shift in the content of residual streams and the functionality of the attention heads.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Check the claims and accurately position the contribution of this work.\n\n2. Justify the choice of using the vanilla Llama-2 7B for analysis. Possibly consider a more convincing analysis on CoT-tuned LLMs, such as Llama-2 Alpaca or Vicuna.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "7DhEqwzm8H", "reviewer": "Reviewer_Fxq8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I find particularly interesting the finding regarding the difficulty of identifying subsets of heads associated with a particular subtask (but that is according to a particular decomposition of the reasoning task, it does not exclude that other ways of thinking about the problem would not find such an assignment).\n\nI find the other results more difficult to interpret. Nonetheless, this is part of a growing body of observations to make sense of what is going on in Transformers trained as LLMs, with a focus on reasoning, which is original and interesting.\n\nThe paragraph on the fictional ontology took me a while to understand. The main message should be that we set up a set of invented facts and entities on which the LLM has to reason. However, these should not contradict the facts seen by the LLM during training. They should be hypothetical facts which could be true in principle and thus indeed constitute a fiction. However, if they contradicted the training set facts, it would be a problem in that those training set facts would bias the reasoning done by the LLM. Also in this paragraph we introduce the term \"substasks Ti\" without really explaining what is intended. An example at this point would be neeeded.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors report the results of a series of experiments to better understand characteristics of computations taking place during chain-of-thought reasoning with an LLM (LLaMA-2). The experiments are performed with questions requiring to reason about \"fictional\" characters and facts to avoid contaminating the answers with pre-learned broad knowledge of facts and only rely on the general reasoning abilities of the LLM. A notable result is they do not find a decomposition of the required \"natural\" decomposition of the reasoning problem into tasks that mapped to any subset of heads: instead heads seemed to participate in many tasks. However, they found that nonlinear probes (but not linear probes) could detect the required movement of information and observed important differences between early and late layers.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please have the text revised by a human or an AI system to correct numerous little grammatical errors.\n\nSee the weak points made in the above section, which should be addressed.\n\naccuracy of 0.9: clarify if it is 90%.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3r4ny3AQay", "reviewer": "Reviewer_9pmv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- The authors apply mechanistic interpretability to a relatively large-scale model (relatively compared to other mech. interp. work), which is an exciting contribution. \n- Clever way of focusing on attention heads over MLPs by looking at a dataset with a fictional ontology\n- The authors do a lot of analyses and present a lot of relatively well-motivated findings (particularly the results in section 5 and 6 are quite convincing due to the clear patterns found). \n\n**Weaknesses**\n\n- I find the descriptions of the methods in the paper very confusing, and not enough information is provided to interpret the results. I am left with many questions that I don't know the answer to but are necessary to properly follow what is going on and decide whether the claims are sound. Mainly for that reason I decided to write a relatively detailed contributions section, so in case I am misunderstanding something it would be great if the authors can point that out.\n- Similarly to the above weakness, there is a lot of information put in the appendix without which you simply cannot understand what is going on in the paper, like the importance scores for the different types of attention heads.\n- The paper is studying a relative simple form of reasoning, namely deductive over a fictional ontology, only studies one model family, and only one model size. There is not limitations section present in the paper or a section that discusses how the results might generalise to other model families and model sizes, or what the results might tell us about chain-of-thought reasoning more generally.\n- This is a personal preference so feel free to ignore this, but the amount of hyperboles in the first two paragraphs of the introduction is a bit much. Personally, I find that overusing hyperboles makes a text come across less strong.\n\nI'm going to pose my questions below, and I'd like to work with the authors to understand the method better so I can be more convinced of the findings:\n\nQuestions:\n1. Why do you decide which attention heads to knockout only using the false ontologies?\n2. I very much struggle to understand what is going on in section 4. For example, I do not understand how you construct Figure 2. Part of my confusion is caused by the fact that the Y-axis in Figure 2 says accuracy is red and fraction of heads involved is blue, whereas the caption says the opposite. Additionally, in the text you claim subtask 3 is a decision making subtask, but it seems like in Figure 1 subtask 3 is a copy task. Can you tell me if the following is correct for my own understanding and otherwise correct me where I'm wrong:\nFor example for the decision heads figure (leftmost), you take all the decision subtasks (i.e. 0, 2, 4, 6, 8 in Figure 1), you calculate the importance scores for each attention head, and you select a range of importance scores within which if you keep those attention heads active, you get a relative accuracy of at least 0.9 for the subtask you look at (from 0 to 10). This leads to a different number of attention heads active for each subtask 0-10, and a different absolute accuracy corresponding to the relative accuracy of the decision heads on each those subtasks. This means that if you get an absolutely high accuracy (high blue bar) in the figure for a low fraction of heads involved (low red bar), those attention heads are important for that subtask\n3. In Section 5, how are you getting the specific representations that are used as inputs to the classifier? So if you have the following example \"A is B, B is C, then A is C\", which activations do you take, those when the model is predicting C (the last token)? And which occurrence of B do you take as input to the classifier then, the first or second? \n4. why do you call deductive reasoning induction? As far as I know, the type of reasoning the dataset you use requires is purely deductive.\n\nI have more questions but I'm going to start with this and see how far we get in discussion.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper analyses the mechanisms involved in the generation of a chain-of-thought for deductive reasoning tasks in a fictional ontology by Llama-2 7B in a few-shot setup. The authors present several interesting results that shed led on the inner-workings of an LLM while performing reasoning tasks with CoT prompting; namely that there is a functional difference between the bottom half of the layers and top half of layers, with the attention heads in the former moving information around and the latter writing answer tokens. Additionally, the authors find that the model utilises multiple pathways for answer writing. The main contribution is shedding light on the inner workings of a model at scale (7B), where most mechanistic interpretability findings in existing literature are about smaller models.\n\nThe tasks the authors investigate is a deductive reasoning task that can be naturally split in 3 subtasks: **deciding** which premise to start with, **copying** the relevant tokens from that premise, and **induction** of a new statement. For example, the context might be *\"Tumpus is bright. Lempus is tumpus. Max is lempus. True of false: Max is Bright\"* and the CoT to generate is then *\"Max is lempus. Lempus is tumpus. Max is tumpus. Tumpus is bright. Max is bright. True\"*. The contribution here is the insight to use a fictional ontology to be able to solely focus on analysing attention heads, as prior work has shown that the MLPs are mostly used for a key-value-style lookup of facts learned from pretraining (which isn't possible with the fictional ontology). \n\n**More detailed description of method and findings**\nThe authors use different techniques from mechanistic interpretability to investigate the contribution of specific attention heads to particular subtasks in the deductive reasoning task. They score each attention head in the Transformer three times, and each score reflects the relative importance of that attention head for one of the three subtasks. An attention head can be deemed a \"decision head\" if the relevant importance score is within a certain threshold. One can \"knock out\" heads that are not \"decision heads\", for example, and leave a set of heads that are task-specific. You can do predictions with those on the subtasks required for generating the CoT. Some attention heads are \"task-only heads\", when they are only one type of task-head (e.g. they are only not knocked out for the induction task). The authors do many analyses in this way to find the following things:\n- The three types of heads do not necessarily only get high accuracy on their respective subtasks (e.g. decision heads are not better at decision making than other tasks)\n- Induction heads get high accuracy on all the subtasks required for generating the CoT\n- The authors hypothesise that induction heads need to mix information, decide which induction to make, and then copy the required tokens, and that's why they are important for all subtasks. They show that indeed the information mixing happens, mostly between layer 10-15\n- The authors find that mostly the deeper layers' attention heads are context-abiding (meaning that applying an unembedding layer to their activations leads to a bigram that is actually present in the context)\n- The authors further find that there are specific answer-writing heads that directly write the answer to a subtask into the residual stream (also found by applying unembedding to their activations which produces the right answer). There are multiple of these, and they implement different algorithms (found by following their attention patterns to construct a tree)", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Requested changes that are necessary for acceptance**\n- Figure 2 is quite difficult to parse not only due to my own confusions, but also due to the fact that the 0-9 subtasks are not annotated according to which category the task refers to. Please annotate this somehow in Figure 2, and everywhere else.\n- Some of the appendix information needs to be moved to the main text, for example I would propose to write an intuitive summary of the importance scores for each subtask in the main text, because it's kind of hard to interpret the rest of the paper without it. Appendix C.4 also really needs to be moved to the main text.\n- Some information simply seems to be missing from the whole paper, like how exactly do you decide which attention heads to keep for Figure 2? \n- Figure 3 is also difficult to interpret, especially using a coloring determining accuracy over a histogram of bars. I would suggest either using a different visualisation or at least add a guiding interpretation in the caption text.\n- In my experience, writing a strong limitations section only makes a paper's conclusions come across stronger, and I advise the authors to do this as well, especially as there are quite a few differences between their setup and general CoT reasoning by modern LLMs. Additionally, the finding that MLPs are mostly key-value lookups from pretraining (roughly) is not a certainty and I would like to see this discussed in a limitations section.\n\n**Writing style suggestions and other less important suggestions**\n- It took me quite a while to parse what the 0-9 subtasks refer to, could you note when you mention them first that they can be seen in Figure 1? Additionally, I would strongly consider using a different word for these 10 subtasks and the 3 subtasks that are task categories, because using subtasks for both is confusing.\n- Less hyperboles in the abstract and introduction\n- Unimportant nit: I think LLaMA-1 was spelled LLaMA, and Llama-2 is Llama.\n- Personal preference: I think the submission could benefit from a representative figure moved up to the top of page 2 with some standalone explanation of the findings", "broader_impact_concerns": "N.A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Overall, the reviewers were satisfied with the paper's methodology and results. Mechanistic interpretability (MI) has so far been applied to relatively small-scale models; this is the first application of MI to an LLM.  The reviewers had some concerns about the level of methodological detail and the discussion of related work in the original paper. However, the authors have done a good job of revising the paper and have adequately addressed these concerns. Given this, I recommend acceptance.
Decision: Accept (poster)