The Pulse
TokenGate riporta un prefill fino a 11,3 volte più veloce sui contesti lunghi
Un gruppo di ricerca presenta TokenGate, un metodo di selezione dei token appreso durante l’addestramento per accelerare il prefill dei modelli linguistici. Lo studio riporta un tempo al primo token fino a 11,3 volte più rapido e una memori
AI.info Team · Scritto con l’IA · verificato ·
In diverse attività, TokenGate ottiene un tempo al primo token (TTFT) da 1,4 a 11,3 volte più rapido e una memoria di picco da 1,5 a 3,4 volte inferiore, senza perdita di accuratezza nel prefill dei transformer.
Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia e Carlos Guestrin, TokenGate
Nello studio del 19 ottobre, gli autori descrivono TokenGate come un modo per consentire ai modelli linguistici di saltare le parti meno utili dei prompt lunghi durante il prefill, l’elaborazione iniziale che precede la generazione di una risposta. Il risultato principale è un intervallo, non un singolo valore: il miglioramento riportato per il tempo al primo token va da 1,4 a 11,3 volte, insieme a una riduzione da 1,5 a 3,4 volte della memoria di picco durante il prefill.
TokenGate impara quali token mantenere
La maggior parte dei modelli transformer elabora ogni token in ingresso attraverso tutti i livelli. TokenGate aggiunge un piccolo componente di valutazione a un livello intermedio: ordina i token, ne mantiene un numero selezionato e invia ai livelli superiori soltanto le rappresentazioni di quei token. I livelli precedenti continuano a elaborare l’intero prompt.
Imparare a effettuare questa selezione pone un problema tecnico: scegliere i K token con il punteggio più alto non è un’operazione differenziabile e quindi, di per sé, non fornisce un gradiente standard per addestrare il selettore. La token-gated attention dello studio affronta il problema usando punteggi di importanza dei token per ripesare i logit dell’attenzione, così da poter addestrare insieme il modello e il meccanismo di selezione. Secondo gli autori, l’addestramento usa un obiettivo di predizione del token successivo e richiede un tempo totale paragonabile al normale fine-tuning efficiente in termini di parametri.
Il risultato di 11,3 volte deriva da un test specifico
Lo studio valuta il completamento di codice, la risposta a domande che richiedono più passaggi, la sintesi di più documenti, l’estrazione di reazioni ai farmaci da testi biomedici e il recupero needle-in-a-haystack. Gli intervalli principali relativi a velocità e memoria confrontano TokenGate con il prefill standard sull’intero contesto, a parità di accuratezza; non significano che ogni attività o impiego raggiungerà il guadagno massimo.
L’appendice dello studio attribuisce il miglioramento di 11,3 volte del tempo al primo token a un test needle-in-a-haystack da 16.000 token con Qwen3-1.7B. Nelle stesse condizioni, gli autori riportano una riduzione di 3,42 volte della memoria di picco. Sottopongono inoltre modelli addestrati su contesti di quella lunghezza a input needle-in-a-haystack più lunghi, riportando un prefill 12,8 volte più rapido con 32.000 token e 13 volte più rapido sia con 64.000 sia con 128.000 token, senza perdita di accuratezza rispetto al prefill sull’intero contesto in quei test.
Rispetto ai precedenti metodi di riferimento più efficaci per la selezione dei token, lo studio riporta un prefill fino a 6,4 volte più rapido e una memoria di picco fino a 2 volte inferiore. Questo confronto è distinto dall’intervallo da 1,4 a 11,3 volte rispetto al prefill sull’intero contesto. La valutazione comprende metodi come FastKV, SlimInfer, DASH e GemFilter, oltre agli approcci di attenzione sparsa FlexPrefill e MInference.
I modelli ibridi ampliano le possibilità di selezione dei token
TokenGate funziona anche su Qwen3.5-2B, un modello ibrido che combina livelli di attenzione completa con livelli Gated-DeltaNet. Gli autori riportano che collocare TokenGate al terzo livello permette di eguagliare l’accuratezza ottenuta sull’intero contesto nel loro test needle-in-a-haystack da 32.000 token, riducendo di 3,8 volte il tempo al primo token. A parità di accuratezza, riportano una latenza 2,1 volte inferiore rispetto alla migliore esecuzione di FastKV.
FastKV applica la selezione dei token ai livelli di attenzione completa, mentre TokenGate può essere collocato anche ai livelli Gated-DeltaNet, secondo lo studio. Spostando il meccanismo di selezione al primo o al secondo livello, l’accelerazione riportata sale fino a 7,8 volte, ma l’accuratezza diminuisce: anziché eguagliare il risultato ottenuto sull’intero contesto, resta entro il 6% di quel risultato. Questo compromesso rende la scelta del livello una parte pratica del metodo, non un semplice dettaglio implementativo.
L’elaborazione iniziale continua a porre un limite alla memoria
TokenGate richiede un adattamento del modello prima dell’impiego; non è un semplice interruttore da attivare durante l’inferenza senza addestramento. Lo studio riporta che ogni passo di addestramento dura circa 1,3 volte quello del metodo di riferimento LoRA, anche se le esecuzioni convergono in meno passi in quattro delle cinque attività e il tempo totale di addestramento varia a seconda dell’attività.
Anche i benefici in termini di memoria hanno un limite intrinseco: i livelli sotto il meccanismo di selezione elaborano ancora il prompt completo e, secondo gli autori, questa elaborazione dell’intero contesto determina in larga misura la memoria di picco. I risultati riportati per TokenGate dipendono quindi da dove inizia la selezione, da quanti token vengono mantenuti e dall’attività misurata, non soltanto dall’accelerazione massima indicata nell’abstract.