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

TL;DR

Jot tackles decoding inefficiency in diffusion language models by introducing a training-free, per-token early stopping mechanism. It uses a per-position confidence ratio $r^i = p_1^i/(p_2^i + \epsilon)$ and a spatially modulated threshold $\tau^i = \tau_{\max} - (\tau_{\max} - \tau_{\min})\cdot \phi^i$, where $\phi^i$ reflects proximity to already-resolved tokens and $w^i$ aggregates local context with a decay $\gamma$. Across Dream-7B-Instruct and LLaDA-8B-Instruct on GSM8K, MMLU, HellaSwag, and HumanEval, Jot delivers up to $19.6\times$ speedups with minimal quality loss, and ablations show the value of threshold tuning and spatial modulation. The method is training-free, orthogonal to other acceleration techniques, and demonstrates strong practical potential for speeding up diffusion-based text generation in diverse tasks.

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 state-of-the-art efficiency gains while preserving generation quality.

Figures (3)

• Figure 1: Confidence ratio dynamics for a GSM8K sample. The baseline (dashed) shows a smooth decline as the sampler gradually unmasks tokens. Jot with $\tau = 90$ rapidly finalizes high-confidence tokens, then exhibits spikes during the reasoning phase before converging to a similar pattern as the baseline.
• Figure 2: Confidence ratio dynamics for another GSM8K sample. Conservative thresholds ($\tau \geq 30$) allow the model to reach the reasoning phase and produce answer patterns similar to the baseline, while aggressive thresholds ($\tau = 3$) exit prematurely.
• Figure 3: Aggregate speed-quality trade-off on Dream-7B, averaged across all benchmarks. Jot achieves the best Pareto positioning: $7.2\times$ average speedup while retaining $98.3\%$ of baseline quality.