NVIDIA, insieme a ricercatori del MIT, ha presentato una nuova tecnica denominata TriAttention, progettata per ridurre drasticamente il consumo di memoria nei modelli LLM intervenendo su uno dei principali limiti strutturali: la gestione del KV cache. L’innovazione punta a mitigare il collo di bottiglia della memoria durante l’inferenza, migliorando al contempo le prestazioni senza compromettere l’accuratezza dei risultati.
Il KV cache rappresenta uno degli elementi fondamentali nel funzionamento dei modelli transformer. Durante l’elaborazione sequenziale del testo, il modello memorizza chiavi e valori associati ai token precedenti per mantenere il contesto della conversazione o del documento. Tuttavia, questo meccanismo comporta una crescita lineare della memoria utilizzata, che aumenta rapidamente con l’estensione del contesto. Nei modelli di grandi dimensioni, l’uso del KV cache può saturare la memoria GPU, limitando la capacità di gestire documenti lunghi o task di ragionamento complessi.
La tecnica TriAttention affronta direttamente questo problema introducendo un nuovo metodo di selezione delle informazioni rilevanti all’interno del KV cache. L’approccio si basa sull’idea che non tutti i token memorizzati contribuiscano allo stesso modo alla generazione della risposta. Utilizzando un filtro matematico specifico, il sistema identifica i dati meno rilevanti e li elimina in modo selettivo, riducendo la memoria necessaria senza degradare la qualità del modello. Questo processo consente di mantenere solo le informazioni più significative per il contesto, ottimizzando l’uso delle risorse hardware.
Uno degli aspetti tecnici più interessanti riguarda la capacità del metodo di operare anche dopo l’applicazione del meccanismo RoPE, utilizzato per codificare le informazioni di posizione nei transformer. In questa fase, i dati risultano mescolati e diventa difficile distinguere gli elementi realmente importanti. TriAttention introduce un filtro matematico che consente di ricostruire i pattern rilevanti anche in questo stato, individuando le caratteristiche chiave da mantenere. In questo modo, il sistema evita la perdita di informazioni critiche, problema che affliggeva precedenti tecniche di compressione.
Il metodo utilizza un approccio basato su funzioni trigonometriche per analizzare i pattern di distanza tra i token. Questo consente al modello di prevedere quali informazioni avranno maggiore rilevanza nelle fasi successive di inferenza. Invece di confrontare ogni elemento con tutti gli altri, il sistema applica un filtro matematico che identifica rapidamente i token meno utili e li elimina. Questo meccanismo riduce il carico computazionale e migliora l’efficienza dell’elaborazione.
I risultati sperimentali evidenziano un miglioramento significativo. Nei test su benchmark complessi di ragionamento matematico, la tecnica ha mantenuto lo stesso livello di accuratezza del modello standard riducendo il consumo del KV cache di circa 10,7 volte. Inoltre, in configurazioni orientate alla velocità, è stato osservato un aumento delle prestazioni fino a 2,5 volte rispetto all’approccio tradizionale full attention. Questi risultati indicano che l’ottimizzazione della memoria può avvenire senza sacrificare la qualità del ragionamento.
Un ulteriore elemento rilevante riguarda la gestione di contesti molto lunghi. Il metodo è stato testato su sequenze di oltre 32.000 token, equivalenti a documenti di dimensioni simili a un intero libro, mantenendo stabilità e precisione. Questa capacità rappresenta un passo importante per applicazioni che richiedono analisi di grandi volumi di testo, come documentazione tecnica, codice o dataset complessi.
Riducendo la memoria necessaria per il KV cache, modelli di dimensioni elevate possono essere eseguiti su GPU con capacità inferiore. In alcuni scenari, modelli da decine di miliardi di parametri possono essere gestiti su una singola GPU consumer con circa 24 GB di VRAM, aprendo la possibilità di eseguire sistemi avanzati anche su infrastrutture meno costose o su workstation locali.
È importante sottolineare che la riduzione del KV cache non elimina completamente i requisiti di memoria complessivi, poiché il modello e altri componenti continuano a occupare spazio. Tuttavia, l’ottimizzazione consente di utilizzare la stessa quantità di VRAM per contesti più lunghi e task più complessi, ampliando le capacità operative dei modelli linguistici. Questo aspetto è particolarmente rilevante per l’adozione di AI on-device e per lo sviluppo di agenti locali ad alte prestazioni.