Grouped Query Attention (GQA)

Внимание с группировкой запросов (GQA) — разновидность многоголового механизма внимания для трансформеров, в которой группы голов запросов используют общие головы ключей и значений. Например, в Llama 3.1 8B на каждом слое 32 головы запросов сопоставлены с 8 головами ключей и значений. GQA занимает промежуточное положение между многоголовым вниманием (MHA), где у каждой головы запросов собственные ключи и значения, и вниманием с несколькими запросами (MQA), где все головы запросов используют одну общую голову ключей и одну голову значений. Джошуа Эйнсли и его коллеги из Google Research представили GQA в работе «GQA: обучение обобщённых трансформеров с несколькими запросами на контрольных точках многоголовых моделей» на EMNLP 2023 и показали, что модели GQA обеспечивают качество, близкое к MHA, при скорости декодирования, близкой к MQA. GQA используют такие открытые большие языковые модели (LLM), как Llama 2 70B, семейство Llama 3, Mistral 7B и Gemma 2; основанные на них визуально-языковые модели (VLM), например PaliGemma 2 от Google на базе Gemma 2, также используют этот механизм в языковых декодерах.

Как работает GQA#

В механизме самовнимания каждый токен для каждой головы преобразуется в запрос, ключ и значение, а каждая голова сопоставляет запросы со своими ключами. GQA сохраняет полный набор голов запросов, но использует меньше голов ключей и значений; каждая группа голов запросов работает с общими ключами и значениями своей группы. Если есть 8 голов запросов и 2 головы ключей и значений, каждая такая голова обслуживает 4 головы запросов:

Вычисление внимания внутри каждой головы остаётся неизменным и представляет собой масштабированное скалярное произведение. Если группа одна, слой превращается в MQA; если число групп равно числу голов запросов — в MHA. PyTorch поддерживает такую конфигурацию напрямую через аргумент enable_gqa функции scaled_dot_product_attention:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

Выводятся Output (1, 8, 16, 64), same as copied heads: True и 256 with GQA, 1024 with MHA: все 8 голов запросов формируют выходные данные, используя вчетверо меньше ключей и значений, чем хранится в MHA.

Как GQA сокращает KV-кеш#

Во время генерации текста модель хранит ключи и значения всех предыдущих токенов в KV-кеше и считывает весь кеш для каждого нового токена, поэтому скорость декодирования ограничена скорее пропускной способностью памяти, чем вычислениями. Размер кеша зависит от числа голов ключей и значений. В Llama 3.1 8B 32 слоя, в каждом из которых по 8 голов ключей и значений размерностью 128; при 16-битной точности модель хранит 128 КиБ на токен, или 16 ГиБ для полного контекстного окна из 131 072 токенов. Для MHA этой же модели потребовалось бы 512 КиБ на токен и 64 ГиБ для максимальной длины. Освободившуюся память можно использовать для более длинного контекста, увеличения батчей при непрерывной пакетной обработке и снижения задержки инференса на токен. Интерфейс Ultralytics LLM принимает base_url для любого совместимого с OpenAI сервера, например собственного сервера Llama 3.1 8B. Благодаря 8 головам ключей и значений такой сервер может хранить в заданном объёме памяти кеша вчетверо больше токенов диалога, чем при конфигурации MHA с 32 головами.

Преобразование моделей в GQA#

Для GQA не нужно обучать модель с нуля. Авторы преобразовали существующую контрольную точку MHA, усреднив значения голов ключей и значений внутри каждой группы, а затем продолжили предобучение, затратив около 5% исходных вычислительных ресурсов. Этот этап они называют дообучением. MQA — крайний вариант с одной общей головой — представила Ноам Шазир в 2019 году. Это самая быстрая конфигурация, но она снижала качество и могла приводить к нестабильности обучения; промежуточное число групп позволяет восстановить большую часть качества.

GQA и родственные методы#

• Многоголовое латентное внимание (MLA): Представлено вместе с DeepSeek-V2; сжимает ключи и значения каждого токена в небольшой общий латентный вектор.
• PagedAttention: Размещает KV-кеш в блоках фиксированного размера, чтобы избежать фрагментации; число голов, которые заполняют кеш, на это не влияет.
• Внимание с скользящим окном: Ограничивает число кешируемых токенов на слой; в Mistral 7B этот механизм сочетается с GQA.
• Flash attention: Ядро GPU, которое быстрее вычисляет точное внимание и изначально поддерживает конфигурации GQA.