기술 요약: d-TreeRPO

문제 제기

확산 기반 대규모 언어 모델 (dLLM) 은 autoregressive(AR) 모델에 대한 병렬 디코딩 대안을 제공하여 추론 속도를 높입니다. 그러나 dLLM 의 추론 능력을 향상시키기 위해 강화 학습 (RL) 을 적용하는 것은 기존 방법들이 해결하지 못하는 두 가지 중요한 신뢰성 병목 현상에 직면해 있습니다:

1. 희소 보상과 검증 불가능성: 기존 dLLM 정책 최적화 방법은 종종 최종 답변의 이진 정확도 같은 거시적 결과 보상이나 검증되지 않은 과정 신호에 의존합니다. 이러한 희소성은 정확한 이점 (advantage) 계산을 방해하며, 검증되지 않은 신호는 보상 해킹 (reward hacking) 을 초래할 수 있습니다.
2. 불안정한 확률 추정: AR 모델이 단일 순전파 (forward pass) 에서 체인 룰을 통해 시퀀스를 분해하여 정확한 확률을 계산하는 것과 달리, dLLM 은 "임의 순서 (any-order)" 탈잡음 과정을 통해 텍스트를 생성합니다. 실제 행동 로그 확률은 모든 가능한 디코딩 순서에 대한 기대값을 필요로 하므로 계산이 불가능합니다. 기존 근사 방법들 (예: ELBO 기반 하한 또는 단일 단계 순전파) 은 분석되지 않은 근사 오차를 겪어 정책 경사 (policy-gradient) 업데이트의 신뢰성을 훼손합니다.

방법론: d-TreeRPO

저자들은 dLLM 을 위해 특별히 설계된 신뢰할 수 있는 RL 프레임워크인 d-TreeRPO를 제안합니다. 이 방법은 세 가지 핵심 구성 요소를 통해 aforementioned 병목 현상을 해결합니다:

1. 세분화된 보상을 위한 트리 구조 롤아웃

보상 희소성을 극복하기 위해, d-TreeRPO 는 롤아웃 과정을 독립적인 샘플들의 평탄한 시퀀스가 아닌 트리 구조로 구성합니다.

• 구조: $N$개의 탈잡음 단계는 $H$개의 트리 단계로 그룹화됩니다. 각 비리프 (non-leaf) 노드에서 $B$개의 독립 샘플링이 수행되어 $B^H$개의 리프 노드를 가진 트리를 생성하며, 이 리프 노드들은 말단 상태를 나타냅니다.
• 하향식 이점 계산: 리프 노드에는 검증 가능한 결과 보상 (예: 퍼즐 해결 정확도) 이 할당됩니다. 부모 노드의 상태 가치는 자식 노드 값들의 평균으로 하향식 (bottom-up) 으로 계산됩니다.
• 세분화된 이점: 부모 노드 $p$에서 자식 노드 $c$로의 전이에 대한 이점은 $A_p^c = V_c - V_p$로 정의됩니다. 이는 균일한 결과 보상보다 더 세분화된 단계별 검증 가능 보상 신호를 제공합니다.

2. 확률 추정에 대한 이론적 분석

이 논문은 dLLM 에서 편향되지 않은 확률 추정의 계산 불가능성을 다룹니다.

• 이론적 증명: 저자들은 (Theorem 1) 단일 시간 순전파 확률 추정과 편향되지 않은 기대 확률 사이의 간격이 모델의 예측 신뢰도가 증가함에 따라 감소함을 증명합니다. 구체적으로, 모델이 더 결정론적이 될수록 추정 오차 상한이 좁아집니다.
• 함의: 높은 신뢰도는 추정 정확도를 향상시키지만, 훈련 초기에 지나치게 결정론적인 정책은 탐색을 방해할 수 있으므로 탐색 - 활용 (exploration-exploitation) 트레이드오프를 야기합니다.

3. 시간 계획 자기 증류 손실

탐색 (훈련 초기) 과 정확한 확률 추정 (훈련 후기) 사이의 트레이드오프를 균형 있게 조절하기 위해, d-TreeRPO 는 시간 계획 자기 증류 목적 함수를 도입합니다.

• 메커니즘: 이 손실은 정책이 양의 이점을 가진 자식 노드들로 구성된 타겟 분포와 정렬되도록 장려합니다.
• 계획:
  • 온도 ($\tau(t)$): 초기에는 다양성을 허용하기 위해 높게 시작하여 시간이 지남에 따라 타겟 분포를 날카롭게 만들기 위해 감소합니다.
  • 가중치 ($\lambda(t)$): 초기 탐색 단계에서 손실을 최소화하기 위해 0 에 가깝게 시작하여 후기 단계에서 결정론을 강제하고 추정 오차를 줄이기 위해 점진적으로 증가합니다.
• 목적: 이 손실은 표준 GRPO 스타일 정책 경사 목적 함수에 추가되어, RL 업데이트에 사용되는 단일 통과 확률 추정의 정확도를 향상시키기 위해 안전할 때만 모델의 결정론성을 강화합니다.

주요 기여

1. d-TreeRPO 프레임워크: 트리 구조 롤아웃을 활용하여 세분화되고 검증 가능한 단계별 보상과 이점을 제공하는 dLLM 을 위한 새로운 정책 최적화 알고리즘.
2. 이론적 통찰 및 자기 증류: 예측 신뢰도와 추정 오차 감소를 연결하는 이론적 증명과 이를 바탕으로 설계된, 탐색 - 활용 트레이드오프를 동적으로 관리하는 시간 계획 자기 증류 손실.
3. 실증적 검증: 여러 추론 벤치마크에서 기존 베이스라인 (Diffu-GRPO, VRPO, wd1 포함) 보다 d-TreeRPO 가 우수한 성능을 보임을 입증한 광범위한 실험.

실험 결과

저자들은 네 가지 추론 벤치마크 (Sudoku, Countdown, GSM8K, Math500) 에서 LLaDA-8B-Instruct와 LLaDA-MoE-7BA1B-Instruct에 대해 d-TreeRPO 를 평가했습니다.

• 성능 향상: 베이스 모델과 비교하여 d-TreeRPO 는 다음과 같은 유의미한 개선을 달성했습니다:
  • Sudoku: +86.2% (LLaDA-8B)
  • Countdown: +51.6% (LLaDA-8B)
  • GSM8K: +4.5% (LLaDA-8B)
  • Math500: +5.3% (LLaDA-8B)
• 절대성 연구 (Ablation Studies):
  • 희소 보상: 최종 보상을 모든 토큰에 브로드캐스팅하는 희소 결과 보상을 사용하는 변형은 성능이 현저히 낮아, 트리 구조를 통한 세분화된 신용 부여의 필요성을 확인했습니다.
  • 자기 증류: 자기 증류 손실을 제거하면 성능 저하와 더 높은 확률 추정 오차가 발생했습니다. 반면, 높은 엔트로피를 유지하기 위한 다양성 증진 손실은 가장 낮은 성능과 불안정한 훈련을 초래했습니다.
  • 계획: 역방향 계획 (초기에는 강한 증류, 후기에는 약한 증류) 은 초기 성과 이후 성능 감소를 초래하여 순방향 계획 설계의 타당성을 입증했습니다.
• 효율성: 트리 롤아웃으로 인해 d-TreeRPO 는 일부 베이스라인 (예: Diffu-GRPO) 보다 계산 비용이 높지만, 보상 측면에서 더 빠르게 수렴하며 우수한 최종 성능을 달성하여 실용적인 선택지가 됩니다.

중요성과 주장

이 논문은 d-TreeRPO 가 보상 희소성과 확률 추정 편향이라는 이중 문제를 해결함으로써 dLLM 에서 RL 을 위한 더 신뢰할 수 있는 기반을 확립한다고 주장합니다. 트리 구조 롤아웃을 활용함으로써 이 방법은 이점 신호가 검증 가능하고 세분화되도록 보장합니다. 또한, 확률 추정에 대한 신뢰도의 필요성을 이론적으로 근거화하고 시간 계획 자기 증류 메커니즘을 구현함으로써, dLLM 고유의 탐색 - 활용 트레이드오프를 효과적으로 관리합니다. 그 결과들은 확산 모델에 대한 신뢰할 수 있는 정책 최적화를 위해서는 단순한 결과 보상과 단일 통과 근사를 넘어선 구조화된 검색과 동적 신뢰도 조절을 채택해야 함을 시사합니다.