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

Yuna Kim

doi:10.54097/w9c56t64

Authors

Yuna Kim KAIST, School of Computing, South Korea

DOI:

https://doi.org/10.54097/w9c56t64

Keywords:

Diffusion language models, Speculative decoding, Long document generation, Uncertainty estimation, Masked denoising, Context compression

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.

Downloads

Download data is not yet available.

References

[1] Hoogeboom, E., Nielsen, D., Jaini, P., Forre, P., & Welling, M. (2021). Argmax flows and multinomial diffusion: Learning categorical distributions. In Advances in Neural Information Processing Systems 34 (NeurIPS).

[2] Sahoo, S. S., Arora, S., Goyal, A., & Arora, S. (2024). Simple and effective masked diffusion language models. arXiv, arXiv:2406.07524. https://doi.org/10.48550/arXiv.2406.07524

[3] Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., & Li, C. (2025). Large language diffusion models. arXiv, arXiv:2502.09992. https://doi.org/10.48550/arXiv.2502.09992

[4] Austin, J., Johnson, D. D., Ho, J., Tarlow, D., & van den Berg, R. (2021). Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems 34 (NeurIPS).

[5] Teng, D., Rhee, M., Qin, Y., Zi, B., & Liu, W. (2026). SW-SpeedDLM: Sliding-window speculative decoding for diffusion language models under long-context constraints. Mathematics, 14(12), Article 2137. https://doi.org/10.3390/math14122137

[6] Lou, A., Meng, C., & Ermon, S. (2023). Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv, arXiv:2310.16834. https://doi.org/10.48550/arXiv.2310.16834

[7] Shi, J., Han, K., Wang, Z., Doucet, A., & Titsias, M. K. (2024). Simplified and generalized masked diffusion for discrete data. arXiv, arXiv:2406.04329. https://doi.org/10.48550/arXiv.2406.04329

[8] Christopher, J. K., Bartoldson, B. R., Ben-Nun, T., Cardei, M., Kailkhura, B., & Fioretto, F. (2024). Speculative diffusion decoding: Accelerating language generation through diffusion. arXiv, arXiv:2408.05636. https://doi.org/10.48550/arXiv.2408.05636

[9] Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast inference from transformers via speculative decoding. In Proceedings of the International Conference on Machine Learning (ICML). PMLR.

[10] Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., & Jumper, J. (2023). Accelerating large language model decoding with speculative sampling. arXiv, arXiv:2302.01318. https://doi.org/10.48550/arXiv.2302.01318

[11] Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., & Dao, T. (2024). Medusa: Simple LLM inference acceleration framework with multiple decoding heads. arXiv, arXiv:2401.10774. https://doi.org/10.48550/arXiv.2401.10774

[12] Sun, Z., Suresh, A. T., Ro, J. H., Beirami, A., Jain, H., & Yu, F. (2023). SpecTr: Fast speculative decoding via optimal transport. In Advances in Neural Information Processing Systems 36 (NeurIPS).

[13] Fu, Y., Bailis, P., Stoica, I., & Zhang, H. (2024). Break the sequential dependency of LLM inference using lookahead decoding. arXiv, arXiv:2402.02057. https://doi.org/10.48550/arXiv.2402.02057

[14] Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE: Speculative sampling requires rethinking feature uncertainty. arXiv, arXiv:2401.15077. https://doi.org/10.48550/arXiv.2401.15077

[15] Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The long-document transformer. arXiv, arXiv:2004.05150. https://doi.org/10.48550/arXiv.2004.05150

[16] Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontañón, S., Pham, P., Ravula, A., Wang, Q., Yang, L., & Ahmed, A. (2020). Big Bird: Transformers for longer sequences. In Advances in Neural Information Processing Systems 33 (NeurIPS).

[17] Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2023). Efficient streaming language models with attention sinks. arXiv, arXiv:2309.17453. https://doi.org/10.48550/arXiv.2309.17453

[18] Teng, D. (2025). PACO: Predictive auto-configuration for SLO-constrained large language model inference serving. Innovation and Technology Studies, 2(1), 1-10.

[19] Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., Wang, Z., & Chen, B. (2023). H2O: Heavy-hitter oracle for efficient generative inference of large language models. In Advances in Neural Information Processing Systems 36 (NeurIPS).

[20] Li, Y., Huang, Y., Yang, Y., Zhang, J., Luo, M., Wang, H., & Zhang, Y. (2024). SnapKV: LLM knows what you are looking for before generation. arXiv, arXiv:2404.14469. https://doi.org/10.48550/arXiv.2404.14469

[21] Cai, Z., Zhang, Y., Gao, L., & Liu, T. (2024). PyramidKV: Dynamic KV cache compression based on pyramidal information funneling. arXiv, arXiv:2406.02069. https://doi.org/10.48550/arXiv.2406.02069

[22] Wang, Z., Yang, J. S., Shang, W., & Ding, J. (2026). FairPromote: Explainable and fairness-aware talent promotion prediction via adversarial debiasing and SHAP-based interpretation. IEEE Access, 14, 1-16.

[23] Zhang, F., Guo, Z., Ding, J., Yang, J., & Liu, W. (2026). Adaptive sensor fusion for robust perception in dense fog: A gated vision and LiDAR integration framework. Sensors, 26(12), Article 3728. https://doi.org/10.3390/s26123728

[24] Zi, B. (2024). Large language models for enterprise workflow automation in financial operations. Innovation and Technology Studies, 1(1), 24-29.

[25] Ding, J., Shen, Z., & Liu, W. (2026). Game-theoretic cost-sensitive adversarial training for robust cloud intrusion detection against GAN-based evasion attacks. Applied Sciences, 16(8), Article 3944. https://doi.org/10.3390/app16083944

[26] Zi, B. (2024). Cloud-native distributed systems for real-time payment intelligence. AI and Data Science Journal, 1(1), 51-56.

[27] Teng, D. (2025). TEAS: Token- and energy-aware autoscaling for cost-efficient LLM serving. AI and Data Science Journal, 6(3), 1-10.

[28] Zhang, H. (2025). Physics-informed neural networks for high-fidelity electromagnetic field approximation in VLSI and RF EDA applications. Journal of Computing and Electronic Information Management, 18(2), 38-46.

[29] Jiao, Y., Fan, H., Yue, X., Ping, W., Sun, T., & Wang, J. (2026). Dynamic heterogeneous graph contrastive learning for uncovering collusive financial fraud. Scientific Reports, 16, Article 14567. https://doi.org/10.1038/s41598-026-58938-5

[30] Ping, W., Jiao, Y., Fan, H., & Zhang, X. (2026). Multimodal fraud detection in financial statements: A trimodal attention network with contrastive evidence chain construction. IEEE Access, 14, 1-17.