PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashDLM: Accelerating Diffusion Language Model Inference via Efficient KV Caching and Guided Diffusion

Zhanqiu Hu, Jian Meng, Yash Akhauri, Mohamed S. Abdelfattah, Jae-sun Seo, Zhiru Zhang, Udit Gupta

arXiv:2505.21467v2cs.CL

TL;DR

Diffusion language models provide parallel, bidirectional generation but face high inference cost, latency, and long-context challenges. The paper introduces training-free FreeCache and Guided Diffusion to reduce KV computation and denoising overhead, achieving substantial speedups with negligible accuracy degradation. These techniques bring DLM latency closer to, and sometimes beyond, similarly sized autoregressive models while supporting broader long-context use.

  • Problem

    Diffusion language models have comparable quality to autoregressive models but suffer high computation latency, limited long-context ability, and inefficient inference.

  • Method

    The paper combines FreeCache, which approximates stable KV states without training, with Guided Diffusion, which uses a lightweight autoregressive model to guide token unmasking.

  • Results

    Across reasoning tasks, the method achieves an average 12.14× speedup with negligible accuracy degradation and comparable or better generation speed than similarly sized autoregressive models.

  • Takeaways & Limitations

    The approach enables practical long-context and high-throughput DLM applications while retaining parallel generation and bidirectional conditioning.

  • Takeaways & Limitations

    FreeCache trades accuracy and memory for speed, while Guided Diffusion adds an autoregressive guider whose alignment affects effectiveness.

Abstract

from arXiv · show

Diffusion language models offer parallel token generation and inherent bidirectionality, promising more efficient and powerful sequence modeling compared to autoregressive approaches. However, state-of-the-art diffusion models (e.g., Dream 7B, LLaDA 8B) suffer from slow inference. While they match the quality of similarly sized autoregressive (AR) models (e.g., Qwen2.5 7B, Llama3 8B), their iterative denoising requires multiple full-sequence forward passes, resulting in high computational costs and latency, particularly for long input prompts and long-context scenarios. Furthermore, parallel token generation introduces token incoherence problems, and current sampling heuristics suffer from significant quality drops with decreasing denoising steps. We address these limitations with two training-free techniques. First, we propose FreeCache, a Key-Value (KV) approximation caching technique that reuses stable KV projections across denoising steps, effectively reducing the computational cost of DLM inference. Second, we introduce Guided Diffusion, a training-free method that uses a lightweight pretrained autoregressive model to supervise token unmasking, dramatically reducing the total number of denoising iterations without sacrificing quality. We conduct extensive evaluations on open-source reasoning benchmarks, and our combined methods deliver an average of 12.14x end-to-end speedup across various tasks with negligible accuracy degradation. For the first time, diffusion language models achieve a comparable and even faster latency as the widely adopted autoregressive models. Our work successfully paved the way for scaling up the diffusion language model to a broader scope of applications across different domains.

1 INTRODUCTION

The paper addresses the high latency, limited scalability, and caching difficulty of diffusion language models by introducing two training-free acceleration methods. FreeCache approximates stable KV states, while Guided Diffusion uses a lightweight autoregressive model to guide token unmasking; together they improve speed with negligible accuracy loss.

  • DLMs achieve comparable quality to autoregressive models but suffer extensive computation latency and limited long-context ability.
  • The central problem is how to accelerate state-of-the-art DLMs while maintaining model performance.
  • FreeCache provides training-free KV caching by approximating earlier clean-token states whose future-token influence diminishes across denoising steps.
  • Guided Diffusion uses a lightweight autoregressive model to select plausible tokens for unmasking, reducing denoising overhead while preserving DLM reasoning.
  • 12.14× and 13.29× average speedups are achieved on Dream-7B-Instruct and LLaDA-8B-Instruct, respectively, with negligible accuracy drops.
  • The methods support long-context diffusion beyond 1024 tokens without hurting model performance and show strong results across simple and complex reasoning tasks.
  • Together, the techniques make DLMs feasible for long-context and high-throughput applications while retaining parallel generation and bidirectional conditioning.

2 RELATED WORK

Related work studies the computational and memory costs of diffusion and autoregressive language-model inference. Existing directions include cached diffusion, KV-cache optimization, and speculative decoding, but their effectiveness depends on training overhead, system design, or draft-target agreement.

  • Diffusion language models offer high scalability but remain computationally expensive and have worse latency than conventional autoregressive models.
  • KV-cache research addresses long-context autoregressive decoding through attention-system designs and compression methods.
  • Speculative decoding accelerates autoregressive decoding with a lightweight assistant that drafts multiple tokens before target-model verification.

3 METHODOLOGY

The methodology targets the computational cost and token incoherence of iterative diffusion generation. It combines FreeCache, which reuses stable KV projections through shrinking computation windows, with Guided Diffusion, which uses autoregressive agreement to select tokens for parallel unmasking.

  • Diffusion generation: Diffusion language models iteratively select and unmask token subsets, requiring repeated denoising passes until the sequence is fully unmasked.At each iteration, the model predicts logits, selects positions for unmasking, and updates the remaining mask set.
  • DLM inference bottleneck: Full-sequence computation at every denoising step creates a major latency bottleneck compared with autoregressive caching.DLMs process the full length L in transformer modules at each step, whereas AR decoding uses a prefix of length l per token.
  • FreeCache: FreeCache exploits the temporal stability of clean-token KV projections by progressively shrinking the active computation window as blocks become finalized.It recomputes the current and subsequent blocks, freezes completed blocks, and excludes them from later generation steps.
  • FreeCache: FreeCache delivers up to 5× speedup over vanilla Dream-7B while preserving minimal accuracy loss, including 4.42× for Dream-7B-Instruct and 6.32× for LLaDA-8B-Instruct.The reported results indicate that caching accelerates inference while maintaining accuracy close to the baseline.
  • Guided Diffusion: Guided Diffusion combines one-step diffusion proposals with autoregressive predictions, unmasking matched tokens to improve semantic coherence during parallel generation.Agreement between the diffusion model and a frozen autoregressive model determines how many tokens can be safely unmasked; if no agreement exists, only the first token is accepted.
  • Guided Diffusion: The guidance strategy adaptively controls parallel unmasking using model agreement rather than heuristic confidence thresholds, without additional training or fine-tuning.The process repeats until all tokens are unmasked, while preserving the diffusion model’s parallel generation and using the autoregressive model as a coherence guide.

4 EXPERIMENTS

The experiments evaluate FreeCache and Guided Diffusion across reasoning, question-answering, and long-context tasks using accuracy and wall-clock latency. The methods provide substantial speedups with negligible accuracy degradation and remain compatible with autoregressive guiders.

  • Experimental setup: Experiments evaluate accuracy and wall-clock latency against vanilla Dream and mainstream autoregressive models across multiple task types.Dream-Instruct-7B is the primary diffusion model, evaluated on a single NVIDIA RTX 6000 Ada GPU.
  • Comparison with the baseline diffusion language model: 12.48× average speedup is achieved on complex reasoning tasks with minimal quality degradation after combining FreeCache and Guided Diffusion.FreeCache alone provides 3.11× speedup on tasks such as MMLU-PRO and GPQA.
  • Long-context diffusion: 34.1× end-to-end average speedup is achieved on long-context math problems such as GSM8K with 8-shot prompting.The evaluation uses up to 1024 generated tokens for selected reasoning, math, and question-answering tasks.
  • Comparison with the baseline diffusion language model: Guided unmasking recovers accuracy lost from approximated caching, including during long-context generation of 1024 tokens.Increasing the guider from 1.5B to 7B yields only marginal accuracy improvements with minimal latency overhead.
  • Comparison with the baseline diffusion language model: A dedicated math autoregressive guider, such as Qwen-2.5-Math-1.5B, further improves Dream-Instruct-7B reasoning on GSM8K.The method supports different autoregressive guider variants across knowledge domains.
  • Comparison with the vanilla autoregressive LLM models: Guided Diffusion provides significant acceleration while maintaining latency on par with autoregressive models on GSM8K.The generation quality remains governed by the diffusion model’s reasoning capacity, while the guider directs unmasking.
  • Model GPU memory usage: Guided unmasking introduces minimal memory overhead relative to the combined memory consumption of the autoregressive and baseline diffusion models.GPU memory usage is reported for LLaDA-8B-Instruct and Dream-7B-Instruct.

5 CONCLUSION AND DISCUSSION

The paper concludes that FreeCache and Guided Diffusion provide training-free acceleration for diffusion language models, with comparable performance to autoregressive models. It also identifies accuracy, memory, system complexity, and guider alignment as remaining trade-offs.

  • Conclusion: 12.14× average speedup is achieved across reasoning tasks with negligible accuracy degradation compared with baseline diffusion models.The approach uses off-the-shelf diffusion models without additional training or fine-tuning.
  • Conclusion: FreeCache enables training-free caching, while Guided Diffusion guides direct diffusion to improve quality and accelerate end-to-end generation.Together, the methods target diffusion-language-model inference without requiring a dedicated training process.
  • Conclusion: The method achieves performance comparable to autoregressive language models while maintaining similar accuracy.This comparison is presented as a central conclusion of the paper.
  • Discussion: FreeCache trades accuracy and memory for speed because KV approximation introduces small errors and higher memory use.Guided Diffusion additionally requires an autoregressive guider, whose alignment affects effectiveness.

B.1 CODE & DATA RELEASE PLAN

The implementation of the proposed methods is publicly available through an anonymous repository.

  • Code release: The method implementation is available at the paper’s anonymous 4open.science repository.The release passage provides the repository URL.

C.1 ADDITIONAL EXPERIMENTAL RESULTS OF LLADA-8B-INSTRUCT

Additional experiments validate the proposed acceleration methods on LLaDA-8B-Instruct using GSM8K with 8-shot chain-of-thought prompting. Guided Diffusion avoids block-length control by using the diffusion model for drafting and the autoregressive model for unmasking guidance.

  • Additional experimental results: LLaDA-8B-Instruct shows consistent speedup from FreeCache and guided unmasking on GSM8K with 8-shot chain-of-thought.This extends validation beyond the Dream-Instruct-7B diffusion model.
  • Additional experimental results: LLaDA’s semi-autoregressive baseline still processes the full sequence at every denoising step despite left-to-right block unmasking.Its user-defined block length controls block size but does not make generation fully autoregressive.
  • Additional experimental results: Guided Diffusion needs no block length because the diffusion model drafts all tokens while the autoregressive model guides which tokens are unmasked.Its effective block length is the number of diffused draft tokens produced at each drafting step.

C.2 ADDITIONAL EXPERIMENTAL RESULTS ON DREAM-V0-INSTRUCT-7B

Additional Dream-v0-Instruct-7B results show that FreeCache and Guided Diffusion provide substantial speedup and accuracy gains on HellaSwag relative to the baseline.

  • Additional HellaSwag Results: FreeCache and Guided Diffusion achieve substantial speedup and accuracy gains over the baseline.
  • Additional HellaSwag Results: HellaSwag results report accuracy, latency, and speedup relative to the Dream-v0-Instruct-7B baseline.The results are presented for Dream-v0-Instruct-7B with FreeCache and Guided Diffusion.

C.3 MODEL GPU MEMORY USAGE

The section reports GPU memory usage for Guided Diffusion with LLaDA-8B-Instruct and Dream-v0-Instruct-7B, finding minimal overhead from guided unmasking.

  • Measurement Setup: GPU memory usage is reported in gigabytes for LLaDA-8B-Instruct and Dream-v0-Instruct-7B diffusion models.The measurements are listed in the Guided Diffusion memory-usage tables.
  • Memory Overhead: Guided Diffusion introduces minimal memory overhead compared with the combined memory consumption of the single autoregressive and baseline diffusion models.
  • Measurement Setup: The LLaDA-8B-Instruct memory results are presented in Table 9.

C.4 IMPLEMENTATION DETAILS

The implementation uses off-the-shelf diffusion and autoregressive models with shared configurations, while stochastic guided unmasking applies a confidence threshold to token decisions and evaluations use task-specific reasoning prompts.

  • Configuration: The method uses off-the-shelf diffusion and autoregressive models with the same configuration set across tasks.The authors retain one configuration set to demonstrate generality and practicality.
  • Guided Unmasking: Stochastic Guided Unmasking unmasks a token when its diffusion-output maximum logit exceeds τ × max(loga), otherwise leaving it masked.τ is the Guidance Confidence Threshold, with default value 0.5; guidance logits use only the AR model's Top-K logits.
  • Configuration: The Dream-v0-Instruct-7B and LLaDA-8B-Instruct diffusion configurations are detailed in Table 11.
  • Prompt Templates: The evaluation prompts require step-by-step reasoning followed by a final answer in a specified format.The templates cover commonsense, question-answering, reasoning, and multiple-choice tasks.
  • Prompt Templates: MMLU-PRO prompts are adopted from the benchmark's official repository.
Loading 2505.21467v2…