Una strategia semplice ed efficace basata sulla norma L_2 per la compressione della cache KV

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

June 17, 2024

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

cs.AI

Abstract

L'implementazione di modelli linguistici di grandi dimensioni (LLM) è spesso ostacolata dagli estesi requisiti di memoria della cache Chiave-Valore (KV), specialmente con l'aumento della lunghezza del contesto. Gli approcci esistenti per ridurre le dimensioni della cache KV implicano il fine-tuning del modello per apprendere una strategia di compressione o l'utilizzo dei punteggi di attenzione per ridurre la lunghezza della sequenza. Analizziamo le distribuzioni di attenzione nei modelli basati su Transformer a solo decodifica e osserviamo che i modelli di allocazione dell'attenzione rimangono consistenti nella maggior parte dei livelli. Sorprendentemente, troviamo una chiara correlazione tra la norma L_2 e i punteggi di attenzione sulle coppie KV memorizzate, dove una bassa norma L_2 di un'embedding chiave solitamente porta a un punteggio di attenzione elevato durante la decodifica. Questa scoperta indica che l'influenza di una coppia KV è potenzialmente determinata dall'embedding chiave stessa prima di essere interrogata. Sulla base di questa osservazione, comprimiamo la cache KV in base alla norma L_2 delle embedding chiave. I nostri risultati sperimentali mostrano che questa semplice strategia può ridurre le dimensioni della cache KV del 50% su compiti di modellazione linguistica e di ricerca dell'ago nel pagliaio e del 90% su compiti di recupero di passkey senza perdere accuratezza.

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.