ThinK: Cache de Chave Mais Fina Através de Poda Orientada por Consulta

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

Autores: Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo

cs.AI

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) revolucionaram o campo do processamento de linguagem natural, alcançando um desempenho sem precedentes em uma variedade de aplicações, aproveitando tamanhos de modelo maiores e comprimentos de sequência aumentados. No entanto, o aumento associado nos custos computacionais e de memória apresenta desafios significativos, especialmente na gestão de sequências longas devido à complexidade quadrática do mecanismo de atenção do transformador. Este artigo concentra-se no cenário de longo contexto, abordando as ineficiências no consumo de memória cache KV durante a inferência. Ao contrário das abordagens existentes que otimizam a memória com base nos comprimentos das sequências, descobrimos que a dimensão do canal do cache KV apresenta uma redundância significativa, caracterizada por uma distribuição de magnitude desequilibrada e uma estrutura de baixa classificação nos pesos de atenção. Com base nessas observações, propomos o ThinK, um novo método de poda de cache KV dependente de consulta projetado para minimizar a perda de peso de atenção enquanto poda seletivamente os canais menos significativos. Nossa abordagem não apenas mantém ou aprimora a precisão do modelo, mas também alcança uma redução nos custos de memória em mais de 20% em comparação com métodos convencionais de evicção de cache KV. Avaliações extensivas nos modelos LLaMA3 e Mistral em vários conjuntos de dados de sequências longas confirmam a eficácia do ThinK, estabelecendo um novo precedente para a implantação eficiente de LLM sem comprometer o desempenho. Também delineamos o potencial de estender nosso método para a poda de cache de valor, demonstrando a versatilidade e ampla aplicabilidade do ThinK na redução tanto dos custos de memória quanto dos custos computacionais.

English

Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications by leveraging increased model sizes and sequence lengths. However, the associated rise in computational and memory costs poses significant challenges, particularly in managing long sequences due to the quadratic complexity of the transformer attention mechanism. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. Unlike existing approaches that optimize the memory based on the sequence lengths, we uncover that the channel dimension of the KV cache exhibits significant redundancy, characterized by unbalanced magnitude distribution and low-rank structure in attention weights. Based on these observations, we propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in memory costs by over 20% compared with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and Mistral models across various long-sequence datasets confirm the efficacy of ThinK, setting a new precedent for efficient LLM deployment without compromising performance. We also outline the potential of extending our method to value cache pruning, demonstrating ThinK's versatility and broad applicability in reducing both memory and computational overheads.