TokenGate: adattare gli LLM a un prefill rapido ed efficiente nell’uso della memoria con la potatura appresa dei token

Abstract degli autori

I lettori esperti esaminano rapidamente input di grandi dimensioni, come repository di codice e documenti lunghi, per leggere con attenzione solo ciò che conta. I modelli linguistici, invece, elaborano tutti i token dell’input attraverso ogni livello del modello, con costi di elaborazione elevati per gli input a contesto lungo. Presentiamo TokenGate, un piccolo componente facilmente integrabile che estende modelli già addestrati e consente una potatura dei token apprendibile end-to-end sugli input a contesto lungo, per un prefill rapido ed efficiente nell’uso della memoria. Per rendere possibile l’apprendimento end-to-end nonostante l’operatore top-K non differenziabile usato per la potatura dei token, introduciamo l’attenzione con gating dei token, un meccanismo ispirato a Mixture-of-Experts che ripesa i logit dell’attenzione in base ai punteggi di importanza dei token. In diverse attività, TokenGate ottiene un tempo di generazione del primo token (TTFT) più rapido di $1.4$-$11.3\times$ e un picco di utilizzo della memoria inferiore di $1.5$-$3.4\times$, senza perdita di accuratezza nel prefill dei transformer. Rispetto ai precedenti metodi di potatura dei token con le prestazioni migliori, TokenGate offre un prefill fino a $6.4\times$ più rapido e un picco di utilizzo della memoria inferiore di $2\times$. TokenGate estende nativamente sia i transformer standard, come Qwen3, sia i modelli ad attenzione ibrida, come Qwen3.5, con un adattamento semplice e leggero, paragonabile al fine-tuning standard efficiente in termini di parametri. Il nostro codice è disponibile all’indirizzo: https://github.com/tokengate-models/TokenGate.