Reliable OOD Virtual Screening with Extrapolatory Pseudo-Label Matching

Machine learning (ML) models are increasingly deployed for virtual screening in drug discovery, where the goal is to identify novel, chemically diverse scaffolds while minimizing experimental costs. This creates a fundamental challenge: the most valuable discoveries lie in out-of-distribution (OOD) regions beyond the training data, yet ML models often degrade under distribution shift. Standard novelty-rejection strategies ensure reliability within the training domain but limit discovery by rejecting precisely the novel scaffolds most worth finding. Moreover, experimental budgets permit testing only a small fraction of nominated candidates, demanding models that produce reliable confidence estimates. We introduce EXPLOR (Extrapolatory Pseudo-Label Matching for OOD Uncertainty-Based Rejection), a framework that addresses both challenges through extrapolatory pseudo-labeling on latent-space augmentations, requiring only a single labeled training set and no access to unlabeled test compounds, mirroring the realistic conditions of prospective screening campaigns. Through a multi-headed architecture with a novel per-head matching loss, EXPLOR learns to extrapolate to OOD chemical space while producing reliable confidence estimates, with particularly strong performance in high-confidence regions, which is critical for virtual screening where only top-ranked candidates advance to experimental validation. We demonstrate state-of-the-art performance across chemical and tabular benchmarks using different molecular embeddings.

💡 Research Summary

The paper tackles a central problem in ligand‑based virtual screening (LBVS): discovering novel, chemically diverse scaffolds that lie outside the distribution of the available training data while still providing reliable confidence estimates for the top‑ranked candidates that will be experimentally tested. Conventional approaches either rely on novelty‑rejection, which discards out‑of‑distribution (OOD) compounds and thus hampers discovery, or they focus on global metrics (AUROC, AUPRC) that do not reflect the early‑recognition nature of screening campaigns. To bridge this gap, the authors introduce EXPLOR (Extrapolatory Pseudo‑Label Matching for OOD Uncertainty‑Based Rejection), a framework that works with a single labeled source dataset and no unlabeled test set, matching realistic prospective screening conditions.

EXPLOR consists of three stages. First, it builds a diverse ensemble of K pseudo‑labelers. Each pseudo‑labeler is trained on a different random subset of features and instances, encouraging specialization and diversity. Second, it expands the effective training support by perturbing inputs in a learned latent space. An autoencoder provides an encoder φ and decoder γ; latent vectors z = φ(x) are radially scaled by (1+|ε|) where ε follows a half‑normal distribution, yielding expanded samples x′ = γ(z′). These samples lie beyond the original data manifold and may overlap with OOD regions. Since true labels are unavailable for expanded points, the pseudo‑labelers supply K soft targets. Third, a multi‑headed neural network is trained such that each head h_k matches the predictions of one pseudo‑labeler g_k via a per‑head matching loss, while an additional regularizer encourages agreement across heads to stabilize confidence estimates.

This design yields two key benefits. The latent‑space expansion forces the model to learn extrapolative behavior, directly exposing it to OOD‑like inputs during training. The multi‑head, per‑head matching preserves the diversity of labeling functions, reducing reliance on any single brittle predictor, and the agreement regularizer prevents excessive variance, resulting in well‑calibrated high‑confidence predictions. Importantly, the framework operates on real‑valued vector representations (e.g., Morgan fingerprints, pretrained embeddings) and does not require modality‑specific augmentations or multiple source domains.

The authors evaluate EXPLOR on several chemical OOD benchmarks (scaffold split, temporal split) and on tabular datasets, comparing against supervised ERM, semi‑supervised methods (Mean Teacher), and domain‑generalization techniques (IRM, GroupDRO). They employ both traditional global metrics and screening‑relevant metrics: a truncated precision‑recall area (AUPRC@R < τ) that measures performance at low recall (the region where only a few top candidates are tested), as well as calibration measures (ECE, NLL). EXPLOR consistently outperforms baselines in the high‑confidence regime, showing markedly higher precision among the top 1–5 % of predictions, better calibration, and reduced variance across random seeds. Qualitative analysis reveals that EXPLOR identifies a broader set of OOD actives rather than overfitting to a narrow region of the training manifold.

In summary, EXPLOR offers a practical, modality‑agnostic solution for single‑source LBVS under distribution shift. By combining diverse pseudo‑labelers, latent‑space extrapolation, and multi‑head matching, it achieves controlled OOD extrapolation while delivering reliable uncertainty estimates where they matter most—among the few compounds that will be purchased and assayed. The work highlights the importance of early‑recognition‑aligned evaluation for drug discovery and opens avenues for further improvements such as meta‑learning of pseudo‑labelers, non‑linear latent augmentations, and prospective validation in real screening campaigns.