Ada-KV: Оптимизация вытеснения KV-кэша путём адаптивного распределения бюджета для эффективного вывода LLM

Основная информация

• ID статьи: 2407.11550
• Название: Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
• Авторы: Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, S. Kevin Zhou
• Классификация: cs.CL cs.AI
• Дата публикации/конференция: 39-я конференция по нейронным системам обработки информации (NeurIPS 2025)
• Ссылка на статью: https://arxiv.org/abs/2407.11550

Аннотация

Большие языковые модели (LLM) демонстрируют превосходные результаты в различных областях, однако сталкиваются с проблемами эффективности из-за растущих требований к кэшу Key-Value (KV) при выводе на длинных последовательностях. Недавние исследования снижают размер KV-кэша путём вытеснения во время выполнения большого количества некритических элементов кэша, сохраняя при этом качество генерации. Однако эти методы обычно распределяют бюджет сжатия равномерно между всеми головками внимания, игнорируя уникальные паттерны внимания каждой головки. В данной работе устанавливается теоретическая верхняя граница потерь между выходами внимания до и после вытеснения, объясняются цели оптимизации предыдущих методов вытеснения кэша и направляется оптимизация адаптивного распределения бюджета. На основе этого авторы предлагают Ada-KV — первую стратегию адаптивного распределения бюджета на уровне головки. Метод обладает преимуществом plug-and-play, позволяя беспрепятственно интегрироваться с существующими методами вытеснения кэша.

Исследовательский контекст и мотивация

Описание проблемы

По мере увеличения длины последовательности, обрабатываемой большими языковыми моделями (например, GPT поддерживает 128K, Claude3 поддерживает 200K, Gemini-Pro-1.5 поддерживает 2M токенов), требования к памяти для KV-кэша растут экспоненциально. Для LLM с 8B параметров обработка одной последовательности из 2M токенов может потребовать до 256GB кэша, что серьёзно влияет на эффективность памяти GPU и время выполнения вычислений.

Ограничения существующих методов

Существующие методы вытеснения кэша подразделяются на две основные категории:

1. Методы вытеснения скользящего окна: просто сохраняют начальные и недавние элементы кэша, но значительно снижают качество генерации
2. Методы Top-k вытеснения: выбирают критические элементы кэша на основе весов внимания, но распределяют бюджет равномерно между всеми головками внимания

Ключевая проблема заключается в том, что существующие методы игнорируют уникальные характеристики различных головок внимания: некоторые головки имеют разреженные сконцентрированные паттерны внимания, тогда как другие головки имеют более рассеянное распределение внимания.

Исследовательская мотивация

Анализируя модель Llama-3.1-8B-Instruct, авторы обнаружили, что большинство головок внимания сохраняют почти все веса внимания, используя лишь небольшую долю кэша (например, top 5%), тогда как рассеянные головки требуют большей доли кэша. Этот неравномерный паттерн концентрации внимания обеспечивает теоретическую основу для адаптивного распределения бюджета.

Основные вклады

1. Стратегия адаптивного распределения бюджета: предложена первая стратегия адаптивного распределения бюджета на уровне головки Ada-KV, которая динамически регулирует распределение бюджета в соответствии с уникальными паттернами внимания каждой головки внимания
2. Установление теоретической базы: создана теоретическая база для вытеснения кэша, определены потери при вытеснении и выведена их верхняя граница, объясняются цели оптимизации существующих методов и направляется разработка Ada-KV
3. Совместимость plug-and-play: Ada-KV обладает характеристикой plug-and-play, позволяя беспрепятственно интегрироваться в существующие методы вытеснения кэша, и сохраняет вычислительную эффективность благодаря эффективной реализации CUDA-ядра
4. Комплексная экспериментальная проверка: проведена комплексная оценка на 29 наборах данных из Ruler и LongBench, демонстрирующая значительные улучшения как в сценариях с учётом вопроса, так и без учёта вопроса

Подробное описание метода

Определение задачи

Для многоголовочного слоя самовнимания выбрать элементы KV-кэша, которые следует сохранить в условиях ограничения бюджета, минимизируя потери между выходом внимания после вытеснения и исходным выходом.

Теоретическая база

Определение потерь при вытеснении L1

Авторы количественно определяют потери при вытеснении как L1-расстояние между выходами механизма самовнимания до и после вытеснения:

$\text{L1 Eviction Loss} = ||y - \hat{y}||_1$

где $y$ и $\hat{y}$ — выходы внимания до и после вытеснения соответственно.

Вывод верхней границы потерь

Теорема 3.1: Потери при вытеснении L1 могут быть ограничены верхней границей $\epsilon$:

$\text{L1 Eviction Loss} \leq \epsilon = 2hC - 2C\sum_{i \in [1,h]}\sum_{j \in [1,n]} I_i^j A_i^j$

где $C = \max\{\|V_iW_i^O\|_\infty\}$ — константа, $I_i^j$ — переменная индикатора решения вытеснения, $A_i^j$ — вес внимания.

Теорема 3.2: Метод вытеснения кэша Top-k при заданном распределении бюджета минимизирует верхнюю границу потерь:

$\epsilon^* = 2hC - 2C\sum_{i \in [1,h]}\sum_{A_i^j \in \text{Top-k}(A_i, k=B_i)} A_i^j$

Алгоритм Ada-KV

Алгоритм 1: Адаптивное распределение бюджета

Входные данные: общий бюджет B, веса внимания каждой головки {A_i}
Выходные данные: распределённый бюджет {B_i^*}
1. Объединить веса внимания всех головок: A = Cat({A_i})
2. Выбрать top B весов из A: Top-k(A, k=B)
3. Подсчитать количество выбранных весов для каждой головки: {f_i}
4. Установить распределённый бюджет: {B_i^* = f_i}

Теоретические преимущества

Теорема 3.3: Адаптивное распределение бюджета достигает минимальной верхней границы потерь:

$\epsilon^{**} = \min_{\{B_i\}} \epsilon^*$

Интеграция с существующими методами

Авторы демонстрируют интеграцию Ada-KV с двумя методами SOTA:

Ada-SnapKV и Ada-Pyramid

Через алгоритм 2 Ada-KV может беспрепятственно интегрироваться в SnapKV и Pyramid:

1. Вычислить веса внимания в окне наблюдения
2. Использовать алгоритм Ada-KV для распределения бюджета
3. Применить параметр защиты α = 0.2 для предотвращения чрезмерно разреженного распределения
4. Выполнить решение вытеснения Top-k

Технические инновации

1. Глобальная перспектива оптимизации: рассмотрение распределения бюджета на уровне головки как глобальной задачи оптимизации, а не локальной
2. Теоретически направляемое проектирование: алгоритм разработан на основе строгого теоретического анализа
3. Гарантия вычислительной эффективности: сохранение вычислительной эффективности благодаря FlashAttention переменной длины и сглаженной структуре кэша
4. Совместимость с GQA: поддержка Group Query Attention для дополнительного сжатия кэша

Экспериментальная установка

Наборы данных

• Эталон Ruler: 13 задач на длинных последовательностях, в основном варианты тестов Needle-in-a-Haystack, оценка на длине 16K
• Эталон LongBench: 16 наборов данных, охватывающих QA на одном документе, QA на нескольких документах, суммаризацию, обучение с несколькими примерами, синтетические задачи и генерацию кода

Базовые модели

• Llama-3.1-8B-Instruct
• Mistral-7B-instruct-v0.2

Метрики оценки

В зависимости от типа задачи используются соответствующие метрики: F1-оценка (задачи QA), Rouge-L (суммаризация), точность (классификация), редакционное сходство (задачи кодирования)

Методы сравнения

• Базовые методы: SnapKV, Pyramid, StreamingLLM
• Улучшенные версии: Ada-SnapKV, Ada-Pyramid

Экспериментальные сценарии

• Сжатие с учётом вопроса: стандартный сценарий, когда вопрос известен
• Сжатие без учёта вопроса: более сложный сценарий реального применения

Результаты экспериментов

Основные результаты

Тестирование на эталоне Ruler

В сценарии без учёта вопроса с использованием Llama-3.1-8B-Instruct:

• Бюджет 80%: Ada-SnapKV повышает оценку SnapKV с 87.59 до 92.67
• Бюджет 20%: Ada-SnapKV повышает оценку SnapKV с 44.02 до 53.29

Тестирование на эталоне LongBench

В сценарии без учёта вопроса:

• Ada-SnapKV и Ada-Pyramid постоянно улучшают качество генерации при всех фиксированных установках бюджета
• При бюджете 2048 достигается почти безубыточная производительность

Анализ подзадач

В сложной задаче Needle-in-a-Haystack:

• Задача S-NIAH-3 (бюджет 80%): Ada-SnapKV повышает SnapKV с 62.4 до 97.6
• Задача MK-NIAH-2 (бюджет 80%): Ada-SnapKV повышает SnapKV с 85.2 до 99.6

Вычислительная эффективность

Ada-SnapKV при фиксированном бюджете 1024:

• Пиковое использование памяти сопоставимо с исходным SnapKV
• Задержка декодирования сопоставима с исходным SnapKV
• Оба значительно превосходят сценарий с полным кэшем

Проверка широкого применения

Стратегия Ada-KV была принята несколькими последующими работами:

• CriticalKV + Ada-KV: повышение с 42.99 до 43.77 при 20% кэша
• DefensiveKV + Ada-KV: повышение с 43.78 до 46.68 при 20% кэша

Связанные работы

Методы вытеснения кэша

• Методы скользящего окна: StreamingLLM и др., простые, но с большими потерями качества
• Методы Top-k: H2O, SnapKV, Pyramid и др., выбирают критические элементы на основе весов внимания

Методы разреженного внимания

Концептуально связаны с вытеснением кэша, но используют другие подходы:

• Вытеснение кэша: сохранение подмножества KV-кэша
• Разреженное внимание: сохранение всех записей, но выборочное использование

Другие связанные технологии

• Квантизация KV-кэша: снижение точности отдельных элементов
• Спекулятивный декодинг: использование моделей с уменьшенным кэшем для генерации черновиков
• Постраничное внимание: эффективное управление памятью

Заключение и обсуждение

Основные выводы

1. Ada-KV впервые предложила стратегию адаптивного распределения бюджета на уровне головки, значительно улучшив производительность существующих методов вытеснения кэша
2. Теоретический анализ устанавливает строгую базу для вытеснения кэша, направляя разработку алгоритма
3. Сценарий сжатия без учёта вопроса выявляет ограничения существующих методов, которым следует уделить больше внимания

Ограничения

1. Текущее распределение на уровне головки ограничено одним слоем, не расширяется на распределение между слоями
2. Параметр защиты α требует компромисса производительности при различных бюджетах
3. Теоретический анализ основан на L1-расстоянии, что может не полностью отражать фактическое качество генерации

Направления будущих исследований

1. Расширение механизма распределения на уровне головки на сценарии между слоями
2. Разработка соответствующего теоретического анализа между слоями
3. Объединение анализа важности головки во время обучения
4. Совместная оптимизация с другими методами оптимизации (квантизация, разреженное внимание)

Глубокая оценка

Преимущества

1. Прочный теоретический вклад: установлена полная теоретическая база, логика от вывода верхней границы потерь до разработки алгоритма ясна
2. Простой и эффективный метод: алгоритм прост и понятен, характеристика plug-and-play облегчает его принятие
3. Комплексная и достаточная экспериментальная проверка: комплексная оценка на 29 наборах данных, включая часто упускаемый сценарий без учёта вопроса
4. Высокая практическая ценность: принята несколькими последующими работами, что подтверждает ценность и влияние метода

Недостатки

1. Разрыв между теорией и практикой: хотя теоретически минимизируются потери верхней границы, это не гарантирует минимизацию фактических потерь
2. Чувствительность к гиперпараметрам: выбор параметра защиты α требует эмпирической настройки
3. Ограничения масштабируемости: в настоящее время рассматривается только распределение бюджета в одном слое
4. Ограничения оценки: основная оценка проводится на моделях среднего размера, эффективность на крупномасштабных моделях требует проверки

Влияние

1. Академический вклад: предоставляет новое направление исследований для области оптимизации KV-кэша
2. Практическая ценность: характеристика plug-and-play облегчает развёртывание в реальных системах
3. Воспроизводимость: предоставляет открытый исходный код и подробные детали реализации
4. Вдохновляющее значение: предоставляет теоретическую базу и методологическое руководство для последующих исследований

Применимые сценарии

1. Вывод на длинных последовательностях: особенно подходит для приложений, требующих обработки длинного контекста
2. Среды с ограниченными ресурсами: оптимизация эффективности вывода при ограниченной памяти GPU
3. Системы реального времени: балансировка качества и эффективности в онлайн-сервисах
4. Многораундовые диалоги: сценарий сжатия без учёта вопроса особенно подходит для систем диалога

Библиография

Статья цитирует 64 связанные работы, включая в основном:

• Фундаментальные работы по большим языковым моделям (GPT-4, Claude, Gemini и др.)
• Методы оптимизации KV-кэша (H2O, SnapKV, Pyramid и др.)
• Оптимизацию механизма внимания (FlashAttention, разреженное внимание и др.)
• Эталоны обработки длинных последовательностей (Ruler, LongBench и др.)

Общая оценка: Это высококачественная исследовательская статья, достигшая хорошего баланса между теоретическим вкладом и практической ценностью. Метод Ada-KV прост и эффективен, теоретический анализ строг, экспериментальная проверка полна. Статья не только решает важные ограничения существующих методов, но и предоставляет ценную базу и направления для будущих исследований.