Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Кэш длинноконтекстной LLM можно сокращать точнее, если учитывать, на каком расстоянии от запроса лежат сохранённые токены. Группа Xianpeng Shang показала в препринте, который не прошёл рецензирование и содержит их собственные замеры, что Distance-KV уменьшает память кэша Llama-3.1-8B-Instruct на 65,4% при длинном контексте. Метод заранее строит схему сокращения и не тратит время на оценку важности токенов при каждом запросе.
Одна голова внимания по-разному работает на разных расстояниях
При генерации модель хранит кэш ключей и значений (KV-кэш) для уже обработанных токенов. Чем длиннее контекст, тем больше памяти занимает этот кэш и тем дольше каждый новый токен обращается к сохранённым состояниям.
Обычные методы сокращают кэш по важности отдельных токенов, слоёв или голов внимания. Подход на уровне голов предполагает, что каждой голове можно назначить одну роль: сохранять её кэш полностью, ограничить его или оставить только ближайшую часть контекста.
Distance-KV опирается на более мелкое разделение. Одна и та же голова может хорошо находить данные рядом с запросом, но хуже работать на дальней части контекста. Другая, наоборот, помогает извлекать удалённые фрагменты. Некоторые головы достигают максимума посередине, поэтому простого деления на «ближние» и «дальние» недостаточно.
Метод делит контекст на интервалы по относительному расстоянию и отдельно решает, какие интервалы сохранить для каждой головы в каждом слое. Получается двоичная схема: конкретное сочетание слоя, головы и диапазона расстояний либо остаётся в кэше, либо удаляется.
Схему обучают заранее на синтетических задачах поиска. Параметры самой LLM при этом заморожены: меняются только переключатели, которые управляют сохранением кэша. Для каждой модели и выбранного бюджета памяти нужна своя схема, но после обучения её можно применять к разным входным данным и запросам.
Во время работы модель сначала полностью обрабатывает исходный контекст. Затем Distance-KV один раз удаляет ненужные состояния и уплотняет оставшийся кэш, сохраняя исходные позиции токенов. Начало контекста и его недавняя часть остаются целиком; новые сгенерированные токены тоже не сокращаются.
Статическая схема выиграла у других способов сжатия
Метод проверили на Llama-2-7B-32K-Instruct, Llama-3.1-8B-Instruct и Qwen2.5-7B-Instruct. В набор тестов вошли Needle-in-a-Haystack, LongBench, RULER и SCBench: они покрывают поиск фрагментов, вопросы по длинным документам, устойчивость при росте контекста и повторное использование общего контекста в диалоге.
Сравнение включало полный кэш Dense и методы StreamingLLM, DuoAttention, MoA и KVzip-CI. Все способы сжатия выбирали состояния до обработки запроса и не использовали его содержание, поэтому Distance-KV не получил преимущества от более позднего доступа к вопросу.
На RULER при контексте 128K Distance-KV опередил сильнейший конкурирующий метод сжатия на 9,3 пункта. В тестах с несколькими запросами одна фиксированная схема продолжала работать без повторной оценки кэша и дала лучший средний результат среди методов сжатия на обеих проверенных моделях.
В основном эксперименте схема сохраняла 20% логических сочетаний голов и диапазонов. На Llama-3.1-8B-Instruct декодирование оказалось в 1,66 раза быстрее Dense. Перемешивание диапазонов между головами ухудшало качество при том же объёме физического кэша: важна не только доля сохранённых состояний, но и точное соответствие головы конкретному расстоянию.
Когда Distance-KV меняет архитектурный план
Работа предлагает практичный вариант для сервисов, где один длинный контекст используют многократно: помощников по репозиторию, анализа документов и диалогов с общей базой материалов. Схему можно подготовить вместе с моделью, загрузить при запуске и применять ко всем запросам без отдельного оценщика важности.
Главное архитектурное следствие — сокращение происходит после полного первичного прохода по контексту. Distance-KV уменьшает память и задержку последующей генерации, но не отменяет стоимость первоначальной обработки длинного документа. Если узкое место находится именно в загрузке контекста, результаты работы сами по себе не обосновывают смену подхода.
Для систем с GQA фактический объём кэша также нельзя вывести прямо из логического бюджета. Несколько голов запросов используют одну физическую голову кэша, поэтому состояние приходится сохранять, если оно нужно хотя бы одной из них. Планировать память следует по физическому кэшу после объединения решений, а не по параметру обучения схемы.
Distance-KV не выглядит универсальным модулем, который можно без настройки перенести между моделями. Схема зависит от конкретной LLM и бюджета, а проверки охватывают три модели сопоставимого размера. Для команды это скорее кандидат на отдельный этап оптимизации развёртывания: сначала выбрать модель и рабочую длину контекста, затем обучить схему и проверить качество на собственных длинных документах.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
