ESCALA: Otimização da Compressão de Cache Chave-Valor na Geração de Contexto Longo

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

Resumo

O cache Chave-Valor (KV) tornou-se um gargalo para os LLMs na geração de contexto longo. Apesar dos inúmeros esforços nessa área, a otimização para a fase de decodificação é geralmente ignorada. No entanto, acreditamos que essa otimização é crucial, especialmente para tarefas de geração de saída longa com base nas seguintes duas observações: (i) Compressão excessiva durante a fase de pré-preenchimento, que requer um contexto completo específico, prejudica a compreensão da tarefa de raciocínio; (ii) Desvio de elementos pesados ocorre nas tarefas de raciocínio com saídas longas. Portanto, SCOPE, um framework simples, porém eficiente, que realiza a otimização do cache KV separadamente durante as fases de pré-preenchimento e decodificação, é introduzido. Especificamente, o cache KV durante a fase de pré-preenchimento é preservado para manter as informações essenciais, enquanto uma estratégia inovadora baseada em deslizamento é proposta para selecionar os elementos pesados essenciais para a fase de decodificação. O uso de memória e a transferência de memória são ainda otimizados usando estratégias adaptativas e descontínuas. Experimentos extensivos no LongGenBench mostram a eficácia e generalização do SCOPE e sua compatibilidade como um plug-in para outros métodos de compressão KV apenas de pré-preenchimento.

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.