Optimizing OOD Detection in Molecular Graphs: A Novel Approach with Diffusion Models

Since graph neural networks can use the topological structure and node properties of graphs for representation learning, they have become the most powerful method for processing graph data (Feng et al., 2022; Zhao et al., 2021; Bevilacqua et al., 2021; Zhou et al., 2020), especially molecular graphs (Wang et al., 2022; Wieder et al., 2020). GCN (Kipf and Welling, 2016), the simplest but most efficient method, has been proved to be equivalent to the first-order approximation filter on graphs (He et al., 2021) and thus performs well in node classification (Hamilton et al., 2017) and link prediction (Cai et al., 2021). On graph instance-related tasks, GIN (Xu et al., 2018) proves that GNN is as powerful as the 1-WL test and leverages an injective summation operation to increase performance. More and more researchers have proposed more representational methods, but they all ignore the performance and trustworthiness issues brought by OOD distribution (Wang et al., 2021; Wu et al., 2022).

Graph generative models aim to learn the distribution of the graph data and sample from it to generate novel graphs (Zhu et al., 2022), especially for molecular graphs since it is related to many science issues (Li et al., 2017; Reiser et al., 2022; Ingraham et al., 2019). Some graph generation methods are inspired by auto-regressive models, such as VAE-based (Mitton et al., 2021) or normalizing flow-based models (Kuznetsov and Polykovskiy, 2021). However, they are limited by the high computational cost and inability to model permutation invariance of graph (Jo et al., 2022). Inspired by the diffusion models in computer vision (Song et al., 2021), the same insight on graphs has developed in recent years (Niu et al., 2020; Vignac et al., 2022; Chen et al., 2023). Although diffusion models achieve state-of-the-art performance, they still suffer from inefficiencies caused by slow denoising processes (Limnios et al., 2023).

Recently, many studies focus on graph OOD detection due to its importance. GOOD-D is the pioneering work for unsupervised OOD graph detection, which performs hierarchical contrastive learning to capture latent ID patterns and detects OOD graphs based on their semantic inconsistency (Liu et al., 2023b). GraphDE determines ID and OOD by inferring the environment variables of the graph generation process (Li et al., 2022). AAGOD aims to learn a parameterized amplifier matrix to emphasize the key patterns which helpful for graph OOD detection, thereby enlarging the gap between OOD and ID graphs (Guo et al., 2023). Anomaly graph detection can also be seen as a special case of OOD detection, since anomaly graphs with anomaly structures and features can be caused by distribution shifts and many methods have been proposed to solve it (Qiu et al., 2022; Ma et al., 2022). All of the above methods require redesigning or training well-performing GNNs on the ID datasets and inevitably lead to a trade-off between OOD detection and ID prediction.

For OOD detection task, we aim to design a detector $g$ to distinguish whether the input graph $G$ is an OOD sample or not:

where $J$ denotes a judging function to score the input molecules and $\tau$ denotes threshold for identifying the OOD samples. A desired OOD detector should assign judge scores with the maximum gap between ID and OOD samples. This target can be described as the following optimization:

Supposing the judge score distributions of ID and OOD have significant divergence, we can distinguish them with a simple intermediate threshold. For reconstruction-based OOD detection as shown in Fig. 2, the similarity between the input and the output molecules of diffusion model $\mathrm{F_{M}}$ is often adopted as the judge function:

where $\mathrm{F_{M}}(G)$ is the reconstructed output and $\operatorname{sim}(\cdot)$ is the similarity function. OOD inputs correspond to the lower reconstruction quality and therefore the lower similarity, while the similarity measurement is higher for the ID inputs.

The typical GNNs are based on message passing paradigm. Specifically, the final representation of graph $G$ for a $L$-layer GNNs is:

where $m^{(0)}_{v}=X_{v}$ is raw node feature, $N(v)$ represents a set of neighbor nodes with respect to node $v$, and $\operatorname{MP}$ is the message passing process that aggregates neighborhood features (e.g., sum, mean, or max) and combines them with the local node. GNNs iteratively perform $\operatorname{MP}$ to learn the effective node representations and utilize function $\operatorname{Pooling}$ to map all the node representations into the graph representations, which is a single vector.

The generative method based on the diffusion model consists of a forward diffusion process and a reverse denoising process. At the forward process, the model progressively adds noise to the original data until a standard normal distribution. At the reverse process, the model learns the score function (i.e., a neural network) to remove the perturbed noise with the same amount of steps (Croitoru et al., 2023; Liu et al., 2023c; Song et al., 2021).

Given a graph ${G}=(A,X)$, we can use continuous time $t\in[0,T]$ to index the diffusion trajectory $\left\{{G_{t}}=(A_{t},X_{t})\right\}^{T}_{t=1}$, such that $G_{0}$ is the original input graph and $G_{T}$ approximately follows the normal distribution. The forward process transforms $G_{0}$ to $G_{T}$ through a stochastic differential equation (SDE):

where $\mathbf{w}$ is standard Wiener process (Jo et al., 2022), $\mathbf{f}_{t}(\cdot):\mathcal{G}\to\mathcal{G}$ is linear drift coefficient, $g(t):\mathbb{R}\to\mathbb{R}$ is a scalar function which represents the diffusion coefficient. $\mathbf{f}_{t}({G}_{t})$ and $g(t)$ relate to the amount of noise $\mathrm{d}\mathbf{w}$ added to the graph at each infinitesimal step $\mathrm{d}t$. In order to generate graphs that follow the distribution of $G_{0}$, we start from $G_{T}$ and utilize a reverse-time SDE for denoising from $T$ to $0$:

where $\mathrm{S}_{\theta}(G_{t},t)$ is score function to estimate the scores of perturbed graphs $\nabla_{G_{t}}\log p_{t}(G_{t})$ and $p_{t}(G_{t})$ is the marginal distribution under the forward process at time $t$. In practice, two GNNs are utilized as the score function to denoise both node features and graph structures. $\bar{\mathbf{w}}$ is a reverse time standard Wiener process.

Inspired by the generative methods (Liu et al., 2023d; Gao et al., 2023), we design a vanilla graph reconstruction model (GR-MOOD) for molecular graph OOD detection. GR-MOOD is pre-trained on a large-scale compound dataset (e.g., QM9 or ZINC) and fine-tuned on $D_{\mathrm{train}}$. Considering input graph $G\in D_{\mathrm{test}}$, we utilize GR-MOOD to perturb and reconstruct it via:

where $\theta$ is the parameters of GR-MOOD, and $T$ is the iteration numbers. Function $\operatorname{diffuse}(\cdot)$ applies Eq. (8) to introduce perturbations that transform $G$ into a noised state $G_{\mathrm{o}}$, while function $\operatorname{denoise}(\cdot)$ utilizes Eq. (9) to reverse the process, effectively denoising $G_{\mathrm{o}}$ to generate reconstruction graph $\hat{G}$.

Upon acquiring the reconstruction graph $\hat{G}$, we utilize a GNN well-trained on the ID dataset to encode both the feature and structure information of $G$ and $\hat{G}$, whose representations are denoted as $z$ and $\hat{z}$, respectively. The cosine similarity between them is treated as OOD judge score and is defined in Eq. (5):

To validate GR-MOOD effectiveness, we conduct experiments on two DrugOOD datasets (Ji et al., 2023). As shown in Fig. 3, the performance of GR-MOOD is comparable (e.g., AUROC and AUPR) or even outperforming (e.g., the smaller score of FPR95 is better) than the SOTA method of GOOD-D (Liu et al., 2023b). The underlying principle is that since GR-MOOD is trained to reconstruct graphs that align with the ID distribution, OOD samples, due to their inherent dissimilarity from the ID distribution, will typically undergo poorer reconstruction when being processed. Such discrepancy is quantified as a lower judge score, which signals the presence of an OOD sample. This mechanism highlights the critical role of diffusion model based reconstruction method in identifying graphs that do not conform to the expected distribution, thereby providing a quantitative basis for distinguishing between ID and OOD samples.

To address the limitations of GR-MOOD, we propose a novel approach based upon diffusion model, PGR-MOOD (Prototypical Graph Reconstruction for Molecular OOD Detection). The innovation of PGR-MOOD has three aspects: A strong similarity function, a prototypical graphs generator, and an efficient and scalable OOD detector. The architecture of PGR-MOOD is shown in Fig. 5.

Specifically, $\operatorname{FGW}$ treats the graph associated with topology and node feature as a probability distribution. It allows for the computation of costs between two distributions with optimal coupling, serving as a distance measure between graphs. For two graphs represented in OT format, $G_{1}=(A_{1},X_{1},\mu_{1})$ and $G_{2}=({{A_{2}}},{{X_{2}},\mu_{2}})$, their $\operatorname{FGW}$ distance is defined as:

where $A_{1}(i,j)$ represents the element of the $i$-th row and $j$-th column in $A_{1}$, $X_{1}(i)$ represents the $i$th row vector of $X$, $\alpha\in[0,1]$ is a parameter to balance the structure term and the feature term, $\Pi(\mu_{1},\mu_{2})=\{\pi\in R_{+}^{m\times n}\mathrm{s.t.},\sum_{i=1}^{m}\pi_{i,j}=\mu_{2}(j),\sum_{j=1}^{n}\pi_{i,j}=\mu_{1}(i)\}$ is the set of all admissible couplings between $\mu_{1}$ and ${\mu_{2}}$. $\operatorname{FGW}(\cdot)$ metric exhibits optimal performance in directly discerning both structural variances and feature disparities between graphs.

To satisfy Property ➀, Eq. (13) is utilized as the distance metric, and the loss function $\mathcal{L}_{\mathrm{ID}}$ is formulated to guide the denosing process at the generator:

Similarly to comply with Property ➁, we introduce loss function $\mathcal{L}_{\mathrm{OOD}}$ to enhance the distance between $\overline{G}$ from OOD samples:

Note that OOD graphs $G_{\mathrm{out}}$ are unreachable during the training phase, precluding the direct formulation of $\mathcal{L}_{\mathrm{OOD}}$. Consequently, it becomes imperative to synthesize graphs as proxies for the absent OOD samples. Recalling the pre-trained diffusion model $\mathrm{F_{M}}$ in Eq. (9), it adopts socre function $\mathrm{S}_{\theta}$ to generate graph. The parameter weights of $\mathrm{S}_{\theta}$ is given by $\theta_{\mathrm{M}}=\{\theta_{\mathrm{M}}^{(l)}\}_{l=1}^{L}$, where $\theta_{\mathrm{M}}^{(l)}$ represent the parameters of the $l$-th score function. We propose to directly perturb parameters $\theta_{\mathrm{M}}$ for generating OOD graphs $G_{\mathrm{out}}$:

where $\alpha>0$ is perturbation strength, $I$ is identity matrix, and $P^{(l)}$ is perturbation matrix. By perturbing the parameters $\theta_{\mathrm{M}}$, a new score function $\mathrm{S}_{\tilde{\theta}}(\cdot)$ is derived. Experimental observations (w/o $\mathcal{L}_{\mathrm{OOD}}$ of Table 2) reveal that $\mathrm{S}_{\tilde{\theta}}(\cdot)$ can induce a deviation in the denoising trajectory away from the original data distribution, thereby enabling the diffusion model to generate $G_{\mathrm{out}}$ during the training phase. In light of these researches, a composite loss function $\mathcal{L}_{\mathrm{guide}}$ is formulated by integrating both $\mathcal{L}_{\mathrm{OOD}}$ and $\mathcal{L}_{\mathrm{ID}}$:

It is leveraged to guides the training of Prototypical Graphs Generator $\mathrm{F_{PG}}$, which has the same architecture and initial parameters $\theta$ with $\mathrm{F_{M}}$, to generate prototypical graph $\overline{G}$. The generation of $\overline{G}$ by $\mathrm{F_{PG}}$ unfolds in two phases: Firstly, in contrast to generating directly from Gaussian noise, a graph $G_{0}$ from $D_{\mathrm{train}}$ is randomly chosen as the start point of generation. We then add $T$-step noise according to Eq. (8) to get the final noise graph $G_{T}$ (i.e., $G_{0}\to G_{T}$). Secondly, $\mathcal{L}_{\mathrm{guide}}$ guides the denoising step of diffusion model to generate prototype graph $\overline{G}$:

where $t$ is the indicator of the denoise step and varies from $T$ to $0$. The prototype graph $\overline{G}$ generated by the above equation can be viewed as the reconstruction of both ID and OOD graphs, but has better discrimination than the reconstruction generated in GR-MOOD. To further reduce the computation, rather than utilizing the entirety of $D_{\mathrm{train}}$, a fixed batch-size dataset $D_{\mathrm{\mathrm{batch}}}$ is employed for the computation of $\mathcal{L}_{\mathrm{ID}}$. Each $D_{\mathrm{batch}}$ can generate one $\overline{G}$, and they are combined to formulate a list $PL=\{\overline{G}^{(i)}\}^{I}_{i=1}$, $I=\lceil\frac{|D_{\mathrm{train}}|}{|D_{\mathrm{batch}}|}\rceil$.

Q: Whether PGR-MOOD achieves the best performance on the OOD detection in molecular graphs? Yes, we utilize the new loss function $\mathcal{L}_{\mathrm{guide}}$ to guide the diffusion model to generate prototypical graphs that are more representative of all ID samples, and more easily detect OOD samples.

Q: Whether PGR-MOOD can enlarge the judge score gap between ID and OOD graphs? Yes, we calculate the similarity between the prototypical graphs and test graphs, which has a massive difference for ID and OOD. A more significant gap between ID and OOD graphs corresponds to a better graph OOD detector. We present the scoring distributions on two datasets in Fig. 6. The ID and OOD are perfectly separated into two distinct distributions, so we can use a simple threshold for OOD detection and achieve SOTA performance.

Q: Whether each module in PGR-MOOD contribute to effectively discriminating OOD molecular graphs? Yes, we conduct experiments on four datasets to verify the role of $\mathcal{L}_{\mathrm{ID}}$, $\mathcal{L}_{\mathrm{OOD}}$, and $\operatorname{FGW}$ modules in PRG-MGOD. The results are shown in Table 2.