LLM 추론을 활용한 분자 최적화: 참조 기반 정책 최적화

Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

대규모 언어 모델(LLM)은 추론 작업에서 지도 학습 미세 조정(SFT) 및 검증 가능한 보상을 이용한 강화 학습(RLVR)을 통해 상당한 이점을 얻습니다. 그러나 각 데이터 포인트가 일반적으로 단일 최적화된 참조 분자만을 제공하고 단계별 최적화 경로를 포함하지 않는 지시 기반 분자 최적화 작업에서는 이러한 방법들이 효과적이지 않습니다. 우리는 참조 분자에 대한 '답변만' 포함하는 SFT 학습이 추론 능력을 저하시키며, RLVR이 모델의 효과적인 탐색 부족으로 인해 유사성 제약 조건 하에서 희소한 피드백을 제공하여 학습 속도를 늦추고 최적화를 제한한다는 것을 밝혀냈습니다. 우리는 새로운 분자를 탐색하는 것을 장려하면서 동시에 참조 분자를 활용하기 위해 참조 기반 정책 최적화(RePO)를 제안합니다. RePO는 경로 데이터 없이 참조 분자로부터 학습하는 최적화 방법입니다. 각 업데이트 단계에서 RePO는 모델로부터 후보 분자와 해당 중간 추론 경로를 샘플링하고, RL 방식으로 유사성 제약 조건 하에서 속성 만족도를 측정하는 검증 가능한 보상을 사용하여 모델을 학습합니다. 동시에, RePO는 정책의 중간 추론 경로를 컨텍스트로 유지하고, 지도 학습 방식으로 답변 부분만 학습하여 참조 정보를 활용합니다. RL 항은 탐색을 촉진하고, 가이드 항은 유효한 분자 편집이 많을 때 출력 결과를 참조 분자에 연결하여 보상 희소성을 완화하고 학습을 안정화합니다. 다양한 분자 최적화 벤치마크에서 RePO는 SFT 및 RLVR 기준 모델(예: GRPO)보다 일관되게 우수한 성능을 보입니다. RePO는 최적화 지표(성공률 × 유사도)를 개선하고, 상충되는 목표 간의 균형을 향상시키며, 새로운 지시 스타일에도 더 잘 일반화됩니다. 저희의 코드는 다음 주소에서 공개적으로 이용할 수 있습니다: https://github.com/tmlr-group/RePO.

Large language models (LLMs) benefit substantially from supervised fine-tuning (SFT) and reinforcement learning with verifiable rewards (RLVR) in reasoning tasks. However, these recipes perform poorly in instruction-based molecular optimization, where each data point typically provides only a single optimized reference molecule and no step-by-step optimization trajectory. We reveal that answer-only SFT on the reference molecules collapses reasoning, and RLVR provides sparse feedback under similarity constraints due to the model's lack of effective exploration, which slows learning and limits optimization. To encourage the exploration of new molecules while balancing the exploitation of the reference molecules, we introduce Reference-guided Policy Optimization (RePO), an optimization approach that learns from reference molecules without requiring trajectory data. At each update, RePO samples candidate molecules with their intermediate reasoning trajectories from the model and trains the model using verifiable rewards that measure property satisfaction under similarity constraints in an RL manner. Meanwhile, it applies reference guidance by keeping the policy's intermediate reasoning trajectory as context and training only the answer in a supervised manner. Together, the RL term promotes exploration, while the guidance term mitigates reward sparsity and stabilizes training by grounding outputs to references when many valid molecular edits exist. Across molecular optimization benchmarks, RePO consistently outperforms SFT and RLVR baselines (e.g., GRPO), achieving improvements on the optimization metric (Success Rate $\times$ Similarity), improving balance across competing objectives, and generalizing better to unseen instruction styles. Our code is publicly available at https://github.com/tmlr-group/RePO.