🤖 machine learning

Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

이 논문은 디퓨전 대규모 언어 모델의 정책 분포를 최적의 보상 기울어진 타겟에 정렬함으로써 해당 모델의 추론 능력을 크게 향상시키고, 지도 미세 조정 없이도 기존 베이스라인 대비 상당한 정확도 개선을 달성하는 새로운 강화 학습 프레임워크인 분포 매칭 정책 최적화(DMPO)를 소개한다.

원저자: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

게시일 2026-07-14

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

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

당신에게 이야기를 쓰고, 수학 문제를 풀고, 게임을 즐기는 아주 똑똑한 로봇이 있다고 상상해 보세요. 오늘날 대부분의 로봇은 사람이 책을 읽을 때처럼 왼쪽에서 오른쪽으로 한 단어씩 읽는 방식으로 작동합니다. 이들은 앞서 나갈 수도, 뒤를 돌아볼 수도 없습니다. 그저 이전 단어들을 바탕으로 다음 단어를 추측해야 할 뿐이죠. 이는 마치 눈앞에 있는 퍼즐 조각 하나만 보고 거대한 퍼즐을 채우려는 것과 같아서 매우 느립니다.

여기에 **디퓨전 대규모 언어 모델(dLLMs)**이 등장합니다. 이들은 새로운 세대의 주인공입니다. 단어 하나하나를 써 내려가는 대신, 이들은 혼란스럽고 뒤섞인 "미스터리 토큰"(마치 모든 조각이 안개에 가려진 퍼즐과 같은 상태)에서 시작하여, 정답을 한꺼번에 혹은 원하는 순서대로 서서히 정돈하며 드러냅니다. 이는 이들을 잠재적으로 훨씬 더 빠르게 생각할 수 있게 만듭니다.

하지만 여기에는 함정이 있습니다. 이 디퓨전 로봇들은 빠르긴 하지만, 고등 수학이나 논리 퍼즐 같은 까다로운 추론 작업에서는 항상 가장 똑똑한 것은 아닙니다. 이들을 더 똑똑하게 만들기 위해 과학자들은 보통 **강화 학습(Reinforcement Learning, RL)**이라는 기술을 사용합니다. 이것은 로봇이 좋은 답을 내놓으면 점수를 얻고, 나쁜 답을 내놓으면 점수를 잃는 비디오 게임과 같습니다.

문제점: "모드 탐색(Mode-Seeking)"의 함정

이 로보들을 훈련시키는 기존 방식(GRPO와 같은 방법)은 마치 자신이 옳다고 생각하는 단 하나의 정답만을 공부하는 학생과 같았습니다. 만약 로봇이 점수를 얻기 쉬운 "안전한" 방법을 찾아내면, 더 이상의 탐색을 멈춰버립니다. 로봇은 정체기에 빠지게 되며, 똑같이 훌륭할 수도 있는 다른 기발한 해결책들을 무시하게 됩니다. 논문에서 저자들은 이를 **"모드 탐색(mode-seeking)"**이라고 부릅니다. 이는 마치 로봇이 수학 문제를 오직 한 가지 특정 방식으로만 푸는 법을 배우고, 그 방식이 막히면 당황하는 것과 같습니다. 또한, 처음 찾은 답과는 다르지만 올바른 "무질서한" 답변들을 무시하는 경향이 있습니다.

이 논문은 단순히 최고 점수만을 쫓는 기존 방식이 디퓨전 모델에는 결함이 있다고 주장합니다. 단 하나의 최선 경로에만 집중함으로써, 여러 다양한 경로를 동시에 탐색할 수 있는 디퓨전 모델만의 독특한 능력을 놓치게 된다는 것입니다.

해결책: DMPO ( "지도 매칭" 로봇)

저자들은 **분포 매칭 정책 최적화(Distribution Matching Policy Optimization, DMPO)**라고 불리는 새로운 방법을 제안합니다.

로봇에게 "그저 가장 높은 점수를 찾아라"라고 말하는 대신, DMPO는 "여기에 네가 제시할 수 있는 모든 좋은 답변들의 전체 지도가 있다. 너의 임무는 그 전체 지도를 학습하여 일치시키는 것이다"라고 말합니다.

당신이 강아지에게 공 가져오기를 가르치는 상황을 상상해 보세요.

기존 방식: 당신이 공을 던지면, 강아지는 공이 떨어진 단 한 곳으로 달려갑니다. 만약 공이 덤불 속에 떨어졌다면, 강아지는 오직 덤불 속에서만 공을 가져오는 법을 배웁니다.
DMPO 방식: 당신은 공이 떨어질 수 있는 모든 가능한 장소(덤불, 잔디, 모래, 물 등)가 그려진 지도를 보여주며 이렇게 말합니다. "이 중 어떤 곳이든, 그 장소가 얼마나 좋은지에 따라 가중치를 두어 공을 가져오는 법을 배워라." 그러면 강아지는 마당의 한 구석만이 아니라 전체를 탐색하며 유연하게 학습하게 됩니다 됩니다.

기술적인 용어로, DMPO는 **가중 디노이징 교차 엔트로피(Weighted Denoising Cross-Entropy, WDCE)**라는 특별한 수학적 기법을 사용합니다. 이를 통해 로봇은 매번 모든 것을 처음부터 다시 생성할 필요 없이, 자신의 과거 시도들(완벽하지 않았던 것들까지도)로부터 학습할 수 있습니다. 이는 마치 로봇이 자신의 옛 게임들을 반복해서 공부하며, 혼란을 겪지 않고 좋았던 움직임과 나빴던 움직임 모두로부터 배우는 "리플레이 버퍼"를 가진 것과 같습니다.

핵심 비결: "가중치 베이스라인(Weight Baseline)"

저자들은 작은 규모의 예시(작은 배치 크기)로 훈련할 때 발생하는 까다로운 문제를 발견했습니다. 때때로 로봇은 단지 자신이 본 것이 그것뿐이라는 이유로 나쁜 답변에 보상을 주는 실수를 범하곤 했습니다.

이를 해결하기 위해 그들은 **가중치 베이스라인 차감(Weight Baseline Subtraction)**이라는 영리한 트릭을 발명했습니다.
이것은 선생님이 시험을 채점하는 것과 비슷합니다. 학생이 문제를 맞히면 금메달을 줍니다. 하지만 선생님이 그 문제 하나만 봤다면, 학생이 한 모든 행동이 훌륭하다고 생각할 수도 있습니다. 여기서 "베이스라인"은 "평균적인 기준"과 같습니다. 선생님은 학생의 점수에서 "평균"을 뺍니다.

만약 학생이 평균보다 더 잘했다면, 큰 금메달을 받습니다.
만약 학생이 평균보다 못했다다면, 설령 기술적으로 점수를 얻었더라도 "벌점(음수 가중치)"을 받게 됩니다.

이는 로봇이 평범한 답변에 너무 들뜨지 않게 하며, 진정으로 탁월한 답변을 향해 계속 나아가도록 보장합니다.

결과: 얼마나 더 좋아졌나?

저자들은 매우 어려운 추론 벤치마크들을 통해 이 새로운 방법을 테스트했습니다. 그들은 두 가지 방식으로 사전 학습된(pretrained) 모델(LLaDA-Instruct 및 Dream-Instruct)에 DMPO를 적용했습니다.

직접 적용: 사전 학습된 베이스 모델에 별도의 "숙제"(추론 데이터셋에 대한 지도 미세 조정 또는 SFT) 없이 DMPO를 직접 적용했습니다. 이 "R1-Zero-like" 접근 방식은 DMPO의 순수한 잠재력을 명확히 보여주기 위해 사용되었습니다.
강화된 적용: 이미 SFT를 거친 모델에도 DMPO를 적용하여, 이것이 이미 훈련된 모델을 위한 강력한 업그레이드 도구임을 보여주었습니다.

두 설정 모두에서 결과는 매우 인상적이었습니다.

GSM8K라는 수학 퍼즐에서, 사전 학습된 베이스 모델에 직접 적용된 모델은 엄청난 개선을 보였습니다. 구체적으로, DMPO는 기존의 비-DMPO 강화 학습 베이스라인보다 최대 39.63 퍼센티지 포인트 높은 정확도를 보였으며, 베이스 모델 자체보다는 무려 67.97 퍼센티지 포인트나 높은 성과를 거두었습니다.
**스도쿠(Sudoku)**라는 논리 퍼즐에서, 그들은 거대한 상승을 목격했습니다. 한 버전의 모델은 성공률이 **16.41%**에서 **80.86%**로 급증했습니다 (64포인트 이상의 차이!).
결정적으로, DMPO는 이미 SFT를 거친 모델에 적용했을 때도 상당한 성능 향상을 지속적으로 보여주었으며, 이는 DMPO가 기초부터 시작하든 기존 훈련을 바탕으로 하든 작동하는 견고한 방법임을 입증했습니다.

이 논문은 DMPO가 단순한 미세한 조정이 아니라 근본적인 변화라고 제안합니다. DM포는 모델이 오프-폴리시(off-policy), 즉 오래된 데이터로부터 효율적으로 학습할 수 있게 하며, 포워드-온리(forward-only), 즉 다른 로봇들을 느리게 만드는 값비싼 역방향 계산을 수행할 필요가 없게 만듭니다.

요점

저자들은 이 방법이 테스트된 특정 모델들(LLaDA, Dream 등)과 특정 추론 작업들에 효과적이라고 확신합니다. 그들은 표준 데이터셋에 대한 엄격한 실험을 통해 이러한 개선을 측정했으며, DMPO가 사전 학습 직후의 모델이나 이미 지도 미세 조정(SFT)을 거친 모델 모두에 적용되어 성능을 높일 수 있음을 보여주었습니다. 다만, 아직 모든 유형의 모델이나 작업에 대해 테스트하지는 못했으므로, 결과가 강력하긴 하지만 전체 잠재력은 여전히 탐구 중이라는 점을 인정하고 있습니다.

요약하자면, DMPO는 디퓨전 로봇이 단 하나의 "완벽한" 정답에 집착하는 것을 멈추고, 대신 가능한 모든 좋은 해결책의 지형을 이해하도록 가르침으로써, 그들을 더 똑똑하고, 빠르고, 창의적인 문제 해결사로 만듭니다.

기술 요약: 분포 매칭 정책 최적화(DMPO)를 통한 확산 LLM(Diffusion LLM)의 추론 능력 향상

문제 정의

확산 대규모 언어 모델(dLLM)은 임의 순서 생성(any-order generation)을 가능하게 하고 잠재적으로 더 높은 추론 처리량을 제공함으로써 자기회귀 LLM(AR-LLM)의 유망한 대안을 제시합니다. 그러나 dLLM을 AR-LLM의 추론 능력 수준으로 확장하는 것은 여전히 큰 과제입니다. 기존의 강화 학습(RL) 미세 조정 방법인 GRPO(Group Relative Policy Optimization)는 다음과 같은 두 가지 주요 요인으로 인해 dLLM에 적합하지 않습니다:

계산적 불일치: dLLM은 양방향 프로세싱에 의존하므로, 표준 정책 경사(policy gradient) 방법에서 요구되는 시퀀스 로그 확률(sequence log-probabilities)을 추정하는 것이 AR 모델에 비해 계산 비용이 많이 들고 종종 실행 불가능할 정도로 복ert합니다.
알고리즘적 한계: 현재 dLLM을 위한 RL 접근 방식은 토큰 위치 간의 상관관계를 무시하는 휴리스틱 근사(예: 평균장 분해)에 의존하는 경우가 많습니다. 또한, 이들은 일반적으로 보상 최대화에 초점을 맞추는데, 이는 단일 보상 모드(single reward mode)를 최적화하는 경향이 있습니다. 이러한 "모드 탐색(mode-seeking)" 동작은 보상 해킹(reward hacking)을 초-래하고, dLLM이 이론적으로 생성할 수 있는 다양하고 고품질인 추론 경로의 탐색을 저해할 수 있습니다.

방법론: 분포 매칭 정책 최적화 (DMPO)

본 논문은 확률적 최적 제어(Stochastic Optimal Control, SOC)에 기반한 원칙적인 RL 미세 조정 프레임워크인 **분포 매칭 정책 최적화(DMPO)**를 제안합니다. 기대 보상을 최대화하는 전통적인 방법과 달리, DMPO는 모델의 정책 분포를 보상에 의해 기울어진(reward-tilted) 최적의 타겟 분포와 일치시키는 것을 목표로 합니다.

핵심 프레임워크

목표는 유도된 시퀀스 분포 $p_\theta(o|q)$ 가 다음과 같이 정의된 최적 분포 $p^*(o|q)$ 와 일치하도록 하는 정책 $\pi_\theta$ 를 학습하는 것입니다:
$p^*(o|q) = \frac{1}{Z(q)} p_{ref}(o|q) e^{r(q,o)/\alpha}$
여기서 $p_{ref}$ 는 사전 학습된 정책이고, $r$ 은 보상 함수이며, $\alpha$ 는 온도 파라미터입니다.

이를 달erm하기 위해, 저자들은 타겟과 현재 정책 사이의 전방 쿨백-라이블러 발산(Forward Kullback-Leibler divergence) (교차 엔트로피와 동일)을 최소화하는 문제로 정식화합니다:
$\min_\theta \mathbb{E}_{q \sim D} \mathbb{E}_{p^*(o|q)} \left[ \log \frac{p^*(o|q)}{p_\theta(o|q)} \right]$
이 "질량 덮기(mass-covering)" 목적 함수는 모델이 단일 모드로 붕괴하기보다 다양한 고보상 궤적을 탐색하도록 장려합니다.

주요 기술 구성 요소

가중 디노이징 교차 엔트로피 (Weighted Denoising Cross-Entropy, WDCE):
$p^*$ 는 미지의 분할 함수(partition function) $Z(q)$ 와 직접 샘플링의 불가능성으로 인해 다루기 어렵기 때문에(intractable), DMPO는 **중요도 샘플링(Importance Sampling)**을 채택합니다. 이는 참조 정책 $p_v$ (모델의 stop-gradient 복사본)로부터의 샘플을 $p^*$ 로부터의 가중 샘플로 취급합니다.
결과물인 WDCE 손실 함수는 dLLM의 독특한 전방(노이징) 프로세스를 활용합니다. 이 함수는 최종 클린 샘플과 그에 따른 중요도 가중치만을 사용하여 손실을 계산하며, 중간 롤아웃 궤적은 버립니다. 이로 인해 WDCE는 **오프-정책(off-policy)**이자 전방 전용(forward-only) 목적 함수가 됩니다.
가중치 베이스라인 차감 (Weight Baseline Subtraction):
작은 훈련 배치 사이즈를 사용할 때 발생하는 WDCE의 불안정성이 중요한 구현 과제로 식별되었습니다. 작은 배치에서는 "나쁜" 응답(여전히 $p^*$ 하에서 0이 아닌 확률을 가짐)에 대한 양수 가중치가 지배적이 되어 잘못된 그래디언트 업데이트를 유발할 수 있습니다.
이를 해결하기 위해 저자들은 중요도 가중치에서 베이스라인 가중치 $w_{base}$ 를 차감하는 방식을 제안합니다:
$w_{real} = w - w_{base}$
세 가지 $w_{base}$ 전략이 제안됩니다:

그룹 가중치 베이스라인 (Group Weight Baseline): 수렴을 장려하기 위한 상수 값 (1).
개별 가중치 베이스라인 (Individual Weight Baseline): 음의 보상 기울기(negative reward tilt)에 기반함.
모델 가중치 베이스라인 (Model Weight Baseline): 현재 모델 자체의 중요도 가중치에 기반함.
이 기술은 효과적으로 낮은 보상을 가진 샘플에 대해 음의 그래디언트를 주입하여, 낮은 배치 상황에서 모델이 나쁜 응답을 촉진하는 것을 방지합니다.

가중 직접 판별 최적화 (Weighted Direct Discriminative Optimization, WDDO):
교차 엔트로 엔트로피의 대안으로서, 본 논문은 나쁜 샘플에 대한 음의 그래디언트를 자연스럽게 통합하여 분포 매칭을 더욱 강화하는 GAN 스타일의 목적 함수를 탐구합니다.

주요 기여

새로운 RL 프레임워크: 보상 최대화 패러다임에서 분포 매칭 패러다임으로 전환하여 dLLM을 위해 특별히 설계된 분포 매칭 프레임워크인 DMPO를 제안했습니다.
이론적 근거: 중요도 샘플링과 확률적 최적 제어를 사용하여 WDCE 손실을 유도함으로써, 값비싼 시퀀스 로그 확률 추정 없이도 오프-정책 훈련이 가능한 이론적으로 견고한 방법을 제공했습니다.
실질적 솔루션: 전방 KL 목적 함수에서의 소규모 배치 실패 모드를 식별하고, 가중치 베이스라인 차감 및 WDDO를 도입하여 훈련을 안정화하고 샘플 효율성을 개선했습니다.
효율성: WDCE가 전방 프로세스만으로 훈련을 가능하게 하여 롤아웃의 재사용(off-policy)을 가능하게 하고, 빠른 추론 기술(예: KV-캐시)과의 호환성을 제공함으로써 궤적 기반 방법론에 비해 훈련 속도를 크게 가속화함을 입증했습니다.

실험 결과

저자들은 LLaDA-8B-Instruct 및 Dream-7B 모델을 대상으로 네 가지 추론 벤치마크(GSM8K, MATH500, Countdown, Sudoku)에서 DMPO를 평가했습니다.

성능 향상: DMPO는 베이스 모델 및 기존 RL 베이스라인(diffu-GRPO, cGRPO, d1 포함)을 일관되게 능가했습니다.
- Sudoku에서 DMPO-LLaDA는 가장 우수한 비-DMPO RL 베이스라인보다 39.63 퍼센티지 포인트 향상되었으며, 베이스 모델보다는 67.97 퍼센티지 포인트 향상되었습니다.
- Countdown에서는 강력한 베이스라인 대비 최대 16.40 퍼센티지 포인트의 향상이 관찰되었습니다.
SFT 불필요: 이 방법은 정제된 데이터셋에 대한 지도 학습 미세 조정(SFT) 없이 이러한 결과를 달성하여, 사전 학습된 모델로부터 직접 추론 능력을 확장할 수 있음을 보여주었습니다.
훈련 효율성: 오프-정책 특성과 전방 전용 계산 덕분에, DMPO는 GRPO 유형의 방법보다 2~3배 높은 샘물 효율성과 최대 8배 빠른 롤아웃 샘플링(fast-dLLM 추론과 결합 시)을 달성했습니다. 또한 d1의 보상 수준에 도 только 31%의 훈련 컴퓨팅 예산만으로 도달했습니다.
안정성: DMPO는 매우 오래된(stale) 데이터(동일한 롤아웃 배치를 최대 24회 재사용)에서도 안정성을 유지했는데, 이는 DMPO에서 사용되는 시퀀스 수준 중요도 가중치의 안정성 덕분입니다.

의의 및 주장

본 논문은 DMPO가 추론 작업을 위한 dLLM을 확장하는 데 있어 중요한 진전임을 주장합니다. 보상 최대화 패러다임과 궤적 기반 RL의 계산적 병목 현상에서 벗어남으로써, DMPO는 확산 모델의 고유한 특성을 완전히 활용합니다:

다양성: 전방 KL 목적 함수는 모드 붕괴를 방지하여 dLLM이 다양한 추론 경로를 탐색할 수 있게 합니다.
효율성: WDCE의 오프-정책 및 전방 전용 특성은 dLLM의 높은 처리량 추론 능력과 결면하여, 대규모 모델에 대해 계산적으로 실행 가능한 확장 가능한 RL 미세 조정을 위한 경로를 제공합니다.
일반화 가능성: 이 프레임와는 LLaDA 및 Dream 아키텍처 모두에서 성공을 거둔 것처럼, 다양한 dLLM 백본에 적용될 수 있는 원칙적인 접근 방식으로 제시됩니다.

저자들은 현재의 연구가 특정 모델과 벤치마크에 집중되어 있지만, 분포 매칭 프레임워크가 생성 모델을 개선하고 미래 연구를 위한 더 효과적인 가중치 베이스라인 기술을 설계하는 데 새로운 방향을 열어줄 것이라고 결론짓습니다.