PaperlayerView exact revision on arXiv

Source-linked AI summary

CAI-DLLM: Convergence Aware Inference for Diffusion Language Models

Farhana Amin, Sabiha Afroz, Dimitrios S. Nikolopoulos

arXiv:2608.22646v1cs.AI

TL;DR

Diffusion language models reduce sequential dependence through parallel token updates but still incur repeated full-sequence denoising, including for stabilized tokens. CAI-DLLM uses first-step confidence to adapt token budgets, commitment schedules, and block thresholds without retraining or extra predictors. Across evaluated benchmarks, it delivers substantial speedups and energy reductions, with accuracy trade-offs on some tasks.

  • Problem

    Diffusion language models still repeatedly recompute stabilized tokens during denoising, while existing methods use uniform computation despite uneven token convergence.

  • Method

    CAI-DLLM is a training-free inference method that uses first-step confidence to guide adaptive parallel decoding, block-level schedules, token budgets, and low-yield early exit.

  • Results

    Across seven benchmarks on LLaDA-8B-Instruct and Dream-7B-Instruct, CAI-DLLM achieves speedups up to 44.8× and energy reductions up to 95.3%.

  • Takeaways & Limitations

    First-step confidence can substantially reduce redundant denoising computation while leaving the original model unchanged.

  • Takeaways & Limitations

    The model-specific confidence-gating choice is selected using the full GSM8K test set rather than a held-out validation set, and speedups can involve accuracy trade-offs.

Abstract

from arXiv · show

Diffusion language models can generate many tokens in parallel, but they still require repeated denoising steps during inference. This makes generation costly, especially when the model continues to recompute tokens that are already stable. To address these limitations, we propose CAI-DLLM, a training-free inference method that uses first-step confidence to guide denoising and reduce inference time. Specifically, CAI-DLLM commits easy tokens earlier, allocates more denoising steps to harder tokens, and adjusts decoding schedules across output blocks. As it relies only on first-step confidence signals, it does not require retraining, extra predictors, or weight updates. We evaluate CAI-DLLM on LLaDA-8B-Instruct and Dream-7B-Instruct across math, code, reasoning, commonsense, and long-context tasks. CAI-DLLM achieves up to 18.2x wall clock inference speedup on LLaDA GSM8K while improving accuracy from 76.27% to 77.41%, and up to 13.1x speedup on Dream HumanEval while achieving higher pass@1 than no-cache inference, 48.17% compared with 46.95%. On harder reasoning tasks, speedups reach 44.8x, with a largest accuracy drop of 4.4 points, while energy consumption is reduced by up to 95.3%.

1 Introduction

Diffusion language models parallelize token updates but still repeat full-sequence denoising, including for stabilized tokens. CAI-DLLM uses free first-step confidence to vary token and block computation, committing easy tokens earlier and refining harder ones longer.

  • Diffusion language models update many masked positions in parallel, but every denoising step still processes the full sequence.This retains substantial inference cost even after some tokens stabilize.
  • Existing efficiency methods still apply the same denoising budget to every token, wasting steps on easy tokens and underserving hard tokens.
  • First-step confidence provides a free estimate of token difficulty because it is produced during the required first forward pass.High-confidence tokens tend to stabilize early, whereas low-confidence tokens require more refinement.
  • CAI-DLLM is a training-free method that assigns token-specific step budgets, block-level commit schedules, and low-yield early exits.It uses first-step confidence to commit easy tokens early and allocate more steps to hard tokens.
  • CAI-DLLM introduces block-level adaptive threshold scheduling, with thresholds decreasing across output blocks as more context becomes available.

2 Background and Motivation

Masked diffusion decoding updates multiple positions iteratively, but fixed policies overlook uneven convergence across layers, blocks, and tokens. The paper motivates first-step confidence as an inexpensive signal for adaptive token commitment and denoising allocation.

  • Masked diffusion models progressively replace masked positions and can update multiple tokens in one forward pass, unlike sequential autoregressive generation.
  • A fixed decoding policy is inefficient because denoising does not progress uniformly across the sequence.The analysis uses four 64-token blocks generated by LLaDA-8B-Instruct on GSM8K.
  • Early transformer layers change substantially less than deeper layers during denoising.
  • First-step confidence separates tokens by convergence difficulty, with high-confidence tokens usually stabilizing early and low-confidence tokens often needing many more steps.A token is defined as stable when its top prediction remains unchanged for three consecutive steps.
  • The method uses first-step confidence as a control signal because it is already computed during standard decoding and requires no additional model or forward pass.

3 Related Work

Prior work accelerates diffusion decoding through caching, hybrid decoding, speculative verification, or layer-level token skipping. CAI-DLLM instead targets the number of denoising steps assigned to tokens while preserving the original model and decoding process.

  • Autoregressive KV caching reduces repeated attention computation but does not remove sequential generation, motivating parallel diffusion decoding.
  • Diffusion language models can match autoregressive models on reasoning, math, and code, but uniform fixed-step treatment remains their core inefficiency.
  • Diffusion caching methods reduce repeated computation through cache reuse or cache sparsity, including stabilization-aware and prompt-response-specific policies.
  • Figure 1 organizes uneven denoising work across layers, output blocks, and token difficulty using first-step confidence thresholds.
  • Hybrid and speculative approaches modify decoding structure or add draft-and-verify computation, whereas CAI-DLLM keeps the original model and decoding process unchanged.
  • Layer-level token skipping reduces computation within denoising steps, while CAI-DLLM reduces the number of denoising steps assigned to each token.

4 Method

CAI-DLLM uses first-step confidence to control masked diffusion decoding, assigning computation according to token difficulty and adapting schedules across blocks. It also detects low-yield late denoising to exit early.

  • 4.1 Diffusion Decoding Setup: CAI-DLLM uses the required first forward pass to estimate token difficulty without extra model calls.High-confidence tokens are treated as easier, while low-confidence tokens are treated as harder.
  • 4.2 First-Step Confidence: High first-step confidence tokens stabilize within 5 steps, whereas low-confidence tokens can continue changing for 30 steps or more.This pattern is reported as consistent across prompts and tasks.
  • 4.3 Adaptive Scheduling: The method assigns per-token budgets, commits easy tokens early, gives hard tokens more refinement, and supports block-adaptive threshold schedules.Later blocks use lower end thresholds because they converge more easily with more committed context.
  • 4.4 Token Budgets and Position Aware Thresholds: A grinding-phase detector exits blocks after several consecutive low-yield steps, while force-commit behavior prevents very uncertain tokens from being forced.The detector tracks newly committed tokens, with a patience window and yield threshold.
  • 4.4 Token Budgets and Position Aware Thresholds: Token positions affect computation: positions p < 32 commit 37% faster on average than positions p ≥48 on LLaDA GSM8K.This motivates position-aware thresholds and budgets.

5 Implementation Details

The evaluation integrates CAI-DLLM with ES-dLLM’s codebase and KV-cache backend, using fixed task settings and common baselines for comparison. Speedup is measured by wall-clock time relative to no-cache decoding.

  • Setup: CAI-DLLM is evaluated on LLaDA-8B-Instruct and Dream-7B-Instruct using a single NVIDIA H200 SXM5 GPU.Generation lengths are 256 tokens for reasoning and commonsense tasks and 512 tokens for code and long-context tasks.
  • Decoding settings: The experiments use block length 64 for most tasks and 32 for LongBench, with all methods evaluated under the same generation settings.The shared settings are intended to support fair comparison.
  • Decoding settings: No-cache decoding uses 64 denoising steps per block without KV caching or parallel decoding.This establishes the reference decoding configuration.
  • Baselines and metrics: Baselines include no-cache decoding, DualCache from Fast-dLLM, and ES-dLLM.Metrics are exact match for several reasoning tasks, pass@1 for code tasks, and F1/ROUGE-L for LongBench.
  • Baselines and metrics: Speedup is defined as wall-clock time relative to no-cache decoding, while energy is computed as average GPU power multiplied by wall-clock generation time.Throughput is measured in output tokens per second after one warm-up pass.

6 Results

Across seven benchmarks on Dream-7B and LLaDA-8B, CAI-DLLM consistently improves throughput, with larger gains on longer generations, while trading efficiency against accuracy on harder tasks.

  • Main results: 44.8× speedup is reached on LLaDA BBH, the largest reported gain among longer-generation tasks.Other long-task gains include 28.5× on LLaDA MBPP, 23.0× on Dream BBH, 20.0× on Dream MBPP, and 19.0× on Dream LongBench.
  • Main results: 13.1× speedup accompanies Dream HumanEval pass@1 improvement from 46.95% to 48.17%, treated as accuracy parity given ±3.76 standard error.
  • Main results: 44.8× speedup on LLaDA BBH accompanies an accuracy decrease from 56.75 to 52.33, illustrating the efficiency-quality trade-off on harder reasoning tasks.Dream GSM8K similarly drops from 79.68 to 77.26 while reaching 18.7× speedup.
  • Main results: 19.04× speedup on Dream LongBench yields an average score of 23.11 versus 24.63 for no-cache decoding.CAI-DLLM provides more than twice the throughput of cached methods in this setting, although cached methods slightly reduce average score.
  • Main results: 95.3% energy reduction on LLaDA GSM8K lowers consumption from 3559 Wh to 169 Wh.
  • Ablation study: 614.6 TPS with pass@1 of 46.34% is obtained for Dream when confidence gating is disabled, while LLaDA reaches 435.3 TPS with pass@1 of 35.37% using token budgets.For LLaDA, APD supplies most throughput gain; for Dream, disabling gating recovers accuracy relative to gated budgets.
  • Ablation study: 39–56% denoising-step reduction occurs on GSM8K and 38–47% on HumanEval relative to a fixed budget.The ablation evaluates CAI-DLLM components on HumanEval for both models.

7 Conclusion

CAI-DLLM is a training-free diffusion-language-model inference method that uses first-step confidence to reduce redundant denoising computation. Across seven benchmarks and two models, it delivers speedups up to 44.8× and energy reductions up to 95.3% while leaving the original model unchanged.

  • CAI-DLLM uses first-step confidence to guide token commitment, block schedules, and token budgets without extra models, retraining, or weight updates.
  • 44.8× speedup and 95.3% energy reduction are achieved across seven benchmarks on LLaDA-8B-Instruct and Dream-7B-Instruct.The largest gains occur on longer generation tasks.

Limitations

The study reports limitations involving model-specific gating selection, accuracy trade-offs, hardware dependence, and reduced evaluation sample sizes.

  • Gating choices are selected from full-test-set GSM8K ablations rather than a held-out validation set, limiting confidence in deployment choices for new models.The authors recommend selecting gating on a small held-out validation subset of a representative task.
  • Accuracy can decline alongside speed gains, including drops from 56.75 to 52.33 on BBH and from 40.00 to 37.00 on MBPP for LLaDA-8B.
  • Results may differ on other hardware because all experiments use a single NVIDIA H200 GPU.Some evaluations also use reduced sample sizes, potentially increasing variance.

Ethical Considerations

The evaluation uses controlled settings and reports substantial efficiency gains, while documenting model-specific tuning, benchmark scope, and accuracy trade-offs. The method preserves unchanged model weights and reduces inference compute and energy.

  • Ethical Considerations: CAI-DLLM is training-free and leaves model capabilities, weights, and decoding structure unchanged.It targets inference cost and energy for existing diffusion language model deployments.
  • Evaluation controls: All compared methods use identical generation settings, evaluation splits, and hardware/software controls for fair comparison.The setup fixes generation length, block length, batch size, temperature, and evaluation split for each benchmark.
  • Model-specific configuration: Gating is enabled for LLaDA but disabled for Dream because Dream confidence scores are more uniform and forced commitment is counterproductive.The gating choice is paired with adaptive threshold and block schedules in both cases.
  • Benchmark scope: The evaluation uses benchmark-specific metrics and selected subsets for MathQA and LongBench to control GPU cost.LongBench is a selected-task evaluation on Dream-7B, with inputs truncated to 4,000 tokens.
  • Accuracy trade-offs: The largest reported accuracy difference is a 4.4-point drop on LLaDA BBH, with the gap exceeding estimated sampling uncertainty.The authors relate this result to early commitment preserving incorrect intermediate choices in multi-step reasoning.
  • Efficiency: 95.0% less inference compute is reported for CAI-DLLM than no-cache decoding across the full evaluation.CAI-DLLM uses 3.9 GPU-hours versus 77.4 hours for no-cache, saving 73.5 GPU-hours.

C Additional Background and Motivation Analysis

The appendix analyzes uneven denoising behavior across layers, schedules, hardware, and memory. These observations support adaptive computation while showing that CAI-DLLM adds little bookkeeping memory for LLaDA.

  • Attention sparsity: Deep layers become more selective during denoising, attending to fewer positions with higher top-10 attention mass.Shallow layers attend broadly, whereas later layers concentrate on fewer positions.
  • Attention sparsity: A uniform compute policy is wasteful because early layers process tokens broadly while deep layers focus selectively.The observed layer variation motivates layer-aware or token-aware inference strategies.
  • APD schedule sensitivity: The selected cosine confidence schedule provides the best accuracy-speed trade-off among tested APD configurations.Fixed thresholds and overly aggressive schedules can reduce accuracy.
  • Hardware behavior: At batch size 8, compute utilization exceeds 150–350% of baseline while bandwidth utilization remains below 16%.This identifies compute, rather than memory bandwidth, as the hardware bottleneck.
  • Memory usage: CAI-DLLM uses 32,787 MB on LLaDA, only 480 MB or 1.5% more than DualCache.The added O(L) confidence, budget, and threshold data are negligible relative to model weights and KV cache.
  • Memory usage: Dream uses 36,533 MB with CAI-DLLM, 5,822 MB more than DualCache but 7,500 MB below no-cache decoding.The larger memory use comes from more aggressive parallel output buffering rather than controller bookkeeping.

D.2 Detailed LongBench Results

On Dream-7B LongBench, CAI-DLLM more than doubles throughput while incurring an average quality drop comparable to other caching methods. Task sensitivity varies, with QMSum improving and MultifieldQA-en declining more.

  • Overall LongBench results: 19.04× speedup raises average Dream-7B LongBench throughput from 2.7 to 51.4 tokens per second.The average score decreases from 24.63 to 23.11, a 1.52-point drop.
  • Overall LongBench results: CAI-DLLM provides more than double the throughput of DualCache and ES-dLLM while showing similar average quality degradation.The average score drop is 1.52 points for CAI-DLLM, compared with 1.10 for DualCache and 0.94 for ES-dLLM.
  • Task-level variation: QMSum slightly improves over both cached baselines, with CAI-DLLM scoring 14.66 versus 14.59 and 14.38.The result suggests confidence-aware commitment need not reduce summarization quality.
  • Task-level variation: MultifieldQA-en declines to 43.72 versus 47.95 for ES-dLLM, indicating greater sensitivity to early token commitment.LongBench uses F1 for MultifieldQA-en and NarrativeQA, and ROUGE-L for GovReport and QMSum.
  • Energy: Energy falls from 3559.1 Wh to 168.7 Wh on LLaDA GSM8K, a 95.3% reduction.The measurements use GPU power and wall-clock generation time, excluding CPU, DRAM, and system-level power.

E Throughput Benchmark Results

CAI-DLLM achieves the highest throughput across tested generation lengths and both models, with larger gains for longer outputs. The benchmark uses tokens-per-second speedup over no-cache decoding under a single H200 setup.

  • Evaluation setup: The evaluation measures wall-clock throughput over generation lengths L ∈ {64, 128, 256, 512} on both instruction-tuned diffusion models.All experiments use one NVIDIA H200 GPU under the same hardware and software setup.
  • Throughput across lengths: CAI-DLLM achieves the highest throughput in every tested setting across LLaDA-Instruct and Dream-Instruct.Speedup is computed as the ratio of a method’s TPS to the no-cache TPS.
  • Throughput across lengths: On LLaDA-Instruct, speedup increases from 2.07× at L = 64 to 4.59× at L = 512.Longer generation lengths provide more room for adaptive denoising-step reduction.
  • Throughput across lengths: On Dream-Instruct at L = 512, CAI-DLLM reaches 15.29× speedup over no-cache, or 692.6 versus 45.3 TPS.DualCache and ES-dLLM do not improve throughput at short lengths and can slightly reduce it.
Loading 2608.22646v1…