Improving Subgraph Recognition with Variational Graph Information Bottleneck

Given the input data $X$ and its label $Y$, the information bottleneck (IB) [43] principle learns the minimal sufficient representation $Z$ by optimizing the IB objective: $\min_{Z}{-I(Z,Y)+\beta I(Z,X)}.$ Here $\beta$ is the Lagrangian parameter to balance the two terms. Inspired by it, Yu et al. propose Graph Information Bottleneck (GIB) principle to recognize an informative yet compressed subgraph from the original graph [52]. The GIB objective is as follows:

$\displaystyle\min_{G_{sub}}{-I(G_{sub},Y)+\beta I(G_{sub},G)}.$

(1)

The first term in Eq. 1 encourages $G_{sub}$ to be informative to the graph label $Y$. And the second term minimizes the mutual information of $G$ and $G_{sub}$ so that $G_{sub}$ only receives limited information from the input graph $G$. The subgraph found by GIB is denoted as the IB-subgraph: $G_{sub}^{*}=\mathop{\arg\min}_{G_{sub}}{-I(G_{sub},Y)+\beta I(G_{sub},G)}$. The GIB objective cannot be directly optimized since the mutual information is intractable to compute. Hence, it first estimates the mutual information with MINE [4], and use the estimated value as a proxy in the optimization of GIB. This bilevel process is inefficient in optimization since it is time-consuming to estimate the mutual information. Meanwhile, inaccurate estimation also leads to an unstable training process and degenerated results.

Rather than directly evaluate the compression quality of $G_{sub}$ with $I(G,G_{sub})$, we inject noise into the node representations of $G$ for an alternative as shown in Fig. 2.

For an input graph $G\in\mathbb{G}$ with node feature matrix $X$, adjacent matrix $A$ and degree matrix $D$, we first generate the node representations with a $l$-layer GNN:

	$\displaystyle H$	$\displaystyle=\mathrm{GNN}(A,X;W^{1},\cdots,W^{l})$		(2)
		$\displaystyle=\underbrace{\sigma(D^{-\frac{1}{2}}AD^{-\frac{1}{2}}\cdots\sigma(D^{-\frac{1}{2}}AD^{-\frac{1}{2}}}_{l-\mathrm{layer}}XW^{1})\cdots W^{l})$

where $H=[h_{1},h_{2},\cdots,h_{n}]^{T}$ is the node representation matrix. $W^{1},\cdots,W^{l}$ are the parameters at different layers.

Then we damp the information in $G$ by injecting noises into node representations with a learned probability. Let $\epsilon$ be the noise sampled from a parametric noise distribution. We assign each node a probability of being replaced by $\epsilon$. Specifically, for the i-th node, we learn the probability $p_{i}$ with a Multi-layer Perceptron (MLP). Then, we add a Sigmoid function on the output of MLP to ensure $p_{i}\in[0,1]$:

$\displaystyle p_{i}=\mathrm{Sigmoid}(\mathrm{MLP}(h_{i})).$

(3)

We then replace the node representation $h_{i}$ by $\epsilon$ with probability $p_{i}$:

$\displaystyle z_{i}=\lambda_{i}h_{i}+(1-\lambda_{i})\epsilon,$

(4)

where $\lambda_{i}\sim\mathrm{Bernoulli}(p_{i})$. The transmission probability $p_{i}$ controls the information sent from $h_{i}$ to $z_{i}$. If $p_{i}=1$, then all the information in $h_{i}$ are transfered to $z_{i}$ without loss. On the contrary, when $p_{i}=0$, then $z_{i}$ contains no information from $h_{i}$ but only noise. Compared with dropping nodes for compression in GIB, this method allows to flexibly adjust the amount of information from $h_{i}$ to $z_{i}$ by changing $p_{i}$. We denote the $G_{N}=\{A,Z\}$ as the perturbed graph. At graph level, the transmissioin probability of each node determines the information sent from $G$ to $G_{N}$ in a similar way. Therefore, we can compress the information of $G$ into $G_{N}$ with a set of $p_{i}$. We hope $p_{i}$ is learnable so that we can selectively preserve the information in $G_{N}$. However, $\lambda_{i}$ is a discrete random variable and we can not directly calculate the gradient of $p_{i}$. Therefore, we employ the concrete relaxation [18, 10] for $\lambda_{i}$:

$\displaystyle\hat{\lambda}_{i}=\mathrm{Sigmoid}(\frac{1}{t}\log{\frac{p_{i}}{1-p_{i}}}+\log{\frac{u}{1-u}}),$

(5)

where $t$ is the temperature parameter and $u\sim\mathrm{Uniform}(0,1)$. Another key component in the noise injection is the specification of the injected noise. Notice the arbitrary noise is detrimental to the semantic of the input graph. This will lead the prediction of the perturbed graph to be different from the graph property. Moreover, appropriately selected noise will endow the whole objective with a variational upper bound. Please refer to Section 4.3 for detailed discussions.

The perturbed graph $G_{N}$ is employed to distill the actionable information of $G$ for predicting its label $Y$. Specifically, we compress the information of $G$ via noise injection to obtain $G_{N}$. Meanwhile, we hope that $G_{N}$ is maximally informative to $Y$, which leads to a novel Variational Graph Information Bottleneck (VGIB) framework:

$\displaystyle\min\limits_{G_{N}}-I(G_{N},Y)+\beta I(G_{N},G).$

(6)

The first term encourages $G_{N}$ to be sufficient for predicting the graph label $Y$, and the second term constrains the information that $G_{N}$ receives from $G$. These two terms require us to inject noise into $G$ selectively so that $G_{N}$ receives actionable information as much as possible. The intuition is that injecting noises into the IB-subgraph of $G$ is more harmful to the functionality of $G$ than that into the label-irrelevant substructures. In that sense, the nodes in the IB-subgraph are less likely to be injected with noise. Therefore, we can select the IB-subgraph from $G_{N}$ by this criterion after training VGIB. We introduce the following lemma before justifying the above formulation.

Lemma 7 indicates when setting $\beta=1$ in Eq. 1, the GIB objective upper bounds the mutual information of $G_{sub}$ and $G_{n}$. That is to say, optimizing the GIB objective encourages $G_{sub}$ to be less related to the label-irrelevant substructure $G_{n}$. $I(G_{n},G_{sub})$ is minimal when $G_{sub}$ is the IB-subgraph. The proof of Lemma 7 is in the Supplementary Materials. We next give the theorem that minimization of VGIB objective in Eq. 6 also leads to the irrelevance of $G_{sub}$ and $G_{n}$.

Please refer to the Supplementary Materials for the proof of Theorem 8. VGIB differs from GIB mainly in noise injection. In the next section, we will show that this process leads to convenience in optimization.

We first examine the first term $I(G_{N},Y)$ in Eq. 6, which encourages $G_{N}$ is informative of graph label $Y$.

	$\displaystyle-I(G_{N},Y)$	$\displaystyle\leq\mathrm{E}_{Y,G_{N}}-\log{q_{\theta}(Y|G_{N})}$		(9)
		$\displaystyle:=\mathcal{L}_{\mathrm{cls}}(G_{N},Y),$

where $q_{\theta}(Y|G_{N})$ is the variational approximation to $p(Y|G_{N})$. $q_{\theta}(Y|G_{N})$ outputs the label distribution of $G_{N}$ and can be modeled as a classifier. $\mathcal{L}_{\mathrm{cls}}$ is the classification loss. We choose the cross-entropy loss and mean square loss for categorical $Y$ and continuous $Y$, respectively.

For the second term $I(G_{N},G)$ in Eq. 6, we first obtain the graph representation $z_{N}$ of $G_{N}$ via a readout function. We employ the sufficient encoder assumption [42] that the information of $z_{N}$ is lossless in the encoding process, leading to $I(z_{N},G)\approx I(G_{N},G)$. By choosing the distribution of noise and the readout function as the Gaussian distribution, $I(G_{N},G)$ has a tractable variational upper bound.

where $A_{G}=\sum_{j=1}^{m_{G}}(1-\lambda_{j})^{2}$ and $B_{G}=\frac{\sum_{j=1}^{m_{G}}\lambda_{j}(h_{j}-\mu_{h})}{\sigma_{h}}$. Please refer to the Supplementary Materials for the proof of Proposition 10. One can efficiently estimate Eq. 9 and Eq. 10 with the batched data in the training set. The overall loss is:

$\displaystyle\mathcal{L}=\mathcal{L}_{\mathrm{cls}}(G_{N},Y)+\beta\mathcal{L}_{\mathrm{MI}}(Z_{N},G)$

(11)

In this experiment, we extract the substructures which have the most similar properties to the original molecules. We consider four properties: QED, HLM-CLint, MLM-CLint, and RLM-CLint. QED measures the probability of a molecule being a drug within the range of $[0,1.0]$. HLM-CLint, MLM-CLint, and RLM-CLint are estimated values of in vitro human, mouse and rat liver microsome metabolic stability, respectively (base 10 logarithm of mL/min/g)²²2We evaluate QED values of molecules with the toolkit on https://www.rdkit.org/. Moreover, we obtain HLM-CLint, MLM-CLint, and RLM-CLint value of molecules on https://drug.ai.tencent.com/.. We collect molecules with QED$\geq 0.85$, HLM-CLint$\geq 2$, MLM-CLint$\geq 2$, and RLM-CLint$\geq 2$ from ZINC250K [17], and we individually build datasets with the selected molecules for the four properties. For each property, we use $85\%$, $5\%$, and $10\%$ of the molecules for training, validating, and testing, respectively. Please refer to the Supplementary Materials for the statistics of datasets.

We compare the proposed VGIB with GIB [52] and the attention-based method [25]. For VGIB, we first learn the node representation with a GCN and inject noise into each node as shown in Eq. 4 and Eq. 3. Then we obtain the perturbed graph representation by pooling the noisy node representations with a readout function. We thereafter optimize the loss function in Eq. 11 and collect the nodes with $p_{i}\geq 0.5$ in Eq. 3 to obtain the IB-subgraph. We simultaneously supervise the classifier in Eq. 9 with the representation of the input graph and its label to enhance the informativeness of the subgraph. For GIB, we follow the existing method [52] and optimize the model in a bilevel optimization scheme. As for the attention-based method, we attentively pool the node representations with the attention scores for label prediction. Furthermore, we select the nodes with top $50\%$ and $70\%$ attention scores to form the predictive subgraphs. For each method, we adopt a 2-layer GCN with 16 hidden dimensions for a fair comparison. We run experiments on one TITAN RTX GPU. If the found subgraph is disconnected, we choose its largest connected part to ensure chemical validity.

We report the mean and standard deviation of absolute property divergence between the input molecules and the found subgraphs in Table 1. We quantitatively compare the performance of different methods on the interpretation of graph properties. The proposed VGIB method shows favorable performance against GCN+GIB. For the attention-based method, the performance is sensitive to the selected value of the threshold since the results of GCN+Att05 vary from those of GCN+Att07. Therefore, one needs to finetune the threshold for different tasks. In contrast, our VGIB is free from a manually selected threshold thanks to the information-theoretic objective.

We then compare the training time of different methods on the QED dataset. We train different methods 3 times and report the average training time in Fig. 2. It is shown that the attention-based methods achieve the fastest training due to their simple architectures. Although VGIB is slower than the attention-based methods, it achieves significant performance gain in finding subgraphs with more similar properties to the input molecules. Compared with GIB, VGIB trains over ten times faster than GIB with better performance. The reason is that estimating the mutual information in GIB is time-consuming. However, VGIB a enjoys tractable upper bound of its objective, which is easy to optimize.

As GCN model becomes prevalent in various tasks [14, 54], there is increasing concerns on its explainability since GCN are treated as black-box [49]. In this section, we employ the proposed VGIB to explain the prediction of GCN for molecule classification on ZINC250K dataset. We consider four properties: QED, HLM-CLint, RLM-CLint and DRD2 construct the dataset for each property. For QED, we label the molecules with QED$\geq$0.85 /$<$0.85 as 1/0. For DRD2, we label the molecules with DRD2$\geq$0.50 /$<$0.50 as 1/0. For HLM-CLint and RLM-CLint, we label the molecules with the property value greater than 2.0 as 1 or 0 otherwise. For each property, we train the GCN on the training set and employ VGIB to generate post-hot explanation on the test set. Please refer to Supplementary Materials for more details on the dataset splits and the training process.

We compare VGIB with various explanation model including GNNExplainer [49], PGExplainer [28], GraphMask [37], IGExplainer [41], GraphGrad-CAM [31] and GIB [52]. These methods interpret the prediction of GCN with the node importance score ranging in $[0,1]$, where 1 indicates the node is the most important to the molecule classification and 0 otherwise. Please refer to Supplementary Materials for more details on the baseline methods.

We employ the fidelity score to evaluate how the explanation is faithful to the GCN model [55]. Specifically, let $y_{i}$ and $\hat{y}_{i}$ be the ground-truth and the prediction of the i-th input molecule. Define $k$ as the sparsity score of the explanatory subgraph. The explanatory subgraph is obtained by choosing the nodes with top $k$% scores from the input molecule, and we denote its prediction as $\hat{y}_{i}^{k}$. The Fidelity- score is computed as follows:

$\displaystyle Fidelity-=\frac{1}{N}\sum_{i=1}^{N}\mathds{1}(y_{i}=\hat{y}_{i})-\mathds{1}(y_{i}-\hat{y}_{i}^{k})$

(12)

where $\mathds{1}(y_{i}=\hat{y}_{i})$ is the indicator function which outputs 1 if $y_{i}=\hat{y}_{i}$ and 0 otherwise. Fidelity- score measures how the prediction of the explanatory subgraph is close to the input molecule. The lower value of Fidelity- score indicates more faithful explanation. Similarly, we define $\hat{y}_{i}^{1-k}$ as the prediction of the complementary subgraph, which is obtained by removing the explanatory subgraph from the input. The Fidelity+ score is defined as follows:

$\displaystyle Fidelity+=\frac{1}{N}\sum_{i=1}^{N}\mathds{1}(y_{i}=\hat{y}_{i})-\mathds{1}(y_{i}-\hat{y}_{i}^{1-k})$

(13)

The Fidelity+ score indicates more important nodes are identified in the explanation.

Table 3 shows the fidelity scores of the explanations produced by different methods. The sparsity score is set to be $k=0.5$ for all the explanatory subgraphs. As shown in table 3, VGIB achieves best fidelity scores on all the properties. This shows that VGIB generate faithful explanation to the predictions of the GCN. Moreover, to comprehensively evaluate different methods, we set $k\in\{0.30,0.35,0.40,0.45,0.50,0.55,0.60\}$ and compare the performance under different sparsity scores. As shown in Figure 4, VGIB generate most faithful explanations to the predictions of GCN in terms of fidelity scores by under different sparsity scores. This shows that VGIB can identify most important nodes to the predictions of GCN.

In this subsection, we aim to find out whether the found subgraph can improve the performance of baselines on graph classification or not. We evaluate different methods on MUTAG [36], PROTEINS [6], DD, IMDB-BINARY, REDDIT-BINARY, and COLLAB [35] datasets, which are wildly used for graph classification. Please refer to Supplementary Materials for the statistics of datasets.

We consider four GNN baselines including GCN [22], GAT  [44], GraphSAGE  [14] and GIN  [47]. We use the mean and sum pooling as a readout function for the baseline methods to obtain the graph representation for prediction. Then, similar to GIB [52], we plug VGIB into these baselines. Specifically, we adopt the baseline models to extract node representation. Then, we recognize the IB-subgraph by optimizing the VGIB objective. We pool the node representations in the IB-subgraph for classification with the same readout function as the baselines. For a fair comparison, we adopt a 2-layer network architecture and 16 hidden dimensions for different methods. We train these methods for 100 epochs and test the models with the smallest validation loss. We report the mean and standard deviation of accuracy across 10 folds ³³3We follow the protocol in https://github.com/rusty1s/pytorch_geometric/tree/master/benchmark/kernel.

The experimental results are summarized in Table 4. Compared with the baselines, VGIB can discover an informative yet compressed subgraph of the input. Therefore, it relieves the perturbation of noise structures and redundant information and boosts the performance of the baselines. VGIB also outperforms the GIB-based methods on most of the datasets.