Sample Efficient Generative Optimization for Molecular Design

기술 요약: 분자 설계를 위한 샘플 효율적 생성 최적화 (SEGO)

문제 정의
신약 개발, 재료 설계 및 촉매 작용에서의 분자 최적화는 엄격한 평가 예산 하에서 방대한 화학 공간을 탐색해야 합니다. 고충실도 오라클(예: 실험적 합성 및 테스트)은 비용이 많이 들고 시간이 오래 걸리므로, 최적화 방법의 "샘플 효율성(sample efficiency)"—즉, 강력한 후보를 찾는 데 필요한 평가 횟수—이 실질적인 유용성을 결정하는 주요 요인이 됩니다. 기존 방식들은 트레이드오프에 직면해 있습니다. 심층 생성 모델(예: RL 기반 SMILES 생성기)은 수렴을 위해 수백에서 수천 번의 오라클 호출이 필요한 경우가 많은 반면, 베이지안 최적화(BO)는 높은 샘플 효율성을 제공하지만 고정된 라이브러리에 국한되거나 분자 표현 방식의 선택에 민감합니다. 기존의 하이브리드 접근 방식들은 타겟 분포가 알려져 있지 않고 직접적인 실험 피드백이 유일한 가용한 오라클인 시나리오에서, 생성적 탐색의 유연성과 BO의 데이터 효율성 사이의 균형을 맞추는 데 어려움을 겪고 있습니다.

방법론: SEGO 프레임워크
저자들은 개방형 분자 공간에서 샘플 효율적인 탐색을 수행하기 위해 강화 학습(RL)과 베이지안 최적화(BO)를 통합한 탠덤 프레임워크인 **샘플 효율적 생성 최적화(SEGO)**를 제안합니다. SEGO는 세 가지 핵심 구성 요소를 포함하는 반복 루프를 통해 작동합니다:

1. 생성 엔진 (Saturn): RL을 통해 학습된 Mamba 아키텍처 기반의 자기회귀(autoregressive) SMILES 생성기입니다. 이 모델은 "증강 메모리(Augmented Memory)"를 활용하여, 일반화 능력을 향상시키고 모드 붕괴(mode collapse)를 방지하기 위해 무작위화된(SMILES 증강을 통한) 상위 보상 SMILES 문자열의 리플레이 버퍼를 유지합니다.
2. 베이지안 최적화 엔진 (GOLLuM): LLM으로부터 유도된 SMILES 임베딩 상에서 직접 작동하는 딥 커널 가우시안 프로세스(Deep Kernel GP) 기반의 대리 모델(surrogate model)입니다. 이를 통해 수작업으로 설계된 디스크립터 없이도 온더플라이(on-the-fly) 특징 추출이 가능합니다.
3. SEGO 루프:
   • 앵커링된 내부 루프 (Anchored Inner Loop): 각 반복 단계에서, "진성 오라클(true-oracle)"에 앵커링된 에이전트($\pi_{Agent}$)로부터 "대리 모델 가이드(surrogate-guided)" 내부 에이전트($\pi_{Inner}$)가 인스턴스화됩니다. 내부 에이전트는 GP 사후 평균(높은 보상이 존재할 것이라고 대리 모델이 가설을 세운 지점)을 대상으로 훈련되어, 유망한 영역 내의 표적 후보군을 생성합니다.
   • 선택 (Selection): 생성된 라이브러리는 딥 커널 GP를 사용하여 특징화되며, 획득 함수(Expected Improvement)가 가장 유망한 후보를 선택합니다.
   • 피드백 및 업데이트: 선택된 후보는 진성 오라클에 의해 평가됩니다. 결과는 다음을 위해 사용됩니다:
     • GP 대리 모델 업데이트 (가설 정교화).
     • 진성 보상을 통한 앵커링된 에이전트($\pi_{Agent}$) 업데이트.
     • 증강 최적화(Augmented Optimization) 수행: 선택된 분자는 $N$개의 무작위화된 SMILES 변형으로 증강되며, GP는 모든 표현형을 동시에 학습합니다. 이는 SMILES의 비단사적(non-injective) 특성을 활용하여 잠재 공간을 정렬함으로써, 특히 저데이터 환경에서 예측 정확도를 향상시킵니다.

주요 기여

• 대리 모델 가이드 생성: 고정된 라이브러리에서 선택하는 표준 BO와 달리, SEGO는 대리 모델이 높은 보상이 포함되어 있다고 가설을 세운 특정 영역을 제안하도록 생성 모델을 유도합니다.
• 앵커링 메커니즘: 이 프레임워크는 진성 오라클 피드백을 통해 학습된 에이전트를 유지하여 안정성을 확보하는 동시에, 임시 내부 에이전트를 사용하여 대리 모델의 지형을 탐색합니다. 이는 대리 모델이 초기에 불확inc한 경우 에이전트가 생산적이지 않은 영역으로 이탈하는 것을 방지합니다.
• SMILES 증강 대리 모델 업데이트: 동일한 분자의 여러 증강된 표현형을 통해 GP를 훈련함으로써, SEGO는 잠재 공간의 정렬과 더 적은 오라클 호출로도 대리 모델의 정확도를 개선합니다.
• 이성질체 붕괴 완화: 저자들은 현재의 최고 성과물(best hit)의 구성 이성질체(constitutional isomers)만을 발견하는 것이 아니라 새로운 분자식을 발견하도록 장려하기 위해, 기존의 챔피언 임계값을 높이는 수정된 Expected Improvement 획득 함수를 도입했습니다.

결과

• 절제 연구 (Ablation Studies): 앵커링 초기화 전략은 사전 확률(prior)로부터 재초기화하거나 반복 단계 전반에 걸쳐 연속적인 에이전트를 유지하는 방식보다 성능이 크게 우수했으며, 이는 유망한 화학 영역에서의 필수적인 "웜 스타트(warm start)"를 제공했습니다. 증강 최적화(5배~20배)는 통계적으로 유의미한 다양성 손실 없이 수율을 높이고 오라클 부담을 줄이는 것으로 나타났습니다.
• 구성 요소 비교: 엄격한 100회 호출 예산이 주어진 다중 속성 최적화 작업에서, 단독 생성 엔진(Saturn)과 단독 BO 엔진(GOLLuM)은 히트를 찾지 못하거나 성능이 저조했습니다. 반면 SEGO는 63개의 고유한 히트를 성공적으로 식별하여 하이브리드 접근 방식의 필요성을 입증했습니다.
• PMO 벤치마크: PMO(Practical Molecular Optimization) 벤치마크에서 SEGO는 다른 방법들이 일반적으로 사용하는 1,000회의 호출 대비 단 **10%의 오라클 예산(100회 호출)**만 사용하여 최첨단 성능(SOTA)을 달al했습니다. 또한 경쟁력 있는 Top-1 및 Top-10 점수를 기록했습니다.
• 다중 파라미터 도킹: 특정 도킹 작업에서 SEGO는 기존 최첨단 방식들이 요구하는 호출 횟수의 약 절반 정도만 사용하여 10개의 히트에 도달했습니다.

의의 및 주장
본 논문은 SEGO가 분자 최적화를 직접적인 실험 피드백에 의해 주도되는 캠페인에 더 가깝게 만든다고 주장합니다. 오라클 호출 횟수를 획기적으로 줄임으로써, 이 방법은 "해킹(hacking)"되기 쉬운 저충실도 in silico 대리물에 의존하는 대신, 고충실도 시뮬레이션이나 실제 실험을 최적화 오라클로 사용하는 것을 가능하게 합니다. 저자들은 SEGO를 고정된 라이브러리의 한계와 순수 생성 방식의 데이터 비효율성을 극적으로 극복하여, 최소한의 실험 비용으로 개방형 화학 공간에서 고품질의 후보를 발견할 수 있는 해결책으로 포지셔닝합니다.