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

TL;DR

Diffusion language models struggle with standard KV caching due to arbitrary decoding orders, incurring heavy hidden-state recomputation. SPA-Cache tackles this by (i) using a singular proxy based on a low-rank projection to identify update-critical tokens and (ii) applying an adaptive, layer-wise budget that concentrates updates on volatile layers while caching stable ones; these strategies are supported by stability bounds linking value-state drift to attention behavior and by a low-rank approximation guarantee. Empirically, SPA-Cache yields up to $8\times$ throughput gains over vanilla decoding and up to $28\times$ when combined with parallel decoding, while maintaining comparable generation quality across seven benchmarks and two diffusion LMs. This work provides a principled, training-free caching approach that makes diffusion-based, high-throughput decoding more practical for real-world deployments.

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.

Figures (6)

• Figure 1: Adjacent-Step Similarities for LLaDA-8B-Instructllada. While input states exhibit uniformly high similarity, our singular proxy at early stage of the layer can efficiently uncover the drifts in the final FFN output. This provides a clear guidance for selective hidden state updates.
• Figure 2: Distribution of Drift Across Layers. Rather than using a fixed threshold (dashed line) to allocate update budgets uniformly across layers, our SPA-Cache employs a dynamic threshold (solid line) that allocates more updates on the most volatile layers and fewer on the more stable layers.
• Figure 3: Overview of a Single SPA-Cache Layer.Phase 1: Input states are projected by $f_{\mathrm{proxy}}$ to produce update scores. A Top-$k$ selection identifies the update index set $\mathcal{I}$, which identifies positions whose cached states from the previous steps should be refreshed. Phase 2: Queries are generated only from the selected states and attend to a key-value cache that is selectively updated at indices $\mathcal{I}$. Phase 3: The resulting sparse attention outputs are fed into the MLP for the layer output. The output states is updated at $\mathcal{I}$, while non-selected positions reuse cached features.
• Figure 4: Component-wise Latency Decomposition. Value proxy reduces Attention and FFN overhead, but it introduces additional computational cost during update identification. Our proposed singular proxy further reduces this identification cost.
• Figure 5: Empirical Evidence of Anisotropy Problem in Attention Outputs. The density plots compare the cosine similarity distributions of value states (blue) versus attention outputs (orange). While value states exhibit an isotropic distribution centered near zero (orthogonality), the attention outputs show a significant positive shift in mean similarity.

Theorems & Definitions (15)

• Theorem 3.1: Bound on Attention Output Similarity
• Theorem 3.2: Bound on FFN Output Divergence
• Remark 3.3
• Theorem 3.4: Similarity Preservation in the Truncated Projection
• Remark 3.5
• proof
• proof
• Proposition 1.1
• proof
• Proposition 1.2