기본 머신러닝 포스 필드 기반 등변 확산 모델 정렬

초록

Elign은 사전 학습된 대규모 머신러닝 포스 필드(MLFF)를 활용해 E(3)‑등변 확산 모델을 사후 미세조정함으로써, DFT 수준의 에너지와 힘을 만족하는 3D 분자 구조를 빠르게 생성한다. 물리적 보상은 강화학습 형태로 정의하고, 에너지와 힘을 별도로 정규화하는 FED‑GRPO 알고리즘으로 정책을 최적화한다. 결과적으로 추론 단계에서는 추가 계산 없이 원본 확산 모델과 동일한 속도로 고품질 구조를 얻는다.

상세 분석

본 논문은 두 가지 차원의 비용 절감을 목표로 하는 새로운 포스트‑트레이닝 프레임워크 Elign을 제안한다. 첫 번째는 고비용 양자화학(DFT) 계산을 대체하기 위해 사전 학습된 ‘기초’ 머신러닝 포스 필드(MLFF)를 이용하는 것이다. 최신 기초 MLFF는 방대한 QM 데이터셋을 기반으로 학습돼 다양한 화학 공간에 대해 근사적인 포텐셜 에너지와 원자 힘을 제공한다. 따라서 Elign은 DFT 대신 이 경량 모델을 물리적 신호원으로 삼아 에너지·힘 보상을 생성한다.

두 번째 차원은 런타임에서 반복적으로 물리적 평가를 수행하는 기존 가이드 기반 확산 모델의 비효율성을 없애는 것이다. 이를 위해 역확산 과정을 마코프 결정 과정(MDP)으로 형식화하고, 정책(denoising network)을 강화학습으로 미세조정한다. 핵심은 ‘Force‑Energy Disentangled Group Relative Policy Optimization(FED‑GRPO)’이다. FED‑GRPO는 (1) 에너지 보상 = ‑Eϕ(z)와 (2) 힘 보상 = ‑‖Fϕ(z)‖²를 각각 z‑score 정규화한 뒤, 시간 단계별로 독립적으로 결합한다. 이렇게 하면 에너지와 힘의 스케일 차이·상관관계에 의해 학습이 왜곡되는 것을 방지하고, 각 보상이 그룹 정규화(즉, 동일한 시간 단계 내에서 평균·표준편차 기준 정규화)되어 E(3)‑등변성을 유지한다.

또한 논문은 정책 업데이트 시 KL‑제한(TRUST‑REGION) 형태의 정규화를 도입해 사전 학습된 베이스 모델에서 크게 벗어나지 않도록 한다. 이는 LLM의 사후 정렬(posterior alignment)과 유사한 접근법으로, 기존 확산 모델이 학습한 데이터 분포를 크게 손상시키지 않으면서 물리적 선호도를 반영한다.

실험에서는 QM9와 GEOM‑Drugs 데이터셋을 사용해, Elign이 사후 가이드 없이도 DFT 수준의 에너지와 낮은 RMS 힘을 달성함을 보였다. 특히, 동일한 샘플링 속도(추론 시 추가 연산 없음)에도 불구하고, 기존 런타임 가이드 방식보다 에너지·힘 오류가 평균 15‑20 % 정도 감소했다. 이는 물리적 안정성을 확보하면서도 대규모 분자 생성 파이프라인에 바로 적용할 수 있음을 의미한다.

요약하면, Elign은 (1) 기초 MLFF를 통한 물리적 신호의 비용 절감, (2) 강화학습 기반 정책 정렬을 통한 런타임 비용 제거, (3) 에너지·힘을 독립적으로 정규화하는 FED‑GRPO라는 새로운 최적화 기법을 도입함으로써, 고품질 3D 분자 생성 모델의 실용성을 크게 향상시킨다.