ALCANCE: Optimización de la Compresión de la Caché de Clave-Valor en la Generación de Contextos Largos

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

December 18, 2024

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

cs.AI

Resumen

La caché de Clave-Valor (KV) se ha convertido en un cuello de botella para los LLMs en la generación de contexto largo. A pesar de los numerosos esfuerzos en esta área, la optimización para la fase de decodificación generalmente se ignora. Sin embargo, creemos que dicha optimización es crucial, especialmente para tareas de generación de salidas largas basadas en las siguientes dos observaciones: (i) La compresión excesiva durante la fase de precarga, que requiere un contexto completo específico, perjudica la comprensión de la tarea de razonamiento; (ii) Se produce una desviación de los elementos más importantes en las tareas de razonamiento con salidas largas. Por lo tanto, se introduce SCOPE, un marco simple pero eficiente que realiza la optimización de la caché KV de forma separada durante las fases de precarga y decodificación. Específicamente, la caché KV durante la fase de precarga se conserva para mantener la información esencial, mientras que se propone una estrategia novedosa basada en deslizamiento para seleccionar los elementos más importantes esenciales para la fase de decodificación. El uso de memoria y la transferencia de memoria se optimizan aún más utilizando estrategias adaptativas y discontinuas. Experimentos extensos en LongGenBench muestran la efectividad y generalización de SCOPE y su compatibilidad como un complemento para otros métodos de compresión KV solo para precarga.

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.