🤖 machine learning

DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Diffusion Language Models

이 논문은 확산 언어 모델(Diffusion Language Models)의 토큰 생성 순서를 최적화하기 위해, 기존의 무작위 마스킹이나 신뢰도 기반 방식의 한계를 극복하고 Doob h-변환 기반의 프로세스 보상 모델(DPRM)을 도입하여 언어, 단백질, 분자 설계 등 다양한 도메인에서 생성 성능을 향상시키는 플러그인 모듈을 제안합니다.

원저자: Dake Bu, Wei Huang, Andi Han, Hau-San Wong, Qingfu Zhang, Taiji Suzuki, Atsushi Nitanda

게시일 2026-04-28

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Dake Bu, Wei Huang, Andi Han, Hau-San Wong, Qingfu Zhang, Taiji Suzuki, Atsushi Nitanda

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

1. 배경: 퍼즐 맞추기 방식의 차이

기존의 AI(언어 모델)는 보통 '왼쪽에서 오른쪽으로' 글을 써 내려가는 방식(자기회귀 방식)을 사용합니다. 마치 우리가 일기를 쓸 때 첫 문장부터 차례대로 쓰는 것과 같죠.

하지만 최근 등장한 **'확산 모델(Diffusion Model)'**은 방식이 다릅니다. 마치 **'흩어진 퍼즐 조각을 한꺼번에 맞추는 것'**과 같습니다. 처음에는 아무것도 안 보이는 빈 판이었다가, AI가 여기저기 빈칸을 하나씩 채워나가며 점점 완성된 그림을 만들어가는 방식이죠.

문제는 여기서 발생합니다. "퍼즐을 맞출 때, 어떤 조각부터 먼저 맞춰야 가장 빠르고 정확하게 완성할 수 있을까?"라는 문제입니다.

2. 기존 방식의 한계: "눈앞의 것만 보는 근시안"

기존의 AI들은 보통 두 가지 방식으로 퍼즐을 맞췄습니다.

무작위 방식: 그냥 아무 조각이나 집어넣기 (너무 비효율적입니다).
자신감 방식 (Confidence-driven): "내가 지금 이 칸에 이 색깔을 넣을 확률이 제일 높아!"라고 확신이 드는 곳부터 채우기.

하지만 **'자신감 방식'**에는 치명적인 약점이 있습니다. 바로 **'근시안적(Myopic)'**이라는 점입니다. 당장 눈앞의 조각을 맞추는 데는 자신감이 넘치지만, 그 조각을 먼저 맞췄을 때 나중에 전체 그림이 망가질지 아닐지는 고려하지 못합니다. 마치 요리를 할 때, 당장 볶기 쉬운 재료부터 넣었다가 나중에 전체 맛의 조화를 망치는 것과 같습니다.

3. DPRM의 등장: "미래를 내다보는 전략가"

이 논문에서 제안하는 DPRM은 이 문제를 해결하기 위해 **'미래의 보상(Reward)'**을 계산에 넣습니다.

비유하자면, DPRM은 단순한 퍼즐 맞추기꾼이 아니라 **'전략적인 설계자'**입니다.

기존 AI: "이 조각은 내가 맞출 확률이 90%니까 이것부터 하자!"
DPRM: "이 조각을 지금 맞추면 당장은 90% 확률로 맞출 수 있지만, 나중에 전체 그림(최종 결과물)의 완성도를 떨어뜨릴 것 같아. 차라리 지금은 조금 불확실하더라도, 나중에 전체 그림을 멋지게 완성하는 데 결정적인 역할을 할 조각부터 먼저 맞추자!"

DPRM은 **'도브 h-변환(Doob h-transform)'**이라는 수학적 원리를 사용하여, 현재의 선택이 미래의 최종 결과(보상)에 어떤 영향을 미칠지를 미리 예측하고, 그 예측에 따라 퍼즐을 맞추는 순서를 조정합니다.

4. DPRM의 똑똑한 학습법: "처음엔 배우고, 나중엔 실전처럼"

하지만 미래를 완벽하게 예측하는 건 매우 어렵고 계산량도 엄청납니다. 그래서 DPRM은 아주 영리한 **'2단계 전략'**을 씁니다.

1단계 (워밍업): 처음에는 기존 방식처럼 **'자신감'**이 높은 조각부터 맞추며 기본기를 다집니다. (안전하게 시작하기)
2단계 (실전 모드): 어느 정도 익숙해지면, 점차 **'미래 보상'**을 고려하는 방식으로 전환합니다. (전략적으로 움직이기)

마치 초보 운전자가 처음에는 핸들 조작과 페달 밟기에 집중하다가(자신감), 숙련자가 되면 내비게이션을 보며 전체 경로와 교통 흐름을 고려해 운전하는 것(보상)과 같습니다.

5. 결과: "언어부터 생명공학까지, 모든 분야의 마스터"

연구진은 이 기술을 다양한 분야에 적용해 보았습니다.

언어/수학: 어려운 수학 문제를 풀 때 훨씬 더 정확한 답을 찾아냈습니다.
단백질 설계: 생명체의 핵심인 단백질 구조를 만들 때, 훨씬 더 안정적이고 실제와 유사한 구조를 설계했습니다.
DNA/분자 설계: 신약 개발이나 유전자 설계 같은 복잡한 과학적 과제에서도 성능을 크게 높였습니다.

요약하자면...

DPRM은 AI에게 "당장 쉬운 것부터 하지 말고, 전체를 위해 무엇이 중요한지 생각하며 순서를 정해라"라고 가르치는 '전략적 순서 결정 장치'입니다. 이 기술 덕분에 AI는 단순히 빈칸을 채우는 것을 넘어, 전체적인 조화와 목적을 고려하며 훨씬 더 똑똑하게 결과물을 만들어낼 수 있게 되었습니다.

[기술 요약] DPRM: 확산 언어 모델을 위한 플러그인 토큰 순서 지정 모듈

1. 문제 정의 (Problem Statement)

기존의 자기회귀(Autoregressive) 언어 모델은 텍스트를 왼쪽에서 오른쪽으로 순차적으로 생성하는 고정된 순서를 가집니다. 반면, **확산 언어 모델(Diffusion Language Models, DLMs)**은 토큰을 생성할 때 고정된 순서가 없으며, 어떤 토큰을 먼저 드러낼지(reveal) 결정하는 **'토큰 순서 지정(Token Ordering)'**이 알고리즘의 핵심적인 선택 사항이 됩니다.

현재 DLM 시스템은 주로 두 가지 방식을 사용합니다:

무작위 마스킹(Random Masking): 학습 시 무작위로 토큰을 가림. (학습-테스트 불일치 문제 발생)
신뢰도 기반 순서 지정(Confidence-driven Ordering): 모델이 가장 확신하는 토큰부터 먼저 생성. (근시안적 탐색으로 인해 전역적으로 가치 있는 경로를 놓치는 'Quality-Exploration Trade-off' 문제 발생)

본 논문은 **"단순히 모델의 확신도(Confidence)에만 의존하는 것이 최선인가?"**라는 질문을 던지며, 전역적인 보상(Reward)을 고려한 최적의 순서 지정 정책을 제안합니다.

2. 방법론 (Methodology)

논문은 **DPRM (Doob $h$ -transform Process Reward Model)**이라는 플러그인 모듈을 제안합니다. 이 모듈은 기존 모델의 구조, 학습 목적 함수, 데이터 파이프라인을 전혀 수정하지 않고 오직 **'토큰을 드러내는 순서(Ordering Policy)'**만 변경합니다.

핵심 개념: Doob $h$ -transform

DPRM은 수학적으로 Doob $h$ -transform 이론에 기반합니다. 이는 현재의 상태에서 미래에 높은 보상을 받을 확률이 높은 경로로 전이 확률을 기울이는(tilt) 방식입니다. 즉, 단순히 "지금 이 토큰이 확실한가?"를 묻는 대신, **"이 토큰을 지금 드러내는 것이 최종적으로 높은 보상을 받는 경로로 이어질 것인가?"**를 계산합니다.

단계적 온라인 근사 (Progressive Online Approximation)

이론적인 정확한 보상 계산은 비용이 너무 많이 들기 때문에, DPRM은 다음과 같은 실용적인 단계적 접근법을 사용합니다:

1단계 (Warmup): 초기에는 기존의 신뢰도 기반(Confidence-based) 순서를 사용하여 안정적으로 학습을 시작합니다.
2단계 (Transition): 학습이 진행됨에 따라, 온라인으로 추정된 **프로세스 보상(Process Reward)**을 결합합니다.
버킷화(Bucketization): 계산 효율성을 위해 '단계(Phase)'와 '신뢰도 구간(Confidence Bucket)'별로 보상 통계를 관리하여, 유사한 상황에서의 경험을 공유합니다.

3. 주요 기여 (Key Contributions)

플러그인 모듈 설계: 기존 아키텍처를 건드리지 않고 순서 지정 규칙만 교체하여 성능을 높이는 범용적인 프레임워크를 제안했습니다.
이론적 보장 (Theoretical Guarantees):
- DPRM의 근사치가 정확한 Gibbs reveal law로 수렴함을 증명 ( $O(1/N)$ 수렴 속도).
- 온라인 버킷 컨트롤러가 통계적 속도로 정확한 점수를 추적함을 증명.
- 무작위 방식이나 신뢰도 방식보다 샘플 복잡도(Sample Complexity) 측면에서 지수적(Exponential) 이점이 있음을 수학적으로 입증.
광범위한 검증: 자연어 처리뿐만 아니라 단백질, 분자, DNA 설계 등 과학적 도메인까지 아우르는 7가지 설정에서 실험을 수행했습니다.

4. 실험 결과 (Results)

DPRM은 다양한 도메인에서 일관된 성능 향상을 보였습니다:

자연어 추론 (Language Reasoning):
- Pretraining (PUMA): GSM8K 벤치마크 성능이 29.34에서 34.27로 +16.8% 향상.
- Post-training (DMPO): 어려운 수학 문제(MATH Hard)와 카운트다운(Countdown Hard) 문제에서 유의미한 성능 향상.
- Test-time Scaling (Prism): 추론 시 탐색 효율을 높여 정확도 개선.
과학적 생성 (Scientific Generation):
- 단백질 (Protein): 구조적 정확도(RMSD, TM-score) 개선.
- 단일 세포 (Single-cell): 유전자 발현 복구 정확도 대폭 향상.
- 분자 및 DNA (Molecular/DNA): 약물 설계의 유효성(Validity) 및 DNA 조절 서열의 보상 최적화 성공.

5. 의의 (Significance)

본 논문은 **'토큰 순서 지정(Token Ordering)'**이 확산 언어 모델(DLM) 및 마스크 기반 이산 확산 시스템에서 성능을 제어할 수 있는 **근본적인 축(Fundamental Control Axis)**임을 밝혀냈습니다.

특히, 모델이 단순히 눈앞의 확신도에 매몰되지 않고 미래의 보상을 예측하여 탐색 경로를 결정하게 함으로써, 복잡한 추론과 정밀한 과학적 설계가 필요한 영역에서 DLM의 잠재력을 극대화할 수 있는 길을 열었습니다.