그룹 상대 정책 최적화 기반 분자 최적화 모델 GRXForm

초록

본 논문은 사전 학습된 Graph Transformer를 기반으로, 시작 구조마다 보상을 정규화하는 Group Relative Policy Optimization(GRPO)을 적용해 단일 전방 패스로 최적 분자를 생성하는 amortized molecular optimization 프레임워크 GRXForm을 제안한다. GRPO는 동일 시작 구조에서 샘플링된 G개의 궤적 평균을 동적 베이스라인으로 사용해 스케폴드 난이도 차이에 따른 보상 편차를 감소시키고, 학습 안정성을 높인다. 실험 결과, GRXForm은 인스턴스 최적화 기법과 동등하거나 우수한 다목적 최적화 성능을 보이며, 추론 시 오라클 호출 없이 새로운 스케폴드에 일반화한다.

상세 분석

GRXForm은 기존 GraphXForm의 구조를 그대로 유지하면서, 목표 지향적인 미세 조정을 위해 강화학습(RL) 방식을 도입한다. 핵심 혁신은 Group Relative Policy Optimization(GRPO)이다. 전통적인 REINFORCE나 Actor‑Critic 방식은 전역 베이스라인을 사용해 보상 평균을 빼지만, 분자 설계에서는 시작 스케폴드마다 최적화 난이도가 크게 달라 전역 베이스라인이 편향된 그라디언트를 초래한다. GRPO는 각 시작 구조 S_i에 대해 Stochastic Beam Search를 이용해 G개의 완성 궤적 O_{i,1…G}을 생성하고, 이들의 평균 μ_i를 동적 베이스라인으로 사용한다. 보상 r_{i,j}에 대해 상대 이득 A_{i,j}=r_{i,j}−μ_i를 계산함으로써 “어려운” 스케폴드에서도 상대적으로 높은 점수를 받은 궤적이 양의 신호를 얻고, “쉬운” 스케폴드에서는 과도한 보상 상승을 억제한다. 이 방식은 보상 분산을 크게 감소시켜 정책 그라디언트 추정의 편차를 낮추고, 학습 초기에 발생할 수 있는 모드 붕괴를 방지한다.

정책 네트워크는 Decoder‑only Graph Transformer이며, 원자·결합 추가를 3단계 계층형 액션 스페이스(작업 선택, 타깃 선택, 결합 차수)로 분해한다. 각 단계마다 화학적 원자 가치를 보장하는 유효성 마스크를 적용해 valence 위반을 사전에 차단한다. 사전 학습 단계에서는 ChEMBL 35 데이터를 사용해 teacher‑forcing 방식으로 화학적 plausibility를 학습하고, 이는 추후 RL 미세 조정 시 탐색 공간을 현실적인 화학 구조에 한정한다.

학습 파이프라인은 (1) 사전 학습된 GraphXForm 로드, (2) 목표 속성(예: QED, LogP, 다목적 스코어) 정의, (3) 각 배치의 시작 구조에 대해 SBS로 G개의 후보 완성 생성, (4) GRPO 기반 상대 보상 계산 후 정책 업데이트 순으로 진행된다. 중요한 점은 추론 단계에서 전혀 오라클 호출이 필요 없다는 것이다. 학습 시에만 오라클(또는 사전 훈련된 속성 예측기) 호출이 발생하므로, 대규모 라이브러리 생성이나 사용자‑인‑루프 설계에 적합한 amortized 접근법이 된다.

실험에서는 PMO(Practical Molecular Optimization) 벤치마크의 10,000 오라클 호출 제한 하에서 GRXForm이 기존 인스턴스 최적화 기법(GA, GenMol 등)과 비교해 비슷하거나 더 높은 평균 스코어를 달성했으며, 특히 새로운 스케폴드(훈련에 포함되지 않은)에서도 성능 저하가 미미했다. 다목적 최적화 실험에서는 파레토 프론트가 넓게 형성되어, 단일 목표 최적화보다 복합적인 약물 설계 요구를 충족시키는 능력을 보여준다.

한계점으로는 (1) GRPO가 그룹 크기 G와 SBS 폭에 민감해 하이퍼파라미터 튜닝이 필요하고, (2) 고신뢰도 물리‑화학 시뮬레이션(예: 도킹, 자유 에너지 계산)과 결합할 경우 학습 비용이 크게 증가한다는 점을 언급한다. 향후 연구에서는 메타‑학습이나 온라인 적응형 베이스라인을 도입해 G와 SBS를 자동 조정하거나, 멀티‑태스크 학습을 통해 다양한 오라클을 동시에 학습하는 방안을 제시한다.

요약하면, GRXForm은 그룹 기반 상대 보상 정규화를 통해 시작 구조별 난이도 차이를 효과적으로 보정하고, 단일 전방 패스로 고품질 분자를 생성하는 실용적인 amortized molecular optimization 솔루션을 제공한다.