SnapKV: LLM Знает, Что Вы Ищете Перед Генерацией

SnapKV: LLM Knows What You are Looking for Before Generation

April 22, 2024

Авторы: Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, Deming Chen

cs.AI

Аннотация

Большие языковые модели (LLM) достигли значительного прогресса в обработке обширных контекстов, где ключево-значимый (KV) кэш играет важную роль в улучшении их производительности. Однако увеличение размера KV кэша в ответ на увеличение длины ввода создает вызовы для эффективности памяти и времени. Для решения этой проблемы в данной статье представлен SnapKV - инновационный подход, не требующий настройки, который эффективно минимизирует размер KV кэша, сохраняя при этом сопоставимую производительность в реальных приложениях. Мы обнаружили, что каждая голова внимания в модели последовательно фокусируется на конкретных особенностях внимания к подсказке во время генерации. В то же время этот устойчивый паттерн можно получить из `окна наблюдения', расположенного в конце подсказок. Основываясь на этом понимании, SnapKV автоматически сжимает KV кэши, выбирая кластеризованные важные позиции KV для каждой головы внимания. Наш подход значительно снижает растущие вычислительные накладные расходы и объем памяти при обработке длинных последовательностей ввода. Конкретно, SnapKV достигает стабильной скорости декодирования с увеличением скорости генерации в 3,6 раза и улучшением эффективности памяти в 8,2 раза по сравнению с базовым уровнем при обработке вводов из 16K токенов. В то же время он сохраняет сопоставимую производительность с базовыми моделями на 16 длинных последовательностях данных. Более того, SnapKV может обрабатывать до 380K токенов контекста на одном графическом процессоре A100-80GB с использованием реализации HuggingFace с незначительными изменениями, проявляя лишь незначительное падение точности в тесте "Игла в стоге сена". Дальнейшие всесторонние исследования указывают на потенциал SnapKV для практических приложений.

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.