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

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

1. Het Probleem

Diffusion-based Large Language Models (dLLMs) bieden een veelbelovend alternatief voor autoregressieve (AR) modellen door parallelle token-generatie via iteratief ontdoening (denoising). Echter, hun architectuur kent een fundamenteel nadeel voor inferentie-efficiëntie:

• Bidirectionele Attention: In tegenstelling tot AR-modellen die causale attention gebruiken (waarbij eerdere Key-Value (KV) paren geldig blijven), gebruiken dLLMs bidirectionele attention. Dit betekent dat het "ontmaskeren" (unmasking) van zelfs één enkel token de representaties van alle posities in de sequentie verandert.
• Gebrek aan Lossless KV Caching: Vanwege deze afhankelijkheid kunnen standaard KV-caching technieken niet worden toegepast zonder verlies aan nauwkeurigheid. Traditionele dLLM-inferentie vereist daarom een volledige forward pass over de hele sequentie bij elke denoising-stap, wat leidt tot een kostprijs die lineair schaalt met de contextlengte en het aantal stappen.
• Beperkingen van Bestaande Oplossingen: Bestaande benaderingen voor benaderende KV-caching (zoals Fast-dLLM, d2Cache, Elastic-Cache) proberen dit te oplossen door selectief cache-states te updaten. Deze methoden hebben echter twee grote nadelen:
  1. Ze vereisen vaak dure berekeningen (zoals attention rollout of per-layer cosine similarity vergelijkingen) om te beslissen wanneer te herschrijven.
  2. De overhead van deze beslissingen schaalt mee met de contextlengte, modeldiepte of verborgen dimensie, wat de potentiële snelheidswinst tenietdoet.

2. Methodologie: EntropyCache

EntropyCache is een trainingsvrije (training-free) methode die een goedkope, constante-kost signaal gebruikt om te beslissen wanneer een volledige herschrijving (recomputation) nodig is. De methode rust op twee empirische observaties:

Observatie 1: Entropie voorspelt KV-cache drift.
De auteurs ontdekten dat de maximale entropie van de nieuw gedecodeerde token-verdelingen sterk correleert met de grootte van de drift in de KV-cache.

• Principe: Wanneer een token wordt ontsluierd vanuit een hoge-entropie (onzeker) verdeling, injecteert dit veel nieuwe informatie in de sequentie, wat leidt tot grote veranderingen in de interne representaties (KV-drift). Bij lage entropie (zekerheid) is de drift minimaal.
• Voordeel: Het berekenen van deze entropie kost slechts $O(V)$ tijd (waarbij $V$ de vocabulairegrootte is) en is onafhankelijk van de contextlengte of modelgrootte.

Observatie 2: Feature-volatiliteit duurt langer dan één stap.
Het volgen van individuele token-trajectoiren toont aan dat sommige tokens na het ontsluieren nog meerdere denoising-stappen onstabiel blijven.

• Implicatie: Bestaande dynamische methoden die alleen de direct vorige stap herschrijven, zijn ontoereikend. EntropyCache lost dit op door de $k$ meest recent gedecodeerde tokens te herschrijven, niet alleen de huidige masker-tokens.

Het Algorithmische Proces:
Bij elke denoising-stap $t$:

1. Forward Pass: Er wordt een volledige of partiële forward pass uitgevoerd. Bij een partiële pass worden alleen de huidige masker-tokens en de $k$ meest recente gedecodeerde tokens herschreven; de rest wordt uit de cache gehaald.
2. Decode & Entropy Evaluatie: Nieuwe tokens worden gedecodeerd en de maximale entropie ($E_{t+1}$) van hun verdelingen wordt berekend.
3. Beslissing (Skip-or-Recompute):
   • Als $E_{t+1} > \tau$ (drempelwaarde): Er wordt een volledige forward pass getriggerd om de hele KV-cache te verversen.
   • Als $E_{t+1} \leq \tau$: Er wordt een partiële forward pass uitgevoerd. De cache wordt hergebruikt voor de meeste tokens, met uitzondering van de huidige maskers en de $k$ recentste tokens.
4. Overhead: De beslissingslogica kost slechts $O(V)$ berekeningen en $O(k)$ geheugen, volledig onafhankelijk van de contextlengte ($L$) of modeldiepte.

3. Belangrijkste Bijdragen

1. Nieuwe Signaal voor Cache-Stalheid: Identificatie van de maximale entropie van gedecodeerde tokens als een lichtgewicht, empirisch onderbouwde proxy voor KV-cache-veroudering, die dure per-layer analyses vervangt.
2. Multi-stap Feature Volatiliteit: Empirisch bewijs dat feature-instabiliteit meerdere stappen aanhoudt na het ontsluieren, wat leidt tot de strategie van het herschrijven van een venster van recente tokens ($k$) in plaats van slechts één stap.
3. Efficiënte Implementatie: Integratie van deze inzichten in EntropyCache, een methode die constante beslissingskosten biedt en onafhankelijk is van schaalparameters van het model.

4. Resultaten

De methode is geëvalueerd op twee dLLM-architecturen: LLaDA-8B-Instruct en Dream-7B-Instruct, over diverse benchmarks (GSM8K, MATH500, MBPP, HumanEval, en Chain-of-Thought taken).

• Snelheidswinst:
  • Op standaard benchmarks bereikte EntropyCache een 15.2x tot 26.4x snelheidswinst ten opzichte van de basislijn (vanilla dLLM inferentie).
  • Op Chain-of-Thought (CoT) benchmarks (langere contexten) was de winst zelfs 22.4x tot 24.1x.
  • Dit overtreft bestaande methoden zoals d2Cache en Elastic-Cache aanzienlijk in throughput.
• Nauwkeurigheid:
  • EntropyCache behaalde consistent concurrerende nauwkeurigheid, vaak gelijk aan of beter dan de basislijn, terwijl andere caching-methoden vaak in nauwkeurigheid inzaken (vooral op HumanEval).
  • Op HumanEval behaalde het de enige "lossless" performance vergeleken met de basislijn, terwijl het toch een snelheidswinst van 8.9x tot 30.7x boekte.
• Overhead:
  • De beslissingslogica van EntropyCache neemt slechts 0.5% van de totale inferentietijd in beslag, vergeleken met 9.2% voor d2Cache en 14.5% voor Elastic-Cache.

5. Betekenis en Impact

EntropyCache adresseert een cruciale bottleneck in de adoptie van diffusion-based taalmodellen: de hoge inferentiekosten door het gebrek aan efficiënte caching.

• Schalbaarheid: Door de beslissingskosten te koppelen van de contextlengte en modelgrootte, maakt EntropyCache dLLMs schaalbaar voor lange contexten en grote modellen zonder dat de overhead exponentieel toeneemt.
• Toepasbaarheid: De methode is trainingsvrij en kan direct worden toegepast op bestaande dLLM-architecturen.
• Toekomstperspectief: De paper suggereert dat het principe van entropie-gestuurde overslaan ook toepasbaar is op andere diffusion-varianten, zoals blockwise diffusion modellen, wat de weg vrijmaakt voor bredere efficiëntieverbeteringen in generatieve AI.

Kortom, EntropyCache biedt een elegante en uiterst efficiënte oplossing om de snelheid van diffusion language modellen te verhogen zonder in te leveren op de kwaliteit van de gegenereerde tekst.