ThinK: Dunnere Sleutelcache door Query-Gestuurde Snoei

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

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

cs.AI

Samenvatting

Grote Taalmodellen (LLM's) hebben een revolutie teweeggebracht in het vakgebied van natuurlijke taalverwerking, waarbij ze ongekende prestaties hebben bereikt in diverse toepassingen door gebruik te maken van grotere modelgroottes en langere sequentielengtes. De daarmee gepaard gaande stijging in reken- en geheugenkosten brengt echter aanzienlijke uitdagingen met zich mee, met name bij het beheren van lange sequenties vanwege de kwadratische complexiteit van het transformer-attentiemechanisme. Dit artikel richt zich op het scenario van lange contexten en behandelt de inefficiënties in het geheugengebruik van de KV-cache tijdens inferentie. In tegenstelling tot bestaande benaderingen die het geheugen optimaliseren op basis van de sequentielengtes, ontdekken wij dat de kanaaldimensie van de KV-cache aanzienlijke redundantie vertoont, gekenmerkt door een onbalans in de grootteverdeling en een laag-rangstructuur in de aandachtswaarden. Op basis van deze observaties stellen wij ThinK voor, een nieuwe query-afhankelijke KV-cache pruning-methode die is ontworpen om het verlies van aandachtswaarden te minimaliseren terwijl de minst significante kanalen selectief worden verwijderd. Onze aanpak behoudt of verbetert niet alleen de modelnauwkeurigheid, maar reduceert ook de geheugenkosten met meer dan 20% vergeleken met standaard KV-cache verwijderingsmethoden. Uitgebreide evaluaties op de LLaMA3- en Mistral-modellen over diverse lange-sequentie datasets bevestigen de effectiviteit van ThinK, waarmee een nieuwe standaard wordt gezet voor efficiënte LLM-implementatie zonder in te leveren op prestaties. We schetsen ook het potentieel van het uitbreiden van onze methode naar value-cache pruning, wat de veelzijdigheid en brede toepasbaarheid van ThinK aantoont in het verminderen van zowel geheugen- als rekenoverhead.

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.