Open Materials Generation with Inference-Time Reinforcement Learning

이 논문은 명시적인 스코어 계산을 요구하지 않고도 효율적이며 목표 속성에 정렬된 결정 구조 예측을 가능하게 하기 위해, 연속 시간 생성 모델의 속도장(velocity fields)에서 직접 작동하는 새로운 정책 경사 강화 학습 프레임워크인 OMatG-IRL을 소개한다.

핵심

당신이 새로운 유형의 초고층 빌딩을 위한 완벽한 건축 블록을 설계하려는 숙련된 건축가라고 상상해 보십시오. 재료 과학의 세계에서 이 "블록"은 바로 결정(crystal)입니다. 오랫동안 컴퓨터는 기존의 수백만 가지 사례를 연구함으로써 이 블록들이 어떻게 생겼는지를 배우는 데 능숙했습니다. 컴퓨터는 실제와 매우 유사하면서도 안정적인 새로운 결정 구조를 생성할 수 있습니다.

하지만 문제가 하나 있습니다. 컴퓨터는 형태를 복제하는 데는 뛰어나지만, "이 결정을 매우 단단하게 만들어줘"라거나 "전기를 더 잘 통하게 해줘"와 같은 구체적인 지시를 따르는 데는 서툽니다. 이것은 마치 로봇에게 완벽한 집을 그릴 줄은 알지만, "불에 타지 않는 집을 그려줘"라고 요청하면 그저 똑같은 집을 다시 그리는 것과 같습니다. 왜냐하면 로봇은 특정 목표를 우선순위에 두는 법을 모르기 때문입니다.

이 논문은 이를 해결하기 위해 OMatG-IRL이라는 새로운 방법을 소개합니다. 이 방법이 어떻게 작동하는지 쉬운 개념으로 나누어 설명하겠습니다.

1. 문제점: "스코어(Score)" vs "속도(Velocity)"

기와 같은 형태를 생성하는 대부분의 고급 AI 모델은 두 가지 방식 중 하나로 작동합니다.

• "스코어(Score)" 방식: AI는 더 나은 모양을 얻기 위해 정확히 어느 방향으로 움직여야 하는지를 알려주는 "스코어"(언덕 위의 경사도와 같은 것)를 학습합니다. 이는 마치 GPS가 "목적지에 가까워지려면 왼쪽으로 도세요"라고 알려주는 것과 같습니다.
• "속도(Velocity)" 방식: AI는 무작위적인 노이즈 덩어리로부터 결정 모양으로 이동하기 위한 "속도"(방향과 속도)를 학습합니다. 이는 산에서 바다로 흐르는 강물과 같습니다. AI는 흐름의 방향은 알지만, 반드시 "스코어"나 언덕의 정확한 수학적 경사도를 알고 있는 것은 아닙니다.

문제는 AI에게 특정 목표를 따르도록 가르치는 가장 강력한 도구들(이를 강화 학습이라 부릅니다)이 대개 "스코어" 방식을 필요로 한다는 점입니다. 만약 당신에게 "속도" 방식만 있다면, 에너지 효율성과 같은 특정 속성을 최적화하도록 AI를 가르치기가 쉽지 않습니다.

2. 해결책: 강의 흐름을 다르게 만드는 법을 가르치기

저자들은 영리한 우회 방법을 만들어냈습니다. 그들은 비록 "속도"(강물의 흐름)만 가지고 있더라도, 흐름에 아주 작은 무작위성(노이즈)을 추가함으로써 AI가 새로운 목표를 따르도록 가르칠 수 있다는 사실을 깨달았습니다.

다음과 같이 생각해 보십시오:

• AI가 가장 낮은 지점(가장 안정적인 결정)을 찾기 위해 언덕 아래로 구슬을 굴리고 있다고 가정해 봅시다.
• 보통 구슬은 AI가 설계한 경로를 따라 똑바로 굴러 내려갑니다.
• OMatG-IRL은 구슬을 경로에서 약간 벗어나게 밀어내는 미세하고 통제된 "산들바람"을 추가합니다.
• 이 바람 때문에 구슬은 가끔 원래 경로와 약간 다른 곳으로 굴러갑니다. 이때 컴퓨터는 확인합니다: "이 새로운 지점이 에너지가 더 낮았나? 더 좋은 결정이었나?"
• 만약 답이 "예"라면, AI는 학습합니다: "좋아, 다음번에는 구슬을 그 방향으로 조금 더 밀어보자."

이를 통해 AI는 복잡한 "스코어" 지도 없이도 실험을 통해 자신의 실수와 성공으로부터 배울 수 있습니다. 즉, 흐름 자체를 가지고 실험하며 배우는 것입니다.

3. "시간 여행"의 기술 (속도 어닐링, Velocity Annealing)

논문은 또한 AI가 결정을 생성하는 속도에 대해 놀라운 사실을 발견했습니다. 일반적으로 완벽한 결정을 얻기 위해 AI는 수백 번의 작고 느린 단계(가파른 계단을 조심스럽게 내려가는 것과 같은 과정)를 거쳐야 합니다. 이는 시간이 오래 걸립니다.

저자들은 이 새로운 학습 방법을 사용하여 AI에게 **새로운 속도 일정(schedule)**을 가르쳤습니다. 단순히 전체 시간 동안 천천히 걷는 대신, AI는 다음과 같이 학습했습니다:

1. 특정 속도로 시작한다.
2. 적절한 순간에 속도를 높이거나 늦춘다.
3. 훨씬 더 짧은 시간 안에 일을 끝낸다.

이는 보통 10마일을 조깅하던 러너에게 갑자기 마지막 1마일을 완벽하게 전력 질주하도록 가르치거나, 특정 페이스로 달릴 때만 통하는 지름길을 이용하도록 가르치는 것과 같습니다. 그 결과? AI는 동일한 정확도를 유지하면서도 훨씬 더 높은 품질의 결정을 10배 더 빠르게(또는 그 이상) 생성할 수 있게 되었습니다.

4. 결정(Crystals)에 있어 이것이 중요한 이유

**결정 구조 예측(CSP)**이라는 구체적인 작업—즉, 탄소와 산소 같은 재료 목록을 AI에게 주고 최적의 결정을 만들라고 요청하는 작업—에서 저자들은 다음과 같은 성과를 보여주었습니다:

• AI가 더 낮은 에너지(이는 자연계에 존재할 가능성이 높은, 더 안정적인 상태를 의미함)를 가진 결정을 만들도록 가르칠 수 있었습니다.
• 이 과정에서 다른 방식들이 요구하는 복잡한 "스코어"를 계산할 필요가 없었습니다.
• 결정의 다양성을 높게 유지하면서도(AI가 정답 하나만을 암기하지 않도록 함) 이 작업을 수행했습니다.
• 과정을 훨씬 빠르게 만들어, 결정을 생성하는 데 필요한 단계를 수백 단계에서 수십 단계로 단축했습니다.

요약

이 논문은 더 나은 재료를 설계하도록 AI를 훈련시키는 새로운 방법을 제시합니다. 이는 마치 자연스럽게 특정 방향으로 흐르는 강물을 대상으로, 전체 지도가 없어도 더 나은 목적지를 찾기 위해 가끔씩 경로를 변경하도록 가르치는 것과 같습니다. 이를 통해 과학자들은 이전보다 더 빠르고, 더 구체적인 특성을 가진 새로운 재료를 설계할 수 있게 되었습니다.

기술 요약

기술 요약: 추론 시점 강화 학습을 이용한 개방형 재료 생성 (OMatG-IRL)

1. 문제 정의

연속 시간 생성 모델은 안정적인 결정 구조를 예측할 수 있는 역재료 설계(inverse materials design)의 강력한 도구로 부상했습니다. 그러나 명시적인 목표 속성(예: 특정 기계적, 전자적 또는 에너지적 목적 함수)을 생성 과정에 통합하는 것은 여전히 어려운 과제로 남아 있습니다. 정책 경사(Policy-Gradient) 강화 학습(RL)은 다운스트림 목적 함수와 생성 모델을 정렬하기 위한 원칙적인 메커니즘을 제공하지만, 이를 플로우 기반 모델(flow-based models)에 적용하는 데에는 기술적 제약이 존재해 왔습니다.

표준 정책 경사 RL 방법은 일반적으로 정책 비율(policy ratios)을 계산하고 업데이트를 수행하기 위해 스코어(로그 확률 밀도의 기울기)에 접근할 수 있어야 합니다. 그러나 Stochastic Interpolants(SI)나 Flow Matching을 활용하는 많은 현대적 플ow 기반 모델들은 **속도장(velocity fields)**만을 학습하며 스코어를 명시적으로 계산하거나 저장하지 않습니다. 결과적으로, 이러한 모델들은 표준 RL 프레임워크를 적용하기 어려웠으며, 이는 모델이 훈련 분포의 내재된 안정성을 넘어 특정 비암시적 목적 함수를 최적화하는 능력을 제한해 왔습니다.

2. 방법론: OMatG-IRL

저자들은 스코어의 명시적 계산 없이 연속 시간 생성 모델의 학습된 속도장에 직접 작동하도록 설계된 정책 경사 RL 프레임워크인 **Open Materials Generation with Inference-Time Reinforcement Learning (OMatG-IRL)**을 소개합니다.

핵심 메커니즘

OMatG-IRL은 표준 결정 구조 예측(CSP) 평가 지표가 기초가 되는 상미분 방정식(ODE) 역학에 도입된 작은 확률적 섭동(stochastic perturbations)에 강건하다는 경험적 관찰을 활용합니다. 이 방법은 다음과 같이 진행됩니다:

1. 대리 확률 과정 (Surrogate Stochastic Process): 속도장 $\hat{v}_\theta(t, x_t)$만을 학습하는 모델을 위해, 결정론적 ODE 적분에 작은 노이즈 스케줄 $\sigma_{ref}(t)$를 추가합니다. 이는 사전 훈련된 모델의 기본 성능을 유지하면서 필요한 탐색을 가능하게 하는 대리 확률 미분 방정식(SDE)을 생성합니다.
   $$x_{t+\Delta t} = x_t + \hat{v}_{\theta_{ref}}(t, x_t)\Delta t + \sigma_{ref}(t)\sqrt{\Delta t}\xi$$
   이 대리 과정은 KL 정규화를 위한 참조 정책(reference policy)을 정의합니다.

2. 추론 시점 탐색 (Inference-Time Exploration): RL 과정 중, 모델은 강화된 속도장 $\hat{v}_\theta(t, x_t)$와 잠재적으로 탐색을 강화하기 위한 다른 노이즈 스케줄 $\sigma(t)$를 사용하여 탐색을 수행합니다.

3. 정책 최적화 (GRPO): 프레임워크는 그룹 상대 정책 최적화(Group Relative Policy Optimization, GRPO)를 채택합니다. 주어진 조성(composition)에 대해 여러 궤적(trajectories)을 생성합니다. 종단 보상(예: 원자당 음의 에너지)을 계산하고, 그룹 상대적 이득(group-relative advantages)을 산출하여 정책을 업데이트합니다. 이 방식은 학습된 가치 함수(value function)를 필요로 하지 않으며, 이질적인 보상 규모 전반에서 최적화를 안정화합니다.

4. 속도 어닐링 학습 (Velocity-Annealing Learning): OMatG-IRL의 새로운 응용 분야는 시간 의존적 속도 어닐링 스케줄 $s_\theta(t)$를 학습하는 것입니다. 수작업으로 만든 어닐링 스케줄 대신, 모델은 고정된 속도장에 대한 잔차 수정(residual correction)을 학습합니다:
   $$x_{t+\Delta t} = x_t + [1 + s_\theta(t)]\hat{v}_{\theta_{ref}}\Delta t + \sigma(t)\hat{v}_{\theta_{ref}}\sqrt{\Delta t}\xi$$
   이를 통해 모델은 샘플링 효율을 개선하기 위해 속도장을 적응적으로 재조정(rescale)할 수 있습니다.

적용 가능성

본 프레임워크는 다음과 같이 유연하게 설계되었습니다:

• 속도 기반 (Velocity-Based): 속도장만을 학습하는 모델에서 작동합니다 (스코어 불필요).
• 스코어 기반 (Score-Based): 속도와 디노이저(denoiser, 스코어)를 모두 예측하는 모델에도 적용 가능하며, 두 구성 요소를 공동으로 업데이트할 수 있습니다.

3. 주요 기여

• CSP에 대한 최초의 RL 적용: 본 연구는 조성이 고정되어 있고 구조가 생성되는 결정 구조 예측(CSP) 작업에 정책 경사 RL을 적용한 첫 사례입니다.
• 플로우 모델을 위한 스코어 프리(Score-Free) RL: OMatG-IRL은 속도장만을 학습하는 플로우 기반 생성 모델에 대한 RL을 가능하게 하여, 기존에 스코어 기반 확산 모델(diffusion models)로 RL을 제한했던 제약을 극복했습니다.
• 다양성 보상 없는 에너지 기반 강화: 모드 붕괴(mode collapse)를 방려하기 위해 명시적인 다양성 보상이 필요한 De Novo Generation (DNG) 작업과 달리, CSP 작업은 조성 조건부(composition conditioning)를 통해 자연스럽게 다양성을 유지합니다. 저자들은 에너지 기반 목적 함수가 추가적인 다양성 패널티 없이도 효과적으로 강화될 수 있음을 입증했습니다.
• 학습된 어닐링 스케줄: 본 논문은 RL을 통해 시간 의당적 속도 어닐링 스케줄을 학습하는 방법을 도입하여 수작업 헤우리스틱을 대체했습니다.

4. 실험 결과

저자들은 MP-20 데이터셋(Materials Project)을 사용하여 OMatG 프레임워크 상에서 OMatG-IRL을 평가했습니다.

• 에너지 강화: 스코어 기반 및 속도 기반 변형 모델 모두 원자당 상대 에너지를 성공적으로 강화하였으며, 사전 훈련된 베이스라인 대비 약 0.5 eV/atom의 감소를 달성했습니다.
• 성능 동등성: 스코어 계산이 필요 없는 속도 기반 접근 방식이 스코어 기반 방식과 대등한 성능을 보여줌으로써, 대리 확률 과정의 효용성을 검증했습니다.
• 샘플링 효율성:
  • RL 프레임워크를 통해 적분 단계를 획기적으로 줄이면서도 정확한 CSP가 가능했습니다.
  • 특히, Velocity-Annealing OMatG-IRL 변형 모델은 $N_t = 950$ 단계의 적분이 필요한 베이스라인 모델의 성능을 단 $N_t = 100$ 단계만으로 회복했습니다.
  • 놀랍게도, 학습된 어닐링 스케줄은 단계 수를 $N_t = 10$까지 공격적으로 줄였을 때도 견고함을 유지한 반면, 수작업 기반 베이스라인은 급격히 실패했습니다.
• 강건성: 방법론은 매치 레이트(match rate)를 유지하고 RMSE를 낮추는 동시에, 생성의 계산 비용을 크게 낮추었습니다.

5. 의의 및 주장

저자들은 OMatG-IRL이 다음과 같은 측면에서 결정질 재료의 역설계(inverse design)에 있어 중요한 진전을 이루었다고 주장합니다:

1. 플로우 모델을 위한 RL의 민주화: 명시적인 스코어 계산에 대한 의존성을 제거함으로써, 본 프레임워크는 Flow Matching 및 일반적인 Stochastic Interpolants를 포함한 더 넓은 범위의 연속 시간 생성 모델로 RL의 이점(특정 다운스트림 목적 함수 최적화)을 확장합니다.
2. 효율성 개선: 최적의 속도 어닐링 스케줄을 학습하는 능력은 훨씬 적은 적분 단계만으로도 정확한 구조 예측을 가능하게 하여, 재료 스크리닝의 계산 병목 현상을 직접적으로 해결합니다.
3. 태스크 특화 최적화: 본 연구는 RL이 구조적 다양성을 해치지 않으면서도 물리적 목적(에너지 최소화 등)을 효과적으로 최적화할 수 있음을 보여주며, 이는 목표한 특성을 가진 재료를 발견하기 위한 더 직접적인 경로를 제공합니다.

저자들은 한계점으로, 대리 확률 과정이 정확히 주변 분포를 보존(marginal-preserving)하지는 않는다는 점(단, 노이즈가 작을 경우 차이가 미미함)과 현재의 에너지 기반 보상이 구조 매칭 지표(매치 레이트 등)를 직접 최적화하지는 않지만 두 지표가 상관관계가 있다는 점 등을 언급했습니다. 코드는 업데이트된 Open Materials Generation (OMatG) 프레임워크의 일부로 공개되었습니다.