Una Estrategia Simple y Efectiva Basada en la Norma L_2 para la Compresión de la Caché 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

Resumen

El despliegue de modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) a menudo se ve obstaculizado por los extensos requisitos de memoria de la caché de Clave-Valor (KV), especialmente a medida que aumentan las longitudes de contexto. Los enfoques existentes para reducir el tamaño de la caché KV implican ajustar el modelo para aprender una estrategia de compresión o aprovechar las puntuaciones de atención para reducir la longitud de la secuencia. Analizamos las distribuciones de atención en modelos basados en Transformers de solo decodificación y observamos que los patrones de asignación de atención se mantienen consistentes en la mayoría de las capas. Sorprendentemente, encontramos una clara correlación entre la norma L_2 y las puntuaciones de atención sobre los pares KV almacenados en caché, donde una norma L_2 baja de una incrustación de clave generalmente conduce a una puntuación de atención alta durante la decodificación. Este hallazgo indica que la influencia de un par KV está potencialmente determinada por la propia incrustación de clave antes de ser consultada. Basándonos en esta observación, comprimimos la caché KV en función de la norma L_2 de las incrustaciones de clave. Nuestros resultados experimentales muestran que esta estrategia simple puede reducir el tamaño de la caché KV en un 50% en tareas de modelado de lenguaje y de búsqueda de aguja en un pajar, y en un 90% en tareas de recuperación de claves de acceso, sin perder precisión.

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.