🤖 machine learning

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

The paper introduces SPA-Cache, a novel caching framework for Diffusion Language Models that utilizes a low-dimensional singular proxy for efficient update identification and an adaptive budget allocation strategy to overcome non-causal limitations, achieving up to an 8x throughput improvement over vanilla decoding and a 2–4x speedup over existing baselines.

Original authors: Wenhao Sun, Rong-Cheng Tu, Yifu Ding, Zhao Jin, Jingyi Liao, Yongcheng Jing, Dacheng Tao

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Wenhao Sun, Rong-Cheng Tu, Yifu Ding, Zhao Jin, Jingyi Liao, Yongcheng Jing, Dacheng Tao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Rewrite-Everything" Dilemma

Imagine you are writing a story, but instead of writing it word-by-word from left to right (like a normal person), you are writing it in a random order. You might write the ending first, then jump to the middle, then go back to the beginning. This is how Diffusion Language Models (DLMs) work. They are flexible and can fill in blanks anywhere, which is great for creativity and complex tasks.

However, there is a huge downside. Because you are jumping around, you can't just remember what you wrote five minutes ago and move on. Every time you add a new word, the whole story changes slightly. To get the next word right, the computer has to re-read and re-calculate the entire story from scratch every single time.

The Old Way (Autoregressive): Like reading a book. You remember the last page, turn it, and read the next one. Fast and easy.
The Diffusion Way: Like trying to solve a puzzle where every time you place one piece, the picture on the other pieces shifts. You have to re-scan the whole puzzle board every time you move a piece. This is incredibly slow and expensive.

The Solution: SPA-Cache

The authors created a system called SPA-Cache to speed this up. Think of it as a smart "Save Game" feature for this chaotic puzzle. Instead of re-calculating the whole story, the system tries to figure out: "Which parts of the story actually changed, and which parts are still the same?"

If a part of the story hasn't changed, the computer skips it and just uses the old memory (the "cache"). If a part has changed, it re-calculates only that specific part.

The paper introduces two main tricks to make this work efficiently:

1. The "Low-Res Snapshot" (Singular Proxies)

To decide what to re-calculate, the computer needs to check if the story has changed.

The Old Problem: Previously, computers tried to check the entire high-definition version of the story to see if it changed. This was like trying to spot a typo by reading every single letter of a 500-page book in high definition. It took too much time, canceling out the speed gains.
The New Trick (Singular Proxy): The authors realized they don't need the high-definition version to spot a change. They can use a low-resolution "snapshot" (a simplified, compressed version) to check for changes.
- Analogy: Imagine you are checking if a painting has been altered. Instead of examining every single brushstroke under a microscope (high cost), you step back and look at a blurry, low-res photo of the painting. If the blurry photo looks the same, the painting hasn't changed. If the blurry photo looks different, then you know you need to check the details.
- Result: This "snapshot" check is incredibly fast, allowing the system to quickly identify which parts of the story need re-writing without slowing down the whole process.

2. The "Smart Budget" (Adaptive Caching)

Once the system knows what to check, it needs to decide how much to re-calculate.

The Old Problem: Previous methods used a "one-size-fits-all" rule. They would say, "Re-calculate 25% of the story, no matter what."
- The Issue: Some parts of the story are very stable (like the setting or character names) and rarely change. Others are chaotic (like the plot twists) and change constantly. Re-calculating the stable parts is a waste of energy, while re-calculating the chaotic parts too little leads to mistakes.
The New Trick (Adaptive Budget): The system now acts like a smart manager who looks at the story and says: "This chapter is boring and stable? Let's only update 5% of it. This chapter is full of action and changing fast? Let's update 40% of it."
- Analogy: Think of a construction crew. If a building's foundation is solid and won't move, you don't send a crew to check it every day. But if a roof is leaking and shifting in the wind, you send a team to fix it immediately. SPA-Cache sends its "fixing crew" only where the "wind" is blowing hardest.

The Results: Speeding Up the Chaos

By combining these two tricks, the paper shows that SPA-Cache makes Diffusion Language Models significantly faster:

8x Faster: It is up to 8 times faster than the standard, slow way of running these models.
2–4x Faster than other tricks: It beats previous attempts to speed up these models by 2 to 4 times.
No Quality Loss: Despite skipping calculations, the quality of the story (the answers the model gives) remains just as good as if it had done all the work.

Summary

The paper solves the "slow puzzle" problem of Diffusion Language Models by teaching the computer to:

Use a quick, blurry snapshot to spot changes instead of a slow, detailed scan.
Spend its energy wisely, focusing only on the parts of the story that are actually changing, while ignoring the parts that are stable.

This makes these flexible, non-linear AI models practical for real-world use without sacrificing their unique ability to think in any order.

Technical Summary: Singular Proxies for Adaptive Caching in Diffusion Language Models

Problem Statement

Diffusion Language Models (DLMs) offer a flexible alternative to autoregressive (AR) models by enabling arbitrary-order decoding and parallel generation. However, their non-causal nature and reliance on bidirectional attention render standard Key-Value (KV) caching, a cornerstone of AR efficiency, incompatible. In DLMs, hidden states cannot be simply reused without recomputation to capture evolving global dependencies, leading to prohibitive inference latency ( $O(T \times N^2)$ complexity).

Existing attempts to mitigate this via caching face two primary limitations:

Costly Identification: Current methods rely on high-dimensional similarity computations (e.g., monitoring Value state drift) to identify which tokens require updates. This identification overhead often offsets the gains from sparse updates as sequence lengths scale.
Rigid Budget Allocation: Existing approaches typically apply a uniform update budget across all model layers. This fails to account for the heterogeneous dynamics of hidden states, where some layers are stable (requiring few updates) and others are volatile (requiring frequent updates), leading to either computational redundancy or performance degradation.

Methodology: SPA-Cache

The authors propose SPA-Cache (Singular Proxies for Adaptive Caching), a framework that jointly optimizes update identification and budget allocation. The workflow for a single Transformer layer is divided into three phases:

1. Update Identification via Singular Proxies

Instead of computing similarity in the full high-dimensional space ( $d$ ), SPA-Cache introduces a singular proxy to identify update-critical tokens in a low-dimensional subspace ( $r \ll d$ ).

Theoretical Basis: The authors prove that the stability of Value states is a sufficient condition for the stability of attention outputs and subsequent FFN outputs. Thus, monitoring Value states provides a conservative guarantee for output stability.
Dimensionality Reduction: To avoid the computational bottleneck of high-dimensional projection, the method applies Singular Value Decomposition (SVD) to the Value projection matrix $W$ . It retains only the top $r$ singular vectors and values to create a truncated projection $W_r$ .
Mechanism: Input states are projected into this low-dimensional manifold. Cosine similarity is computed between the current singular proxy and the cached proxy. Tokens with the lowest similarity (highest drift) are selected for recomputation.
Efficiency: This reduces the identification complexity from $O(d)$ to $O(r)$ and the projection overhead from $O(d^3)$ to $O(rd^2)$ .

2. Adaptive Budget Allocation

Recognizing that hidden state drift varies significantly across model layers, SPA-Cache replaces static, uniform update ratios with a dynamic, layer-wise adaptive strategy.

Observation: Drift is low in early layers (embedding transformations), peaks in middle layers (complex semantic transformations), and stabilizes in late layers.
Implementation: The update ratio $\rho(l)$ for layer $l$ is parameterized using a piecewise Gaussian function. This allocates a higher update budget to volatile middle layers and aggressively caches stable early and late layers, maximizing throughput without sacrificing quality.

3. Sparse Execution

The selected tokens (indices $I$ ) are used to generate sparse Queries, Keys, and Values. The Attention mechanism computes outputs using these sparse queries against a partially updated KV cache. The Feed-Forward Network (FFN) processes only the sparse attention outputs, and the final layer output is updated only at the selected indices, leaving other positions to reuse cached features.

Key Contributions

Theoretical Justification for Value Proxies: The paper provides formal proofs (Theorems 3.1 and 3.2) demonstrating that Value state similarity bounds the divergence of attention and FFN outputs, validating the use of Value states as a principled identifier for update selection.
Singular Proxy for Low-Overhead Identification: By leveraging truncated SVD, the authors introduce a method to identify critical tokens in a low-dimensional subspace, significantly reducing the computational cost of the identification process itself.
Adaptive Budget Allocation: The introduction of a dynamic update strategy that aligns computational resources with layer-wise heterogeneity, avoiding the inefficiencies of uniform allocation.
Orthogonality to Parallel Decoding: SPA-Cache is designed to be orthogonal to sequence-level optimizations and can be integrated with parallel decoding frameworks to amplify throughput gains.

Experimental Results

The authors evaluated SPA-Cache on two DLMs (LLaDA-8B and Dream-v0-Instruct-7B) across seven benchmarks (including GSM8K, MMLU-PRO, MBPP, and HumanEval).

Throughput Improvements: SPA-Cache achieves up to an 8× throughput improvement over vanilla decoding and a 2–4× speedup over existing caching baselines (dLLM-Cache, Fast-dLLM).
Specific Benchmarks:
- On MBPP with LLaDA-8B, SPA-Cache attained an 8.0× speedup.
- On MMLU-PRO, it achieved 6.0× (LLaDA-8B) and 6.2× (Dream-7B) speedups.
Integration with Parallel Decoding: When combined with parallel decoding methods, SPA-Cache yielded a total 28× speedup over vanilla decoding on the BBH benchmark, while maintaining comparable accuracy.
Ablation Studies:
- Replacing full-dimensional Value identifiers with the singular proxy ( $r=128$ ) improved Tokens Per Second (TPS) from ~165 to ~179 without accuracy loss.
- Adaptive budget allocation further increased TPS to ~189 by reducing the average update ratio from 25% to 16%, whereas a uniform 16% ratio caused accuracy degradation.

Significance and Claims

The paper claims that SPA-Cache establishes a new efficiency baseline for Diffusion Language Models, positioning them as a computationally viable alternative for high-throughput, real-world applications. By fundamentally addressing the inference bottleneck through principled theoretical analysis and adaptive resource allocation, the method enables DLMs to scale effectively without compromising generation fidelity. The authors emphasize that their approach is training-free and compatible with existing parallel decoding strategies, paving the way for the practical deployment of discrete diffusion models.