Групповой запрос внимания

Групповой запрос внимания (англ. Group Attention Request, GAR) — это механизм в архитектуре трансформеров, применяемый в некоторых больших языковых моделях (LLM) для улучшения эффективности обработки последовательностей данных, в частности, для снижения вычислительных затрат и требований к памяти при работе с длинными контекстами. Является вариацией многоголового внимания (Multi-Head Attention, MHA), при которой головы внимания группируются, и внутри каждой группы вычисления выполняются с общими ключами (keys) и значениями (values), что позволяет уменьшить количество уникальных операций.

История

Механизм группового запроса внимания был предложен как развитие идей, заложенных в архитектуре трансформеров, впервые описанной в статье «Attention Is All You Need» (2017). Классическое многоголовое внимание (MHA) требует для каждой головы отдельного набора матриц для запросов (queries), ключей (keys) и значений (values), что приводит к линейному росту вычислительных затрат и объёма памяти с увеличением числа голов. В 2022 году была предложена архитектура многоквартирного внимания (Multi-Query Attention, MQA), где все головы используют общие ключи и значения, что значительно ускоряет инференс, но может снижать качество модели.

Групповой запрос внимания (GAR) был впервые представлен в 2023 году в работе «GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints» (Ainslie et al., Google Research). Авторы предложили компромиссный подход: головы делятся на группы, каждая из которых имеет собственный набор ключей и значений, но внутри группы эти параметры общие. Это позволило сохранить часть вычислительной эффективности MQA, избежав существенного падения качества, характерного для полного разделения.

Принцип работы

В основе GAR лежит идея уменьшения количества уникальных проекций ключей и значений по сравнению с числом голов внимания. В стандартном MHA для каждой из \( h \) голов вычисляются собственные \( Q_i, K_i, V_i \) матрицы. В GAR головы объединяются в \( g \) групп, где \( g < h \) (обычно \( g \) — это делитель \( h \), например, 2, 4 или 8). Внутри каждой группы все головы используют один и тот же набор ключей и значений, но каждая голова сохраняет свой собственный запрос.

Формальное описание

Пусть \( h \) — общее количество голов внимания, \( g \) — количество групп. Тогда для каждой группы \( j \) (\( j = 1, \dots, g \)) вычисляются:

• Общая матрица ключей \( K_j \in \mathbb{R}^{d_k \times L} \)
• Общая матрица значений \( V_j \in \mathbb{R}^{d_v \times L} \)

где \( L \) — длина последовательности, \( d_k \) и \( d_v \) — размерности ключей и значений соответственно.

Для каждой головы \( i \) в группе \( j \) вычисляется индивидуальная матрица запросов \( Q_{i,j} \in \mathbb{R}^{d_q \times L} \). Выход головы определяется как:

\[ \text{Attention}(Q_{i,j}, K_j, V_j) = \text{softmax}\left(\frac{Q_{i,j} K_j^T}{\sqrt{d_k}}\right) V_j \]

Результаты всех голов конкатенируются и проецируются через выходную матрицу.

Сравнение с MHA и MQA

Применение

GAR используется в основном в больших языковых моделях, где критичны скорость инференса и объём оперативной памяти, особенно при работе с длинными последовательностями (например, в задачах генерации текста, перевода, анализа документов).

Примеры моделей, использующих GAR

• LLaMA 2 (Meta — организация признана экстремистской и запрещена в РФ) — в некоторых конфигурациях применяет GAR для оптимизации.
• Gemini (Google) — использует GAR в архитектуре для баланса между производительностью и качеством.
• Mistral — в моделях Mistral 7B и Mixtral 8x7B применяется GAR с 8 группами, что позволило достичь высокой эффективности при сохранении качества, сравнимого с MHA.
• Qwen (Alibaba) — в ряде версий Qwen-7B и Qwen-14B используется GAR.

В России GAR также находит применение в разработках, связанных с NLP, например, в моделях на базе архитектуры трансформеров, создаваемых в рамках проектов Yandex, Sber AI и других компаний, хотя публичные детали реализации часто не раскрываются.

Преимущества и недостатки

Преимущества

• Снижение вычислительной сложности: Уменьшение количества операций умножения матриц для ключей и значений по сравнению с MHA.
• Экономия памяти: Кэш ключей и значений (KV cache) при инференсе занимает меньше места, что критично для длинных контекстов.
• Скорость инференса: Ускорение генерации токенов за счёт уменьшения объёма пересылаемых данных между памятью и вычислительными блоками.
• Сохранение качества: При правильном выборе числа групп (\( g \)) качество модели близко к MHA, в отличие от MQA.

Недостатки

• Сложность настройки: Выбор оптимального числа групп требует экспериментов и может зависеть от конкретной задачи.
• Потенциальное снижение точности: При слишком малом числе групп (например, \( g = 1 \), что эквивалентно MQA) качество может упасть.
• Необходимость модификации архитектуры: Внедрение GAR требует изменения исходной реализации трансформера, что может быть нетривиально для существующих моделей.

Интересные факты

• GAR часто называют «гибридным» подходом между MHA и MQA, так как он объединяет их преимущества.
• В некоторых реализациях (например, в модели Falcon) используется вариант GAR с фиксированным числом групп, равным 4, что стало стандартом для многих современных LLM.
• Исследования показывают, что GAR особенно эффективен на этапе инференса (генерации), когда кэш KV активно используется, но менее критичен на этапе обучения.
• В 2024 году появились работы, предлагающие динамическое изменение числа групп в зависимости от длины последовательности или типа задачи, что может стать следующим шагом в развитии механизма.

Критика

Основная критика GAR связана с тем, что он не решает фундаментальной проблемы квадратичной сложности внимания относительно длины последовательности. Хотя GAR уменьшает константу в вычислительной сложности, асимптотическая сложность остаётся \( O(L^2) \) для полного внимания. Для преодоления этого ограничения GAR часто комбинируют с другими методами, такими как скользящее окно внимания (sliding window attention) или sparse attention.

Также отмечается, что эффективность GAR сильно зависит от аппаратной реализации: на некоторых GPU (например, с большим объёмом кэша) выигрыш может быть менее заметен, чем на специализированных ускорителях.

Источники

1. Ainslie, J., et al. «GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.» arXiv preprint arXiv:2305.13245, 2023.
2. Shazeer, N. «Fast Transformer Decoding: One Write-Head is All You Need.» arXiv preprint arXiv:1911.02150, 2019.
3. Vaswani, A., et al. «Attention Is All You Need.» Advances in Neural Information Processing Systems, 2017.
4. Touvron, H., et al. «LLaMA: Open and Efficient Foundation Language Models.» arXiv preprint arXiv:2302.13971, 2023.
5. Jiang, A. Q., et al. «Mistral 7B.» arXiv preprint arXiv:2310.06825, 2023.