Простая и эффективная стратегия сжатия кэша KV на основе нормы L_2

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

June 17, 2024

Авторы: Alessio Devoto, Yu Zhao, Simone Scardapane, Pasquale Minervini

cs.AI

Аннотация

Развертывание больших моделей языков (LLM) часто затрудняется из-за обширных требований к памяти кеша ключ-значение (KV), особенно при увеличении длины контекста. Существующие подходы к уменьшению размера кеша KV включают либо донастройку модели для изучения стратегии сжатия, либо использование оценок внимания для сокращения длины последовательности. Мы анализируем распределения внимания в моделях на основе трансформеров только с декодером и наблюдаем, что образцы распределения внимания остаются последовательными на большинстве слоев. Удивительно, мы обнаруживаем четкую корреляцию между L_2 и оценками внимания для кешированных пар KV, где низкий L_2 встраивания ключа обычно приводит к высокой оценке внимания во время декодирования. Это открытие указывает на то, что влияние пары KV потенциально определяется самим встраиванием ключа до его запроса. Основываясь на этом наблюдении, мы сжимаем кеш KV на основе L_2 встраиваний ключей. Наши экспериментальные результаты показывают, что эта простая стратегия может уменьшить размер кеша KV на 50% при языковом моделировании и задачах поиска иголки в стоге и на 90% при задачах поиска ключа без потери точности.

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.