기술 요약: GDSD: 확산 언어 모델을 위한 유도된 디노이저 자기 증류로서의 강화 학습

1. 문제 제기

확산 대규모 언어 모델 (dLLMs) 은 여러 토큰을 동시에 생성하고 엄격한 왼쪽에서 오른쪽으로의 생성 순서를 깨뜨림으로써 자기회귀 모델 (ARM) 대비 상당한 추론 효율성 향상을 제공합니다. 그러나 dLLMs 는 현재 최첨단 ARMs 에 비해 생성 품질이 뒤처져 있어, 강화 학습 (RL) 과 같은 효과적인 미세 조정 방법이 필요합니다.

dLLMs 에 RL 을 적용하는 데 있어 핵심적인 장애물은 정책 가능도 (policy likelihood) 의 계산 불가능성입니다. 기존 해결책은 일반적으로 두 가지 범주로 나뉩니다:

1. 궤적 가능도 (Trajectory Likelihood) 방법: 전이 확률을 누적하여 정확한 가능도를 추정합니다. 편향은 없으나 계산 비용이 많이 들며 사전 학습 목표와 불일치합니다.
2. 시퀀스 가능도 (ELBO) 방법: 무작위로 마스킹된 시퀀스에서 추정된 증거 하한 (ELBO) 을 실제 가능도의 대리 변수로 사용합니다. 이는 계산 효율성이 높고 사전 학습과 정렬되어 있습니다.

핵심 문제: 본 논문은 ELBO 기반 RL 방법이 학습 - 추론 불일치 (Training-Inference Mismatch, TIM) 편향을 겪는다고 주장합니다. 이 편향은 두 가지 원인에서 발생합니다:

• 가능도 간극 (Likelihood Gap): ELBO 는 실제 가능도에서 무시할 수 없는 간극을 가진 하한이므로, 정책 경사 (PG) 또는 PPO 목적 함수에서 사용되는 중요도 샘플링 비율에 편향과 분산을 도입합니다.
• 샘플링 분포 불일치: dLLM 추론은 반복적인 재마스킹과 블록 단위 디코딩을 사용하므로, ELBO 를 통해 추정된 학습 정책 분포 ($\hat{\pi}_{old}$) 와는 다른 샘플링 분포 ($\pi^{rm}_{old}$) 를 유도합니다.

이러한 불일치는 중요도 샘플링 보정이 비효율적이게 만들어 성능 저하를 초래하며, 심각한 경우 치명적인 학습 붕괴로 이어집니다.

2. 방법론: 유도된 디노이저 자기 증류 (GDSD)

저자들은 dLLMs 를 위한 RL 을 오프-폴리시 자기 증류 문제로 재정의하여 가능도 추정의 필요성을 우회하고 TIM 편향을 피하는 유도된 디노이저 자기 증류 (GDSD) 프레임워크를 제안합니다.

2.1 이론적 기반

역 KL 정규화 하의 RL 에서 최적 정책 $\pi^*$는 폐쇄형 해를 가집니다:
$$\pi^*(x_0) \propto \pi_{old}(x_0)^{1-\beta} \cdot \pi_{ref}(x_0)^\beta \cdot \exp(\psi A(x_0))$$
여기서 $A(x_0)$는 이점 함수 (advantage function) 이고, $\beta$는 정규화 계수이며, $\psi$는 유도 계수입니다.

저자들은 이 최적 정책이 마스킹 확산 모델을 위한 유도된 디노이싱 분포 (즉, "교사" 디노이저 $p^*$) 를 유도함을 보여줍니다:
$$p^*(x_0|x_t) = p^{ref}_{old}(x_0|x_t) \cdot \exp(\psi A(x_0) - A_t(x_t))$$
여기서 $p^{ref}_{old}$는 이전 정책과 참조 정책의 기하학적 혼합이며, $A_t(x_t)$는 정규화 상수입니다.

2.2 GDSD 목적 함수

가능도 기반 목적 함수를 최적화하는 대신, GDSD 는 로지트 (logits) 를 일치시킴으로써 현재 디노이저 $p_\theta$를 교사 $p^*$로부터 직접 증류합니다. 이 목적 함수는 학생과 교사의 로지트 간의 제곱 $L_2$ 거리를 최소화합니다:
$$\mathcal{L}_{GDSD} = \mathbb{E} \left[ \left( \log p_\theta(x_0|x_t) - \log p^*(x_0|x_t) \right)^2 \right]$$

2.3 정규화 없는 재형성

이를 구현하는 데 있어 주요 과제는 교사 분포의 계산 불가능한 정규화 상수 $Z_t$(또는 $A_t$) 입니다. 저자들은 이를 제거하기 위한 두 가지 해결책을 제안합니다:

1. 직접 일치 (Direct Matching): Softmax 연산자의 이동 불변성을 활용하면, 학생 로지트가 상수만큼 더해진 비정규화 교사 로지트와 일치하는 경우 해당 상수를 무시할 수 있습니다.
2. 토큰 레벨 로지트 중앙화 (Token-Level Logit Centralization, TLC): 반복적인 업데이트 중 통제되지 않는 로지트 스케일 변화를 방지하기 위해, 저자들은 로지트를 중앙화할 것을 제안합니다. 각 토큰 위치에서 어휘 전체에 대한 평균 로지트를 빼면 목적 함수에서 정규화 상수가 완전히 제거됩니다.

TLC 를 포함한 최종 실용적 목적 함수는 다음과 같습니다:
$$\mathcal{L}_{GDSD} = \mathbb{E} \left[ \left( \log \bar{p}_\theta(x_0|x_t) - \log \bar{p}_{old}(x_0|x_t) - \psi A(x_0) \right)^2 + \beta \left( \log \bar{p}_\theta(x_0|x_t) - \log \bar{p}_{ref}(x_0|x_t) \right)^2 \right]$$
여기서 $\bar{p}$는 로지트 중앙화 모델을 나타냅니다.

3. 주요 기여

1. 자기 증류로의 환원: 본 논문은 dLLMs 를 위한 역 KL 정규화 RL 이 폐쇄형 최적 정책에서 유도된 자기 교사로부터 유도된 디노이저를 증류하는 것과 동등함을 확립합니다.
2. GDSD 프레임워크: 중요도 샘플링 비율이나 ELBO 대리 변수에 의존하지 않으므로 설계상 TIM 편향을 제거하는 새로운 가능도 없는 자기 증류 방법입니다.
3. 정규화 없는 최적화: 토큰 레벨 로지트 중앙화 (TLC) 의 도입은 계산 불가능한 분할 함수를 계산하지 않고도 효율적인 온라인 RL 을 가능하게 합니다.
4. 이론적 통합: 본 논문은 기존 ELBO 기반 방법들 (예: wd1, DMPO, SPG, UniGRPO) 이 서로 다른 발산 측정 (Forward-KL 또는 Reverse-KL) 을 사용하는 증류의 특정 사례로 볼 수 있음을 보여주지만, GDSD 가 피하는 데이터 비효율성이나 TIM 편향을 겪는다고 주장합니다.

4. 실험 결과

저자들은 여섯 가지 벤치마크에서 LLaDA-8B와 Dream-7B에 대해 GDSD 를 평가했습니다:

• 수학적 추론: GSM8K, MATH500
• 계획 수립: Countdown, Sudoku
• 코딩: HumanEval, MBPP

주요 발견:

• 성능 향상: GDSD 는 일관되게 최첨단 ELBO 기반 방법들을 능가했습니다. Dream-7B 모델에서 GDSD 는 계획 수립 작업 (Sudoku 및 Countdown) 에서 최대 **+19.6%**의 테스트 정확도 향상을 달성했습니다. LLaDA-8B에서는 계획 수립, 수학, 코딩 벤치마크 전반에 걸쳐 **+0.6% 에서 +5%**까지의 향상이 관찰되었습니다.
• 학습 안정성: GDSD 는 진동이나 붕괴를 보인 ELBO 기반 베이스라인 (예: SPG, ESPO) 에 비해 훨씬 더 안정적인 학습 보상 동역학을 보여주었습니다.
• 애블레이션 연구: 토큰 레벨 로지트 중앙화 (TLC) 구성 요소는 학습을 더욱 안정화시키고, 일부 경우 (Sudoku 등) 에서는 모델이 지역 최소값을 탈출하여 더 높은 최종 보상을 달성하는 데 도움이 되는 것으로 나타났습니다.

5. 의미 및 주장

본 논문은 직접 디노이저 자기 증류가 dLLMs 를 위한 ELBO 기반 RL 보다 더 안정적이고 효과적인 대안이라고 주장합니다. 가능도 대리 변수에 대한 의존성을 제거함으로써 GDSD 는 중요도 샘플링 접근법 내재적 구조 편향을 피합니다. 관찰된 성능 향상은 단순한 하이퍼파라미터 튜닝 때문이 아니라, dLLMs 의 실제 추론 분포와 학습 목적 함수를 정렬하는 더 원칙적인 최적화 경로에서 비롯된다는 결과가 시사합니다. 저자들은 ELBO 가 dLLMs 를 통한 RL 에 필수적이지 않으며, 이를 제거하면 더 강력한 정렬 및 추론 능력이 도출된다고 결론지었습니다.