Causal Interpretation of Self-Attention in Pre-Trained Transformers

Summary

The paper establishes a formal link between a (pre-trained) self-attention layer and a causal graph underlying a sequence of symbols. Specifically, it shows that self-attention learns to represent pairwise associations for which there exists a linear-Gaussian Structural Causal Model that has the exact same pair-wise associations between its nodes. The paper then proceeds to derive an algorithm for recovering the causal structure of the SCM solely from the weights of a pre-trained attention layer

Strengths

The Attention-Based Causal-Discovery (ABCD) method, which is developed for learning an equivalence class of the causal graph that underlies a given input sequence is reasonable. It relies on partial correlarion CI-testing, which is sound. Then the paper proceeds to apply ABCD to reasoning about predictions by generating causal explanations from pre-trained self-attention models, such as BERT (8). Specifically, the algorithm produces explanations through selecting a subset of the input symbols that are claimed to be solely responsible for BERT prediction. ABCD is shown empirically to be much stronger than the common approach of using the attention matrix to learn about input-output relations for providing explanations.

Weaknesses

The Attention-Based Causal-Discovery (ABCD) method relies on an unsupervised training assumption. This restricts applicability to only models of BERT type, you cannot deal e.g. with decoder-only models. In addition, it is extremely troubling that the authors do not perform any comparative evaluation against standard methods in the field. They limit their examination to standard self-attention and a variant of it. Even worse, they have avoided use of any standard benchmark in the field. Eventually, it is not clear how useable the generated causal explanations are, and how easy it is to exploit them in an application setting.

Questions

1. I want to see a serious reconsideration of all the empirical evidence the paper provides. 2. How can the causal explanations be used in an easy and/or automated setting?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

n/a

Summary

The ABCD paper starts off by motivating the need for understanding the complex dependencies that underly current Transformer architectures, pointing to causality as the missing puzzle piece. Following brief introductions of both SCMs and Self-attention, the authors then provide covariance matrix derivations for both linear Gaussian SCM and attention such that they are able to equate them eventually. By introducing another assumption on how the attention was trained, they propose a theorem that provides a clear connection between causal discovery and attention in that one can make CI tests based on the attention matrix essentially. Finally, the authors propose a complete pipeline for going from e.g. a pre-trained Transformer model to causal explanations based on their ABCD approach computed from the last attention layer. As put by the authors, this can be considered as a "one-shot" causal discovery paradigm.

Strengths

1. A unique perspective and valuable contribution to the causal discovery literature in that it expands the horizon on what kind of connections can be established between existing models to foster new understanding of said models. The derivation of the covariances and then equating them is both simple and effective, and then Thm.1 seems like a nice result for understanding the final correspondence between attention and CI-testing for constraint-based causal discovery. 2. Presenting a complete pipeline is a great bonus (since it is a difficult task) and shows an immediate practical approach to the previously presented theory. The explanation approach is convincing and practically applicable since the theory justifies the use of the attention matrix.

Weaknesses

1. While the causality assumptions Markov and faithfulness can be considered standard, as soon as moving to the realm of data that nowadays Transformers are confronted with we run the risk of violating them. Furthermore, the assumption of a linear Gaussian SCM is quite restrictive and should be highlighted more, especially under consideration of Transformers. Finally, Assumption 1 right before Theorem 1 is only vaguely comprehensible IMHO. I'd wish for the authors to be more specific about what property is required for the deepest attention layer such that the equating to CB-CD works. This could be done through formalization of that assumption or if not possible, then through an example. 2. Thm.1 is incomplete writing-wise i.e., sentence cuts off at end of page 4. Starting from p.4 the paper generally lend the impression of being not as carefully written as the pages before. Please consider a careful readthrough post-submission. 3. Missing discussion of explanation methods based on causal graphs. Two very different approaches include Causal Shapely Values by Heskes et al. 2020 and Causal Explanation of SCM by Zečević et al. 2021. The first is a rather classical approach, where numerical attributions are being computed and in Sec.5 they present an approach for causal chain graphs (a special case like your Linear-Gaussian case), and the second discusses a traversal algorithm for generating immediate textual representations based on more general SCM (although only showcased for a linear SCM setting).

Questions

1. Can you comment on Theorem 1 as mentioned in the weaknesses section? 2. The assumption of linear Gaussian SCM can severely limit the method. What was the reasoning behind the assumption? 3. What property is required for the deepest attention layer such that the equating to CB-CD works. Overall, a nice contribution but not devoid of weaknesses. I will be willing to raise my score based on the rebuttal.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

No concerns here.

Summary

In this work, authors, interpret self attention in transformers as linear gaussian SCM over output embeddings. SCM over input embeddings is same as SCM over output embeddings, since deepest attention layer is trained to predict input symbol. Existing conditional independence based testing algorithm is applied on the weights of the deepest attention layer of the pretrained transformer in order to estimate causal structure. Furthermore, the authors demonstrate the applicability of the proposed approach in deducing causal explanations for attention based recommendations.

Strengths

Authors proposed novel way of viewing self attention in transformers as weight matrix of linear gaussian SCM Pretrained model is used to obtain causal structure, which would be of great use to gain causal perspectives for existing model Estimating causal structure from the attention layer and using it for causal explanantions is promising.

Weaknesses

Theorem1, in main paper has few words missing, however, it's available in supplementart material

Questions

In the paper, the covariance matrix of value is assumed to be identity using central limit theorem, however, in practice, it wouldnot be identity as a result the graph learnt wouldnot be faithful. This will act as one more reason of inaccuracy. Is it possible for you to throw some light on this ?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

1) The accuracy of the proposed method as explained by the authors is dependent on the prediction of the transformers and the conditional independence testing accuracy 2) Authors assumed learnt graph will be faithful, since exogeneous variables V under central limit theorem will have identity covariance, however in practice the faithfullness assumption will be violated, this will act as one more reason of inaccuracy

Summary

The paper proposes Attention-Based Causal Discovery (ABCD), a causal interpretation for the attention matrix of transformers (e.g., LLMs like BERT). The authors suggest CLEANN, an algorithmic approach that first finds a causal structure of the input symbols (e.g., tokens (words) of a sentence as input to BERT). This causal structure is unique to each input (sentence). This can then be used to get causal explanation for different taks (in their experiments classification or recommendation) based on the representations of the input sentences (e.g., the word "but" and "bad" in the input sentence caused the negative sentiment classification).

Strengths

Thank you for your submission. I enjoyed reading your paper. Clarity & Presentation: I found the paper to be clearly structured and written. On clarity and presentation, I would like to positively highlight the introduction to each (sub) section gave an orientation that helped to clearly guide the reader, what the next chapter is about and how it fits in the bigger picture. The figures (e.g., Fig2a, Fig 4) helped a lot to understand the method/pipeline. Novelty: I found the causal interpretation and its application to causal explanation of predictions apealing and novel. Finding a causal explanation to downs-stream tasks of LLMs for example, can be very helpful. However, I am not an expert in the field of LLMs and I do not feel qualified to fully assess the novelty and hence the contribution of the work in the general field. Evaluation: I found the experimental section clearly structured and interesting to read.

Weaknesses

I have a few comments for improvements. * Section 5, Table 1: What do the bold numbers indicate, what the +- (i.e., over how many runs did you do it, were these runs of CLEANN or different data samples)? This would be helpful to comment in the caption. Am I right in understanding tha CLEANN does yield lower average explanations, but the numbers suggest that it is not a significant? Over how many runs did you do it? * Section 5, Figure 2 (b, c): I find it hard to understand and interpret Figure 2(b,c) despite your explanation in ln. 255ff. Also, the labels are very small and hard to read. * Section 5, Figure 3: I am missing an explanation/interpretation of Figure 3 in the text. Did I overlook it? * I am missing a related work section. Minor: * Section 1: While I enjoyed reading the empirical evaluation, I found the causal explaining application could be stronger highlighted in the introduction. * Section 2.1, line 67: While I appreciate the introduction to SCMs, I was a bit confused about the Causal Markov and Faithful properties at the end of the section. They seem very disconnected of the rest of the text. It would help, if there would be a short explanation, why they are stated / needed for the paper. * Section 2.2 is very brief. In the light of Section 3 (ln. 100 ff.), where the self-attention mechanisms is further detailed, it makes sense. However, it may be helpful to just point to that at the end of the section 2.2. * Section 5: Caption of Table 1 - by review length, you mean the length of the sentence, right? It would have helped me to point that out. * Section 5: I find it weird that Fig. 6 is discussed before Fig. 5. A natural way would be to discuss them the other way around. * I got a bit confused by the difference between ABCD and CLEANN. Could you kindly clarify? Typos: * ln 222 "(1) In addition, it" -> "(1) It" * ln 262 "influences on the " -> "influences the"

Questions

I have a few questions: * Section 2.1, ln. 50, Definition of an SCM. I have previously seen an SCM being defined often $M = (F P(U))$, or $M = (X, U , F)$, but I have not seen both U or P(U). Is yours the proper formal definition? * Section 2.1, ln. 59, your understanding of a linear-Gaussian SCM is the same as an Additive Noise Model (ANM)? * Section 5, Fig. 6: I do not understand the x-axis. Could you explain this again? * Are there other applications of ABCD beyond providing causal explanations?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Section 6 Discussion: ln. 305 mentions "under certain assumptions", but these assumptions are not re-stated here. It would be helpful to have them restated. In this regard, I am missing a discussion of the limitations regarding the (causal) assumptions. I am also wondering, how robust your causal explanations are? I believe, when providing causal explanations to individuals, it is important to discuss any limitations regarding the sensitivity to assumptions and to comment the potential negative societal impact, if assumptions are not met.

Thank you for the response

I would like to thank the authors for their detailed response. I am overall positive about the paper and would like to keep my score.