Grouped Query Attention (GQA)

Grouped-Query Attention (GQA) ist eine Form der Mehrkopf-Attention für Transformer, bei der Gruppen von Query-Köpfen jeweils einen Key-Kopf und einen Value-Kopf gemeinsam nutzen. Llama 3.1 8B verknüpft beispielsweise in jeder Schicht 32 Query-Köpfe mit 8 Key-Value-Köpfen. GQA liegt zwischen Multi-Head Attention (MHA), bei der jeder Query-Kopf eigene Keys und Values hat, und Multi-Query Attention (MQA), bei der alle Query-Köpfe einen einzigen Key-Kopf und Value-Kopf gemeinsam nutzen. Joshua Ainslie und seine Kolleginnen und Kollegen bei Google Research stellten GQA in GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints auf der EMNLP 2023 vor. Sie zeigten, dass GQA-Modelle eine ähnliche Qualität wie MHA erreichen und annähernd so schnell decodieren wie MQA. Offene große Sprachmodelle (LLMs) wie Llama 2 70B, die Llama-3-Familie, Mistral 7B und Gemma 2 verwenden GQA. Darauf aufbauende Bild-Sprach-Modelle (VLMs), etwa Googles PaliGemma 2 auf Gemma 2, übernehmen GQA in ihren Sprachdecodern.

Funktionsweise von GQA#

Bei der Self-Attention wird jedes Token für jeden Kopf in einen Query, einen Key und einen Value projiziert. Jeder Kopf vergleicht seine Queries mit seinen Keys. GQA behält alle Query-Köpfe bei, projiziert jedoch weniger Key- und Value-Köpfe. Jede Gruppe von Query-Köpfen verwendet die gemeinsam genutzten Keys und Values ihrer Gruppe. Bei 8 Query-Köpfen und 2 Key-Value-Köpfen versorgt jeder Key-Value-Kopf 4 Query-Köpfe:

Die Berechnung der Attention innerhalb jedes Kopfes bleibt die unveränderte skalierte Skalarprodukt-Attention. Mit einer Gruppe wird die Schicht zu MQA, mit genauso vielen Gruppen wie Query-Köpfen zu MHA. PyTorch unterstützt dieses Layout direkt über das Argument enable_gqa von 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")

Ausgegeben werden Output (1, 8, 16, 64), same as copied heads: True und 256 with GQA, 1024 with MHA: Alle 8 Query-Köpfe erzeugen Ausgaben mit einem Viertel der Keys und Values, die MHA speichern würde.

Warum GQA den KV-Cache verkleinert#

Während der Textgenerierung speichert ein Modell die Keys und Values jedes vorherigen Tokens in seinem KV-Cache und liest bei jedem neuen Token den gesamten Cache aus. Daher wird die Decodiergeschwindigkeit stärker durch die Speicherbandbreite als durch die Rechenleistung begrenzt. Die Cache-Größe wächst mit der Anzahl der Key-Value-Köpfe. Llama 3.1 8B hat 32 Schichten mit jeweils 8 Key-Value-Köpfen und 128 Dimensionen. Bei 16-Bit-Genauigkeit speichert es daher 128 KiB pro Token beziehungsweise 16 GiB bei seinem vollständigen Kontextfenster von 131.072 Tokens. Mit MHA würde dasselbe Modell 512 KiB pro Token und 64 GiB bei voller Länge benötigen. Der freigewordene Speicher ermöglicht längere Kontexte, größere Batches für kontinuierliche Batchverarbeitung und eine geringere Inferenzlatenz pro Token. Die Schnittstelle LLM von Ultralytics akzeptiert ein base_url für jeden OpenAI-kompatiblen Server, etwa einen selbst gehosteten Llama 3.1 8B. Mit seinen 8 Key-Value-Köpfen kann der Server bei gleichem Cache-Speicher viermal so viele Gesprächstokens speichern wie mit einem MHA-Layout mit 32 Köpfen.

Modelle in GQA umwandeln#

Für GQA ist kein Training von Grund auf nötig. Die Autorinnen und Autoren wandelten einen bestehenden MHA-Checkpoint um, indem sie die Key- und Value-Köpfe jeder Gruppe mittelten. Anschließend setzten sie das Vortraining mit etwa 5 % des ursprünglichen Rechenaufwands fort. Diesen Schritt nennen sie Weitertraining. MQA, die Extremvariante mit einem gemeinsam genutzten Kopf, wurde 2019 von Noam Shazeer vorgestellt. Sie ist das schnellste Layout, verschlechterte jedoch die Qualität und konnte das Training instabil machen. Eine mittlere Anzahl von Gruppen stellt den Großteil dieser Qualität wieder her.

GQA und verwandte Verfahren#

• Multi-Head Latent Attention (MLA): Zusammen mit DeepSeek-V2 vorgestellt. Komprimiert die Keys und Values jedes Tokens zu einem kleinen gemeinsamen latenten Vektor.
• PagedAttention: Teilt den KV-Cache in Blöcke fester Größe auf, um Fragmentierung zu vermeiden, unabhängig davon, wie viele Köpfe ihn füllen.
• Attention mit gleitendem Fenster: Begrenzt die Anzahl der pro Schicht gespeicherten Tokens. Mistral 7B kombiniert dieses Verfahren mit GQA.
• Flash Attention: Ein GPU-Kernel, der exakte Attention schneller berechnet und GQA-Layouts nativ unterstützt.