Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models
Authors: Yifan Hou
Jiaoda Li
Yu Fei
Alessandro Stolfo
Wangchunshu Zhou
Guangtao Zeng
Antoine Bosselut
Mrinmaya Sachan
Venue: EMNLP 2023 Main (2023)
Keywords: model interpretation
reasoning
attention mechanism
large language model

Abstract: Recent work has shown that language models (LMs) have strong multi-step (i.e., procedural) reasoning capabilities. However, it is unclear whether LMs perform these tasks by cheating with answers memorized from pretraining corpus, or, via a multi-step reasoning mechanism. In this paper, we try to answer this question by exploring a mechanistic interpretation of LMs for multi-step reasoning tasks. Concretely, we hypothesize that the LM implicitly embeds a reasoning tree resembling the correct reasoning process within it. We test this hypothesis by introducing a new probing approach (called MechanisticProbe) that recovers the reasoning tree from the model’s attention patterns. We use our probe to analyze two LMs: GPT-2 on a synthetic task (k-th smallest element), and LLaMA on two simple language-based reasoning tasks (ProofWriter \& AI2 Reasoning Challenge). We show that MechanisticProbe is able to detect the information of the reasoning tree from the model’s attentions for most examples, suggesting that the LM indeed is going through a process of multi-step reasoning within its architecture in many cases.

Towards Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models
Yifan Hou', Jiaoda Li', Yu Fei", Alessandro Stolfo', Wangchunshu Zhou', Guangtao Zeng', Antoine Bosselut", Mrinmaay Sachan'
ETH Zürich, ²UC Irvine, $ AlWaves, "SUTD, SEPFL.
fan hou, jiaoda. alessandro.stol
saaayaa-achaaeeinf .etz..ch, ch, 'yu. fei@uci eddu,
"hhnssu@aiwaves.ccmm
auutd.du.... antooine. wow.iee.sssseeeee........
Abstract
Zebra Puzzle:
1. There are five houses.
Englishman lives red house.
Recent work has shown that language models (LMs) have strong multi-step (i.c.. procedural) reasoning capabilities. However, whether LMs perform these tasks by cheating with answers memorized from pretraining cor- pus, or, via a multi-step reasoning mechanism. In this paper, try to answer this question by exploring a mechanistic interpretation LMs for multi-step reasoning tasks. Concretely, we hypothesize that the LM implicitly embeds reasoning tree resembling the correct reason- ing process within it. We test this hypothesis by introducing new probing approach (called Mechanis sticProbe) that recovers the reason- ing tree from the model's attention patterns. We use our probe to analyze two LMs: GPT- 2 on synthetic task (k-th smallest element), and LLaMA on two simple language-based rea- soning tasks (ProofWriter & A12 Reasoning Challenge). We show that Mechanistion stiPrrobe is able to detect the information of the reason- ing tree from model's attentions for most examples, suggesting that the LM indeed is go- ing through a process of multi-step reasoning within its architecture in many cases.?
Human
Now, who drinks water? Wha owns the zebra?
unclear
Reasoning
Recall:
that
remember d! The anower
ortnis water and
The Norwegian drinks The Japanese owns the zebra
Figure example showing how LMS solve reason- ing tasks. It is unclear if LMs produce the answer by reasoning or by recalling memorized information.
language-based reasoners (Creswell and Shanahan, 2022; Creswell et al., 2022; Chen et al., 2023). A recent line of work tests the behavior of LMs by designing input-output reasoning exam- ples (Zhang al., 2023; Dziri et al., 2023). How- ever, it is expensive and challenging to construct such high-quality examples, making it hard to generalize these analyses to other tasks/models. Another line of work. mechanistic interpretabil- ity (Merullo et al., 2023; Wu et al., 2023; Nanda et al., 2023; Stolfo et al., 2023; Bayazit et al., 2023), directly analyzes the parameters of LMs, which be easily extended different tasks. Inspired by re- cent work (Abnar and Zuidema, 2020; Voita et al., 2019; Manning et al., 2020; Murty et al., 2023) that uses attention patterns for linguistic phenomena prediction, we propose an attention-based mecha- nistic interpretation to expose how LMs perform multi-step reasoning tasks (Dong et al., 2021).
1 Introduction
Large language models (LMs) have shown impres- sive capabilities of solving complex reasoning prob- lems (Brown et al., 2020; Touvron al., 2023). Yet, what is the underlying "thought process" of these models is still unclear (Figure 1). Do they cheat with shortcuts memorized from pretraining corpus (Carlini et al., 2022; Razeghi et al., 2022; Tang et al., 2023)? Or, do they follow rigorous reasoning process and solve the problem proce- durally (Wei et al., 2022: Kojima al., 2022)? Answering this question not only critical to our understanding of these models but is also criti- cal for the development of next-generation faithful
We assume that the reasoning process for an- swering multi-step reasoning question can be rep- resented as a reasoning tree (Figure 2). Then, we investigate if the LM implicitly infers such a tree when answering the question. To achieve this, we designed a probe model, Mechani s crobe,,, that recovers the reasoning tree from the LM's attention patterns. To simplify the probing problem and gain a more fine-graine aine-graned understanding, we decompose the problem of discovering reasoning trees into two
OOu code, as well as analysis results, are available at https:/githuh.com/yffaa wtw.ooyyyyyaaaa//ccaciiiiieie Prroan

---

Reasoning tree
5
Bob Fongerous. True Sabe?
(.Whch event occun dayy syou?
21
the is kind event
Jy serreone them rater herate
5y the nany a everr (a)
Syy s2 Se
Seg 336
Sg B sormecee is kind then hershe nos dargarous
by the mang and setting is the event that ocoun
statements
Input
smallest element (R=5)
ProofWriter
A12 Reasoning Challenge
Figure Illustration of our MechanisticProbe with one example from each the three reasoning tasks considered in this work. We are given a number of input statements: S S2,.. and question: Q. EechanisticPron recovers the reasoning tree describing the ground-truth reasoning process to answer the ques- tion. MechanisticPron
works in two stages: in the first stage, MechanisticProbe detects if the LM can select the set of useful statements required in reasoning, and then in the second stage, AaccanisstiProon detects if the LM can predict the reasoning tree given the useful statements.
subproblems: 1) identifying the necessary nodes in the reasoning tree; 2) inferring the heights of iden- tified nodes. We design our probe using two simple non-parametric classifiers for the two subproblems. Achieving high probing scores indicates that the LM captures the reasoning tree well.
2.1 Reasoning Formulation
In our work, the LM is asked to answer a ques- tion Q given set of statements denoted by S = [S1, S2. Some of these statements may not be useful for answering Q. To obtain the answer, the LM should perform reasoning using the statements in multiple steps. We assume that this process can be represented by a reasoning tree G. We pro- vide specific examples in Figure for the three reasoning tasks used our analysis.? This is a very broad formulation and includes settings such as theorem proving (Loveland, 1980) where the statements could be facts or rules. In our analyses, study the tasks described below.
We conduct experiments with GPT-2 (Radford et al., 2019) on a synthetic task (finding the k-th smallest number in a sequence of numbers) and with LLaMA (Touvron et al., 2023) on two natu- ral language reasoning tasks: Proof (Tafjord al., 2021) and AI2 Reasoning Challenge (i.e., ARC: Clark et al., 2018). For most examples, successfully detect reasoning trees from attentions of (finetuned) GPT-2 and (few-shot finetuned) LLaMA using MechanisticProbe We also ob- serve that LMs find useful statements immediately at the bottom layers and then do the subsequent reasoning step by step ($4).
2.2 Reasoning Tasks
k-th smallest element. In this task, given a list of m numbers (m = 16 by default) in any order, the LM is asked to predict (i.e., generate) the k-th smallest number in the list. For simplicity, only consider numbers that can be encoded as token with GPT-2's tokenizer. We select m numbers ran- domly among them to construct the input number list. The reasoning trees have a depth (Figure 2 left). The root node is the k-th smallest number and the leaf nodes are top-k numbers.
To validate the influence on the LM's predictions of the attention mechanisms that we identify, we conduct additional analyses. First, we prune the attention heads identified by MechanisticProba and observe a significant accuracy degradation ($5). Then, we investigate the correlation between our probing scores and the LM's performance and robustness ($6). Our findings suggest that LMs exhibit better prediction accuracy and tolerance to noise on examples with higher probing scores. Such observations highlight the significance of ac- curately capturing the reasoning process for the efficacy and robustness of LMs.
For this task, we select GPT-2 (Radford et al.. 2019) as the LM for the analysis. We randomly generate training data finetune GPT-2 and ensure that the test accuracy on the reasoning task larger than 90%. For each k, we finetune an independent GPT-2 model. More details about finetuning (c.g.. hyperparameters) are in Appendix B.
2 Reasoning with LM
ProofWriter. The ProofWriter dataset (Tafjord et al., 2021) contains theorem-proving problems.
In this section, we formalize the reasoning task and introduce the three tasks used in our analysis: k-th smallest element, ProofWriter, and ARC.
Note that we can leave out the question if Q remains the same for all examples (c.g., Figure 2 left).

---

In this task, given a of statements (verbalized rules and facts) and a question, the LM asked to determine if the question statement is true or false. Annotations of reasoning trees are also provided in the dataset. Again, each tree has only one node at any height larger than O. Thus, know- ing the node height is sufficient to recover G. To simplify our analysis, we remove examples anno- tated with multiple reasoning trees. Details are in Appendix C.3. Furthermore, to avoid tree ambigu- ity (Appendix C.4), we only keep examples with reasoning trees of depth upto 1, which account for 70% of the data in ProofWriter.
Layer 1
Zurich
41
i
located
M
Switzerland
A(1,))
A(2,
ALL
Figure 3: An illustration of the attention mechanism in LM. The example sentence here is: Zurich is located in Switzerland. Token fy is the last token used for the prediction the next token f|1++ The arrows here show the attentions, where the red arrows denote the attentions to the last token.
A12 Reasoning Challenge (ARC). The ARC dataset contains multiple-choice questions from middle-school science exams (Clark et al., 2018). However, the original dataset does not have reason- ing tree annotations. Thus, we consider a subset of the dataset provided by Ribeiro et al. (2023), which annotates around 1000 examples. Considering the limited number of examples, we do not include analysis of finetuned LMs and mainly focus on the in-context learning setting. More details about the dataset are in Appendix C.3.
Let us consider the simplest of the three tasks (k-th smallest element). Each statement S; here is encoded taining all statements in S) to the LM comprises set of tokens as T (t1, t.... We denote the hidden representations of the token 4 at layer by z/. We further denote the attention of head h be- tween 2/-1 and z) as A(I, (t,,).) As shown in Figure 3, the attention matrix at layer I of head h is denoted as A(1, h), the lower triangular matrix.* The overall attention matrix then can be denoted by A - {L(L, A)]1 s IS L; S h S H).
token f4- The input text (con-
For both ProofWriter and ARC tasks, we select LLaMA (7B) (Touvron et al., 2023) as the LM for analysis. The tasks are formalized as classifi- cations: predicting the answer token (c.g., true or false for ProofWriter). We compare two settings: LLaMA with 4-shot in-context learning setting, and LLaMA finetuned with supervised signal (i.e., LLaMArr). We partially finetune LLaMA on at- tention parameters. Implementation details about in-context learning and finetuning on LLaMA can be found in Appendix C.I.
Our probing task detect G from A, i.e. mod- cling A). However, note that the size of is very large L x H x [ and it could con- tain many redundant features. For example, if the number of tokens in the input is 100, the attention matrix A for LLaMA contains millions of attention weights. It is impossible probe the information directly from A. In addition, the tree prediction task is difficult (Hou and Sachan, 2021) and we want our probe to simple (Belinkov, 2022) as we want to provide reliable interpretation about the LM rather than learn the task itself. Therefore, we introduce two ways to simplify attentions and the probing task design.
3 MechanisticProba
this section, we introduce MechanisticProbe a probing approach that analyzes how LMs solve multi-step reasoning tasks by recovering reasoning trees from the attentions of LMs.
3.2 Simplification of A
3.1 Problem Formulation
In general, we propose two ways to simplify A. Considering that LLaMA is large LM, we propose extra two ways to further reduce the number of considered attention weights in A for it.
Our goal is to shed light how L.S solve proce- dural reasoning tasks. Formally, given an LM with I. layers and H attention heads, we assume that the LM can handle the multi-step reasoning task with sufficiently high accuracy.
Focusing the last token. For causal LMs, the representation of the last token in the last layer 2/r1
'Note that most reasoning tasks with LMs can be formal- ized as the single-token prediction format (e.g., multi-choice question answering).
"Note that in this work we only consider causal L.ss.

---

is used predict the next token 41r1+1 (Figure 3). Thus, we simplify A by focusing on the attentions on the last input token, denoted as Aimp- This re- duces the size of attentions to Lx H x [T] Findings in previous works support that simp is sufficient to reveal the focus of the prediction token (Brunner et al., 2020; Geva et al., 2023). In our experimental setup, also find that analysis Asip gives similar results to that A (Appendix B.4).
Mean pooling
Max pooling
Layer 1 to
A
Stmp
Figure 4: An illustration of the way extend our atten- tion simplification method. In this example, S1 and Q are both composed of 2 tokens. Thus, there are 2 x 2 attentions. We mean pool all the tokens within the state- ments (remaining 1 x attentions), and max pool all the tokens in the question (remaining 1 attention).
Attention head pooling. Many existing attention-based analysis methods use pooling (e.g., mean pooling or max pooling) on attention heads for simplicity (Abnar and Zuidema, 2020; Manning et al., 2020; Murty et al., 2023). We follow this idea and take the mean value across all attention heads for our analysis. Then, the size of Asimp is further reduced to [T].
ProofWriter and ARC respectively. For finetuned LLaMA, 18 (out of 32) layers are removed for ProofWriter. More details about the attention prun- ing can be found in Appendix C.2.
Ignoring attention weights within the statement (LLaMA). For LLaMA the two natural lan- guage reasoning tasks, statement S; could contain multiple tokens, i..., [T] >> (S]. Thus, the size of simp can still be large. To further simplify Asip under this setting, we regard all tokens of a statement as a hypernode. That is, we ignore attentions within hypernodes and focus on atten- tions across hypernodes. As shown in Figure 4, we get A cross via mean pooling on all tokens of simp statement and max pooling on all tokens of Q as:$
3.3 Simplification of the Probing Task
We simplify the problem of predicting G by break- ing it down into two classification problems: clas- sifying if the statement is useful or not, and classi- fying the height of the statement in the tree.
P(GIAsimp) = P(V Asipp)) P(G|V, Asimp).
Here, V is the of nodes in G. P(iIA)imp) (binary classification) measures if LMs can select useful statements from the input based on atten- tions, revealing LMs correctly focus on useful statements. Given the set of nodes in G, the second probing task is to decide the reasoning tree. We model P,GIV, Asimp) as the multiclass classifica- tion for predicting the height of each node. For example, when the reasoning tree depth is 2, the height label set {0, 1, Note that for the three reasoning tasks considered by us, G is always a simple tree that has multiple leaf nodes but one intermediate node at each height.
A cruss (1, h)M max mean (A(,A),,))) ty
ES,
The size of simplified cross-hypernode attention matrix Agross is further reduced to L x (IS) + 1).
Pruning layers (LLaMA). Large LMs (e.g.. LLaMA) are very deep (i.e., L large), and are pretrained for a large number of tasks. Thus, they have many redundant parameters for performing the reasoning task. Inspired by Rogers et al. (2020), we prune the useless layers of LLaMA for the rea- soning task and probe attentions of the remaining keep minimum num- ber of layers that maintain the LM's performance on held-out development set, and deploy our analysis on attentions of these layers. For 4-shot LLaMA, 13/15 (out of 32) layers are removed for
3.4 Probing Score
layers. Specifically,
To limit the amount of information the probe learns about the probing task, classifier: k-nearest neighbors (kNN) to perform the two classification tasks. We use SFI (V (/Assm)) and SF1 (GVV, A simp) to denote their F1-Macro scores. To better interpret the probing results, we
non non-parametric
WWe take mean pooling for statements since max pooling cannot differentiate statements with many overlapped words. We take max pooling for Q since it is more diffcrentiable for the long input text. In practice, users can select different pooling strategies based on their own requirements.
"In practice, many reasoning graphs can be formalized as the chain-like tree structure (Wei et al., 2022). We leave the analysis with more complex G for future work.

---

introduce random probe baseline as the control task (Hewitt and Liang, 2019) and instead of look- ing at absolute probing values, we interpret the probing scores compared to the score of the ran- dom baseline. The probe scores are defined as:
or A cross (LLaMA). simp
GPT-2 on k-tth smallest element (Asimp). We consider two versions of GPT-2: a pretrained ver- sion and a finetuned version (GPT-2rr). We report our two probing scores (Eq. and Eq. 2) with dif- ferent in Table The unnormalized F1-macro scores (i.e., SFI ( and SFI Aiim)) can be found in Appendix B.2.
Spi SFI (V) Asimp)
1 - SFI ( (VI Aamit)
Sp2 = SFI VG A simp) SFI(G]V A
(2)
- (F Arand)
Table I:
Probing scores for GPT-2 models on syn- thetic reasoning tasks with different k. We also provide the test accuracy of these finetuned models for refer- ence. Note that when k = 1, the depth of G is 0, and SFI (GIV, Ammp)) always equal 1. Thus, we leave these results blank. Results show that can clearly detect G from attentions of GPT-2pT.
where Arand is the simplified attention matrix given by randomly initialized LM. After normaliza- tion, we have the range of our probing scores: SpI, Sp2 E [0, 1]. Small values mean that there is no useful information about G in attention, and large values mean that the attention patterns indeed contain much information about G.
Test Aoc.
Spi
Spa
GPT-2
GPT-2m
GPT-2
GPT-2rr
GPT-2
GPT-2p
4 Mechanistic Probing of LMs
99.63
7.09
92.94
We use the probe to analyze how LMs perform the reasoning tasks. We first verify the usefulness of at- tentions in understanding the reasoning process of LMs by visualizing Asimp ($4.1). Then, we use the probe quantify the information of G contained in Asimp (84.2). Finally, we report layer-wise probing results to understand if LMs are reasoning proce- durally across their architecture ($4.3).
2
99.42
5.88
93.71
20.65
98.05
3
98.23
13.48
91.62
13.59
95.76
4
94.89
20.73
88.34
8.04
92.32
0.00
5
93.38
24.15
87.54
13.81
88,17
6
92.62
24.82
86.89
11.18
95.06
7
91.37
25.87
76.61
<1
91.56
8
91.29
22.30
78.73
1.06
93.93
These results show that without finetuning, GPT- 2 is incapable of solving the reasoning task, and we can detect little information about G from GPT- 2's attentions. has high test accuracy the reasoning task, MechanisticProbe can casily recover the reason- ing tree G from Asimp- This further confirms that GPT-2FT solves this synthetic reasoning task fol- lowing G in Figure 2 (left).
4.1 A Attention Visualization
8 However, for GPT-2FT, which
We first analyze Asmp on the k-th smallest element task via visualizations. We permute Asipp arrang- ing the numbers in ascending order and denote this permulation as 7 (Asimp). We show visualiza- tions of IE[T( [| on the test data in Figure 5. We observe that when GPT-2 tries to find the k-th smallest number, the prediction token first focuses on top-k numbers in the list with bottom layers. Then, the correct answer found in the top layers. These findings suggest that GPT-2 solves the rea- soning task in two steps following the reasoning tree G. We further provide empirical evidence in Appendix B.4 to show that amalysis on Asimp gives similar conclusions compared to that on A.
LLaMA on ProofWriter and ARC (A (rmee)) ) Similarly, we use MechanisticProbe to probe LLaMA on the two natural language reasoning tasks. For efficiency, we randomly sampled 1024 examples from the test sets for our analysis. When depth = O, LLaMA only needs to find out the use- ful statements for reasoning (Sp1). When depth=1, LLaMA needs determine the next reasoning step. Probing results with different numbers of in- put statements (i.e., (S) are in Table 2. Unnor- malized classification scores can be found in Ap- pendix C.5. It can be observed that all the prob- ing scores are much larger than 0, meaning that the attentions indeed contain information about G.
4.2 Probing Scores
Next, we use our MechanisticProbe to quantify the information of G contained in Asimp (GPT-2)
We have additional empirical explorations in the Ap- pendix. Appendix B.6 shows that different finetuning methods would not influence probing results. Appendix B.7 explores how the reasoning task difficulty and LM capacity influence the LM performance. Appendix B.8 researches the relation- ships between GPT-2 performance and our two probing scores.
VVisualization of E (T( Ar(A) for (pretrained) GPT-2 can he found in Appendix B.3. It shows that GPT-2 can somehow find out the largest number from the number list.

---

1990
(a) ==1
(b) A:
(c) =
(d) 1:
14-
120
um
oo-
um
um
5
&
(g) 7
(h)
k
Figure 5: Visualization of (Ammp) for kt ranging from to 8. Note that 16 here, and when &: 8, the reasoning task becomes finding k-th largest number. The visualizations are roughly the same without loss of generality. The r-axis represents the layer index and the g-axis represents the size (rank) of the number. The value of each cube is the attention weight (better view in color). Results show that layers of finetuned GPT-2 have different functions: bottom layers focusing on top-k numbers and top layers focusing on k-smallest numbers.
Table Probing results LLaMA on natural language
statements from attentions easily.
reasoning. Note that when depth 0, SF1 (G V, Asimp) is always equal to I. When depth = and |S| = 2, imput statements are useful. SFI (V (A)))) is always equal to 1. We leave these results blank. Regarding ARC, we only report the scores under the 4-shot set- ting. Results show that attentions of LLaMA contain some information about selecting useful statements for ProefWriter but much information for ARC. Regarding determining reasoning steps, attentions of LLaMMA con- tain much information for both ProofWriter and ARC.
Looking at Sp2, we notice that our probe can easily determine the height of useful statements based on attentions on both ProofWriter and ARC datasets. By comparing the probing scores on ProofWriter, we find that LLaMAFT always has higher prohing scores than 4-shot LLaMA, imply- ing that finetuning with supervised signals makes the LM follow the reasoning tree more clearly. We also notice that 4-shot LLaMA affected more by the number of useless statements than LLaMAFT. indicating lack of robustness of rea- soning in the sho- setting.
ProofWriter
Depth
Spn
LLaMA LLaMApT
ILMA
LLaMAyr
LLAMA LLaMArr
81.72
57.21
49.08
92.12
44.83
48.14
91.34
96
4.3
Layer-wise Probing
92.19
27.39
40.09
83.75
96.11
100
12.70
77.38
89:55
17.18
21.07
77.85
89.31
After showing that LMs perform reasoning follow- ing oracle reasoning trees, we investigate how this reasoning happens inside the LM layer-by-layer.
85.38
11.10
13.84
79.99
91.1
17.33
90.32
94.42
07.49
61.73
55:41
96.49
$5.00
GPT-2 on k-th smallest element. In order to use our probe layer-by-layer, we define the set of simplified attentions before layer as Asimp( (ii)p((')) I). Then, report our prob- ing scores (p) and Sp) on these partial atten- tions from layer to layer 12. We denote these layer-wise probing scores as (FI(V) Asimp(: 1)) and Sp (G) V, Asimp:1)).
Looking Spı ProofWriter, when the number of input statements is small, Mechanistin crrobe can clearly decide the useful statements based on attentions. However, it becomes harder when there are more useless statements (i.c., |S] large). How- ever, for ARC, our probe can always detect useful

---

Prohon SOIES
Proking Pokingsores sores
Wionkar
0.8
0.8
0.5
0.5
0.8
D:4
D.G
0:2
MBd
0:2
0-4
D.G
20
20
20
Tayer number
Tayer number
Layer number
Layer index
Layer indox
Layer Index
(a) k 1
(b) 1: = 2
(c) l: =
(a) [S] =
(b)
S]
8
(c) |S| = 12
scotss
Probing scoes
0.8
0.8
Spx (4) Spx
0.4
0.0
0.2
0.4
0.2
0 5 10 15 20
0 5 10
15 20
10
Layer number
Layer number
Layer incess
Layer index
Layer index
(d)
[S]
16
(c)
|S| 20
(d) 1:
(e)
(1) l:
0.8
Sp (1a(1)
Pathing
0-6
0.4
U.5
0.6
beight beight
0.2
0.53
10
5
10
0.5
Layer index
Layer inders
5
10
15
20
&
(h) 1: 8
Layer
(i) Height probing
Figure 6: Two probing scores partial attentions across 12 layers. The æ-axis represents the layer index of GPT- 2pr, and the y-axis represents the probing scores. Simi- larly, we test different & from I to 8 (m = 16 by default). Results show that each layer of GPT-2pr focuses on dif- ferent steps of the synthetic reasoning task.
Figure 7: Layer-wise probing results ProofWriter. r-axis represents the number of layers used for prob- ing. We can find that Spi scores reach the plateau quickly and (y) scores increase smoothly till the mid- dle layers. This indicates that useful statement selection is mainly finished in the bottom layers and the reason- ing step is decided in the middle layers. Results of SFi Vineigull A Ammp (: f)) shows that statements at height 0 are identified by LLaMA in bottom layers, and state- ment at height are identified later in middle layers.
Figure 6 shows the layer-wise probing scores for each & for GPT-2FT models. Observing Sp1, i.e., selecting top-k numbers, we notice that GPT- 2FT quickly achieves high scores in initial layers and then, Spi increases gradually. Observing Sp2. i.e. selecting the k-th smallest number from top-k numbers, we notice that GPT-2FT does not achieve high scores until layer 10. This reveals how GPT- 2FT solves the task internally. The bottom layers find out the top-k numbers, and the top layers select the k-th smallest number among them. Results in Figure also support the above findings.
Probing scores
Sen (
Spa(O)
10 15
Layar
Figure 8: Layer-wise probing results ARC (depth=1). 1). Similar to that of ProofWriter, we can find that (I) scores reach the plateau quickly and Sp)()) scores in- crease smoothly till the middle layers.
LLaMA on ProofWriter and ARC (A Similarly, we report layer-wise probing scores (P) and Sp)(l) for LLaMA under the 4-shot set- ting. We further report SF1 (Vheightl Ammp (: ()) for nodes with height- 0 and height = 1 to show if the statement at height 0 is processed in LLaMA before the statement at height 1.
simp )
Then, it focuses on predicting the reasoning steps given the useful statements. Figure 7(f) shows how SpI (Vieighu A smmp 1)) varies across layers. LLaMA identifies the useful statements from all the input statements (height 0) immediately in layer 2. Then, LLaMA gradually focuses on these state- ments and builds the next layer of the reasoning tree (height 1) in the middle layers.
The layer-wise probing results for ProofWriter are shown in Figure 7(a-e). We find that similar to GPT-2, probing results for 4-shot LLaMA reach a plateau at an early layer (layer 2 for Spi(1) (I)). and middle layers for Sp2(1). This obserrva- tion holds as we vary|S] from to 20. This shows that similar to GPT-2, LLaMA first tries to identify useful statements in the bottom layers.
The layer-wise probing results for ARC are in Figure 8. Similar to ProofWriter, the Spi scores on ARC for 1-shot LLaMA reach a plateau an early layer (layer 2), and at middle layers for Sp2(1). This also shows that LLaMA tries to identify useful statements in the bottom layers and focuses on the next reasoning steps in the higher layers.

---

5 Do LMs Reason Using Asimp?
0.6
0.5
0.5
Our analysis so far shows that LMs encode the reasoning trees in their attentions. However, as argued by Ravichander et al. (2021); Lasri et al. (2022); Elazar et al. (2021), this information might be accidentally encoded but not actually used by
50
50
50
Pruned heads
Praned beads (5)
Pruned heads (R))
(a) A: 1
(b) h: 2
(c) &
3
the LM for inference. Thus, we design a causal analysis for GPT-2 on the k-th smallest element task to show that LMs indeed perform reasoning following the reasoning tree. The key idea is to prove that the attention heads that contribute to Asimp are useful to solve the reasoning task, while those heads that are irrelevant (i.e., independent) to simp are not useful in the reasoning task.
0.6
0.5
0.5
50
G0
50
Praned heads (16)
Praned beacs (%)
Pruned heads (5R)
(d) & = 4
(e)
A
5
(f) A
6
Test aosuras
Random
Podition estogy
Sinee enimury
D 50
0
50
Proned heads (1)
Pruned heads (8))
Intuitively, for the k-th smallest element task, attention heads that are sensitive to the number size (rank) are useful, while heads that are sensitive to the input position are not useful. Therefore, for each head, we calculate the attention distribution on the test data to see if the head specially focuses on numbers with particular size or position. We use the entropy of the attention distribution measure this": small entropy means that the head focuses particularly on some numbers. We call the entropy with respect to number size as size entropy and that with respect to input position as position entropy. The entropy of all the heads in terms of number size and position can be found in Appendix B.5.
(g) h = 7
(h) A
Figure 9: Test accuracy on the reasoning task for differ- ent pruning rates. The r-axis represents pruning rates: ranging from 0% to 90%, and the y-axis represents the test accuracy. We prune heads in the ascending order of their size/position entropy. Results show that heads with small size entropy are essential to the test accuracy while those with small position entropy are useless.
size and is independent of the number position. MechanisticProbe detects the information of G from attentions. Thus, our probing scores would be affected by the heads with small size entropy but would not be affected by heads with small po- sition entropy. Then, we can say that changing our probing scores (via pruning heads in terms of size entropy) would cause the test accuracy change. Therefore, we say that there a causal relationship between our probing scores and LM performance, and LMS perform reasoning following the reason- ing tree in Asimp-
We prune different kinds of attention heads in order of their corresponding entropy values and report the test accuracy on the pruned GPT-2. 10 Results with different it are shown in Figure 9. We find that the head with a small size entropy is essen- tial for solving the reasoning task. Dropping 10% of this kind of head, leads to a significant drop in performance on the reasoning task. The heads with small position entropy are highly redundant. Drop- ping 40% of the heads with small position entropy does not affect the test accuracy much. Especially when & 1, dropping 90% position heads could still promise a high test accuracy.
6 Correlating Probe Scores with Model Accuracy and Robustness
Our results show that LMs indeed reason mecha- nistically. But, is mechanistic reasoning necessary for LM performance robustness? We attempt to answer this question by associating the probing scores with the performance and robustness of LMs. Given that finetuned GPT-2 has a very high test ac- curacy on the synthetic task and LLaMA does not perform as well on ARC, we conduct our analysis mainly with LLaMA on the ProofWriter task.
These results show that heads with small size entropy are fairly important for GPT-2 to find k-th smallest number while those with small position entropy are useless for solving the task. Note that the reasoning tree G is defined on the input number
"We regard the attentions BB| B[A h.] as probabilities, and calculate the entropy of head at layer| using them. "We gradually prune heads, and the order of pruning heads is based on the size/position entropy. For example, for the size entropy pruning with 50% pruned heads, we remove attention heads that have the top 50% smallest size entropy.
A ccuracy. We randomly sample 64 to 128 exam- ples from the dataset and test 4-shot LLaMA on these examples. We calculate their test accuracy

---

Pearson correlation coefficient between our probing scores and test accuracy for 4-shot LLaMA.
Table 3:
Belinkov, 2019; DeRose et al., 2021; Bibal et al., 2022). A direct way for interpreting an attention- based model is to visualize attentions (Samaran et al., 2021; Chefer et al., 2021). But the irrelevant, redundant, and noisy information captured by at- tentions makes it hard to find meaningful patterns, Alternatively, accumulation attentions that quan- tify how information flows across tokens can be used for interpretation (Abnar and Zuidema, 2020; Eberle al., 2022). However, for casual LMs, in- formation flows in one direction, and causes an over over-smoothing problem when the model deep." !2 To tackle this, other works propose new metrics and analyze attentions using them (Lthhyaaajh and Ju- rafsky, 2021; Liu et al., 2022). However, these metrics are proposed under specific scenarios, and these are not useful for detecting the reasoning process in LMs. We address this challenge by de- signing more structured probe that predicts the reasoning tree LMs.
Pearson corvelation coefficient p (x100%)
(Sp1, Spz)
0.01
P (Test accuracy, Spt)
27.42
(Test accuracy, Spz)
71.13
and the two probing scores Sp1 and Sp2. We repeat this experiment 2048 times. Then, we calculate the correlation between the probing scores and test accuracy. From Table 3, we find that test accuracy is closely correlated with Sp2. This implies that when we can successfully detect reasoning steps of useful statements from LM's attentions, the model is more likely to produce a correct prediction.
Following the same setting, also associate the probing scores with LM robustness. In order to quantify model robustness, we randomly corrupt one useless input statement for each ex- ample, such that the prediction would remain un- changed." 11 We measure robustness by the decrease in test accuracy after the corruption.
Robustness.
Mechanistic Interpretability. Mechanistic inter- pretation explains how LMs work by reverse en- gineering, i.e., reconstructing LMs with different components (Räuker et al., 2022). recent line of work provides interpretation focusing on the LM's weights and intermediate representations (Olah et al., 2017, 2018, 2020). Another line of work interprets LMS focusing on how the information is processed inside LMs (Olsson et al., 2022; Nanda et al., 2023). Inspired by them, Qiu et al. (2023) attempts to interpret how LMs perform reasoning. However, existing explorations do not cover the research problem we discussed.
That accunacy change
Test uocuracy decrosse
0.4
0.5
0.6
0.7
Sys (ninges
tomasimu))
Figure 10: Histogram for LM prediction robustness and Sp2. We measure robustness by the change in accuracy the corrupted dataset. We present results with & bins. Results show that LLaMA is more robust to noise on examples with higher probing scores Sp2-
Conclusion
Figure 10 shows that if Sp2 is small (less than 0.7), the prediction of LLaMA could be easily influ- enced by the noise (test accuracy decreases around 10%). However, if the probing Sp2 is high, LLaMA is more robust, i.c., more confident in its correct pre- diction (test accuracy increases around 4%). This provides evidence that if the LM encodes the gold reasoning trees, its predictions are more reliable (i.e., robust to noise in the input).
In this work, we raised the question of whether LMs solve procedural reasoning tasks step-by-step within their architecture. In order to answer this question, we designed a new probe that detects the oracle reasoning tree encoded in the LM architec- ture. We used the probe to analyze GPT-2 on a synthetic reasoning task and the LLaMA model on two natural language reasoning tasks. Our em- pirical results show that can often detect the information in the reasoning tree from the LM's attention patterns, lending support to the claim that LMs may indeed reasoning "mechanistically..
7 Related Work
A Attention- based analysis of LMs. Attention has been popularly used to interpret LMs (Vig and
I1 We do by adding negation: corrupting S; as !"That"o
12The issue is that all token representations are dominated by the first token. Detailed discussion is in Appendix A.
S: o is false"]. where means string concatenation.

---

Limitations
knowledge-critical subnetworks in pretrained lan- guage models.
One key limitation of this work is that we con- sidered fairly simple reasoning tasks. We invite future work to understand the mechanism behind LM-based reasoning by exploring more challeng- ing tasks. We list few other limitations of our work below:
Yonatan Belinkov, 2022. Probing classifiers:
Promises, shortcomings, and advances. Computational Linguis- rics, 48(1 48(1):207-219
Adrien Bibal, Rémi Cardon, David Alfter, Rodrigo Wilkens, Xiaoou Wang, Thomas François, and Patrick Watrin. 2022. Is attention explanation? introduction to the debate. In Proceedings of the 60th Annual Meeting of the Association for Compu- lational Linguistics (Volume I: Long Papers), pages 3889-3900, Dublin, Ireland. Association for Compu- tational Linguistics.
Mutli-head attention. In this work, most of our analysis takes the mean value of attentions across all heads. However, we should notice that attention heads could have different functions, especially when the LM is shallow but wide (e.g.. with many attention heads, and very high-dimensiona hidden states). Shallow models might still be able to solve procedural reasoning tasks within few layers, but the functions of the head could not be ignored.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss. Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. CoRR, abs/2005. 14165.
Auto-regressive reasoning tasks. In our analy- sis, formalize the reasoning task as classifica- tion, i.e., single-token prediction. Thus, the analy- sis could only be deployed on selected reasoning tasks. Some recent reasoning tasks are difficult and can only be solved by LMs via chain-of- ththugh prompting. We leave the analysis of reasoning un- der this setting for future work.
Gino Brunner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Wat- tenhofer. 2020. On identifiability transformers. 8th International Conference on Learning Represen- tations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net
A Acknowledgements
We are grateful to anonymous reviewers for their insightful comments and suggestions. Yifan Hou is supported by the Swiss Data Science Center PhD Grant (P22-05) and Alessandro Stolfo is supported by armasuisse Science and Technology through a CYD Doctoral Fellowship. Antoine Bosselut gratefully acknowledges the support of the Swiss National Science Foundation (No. 215390), Inno- suisse (PFFS-21-20) the EPFL Science Seed Fund, the EPFL Center for Imaging, Sony Group Corpo- ration, and the Allen Institute for AI. Mrinmaya Sachan acknowledges support from the Swiss Na- tional Science Foundation (Project No. 197155). Responsible AI grant by the Haslerstiftung; and an ETH Grant (ETH-19 21-1).
Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian Tramer, and Chiyuan Zhang, 2022. Quantifying memorization across neural lan- guage models. CoRR, abs/2202.076466.
Hila Chefer, Shir Gur, and Lior Wolf. 2021. Generic attention-model explainability for interpreting bi- modal and encoder-decoder transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021, pages 387-396. IEEE.
Zeming Chen, Gail Weiss, Eric Mitchell, Asli Celiky- ilmaz, and Antoine Bosselut. 2023. RECKONING: reasoning through dynamic knowledge encoding. CoRR, abs/2305.0623
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question an- swering? try arc, the A12 reasoning challenge. CORR, abs/1803.05457.
References
Samira Abnar and Willem Zuidema. 2020. Quantify- ing attention flow in transformers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4190-4197, On- line. Association for Computational Linguistics.
Antonia Creswell and Murray Shanahan. 2022. Faith- ful reasoning using large language models. CoRR, abs/2208.1422
Antonia Creswell, Murray Shanahan, and Irina Higgins. 2022. Selection-inference:: Exploiting large language models for interpretable logical reasoning. CoRR. abs/2205.0977 099102
Deniz Bayazit, Negar Foroutan, Zeming Chen, Gail Weiss, and Antoine Bosselut. 2023. Discovering

---

Joseph F. DeRose, Jiayao Wang, and Matthew Berger. 2021. Attention flows: Analyzing and comparing attention mechanisms in language models. JEFE Trans. Vis. Comput. Graph., 27(2):1100-117
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yu- taka Matsuo, and Yusuke Iwasawa. 2022. Large lan- guage models are zero-shot reasoners. In NeurlPS.
Karim Lasri, Tiago Pimentel, Alessandro Lenci, Thierry
Yue Dong, Chandra Bhagavatula, Ximing Lu, Jena D. Hwang, Antoine Bossclut, Jackie Chi Kit Cheung, and Yejin Choi. 2021. On-the-fly attention modula- tion for neural generation. In Findings of the Associ- ation for Computational Linguisties: ACL/JJCNPP 2021, Online Event, August 1-6, 2021, volume ACL/MCNLLP 2021 of Findings of ACL. pages 1261- 1274, Association for Computational Linguistics.
Poibeau, and Ryan Cotterell. 2022. Probing for the usage of grammatical number. In Proceedings of the 60th Annual Meeting of the Association for Compu- tational Linguistics (Volume I: Long Papers), pages 2818-8831, Dublin, Ireland. Association for Compu- tational Linguistics.
Yibing Liu, Haoliang Li, Yangyang Guo. Chenqi Kong- Jing Li, and Shiqi Wang. 2022. Rethinking attention- model explainability through faithfulness violation test. In International Conference on Machine Learn- ing, ICMI. 2022, 17-23 July 2022, Baltimore, Mary- land, USA, volume 162 of Proceedings of Machine Learning Research, pages 13807-13824 PMLR.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang. Bill Yuchen Lin, Peter West, Chan- dra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. 2023. Faith and fate: Limits of transformers on compo- sitionality. CoRR, abs/2305. 18654.
Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. In 7th International Conference on Learning Represemations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenRe- view.net.
Oliver Eberle, Stephanie Brandl, Jonas Pilot, and An- ders Søgaard. 2022. Do transformer models show similar attention patterns to task-specific human gaze? In Proceedings of the 60th Annual Meet- ing of the Association for Computational Linguistics (Volume I: Long Papers), pages -299,430,, Dublin, Ireland. Association for Computational Linguistics.
Donald W. Loveland. 1980. Automated theorem prov- ing. a logical basis. Journal of Symbulic Logic. 45(3):229-630.
Christopher D. Manning, Kevin Clark, John Hewitt, Urvashi Khandelwal, and Omer Levy. 2020. Emer- gent linguistic structure in artificial neural networks trained by self-supervision. Proc. Natl. Acad. Sci. USA. 1 117(48):300463000 30054..
Yanai Elazar, Shauli Ravfogel, Alon Jacovi, and Yoav Goldberg. 2021. Ammesic probing: Behavioral expla- nation with amnesic counterfactuals, Transactions of the Association for Computational Linguisties, 9:160- 175.
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick, 2023. Language models implement simple word2vec-style vector arithmetic.
Kawin Ethayarajh and Dan Jurafsky. 2021. Attention flows are shapley value explanations. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1 1ith International Joint Conference on Natural Language Processing (Volume 2: Short Papers), pages 49-54, Online, As- sociation for Computational Linguistics.
Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. 2023. Grokking of hierarchi- cal structure in vanilla transformers. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 439-448, Toronto, Canada. Association for Computational Linguistics.
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson, 2023. Dissecting recall of factual asso- ciations in auto-regressive language models. CoRR, abs/2304. 14767.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. 2023. Progress mea- sures for grokking via mechanistic interpretabillity In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. Open
John Hewitt and Percy Liang. 2019. Designing and in- terpreting probes with control tasks. In Proceedings of the 2019 Conference on Empirical Methods in Nut- ural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLL pages 3733-2745 Hong Kong, China. Association for Computational Linguistics.
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. 2020. Zoom in: An introduction to circuits. Distill, 5(3):c00024- 001.
Yifan Hou and Mrinmaya Sachan. 2021. Bird's eye: Probing for linguistic graph structures with a sim- ple information- thhorreti approach. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1 1th International Joint Conference Natural Language Processing (Volume I: Long Papers), pages 1844-1859, Online. Association for Computational Linguistics.
Chris Olah, Alexander Mordvintsev, and Ludwig Schu- bert. 2017. Feature visualization. Distill, 2(11)
Chris Olah, Arvind Satyanarayan, Ian Johnson, Shan Carter, Ludwig Schubert, Katherine Ye, and Alexan- der Mordvintsev. 2018. The building blocks of inter- pretability. Distill, 303)e10.

---

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. 2022. In-context learning and induction heads. arXiv preprint urXiv2209.1/892
vision-and-lanee lannnuun-- transformers. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1 11th International Joint Conference on Natural Language Processing: Student Research Workshop, pages 81-86, Online. Association for Computational Linguisties.
Linlu Qiu, Liwei Jiang, Ximing Lu, Melanie Sclar, Valentina Pyatkin, Chandra Bhagavatula, Bailin Wang, Yoon Kim, Yejin Choi, Nouha Dziri, et al. 2023. Phenomenal yet puzzling: Testing in- ductive reasoning capabilities of language mod- els with hypothesis refinement. arXiv preprint arXiv:2310.08559.
Alessandro Stolfo, Yomatan Belinkov, and Mrinmaya Sachan. 2023. Understanding arithmetic reasoning in language models using causal mediation analysis. arXiv preprint 2.55::2300555555.
Oyvind Tafjord, Bhavana Dalvi, and Peter Clark. 2021. ProofWriter Generating implications, proofs, and abductive statements over natural language. In Find- ings of the Association for Computational Linguis- tics: ACL-IICNLP 2021, pages 3621-3634 Online. Association for Computational Linguistics.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Nla Sutskever. 2019. Language models are unsupervised multitask learners. OpenAl blog.
Tilman Räuker. Anson Ho. Stephen Casper, and Dylan Hadfield- 2022. Toward transparent Al: A survey on interpreting the inner structures of deep neural networks. CoRR, abs/2207.13243.
Xiaojuan Tang, Zilong Zheng, Jiaqi Li, Fanxu Meng, Song-Chus Chun Zhu, Yitao Liang, and Muhan Zhang. 2023. Large language models are in-context semantic reasoners rather than symbolic reasoners. CoRR, 14822
Abhilasha Ravichander, Yonatan Belinkov, and Eduard Hovy. 2021. Probing the probing paradigm: Does prohing accuracy entail task relevance? In Proceed- ings of the 16th Conference of the European Chap- ter of the Association for Computational Linguisties: Main Volume, pages 3363-3377, Online. Association for Computational Linguistics.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière. Naman Goyal. Eric Hambro. Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models. CoRR, 0ts/2302..3377.
Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. 2022. Impact of pretraining term frequencies on few-shot numerical reasoning. In Findings of the Association for Computational Lin- guistics: EMNLP 2022, pages 840-854, Abu Dhabi, United Arab Emirates. Association for Computa- tional Linguistics.
Jesse Vig and Yonatan Belinkov. 2019. Analyzing the structure attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 63-76, Florence, Italy. As- sociation for Computational Linguistics.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sen- nrich, and Ivan Titov. 2019. Analyzing multi-head self-attention: Specialized heads do the heavy lift- ing, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pages 5797-5808, Florence, Italy. Association for Computational Linguistics.
Danilo Neves Ribeiro, Shen Wang, Xiaofei Ma, Henghui Zhu, Rui Dong, Deguang Kong, Juli- ette Burger, Anjelica Ramos, Zhiheng Huang, William Yang Wang, George Karypis, Bing Xiang, and Dan Roth. 2023. STREET: multi-task struc- tured reasoning and explanation benchmark. In The Eleventh International Conference on Learning Rep- resentations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net.
Jason Wei, Xuezhi Wang. Dale Schuurmans. Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou 2022. Chain-of-thought prompt- ing elicits reasoning in large language models. In NeurIPS.
Anna Rogers, Olga Kovaleva, and Anna Rumshisky. 2020. A primer in BERTology: What we know about how BERT works. Transactions of the Association for Computational Linguistics, 8:842-866.
Zhengxuan Wu, Atticus Geiger, Christopher Potts, and Noah D. Goodman. 2023. Interpretability at scale: Identifying causal mechanisms in alpaca. CoRR. Abs/2305.08809.
Sebastian Ruder, Jonas Pfeifferr and Ivan Vulić. 2022. Modular and parameter-efficient fine-tuning for NLP models. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing: Tutorial Abstracts, pages 23-29, Abu Dubai, UAE, Association for Computational Linguistics.
Shizhuo Dylan Zhang, Curt Tigges, Stella Biderman, Maxim Raginsky, and Talia Ringer. 2023. Can trans- formers learn to solve problems recursively? CoRR, abs/2305.14699
Jules Samaran, Noa Garcia, Mayu Otani, Chenhui Chu, and Yuta Nakashima. 2021. Attending self-attention: A case study of visually grounded supervision in

---

A Proof of The First-Token Domination
B Supplementary about GPT-2
Proof. Without loss of generality, assume the LM have L layers with only 1 attention head (h H 1), and the attention weight matrix layer A(I, h). We consider the input that have more than token. We model attention as flows following the setting of Abnar and Zuidema (2020). Then, the attention accumulation Accum() of the token z! in the layer + 1 can be written
We provide some more details on our experiments on GPT-2 as well as LLaMA to help in reproducibil- ity. First of all, we fix the random seed (42) and use the same random seed for all experiments, in- cluding LM finetuning and interpretations. In addi- tion, make sure the random seed is unbiased, we further re-run the same experiment with different random seeds. All our experiments have roughly the same results as those using other seeds. Second, we design our analysis as simply as possible to ensure that there is as little random influence (i.e., confounder) as possible. For MechanisticProbe, we select the kNN classifier. For LLaMA, we run analysis experiments on a 4- shot in-context learning setting.
Accum 2 +1)= =
Σ
(t)(1,h) Accum(z)) (z)),
15jSM
where we have
Accum(z)) (xcummzz)=
Σ
an)(1,h)+t . ty.
(4)
Third, we report as many intermediate and sup- plement results as possible. In the Appendix, there are many other interesting findings. However, due to the space limit, we cannot present them in the main paper. We hope our findings are helpful to the community to better understand LMs.
18JS
Since (1,(1, is the attention weight for casual LM, we have
{ 0< (aj <1 i2j Σ, (1,hh)
=0 if i i<j
(5)
B.1
GPT-2 Finetuning
The finetuning settings for all GPT-2 models are roughly identical. We generate 0.98 million se- quences of numbers as the training data and 10, 000 in data for validation and testing. Note that the col- lision probability is extremely small, thus we can assume that there is no data lcakage. The epoch number is set as 2, and the batch size here is 256. We use the AdamW (Loshchilov and Hutter, 2019) optimizer with weight decay 1e - 3 from Hugging- fà for finetuning. The learning rate is le
Note that attention is normalized by Softmax func- tion in LMs. The minimum attention weight is non-zero, and we assume there exist a constant 0 such as < (1, < 1 i j. Now we define the information ratio IR (tg) as the informa- tion of token tj stored in the hidden representation z/. Consider that in each layer, token t would propagate its information to all tokens the next layer with at least amount. Then, by tracing the information flow from other tokens j > 1, we have
B.2 Original Probing Scores
1 - IR (R) s (1 () 1R (4)).
(6)
The original classification scores (F1-Macro) for two probing tasks can be found in Table 4. Here, random means we randomly initialize GPT-2 and use its attentions for probing as the random base- line. The other two preirained and finetuned are the pretrained GPT-2 model and GPT-2 model after finetuning with supervised signals...
Using the chain rulc, we have
(7) 1 1- IR 12 (4) < (1 (1- (4)( (I)()))))(
which means
IR,: (t)) > 1 - (1 (1- -1- IR_ (t1),, (8)
B.3 Visualization of E[T [A((mA))p)) for GPT-2
We visualize the Asimp for pretrained GPT-2 with- out finetuning in Figure 11. We can find that even if pretrained GPT-2 cannot solve the synthetic rea- soning task (finding k-th smallest number from a list). It can still somehow differentiate the size of
IR. (t)( 1- (1-
With the inequality above, we know that if the LM is deep, i.c., L is large, we have IR,: (t1) increase exponentially in terms of layer L, which means that IR_. (t1) z 1 with large I in general.
hhttss//huggingtfcc..o//

---

Table The original classification F1-Macro scores of GPT-2 for probing tasks.
3-
Probing
(Aaeg)
(C])V, M Awsy)
GPT-2
random
pretrained
Gietuned
random
peotrained
finctuned
Token
a (ler))
2=
48.60
$1.61
96.77
73.18
78.72
99.47
10
k=3
47.62
$4.68
95.61
61.45
66.89
98.36
58
55.71
59.28
13
48,50
60.94
93.58
54.36
55.48
94.06
Am6
48.06
61.40
50.44
55,98
97.55
15
A=
19.28
62
51.75
51.11
97.55
6 Layer
7
8
9
10
31
32
Rm
49.74
60.95
89.31
50.54
$1,07
97.00
Figure 12: Visualization of A]. The a-axis represents layers and the y-axis represents input token positions. We take the mean pooling of 60% attention heads with large position entropy. The leaf node (i.e., the largest number) always position 8, and the root node (i.c., the second largest number) is always position 12. We track their attentions and visualize them by blue and red lines.
numbers. The largest number of the list often has slightly larger attentions in layer 11.
Layers
12
38075
38790
typical tendency.
1
2622
3-
Figure 11: Visualization of E[*(A,img) for pretrained GPT-2 (w/o any finetuning). Similarly, we take mean pooling for different attention heads.
Token
11
B.4 Visulization of A
To avoid disturbance, remove 10% position heads (heads with small position entropy) and take mean pooling on all left heads. visualize to directly show that Aimp contains sufficient essen- tial information of A. We consider a special case when & = 2, the largest number is at position 8 and the second largest number is at position 12. 14
14-
16
9
10
12
Layer
Figure 13: Visualization of E[A] on normal data. The æ-axis represents layers and the y-axis represents input token positions. We take mean pooling of 60% attention heads with large position entropy.
The attention E A[A] visualized as in Figure 12. From Figure 12, we can get similar conclusion as on the visualization of E. [imp] In the bottom layers, most hidden representations focus on top-2 numbers. In the top layers, most hidden represen- tations focus on the second smallest number. This result proves that our analysis on Asimp is as rea- sonable that on A.
B.5 Attention Head Entropy
From Figure 14, we can find that most heads belong to either position head or size head. Note that we generate input data randomly, thus, the number size and input position independent of each other. Thus, one head cannot be both position head and size head.
We also provide the visualization of E[A] on normal test data (i.e., the input position indepen- dent the size) in Figure 13. We can find that the attention distribution is even, and there is no
B.6 Do Finetuning Methods Matter?
To show that our analysis is robust, we explore the attention of GPT-2 with different ways param-
We modify our test data only to satisfy this condition.

---

Layers
sth
thm
sin-
(a) Entropy terms num-(b) Entropy in terms num- her size (rank)
ber input position
(a) Full
(b) Partial (Attention)
Figure 14: Entropy for all attention heads of finetuned GPT-22 z-axis represents layers and y-axis represents the head index. The cube (better view color) shows the entropy value.
Larerr
Layar
eter efficient finetuning (Ruder et al., 2022). We report our two probing scores Sp and Sp2 with ious ways of finetuming in Figure 15 (We consider the condition when k and m 16). We find full GPT-2
10
CM
(c) Partial (MI.P)
(d) Partial (Bias)
that probing scores of finetuning model are similar to that of partially finetuning on attention parameters and MLP (multilayer percep- tron) parameters. These consistent results ensure the general usage of our MechanisticPrre on current large LMs with partial finetuning.
Figure 16: Visualization of TB[ (Amm)) with different ways of finetuning. The x-axis represents layers and the y-axis represents size ranking. We take mean pooling attention heads. From Figures (a-d) GPT-2 with full finetuning, GPT-2 with partial finetuning on atten- tion parameters, GPT-2 with partial finetuning on MLP parameters, and GPT-2 with partial finetuning bias parameters. We can find that different finetuning meth- ods would not affect the attention patterns much.
Syi
Sez
Proking scoes
0:99997
0,9g1.90
1.08
terns that of full model finetuning. Regarding the bias tuning, it slightly different. We speculate that this is because the LM does not learn the k-th smallest element task well (with much lower test accuracy). Generally, without significant perfor- mance drops, these results can support the general usage of our probing method Mechani wtwnnnssttiPoob under different finetuning settings.
0.0
0.21
Attetion
MLP
Blas
Finetuning methods
Figure 15: Probing scores GPT-2 with different fine- tuning methods. Here, w/o denotes the baseline when GPT-2 model is not finetuned. Full denotes finetuning with parameters. Other models are partial finetuned with corresponding parameters. With the exception of probing scores of the Bias tuned model, other par- tial finetuning methods have roughly same probing scores compared to that full finetuning, This indicates that Mechanis MecchnisticPreed sticProbe can provide consistent analy- sis for LMs finetuned in different ways.
B.7 Task Difficulty Model Capacity
This subsection discusses which factors can let GFT-2 handle reasoning tasks with more leaf nodes in G, i.e., large k. We explore the model capacity and reasoning task difficulty. Specifically, we ex- tend the list length m from 16 to 64, and maximum k from 8 to 32. For each k. we finetune LMs with the same finetuning settings and evaluate them test data with accuracy. For model capacity, we compare three versions of GPT-2 with different sizes: Distilled GPT-2, GPT-2 (small), and GPT-2 Medium. For task difficulty, we construct synthetic data by selecting m 64 numbers from 256, 384,
The direct visualization of [Asipp] for dif- ferent finetuning methods can be found Figure 16(a- d). And the test accuracy of these finetuned mod- els are: 09.42, 99.31, 94.11, and 76.84. We can find that finetuning attention parameters or MLP parameters can obtain quite similar attention pat-

---

Till now, we have experimented on a relatively casy task (m 16). In this subsection, we in- crease the difficulty of the task by extending the input number list from m = 16 numbers to m 64 numbers. We report the test accuracy as well as the two probing scores in Figure 18, varying l from I to 32. As expected, the test accuracy decreases smoothly from near 100% to around 15%. Interest- ingly, SpI and Sp2 do not decrease with the same speed. Even when the model has very low ac- curacy, Spi still maintains at a high score (above 30%), while Sp2 quickly jumps to 0 when is around 20. This suggests that GPT-2pT solves the two steps sequentially, and step 2 fails first when the task goes beyond the capacity of the model.
Diviilld (IPT-2
frrm from numbers from mumbers
GPT-22 (Smal))
GPT-2 Medium
Teat uocuracy (R)
100
20
30
A
(a) Model capacity
(b) Procedural task difficulty
Test accuracies for LMs under different con- ditions. The z-axis represents (finding the k-th small- est number) and the y-axis represents the test accuracy. The left figure explores the performance of LMs with different capacities, and the right figure explores the performance of GPT-2 with different task difficulties.
Figure 17:
C Supplementary about LLaMA
C.1 Settings for 4-shot and Finetuned LLaMA
and 512 distinct numbers.
We report their test accuracies in Figure 17. Note that here the setting remains the same: for each reasoning task (i.e., k), we finetune an individual model. From Figure 17()), we find that LMs with large model capacities can better solve procedural tasks with more complex G (i.c., more leaf nodes in G). But it does not mean that small LMs fail in this case. If we can reasonably reduce the task difficulty (e.g., decompose the procedural task), small LMs are still able to handle that task with complex G (Figure 17(b)).
For the in-context learning of LLaMA, we con- struct the input prompt as simple possible. Given of statements [St, S2, ...]. the question state- ment Q, and the answer label A (e.g.. "True" or "False"), the prompt templates for ProofWriter and ARC are:
[S1, S2, . + [Q] + True or False? + [A],
[S1, S2. + [Q] + The answer is: + [A].
The test accuracy of 0-shot, 2-shot, 4-shot, and 8- shot prompting LLaMA can be found in Table 5, We select 4-shot in-context learning setting in our analysis due to its best performance.
B.8 What if Reasoning Tasks Become Harder?
Table 5: Test accuracy of LLaMA.
ProofWriter
0.0
Sy2
Acc. (%)
depth=0
depth=1
depth=2
depth=3
Protengoom
0-shot LLaMA
50.07
49.61
49.74
49.83
2-shot LLaMA
75.58
75.83
74.31
73.74
4-shot LLaMA
81.72
78.33
75.58
74.64
0.5
8-shot LLaMA
54.26
52.96
52.82
52.42
Finetuned I.LMMA
100
100
100
100
ARC
4-shot LLaMA
56.32
53.40
Figure 18: Prohing scores and test accuracy of GPT-2FT on more difficult reasoning tasks. We finetune GPT- 2 models to find the j-th smallest number from long number list (m 64, and ranges from to 32). Results show that when the accuracy is low, GPT-2pT would still know how to select top-k numbers more less. But they are unable to find the k-th smallest number from top-k numbers anymore.
Regarding the finetuning of LLaMA (i.e., par- tially finetuning on attention parameters), most set- tings are similar to that of GPT-2. The epoch num- ber is set as 2, and the batch size 256. We use the Adam W optimizer with weight decay 1e 5 for finetuning, and the warmup number is 500. The

---

learning rate is le 6. Test accuracy of finetuned models can be found in Table 5 as well.
can be found in Table 6. We can find that there are less than 2000 4/5-depth examples in test data.
C.2 Layer (Attention) Pruning
Table 7: Data statistics of annotated ARC in terms of different depth.
For the layer pruning, we use the greedy search strategy. Specifically, we remove all attentions in layers from top to bottom." 15 If the performance decrease on test data is small (less than 5% in to- tal), the attention in that layer is dropped. For 4-shot LLaMA on PriefWriter, 13 (out of 32) top layers removed, and 15 (out of 32) top layers are removed for ARC. For finetuned LLaMA ProofWriter, middle layers (layer and layer 13) and 16 top layers are removed. After removing all attentions in these layers, the performance de- creases are around 2% for both in-context learning and finetuning settings.
# of examples
Training
Development
Test
depth =
331
51
87
depth
380
64
95
depth 3
277
28
67
depth =
175
18
depth
150
26
40
ARC. We follow the original data split for train- ing, development, and test sets (Ribeiro et al., 2023). Note that the number of examples in ARC is quite small. Thus, in analysis, do not experiments only on test data. We simply merge all data for the analysis.
C.3 Statistics of Cleaned ProofWriter and Annotated ARC
Table Data statistics cleaned ProofWriter in terms of different depth.
C.4 Reasoning Tree Ambiguity Example
We consider simple case to explore if the rea- soning tree ambiguity issues happens in LLaMA. We sample 1024 examples whose annotations 2-depth reasoning trees are
of examples
Training
Development
Test
depth
84,568
12,227
24,270
depth 1
41,718
6,101
12,044
depth 2
25,021
3,712
7,215
S > S2 53 Q.
depth
14,042
2,079
4,132
We give real example from the dataset randomly to illustrate the issue of reasoning tree ambiguity. Consider the three statements of G (from 17 input statements) as
depth 4
6,078
891
1765
depth 5
5,998
874
1756
S1 Erin is cold;
ProofWriter. We follow the original data split for training, development, and test sets. However, the depth split of PrtefWriter is not suitable in our case. The original dataset only considers the largest depth of set of examples (with similar templates) for the split. It means for example in depth 5, there would be many of them with depth smaller than 5. In our case, we classify examples into types from depth to only based on the example's reasoning tree depth. After the depth split, we also remove examples whose reasoning trees have loops multiple annotations. Besides, we remove few examples whose depth annotations are wrong (e.g., annotated depth 5 but only with nodes in G). Statistics of the cleaned and re-split ProofWriter
S2 If someone is cold then they rough;
S3 If someone is rough then they white;
Q: Erin is white (True).
It is intuitive that following the annotated reasoning tree could obtain correct answer. However, there other ways to answer the question. We can first combine S and get a new statement S4 as
S4 If someone is cold then they are white.
Then, the reasoning tree becomes
S S3 S - Q, and we can rewrite it with brackets as
15 Transformer layers without attention degrade to MLPs.
S (S2 > Sz) - Q.

---

There are multiple ways do reasoning for this example, and we do not know which one the LM uses. Thus, in this work, we ignore these kinds of examples with reasoning tree depth larger than 16
C.5 Original Probing Scores
Table The original classification FI-Macro scores LLaMA for probing tasks.
ProofWriter
Probing
(Akee)
Sey[C Aaer)
LLaMA
random
6-shot
finetuned
mandom
4-shot
finetuned
depth
[8] 1
48.10
77.79
33,59
=
63:31
73.36
78.00
54.07
92.03
98.39
depth
S
$1.64
64.33
67.45
49.75
58.73
96.71
$7.65
8
24
18.40
53.25
57.34
51.77
90.51
97.31
depth-1
51.21
98.77
48.73
90.38
depth=2
30.36
98.28
20.38
76.83
The original classification scores (FI-Macro) for two probing tasks on LLaMA can be found Ta- ble 8. Here, random means we randomly initialize LLaMA and use its attentions for probing as the random baseline.
in ProofWriter, there are only 30% examples that have reasoning trees with depth larger than 1.


Reviews:
[{"review_id": "3N8MA8UbJh", "reviewer": "Reviewer_jwyS", "summary": "", "questions": "1. It's strange that one model is only evaluated on one dataset.  Is there a particular reason that makes the authors do this?\n2. It's extremely misleading that the authors call proofwriter a real world task, which really isn't.  Proofwriter is synthetically generated and not even in purely natural language.\n3. I'm skeptical of how well this works in reality.  Based on the evidence shown in the two papers \"Attention is not Explanation\" and \"Attention is not not Explanation\", attention cannot be served as explanation, especially for real language datasets.  The authors should consider add hotpotqa and musique, which are procedural reasoning datasets and in natural language.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "*We thank the reviewer for recognizing the motivation of our work. Targeting an interesting question, we propose a new attention-based analysis method for LLMs, and provide new insights about how LLMs procedurally perform multi-step reasoning step-by-step within their architecture. However, we would like to respectfully point out that there may be a few misunderstandings about our method and experiment design. Responses are given below point by point.*\n\n---\n\n**Point 1 & Question 1: The experiment setting is a bit strange.**\n    \n    The experiments are limited, … And only one model is tested for one dataset.\n    It's strange that one model is only evaluated on one dataset. Is there a particular reason that makes the authors do this?\n\n**Response:**\n>Regarding the k-th smallest number task, we did experiments on other models with different model capacities: GPT-2 medium, and Distilled GPT-2 (see Appendix B.5). And similar results were obtained. We will add the results of LLaMA in the camera-ready version. \n\n>For the ProofWriter task, in order to obtain meaningful analysis, the LLMs need to have a reasonable performance on the task. At the same time, due to the limitation of computing resources, we can only afford to conduct our experiment with LLMs with moderate numbers of parameters.    \n>We have experimented on GPT-J (6B) and GPT-Neo (2.7B). However, their test accuracies are \\~50% (GPT-J w/ in-context learning) and \\~60% (GPT-J w/ finetuning); \\~50% (GPT-Neo w/ in-context learning) and \\~60% (GPT-Neo w/ finetuning), which is only slightly better than random guess (\\~50%). LLaMA 7B is the smallest model that can achieve >80% test accuracy under the in-context learning setting (see Table 2 in the paper). \n\n---\n\n\n**Point 1 & Question 2 & Question 3: Experiment datasets are not real-world data.**\n\n    For datasets, k-th smallest number task is too simplified, and proofwriter is a totally synthetically generated dataset.\n    It's extremely misleading that the authors call proofwriter a real world task, which really isn't. Proofwriter is synthetically generated and not even in purely natural language.\n    The authors should consider add hotpotqa and musique, which are procedural reasoning datasets and in natural language.\n\n**Response:** \n>In our experiments, we first proposed a synthetic task to give readers preliminary insights into how LLMs perform reasoning. Then, we introduced a more practical reasoning task. We think ProofWriter is a suitable task for our analysis for the following reasons:    \n>>1. We agree that “real-world task” might not be accurate for ProofWriter and one could argue that the sentences are not “natural” enough, we will rephrase this in the paper to avoid misunderstanding. However, the task still requires logical reasoning and an understanding of certain concepts in human language.     \nFor example, *“Charlie is blue; If someone is blue then they are kind -> Charlie is kind.”*. The LLMs need to know that Charlie falls within the scope of “someone” and understand the if-then statement.\n2. Reasoning analysis on synthetic data (i.e., ProofWriter) is commonplace and sometimes desirable because it allows us to study some phenomena not easy to do in natural text [2]. For instance, ProofWriter comes with the annotation of reasoning trees, which provides a ground truth to be compared with the extracted tree from LLMs. The lack of such reasoning trees in many real-world datasets (e.g., HotpotQA) makes it inadequate for such analysis. \n3. Various real-world datasets could have already been used for pretraining LLMs (see Tables 1&9 in [1]). For example, both HotpotQA and MuSiQue are based on Wikipedia. This could cause data leakage and LLMs would learn shortcuts to solve these multi-step reasoning tasks [3]. This makes them unsuitable to be used to answer our main research questions.\n\n[1] Longpre, Shayne, et al. \"A Pretrainer's Guide to Training Data: Measuring the Effects of Data Age, Domain Coverage, Quality, & Toxicity.\" arXiv preprint arXiv:2305.13169 (2023).    \n[2] Saparov, Abulhair, and He He. \"Language models are greedy reasoners: A systematic formal analysis of chain-of-thought.\" arXiv preprint arXiv:2210.01240 (2022).    \n[3] Trivedi, Harsh, et al. \"MuSiQue: Multihop Questions via Single-hop Question Composition.\" Transactions of the Association for Computational Linguistics 10 (2022): 539-554.\n\n---\n\n**Point 2 & Question 3: Using attention for interpretation could be problematic.** \n\n    The probing approach's reliance on attention patterns may raise concerns about its ability to provide a definitive mechanism to interpret LLMs.\n    I'm skeptical of how well this works in reality. Based on the evidence shown in the two papers \"Attention is not Explanation\" and \"Attention is not not Explanation\", attention cannot be served as explanation, especially for real language datasets. \n\n**Response:** \n>Thanks for the references, we will add them in the paper. First of all, we would like to respectfully point out that the two references only discussed the attention mechanism of RNN/LSTM. Analyzing transformers (“attention is all you need”) via attention is popular and many attention-based interpretation methods/findings have been proposed in recent years.\n\n>We understand the concern that the attention mechanism could fail in model interpretation. However, as said in [4], we do not think [5] disproves the usefulness of attention mechanisms for explainability. More specifically, [5] claims “attention distributions often have no effect on model outputs, while we include two experiments (Section 4.3 & Appendix B.4) to verify that the attention patterns do signify the model’s reasoning ability:\n>>1. Section 4.3. We explore the relationships between our proposed probing scores and an LLM’s performance. Results imply that when we can successfully detect reasoning steps of useful statements from LLM’s attention, the model is more likely to produce a correct prediction.    \n2. Appendix B.4. We classify attention heads into two groups: those focusing on number size (called size head) and others. The size heads contribute to the attention map for probing, while other heads can be regarded as noise for our analysis. We find that dropping size heads (less than 10% heads) has a much bigger impact than dropping other heads (>40% heads) (Figure 12). This supports the causation relationship between attention patterns and test accuracy. \n\n[4] Wiegreffe, Sarah, and Yuval Pinter. \"Attention is not not explanation.\" arXiv preprint arXiv:1908.04626 (2019).    \n[5] Jain, Sarthak, and Byron C. Wallace. \"Attention is not explanation.\" arXiv preprint arXiv:1902.10186 (2019).\n\n---\n\n**Point 3: Scalability.**\n \n    Unclear how this method scales to the actual large models people use these days\n\n**Response:** \n>The models we experiment on have a wide range of applications in real-world scenarios. Both GPT-2 and LLaMA have been widely used in both academia and industry.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The problem studied is interesting and important.\n2. Some interesting analysis has been done.", "reasons_to_reject": "1. The experiments are limited, w.r.t. both datasets and models tested.  For datasets, k-th smallest number task is too simplified, and proofwriter is a totally synthetically generated dataset.  And only one model is tested for one dataset.\n2. Although the authors demonstrate that attention patterns can help recover reasoning trees, the interpretability of attention patterns is a subject of ongoing debate. (See questions) The probing approach's reliance on attention patterns may raise concerns about its ability to provide a definitive mechanism to interpret LLMs.\n3. Unclear how this method scales to the actual large models people use these days", "questions_to_authors": "1. It's strange that one model is only evaluated on one dataset.  Is there a particular reason that makes the authors do this?\n2. It's extremely misleading that the authors call proofwriter a real world task, which really isn't.  Proofwriter is synthetically generated and not even in purely natural language.\n3. I'm skeptical of how well this works in reality.  Based on the evidence shown in the two papers \"Attention is not Explanation\" and \"Attention is not not Explanation\", attention cannot be served as explanation, especially for real language datasets.  The authors should consider add hotpotqa and musique, which are procedural reasoning datasets and in natural language.", "paper_topic_and_main_contributions": "Summary: The paper investigates if large language models (LLMs) such as GPT-2 and LLaMA perform multi-step reasoning tasks by memorizing answers or through a procedural reasoning mechanism. The authors propose a new probing approach, MechanisticProbe, that recovers reasoning trees from the model's attention patterns for a given problem. They test the probe on a synthetic task (k-smallest element) and a real-world task (ProofWriter) and find that reasoning trees can be largely recovered, suggesting that LLMs go through multi-step reasoning within their architecture. Higher probing scores correlate with better performance and noise tolerance, highlighting the importance of capturing the correct reasoning process.\n\nContributions:\n1. The proposal of a new probing approach, MechanisticProbe, which recovers reasoning trees from LLMs' attention patterns, helping to gain insights into the multi-step reasoning process of these models.\n2. The successful application of MechanisticProbe on GPT-2 and LLaMA, on synthetic and real-world tasks, providing evidence to suggest that LLMs indeed go through a multi-step reasoning process within their architecture.\n3. The discovery that higher probing scores correlate with better model performance and robustness against noise, implying the importance of capturing the correct reasoning process for LLM accuracy and robustness.", "excitement": 4, "reproducibility": 2, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "* The authors should cite \"Attention is not Explanation\" and \"Attention is not not Explanation\""}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "1. The authors don't need to repeatedly cite same papers in one paragraph.  This makes the paper hard to read.  For example, not sure the connection is tight to Wei et al., 2022; Kojima et al., 2022 on line 32-35?"}}, {"review_id": "C8HRhQt8Ow", "reviewer": "Reviewer_3L4M", "summary": "", "questions": "A. Is there a threshold for probing scores that indicate the tree exists? The S_p1 scores for the two experiments are quite different.\n\nB. Do the two subtasks (sec. 2.3) represent a two-step solution to reasoning tasks? The first step is to select related statements, and the second step is to reach the conclusion given these statements.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "*We thank the reviewer for recognizing the novelty and contribution of our work. According to our knowledge, we are the first to propose an attention-based method for a mechanistic understanding of the LLM’s reasoning ability. We present comprehensive experiments verifying our analysis method and provide new insights for understanding the multi-step reasoning ability of LLMs. Responses to proposed concerns are given below point by point.*\n\n---\n\n**Point 1:** \n\n    The probing task lacks general applicability. It is based on several assumptions, like the task format (footnote 2), tree structure (line 147-149), and node dependence (footnote 5). Even for the two tasks presented, their problem setups are very different.\n\n**Response:** \n>We presented two tasks that have the same task format: classification. For the synthetic task, each number is a label, and LMs do classification across all tokens in the vocabulary (for simplicity). For the ProofWriter, there are only two labels (true and false). LMs do classification across these two label tokens.    \n>>1. *Task format:* A variety of reasoning tasks can be formalized as multiple-choice questions (e..g, MMLU, GSM8K), which can also be easily adapted into our format.        \n>>2. *Tree structure:* Even if the reasoning structure is very complex, most of them in practice are chain-like. We consider reasoning structures as simple trees, which cover most reasoning examples in reality. Specifically, most the reasoning examples that can be solved by current popular chain-of-thought prompting methods are chain-like.    \n>>3. *Node dependence:* This assumption is not necessary. We predict node depths independently because it suffices to recover the reasoning tree. Otherwise, there are many existing tools available [1] that consider the node dependence for extracting trees from the attention patterns. Given that our work is the first such attempt, we make several simplifying assumptions.\n\n[1] Hewitt, John, and Christopher D. Manning. \"A structural probe for finding syntax in word representations.\" Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers). 2019.\n\n---\n\n**Point 2 and Point 3:**\n\n    Given its small size and early training methods, GPT-2 may not qualify as a LLM that could do multi-step reasoning. The GPT-2 experiment, despite promising results, does not strongly support the claims.\n    How to compute F1-Macro scores (line 160-161) is not discussed.\n\n**Response:** \n>We will change the description (LLM to LM) to avoid misunderstanding, and add more details about F1-Macro score computation in the paper. Our approach is agnostic to model size.\n\n---\n\n**Question A:** \n\n    Is there a threshold for probing scores that indicate the tree exists? \n\n**Response:** \n>We use the random baseline as a threshold (i.e., S_p1/S_p2=0). We also computed the p-value using a permutation test. The p-value is always is 0, so we didn’t include it in the results. \n\n---\n\n**Question B:** \n\n    Do the two subtasks (sec. 2.3) represent a two-step solution to reasoning tasks?...\n\n**Response:** \n>Yes. The first subtask/step selects the k leaves in the reasoning tree, and the second subtask/step selects the root from the leaves.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The idea of discovering trees from the attention of LLMs is novel and timely. \n2. This paper presents comprehensive experiments and analyses to verify such trees, and show their impact on performance and robustness.", "reasons_to_reject": "1. The probing task lacks general applicability. It is based on several assumptions, like the task format (footnote 2), tree structure (line 147-149), and node dependence (footnote 5). Even for the two tasks presented, their problem setups are very different.\n2. Given its small size and early training methods, GPT-2 may not qualify as a LLM that could do multi-step reasoning. The GPT-2 experiment, despite promising results, does not strongly support the claims.\n3. How to compute F1-Macro scores (line 160-161) is not discussed.", "questions_to_authors": "A. Is there a threshold for probing scores that indicate the tree exists? The S_p1 scores for the two experiments are quite different.\n\nB. Do the two subtasks (sec. 2.3) represent a two-step solution to reasoning tasks? The first step is to select related statements, and the second step is to reach the conclusion given these statements.", "paper_topic_and_main_contributions": "This paper analyzes multi-step reasoning of LLM by introducing a mechanistic probing task. LLMs perform well on multi-step reasoning tasks, but it is unclear whether they are truly reasoning or just using memorized content. This paper assumes that LLMs are indeed reasoning, and the reasoning progress can be expressed as a tree. Therefore, the paper presents a probing task that predicts the tree given the attention of the model. The task is further simplified and produces two scores that represent the likelihood of the tree. Experiments on GPT-2 and LLaMA suggest the existence of such trees. The LLaMA experiment also shows higher scores are correlated with better performance and robustness.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "Line 94-97: Incomplete sentence."}, "extra_text": {}}, {"review_id": "kF8xQWGLm2", "reviewer": "Reviewer_F7aK", "summary": "", "questions": "A. Can the authors provide some additional explanation/hypothesis as to why the model fails on longer inputs, despite being fine-tuned on them? Is this length beyond the sequence length that the model has been trained on? \n\nB. In the kth-smallest element experiment, it is unclear if the gains in performance/attention patterns are learned exclusively during finetuning. Would finetuning a randomly initialized GPT-2 achieve similar results?\n\nC. Can the authors expand more on the kNN classification of tokens/sentences as useful (similarly for the height classification)? Is the training data being used to identify the neighbors?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "*We thank the reviewer for recognizing the motivation and contribution of our work. Understanding the inner workings of LLMs when they do reasoning is very important at this time in the development of language models. To provide reliable insights, we carefully design experiments and try to be as comprehensive as possible. Concerns are responded to below point by point.*\n\n---\n\n**Point 1: Scalability**\n \n    …and it is unclear how this method can be scaled to more complex reasoning tasks.\n\n**Response:** \n>Our method can be applied to many more complex reasoning tasks as long as they can be formalized as multiple-choice questions (e..g, MMLU, GSM8K), and the reasoning tree is chain-like (most of the reasoning examples that can be solved by current popular chain-of-thought prompting methods are chain-like.).\n\n---\n\n**Point 2: Correlation does not imply causation.**\n\n    Attributing model behavior to attention is not always reliable…\n\n    Additionally, correlation does not imply causation. As the authors have pointed out, just because attention encodes the correct reasoning tree, it does not mean that these trees are being used by the model, regardless of its correlation with the probing scores. I think a causal intervention approach would instill more confidence in believing that the model actually leverages the reasoning tree encoded in its attention for making predictions.\n\n**Response:**\n\n>We thank the reviewer for the constructive comments. In section 4.3, we explore the relationship between our proposed probing scores and the LLM’s performance. Results imply that when we can successfully detect reasoning steps from LLM’s attention, the model is more likely to produce a correct prediction.\n\n>As for causal intervention, we classify attention heads into two groups: those focusing on number size (called size head) and others. The size heads contribute to the attention map for probing, while other heads can be regarded as noise in our analysis. We find that dropping size heads (less than 10% heads) has a much bigger impact than dropping other heads (>40% heads) (Figure 12). This supports the causation relationship between attention patterns and test accuracy.  We will also try to run experiments on ProofWriter and add the results as suggested.\n\n---\n\n**Question A:** \n\n    Can the authors provide some additional explanation/hypothesis as to why the model fails on longer inputs, despite being fine-tuned on them? Is this length beyond the sequence length that the model has been trained on?\n\n**Response:** \n>The input length is fixed, it is k that is increased. We suppose that the model failed due to the increase in task difficulty. Even if we finetune LMs, limited model capacity could make it hard to learn such difficult tasks well. We will add discussions in section 3.4 to provide more intuitions.\n\n---\n\n**Question B:** \n\n    In the kth-smallest element experiment, it is unclear if the gains in performance/attention patterns are learned exclusively during finetuning. Would finetuning a randomly initialized GPT-2 achieve similar results?\n\n**Response:** \n>No. We find that finetuning GPT-2 from scratch cannot converge even with millions of supervised data points (test accuracy < 1%). Note that the synthetic task for LMs is not that easy: LMs treat numbers as tokens, and they should first differentiate the size of numbers before giving the correct k-th smallest number. In Appendix B.3, we probe the pretrained GPT-2 (without any finetuning). We find that GPT-2 could learn to differentiate the size of numbers during pretraining. \n\n---\n\n**Question C:** \n\n     Can the authors expand more on the kNN classification of tokens/sentences as useful (similarly for the height classification)? Is the training data being used to identify the neighbors?\n\n**Response:** \n>We run the KNN classifier under the cross-validation setting. Part of the data is used as training data. Then, during testing, the classifier would find the nearest neighbor of the test data and assign its neighbor’s label (useful / not useful) to it. We would add more descriptions about the KNN classifier in the paper for clarification.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The paper is well-written and the experiments are carefully designed. With LLMs being prevalent across different fields and applications, it is important to understand their inner workings. This work sheds light on some interesting insights into the internals of model behavior that would be useful for the community to build upon.", "reasons_to_reject": "Attributing model behavior to attention is not always reliable and it is unclear how this method can be scaled to more complex reasoning tasks. Additionally, correlation does not imply causation. As the authors have pointed out, just because attention encodes the correct reasoning tree, it does not mean that these trees are being used by the model, regardless of its correlation with the probing scores. I think a causal intervention approach would instill more confidence in believing that the model actually leverages the reasoning tree encoded in its attention for making predictions. Nevertheless, I believe the paper presents some interesting findings about models’ attention mechanisms.", "questions_to_authors": "A. Can the authors provide some additional explanation/hypothesis as to why the model fails on longer inputs, despite being fine-tuned on them? Is this length beyond the sequence length that the model has been trained on? \n\nB. In the kth-smallest element experiment, it is unclear if the gains in performance/attention patterns are learned exclusively during finetuning. Would finetuning a randomly initialized GPT-2 achieve similar results?\n\nC. Can the authors expand more on the kNN classification of tokens/sentences as useful (similarly for the height classification)? Is the training data being used to identify the neighbors?", "paper_topic_and_main_contributions": "This work explores the internal multi-step reasoning process in LLMs. It introduces a novel probing approach called MechanisticProbe to recover the reasoning tree that resembles the correct reasoning process. The probe is used to analyze two LLMs (GPT-2 and LLaMA) on a synthetic task (finding the kth-smallest element) and a real-world reasoning task (ProofWriter) respectively. Using this probe, the authors find that the model’s attention patterns are indicative of the reasoning process represented by the reasoning tree. They also find that LLMs identify useful tokens/statements in the early layers and then reason over it step-by-step.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "- Replace k-smallest with kth-smallest as k-smallest element confuses the reader with the top-k subtask.\n\n- missing x-axis labels in Figure 9\n\n- typos in lines 056, 209, 237"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers were generally on the same page that the current scope of the work is somewhat narrow (on a synthetic task and on the naturalistic but synthetically-generated ProofWriter task), yet the overall findings about being able to mechanistically recover the implicit tree-structured reasoning from the model is valuable enough to be of interest. Unlike most explainability work that tries to obtain natural language multi-step explanations, the mechanistic interpretability angle of this work is novel in the multi-step reasoning space and provides useful insights.

The authors adequately addressed many of the concerns raised during the review period. The findings seem valuable to the community, even if currently somewhat narrow in scope. The reviewers have suggested specific other datasets (e.g., Musique) to push this line of work further.
Decision: Accept-Main