Uma Estratégia Simples e Eficaz Baseada na Norma L_2 para Compressão do Cache KV

A Simple and Effective L_2 Norm-Based Strategy for KV Cache Compression

June 17, 2024

Autores: Alessio Devoto, Yu Zhao, Simone Scardapane, Pasquale Minervini

cs.AI

Resumo

A implantação de grandes modelos de linguagem (LLMs) é frequentemente dificultada pelos extensos requisitos de memória do cache de Chave-Valor (KV), especialmente à medida que os comprimentos de contexto aumentam. As abordagens existentes para reduzir o tamanho do cache KV envolvem tanto o ajuste fino do modelo para aprender uma estratégia de compressão quanto a utilização de pontuações de atenção para reduzir o comprimento da sequência. Analisamos as distribuições de atenção em modelos baseados em Transformers de apenas decodificação e observamos que os padrões de alocação de atenção permanecem consistentes na maioria das camadas. Surpreendentemente, encontramos uma clara correlação entre a norma L_2 e as pontuações de atenção sobre os pares KV armazenados em cache, onde uma norma L_2 baixa de uma incorporação de chave geralmente resulta em uma pontuação de atenção alta durante a decodificação. Essa descoberta indica que a influência de um par KV é potencialmente determinada pela própria incorporação da chave antes de ser consultada. Com base nessa observação, comprimimos o cache KV com base na norma L_2 das incorporações de chave. Nossos resultados experimentais mostram que essa estratégia simples pode reduzir o tamanho do cache KV em 50% em tarefas de modelagem de linguagem e de "agulha no palheiro" e em 90% em tarefas de recuperação de senha sem perda de precisão.

English

The deployment of large language models (LLMs) is often hindered by the extensive memory requirements of the Key-Value (KV) cache, especially as context lengths increase. Existing approaches to reduce the KV cache size involve either fine-tuning the model to learn a compression strategy or leveraging attention scores to reduce the sequence length. We analyse the attention distributions in decoder-only Transformers-based models and observe that attention allocation patterns stay consistent across most layers. Surprisingly, we find a clear correlation between the L_2 and the attention scores over cached KV pairs, where a low L_2 of a key embedding usually leads to a high attention score during decoding. This finding indicates that the influence of a KV pair is potentially determined by the key embedding itself before being queried. Based on this observation, we compress the KV cache based on the L_2 of key embeddings. Our experimental results show that this simple strategy can reduce the KV cache size by 50% on language modelling and needle-in-a-haystack tasks and 90% on passkey retrieval tasks without losing accuracy.