ThinK: Cache delle Chiavi Sottile tramite Potatura Guidata dalle Query

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

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

cs.AI

Abstract

I Large Language Models (LLM) hanno rivoluzionato il campo dell'elaborazione del linguaggio naturale, raggiungendo prestazioni senza precedenti in una varietà di applicazioni grazie all'aumento delle dimensioni dei modelli e delle lunghezze delle sequenze. Tuttavia, l'aumento associato dei costi computazionali e di memoria pone sfide significative, in particolare nella gestione di sequenze lunghe a causa della complessità quadratica del meccanismo di attenzione dei transformer. Questo articolo si concentra sullo scenario di contesto lungo, affrontando le inefficienze nel consumo di memoria della cache KV durante l'inferenza. A differenza degli approcci esistenti che ottimizzano la memoria in base alle lunghezze delle sequenze, scopriamo che la dimensione del canale della cache KV presenta una ridondanza significativa, caratterizzata da una distribuzione di magnitudine squilibrata e da una struttura a basso rango nei pesi di attenzione. Sulla base di queste osservazioni, proponiamo ThinK, un nuovo metodo di pruning della cache KV dipendente dalle query, progettato per minimizzare la perdita dei pesi di attenzione mentre selettivamente elimina i canali meno significativi. Il nostro approccio non solo mantiene o migliora l'accuratezza del modello, ma riduce anche i costi di memoria di oltre il 20% rispetto ai metodi tradizionali di evizione della cache KV. Valutazioni estensive sui modelli LLaMA3 e Mistral su vari dataset di sequenze lunghe confermano l'efficacia di ThinK, stabilendo un nuovo precedente per il dispiegamento efficiente di LLM senza compromettere le prestazioni. Delineiamo inoltre il potenziale di estendere il nostro metodo al pruning della cache dei valori, dimostrando la versatilità e l'ampia applicabilità di ThinK nella riduzione sia della memoria che dei sovraccarichi computazionali.

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.