Random Attention: случайное вытеснение KV-кэша ускоряет вывод рассуждающих моделей на 32, 43% без потери качества

Длинные цепочки рассуждений повышают качество ответов языковых моделей на сложных задачах, но именно они превращают KV-кэш, память, где модель хранит промежуточные данные внимания для уже сгенерированных токенов, в узкое место: цепочка растёт, кэш переполняется, и часть токенов приходится вытеснять. Все существующие методы сжатия KV-кэша устроены одинаково: каждому токену в кэше присваивается оценка того, насколько он понадобится модели позже, и остаются только токены с самой высокой оценкой.

Хэн Ван с соавторами предлагают метод Random Attention, который отказывается от самой идеи такой оценки. Токены исходного запроса (промпта) метод не трогает и оставляет в кэше целиком, а всё, что накопилось сверх этого, вытесняет полностью случайно, независимо в каждой голове внимания, вообще не вычисляя никаких оценок важности. На четырёх моделях и шести задачах, требующих рассуждений, случайное вытеснение по качеству не уступает лучшему из существующих методов оценки, и при этом в реальном развёртывании на движке вывода vLLM даёт на 32, 43% более высокую пропускную способность, чем этот метод-конкурент: весь выигрыш, за счёт того, что не нужно тратить вычисления на подсчёт оценок.

Контролируемые эксперименты авторов объясняют, почему случайного отбора достаточно. Во-первых, самый уязвимый участок кэша, это промпт: большая часть разницы в качестве между разными методами оценки объясняется тем, удалось ли им сохранить промпт в кэше, а не тем, насколько умно они выбирали остальные токены. Во-вторых, если промпт защищён, сама цепочка рассуждений устойчива к случайному вытеснению, у неё есть избыточность на двух уровнях: в тексте (модель по ходу рассуждения сама повторяет то, что ей ещё понадобится) и между головами внимания (каждая голова хранит свою отдельную копию цепочки). Поэтому случайная выборка почти всегда оставляет достаточно копий нужной информации, и оценивать токены, чтобы их найти, не требуется.

Код метода авторы выложили в открытом доступе на GitHub, в аккаунте SalesforceAIResearch.

Ключевые факты

Random Attention оставляет в KV-кэше без изменений только токены промпта, а весь остальной кэш вытесняет полностью случайно, отдельно в каждой голове внимания, не вычисляя никаких оценок важности токенов.
На четырёх моделях и шести задачах, требующих рассуждений, случайное вытеснение по качеству не уступает лучшему из существующих методов оценки токенов.
В реальном развёртывании на vLLM метод даёт на 32, 43% более высокую пропускную способность, чем этот метод-конкурент, экономия целиком за счёт отказа от подсчёта оценок.
Авторы объясняют результат избыточностью цепочки рассуждений: модель сама повторяет нужную информацию по ходу текста, а каждая голова внимания хранит собственную копию цепочки.
Код метода выложен в открытом доступе на GitHub, в аккаунте SalesforceAIResearch.

Почему это важно

KV-кэш, это память, в которой модель хранит представления уже обработанных токенов, чтобы не пересчитывать их заново; при длинных цепочках рассуждений эта память быстро становится главным узким местом по объёму и по скорости вывода. Вся индустрия методов сжатия KV-кэша построена на одной идее: оценивать каждый токен по тому, насколько он понадобится модели в будущем, и оставлять только токены с лучшей оценкой. Эта работа показывает, что сам сигнал такой оценки даёт почти нулевой вклад в качество: если просто защитить промпт и вытеснять остальной кэш случайно, вообще без оценок, результат на четырёх моделях и шести задачах на рассуждение не уступает лучшему из существующих методов, а скорость вывода в реальном развёртывании на vLLM выше на 32, 43%. Это ставит под вопрос смысл целого класса алгоритмов, которые тратят вычисления на оценку токенов, если та, судя по всему, не даёт ощутимой пользы.

Кому это важно

Инженерам, которые обслуживают в проде рассуждающие модели с длинными цепочками мыслей и упираются в память и скорость из-за разросшегося KV-кэша. Командам, использующим vLLM или похожие движки вывода, где пропускная способность напрямую определяет стоимость обработки запросов. Исследователям, которые занимаются сжатием KV-кэша и отбором токенов по важности, результат ставит под сомнение исходную предпосылку целого направления. Компаниям, которые считают стоимость вывода рассуждающих моделей и заинтересованы в более простом и быстром способе сократить память кэша без потери качества ответов.

Как это применить

Рецепт метода нарочно простой: токены исходного запроса (промпта) в KV-кэше не трогать вообще, а всё, что накопилось в кэше сверх этого, вытеснять случайно и независимо в каждой голове внимания, без вычисления какой-либо оценки важности токенов. Именно отказ от подсчёта оценок и даёт измеренный прирост пропускной способности на 32, 43% в реальном развёртывании на vLLM, а не только теоретическую экономию. Авторы уже выложили код в открытом доступе на GitHub, в аккаунте SalesforceAIResearch, так что метод можно опробовать напрямую, не восстанавливая его по одному описанию из статьи. Цена, лицензия на код и требования к оборудованию для развёртывания в источнике не указаны.

Можно ли доверять

Это исследовательская работа, материал из ленты научных статей Hugging Face Papers, а не анонс продукта. В её пользу говорит то, что результат проверен не на одной модели и не на одной задаче, а на четырёх моделях и шести задачах на рассуждение, и что авторы сразу выложили код в открытом доступе, это позволяет независимо перепроверить заявленные цифры. Ограничение в том, что сам текст не называет ни эти четыре модели и шесть задач, ни конкретное название метода-конкурента, обозначенного просто как «сильнейший из существующих методов вытеснения», оценить по этому тексту, насколько такая проверка репрезентативна, невозможно. Не приведены и абсолютные цифры пропускной способности, только относительный прирост в 32, 43%, так что итоговая практическая выгода зависит от того, с какой базовой скоростью сравнивали.

Риски и подводные камни

Вывод о том, что случайного отбора достаточно, опирается на избыточность именно длинных цепочек рассуждений: модель сама повторяет нужное по ходу текста, а разные головы внимания хранят по своей копии. Для коротких запросов или задач без развёрнутого рассуждения такой избыточности может не быть, и перенос результата на них в тексте не подтверждён. Схема работает только при условии, что промпт в кэше защищён отдельно и не вытесняется наравне с остальным содержимым, источник прямо называет промпт самым уязвимым местом кэша, и реализация без этой защиты рискует потерять в качестве. Наконец, случайное вытеснение именно сравнялось по качеству с лучшим существующим методом, а не превзошло его, выигрыш здесь только в скорости, и если позже появится более точный метод оценки, паритет по качеству у Random Attention ничем не гарантирован.