Globally Interpretable Graph Learning via Distribution Matching

Instance-level methods provide input-dependent explanations for each individual graph (Agarwal et al., 2023; Yuan et al., 2022). Given an input graph, these methods explain GNNs by extracting a small interpretive subgraph. Existing solutions can be categorized as gradient-based (Baldassarre and Azizpour, 2019; Pope et al., 2019), attention-based (Miao et al., 2022), perturbation-based (Ying et al., 2019; Luo et al., 2020b), decomposition-based (Huang et al., 2022), and surrogate-based methods (Vu and Thai, 2020). Gradient-based method directly uses the gradients as the approximations of feature importance. Attention-based methods use the attention mechanism to identify important subgraph as interpretation. Perturbation-based methods optimize a subgraph mask to captures the important nodes and edges. Surrogate-based explanation methods use data sampling to filter out unimportant features and an explainable small model — such as a probabilistic graphical model — is fitted on the filtered data as a topological explanation. Decomposition-based methods decompose predictive scores to represent how importance the input contributes to the predicted results. Again, instance-level methods are based on each input instance. Although they are helpful for getting an explanation for every single graph, they can hardly capture the commonly important features that are shared by graph instances for each class. Therefore, it is necessary to have both instance-level and model-level interpretations for GNNs.

Model-level interpretation aims at capturing the global behaviour of the model as a whole, such that a robust overview of the model can be summarized from individual noisy local explanations. This type of interpretation on graph learning is less studied. XGNN (Yuan et al., 2020) leverages a reinforcement learning technique to sequentially generate edges based on the prediction reward. However, this approach requires domain expert knowledge to design valid reward function for different inputs, which is not always available. GNNInterpreter (Wang and Shen, 2022) learns a probabilistic generative graph distribution and identifies the key graph pattern when GNN tries to make a certain prediction. GLGExplainer (Azzolin et al., 2022) generates explanations as Boolean combinations of learned graphical concepts, represented as clusters of local explanations. While these methods identify intuitive class-related patterns that can be recognized by the model (with high predictive accuracy), they usually ignores the training utility of these explanations. Ideally, high-quality interpretations capturing class discriminative patterns from the training data should be able to train a similar model. From this perspective, in this work, we define model fidelity as a new metric, and propose a novel globally interpretable graph learning framework that explains by matching the distribution along the model training trajectory.

Dataset condensation aims to synthesizing a compact training dataset to distill massive training data. Multiple techniques are proposed to maintain the utility of dataset for model training, including gradient matching (Zhao et al., 2020) with data regularity (Kim et al., 2022), trajectory matching (Cazenavette et al., 2022), and distribution matching (Zhao and Bilen, 2023). While majority of study focuses on i.i.d. data (image), recently these techniques are adapted to condense large graphs into small and highly-informative graphs (Jin et al., 2022b, a; Hashemi et al., 2024; Zheng et al., 2023; Xu et al., 2023). While data condensation and global interpretation could share similar techniques, they have very different goals: data condensation aims to boost model performance using as small data as possible, while global interpretation is to faithfully identify the key patterns dominating the target model’s behavior.

We focus on explaining GNNs’ global behavior for the graph classification task. A graph classification dataset with $N$ graphs can be denoted as $\mathcal{G}=\{G^{(1)},G^{(2)},\dots,G^{(N)}\}$ with a corresponding ground-truth label set $\mathcal{Y}=\{y^{(1)},y^{(2)},\dots,y^{(N)}\}$. Each graph consists of two components, $G^{(i)}=(\textbf{A}^{(i)},\textbf{X}^{(i)})$, where $\textbf{A}^{(i)}\in\{0,1\}^{n\times n}$ denotes the adjacency matrix and $\textbf{X}^{(i)}\in\mathbb{R}^{n\times d}$ is the node feature matrix. The label for each graph is chosen from a set of $C$ classes $y^{(i)}\in\{1,\dots,C\}$, and $y^{(i)}_{c}$ denotes that the label of graph $G_{i}$ is $c$, that is $y^{(i)}=c$. A set of graphs that belong to class $c$ could be further represented as $\mathcal{G}_{c}=\{G^{(i)}|y^{(i)}=c\}$.

A GNN model $\Phi(\cdot)$ is a concatenation of a feature extractor $f_{\bm{\theta}}(\cdot)$ parameterized by $\bm{\theta}$ and a classifier $h_{\bm{\psi}}(\cdot)$ parameterized by $\bm{\psi}$, where $\Phi(\cdot)=h_{\bm{\psi}}(f_{\bm{\theta}}(\cdot))$. The feature extractor $f_{\bm{\theta}}:\mathcal{G}\rightarrow\mathbb{R}^{d^{\prime}}$ takes in a graph and embeds it to a low-dimentional space with $d^{\prime}\ll d$. The classifier $h_{\bm{\psi}}:\mathbb{R}^{d^{\prime}}\rightarrow\{1,\dots,C\}$ further outputs the predicted class given the graph embedding.

We now provide a general form for the global interpretation problem. The idea is to generate a small set of compact graphs that can explain the high-level behavior of the GNN model, e.g., what patterns lead the model to discriminate different classes. Specifically, given a GNN model $\Phi^{*}$, exsiting global interpretation method aims to generate interpretive graphs that have the maximal predicted probability for a certain class $y_{c}$. Formally, this problem can be defined as:

where $\mathcal{S}$ is one or multiple compact interpretive graphs capturing key graph structures and node characteristics for interpretation, and $\mathcal{L}(\cdot,\cdot)$ is the loss (e.g., cross-entropy loss) of predicting $\mathcal{S}$ as label $y_{c}$. Existing global interpretation techniques can fit in this form but differ in the generation procedure of $\mathcal{S}$. For instance, in XGNN (Yuan et al., 2020), $\mathcal{S}$ is defined as a set of completely synthesized graphs with each edge generated by a reinforcement learning strategy. The goal of the reward function is to maximize the probability of predicting $\mathcal{S}$ as a certain class $y_{c}$. In GNNInterpreter (Wang and Shen, 2022), $\mathcal{S}$ is generated by sampling from an estimated graph distribution. In GLGExplainer (Azzolin et al., 2022), $\mathcal{S}$ is generated by a Boolean logic function. Despite their difference in generation techniques, they stand on a common ground as a model consumer: they can only access and inspect the final pre-trained model $\Phi^{*}$ to explain its behavior.

If standing from the perspective of model provider, such a problem formulation may not fully leverage all accessible information, such as the the whole training trajectory, leading to limited interpretation capability. Specifically, we consider interpretation quality from the following two aspects:

• •

  Predictive Accuracy reflects whether the extracted interpretative patterns are really class-relevant. It is calculated as the model accuracy on generated interpretive graphs. Existing works mainly focus on this aspect (Yuan et al., 2020; Wang and Shen, 2022; Azzolin et al., 2022).

• •

  Model Fidelity measures whether the interpretive graphs are class discriminative enough to train a similar model. It is calculated as the cosine similarity between the predictive probabilities of the target model and that of the surrogate model (trained by interpretative graphs) on a same set of instances. This aspect however has never been inspected in prior studies.

As shown in Figure 1, existing works following this formulation provide a limited model fidelity. This observation motivates us to rethink the global interpretation problem from the model provider’s perspective and design a globally interpretable learning framework.

Our goal is to generate global explanations that can not only be accurately predicted as the corresponding class, but also lead to a high-fidelity model. In order to achieve this goal, we propose to optimize the explanations in the model developing stage, such that the training trajectory information can be leveraged. We thus propose a novel research problem: how to provide global interpretation for a model training procedure, such that training on such interpretation can recover a similar model? We frame this problem as globally interpretable graph learning, which can be defined as follows:

where $\mathcal{T}=[0,\dots,T-1]$ is the normal training iterations for the target GNN model, and $\mathtt{opt-alg}$ is a specific model update algorithm (e.g., gradient descent) with a fixed number of steps $\varsigma$. $\mathcal{L}_{\text{CE}}(\Phi)=\mathbb{E}_{G,y\sim\mathcal{G,Y}}[\mathscr{l}(\Phi(G),y)]$ is the cross-entropy loss used for normal GNN model training.

This formulation of globally interpretable graph learning states that the interpretable patterns $\mathcal{S}$ should be optimized based on the whole training trajectory $\Phi_{0}\rightarrow\Phi_{1}\rightarrow\cdots\rightarrow\Phi_{T-1}$ of the model. This stands in sharp contrast to other global interpretation where only the final model $\Phi^{*}=\Phi_{T-1}$ is considered. The training trajectory reveals more information about model’s training behavior to form a constrained model space, such as essential graph patterns that dominate the training of this model.

To realize globalinterpretation as demonstrated in Eq. (2), we now introduce the exact form of the objective function for optimizing interpretive graphs that encapsulate the model’s learning behavior from the data. Recall that a GNN model is a combination of feature extractor and a classifier. The feature extractor $f_{\bm{\theta}}$ usually carries the most essential information about the model, while the classifier is a rather simple multi-perceptron layer. Since the feature extractor plays the majority role, a natural idea for generating interpretation is to match its distribution with training graphs in the model’s feature space. We name this interpertation principle as Graph Distribution Matching (GDM).

The proposed framework is illustrated in Figure 2, for each training step $t$, we update interpretive graphs by aligning with the training graphs in the GNN model’s feature space via distribution matching. Along the whole training trajectory, we keep updating interpretive graphs in a curriculum learning manner to capture the model’s training behavior. It is worth noting that such a distribution matching scheme has shown success in distilling rich knowledge from training data to synthetic data (Zhao and Bilen, 2023), which preserve sufficient discriminative information for training the underlying model. This justifies our design of distribution matching for interpretation.

Optimizing each interpretive graph is essentially optimizing its adjacency matrix and node feature. Denote a interpretive graph as $S=(\textbf{A}_{s},\textbf{X}_{s})$, with $\textbf{A}_{s}\in\{0,1\}^{m\times m}$ and $\textbf{X}_{s}\in\mathbb{R}^{m\times d}$. To generate solid graph explanations using Eq. (5), we introduce several practical constraints on $\textbf{A}_{s}$ and $\textbf{X}_{s}$. The constraints are applied on each interpretive graph, concerning discrete graph structure, matching edge sparsity, and feature distribution with the training data.

Specifically, for each graph $G=(\textbf{A, X})$ with $n$ nodes, we estimate the graph-level feature distribution as $\bar{\textbf{x}}=\sum^{n}_{i=1}\textbf{X}_{i}/n\in{\mathbb{R}^{d}}$, which is essentially a mean pool of the node features. For each class $c$, we then define the following feature matching loss:

where we empirically measure the class-level feature distribution by calculating the average of graph-level features. By minimizing the feature distribution distance in Eq. (8), even rare features can have a chance to be distilled in the interpretive graphs.

Integrating the practical constraints discussed in Section 3.5 with the distribution matching based interpretation framework in Eq. (5), we now obtain the final objective for interpretation optimization, which essentially is determined by the Bernoulli parameters for sampling discrete adjacency matrices and the node feature matrices. Formally, we aims to solve the following optimization problem:

where we use $\alpha$ and $\beta$ to control the strength of regularizations on feature distribution matching and edge sparsity respectively. Algorithm 1 details the steps for solving this optimization problem.

Dataset The interpretation performance is evaluated on the following synthetic and real-world datasets for graph classification, whose statistics can be found in Appendix A.2.

• •

  Real-world data includes: MUTAG (Debnath et al., 1991) consists of chemical compounds with atoms as nodes and chemical bonds as edges, labeled by whether it has a mutagenic effect on a bacterium. Graph-Twitter (Socher et al., 2013) includes Twitter comments for sentiment classification with three classes. Each comment sequence is presented as a graph, with word embedding as node feature. Graph-SST5 (Yuan et al., 2021) is a similar dataset with reviews, where each review is converted to a graph labeled by one of five rating classes.

• •

  Synthetic data includes: Shape contains four classes, i.e., Lollipop, Wheel, Grid, and Star. Each class has the same number of synthesized graphs with a random number of nodes. BA-Motif (Luo et al., 2020b) uses Barabasi-Albert (BA) graph as base graphs, among which half graphs are attached with a “house” motif and the rest with “non-house” motifs. BA-LRP (Schnake et al., 2021) based on Barabasi-Albert (BA) graph includes one class being node-degree concentrated graphs, and the other degree-evenly graphs. These datasets do not have node features, thus we use node index as the surrogate feature.

This evaluation aims to answer the first question Q1. Meanwhile, we also report the commonly adopted predictive accuracy.

We qualitatively visualize the global interpretations provided by GDM to verify that GDM can capture human-intelligible patterns, which indeed correspond to the ground-truth rules for discriminating classes. Table 5 shows examples in BA-Motif, MUTAG, BA-LRP and Shape datasets, and more results and analyses on other datasets can be found in Appendix A.6. The qualitative results show that the global explanations successfully identify the discriminative patterns for each class. If we look at BA-Motif dataset, for the house-shape class, the interpretation has captured such a pattern of house structure, regardless of the complicated base BA graph in the other part of graphs; while in the non-house class with five-node cycle, the interpretation also successfully grasped it from the whole BA-Motif graph. Regarding the Shape dataset, the global interpretations for all the classes are almost consistent with the ground-truth graph patterns, i.e., Wheel, Lollipop, Grid and Star shapes are also reflected in the interpretation graphs. Note that the difference for interpretative graphs of Star and Wheel are small, which provides a potential explanation for our fidelity results in Table 2, where pre-trained GNN models cannot always distinguish interpretative graphs of Wheel shape with interpretative graphs of Star shape.