LowRAM: chiavi a basso rango e valori ammortizzati per la compressione componibile della cache KV in fase di test

Abstract degli autori

La decodifica autoregressiva si basa sulla cache KV, il cui ingombro in memoria cresce linearmente con la lunghezza del contesto e limita le dimensioni del batch e il throughput. I metodi esistenti di post-training sfruttano la ridondanza tra token, layer e dimensioni delle feature, oppure riducono la precisione numerica. Individuiamo un asse complementare: all’interno di ogni head di attenzione, i valori sono in larga misura prevedibili a partire dalle rispettive chiavi. Sfruttiamo questa proprietà con l’ammortamento della cache dei valori, che adatta al test time un predittore leggero per ciascun head e memorizza i residui solo per i token che ricostruisce con minore accuratezza, ottenendo una compressione dei valori $4\times$ quasi senza perdita. Il predittore può essere inizializzato tramite meta-learning, riducendo notevolmente l’overhead di prefill, oppure addestrato da zero in contesti in cui è ammesso un costo una tantum, come l’inferenza multi-query. Combiniamo l’ammortamento dei valori con la compressione delle chiavi a basso rango tra i layer per ottenere LowRAM, una pipeline di compressione asimmetrica che mantiene prestazioni prossime alla perdita zero su tre famiglie di modelli. Ricostruendo solo un sottoinsieme selezionato di coppie KV a ogni passo di decodifica, LowRAM può mantenere minimo l’overhead di decodifica e raggiungere un throughput di decodifica di picco fino a $5\times$ superiore con una compressione $4\times$. LowRAM si combina inoltre con quantizzazione ed eviction, raggiungendo una compressione di $26.9\times$ con il $94.5\%$ delle prestazioni di riferimento su contesti lunghi. Questi risultati mostrano che considerare chiavi e valori come elementi strutturalmente distinti consente una compressione della cache KV efficiente, aggressiva e componibile.