Kronecker Product Attention сжимает KV-кэш до 10 раз и догоняет MLA по качеству

Память под KV-кэш давно стала главным ограничителем авторегрессионной генерации: она растёт линейно с длиной контекста, размером батча и глубиной модели. Новый препринт на arXiv предлагает механизм Kronecker Product Attention (KPA), который, по заявлению авторов, сжимает этот кэш до 10 раз и при этом не проигрывает по качеству MHA, MQA, GQA, MLA и TPA. Сравнение проведено на четырёх масштабах моделей и девяти downstream-бенчмарках при одинаковых условиях обучения.

Как устроен KPA

Главная идея: ключи и значения не хранятся целиком, а раскладываются в сумму кронекеровых произведений компактных матричных факторов. Сами факторы зависят от токена, то есть вычисляются из входа, а не фиксируются в весах. В кэше остаются только эти факторы, а полные тензоры восстанавливаются при вычислении внимания.

У существующих подходов компромиссы разные. MQA и GQA экономят память, объединяя головы, и теряют в выразительности. MLA сжимает кэш через латентное представление, но требует отдельной «развязанной» ветки для позиционного кодирования RoPE. TPA использует факторизацию через внешнее произведение векторов, то есть жёсткую структуру ранга один. KPA, как утверждают авторы, строго обобщает TPA: векторное внешнее произведение получается как вырожденный случай кронекерова разбиения, а невырожденные матричные факторы способны описывать более богатые блочные паттерны.

Отдельный технический результат касается позиций. Авторы вводят структурированный ротационный оператор, согласованный с кронекеровым представлением. Он сохраняет свойство относительной позиции RoPE, но вращает только компактные факторы. Поэтому дополнительная ветка для позиционного кодирования, как в MLA, не требуется.

Что заявлено в статье

• Сжатие KV-кэша: до 10x.
• Сравнение с MHA, MQA, GQA, MLA и TPA: результат на уровне или выше.
• Масштабы: четыре размера моделей, девять downstream-бенчмарков.
• Условия: идентичное обучение для всех архитектур.
• Позиционное кодирование: RoPE-совместимый оператор на компактных факторах, без отдельной ветки.
• Настройка: «кронекерово разбиение голов» (Kronecker head split) задаётся как гиперпараметр.

Последний пункт, возможно, самый практичный. Разбиение определяет компромисс между степенью сжатия и точностью, и менять его можно на этапе проектирования без смены архитектуры. Команде не нужно переписывать модель, чтобы переехать с более мягкой компрессии на более агрессивную. Достаточно переобучить с другим значением.

Что делать инженерам и где осторожность

Пока это препринт v1 без независимой проверки. В аннотации не указаны ни конкретные размеры моделей, ни длины контекста, ни скорость декодирования на реальном железе. Выигрыш по памяти не всегда равен выигрышу по задержке: восстановление ключей и значений из факторов добавляет вычисления, а их цена зависит от ядер и GPU. Заявленное «до 10x» относится к лучшему режиму, а в какой конфигурации оно достигается и какой ценой по качеству, станет ясно из полного текста и кода, если авторы его опубликуют.

Командам, которые строят инференс длинного контекста или обслуживают большие батчи, KPA стоит держать в списке кандидатов рядом с MLA. Проверять нужно три вещи: сколько памяти освобождается на ваших длинах контекста, как ведёт себя качество при разных разбиениях и какая получается пропускная способность с вашими ядрами внимания. Тем, кто уже обучил модель на GQA или MLA, переход потребует переобучения, так что имеет смысл закладывать KPA в следующий цикл pre-training, а не пытаться заменить механизм в готовых весах.