Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

1. Il Problema

L'inferenza dei Modelli Linguistici su Grande Scala (LLM) con contesti lunghi è limitata dall'ingente footprint di memoria richiesto dalla Cache Key-Value (KV). Per mitigare questo problema, sono state sviluppate tecniche di "pruning" (potatura) dei token, che selezionano un sottoinsieme di token importanti da mantenere nella cache KV, scartando gli altri.

Le attuali metodologie di pruning a livello di strato (layer-wise token pruning), come FastKV o GemFilter, soffrono di una limitazione critica: scelgono i token da mantenere in strati predefiniti e fissi, indipendentemente dal compito specifico o dalla difficoltà del contesto.

• Flessibilità insufficiente: Come illustrato nel paper, compiti semplici (es. domande di risposta su testi) possono beneficiare di una selezione precoce (strati intermedi), mentre compiti difficili (es. recupero di informazioni chiave in un vasto contesto, o KV retrieval) richiedono che la selezione avvenga in strati più profondi per permettere al modello di affinare l'attenzione.
• Conseguenze: L'uso di strati fissi porta a una significativa degradazione delle prestazioni nei compiti difficili o richiede di aumentare il budget della cache KV, annullando i benefici di risparmio di memoria.

2. Metodologia Proposta: ASL (Adaptive Selection Layer)

Gli autori propongono ASL, un metodo senza addestramento (training-free) che seleziona dinamicamente lo strato ottimale per il pruning dei token basandosi sui pattern di attenzione osservati durante l'inferenza.

Meccanismo Chiave: Varianza dei Rank dei Token

L'idea centrale è monitorare la varianza dei rank dei token ordinati in base al loro punteggio di attenzione attraverso gli strati consecutivi.

1. Osservazione: Nei primi strati, i punteggi di attenzione sono distribuiti uniformemente. Man mano che si scende negli strati più profondi, l'attenzione si concentra stabilmente su un sottoinsieme fisso di token.
2. Calcolo della Varianza:
   • ASL calcola i punteggi di attenzione medi (pooled) su una finestra di token recenti.
   • Ordina i token in base a questi punteggi e calcola la varianza dei loro rank su un intervallo di $L_{obs}$ strati recenti.
   • Viene calcolata una varianza relativa normalizzando la varianza corrente rispetto alla varianza iniziale (allo strato $L_{min}$).
3. Criterio di Selezione: Quando la varianza relativa scende sotto una soglia predefinita ($\tau$), ciò indica che l'attenzione si è stabilizzata su un sottoinsieme specifico di token. In quel momento, lo strato corrente viene selezionato come strato di selezione.
4. Esecuzione:
   • Prefilling: ASL opera durante la fase di prefilling. Una volta determinato lo strato di selezione, i token vengono selezionati "in una sola passata" (one-shot) e propagati agli strati successivi.
   • Integrazione: ASL può essere combinato con metodi esistenti (es. SnapKV per la compressione prima della selezione, o GemFilter per una strategia a due passaggi) per ottimizzare sia la fase di prefilling che quella di decoding, rispettando un budget KV specifico.

3. Contributi Chiave

1. Identificazione della rigidità: Dimostrazione empirica che gli strati di selezione fissi sono inadeguati per compiti di diversa difficoltà, causando cali di accuratezza significativi in scenari complessi.
2. Proposta di ASL: Un metodo adattivo che determina lo strato di selezione monitorando la stabilità dei rank dei token, permettendo una riduzione della cache KV consapevole del compito (task-aware).
3. Integrazione Modulare: ASL è progettato per funzionare in sinergia con tecniche di riduzione della cache esistenti, ottimizzando sia la fase di prefilling (TTFT - Time to First Token) che di decoding (TPOT - Time Per Output Token).
4. Validazione Estensiva: Dimostrazione attraverso esperimenti che ASL offre un miglior compromesso tra accuratezza ed efficienza rispetto agli stati dell'arte, specialmente nei compiti difficili.

4. Risultati Sperimentali

Il metodo è stato valutato su tre benchmark principali: InfiniteBench, RULER e Needle in a Haystack (NIAH), utilizzando modelli come Llama-3.1-8B e Qwen2.5-7B con contesti fino a 256k token.

• Accuratezza:
  • ASL supera i metodi SOTA (come FastKV, GemFilter, PyramidInfer) in termini di accuratezza media, specialmente nei compiti difficili come il KV retrieval e il NIAH multi-key.
  • In InfiniteBench, ASL (configurazione one-shot) ottiene il punteggio medio più alto su Llama-3.1-8B, mentre ASL_2pass (due passaggi) mostra prestazioni eccellenti su Qwen2.5-7B.
  • Nei test NIAH, ASL e ASL_2pass raggiungono il 100% di accuratezza su tutte le lunghezze di contesto, superando di gran lunga GemFilter che fatica con contesti lunghi.
• Efficienza (Velocità):
  • TTFT (Prefilling): ASL è leggermente più lento di FastKV perché tende a selezionare i token in strati più profondi (necessari per compiti difficili), ma rimane significativamente più veloce dell'attenzione completa.
  • TPOT (Decoding): Una volta selezionati i token, la velocità di decoding è paragonabile o superiore ai metodi SOTA, poiché il numero di token da elaborare è ridotto al budget KV (es. 2048 token).
  • Throughput: Sebbene ASL sacrifichi leggermente il throughput medio rispetto a metodi più aggressivi come GemFilter, il divario si riduce drasticamente per query lunghe o in scenari reali dove il tempo di decoding domina.
• Consumo di Memoria:
  • L'overhead di memoria introdotto da ASL (per il calcolo delle varianze e la cache dei rank) è trascurabile. Tutti i metodi di pruning riducono l'uso di memoria da ~18GB (Full KV) a ~0.3GB, mantenendo un footprint quasi identico tra i vari metodi.

5. Significato e Impatto

Il lavoro di ASL rappresenta un passo avanti significativo verso l'inferenza efficiente di LLM su contesti ultra-lunghi.

• Adattività Dinamica: Sposta il paradigma da una selezione statica a una dinamica, riconoscendo che la "finestra di attenzione" ottimale varia in base alla complessità semantica del compito.
• Efficienza Senza Addestramento: Offre un miglioramento delle prestazioni senza richiedere il ri-addestramento del modello, rendendolo immediatamente applicabile a modelli LLM esistenti.
• Bilanciamento Ottimale: Risolve il dilemma tra velocità di inferenza e accuratezza, permettendo di mantenere un'alta precisione anche in compiti complessi senza dover aumentare eccessivamente il consumo di memoria.

In sintesi, ASL dimostra che l'adattività basata sull'analisi statistica dei pattern di attenzione è la chiave per sbloccare le potenzialità dei modelli LLM su contesti estremamente lunghi, superando i limiti delle strategie di pruning rigide.