Rationale refers to a subset of the input features that play a crucial role in model predictions for downstream tasks [8,18,37,49,55,77,84]. In the context of graph machine learning, graph rationale is defined as a subgraph of the input graph containing the most taskrelevant semantics. The application of graph rationale is broad; for example, it can greatly enhance model performance for graph-level tasks [33,55,61,67,81] by identifying the key components of the input graph. Additionally, the discovery of rationales can improve model explainability [37], as it highlights the parts of the input graph that significantly contribute to the final prediction.
Existing graph rationalization solutions [37,55] employ a trainable augmenter to execute the rationale/environment decomposition. In this process, a node/edge mask is generated by the augmenter to decompose the given graph into a rationale graph and an environment graph. Inspired by the content-style decomposition [27], the key idea of graph rationalization is to preserve the utility of the graph rationale when combined with changing environment graphs (see Figure 1). To achieve this, a technique named intervention is used, where the environment graph interacts with the rationale graph.
The intervention mechanism (named intervener) is essential in the graph rationalization process, as it must accurately represent the interaction between the rationale and the environment. Intuitively, the intervener should work in an adversarial behavior against the augmenter mentioned above, a point largely overlooked in the existing literature. If the intervener is more powerful, it can capture more detailed interactions between the rationale and environment subgraphs. Given such a powerful intervener, the augmenter is compelled to minimize these interactions between the graph rationale and the environment to obtain a "purer" graph rationale.
Unfortunately, existing works develop interveners in a coarse and non-parametric manner. After performing rationale/environment decomposition on the graph data, they compute graph-level embeddings for the rationale and environment subgraphs. The intervention is then designed as an interaction between these graph-level embeddings. For example, [37] adds the environment embedding into the rationale embedding as the intervened rationale embedding; [55] defines the intervened prediction as the Hadamard product between the predictions based on the rationale subgraph and the environment subgraph. We argue that such a graph-level non-parametric intervention is insufficient to fully represent the interaction between the rationale and environment graphs effectively.
In response to this limitation, we propose a fine-grained, parametric intervention named FIne-grained Graph rationalization (FIG). FIG draws inspiration from the self-attention module in the Transformer model, which captures interactions between input tokens. Building upon insights from Transformer [51] and its linear variant Linformer [54], FIG formulates the interaction between the rationale and environment subgraphs at the node-level or the virtual node-level. The two variants are named FIG-N and FIG-VN. Furthermore, to maximize the effectiveness of the intervention, we formulate a min-max game involving the node encoder, augmenter, intervener, and predictor, forcing the rationale subgraph to be as informative as possible.
We conduct comprehensive experiments on 7 graph-level benchmarks to evaluate the proposed approach and compare FIG-N/VN against 13 state-of-the-art baseline methods. The results demonstrate that FIG and its variants outperform the baseline methods, validating their superior performance. Our primary contributions in this paper are summarized as follows: (1) we address the graph rationalization problem via a fine-grained node/virtual node-level model, FIG ; (2) a min-max game is proposed to boost the effectiveness of the proposed intervener; (3) extensive experiments covering 13 baseline methods and 7 real-world datasets demonstrate the efficacy of the proposed method.
Section 2.1 introduces the notations used throughout this paper. Then, the classic Transformer architecture is introduced in Section 2.2, whose self-attention module is an important building block of our model. Last but not least, we introduce the overall ideas of existing graph rationalization works in Section 2.3.
We adopt the following notation conventions: bold uppercase letters for matrices and tensors (e.g., A), bold lowercase letters for column vectors (e.g., u), lowercase and uppercase letters in regular font for scalars (e.g., 𝑑, 𝐾), and calligraphic letters for sets (e.g., T ). To index vectors/matrices/tensors, we follow the syntax from NumPy (0-based). Specifically, A[𝑝, :] and A[:, 𝑞] represent the 𝑝-th row and the 𝑞-th column of matrix A respectively; A[𝑝, 𝑞] represents the entry at the 𝑝-th row and the 𝑞-th column. Similarly, u[𝑝] denotes the 𝑝-th entry of vector u. In addition, the slicing syntax for vectors/matrices/tensors is used. For example, for a matrix A, A[𝑖 : 𝑗, :] denotes rows from the 𝑖-th row (included) to the 𝑗-th row (excluded) and A[:, : 𝑘] denotes all the columns before the 𝑘-th column. ⊤ denotes the transpose of matrices and vectors. ⊙ represents the Hadamard product, and • denotes function composition. We use || for the concatenation operation, and the specific dimension of concatenation will be clarified based on the context.
An attributed graph can be represented as G = (A, X, E), where A ∈ R 𝑛×𝑛 is the adjacency matrix, X ∈ R 𝑛×𝑑 𝑋 is the node feature matrix, and E ∈ R 𝑛×𝑛×𝑑 𝐸 is the edge feature tensor. Here, 𝑛 denotes the number of nodes, and 𝑑 𝑋 (𝑑 𝐸 ) represents the dimensions of node (edge) features, respectively. This paper assumes the node and edge feature dimensions are the same (i.e., 𝑑 𝑋 = 𝑑 𝐸 = 𝑑) for brevity; if they differ, a fully connected layer can map them into a shared feature space. Our main focus in this paper is on graph property prediction tasks. The ground truth of a graph is represented by 𝑦.
The core modules of the Transformer architecture [51] are the self-attention layer and the feed-forward network layer. Given the input as a sequence of symbol representations H ∈ R 𝑛×𝑑 𝐻 , it is first transformed into the query, key, and value matrices as
where
For brevity, we set 𝑑 𝐻 = 𝑑 𝑄 = 𝑑 𝐾 = 𝑑 𝑉 = 𝑑. Then, the self-attention module is,
Typically, the non-linearity 𝜎 is Softmax. The feed-forward network (FFN) updates the representations H as:
Optional techniques such as normalization [3,24], dropout [47], and multi-head attention [51] are omitted here for brevity. While Transformers were initially designed for sequence or set data with positional encoding, numerous techniques have since been introduced to adapt Transformers for graph data. Interested readers are referred to [41] for a detailed taxonomy. Interestingly, it is well-known in both the graph learning [9] and natural language processing communities [54,78] that, from the message-passing perspective, the key idea of the Transformer architecture is to construct a weighted complete graph, whose adjacency matrix is P = 𝜎 QK ⊤ √ 𝑑 with complexity 𝑂 (𝑛 2 ).
The graph rationale is a subgraph that encodes most downstream task-relevant semantics. A typical example is the functional groups in polymer graphs [37,55], which fundamentally determines the chemical property of polymers. Mathematically, a given graph is decomposed into a rationale graph and an environment graph:
Commonly, the graph embeddings on G ra and G env are computed as h ra and h env . To ensure the rationale graph is invariant w.r.t. the prediction results when combined with different environments, a utility loss is minimized given the rationale embedding h ra intervened by the environment embedding henv , i.e.,
Here, henv could be either from the same graph (i.e., henv = h env ), or could be environment embeddings from other graphs, such as those in the batch.
A key difference among existing methods lies in the intervention operation, of which we mention two: (1) GREA [37] designs the intervention as sum, i.e., h ra + henv and (2) DIR [55] designs an element-wise intervention: 𝜃 pred (h ra ) ⊙ Sigmoid(𝜃 pred ( henv )), where ⊙ is the Hadamard product and 𝜃 pred is a predictor. The above intervention is graph-level because h ra and fenv are graph embeddings. In Figure 2a, an overview of the GREA [37] is presented. As a comparison, we aim to design the intervention at finer grains (e.g., node level) to handle the interaction between rationale and environment graphs, which will be detailed as follows.
This section introduces the proposed method, FIG. Figure 2b provides an overview of FIG, highlighting its four parametric modules: the encoder, augmenter, intervener, and predictor. Encoder. The encoder, denoted as 𝜃 enc : G → R 𝑛×𝑑 , accepts a graph data as input and outputs a node embedding matrix. There are various graph encoders available, such as graph neural networks (GNNs) [57] and graph Transformers [41]. From the methodology perspective, the encoder module is not the main contribution of this paper, so in this section, we do not specify a specific graph encoder 𝜃 enc . Predictor. The predictor, denoted as 𝜃 pred : R 𝑑 → R 𝑐 takes as input a graph embedding and outputs a prediction vector/scalar. For graph regression tasks, 𝑐 = 1; for graph classification tasks, 𝑐 is the number of classes. A typical predictor is a multi-layer perceptron (MLP) with appropriate activation functions. Details of the encoder and predictor in our implementation are presented in Section 4.
In subsequent subsections, we will elaborate on the augmenter and intervener, two essential modules. Their detailed designs derive two variants of the proposed FIG.
Node-level augmenter. The augmenter is a critical module of the proposed FIG. For the node-level variant, termed FIG-N, the augmenter's primary function is decomposing the node set into two distinct subsets: rationale nodes and environment nodes. This decomposition is operated by parameterizing the node-level augmenter as a learnable node partitioner, denoted by 𝜃 aug-N ,
whose input is the node embedding matrix H ∈ R 𝑛×𝑑 , and its output is a partition vector m ∈ [0, 1] 𝑛 . MLP is a multi-layer perceptron. Each entry within m, such as m[𝑖], denotes the probability of the 𝑖-th node being categorized as a rationale node.
For the partition vector m, its top-𝐾 entries are indexed as idx ra = argtopK(m) which is used to index the rationale nodes from the node embedding matrix H; naturally, the remaining nodes are categorized as the environment nodes whose indices are idx env = {1, . . . , 𝑛}/idx ra . 𝐾 is a hyper-parameter whose impact is studied in Section 4.5. Also, in our implementation, we use a soft argtopK operation to maintain differentiability, whose details are in Section B.
Using the indices mentioned above, rationale and environment embeddings, denoted as H ra and H env , respectively, can be extracted from the node embedding matrix H:
Node-level intervener. The design of the fine-grained intervener draws inspiration from the Transformer architecture [51]. Explicitly, the node-level intervener 𝜙 is presented as,
In this representation, the operator || concatenates along the first dimension of the matrices H ra and H env . We dub the Eqs. ( 1)-( 3) as Transformer and P is the intermediate attention matrix from the self-attention layer (Eq. (6b)). Here, the self-attention module models the interactions between the rational nodes H ra and the environment nodes H env . 𝜙 includes all the parameters of the Attn (Eq. (6b)) and FFN (Eq. ( 3)) modules. In some contexts where the attention matrix P is not explicitly used as an output, input/output of the intervener 𝜙 can be presented as
The utility loss is computed as
where L task is the task-specific objective; e.g., it is the mean squared error for regression or the cross-entropy for classification. Notice that we select the top-𝐾 rows of the output of 𝜙, i.e., the intervened rationale part and ideally, we expect
As Figure 1 shows, invariant rationale discovery is to find the graph rationale so that the utility loss is minimized given changing environments. Thus, the utility objective is
where Henv is the node embeddings from the changing environments. In practical implementations, Henv is the environment node embeddings from other graphs in the mini-batch. Additionally, to fully utilize the rich interactions from the fine-grained intervention module, we apply the following regularization term, where P ∈ R 𝑛×𝑛 is the self-attention matrix from Eq. (6b),
The binary s vector is used to designate whether a particular row of the matrix H ra ||H env originates from the rationale nodes or the environment nodes. The underlying notion of the regularization term Eq. ( 8) is to impose penalties on interactions between the rationale nodes and the environment nodes. Namely, these two terms s ⊤ P(1-s) and (1-s) ⊤ Ps denote the total weights on the links (i.e., cut) between the rationale and environment subgraphs. To handle the changing environments, we introduce an additional regularization term on the changing environments as L reg (H ra || Henv ) where Henv is the environment node embeddings from another graph within the same mini-batch. Then, the total regularization term is
and the final objective function is L util + 𝛽L reg . To fully harness the capabilities of the fine-grained parametric intervener, it is crucial to note-as highlighted in the introduction-that the behavior of the intervener 𝜙 operates in an adversarial fashion to the other modules. As a result, we formulate a min-max game that involves 𝜃 = {𝜃 enc , 𝜃 aug-N , 𝜃 pred } and 𝜙 as, min
Here, the intervener 𝜙 is trained to decrease the utility of the graph rationale by facilitating interactions between the rationale nodes and the environment nodes. Conversely, the encoder, augmenter, and predictor (i.e., 𝜃 ) are optimized in an opposing manner to the intervener's objectives.
Remark 1. The min-max objective is necessary. If it is replaced with the min objective, trivial solutions exist. E.g., if feature dimension 𝑑 3 concatenate H ra ||H env and H ra || Henv ; 4 compute L util and L reg via Eqs. (7), (9), and (10); 5 update 𝜃 via gradient descent with 𝜕 ( L util +𝛽 L reg ) 𝜕𝜃 ; 6 update 𝜙 via gradient ascent with 𝜕 ( L util +𝛽 L reg ) 𝜕𝜙 ; is large enough, 𝜙 could be an identity function with P = I, so that (1) L reg = 0 and (2) we would always have 𝜙 (H ra ||H env ) [: 𝐾] = H ra . However, in that case, the pipeline degenerates to a regular graph classification/regression model, with no intervener. Complexity of FIG-N. As the encoder 𝜃 enc and the predictor 𝜃 pred are off-the-shelf, the FIG-N introduces two new modules: the nodelevel augmenter, 𝜃 aug-N , and the Transformer-based intervener, 𝜙. Notably, despite these additions, the increase in the number of parameters remains modest. The parameters for 𝜃 aug-N originate from the MLP defined in Eq. (4). In a configuration where the MLP has 3 layers with a feature dimension of 𝑑, the parameter count is 𝑂 (2𝑑 2 ). The intervener 𝜙, driven by the Transformer layer in Eq. (6a), has its parameters confined to 𝑂 (3𝑑 2 + 2𝑑 2 ) = 𝑂 (5𝑑 2 ), owing to its query, key, value projection matrices and the feedforward net (FFN from Eq. (3), typically a 3-layered MLP). A step-by-step algorithm for FIG-N is in Algorithm 1. In test phase, the output of 𝜃 pred • Readout • 𝜙 (H ra ||H env ) is evaluated. Algorithm 2: FIG-VN single training step for graph G Input : a labeled graph (G, 𝑦), a sampled graph G from the same batch as G, 𝜃 = {𝜃 enc , 𝜃 aug-VN , 𝜃 pred }, 𝜙; Output : updated 𝜃 and 𝜙; 1 compute H = 𝜃 enc (𝐺) and H = 𝜃 enc ( G); 2 compute (H ra , H env ) = 𝜃 aug-VN (H), ( Hra , Henv ) = 𝜃 aug-VN ( H) via Eqs. (12), (13), and (14); 3 concatenate H ra ||H env and H ra || Henv ; 4 compute L util and L reg via Eqs. (7), (9), and (10); 5 update 𝜃 via gradient descent with 𝜕 ( L util +𝛽 L reg ) 𝜕𝜃 ; 6 update 𝜙 via gradient ascent with 𝜕 ( L util +𝛽 L reg ) 𝜕𝜙 ; 3.2 Virtual Node-Level Variant: FIG-VN In the previously introduced FIG-N, its augmenter decomposes the nodes into rationale nodes and environment nodes via a trainable node partitioner 𝜃 aug-N so that the interaction is conducted at the node level (Eqs. (6a)) whose dense attention matrix's complexity is quadratic in terms of the node numbers. This section extends this idea to extract the graph rationale at the virtual node level, which has a lower computation complexity (linear in terms of the node numbers) compared to FIG-N and provides an intermediate intervention granularity between the node-level model (FIG -N) and the graph-level model (GREA [37]).
Virtual node-level augmenter. Our idea is partly inspired by the speedup technique from Linformer [54], which reformulates both the attention matrix and node (token) embedding matrix to dimensions of R 𝑛×𝑟 and R 𝑟 ×𝑑 , respectively. This reformulation ensures that their multiplication scales linearly with the number of nodes (tokens) 𝑛 because 𝑟 ≪ 𝑛 and 𝑑 represents the feature dimension.
Building upon this insight, given node embeddings H from the encoder, the virtual node embeddings are:
Here, the row-wise applied Softmax function, along with Softmax(W N-VN ) ∈ R 𝑟 ×𝑛 , yields a trainable matrix assigning 𝑛 nodes to 𝑟 virtual nodes, where 𝑟 acts as a tunable hyper-parameter. More specifically, Softmax(W N-VN ) can be viewed as the adjacency matrix between 𝑟 virtual nodes and 𝑛 original nodes; the virtual node embedding, H VN , is the aggregation of original node embeddings after 1-step message passing. In experiments, we set 𝑟 = 8. As all the virtual node embeddings are learned, a subset of the 𝑟 virtual nodes can be designated as rationale virtual nodes, whose rationality is data-driven by the intervention procedure discussed in subsequent subsections. For brevity, the initial 𝐾 virtual nodes are deemed as rationale virtual nodes, while the last 𝑟 -𝐾 nodes are considered the environment virtual nodes; their embeddings are:
Like the FIG-N, here 𝐾 is a hyperparameter whose impact is studied in Section 4.5. The parameter of 𝜃 aug-VN is only W N-VN .
Virtual node-level intervener. This section discusses the design of a virtual node-level intervener, similar to the intervener introduced in Section 3.1. The main difference is that the intervention here functions on the virtual nodes rather than the given real ones.
We recall the rationale and environment virtual node embeddings
Thanks to the property of the Transformer that it can process sets with variable size, the design of the virtual node-level intervener 𝜙 is similar to the nodelevel intervener as H inter , P = Transformer(H ra ||H env ) or short as H inter = 𝜙 (H ra ||H env ) if the attention matrix P is not used. Notably, for FIG-VN, P ∈ R 𝑟 ×𝑟 describes the interaction among the 𝑟 virtual nodes, whose complexity is only 𝑂 (𝑟 2 ). FIG-VN optimization objective. The output of 𝜃 pred • Readout • 𝜙 (•) is used for minimizing the utility loss L util = L util (H ra ||H env )+ 𝛼 L util (H ra || Henv ), where L util (H ra ||H env ) = L task (𝜃 pred • Readout • 𝜙 (H ra ||H env ) [: 𝐾], y) and L util (H ra || Henv ) is defined similarly. For modeling the changing environment, Henv is the virtual node embeddings from other graphs in the mini-batch. Additionally, the previously proposed regularization term Eq, (8) can be extended to the virtual node-level variant: L reg (H ra ||H env ) = s ⊤ P(1s) + (1s) ⊤ Ps. The total regularization term, considering the changing environment Henv , is L reg = L reg (H ra ||H env ) + L reg (H ra || Henv ). As the P depicts interactions among virtual nodes, we construct the rationale/environment indicator vector s analogously to Eq. (9). Put everything together, and the optimization objective of FIG-VN is min 𝜃 max 𝜙 L util + 𝛽L reg , where 𝜃 = {𝜃 enc , 𝜃 aug-VN , 𝜃 pred }. Complexity of FIG-VN. As we mentioned the encoder 𝜃 enc and the predictor 𝜃 pred are off-the-shelf. Thus, the extra modules introduced by the FIG-VN are the virtual node-level augmenter 𝜃 aug-VN and the Transformer-based intervener 𝜙. The parameters for 𝜃 aug-VN originate from the matrix W N-VN , as defined in Eq. (12), whose complexity is 𝑂 (𝑛𝑟 ), where 𝑛 denotes the number of nodes. For practical implementation purposes, 𝑛 is pre-set; it is set to 10× the average size of graphs from the dataset, and we truncate the input graphs if their sizes are larger than 10× the average size. The intervener 𝜙 parameters originate from the Transformer layer, outlined in Eq. (6a). The number of parameters here is 𝑂 (5𝑑 2 ), owing to its query, key, value projection matrices, and the feed-forward net (Eq. (3), typically a 3-layered MLP). A step-by-step algorithm for FIG-VN is in Algorithm 2. In test phase, the output of 𝜃 pred • Readout • 𝜙 (H ra ||H env ) is evaluated.
In this section, the datasets and baseline methods are detailed first. Subsequent subsections evaluate the effectiveness of FIG, supplemented by efficiency studies, ablation studies, sensitivity studies, training convergence, and a visualization of the attention matrix. The detected rationales are visualized in the Appendix, Section 4.6.
In this paper, 7 publicly-accessible real-world datasets are used: (1) graph classification datasets molhiv [23], moltox21 [23], molbace [23], molbbbp [23] and (2) graph regression datasets ZINC [14], AQSOL [14], and mollipo [23]. We strictly follow the metrics and dataset split recommended by the given benchmarks. To be concrete, the area under the ROC curve (AUC) is the metric for datasets molhiv, moltox21, molbace, molbbbp; root-mean-square deviation (RMSE) is the metric for dataset mollipo; mean absolute error (MAE) is the metric for datasets ZINC and AQSOL. The detailed statistics of the datasets are given in Table 5 (Appendix). We report the average result with the standard deviation in 10 runs 1 .
Our baseline methods include (1) 4 graph neural network models: GIN [59], GAT [52], GATv2 [7], and GatedGCN [6] (2) 7 graph Transformers: GT [13], GraphiT [40], SAN [29], SAT [9], Graphormer [74], GraphTrans [56], GPS [46], and (3) 2 graph rationale discovery methods: DIR [55] and GREA [37]. 1. To ensure a fair comparison, some pre-trained models, such as the pre-trained Graphormer [74], are excluded. As DIR [55] is designed to conduct 1 The code is available at https://github.com/pricexu/FIG interventions on the label prediction vectors, it cannot be directly applied to graph regression tasks.
We have several observations. First
A detailed comparison regarding the number of parameters and the FLOPs (floating point operations) is presented in Table 2, where we list 4 typical 5-layered encoders (GIN, SAT, GraphTrans, and GPS), and our proposed node-/virtual node-level augmenter, intervener modules (i.e., {𝜃 aug-N , 𝜙 } and {𝜃 aug-VN , 𝜙 }). The comparison shows that our proposed interveners are lightweight and only increase very minor computational costs.
We also compare the training efficiency ( iterations/second) of 3 working with different encoders (GIN and GPS). The batch size is set as 32. We note that the inclusion of our proposed augmenter and intervener, represented as {𝜃 aug-N , 𝜙 } or {𝜃 aug-VN , 𝜙 }, introduces a slight reduction in training speed. That is because the proposed parametric augmenter and intervener increase the data pipeline steps, as presented in figure 2b, and enlarge the computational graph for auto-gradient tools, such as PyTorch. Fortunately, the parameter count of the parametric augmenter and intervener is low, ensuring that the overall training speed of the model is not dramatically affected.
We conducted an ablation study on the proposed models, FIG-N and FIG-VN. We designed two ablated variants as baselines: (
𝜃 enc • 𝜃 pred which is a pure composition of the encoder 𝜃 enc and the predictor 𝜃 pred without any rationale discovery module. Many of the existing graph classifiers are in this form, and here we select the GraphGPS [46], which also serves as the backbone of our 10)) from the objective function. Our results in Table 4 highlight that (1) equipped the proposed Transformer-based intervener, the model's performance improves across all the datasets; e.g., the AUC is improved from 79.6% to 84.3% (FIG-N) on the molbace dataset. ( 2) With the proposed regularization term, the model's performance can be improved further; e.g., the AUC of the FIG-N is improved from 72.3% to 73.8% on the molbbbp dataset.
In this section, we carefully study the impact of hyperparameter 𝐾 (from Eq. (5a), (5b), (13), and ( 14)), which determines the ratio of the rationale and environment subgraphs. In our implementation, we set
. We evaluate the model performance across varying K on the molbace and molbbbp datasets in Figure 3. We note that the model performance degrades if most nodes/virtual nodes are marked as the environment component. Similar performance degradation is observed if too many nodes/virtual nodes are marked as the rationale nodes (e.g., K = 0.875). That is because for a small K (e.g., 0.25), too few nodes/virtual nodes are involved for the downstream
(a) molbace (b) molbbbp tasks; for a large K (e.g., 1), the model degenerates to a vanilla graph encoder. The best performance is observed when K is 0.75 or 0.675.
We
In this section, we aim to evaluate the significance of the regularization term by visualizing the attention matrix P of the intervener 𝜙 in Figure 5. For clarity in visualization, FIG-VN is chosen because its number of virtual nodes is predefined. Specifically, we set the number of virtual nodes 𝑟 to 16 with 𝐾 at 10, designating 10 virtual nodes to rationales and the remainder as environments. All the visualization results are obtained from the molbace dataset. It is worth noting that the attention matrix P is normalized row-wise by Softmax. From our observations, we highlight two insights:
• Interestingly, even in the absence of the regularization term, in Figure 5(a), interactions between rationales and environments appear significantly weaker than those within the rationales themselves. One potential explanation is the changing of the environment. In optimizing the utility loss L util = L task (H ra ||H env ) + 𝛼 L task (H ra || Henv ), the everchanging environment ( Henv ) might lead the model to minimize interactions between rationales and environments so that the utility of the rationale can be preserved. • The first observation aligns with the motivation to introduce the regularization term, which aims to penalize rationaleenvironment interactions. When the proposed regularization term (Eq. ( 8)) is implemented, in Figure 5(b), there is a noticeable decrease in rationale-environment interactions (the off-diagonal blocks in Figure 5). As our earlier ablation study demonstrated, this leads to improved model performance.
This section introduces two topics related to this paper: (1) invariant learning on graphs and (2) Transformer on graphs. Invariant Learning on Graphs. Invariant learning, which has extensive applications in many fields such as computer vision [2,8,20,27], is gaining more attention in the community of graph machine learning [4, 16, 34, 36, 39, 43, 50, 53, 58, 62-65, 68-72, 75, 76, 79, 80, 82, 83]. OOD-GNN [31] applies random Fourier features to eliminate the statistical dependence between relevant and irrelevant graph representations. [5] studies the graph invariant representation learning with a particular focus on the size discrepancies between the training/test graphs. DIR [55] decomposes the given graph into a rationale subgraph and an environment subgraph; after that, it uses a graph classifier to conduct prediction on the above two graphs respectively, and its intervention is via the Hadamard product between the prediction vectors. Similarly, GREA [37] conducts the rationale/environment decomposition at the node level, and its intervention operation is to directly add the environment graph embedding into the rationale graph embedding. In a similar vein, GIL [32] decomposes the given graph into the invariant and variant subgraphs; based on that, it infers the environment label, as input of the invariance regularization term [28]. Furthermore, invariant learning can also benefit graph contrastive learning [25,26,48,85,86] to enhance data augmentation. Using invariant learning to address the OOD generalization challenges on graph [12,19,21,30,49,73,77,87] is promising, but our paper does not concentrate on this setting, and we leave it as future work.
Transformer on Graphs. Transformer [1,11,38,44,45,51] has achieved significant success in various domains, including natural language processing [51], computer vision [22], and more [10,15,17,35,42,60,66]. In recent years, there has been a surge of interest in enhancing graph machine learning methods by leveraging the capabilities of Transformers. For example, GraphTrans [56] concatenates the Transformer architecture after various messagepassing neural networks; GPS [46] proposes a powerful layer which operates both a graph convolution layer and a Transformer layer in parallel; GT [13] generalizes the attention module on graphs via concatenating the spectral vectors with the raw node features and computing the attention weights on all the existing edges; Graphormer [74] encodes both the node centrality and node-pair shortest path into the Transformer architecture. A comprehensive survey about the graph Transformer can be found in [41].
This paper studies the invariant rationale discovery problem on graphs. Distinct from existing methods, our solutions (FIG-N and FIG-VN) are designed at more fine-grained levels, specifically at the node and virtual node levels, so that they can better model the interactions between the rationale and environment subgraphs. More importantly, we formulate the intervener and the other model modules in a min-max game, which can significantly improve the quality of the extracted graph rationales. Comprehensive experiments on 7 real-world datasets illustrate the effectiveness of the proposed method compared to 13 baseline methods. Our evaluation of the proposed augmenters' and interveners' efficiency shows that they can largely retain the overall training efficiency.
A Hardware and Implementations.
PyTorch and PyTorch-geometric. All the efficiency study results are from one NVIDIA Tesla V100 SXM2-32GB GPU on a server with 96 Intel(R) Xeon(R) Gold 6240R CPU @ 2.40GHz processors and 1.5T RAM. We directly use those statistics when baseline methods have pre-existing results for specific datasets. In cases where such results are absent, we implement the models based on either the available code or details described in the associated publications. We first introduce some shared implementations among FIG-N/VN and baseline methods. The batch size is set as 32, and the weight decay is set as 0. The hidden dimension is set as 300. The dropout rate for both the encoder and the intervener is searched between {0, 0.5}. The pooling is searched between {mean, sum}. The normalization is searched between the batch normalization [24] and layer normalization [3]. The learning rate is initialized as 0.0001 The statistics of the datasets used are presented in Table 5. Implementation of FIG-N/VN. The encoder is set as GPS [46] on ZINC, AQSOL, mollipo, molhiv, molbace, and set as GraphTrans [56] on moltox21 and molbbbp. We follow the typical design for the predictor module as a 3-layered MLP with ReLU activation in the intermediate layers. In our implementation, we set 𝛽 = Implementation of baseline methods. We search the number of layers of GIN [59], GAT [52], GATv2 [7], GatedGCN [6], DIR [55], and GREA [37] between {2, 3, 5, 10} and report the best performance, considering configurations both with and without a virtual node connecting to all the given nodes.
Regarding the Transformer-based baselines (GT [13], GraphiT [40], SAN [29], SAT [9], Graphormer [74], GraphTrans [56], GPS [46]), for the (absolute or relative) positional encoding, we adhere to the suggestions made in their original papers. We also searched the number of layers between {2, 3, 5, 10}.
Our GIN, GAT, GATv2, and GatedGCN implementations are from the PyTorch-geometric package. Our implementations of GT, GraphiT, SAN, SAT, Graphormer, GraphTrans, GPS, DIR, and GREA are adapted from publicly available code.
The hard argtopK breaks the differentiability of the model so that 𝜃 aug-N has no gradient. Here, we present a soft argtop-𝐾 trick which is inspired by the soft top-K trick 2 . Overall, the key ideas are as follows, (1) The index vector idx ra ∈ N 𝐾 + can be viewed as a list of onehot index encoding idx ∈ {0, 1} 𝐾 ×𝑛 whose every row is a one-hot vector. If the 𝑖-th row's 𝑗-th element is 1, it means the 𝑗-th element in m is the 𝑖-th largest element in m. Then, we can use the matrix multiplication to index the matrix H:
(2) We aim to find a soft and differentiable matrix to approximate idx. The trick is to use the fact that every row of idx is onehot, which can be approximated by the output of the softmax function. Hence, the implementation is to call the softmax 𝐾 times repeatedly. (3) The parameterization trick can be used: y_hard -y_soft.detach() + y_soft, whose main idea is to ensure that (1) the forward process uses the hard indexing and (2) the backpropagation updates the soft indices. We provide the exemplar code as follows.
https://github.com/ZIB-IOL/merlin-arthur-classifiers/blob/main/soft-topk.py