DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference

1. Enunciazione del Problema

I Large Language Models (LLM) affrontano un collo di bottiglia critico di memoria durante l'inferenza a contesto lungo. Sebbene la generazione autoregressiva si basi sulla cache Key-Value (KV) per evitare di ricalcolare i token passati, l'impronta di memoria di questa cache cresce linearmente con la lunghezza della sequenza ($O(L \cdot N \cdot d)$, dove $L$ è il numero di strati, $N$ è la lunghezza della sequenza e $d$ è la dimensione nascosta).

Le soluzioni esistenti post-allenamento utilizzano principalmente la potatura della cache KV, che scarta i token con punteggi di attenzione bassi per rientrare in un budget di memoria fisso. Tuttavia, una limitazione fondamentale dei metodi attuali (ad esempio H2O, StreamingLLM) è l'assunzione di potatura uniforme: applicano lo stesso rapporto di potatura attraverso tutti gli strati del transformer. Gli autori sostengono che questo sia subottimale perché assume implicitamente che tutti gli strati contribuiscano equamente alle prestazioni del modello, ignorando la sensibilità eterogenea di diversi strati alla rimozione dei token.

2. Metodologia: DepthKV

Il framework proposto, DepthKV, sfida l'assunzione di potatura uniforme introducendo un'allocazione della cache KV dipendente dallo strato. Invece di scartare una percentuale fissa di token per strato, DepthKV rialloca un budget KV globale fisso attraverso gli strati in base alla loro specifica sensibilità alla potatura.

A. Pre-studio: Analisi della Sensibilità degli Strati

Prima di progettare la strategia di allocazione, gli autori hanno condotto uno studio di ablazione strato per strato:

• Metodo: Hanno potato i token in un solo strato specifico alla volta, mantenendo gli altri intatti, misurando il conseguente calo delle prestazioni (ROUGE-1).
• Risultato: Il degrado delle prestazioni è altamente non uniforme. Certi strati "critici" causano crolli massicci quando potati, mentre altri mostrano un impatto trascurabile o addirittura ridondanza.
• Amplificazione del Contenuto: Hanno identificato gli "strati di amplificazione del contenuto". La potatura di questi strati non solo riduce l'accuratezza, ma sopprime significativamente la lunghezza e l'informatività dell'output (misurato tramite YapScore), indicando che questi strati sono cruciali per sostenere la generazione.
• Validazione Statistica: Test di permutazione hanno confermato che la variazione nell'importanza degli strati è statisticamente significativa attraverso i modelli (Gemma, LLaMA, Qwen) e i dataset.

B. Metriche di Rappresentazione per l'Importanza

Per automatizzare l'allocazione senza un'ablazione esaustiva, gli autori hanno analizzato le proprietà di rappresentazione per trovare un proxy per l'importanza dello strato.

• Selezione della Metrica: Hanno valutato sei metriche (Entropia Spettrale, Rank Effettivo, Curvatura, DiME, LiDAR e InfoNCE).
• Risultato Chiave: InfoNCE (una perdita contrastiva che misura la robustezza della rappresentazione alle perturbazioni di input) ha mostrato la correlazione negativa più forte con il degrado delle prestazioni.
  • Interpretazione: Gli strati con valori più bassi di InfoNCE (rappresentazioni più robuste) sono più critici. La loro potatura causa cali di prestazioni più grandi.
  • Utilizzo: DepthKV utilizza i punteggi InfoNCE per classificare gli strati e determinare l'allocazione del budget.

C. Strategie di Allocazione

DepthKV opera sotto un rapporto di potatura globale fisso ($\rho$) ma lo distribuisce in modo non uniforme. Vengono proposte tre strategie:

1. Protezione degli Strati Centrali (MLP): Un approccio euristico che preserva un sottoinsieme fisso di strati centrali (ad esempio, gli strati $\lfloor L/2 \rfloor$ e $\lfloor L/2 \rfloor + 1$) mentre pota gli altri in modo uniforme.
2. Allocazione Guidata dalla Metrica (MGA): Il metodo principale proposto. Calcola un punteggio di importanza $s^{(l)}$ basato su InfoNCE. Gli strati con maggiore robustezza (InfoNCE più basso) ricevono un budget KV più ampio (rapporto di potatura più basso), mentre gli strati meno critici vengono potati più aggressivamente.
3. Allocazione Metrica degli Strati Centrali (MLMA): Un approccio ibrido che protegge strutturalmente gli strati centrali mentre distribuisce il budget rimanente tra gli altri strati in base ai punteggi InfoNCE.

3. Contributi Chiave

1. Insight Teorico: Ha dimostrato che gli strati del transformer mostrano una sensibilità statisticamente significativa e non uniforme alla potatura della cache KV, invalidando l'assunzione di potatura uniforme.
2. Scoperta dell'Amplificazione del Contenuto: Ha identificato strati specifici la cui potatura sopprime la generazione di contenuti, collegando direttamente la robustezza della rappresentazione alla qualità e alla lunghezza dell'output.
3. Proposta di Framework: Ha introdotto DepthKV, un framework di allocazione dipendente dallo strato che ridistribuisce un budget di memoria fisso in base all'importanza dello strato (utilizzando specificamente InfoNCE).
4. Validazione Empirica: Ha dimostrato che DepthKV supera costantemente la potatura uniforme attraverso compiti e modelli diversi senza richiedere ulteriore allenamento.

4. Risultati Sperimentali

Gli autori hanno valutato DepthKV su tre famiglie di modelli (Gemma-7B, LLaMA-3.1-8B, Qwen2.5-7B) attraverso sette dataset:

• Compiti: Riassunto di documenti lunghi (arXiv, PubMed, GovReport, LegalCase), QA basato su documenti (HotpotQA, Qasper) e Ragionamento Matematico (GSM-$\infty$).
• Configurazione: Riduzione KV globale fissa del 60% (mantenimento del 40%).
• Prestazioni:
  • Riassunto: DepthKV (in particolare MGA) ha migliorato significativamente ROUGE-1 e la similarità SBERT rispetto alla potatura uniforme (ad esempio, +3.0 ROUGE-1 su arXiv per Gemma).
  • QA e Ragionamento: Le varianti di DepthKV hanno ottenuto punteggi Exact Match (EM) e F1 più alti. Ad esempio, su Qasper, MLMA-6L ha superato tutti i baseline.
  • Robustezza: I benefici si sono estesi ai compiti di ragionamento (GSM-$\infty$), dove tutte le varianti di DepthKV hanno superato la potatura uniforme.
  • LLM-as-a-Judge: Le valutazioni utilizzando Prometheus hanno confermato che gli output di DepthKV erano più corretti, completi e concisi rispetto ai baseline di potatura uniforme.
• Ablazione: Lo studio ha confermato che come viene distribuito il budget (allocazione strato per strato) è più importante dello specifico stimatore di importanza dei token (ad esempio, H2O con V vs. senza V).

5. Significato e Limiti

Significato:

• Efficienza: DepthKV offre una soluzione semplice e post-allenamento al collo di bottiglia della memoria a contesto lungo, consentendo finestre di contesto più lunghe sull'hardware esistente senza ri-allenamento.
• Generalizzabilità: Il metodo è agnostico rispetto all'architettura ed efficace attraverso diverse famiglie di modelli e tipi di compiti.
• Cambio di Paradigma: Sposta il campo dalla potatura a "livello di token" all'allocazione delle risorse a "livello di strato", riconoscendo l'eterogeneità funzionale degli strati del transformer.

Limiti:

• Non Consapevole della Query: DepthKV stima l'importanza dei token durante la fase di prefill basandosi sull'attenzione storica. Non si adatta dinamicamente alla specifica query di decodifica, potenzialmente perdendo token che diventano critici solo per specifiche query future.
• Granularità a Livello di Testa: Il framework attuale alloca i budget a livello di strato ma aggrega l'attenzione attraverso le teste. Non tiene conto della variabilità intra-strato dove specifiche teste di attenzione potrebbero essere più critiche di altre.

Conclusione

DepthKV stabilisce che non tutti gli strati del transformer sono creati uguali nell'inferenza a contesto lungo. Riallocando i budget di memoria per preservare gli strati più sensibili e robusti (identificati tramite InfoNCE), il metodo raggiunge compromessi prestazioni-per-memoria superiori rispetto alla potatura uniforme, offrendo una via pratica per scalare le finestre di contesto degli LLM.