通过滑动窗口修剪与缓存,实现扩散语言模型的高效推理。
Window-Diffusion: Accelerating Diffusion Language Model Inference with Windowed Token Pruning and Caching
- 用滑动窗口划分待解码词元,仅计算局部活跃词元和缓存的缓冲词元
- 在相同计算预算下,推理速度提升最高达99倍,生成质量损失小
- 适用于预训练扩散语言模型,无需重新训练,适合实际部署
扩散语言模型(DLM)通过迭代去噪生成文本,但每轮推理需对全序列进行注意力计算,导致大量冗余计算。块级扩散虽可降低开销,但通常依赖重训练和受限更新顺序,难以直接应用于预训练DLM。我们发现DLM推理中存在显著的结构局部性:解码由少量前缀局部活跃词元驱动;远距离未解码上下文影响迅速衰减;已解码词元表现出阶段性时间稳定性,中间表示可复用,仅短暂解码后存在瞬态变化。基于此,我们提出Window-Diffusion,一种基于窗口的词元剪枝与缓存推理方法。维护一个随去噪过程向右滑动的局部计算窗口,将未解码词元分为三类:(i) 在线计算的活跃词元,(ii) 缓存并定期刷新键值状态的缓冲词元,(iii) 窗外被剪枝的远场词元。每阶段仅在窗口内计算活跃与缓冲词元,远场词元被忽略。在LLaDA和Dream上的实验表明,在匹配计算预算下,该方法实现最高99倍的推理加速,同时基本保持生成性能。
原文摘要 · Abstract (English)
Diffusion language models (DLMs) generate text through iterative denoising, but inference requires full-sequence attention at every iteration, resulting in substantial redundant computation on masked tokens. Block-wise diffusion can reduce this cost, yet it typically relies on retraining and constrained update orders, limiting its direct applicability to pretrained DLMs. Our token-level analysis reveals pronounced structural locality in DLM inference. Decoding is driven by a small set of prefix-localized active tokens; the influence of distant undecoded context diminishes rapidly, and decoded tokens exhibit stage-wise temporal stability, enabling reuse of intermediate representations except for a brief post-decode transient. Motivated by these observations, we propose \textbf{\placeholder}\footnote{The source code is available at https://github.com/vhicrgit/Window-Diffusion.}, a window-based token pruning and caching method for inference. We maintain a local computation window that slides rightward as denoising progresses, and partition undecoded tokens into: (i) \textit{active tokens} that are computed online, (ii) \textit{buffer tokens} whose KV states are cached and periodically refreshed, and (iii) \textit{far-field tokens} that are pruned outside the window. Computation is restricted to active and buffer tokens within the window, while far-field tokens are omitted at each stage. Experiments on LLaDA and Dream show that, under matched compute budgets, our method achieves up to $99\times$ inference speedup while largely preserving generation performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。