Just on Time: Token-Level Early Stopping for Diffusion Language Models
This addresses the problem of slow text generation for users of diffusion language models, offering an incremental improvement in efficiency.
The paper tackles the computational inefficiency of diffusion language models by introducing a training-free, token-level early stopping method that identifies convergence independently at each position, achieving state-of-the-art efficiency gains while preserving generation quality across diverse benchmarks.
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 state-of-the-art efficiency gains while preserving generation quality.