Singular Proxies for Adaptive Caching in Diffusion Language Models
Wenhao SUN ⋅ Rong-Cheng Tu ⋅ Yifu Ding ⋅ Zhao Jin ⋅ Jingyi Liao ⋅ Yongcheng Jing ⋅ Dacheng Tao
Abstract
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 caching approaches reduce this cost by selective hidden state updates; however, they are still limited by (i) computationally costly token-wise update identification heuristics and (ii) rigid, uniform budget allocation that fails to account for heterogeneous hidden- tate dynamics. To address these challenges, we present SPA-Cache that jointly optimizes update identification and budget allocation. 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, motivated by the layer-wise heterogeneity in hidden state dynamics, 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 models decoding and a $2$-$4\times$ speedup over existing caching baselines.
Lay Summary
Diffusion Language Models (DLMs) are great because they don't have to generate text strictly left-to-right, but they are incredibly slow because you can't use standard KV-caching. You have to recompute the network's hidden states at every single denoising step. Existing fixes try to cache some states, but finding which tokens to update is too slow, and they treat every layer the same. SPA-Cache fixes this by finding token updates in a cheap, low-dimensional subspace and dynamically skipping updates for stable layers. It gives DLMs a $2\times$ to $4\times$ speedup over existing caches and up to $8\times$ throughput over vanilla decoding.
Successful Page Load