arXiv:2602.02544cs.LGcs.AI2026-02中稿 · ICML

提出SPA-Cache,让扩散语言模型推理更快更省资源。

SPA-Cache: Singular Proxies for Adaptive Caching in Diffusion Language Models

  • 用低维奇异代理识别关键更新词元,降低计算开销。
  • 动态分配缓存预算,稳定层少更新,性能不下降。
  • 相比原方法提速8倍,比现有方案快2至4倍。

扩散语言模型(DLMs)虽能灵活处理任意顺序生成,但其非因果特性导致无法使用标准的键值缓存,必须在每一步解码时重新计算隐藏状态,代价高昂。现有缓存方法通过选择性更新隐藏状态来降低成本,但仍受限于:(i) 高成本的逐词元更新识别启发式策略,(ii) 固定统一的缓存预算分配,无法适应不同层隐藏状态动态差异。为此,我们提出SPA-Cache,联合优化更新识别与缓存预算分配。首先,我们构建一个低维奇异代理,可在低维子空间中高效识别需更新的关键词元,显著降低识别开销;其次,引入自适应策略,对状态稳定的层分配更少更新,而不影响生成质量。两项贡献共同大幅提升DLM效率,在吞吐量上相较原始解码提升最高达8倍,较现有缓存基线提速2至4倍。

原文摘要 · Abstract (English)

While Diffusion Language Models (DLMs) offer a flexible, arbitrary-order alternative to the autoregressive paradigm, their non-causal nature precludes standard KV caching, forcing costly hidden state recomputation at every decoding step. Existing DLM caching approaches reduce this cost by selective hidden state updates; however, they are still limited by (i) costly token-wise update identification heuristics and (ii) rigid, uniform budget allocation that fails to account for heterogeneous hidden state dynamics. To address these challenges, we present SPA-Cache that jointly optimizes update identification and budget allocation in DLM cache. First, we derive a low-dimensional singular proxy that enables the identification of update-critical tokens in a low-dimensional subspace, substantially reducing the overhead of update identification. Second, we introduce an adaptive strategy that allocates fewer updates to stable layers without degrading generation quality. Together, these contributions significantly improve the efficiency of DLMs, yielding up to an $8\times$ throughput improvement over vanilla decoding and a $2$--$4\times$ speedup over existing caching baselines.

扩散模型缓存优化推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。