В качестве отправной точки берут метод KVZip, который дублирует входной промпт следующим образом
user: <prompt>
Repeat the previous context exactly.
assistant: <prompt>И на основе внимания отбирает наиболее важные пары ключей и значений. Данный метод хорош с точки зрения качества, но требует двух проходов по сети, и годится только для сжатия префикса.
Авторы предлагают модификацию KVZip+, которая дополнительно еще учитывает тот факт, что внимание домножается на value и выходную проекцию O - тем самым получая более точную оценку на важность данного токена.
Конечный метод KVZap предлагает обучать небольшую нейросеть (линеный слой или MLP) предсказывать оценку из KVZip+ и использовать ее для выбрасывания токенов. Предсказание стоит заметно меньше, чем выгрузка большого KV-кэша на каждом шаге генерации. На инференсе задается некий порог важности, выше которого токен принимается, а ниже выбрасывается.
Для обучения предсказателей собирают выборку из 1.2M пар hidden / метрика из датасета Nemoton-CC. Кроме того для гарантированного сохранения ближайших токенов сохраняют всегда 128 последних токенов.
🧪 Эксперименты
Метод валидируют на Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-32.
Качество оценивают на бенчах длинного контекста - RULER/LongBench и AIME25 (ризонинг с длинным контекстом).
По квадратичной ошибке двуслойная MLP всегда лучше линейного слоя, и превосходство имеет место в большинстве случаев и на бенчах, кроме Llama.
KVZap по качеству равен или чуть хуже KVZip, но опережает заметно, все остальные бейзлайны.
В ablation показывают, что фиксированный порог лучше, чем topk pruning, потому что разные тексты могут иметь разную информационную нагрузку. Скользящее окно важно для хорошего качества.
В итоге удается достичь сжатия в 2-4 раза без просадки в качестве.
💡 Выводы
Выглядит как вполне рабочий и сравнительно легко применимый на практике подход с небольшим пост обучением. Из бейзлайнов не хватает очень похожего по смыслу и мотивации AutoJudge, где важность токенов определяется на основе конечной задачи. В то же время, данный подход, по всей видимости, более обобщаем на другие задачи.