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

Motivation

• Known: Bayesian optimization과 active learning은 expensive evaluation 하에서 discrete chemical space를 탐색하는 표준 도구로 자리잡았으며, surrogate 모델을 통해 evaluation 횟수를 줄이는 방식이 널리 사용되어 왔다. 전통적으로는 acquisition function을 전체 라이브러리 X에 대해 exhaustive하게 평가하는 방식이 표준이었다.
• Gap: make-on-demand 화학의 발전으로 라이브러리 규모가 수십억~수조 개 화합물로 커지면서, evaluation cost뿐 아니라 surrogate model의 full-library forward pass 자체(inference cost)가 새로운 병목이 되었으나 기존 BO/active learning 파이프라인은 이 inference cost를 무시할 수 있다는 가정에 기반해 있어 이러한 초대형 라이브러리 환경에 적용하기 어렵다.
• Why: 신약 발굴의 hit identification 단계는 비용과 시간 측면에서 파이프라인 전체의 핵심 병목이며, 다운스트림 단계로 갈수록 비용이 급격히 증가하므로 초기 computational screening의 질이 결과에 불균형적으로 큰 영향을 미친다. 따라서 수조 규모 라이브러리에서도 실용적으로 작동하는 virtual screening 방법론은 실질적 파급력이 크다.
• Approach: 화학 공간을 partition으로 나누고 이를 multi-armed bandit의 arm으로 취급하여, 유망한 partition에 surrogate inference와 evaluation 자원을 optimism-under-uncertainty 원칙(예: UCB1)에 따라 적응적으로 집중 배분함으로써 전체 라이브러리에 대한 exhaustive inference를 회피한다.

Achievement

Evaluation of Bandit algorithms, and observed tradeoff between performance and inference cost. (Left) Optimization

1. BOBA 프레임워크 제안: bandit 기반 partition 할당과 partition 내 surrogate-guided 최적화를 결합해 전체 라이브러리 inference를 제거하는 새로운 최적화 프레임워크를 도입했다.
2. 실제 synthesis-on-demand 라이브러리 벤치마크: optimism-under-uncertainty 기반 bandit(UCB1 등)과 의미 있는 화학공간 partitioning이 효과적인 자원 배분에 필수적임을 실험적으로 입증했다.
3. 성능-비용 tradeoff 규명: screening 성능과 surrogate inference 비용 사이에 조절 가능한 tradeoff가 존재하며, 이는 라이브러리 규모가 커질수록(최대 약 10^8 분자까지 스케일링 실험) 더 유리해짐을 보였다.
4. 이론적 분석: partition 개수가 inference 절감과 bandit regret 간 균형을 결정하는 핵심 제어 변수임을 간단한 이론적 분석으로 규명했다.

How

Influence of static vs. dynamic partitioning. Opti-

• Virtual screening을 discrete optimization 문제로 정식화하고, evaluation cost와 inference cost를 명시적으로 구분
• 화학 공간 X를 topological feature 등 분자 표현 ϕ(x)에 기반해 K개의 disjoint partition {X1,...,XK}으로 분할 (static/dynamic partitioning 비교)
• 각 partition을 MAB의 arm으로 취급, ε-greedy, softmax sampling, UCB1 등 다양한 bandit policy를 비교
• 각 라운드에서 선택된 arm(partition) 내부에서만 surrogate 모델로 acquisition function을 계산하여 evaluation 후보 선정
• 실제 synthesis-on-demand 라이브러리 데이터를 사용한 systematic benchmark 및 최대 10^8 분자 규모의 scaling 실험 수행
• partition 개수에 따른 inference 절감과 regret 간 tradeoff에 대한 이론적 분석 수행

Originality

• 기존 BO/active learning이 가정하던 "surrogate inference는 무시 가능하다"는 전제를 명시적으로 깨고, inference cost를 evaluation cost와 별개의 제약으로 정식화한 문제 설정 자체가 독창적이다.
• Chemical space partitioning을 MAB의 arm으로 매핑하여 global-local search를 분리하는 아이디어(BOBA)가 새롭다.
• 화학적으로 의미 있는 partitioning(구조 기반 feature space)과 optimism 기반 bandit을 결합해야만 효과적임을 실증적으로 규명한 점이 기존 random exploration 기반 접근과 차별화된다.

Limitation & Further Study

• 발췌된 본문에서 partition 개수 K의 최적 선택 방법이나 partition 경계의 chemical validity에 대한 심층적 논의가 제한적으로 보인다.
• Scaling 실험이 약 10^8 분자 규모까지만 수행되어 실제 목표인 수십억~수조 규모 라이브러리에서의 실증적 검증은 아직 부재하며, 이는 추정(extrapolation)에 의존한다.
• 사용된 surrogate 모델의 종류나 docking score의 노이즈 특성이 partition 간 이질성에 미치는 영향에 대한 강건성 분석이 추가로 필요해 보인다.
• 후속 연구로 실제 수조 규모 라이브러리에서의 end-to-end 검증, 동적 partition 재구성 전략의 정교화, 다양한 target에 대한 일반화 검증이 요구된다.

Evaluation

총평: Surrogate inference 자체를 병목으로 인식하고 이를 bandit 기반 partition 할당으로 해결하는 문제 설정과 접근이 신선하며, 초대형 chemical library 스크리닝에 실질적 파급력을 지닌 견실한 연구이다.