CForce: Boosting Parallel Decoding for dLLMs via Consistency Forcing

AI summary

Overview

Research area: Machine learning, specifically diffusion large language models (dLLMs) and their inference acceleration.

Technical level: Advanced. The paper assumes familiarity with masked diffusion language models, KL divergence, consistency distillation, and threshold-based parallel decoding.

Scope: The paper introduces Consistency Forcing (CForce), a training-time distillation method that makes early, low-context denoising predictions in diffusion LLMs agree with later, more informed predictions from the same model, thereby enabling faster high-parallelism decoding.

What This Paper Is About

Diffusion language models generate text by starting from a fully masked sequence and filling in multiple masked positions per forward pass. When decoding is made aggressive (many tokens committed at once), the model's early predictions are made with very little context and are often wrong, and those errors then propagate through the rest of the denoising trajectory. CForce addresses this by training the model so that its predictions at early denoising stages match predictions the same model would make at later stages, using trajectories the model itself produced during decoding. The goal is a better speed-quality trade-off: more tokens committed per forward pass without losing accuracy.

Key Contributions

1. Reframing dLLM acceleration around early-stage reliability. The paper argues that faster dLLM decoding depends not only on reducing denoising steps but on improving the dependability of predictions made in early, low-context stages under aggressive threshold decoding.

2. Consistency Forcing (CForce), a distillation method with no frozen teacher. The model is trained on pre-collected self-rollout threshold-decoding trajectories, partitioned into stages. Each earlier stage is trained to match a stop-gradient prediction from the same model at the next stage, using Confidence Adaptive KL Divergence (CAD) for still-masked positions and a cross-entropy (CE) anchor for newly revealed positions, plus a curriculum over reveal difficulty.

3. Extension to edit-capable dLLMs. For models that support token-to-token (T2T) refinement after mask-to-token (M2T) commitment, the formulation is broadened so that later T2T corrections supervise earlier M2T predictions, with the CAD domain extended from the still-masked set to the union of masked positions and the stage-level edit set, and the CE anchor extended to all student-masked positions.

4. A theoretical justification and empirical validation. The paper proves an upper bound on early-stage prediction error in terms of adjacent-stage distributional drift plus reveal-boundary token error, and shows experimentally that the same recipe improves both edit-capable and non-edit LLaDA variants.

Main Findings

• Edit-capable model gains are the strongest result. On LLaDA2.1-mini, Consistency Forcing raises average Tokens Per Forward (TPF) from 6.94 to 9.08 while also improving average accuracy from 85.57 to 86.41. Gains are largest on code benchmarks (MBPP score 77.75 to 81.97, HumanEval 85.98 to 86.59) and modest on math (GSM8K 93.56 to 92.27, MATH500 85.00 to 84.80).

• Non-edit model shows an explicit speed-quality trade-off. On LLaDA2.0-mini, average TPF rises from 3.60 to 6.42 for CForce versus 5.07 for the CAP baseline. Average score is lower than the base model but higher than CAP.

• Few-step generation at fixed budgets. Under constrained decoding at a fixed TPF of 8 or 4 on the LLaDA2.0 series, CForce obtains the best average score at both budgets, with the largest margin at TPF = 8.

• Stage size matters. Ablating the number of newly revealed tokens per saved stage on LLaDA2.0-mini-CForce: S=4 gives 80.62 average score and 5.64 average TPF; S=8 (the default) gives 80.88 and 6.42; S=16 gives 75.05 and 7.97. Larger transitions are harder consistency targets.

• Curriculum learning stabilizes training. Removing the gradual exposure schedule (which increases the fraction of newly revealed positions exposed by the later state from 10% to 100% over training) lowers average score from 80.88 to 79.87 and average TPF from 6.42 to 5.97.

• Same-model targets beat a frozen teacher. The same-model stop-gradient target outperforms a frozen-teacher variant on both average score (+1.39) and TPF (+0.51), suggesting an evolving target better tracks the student's own decoding distribution.

• CAD outperforms either KL alone. Forward KL yields 85.75 average score but only 5.07 average TPF; reverse KL reaches 8.72 average TPF but collapses to 65.57 average score; CAD sits in between at 80.88 score and 6.42 TPF.

• The CE anchor helps the score side. Setting λ_CE = 0 lowers average score from 80.88 to 79.02 with TPF nearly unchanged; λ_CE = 2.0 gives 80.22 score and 6.52 TPF but drops noticeably on HumanEval. The default is λ_CE = 0.1.

• Theoretical bound. Early-stage prediction error measured by total variation distance is upper-bounded by a sum of square roots of forward KL terms across adjacent stages plus the negative log-probability of the correct token immediately before reveal. The forward KL part of CAD targets the first term; the CE loss targets the second.

Methodology in Plain English

The researchers start with a pretrained dLLM and use it to generate its own decoding trajectories on prompts from OpenMath-Instruct-2 and OpenCodeInstruct, discarding responses longer than 4096 tokens. Rather than treating every single native decoding step as a training example (which would give very weak signal, since consecutive steps often differ by only one or two revealed tokens), they save a checkpoint state only every time the cumulative number of newly revealed tokens reaches a fixed stage size S.

Each adjacent pair of saved stages then becomes a training example. The earlier stage is the student input; the later stage is used as a target, computed with stop-gradient so gradients do not flow back into it and no separate frozen teacher model is needed. Positions still masked in both stages are aligned using Confidence Adaptive KL Divergence, which mixes forward KL (stable drift control) with a confidence-weighted reverse KL (sharper, mode-seeking predictions when the later stage is confident). Positions newly revealed between the two stages get a cross-entropy anchor against the clean target token. A curriculum gradually increases how much of the later-stage reveal is shown, starting at 10% and ramping to 100% over training, so the model first learns small consistency steps before facing large context gaps. For edit-capable models, the same machinery also uses positions whose visible tokens changed between stages as extra supervision.

Training uses full-parameter fine-tuning with AdamW, learning rate 1.0 × 10⁻⁵, a cosine schedule, stage size 8, global batch size 64, 5 epochs, block size 32, and top-20 token truncation with renormalization during KL distillation, on 64 NVIDIA H800 GPUs. Training thresholds: τ = 0.95 for LLaDA2.0, τ = 0.85 and τ_edit = 0.5 for LLaDA2.1. Evaluation uses SGLang with block size 32 and maximum generation length 4096, with τ = 0.9 for the LLaDA2.0 family and τ = 0.7 with τ_edit = 0.5 for the LLaDA2.1 family. Metrics are Score (exact match on GSM8K and MATH500, pass@1 on MBPP and HumanEval), TPF, and AUP (Accuracy Under Parallelism).

Why This Matters

Impact on research. The paper offers a different angle on dLLM acceleration: instead of only reducing sampling steps or improving token selection at inference time, it trains the model to be reliable in the low-context states its own decoding path actually visits. It also provides a theoretical bound linking early-stage error to adjacent-stage drift and reveal-boundary error, and shows that edit-capable models gain the most because later token-to-token corrections can supervise earlier mask-to-token drafting.

Real-world applications:

• Fast code generation and completion, where the paper reports its largest gains (MBPP and HumanEval).
• Mathematical reasoning assistants that need low-latency multi-step answers.
• Interactive chat and agentic systems where per-token latency dominates user experience.
• Batch inference services where raising tokens per forward pass directly reduces cost.

Industry relevance. Higher tokens per forward pass at comparable quality means fewer forward passes, which translates into lower serving cost and latency. The code is released at https://github.com/inclusionAI/dFactory, and the work was supported by Ant Group Research Fund, indicating direct industrial interest in serving dLLMs efficiently.

Future Directions

• Online trajectory collection. The authors state that CForce currently relies on trajectories collected in advance from a pretrained initialization; as training changes the model, these cached trajectories may drift from the ones the updated model would actually visit. They plan to explore online trajectory collection and policy-updating strategies.

• Scaling the approach. Experiments cover two mini-scale LLaDA variants; whether the gains hold at larger model scales is not reported.

• Tuning the stage-size trade-off. The ablation shows a clear tension: S=4 gives the best score, S=16 the best parallelism, and S=8 the chosen balance. Adaptive or per-example stage sizing is a natural open question.

• Broader edit-capable designs. The edit-capable extension produced the largest gains; applying the same T2T-supervises-M2T idea to other editing or revision-based decoding schemes is unexplored.

Target Audience

Researchers and engineers working on efficient language model inference, discrete diffusion models, and knowledge distillation. Also relevant to practitioners deploying diffusion LLMs in latency- or cost-sensitive production settings, and to readers interested in how consistency-model ideas from continuous diffusion translate to discrete, masked text generation. A background in probabilistic modeling and transformer inference is helpful, as the paper uses KL divergences and total variation bounds directly.

Authors’ abstract

Diffusion large language models (dLLMs) accelerate language generation by predicting multiple masks in a single forward pass. However, existing dLLMs can suffer from unreliable predictions in early denoising stages under aggressive parallelism strategies, leading to errors that can propagate to later stages. To tackle this issue, we present Consistency Forcing (CForce) for dLLMs, a distillation method to force the mask predictions of early stages to align with those of later stages. CForce trains the model on pre-collected self-rollout trajectories, thereby improving training-inference alignment. We introduce Confidence Adaptive KL Divergence as a distillation objective to conjoin the merits of forward and reverse KL. We further provide a theoretical analysis for the consistency objective to explain why CForce can approximately minimize the prediction error of early stages. Critically, the same formulation applies to both mask-to-token decoding and edit-capable decoding; in the edit-capable case, later token-to-token refinements provide additional supervision for earlier masked-state predictions. Experiments on non-edit and edit-capable LLaDA models show improved speed-quality trade-offs, especially under high-parallelism decoding budgets. Code is available at: https://github.com/inclusionAI/dFactory.