SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Raanan Y. Rohekar
raanan.yehezkel@intel.com
Yaniv Gurwicz
yaniv.gurwicz@intel.com
Shami Nisimov
shami.nisimov@intel.com
We propose a causal interpretation of self-attention in the Transformer neural network architecture. We interpret self-attention as a mechanism that estimates a structural equation model for a given input sequence of symbols (tokens). The structural equation model can be interpreted, in turn, as a causal structure over the input symbols under the specific context of the input sequence. Importantly, this interpretation remains valid in the presence of latent confounders. Following this interpretation, we estimate conditional independence relations between input symbols by calculating partial correlations between their corresponding representations in the deepest attention layer. This enables learning the causal structure over an input sequence using existing constraint-based algorithms. In this sense, existing pre-trained Transformers can be utilized for zero-shot causal-discovery. We demonstrate this method by providing causal explanations for the outcomes of Transformers in two tasks: sentiment classification (NLP) and recommendation.
Causality plays an important role in many sciences such as epidemiology, social sciences, and finance (25; 34). Understanding the underlying causal mechanisms is crucial for tasks such as explaining a phenomenon, predicting, and decision making. An automated discovery of causal structures from observed data alone is an important problem in artificial intelligence. It is particularity challenging when latent confounders may exist. One family of algorithms (35; 7; 8; 46; 30), called constraint-based, recovers in the large sample limit an equivalence class of the true underlying graph. However, they require a statistical test to be provided. The statistical test is often sensitive to small sample size, and in some cases it is not clear which statistical test is suitable for the data. Moreover, these method assume that data samples are generated from a single causal graph, and are designed to learn a single equivalence class. There are cases where data samples may be generated by different causal mechanisms and it is not clear how to learn the correct causal graph for each data sample. For example, decision process may be different among humans. Data collected about their actions may not be represented by a single causal graph. This is the case for recommender systems that are required to provide a personalized recommendation for a user, based on her past actions. In addition, it is desirable that such automated systems will provide a tangible explanation to why the specific recommendation was given.
Recently, deep neural networks, based on the Transformer architecture (39) have been shown to achieve state-of-the-art accuracy in domains such as natural language processing (9; 3; 45), vision (10; 19), and recommender systems (36; 18). The Transformer is based on the attention mechanism, where it calculates context-dependent weights for a given input (32). Given an input consisting of a sequence of symbols (tokens), the Transformer is able to capture complex dependencies, but it is unclear how they are represented in the Transformer nor how to extract them.
*All authors contributed equally.
In this paper we bridge between structural causal models and attention mechanism used in the Transformer architecture. We show that self-attention is a mechanism that estimates a linear structural equation model in the deepest layer for each input sequence independently, and show that it represents a causal structure over the symbols in the input sequence. In addition, we show that an equivalence class of the causal structure over the input can be learned solely from the Transformer's estimated attention matrix. This enables learning the causal structure over a single input sequence, using existing constraint-based algorithms, and utilizing existing pre-trained Transformer models for zero-shot causal discovery. We demonstrate this method by providing causal explanations for the outcomes of Transformers in two tasks: sentiment classification (NLP) and recommendation.
In recent years, several advances in causal reasoning using deep neural networks were presented. One line of work is causal modeling with neural networks (42; 12). Here, a neural network architecture follows a given causal graph structure. This neural network architecture is constructed to model the joint distribution over observed variables (12) or to explicitly learn the deterministic functions of a structural causal model (42). It was recently shown that a neural network architecture that explicitly learns the deterministic functions of an SCM can be used for answering interventional and counterfactual queries (43). That is, inference in rungs 2 and 3 in the ladder of causation (26), also called layers 2 and 3 in the Pearl causal hierarchy (2). Nevertheless, this architecture requires knowing the causal graph structure before training, and supports a single causal graph.
In another line of work, an attention mechanism is added before training, and a causal graph is inferred from the attention values at inference (21; 17). Attention values are compared against a threshold, and high values are treated as indication for a potential causal relation. Nauta et al. (21) validate these potential causal relations by permuting the values of the potential cause and measuring the effect. However, this method is suitable only if a single causal model underlay the dataset.
First, we provide preliminaries for structural causal models and the attention mechanism in the Transformer-based neural architecture. Throughout the paper we denote vectors and sets with bold-italic upper-case (e.g.,  $V$ ), matrices with bold upper-case (e.g.,  $\mathbf{A}$ ), random variables with italic upper-case (e.g.,  $X$ ), and models in calligraphic font (e.g.,  $\mathcal{M}$ ). Different sets of letters are used when referring to structural causal models and attention (see Appendix A-Table 1 for a list).
3.1 Structural Causal Models
A structural causal model (SCM) represents causal mechanisms in a domain (24; 35; 28). An SCM is a tuple  $\{\pmb{U},\pmb{X},\mathcal{F},P(\pmb{U})\}$ , where  $\pmb{U} = \{U_1,\dots ,U_m\}$  is a set of latent exogenous random variables,  $\pmb{X} = \{X_{1},\ldots ,X_{n}\}$  is a set of endogenous random variables,  $\mathcal{F} = \{f_1,\dots ,f_n\}$  is a set of deterministic functions describing the values  $\pmb{X}$  given their direct causes, and  $P(U)$  is the distribution over  $\pmb{U}$ . Each endogenous variable  $X_{i}$  has exactly one unique exogenous cause  $U_{i}$  ( $m = n$ ). The value of an endogenous variable  $X_{i},\forall i\in [1,\dots ,n]$  is
$$
X _ {i} \leftarrow f _ {i} \left(\boldsymbol {P} \boldsymbol {a} _ {i}, U _ {i}\right) \tag {1}
$$
(the left-arrow indicates an assignment due to cause-effect relation), where  $\mathbf{P}\mathbf{a}_i$  is the set of direct causes (parents in the causal graph) of  $X_{i}$ . A graph  $\mathcal{G}$  corresponding to the SCM consists of a node for each variable, and a directed edge for each direct cause-and-effect relation, as evident from  $\mathcal{F}$ . If the graph is directed and acyclic (DAG) then the SCM is called recursive.
In a special subclass of SCMs, referred in this paper linear-Gaussian SCM, a node value is a linear combination of its parents, and exogenous nodes are independently and normally distributed additive noise. Let  $\mathbf{G}$  be a weight matrix, where  $\mathbf{G}(i,j)$  is the weight of parent node  $X_{j}$  determining the child node  $X_{i}$  and  $\mathbf{G}(i,k) = 0$  if  $X_{k}$  is not a parent of  $X_{i}$ , and  $\Lambda$  be a diagonal matrix, where  $\Lambda (i,i)$  is the weight of exogenous node  $U_{i}$ . Then,  $\forall i\in [1,\dots ,n]$
$$
X _ {i} \leftarrow \mathbf {G} (i, \cdot) \boldsymbol {X} + \boldsymbol {\Lambda} (i, i) U _ {i}, \tag {2}
$$
and in matrix form  $\mathbf{X} = \mathbf{G}\mathbf{X} + \boldsymbol{\Lambda}\mathbf{U}$ . In the case of a recursive SCM, with nodes sorted topologically (ancestors before their descendants), the weight matrix  $\mathbf{G}$  is strictly lower triangular (zeros on the diagonal). In this paper we learn the causal graph, described by the non-zero elements of  $\mathbf{G}$ , using constraint-based causal discovery approach, which generally requires assuming the following.
Definition 1 (Causal Markov). In a causally Markov graph, a variable is independent of all other variables, except its effects, conditional on all its direct causes.
Definition 2 (Faithfulness). A distribution is faithful to a graph if and only if every independence relation true in the distribution is entailed by the graph.
3.2 Self-Attention Mechanism
Attention is a mechanism that predicts context-dependent weights (32) for the input sequence. In this paper we analyse the attention mechanism used by Vaswani et al. (39). Here, the input is a matrix  $\mathbf{Y}$  where row vector  $\mathbf{Y}(i,\cdot)$  is the embedding of symbol  $s_i$  in an input sequence  $\boldsymbol{s} = [s_1,\dots,s_n]$ . The attention matrix is calculated by  $\mathbf{A} = \text{softmax}(\mathbf{YW}_{QK}\mathbf{Y}^{\top})$ , such that the rows of  $\mathbf{A}$  sum to  $1^{1}$ . Values matrix is computed by  $\mathbf{V} = \mathbf{Y}\mathbf{W}_V$ , where row  $\mathbf{V}(i,\cdot)$  is the value vector of input symbol  $s_i$ . We treat  $\mathbf{V}$  as a set of projection of  $\mathbf{Y}$  on column vectors  $\mathbf{W}_V$ . Finally, the output embeddings are  $\mathbf{Z} = \mathbf{A}\mathbf{V}$ , where the  $i$ -th output embedding  $\mathbf{Z}_i$  is the  $i$ -th row of  $\mathbf{Z}$ .
4 A Link between Pre-trained Self-Attention and a Causal Graph underlying an Input Sequence of Symbols
In this section we first describe self-attention as a mechanism that encodes correlations between symbols in an unknown structural causal model (Section 4.1). After establishing this relation, we present a method for recovering an equivalence class of the causal graph underlying the input sequence (Section 4.2). We also extend the results to the multi-head and multi-layer architecture (Section 4.3).
4.1 Self-Attention as Correlations between Nodes in a Structural Causal Model
We now describe how self-attention can be viewed as a mechanism that estimates the values of observed nodes of an SCM. We first show that the covariance over the outputs of self-attention is similar to the covariance over observed nodes of an SCM. Specifically, we model relations between symbols in an input sequence using a linear-Gaussian SCM at the output of the attention layer. In an SCM, the values over endogenous nodes, in matrix form, are  $\mathbf{X} = \mathbf{G}\mathbf{X} + \boldsymbol{\Lambda}\mathbf{U}$ , which means
$$
\boldsymbol {X} = (\mathbf {I} - \mathbf {G}) ^ {- 1} \boldsymbol {\Lambda} \boldsymbol {U}. \tag {3}
$$
Since  $\mathbf{G}$  is a strictly lower-triangular weight matrix,  $(\mathbf{I} - \mathbf{G})^{-1}$  is a lower unitriangular matrix, which is equal to the sum of a geometric series
$$
\left(\mathbf {I} - \mathbf {G}\right) ^ {- 1} = \sum_ {k = 0} ^ {n - 1} \mathbf {G} ^ {k}. \tag {4}
$$
Thus, the  $(i,j)$  element represents the sum of directed-paths' weights, for all directed paths from  $X_{j}$  to  $X_{i}$  having length up to  $n - 1$ . The weight of a directed path is the product of the weights of edges along the path. In case some of the nodes are latent confounders, then the matrix in Equation 4 may not be triangular, due to spurious associations. Equation 3 represents a system with input  $\mathbf{U}$ , output  $\mathbf{X}$  and weights  $(\mathbf{I} - \mathbf{G})^{-1}\mathbf{\Lambda}$ . The covariance matrix of the output is
$$
\begin{array}{l} \mathbf {C} _ {\boldsymbol {X}} = \mathbb {E} [ (\boldsymbol {X} - \boldsymbol {\mu} _ {\boldsymbol {X}}) (\boldsymbol {X} - \boldsymbol {\mu} _ {\boldsymbol {X}}) ^ {\top} ] = \\ = \mathbb {E} [ (\mathbf {I} - \mathbf {G}) ^ {- 1} \boldsymbol {\Lambda} (\boldsymbol {U} - \boldsymbol {\mu} _ {\boldsymbol {U}}) (\boldsymbol {U} - \boldsymbol {\mu} _ {\boldsymbol {U}}) ^ {\top} ((\mathbf {I} - \mathbf {G}) ^ {- 1} \boldsymbol {\Lambda}) ^ {\top} ] = \tag {5} \\ = (\mathbf {I} - \mathbf {G}) ^ {- 1} \boldsymbol {\Lambda} \mathbb {E} [ (\boldsymbol {U} - \boldsymbol {\mu} _ {\boldsymbol {U}}) (\boldsymbol {U} - \boldsymbol {\mu} _ {\boldsymbol {U}}) ^ {\top} ] ((\mathbf {I} - \mathbf {G}) ^ {- 1} \boldsymbol {\Lambda}) ^ {\top} = \\ = \left(\left(\mathbf {I} - \mathbf {G}\right) ^ {- 1} \boldsymbol {\Lambda}\right) \mathbf {C} _ {\boldsymbol {U}} \left(\left(\mathbf {I} - \mathbf {G}\right) ^ {- 1} \boldsymbol {\Lambda}\right) ^ {\top}, \\ \end{array}
$$
where  $\pmb{\mu}_{\pmb{X}} = (\mathbf{I} - \mathbf{G})^{-1}\pmb{\Lambda}\pmb{\mu}_{\pmb{U}}$ , and  $\mathbf{C}_U$  is the covariance matrix of exogenous variables  $\pmb{U}$ .
<sup>1</sup>In the paper by Vaswani et al. (39) the weight matrix is  $\mathbf{W}_{QK} = \mathbf{W}_Q\mathbf{W}_K^\top /\sqrt{d_K}$ , where the weight matrices  $\mathbf{W}_Q$  and  $\mathbf{W}_K$  are learned explicitly and  $d_{K}$  is the number of columns in these matrices.
An attention layer in a Transformer architecture, estimates an attention matrix  $\mathbf{A}$  and a values matrix  $\mathbf{V}$  from embeddings  $\mathbf{Y}$  of an input sequence of symbols  $s = [s_1,\dots ,s_n]^\top$ . The output embeddings is calculated by  $\mathbf{Z} = \mathbf{A}\mathbf{V}$ . It is important to note that during self-attention training, an inductive bias is used in the form of learning separate representations for the attention matrix  $\mathbf{A}$  and values matrix  $\mathbf{V}$ . It can be seen that the  $j$ -th element of the  $i$ -th row-vector,  $\mathbf{V}(i,j)$  is the projection of the symbol  $s_i$  input embedding  $\mathbf{Y}(i,\cdot)$  on column vector  $\mathbf{W}_V(\cdot ,j)$ . In other words, an embedding is projected on the direction of vector  $\mathbf{W}_V(\cdot ,j)$ . Thus, each column of  $\mathbf{V}$  can be seen as a projection of the input embeddings. Moreover, the attention matrix is shared with all projections for calculating the output embeddings. That is, for the  $j$ -th projection the output embedding is
$$
\mathbf {Z} (\cdot , j) = \mathbf {A V} (\cdot , j). \tag {6}
$$
Denote  $\mathbf{Z}_j \equiv \mathbf{Z}(\cdot, j)$  and  $\mathbf{V}_j \equiv \mathbf{V}(\cdot, j)$  corresponding to projection  $j$ . Covariance of  $\mathbf{Z}_j$  is
$$
\begin{array}{l} \mathbf {C} _ {\boldsymbol {Z} _ {j}} = \mathbb {E} \left[ \left(\boldsymbol {Z} _ {j} - \boldsymbol {\mu} _ {\boldsymbol {Z} _ {j}}\right) \left(\boldsymbol {Z} _ {j} - \boldsymbol {\mu} _ {\boldsymbol {Z} _ {j}}\right) ^ {\top} \right] \\ = \mathbf {A} \mathbb {E} \left[ \left(\boldsymbol {V} _ {j} - \boldsymbol {\mu} _ {\boldsymbol {V} _ {j}}\right) \left(\boldsymbol {V} _ {j} - \boldsymbol {\mu} _ {\boldsymbol {V} _ {j}}\right) ^ {\top} \right] \mathbf {A} ^ {\top} = \tag {7} \\ = \mathbf {A} \mathbf {C} _ {V _ {j}} \mathbf {A} ^ {\top}, \\ \end{array}
$$
where  $\mathbf{C}_{V_j}$  is the covariance matrix of input embeddings projected on vector  $\mathbf{W}_V(\cdot ,j)$ . Since columns of  $\mathbf{V}$  describe different projections corresponding to independent columns of  $\mathbf{W}_V$ , and from the central limit theorem it follows that  $\mathbf{C}_{\mathbf{V}}\rightarrow \mathbf{I}$ , which leads to
$$
\mathbf {C} _ {\mathbf {Z}} = \mathbf {A} \mathbf {A} ^ {\top}. \tag {8}
$$
How does attention matrix  $\mathbf{A}$  and a values vector  $V_{j}$  translate into a structural causal model? We interpret the attention matrix  $\mathbf{A}$  and values vector  $\mathbf{V}$  as components of an SCM, such that the covariance matrix of its endogenous variables is equal to the covariance matrix of the attention layer's output,  $\mathbf{C}_{\mathbf{X}} = \mathbf{C}_{\mathbf{Z}}$ . From Equation 5 and Equation 8,
$$
\mathbf {A} \mathbf {A} ^ {\top} = \left(\left(\mathbf {I} - \mathbf {G}\right) ^ {- 1} \boldsymbol {\Lambda}\right) \mathbf {C} _ {\boldsymbol {U}} \left(\left(\mathbf {I} - \mathbf {G}\right) ^ {- 1} \boldsymbol {\Lambda}\right) ^ {\top}. \tag {9}
$$
The attention matrix  $\mathbf{A}$  measures the total effect each variable has on another variable, for any projection  $\mathbf{V}(\cdot ,j)$  of the input embeddings. Similarly, matrix  $(\mathbf{I} - \mathbf{G})^{-1}\mathbf{\Lambda}$  in an SCM measures the effect one variable has on another variable through all directed paths in the causal graph (Equation 4) for any value of the exogenous variables  $U$ . Thus, form Equation 9 matrix  $\mathbf{A}$  is analogous to  $(\mathbf{I} - \mathbf{G})^{-1}\mathbf{\Lambda}$  and  $\mathbf{V}(\cdot ,j)$  to value assignment for  $U$ , which serves as context. Both  $\mathbf{A}$  and  $(\mathbf{I} - \mathbf{G})^{-1}\mathbf{\Lambda}$  are not required to be triangular as latent confounders may introduce spurious associations, such that variables might be required to estimate the values of variables earlier in the topological order. In addition, the symbols in the sequence  $S$  are not assumed to be topologically ordered.
4.2 Attention-based Causal-Discovery (ABCD)
In Section 4.1 we show that self-attention learns to represent pairwise associations between symbols of an input sequence, for which there exists a linear-Gaussian SCM that has the exact same pair-wise associations between its nodes. However, can we recover the causal structure  $\mathcal{G}$  of the SCM solely from the weights of a pre-trained attention layer? We now present the Attention-Based CausalDiscovery (ABCD) method for learning an equivalence class of the causal graph that underlies a given input sequence.
Causal discovery (causal structure learning) from observed data alone requires placing certain assumptions. Here we assume the causal Markov (24) and faithfulness (35) assumptions. Under these assumptions, constraint-based methods use tests of conditional independence (CI-tests) to learn the causal structure (27; 35; 8; 7; 30; 31; 22). A statistical CI-test is used for deciding if two variables are statistically independent conditioned on a set of variables. Commonly, partial correlation is used for CI-testing between continuous, normally distributed variables with linear relations. This test requires only a pair-wise correlation matrix (marginal dependence) for evaluating partial correlations (conditional dependence). We evaluate the correlation matrix from the attention matrix. From Equation 8 the covariance matrix of output embeddings is  $\mathbf{C_Z} = \mathbf{A}\mathbf{A}^\top$ , and (pairwise) correlation coefficients are  $\rho_{i,j} = \mathbf{C_Z}(i,j) / \sqrt{\mathbf{C_Z}(i,i)\mathbf{C_Z}(j,j)}$ . Unlike kernel-based CI tests (1; 11; 13; 14; 37; 48), we do not need to explicitly define or estimate the kernel, as it is readily
available by a single forward-pass of the input sequence in the Transformer. This implies the following. Firstly, our CI-testing function is inherently learned during the training stage of a Transformer, by that enjoying the efficiency in learning complex models from large datasets. Secondly, since attention is computed for each input sequence uniquely, CI-testing is unique to that specific sequence. That is, conditional independence is tested under the specific context of the input sequence.
Finally, the learned causal graph represents an equivalence class in the form of a partial ancestral graph (PAG) (29; 47), which can also encode the presence of latent confounders. A PAG represents a set of causal graphs that cannot be refuted given the data. There are three types of edge-marks (at some node  $X$ ): an arrow-head  $\rightarrow X$ , a tail  $\rightarrow X'$ , and circle  $\rightarrow 0X'$  which represent an edge-mark that cannot be determined given the data. Note that reasoning from a PAG is consistent with every member in the equivalence class it represents.
What is the relation between the equivalence class learned by ABCD and the causal graph underlying the input sequence  $s = [s_1, \ldots, s_n]$ ?
Definition 3 (SCM-consistent probability-space). Let  $\Omega_{n} = \{[s_{1},\ldots ,s_{n}]:s_{i}\in S\}$  be a sample space, where  $S$  is a countable set of symbols. Let  $\mathcal{F}_n$  be an event space, and  $P_{n}$  be a probability measure. For a finite  $n\in \mathbb{N}$ , the probability space  $(\Omega_n,\mathcal{F}_n,P_n)$  is called SCM-consistent if every sequence  $s\in \Omega_{n}$ , was generated by a corresponding SCM with graph  $\mathcal{G}_{s}$ .
Assumption 1. Attention mechanism  $\mathcal{M}$  is trained on sequences sampled from an SCM-consistent probability space, such that each input symbol is uniquely represented by its output embedding.
Many training methods, such as MLM (9), satisfy Assumption 1 as they train the network to predict input symbols using their corresponding output embeddings in the deepest layer. This assumption is required for relating the causal graph learned from the deepest attention layer to the causal graph underlying the input sequence.
Theorem 1. If  $s$  is a sequence from an SCM-consistent probability space used to train  $\mathcal{M}$ ,  $\mathcal{G}_s$  is the causal graph of the SCM from which  $s$  was generated, and  $\mathbf{A}$  the attention matrix for input sequence  $s$  in the deepest layer, then  $\mathcal{G}_s$  is in the equivalence class learned by a sound constraint-based algorithm that uses covariance matrix  $\mathbf{C_Z} = \mathbf{A}\mathbf{A}^\top$  for partial correlation CI-testing.
Corollary 1. Nodes  $X_{i}$  and  $X_{j}$  are  $d$ -separated in  $\mathcal{G}_{\mathbb{Z}}$  conditioned on  $\pmb{X}^{\prime} \subset \pmb{X}$  if and only if symbols  $s_{i}$  and  $s_{j}$  are  $d$ -separated in  $\mathcal{G}_{\mathbf{s}}$  conditioned on  $\pmb{s}^{\prime}$ , a subset containing symbols corresponding to  $\pmb{X}^{\prime}$ .
Corollary 1 is useful for cases in which the graph over the input symbols  $\mathcal{G}_s$ , as well as the structural equation model described by  $\mathcal{G}_Z$ , does not have a causal meaning. In these cases, one can reason about conditional independence relations among input symbols.
4.3 Multiple Attention Layers and Multi-Head Attention as a Structural Causal Model
Commonly, the encoder module in the Transformer architecture (39) consists of multiple self attention layers executed sequentially. A non-linear transformation is applied to the output embeddings of one attention layer before feeding it as input to the next attention layer. The non-linear transformation is applied to each symbol independently. In addition, each attention layer consists of multiple self-attention heads, processing the input in parallel using head-specific attention matrix. The output embedding of each symbol is then linearly combined along the heads. It is important to note that embedding of each symbol is processed independently of embeddings of other symbols. The only part in which an embedding of one symbol is influenced by embeddings of other symbols is the multiplication by the attention matrix.
Thus, a multi-layer, multi-head, architecture can be viewed as a deep graphical model, where exogenous nodes of an SCM are estimated by a previous attention layer. The different heads in an attention layer learn different graphs for the same input, where their output embeddings are linearly combined. This can be viewed as a mixture model. We recover the causal graph from the last (deepest) attention layer, and treat earlier layers as context estimation, which is encoded in the values of the exogenous nodes (Figure 1).
Figure 1: A network with three attention layers. Early attention layers (0 and 1) estimate context, which is encoded in the values of exogenous nodes for the last layer (2).
4.4 Limitations of ABCD