Organizations: Department of Artificial Intelligence, Xi’an Jiaotong University, Xi’an, Shaanxi, China · College of Computing and Data Science, Nanyang Technological University, Singapore · School of Computer Science and Technology, Harbin Institute of Technology, Harbin, China
Diffusion language models generate text through iterative denoising under a uniform refinement rule applied to all tokens. However, tokens stabilize at different rates in practice, leading to substantial redundant refinement and motivating refinement control over the denoising process. Existing approaches typically assess refinement necessity from instantaneous, step-level signals under a fixed decoding process. In contrast, whether a token has converged is defined by how its prediction changes along its future refinement trajectory. Moreover, changing the refinement rule reshapes future refinement trajectories, which in turn determine how refinement rules should be formulated, making refinement control inherently dynamic. We propose Progressive Refinement Regulation (PRR), a progressive, trajectory-grounded refinement control framework that derives a token-level notion of empirical convergence progress from full decoding rollouts. Based on this signal, PRR learns a lightweight token-wise controller to regulate refinement via temperature-based distribution shaping under a progressive self-evolving training scheme. Experiments show that PRR substantially accelerates diffusion language model decoding while preserving generation quality.
Figures & tables
Figure 1: Empirical visualization of refinement trajectory changes under refinement control. Blue denotes unmasked tokens, while yellow marks redundant refinement where the current prediction already matches the final unmasked value; the dashed line indicates the final decoding step. Starting from uniform top-2 refinement, inducing a first-stage refinement rule produces different refinement trajectories rather than merely shortening the refinement steps. Further refinement control induced from the first-stage reshaped trajectories reorganizes the decoding trajectories again, demonstrating the dynamic nature of refinement control.
Figure 2: (a) Progressive self-evolving training of PRR. At each stage, a refinement controller regulates diffusion decoding to induce refinement trajectories, which are then used to construct supervision for the next-stage controller. (b) Empirical stability signal construction. For a given token, the refinement trajectory records its predictions across denoising steps. The empirical stability target yi,t is computed from distance-weighted suffix agreement with the final decoded token (Eq. 1), quantifying whether the current prediction has aligned with the final outcome and how persistently this alignment holds.
Figure 3: Accuracy versus number of function evaluations (NFE) on HumanEval, MBPP, GSM8K, and IFEval under Dream-7B and LLaDA-8B backbones.
Method
Accuracy ( ↑ )
NFE ( ↓ )
GSM8K
HumanEval
MBPP
IFEval
MATH
GSM8K
HumanEval
MBPP
IFEval
MATH
Dream-7B
Vanilla
73.62
55.49
54.00
43.06
38.82
256
512
512
512
512
Dynamic-Sampler
72.93
57.32
54.60
43.26
38.62
138.68
237.82
107.67
318.57
161.23
EB-Sampler
73.69
50.61
53.00
43.05
38.30
141.50
238.87
114.57
281.41
153.94
PRR(Ours)
74.15
59.15
55.60
43.76
39.02
138.02
238.08
100.84
288.50
149.78
Table 1: Main results on Dream-7B and LLaDA-8B.Accuracy ( ↑ ): higher is better. NFE ( ↓ ): lower is better. PRR consistently achieves higher accuracy under similar or lower decoding budgets, improving the accuracy-efficiency trade-off across reasoning and code-generation benchmarks.
Figure 4: Token-level unmasking schedule during PRR-regulated decoding on one example. At each decoding step, blue bars mark tokens selected by the top-1 refinement rule, while purple bars indicate additional tokens refined under PRR’s regulation (i.e., the expanded unmasking set after regulation). Yellow stars denote the positions newly unmasked at each step.
Figure 5: Trajectory-grounded convergence progress labels and controller predictions. Left: empirical convergence progress labels derived from full decoding rollouts (Eq. (1)). Light-blue masks denote positions that have already been decoded, and the colormap shows convergence progress over positions that remain under refinement. Middle: predictions of the early-stage controller. Right: predictions of the final controller. The top row corresponds to vanilla refinement trajectories, while the bottom row shows PRR-regulated trajectories from the (N−1) -th training stage. Panel annotations report decoded-token ratios and prediction metrics over the token-step grid.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Training loss evolution across successive refinement heads. Training loss of the controller across refinement heads on four benchmarks. Loss increases as supervision is drawn from increasingly self-induced refinement trajectories, while remaining bounded across stages without divergence.
Figure 7: Accuracy evolution across successive refinement heads. Controller accuracy remains stable across training stages, with only mild fluctuations.
Figure 8: Training dynamics of the progressive self-evolving regulator on GSM8K and HumanEval. We report accuracy, precision, recall, F1, and loss across successive training stages. While loss increases substantially, all predictive metrics remain stable.
Figure 9: Effect of progressive self-evolution and trust-region regularization. We report accuracy as a function of NFE under different supervision update strategies: one-shot self-evolution (100%), progressive self-evolution with mixing ratios of 10% and 1%, and progressive self-evolution (1%) without trust-region regularization. Each curve shows the resulting accuracy-efficiency frontier under a fixed supervision evolution scheme.
Figure 10: Effect of refinement regulation strength α on the quality-efficiency trade-off. We report accuracy as a function of the number of function evaluations (NFE) under α∈{0.5,1.0,1.5} . Larger α enforces stronger temperature-based refinement regulation, leading to earlier token unmasking and lower NFE, while smaller α yields more conservative refinement. Each curve traces the resulting accuracy-NFE frontier induced by a fixed α .
Diffusion language model decoding generates text through iterative parallel refinement, balancing fewer denoising steps with accuracy, coherence, and control.
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.
Diffusion language models generate text through iterative denoising, offering a powerful alternative to autoregressive generation. However, discrete language spaces lack a natural neighborhood structure for defining effective perturbations, so some artificial corruption schemes are proposed in the forward process. Such prescribed forward processes often produce states that are mathematically convenient but misaligned with drafts and errors encountered during generation, resulting in degraded sample quality. To address this limitation, we propose FReDA, a forward-free diffusion language model that eliminates the need for a hand-designed forward process. We formulate diffusion language modeling as recursive distribution refinement, in which model-generated drafts serve as implicit intermediate states, and the learned refinement model progressively moves the draft distribution toward the target distribution. Concretely, FReDA refines drafts by proposing candidate draft sequences and either directly performing self-refinement or selecting among parallel candidates via best-of-N refinement. With this design, FReDA is neighborhood-agnostic, model-complexity-aware, and compatible with flexible refinement parameterizations. Extensive evaluations in the sub-8B regime show that FReDA-4B outperforms larger diffusion base models on reasoning and coding benchmarks, achieving absolute gains of up to 15%, while reaching a 1.5-1.8x average speedup over diffusion baselines and scaling effectively with additional refinement computation.
Diffusion language models (DLMs) support bidirectional denoising and iterative refinement, but generation remains largely controlled by token-level masking schedules. This leaves draft-level refinement implicit and limits explicit control over when to insert missing content, delete redundant spans, revisit uncertain regions, or stop editing. We introduce Diffusion Editing Language Adapter (DELA), a plug-in framework that enables controllable draft refinement with a frozen DLM. DELA operates on clean, variable-length drafts using four explicit editing actions: Insert, Delete, Remask, and Stop. The controller determines when and where to edit the draft, while the frozen DLM fills masked regions through denoising, separating high-level edit control from token-level generation. A naive single-phase editor would be computationally expensive because arbitrary interior edits can disrupt prefix stability and limit reliable KV-cache reuse. To address this, DELA uses a two-phase decoding procedure: a fast, append-only drafting phase that supports KV-cache reuse, followed by a global refinement phase that enables targeted non-monotonic edits to the full draft. We train the controller using a staged curriculum that begins with learning local edit priors from synthetic recovery tasks and mined actions that improve draft quality. Subsequent stages refine edit selection through rollout-based outcome preferences and alignment with downstream task feedback. Experiments on mathematical reasoning and code generation show that DELA achieves gains of up to 5.5 percentage points over direct diffusion decoding while keeping the backbone frozen. DELA thus provides a modular approach to explicit draft-level control in diffusion language models.
Xinsong Feng, Chenan Wang, Leshu Tang +4
Georgia Institute of Technology · William & Mary · UCLA +1