Target-Aware Bandit Allocation for Scalable Surrogate Optimization in Chemical Space

Identifying high-utility candidates from massive discrete spaces under expensive evaluations is a recurring challenge across the sciences, with structure-based drug discovery as a prominent example. While surrogate-based optimization can increase sample efficiency by reducing the number of expensive evaluations, modern molecular libraries have reached billions to trillions of compounds, making full-library surrogate inference itself a major computational bottleneck. We introduce BOBA, a bandit-guided surrogate optimization framework that eliminates full-library inference by adaptively allocating computation across partitions of the action space. By treating partitions as arms in a multi-armed bandit, BOBA concentrates inference and evaluations on empirically promising partitions while maintaining principled exploration. Experiments on real-world synthesis-on-demand libraries demonstrate that optimism-under-uncertainty bandits, combined with meaningful action space partitioning, are essential for effective allocation of inference and evaluations. Our findings reveal a tunable tradeoff between screening performance and surrogate inference cost, which supports practical optimization over current libraries, and establishes a viable route to ultra-large library virtual screening.

Drug discovery usually starts by searching through massive collections of possible molecules to find a few that may bind to a disease-related protein. Machine learning can help this process by learning from a small number of expensive experiments, and then suggesting better molecules to test next. However, modern molecule collections can contain billions or even trillions of candidates, so even asking a machine-learning model to score every molecule can become too slow. Here, we introduce BOBA, a method that avoids scoring the whole collection of molecules. BOBA first divides the molecule collection into smaller groups. It then uses a decision-making strategy to decide which groups look most promising, while still checking uncertain groups when needed. Within the selected group, a machine-learning model scores all molecules, and chooses the ones for the next expensive experiment. In experiments on large molecule libraries, BOBA found many of the best-scoring molecules while using far fewer model predictions than standard techniques. The results show that both good grouping of molecules and uncertainty-aware decision making across groups are important. This makes BOBA a practical step toward machine-learning-guided search in extremely large chemical libraries, with possible uses in drug discovery and related scientific discovery scenarios.