Open Materials Generation with Inference-Time Reinforcement Learning

Motivation

• Known: score-based diffusion model은 score를 명시적으로 학습하므로 policy-gradient RL(예: MatInvent, Chemeleon2)을 de novo generation(DNG) task에 적용해 target property를 정렬할 수 있음이 알려져 있다. 반면 flow-matching 및 stochastic interpolants(SI) 프레임워크는 velocity field만 학습하기 때문에 RL 적용이 제한적이었다.
• Gap: Flow-GRPO 등 기존 접근은 Gaussian base distribution과 linear interpolant path라는 제한된 설정에서만 velocity와 score를 연결할 수 있어, OMatG의 넓은 SI design space 대부분에는 RL을 적용할 수 없었다. 또한 RL을 crystal structure prediction(CSP) task에 적용한 사례는 전무했다.
• Why: RL은 likelihood 기반 학습의 근본적 한계를 극복하고 black-box reward function을 통해 생성모델을 downstream objective(안정성, 물성 등)에 직접 정렬시킬 수 있는 유연한 프레임워크이므로, 이를 velocity-field 기반 생성모델과 CSP task로 확장하는 것은 역물질설계(inverse materials design) 파이프라인의 실용성을 크게 높인다.
• Approach: 저자들은 사전학습된 velocity field에 작은 noise schedule을 추가한 surrogate SDE를 정의하여 결정론적 ODE 생성 과정을 확률적 정책(policy)으로 변환하고, GRPO 및 PPO 스타일의 clipped objective를 사용해 velocity field 자체를 강화학습으로 미세조정한다.

Achievement

Figure 4. Test-set evaluation metrics for velocity-annealing OMatG

1. Score-free policy-gradient RL 프레임워크: velocity field에 직접 작동하며 명시적 score 계산 없이 policy-gradient RL을 수행하는 OMatG-IRL을 제안하고, score-based 모델로도 자연스럽게 확장 가능함을 보였다.
2. CSP task에 대한 RL 최초 적용: 에너지 기반 목적함수를 명시적 diversity reward 없이도 효과적으로 강화할 수 있음을 보였으며, 이는 DNG와 달리 composition conditioning에서 다양성이 자연스럽게 발생하기 때문임을 규명했다.
3. Score 기반 대비 경쟁력 있는 성능: velocity-based와 score-based OMatG-IRL을 직접 비교하여 두 접근이 유사한 강화 성능을 달성함을 실험적으로 입증했다.
4. Velocity-annealing 스케줄 학습: 시간 의존적 velocity-annealing schedule을 RL로 학습하여 handcrafted annealing scheme을 대체하고, 적분 스텝 수를 한 자릿수(order-of-magnitude) 줄이면서도 정확한 CSP를 달성해 샘플링 효율과 생성 시간을 크게 개선했다.

How

Figure 1. Inference-time RL for CSP in velocity-based OMatG-IRL. (a) The deterministic base ODE with pretrained velocity

• OMatG의 equivariant crystal representation({A, X, L})과 stochastic interpolants(SI) 기반 생성 과정을 기반으로 프레임워크 구축
• 사전학습된 deterministic velocity field bθref(t, xt)에 작은 noise schedule σref(t)를 추가하여 evaluation metric을 거의 변화시키지 않는 surrogate SDE 정의 (Fig.1a)
• 이 frozen surrogate를 KL 정규화를 위한 reference policy로 사용하고, 강화 대상 velocity field bθ(t, xt)와 (다를 수 있는) noise schedule σ(t)로 stochastic exploration 수행 (Fig.1b)
• 동일 조건화 하에서 다수의 stochastic-policy rollout으로부터 terminal reward ri = r(xi_t=1)를 얻어 GRPO 방식으로 비교 (Fig.1c)
• reward를 GRPO group advantage로 변환하고 KL 정규화와 함께 PPO 스타일 clipped objective에 사용하여 policy 업데이트 (Fig.1d)
• score-based SDE와 perturbed velocity 접근을 비교 실험(Fig.2), validation metric의 학습 곡선 비교(Fig.3), velocity-annealing 실험(Fig.4, Fig.5) 수행

Originality

• Score 계산 없이 velocity field에 직접 policy-gradient RL을 적용하는 최초의 프레임워크 제안으로, Flow-GRPO의 제약(Gaussian base, linear interpolant)을 벗어나 OMatG의 전체 SI design space에 RL 적용을 가능하게 함
• 사전학습 생성모델의 baseline 성능을 보존하면서 탐색과 policy-gradient 추정을 동시에 가능하게 하는 stochastic perturbation 기법(surrogate SDE) 도입
• CSP task에 RL을 최초로 적용하여, DNG와 달리 diversity reward 없이도 composition conditioning만으로 다양성이 자연스럽게 보존됨을 규명
• RL을 이용해 handcrafted annealing scheme을 대체하는 학습형 velocity-annealing schedule을 제시, 샘플링 효율의 실질적 개선을 달성

Limitation & Further Study

• 실험이 특정 데이터셋과 특정 energy-based reward에 한정되어 있어, 보다 다양한 target property(전자적, 자기적 물성 등)에 대한 일반화 가능성 검증이 추가로 필요함
• surrogate SDE의 noise schedule σref(t) 및 σ(t) 선택이 성능에 미치는 민감도에 대한 체계적 분석이 제한적임
• DNG task로의 확장 및 diversity reward 없는 설정이 다른 데이터셋/조성 분포에서도 유지되는지에 대한 추가 검증 필요
• score-based와 velocity-based 접근의 비교가 제한된 벤치마크에 국한되어, 대규모/다양한 벤치마크에서의 일반성 확인이 후속 연구로 남음

Evaluation

총평: Flow-matching 기반 생성모델에 score 계산 없이 policy-gradient RL을 적용 가능케 한 방법론적 기여가 뚜렷하며, CSP task에 대한 RL 최초 적용과 velocity-annealing을 통한 샘플링 효율 개선이라는 실용적 성과도 인상적이다. 다만 실험 범위가 제한적이어서 다양한 물성 및 데이터셋에 대한 일반화 검증이 후속 연구로 필요하다.