OMVANG: Optimalisatie van Key-Value Cache Compressie bij Generatie met Lange Context

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

December 18, 2024

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

cs.AI

Samenvatting

De Key-Value (KV) cache is een knelpunt geworden voor LLM's voor het genereren van lange contexten. Ondanks de talrijke inspanningen op dit gebied, wordt optimalisatie voor de decoderingsfase over het algemeen genegeerd. Wij zijn echter van mening dat dergelijke optimalisatie cruciaal is, vooral voor taken waarbij lange uitvoer gegenereerd moet worden, gebaseerd op de volgende twee observaties: (i) Overmatige compressie tijdens de voorvul fase, die specifieke volledige context vereist, belemmert het begrip van de redeneertaak; (ii) Afwijking van zware invloeden doet zich voor bij redeneertaken met lange uitvoer. Daarom wordt SCOPE geïntroduceerd, een eenvoudig maar efficiënt raamwerk dat afzonderlijk KV cache optimalisatie uitvoert tijdens de voorvul en decoderingsfases. Specifiek wordt de KV cache tijdens de voorvul fase behouden om de essentiële informatie te behouden, terwijl een nieuwe strategie op basis van schuiven wordt voorgesteld om essentiële zware invloeden te selecteren voor de decoderingsfase. Geheugengebruik en geheugenoverdracht worden verder geoptimaliseerd met adaptieve en discontinue strategieën. Uitgebreide experimenten op LongGenBench tonen de effectiviteit en generalisatie van SCOPE en de compatibiliteit ervan als plug-in voor andere alleen-voorvul KV-compressiemethoden.

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.