Een Eenvoudige en Effectieve L_2 Norm-Gebaseerde Strategie voor KV Cache Compressie

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

June 17, 2024

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

cs.AI

Samenvatting

De inzet van grote taalmodellen (LLMs) wordt vaak belemmerd door de uitgebreide geheugeneisen van de Key-Value (KV) cache, vooral naarmate de contextlengtes toenemen. Bestaande benaderingen om de KV-cachegrootte te verkleinen, omvatten ofwel het finetunen van het model om een compressiestrategie te leren, ofwel het gebruikmaken van attentiescores om de sequentielengte te reduceren. Wij analyseren de aandachtverdelingen in decoder-only Transformers-gebaseerde modellen en merken op dat de aandachtstoewijzingspatronen consistent blijven over de meeste lagen. Verrassend genoeg vinden we een duidelijke correlatie tussen de L_2 en de attentiescores over gecachete KV-paren, waarbij een lage L_2 van een key-embedding meestal leidt tot een hoge attentiescore tijdens het decoderen. Deze bevinding geeft aan dat de invloed van een KV-paar mogelijk wordt bepaald door de key-embedding zelf voordat deze wordt bevraagd. Op basis van deze observatie comprimeren we de KV-cache op basis van de L_2 van key-embeddings. Onze experimentele resultaten tonen aan dat deze eenvoudige strategie de KV-cachegrootte kan verkleinen met 50% bij taalmodellering en needle-in-a-haystack taken, en met 90% bij passkey-retrieval taken, zonder verlies van nauwkeurigheid.

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.