Diffusion Language Model을 위한 Unmasking 정책 학습

Diffusion Large Language Models(dLLMs)는 여러 토큰을 병렬로 디코딩할 수 있어 추론 효율성 측면에서 Autoregressive 모델의 유망한 대안으로 꼽힙니다. 하지만 각 단계에서 어떤 토큰을 먼저 복원(Unmask)할지 결정하는 샘플링 과정이 성능의 핵심적인 설계 요소로 작용합니다.

기존에는 신뢰도 임계값(Confidence Thresholding)과 같은 휴리스틱 전략을 사용해 왔으나, 이는 수동 튜닝이 필요하고 블록 크기가 커질수록 성능이 저하되는 한계가 있었습니다. 본 연구에서는 이러한 샘플링 과정을 **마르코프 결정 과정(MDP)**으로 정식화하고, 강화 학습을 통해 최적의 샘플링 정책을 학습하는 방법을 제안합니다.

제안된 방식의 특징은 다음과 같습니다:

• Lightweight Policy: 단일 레이어 Transformer를 기반으로 하여 dLLM의 토큰 신뢰도를 Unmasking 결정으로 매핑합니다.
• 성능 입증: 실험 결과, 제안된 정책은 Semi-autoregressive(Block) 생성 방식에서 기존 SOTA 휴리스틱과 대등한 성능을 보였으며, Full-diffusion 설정에서는 이를 능가하는 성능을 기록했습니다.