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

1. Problemstellung

Diffusion-basierte Large Language Models (dLLMs) stellen eine vielversprechende Alternative zu autoregressiven (AR) Modellen dar, da sie durch parallele Token-Generierung über mehrere Denoisingschritte hinweg potenziell einen höheren Durchsatz erreichen können. Allerdings leiden dLLMs unter einer fundamentalen architektonischen Einschränkung: Sie verwenden bidirektionale (nicht-kausale) Aufmerksamkeit.

Im Gegensatz zu AR-Modellen, bei denen K-V-Paare (Key-Value) für bereits generierte Tokens unverändert bleiben und somit effizient im Cache gespeichert werden können (lossless KV Caching), ändert sich bei dLLMs die Repräsentation aller Positionen, sobald auch nur ein einzelner Token „entschlüsselt" (unmasked) wird. Dies macht herkömmliches KV Caching unmöglich und zwingt das Modell dazu, bei jedem Denoisingschritt einen vollständigen Forward-Pass über die gesamte Sequenz durchzuführen. Der Rechenaufwand skaliert somit linear mit der Kontextlänge und der Anzahl der Schritte.

Bestehende Ansätze zur Approximation (z. B. Fast-dLLM, d2Cache, Elastic-Cache) versuchen, dies durch selektive Aktualisierung von Cache-Zuständen zu lösen. Diese Methoden leiden jedoch oft unter einem hohen Entscheidungs-Overhead, der mit der Kontextlänge oder der Modelltiefe skaliert, oder sie opfern Genauigkeit zugunsten von Geschwindigkeit.

2. Methodik: EntropyCache

Die Autoren schlagen EntropyCache vor, eine trainingsfreie KV-Caching-Methode, die zwei empirische Beobachtungen nutzt, um zu entscheiden, wann ein vollständiger Forward-Pass notwendig ist und wann der Cache wiederverwendet werden kann.

Kernidee

Anstatt komplexe Berechnungen (wie Attention Rollouts oder Layer-Vergleiche) durchzuführen, nutzt EntropyCache die maximale Entropie der neu entschlüsselten Token-Verteilungen als Signal für die „Staleness" (Veraltetheit) des Caches.

Die zwei empirischen Beobachtungen

1. Entropie korreliert mit KV-Cache-Drift: Die maximale Entropie der Token-Verteilung eines Schritts ist ein starker Prädiktor dafür, wie stark sich die internen K-V-Repräsentationen im nächsten Schritt ändern werden. Hohe Entropie (hohe Unsicherheit) führt zu großen Drifts, niedrige Entropie zu minimalen Änderungen.
2. Persistente Feature-Volatilität: Die Instabilität von Features nach dem Entschlüsseln eines Tokens hält oft über mehrere Denoisingschritte an, nicht nur für einen einzigen Schritt, wie es frühere dynamische Methoden annahmen.

Der Algorithmus

Der Prozess läuft in jedem Denoisingschritt $t$ in drei Phasen ab:

1. Forward Pass (Vollständig oder Teilweise):
   • Wenn ein „Skip"-Flag gesetzt ist (basierend auf dem vorherigen Schritt), wird nur ein partieller Forward-Pass durchgeführt. Dabei werden nur die aktuellen Mask-Token und eine kleine Menge der k kürzlich entschlüsselten Tokens neu berechnet. Der Rest wird aus dem Cache gelesen.
   • Wenn das Flag nicht gesetzt ist, wird ein vollständiger Forward-Pass durchgeführt, um den Cache zu aktualisieren.
2. Entschlüsselung & Entropie-Bewertung:
   • Nach dem Decoding der neuen Token wird die maximale Entropie $E_{t+1}$ über die vorhergesagten Verteilungen berechnet.
   • Die Berechnung kostet nur $O(V)$ (wobei $V$ die Vokabulargröße ist) und ist unabhängig von der Kontextlänge oder Modelltiefe.
3. Skip-Entscheidung & Auswahl:
   • Schwellenwert-Check: Wenn $E_{t+1} > \tau$ (Schwellenwert), wird für den nächsten Schritt ein vollständiger Forward-Pass ausgelöst.
   • Recomputation der jüngsten Tokens: Wenn $E_{t+1} \le \tau$, wird der Cache wiederverwendet, aber um die Genauigkeit zu erhalten, werden zusätzlich zu den Mask-Token auch die k kürzlich entschlüsselten Tokens (basierend auf einer Historie) neu berechnet. Dies kompensiert die beobachtete multi-stufige Volatilität.

3. Hauptbeiträge

• Leichtgewichtiges Signal: Identifikation der maximalen Token-Entropie als effizienter Proxy für KV-Cache-Staleness. Dies ersetzt teure Layer-weite Berechnungen durch eine konstante $O(V)$-Operation.
• Erkenntnis zur Volatilität: Empirischer Nachweis, dass Feature-Instabilität über mehrere Schritte anhält, was zu einer Strategie führt, die nicht nur den aktuellen, sondern auch die $k$ jüngsten Tokens neu berechnet.
• Training-freie Implementierung: EntropyCache erfordert kein Fine-Tuning und ist direkt auf existierende dLLMs anwendbar.
• Konstanter Overhead: Der Entscheidungs-Overhead ist unabhängig von Kontextlänge ($L$), Modelltiefe ($\ell$) und versteckter Dimension ($d$).

4. Ergebnisse

Die Methode wurde auf den Modellen LLaDA-8B-Instruct und Dream-7B-Instruct evaluiert.

• Beschleunigung:
  • Auf Standard-Benchmarks (GSM8K, MATH500, MBPP, HumanEval) wurde eine 15,2-fache bis 26,4-fache Beschleunigung gegenüber dem Basis-Modell erreicht.
  • Auf Chain-of-Thought (CoT) Benchmarks (längere Kontexte) lag die Beschleunigung bei 22,4-fach bis 24,1-fach.
• Genauigkeit: EntropyCache erreicht eine Genauigkeit, die mit dem Baseline-Modell (ohne Caching) vergleichbar ist und andere Caching-Methoden (wie d2Cache oder Elastic-Cache) in der Kombination aus Durchsatz und Genauigkeit übertrifft. Insbesondere auf HumanEval konnte es die Genauigkeit des Baselines beibehalten, während andere Methoden Einbußen hatten.
• Overhead: Der Entscheidungs-Overhead macht nur 0,5 % der gesamten Inferenzzeit aus (verglichen mit 9,2 % bei d2Cache und 14,5 % bei Elastic-Cache).

5. Bedeutung und Ausblick

EntropyCache adressiert einen der größten Engpässe bei der praktischen Anwendung von Diffusion-LLMs: den hohen Rechenaufwand durch fehlendes effizientes KV Caching.

• Skalierbarkeit: Da der Overhead nicht mit der Kontextlänge skaliert, eignet sich die Methode besonders gut für lange Generationen und komplexe Aufgaben (wie CoT), wo andere dynamische Caching-Methoden an ihre Grenzen stoßen.
• Ressourceneffizienz: Durch die drastische Reduktion der Inferenzkosten (bis zu 26x schneller) wird der Einsatz von Diffusion-Modellen in ressourcenbeschränkten Umgebungen oder für Echtzeitanwendungen praktikabler.
• Zukunftsperspektive: Das Prinzip der entropiegesteuerten Überspringung könnte auch auf andere Architekturen wie Block-Diffusion-Modelle (z. B. Fast-dLLMv2) übertragen werden, um die Rechenlast innerhalb von Sub-Blöcken zu optimieren.

Zusammenfassend bietet EntropyCache einen eleganten, theoretisch fundierten und empirisch validierten Ansatz, um die Lücke zwischen der theoretischen Parallelisierbarkeit von Diffusion-Modellen und der praktischen Effizienz von KV-Caching zu schließen.