d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation

이 논문은 특화된 궤적 가능도 추정기(d2-AnyOrder 및 d2-StepMerge)를 채택하여 논리 및 수학적 벤치마크에서 추론 능력을 크게 향상시키고 새로운 최첨단 성능을 달성한 마스크 확산 언어 모델을 위한 새로운 강화 학습 프레임워크인 d2를 소개한다.

핵심

매우 재능 있는 예술가가 빈 캔버스에서 시작하여 세부 사항을 점진적으로 추가하며 걸작을 완성해 나가는 모습을 상상해 보세요. 이것이 바로 **확산 언어 모델(Diffusion Language Models, DLMs)**이 작동하는 방식입니다. 단어를 하나씩 쌓아 올리는 전통적인 AI 작가(벽돌을 쌓는 것과 같은 방식)와 달리, 이 모델들은 "노이즈"(오래된 TV의 정전기 같은 상태)에서 시작하여 단계적으로 노이즈를 제거(denoise)하며 일관된 이야기나 해결책을 드러냅니다.

이 논문은 이러한 예술가들에게 어려운 퍼즐(수학 문제나 논리 게임 등)을 훨씬 더 잘 풀 수 있도록 가르치는 새로운 훈련 방법인 d2를 소개합니다. 다음은 이 방법이 어떻게 작동하는지에 대한 설명을 쉬운 비유를 사용한 브레이크다운입니다.

문제점: "눈먼" 비평가

AI에게 추론을 가르치기 위해 연구자들은 **강화 학습(Reinforcement Learning, RL)**을 사용합니다. 이것은 AI가 퍼즐을 풀려고 시도하면 "비평가"(보상 시스템)가 점수를 주는 게임과 같습니다. 만약 AI가 좋은 점수를 받으면, AI는 그 행동을 다시 하도록 학습합니다.

하지만 확산 모델에는 함정이 있습니다. 효과적으로 학습하려면, AI는 자신이 생성한 특정 단어 시퀀스가 얼마나 발생할 가능성이 높았는지 정확히 알아야 합니다.

• 전통적인 AI (자기회귀 모델, Autoregressive): 이것은 편지를 쓰는 것과 같습니다. 첫 번째, 두 번째, 세 번째로 무엇을 썼는지 정확히 알 수 있습니다. 따라서 "가능성(likelihood)"을 계산하는 것이 쉽고 빠릅니다.
• 확산 AI (Diffusion AI): 이것은 돌덩어리에서 조각상이 나타나는 과정을 지켜보는 것과 같습니다. AI는 수천 번의 미세한 조정을 거칩니다. 최종 조각상의 정확한 확률을 알기 위해서는 이론적으로 조각 과정의 모든 프레임을 되감아서 매 프레임을 다시 재생해야 합니다. 모든 훈련 단계마다 이 작업을 수행하는 것은 계산적으로 불가능합니다(너무 느리고 비용이 많이 듭니다).

이 "점수"를 정확하게 계산할 수 없었기 때문에, 이전의 방식들은 추측에 의존했고, 이는 부족한 추론 능력으로 이어졌습니다.

해결책: d2 (스마트한 점수 기록원)

저자들은 매우 효율적인 점수 기록원 역할을 하는 d2 프레임워크를 만들었습니다. 이들은 AI의 여정(가능성)을 매번 전체 영상을 다시 재생하지 않고도 파악해 냅니다. 그들은 두 가지 서로 다른 유형의 예술가를 위한 두 가지 도구를 구축했습니다.

도구 1: d2-AnyOrder ("마법 거울")

어떤 확산 모델들은 특별합니다. 이들은 어떤 순서로든 최종 이미지를 드러낼 수 있습니다. 눈을 먼저 그리고 코를 그린 다음 배경을 그릴 수도 있고, 그 반대도 가능하며, 어떤 순서든 상관없이 작동합니다.

• 비유: 당신에게 마법 거울이 있다고 상상해 보세요. 그림이 그려지는 과정을 단계별로 지켜보는 대신, 완성된 그림을 보고 예술가가 모든 단계에서 무엇을 말했을지 한 번에 즉시 볼 수 있습니다.
• 결과: d2-AnyOrder라고 불리는 이 도구는 단 한 번의 시선(한 번의 컴퓨터 패스)만으로 완벽한 점수를 계산합니다. 이는 믿을 수 없을 정도로 정확하며 AI가 매우 빠르게 학습하도록 만듭니다.

도구 2: d2-StepMerge ("청킹 전략")

가장 유명한 LLaDA와 같은 대부분의 인기 있는 확산 모델들은 아무 순서로든 그릴 수 없으며, 특정 순서를 따라야 합니다. 이들에게는 "마법 거울"이 없습니다.

• 비유: 영화를 한꺼번에 볼 수 없다면, 영화를 **청크(덩어리)**로 나눕니다. 모든 프레임을 다 보는 대신, 영화를 10초 단위의 세그먼트로 나누어 봅니다. 각 세그먼트에 대한 점수를 추정하고 이를 모두 더합니다.
• 결과: d2-StepMerge라고 불리는 이 도구는 근사치(approximation)를 구하는 방식입니다. 완벽하지는 않지만, 저자들은 수학적으로 오차가 작고 통제 가능하다는 것을 증로했습니다. 이는 약간의 정확도를 희생하는 대신 엄청난 속도 이득을 취하며, 이를 통해 표준 모델들에 실용적으로 적용할 수 있게 합니다.

결과: 초보자에서 그랜드마스터까지

연구진은 이 도구들을 AI에게 가장 어려운 퍼즐들에 테스트했습니다:

1. 논리 게임: 스도쿠(Sudoku)와 카운트다운(Countdown, 수학 게임).
2. 수학 문제: GSM8K 및 MATH500 (수학적 추론을 위한 표준 벤치마크).

결과:

• 추가적인 "치트 시트"(지도 학습 기반 미세 조정) 없이도: d2 프레임워크는 이전의 어떤 방법보다 모델이 더 잘 추론하도록 가르쳤습니다.
• 경쟁 상대를 압도함: 스도쿠와 카운트다운 테스트에서 d2로 훈련된 모델은 각각 **92%**와 **56%**에 가까운 점수를 기록하며, 약 22%와 42%에 머물러 있던 기존의 최고 방법들을 압도했습니다.
• 수학 숙련도: 복잡한 수학 벤치마크에서 d2는 새로운 "SOTA(State-of-the-Art, 최고 수준)" 기록을 세웠으며, 확산 모델이 전통적인 텍스트 생성 모델만큼이나 추론 능력이 뛰어날 수 있음을 입증했습니다. 동시에 더 빠르고 제어 가능하다는 장점도 보여주었습니다.

요약

이 논문의 핵심은 다음과 같습니다: "우리는 확산 AI 모델이 명확하게 생각하는 법을 가르칠 방법을 찾아냈다."

• 기존의 AI 작업 채점 방식은 너무 느리고 부정확하다는 것을 깨달았습니다.
• 그들은 d2라는 새로운 채점 시스템을 발명했습니다. 이는 완벽하게 정확하거나(특수 모델용), 혹은 스마트하게 효율적(표준 모델용)입니다.
• 이 새로운 시스템을 사용하여, AI 모델은 정답을 미리 배우지 않고도 논리 및 수학 퍼즐을 훨씬 더 잘 풀 수 있도록 학습했습니다.

요약하자면, d2는 확산 모델이 자신의 사고 과정을 정확하게 "되돌아볼" 수 있는 능력을 부여하여, 실수를 통해 배우고 진정한 추론 전문가가 될 수 있도록 해줍니다.

기술 요약

기술 요약: d2: 트래젝토리 가능도 추정을 통한 확산 언어 모델의 추론 능력 향상

문제 정의

확산 언어 모델(Diffusion Language Models, DLMs), 특히 마스크드 DLM(masked DLMs)은 제어 가능성과 병렬 생성 능력 덕분에 자기회귀(autoregressive, AR) 모델의 강력한 대안으로 부상했습니다. 그러나 RL을 통해 이들의 추론 능력을 강화하는 것은 여전히 까 remains 어려운 과제입니다. RL이 AR 모델에서 추론을 유도하는 표준이 된 반면, DLM에 이를 적용하는 것은 쉽지 않습니다. 핵심적인 어려움은 정책 경사(policy gradient) 공식화에 있습니다. 토큰 위치에 따라 가능도가 깔끔하게 분해되는 AR 모델과 달리, 확산 트래젝토리(diffusion trajectory)의 정확한 가능도를 계산하는 것은 계산적으로 불가능합니다. 나이브하게 이러한 가능도를 계산하려면 $T$번의 순방향 패스(forward pass, 여기서 $T$는 확산 단계 수)가 필요하며, 이는 GRPO(Group Relative Policy Optimization)와 같은 표준 정책 경사 방법을 계산적으로 매우 비효과적으로 만듭니다. 기존의 근사 방식들은 상당한 편향(bias)을 유발하여 최적이 아닌 정책 업데이트를 초래하는 경우가 많습니다.

방법론

저자들은 마스크드 DLM에 특화된 추론 프레임워크인 d2를 소개합니다. 이 프레임워크는 샘플링 트래젝토리의 가능도를 정확하게 추정하는 것을 명시적으로 요구하는, GRPO 목적 함수로부터 유도된 새로운 정책 경사 알고리즘을 중심으로 합니다. 계산상의 병목 현상을 해결하기 위해, 본 논문은 서로 다른 모델 클래스에 맞춤화된 일련의 추정기(estimators)를 제안합니다.

1. 이론적 기초

저자들은 시간 잠재 변수(time latents)에 대해 가능도를 한계화(marginalizing)하고 중요도 샘플링(importance sampling)을 사용하여 마스크드 DLM을 위한 GRPO 스타일의 목적 함수를 유도합니다. 결과적인 경사 분해(Theorem 3.1)는 현재 정책 $\pi_\theta$와 오래된 정책 $\pi_{old}$ 사이의 트래젝토리 가능도 비율에 의존합니다. 핵심 과제는 $\pi(x_{0:T})$를 효율적으로 평가하는 것입니다.

2. 추정기(Estimators)

d2 프레임워크는 모델의 아키텍처에 따라 두 가지 특정 추정기를 제공합니다.

• d2-AnyOrder (정확한 추정기):

  • 대상: **임의 순서 디코딩(any-order decoding, AO-dLLMs)**을 지원하는 DLM. 이 모델들은 어텐션 마스크가 특정 인과성 제약(독립 마스크 및 순서 인과성)을 준수한다면 어떤 순서로든 토큰을 디코딩할 수 있습니다.
  • 메커니즘: 특정 위치 인코딩(positional encodings)과 어텐션 마스크를 가진 길이 $2L$의 시퀀스(클린 토큰과 마스크 토큰을 결전합)를 구성합니다. 이를 통해 모델은 단 한 번의 순방향 패스만으로 모든 토큰의 정확한 조건부 가능도를 계산할 수 있습니다.
  • 한계: 표준 마스크드 DLM(예: LLaDA)은 기본적으로 임의 순서 디코딩을 지원하지 않으며, 이 추정기가 편향되지 않기 위해 필요한 어텐션 속성을 보장하는 학습 목표를 가지고 있지 않습니다.

• d2-StepMerge (근사 추정기):

  • 대상: 임의 순서 디코딩을 지원하지 않는 표준 마스크드 확산 모델(MDMs).
  • 메커니즘: 블록 합성 가능도(block composite likelihood)에서 영감을 얻은 이 추정기는 $T$-단계 트래젝토리를 $N$개의 연속적인 시간 세그먼트로 나누어 트래젝토리 가능도를 근사합니다. 모든 단계를 평가하는 대신, 각 세그먼트 내에서 마스크가 해제된 토큰의 가능도를 세그먼트당 한 번의 순방향 패스를 통해 평가합니다.
  • 트레이드오프: 순방향 패스의 횟수를 $T$에서 $N$으로 줄여줍니다. 논문은 근사 오차(KL 발산)에 대한 이론적 상한선(Theorem 4.1)을 제공하며, $N$이 증가함에 따라 오차가 단조 감소함을 보여줌으로써 계산량과 편향 사이의 관계를 정량화합니다.

주요 기여

1. d2 프레임워크: 정확한 트래젝토리 가능도 추정에 의존하는 GRPO 목적 함수를 도출한, 마스크드 DLM을 위한 원칙적인 RL 사후 학습(post-training) 프레임워크입니다.
2. d2-AnyOrder: AO-dLLM을 위한 단일 패스 정확한 가능도 추정기로, 편향되지 않은 가능도 추정이 RL 환경에서 기존 베이스라인보다 크게 우수함을 입증했습니다.
3. d2-StepMerge: 분석적으로 경계가 정해진 근사 오차를 가진, 표준 MDM을 위한 결정론적 다중 패스 추정기로, diffu-GRPO와 같은 기존 방법보다 우수한 효율성-정확도 트레이드오프를 제공합니다.
4. 실험적 검증: 지도 학습 기반의 생각의 사슬(supervised CoT) 미세 조정 없이도 추론 벤치마크에서 최첨단 성능을 달s하는 광범위한 실험을 수행했습니다.

실험 결과

저자들은 논리적 추론(Countdown, Sudoku) 및 수학적 추론(GSM8K, MATH500) 벤치마크, 그리고 코딩 작업 및 독성 제어(toxicity steering)에 대해 d2를 평가했습니다.

• AO-dLLM에서의 성능: 임의 순서 디코딩을 지원하는 모델(예: Eso-LM, Set Diffusion, 커스텀 Any-Order Causal LLaDA)에 적용했을 때, d2-AnyOrder는 DDPO 및 diffu-GRPO와 같은 베이스라인을 크게 앞질렀습니다. 예를 들어, Qwen3-1.7B 베이스 모델의 GSM8K에서 d2-AnyOrder는 67%의 정확도를 기록하여 diffu-GRPO의 63%보다 높은 성과를 보였습니다.
• 표준 MDM에서의 성능: 표준 마스크드 DLM인 LLaDA-8B-Instruct에 적용했을 때, d2-StepMerge($N=16$)는 다음과 같이 새로운 SOTA 결과를 세웠습니다:
  • Sudoku: 91.9% (SFT를 사용한 d1의 22.1% 대비)
  • Countdown: 56.6% (d1의 42.2% 대비)
  • GSM8K: 85.0% (d1의 82.1% 대비)
  • MATH500: 41.6% (wd1의 39.0% 대비)
    특히, 이러한 결과는 추가적인 CoT 데이터에 대한 지도 학습 미세 조정 없이 달성되었습니다.
• 일반화: 이 방법은 다른 아키텍처(Dream 7B)와 작업(HumanEval, MBPP)으로 일반화되었으며, diffu-GRPO보다 개선된 길이 일반화(length generalization) 능력을 보여주었습니다.
• 절제 연구(Ablation): 실험을 통해 정확한 가능도(d2-AnyOrder)가 근사치보다 더 높은 충실도의 신호를 제공하며, d2-StepMerge가 최적의 $N$(예: 16)을 사용할 때 계산 비용과 정확도 사이의 균형을 극단적인 값들보다 더 잘 맞춘다는 것을 확인했습니다.

의의 및 주장

본 논문은 정확한 트래젝토리 가능도 추정이 DLM의 추론을 위한 RL 성공의 핵심이라고 주장합니다. 저자들은 이전 방법들이 편향되거나 계산 비용이 많이 드는 가능도 추정에 의존했기 때문에 최적의 성능을 달성하지 못했다고 설명합니다.

• 이론적 기여: 본 연구는 마스크드 DLM을 위한 최초의 엄밀한 GRO 목적 함수 유도를 제공하며, 스텝 머징(step-merging) 전략의 근사 오차에 대한 이론적 경계를 설정합니다.
• 실질적 영향: d2는 적절한 RL 학습을 통해 확산 모델이 광범적인 지도 CoT 데이터 없이도 자동 회귀 모델과 경쟁하거나 이를 능가하는 추론 능력을 갖출 수 있음을 입증합니다.
• 향-후 전망: 저자들은 d2-AnyOrder가 우수하지만, 현재는 임의 순서 디코딩을 위해 특별히 학습된 모델에 국한되어 있다는 점을 언급했습니다. 이들은 정확한 추정기를 온전히 활용하기 위해 더 강력한 범용 임의 순서 dLLM을 개발하는 것이 자연스러운 다음 단계라고 보고 있습니다.

결론적으로 d2는 논리 및 수학적 추론 작업에서 DLM의 새로운 SOTA를 설정하며, 추론 잠재력을 끌어올리는 데 있어 가능도 추정 전략의 효용성을 입증했습니다.