SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sangwoo Seo $^{1}$  Sungwon Kim $^{1}$  Chanyoung Park $^{1*}$
1KAIST
{sangwooseo@kaist.ac.kr, swkim@kaist.ac.kr, cy.park@kaist.ac.kr}
The success of Graph Neural Networks (GNNs) has led to a need for understanding their decision-making process and providing explanations for their predictions, which has given rise to explainable AI (XAI) that offers transparent explanations for black-box models. Recently, the use of prototypes has successfully improved the explainability of models by learning prototypes to imply training graphs that affect the prediction. However, these approaches tend to provide prototypes with excessive information from the entire graph, leading to the exclusion of key substructures or the inclusion of irrelevant substructures, which can limit both the interpretability and the performance of the model in downstream tasks. In this work, we propose a novel framework of explainable GNNs, called interpretable Prototype-based Graph Information Bottleneck (PGIB), that incorporates prototype learning within the information bottleneck framework to provide prototypes with the key subgraph from the input graph that is important for the model prediction. This is the first work that incorporates prototype learning into the process of identifying the key subgraphs that have a critical impact on the prediction performance. Extensive experiments, including qualitative analysis, demonstrate that PGIB outperforms state-of-the-art methods in terms of both prediction performance and explainability. The source code of PGIB is available at https://github.com/sang-woo-seo/PGIB.
With the success of Graph Neural Networks (GNNs) in a wide range of deep learning tasks, there has been an increasing demand for exploring the decision-making process of these models and providing explanations for their predictions. To address this demand, explainable AI (XAI) has emerged as a way to understand black-box models by providing transparent explanations for their predictions. This approach can improve the credibility of models and ensure transparency in their decision-making processes. As a result, XAI is actively being used in various applications, such as medical, finance, security, and chemistry [3, 26].
In general, explainability can be viewed from two perspectives: 1) improving the interpretability by providing explanations for the model's predictions, and 2) providing the reasoning process behind the model prediction by giving explanations for the model's training process. Improving the interpretability in GNNs involves detecting important substructures during the inference phase, which is useful for tasks such as identifying functional groups (i.e., important substructures) in molecular chemistry [29, 33, 23, 12]. On the other hand, it is also important to provide the reasoning process for why the model predicts in a certain way, which requires an in-depth analysis of the training phase, so as to understand the model in a more fundamental level. Through this reasoning process, we can visualize and analyze how the model makes correct or incorrect decisions, thus obtaining crucial information for improving its performance.
Figure 1: Comparison of the learned prototypes between ProtGNN and PGIB.
In recent years, there has been a growing interest in exploring the reasoning process to provide greater transparency and explainability in deep learning models. These approaches can be generally classified into two main categories: 1) post-hoc approaches, and 2) built-in approaches. Post-hoc approaches focus on exploring the model outputs by visualizing the degree of activation of neurons through measuring their contribution based on gradients of the model predictions. For instance, recent works [14, 9, 34] use techniques such as saliency maps and class activation maps to visualize the activated areas during the model's prediction process. However, these approaches require a separate explanatory model for each trained model, resulting in the need for a new explanatory model for additional training data and different models [2, 15]. In order to address the aforementioned challenges, built-in approaches aim to integrate the generation of explanations into the model training process. One such approach is prototype learning, which involves learning prototypes that represent each class of the input data, which are then compared with new instances to make predictions. ProtGNN [35], for instance, measures the similarity between the embedding of an input graph and each prototype, providing explanations through the similarity calculation and making predictions for the input graph based on its similarity with the learned prototypes. More precisely, ProtGNN projects each learned prototype onto the closest training graph, enabling it to provide explanations of primary structures for its prediction.
However, since ProtGNN compares the graph-level embedding of an input graph with the learned prototypes, the model overlooks the key substructures in the input graph while also potentially including uninformative substructures. This not only results in a degradation of the interpretability of the reasoning process, but also limits the performance on the downstream tasks. Figure 1(a) shows the prototype graphs in the training set (i.e.,  $\mathcal{G}_p$ , denoted as bold edges) detected by ProtGNN for an input molecule (i.e.,  $\mathcal{G}$ ) that belongs to the "mutagenic" class. Despite the  $\mathrm{NO}_2$  structure being the key functional group for classifying a given molecule as "mutagenic,"  $\mathcal{G}_p$  detected by ProtGNN tends to include numerous ring structures (i.e., uninformative substructures) that are commonly found throughout the input graph, and exclude  $\mathrm{NO}_2$  structures (i.e., key substructures) in learned prototype graphs, which is mainly due to the fact that the input graph  $\mathcal{G}$  is considered in the whole graph-level. As a result, it is crucial to identify a key subgraph within the input graph that holds essential information for the learning of prototypes, which in turn enhances both the explanation of the reasoning process and the performance on the downstream tasks. Among the various solutions for detecting important subgraphs, the Information Bottleneck (IB) has emerged as one of the most effective methods [27, 31], and it has been demonstrated that key subgraphs detected based on IB can contribute to performance improvement in various tasks such as relational learning [11] and structure learning [19]. We aim to approach the IB principle from the perspective of prototypes to convey important substructure information to the prototypes.
To this end, we propose a novel framework of explainable GNNs, called Interpretable Prototype-based Graph Information Bottleneck (PGIB). The main idea is to incorporate prototype learning within the Information Bottleneck (IB) framework, which enables the prototypes to capture the essential key subgraph of the input graph detected by the IB framework. Specifically, PGIB involves prototypes (i.e.,  $\mathcal{G}_p$ ) in a process that maximizes the mutual information between the learnable key subgraph (i.e.,  $\mathcal{G}_{sub}$ ) of the input graph (i.e.,  $\mathcal{G}$ ) and target information (i.e.,  $Y$ ), which allows the prototypes to interact with the subgraph. This enables the learning of prototypes  $\mathcal{G}_p$  based on the key subgraph  $\mathcal{G}_{sub}$  within the input graph  $\mathcal{G}$ , leading to a more precise explanation of the reasoning process and improvement in the performance on the downstream tasks. To the best of our knowledge, this is the first work that combines the process of optimizing the reasoning process and interpretability by identifying the key subgraphs that have a critical impact on the prediction performance. In Figure 1(b), PGIB is shown to successfully detect the key subgraph  $\mathcal{G}_{sub}$  that includes  $\mathrm{NO}_2$  from the input graph  $\mathcal{G}$ , even when the ring structures are dominant in  $\mathcal{G}$ . It is important to note that PGIB is highly efficient in detecting  $\mathcal{G}_{sub}$  from  $\mathcal{G}$  since PGIB adopts a learnable masking technique, effectively resolving the time complexity issue. Last but not least, since the number of prototypes for each
class is determined before training, some of the learned prototypes may share similar semantics, which negatively affects the model interpretability for which the small size and low complexity are desirable [6, 25, 17]. Hence, we propose a method for effectively merging the prototypes, which in turn contributes to enhancing both the explanation of the reasoning process and the performance on the downstream tasks.
We conducted extensive experiments to evaluate the effectiveness and interpretability of the reasoning process of PGIB in graph classification tasks. Our results show that PGIB outperforms recent state-of-the-art methods, including existing prototype learning-based and IB-based methods. Moreover, we evaluated the ability of PGIB in capturing the label information by evaluating the classification performance using only the detected subgraph  $\mathcal{G}_{sub}$ . We also conducted a qualitative analysis that visualizes the subgraph  $\mathcal{G}_{sub}$  and prototype graph  $\mathcal{G}_p$ , suggesting the ability of PGIB in detecting the key subgraph. Overall, our results show that PGIB significantly improves the interpretability of both  $\mathcal{G}_{sub}$  and  $\mathcal{G}_p$  in the reasoning process, while simultaneously improving the performance in downstream tasks.
In summary, our main contributions can be summarized as follows: 1) We propose an effective approach, PGIB, that not only improves the interpretability of the reasoning process, but also the overall performance in downstream tasks by incorporating the prototype learning in a process of detecting key subgraphs based on the IB framework. 2) We provide theoretical background with our method that utilizes interpretable prototypes in the process of optimizing  $\mathcal{G}_{sub}$ . 3) Extensive experiments, including qualitative analysis, demonstrate that PGIB outperforms state-of-the-art methods in terms of both prediction performance and explainability.
In this section, we introduce notations used throughout the paper followed by the definitions of IB and IB-Graph.
Notations. We use  $\mathcal{G} = (\mathcal{V},\mathcal{E},\mathbf{A},\mathbf{X})$  to denote a graph, where  $\mathcal{V}$ ,  $\mathcal{E}$ ,  $\mathbf{A}$  and  $\mathbf{X}$  denote the set of nodes and edges, the adjacency matrix and node features, respectively. We assume that each node  $v_{i}\in \mathcal{V}$  is associated with a feature vector  $\mathbf{x}_i$ , which is the  $i$ -th row of  $\mathbf{X}$ . We use  $\{(\mathcal{G}_1,y_1),(\mathcal{G}_2,y_2),\dots ,(\mathcal{G}_N,y_N)\}$  to denote the set of  $N$  graphs with its corresponding labels. The graph labels are given by a set of  $K$  classes  $\mathcal{C} = \{1,2,\ldots ,K\}$ , and the ground truth label of a graph  $\mathcal{G}_i$  is denoted by  $y_{i}\in \mathcal{C}$ . We use  $\mathcal{G}_{sub}$  to denote a subgraph of  $\mathcal{G}$ , and use  $\bar{\mathcal{G}}_{sub}$  to denote the complementary structure of  $\mathcal{G}_{sub}$  in  $\mathcal{G}$ . We also introduce the prototype layer, which consists of a set of prototypes  $\mathcal{Z}_p = \left\{\mathbf{z}_{\mathcal{G}_p}^1,\mathbf{z}_{\mathcal{G}_p}^2,\dots ,\mathbf{z}_{\mathcal{G}_p}^M\right\}$ , where  $M$  is the total number of prototypes, and each prototype  $\mathbf{z}_{\mathcal{G}_p}^m$  is a learnable parameter vector that serves as the latent representation of the prototypical part (i.e.,  $\mathcal{G}_p$ ) of graph  $\mathcal{G}$ . We allocate  $J$  prototypes for each class, i.e.,  $M = K\times J$ .
Graph Information Bottleneck. The mutual information between two random variables  $X$  and  $Y$ , i.e.,  $I(X;Y)$ , is defined as follows:
$$
I (X; Y) = \int_ {X} \int_ {Y} p (x, y) \log \frac {p (x , y)}{p (x) p (y)} d x d y \tag {1}
$$
Given the input  $X$  and its associated label  $Y$ , the Information Bottleneck (IB) [21] aims to obtain a bottleneck variable  $Z$  by optimizing the following objective:
$$
\min  _ {Z} - I (Y; Z) + \beta I (X; Z), \tag {2}
$$
where  $\beta$  is the Lagrange multiplier used to control the trade-off between the two terms. IB principle has recently been applied to learning a bottleneck graph, named IB-Graph, for a given graph  $\mathcal{G}$ , which retains the minimal sufficient information in terms of  $\mathcal{G}$ 's properties [31]. This approach is motivated by the Graph Information Bottleneck (GIB) principle, which seeks to identify an informative yet compressed subgraph  $\mathcal{G}_{sub}$  from the original graph  $\mathcal{G}$  by optimizing the following objective:
$$
\min  _ {\mathcal {G} _ {s u b}} - I (Y; \mathcal {G} _ {s u b}) + \beta I (\mathcal {G}; \mathcal {G} _ {s u b}), \tag {3}
$$
where  $Y$  is the label of  $\mathcal{G}$ . The first term maximizes the mutual information between the graph label and the compressed subgraph, which ensures that the compressed subgraph contains as much information as possible about the graph label. The second term minimizes the mutual information between the input graph and the compressed subgraph, which ensures that the compressed subgraph contains minimal information about the input graph.
Figure 2: The architecture of our proposed PGIB. PGIB generates a subgraph  $\mathcal{G}_{sub}$  by injecting noise to identify core subgraphs, and it is used to compute similarity scores between prototypes in the prototype layer. The trained prototypes play a crucial role in visualizing the reasoning processes during training in an interpretable manner. PGIB also involves merging pairs of similar prototypes to decrease the number of prototypes. Finally, the integrated prototypes are utilized to predict the graph labels in the fully connected layer.
In this section, we present our proposed method, called PGIB. We introduce Prototype-based Graph Information Bottleneck (Section 3.1), each layer in the architecture (Section 3.2 - 3.4), and the interpretability stabilization process (Section 3.5), which enhances the interpretability and the tracking capabilities of the reasoning process during the model training.
Model Architecture. Figure 2 presents an overview of PGIB. We first generate the representations of nodes in the input graph  $\mathcal{G}$  using a GNN encoder. Then, the node representations are passed to the subgraph extraction layer that assigns each node in  $\mathcal{G}$  to either  $\mathcal{G}_{sub}$  or  $\bar{\mathcal{G}}_{sub}$ . Next, we compute the similarities between the embedding  $\mathbf{z}_{\mathcal{G}_{sub}}$  and the set of prototypes  $\mathcal{Z}_p = \left\{\mathbf{z}_{\mathcal{G}_p}^1,\mathbf{z}_{\mathcal{G}_p}^2,\dots ,\mathbf{z}_{\mathcal{G}_p}^M\right\}$  in the prototype layer. Finally, we merge the prototypes that are semantically similar, which are then used to generate the final prediction.
3.1 Prototype-based Graph Information Bottleneck
PGIB is a novel explainable GNN framework that incorporates the prototype learning within the IB framework, thereby enabling the prototypes to capture the essential key subgraph of the input graph detected by the IB framework. More precisely, we reformulate the GIB objective shown in Equation 3 by decomposing the first term, i.e.,  $I(Y; \mathcal{G}_{sub})$ , with respect to the prototype  $\mathcal{G}_p$  using the chain rule of mutual information in order to examine the impact of the joint information between  $\mathcal{G}_{sub}$  and  $\mathcal{G}_p$  on  $Y$  as follows:
$$
\min  _ {\mathcal {G} _ {s u b}} \underbrace {- I (Y ; \mathcal {G} _ {s u b} , \mathcal {G} _ {p}) + I (Y ; \mathcal {G} _ {p} | \mathcal {G} _ {s u b})} _ {\text {S e c t i o n 3 . 3}} + \beta \underbrace {I (\mathcal {G} ; \mathcal {G} _ {s u b})} _ {\text {S e c t i o n 3 . 2}}. \tag {4}
$$
Please refer to Appendix A.1 for a detailed proof of Equation 4. In the following sections, we describe how each term is optimized during training.
3.2 Subgraph Extraction Layer (Minimizing  $I(\mathcal{G};\mathcal{G}_{sub})$
The goal of the subgraph extraction layer is to extract an informative subgraph  $\mathcal{G}_{sub}$  from  $\mathcal{G}$  that contains minimal information about  $\mathcal{G}$ . We minimize  $I(\mathcal{G};\mathcal{G}_{sub})$  by training the model to inject noise into insignificant subgraphs  $\bar{\mathcal{G}}_{sub}$ , while injecting less noise into more informative ones  $\mathcal{G}_{sub}$  [32]. Specifically, given the representation of node  $v_{i}$ , i.e.,  $\mathbf{h}_i$ , we compute the probability  $p_i$  with an MLP followed by a sigmoid function, which is then used to replace the representation  $h_i$  to obtain the final representation  $\mathbf{z}_i$  as follows:
$$
p _ {i} = \operatorname {S i g m o i d} \left(\operatorname {M L P} \left(h _ {i}\right)\right)
$$
$$
z _ {i} = \lambda_ {i} h _ {i} + (1 - \lambda_ {i}) \epsilon , \text {w h e r e} \lambda_ {i} \sim \operatorname {B e r n o u l l i} \left(p _ {i}\right) \text {a n d} \epsilon \sim \mathcal {N} \left(\mu_ {\mathbf {h} _ {i}}, \sigma_ {\mathbf {h} _ {i}} ^ {2}\right). \tag {5}
$$
That is, the learned probability  $p_i$  enables selective preservation of information in  $\mathcal{G}_{sub}$ , and based on this probability, the quantity of information transmitted from  $\mathbf{h}_i$  to  $\mathbf{z}_i$  can be flexibly adjusted to compress the information from  $\mathcal{G}$  to  $\mathcal{G}_{sub}$ . This approach not only retains interpretability within the
subgraph itself, but also potentially facilitates the learning of prototypes that are introduced in the next step. Following [32], we minimize the upper bound of  $I(\mathcal{G};\mathcal{G}_{sub})$  as follows:
$$
I (\mathcal {G}; \mathcal {G} _ {s u b}) \leq \mathbb {E} _ {\mathcal {G}} \left(- \frac {1}{2} \log A + \frac {1}{2 | \mathcal {V} _ {\mathcal {G}} |} A + \frac {1}{2 | \mathcal {V} _ {\mathcal {G}} |} B ^ {2}\right) =: \mathcal {L} _ {\mathrm {M I}} ^ {1} (\mathcal {G}, \mathcal {G} _ {s u b}), \tag {6}
$$
where  $A = \sum_{i=1}^{|\mathcal{V}_{\mathcal{G}}|}(1 - \lambda_i)^2$  and  $B = \frac{\sum_{i=1}^{|\mathcal{V}_{\mathcal{G}}|}\lambda_i(\mathbf{h}_i - \mu_{\mathbf{h}_i})}{\sigma_{\mathbf{h}_i}}$ . Thus, minimizing  $\mathcal{L}_{\mathrm{MI}}^1$  allows us to minimize the upper bound of  $I(\mathcal{G};\mathcal{G}_{sub})$ . After noise injection, we compute the embedding  $\mathbf{z}_{\mathcal{G}_{sub}}$  through a graph readout function such as max pooling and sum pooling. For further details and analysis of the different graph readout functions, please refer to Appendix A.4.2.
3.3 Prototype Layer (Minimizing  $-I(Y; \mathcal{G}_{sub}, \mathcal{G}_p) + I(Y; \mathcal{G}_p | \mathcal{G}_{sub})$
The prototype layer involves allocation of a fixed number of prototypes for each class. The prototypes are required to capture the most significant graph patterns that can aid in the identification of the graphs within each class. To begin with, we define the similarity score between the prototype  $\mathbf{z}_{\mathcal{G}_p}$  and the embedding  $\mathbf{z}_{\mathcal{G}_{sub}}$  obtained from noise injection as follows:
$$
g \left(\mathbf {z} _ {\mathcal {G} _ {s u b}}, \mathbf {z} _ {\mathcal {G} _ {p}}\right) = \log \left(\frac {\left\| \mathbf {z} _ {\mathcal {G} _ {s u b}} - \mathbf {z} _ {\mathcal {G} _ {p}} \right\| _ {2} ^ {2} + 1}{\left\| \mathbf {z} _ {\mathcal {G} _ {s u b}} - \mathbf {z} _ {\mathcal {G} _ {p}} \right\| _ {2} ^ {2} + \epsilon}\right), \tag {7}
$$
where  $\mathbf{z}_{\mathcal{G}_p}$  is the prototype and shares the same dimension as  $\mathbf{z}_{\mathcal{G}_{sub}}$
3.3.1 Minimizing  $-I(Y;\mathcal{G}_{sub},\mathcal{G}_p)$
We derive the lower bound of  $I(Y; \mathcal{G}_{sub}, \mathcal{G}_p)$  as follows:
Proposition 1. (Lower bound of  $I(Y; \mathcal{G}_{sub}, \mathcal{G}_p)$ ) Given significant subgraph  $\mathcal{G}_{sub}$  for a graph  $\mathcal{G}$ , its label information  $Y$ , prototype graph  $\mathcal{G}_p$  and similarity function  $\gamma$ , we have
$$
\begin{array}{l} I (Y; \mathcal {G} _ {s u b}, \mathcal {G} _ {p}) = \mathbb {E} _ {Y, \mathcal {G} _ {s u b}, \mathcal {G} _ {p}} [ \log p (Y | \mathcal {G} _ {s u b}, \mathcal {G} _ {p}) ] - \mathbb {E} _ {Y} [ \log p (Y) ] \\ \geq \mathbb {E} _ {Y, \mathcal {G} _ {s u b}, \mathcal {G} _ {p}} \left[ \log p (Y | \gamma (\mathcal {G} _ {s u b}, \mathcal {G} _ {p})) \right] - \mathbb {E} _ {Y} [ \log p (Y) ] \\ \geq \mathbb {E} _ {Y, \mathcal {G} _ {s u b}, \mathcal {G} _ {p}} \left[ \log q _ {\theta} \left(Y | \gamma \left(\mathcal {G} _ {s u b}, \mathcal {G} _ {p}\right)\right) \right] \\ =: - \mathcal {L} _ {\mathrm {c l s}} \left(q _ {\theta} (Y | \gamma (\mathcal {G} _ {s u b}, \mathcal {G} _ {p})) \right. \\ \end{array}
$$
where  $q_{\theta}\left(Y|\gamma \left(\mathcal{G}_{sub},\mathcal{G}_p\right)\right)$  is the variational approximation to the true posterior  $p\left(Y|\gamma \left(\mathcal{G}_{sub},\mathcal{G}_p\right)\right)$ .
Equation 8 demonstrates that the maximization of the mutual information  $I(Y; \mathcal{G}_{sub}, \mathcal{G}_p)$  can be attained by minimizing the classification loss, denoted as  $\mathcal{L}_{cls}$ . This maximization of mutual information between the label  $Y$  and the similarity information  $\gamma(\mathcal{G}_{sub}, \mathcal{G}_p)$  promotes the subgraph  $\mathcal{G}_{sub}$  and prototype  $\mathcal{G}_p$  to possess predictive capabilities concerning the graph label  $Y$ . In practical applications, the cross-entropy loss is chosen for a categorical  $Y$ . For a comprehensive understanding of the derivation process of Equation 8, refer to the Appendix A.2. Note that the similarity between  $\mathcal{G}_{sub}$  and  $\mathcal{G}_p$ , i.e.,  $\gamma(\mathcal{G}_{sub}, \mathcal{G}_p)$ , is computed by the similarity score defined in Equation 7, i.e.,  $g(\mathbf{z}_{\mathcal{G}_{sub}}, \mathbf{z}_{\mathcal{G}_p})$ .
3.3.2 Minimizing  $I(Y; \mathcal{G}_p | \mathcal{G}_{sub})$
We investigate the mutual information, denoted as  $I(Y; \mathcal{G}_p | \mathcal{G}_{sub})$ , from the perspective of the interaction between  $\mathcal{G}_{sub}$  and  $\mathcal{G}_p$ . We decompose  $I(Y; \mathcal{G}_p | \mathcal{G}_{sub})$  into the sum of two terms based on the chain rule of mutual information as follows:
$$
I (Y; \mathcal {G} _ {p} | \mathcal {G} _ {s u b}) = I (\mathcal {G} _ {p}; Y, \mathcal {G} _ {s u b}) - I (\mathcal {G} _ {s u b}; \mathcal {G} _ {p}). \tag {9}
$$
It is important to note that the first term, i.e.,  $I(\mathcal{G}_p;Y,\mathcal{G}_{sub})$ , minimizes the mutual information between  $\mathcal{G}_p$  and the joint variables  $(Y,\mathcal{G}_{sub})$ , which eliminates the information about  $Y$  related to  $\mathcal{G}_{sub}$  from  $\mathcal{G}_p$ . However, since our goal is not to solely minimize  $I(\mathcal{G}_p;Y,\mathcal{G}_{sub})$  but to ensure the interpretability of the prototype  $\mathcal{G}_p$ , including this term leads to diminished interpretability of  $\mathcal{G}_p$ . Consequently, we excluded the first term during training, and only consider the second term, i.e.,  $-I(\mathcal{G}_{sub};\mathcal{G}_p)$ , to simultaneously guarantee the interpretability of both  $\mathcal{G}_{sub}$  and  $\mathcal{G}_p$ . A detailed derivation for Equation 9 is given in Appendix A.3. From now on, we describe approaches for minimizing the second term. Inspired by [11], we introduce two different approaches for minimizing  $-I(\mathcal{G}_{sub};\mathcal{G}_p)$ .
1) Variational IB-based approach. We obtain the upper bound of  $-I(\mathcal{G}_{sub};\mathcal{G}_p)$  using the variational IB-based approach as follows:
$$
- I (\mathcal {G} _ {s u b}; \mathcal {G} _ {p}) \leq \mathbb {E} _ {\mathcal {G} _ {s u b}, \mathcal {G} _ {p}} [ - \log q _ {\phi} (\mathcal {G} _ {p} | \mathcal {G} _ {s u b}) ] := \mathcal {L} _ {\mathrm {M I}} ^ {2} (\mathcal {G} _ {s u b}, \mathcal {G} _ {p}), \tag {10}
$$