Uncertainty-Aware Speculative Decoding for Diffusion Language Models in Long Document Generation

Abstract

Masked diffusion language models (MDLMs) are attractive for long-form generation because they denoise many positions under bidirectional context rather than committing tokens strictly from left to right. The same property makes inference expensive: every reverse step may require dense attention over a long partially masked sequence. Recent sliding-window and speculative wrappers demonstrate that pretrained MDLMs can be accelerated without changing their weights, especially the SW-SpeedDLM framework, which combines segmented sliding-window denoising, cross-segment summary compression, and window-level speculative acceptance. This paper proposes Uncertainty-Aware Speculative Decoding for Diffusion Language Models (UASD-DLM), a complementary inference scheduler for long document generation. Instead of using a fixed speculation depth for every window, UASD-DLM estimates local predictive uncertainty from entropy, margin, and summary-conflict signals, then chooses both the draft depth and context refresh policy dynamically. Low-risk windows receive longer drafts, while high-risk windows are verified earlier to avoid wasted drafting and boundary errors. The acceptance step retains the same rejection-sampling correction as standard speculative decoding, so changing the draft length does not alter the verifier distribution. To avoid unsupported large-model claims, we provide a reproducible CPU proxy benchmark using real text extracted from the uploaded SW-SpeedDLM article and a seeded long-range dependency corpus. In this controlled setting, UASD-DLM improves the speed-memory trade-off over fixed-depth speculative decoding and recovers technical-term accuracy lost by summary-free windowing. The complete code, raw results, and Word manuscript accompany this submission.