Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
PeekKV сокращает кэш ключей и значений (KV-cache), который LLM хранит для уже прочитанного контекста, и подстраивает его размер под конкретный запрос без предварительного профилирования. Хотя препринт не рецензирован и числа получили сами авторы, на Llama-3.1-8B-Instruct метод сохранил 95,7% точности полного кэша при бюджете памяти 10%. Это позволяет пересмотреть статическое распределение памяти в сервисах, которые обрабатывают длинный контекст.
Важные головы меняются после начала ответа
KV-cache хранит ключи и значения для предыдущих токенов, чтобы модель не вычисляла их заново на каждом шаге. Его размер растёт вместе с контекстом, поэтому длинные документы и состояние агентов занимают больше памяти GPU и требуют больше данных при каждом обращении механизма внимания.
Часть методов сокращает кэш равномерно, а часть оставляет больше токенов отдельным головам внимания — компонентам, которые независимо выбирают нужные фрагменты контекста. Важные головы обычно находят заранее на наборе запросов или оценивают во время первичной обработки входного текста.
PeekKV опирается на другое наблюдение: набор активных голов заметно меняется на границе между чтением запроса и генерацией ответа. Поэтому сигнал из входного текста плохо предсказывает, какие части модели понадобятся при дальнейшем выводе.
Первый сгенерированный токен тоже подходит не всегда. Код может начинаться с ограждения блока и ключевого слова, документ — с разметки, обычный текст — с артикля или знака препинания. Такие служебные токены активируют головы, связанные с формой ответа, а не с его содержанием. После их пропуска состав важных голов остаётся достаточно стабильным на следующих содержательных токенах.
Как один токен задаёт бюджет кэша
После первичной обработки запроса модель начинает обычную генерацию с полным кэшем. PeekKV пропускает служебные элементы и останавливает наблюдение на первом токене, который несёт содержание. Для кода тип токена определяет синтаксический анализатор, а для текста и разметки используются заданные правила.
На этом шаге метод оценивает, насколько сильно каждая голова сосредоточена на одной позиции исходного контекста. Сфокусированные головы получают высокий балл, а головы с рассеянным вниманием — низкий. Поскольку модели с групповыми запросами хранят общие ключи и значения для нескольких голов, PeekKV суммирует их баллы внутри группы.
Затем метод распределяет память между группами пропорционально этим оценкам. Каждой группе гарантирован минимальный объём, а максимальный ограничен потолком. В результате несколько важных для текущего запроса групп могут сохранить почти весь контекст, тогда как остальные оставляют лишь небольшую его часть.
Выбор токенов внутри группы строится на изменении выхода внимания при удалении каждого кандидата. Чем сильнее удаление токена меняет результат, тем выше его место в списке сохранения. Эту величину можно вычислить по уже полученным весам внимания и векторам значений, не запуская модель отдельно для каждого токена.
Токены, созданные во время подсматривания, не отбрасываются: после сжатия исходного префикса генерация продолжается с того же места. Обучение и профиль модели не нужны, но до первого содержательного токена память остаётся на уровне полного кэша. PeekKV поэтому сокращает постоянный расход при генерации, но не помогает загрузить контекст, который изначально не помещается на устройстве.
Когда PeekKV меняет планы внедрения
На RULER при бюджете 3% PeekKV набрал 81,1% против 70,1% у ближайшего соперника AdaKV. На LongBench метод сохранил 95,9% точности полного кэша при бюджете 1,5%. В сравнении со StreamingLLM, H2O, SnapKV и PyramidKV он лидировал на каждом использованном бенчмарке; против методов с неравномерным бюджетом преимущество зависело от модели и задачи.
Проверка охватила Llama-3.1-8B-Instruct и Mistral-7B-Instruct-v0.2. RULER проверял поиск информации в длинном контексте, LongBench — ответы на вопросы, суммаризацию и продолжение кода, BFCL v3 — вызов инструментов, HumanEval — генерацию программ. Сравнение проводили с полным кэшем и семью способами его сжатия.
Один токен хуже покрывал задания, где ответ нужно собрать из нескольких разнесённых фрагментов. Дополнительные содержательные токены помогали сохранять окончания таких фрагментов, особенно при малом бюджете. Это оставляет практический выбор: короткое наблюдение уменьшает задержку перед сжатием, более длинное — лучше подходит для поиска нескольких элементов и агрегации.
Для команды, которая уже планирует неравномерно сокращать KV-cache, работа меняет точку измерения: важность голов стоит оценивать после начала ответа, а не только на входном тексте. Это особенно полезно, если сервис обслуживает разные типы запросов и статический профиль приходится собирать отдельно для каждой модели и нагрузки.
Доказательства пока относятся к указанным моделям, однократной генерации и четырём наборам задач; многотурновые сценарии авторы оставили для дальнейшей работы. Внедрение также требует правил, которые отделят служебные токены от содержательных в каждом поддерживаемом формате.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
