Was ist Grouped Query Aufmerksamkeit (GQA)?

Ainslie et al. konzipierten die gruppierte Abfrageaufmerksamkeit als eine Optimierung der Multi-Head-Attention (MHA), dem innovativen Selbstbeobachtungsalgorithmus, der in der wegweisenden Arbeit „Attention is All You Need“ (2017) vorgestellt wurde, das neuronale Netze für Transformatoren einführte. Genauer gesagt wurde er als eine Verallgemeinerung und zurückhaltendere Anwendung von Multi Query Attention (MQA), einer früheren Optimierung von MHA, vorgestellt.

Obwohl die standardmäßige Multi-Head-Attention einen Evolutionssprung im maschinellen Lernen, in der Verarbeitung natürlicher Sprache (NLP) und in der generativen KI katalysiert hat, ist sie extrem anspruchsvoll in Bezug auf Rechenressourcen und Speicherbandbreite. Als die LLMs immer größer und ausgefeilter wurden, wurden diese Anforderungen an die Speichernutzung zu einem Hindernis für den Fortschritt, insbesondere für die autoregressiven LLMs, die nur für die Texterzeugung, die Zusammenfassung und andere generative KI-Aufgaben verwendet werden.

Die anschließende Forschung konzentrierte sich auf Techniken zur Verbesserung oder Optimierung der Multi-Head-Aufmerksamkeit. Einige, wie z. B. Flash Attention und Ring Attention, verbessern wie GPUs, die zum Trainieren und Ausführen von Modellen verwendet werden, Berechnungen und Speicher verarbeiten. Andere, wie GQA und MQA, erkunden Änderungen an der Art und Weise, wie Transformer-Architektur Token verarbeiten.

Die Grouped Query Attention zielt darauf ab, die Kompromisse zwischen der standardmäßigen Multi-Head-Attention und der Multi-Query-Attention auszugleichen. Erstere maximiert die Genauigkeit auf Kosten eines erhöhten Speicherbandbreiten-Overheads und geringerer Geschwindigkeit. Letztere maximiert die Geschwindigkeit und Effizienz auf Kosten der Genauigkeit.

Um zu verstehen, wie Grouped Query Attention Transformer-Modelle optimiert, ist es wichtig, zunächst zu verstehen, wie Multihead-Aufmerksamkeit im Allgemeinen funktioniert. Sowohl GQA als auch MQA verfeinern einfach die Kernmethodik von MHA, anstatt sie zu ersetzen.

Die treibende Kraft hinter LLMs und anderen Modellen, die die Transformer-Architektur verwenden, ist die Selbstbeobachtung, ein mathematischer Framework zum Verständnis der Beziehungen zwischen den verschiedenen Token in einer Sequenz. Die Selbstbeobachtung ermöglicht es einem LLM, Textdaten nicht nur durch statische Basisdefinitionen, sondern auch durch den Kontext anderer Wörter und Phrasen zu interpretieren.

In autoregressiven LLMs, die zur Textgenerierung verwendet werden, hilft der Beobachtungsmechanismus dem Modell, das nächste Token in einer Sequenz vorherzusagen, indem er bestimmt, welche vorherigen Token es in diesem Moment am meisten wert „beachtet zu werden“. Informationen von Token, die er als am relevantesten einstuft, erhalten eine größere Beobachtungsgewichtung, während Informationen von Token, die als irrelevant eingestuft werden, eine Aufmerksamkeitsgewichtung von nahezu 0 erhalten.

Der Multi-Head-Aufmerksamkeitsmechanismus, der Transformer-Modelle animiert, generiert umfangreiche Kontextinformationen, indem er die Selbstaufmerksamkeit viele Male parallel berechnet und dabei Aufmerksamkeitsebenen in mehrere Attention Heads aufteilt.

Die Autoren von „Attention is All You Need“ haben ihren Beobachtungsmechanismus anhand der Terminologie einer relationalen Datenbank formuliert: Abfragen, Schlüssel und Werte. Relationale Datenbanken sollen die Speicher und den Abruf relevanter Daten vereinfachen: Sie weisen jedem Datenelement eine eindeutige Kennung („Schlüssel“) zu, wobei jeder Schlüssel mit einem entsprechenden Wert verknüpft ist. Das Ziel einer relationalen Datenbank besteht darin, jede Abfrage dem entsprechenden Schlüssel zuzuordnen.

Für jedes Token in einer Sequenz erfordert die Multi-Head-Aufmerksamkeit die Erstellung von 3 Vektoren.

• Ein Abfragevektor, Q, der die Informationen darstellt, die das Token „sucht“. Zum Beispiel könnte der Abfragevektor für ein Substantiv eine Suche nach Adjektiven darstellen, die es beschreiben.
  
  
• Ein Schlüsselvektor, K, der die im Token enthaltenen Informationen darstellt. Zur Berechnung der Aufmerksamkeitsgewichte werden Ausrichtungsscores verwendet, die Relevanz des Schlüsselvektors eines jeden Tokens für den Abfragevektor eines jeden anderen Tokens darstellen.
  
  
• Ein Wertevektor, V, der die Kontextinformationen darstellt, die durch die Beachtung gewichteter Beiträge aus den Schlüsselvektoren anderer Token aktualisiert werden.

Die mathematischen Wechselwirkungen zwischen diesen 3 Vektoren, die durch den Attention-Mechanismus vermittelt werden, sind die Art und Weise, wie ein Modell sein kontextspezifisches Verständnis jedes Tokens anpasst.
 

Um jeden dieser 3 Vektoren für ein bestimmtes Token zu generieren, beginnt das Modell mit der ursprünglichen Vektoreinbettung dieses Tokens: eine numerische Kodierung, bei der jede Dimension des Vektors einem abstrakten Element der semantischen Bedeutung des Tokens entspricht. Die Anzahl der Dimensionen in diesen Vektoren ist ein vorgegebener Hyperparameter.

Die Q-, K- und V-Vektoren für jedes Token werden generiert, indem die ursprüngliche Token-Einbettung durch eine lineare Schicht geleitet wird, die der ersten Aufmerksamkeitsschicht vorausgeht. Diese lineare Schicht ist in 3 einzigartige Matrizen von Modellgewichten unterteilt: W_Q, W_K und W_V. Die darin enthaltenen spezifischen Gewichtungswerte werden durch selbstüberwachtes Vortraining an einem riesigen Datensatz von Textbeispielen gelernt.

Multipliziert man die ursprüngliche Vektoreinbettung des Tokens mit W_Q, W_K und W_V, erhält man den entsprechenden Abfragevektor, den Schlüsselvektor bzw. den Wertvektor. Die Anzahl der Dimensionen d, die jeder Vektor enthält, wird durch die Größe der einzelnen Gewichtungsmatrix bestimmt. Q und K haben die gleiche Anzahl von Dimensionen, d_k.

Diese 3 Vektoren werden dann an die Aufmerksamkeitsschicht weitergegeben.

In der Attention-Ebene werden die Q-, K - und V-Vektoren verwendet, um einen Alignment-Score zwischen jedem Token an jeder Position in einer Sequenz zu berechnen. Diese Ausrichtungswerte werden dann mithilfe einer Softmax-Funktion in Aufmerksamkeitsgewichte normalisiert.

Für jedes Token x in einer Sequenz werden die Alignment-Scores berechnet, indem das Punktprodukt des Abfragevektors Q_x dieses Tokens mit dem Schlüsselvektor K jedes der anderen Token berechnet wird: mit anderen Worten, indem sie miteinander multipliziert werden. Wenn sich eine sinnvolle Beziehung zwischen 2 Token in Ähnlichkeiten zwischen ihren jeweiligen Vektoren widerspiegelt, ergibt die Multiplikation einen großen Wert. Wenn die 2 Vektoren nicht aufeinander abgestimmt sind, ergibt die Multiplikation einen kleinen oder negativen Wert. Die meisten Transformer-Modelle verwenden eine Variante, die als Scaled Dot-Product Attention bezeichnet wird, bei der QK skaliert, d. h. mit $\frac{1}{\sqrt{d_k}}$  zur Verbesserung der Trainingsstabilität multipliziert wird.

Diese Abfrage-Schlüssel-Alignment-Scores werden dann in eine Softmax-Funktioneingetippt. Softmax gleicht alle Eingaben auf einen Wert zwischen 0 und 1 aus, sodass sie alle auf 1 addieren. Die Ausgaben der Softmax-Funktion sind die Aufmerksamkeitsgewichte, die jeweils den Anteil (aus 1) der Aufmerksamkeit von Token x darstellen, der jedem der anderen Token zukommt. Wenn das Aufmerksamkeitsgewicht eines Tokens nahe bei 0 liegt, wird es ignoriert. Eine Aufmerksamkeitsgewichtung von 1 würde bedeuten, dass ein Token die gesamte Aufmerksamkeit von x erhält und alle anderen ignoriert werden.

Schließlich wird der Wertvektor für jeden Token mit seinem Aufmerksamkeitsgewicht multipliziert. Diese aufmerksamkeitsgewichteten Beiträge jedes vorherigen Tokens werden zusammengezählt und zur ursprünglichen Vektoreinbettung für Token x addiert. Damit wird die Einbettung von Token x nun aktualisiert, um den Kontext widerzuspiegeln, der von den anderen Token in der Sequenz bereitgestellt wird, die für sie relevant sind.

Die aktualisierte Vektoreinbettung wird dann an eine andere lineare Schicht mit einer eigenen Gewichtungsmatrix W_Z gesendet, wo der kontextaktualisierte Vektor wieder auf eine konsistente Anzahl von Dimensionen normalisiert und dann an die nächste Aufmerksamkeitsebene gesendet wird. Jede progressive Aufmerksamkeitsebene erfasst größere kontextuelle Nuancen.

Die Verwendung von Durchschnittswerten der Attention-gewichteten Beiträge von anderen Tokens, anstatt jeden Teil des Attention-gewichteten Kontexts einzeln zu berücksichtigen, ist zwar mathematisch effizient, aber es führt zu einem Detailverlust.

Um dies zu kompensieren, teilen die Transformer-Netzwerke die ursprüngliche Eingabe des Token-Einbettung in h gleichgroße Teile auf. Sie teilen W_Q,W_K und W_V ebenfalls in h Teilmengen auf, die Query Heads, Key Heads bzw. Value Heads genannt werden. Jeder dieser Heads erhält einen Teil der ursprünglichen Token-Einbettung. Die Vektoren, die von jedem dieser parallelen Tripletts von Abfrageköpfen, Schlüsselköpfen und Wertköpfen erzeugt werden, werden in einen entsprechenden Attention Head eingespeist.Schließlich werden die Ausgaben dieser h parallelen Schaltkreise wieder zusammengefügt, um die vollständige Token-Einbettung zu aktualisieren.

Beim Training lernt jeder Kreislauf unterschiedliche Gewichte, die einen separaten Aspekt der Semantik-Bedeutung erfassen. Dies wiederum hilft dem Modell, die verschiedenen Arten zu verarbeiten, in denen die Implikationen eines Wortes durch den Kontext anderer Wörter in seiner Umgebung beeinflusst werden können.

Die Kehrseite der normalen Multi-Head-Attention ist nicht so sehr das Vorhandensein eines entscheidenden Fehlers, sondern das Fehlen jeglicher Optimierung. MHA war der erste Algorithmus seiner Art und stellt die komplexeste Ausführung seines allgemeinen Mechanismus zur Attendorn-Berechnung dar.

Der größte Teil der Ineffizienz von MHA ist auf die Fülle von Berechnungen und Modellparametern zurückzuführen. In der Standard-MHA hat jeder Query-Head, Key-Head und Value-Head in jedem Aufmerksamkeitsblock seine eigene Matrix von Gewichten. Ein Modell mit 8 Aufmerksamkeitsköpfen in jeder Aufmerksamkeitsschicht – weit weniger als bei den meisten modernen LLMs – würde also 24 eindeutige Gewichtsmatrizen allein für die Q-, K- und V-Köpfe der Schicht erfordern. Dies bedeutet eine enorme Anzahl von Zwischenberechnungen in jeder Schicht.

Eine Folge dieser Konfiguration ist, dass sie rechenintensiv ist. Die Rechenanforderungen für MHA skalieren quadratisch in Bezug auf die Sequenzlänge: Die Verdoppelung der Anzahl von Token in einer Eingabesequenz erfordert die vierfache Komplexität. Dies setzt der Größe von Kontextfenstern in der Praxis enge Grenzen.

MHA belastet auch den Systemspeicher stark. GPUs verfügen nicht über viel internen Speicher, um die Ergebnisse der riesigen Menge an Zwischenberechnungen zu speichern, die bei jedem weiteren Verarbeitungsschritt abgerufen werden müssen. Diese Zwischenergebnisse werden stattdessen im High-Bandwidth-Speicher (HBM) gespeichert, der sich nicht auf dem GPU-Chip selbst befindet. Dies führt jedes Mal, wenn Schlüssel und Werte aus dem Speicher gelesen werden müssen, zu einer geringen Latenz. Als die Transformer-Modelle auf viele Milliarden von Parametern skalieren wurden, wurden die Zeit und der Rechenaufwand für das Training und die Ausführung der Inferenz zu einem Engpass für die Modellleistung.

Für weitere Fortschritte waren Methoden erforderlich, mit denen die Anzahl der Rechenschritte reduziert werden konnte, ohne die Fähigkeit der Transformatoren zu beeinträchtigen, komplizierte sprachliche Muster zu erlernen und zu reproduzieren. In diesem Zusammenhang wurden die MQA und später die GQA eingeführt.

Multi-Query Attention (MQA) ist ein rechnerisch effizienterer Aufmerksamkeitsmechanismus, der Multi-Head Attention vereinfacht, um Speicherverbrauch und Zwischenberechnungen zu reduzieren. Anstatt für jeden Aufmerksamkeitskopf einen eindeutigen Schlüsselkopf und Wertekopf zu trainieren, verwendet MQA auf jeder Ebene einen einzelnen Schlüsselkopf und einen einzelnen Wertekopf. Daher werden Schlüsselvektoren und Wertevektoren nur einmal berechnet; dieser einzelne Satz von Schlüssel- und Wertevektoren wird dann von allen h Aufmerksamkeitsköpfen gemeinsam genutzt.

Durch diese Vereinfachung wird die Anzahl der linearen Projektionen, die das Modell berechnen und im Speicher mit hoher Bandbreite speichern muss, erheblich reduziert. Laut dem Papier aus dem Jahr 2019, in dem MQA vorgestellt wurde, ermöglicht MQA einen 10- bis 100-mal kleineren Schlüssel-Wert-Paarspeicher (oder KV-Cache) und eine 12-mal schnellere Decoder-Inferenz. Die reduzierte Speicherauslastung von MQA beschleunigt auch das Training erheblich, indem eine größere Batchgröße ermöglicht wird.

Trotz seiner Vorteile bringt MQA einige unvermeidliche Nachteile mit sich.

• Leistungsdegradation: Natürlich reduziert die Verringerung der Anzahl eindeutiger, trainierbarer Modellparameter die Fähigkeit des Modells, Wissen und Nuancen zu erkennen. MQA führt zu einem deutlichen Rückgang der Genauigkeit im Vergleich zur Standard-MHA, was es für bestimmte Situationen und Anwendungsfälle ungeeignet macht.
  
  
• Muss von Grund auf neu trainiert werden: Ein Modell, das mit Standard-MHA trainiert wurde, kann nicht einfach an MQA angepasst werden, sondern muss stattdessen von Grund auf mit MQA trainiert werden. Dies bedeutet, dass MQA nicht zur Optimierung bestehender Modelle verwendet werden kann und erhebliche Opportunitätskosten beim Experimentieren mit MQA für neue Modelle mit sich bringt.
  
  
• Redundanzen in der Tensorparallelität: Einer der Hauptvorteile des Trainings von Transformatormodellen auf GPUs ist die Möglichkeit, mehrere komplexe Tensoroperationen parallel auszuführen. K- und V-Werte müssen auf jedem Knoten des GPU-Clusters vorhanden sein, der diese Vorgänge ausführt, was bedeutet, dass sie in der Praxis für jeden Knoten repliziert werden müssen. Dies ist keine optimale Nutzung der Ressourcen, obwohl es immer noch effizienter ist als Standard-MHA.

Die Aufmerksamkeit für gruppierte Abfragen ist eine allgemeinere, flexiblere Formulierung der Aufmerksamkeit für mehrere Abfragen, bei der Abfrageköpfe in mehrere Gruppen unterteilt werden, die jeweils einen Satz von Schlüsseln und Werten teilen, anstatt einen Satz von Schlüsseln und Werten für alle Abfrageköpfe freizugeben.

Nachdem „GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints“ im Mai 2023 veröffentlicht wurde, haben viele LLMs GQA schnell übernommen Meta beispielsweise übernahm GQA erstmals im Juli 2023 für seine Llama-2-Modelle und behielt GQA in den Llama-3-Modellen bei, die im Jahr 2024 veröffentlicht wurden. Mistral KI verwendete GQA für das Modell Mistral 7B, das im September 2023 auf den Markt kam. Auch die IBM Granite 3.0-Modelle verwenden GQA für eine schnelle Inferenz.

Theoretisch kann GQA als eine Verallgemeinerung des Spektrums zwischen Standard-MHA und vollständigem MQA betrachtet werden. GQA mit der gleichen Anzahl von Schlüsselwert-Head-Gruppen wie Attention Heads entspricht der Standard-MHA; GQA mit 1 Head-Gruppe ist das Äquivalent zu MQA.

In der Praxis impliziert GQA fast immer einen Zwischenansatz, bei dem die Anzahl der Gruppen selbst ein wichtiger Hyperparameter ist.

Grouped Query Attention bietet mehrere Vorteile, die dazu geführt haben, dass sie relativ weit verbreitet für führende LLMs eingesetzt wird.

• Effiziente GPU-Nutzung: Die Verteilung der Schlüssel-Wert-Paare durch GQA nutzt die Vorteile der Tensorparallelität und reduziert die Rechenleistung, die durch die Replikation redundanter Werte „verschwendet“ wird.
  
  
• Effektiver Kompromiss: GQA bietet einen idealen Kompromiss zwischen Decoder-Inferenzgeschwindigkeit und Leistung, da es fast so genau wie MHA und gleichzeitig fast so schnell wie MQA ist.
  
  
• Reduzierter Speicherbandbreitenaufwand: Wie MQA reduziert GQA deutlich die Zwischenberechnungen, die zur Inferenzzeit berechnet, speichern und abgerufen werden müssen.
  
  
• Flexibles Training: Im Gegensatz zu MQA erfordert Group Query Attention nicht, dass die Modelle mit diesem Ansatz von Grund auf neu trainiert werden. Modelle, die mit Standard-MHA trainiert wurden, können durch einen Feinabstimmungsprozess namens „Uptraining“ an die Verwendung von GQA angepasst werden.