Pretrained Joint Predictions for Scalable Batch Bayesian Optimization of Molecular Designs

2025-11-13

AI Summary

In drug discovery, molecular design is bottlenecked by low accuracy in binding affinity prediction and poor scalability of Bayesian optimization (BO) for batch synthesis and testing. To address this, we propose Epistemic Neural Networks—a framework grounded in epistemic uncertainty modeling—that integrates structure-aware molecular representations, synthetic-data pretraining, and joint probability density sampling to construct an efficient, scalable surrogate model supporting parallel acquisition function optimization. Our method is the first to enable end-to-end modeling of affinity distributions within batch BO. In a semi-synthetic EGFR inhibitor discovery task, it achieves potent molecule rediscovery in one-fifth the iterations required by baselines; in real-world small-molecule library screening, it delivers up to 10× speedup. The core contribution lies in explicitly embedding epistemic uncertainty into the pretrain-fine-tune paradigm, markedly enhancing robustness and generalization of batch decisions under limited data.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Batched synthesis and testing of molecular designs is the key bottleneck of drug development. There has been great interest in leveraging biomolecular foundation models as surrogates to accelerate this process. In this work, we show how to obtain scalable probabilistic surrogates of binding affinity for use in Batch Bayesian Optimization (Batch BO). This demands parallel acquisition functions that hedge between designs and the ability to rapidly sample from a joint predictive density to approximate them. Through the framework of Epistemic Neural Networks (ENNs), we obtain scalable joint predictive distributions of binding affinity on top of representations taken from large structure-informed models. Key to this work is an investigation into the importance of prior networks in ENNs and how to pretrain them on synthetic data to improve downstream performance in Batch BO. Their utility is demonstrated by rediscovering known potent EGFR inhibitors on a semi-synthetic benchmark in up to 5x fewer iterations, as well as potent inhibitors from a real-world small-molecule library in up to 10x fewer iterations, offering a promising solution for large-scale drug discovery applications.

Problem

• Developing scalable probabilistic surrogates for molecular binding affinity prediction
• Accelerating batch Bayesian optimization for drug discovery applications
• Improving joint predictive distributions using pretrained epistemic neural networks

Innovation

• Uses Epistemic Neural Networks for joint predictions
• Pretrains prior networks on synthetic data
• Leverages structure-informed representations for binding affinity