PaperlayerView exact revision on arXiv

Source-linked AI summary

CForce: Boosting Parallel Decoding for dLLMs via Consistency Forcing

Yuji Ren, Chenkai Xu, Zhuocheng Gong, Jianguo Li, Zhijie Deng

arXiv:2608.13925v1cs.LGcs.AIcs.CL

TL;DR

Aggressive parallel decoding can make early dLLM predictions unreliable, while existing methods do not enforce consistency with later states. CForce distills later predictions along self-rollout trajectories, improving speed-quality trade-offs on edit-capable and non-edit LLaDA models, especially at high parallelism.

  • Problem

    Aggressive parallel decoding makes early dLLM predictions unreliable, and existing methods do not train early-to-late consistency along the model’s decoding trajectory.

  • Method

    CForce uses self-rollout trajectories to align earlier masked-state predictions with later states through Confidence Adaptive KL Divergence, a CE anchor, and staged reveal difficulty.

  • Results

    On edit-capable LLaDA2.1-mini, CForce raises average TPF from 6.94 to 9.08 and average accuracy from 85.57 to 86.41, while improving high-parallelism behavior on non-edit LLaDA2.0-mini.

  • Takeaways & Limitations

    Trajectory-based consistency forcing supports higher-parallelism decoding while improving the speed-quality frontier across edit-capable and conventional dLLMs.

  • Takeaways & Limitations

    CForce relies on trajectories cached before optimization, which may diverge from inference trajectories of the updated model.

Abstract

from arXiv · show

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.

1 Introduction

CForce addresses unreliable early-stage predictions in aggressively parallel dLLM decoding by aligning them with later predictions using self-rollout trajectories and confidence-adaptive distillation. It improves speed–quality trade-offs in both non-edit and edit-capable LLaDA models, with later token-to-token refinements supervising earlier masked-state predictions.

  • Motivation: Aggressive parallelism makes early dLLM predictions unreliable, allowing errors to propagate into later denoising stages.Faster decoding therefore requires improving early-stage reliability, not only reducing denoising steps.
  • Edit-capable decoding: In edit-capable dLLMs, later token-to-token refinements provide richer supervision for earlier mask-to-token predictions before token commitment.Later denoising states can revise earlier drafts through token-to-token corrections.
  • Method: CForce trains on self-rollout trajectories to align early-stage mask predictions with later-stage predictions and improve training–inference alignment.The method partitions each pre-collected trajectory into stages.
  • Theory: CForce’s consistency objective theoretically bounds early-stage prediction error by adjacent-stage distributional drift plus reveal-boundary token error.These terms correspond to the CAD and CE components of the objective.
  • Results: 6.94 to 9.08 average TPF and 85.57 to 86.41 average accuracy are achieved on edit-capable LLaDA2.1-mini with CForce.On non-edit LLaDA2.0-mini, average TPF increases from 3.60 to 6.42 while few-step accuracy improves at fixed TPF budgets under aggressive decoding.

2 Related Work

Related work spans masked diffusion language models, decoding and system-level acceleration, and consistency distillation. These directions motivate improving early predictions through later-state supervision and agreement across diffusion trajectories.

  • Masked Diffusion Language Models: Masked diffusion language models formulate text generation as discrete denoising, with recent dLLMs scaling the paradigm.Confidence-based parallel decoding accelerates generation, while edit-capable models permit later denoising states to revise committed tokens.
  • Masked Diffusion Language Models: Edit-capable decoding enables later denoising states to revise committed tokens, providing supervision for early M2T predictions under aggressive parallel decoding.Early predictions are made with limited context and can benefit from later denoising states.
  • Acceleration Methods for dLLMs: dLLM acceleration methods improve token selection or verification at inference time, reduce computation through caching, or train models for few-step generation.These approaches respectively modify inference procedures, system computation, or model training.
  • Consistency Distillation in Diffusion Models: Diffusion distillation compresses slow multi-step generation into faster samplers or models, while consistency models enforce agreement across trajectory states for one-step or few-step generation.Related consistency objectives have also been studied for discrete generative models.

3 Preliminaries

Masked diffusion language models generate text by denoising masked token sequences, with confidence-based decoding revealing multiple positions in parallel. Block diffusion generates blocks sequentially, while edit-capable variants additionally revise committed tokens through later token-to-token refinement.

  • Masked Diffusion Language Models: Masked diffusion language models formulate text generation as denoising in a discrete token space.The forward process independently replaces clean target tokens with [MASK], and the reverse model predicts categorical distributions for masked positions.
  • Masked Diffusion Language Models: The masked diffusion training objective is a weighted cross-entropy that recovers clean tokens from partially observed contexts and enables parallel denoising at inference.The positive ELBO weight is induced by the absorbing diffusion process.
  • Block Diffusion: Block diffusion partitions target sequences into blocks, denoises each block in parallel, and generates blocks sequentially.Current large-scale dLLMs commonly combine this framework with the masked diffusion objective.
  • Confidence-based Parallel Decoding: Confidence-based decoding commits all positions above threshold τ and otherwise selects the most confident remaining position to ensure progress.Each decoding trajectory starts from an all-mask block state and iteratively produces an all-clean state.
  • Edit-capable Diffusion Decoding: Edit-capable dLLMs add token-to-token refinement after masked-to-token commitment, allowing later states to correct earlier committed tokens.These refinements are especially important when the masked-to-token threshold τ is low and provide later-stage supervision for earlier masked-state predictions.

4 Method

Consistency Forcing constructs coarser self-rollout stage pairs and trains early masked states to match later, more informative predictions. Its CAD-plus-CE objective extends to edit corrections, uses a transition curriculum, and serves as a surrogate for an upper bound on early-stage prediction error.

  • Trajectory Construction: CForce downsamples confidence-based self-rollouts into adjacent stages separated by cumulative M2T reveals, producing pairs with more substantial visible-context changes.Native trajectories from qθ0 are staged when cumulative newly revealed tokens reach a fixed size S; the all-mask and final states are retained.
  • Adjacent-Stage Objective: Each pair uses the earlier stage as student input and the later stage as a stop-gradient target, combining CAD on persistent masks with a CE anchor at reveal positions.CAD controls distributional drift on positions remaining masked, while CE stabilizes token commitments immediately before reveal.
  • Confidence Adaptive KL Divergence: CAD mixes forward and confidence-weighted reverse KL, controlling distributional drift while adding mode-seeking pressure only when later-stage predictions are confident.The confidence signal is cj = maxv pj tar; reverse KL is avoided as an unconditional objective because uncertain targets can collapse prematurely.
  • Edit-Capable dLLMs: For edit-capable decoding, CAD supervises masked positions and stage-level T2T edits, while CE extends from newly revealed positions to all student-masked positions.The edit-capable formulation preserves the decoding algorithm while incorporating later visible-token corrections as additional supervision.
  • Curriculum Transition: The curriculum begins with ρ0 = 0.1 of newly revealed positions exposed in the later context and increases exposure over training to teach larger stage transitions.At step u, the selected subset Γi(u) has size ⌊ρu|∆i|⌋, making the initial target state close to the student state before progressing toward fast-decoding transitions.
  • Theoretical Motivation: The theory bounds early-stage prediction error by adjacent-stage distributional drift plus reveal-boundary token error, matching CAD’s forward KL and the CE anchor.The confidence-weighted reverse KL is non-negative, preserving forward-KL alignment while sharpening predictions for confident later-stage targets.

5 Experiments

Experiments evaluate CForce on non-edit and edit-capable LLaDA models across mathematical reasoning and code-generation benchmarks. CForce improves the edit-capable speed-quality trade-off, while non-edit and ablation results expose the effects of parallelism, curriculum learning, target choice, KL form, and CE anchoring.

  • Experimental Setup: Experiments use LLaDA2.0-mini for non-edit decoding and LLaDA2.1-mini for edit-capable decoding, evaluated on GSM8K, MATH500, MBPP, and HumanEval.Score denotes exact-match accuracy on the mathematical benchmarks and pass@1 on the code benchmarks; TPF and AUP measure parallelism and speed-quality performance.
  • Edit-capable dLLMs: On LLaDA2.1-mini, CForce raises average TPF from 6.94 to 9.08 and average score from 85.57 to 86.41, with strongest gains on code benchmarks.The edit-capable setting benefits from later corrections to decoded tokens feeding back into earlier masked-position predictions through the stop-gradient target.
  • Non-edit dLLMs: On LLaDA2.0-mini, CForce reaches 6.42 average TPF versus 3.60 for the base model and 5.07 for CAP, trading lower average score than the base for higher score than CAP.This result places CForce at a higher-parallelism operating point under non-edit decoding.
  • Ablations: Stage size S = 8 provides the best observed ablation balance, with 80.88 average score and 6.42 average TPF; S = 16 reaches 7.97 average TPF but lowers score to 75.05.A small stage size of S = 4 yields 80.62 average score and 5.64 average TPF.
  • Ablations: Removing curriculum learning lowers average score from 80.88 to 79.87 and average TPF from 6.42 to 5.97, while the same-model target beats the frozen-teacher baseline by +1.39 score and +0.51 TPF.The curriculum gradually increases later-state exposure from 10% to 100%, and the evolving same-model target better tracks the student’s decoding distribution.
  • Ablations: CAD increases average TPF from 5.07 with forward KL to 6.42, whereas reverse KL reaches 8.72 average TPF but drops average score to 65.57; removing the CE anchor lowers score from 80.88 to 79.02.The moderate setting λCE = 0.1 is retained because it improves average score while preserving KL-driven parallelism.

6 Conclusion

Consistency Forcing improves the speed-quality frontier of diffusion language models under aggressive threshold decoding by aligning early low-context states with later, more informative states. It uses self-generated staged trajectories, same-model stop-gradient supervision, Confidence Adaptive KL Divergence, a CE anchor, and a reveal-difficulty curriculum to improve early-stage prediction reliability.

  • Consistency Forcing improves the speed-quality frontier of diffusion language models under aggressive threshold decoding.
  • The method constructs staged trajectories from the model’s own decoding path and aligns earlier low-context states with later, more informative states.Alignment uses same-model stop-gradient supervision.
  • Confidence Adaptive KL Divergence, a CE anchor, and a curriculum over reveal difficulty improve the reliability of early-stage predictions.

Limitations · A Staged Trajectory Construction Algorithm

CForce currently constructs training data from fixed pretrained-model self-rollout trajectories, which may diverge from trajectories of the updated model. Algorithm 1 stages native threshold-decoding traces into a trajectory dataset by recording masked states and edit information at fixed reveal intervals.

  • Limitations: CForce relies on trajectories collected before optimization from a pretrained initialization, creating a possible mismatch with inference trajectories after model updates.The cached trajectories may gradually differ from the trajectories actually visited by the updated model.
  • A Staged Trajectory Construction Algorithm: Algorithm 1 constructs staged self-rollout trajectories from native threshold-decoding traces.The procedure produces a trajectory dataset D.
  • A Staged Trajectory Construction Algorithm: The algorithm initializes the native trajectory at the all-mask state, then tracks stage and revealed-position counters during decoding.It sets s̄0 ← s0, n ← 0, and R ← 0 before processing native decoding steps.
  • A Staged Trajectory Construction Algorithm: At each native decoding step, it computes predictions, selects newly revealed positions, optionally records edits, and updates the native state.The newly revealed positions are selected by Eq. (2), while edit-capable decoding records Et by Eq. (5).
  • A Staged Trajectory Construction Algorithm: Whenever the reveal counter reaches stage size S, the algorithm saves the current staged state and resets the stage counters.It saves s̄n+1 ← st+1, increments n, and resets R ← 0.
  • A Staged Trajectory Construction Algorithm: After full decoding, the procedure saves the final native state if needed and appends the complete staged trajectory to the dataset.The trajectory is represented as T̄ = (s̄0, . . . , s̄N).

B Implementation Details

CForce uses full-parameter AdamW fine-tuning with a cosine learning-rate schedule and fixed training hyperparameters. KL distillation truncates target distributions to their top-20 tokens, while training runs on 64 NVIDIA H800 GPUs.

  • Optimization: Models use full-parameter fine-tuning with AdamW, a learning rate of 1.0 × 10−5, and a cosine learning rate schedule.The optimization setup is shared across models.
  • Training Configuration: Training uses stage size 8, global batch size 64, 5 epochs, CE anchor loss weight λCE 0.1, and block size 32.These settings define the training schedule, loss weighting, and sequence blocking configuration.
  • KL Distillation: KL distillation retains and renormalizes the top-20 tokens ranked by the target distribution pj tar before computing KL terms.Top-k truncation reduces softmax and KL computation while focusing learning on semantically meaningful probability mass.
  • Hardware: All training is conducted on 64 NVIDIA H800 GPUs.The implementation uses distributed training across the full GPU allocation.

C Full Derivation of the Theoretical Analysis

The theoretical analysis bounds early-stage prediction error by adjacent-stage distributional drift plus masked-state error before token reveal. Consistency Forcing targets these terms through forward-KL alignment, CE anchoring, and confidence-weighted reverse-KL sharpening.

  • Derivation: For a token revealed at stage r_j, the triangle inequality decomposes earlier-stage error into adjacent-stage distributional drift and masked-state error before reveal.The revealed position satisfies j ∈ M̄_{r_j−1} and j ∉ M̄_{r_j}, with final token y_j = s̄_j.
  • Derivation: Pinsker’s inequality bounds each adjacent-stage total-variation term using forward KL, while the reveal-boundary term follows from 1 − x ≤ −log x.Substituting these bounds into the triangle-inequality decomposition produces the stated early-stage prediction-error bound.
  • Connection to CForce: Consistency Forcing’s forward KL targets adjacent-stage drift, and its CE anchor captures reveal-boundary error while supplying token-level supervision at commitment boundaries.The confidence-weighted reverse KL additionally encourages mode-seeking predictions when later-stage targets are reliable.
  • Connection to CForce: Optimizing L_CAD preserves forward-KL alignment while sharpening predictions toward the cleaner state’s dominant mode, thereby reducing a tractable surrogate of early-stage prediction error.The objective aligns adjacent-stage distributions and anchors masked-token predictions along the self-rollout trajectory.

D AUP Evaluation Details · E Wall-clock Throughput

The paper constructs AUP from filtered, weighted score–parallelism frontiers using exact decoding configurations, then reports wall-clock throughput under fixed two-GPU conditions. CForce improves throughput over base models in both non-edit and edit-capable settings.

  • D AUP Evaluation Details: AUP uses a sorted frontier of operating points, each pairing TPF with the benchmark task score.The frontier contains multiple score–parallelism points rather than only the single Score/TPF pair shown in Table 1.
  • D AUP Evaluation Details: AUP retains points within five score points of the lowest-TPF point and weights accuracy-preserving parallelism more favorably.The weighting function is W(y) = min(exp(−3(1−y/ymax)), 1), with ymax defined per benchmark among compared methods.
  • D AUP Evaluation Details: Table 5 provides the decoding configurations and benchmark-level Score/TPF pairs used to build LLaDA-family AUP frontiers.LLaDA2.0 models use threshold-decoding points τ = 0.95 and τ = 0.9; LLaDA2.1 models use official quality and speed modes.
  • E Wall-clock Throughput: Table 6 reports wall-clock throughput for GSM8K and HumanEval, with bold indicating the best throughput within each edit-capability group.The table covers both non-edit and edit-capable comparisons.
  • E Wall-clock Throughput: All throughput measurements use two NVIDIA H20 GPUs with tensor parallelism TP = 2 and batch size 1.TPS denotes generated tokens per second.
  • E Wall-clock Throughput: 857.96 average TPS is achieved by non-edit CForce, improving over the base model by 42.82% and over CAP by 7.06%.In edit-capable LLaDA2.1, CForce reaches 1292.19 average TPS, improving over the base model by 22.84%.

F Full Ablation Results … G.3 Example 3 Prompt:

The appendix reports full benchmark-level ablations for LLaDA2.0-mini-CForce and qualitative case studies contrasting CForce with the baseline on arithmetic, profit, and programming tasks. The ablations evaluate stage size, curriculum learning, target type, KL divergence form, and CE anchor weight using Score and TPF.

  • F Full Ablation Results: Tables 7–9 expand the compact ablations into benchmark-level results, reporting Score and TPF for each benchmark group under threshold decoding.The evaluation uses τ = 0.9 and maximum generation length 4096.
  • G Case Study: The case study presents qualitative generation comparisons between LLaDA2.0-mini-CForce (Ours) and LLaDA2.0-mini (Baseline).The examples cover arithmetic reasoning, profit calculation, and a programming task.
  • F Full Ablation Results: The full ablations vary stage size, curriculum learning, target type, KL Divergence Form, and CE Anchor Weight for LLaDA2.0-mini-CForce.The curriculum and target-type table distinguishes frozen-teacher and no-curriculum variants.
  • G.2 Example 2 (Reference: 160): For Tara’s canvas-bag problem, the CForce example concludes that she earned a profit of $160.Its shown reasoning computes 8 packs × 5 bags per pack = 40 canvas bags, with each bag costing $4 and selling for $8.
  • G.3 Example 3 Prompt:: The programming prompt asks for median_numbers to return the median of three numbers and satisfy three specified assertions.The requested code must be enclosed in Python delimiters.
  • G.3 Example 3 Prompt:: Both the CForce and baseline programming examples implement the median by ordering three inputs and returning the middle element as a float.The CForce version uses sorted([a, b, c]), while the baseline constructs a list and returns float(numbers[1]).
Loading 2608.13925v1…