Out-of-Distribution Detection in Molecular Complexes via Diffusion Models for Irregular Graphs

Predictive machine learning models generally excel on in-distribution data, but their performance degrades on out-of-distribution (OOD) inputs. Reliable deployment therefore requires robust OOD detection, yet this is particularly challenging for irregular 3D graphs that combine continuous geometry with categorical identities and are unordered by construction. Here, we present a probabilistic OOD detection framework for complex 3D graph data built on a diffusion model that learns a density of the training distribution in a fully unsupervised manner. A key ingredient we introduce is a unified continuous diffusion over both 3D coordinates and discrete features: categorical identities are embedded in a continuous space and trained with cross-entropy, while the corresponding diffusion score is obtained analytically via posterior-mean interpolation from predicted class probabilities. This yields a single self-consistent probability-flow ODE (PF-ODE) that produces per-sample log-likelihoods, providing a principled typicality score for distribution shift. We validate the approach on protein-ligand complexes and construct strict OOD datasets by withholding entire protein families from training. PF-ODE likelihoods identify held-out families as OOD and correlate strongly with prediction errors of an independent binding-affinity model (GEMS), enabling a priori reliability estimates on new complexes. Beyond scalar likelihoods, we show that multi-scale PF-ODE trajectory statistics - including path tortuosity, flow stiffness, and vector-field instability - provide complementary OOD information. Modeling the joint distribution of these trajectory features yields a practical, high-sensitivity detector that improves separation over likelihood-only baselines, offering a label-free OOD quantification workflow for geometric deep learning.

💡 Research Summary

The paper introduces a novel unsupervised out‑of‑distribution (OOD) detection framework specifically designed for irregular three‑dimensional (3D) molecular graphs, such as protein‑ligand complexes, which combine continuous spatial coordinates with categorical chemical attributes. Traditional OOD methods either rely on supervised discriminative scores or on generative models that treat 2D topological graphs; they struggle with the joint geometry‑chemistry nature of molecular complexes and suffer from complexity bias, where simple OOD inputs receive high likelihoods.

To address these challenges, the authors develop a unified continuous diffusion model. Categorical atom and residue types are embedded onto a spherical latent space, concatenated with the 3D coordinates, and the combined state is diffused using a stochastic differential equation (SDE). A single SE(3)‑equivariant graph neural network (EGNN) predicts denoised coordinates and class logits at each diffusion step. The logits are transformed into a posterior‑mean “clean” embedding, allowing the model to be trained with standard cross‑entropy while preserving smooth continuous dynamics required for diffusion.

The key theoretical tool is the probability‑flow ordinary differential equation (PF‑ODE), the deterministic counterpart of the diffusion SDE that describes the average marginal flow of probability mass. For any input sample, solving the PF‑ODE yields a unique trajectory from the data point (t = 0) to pure Gaussian noise (t = T) and back. The exact log‑likelihood of the sample is obtained by integrating the divergence of the drift term along this trajectory, providing a principled typicality score. However, as shown in prior work, raw likelihoods are vulnerable to complexity bias: low‑complexity OOD structures can be assigned high likelihoods.

To mitigate this, the authors augment the scalar likelihood with eighteen trajectory‑level geometric features extracted from the PF‑ODE path. These include path efficiency, tortuosity, total angular deviation, smoothness, acceleration, vector‑field L2 statistics (mean, std, max), spikiness, total flow energy, Lipschitz estimates (mean and max), center‑of‑mass drift statistics, inter‑molecular distance changes, and ratios that capture the coupling between coordinate and feature diffusion. The intuition is that in‑distribution (ID) samples travel through well‑trained, high‑density regions, resulting in short, smooth, low‑energy trajectories, whereas OOD samples must traverse low‑density, poorly constrained regions, producing erratic, high‑energy, and unstable flows.

Empirical evaluation uses the PDBbind 2020 dataset (19 443 complexes). The authors construct a rigorous split: a standard validation set (minimal shift), the CASF‑2016 benchmark (intermediate shift), and seven strict OOD test sets obtained by withholding entire protein families (e.g., serine/threonine kinases, estrogen receptors, HIV proteases, etc.). Bioinformatic similarity analyses (TM‑score for protein structure, ligand fingerprint similarity, and binding‑pose alignment) confirm that the OOD sets are progressively farther from the training distribution.

Results show that PF‑ODE log‑likelihood alone can separate most OOD families but still exhibits substantial overlap with the training distribution, and it fails dramatically on a low‑complexity OOD set, misclassifying it as ID. When the eighteen trajectory features are combined with the likelihood in a simple classifier (e.g., logistic regression or a shallow MLP), detection performance improves markedly: ROC‑AUC values exceed 0.90 across all OOD sets, and the previously misidentified low‑complexity set is correctly flagged. Moreover, the PF‑ODE likelihood correlates strongly (Pearson ≈ 0.68–0.73) with prediction errors of an independent binding‑affinity predictor (GEMS) trained on the same data, demonstrating that the OOD score can serve as a pre‑emptive reliability estimate for downstream tasks.

The contributions of the work are fourfold: (1) a unified diffusion framework that jointly models continuous geometry and discrete chemical identities; (2) a principled OOD scoring mechanism based on PF‑ODE likelihood augmented with multi‑scale trajectory statistics; (3) the first unsupervised OOD detection study on 3D molecular graphs with strict family‑level hold‑out validation; (4) empirical evidence that the OOD score predicts downstream model errors, enabling confidence‑aware deployment in drug‑discovery pipelines.

In discussion, the authors note that the approach is model‑agnostic (any diffusion model with a tractable PF‑ODE can be used) and label‑free, making it applicable to other scientific domains where data are represented as irregular 3D graphs, such as metal‑organic frameworks, nanoparticle assemblies, or materials‑science simulations. Future work may explore more sophisticated trajectory‑feature learning, integration with active learning loops, or extension to conditional diffusion models for generative design tasks. Overall, the paper demonstrates that diffusion‑based density estimation, when coupled with detailed flow dynamics analysis, provides a robust and interpretable solution to OOD detection in complex geometric deep‑learning settings.