Recent work in machine learning community proposed multiple methods for performing lossy compression (quantization) of large matrices. This quantization is important for accelerating matrix multiplication (main component of large language models), which is often bottlenecked by the speed of loading these matrices from memory. Unlike classical vector quantization and rate-distortion theory, the goal of these new compression algorithms is to be able to approximate not the matrices themselves, but their matrix product. Specifically, given a pair of real matrices $A,B$ an encoder (compressor) is applied to each of them independently producing descriptions with $R$ bits per entry. These representations subsequently are used by the decoder to estimate matrix product $A^\top B$. In this work, we provide a non-asymptotic lower bound on the mean squared error of this approximation (as a function of rate $R$) for the case of matrices $A,B$ with iid Gaussian entries. Algorithmically, we construct a universal quantizer based on nested lattices with an explicit guarantee of approximation error for any (non-random) pair of matrices $A$, $B$ in terms of only Frobenius norms $\|\bar{A}\|_F, \|\bar{B}\|_F$ and $\|\bar{A}^\top \bar{B}\|_F$, where $\bar{A},\bar{B}$ are versions of $A,B$ with zero-centered columns, respectively. For iid Gaussian matrices our quantizer achieves the lower bound and is, thus, asymptotically optimal. A practical low-complexity version of our quantizer achieves performance quite close to optimal. In addition, we derive rate-distortion function for matrix multiplication of iid Gaussian matrices, which exhibits an interesting phase-transition at $R\approx 0.906$ bit/entry, showing necessity of Johnson-Lindestrauss dimensionality reduction (sketching) in the low-rate regime.
- Papier-ID: 2410.13780
- Titel: Optimal Quantization for Matrix Multiplication
- Autoren: Or Ordentlich (Hebrew University of Jerusalem), Yury Polyanskiy (MIT)
- Klassifizierung: cs.IT cs.AI cs.CL cs.LG math.IT
- Veröffentlichungsdatum: Oktober 2024 (arXiv-Preprint)
- Papierlink: https://arxiv.org/abs/2410.13780
Dieses Papier untersucht eingehend die Quantisierungsproblematik bei großflächiger Matrixmultiplikation. Im Gegensatz zur klassischen Vektorquantisierung besteht das Ziel nicht darin, die Matrizen selbst zu approximieren, sondern deren Matrixprodukt. Gegeben zwei reelle Matrizen A und B komprimiert der Encoder jede Matrix unabhängig, wobei jeder Eintrag mit R Bits beschrieben wird. Der Decoder nutzt diese komprimierten Darstellungen, um das Matrixprodukt A⊤B zu schätzen. Das Papier liefert nicht-asymptotische untere Schranken für die approximative mittlere quadratische Abweichung bei Matrizen mit unabhängig identisch verteilten Gaußschen Einträgen, konstruiert universelle Quantisierer basierend auf verschachtelten Gittern und entdeckt ein interessantes Phasenübergangsphänomen bei R ≈ 0,906 Bits/Eintrag, das die Notwendigkeit der Johnson-Lindenstrauss-Dimensionsreduktion bei niedrigen Bitraten anzeigt.
Mit dem Aufstieg von tiefen neuronalen Netzen und großen Sprachmodellen ist die Matrixmultiplikation zum Hauptengpass der Berechnung geworden. Moderne Computerhardware wird häufig durch Speicherbandbreite begrenzt, nicht durch Rechenleistung. Daher ist die verlustbehaftete Kompression von Matrizen zur Verringerung des Speichertransfers ein wichtiges Problem.
Für große Sprachmodelle schätzen die Autoren die erforderlichen Quantisierungsraten:
- In der Generierungsphase benötigt die CPU eine Quantisierungsrate von 1-3 Bits/Eintrag, um Rechenressourcen angemessen zu nutzen
- In der Vorausfüllungsphase benötigen kleine LLMs, die auf schnellen GPUs laufen, eine Quantisierungsrate von etwa 11,7 Bits/Eintrag
- Klassische Vektorquantisierung: Die unabhängige Quantisierung von Matrizen A und B gefolgt von der Berechnung des Produkts der quantisierten Matrizen führt zu O(n²)-Fehlern
- Skizzierungsmethoden: Obwohl sie unverzerrte Schätzungen bieten, bleibt die Varianz O(n²)
- Deterministische Quantisierer: Für Vektoren auf der Sphäre existiert eine untere Schranke von Ω(n²)
- Theoretische untere Schranken: Nicht-asymptotische untere Schranken für die Matrixmultiplikationsquantisierung bei Matrizen mit iid-Gaußschen Einträgen
- Universelle Quantisierer: Konstruktion universeller Quantisierer basierend auf verschachtelten Gittern mit expliziten Fehlergarantien für beliebige Matrizen
- Asymptotische Optimalität: Nachweis, dass die vorgeschlagenen Quantisierer für iid-Gaußsche Matrizen die untere Schranke erreichen und daher asymptotisch optimal sind
- Phasenübergangsphänomen: Entdeckung eines Phasenübergangs bei R ≈ 0,906 Bits/Eintrag, der die Notwendigkeit der Dimensionsreduktion bei niedrigen Bitraten offenbart
- Praktische Algorithmen: Bereitstellung von Implementierungen mit niedriger Komplexität, die nahezu optimale Leistung erreichen
Gegeben Matrizen A ∈ R^(n×a) und B ∈ R^(n×b), besteht das Ziel darin, Encoder f₁: R^(n×a) → 2^(naR) und f₂: R^(n×b) → 2^(nbR) sowie einen Decoder g zu entwerfen, sodass:
E∥A⊤B−g(f1(A),f2(B))∥F2
minimiert wird, wobei jeder Matrixeintrag mit R Bits beschrieben wird.
Das Papier definiert die kritische Rate-Distortion-Funktion:
Γ(R)={1−(1−(2⋅2−2R∗−2−4R∗))R∗R2⋅2−2R−2−4RR<R∗R≥R∗
wobei R* ≈ 0,906 die Lösung der Fixpunktgleichung R = ½log₂(1 + 4R ln 2) ist.
Für ρ-korrelierte Vektoren U, V auf der Sphäre wird Quantisierung mit verschachtelten Gittern Λc ⊂ Λf verwendet:
Kodierungsprozess:
- Unabhängige Dithering-Vektoren Z₁, Z₂ zu U und V addieren
- Quantisierung zum feinen Gitter Λf
- Ausgabe der Nebenklassendarstellung im groben Gitter Λc
Dekodierungsprozess:
- Quantisierungspunkt aus der Nebenklasse wiederherstellen
- Dithering entfernen
- Innenproduktsschätzung berechnen
Vorverarbeitungsschritte:
- Zentrierung: Berechnung von Ā = A - (1/n)1·1^⊤A, B̄ = B - (1/n)1·1^⊤B
- Normquantisierung: Hochpräzise Beschreibung des Mittelwerts und der Norm jeder Spalte
- Zufällige Rotation: Anwendung derselben orthogonalen Matrix S auf Ā und B̄
Quantisierungsschritte:
- Anwendung des Innenproduktsquantisierers auf jede rotierte Spalte
- Verwendung von Zeitanteil-Parametern κ und MMSE-Skalierungsparameter α
- Dithering-Technik: Macht Quantisierungsfehler unabhängig von der Eingabe und vermeidet O(n²)-Fehler deterministischer Quantisierer
- Verschachtelte Gitterstruktur: Bietet endliche Bitrate bei gleichzeitig guter Quantisierungsleistung
- Zeitanteil: Erreicht optimale Leistung bei niedrigen Bitraten durch Dimensionsreduktion
- Zufällige Rotation: Konvertiert beliebige Vektoren in gleichmäßig verteilte Vektoren auf der Sphäre zur vereinfachten Analyse
Datenerzeugung:
- Matrizen A, B ∈ R^(n×n) mit iid N(0,1)-Einträgen
- n = 3 × 2¹¹
Implementierungsdetails:
- Basis-Gitter: D₃-Gitter (3-dimensional)
- Verschachtelungsverhältnis: q = 6
- Lookup-Tabellengröße: < 64KB (passt in L1-Cache)
- Effektive Bitrate: ≈ 3,015 Bits/Symbol
- 3-Bit-Skalarquantisierer (ℓ∞-Normalisierung)
- Theoretische untere Schranke Γ(R)
Leistungsvergleich:
- Vorgeschlagene Methode: 1/n³ ∥Â⊤B - A⊤B∥²F ≈ 0,0593
- 3-Bit-Skalarquantisierung: ≈ 0,1668 (etwa 3-facher Unterschied)
- Theoretische untere Schranke: Γ(3,015) = 0,0304
Schlüsselergebnisse:
- Das auf D₃-Gitter basierende Schema ist deutlich besser als Skalarquantisierung
- Leistung liegt nahe am theoretischen Optimum (etwa 2-facher Unterschied)
- Der Leistungsunterschied wird mit wachsendem n weiter verringert
Kodierungskomplexität: O(n log n) (unter Verwendung der schnellen Hadamard-Transformation)
Dekodierungskomplexität: O(1) (unter Verwendung von Lookup-Tabellen)
Speicheraufwand: Jeder Quantisierer benötigt O(log n) zusätzliche Bits zur Beschreibung von Skalierungsfaktoren
- Monte-Carlo-Matrixmultiplikation (MCMM): Zufällige Stichprobennahme von Zeilen zur Approximation
- Lokalitätssensitives Hashing (LSH): Für niedrigdimensionale Skizzen von Kosinusähnlichkeit
- Einschränkungen: Relativer Fehler wächst mit ∥A∥²F∥B∥²F/∥A⊤B∥²F
- Post-Training-Quantisierung: OPTQ, GPTQ und andere Methoden
- Rotationstechniken: QuIP, QuaRot verwenden Hadamard-Transformationen
- Gitter-Quantisierung: QUIP# verwendet E₈-Gitter für Gewichtsquantisierung
- Verteilte Kompression: Kompression für lineare Funktionsberechnung
- Codebuch-Design: Voronoi-Codes und verschachtelte Gitter-Codes
- Optimalität: Für iid-Gaußsche Matrizen erreicht das vorgeschlagene Schema die informationstheoretische untere Schranke
- Universalität: Bietet explizite Leistungsgarantien für beliebige Matrizen
- Phasenübergangsphänomen: R* ≈ 0,906 Bits/Eintrag ist ein kritischer Schwellenwert
- Praktikabilität: Implementierung mit niedriger Komplexität erreicht nahezu theoretische Leistung
- Gemeinsame Zufälligkeit: Encoder und Decoder müssen Zufallssamen teilen
- Matrixbedingungen: Erfordert begrenzte Matrixeinträge (M = n^(10^22000))
- Hochdimensionale Gitter: Theoretisches Optimum erfordert hochdimensionale "gute" Gitter; in der Praxis werden Produkte niedrigdimensionaler Gitter verwendet
- Deterministische Schemata: Unklar, ob optimale deterministische Schemata ohne gemeinsame Zufälligkeit existieren
- Mehrfach-Matrixprodukte: Erweiterung auf Produkte von k>2 Matrizen
- Alternative Distanzmaße: KL-Divergenz und andere Nicht-MSE-Maße
- Deterministische Quantisierer: Erforschung von Schemata ohne gemeinsame Zufälligkeit
- Anwendungen in tiefen Netzen: Einsatz und Optimierung in praktischen LLMs
- Theoretische Strenge: Vollständige informationstheoretische Analyse mit oberen und unteren Schranken
- Praktischer Wert: Löst tatsächliche Engpässe bei der LLM-Inferenz
- Technische Innovation: Geschickte Kombination von Gitter-Quantisierung, zufälliger Rotation und Zeitanteil
- Universalität: Nicht abhängig von spezifischen Matrixverteilungsannahmen
- Komplexität: Theoretische Analyse ist relativ komplex; praktische Implementierung erfordert mehrere technische Komponenten
- Konstante Faktoren: Obwohl asymptotisch optimal, können Konstanten bei endlichen Stichproben groß sein
- Hardware-Anpassung: Implementierung muss für verschiedene Hardwareplattformen optimiert werden
- Skalierbarkeit: Erweiterung von 2 Matrizen auf mehrere Matrixprodukte ist nicht trivial
Theoretische Beiträge:
- Etablierung der informationstheoretischen Grundlagen der Matrixmultiplikationsquantisierung
- Offenlegung des Phasenübergangsphänomens und der Notwendigkeit der Dimensionsreduktion
- Bereitstellung wichtiger theoretischer Benchmarks für das Feld
Praktische Anwendungen:
- Neue theoretische Richtlinien für LLM-Quantisierung
- Nachfolgearbeit NestQuant hat SOTA-Leistung auf praktischen LLMs erreicht
- Theoretische Grundlagen für das Design von Hardware-Beschleunigern
- Inferenz großer Sprachmodelle: Gemeinsame Quantisierung von Gewichten und Aktivierungen
- Edge-Computing: Matrixoperationen in speicherbegrenzten Umgebungen
- Verteiltes Rechnen: Matrixmultiplikation mit Kommunikationsbeschränkungen
- Wissenschaftliches Rechnen: Großflächige numerische lineare Algebra-Probleme
Das Papier zitiert 44 verwandte Arbeiten, die mehrere Bereiche abdecken, darunter Informationstheorie, Gittertheorie, randomisierte lineare Algebra und Quantisierung neuronaler Netze. Besonders bemerkenswert sind:
- Zamirs Monographie zur Gitter-Kodierung bietet theoretische Grundlagen
- Bahnbrechende Arbeiten von Erez und Zamir zu verschachtelten Gittern
- Neuere LLM-Quantisierungsmethoden wie OPTQ, QuIP und andere
- Ergebnisse aus Zufallsmatrixtheorie und sphärischer Geometrie
Gesamtbewertung: Dies ist ein ausgezeichnetes Papier mit wichtigen Beiträgen sowohl in Theorie als auch in der Praxis, das solide informationstheoretische Grundlagen für das Matrixmultiplikationsquantisierungsproblem bietet und praktische Algorithmen mit nahezu optimaler Leistung demonstriert. Diese Arbeit ist von großer Bedeutung für das Verständnis und die Verbesserung von Quantisierungstechniken in großflächigen Maschinenlern-Systemen.