SnapKV: LLM Weet Wat Je Zoekt Voor Generatie

Samenvatting

Grote Taalmodellen (LLMs) hebben opmerkelijke vooruitgang geboekt in het verwerken van uitgebreide contexten, waarbij de Key-Value (KV) cache een cruciale rol speelt bij het verbeteren van hun prestaties. De groei van de KV cache als reactie op toenemende invoerlengte brengt echter uitdagingen met zich mee voor geheugen- en tijdefficiëntie. Om dit probleem aan te pakken, introduceert dit artikel SnapKV, een innovatieve en fine-tuning-vrije aanpak die de grootte van de KV cache efficiënt minimaliseert, terwijl deze nog steeds vergelijkbare prestaties levert in praktische toepassingen. Wij ontdekken dat elke aandachtskop (attention head) in het model consistent gericht is op specifieke prompt-attentiefuncties tijdens het genereren. Tegelijkertijd kan dit robuuste patroon worden verkregen uit een 'observatie'-venster aan het einde van de prompts. Gebaseerd op dit inzicht comprimeert SnapKV automatisch KV caches door gegroepeerde belangrijke KV-posities te selecteren voor elke aandachtskop. Onze aanpak vermindert aanzienlijk de toenemende rekenkosten en geheugenvoetafdruk bij het verwerken van lange invoerreeksen. Specifiek behaalt SnapKV een consistente decodersnelheid met een 3,6x toename in generatiesnelheid en een 8,2x verbetering in geheugenefficiëntie vergeleken met de baseline bij het verwerken van invoer van 16K tokens. Tegelijkertijd behoudt het vergelijkbare prestaties met baseline-modellen over 16 lange-reeksen datasets. Bovendien kan SnapKV tot 380K contexttokens verwerken op een enkele A100-80GB GPU met behulp van een HuggingFace-implementatie met minimale aanpassingen, waarbij slechts een verwaarloosbare nauwkeurigheidsdaling wordt waargenomen in de Needle-in-a-Haystack-test. Verdere uitgebreide studies suggereren het potentieel van SnapKV voor praktische toepassingen.

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.

SnapKV: LLM Weet Wat Je Zoekt Voor Generatie

SnapKV: LLM Knows What You are Looking for Before Generation

Samenvatting

Support