RetrOrchestrator: A Multi-Step Retrosynthesis Agent Dynamically Orchestrating Single-Step Transition Models

Motivation

• Known: 기존 multi-step retrosynthesis planning은 MDP로 정식화되며, 고정된 단일 SSR model을 transition model로 사용하고 MCTS나 A* 같은 탐색 알고리즘 및 학습된 value function을 통해 "어떻게 탐색할지(how to search)"를 개선하는 데 집중해왔다.
• Gap: 서로 다른 SSR model들은 분자 상태에 따라 성능 편차(skill disparity)가 크게 나타남에도 불구하고, transition space 자체를 어떻게 탐색할지, 즉 어떤 SSR을 언제 사용할지에 대한 연구는 거의 이루어지지 않았다. 특히 분자-model 적합성에 대한 세밀한 supervision이 부재하고 보상이 경로 완료 시에만 sparse하게 주어져 credit assignment가 어렵다는 문제가 존재한다.
• Why: 단일 SSR model에 의존하는 정적 탐색은 특정 화학 공간에서 최적이 아닌 model을 강제로 사용하게 되어 고품질 retrosynthetic route 발견을 저해하므로, 여러 model의 강점을 상황에 맞게 조합하는 능력은 실제 화학 합성 계획의 실용성과 성공률을 크게 좌우한다.
• Approach: Multi-step retrosynthesis planning을 POMDP로 재정의하고 LLM의 진화하는 context를 belief state의 proxy로 활용하여, 각 SSR model을 tool로 취급해 agent가 분자 선택과 model 선택을 동시에 학습하도록 하는 scaffold-aware 강화학습 알고리즘(SA-GRPO)을 제안한다.

Achievement

Figure 1. Relative performance of retrosynthesis agents with dif-

1. 최초의 LLM 기반 retrosynthesis planning agent: multi-step retrosynthesis를 POMDP로 정식화하여 LLM의 context를 implicit belief state로 활용, 여러 SSR model을 동적으로 orchestrate하는 최초의 agent를 제안했다.
2. SA-GRPO 알고리즘: scaffold-level clustering을 이용해 sparse하고 delayed된 reward 상황에서도 각 SSR model에 대한 credit assignment를 가능하게 하고, 쉬운 분자에서의 chemical rule 학습부터 어려운 분자에서의 전략적 planning까지 점진적으로 전이하는 curriculum을 설계했다.
3. Pareto-optimal 성능 달성: Retro*-190 벤치마크에서 94.21%의 state-of-the-art 성공률(off-the-shelf LLM 9.47%, non-LLM state-dependent router 82.63% 대비)을 달성했으며, 해결된 route의 92.49%가 2개 이상의 SSR을 호출해 정책이 단일 specialist나 static router로 붕괴하지 않았음을 입증했다. PDB-600이라는 더 큰 out-of-distribution set에서도 동일한 성능 향상이 유지되며, wall-clock time과 model-query count 양 측면에서 Pareto-optimal을 달성했다.

How

Figure 3. Overview of the RetrOrchestrator training algorithm and pipeline. The bottom panel illustrates the scaffold-aw

• Multi-step retrosynthesis planning을 (S, A, T, R) 대신 belief state를 포함하는 POMDP로 재정의하고, LLM의 context(scenario, tools, rules, history, observation)를 belief state 업데이트의 proxy로 사용
• 각 SSR model(LocalRetro, RetroKNN, Chemformer, RootAligned, MEGAN, Graph2Edits, EditRetro 등)을 agent가 호출 가능한 tool로 등록하여, 매 스텝마다 확장할 분자와 사용할 SSR model을 함께 결정
• Scaffold-aware Group-Relative Policy Optimization(SA-GRPO)을 제안하여 scaffold 유사 분자 상태들에 대한 step-wise advantage를 기반으로 서로 다른 SSR model의 skill disparity를 추론하도록 학습
• 쉬운 분자에서의 chemical rule adherence부터 어려운 분자에서의 전략적 planning으로 점진적으로 전환하는 curriculum 학습 전략 도입
• Retro*-190 벤치마크와 더 큰 out-of-distribution 벤치마크 PDB-600에서 성공률, wall-clock time, model-query count를 기준으로 정적 SSR baseline, ensemble, off-the-shelf LLM, non-LLM state-dependent router와 비교 평가

Originality

• Multi-step retrosynthesis를 MDP가 아닌 POMDP로 재정의하여 SSR model 선택 자체를 탐색 대상에 포함시킨 최초의 시도
• LLM의 context를 belief state의 proxy로 삼아 tool-calling 방식으로 여러 SSR model을 조합하는 agent 설계
• Scaffold 유사성 기반의 group-relative advantage 추정을 통해 sparse reward 환경에서 model별 기여도를 credit assignment하는 SA-GRPO 제안
• 화학적 규칙 준수에서 전략적 planning으로 전환하는 curriculum 학습 설계를 강화학습 훈련에 결합

Limitation & Further Study

• 논문 발췌 내용상 SA-GRPO의 구체적 수식 및 scaffold clustering 방법의 세부 사항, LLM backbone의 규모와 종류에 대한 상세 정보가 충분히 제시되지 않아 재현성 검증에 한계가 있음
• 다양한 SSR model을 tool로 호출하는 과정에서 발생하는 계산 비용(다중 model 추론 오버헤드)이 실제 대규모 배치 상황에서 어떻게 확장되는지에 대한 분석이 제한적일 수 있음
• Retro*-190과 PDB-600이라는 특정 벤치마크에 대한 평가로, 더 다양한 화학 공간이나 실험실 검증(wet-lab validation)으로의 일반화 가능성에 대한 추가 연구가 필요함
• LLM 기반 agent 특성상 학습된 policy의 해석 가능성 및 실패 사례에 대한 심층 분석이 추가되면 신뢰성이 높아질 것으로 보임

Evaluation

총평: SSR model 간 skill disparity라는 실질적으로 중요하지만 간과되어온 문제를 POMDP 정식화와 scaffold-aware 강화학습으로 해결한 독창적이고 실용적인 연구로, retrosynthesis planning 분야에 새로운 방향을 제시한다.