Multi-Query Attention

Das Modell verbindet das aktuelle Textfragment über mehrere parallele Berechnungen mit dem früheren Gespräch. Solche Berechnungen heißen Attention-Köpfe, also Aufmerksamkeitsköpfe. Wenn jeder Kopf seine eigenen Beschreibungen des gesamten früheren Gesprächs speichert, steigen der Speicherbedarf und die Kosten des Auslesens. Lassen sich einige Daten gemeinsam nutzen, während unterschiedliche Arten des Auslesens erhalten bleiben?

Multi-Query Attention (MQA) behält viele Queries — Beschreibungen des Bedarfs der aktuellen Positionen —, gibt ihnen aber einen gemeinsamen Satz von Key und Value. Key dient dem Abgleich, Value enthält weiterzugebende Informationen. Alle sind Listen von Zahlen.

Zwei Köpfe können denselben Key/Value-Satz nutzen und trotzdem die Anteile anders verteilen, weil sie unterschiedliche Queries haben. Gemeinsame Daten erzwingen kein identisches Ergebnis.

Die Einsparung betrifft unter anderem den KV Cache, also gespeicherte Beschreibungen früheren Textes. Die Änderung kann jedoch die Qualität beeinflussen und bedeutet keine proportionale Beschleunigung des gesamten Modells. Grouped-Query Attention nutzt die gemeinsame Verwendung innerhalb kleinerer Gruppen.

Mechanismus und Details

Bei gewöhnlicher Multi-Head Attention hat jeder Kopf eigene Query-, Key- und Value-Projektionen. MQA entfernt die Kopfdimension aus K/V und deren Projektionen und erhält sie für Query. So definiert Noam Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, §3 diese Variante.

Gleicher Speicher, andere Frage

Das eigene Beispiel hat zwei Schlüssel: $k_1=(1,0)$ und $k_2=(0,1)$. Zwei Query-Köpfe erhalten jeweils $q_1=(2,0)$ und $q_2=(0,2)$. Scaled Dot-Product Attention teilt die Produkte durch $\sqrt{2}$ und wendet Softmax an. Der erste Kopf gibt den Schlüsseln Gewichte von etwa 0,804 und 0,196; der zweite kehrt diese Anteile um. Bei gemeinsamen $v_1=(1,0)$ und $v_2=(0,1)$ sind dies zugleich die Koordinaten zweier unterschiedlicher Ergebnisse.

Woher die Einsparung kommt

Der KV-Cache speichert frühere K/V. Bei acht Query-Köpfen speichert MHA acht K/V-Sätze, MQA einen. Bei gleicher Sequenzlänge, Kopfdimension und Präzision sind die Cachetensoren dann achtmal kleiner. Das ist eine Größenberechnung, keine achtfache Beschleunigung des gesamten Modells.

Shazeers Analyse in §2.4.1 und §3.1 verbindet diese Änderung mit dem Speichertransfer während aufeinanderfolgender Generierungsschritte. Allerdings verändern sich Architektur und Repräsentationsmöglichkeiten. Es handelt sich nicht um das verlustfreie Entfernen identischer Kopien aus einem beliebigen trainierten MHA. Grouped-Query Attention behält als Zwischenvariante mehrere K/V-Sätze.

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