CD4LM: Consistency Distillation and aDaptive Decoding for Diffusion Language Models

TL;DR

CD$^{4}$LM tackles latency in diffusion language models by decoupling training from inference and enabling adaptive compute through two components: Discrete-Space Consistency Distillation (DSCD) and Confidence-Adaptive Decoding (CAD). DSCD trains a trajectory-robust, trajectory-invariant student by pairing a nested teacher-subset masking scheme with KL-distillation and a reconstruction objective, effectively projecting the teacher’s denoising trajectory onto the student’s reduced information set. CAD leverages the distilled student to dynamically commit high-confidence tokens within a block-diffusion framework, achieving substantial wall-clock speedups while maintaining or improving accuracy, and yielding emergent hierarchical planning (skeleton-first generation). Across code and mathematics benchmarks (notably GSM8K, HumanEval, and MBPP), CD$^{4}$LM attains a strong Pareto improvement, e.g., a mean speedup of $3.62\times$ with better average accuracy, and on GSM8K up to $5.18\times$ wall-clock acceleration, demonstrating practical impact for scalable, structured reasoning tasks. The approach preserves compatibility with existing backbones and opens avenues for dynamic-context diffusion, on-policy training, and caching integrations.

Abstract

Autoregressive large language models achieve strong results on many benchmarks, but decoding remains fundamentally latency-limited by sequential dependence on previously generated tokens. Diffusion language models (DLMs) promise parallel generation but suffer from a fundamental static-to-dynamic misalignment: Training optimizes local transitions under fixed schedules, whereas efficient inference requires adaptive "long-jump" refinements through unseen states. Our goal is to enable highly parallel decoding for DLMs with low number of function evaluations while preserving generation quality. To achieve this, we propose CD4LM, a framework that decouples training from inference via Discrete-Space Consistency Distillation (DSCD) and Confidence-Adaptive Decoding (CAD). Unlike standard objectives, DSCD trains a student to be trajectory-invariant, mapping diverse noisy states directly to the clean distribution. This intrinsic robustness enables CAD to dynamically allocate compute resources based on token confidence, aggressively skipping steps without the quality collapse typical of heuristic acceleration. On GSM8K, CD4LM matches the LLaDA baseline with a 5.18x wall-clock speedup; across code and math benchmarks, it strictly dominates the accuracy-efficiency Pareto frontier, achieving a 3.62x mean speedup while improving average accuracy. Code is available at https://github.com/yihao-liang/CDLM

Figures (7)

• Figure 1: Overview of the DSCD pipeline.(1) Input & Masking: We employ a teacher-subset masking strategy where the teacher's mask $\mathcal{M}_{\text{T}}$ is strictly sampled from the student's mask $\mathcal{M}_{\text{S}}$ ($\mathcal{M}_{\text{T}} \subseteq \mathcal{M}_{\text{S}}$). This ensures the teacher always conditions on a superset of the student's context, acting as a lower-variance guide. (2) Forward & Optimization: The student $p_\theta$ predicts tokens from the heavily masked view, optimized jointly by a reconstruction loss $\mathcal{L}_{\text{recon}}$ (anchoring to ground truth) and a consistency loss $\mathcal{L}_{\text{cons}}$ (aligning with the frozen teacher's soft targets). (3) Curriculum: A dynamic schedule $\lambda(g)$ governs the transition from pure distillation to supervised refinement.
• Figure 2: Overview of different decoding paradigms. (a) Autoregressive decoding generates tokens left-to-right. (b) Block diffusion decodes a preset number of masked tokens per step within a fixed block schedule. (c) Block diffusion with adaptive decoding: Within each block, we adaptively select a high-confidence set of masked positions to unmask, enabling variable token counts (and selected positions) across steps while preserving the block-wise decoding structure. (d) The same adaptive rule can also be applied to pure diffusion, iteratively denoising the entire sequence starting from all <mask> tokens.
• Figure 3: Illustration of CAD. In each step, the model predicts token distributions for all masked positions in parallel. We compute the confidence score $c_i^{(s)}$ based on the maximum probability and selectively unmask tokens that satisfy the threshold $c_i^{(s)} \ge \gamma_{\text{conf}}$. Low-confidence tokens remain masked for subsequent iterations.
• Figure 4: Accuracy--compute trade-off under different decoding budgets. We plot performance (%) versus the average NFE across benchmarks. The dashed gray curve denotes the LLaDA-8B-Instruct baseline evaluated under increasing NFE budgets, while the solid orange curve denotes our CAD with the given threshold ($\gamma_{\text{conf}}=0.95$). Each marker corresponds to one decoding budget. Across all four benchmarks, our method yields a consistently better trade-off, achieving comparable or higher performance at substantially lower NFE.
• Figure 5: Accuracy-compute Pareto frontier on GSM8K. The dashed gray curve represents the LLaDA baseline. Strict Dominance: Our method (solid curves) consistently pushes the frontier upwards and leftwards. The orange arrow highlights our selected operating point ($\gamma_{\text{conf}}=0.95$, dark blue curve), which achieves a 3.4$\times$ speedup while maintaining accuracy comparable to the baseline. While a higher threshold ($\gamma_{\text{conf}}=0.99$) prioritizes quality, $\gamma_{\text{conf}}=0.95$ strikes an optimal efficiency-quality balance.

Theorems & Definitions (10)

• Definition 1: Information ordering
• Proposition 1: DSCD as projection onto $\mathcal{F}_S$
• proof
• Lemma 1: Conditional Rao-Blackwellization
• proof
• Lemma 2: Monotonicity within the active block
• proof
• Corollary 1: Per-block and total NFE bound
• Proposition 2: Pure diffusion as a special case
• Proposition 3: Reduction to fixed-step block diffusion