모든 노이즈 제거 단계가 동일하지 않다: 마스크 확산 언어 모델의 빠른 생성을 위한 모델 스케줄링
Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models
April 11, 2026
저자: Ivan Sedykh, Nikita Sorokin, Valentin Malykh
cs.AI
초록
마스킹 확산 언어 모델(MDLM)의 최근 발전은 자기회귀 언어 모델과의 성능 격차를 좁혔지만, 샘플링은 여전히 비용이 많이 든다. 생성 과정에서는 대규모 Transformer를 사용한 많은 전체 시퀀스 노이즈 제거 단계가 필요하며, 자기회귀 디코딩과 달리 KV 캐싱의 이점을 활용할 수 없기 때문이다. 본 연구에서는 확산 프레임워크의 유연성을 활용하여 더 작은 MDLM이 노이즈 제거 단계의 일부에서 전체 모델을 대체하는 모델 스케줄링을 탐구한다. OpenWebText와 LM1B에서 훈련된 모델을 대상으로 한 실험에서, 초기와 후기 노이즈 제거 단계는 중간 단계보다 이러한 대체에 상당히 더 강력하여, 무조건 및 접두사 조건 생성 모두에서 생성적 perplexity의 약간의 저하만으로 FLOPs를 최대 17%까지 줄이면서도 샘플 다양성을 유지할 수 있음을 보여준다. 우리는 이러한 결과를 다양한 시간 단계에서의 소규모 및 대규모 모델 간의 손실과 KL 발산에 기반한 단계 중요도 분석, 그리고 coarse step 세그먼트에 대한 철저한 탐색을 통해 뒷받침한다. 두 방법 모두 데이터셋에 관계없이 일관되게 확산 경로의 중간 부분이 가장 민감함을 확인한다. 우리의 결과는 단순하고 아키텍처에 구애받지 않는 스케줄링 규칙이 생성 품질을 대부분 유지하면서 MDLM 샘플링을 크게 가속화할 수 있음을 시사한다.
English
Recent advances in masked diffusion language models (MDLMs) narrow the quality gap to autoregressive LMs, but their sampling remains expensive because generation requires many full-sequence denoising passes with a large Transformer and, unlike autoregressive decoding, cannot benefit from KV caching. In this work, we exploit the flexibility of the diffusion framework and study model scheduling, where a smaller MDLM replaces the full model at a subset of denoising steps. Across models trained on OpenWebText and LM1B, we show that early and late denoising steps are substantially more robust to such replacement than middle steps, enabling up to a 17% reduction in FLOPs with only modest degradation in generative perplexity under both unconditional and prefix-conditional generation, while preserving sample diversity. We support these findings with a step-importance analysis based on loss and KL divergence between small and large models across timesteps, as well as an exhaustive search over coarse step segments, both of which identify the middle of the diffusion trajectory as most sensitive consistently across datasets. Our results suggest that simple, architecture-agnostic scheduling rules can significantly accelerate MDLM sampling while largely preserving generation quality.