ИИтог

Virtual Head Attention: схема внимания с KV-кэшем в 2-4 раза меньше, чем у GQA

исследованияРедакция ИИтогHugging Face🔥56Рейтинг ИИтог56Влияние45Новизна65Релевантность100Неожиданность30Участники10Достоверность70
Virtual Head Attention: схема внимания с KV-кэшем в 2-4 раза меньше, чем у GQA

В статье предложен механизм внимания Virtual Head Attention (VHA) для языковых моделей с длинным контекстом. Авторы исходят из «невозможного треугольника»: при проектировании внимания приходится выбирать между размером KV-кэша, вычислениями (FLOPs) на обучении и префилле и качеством модели. GQA даёт хорошее качество, но требует много памяти под KV-кэш. Альтернативы MLA и MFA кэш сжимают, но MLA теряет в качестве, а MFA заметно увеличивает вычисления на обучении и префилле. VHA, по утверждению авторов, улучшает сразу все три вершины этого треугольника.

Устройство метода: две лёгкие линейные операции. Q Premix применяет к запросам внутри каждой KV-группы почти тождественное преобразование, чтобы вернуть разнообразие «виртуальных» голов при вдвое меньшем числе физических голов запросов. Linear Postmix объединяет признаки разных голов через низкоранговую остаточную структуру вида I + AB^T. Нелинейных операций в VHA нет, поэтому на инференсе Premix можно свернуть в проекцию запросов, а Postmix, в выходную проекцию. Модель при этом превращается в обычную архитектуру GQA-2 и напрямую использует все существующие оптимизации инференса GQA. Всего с двумя KV-головами размер группы в больших моделях легко достигает 64, что, как пишут авторы, позволяет полностью задействовать параллелизм Tensor Core при декодировании.

Результаты. При предобучении на 50 млрд токенов с контекстом 4K на четырёх масштабах VHA стабильно превосходит стандартный GQA на каждом масштабе (от +0,05% до +0,81%; в тексте не уточнено, какая это метрика и в каких единицах). KV-кэш при этом в 4 раза меньше, чем у GQA-8, на плотных моделях и в 2 раза меньше, чем у GQA-4, на MoE-масштабе. Вычисления на обучении и префилле снижены на 7,9% (для модели 1,7B при длине последовательности 4096).

На «индустриальной» MoE-модели 30B-A3B, обученной на 500 млрд токенов с контекстом 8K, VHA опережает MLA и GQA-4 на +0,70 и +2,60% соответственно на наборе из 12 бенчмарков для базовых моделей. В отдельном сопоставлении MoE-моделей 10B-A1B после длинноконтекстного дообучения (mid-training) средний балл RULER на длинах 8K, 128K составил 84,4 у VHA против 73,0 у GQA, причём VHA использует две KV-головы вместо четырёх в слоях с полным вниманием.

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

  • VHA добавляет две линейные операции, Q Premix (преобразование запросов внутри KV-группы) и Linear Postmix (низкоранговое объединение признаков голов), и на инференсе сворачивается в обычную архитектуру GQA-2.
  • По заявлению авторов, KV-кэш в 4 раза меньше, чем у GQA-8, на плотных моделях и в 2 раза меньше, чем у GQA-4, на MoE-масштабе; вычисления на обучении и префилле ниже на 7,9% (модель 1,7B, длина 4096).
  • При предобучении на 50 млрд токенов (контекст 4K) VHA превосходит стандартный GQA на всех четырёх масштабах, на величину от +0,05% до +0,81%.
  • На MoE-модели 30B-A3B (500 млрд токенов, контекст 8K) VHA впереди MLA и GQA-4 на +0,70 и +2,60% по набору из 12 бенчмарков.
  • В сопоставлении MoE-моделей 10B-A1B средний RULER на 8K, 128K: 84,4 у VHA против 73,0 у GQA.

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

KV-кэш, одна из главных статей расхода памяти при работе языковых моделей с длинным контекстом, и его размер обычно приходится менять на качество или на дополнительные вычисления. Авторы называют это «невозможным треугольником» и утверждают, что VHA улучшает размер кэша, вычисления на обучении и префилле и качество одновременно. Если заявленное подтвердится независимо, это будет заметный сдвиг в схеме внимания. Пока это лишь утверждения самой статьи.

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

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

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

Для пользователей готовых нейросетей практической пользы сейчас нет: это препринт об архитектуре, которую нужно закладывать при обучении модели. О выпуске кода или моделей в тексте не сообщается. Команде, обучающей модель, идея применима как замена блока внимания; на инференсе, по словам авторов, Premix сворачивается в проекцию запросов, а Postmix, в выходную проекцию, так что отдельной поддержки в движке инференса не требуется.

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

Все цифры взяты из аннотации и приведены как заявления авторов. В тексте не названы ни авторы, ни организации, ни владелец «индустриальной» модели 30B-A3B; 12 бенчмарков не перечислены, а в каких единицах измерены +0,05…+0,81% и +0,70 / +2,60%, не сказано. Сравнение на RULER (84,4 против 73,0) сделано только с GQA и отдельно от сравнения на 30B-A3B. Пока это препринт без независимой проверки.

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

Улучшения на малых прогонах (50 млрд токенов, контекст 4K) скромные, доли процента, и могут не перенестись на другие задачи и настройки. Снижение вычислений на 7,9% приведено только для модели 1,7B при длине 4096, для других масштабов цифр нет. Измерений задержки и пропускной способности инференса в аннотации нет, только KV-кэш и FLOPs. Кроме того, сравнение идёт с разными вариантами GQA (GQA-8, GQA-4, GQA-2), поэтому кратность сжатия кэша зависит от выбранного базового варианта.