[NeurIPS 2024] EntropyCache: Unlocking 26x Speedup in Diffusion LLMs via Entropy-Guided KV Caching

TL;DR

Diffusion-based Large Language Models (dLLMs) like LLaDA and Dream offer massive parallelism but are hindered by bidirectional attention, which typically requires a full, expensive forward pass at every denoising step. EntropyCache breaks this bottleneck by using the maximum entropy of decoded tokens as a "cheap" trigger for cache updates. By only recomputing when the model is "surprised," it achieves up to 26.4x speedup with a decision overhead of nearly zero (0.5% of runtime).

The Motivation: The Bidirectional Attention Hex

In standard Autoregressive (AR) models (like GPT-4), the past is fixed. Once a token is generated, its Key-Value (KV) pair never changes, making KV caching a "free lunch."

However, dLLMs use bidirectional (non-causal) attention. When a dLLM unmasks a token in the middle of a sentence, that new information propagates to every other token in the sequence. Technically, this "drifts" the entire KV cache, forcing a total recomputation.

Prior works tried to solve this by:

1. Static Caching: Freezing most of the cache (fast but inaccurate).
2. Dynamic Caching: Using "Attention Rollout" to see which tokens moved most. The irony? Calculating who moved most often costs more than the computation saved.

Core Insight: Entropy is a Proxy for Drift

The authors discovered a striking physical intuition: Cache drift isn't random; it follows uncertainty.

As shown in the paper's empirical analysis, "Entropy Spikes" (moments where the model is unsure which token to pick) correlate almost perfectly with "KV Drift Spikes" (moments where internal representations shift significantly).

Figure 1: Notice how the red entropy spikes precede the green cosine distance (drift) jumps.

Methodology: The EntropyCache Algorithm

EntropyCache operates on a simple three-phase loop at each denoising step $t$:

1. Skip-or-Recompute Decision: Calculate the entropy of the newly decoded tokens. If $Max(Entropy)>au$, trigger a full forward pass.
2. Partial Forward Pass: If entropy is low, reuse the old cache for 99% of the tokens. Recompute KV pairs only for the current mask positions and a small budget ($k=64$) of recently decoded tokens.
3. Recent Token Strategy: Why $k$ tokens? The authors found that "feature volatility" persists. A token doesn't stabilize the instant it is unmasked; it "wobbles" for a few steps.

Figure 2: The EntropyCache workflow—Phase 1 (Decision), Phase 2 (Decoding), and Phase 3 (Update).

Why this is a Breakthrough

The computational complexity of the decision is just $O(V)$ (Vocabulary size). Unlike competing methods like d2 Cache ($O(L^2)$) or Elastic-Cache ($O(depthimeslength)$), EntropyCache's overhead does not grow as the conversation gets longer or the model gets deeper.

Experimental Results: Speed without the Sacrifice

Experiments on LLaDA-8B and Dream-7B show that EntropyCache is the first method to cross the "efficiency/accuracy" chasm.

| Method | Avg Accuracy (LLaDA) | Avg Speedup | Decision Overhead | | :--- | :---: | :---: | :---: | | Baseline | 52.20% | 1.0x | 0% | | d2 Cache | 50.97% | 13.1x | 9.2% | | EntropyCache | 52.28% | 15.2x | 0.5% |

In Chain-of-Thought (CoT) benchmarks where sequences are long, EntropyCache reached a staggering 107.3x speedup on BBH tasks.

Figure 3: Accuracy-Throughput tradeoff. EntropyCache (Max Entropy) provides a more robust Pareto frontier than Max Confidence or Entropy Sum.

Critical Analysis & Future Outlook

Strengths:

• Training-Free: Plug-and-play for any existing dLLM.
• Constant Overhead: Finally, a dynamic method that doesn't get bogged down by context length.

Limitations:

• Hyperparameter Sensitivity: The threshold $au$ needs calibration. Setting it too high in coding tasks (like MBPP) can lead to logic errors because code tokens stabilize differently than natural language.

The Verdict: EntropyCache proves that the "intelligence" of a model (its confidence/entropy) can be used to optimize its own "physics" (KV cache management). This opens the door for dLLMs to finally compete with, or even surpass, the inference efficiency of traditional GPT-style autoregressive models.