Elign: Equivariant Diffusion Model Alignment from Foundational Machine Learned Force Fields

Motivation

• Known: score-based diffusion model과 E(3)-equivariant architecture는 분자 conformation 생성에서 우수한 성능을 보였으며, DFT나 semi-empirical potential 같은 물리 기반 guidance를 sampling 중에 적용해 저에너지 구조를 유도하는 다양한 방법이 존재한다.
• Gap: 기존 diffusion model은 semi-empirical relaxation이나 heuristic enumeration으로 만들어진 학습 데이터의 편향을 그대로 재현하여 진짜 열역학적 평형 분포(Boltzmann law)를 따르지 않으며, 물리 기반 guidance는 DFT 등 quantum-chemical 평가 비용이 크고 매 sampling step마다 반복 쿼리가 필요해 run-time 오버헤드가 크다는 두 가지 계산적 병목이 존재한다.
• Why: quantum-chemical 정확도와 계산 효율성을 동시에 만족하는 분자 conformation 생성기는 신약 개발, 재료과학, 계산화학에서 실질적 활용도가 높으며, inference 속도를 희생하지 않으면서 물리적으로 타당한 구조를 생성하는 것은 대규모 스크리닝에 필수적이다.
• Approach: DFT 평가를 pretrained foundational MLFF로 대체하여 첫 번째 amortization을 달성하고, reverse diffusion을 강화학습 문제로 재정식화한 뒤 Group Relative Policy Optimization(GRPO)으로 denoising policy를 post-training 단계에서 fine-tuning함으로써 run-time 쿼리를 제거하는 두 번째 amortization을 수행한다.

Achievement

Figure 5. DFT oracle evaluation of generated molecules from

1. 이중 amortization 프레임워크 제안: 비용이 큰 DFT 계산을 pretrained foundational MLFF로 대체(1차 amortization)하고, 물리적 guidance를 post-training 단계로 이전하여 inference 시 energy evaluation을 완전히 제거(2차 amortization)함.
2. GRPO 기반 disentangled reward 최적화: potential 기반 energy reward와 force 기반 stability reward를 결합하되, MLFF 학습 방식과 유사하게 disentangled 방식으로 최적화하는 GRPO 목적함수를 설계함.
3. DFT gold-standard 검증: QM9과 GEOM-Drugs에서 생성된 conformation을 실제 DFT oracle로 평가한 결과, 기존 run-time guidance 방법 대비 더 낮은 energy와 force, 향상된 안정성을 보임.
4. inference 속도 유지: guidance가 없는 baseline sampler와 동일한 inference 속도를 유지하면서도 물리적 안정성을 개선함을 실험적으로 입증함.

How

Figure 3. Overview of Elign. (a) Rollout branching with shared prefix: starting from a CoM-free Gaussian prior at t = T,

• reverse diffusion 과정을 Markov decision process로 재정식화하여 denoising policy를 강화학습 policy로 취급함.
• Group Relative Policy Optimization(GRPO)을 사용해 rollout branching과 shared prefix 구조를 통해 여러 trajectory를 샘플링하고 group 내 relative reward로 policy를 업데이트함(Figure 3 참조).
• reward는 foundational MLFF가 추정하는 potential energy 기반 reward와 force norm 기반 stability reward 두 가지로 구성되며, 두 reward를 disentangled 방식으로 별도 최적화함.
• E(3)-equivariant score network의 대칭성을 보존하기 위해 center-of-mass subspace로 diffusion을 제한하고, equivariant parameterization을 유지한 채 fine-tuning을 수행함.
• 학습된 policy는 post-training 이후 추가적인 energy/force 평가 없이 unguided sampler와 동일한 속도로 inference를 수행함.

Originality

• MLFF와 diffusion model 간의 denoising objective 유사성을 활용해, MLFF를 물리적 guidance의 amortized oracle로 사용하는 새로운 연결고리를 제시함.
• 기존의 run-time guidance(터미널 단계 reward, post-processing, SPSA 기반 zeroth-order gradient 등)와 달리, 물리적 제약을 post-training 단계에서 policy 자체에 compile하는 방식으로 완전히 새로운 amortization 레벨을 도입함.
• adjoint 기반 stochastic optimal control 방법(Adjoint Matching, Adjoint Schrödinger Bridge Sampler)이 요구하는 reward 미분가능성 가정을 우회하기 위해 GRPO라는 gradient-free 정책 최적화 기법을 diffusion fine-tuning에 도입함.
• energy reward와 force reward를 disentangled 방식으로 결합하는 설계는 기존 단일 reward 기반 guidance와 차별화됨.

Limitation & Further Study

• foundational MLFF 자체의 예측 오차나 out-of-distribution 분자에 대한 신뢰도 문제가 최종 생성 품질에 그대로 전파될 위험이 있음.
• GRPO 기반 fine-tuning이 요구하는 rollout 수 및 학습 비용이 명시적으로 분석되지 않아, post-training 단계의 총 계산 비용 대비 이득이 충분히 정량화되지 않음.
• QM9, GEOM-Drugs 같은 비교적 작은 유기 분자 데이터셋에 대해서만 검증되었으며, 더 큰 분자나 다양한 화학 공간(단백질-리간드 복합체 등)으로의 일반화 가능성은 추가 검증이 필요함.
• disentangled reward 최적화의 이론적 정당성(왜 별도 최적화가 결합 최적화보다 우수한지)에 대한 심층적 분석이 부족함.

Evaluation

총평: MLFF와 diffusion model의 이중 amortization이라는 아이디어가 명확하고 실용적이며, DFT gold-standard 평가를 통한 검증이 설득력 있으나, 계산 비용 분석과 더 큰 화학 공간으로의 확장성 검증이 보완되면 더욱 완성도 높은 연구가 될 것이다.