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

1. Il Problema

I Large Language Models basati sulla Diffusione (dLLM) rappresentano un'alternativa promettente ai modelli autoregressivi (AR), permettendo la generazione parallela di token attraverso processi iterativi di denoising. Tuttavia, soffrono di un limite architetturale fondamentale: utilizzano l'attenzione bidirezionale (non causale).

• Limitazione del KV Caching: Nei modelli AR, l'attenzione causale garantisce che le coppie Chiave-Valore (KV) calcolate in precedenza rimangano valide man mano che nuovi token vengono aggiunti, rendendo il caching delle KV lossless ed efficiente. Nei dLLM, invece, ogni posizione dell'attenzione guarda tutte le altre; di conseguenza, anche la decodifica di un singolo token altera le rappresentazioni di tutte le posizioni.
• Costo Computazionale: L'inferenza "vanilla" (standard) dei dLLM richiede un forward pass completo su tutta la sequenza ad ogni passo di denoising. Il costo scala linearmente con la lunghezza del contesto e il numero di passi, rendendo l'inferenza estremamente costosa.
• Limiti delle Soluzioni Esistenti: I metodi di caching approssimato esistenti (es. Fast-dLLM, d2Cache, Elastic-Cache) tentano di ridurre questo costo riutilizzando selettivamente le KV cache. Tuttavia, i loro meccanismi di decisione su quando o dove ricalcolare hanno un sovraccarico che scala con la lunghezza del contesto o la profondità del modello, limitando i guadagni reali di velocità.

2. Metodologia: EntropyCache

Gli autori propongono EntropyCache, un metodo di caching delle KV senza training (training-free) che utilizza l'entropia dei token decodificati come segnale per decidere quando ricalcolare.

Il metodo si basa su due osservazioni empiriche fondamentali:

1. L'entropia del token decodificato predice la deriva della KV Cache:
   Gli autori hanno osservato una forte correlazione positiva tra l'entropia massima della distribuzione dei token appena decodificati e la deriva (drift) delle rappresentazioni interne (KV cache) nel passo successivo.

   • Quando un token viene decodificato da una distribuzione ad alta entropia (alta incertezza), introduce una grande quantità di "sorprendenza" nella sequenza, causando un forte spostamento nelle rappresentazioni KV.
   • Quando l'entropia è bassa (distribuzione deterministica), il modello aveva già anticipato l'output, e la deriva delle KV è minima.
   • Vantaggio: Calcolare l'entropia massima richiede solo un'operazione $O(V)$ (dove $V$ è la dimensione del vocabolario), indipendente dalla lunghezza del contesto o dalla profondità del modello.

2. La volatilità delle feature persiste per più passi:
   Analizzando le traiettorie dei vettori di valore, gli autori hanno scoperto che alcuni token rimangono instabili per diversi passi di denoising dopo essere stati "sbloccati" (unmasked), non solo per il singolo passo successivo come ipotizzato dai metodi dinamici precedenti.

   • Soluzione: Invece di ricalcolare solo i token del passo corrente, EntropyCache ricalcola un budget fisso di $k$ token più recentemente decodificati insieme ai token mascherati attuali.

Algoritmo di Funzionamento

Ad ogni passo di denoising $t$:

1. Decodifica e Calcolo Entropia: Si decodificano nuovi token e si calcola l'entropia massima $E_{t+1}$ delle loro distribuzioni.
2. Decisione di Skip/Ricalcolo:
   • Se $E_{t+1} > \tau$ (soglia): Si attiva un forward pass completo per aggiornare l'intera KV cache.
   • Se $E_{t+1} \le \tau$: Si esegue un forward pass parziale. Si riutilizzano le KV cache per la maggior parte dei token, ricalcolando solo:
     • I token attualmente mascherati.
     • I $k$ token più recentemente decodificati (per gestire la volatilità a lungo termine).
3. Overhead Costante: La decisione di skip richiede solo $O(V)$ computazioni, rendendo il sovraccarico trascurabile e costante rispetto alla lunghezza della sequenza.

3. Contributi Chiave

• Nuovo Segnale di Decisione: Identificazione dell'entropia massima dei token sbloccati come proxy leggero ed empiricamente fondato per la "stalezza" della KV cache, sostituendo costosi calcoli di rollout dell'attenzione o confronti di stati nascosti.
• Analisi della Volatilità: Dimostrazione che l'instabilità delle feature persiste per più passi dopo la decodifica, giustificando il ricalcolo di una finestra di token recenti ($k$) invece di un singolo passo.
• Metodo Training-Free: Integrazione di queste scoperte in EntropyCache, un metodo che non richiede riaddestramento del modello.
• Efficienza Scalabile: La complessità decisionale è decoupled dalla lunghezza del contesto e dalla scala del modello, a differenza delle soluzioni precedenti.

4. Risultati Sperimentali

Il metodo è stato valutato su due modelli dLLM (LLaDA-8B-Instruct e Dream-7B-Instruct) su benchmark standard e di ragionamento (Chain-of-Thought).

• Velocità (Speedup):
  • Su benchmark standard: 15.2× – 26.4× di accelerazione rispetto all'inferenza base.
  • Su benchmark Chain-of-Thought (CoT): 22.4× – 24.1× di accelerazione.
  • In confronto, i metodi precedenti (d2Cache, Elastic-Cache) ottengono speedup inferiori (es. 13×–18×) con un overhead decisionale molto più alto.
• Accuratezza: EntropyCache mantiene un'accuratezza competitiva, spesso pari o superiore alla baseline lossless, superando altri metodi di caching che tendono a degradare le prestazioni su task complessi (es. HumanEval).
• Overhead: La logica di decisione di EntropyCache occupa solo lo 0.5% del tempo totale di inferenza, contro il 9.2% di d2Cache e il 14.5% di Elastic-Cache.
• Efficienza di Memoria: L'uso di memoria aggiuntiva è minimo ($O(k)$), a differenza di metodi che richiedono buffer dipendenti dalla profondità del modello o dalla lunghezza del contesto.

5. Significato e Impatto

EntropyCache risolve il collo di bottiglia principale nell'inferenza dei modelli di linguaggio basati sulla diffusione, rendendoli competitivi in termini di throughput rispetto ai modelli autoregressivi tradizionali.

• Accessibilità: Riducendo drasticamente i costi di inferenza e l'uso di risorse hardware, rende possibile l'implementazione di dLLM in scenari con risorse limitate.
• Generalità: Il principio di skipping basato sull'entropia è applicabile a diverse strategie di decodifica (es. finestre scorrevoli, blocchi) e potrebbe essere esteso a future architetture di modelli di diffusione.
• Precisione vs. Velocità: Dimostra che è possibile ottenere accelerazioni massive senza sacrificare la qualità della generazione, identificando dinamicamente i momenti critici in cui il ricalcolo completo è necessario per mantenere l'accuratezza semantica.

In sintesi, EntropyCache rappresenta un avanzamento significativo nell'ottimizzazione dell'inferenza per i dLLM, offrendo un compromesso ottimale tra velocità, accuratezza e complessità computazionale.