Response to Official Review by Reviewer Dtpy
We appreciate your invaluable comments.
# Alternatives to Integrated Gradients
>Weaknesses 1
>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]
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.
# Application to a new-age LLM
>Weaknesses 2
>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)
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.
# Computational cost of the proposed method
> Questions 1
>What is the computational [and memory] complexity of generating these maps? This may sound nitpicky but with increasingly large LLMs which barely fit in the accelerator time and memory bounds whether at training or inference time, this can indeed become a factor and consideration in how widely this gets adapted.
Given an input of length $N$ and FF's intermediate dimension $d$, estimating one attention map requires the computation of Integrated Gradients $Nd$ times. In other words, the computational cost increases with input length and model size. Fortunately, our preliminary analysis shows that the practical computational cost (speed and memory use) seems to be proportional or better than the linear increase with respect to input length and model size ($N$ and $d$) due to some implementation tricks and parallelism (e.g., provided by captum https://captum.ai/).
# Relation to mechanistic explanation
> Questions 2
> 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]
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
# Local sparse and global dense attention in recent LLMs
> Questions 3
> 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.
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.