ThinK: Schlankerer Schlüsselspeicher durch Abfragegesteuertes Beschneiden

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben das Feld der natürlichen Sprachverarbeitung revolutioniert, indem sie durch die Nutzung größerer Modellgrößen und Sequenzlängen eine beispiellose Leistung in einer Vielzahl von Anwendungen erzielen. Allerdings stellen die damit verbundenen steigenden Rechen- und Speicherkosten erhebliche Herausforderungen dar, insbesondere bei der Verwaltung langer Sequenzen aufgrund der quadratischen Komplexität des Transformer-Aufmerksamkeitsmechanismus. Dieser Artikel konzentriert sich auf Szenarien mit langem Kontext und befasst sich mit den Ineffizienzen beim Verbrauch des KV-Cache-Speichers während der Inferenz. Im Gegensatz zu bestehenden Ansätzen, die den Speicher basierend auf den Sequenzlängen optimieren, zeigen wir auf, dass die Kanaldimension des KV-Caches eine signifikante Redundanz aufweist, die durch eine unausgeglichene Verteilung der Magnituden und eine niederdimensionale Struktur in den Aufmerksamkeitsgewichten gekennzeichnet ist. Basierend auf diesen Beobachtungen schlagen wir ThinK vor, eine neuartige abfrageabhängige KV-Cache-Pruning-Methode, die darauf abzielt, den Verlust von Aufmerksamkeitsgewichten zu minimieren, während gleichzeitig die unwichtigsten Kanäle selektiv entfernt werden. Unser Ansatz erhält nicht nur die Modellgenauigkeit oder verbessert sie, sondern erzielt auch eine Reduzierung der Speicherkosten um über 20% im Vergleich zu herkömmlichen KV-Cache-Verdrängungsmethoden. Umfangreiche Evaluationen an den Modellen LLaMA3 und Mistral über verschiedene Datensätze mit langen Sequenzen bestätigen die Wirksamkeit von ThinK und setzen einen neuen Maßstab für die effiziente Bereitstellung von LLMs, ohne die Leistung zu beeinträchtigen. Wir skizzieren auch das Potenzial, unsere Methode auf das Pruning des Wertecaches auszudehnen, was die Vielseitigkeit und breite Anwendbarkeit von ThinK bei der Reduzierung sowohl des Speicher- als auch des Rechenaufwands zeigt.

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.