Paper: Interpreting the Inner Mechanisms of Large Language Models in Mathematical Addition (ICLR-2024)

Reviewer comment (Evaluation):
The paper identifies attention heads that take part in the processing of summation but does not look into or explain what each of the “key heads” actually does and what is the mechanism through which it contributes to summation. Therefore, the paper focuses on _localization_ of the heads that partake in summation, rather than _interpreting_ them.

Author response:
**(a)** Thank you for bringing up the insightful comment!
* We agree that delving into the *"causal mechanism"* of addition calculation is a quite promising direction. In this study, we follow the recent research in mechanistic interpretability [1, 2] that adapt the theory of causal mediation analysis to measure the *"causal effect"* of component mediators to the model behavior. We would like to highlight the identification of key heads *implicitly* interpret the model behavior (e.g. addition calculation).
* Before this study, there was a lack of clear interpretation regarding several important questions of addition calculation in LLMs. 
  (i) Are there attention heads consistently involved in addition calculation? 
  (ii) If yes, are they sparsely or densely distributed? 
  (iii) In the case of sparsity, where exactly are these key heads located?  
To shed light on these matters, we performed sufficient experiments using different models and data formats, striving to present the observed phenomena in a clear and understandable manner.  

We hope that our efforts could contribute to making the inner workings of LLMs in addition more understandable to humans.

Reviewer comment (Experiments):
The paper does not look at alternative representations of numbers. For example, in words (“two” instead of 2), Roman numerals (II instead of 2), and other languages (二 or ٢ instead of 2). The lack of such analysis leaves the question open whether these heads simply attend to numerical tokens or whether they are involved in higher-order reasoning about numbers and arithmetic.

Author response:
**(a)** To investigate the effects of numbers represented in different languages, we conduct the experiments as settings a1-a3. After knocking out the key heads identified based on Arabic numbers, the prediction accuracy for samples containing English numbers still decreases noticeably, while there is almost no decline in accuracy for samples containing Chinese and Roman numbers. We assume this is because the training corpus of LLaMA2 mainly consists of English language texts, which could explain the disparity in performance.

Reviewer comment (Experiments):
The paper seems to focus only on single-digit summation. It is unclear whether the results would translate to the summation of larger numbers (or more than two numbers). This is important as prior works have shown that the ability of LLMs to do arithmetic quickly decreases with the increase of the number of digits. It would be interesting to see if your analysis would be able to provide insights into this phenomenon.

Author response:
**(b)** In response to the issue of one-digit addition, we conduct the experiments as settings b1-b3.
* After knocking out the key heads identified based on one-digit addition, the performance for samples containing multi-digit numbers decreases significantly. When scaling to larger numbers, the performance drops more greatly. We assume this is because one-digit addition serves as the fundamental computation *"unit"* for multi-digit addition. The perturbation effects on one-digit addition could be accumulated when scaling to larger numbers. 
* Moreover, the reason why we experiment on one-digit addition in the primary experiments is that the adopted three LLMs (e.g., LLaMA2-7B) tokenize each digit individually (e.g., '42' is tokenized to '4' and '2'). We follow the one-digit nature of LLMs for simplicity in generating large-scale data.

The above studies may provide an explanation for why *"the ability of LLMs to do arithmetic quickly decreases with the increase of the number of digits"*.

Reviewer comment (Evaluation):
It is unclear how to read the attention patterns in Fig. 4. The reviewer questions how the attention can be negative. In fact, it does not seem that these heads attend to all numbers. The first head seems to attend to the completion of “or” with “anges” and the full stop. Both heads seem to attend only to 3 while solving the task would also require attention to 5. Therefore, it is not clear how these heads participate in performing summation.

Author response:
**Ans for Q3:**  
Thanks for your careful review. This is a visualization typo. The attention is bounded in [0, 1]. As for the selected example in Fig. 4, we find that it particularly attends to the token "3" by Head 13.11 and Head 12.22, while Head 17.31 attends to the tokens "5" and "3" simultaneously. We have revised the potentially confusing issue in the revision, by replacing it with another representative example.

Reviewer comment (Experiments):
The paper looks predominantly at attention heads. However, it is well known that a lot of the computation and processing happens in the MLPs. Hence, a full picture of the interoperation of the mechanisms for summation should also include the MLPs.

Author response:
**(b)** Thank you for the valuable suggestion on MLPs. 
* We added one experiment of path patching to measure the causal effect of each MLP layer. The results indicate that: (i) the early MLP layers 0-13 before the key heads (e.g., 13.11) have a slight effect on the output (approximately $\pm$0.5%). (ii) The late MLP layers 14-31 after the key heads exhibit a much larger effect (approximately $\pm$10.0%). These findings may reveal that the key heads are responsible for attending to number tokens, while the following MLP layers process the number tokens.
* We also noticed that previous works [1, 2] tried to explain MLPs for performing computations and retrieving facts. However, it necessitates thorough and rigorous analysis to make the detailed process in MLPs human-understandable, and interpret the collaboration between attention heads and MLPs. We believe that making the calculation fully understandable to humans is an intriguing direction that warrants sustained research efforts.

Reviewer comment (Novelty):
In the Introduction, the paper states _“Contrary to expectations that LLMs may involve diverse attention heads across all layers, our findings reveal that only a small percentage (0.5%) of the attention heads significantly impact the model’s performance.”_ However, the reviewer notes that this is exactly the expectation: attention heads have diverse functions so it is not surprising that only a few of them would be involved in summation.

Author response:
**(a)** We remove *'Contrary to expectations that LLMs may involve diverse attention heads across all layers,'* for more precise expression.

Reviewer comment (Theory):
In the Introduction, the paper states _“Remarkably, only the absence of these heads leads to a notable decline in model performance.”_ The reviewer disputes this, noting that there are surely many other weights that, if perturbed, would result in a significant decline in model performance (e.g. the embedding matrix or the final output projection matrix).

Author response:
**(b)** We change the original sentence to *'Remarkably, the absence of these heads leads to a relatively notable decline in model performance, compared to the absence of randomly-selected heads.'*.

Reviewer comment (Reproducibility):
In Section 4.2, the reviewer asks how the paper decouples the effect of the individual heads. In the implementation of LLAMA there are no separate $W_O$ for each head but a single one that is applied to the concatenation of all the heads. Therefore, it mixes information across heads. The reviewer asks how this issue is resolved.

Author response:
**Ans for Q6:**  
Thanks for pointing out this issue. We follow the implementation in [5] to split $W_O$ into equal size blocks for each head $\left [  W_O^{h_1}, W_O^{h_2}, \cdots \right ]$. As is verified in [6], it's equivalent to running heads independently, multiplying each by its own output matrix, and adding them into the residual stream. 

[5] Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small. In ICLR 2023.  
[6] A Mathematical Framework for Transformer Circuits. In Transformer Circuits Thread 2021.

Reviewer comment (Evaluation):
In Figure 4 left, the reviewer questions how the authors know that the observed effect is because of the heads specialising in numbers and not simply because the test sequences contain numbers. The reviewer would be curious to see how this plot and the rest of the analysis would look if applied to sentences unrelated to numbers and arithmetic. Possibly the heads found to be important for arithmetic would be especially unimportant for other tasks.

Author response:
**Ans for Q7:**  
Thank you for bringing up this concern. We conduct the plot experiment based on the counterfactual sentences. These cases contain no arithmetic, i.e., excluding addition logic, and implicitly have nothing to do with numbers. The findings indicate that the key heads still exhibit a relatively lower attention probability towards non-numeric tokens, similar to the observation depicted in Figure 4. This phenomenon supports the claimed effects of the key heads at a certain level.

Reviewer comment (Experiments):
The robustness of this study could be significantly enhanced if the authors were to conduct analogous experiments on subtraction, multiplication, and division. Such investigations would shed light on whether a select group of attention heads can consistently influence performance across all four mathematical operations.

Author response:
Thanks a lot for bring up this issue! We conduct the path patching experiments on four mathematical tasks using the following representative templates.  

| tasks | template1 | template2 | template3 |
| ------- | ------- | ------- |  ------- | 
|    addition   | "{A} + {B} = " |    "The sum of {A} and {B} is "   | "Question: What is {A} plus {B}? Answer: "|
|    subtraction   | "{A} - {B} = "|    "The difference between {A} and {B} is "   | "Question: What is {A} minus {B}? Answer: " |
|    multiplication   | "{A} * {B} = " |    "The product of {A} and {B} is "   | "Question: What is {A} times {B}? Answer: "|
|    division   | "{A} / {B} = "|    "The ratio between {A} and {B} is "   | "Question: What is {A} over {B}? Answer: " |

* We find that: (i) the sparsity of key heads remains consistent across all four tasks (less than 1.0% of all heads). (ii) The key heads mainly distribute in the middle layers. The phenomena are analogous to the primary findings on the addition task (Section 4.1), demonstrating the potential of extending the observed effects of the addition task to other mathematical tasks. 
* We compare the location of key heads across four mathematical tasks. An interesting finding is that the key heads used in "subtraction" and "addition" tasks overlapped significantly, as did the key heads used in "multiplication" and "division" tasks. Moreover, the four tasks share the heads (e.g., 13.11 and 12.22) that deliver the most significant effects, while they have task-specific heads that only emerge in its own task. These findings suggest that LLMs exhibit behavior aligned with human thinking to some extent, since "subtraction-addition" and "multiplication-division" are opposite mathematical operations.

Thanks for your valuable suggestion, the above results and discussions have been added to the revision (Appendix B).

Reviewer comment (Experiments):
The paper (seems to) limit itself to simple one-digit addition and subtraction (without “carry over one” or “borrow one” examples - reducing its scope to a subset of addition and subtraction. The abstract should explicitly say that the scope is simple one-digit integer addition.

Author response:
**Ans for Q1:**  
**ii)** Thank you for the insightful comment. Following the knockout experiments in Section 4.3, we first generate 200 samples that need to carry over one or borrow one (e.g., "17 + 9 = " and "17 - 9 = "). Then we knock out the identified key heads in addition, resulting in a wrong prediction on over 150 samples. We assume this is because the key heads attend to the addends "7" and "9", regardless of whether it needs to carry over one or borrow one.

Reviewer comment (Evaluation):
The paper touches on subtraction, showing similarities, but a detailed analysis is left for future work.

Author response:
**Ans for Q3:**  
Thanks for your constructive suggestion. 
* We conduct the experiment of path patching based on the subtraction sentences. We find that the identified key heads in the subtraction task are almost the same to those in the addition task, albeit with different magnitude of the effect. This phenomenon could reveal the similar key head *"location"* in addition and subtraction. 
* Moreover, as we show the attention pattern of key heads, we observe that these heads particularly attend to the number tokens regardless of whether they are given addition or subtraction sentences. This phenomenon could reveal the similar key head *"behavior"* in addition and subtraction. 

The above studies further explain why knocking out the key heads identified in addition could affect the accuracy in performing subtraction. We have add the above discussions in the revision (Appendix B).

Reviewer comment (Evaluation):
A discussion of the differences in how each of the LLMs implement one-digit addition would have been interesting e.g. do all the models use roughly the same number of attention heads to implement addition? If no differences were found, then this would be an interesting finding in itself.

Author response:
**Ans for Q4:**  
Thanks for your inspiring comment:
* We find that the quantity of key heads in different LLMs is different from each other. This could be attributed to their different architecture, model size, learning paradigm, training corpus, and more. 
* We also noticed that the adopted three LLMs show different ability in understanding the input context. Given the input addition sentences (e.g., "The sum of {A} and {B} is "), we gather the averaged prediction probabilities of number tokens 1-9. The results are: LLaMA2-7B=99.3%, Qwen-7B=97.6%, and chatGLM2-6B=94.7%. The higher probability of LLaMA2-7B shows that it generates number tokens with a greater confidence level, demonstrating a better understanding of the input context. We assume this explains why LLaMA2-7B has more key heads in performing the addition task.

Reviewer comment (Novelty):
The small scope of this paper limits the reusability of this work.

Author response:
**Ans for Q1:**  
**iii)** We agree that a larger scope is important for the reusability. The above studies suggest the potential of extending the scope of this work to multi-digit and non-trivial scenarios. More analyses and discussions will be updated in the revision.

We believe these novel observations would benefit a lot to the quality of our paper. Thanks again for your constructive suggstions.

Reviewer comment (Experiments):
The addition examples seem to be “simple” one-digit integer addition with a one character answer. There appear to be no “carry over one” examples in the test questions e.g “5 plus 7 is equal to 1_”. If this is so, it reduces the findings scope to some subclasses of addition.

Author response:
**Ans for Q6:**  
Thanks for bring up the concern. We conduct the experiments in the responses to **Q1**, and illustrate the scope of this work could extend to multi-digit and non-trivial scenarios.

Reviewer comment (Experiments):
The subtraction examples all seem to be “simple” one-digit integer subtraction with a one character answer. There appear to be no “borrow one” examples in the test questions e.g “112 minus 5 is equal to 10_”. If this is so, it reduces the findings scope to some subclasses of subtraction.

Author response:
**Ans for Q6:**  
Thanks for bring up the concern. We conduct the experiments in the responses to **Q1**, and illustrate the scope of this work could extend to multi-digit and non-trivial scenarios.

Reviewer comment (Evaluation):
The calculation of the subtraction question “{A} - {B} =” likely has two distinct calculation algorithms: one for when A > B and one for when A < B. Do the authors think that this explains the 52% performance drop when the addition attention heads are ablated?

Author response:
**Ans for Q5:**  
Thanks for your careful reviews. We agree that the calculation algorithms of "{A} > {B}" and "{A} < {B}" are distinct. As mentioned in our response to Q3, we hypothesize that the decrease in performance could be attributed to the shared key heads in both subtraction and addition tasks. The identified key heads in the addition task are also involved in the inner workings of subtraction. More discussions have been incorporated in our revision (Appendix B).

Reviewer comment (Novelty):
The authors didn't include, as related work, some publications that also deal with mathematical reasoning, such as [1] https://arxiv.org/pdf/2305.08809.pdf

Author response:
**Ans for Q1:**  
Thanks for bringing attention to this outstanding work [1], we have added a discussion in the Background section. The studies in [1] scale the methods from causal abstraction to understand how Alpaca (7B) follows a particular instruction: "Please say yes only if it costs between [X] and [Y] dollars, otherwise no.", where it needs to compare the input value with the lower bound [X] and the upper bound [Y]. Different from the task of *number comparison*, we focus on the *number addition* behavior of LLMs and attempt to make their inner workings more understandable to humans.

[1] Interpretability at Scale: Identifying Causal Mechanisms in Alpaca. In ArXiv preprint 2023.

Reviewer comment (Experiments):
Studying only mathematical addition and subtraction seems restrictive. I do note that the authors state at the end however: "_A more thorough study on the subtraction task as well as the validation on more computation tasks (e.g., multiplication and division, etc.) is left for future work._"

Author response:
**Ans for Q2:**  
Thanks for your constructive suggestion. We conduct experiments on four mathematical tasks using the following representative templates.  

| tasks | template1 | template2 | template3 |
| ------- | ------- | ------- |  ------- | 
|    addition   | "{A} + {B} = " |    "The sum of {A} and {B} is "   | "Question: What is {A} plus {B}? Answer: "|
|    subtraction   | "{A} - {B} = "|    "The difference between {A} and {B} is "   | "Question: What is {A} minus {B}? Answer: " |
|    multiplication   | "{A} * {B} = " |    "The product of {A} and {B} is "   | "Question: What is {A} times {B}? Answer: "|
|    division   | "{A} / {B} = "|    "The ratio between {A} and {B} is "   | "Question: What is {A} over {B}? Answer: " |

- Our experimental results show that: 
  - The sparsity of key heads remains consistent across all four tasks (less than 1.0% of all heads); 
  - The key heads mainly distribute in the middle layers.  
The phenomena are analogous to the primary findings on the addition task (Section 4.1), demonstrating the potential of extending the observed effects of the addition task to other mathematical tasks. 
* We compare the location of key heads across four mathematical tasks. An interesting finding is that the key heads used in "subtraction" and "addition" tasks overlapped significantly, as did the key heads used in "multiplication" and "division" tasks. Moreover, the four tasks share the heads (e.g., 13.11 and 12.22) that deliver the most significant effects, while they have task-specific heads that only emerge in its own task. These findings suggest that LLMs exhibit behavior aligned with human thinking to some extent, since "subtraction-addition" and "multiplication-division" are opposite mathematical operations.

The above results and discussions have been incorporated into the revision (Apendix B).

Reviewer comment (Evaluation):
Since addition and subtraction are opposite mathematical operations, is there some kind of similar symmetry observable on the level of attention heads?

Author response:
**Ans for Q3:**  
Thanks a lot for bringing up this insightful question!
* In the above experiments for **Q2**, we find that the identified key heads in the addition task are almost the same to those in the subtraction task, albeit with different magnitude of the effect. This phenomenon may reveal the symmetry of key head *"location"* in addition and subtraction. 
* Moreover, as we show the attention pattern of key heads, we observe that these heads particularly attend to the number tokens regardless of whether they are given addition or subtraction sentences. This phenomenon may reveal the symmetry of key head *"behavior"* in addition and subtraction.