PaperlayerView exact revision on arXiv

Source-linked AI summary

d$^2$Cache: Accelerating Diffusion-Based LLMs via Dual Adaptive Caching

Yuchu Jiang, Yue Cai, Xiangzhong Luo, Jiale Fu, Jiarui Wang, Chonghan Liu, Xu Yang

arXiv:2509.23094v2cs.CL

TL;DR

dLLMs have inferior inference efficiency because bidirectional attention prevents direct reuse of standard ARM KV caches. d2Cache is a training-free, two-stage fine-grained approximate cache that adaptively updates selected token states and reuses the rest, achieving substantial speedups and consistent quality improvements across LLaDA and Dream.

  • Problem

    Bidirectional attention makes dLLM inference inefficient because dLLMs cannot directly benefit from the standard KV cache used by autoregressive models.

  • Method

    d2Cache uses two-stage fine-grained token selection to adaptively update selected KV states while caching the remaining states for reuse.

  • Results

    d2Cache achieves substantial inference speedups and consistent generation-quality improvements across representative dLLMs and datasets.

  • Takeaways & Limitations

    The framework provides a training-free approximate KV-cache approach for accelerating dLLM inference while retaining generation quality.

  • Takeaways & Limitations

    The paper reports that d2Cache maintains comparable performance rather than uniformly improving it across all evaluated settings.

Abstract

from arXiv · show

Diffusion-based large language models (dLLMs), despite their promising performance, still suffer from inferior inference efficiency. This is because dLLMs rely on bidirectional attention and cannot directly benefit from the standard key-value (KV) cache as autoregressive models (ARMs) do. To tackle this issue, we introduce \textit{Dual aDaptive Cache} (d$^2$Cache), which is a training-free approximate KV cache framework for accelerating dLLM inference. d$^2$Cache features a two-stage fine-grained selection strategy to identify tokens and adaptively update their KV states at each decoding step, while caching the KV states of the remaining tokens for reuse. Furthermore, d$^2$Cache naturally offers a more reliable decoding alternative, which can enable quasi left-to-right generation and mitigate premature overconfidence in tokens at the end of the sequence. Extensive experimental results on two representative dLLMs (\ie, LLaDA and Dream) demonstrate that d$^2$Cache not only achieves substantial inference speedups, but also yields consistent improvements in generation quality. The code is available at https://github.com/Kamichanw/d2Cache.

1 INTRODUCTION

dLLMs cannot directly reuse standard ARM KV caches because bidirectional attention makes token updates affect the full sequence. d2Cache addresses this with fine-grained adaptive caching, delivering faster inference and improved generation quality.

  • 1 INTRODUCTION: Bidirectional attention makes dLLM KV states difficult to reuse because changing one masked token alters the context seen by other tokens.This prevents dLLMs from directly benefiting from the standard KV cache used by ARMs.
  • 1 INTRODUCTION: Coarse-grained approximate caches divide tokens into static and dynamic segments, limiting flexibility and potentially updating reusable states or reusing states that should change.These methods may also require complicated tuning.
  • 1 INTRODUCTION: KV dynamics and attention patterns motivate separate fine-grained treatment of tokens, including phase-aware masked-token updates and attention-aware handling of prompt and decoded tokens.The paper reports a three-phase masked-token pattern and uneven attention distribution.
  • 1 INTRODUCTION: d2Cache adaptively selects tokens and updates their KV states at each decoding step while reusing the remaining cached states.Its two-stage selection strategy is training-free and operates at token level.
  • 1 INTRODUCTION: d2Cache achieves substantial inference speedups while consistently improving generation quality across representative dLLMs and datasets.The introduction presents this as the main experimental outcome.

2 RELATED WORK

Diffusion-based language models generate text through iterative denoising over masked sequences rather than autoregressive next-token generation. Approximate KV-cache methods exploit similarity across decoding steps, but existing approaches use coarse sequence segmentation.

  • 2 RELATED WORK: dLLMs iteratively denoise masked sequences, enabling bidirectional context modeling unlike sequential ARM generation.This bidirectional process prevents direct use of the standard ARM KV cache.
  • 2 RELATED WORK: Approximate dLLM KV-cache methods reuse states that remain similar across adjacent decoding steps by partitioning tokens into static and dynamic segments.The cited related methods build on temporal redundancy in KV states.

3 PRELIMINARIES

dLLMs begin with masked response positions and progressively replace selected masks over discrete decoding steps. The paper analyzes token-state dynamics, decoding order, and attention concentration to motivate selective KV updates.

  • 3 PRELIMINARIES: At each decoding step, the dLLM predicts masked positions, assigns confidence scores, and schedules selected positions for replacement until no masks remain.Scheduling may use random sampling or the highest-confidence masked positions.
  • 3 PRELIMINARIES: Masked-token KV states exhibit gradual change early, rapid change immediately before decoding, and stability after decoding.The reported PCA trajectory uses LLaDA-8B-Instruct on GSM8K.
  • 3 PRELIMINARIES: 90% of subsequently decoded tokens lie within distance 10 of the most recently decoded token, indicating localized decoding order.This supports estimating imminent decoding from the density of nearby decoded tokens.
  • 3 PRELIMINARIES: Attention rollout shows that queries consistently attend to a small subset of key positions rather than distributing attention uniformly.This motivates differentiated treatment of token KV states.
  • 3 PRELIMINARIES: Attention allocations remain highly similar across adjacent decoding steps, allowing current-step allocations to approximate the next step’s allocation.KV updates can therefore be restricted to tokens receiving higher attention.

4 D2CACHE: DUAL ADAPTIVE CACHE

d2Cache uses two adaptive selection stages to update only selected token KV states while reusing the remaining cached states, with certainty-guided masked-token selection and attention-aware selection for other tokens.

  • d2Cache adaptively identifies tokens for KV updates at each decoding step while caching the remaining tokens for reuse.The framework is training-free and targets approximate KV caching in dLLMs.
  • 4.1 STAGE 1: CERTAINTY PRIOR-GUIDED SELECTION: Masked-token selection combines prediction confidence with a certainty prior based on the local density of known tokens.The certainty prior reflects structural certainty and favors masked tokens near known tokens.
  • 4.1 STAGE 1: CERTAINTY PRIOR-GUIDED SELECTION: The position-aware certainty density weights nearer known tokens more strongly and selects masked tokens using top-k calibrated scores.The Gaussian weighting scale σ controls the positional scope of the density calculation.
  • 4.1 STAGE 1: CERTAINTY PRIOR-GUIDED SELECTION: Certainty prior-guided decoding can preserve a quasi left-to-right order and mitigate premature overconfidence in sequence termination.The paper reports more reliable decoding performance than default confidence-based decoding.
  • 4.2 STAGE 2: ATTENTION-AWARE SELECTION: For remaining tokens, attention rollout aggregates layerwise attention into global influence scores and selects the smallest token set exceeding threshold p.The selected set U is used so the next decoding step receives a subset rather than the full sequence.

5 EXPERIMENTS

Experiments evaluate d2Cache on two dLLMs across six benchmarks against Vanilla and approximate-cache baselines. The method achieves strong efficiency and quality results, while ablations examine decoding order, update phases, and hyperparameter sensitivity.

  • 5.2 MAIN RESULTS: 3.5× average speedup over Vanilla across all models and datasets, while maintaining comparable or better average scores across six datasets.On Dream-Inst GSM8K, throughput increases from 2.62 to 12.25 tokens per second, a 4.7× speedup.
  • 5.2 MAIN RESULTS: d2Cache achieves the best average throughput, latency, and score across benchmarks, outperforming Vanilla, dLLM-Cache, and Fast-dLLM overall.Compared with Fast-dLLM on Dream-Inst, it provides 1.3× higher inference speedup while maintaining +1.7% average accuracy across six datasets.
  • 5.2 MAIN RESULTS: Certainty prior-guided decoding produces a more natural, controllable left-to-right order than NAR decoding, with smaller σ yielding behavior closer to autoregressive decoding.The comparison uses LLaDA-Inst on 64 randomly sampled GSM8K examples and σ values of 10, 40, and 80.
  • 5.3 ABLATIONS AND ANALYSIS: Updating masked-token KV states only during the rapid-change phase improves throughput over updating during both gradual- and rapid-change phases, with comparable or better scores.The ablation on Dream-Inst identifies full updates during the gradual-change phase as computationally redundant.
  • 5.3 ABLATIONS AND ANALYSIS: Performance improves as k increases but saturates or slightly declines beyond k = 32, making k = 32 the most stable setting across p and sequence lengths.The sensitivity study varies generation lengths of 256, 512, and 1024 on Dream-Inst and HumanEval; p does not monotonically improve performance.

6 CONCLUSION

d2Cache is a training-free approximate KV cache that uses fine-grained token selection to reduce redundant dLLM computation while preserving generation quality. The paper also positions quasi-left-to-right decoding as a way to retain bidirectional context while avoiding premature decoding errors.

  • d2Cache adaptively selects tokens and updates their KV states at each decoding step, reusing the remaining states to reduce redundant computation.Its two-stage selection strategy operates at token level rather than applying fixed update schedules to coarse segments.
  • KV states of masked tokens change most rapidly shortly before decoding and become stable afterward, while prompt and decoded tokens have smaller consecutive-step dynamics.These observations motivate different caching decisions for masked, prompt, and decoded tokens.
  • Attention is concentrated on a small subset of prompt and decoded tokens, supporting reuse of low-attention tokens' KV states.This finding supplies the second basis for d2Cache's dual adaptive caching design.
  • d2Cache differs from Fast-dLLM and dLLM-Cache by selecting individual tokens rather than refreshing fixed blocks or prompt/response segments.The compared methods use block-wise or segment-level update schedules, whereas d2Cache updates selected tokens per decoding step.
  • Quasi-left-to-right decoding can preserve sequential reasoning while retaining bidirectional attention and parallel decoding flexibility.dLLMs can defer uncertain positions and use later context before committing, unlike autoregressive models that must choose the next position immediately.

D.3 LIMITATIONS AND FUTURE WORK

The paper identifies three future-work boundaries: scaling d2Cache to larger dLLMs, reducing attention-rollout cost without harming reasoning, and replacing its Gaussian contextual-contribution score with more adaptive alternatives.

  • d2Cache is evaluated on LLaDA-8B and Dream-7B because they are the only publicly available dense dLLMs considered, leaving larger-scale evaluation for future work.The authors expect larger depth, width, and context length to make bidirectional decoding more costly to maintain.
  • Reducing attention-rollout depth from 28 to 5 slightly improves speed but noticeably degrades GSM8K and Math-500 performance, while code-generation losses are minimal.The comparison motivates more efficient rollout approximations that still identify key tokens.
  • The current Gaussian contextual-contribution score performs well empirically, but more context-adaptive scoring functions may improve performance.Alternative scoring functions remain an explicitly identified future direction.

E.1 EXPERIMENTAL RESULTS ON THE BASE VARIANTS

On the Base variants, d2Cache consistently outperforms other approximate KV cache methods in average inference efficiency and accuracy across six datasets.

  • Across six datasets, d2Cache consistently outperforms other approximate KV cache methods in both average inference efficiency and accuracy on LLaDA-Base and Dream-Base.The evaluation extends beyond the Instruct variants to the corresponding Base models.

E.2 EXPERIMENTAL RESULTS UNDER PARALLEL DECODING SETTINGS

Under parallel decoding, d2Cache remains broadly applicable across the evaluated instruction-tuned models. It reaches substantial acceleration while maintaining performance comparable to the other baselines.

  • Up to 48.3× acceleration over the single-token-per-step baseline is achieved under parallel decoding while performance remains comparable to other baselines.The evaluation uses a threshold of 0.9 and includes LLaDA-Inst and Dream-Inst with semi-AR parallel decoding.

E.3 EXPERIMENTAL RESULTS UNDER LONG-CONTEXT SETTINGS

Under long-context settings, d2Cache maintains substantial speedups without performance loss, while other methods degrade as context length increases.

  • d2Cache maintains substantial speedups without performance loss when generation length reaches 1024.The evaluation uses Dream-Inst under a longer generation length of 1024.
  • Other methods experience severely degraded acceleration as context length increases.The passage attributes this degradation to their coarse-grained nature.
  • Fine-grained two-stage token selection supports d2Cache’s long-context acceleration.The passage identifies this selection strategy as the reason d2Cache retains speedups without performance loss.

E.4 MORE VISUALIZATION RESULTS ON ATTENTION ROLLOUT

Additional attention-rollout visualizations on LLaDA-Inst and GSM8K use the same sample and configuration as Figure 3 and align with the paper’s earlier findings.

  • Figure 8 presents an additional attention-rollout visualization for LLaDA-Inst on GSM8K.
  • The observed attention pattern aligns with findings reported in Section 3.3.

E.5 MORE VISUALIZATION RESULTS ON KV STATE DYNAMICS

Additional KV-state visualizations show that key and value states for a masked token evolve similarly, motivating selective updates and distance-aware certainty decay.

  • Key and value states of the same masked token closely match in trajectory shape, magnitude, and gradual–rapid–stable dynamics.The visualization concerns the 91st masked token during decoding.
  • Updating key and value states only during the rapid-change phase is sufficient for the masked token dynamics described.
  • Decoding a nearby masked token supplies additional contextual information that substantially alters the other token’s representation.The smaller the token distance, the more tightly constrained the context becomes.
  • These dynamics motivate distance-aware decay in the certainty density defined by Equation (3).
  • Figures 8 and 9 provide additional visualizations using samples and configurations corresponding to earlier figures.
Loading 2509.23094v2…