По мере роста спроса на задачи, требующие рассуждений, от больших языковых моделей (LLM) всё чаще требуется генерировать более длинные последовательности или параллельные цепочки рассуждений. Однако производительность во время вывода сильно ограничена объёмом памяти, занимаемым KV-кэшем, а не только количеством генерируемых токенов.
В недавней статье исследователи из NVIDIA и Эдинбургского университета представляют Dynamic Memory Sparsification (DMS) — эффективный с точки зрения данных метод, который сжимает KV-кэши и открывает возможности гипермасштабирования во время вывода без снижения точности модели.
Узкое место: KV-кэш в выводе трансформеров
Модели на основе трансформеров, такие как GPT, LLaMA и Qwen, используют KV-кэши для хранения представлений прошлых токенов для авторегрессивной генерации. Этот кэш растёт линейно с длиной последовательности и шириной (параллельными потоками), потребляя много памяти GPU и приводя к замедлению вывода из-за частого доступа к памяти.
Существующие методы оптимизации KV-кэша либо основаны на эвристиках без обучения, таких как вытеснение токенов на основе весов внимания, либо требуют сложных доработок после обучения, например, динамической компрессии памяти (DMC). У обоих подходов есть существенные недостатки: первый снижает точность, а второй требует больших вычислительных затрат.
Динамическая разреженность памяти DMS: сжатие без компромиссов
Dynamic Memory Sparsification (DMS) устраняет эти ограничения с помощью гибридного подхода: она делает KV-кэш более разреженным, как традиционные методы обрезки, но с минимальными затратами на обучение (~1000 шагов) и отсроченным вытеснением, которое временно сохраняет токены после их маркировки для удаления. Такая схема сохраняет важную контекстную информацию и позволяет избежать резких падений точности.
Основная идея заключается в том, чтобы сделать решения о вытеснении дифференцируемыми во время обучения с помощью механизма выборки на основе Gumbel-sigmoid. Токены, предсказанные для будущего вытеснения, остаются доступными в течение определённого времени, прежде чем будут удалены, что позволяет модели более эффективно усваивать их информационную ценность.
В отличие от DMC, которая требует тысяч шагов обучения и сложной оптимизации на основе градиентов, DMS не вводит дополнительных параметров для каждой головы внимания. Она повторно использует небольшую часть механизма внимания (один нейрон) для прогнозирования вытеснения. Это делает DMS идеальной для доработки существующих моделей без архитектурных изменений.
Эмпирические результаты показывают, что всего за 1 тыс. шагов обучения DMS может достичь сжатия KV-кэша в 8 раз, сохраняя или даже улучшая производительность модели в задачах, требующих рассуждений.
Результаты тестирования: масштабирование производительности без увеличения затрат
Исследовательская группа протестировала DMS на нагруженных рассуждениями бенчмарках, таких как:
* AIME 2024 (продвинутая математика);
* MATH 500 (решение математических задач);
* GPQA Diamond (сложные научные вопросы);
* LiveCodeBench (генерация кода).
На моделях разных размеров — Qwen-R1 1,5B, 7B и 32B — DMS улучшила точность совпадения на 9,1 пункта в AIME, на 7,6 в GPQA и на 9,6 в LiveCodeBench при тех же бюджетах памяти и вычислений.
По сравнению с топовыми базовыми показателями, такими как Quest и TOVA, DMS последовательно превосходила их как по эффективности чтения из KV-кэша (прокси времени выполнения), так и по пиковому использованию памяти, достигая лучших парето-границ.
Универсальность
DMS также эффективна в задачах, не требующих рассуждений. На бенчмарках с коротким контекстом, таких как MMLU, GSM8K и HellaSwag, DMS сохраняла производительность при коэффициентах сжатия до 4 раз с минимальной деградацией (~3,5 пункта). В задачах с длинным контекстом, таких как Needle-in-a-Haystack и Variable Tracking, DMS даже превзошла стандартные модели, что говорит о её потенциале для решения таких проблем, как информационное переполнение в длинных последовательностях.
Заключение
Dynamic Memory Sparsification (DMS) представляет собой практичное и масштабируемое решение для повышения эффективности вывода трансформерных языковых моделей. Благодаря интеллектуальному сжатию KV-кэша с минимальными затратами на переобучение DMS позволяет моделям рассуждать над более длинными последовательностями или параллельно без увеличения времени выполнения или требований к памяти.
Её последовательные достижения в различных задачах, требующих рассуждений, и общих задачах подчёркивают её универсальность и эффективность. Поскольку LLM всё чаще используются в средах с ограниченными ресурсами, DMS предлагает перспективный путь вперёд — балансируя сжатие, точность и простоту интеграции для реальных рабочих нагрузок вывода.