SnapKV: El LLM sabe lo que estás buscando antes de la generación

SnapKV: LLM Knows What You are Looking for Before Generation

April 22, 2024

Autores: Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, Deming Chen

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) han logrado avances notables en el procesamiento de contextos extensos, donde la caché Clave-Valor (KV) desempeña un papel crucial para mejorar su rendimiento. Sin embargo, el crecimiento de la caché KV en respuesta al aumento de la longitud de la entrada plantea desafíos para la eficiencia de memoria y tiempo. Para abordar este problema, este artículo presenta SnapKV, un enfoque innovador y libre de ajuste fino que minimiza eficientemente el tamaño de la caché KV mientras mantiene un rendimiento comparable en aplicaciones del mundo real. Descubrimos que cada cabeza de atención en el modelo se enfoca consistentemente en características específicas de atención del prompt durante la generación. Además, este patrón robusto puede obtenerse a partir de una ventana de `observación' ubicada al final de los prompts. Basándonos en esta idea, SnapKV comprime automáticamente las cachés KV seleccionando posiciones importantes agrupadas para cada cabeza de atención. Nuestro enfoque reduce significativamente la sobrecarga computacional y la huella de memoria al procesar secuencias de entrada largas. Específicamente, SnapKV logra una velocidad de decodificación constante con un aumento de 3.6x en la velocidad de generación y una mejora de 8.2x en la eficiencia de memoria en comparación con la línea base al procesar entradas de 16K tokens. Al mismo tiempo, mantiene un rendimiento comparable a los modelos de referencia en 16 conjuntos de datos de secuencias largas. Además, SnapKV puede procesar hasta 380K tokens de contexto en una sola GPU A100-80GB utilizando la implementación de HuggingFace con cambios mínimos, mostrando solo una caída insignificante en la precisión en la prueba Needle-in-a-Haystack. Estudios adicionales sugieren el potencial de SnapKV para aplicaciones prácticas.

English

Large Language Models (LLMs) have made remarkable progress in processing extensive contexts, with the Key-Value (KV) cache playing a vital role in enhancing their performance. However, the growth of the KV cache in response to increasing input length poses challenges to memory and time efficiency. To address this problem, this paper introduces SnapKV, an innovative and fine-tuning-free approach that efficiently minimizes KV cache size while still delivering comparable performance in real-world applications. We discover that each attention head in the model consistently focuses on specific prompt attention features during generation. Meanwhile, this robust pattern can be obtained from an `observation' window located at the end of the prompts. Drawing on this insight, SnapKV automatically compresses KV caches by selecting clustered important KV positions for each attention head. Our approach significantly reduces the growing computational overhead and memory footprint when processing long input sequences. Specifically, SnapKV achieves a consistent decoding speed with a 3.6x increase in generation speed and an 8.2x enhancement in memory efficiency compared to baseline when processing inputs of 16K tokens. At the same time, it maintains comparable performance to baseline models across 16 long sequence datasets. Moreover, SnapKV can process up to 380K context tokens on a single A100-80GB GPU using HuggingFace implementation with minor changes, exhibiting only a negligible accuracy drop in the Needle-in-a-Haystack test. Further comprehensive studies suggest SnapKV's potential for practical applications.