Scaffold-Aware Generative Augmentation and Reranking for Enhanced Virtual Screening

TL;DR

The paper tackles three persistent challenges in ligand-based virtual screening: extreme class imbalance, structural imbalance among active molecules, and the need for scaffold diversity among top candidates. It introduces ScaffAug, a scaffold-aware framework with three modules: (i) augmentation using scaffold-aware sampling (SAS) and scaffold extension via a graph diffusion model (DiGress) to produce a Generative Diverse Scaffold-Augmented (G-DSA) dataset, (ii) a model-agnostic self-training pipeline with confidence-based pseudo-labeling to safely integrate synthetic data, and (iii) a reranking step based on Maximal Marginal Relevance (MMR) to boost scaffold diversity in the top predictions. Across five WelQrate datasets, ScaffAug with DeepGIN delivers superior performance on ranking metrics such as logAUC and BEDROC, and achieves higher scaffold diversity (SD$_{100}$) while maintaining hit enrichment (EF$_{100}$). The work demonstrates that scaffold-aware generative augmentation, safe integration of synthetic data, and diversity-oriented reranking can substantially improve early discovery efficiency in virtual screening and promote exploration of novel chemical spaces. This framework advances computational drug discovery by effectively leveraging generative AI to balance data distributions and diversify top hits, with practical implications for scaffold hopping and novel active discovery.

Abstract

Ligand-based virtual screening (VS) is an essential step in drug discovery that evaluates large chemical libraries to identify compounds that potentially bind to a therapeutic target. However, VS faces three major challenges: class imbalance due to the low active rate, structural imbalance among active molecules where certain scaffolds dominate, and the need to identify structurally diverse active compounds for novel drug development. We introduce ScaffAug, a scaffold-aware VS framework that addresses these challenges through three modules. The augmentation module first generates synthetic data conditioned on scaffolds of actual hits using generative AI, specifically a graph diffusion model. This helps mitigate the class imbalance and furthermore the structural imbalance, due to our proposed scaffold-aware sampling algorithm, designed to produce more samples for active molecules with underrepresented scaffolds. A model-agnostic self-training module is then used to safely integrate the generated synthetic data from our augmentation module with the original labeled data. Lastly, we introduce a reranking module that improves VS by enhancing scaffold diversity in the top recommended set of molecules, while still maintaining and even enhancing the overall general performance of identifying novel, active compounds. We conduct comprehensive computational experiments across five target classes, comparing ScaffAug against existing baseline methods by reporting the performance of multiple evaluation metrics and performing ablation studies on ScaffAug. Overall, this work introduces novel perspectives on effectively enhancing VS by leveraging generative augmentations, reranking, and general scaffold-awareness.

Figures (5)

• Figure 1: Representative overview of ScaffAug's performance compared to established baseline methods on benchmark datasets. The figure highlights improvements achieved through our scaffold-aware augmentation and reranking framework. Comprehensive experimental analyses, including additional metrics and ablation studies, are presented in Section \ref{['experiments']}.
• Figure 2: ScaffAug Framework. (a) The augmentation module first starts with the scaffold-aware sampling algorithm, which addresses the structural imbalance among active molecules. It clusters active scaffolds based on structure and favors sampling underrepresented scaffolds when constructing the scaffold library. GDM then generates new synthetic samples conditioned on scaffolds in the scaffold library to form the generative diverse scaffold-augmented (G-DSA) dataset. (b) The self-training module incorporates labeled data and augmentation data with a pseudo-labeling strategy. (c) The reranking module leverages prediction scores output by models to produce a more structurally diverse set of top compounds.
• Figure 3: MMR reranking results for Scaffaug$_{\text{DeepGIN}}$. SD$_{100}$ and EF$_{100}$ represent scaffold diversity and enrichment factor for the reranked molecular sets, respectively. Note that the original scores before reranking correspond to $\lambda=1$. The highest EF$_{100}$ score after reranking and its corresponding SD$_{100}$ score at the same $\lambda$ value are highlighted.
• Figure 4: Sampling algorithms affect molecular diversity. We compare the chemical descriptors of original active molecules from the training set, molecules generated using the uniform actives (UA) sampling algorithm, and molecules generated using the scaffold-aware sampling (SAS) algorithm. We convert all molecules to ECFP fingerprints and project them to the same 2D space via UMAP. We then display the number of molecules in each grid, showing the diversity level in different regions of chemical space.
• Figure 5: Comparison of oversampling baseline with ScaffAug variants using UA and SAS. Mean percentage change across five datasets and both split types..