Latent Compaction con addestramento al momento del test

Abstract degli autori

I contesti lunghi fanno crescere le cache chiave-valore (KV), aumentando l’uso della memoria e i costi dell’attenzione durante la generazione. Comprimere queste cache richiede di preservare le informazioni necessarie per rispondere a domande non ancora viste. Presentiamo **Latent Compaction**, un metodo che apprende una cache compatta specifica per il documento tramite addestramento al momento del test. Inizializziamo le voci della cache in base alla loro importanza per la ricostruzione, all’attenzione del teacher e alla diversità nello spazio delle chiavi. L’addestramento adatta dapprima gli output dell’attenzione e i log-normalizzatori a quelli della cache completa, mantenendo fisse le voci selezionate in base alla ricostruzione. Ottimizza poi congiuntamente l’intera cache tramite distillazione end-to-end dal modello con la cache completa. I pesi del modello restano fissi per tutto il processo e l’addestramento usa soltanto prompt generici sul documento, senza accesso alle domande o alle risposte a valle. Gli esperimenti su RULER, LongBench e Phonebook Lookup mostrano miglioramenti sostanziali rispetto ai metodi di compressione di riferimento con lo stesso budget. Su RULER, per esempio, il nostro metodo raggiunge un’accuratezza del 96,7% con una compressione di 32$\times$, contro il 40,0% di KVZip+. Il metodo si trasferisce anche tra modelli di diverse dimensioni e famiglie, con miglioramenti nel recupero delle informazioni sia su Qwen3 sia su MiniCPM5. Questi risultati mostrano che trattare la cache KV come una rete neurale e ottimizzarla con l’addestramento al momento del test può preservare l’accuratezza nel recupero delle informazioni anche con una compressione elevata. Il codice è disponibile all’indirizzo https://www.github.com/THUNLP/latent-compaction