EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models

1. Problem Statement

Diffusion-based Large Language Models (dLLMs) offer a promising alternative to autoregressive (AR) models by enabling parallel token generation through iterative denoising. However, they face a fundamental architectural bottleneck:

• Bidirectional Attention: Unlike AR models which use causal attention (allowing for lossless Key-Value (KV) caching), dLLMs employ bidirectional attention. Unmasking even a single token alters the representations of all positions in the sequence.
• Inefficiency: Consequently, standard dLLM inference requires a full forward pass over the entire sequence at every denoising step, resulting in a cost proportional to the product of sequence length and the number of steps.
• Limitations of Existing Solutions:
  • Static Caching (e.g., Fast-dLLM): Freezes KV pairs outside the current block, sacrificing accuracy for speed.
  • Dynamic Caching (e.g., d2Cache, Elastic-Cache): Dynamically decide when to recompute based on attention rollout or layer-wise drift. However, these methods incur high decision overhead that scales with context length ($L$), model depth ($\ell$), or hidden dimension ($d$), limiting their practical speedup.

2. Methodology: EntropyCache

The authors propose EntropyCache, a training-free KV caching method that minimizes decision overhead while maintaining accuracy. The method is grounded in two empirical observations:

Key Observations

1. Entropy Predicts KV Drift: The maximum entropy of newly decoded token distributions correlates strongly with the magnitude of KV cache drift (measured by cosine distance of value vectors). High entropy indicates high uncertainty and significant representation shifts, signaling the need for a full recomputation.
2. Feature Volatility Persists: Unmasked tokens do not stabilize immediately. Feature volatility often persists for multiple steps after unmasking. Therefore, recomputing only the current step's tokens is insufficient; a window of recently decoded tokens must also be refreshed.

Algorithm Workflow

EntropyCache operates in three phases per denoising step $t$:

1. Forward Pass (Full or Partial):

   • If a "skip" flag is active, the model performs a partial forward pass, recomputing K/V vectors only for:
     • The current mask tokens (being denoised).
     • The $k$ most recently decoded tokens (selected based on a decode history).
   • All other positions reuse cached KV states.
   • If the skip flag is inactive, a full forward pass is executed.

2. Decode & Entropy Evaluation:

   • New tokens are decoded from model logits.
   • The maximum entropy ($E_{t+1}$) of the newly decoded token distributions is calculated:
     $$E_{t+1} = \max_{i \in D_{t+1}} \left( -\sum_{j} p_{ij} \log p_{ij} \right)$$
   • This calculation costs $O(V)$, where $V$ is the vocabulary size, independent of context length or model depth.

3. Skip Decision & Token Selection:

   • Decision: If $E_{t+1} > \tau$ (threshold), the skip flag is set to False (triggering a full recomputation at the next step). Otherwise, the flag remains True.
   • Recent Token Selection: The algorithm maintains a history of when tokens were decoded. It selects the top $k$ most recently decoded tokens to include in the partial recomputation window, ensuring that volatile features are captured.

3. Key Contributions

• Constant-Cost Decision Signal: Replaces expensive per-layer attention rollout or hidden-state comparisons with a single $O(V)$ entropy scalar. The decision overhead is decoupled from context length ($L$), model depth ($\ell$), and hidden dimension ($d$).
• Multi-Step Volatility Handling: Introduces a mechanism to recompute the $k$ most recently decoded tokens, addressing the finding that feature instability persists beyond a single step.
• Training-Free Integration: The method requires no fine-tuning and can be integrated into existing dLLM inference pipelines with minimal code changes.

4. Experimental Results

The method was evaluated on LLaDA-8B-Instruct and Dream-7B-Instruct across standard benchmarks (GSM8K, MATH500, MBPP, HumanEval) and Chain-of-Thought (CoT) benchmarks.

• Speedup:
  • Standard Benchmarks: Achieved 15.2× – 26.4× speedup over vanilla dLLM inference.
  • Chain-of-Thought Benchmarks: Achieved 22.4× – 24.1× speedup, significantly outperforming previous methods (e.g., d2Cache, Elastic-Cache) which struggled with longer contexts.
• Accuracy:
  • EntropyCache maintained competitive accuracy, often matching or exceeding the baseline.
  • Notably, on HumanEval, it was the only method to achieve lossless performance compared to the baseline while delivering an 8.9×–30.7× speedup.
• Overhead:
  • The decision logic accounts for only 0.5% of total inference time.
  • In contrast, d2Cache and Elastic-Cache incur 9.2% and 14.5% overhead, respectively, due to their scaling complexity.
• Efficiency: EntropyCache achieved the highest throughput (tokens/sec) across almost all configurations, particularly in long-generation tasks where cache reuse is maximized.

5. Significance and Impact

• Scalability: By decoupling decision overhead from model size and context length, EntropyCache is uniquely suited for scaling dLLMs to longer contexts and larger models where existing dynamic caching methods become computationally prohibitive.
• Practical Deployment: The low overhead and training-free nature make it immediately applicable to production environments, potentially lowering the energy and hardware costs of deploying diffusion-based LLMs.
• Insight into dLLM Dynamics: The work provides a novel theoretical link between predictive uncertainty (entropy) and internal representation drift, offering a new perspective on how to optimize diffusion inference.

In conclusion, EntropyCache effectively bridges the gap between the theoretical efficiency of diffusion models and the practical constraints of inference latency, offering a robust, high-speed alternative to autoregressive generation for specific use cases.