ThinK: Caché de Claves Más Delgada mediante Poda Dirigida por Consulta

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

Autores: Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Tamaño (LLMs) han revolucionado el campo del procesamiento de lenguaje natural, logrando un rendimiento sin precedentes en una variedad de aplicaciones al aprovechar tamaños de modelo aumentados y longitudes de secuencia extendidas. Sin embargo, el aumento asociado en costos computacionales y de memoria plantea desafíos significativos, especialmente en la gestión de secuencias largas debido a la complejidad cuadrática del mecanismo de atención del transformador. Este artículo se centra en el escenario de contexto largo, abordando las ineficiencias en el consumo de memoria caché KV durante la inferencia. A diferencia de enfoques existentes que optimizan la memoria en función de las longitudes de secuencia, descubrimos que la dimensión del canal de la caché KV muestra una redundancia significativa, caracterizada por una distribución desequilibrada de magnitudes y una estructura de bajo rango en los pesos de atención. Basándonos en estas observaciones, proponemos ThinK, un novedoso método de poda de caché KV dependiente de la consulta diseñado para minimizar la pérdida de peso de atención al podar selectivamente los canales menos significativos. Nuestro enfoque no solo mantiene o mejora la precisión del modelo, sino que también logra una reducción en los costos de memoria de más del 20% en comparación con los métodos convencionales de evicción de caché KV. Evaluaciones exhaustivas en los modelos LLaMA3 y Mistral en diversos conjuntos de datos de secuencias largas confirman la eficacia de ThinK, estableciendo un nuevo precedente para la implementación eficiente de LLM sin comprometer el rendimiento. También esbozamos el potencial de extender nuestro método a la poda de caché de valor, demostrando la versatilidad y amplia aplicabilidad de ThinK en la reducción tanto de la memoria como de la carga computacional.

English

Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications by leveraging increased model sizes and sequence lengths. However, the associated rise in computational and memory costs poses significant challenges, particularly in managing long sequences due to the quadratic complexity of the transformer attention mechanism. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. Unlike existing approaches that optimize the memory based on the sequence lengths, we uncover that the channel dimension of the KV cache exhibits significant redundancy, characterized by unbalanced magnitude distribution and low-rank structure in attention weights. Based on these observations, we propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in memory costs by over 20% compared with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and Mistral models across various long-sequence datasets confirm the efficacy of ThinK, setting a new precedent for efficient LLM deployment without compromising performance. We also outline the potential of extending our method to value cache pruning, demonstrating ThinK's versatility and broad applicability in reducing both memory and computational overheads.