Just on Time: Token-Level Early Stopping for Diffusion Language Models

2026/02/28 by Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi +7
Computer Science · #cs.CL #cs.LG

paper · pdf · doi:10.48550/arxiv.2602.11133

Under review

arxiv created 2026/08/02 · arxiv updated 2026/08/04

Abstract

Diffusion language models generate text through iterative refinement, a process that is often computationally inefficient because many tokens reach stability long before the final denoising step. We introduce a training-free, token-level early stopping approach that identifies convergence independently at each position. Our method leverages lightweight signals derived from the model's predictions and local context to dynamically determine when individual tokens can be finalized. This yields adaptive per-token freezing without task-specific fine-tuning, substantially reducing the total number of diffusion steps required. Across diverse benchmarks, spanning mathematical reasoning, general question answering, and scientific understanding, our approach achieves substantial efficiency gains while preserving generation quality.

Citations

• KLASS: KL-Guided Fast Inference in Masked Diffusion Models
• Diffusion Language Models Know the Answer Before Decoding
• Dream 7B: Diffusion Large Language Models
• Diffusion LLMs Can Do Faster-Than-AR Inference via Discrete Diffusion Forcing
• Accelerating Diffusion Large Language Models with SlowFast Sampling: The Three Golden Principles
• Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
• dKV-Cache: The Cache for Diffusion Language Models
• Large Language Diffusion Models
• Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling
• Simple and Effective Masked Diffusion Language Models
• Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data
• Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution
• DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
• Diffusion-LM Improves Controllable Text Generation
• Training Verifiers to Solve Math Word Problems
• Evaluating Large Language Models Trained on Code
• Structured Denoising Diffusion Models in Discrete State-Spaces
• MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers
• Score-Based Generative Modeling through Stochastic Differential Equations
• Measuring Massive Multitask Language Understanding
• Denoising Diffusion Probabilistic Models
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• HellaSwag: Can a Machine Really Finish Your Sentence?
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding