Grouped-Query Attention

Das Modell liest früheren Text auf mehrere Arten parallel. Jede solche Berechnung ist ein Attention-Kopf, also ein Aufmerksamkeitskopf. Du möchtest den Speicherbedarf dieser Köpfe verringern, aber die vollständige gemeinsame Nutzung eines einzigen Datensatzes könnte das Modell zu stark einschränken. Möglich ist eine Zwischenvariante: Manche Köpfe nutzen einen gemeinsamen Satz, andere einen anderen.

Grouped-Query Attention (GQA) unterteilt Query-Köpfe in Gruppen, die Key und Value gemeinsam nutzen. Query ist eine numerische Beschreibung zum Lesen des Kontexts, Key zum Abgleich und Value zur Weitergabe von Informationen. Jeder Kopf berechnet weiterhin seine eigenen Anteile.

Acht Query-Köpfe, aufgeteilt in zwei Gruppen mit je vier Köpfen, bedeuten zwei Key/Value-Sätze statt acht. Eine einzige Gruppe ergibt MQA; eine Gruppe pro Kopf entspricht gewöhnlichem MHA.

Weniger Sätze verringern bei passenden Dimensionen die Speicherkosten für KV. Das ist ein Kompromiss zwischen Konstruktion und Qualität, kein Schalter, der eine bestimmte Geschwindigkeit garantiert. Die Laufzeit hängt auch von der Hardware und den übrigen Berechnungen ab.

Mechanismus und Details

Für $H$ Query-Köpfe und $G$ Gruppen erhalten wir $G$ K/V-Sätze. Bei $G=1$ ist dies Multi-Query Attention. Bei $G=H$ kehren wir zu Multi-Head Attention zurück. Diese Familie und die Bezeichnung GQA-g beschreiben Ainslie et al., GQA, §2.2.

Berechne den Speicher, bevor du die Geschwindigkeit schätzt

Die eigene Berechnung für den KV Cache nimmt Batchgröße 1, 32 Schichten, 1024 gespeicherte Tokens, die Dimension 128 für jeden Key/Value-Kopf und 2 Byte pro Zahl an. Die K/V-Elemente allein belegen:

$$2\cdot32\cdot1024\cdot G\cdot128\cdot2\ \text{Byte}$$

Die erste Zwei steht für K und V. Bei acht Query-Köpfen ergibt sich 128 MiB für MHA, 32 MiB für zwei KV-Köpfe und 16 MiB für MQA. Ein MiB sind $2^{20}$ Byte. Modellgewichte, Puffer, Speicherausrichtung und Replikation zwischen Geräten zählen wir nicht mit.

Die Verringerung der KV-Kopfanzahl verkürzt hier weder den Kontext noch entfernt sie Query-Köpfe. Sie begrenzt die Zahl separater K/V-Repräsentationen. Die Speicherrechnung sagt weder Qualität noch Laufzeit voraus; beide hängen auch von Training, Hardware und Implementierung ab.

Bei der Checkpoint-Konvertierung mittelten die Autoren die K/V-Projektionen innerhalb der Gruppen und setzten danach das Pre-Training fort (§2.1–2.2). Die Ergebnisse in §3 betreffen konkrete T5-Varianten und Aufgaben. Sie bestimmen weder eine universelle Gruppenanzahl noch garantieren sie Qualitätserhalt beim bloßen Umstellen der Konfiguration.

Ich verwende KI-generierte Inhalte als Teil meines täglichen Lernprozesses.