QKV Projections Require a Fraction of Their Memory

Die Arbeit stellt PAMM vor, eine neuartige Tensor-Kompressionstechnik, die den Speicherbedarf der QKV-Projektionen in LLM-Trainingsprozessen um den Faktor 512 reduziert und dabei die Perplexität erhält oder sogar verbessert.

Das Wesentliche

Das Problem: Der riesige Rucksack

Stell dir vor, ein großes Sprachmodell (ein KI-Modell wie ChatGPT) ist wie ein genialer Architekt, der gerade einen riesigen Wolkenkratzer baut. Um jeden Stein zu setzen, muss der Architekt ständig auf seine Baupläne schauen.

Beim Training dieser KI passiert Folgendes:

1. Der Architekt liest einen Satz (die Eingabe).
2. Er wandelt diesen Satz in eine komplexe Formel um (die sogenannten Q, K und V Projektionen).
3. Er vergleicht diese Formeln, um zu verstehen, welche Wörter zusammengehören (das ist die „Aufmerksamkeit" oder Attention).
4. Das Problem: Damit der Architekt später lernen kann, was er falsch gemacht hat, muss er sich jeden einzelnen Bauplan merken, den er während des Bauens erstellt hat.

Bei großen Modellen sind diese Baupläne so riesig, dass sie den gesamten Speicherplatz (den Arbeitsspeicher des Computers) füllen. Es ist, als würde der Architekt für jeden einzelnen Ziegelstein einen ganzen LKW voller Papierpläne mit sich herumschleppen. Er kann kaum noch neue Steine setzen, weil er nur noch Platz für den LKW hat.

Die Lösung: PAMM – Der „Stichwort"-Trick

Die Forscher aus Israel haben eine clevere Methode namens PAMM (Point-Approximate Matrix Multiplication) entwickelt. Sie nennen es „Punkt-Approximative Matrix-Multiplikation", aber lass uns das mit einer besseren Analogie erklären:

Stell dir vor, der Architekt hat 10.000 verschiedene Skizzen von Fenstern vor sich.

• Der alte Weg: Er fotografiert jede einzelne der 10.000 Skizzen und speichert sie alle auf einer riesigen Festplatte. Das kostet enorm viel Platz.
• Der PAMM-Weg: Der Architekt schaut sich die 10.000 Skizzen an und merkt: „Hey, 9.900 davon sehen fast genau gleich aus! Sie sind nur leicht gedreht oder etwas heller."

Anstatt alle 10.000 Fotos zu speichern, macht er folgendes:

1. Er wählt nur 20 repräsentative Skizzen aus (die sogenannten „Generatoren").
2. Für die restlichen 9.980 Skizzen schreibt er nur einen kurzen Zettel: „Skizze Nr. 50 ist fast wie Skizze Nr. 3, nur 10 % heller."

Das ist der Kern von PAMM: Anstatt die riesigen Datenmengen zu speichern, speichert er nur die wichtigsten Muster und eine kurze Anleitung, wie man den Rest daraus rekonstruieren kann.

Warum funktioniert das? (Die Cluster-Analogie)

Warum sind die Daten überhaupt so ähnlich? Stell dir vor, du hast eine riesige Menge an Fotos von Hunden.

• Die meisten Fotos zeigen Hunde, die auf vier Beinen stehen.
• Nur wenige zeigen Hunde, die tanzen.

Die KI sieht in ihren Daten oft immer wieder die gleichen Muster (z. B. „dieses Wort kommt oft nach jenem Wort"). Die Forscher haben entdeckt, dass man diese Wiederholungen extrem stark komprimieren kann, ohne dass die KI etwas Wichtiges verliert.

Es ist, als würdest du eine Bibliothek mit 1 Million Büchern haben, aber 99 % davon sind nur Kopien desselben Buches mit leicht unterschiedlichen Einbandfarben. Anstatt 1 Million Bücher zu lagern, lagerst du ein Originalbuch und einen Zettel mit der Liste: „Buch 1 bis 1000 sind wie das Original, aber mit rotem Einband."

Die Ergebnisse: Weniger Platz, gleiche Intelligenz

Das Paper zeigt beeindruckende Ergebnisse:

• Platzersparnis: PAMM spart bis zu 97–99 % des Speichers ein, der für diese Zwischenschritte benötigt wird. Das ist, als würde man einen LKW voller Papier in einen kleinen Rucksack verwandeln.
• Qualität: Die KI wird nicht dumm. Sie lernt genauso gut wie vorher, manchmal sogar besser, weil der „Lärm" der überflüssigen Details wegfällt.
• Geschwindigkeit: Es kostet kaum Zeit, diese Komprimierung durchzuführen. Der Architekt muss nicht langsamer bauen.

Zusammenfassung für den Alltag

Stell dir vor, du musst eine lange Reise mit einem Auto machen, das nur 5 Liter Tankkapazität hat, aber du brauchst 500 Liter Benzin für die ganze Strecke.

• Ohne PAMM: Du musst ständig anhalten, um neue Tanks zu organisieren, und kommst kaum voran.
• Mit PAMM: Du stellst fest, dass du das Benzin nicht in 500 einzelnen Kanistern brauchst, sondern in einem einzigen, hochverdichteten Tank, der den gleichen Inhalt hat. Du fährst weiter, ohne anzuhalten, und kommst schneller ans Ziel.

Das Fazit: Die Forscher haben einen Weg gefunden, die KI-Modelle so effizient zu trainieren, dass wir bald viel größere und intelligentere Modelle auf ganz normalen Computern (oder mit weniger teuren Grafikkarten) bauen können, ohne dabei an Intelligenz zu verlieren. Sie haben den „Rucksack" der KI entladen, damit sie endlich frei laufen kann.

Technische Zusammenfassung

1. Problemstellung

Das Training von Large Language Models (LLMs) ist durch hohe Anforderungen an Rechenleistung und vor allem an den GPU-Speicher (VRAM) limitiert. Ein wesentlicher Engpass liegt in der Multi-Head-Attention-Schicht. Während viele Forschungsarbeiten darauf abzielen, die Berechnung der Skalierten-Dot-Product-Attention selbst zu optimieren (z. B. durch FlashAttention), wird der Speicherbedarf der linearen Projektionen (Q, K, V) oft übersehen.

• Der Engpass: Während des Vorwärtsdurchlaufs (Forward Pass) müssen die Eingabeaktivierungen $x$ gespeichert werden, um im Rückwärtsdurchlauf (Backward Pass) die Gradienten für die Gewichtsmatrizen $W_Q, W_K, W_V$ zu berechnen ($\nabla W = X^\top \nabla Z$).
• Die Konsequenz: Diese Aktivierungen können bis zu 20 % des gesamten Spitzen-GPUspeichers der Attention-Blöcke ausmachen. Da Aktivierungen mit der Batch-Größe und der Sequenzlänge skalieren, werden sie zum limitierenden Faktor, insbesondere bei langen Kontexten und großen Batches.
• Bisherige Ansätze: Viele Komprimierungsmethoden nutzen niedrigen Rang (Low-Rank) in der versteckten Dimension (Embedding-Dimension). Das Paper argumentiert jedoch, dass eine viel größere Redundanz in der Sequenzdimension (über die Tokens hinweg) existiert, die bisher ungenutzt blieb.

2. Methodik: Point-Approximate Matrix Multiplication (PAMM)

Die Autoren schlagen PAMM vor, eine neue Tensor-Komprimierungstechnik, die die Aktivierungen der Q-, K- und V-Projektionen drastisch reduziert, ohne die Modellleistung signifikant zu beeinträchtigen.

Kernprinzip

PAMM nutzt die Beobachtung, dass Token-Aktivierungen in Transformer-Modellen oft stark redundant sind (durch Wiederholungen, Padding oder lokale Kontextähnlichkeit). Anstatt alle $b$ Token-Rows zu speichern, werden nur eine kleine Menge repräsentativer „Generator-Punkte" gespeichert.

Der Algorithmus (in zwei Stufen)

Gegeben sei eine Matrix $A \in \mathbb{R}^{b \times n}$ (die Aktivierungen) und eine Matrix $B \in \mathbb{R}^{b \times m}$ (z. B. der Gradient $\nabla Z$). Ziel ist die Approximation von $O = A^\top B$.

1. Komprimierung (Forward Pass / Speichern):

   • Es werden $k$ Generator-Punkte $C = \{C_j\}_{j=1}^k$ zufällig aus den Zeilen von $A$ ausgewählt (ohne Zurücklegen).
   • Für jede Zeile $A_i$ wird der nächste Punkt auf der Linie bestimmt, die von einem der Generatoren aufgespannt wird. Dies entspricht der Projektion von $A_i$ auf $C_j$.
   • Die beste Zuordnung $f(i)$ wird basierend auf der Kosinus-Ähnlichkeit gewählt (Lemma 1: Der Generator mit der höchsten absoluten Kosinus-Ähnlichkeit ist der beste Repräsentant).
   • Ein Skalierungsfaktor $\alpha_i$ wird berechnet.
   • Eine Nachbarbedingung (Neighborhood Condition) mit Toleranz $\varepsilon$ wird geprüft: Wenn der Abstand zwischen $A_i$ und seinem Repräsentanten zu groß ist, wird die Zeile verworfen ($\alpha_i = 0$).
   • Speicher: Statt der vollen Matrix $A$ werden nur gespeichert:
     • Die Generatoren $C \in \mathbb{R}^{k \times n}$ (sehr klein, da $k \ll b$).
     • Die Zuordnung $f \in \mathbb{R}^b$ (Indizes).
     • Die Skalierungsfaktoren $\alpha \in \mathbb{R}^b$.
   • Der Speicherbedarf sinkt von $O(b \cdot n)$ auf $O(k \cdot n + 2b)$.

2. Approximierte Multiplikation (Backward Pass):

   • Statt den vollen Gradienten $\nabla W = A^\top B$ zu berechnen, wird eine rekonstruierte Matrix $\tilde{A}$ verwendet.
   • Durch Ausnutzen der Struktur (dass viele Zeilen denselben Generator teilen) wird die Berechnung umgeformt:
     $$\tilde{O} = \sum_{j=1}^k C_j^\top \cdot \left( \sum_{i: f(i)=j} \alpha_i B_i \right)$$
   • Zuerst wird $B$ in eine kleinere Matrix $\tilde{B} \in \mathbb{R}^{k \times m}$ komprimiert (Aggregation der gewichteten Zeilen von $B$ nach ihren Generatoren).
   • Dann wird die Multiplikation $\tilde{O} = C^\top \tilde{B}$ durchgeführt. Dies ist deutlich billiger als die ursprüngliche Operation.

3. Korrektur:

   • Da einige Zeilen verworfen wurden ($\alpha_i = 0$), wird ein Korrekturfaktor $\beta$ eingeführt, um den Erwartungswert der Schätzung unverzerrt zu halten.

3. Wichtige Beiträge

• Identifikation der Sequenz-Redundanz: Nachweis, dass die Redundanz über die Token-Sequenz (Batch/Sequenz-Dimension) für die Komprimierung von Attention-Aktivierungen viel effektiver ist als Redundanz in der Embedding-Dimension.
• Einfache Zufallsauswahl: Überraschenderweise reicht eine zufällige Auswahl der Generator-Punkte aus, um die Leistung zu erhalten. Komplexe Clustering-Algorithmen sind nicht nötig und wären zu rechenintensiv.
• Kompatibilität: PAMM ist vollständig kompatibel mit bestehenden effizienten Techniken wie FlashAttention, Gradient Checkpointing und Low-Rank Adapters (LoRA). Es greift nicht in die Vorwärtsrechnung oder die Gradienten anderer Schichten ein.
• Theoretische Garantien: Das Paper liefert theoretische Schranken für die Anzahl der Generatoren $k$, die notwendig sind, um die Datenverteilung mit hoher Wahrscheinlichkeit abzudecken (Lemma 2). Es zeigt, dass $k$ nur logarithmisch mit der Batch-Größe wachsen muss.

4. Ergebnisse

Die Autoren evaluieren PAMM auf verschiedenen Modellen (LLaMA-60M bis 7B, RoBERTa, Pixtral-12B) und Aufgaben (Pretraining, Finetuning).

• Speichereinsparung: PAMM reduziert den Speicherbedarf für Q-, K-, V-Aktivierungen um einen Faktor von bis zu 512x (entspricht einer Komprimierung auf ca. 0,2 % des ursprünglichen Speichers). In absoluten Zahlen bedeutet dies oft eine Reduktion von mehreren GB auf wenige MB.
• Modellleistung (Perplexity & Accuracy):
  • Beim Pretraining von LLaMA-Modellen auf dem C4-Datensatz bleibt die Perplexity nahezu unverändert oder verbessert sich sogar leicht bei extremen Kompressionsraten ($r = 1/512$).
  • Beim Finetuning auf dem GLUE-Benchmark (RoBERTa) werden die Ergebnisse mit voller Feinabstimmung fast erreicht, während der Speicherbedarf um über 97 % sinkt.
  • Auch bei Multi-Modal-Modellen (Pixtral-12B mit LoRA) bleibt die Leistung erhalten.
• Durchsatz (Throughput): Der Overhead durch PAMM ist minimal. Bei größeren Modellen (LLaMA-1B und 7B) liegt der Durchsatzverlust unter 2,7 %. Bei kleineren Modellen ist der Overhead höher, aber immer noch akzeptabel (< 20 %).
• Vergleich mit anderen Methoden: PAMM übertrifft andere Komprimierungsmethoden wie CompAct oder Uniform-CRS (Column-Row-Sampling) deutlich. Diese Methoden leiden bei hohen Kompressionsraten unter starken Leistungseinbußen, während PAMM stabil bleibt.

5. Bedeutung und Fazit

Das Paper demonstriert, dass der Speicherbedarf von Attention-Projektionen während des Trainings durch eine intelligente Approximation der Matrixmultiplikation fast vollständig eliminiert werden kann.

• Praktische Relevanz: PAMM ermöglicht das Training größerer Modelle oder mit größeren Batch-Größen auf vorhandener Hardware, ohne die Modellqualität zu opfern.
• Synergie: Da PAMM orthogonal zu FlashAttention und anderen Optimierungen ist, kann es als „Plug-in" in bestehende Trainings-Pipelines integriert werden, um die Gesamteffizienz weiter zu steigern.
• Paradigmenwechsel: Die Arbeit verschiebt den Fokus von der Komprimierung der Parameter oder Optimierer-Zustände hin zur Komprimierung der Aktivierungen in der Sequenzdimension, was als bisher übersehene, aber hochredundante Quelle identifiziert wurde.

Zusammenfassend bietet PAMM eine einfache, theoretisch fundierte und empirisch bewiesene Lösung, um den „Memory Wall" beim Training von LLMs effektiv zu durchbrechen.