Pretrained Joint Predictions for Scalable Batch Bayesian Optimization of Molecular Designs

TL;DR

This paper introduces a scalable probabilistic surrogate model using Epistemic Neural Networks for batch Bayesian optimization in molecular design, significantly reducing the number of iterations needed to discover potent drug candidates.

Contribution

It develops a scalable joint predictive distribution framework with pretraining strategies for ENNs, enhancing batch Bayesian optimization in drug discovery.

Findings

Rediscovered known EGFR inhibitors with 5x fewer iterations.

Identified potent inhibitors from real-world libraries with 10x fewer iterations.

Demonstrated effectiveness on semi-synthetic and real-world benchmarks.

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.