확산 언어 모델을 위한 언마스킹 정책 학습

초록

확산(대형) 언어 모델(dLLM)은 이제 많은 과제에서 자기회귀 모델과 동등한 하류 작업 성능을 보이면서도 추론 과정에서 더 높은 효율성을 약속하고 있습니다. 특히 성공적인 변종 중 하나는 마스크된 이산 확산으로, 특수 마스크 토큰으로 채워진 버퍼가 점진적으로 모델의 어휘 집합에서 샘플링된 토큰으로 대체되는 방식입니다. 여러 토큰을 병렬로 마스크 해제하면 효율성을 높일 수 있지만, 한 번에 너무 많은 토큰을 처리하면 생성 품질이 저하될 위험이 있습니다. 따라서 dLLM의 중요한 설계 측면은 확산 과정의 각 단계에서 대체할 토큰을 선택하는 샘플링 절차입니다. 실제로 최근 연구에서는 확신도 임계값 설정과 같은 경험적 전략이 무작위 마스크 해제보다更高的 품질과 토큰 처리량을 동시에 달성한다는 사실을 발견했습니다. 그러나 이러한 경험적 방법에는 단점이 있습니다: 수동 조정이 필요하며, 버퍼 크기가 커질수록 성능이 저하된다는 점을 우리는 확인했습니다. 본 연구에서는 대신 강화 학습을 사용하여 샘플링 절차를 훈련하는 방법을 제안합니다. 구체적으로, 우리는 마스크된 확산 샘플링을 dLLM이 환경 역할을 하는 마르코프 결정 과정으로 공식화하고, dLLM 토큰 확신도를 마스크 해제 결정에 매핑하는 단일 계층 트랜스포머 기반의 경량 정책 아키텍처를 제안합니다. 우리의 실험 결과는 이러한 훈련된 정책들이 준-자기회귀 생성과 결합했을 때 최신 경험적 전략의 성능에 필적하며, 완전 확산 설정에서는 이를 능가함을 보여줍니다. 또한 우리는 이러한 정책들의 전이 가능성을 검토한 결과, 새로운 기반 dLLM과 더 긴 시퀀스 길이로 일반화될 수 있음을 확인했습니다. 그러나 해당 정책들을 도메인 외 데이터에 적용하면 성능이 저하되며, 우리의 접근 방식으로는 정확도-효율성 트레이드오프의 미세 조정이 어려울 수 있다는 점도 관찰했습니다.

English

Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary. Efficiency can be gained by unmasking several tokens in parallel, but doing too many at once risks degrading the generation quality. Thus, one critical design aspect of dLLMs is the sampling procedure that selects, at each step of the diffusion process, which tokens to replace. Indeed, recent work has found that heuristic strategies such as confidence thresholding lead to both higher quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger buffer sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy architecture based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive generation, while outperforming them in the full diffusion setting. We also examine the transferability of these policies, finding that they can generalize to new underlying dLLMs and longer sequence lengths. However, we also observe that their performance degrades when applied to out-of-domain data, and that fine-grained tuning of the accuracy-efficiency trade-off can be challenging with our approach.

확산 언어 모델을 위한 언마스킹 정책 학습

Learning Unmasking Policies for Diffusion Language Models

초록

Support