💬 NLP

Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

Questo articolo introduce un metodo di fine-tuning per modelli transformer che consente un'inferenza efficiente su contesti lunghi con attenzione sparsa su hardware moderato, permettendo ai modelli di co-adattarsi con varie policy del KV cache, spesso superando gli approcci di attenzione esatta.

Autori originali: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Pubblicato 2026-08-21

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

I moderni sistemi di intelligenza artificiale che generano testi simili a quelli umani si affidano a un meccanismo che funge da memoria a breve termine, permettendo loro di ricordare ciò che è stato detto precedentemente in una conversazione o in un lungo documento. Questa memoria è conservata in un buffer digitale che cresce ad ogni nuova parola che il sistema elabora. Affinché questi sistemi possano funzionare bene su lunghi tratti di testo, questa memoria deve essere vasta, ma l'hardware informatico necessario per contenerla è costoso e limitato. Quando il buffer di memoria si riempie, il sistema deve decidere quali vecchi pezzi di informazione scartare per fare spazio ai nuovi. Se scarta l'informazione sbagliata, il sistema perde la capacità di ragionare o di rispondere alle domande con precisione. Ciò crea un difficile compromesso: mantenere la memoria piccola permette di risparmiare denaro e consente al sistema di funzionare su apparecchiature standard, ma rischia di perdere il contesto necessario per essere intelligente.

I ricercatori hanno a lungo cercato di risolvere questo problema insegnando al sistema a essere selettivo su ciò che conserva, un processo noto come attenzione sparsa (sparse attention). Tuttavia, uno studio recente rivela un difetto critico nel modo in cui questi sistemi sono stati addestrati finora. La maggior parte dei metodi esistenti addestra l'IA utilizzando una memoria perfetta e illimitata e poi cerca di costringerla a operare con una memoria limitata in un secondo momento. I ricercatori hanno scoperto che questo approccio fallisce perché l'IA non ha mai imparato come funzionare sotto i vincoli specifici che effettivamente affronterebbe. Addestrando il modello a dimenticare intenzionalmente e ad adattarsi a una dimensione di memoria fissa fin dall'inizio, il team ha dimostrato che il sistema può performare significativamente meglio rispetto a quelli addestrati con risorse illimitate, anche eseguendosi su un singolo chip di un computer moderatamente potente.

Il team, guidato da scienziati di Amazon Web Services e dell'Università di Amsterdam, ha sviluppato un nuovo modo per perfezionare (fine-tune) questi grandi modelli linguistici. Inveve di utilizzare massicci supercomputer per simulare una memoria perfetta, hanno insegnato ai modelli a co-adattarsi con una specifica politica di gestione della memoria. Immaginate un bibliotecario che viene addestrato a organizzare libri in una biblioteca con scaffali infiniti, per poi sentirsi dire in seguito di lavorare in una stanza minuscola con un singolo scaffale. Il bibliotecario addestrato nella grande biblioteca probabilmente farebbe fatica a dare priorità a ciò che tenere nella stanza piccola. Al contrario, il metodo proposto in questo articolo addestra il bibliotecario direttamente nella stanza piccola, insegnandogli esattamente quali libri tenere e quali scartare in base alle regole di quello spazio specifico. Ciò consente al modello di apprendere il ritmo dei propri limiti, piuttosto che cercare di disimparare le abitudini derivanti dall'avere troppo spazio.

I ricercatori hanno testato questo approccio su un modello con quattro miliardi di parametri, una dimensione sostanziale ma gestibile. Hanno eseguito i loro esperimenti su una singola scheda grafica con 40 gigabyte di memoria, una configurazione accessibile per molte organizzazioni rispetto ai cluster di decine di schede richiesti dai metodi precedenti. Hanno confrontato il loro nuovo metodo di addestramento con lo standard, che utilizza una tecnica chiamata parallelismo di sequenza (sequence parallelism) per suddividere il carico di memoria su più dispositivi costosi. I risultati hanno mostato che i modelli addestrati con il nuovo metodo spesso superavano quelli standard, particolarmente quando il compito richiedeva al sistema di generare risposte specifiche e concise piuttosto che testi lunghi e prolissi. In diversi test riguardanti domande complesse ed estrazione di dati, il metodo standard produceva output troppo lunghi e pieni di numeri casuali e privi di senso, mentre il nuovo metodo imparava a fermarsi al momento giusto e a fornire l'unico valore corretto.

Una parte chiave di questo successo è stata il miglioramento dell' "heavy-hitter oracle", una popolare strategia per decidere quali informazioni mantenere. Questa strategia funziona tracciando a quali pezzi di informazione il modello presta maggiore attenzione nel tempo. I ricercatori hanno scoperto che la versione originale di questa strategia era lenta ed inefficiente. Hanno riscritto il codice sottostante per farlo funzionare molto più velocemente, permettendo al sistema di calcolare questi punteggi di importanza senza rallentare l'intero processo. Questa ottimizzazione ha significato che il sistema poteva prendere decisioni intelligenti su cosa dimenticare in tempo reale, senza necessitare della massiccia potenza computazionale che solitamente accompagna tali compiti. Il team ha anche rilasciato una nuova libreria software open-source per rendere queste tecniche disponibili ad altri, con l'obiettivo di abbassare la barina per chiunque voglia costruire sistemi di IA a lungo contesto senza aver bisogno di una fortuna in hardware.

Lo studio evidenzia che il modo in cui un modello viene addestrato è importante quanto l'hardware su cui gira. Quando i ricercatori hanno costretto il modello ad addestrarsi con gli stessi identici vincoli di memoria che avrebbe affrontato durante l'uso, esso ha imparato a navigare efficacemente tali vincoli. In un test specifico che coinvolgeva dati JSON, il metodo standard è fallito completamente, incapace di trovare i punti dati corretti, mentre il nuovo metodo è riuscito a identificarli circa la metà delle volte. Ciò suggerisce che la capacità di gestire contesti lunghi non è solo una questione di avere più memoria, ma di insegnare al sistema come gestire la memoria che possiede. Le conclusioni indicano che, per molte applicazioni, la strada più efficace non è costruire computer più grandi, ma insegnare al software a essere più efficiente con le risorse che già possiede.

Riassunto Tecnico: Imparare a Dimenticare: Fine-tuning per l'Attention Sparsa a Lungo Contesto

1. Definizione del Problema

I moderni modelli linguistici di grandi dimensioni (LLM) richiedono la capacità di elaborare contesti estremamente lunghi per compiti quali conversazioni multi-turno, l'uso di strumenti (tool use) e il ragionamento chain-of-thought. Mentre le implementazioni standard dei transformer scalano quadraticamente in termini di calcolo e linearmente in memoria con la larghezza del contesto, sono stati proposti metodi di attenzione sparsa per ottenere una scalabilità temporale lineare e una scalabilità della memoria costante. Questi metodi utilizzano una cache Key-Value (KV) di dimensione fissa in cui le informazioni vengono rimosse (evicted) in base a specifiche policy (ad esempio, Heavy-Hitter Oracle, H2O) una volta che la cache è piena.

Tuttavia, esiste un divario significativo tra l'inferenza e l'addestramento in questo dominio:

Vincoli di Addestramento: Il fine-tuning di modelli con policy di attenzione sparsa è notoriamente difficile con budget hardware moderati. La differenziazione automatica standard (autograd) richiede la memorizzazione dell'intero grafo computazionale, inclusi i buffer della cache KV per ogni chunk della sequenza. Per contesti lunghi, questo requisito di memoria supera la memoria GPU disponibile, rendendo spesso necessario l'uso di costosi setup multi-GPU con parallelismo di sequenza o di contesto (SP/CP).
Gap di Co-adattamento: I modelli addestrati con attenzione esatta (usando SP/CP) spesso performano significativamente peggio quando vengono distribuiti con attenzione sparsa durante l'inferenza. Questo perché il modello non viene mai esposto ai vincoli di memoria specifici e alla logica di rimozione che incontrerà durante l'inferenza.
Inefficienza di Implementazione: Le principali policy di attenzione sparsa come H2O sono sottoutilizzate nelle librerie di inferenza allo stato dell'arte (SotA) (ad esempio, vLLM) perché le attuali implementazioni sono lente, mancando di kernel ottimizzati per i requisiti specifici dell'attenzione sparsa (ad esempio, restituire i pesi di attenzione sommati).

2. Metodologia

Il paper propone un nuovo metodo di fine-tuning che permette ai modelli transformer di co-adattarsi con arbitrarie policy della cache KV su hardware moderato (dimostrato su una singola Nvidia A100 40GB).

2.1 Nested Activation Checkpointing con CPU Offloading

Per affrontare il collo di bottiglia della memoria nell'addestramento su sequenze lunghe ( $N \gg N_C$ , dove $N_C$ è la dimensione del cache), gli autori impiegano una strategia di nested activation checkpointing:

Outer Loop (Layer): Il passaggio backward itera sui layer del modello.
Inner Loop (Cell): La sequenza è divisa in "cell" (gruppi di chunk). Per ogni cell, viene eseguito un forward pass per memorizzare i buffer della cache KV necessari nella memoria della CPU.
Meccanismo di Replay: Invece di differenziare attraverso la policy della cache KV (che è spesso non differenziabile o costosa), il forward pass registra tutte le decisioni di rimozione. Durante il backward pass, queste decisioni vengono riprodotte utilizzando "replay caches" che mimano la logica di rimozione del forward pass senza richiedere il calcolo del gradiente attraverso la policy stessa.

2.2 Delta Encoding per i Buffer della Cache KV

L'ottimizzazione della memoria più critica riguarda il modo in cui i buffer della cache KV vengono memorizzati all'interno del grafo autograd.

Recorrenza Lineare: Gli autori osservano che il buffer della cache KV per un nuovo chunk ($keys'$) è una trasformazione lineare del buffer precedente ($keys$), che differisce solo dai valori dei nuovi token e dagli indici specifici degli slot sovrascritti dalla policy di rimozione.
Storage del Delta: Invece di memorizzare tutti i tensori $keys $e$ values$ per ogni chunk nel grafo autograd, il metodo memorizza solo il delta (la differenza) e gli indici degli slot sovrascritti.
Hook per i Tensori Salvati da Autograd: Utilizzando gli autograd_saved_tensors_hooks di PyTorch, gli autori implementano un meccanismo di "packing". Durante il forward pass, i tensori completi vengono sostituiti da riferimenti a tuple (indice, delta). Durante il backward pass, i tensori completi vengono ricostruiti ("unpacked") dai delta e dai buffer successivi. Ciò riduce l'impronta di memoria del grafo autograd di un fattore $k$ (il numero di chunk per cell), portandola a livelli comparabili all'inferenza.

2.3 Miglioramenti della Policy H2O

Il paper introduce miglioramenti specifici alla policy Heavy-Hitter Oracle (H2O):

Indipendenza dal Batch: L'implementazione originale di H2O aggregava i punteggi attraverso le dimensioni del batch. Il nuovo metodo permette di prendere decisioni di rimozione in modo indipendente per ogni posizione del batch.
Normalizzazione: Viene introdotto uno score H2O normalizzato per prevenire il bias verso i token che sono rimasti in cache più a lungo.
Ottimizzazione del Kernel: Per abilitare un addestramento e un'inferenza efficienti, gli autori forniscono codice Triton che restituisce i pesi di attenzione sommati insieme a un kernel FlashInfer Scaled Dot Product Attention (SDPA). Ciò consente a H2O di girare efficientemente senza ricorrere alle lente implementazioni naive di SDPA.

3. Contributi Chiave

Nuovo Metodo di Fine-Tuning: Un metodo per effettuare il fine-tuning di modelli transformer con attenzione sparsa e arbitrarie policy della cache KV su risorse comparabili all'inferenza con attenzione sparsa (ad esempio, una singola A100). Combina nested activation checkpointing, CPU offloading e delta encoding tramite hook autograd.
Implementazione e Ottimizzazione di H2O: Miglioramenti metodologici alla policy H2O (indipendenza dal batch, normalizzazione) e un'implementazione ad alte prestazioni utilizzando kernel SDPA dedicati che restituiscono i pesi sommati, segnando un passo significativo verso latenze competitive con le librerie di inferenza SotA.
Valutazione Esaustiva: Esperimenti sulla suite di benchmark Helmet (finestre di contesto da 64k e 128k) dimostrano che, sebbene le performance varino a seconda del dataset, il metodo proposto permette ai modelli di co-adattarsi con le policy della cache KV, spesso superando i modelli addestrati con il parallelismo di sequenza quando valutati sotto inferenza con attenzione sparsa.
Libreria KeysAndValues: Il rilascio di una libreria open-source (awslabs/keys_values) che fornisce codice performante e facile da usare per l'inferenza a lungo contesto e il fine-tuning con attenzione sparsa.

4. Risultati Sperimentali

Gli autori hanno valutato un modello Qwen3-4B-Instruct sul benchmark Helmet su lunghezze di contesto di 64k e 128k.

Confronto: I modelli sono stati addestrati usando:
- SP: Parallelismo di Sequenza con attenzione esatta (baseline SotA).
- US: Il metodo proposto con la specifica policy della cache KV in atto.
- NO: Il checkpoint base senza fine-tuning.
Risultati:
- Risultati Misti sui Dataset della Tabella 1: Per i dataset della Tabella 1, i risultati sono stati misti e inconcludenti. La baseline SP (attenzione esatta) ha ottenuto i risultati migliori per nq e hotpot qa. Per trivia qa, il checkpoint base senza fine-tuning (no) ha performato meglio, indicando che il fine-tuning non ha aiutato in questo caso specifico. Il metodo proposto (us) ha superato SP su pop qa.
- Prestazioni Forti sui Dataset della Tabella 2: Per i dataset della Tabella 2, us ha superato nettamente sia sp che no.
- Failure Modes di SP: I modelli addestrati con SP hanno mostrato un failure mode in cui gli output erano eccessivamente lunghi e contenevano nonsense casuale quando sottoposti alla rimozione della cache sparsa, probabilmente perché non erano mai stati addestrati a gestire il meccanismo di "dimenticanza". Al contrario, us ha imparato a terminare correttamente e solitamente produce un singolo numero o il target corretto.
- Varianza della Policy: I risultati sono variati significativamente per dataset. Le varianti di H2O hanno performato bene su diversi task, ma la policy lastrec (lr) ha superato H2O su dataset come nq e hotpot qa. Inoltre, per nq e hotpot qa, la baseline SP (attenzione esatta) ha spesso ottenuto i risultati migliori, mentre per pop qa e altri dataset della Tabella 2, il metodo US con policy di attenzione sparsa era superiore.
- Efficienza: L'implementazione di H2O con il nuovo supporto ai kernel rappresenta un grande miglioramento dell'efficienza, facendo un grande passo verso latenze competitive con le librerie SotA, sebbene il paper riconosca che l'attuale latenza dell'attenzione sparsa rimanga significativamente superiore ai metodi che usano il parallelismo di sequenza o di contesto.

- 5. Significato e Claim

Il paper sostiene che l'attenzione sparsa è stata sottoutilizzata nella pratica principalmente a causa dei problemi di latenza e della difficoltà di addestrare i modelli ad adattarsi a specifiche policy di rimozione.

Fattibilità: Gli autori dimostrano che un fine-tuning di alta qualità per l'attenzione sparsa a lungo contesto è fattibile su hardware moderato, eliminando la necessità di enormi cluster di addestramento distribuito esclusivamente per adattarsi a una policy della cache.
Necessità di Co-adattamento: I risultati suggeriscono fortemente che i modelli devono essere addestrati sotto gli stessi vincoli (policy di rimozione) che affronteranno durante l'inferenza per evitare il degrado delle prestazioni.
Direzione Futura: Il paper sostiene che semplificando la struttura della cache KV e chiarendo i requisiti dei kernel SDPA (ad esempio, restituendo i pesi sommati), la comunità può integrare meglio l'attenzione sparsa nelle librerie di inferenza SotA.
Limitazioni: Gli autori riconoscono che, sebbene il loro metodo migliori la latenza, la natura sequenziale delle decisioni di attenzione sparsa potrebbe ancora risultare in una latenza maggiore rispetto a RingAttention o al Parallelismo di Sequenza per budget hardware molto grandi. Prevedono di affrontare questo tramite la fusione dei kernel e l'offloading asincrono della CPU in lavori futuri.

In sintione, il paper fornisce una via pratica per "imparare a dimenticare", consentendo ai modelli di specializzarsi per l'inferenza a lungo contesto su hardware accessibile. Sebbene le prestazioni rispetto al parallelismo di sequenza varino a seconda del dataset, il metodo spesso supera l'addestramento con attenzione esatta in scenari che richiedono l'adattamento a specifiche policy di attenzione sparsa.