Pretrained Joint Predictions for Scalable Batch Bayesian Optimization of Molecular Designs

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.

• The paper introduces a scalable surrogate modeling approach using pretrained prior functions to achieve accurate joint predictions essential for batch selection.
• It integrates chemistry-specific foundation models with Epistemic Neural Networks to capture binding affinity and quantify uncertainty in candidate batches.
• The method demonstrates significant efficiency improvements, outperforming traditional optimization and random sampling in drug discovery benchmarks.

Scalable Joint Predictions for Batch Bayesian Optimization in Molecular Design

Introduction

This paper addresses a primary bottleneck in small-molecule drug discovery: the efficient selection of compound batches for experimental synthesis and testing in Design-Make-Test-Analyze (DMTA) cycles. Bayesian Optimization (BO), particularly its batch implementation, has been a canonical framework for navigating vast chemical spaces under uncertainty. The paper focuses on surrogate modeling for binding affinity given high-dimensional, diverse datasets, examining the need for scalable joint predictive distributions to enable parallel acquisition functions that account for correlations within a candidate batch. The work advances Epistemic Neural Networks (ENNs) by pretraining their additive prior networks on synthetic processes, integrating representations from chemistry-specific foundation models, and benchmarking performance on realistic drug discovery tasks.

Methodology

Batch Bayesian Optimization and Surrogate Requirements

Batch BO requires both:

• Parallel acquisition functions that balance between exploitation and hedged exploration among selected candidate molecules;
• Joint predictive densities $p_e(y_{1:N}|x_{1:N})$ capturing dependencies between candidates to reason about batch-wise uncertainty and facilitate acquisition such as qPO and EMAX.

Traditional surrogates like Gaussian Processes (GPs) excel at modeling joint predictions but scale cubically in data size, intractable for large chemical libraries. Deep ensembles and Bayesian Neural Networks (BNNs) ease scaling via parameter amortization but struggle to produce joint sample paths necessary for parallel acquisition—MC errors decrease slowly and are not sample-efficient for batch sizes typical in industrial workflows.

Epistemic Neural Networks and Pretrained Prior Functions

ENNs propose marginalizing output predictions over a latent "epistemic index" $z \sim p_z(z)$, rather than network weights, yielding parameter-efficient surrogates and facilitating joint predictive sampling. Architecturally, ENNs combine:

• A point-estimate function,
• A trainable correction conditioned on $z$,
• An additive prior network, $f_0(x, z)$, frozen during learning.

Instead of hand-designing the functional form or initializing the prior randomly, the paper proposes pretraining $f_0$ on synthetic datasets sampled from a reference stochastic process (such as GP sample paths). This function-space regularization directly shapes the surrogate's inductive bias to align with domain-specific expectations (e.g., binding affinity distribution shapes) and improves calibration of joint uncertainty.

Algorithmic pretraining proceeds by minimizing squared loss between the prior network output and the reference paths, optionally warped to reflect realistic assay measurement distributions.

Representational Integration

Latent representations from large foundation models (e.g., COATI, Pairformer, ESM) serve as the ENN's input features, allowing the surrogate to exploit rich chemical structure and property relationships learned across billions of molecules and protein targets.

Experimental Results

Synthetic Benchmarks

Variants of the Epinet (ENN architecture) with linear, Random Fourier Features (RFF), and Pretrained prior functions are evaluated on joint negative log-loss (NLL) after training with synthetic warped GP data. The Pretrained variant consistently yields the best joint NLL across input dimensions, indicating more accurate modeling of batch-wise correlation and uncertainty. Notably, all variants perform similarly on marginal NLL, but only well-specified functional priors deliver strong joint predictions—directly impacting acquisition efficacy in real BO tasks.

EGFR Inhibitor Optimization

A semi-synthetic benchmark is constructed by augmenting 13,201 public EGFR inhibitors with decoys generated via COATI-LDM. Surrogate models using COATI embeddings as inputs are trained to maximize binding affinity (pIC50) over 46,000 compounds.

Batch BO is performed with three Epinet surrogates (Linear, RFF, Pretrained prior) and three acquisition strategies (greedy MAP, EMAX, qPO). The Pretrained Epinet combined with parallel acquisition functions retrieves the top pIC50 inhibitors in 5x fewer iterations than greedy baselines and achieves a 9x improvement over random selection for mean IC50 in the top-10 compounds at the final iteration. Performance is distinguished early in optimization, when modeling prior uncertainty is critical due to scarce data. This highlights the sample-efficiency gains enabled by accurately specifying joint predictive distributions.

Real-World Screening with tArray

The platform screens 50,000 compounds against a pipeline target using experimental fold-over (FO) measurements. Ablating the number of sampled particles for joint predictions confirms that Batch BO efficiency (measured as area-under-curve for top FO candidates) degrades rapidly when particle count is reduced. Here, drawing 5,000 joint samples for 50,000 compounds completes in 13 seconds on an A100 GPU, confirming the ENN framework's practical scalability compared to ensemble-based surrogates.

Implications and Future Directions

The demonstrated improvements have immediate practical consequences for drug screening workflows and lead optimization campaigns, enabling laboratory resources to be allocated more effectively by reducing unnecessary synthesis/testing cycles. The framework is amenable to more complex molecular representations (e.g., protein-ligand co-fold information) and alternative property prediction tasks (e.g., ADME). Pretraining prior networks using rich simulation data or physics-based simulators further expands applicability, as functional regularization has been shown to generalize better than weight-space priors for complex chemical objectives.

The findings also imply that function-space specification via pretrained priors empowers BO practitioners to inject detailed domain knowledge without the ambiguity associated with architectural tuning or weight heuristics. As chemical and biological foundation models continue to advance, integrating their representations into the ENN surrogate framework may further improve generalization and accelerate discovery cycles for scenario-specific molecular engineering.

Conclusion

This work presents a scalable surrogate modeling solution for Batch Bayesian Optimization in molecular discovery, validating substantive gains in sample efficiency and optimization performance through the function-space pretraining of ENN prior networks. The incorporation of foundation model representations and rapid joint sampling underpin its practical utility in both simulated and real-world screening pipelines. The results substantiate the claim that careful prior function specification directly enhances batch-aware uncertainty modeling and acquisition, providing a mathematically and computationally sound approach to large-scale chemical optimization tasks.