ТонкийК: Уменьшение размера кэша ключей с помощью обрезки, управляемой запросами

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

Авторы: Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo

cs.AI

Аннотация

Большие языковые модели (LLM) революционизировали область обработки естественного языка, достигая беспрецедентной производительности в различных приложениях за счет увеличения размеров моделей и длины последовательностей. Однако рост вычислительных и памятных затрат представляет существенные вызовы, особенно в управлении длинными последовательностями из-за квадратичной сложности механизма внимания трансформера. В данной статье рассматривается сценарий с длинным контекстом, решая проблемы неэффективного использования памяти KV-кэша во время вывода. В отличие от существующих подходов, оптимизирующих память на основе длины последовательностей, мы обнаруживаем, что размер канала KV-кэша обладает значительной избыточностью, характеризуемой несбалансированным распределением величин и низкоранговой структурой весов внимания. Исходя из этих наблюдений, мы предлагаем ThinK, новый метод обрезки KV-кэша, зависящий от запроса, разработанный для минимизации потерь веса внимания при селективном обрезании наименее значимых каналов. Наш подход не только поддерживает или улучшает точность модели, но также обеспечивает снижение затрат памяти более чем на 20% по сравнению с обычными методами вытеснения KV-кэша. Обширные оценки на моделях LLaMA3 и Mistral на различных длинных последовательностях подтверждают эффективность ThinK, устанавливая новый прецедент для эффективного развертывания LLM без ущерба производительности. Мы также обрисовываем потенциал расширения нашего метода на обрезку кэша значений, демонстрируя универсальность и широкие возможности ThinK в снижении как памятных, так и вычислительных нагрузок.

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.