Paper: Analyzing Feed-Forward Blocks in Transformers through the Lens of Attention Maps (ICLR-2024)

Reviewer comment (Writing):
To enhance the clarity of the proposed method, it would be beneficial to provide a running example that illustrates the step-by-step process. By walking readers through a concrete example, they can more easily grasp the methodology and its application.

Author response:
We agree that providing concrete examples would be helpful to readers. We will add them to the camera-ready version.

Reviewer comment (Novelty):
To better ground the paper in the existing literature, it would be valuable to provide a more detailed and comprehensive literature review. By thoroughly reviewing relevant prior research, the paper can establish its position within the broader academic discourse and highlight its unique contributions.

Author response:
As for the additional literature related to this study, for example, there is research on not only reverse engineering of transformer models (mechanistic interpretability) [Elhage+’21,’22] but also research on comparing transformer models with human behavior [Oh&Schuler’22]. Unfortunately, due to space limitations, it is not immediately possible to add them to the paper; we will try to enhance the Introduction and Related work to be more comprehensive as much as possible in the camera-ready version.  

[Elhage+’21] A Mathematical Framework for Transformer Circuits (Anthropic) https://transformer-circuits.pub/2021/framework/index.html  
[Elhage+’22] Softmax Linear Units (Anthropic) https://transformer-circuits.pub/2022/solu/index.html  
[Oh&Schuler’22] Entropy- and Distance-Based Predictors From GPT-2 Attention Patterns Predict Reading Times Over and Above GPT-2 Surprisal (EMNLP 2022) https://aclanthology.org/2022.emnlp-main.632/

Reviewer comment (Experiments):
What could be the challenges of applying the proposed method in larger models, such as OPT, LLaMA, as mentioned in the future work?

Author response:
Although we used relatively small transformer models to demonstrate our proposed analysis method, analyzing LLMs is an important future work. Our method can be extended to LLMs without any technical modification, but it generally requires a computational cost typically beyond an academic resource. Formally, given an input of length $N$ and a model containing FF's intermediate dimension $d$, estimating one attention map requires $Nd$ times Integrated Gradients computation.

Reviewer comment (Experiments):
The paper suggests that because of the cancelled out effects in surrounding components there is "potential redundancy in transformer layers". However, there is not enough evidence to justify this claim (e.g., by training a transformer model removing some of the components). Without more experiments and results, it is not convincing what the conclusion of this paper is. More importantly, there is no systematic evaluation on the linguistic patterns (e.g., linguistic patterns distribution from different datasets on each layer) apart from some sampled amplified pairs. The results presented in Table 1 and Figure 5 are not evident.

Author response:
> Weaknesses 1  
> The paper suggests that because of the cancelled out effects in surrounding components there is "potential redundancy in transformer layers". However, there is not enough evidence to justify this claim (e.g., by training a transformer model removing some of the components). Without more experiments and results, it is not convincing what the conclusion of this paper is. 

The results in Section 6 revealed that contextualization effects by the Feed-forward network (FF) are canceled by the surrounding components. We agree that this *only suggests the potential redundancy* in the transformer layers. That is why we hedged with the words “potential”, “suggest,” or “imply” in the paper. It would be worthwhile to directly investigate the redundancy of FF, e.g., by observing the model’s behavior change while removing or pruning them. However, we fear that adding such experiments exceeds the focus of a single conference paper. We have refined the manuscript (Section 6) to state the need for further investigation explicitly.

> More importantly, there is no systematic evaluation on the linguistic patterns (e.g., linguistic patterns distribution from different datasets on each layer) apart from some sampled amplified pairs. The results presented in Table 1 and Figure 5 are not evident.

Section 5.2 shows that FF emphasizes specific linguistic compositions, at least across two different models, BERT-base and GPT2 (Figure 5). We agree that the current validation is not fully systematic and this problem can be alleviated by conducting additional experiments on other datasets. We will address this point in the camera-ready version. Nevertheless, it is worth noting that the essential contribution of this study is in part that it provides the new interpretation method and we hope that the acceptance of our study will facilitate more extensive linguistic analyses as you suggested in this community.

Reviewer comment (Theory):
Why is b set to a zero vector in equation 12?

Author response:
> Questions 1  
> Why is b set to a zero vector in equation 12?

There are mainly two reasons:
1. Using zero vector is one of the most common treatments employed as the neutral baseline for Integrated Gradients [Bastings&Filippova’20; Nayak+’21].
2. Since the activation function $g$ in FF satisfies $g(0) = 0$, we can ignore a particular term (output relative to baseline input) when applying Integrated Gradients with the zero baseline vector. This allows Integrated Gradients to fully additively decompose the output into the inputs (Appendix B.2), which is desirable for combining with the norm-based approach.

[Bastings&Filippova’20] The elephant in the interpretability room: Why use attention as explanation when we have saliency methods? (BlackboxNLP 2020) https://aclanthology.org/2020.blackboxnlp-1.14/  
[Nayak+’21] Using Integrated Gradients and Constituency Parse Trees to explain Linguistic Acceptability learnt by BERT (ICON 2021) https://aclanthology.org/2021.icon-main.11/

Reviewer comment (Evaluation):
Why do you think the patterns of changes between BERT and GPT-2 are quick different in each component (e.g., from Fig. 3)? What does this entail for different architectures?

Author response:
> Questions 2  
> Why do you think the patterns of changes between BERT and GPT-2 are quick different in each component (e.g., from Fig. 3)? What does this entail for different architectures?

At least there are two directions in the differences of BERT–GPT2 results: (i) module level and (ii) layer level.  
At the module level, contextualization changes by RES are remarkably different between BERT and GPT-2 (Figures 3, 9, and 10). There was almost no contextualization change in BERT, while there was the change in the GPT-2. We suspect that this is due to their different architectures: Post-LN model (BERT) and Pre-LN model (GPT-2). Given that RES in the Pre-LN model bypasses more modules than Post-LN, it is somewhat intuitive that RES in GPT-2 (Pre-LN) overwrites the internal representation more strongly than that in BERT (Post-LN).  
At the layer level, the contextualization change by FF in BERT is pronounced in the *middle to latter* layers, which is consistent across many variants of BERT (Figures 9 and 10). On the other hand, the change by FF in GPT-2 is pronounced in the *earlier* layers. Understanding this difference may follow the understanding of what each layer does in these models, and the exact interpretation of the model internals is still challenging [Tenney+'19; Elhage+’21]. We hope our study encourages such a deeper level of model understanding. We will add these discussions to the manuscript.

[Tenney+‘19]: BERT Rediscovers the Classical NLP Pipeline (ACL 2019) https://aclanthology.org/P19-1452/  
[Elhage+’21]: A Mathematical Framework for Transformer Circuits (Anthropic) https://transformer-circuits.pub/2021/framework/index.html

Reviewer comment (Evaluation):
Are results from Section 5 averaged across positions? Would position representation bias your findings?

Author response:
> Questions 3  
> Are results from Section 5 averaged across positions? Would position representation bias your findings?

We conducted two types of analyses: (i) the macro analysis (Section 5.1), which measures correlations between the attention patterns before and after a particular component; and (ii) the micro analysis (Section 5.2), which analyzes the amplification of attentions between specific types of token pairs. We guess that your concerns are about the latter micro analysis.  
The FF-amp scores are averaged across the targeted pairs in the corpus; that is, their positions are aggregated following their position distribution in the natural corpus. Further ablation of the position biases (e.g., analyzing with shifting or shuffling input positions) will be an interesting future work.

Reviewer comment (Evaluation):
How is the micro contextualization change (by subtracting a pre-FF attention map) in Section 5.2 different from measuring correlations?

Author response:
>Questions 4  
>How is the micro contextualization change (by subtracting a pre-FF attention map) in Section 5.2 different from measuring correlations?

In Section 5.2, we aimed to investigate which type of word-to-word interactions are controlled by FFs. For this goal, the single score of correlation between attention maps is not suited since it just summarized the degree of overall similarity between the patterns. Thus, instead of measuring the correlation, we computed the difference of attention maps and observed which cells (word-word pairs) got a large change in Section 5.2.

Reviewer comment (Experiments):
It would have been nice if the authors could have discussed and potentially also experimented with atleast one alternative formulation to IG, or atleast one of its variants [They do mention other formulations in B.1 Appendix but did not see further broaching of this angle beyond this]

Author response:
Yes, we can use another formulation, such as Shapley values (SV) (Appendix B.1), instead of Integrated Gradients (IG) in our proposed method. At least in the case of using SV, it is required to calculate an expected output change over **all permutations** of input features; that is, SV increases in computational cost more rapidly than IG as input becomes longer. This could especially be a severe problem in modern LLM use, such as providing long prompts/few-shot instances to models and/or considering verbose outputs generated by chain-of-though style reasoning; thus, we used IG as a practical choice.  
In the camera-ready version, we will add this discussion and an experimental comparison of IG and SV for small models in Appendix.

Reviewer comment (Experiments):
[Doesn't apply after authors comprehensively addressed this on 20th Nov] A marginal weakness but one nonetheless [and this is alluded to in future work], would have been nice to see this for a new-age LLM, of which some variants are available at lower or comparable parametrizations to GPT2-117M (e.g. OPT-125M)

Author response:
We are starting to experiment on OPT-125M. We may be able to share results during the discussion period. In any case, we plan to add it to the camera-ready version.

Reviewer comment (Novelty):
I know mechanistic explanation is a somewhat orthogonal paradigm of interpreting large transformer architectures, but it would be nice to have some comments on how this can relate or synergize with that paradigm [if at all]

Author response:
We recognize that mechanistic interpretability/explanation is an attempt to reverse engineer neural networks at the algorithmic level. We believe that our focus, “how the information of a particular token propagates to surrounding tokens in the model (similar to MOV or ADD commands in assembler),” can be an important perspective and provide a contribution there. In particular, our study may facilitate understanding MLPs (feed-forward networks), which has been difficult to interpret mechanistically [Elhage+’22]. We will explicate this point in the camera-ready version (currently, the paper has space limitations, and we will modify the paper after the discussions converge).

[Elhage+’22] Softmax Linear Units (Anthropic) https://transformer-circuits.pub/2022/solu/index.html

Reviewer comment (Experiments):
What is the effect of banded local attention [alternating banded local sparse and global attention] are a common part of the architecural recipe in many GPT3 or later LLMs so this would be a valuable insight to have.

Author response:
Thanks for your insightful comment. We will mention this as a future work in the camera-ready version. We hypothesize that the FF affects local and global attention differently; for example, a similar observation as this study (Figure 5) may be observed for local contextualization but, global attention may be affected differently, such as FF selectively expanding specific contextual information that is related to the target token.

Reviewer comment (Writing):
I think it might help the reader with a description or definition of contextualization.

Author response:
We have clarified what is referred to by “contextualization” in the second paragraph of the Introduction (Section 1) and the second paragraph of the Background (Section 2).