기술 요약: DACA-GRPO

문제 정의

확산 기반 대규모 언어 모델 (dLLMs) 은 고처리량 병렬 텍스트 생성을 가능하게 함으로써 autoregressive 모델에 대한 매력적인 대안을 제공합니다. 그러나 강화 학습 (RL), 구체적으로 그룹 상대 정책 최적화 (GRPO) 를 dLLMs 에 적용하는 것은 현재 방법론이 해결하지 못하는 두 가지 근본적인 약점에 직면해 있습니다:

1. 시간적 신용 할당의 부재: 토큰이 순차적으로 생성되는 autoregressive 모델과 달리, dLLMs 는 각 비음성화 (denoising) 단계에서 모든 마스킹된 위치를 동시에 예측합니다. 기존 RL 방법론은 전체 생성 결과에 균일한 스칼라 이득 (advantage) 을 할당하여 비음성화 경로를 무시합니다. 이는 모든 단계를 동등하게 중요하게 취급하여, "결정적" 단계 (모델이 수학적 연산자나 코드의 제어 흐름과 같은 핵심 구조적 선택을 확정하는 단계) 와 "일상적" 단계 (단순한 채우기) 를 구분하지 못하게 합니다. 결과적으로 모델의 진화하는 이해를 인코딩하는 중간 예측들이 폐기되어 비효율적인 기울기 업데이트를 초래합니다.
2. 편향되고 분산이 큰 가능도 (Likelihood) 추정: dLLMs 의 표준 정책 최적화는 생성 결과의 로그 가능도 추정에 의존합니다. 주류 접근법은 모든 생성 토큰을 마스킹하고 토큰 간 컨텍스트가 없는 상태 (평균장 근사, mean-field approximation) 에서 이를 예측합니다. 이는 체계적인 편향 (컨텍스트 없이 예측하는 것이 더 어렵다) 과 높은 분산 (단일 무작위 마스킹 패턴이 일반적인 난이도를 대표하지 못할 수 있음) 을 유발하여 정책 기울기를 왜곡합니다.

방법론: DACA-GRPO

저자들은 식별된 약점들을 해결하기 위해 상호 보완적인 두 가지 메커니즘을 통합한 경량 플러그 - 앤 - 플레이 프레임워크인 DACA-GRPO(GRPO 를 위한 비음성화 인지 신용 할당, Denoising-Aware Credit Assignment for GRPO) 를 제안합니다.

1. 비음성화 진행 점수 (Denoising Progress Scores, DPS)

DPS 는 추가적인 순전파 (forward passes) 없이 비음성화 단계 전반에 걸친 모델의 믿음 변화를 정량화함으로써 시간적 신용 할당의 부재를 해결합니다.

• 메커니즘: 각 비음성화 단계 $k$에서 모델은 모든 마스킹된 위치에 대한 분포를 예측합니다. DPS 는 다음 단계 ($k+1$) 에서 여전히 마스킹된 위치들에 대한 모델의 평균 로그 확률 변화를 측정하여 "진행 델타"($\delta_k$) 를 계산합니다. 구체적으로, $\delta_k = S(k+1, k+1) - S(k, k+1)$이며, 여기서 $S$는 단계 $k+1$에서 여전히 마스킹된 위치 집합에 대한 평균 로그 확률입니다.
• 해석: 큰 양의 $\delta_k$는 단계 $k$에서 토큰을 드러내는 것이 모델의 이해를 크게 향상시켰음을 (결정적 단계) 나타냅니다. 작거나 음의 $\delta_k$는 일상적인 진행이나 혼란을 시사합니다.
• 구현: 이러한 원시 델타는 비정상적인 스케일을 처리하기 위해 배치 내에서 단계별로 정규화 (z-score) 됩니다. 정규화된 점수는 RL 손실의 토큰별 이득을 조절합니다: $\tilde{A}_i = A \cdot (1 + \lambda \cdot \bar{\delta}_{\text{birth}(i)})$. 이는 결정적 전환 동안 생성된 토큰에 대한 학습 신호를 증폭시키고 일상적 단계에 대해서는 이를 감쇠시킵니다.
• 비용: 생성 과정에서 자연스럽게 계산된 로지츠 (logits) 를 재사용하므로 추가적인 순전파는 0 회입니다.

2. 계층화 마스킹 가능도 (Stratified Masking Likelihood, SML)

SML 은 토큰들이 상당한 컨텍스트와 함께 예측되도록 보장함으로써 가능도 추정에서의 평균장 편향을 해결합니다.

• 메커니즘: 모든 토큰을 마스킹하는 (0% 컨텍스트) 대신, SML 은 출력 토큰 위치를 $K$개의 불교집합 계층 (disjoint strata) 으로 분할합니다. 각 계층 $S_k$에 대해 모델은 다른 모든 계층을 가시적으로 유지한 채 $S_k$ 내의 토큰을 예측합니다.
• 추정: 시퀀스 수준의 가능도는 모든 계층에 걸친 로그 확률의 합입니다: $\log p_{\text{SML}} = \sum_{k=1}^K \sum_{n \in S_k} \log p_\theta(o_n | o_{\setminus S_k}, q)$.
• 효과: 각 토큰은 시퀀스의 $(K-1)/K$만큼의 컨텍스트 (예: $K=4$인 경우 75% 컨텍스트) 와 함께 예측되어, 0 컨텍스트 가정으로 인한 편향을 현저히 줄입니다.
• 통합:
  • **비비율 기반 방법 (ratio-free methods, 예: wd1)**의 경우, SML 은 보조 손실 항으로 추가됩니다.
  • **비율 기반 방법 (ratio-based methods, 예: Diffu-GRPO, GDPO)**의 경우, SML 은 생성 체제와의 정렬을 유지하면서 편향을 줄이기 위해 완전히 마스킹된 예측과 계층화 예측을 평균화하여 중요도 비율 (importance ratio) 에 통합됩니다.

주요 기여

1. 약점 식별: 본 논문은 현재 dLLM RL 방법론의 주요 병목 현상으로 시간적 신용 할당의 부재와 평균장 가능도 추정의 체계적 편향을 공식적으로 식별합니다.
2. 비음성화 진행 점수 (DPS): 추가적인 계산 비용 없이 중간 예측에서 토큰별 중요도 가중치를 추출하는 새로운 메커니즘으로, 단계별 신용 할당을 가능하게 합니다.
3. 계층화 마스킹 가능도 (SML): 풍부한 컨텍스트로 토큰을 평가함으로써 평균장 편향을 줄이는 위치 기반 계층화 추정기로, 기존 분산 감소 기법을 보완합니다.
4. 통합 프레임워크: 다양한 GRPO 변형 (Diffu-GRPO, wd1, GDPO) 과 호환되는 플러그 - 앤 - 플레이 프레임워크 (DACA-GRPO) 로 DPS 와 SML 을 통합합니다.

실험 결과

저자들은 LLaDA-8B-Instruct 기본 모델을 대상으로 일곱 가지 벤치마크에 DACA-GRPO 를 평가했습니다:

• 수학적 추론: MATH-500, GSM8K.
• 코드 생성: MBPP, HumanEval.
• 제약 조건 충족: Countdown, Sudoku.
• 제약 생성: JSON 스키마 준수.

주요 발견:

• 일관된 개선: DACA-GRPO 는 모든 벤치마크와 기본 방법론에서 일관된 향상을 달성했습니다.
  • 수학 추론: 최대 5.6 퍼센트 포인트 (pp) 향상.
  • 코드 생성: 최대 7.4 pp 향상.
  • 제약 조건 충족: 최대 36.3 pp 향상 (특히 Sudoku 의 경우, wd1 + DPS 에서 셀당 정확도가 약 34% 에서 약 70% 로 급증).
  • JSON 준수: 최대 5.9 pp 향상.
• 상호 보완성: DPS 와 SML 은 상호 보완적입니다. DPS 는 결정적 단계를 식별함으로써 구조화되고 짧은 생성 작업 (예: Sudoku, Countdown) 에서 가장 큰 향상을 가져옵니다. SML 은 토큰 간 의존성이 강하고 평균장 편향이 심각한 장문 작업 (예: 수학, 코드) 에서 가장 큰 향상을 제공합니다.
• 적용 범위: SML 은 긴 출력에는 효과적이지만, 부분적 컨텍스트가 답을 쉽게 결정하는 짧고 엄격하게 제약된 작업 (Sudoku 등) 에서는 모드 붕괴 (mode collapse) 를 유발합니다. 이러한 경우 DPS 만으로도 충분하며 매우 효과적입니다.

중요성과 주장

본 논문은 확산 생성 과정에서 "낭비"되던 중간 예측들이 적절히 활용될 때 의미 있는 훈련 개선을 위한 충분한 신호를 포함하고 있다고 주장합니다. DACA-GRPO 의 주요 중요성은 다음과 같습니다:

1. 효율성: 기존 계산 (로지츠) 을 활용하여 신용 할당을 개선하며, 상당한 오버헤드 없이 구현됩니다 (DPS 는 월시클 시간의 1% 미만을 추가; SML 은 $K$에 비례하는 modest 한 비용 추가).
2. 일반화 가능성: 누락된 신용 할당 구조는 다양한 손실 공식 (비율 기반 대 비비율 기반) 과 도메인 (수학, 코드, 제약 조건) 전반에 걸쳐 공유되는 병목 현상으로 보입니다.
3. 실질적 영향: Sudoku 와 같은 구조화된 작업에서의 극적인 개선은 현재 균일한 신용 할당이 토큰 간 의존성이 긴밀한 작업을 심각하게 소홀히 하고 있음을 시사합니다.

저자들은 DPS 가 보편적인 추가 기능이지만, SML 은 특히 장문 생성에서 가능도 품질이 제약 조건이 되는 방법론과 작업에 우선적으로 적용되어야 한다고 결론지었습니다. 또한 짧은 제약된 출력에 대한 SML 의 적용 한계를 인정하고, 적응형 스트라이드 선택 및 계층화 추정을 사전 학습 목표까지 확장하는 향후 연구를 제안합니다.