Causal-Based Supervision of Attention in Graph Neural Network: A Better and Simpler Choice towards Powerful Attention

1 Introduction

Graph-structured data is widely used in real-world domains, such as social networks [Zhang and Chen, 2018], recommender systems [Wu et al., 2022b], and biological molecules [Gilmer et al., 2017]. The non-euclidean nature of graphs has inspired a new type of machine learning model, Graph Neural Networks (GNNs) [Kipf and Welling, 2016; Defferrard et al., 2016; Du et al., 2022]. Generally, GNN iteratively updates features of the center node by aggregating those of its neighbors and has achieved remarkable success across various graph analytical tasks. However, the aggregation of features between unrelated nodes has long been an obstacle for GNN, keeping it from further improvement.

Recently, Graph Attention Network (GAT) [Veliˇckovi´c et al., 2017] pioneered the adoption of the attention mechanism, a well-established method with proven effectiveness in deep learning [Vaswani et al., 2017], into the neighborhood aggregation process of GNNs to alleviate the issue. The key concept behind GAT is to adaptively assign importance to each neighbor during the aggregation process. Its simplicity and effectiveness have made it the most widely used variant of GNN. Following this line, a myriad of attention-based GNNs have been proposed and have achieved state-of-the-art performance in various tasks [Sun et al., 2021; Zhang et al., 2022; Ying et al., 2021; Brody et al., 2021].

Nevertheless, despite the widespread use and satisfying results, in the past several years, researchers began to rethink if the learned attention functions are truly effective [Kim and Oh, 2022; Wang et al., 2019; Wu et al., 2022a; Knyazev et al., 2019; Liu et al., 2021; Wang et al., 2021]. As we know, most existing attention-based GNNs learn the attention function in a weakly-supervised manner, where the attention modules are simply supervised by the final loss function, without a powerful supervising signal to guide the training process. And the lack of direct supervision on attention might be a potential cause of a less robust and generalizable attention function against real-world noisy graphs [Sui et al., 2022; Knyazev et al., 2019; Li et al., 2022; Wang et al., 2019; Ye and Ji, 2021; Huang et al., 2023]. To address this problem, existing work enhances the quality of attention through auxiliary regularization terms (supervision). However, concerns have been raised that these methods often rely heavily on human-specified prior assumptions about a specific task, which limits their generalizability [You et al., 2020; Wu et al., 2022a]. Additionally, the auxiliary regularization is formulated independently of the primary prediction task, which may disrupt the original optimization target and cause the model to "switch" to a different objective function during training [Kim and Oh, 2022; You et al., 2020].

Recently, causal inference [Pearl, 2009] has attracted many researchers in the field of GNNs by utilizing structural causal model (SCM) [Cinelli et al., 2019] to handle distribution shift [Zhao et al., 2022] and shortcut learning [Feng et al., 2021]. In this paper, we argue that the tool of causal inference has also shed light on a promising avenue that could supervise and improve the quality of GNN's attention directly, while in the meantime we will not make any assumptions about specific tasks or models, and the supervision signal for attention implicitly aligns well with the primary task. Before going any deeper, we first provide a general schema for the SCM of attention-based GNNs in Figure 1, which uses nodes to represent variables and edges to indicate causal relations between variables. As we can see, after a high-level abstraction, there are only three key factors in SCM, including the node features X, attention maps A, and the model's final prediction Y . Note that in causal language, X, A, and Y also denotes the context, treatment, and outcome respectively. For edges in SCM, the link X → A represents that the attention generation relies on the node's features (i.e., context decides treatment). And links (X,A) → Y indicate that the model's final prediction is based on both the node's features X and the attention A (i.e., the final outcome is jointly determined by both context and treatment).

In order to provide a direct supervision signal and further enhance the learning of attention functions, the first step would be finding a way to measure the quality of attention (i.e., quantifying what to improve). Since there are no unified criteria on the way of measurement, researchers usually propose their own solution according to the tasks they are facing, and this is very likely to introduce the unfavorable human-intervened prior assumptions [You et al., 2020]. For example, CGAT [Wang et al., 2019] believes that better attention should focus more on one-hop neighbors. While this assumption surely works on homophilous graphs, it will suffer from huge performance degradation in heterophilous scenarios. Our method differs a lot from existing work in that we introduce SCM to effectively decouple the direct causal effect of attention on the final prediction (i.e., link A → Y ), and use such causal effect as a measurement for the quality of attention. In this way, it is the model and data that decide if the attention works well during training instead of human-predefined rules. And this has been shown to be non-trivial in various machine learning fields because what might seem reasonable to a human might not be considered the same way by the model [Kumar et al., 2010; Wang et al., 2022b]. Another drawback of existing attention regularization methods, as previously mentioned, is the deviation from primary tasks. SuperGAT [Kim and Oh, 2022] uses link prediction to improve the attention quality for node classification, but as the author claims in the paper, there is an obvious trade-off between the two tasks. In this paper, we alleviate this problem by directly maximizing the causal effect of attention on the primary task (i.e., strengthening the causal relation A → Y ). Under mild conditions, we can deem the overall optimization is still towards the primary objective, except that we additionally provide a direct and powerful signal for the learning of attention in a fully-supervised manner.

In summary, this paper presents a Causal Supervision for Attention in graph neural networks (abbreviated as CSA in the following paragraphs). CSA has strong applicability because no human-intervened assumptions are made on the target models or training tasks. And the supervision of CSA can be easily and smoothly integrated into optimizing the primary task to performing end-to-end training. We list the main contributions in this paper as follows:

• We explore and provide a brand-new perspective to directly boost GNN's attention with the tool of causality. To the best of our knowledge, this is a promising direction that still remains unexplored.
• We propose CSA, a novel causal-based supervision framework for attention in GNNs, which can be formulated as a simple yet effective external plug-in for a wide range of models and tasks to improve their attention quality.
• We perform extensive experiments and analysis on CSA and the universal performance gain on standard benchmark datasets validates the effectiveness of our design.

2 RelatedWork

Attention-based Graph Neural Networks. Modeling pairwise importance between elements in graph-structured data dates back to interaction networks [Battaglia et al., 2016; Hoshen, 2017] and relational networks [Santoro et al., 2017]. Recently GAT [Veliˇckovi´c et al., 2017] rose as one of the representative work of attention-based GNNs using selfattention [Vaswani et al., 2017]. The remarkable success of GAT in multiple tasks has motivated many works focusing on integrating attention into GNN [Thekumparampil et al., 2018; Zhang et al., 2018; Wang et al., 2022a; Zhang et al., 2020; Gao and Ji, 2019; Hou et al., 2022]. Lee et al. have also conducted a comprehensive survey [Lee et al., 2019] on various types of attention used in GNNs.

Causal Inference in Graph Neural Network. Causality [Pearl, 2014] provides researchers new methodologies to design robust measurements, discover hidden causal structures and confront data biases. A myriad of studies has shown that incorporating causality is beneficial to graph neural network in various tasks. [Zhao et al., 2022] makes use of counterfactual links to augment data for link prediction improvement. [Sui et al., 2022] performs interventions on the representations of graph data to identify the causally attended subgraph for graph classification. [Feng et al., 2021] on the other hand, applies causality to estimate the causal effect of node's local structure to assist node classification.

Improving Attention in GAT. There is a great number of work dedicated to improving attention learning in GAT. [Kim and Oh, 2020] enhances attention by exploiting two attention forms compatible with a self-supervised task to predict edges. [Brody et al., 2021] introduces a simple fix by modifying the order of operations in GAT. [Wang et al., 2019] develops an approach using constraint on the attention weights according to the class boundary and feature aggregation pattern. In addition, causality also plays a role in boosting the attention of GATs recently. [Wu et al., 2022a] estimates the causal effect of edges by intervention and regularizes edges' attention weights according to their causal effects.

3 Preliminaries

We start by introducing the notations and formulations of graph neural networks and their attention variant. Let G = {V, E} represents a graph where V = {vi}ni=0 is the set of nodes and E ∈ V×V is the set of edges. For each node v ∈ V, it has its own neighbor set N(v) = {u ∈ V | (v, u) ∈ E} and its initial feature vector x0v ∈ Rd0 , where d0 is the original feature dimension. Generally, GNN follows the messagepassing mechanism to perform feature updating, where each node's feature representation is updated by aggregating the representations of its neighbors and then combining the aggregated messages with its ego representation [Xu et al., 2018]. Let mlv ∈ Rdl and xlv ∈ Rdl be the message vector and representation vector of node v at layer l, we formally define the updating process of GNN as:

where AGGREGATE and COMBINE are aggregation functions (e.g., mean, LSTM) and combination function (e.g., concatenation), respectively. The design of these two functions is what mostly distinguishes one type of GNN from the other. GAT [Veliˇckovi´c et al., 2017] augments the normal aggregation with the introduction of self-attention. The core idea of self-attention in GAT is to learn a scoring function that computes an attention score for every node in N(v) to indicate their relational importance to node v. In layer l, such process is defined by the following equation:

where (al ,W l ), σ are learnable matrices and activation function (e.g., LeakyReLU) respectively, and ∥ denotes vector concatenation. The attention scores are then normalized across all neighbors vj ∈ N(vi) using softmax to ensure consistency:

Finally, GAT computes a weighted average of the features of the neighboring nodes as the new feature of vi, which is demonstrated as follows:

4 Casual-based Supervision on GNN's Attention

In this section, we first introduce how the causal effect of attention can be derived from the structural causal model of attention-based GNNs. Specifically, this is done with the help of the widely used counterfactual analysis in causal reasoning. After that, with the obtained causal effects, we elaborate three candidate schemes to incorporate with the training of attention-based GNNs to improve their quality of attention.

4.1 Causal Effects of Attention

As previously mentioned, the first step towards improving attention lies in measuring the quality of existing attention. However, since deep learning models usually exhibit as black boxes, it is generally infeasible to directly assess their attention qualities. Existing works mainly address this issue by introducing human priors to build pre-defined rules for some specific models and tasks. Yet, it has been a long debate on whether human-made rules share consensus with deep learning models during training [Kumar et al., 2010; Ribeiro et al., 2016]. Fortunately, the recent rise of causal inference technology has offered effective tools to help us think beyond the black box and analyze causalities between model variables, which leads us to an alternative way to directly utilize the causal effect of attention to measure its quality. Since the obtained causal effects are mainly affected by the model itself, it is a more accurate and unbiased measurement of how well the attention actually learns.

We first give a brief review of the formulation of attentionbased graph neural network in causal languages, as shown in Figure 2(a). The generated attention map A is directly affected by node feature X. And the model prediction Y is jointly determined by both X and A. We denote the inferring process of the model as:

which indicates that model will give value Yx,a if the value of X and A are set to x and a respectively. In order to pursue the attention's causal effect, we introduce the widely-used counterfactual analysis [Pearl, 2022] in causal reasoning.

The core idea of counterfactual causality lies in asking: given a certain data context (node feature X), what the outcome (model prediction Y ) would have been if the treatment (attention map A) had not been the observed value? To answer the imaginary question, we have to manipulate the values of several variables to see the effect, and this is formally termed as intervention in causal inference literature, which can be denoted as do(·). In do(·) operation, we compulsively select a counterfactual value to replace the original factual value of the intervened variable. And once a variable is intervened, its all incoming links in the SCM will be cut off and its value is independently given, while other variables that are not affected still maintain the original value. In our case, for example, do(A = a∗ ) means we demand the attention A to take the non-real value a∗ (e.g., reverse/random attention) so that the link X → A is cut-off and A is no longer be affected by its causal parent X. This process is illustrated in Figure 2(b) and the mathematical formulation is given as:

which indicates that after do(·) operation which changes the value of attention to be a∗ , the output value of the model also changes to Yx,a∗. Finally, let us consider a case where we assign a dummy value ˜a to the attention map so that for each ego node, all its neighbors share the same attention weights, the feature aggregation of the graph attention model will then degrade to an unweighted average. In this case, according to the theory of causal inferences [VanderWeele, 2016], the Total Direct Effect (TDE) of attention to model prediction can be obtained by computing the differences of model outcome Yx,a and Yx,˜a, which is formulated as follows:

It is worth noting that the induction of attention's causal effect does not exert any assumptions and constraints, which lays a solid foundation for our wide applicability on any graph attention models.

4.2 Supervision on Attention with Causal Effects

We have already demonstrated the derivation of attention's causal effect in the previous section. In this part, we will discuss how to utilize the obtained causal effect for attention quality improvement. Previous works that make use of an auxiliary task to regularize attention usually suffered from performance trade-off between the primary task and auxiliary task. In this work, we alleviate this problem by directly maximizing the causal effect of attention on the primary task. The overall schema of this part is shown in Figure 3.

Consider a simple case where we conduct the node classification task with a standard L-layer GAT. For each layer l, we have the node representations Xl-1 ∈ Rn×dl-1 from the previous layer as input. Then, we perform feature aggregation and updating with factual attention map Al to obtain the factual output feature Xl = f(Xl-1 ,Al ). Similarly, when we intervene the attention maps of layer l (e.g., assigning dummy values using do(·) operation), we can get a counterfactual output feature ˆXl . We further employ a learnable matrix Wl ∈ Rc×dl (c denotes the number of classes) to get the node's factual predicted label Y l pred and counterfactual predicted label ˆ Y l pred using the corresponding features from layer l. Therefore, the causal effect of attention at layer l is obtained as: Y l pred - ˆ Y l pred . To this end, we can use the causal effect as a supervision signal to explicitly guide the attention learning process. The new objective of the CSA-assisted GAT model can be formulated as:

where y is the ground-truth label, Lce is the cross-entropy loss, λl is the coefficient to balance training, and Lothers represents the original objective such as standard classification loss. Note that Equation.(4) is a general form of CSA where we compute additional losses for each GAT layer to supervise attention directly. However in practice it is not necessary, and we found that simply selecting one or two layers is enough for CSA to bring satisfying performance improvement.

Moreover, since our aim is to boost the quality of attention, it is not necessary to estimate the correct causal effect of attention using dummy values. Instead, a strong counterfactual baseline might even be helpful for the attention quality improvement. We hereby further propose three heuristic counterfactual schemes and test them in our experiments. We note that the exact form of how counterfactual is achieved is not limited, and our goal here is just to set the ball rolling.

Scheme I: In the first scheme, we utilize the uniform distribution to generate the counterfactual attention map. Specifically, the counterfactual attention is produced by

where e and f are the lower and upper boundaries. In this case, the generated counterfactual could vary from very bad (i.e., mapping all unrelated neighbors) to very good (i.e., mapping all meaningful neighbors). This is a similar process to the Randomized Controlled Trial [Stolberg et al., 2004] where all possible treatments are enumerated. We hope that maximizing causal effects computed over all possible treatments can lead to a robust improvement of attention.

Scheme II: Scheme I is easy and straightforward to apply. However due to its randomness, a possible concern is that if most of the generated counterfactual attentions are inferior to the factual one, then we will only have very small gradient on attention improvements. Therefore we are actually motivated to find "better" counterfactuals to spur the factual one to evolve. Heuristically, given that MLP is a strong baseline one several datasets (e.g., Texas, Cornell, and Wisconsin), we employ an identical mapping to generate the counterfactual attention, which only attends to the ego-node instead of neighbors. Specifically, the counterfactual attention map is equal to the identity matrix I:

Scheme III: Our last schema can be considered as an extension of Schema II. Since the fast development of GAT family has introduced to us some variants that already outperform MLP in many datasets, using counterfactuals derived from the behavior of MLP does not seem to be a wise choice for these GAT variants to improve their attentions. Inspired by the self-boosting concept [Pi et al., 2016] widely used in machine learning, we leverage the historical attention map as the counterfactual to urge the factual one keep refining itself. The specific formulation is written as follows:

where Ahist denotes the historical attention map (e.g., the attention map from the last update iteration).

5 Experiment

In this section, we conduct extensive node classification experiments to evaluate the performance of CSA. Specifically, we 1) validate the effectiveness of CSA on three popular GAT variants using a wide range of datasets, including both homophily and heterophily scenarios; 2) compare CSA with other attention improvement baselines of GAT to show the superiority of our method; 3) show that CSA can induce better attention that improve the robustness of GAT; 4) test CSA's sensitivity to hyper-parameters; 5) analyze the influences of CSA in feature space; and 6) examine the performances of some special cases of CSA.

5.1 Datasets

For heterophily scenario, we select seven standard benchmark datasets: Wisconsin, Cornell, Texas, Actor, Squirrel, Chameleon, and Crocodile. Wisconsin, Cornell, and Texas collected by Carnegie Mellon University are published in WebKB 1 . Actor [Pei et al., 2020] is the actor-only subgraph sampling from a film-director-actor-writer network. Squirrel, Chameleon, and Crocodile are datasets collected from the English Wikipedia [Rozemberczki et al., 2021]. We summarize the basic attributions for each dataset in Table 1. H(G) is the homophily ratio [Zhu et al., 2020], where H(G) → 1 represents extreme homophily and vice versa.

1 http://www.cs.cmu.edu/ webkb/

For homophily scenario, two large datasets released by Open Graph Benchmark (OGB) 2 [Hu et al., 2020]: ogbnproducts and ogbn-arxiv, are included in our experiments, together with two small-scale homophily graph datasets: Cora and Citeseer [McCallum et al., 2000]. Similarly, the attribution of the dataset is summarized in Table 2.

5.2 Experimental Setup

We employ popular node classification models in our experiments as the baselines: GCN [Kipf and Welling, 2016], GAT [Veliˇckovi´c et al., 2017], SGCN [Wu et al., 2019], FAGCN [Bo et al., 2021], GPR-GNN [Chien et al., 2020], H2GCN [Zhu et al., 2020], WRGAT [Suresh et al., 2021], APPNP [Gasteiger et al., 2018] and UniMP [Shi et al., 2020]. We also present the performance of MLPs, serving as a strong non-graph-based baseline. Due to page limit, we only select four models: GAT, FAGCN, WRGAT and UniMP to examine the effectiveness of CSA. These models ranges from the classic ones to the latest ones, and are considered as representatives for state-of-the-art node classification models. One thing to be noted here is that for all these models, we implement CSA only in their first layers to avoid excessive computational cost.

In our experiments, each GNN is run with the best hyperparameters if provided. We set the same random seed for each model and dataset for reproducibility. The reported results are in the form of mean and standard deviation, calculated from 10 random node splits (the ratio of train/validation/test is 48%/32%/20% from [Pei et al., 2020]). Our experiments are conducted on a GPU server with eight NVIDIA DGX A100 graphics cards, and the codes are implemented using Cuda Toolkit 11.5, PyTorch 1.8.1 and torch-geometric 2.0.1.

2 https://ogb.stanford.edu/

5.3 Performance Analysis

Table 1 and Table 2 provide the test accuracy of different GNNs in different variants of CSA in the supervised node classification task. A graph's homophily level is the average of its nodes' homophily levels. CSA achieves the best in terms of the vanilla one across all datasets. In particular, the highest improved datasets are Texas, Wisconsin, and Cornell. By observing the performance of MLPs, we can see that the common ground of those three datasets contains distinguishable features and a large part of non-homophilous edges. In the meanwhile, the performance of CSA is proportional to the modeling ability. The mechanism behind CSA is to extend the causal effect between nodes representation and final prediction. Therefore, CSA owns limited performance when the node's representations are chaotic. Our experiments highlight that I) The model, which is already better than MLP, does not improve much in CSA-II, while CSA-III improves it relatively more. This is because in those datasets, the graph structure can provide meaningful information, so that CSAIII have more advantages. II) The dataset, which has distinctive features indicated by the performance of MLPs, is more satisfied CSA-II. Similarly, in this scene, the features can be more informative. III) The random strategy (CSA-I) relatively inferior to others, since the distribution is hard to control and tend to generate worst attention map, whereby the regularization is vanished.

(a) Texas Accuracy

(b) Cornell

5.4 Comparison with Attention Promotion Baselines

We here apply multiple attention promotion baselines: CAL [Sui et al., 2022], CAR [Wu et al., 2022a], Super [Kim and Oh, 2022], Constraint interventions [Wang et al., 2019] and the result is shown in Figure 4. Among them, CAL is a method for graph property prediction that relies on an alternative formulation of causal attention with interventions on implicit representations. We adapted CAL for node classification by removing its final pooling layer. Super is well-known as SuperGAT, a method seeking to constrain node feature vectors through a semi-supervised edge prediction task. CAR aligns the attention mechanism with the causal effects of active interventions on graph connectivity in a scalable manner. Constraint method has two auxiliary losses: graph structurebased constraint and class boundary constraint. The results on four datasets are shown in Figure 4. While CAL, CAR, and CSA have related goals of enhancing graph attention using concepts from causal theory, CAL uses abstract perturbations on graph representation to perform causal interventions, and CAR employs an edge intervention strategy that enables causal effects to be computed scalable, while our method does not exert any assumptions and constraints on GATs, compared with CAL and CAR. Therefore, CSA tends to own good generalization ability. In terms of SuperGAT and Constraint method, since there is a trade-off between node classification and regularization. For example, SuperGAT implies that it is hard to learn the relational importance of edges by simply optimizing graph attention for link prediction.

5.5 CSA Provides Robustness

In this section, we systematically study the performance of CSA on robustness against the input perturbations including feature and edge perturbations. Following [Stadler et al., 2021], we conduct node-level feature perturbations by replacing them with the noise sampled from the Bernoulli distribution with p = 0.5 and edge perturbations by stochastically generating the edges. According to the performance shown in Figure 5, CSA produces robust performance on input perturbations. Figure 5 demonstrate that CSA in higher noise situations achieves more robust results than in lower noise scenes on both node and edge perturbations with perturbation percentages ranging from 0% to 40%.

5.6 Hyper-parameter Analysis

We analyze the sensitivity of λ and plot node classification performance in Figure 6. For λ, there is a specific range that maximizes test performance in all datasets. The performance in Texas is the highest when λ is 0.4, but the difference is relatively small compared to Cora. We observe that there is an optimal level of causal supervision for each dataset, and using too large λ degrades node classification performance. Since Cora owns friendly neighbors, its performance is less sensitive than Texas. Based on this, we can also see that Texas relatively needs a larger regularization.

5.7 Pairwise Distance among Classes

To further evaluate whether the good performance of CSA can be contributed to the mitigation of lacking supervision, we visualize the pairwise distance of the node representations among classes learned by CSA and vanilla GAT. Following [Stadler et al., 2021], we calculate the Mean Average Distance (MAD) with cosine distance among node representations for the last layer. The larger the MAD is, the better the node representations are. Results are reported in Figure 7. It can be observed that the node representations learned by CSA keep having a large distance throughout the optimization process, indicating relief of lacking supervision issues. On the contrary, GAT suffers from severely indistinguishable representations of nodes.

5.8 Comparison with Heuristic Causal Strategies.

To validate the effectiveness of CSA, we compare it with two heuristic causal designs (Last and Pure) that 1) directly estimate the total causal effect by subtracting between the model's and the counterfactual's output in the final layer; 2) replace the attention with the static aggregate weights (i.e., each node allocates the same weight). The results are shown in Table 3. We observe that their performance outperforms vanilla one, but is still inferior to ours. In terms of Last, the major difference is whether to explicitly estimate causal effect or not. In our framework, we plug the MLPs into the hidden layers to precisely estimate the causal effect for each layer. Regarding Pure, our strategy can provide more strong baselines, leading to better regularization.

6 Conclusion

We introduced CSA, a counterfactual-based regularization scheme that can be applied to graph attention architectures. Unlike other causal approaches, we first built the causal graph of GATs in a general way and do not impose any assumptions and constraints on GATs. Subsequently, we introduce an efficient scheme to directly estimate the causal effect of attention in hidden layers. Applying it to both homophilic and heterophilic node-classification tasks, we found accuracy improvements and robustness in almost all circumstances. We performed three variants of counterfactual attention strategies and found that they can adapt to different situations, respectively.

Contribution Statement

Hongjun Wang and Jiyuan Chen contributed equally to the work. Lun Du and Xuan Song are the corresponding authors. This work is done during Hongjun Wang's internship at Microsoft Research Asia.