NVIDIA представила Dynamic Memory Sparsification для 8× сжатия KV кеша в трансформерах

Проблема KV кеша в больших языковых моделях

Трансформерные языковые модели, такие как GPT, LLaMA и Qwen, используют KV кеш для хранения представлений предыдущих токенов при автогрессивной генерации текста. Однако с ростом длины последовательности и числа параллельных потоков размер кеша растет линейно, что ведет к большому потреблению памяти на GPU и замедлению инференса из-за частого доступа к памяти.

Ограничения существующих методов оптимизации KV кеша

Существующие методы оптимизации либо удаляют токены по эвристикам, например, на основе весов внимания, что снижает точность модели, либо применяют сложные постобучающие методы, такие как Dynamic Memory Compression (DMC), которые требуют больших вычислительных ресурсов.

Новый подход: Dynamic Memory Sparsification (DMS)

Исследователи из NVIDIA и Университета Эдинбурга разработали Dynamic Memory Sparsification (DMS) — эффективный и удобный для ретрофита метод сжатия KV кеша без потери точности. DMS реализует разреживание кеша с минимальными затратами на обучение (около 1000 шагов). При этом используется отложенное удаление токенов: после их пометки на удаление они некоторое время остаются в кеше, что сохраняет важный контекст и предотвращает резкое падение точности.

Решения об удалении токенов дифференцируемы во время обучения и основаны на механизме выборки Gumbel-сигмоид. Токены, помеченные на удаление, остаются доступными в течение скользящего окна, что помогает модели лучше усвоить их информацию.

Эффективный ретрофит и минимальное обучение

В отличие от DMC, DMS не требует дополнительных параметров на каждую голову внимания. Используется лишь небольшой элемент механизма внимания (один нейрон) для предсказания удаления токенов. Это упрощает интеграцию DMS в существующие модели без архитектурных изменений.

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

Всего за 1000 шагов обучения DMS достигает сжатия KV кеша в 8 раз, сохраняя или улучшая производительность на задачах, требующих рассуждений. Использованы следующие бенчмарки:

AIME 2024 (продвинутая математика)
MATH 500 (математическое решение задач)
GPQA Diamond (сложные научные вопросы)
LiveCodeBench (генерация кода)

На моделях Qwen-R1 разного размера (1.5B, 7B, 32B) DMS улучшил точность на AIME на 9.1 балла и показал лучшие результаты по сравнению с лидирующими методами Quest и TOVA, при том что затраты памяти и времени работы не увеличились.

Универсальность метода

DMS хорошо работает и на задачах с коротким контекстом, таких как MMLU, GSM8K и HellaSwag, поддерживая производительность при сжатии до 4× с минимальным снижением (~3.5 балла). На задачах с длинным контекстом, например Needle-in-a-Haystack и Variable Tracking, DMS даже превзошел исходные модели, что говорит о способности снижать потерю информации в длинных последовательностях.

Практическое значение

Dynamic Memory Sparsification предлагает масштабируемое и эффективное решение для повышения производительности инференса трансформерных языковых моделей. Умное сжатие KV кеша с минимальным дополнительным обучением позволяет моделям обрабатывать более длинные или параллельные цепочки рассуждений без увеличения времени работы и потребления памяти. Это делает DMS перспективным для использования в условиях ограниченных ресурсов.

Для подробностей ознакомьтесь с оригинальной статьей исследователей NVIDIA и Университета Эдинбурга.