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

TL;DR

SPA-Cache introduces a novel caching method for diffusion language models that adaptively identifies critical tokens and allocates update budgets, significantly enhancing decoding efficiency and throughput.

Contribution

It proposes a low-dimensional singular proxy for critical token identification and an adaptive update strategy, improving DLM caching efficiency over prior methods.

Findings

Up to 8x throughput improvement over vanilla decoding

2-4x speedup over existing caching baselines

Effective identification of update-critical tokens with low overhead

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 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.