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

AI summary

Overview

Research area: Efficient inference for Diffusion Language Models (DLMs) — specifically caching mechanisms for bidirectional transformer decoders.

Technical level: Advanced. The paper assumes familiarity with transformer internals (KV caching, attention, FFN blocks), diffusion language modeling, and low-rank matrix theory (SVD), though its high-level argument can be followed without the proofs.

Scope: A caching framework that jointly optimizes which hidden states to recompute and how much recomputation budget each layer receives, delivering up to 8× throughput gains over vanilla DLM decoding.

What This Paper Is About

Diffusion Language Models generate text in arbitrary order using bidirectional attention, which breaks the standard KV cache that makes autoregressive models fast — every decoding step forces a full-sequence recomputation with O(T × N²) cost. SPA-Cache attacks this by identifying only the tokens whose hidden states have actually drifted (using a cheap low-dimensional projection of the Value states) and by varying the update budget per layer to match each layer's natural stability profile.

Key Contributions

1. Theoretical justification for Value-state identification. The authors prove that if the Value state of a token remains similar across decoding steps, the resulting attention output (Theorem 3.1) and FFN output (Theorem 3.2) are provably bounded in divergence — turning a prior empirical heuristic into a principled identifier.

2. Singular proxy identifier. By applying truncated SVD to the Value projection matrix, the paper replaces the expensive full d-dimensional similarity check (d = 4096 for LLaDA-8B) with a low-rank proxy (r = 128), cutting projection cost from O(d³) to O(rd²) and similarity cost from O(d) to O(r). Theorem 3.4 bounds the similarity distortion introduced by truncation.

3. Adaptive layer-wise budget allocation. Instead of applying a uniform update ratio (e.g., 25%) to every layer, SPA-Cache parameterizes a per-layer update ratio with a piecewise Gaussian that concentrates updates on the volatile middle layers and aggressively caches stable early and late layers.

4. Empirical validation and composability. Evaluated on LLaDA-8B-Instruct and Dream-7B across seven benchmarks, SPA-Cache delivers up to 8× speedup over vanilla decoding, 2–4× over prior caching baselines, and up to 28× when composed with parallel decoding.

Main Findings

• Value states are the best identifier. A head-to-head comparison on GSM8K shows Value-state proxies preserve accuracy (78.59%) far better than Query (77.21%), Key (76.83%), attention input (77.29%), or attention output (73.92%), which suffers from feature-space anisotropy.

• Low-rank proxy preserves accuracy at a fraction of the cost. Using rank r = 128 versus the full d = 4096 dimension raises TPS from 164.88 to 179.43 with no meaningful accuracy change; accuracy only degrades when r drops below 128.

• Adaptive allocation saves compute without hurting quality. Layer-wise adaptive budgeting lowers the average update ratio from 25% to 16% while pushing TPS to 189.13 and keeping accuracy at 78.24%. Forcing a uniform 16% instead drops accuracy to 75.56%, confirming that the drift profile is genuinely heterogeneous.

• Large, consistent speedups across benchmarks. SPA-Cache reaches 6.0× and 6.2× speedups on MMLU-pro for LLaDA-8B and Dream-7B, and its peak result (8.0×) is on MBPP with LLaDA-8B — all while matching vanilla accuracy within reported error bars.

• Compatible with parallel decoding. Stacking SPA-Cache on top of Fast-dLLM's parallel decoding scheme yields a 27.9× speedup on BBH and 24.9× on MBPP, with only marginal accuracy trade-offs from the parallel sampling itself.

• Layer drift follows a bell shape. Analysis over 100 samples from GSM8K, MMLU-pro, and MBPP shows drift is low in early layers, peaks in the middle, and stabilizes near the output — motivating the piecewise Gaussian parameterization.

Methodology in Plain English

The method operates inside each transformer layer and proceeds in three phases:

1. Decide what to recompute. Each token's input state is passed through a small SVD-truncated version of the Value projection matrix, producing a compact 128-dimensional "fingerprint." This fingerprint is compared (cosine similarity) against the fingerprint stored from the previous decoding step. Tokens whose fingerprint has changed the most — the ones with the lowest similarity scores — are flagged for recomputation.

2. Recompute only the flagged tokens. Only those selected tokens get fresh Query, Key, and Value projections. The new K and V vectors are scattered back into the running KV cache, and attention is computed between the sparse queries and the partially refreshed cache.

3. Update FFN output selectively. The sparse attention outputs pass through the feed-forward layer, and only the flagged positions have their cached output states overwritten. Everything else reuses cached features.

The per-layer update fraction is not fixed. Based on the observed bell-shaped drift curve, the authors fit a piecewise Gaussian to the update ratio as a function of layer index, setting a peak ratio (ρ_p = 0.25 by default) in the middle of the network and decaying it toward the first and last layers. The full pipeline therefore scales computational cost with the number of updated tokens k rather than the full sequence length N.

Why This Matters

Research impact. The paper converts a heuristic (Value states track drift) into a provable property of transformer layers, and shows that low-rank subspaces preserve enough of that signal for reliable identification. This offers a general template for other caching, pruning, or sparse-activation methods in non-causal architectures, including multimodal diffusion transformers and mixture-of-experts models (which the authors note are built from standard FFNs).

Real-world applications.

• High-throughput chat and assistant systems built on diffusion-based LLMs, where per-token latency currently disqualifies DLMs from production deployment.
• Code generation tools that need long, structured outputs, since MBPP and HumanEval are among the benchmarks with the largest speedups.
• Mathematical and multi-step reasoning assistants, which typically require long generation chains and benefit from caching across many decoding steps.
• Edge or single-GPU deployments where inference throughput is the bottleneck and adding more hardware is not an option.

Industry relevance. DLM inference latency is the single largest barrier separating open-source diffusion LMs from proprietary high-throughput prototypes like Gemini Diffusion and Mercury. SPA-Cache narrows that gap by an order of magnitude and is orthogonal to sequence-level optimizations, making it an attractive drop-in addition to existing serving stacks.

Future Directions

• Extending the theory to mixture-of-experts layers. The FFN divergence bound (Theorem 3.2) is argued to hold for MoE blocks, but empirical validation on MoE-based DLMs is left open.

• Joint optimization with parallel decoding schedules. The composition with Fast-dLLM is currently a pipeline of two independent methods; co-designing the caching budget with the number of tokens decoded per step could yield further gains.

• Learning the layer-wise budget curve instead of fitting it. The piecewise Gaussian (ρ₁, ρ_p, ρ_L, l_p) is hand-tuned per model. A calibration-free or learned alternative would improve portability across architectures and tasks.

• Scaling to longer contexts and multimodal DLMs. The paper evaluates up to benchmark-length sequences on text-only models; whether the drift profile and the singular proxy remain this effective at much longer contexts or with cross-modal attention is untested.

Target Audience

Researchers and engineers working on efficient LLM inference, especially those focused on diffusion language models, KV-cache alternatives, or sparse activation. The paper is also useful for readers interested in the theoretical characterization of hidden-state dynamics inside transformers, and for practitioners evaluating whether DLMs are ready for latency-sensitive production deployments.

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