SnapKV: LLM weiß, wonach Sie suchen, bevor die Generierung erfolgt.

SnapKV: LLM Knows What You are Looking for Before Generation

April 22, 2024

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben bemerkenswerte Fortschritte bei der Verarbeitung umfangreicher Kontexte erzielt, wobei der Schlüssel-Wert (KV) Cache eine wichtige Rolle bei der Verbesserung ihrer Leistung spielt. Das Wachstum des KV-Caches als Reaktion auf die zunehmende Eingabegröße stellt jedoch Herausforderungen für die Speicher- und Zeiteffizienz dar. Um dieses Problem anzugehen, stellt dieser Artikel SnapKV vor, einen innovativen und feinabstimmungsfreien Ansatz, der die Größe des KV-Caches effizient minimiert, während er dennoch vergleichbare Leistungen in realen Anwendungen erbringt. Wir entdecken, dass jeder Aufmerksamkeitskopf im Modell während der Generierung konsistent auf bestimmte promptbezogene Aufmerksamkeitsmerkmale fokussiert. Gleichzeitig kann dieses robuste Muster aus einem 'Beobachtungs'-Fenster am Ende der Prompts gewonnen werden. Basierend auf dieser Erkenntnis komprimiert SnapKV automatisch KV-Caches, indem es für jeden Aufmerksamkeitskopf gruppierte wichtige KV-Positionen auswählt. Unser Ansatz reduziert signifikant den wachsenden Rechenaufwand und den Speicherbedarf bei der Verarbeitung langer Eingabesequenzen. Speziell erreicht SnapKV eine konsistente Dekodierungsgeschwindigkeit mit einer 3,6-fachen Steigerung der Generierungsgeschwindigkeit und einer 8,2-fachen Verbesserung der Speichereffizienz im Vergleich zur Basislinie bei der Verarbeitung von 16K Tokens. Gleichzeitig erzielt er vergleichbare Leistungen wie Basislinienmodelle über 16 Datensätze langer Sequenzen. Darüber hinaus kann SnapKV mit geringfügigen Änderungen bis zu 380K Kontext-Token auf einer einzelnen A100-80GB GPU unter Verwendung der HuggingFace-Implementierung verarbeiten und zeigt nur einen vernachlässigbaren Genauigkeitsverlust im Needle-in-a-Haystack-Test. Weitere umfassende Studien deuten auf das Potenzial von SnapKV für praktische Anwendungen hin.

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.