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

2026-06-25

AI Summary

This work addresses the prohibitive computational cost of surrogate model inference across ultra-large discrete chemical spaces, which hinders efficient identification of high-value molecules. The authors propose BOBa, a novel framework that, for the first time, integrates a multi-armed bandit mechanism into trillion-scale virtual screening. By partitioning the chemical library into target-aware subspaces modeled as “arms” and employing an uncertainty-based upper confidence bound strategy, BOBa dynamically allocates computational resources to perform surrogate inference and evaluation only in high-potential regions. Validated on real synthesizable molecular libraries, the method significantly reduces computational overhead while maintaining strong screening performance, achieving a tunable trade-off between exploration and exploitation as well as between cost and efficacy, thereby establishing a new paradigm for efficient optimization over massive chemical libraries.

Abstract

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.

Problem

• surrogate optimization
• chemical space
• expensive evaluations
• large-scale screening
• computational bottleneck

Innovation

• bandit allocation
• surrogate optimization
• chemical space
• action space partitioning
• sample efficiency