AMBITO: Ottimizzazione della compressione della cache chiave-valore nella generazione di contesti lunghi

SCOPE: Optimizing Key-Value Cache Compression in Long-context Generation

December 18, 2024

Autori: Jialong Wu, Zhenglin Wang, Linhai Zhang, Yilong Lai, Yulan He, Deyu Zhou

cs.AI

Abstract

La cache chiave-valore (KV) è diventata un collo di bottiglia per le LLMs nella generazione di contesti lunghi. Nonostante i numerosi sforzi in questo settore, l'ottimizzazione della fase di decodifica viene generalmente ignorata. Tuttavia, riteniamo che tale ottimizzazione sia cruciale, specialmente per compiti di generazione di output lunghi basati sulle seguenti due osservazioni: (i) Eccessiva compressione durante la fase di precaricamento, che richiede un contesto completo specifico, compromette la comprensione del compito di ragionamento; (ii) Si verifica una deviazione dei dati più significativi nei compiti di ragionamento con output lunghi. Pertanto, SCOPE, un framework semplice ma efficiente che esegue separatamente l'ottimizzazione della cache KV durante le fasi di precaricamento e decodifica, viene introdotto. In particolare, la cache KV durante la fase di precaricamento viene preservata per mantenere le informazioni essenziali, mentre viene proposta una nuova strategia basata su scorrimento per selezionare i dati più significativi essenziali per la fase di decodifica. L'utilizzo della memoria e il trasferimento della memoria vengono ulteriormente ottimizzati utilizzando strategie adattive e discontinue. Estesi esperimenti su LongGenBench mostrano l'efficacia e la generalizzazione di SCOPE e la sua compatibilità come plug-in per altri metodi di compressione KV solo per il precaricamento.

English

Key-Value (KV) cache has become a bottleneck of LLMs for long-context generation. Despite the numerous efforts in this area, the optimization for the decoding phase is generally ignored. However, we believe such optimization is crucial, especially for long-output generation tasks based on the following two observations: (i) Excessive compression during the prefill phase, which requires specific full context impairs the comprehension of the reasoning task; (ii) Deviation of heavy hitters occurs in the reasoning tasks with long outputs. Therefore, SCOPE, a simple yet efficient framework that separately performs KV cache optimization during the prefill and decoding phases, is introduced. Specifically, the KV cache during the prefill phase is preserved to maintain the essential information, while a novel strategy based on sliding is proposed to select essential heavy hitters for the decoding phase. Memory usage and memory transfer are further optimized using adaptive and discontinuous strategies. Extensive experiments on LongGenBench show the effectiveness and generalization of SCOPE and its compatibility as a plug-in to other prefill-only KV compression methods.