기술 요약: 확산 언어 모델 정렬을 위한 엔트로피 인식 보상 안내

문제 제기

보상 안내 (또는 사후 표본 추출) 는 연속 확산 모델에서 매우 효과적으로 입증되었습니다. 여기서 보상 모델로부터의 미분 가능한 피드백은 생성 단계의 반복적 정제를 가능하게 합니다. 그러나 이 패러다임을 **이산 확산 대규모 언어 모델 (dLLM)**에 적용하는 것은 근본적인 도전 과제를 제시합니다: dLLM 은 미분 불가능한 이산 토큰을 생성합니다.

기존 접근법들은 이 격차를 메우려 시도하지만, 보상 모델의 신뢰성과 기울기 정확도 사이의 중요한 트레이드오프에 직면합니다:

1. 연속 완화 (기대값): 보상 모델에 "부드러운" 임베딩 (토큰 확률의 기대값) 을 입력하면 기울기 흐름이 가능해집니다. 그러나 보상 모델은 일반적으로 하드 토큰으로 훈련되므로, 부드러운 외래어 (out-of-vocabulary) 입력은 신뢰할 수 없는 보상 신호와 성능 저하를 초래합니다.
2. 스트레이트-쓰루 추정기 (STE): APS(Rout et al., 2025c) 와 같은 방법은 샘플링된 하드 토큰에서 보상을 평가하여 (신뢰성 확보) 입력이 부드러운 것처럼 기울기를 전파합니다. 이는 보상이 평가되는 위치와 기울기가 적용되는 위치 사이의 불일치를 초래하여, 특히 모델의 예측 엔트로피가 높은 경우 (즉, 초기 탈잡음 단계에서) 근사 오차를 유발합니다.

방법론: EntRGi

저자들은 **EntRGi(엔트로피 인식 보상 안내)**를 소개합니다. 이는 토큰 단위로 연속 토큰 완화와 샘플링된 하드 토큰 사이를 동적으로 보간하는 메커니즘입니다. 핵심 혁신은 확산 모델 자체의 토큰별 예측 엔트로피를 사용하여 보간 가중치를 결정하는 것입니다.

메커니즘

탈잡음 과정 중 각 마스킹된 위치 $l$에서 EntRGi 는 보상 모델을 위한 입력 임베딩 $e_l^{EntRGi}$를 볼록 결합으로 구성합니다:
$$e_l^{EntRGi} = (1 - w_l)e_l^{soft} + w_l e_l^{hard}$$
여기서:

• $e_l^{soft}$는 모델의 확률 분포에서 유도된 기대 부드러운 임베딩입니다.
• $e_l^{hard}$는 샘플링된 하드 토큰의 임베딩입니다.
• $w_l$은 토큰의 엔트로피 $H(q_l)$에서 유도된 적응형 가중치이며, 구체적으로 $w_l = H(q_l) / \log K$입니다.

적응 논리:

• 낮은 엔트로피 (높은 확신): 모델이 확신할 때 (낮은 엔트로피), $w_l \to 0$이 됩니다. 입력은 주로 부드러운 임베딩이 되어, 부드러운 임베딩과 하드 임베딩이 서로 가까워 근사 오차가 크지 않기 때문에 깨끗한 기울기 흐름을 허용합니다.
• 높은 엔트로피 (높은 불확실성): 모델이 불확실할 때 (높은 엔트로피), $w_l \to 1$이 됩니다. 입력은 하드 토큰 임베딩 쪽으로 이동합니다. 이는 보상 모델이 현실적이고 분포 내 (in-distribution) 인 입력을 받도록 보장하여, 기울기 근사 오차가 더 크더라도 보상 신뢰성을 유지합니다.

이 접근법은 모델이 확신할 때는 부드러운 임베딩을 신뢰하고 불확실성이 높을 때는 이산 토큰으로 되돌아가는 방식으로 긴장 관계를 해결합니다.

사후 훈련: RGRL

본 논문은 또한 EntRGi(또는 APS) 를 활용하여 보상 안내 완료본을 생성하는 **RGRL(보상 안내 강화 학습)**이라는 사후 훈련 레시피를 제안합니다. 이러한 완료본은 dLLM 을 자기 증류 (self-distillation) 를 통해 미세 조정하는 데 사용되며, diffu-GRPO 와 같은 표준 RL 방법에서 사용하는 스칼라 보상에만 의존하는 대신 밀집된 보상 - 기울기 피드백 루프를 생성합니다.

주요 기여

1. EntRGi: 기울기 정확도와 보상 모델 신뢰성을 적응적으로 균형 잡는 이산 확산 LLM 을 위한 새로운 엔트로피 인식 보상 안내 메커니즘.
2. 테스트 시간 적응 우수성: 엔트로피가 더 높은 더 높은 샘플링 온도에서 특히 여러 벤치마크에 걸쳐 EntRGi 가 이전 최첨단 (APS) 및 기타 기준선 (기대값, Best-of-N) 보다 우수함을 보여주는 실증적 검증.
3. RGRL 사후 훈련: 보상 안내 샘플을 미세 조정에 활용하는 새로운 훈련 레시피. 저자들은 RGRL-EntRGi 가 dLLM 을 위해 널리 채택된 RL 알고리즘인 diffu-GRPO 를 크게 능가하며, 특정 데이터셋에서 보상 점수에서 최대 70% 의 상대적 개선을 달성함을 입증했습니다.
4. 토크나이저 불일치에 대한 강건성: dLLM 과 보상 모델이 토크나이저가 불일치하는 경우 (예: LLaDA 대 Qwen/Llama) 도 이 방법은 효과적으로 유지됩니다. 이는 기존 기울기 기반 방법들이 종종 어려움을 겪는 일반적인 현실 세계 시나리오입니다.

실험 결과

저자들은 5 개의 다중 스킬 데이터셋과 4 개의 보상 모델을 통해 7B 파라미터 모델(Dream-v0-Instruct-7B 및 LLaDA-8B-Instruct) 에서 EntRGi 를 평가했습니다.

• 추론 시간 성능: EntRGi 는 APS 및 기대값에 비해 일관되게 더 높은 Top@1 및 Avg@4 보상 점수를 달성했습니다. 특히, APS 는 더 높은 온도 ($\tau=0.7$) 에서 성능이 크게 저하되는 반면, EntRGi 는 성능을 유지하거나 개선하여 높은 엔트로피 상태를 효과적으로 처리할 수 있음을 보여주었습니다.
• 사후 훈련 효율성: RGRL-EntRGi 는 diffu-GRPO 에 비해 더 빠른 수렴과 더 높은 최종 보상 점수를 보여주었습니다. 이러한 이득은 초기 보상이 낮은 데이터셋 (예: WildChat-IF) 에서 가장 두드러졌으며, 이 방법이 어려운 정렬 작업에 특히 효과적임을 시사합니다.
• 애블레이션 연구: 분석은 엔트로피 기반 가중치가 L2-노름 거리와 같은 대안 휴리스틱보다 우수하며, 엔트로피가 가장 높은 초기 탈잡음 단계에서 근사 오차를 줄인다는 것을 확인했습니다.

의의

본 논문은 EntRGi 가 확산 언어 모델의 "이산 대 미분 가능" 갈등에 대한 원칙적인 해결책을 제공한다고 주장합니다. 모델 불확실성을 피드백 통합을 조절하는 신호로 활용함으로써 EntRGi 는 다음을 가능하게 합니다:

1. 강력한 추론 시간 조향: dLLM 이 재훈련 없이 작업별 목표를 최적화할 수 있도록 하며, 연속 확산 모델과 비교 가능한 신뢰성을 제공합니다.
2. 실현 가능한 사후 훈련 대안: dLLM 을 위한 지도 미세 조정 및 스칼라 보상 RL 보다 더 효과적인 밀집 보상 - 기울기 피드백을 활성화합니다.

저자들은 이 작업을 새로운 이산 확산 언어 모델 클래스에 대해 보상 안내 적응 및 사후 훈련을 강건하고 확장 가능하게 만드는 데 필요한 단계로 위치시킵니다.