🤖 machine learning

JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates

JacQuant stellt ein neuartiges quantisierungsbewusstes Trainingsframework vor, das den spröden Straight-Through-Schätzer durch einen leichten, erlernten Jacobian-Surrogat ersetzt, um das Training zu stabilisieren und bei der Quantisierung von LLMs mit extrem niedriger Bitbreite eine höhere Genauigkeit zu erzielen, ohne die Vorwärtsquantisierer zu modifizieren oder erhebliche Laufzeitkosten zu verursachen.

Ursprüngliche Autoren: Kai Yi, Vignesh Vivekraja, Harshit Khaitan, Steven Li

Veröffentlicht 2026-05-26

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Kai Yi, Vignesh Vivekraja, Harshit Khaitan, Steven Li

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Das große Problem: Der "kaputte Kompass" beim KI-Training

Stellen Sie sich vor, Sie versuchen, einem Roboter beizubringen, ein Labyrinth zu navigieren. Um Platz und Energie zu sparen, entscheiden Sie sich, dem Roboter eine sehr niedrig aufgelöste Karte zu geben (dies wird als Quantisierung bezeichnet). Anstatt die genaue Entfernung zu einer Wand zu kennen, weiß der Roboter nur, ob er sich in "Zone A", "Zone B" oder "Zone C" befindet.

Das Problem entsteht, wenn sich der Roboter genau auf der Linie zwischen Zone A und Zone B befindet.

Die Realität: Wenn der Roboter sich ein winziges Stück bewegt, bleibt er möglicherweise in Zone A. Bewegt er sich ein winziges Stück in die andere Richtung, springt er möglicherweise in Zone B. Die Karte ist "blockig" und ändert sich nicht sanft.
Die alte Methode (STE): Seit Jahren nutzte der Trainingsalgorithmus, wenn der Roboter versuchte, aus seinen Fehlern zu lernen, eine Abkürzung namens Straight-Through Estimator (STE). Er tat so, als wäre die Karte glatt und kontinuierlich. Er sagte dem Roboter: "Sie haben sich 1 Zoll bewegt, also hat sich die Karte um 1 Zoll geändert."
Das Ergebnis: Dies ist so, als würde man dem Roboter einen kaputten Kompass geben. Wenn sich der Roboter in der Nähe einer Grenze befindet, dreht sich der Kompass wild oder zeigt in die falsche Richtung, weil die Karte tatsächlich nicht glatt ist. Dies macht das Training instabil, insbesondere wenn die Karte sehr niedrig aufgelöst ist (wie bei 1 oder 2 Bits).

Die Lösung: JacQuant (Der "intelligente Sensor")

Die Autoren dieses Papiers, die bei Meta AI arbeiten, haben JacQuant eingeführt. Anstatt so zu tun, als wäre die Karte glatt, lehren sie den Roboter einen leichtgewichtigen Sensor, der genau versteht, wie die blockige Karte auf Bewegungen reagiert.

So funktioniert es, Schritt für Schritt:

1. Der "Surrogat"-Sensor

Stellen Sie sich vor, der Roboter hat einen kleinen, billigen Sensor an seinem Fuß befestigt. Dieser Sensor verändert die Karte nicht; die Karte bleibt blockig. Aber dieser Sensor misst: "Wenn ich meinen Fuß ein winziges Stück wackele, ändert sich die Zone tatsächlich?"

Befindet sich der Roboter in der Mitte einer Zone, sagt der Sensor: "Ja, eine kleine Bewegung verändert die Dinge." (Empfindlichkeit = 1).
Ist der Roboter gegen eine Wand gedrückt (Sättigung) oder genau an einer Grenze, wo eine Bewegung die Zone nicht ändert, sagt der Sensor: "Nein, eine kleine Bewegung bringt nichts." (Empfindlichkeit = 0).

Dieser Sensor wird als Learned Jacobian Surrogate (gelernter Jacobischen-Surrogat) bezeichnet. Es ist eine einfache mathematische Karte, die dem Trainingsalgorithmus mitteilt, wie "empfindlich" der aktuelle Ort ist.

2. Den Kompass reparieren

Wenn der Roboter einen Fehler macht, betrachtet der Trainingsalgorithmus die Messung des Sensors, bevor er ein Korrektursignal sendet.

Alter Weg (STE): "Sie haben einen Fehler gemacht! Drücken Sie stark in diese Richtung!" (Selbst wenn der Roboter gegen eine Wand gedrückt ist, drückt er stark, was dazu führt, dass er nutzlos hin und her springt).
JacQuant-Weg: Der Sensor sagt: "Hey, Sie sind gegen eine Wand gedrückt; Bewegung wird nicht helfen." Daher dämpft der Algorithmus den Druck. Er sagt: "Okay, drücken Sie hier nicht so stark; versuchen wir eine andere Richtung."

Dies verhindert, dass der Roboter in der Nähe der Grenzen wild hin und her springt, und hilft ihm, viel schneller den besten Pfad zu finden.

3. Die "amortisierten" Kosten (Warum es günstig ist)

Sie denken vielleicht: "Das Messen dieser Empfindlichkeit klingt teuer!"
Die Autoren haben dies gelöst, indem sie (auf eine clevere Weise) faul waren. Sie überprüfen den Sensor nicht jede einzelne Sekunde.

Sie überprüfen den Sensor gelegentlich (wie einmal alle 100 Schritte).
Sie verwenden diese Messung für eine lange Zeit, bis sich der Roboter weit genug bewegt hat, um eine neue Überprüfung zu benötigen.
Dies wird als Amortisation bezeichnet. Es ist wie das Überprüfen der Wettervorhersage einmal am Morgen und die Verwendung dieser Informationen den ganzen Tag über, anstatt alle 5 Minuten nach draußen zu treten. Die Kosten sind so gering (weniger als 2 % zusätzliche Zeit), dass es sich wie kostenlos anfühlt.

Die Ergebnisse: Glatter, schneller, besser

Das Papier testete dies an Large Language Models (LLMs) wie LLaMA und Qwen, speziell beim Zusammenpressen auf ultra-niedrige Präzision (1 oder 2 Bits).

Die Analogie: Stellen Sie sich vor, Sie versuchen, ein riesiges, hochauflösendes Gemälde in einen winzigen, pixeligen Rahmen zu passen. Die alte Methode (STE) ließ das Bild verschwommen und wackelig aussehen. JacQuant wirkt wie ein intelligenter Filter, der genau weiß, welche Pixel scharf zu halten sind und welche geglättet werden müssen, was zu einem viel klareren Bild führt.
Die Beweise: Bei ihren Tests zeigten Modelle, die mit JacQuant trainiert wurden:
- Sie lernten schneller (konvergierten schneller).
- Sie machten weniger Fehler (niedrigere "Perplexität", ein Maß dafür, wie verwirrt die KI ist).
- Sie beantworteten Fragen besser (höhere Genauigkeit bei Denkaufgaben).
- Sie blieben stabil (stürzten nicht ab oder oszillierten nicht in der Nähe der "Wände" der Quantisierungszonen).

Zusammenfassung

JacQuant ist eine neue Art, KI-Modelle zu trainieren, die so stark komprimiert wurden, dass sie sehr klein sind. Anstatt blind zu raten, wie das komprimierte Modell lernt (unter Verwendung des alten "Straight-Through Estimator"), lernt es eine einfache, günstige "Empfindlichkeitskarte". Diese Karte sagt der KI genau, wie sehr sie ihren eigenen Bewegungen in der Nähe der Ränder ihrer komprimierten Welt vertrauen kann. Das Ergebnis ist ein stabilerer, genauerer und effizienterer Trainingsprozess für die kleinsten, speicherschonendsten KI-Modelle.

Technische Zusammenfassung: JacQuant

Problemstellung

Quantization-aware Training (QAT) ist ein Standardansatz für den Einsatz von Large Language Models (LLMs) unter strengen Speicher- und Latenzbedingungen, insbesondere in Ultra-Low-Bit-Bereichen (z. B. Gewichte $\le 2$ Bit). Moderne QAT-Pipelines verlassen sich jedoch fast ausschließlich auf den Straight-Through Estimator (STE). Der STE approximiert den nicht differenzierbaren Quantisierer, indem er Gradienten so durchleitet, als wäre der Quantisierer die Identitätsfunktion.

Die Arbeit identifiziert einen kritischen Fehler in diesem Ansatz: Harte Quantisierer sind stückweise konstant mit Bin-Grenzen sowie Sättigungs- bzw. Beschneidungsbereichen. In Ultra-Low-Bit-Szenarien liegen viele Parameter nahe diesen Grenzen oder erreichen die Beschneidungsgrenzen. Die Annahme des STE, eine Identitäts-Jacobi-Matrix zu besitzen, erzeugt eine persistente Diskrepanz, wodurch Gradienten in gesättigten Bereichen übermäßig propagiert werden, in denen der Quantisierer lokal unempfindlich ist. Dies führt zu einem brüchigen Training, Instabilität in der Nähe von Bin-Grenzen und Lösungen, die nur schwach mit dem tatsächlichen Verhalten des niedrigpräzisen Modells übereinstimmen.

Methodik: JacQuant

JacQuant schlägt ein Framework vor, das die feste Identitäts-Jacobi-Matrix des STE durch eine gelernte, leichte Ersatz-Jacobi-Matrix ersetzt, bezeichnet als $B(W)$ . Die Kernphilosophie besteht darin, QAT als ein „fehlende-Jacobi-Matrix"-Problem zu behandeln: Während der Vorwärtsdurchlauf den eingesetzten harten Quantisierer $Q_\Delta(\cdot)$ beibehält, lernt der Rückwärtsdurchlauf, wie Störungen in Vollpräzisionsgewichten ihre quantisierten Proxy-Werte beeinflussen.

Kernmechanismus

Anstatt des Standard-STE-Updates $\nabla_W \ell \approx \nabla_{\hat{W}} \ell$ verwendet JacQuant:
$\nabla_W \ell \approx B(W) \nabla_{\hat{W}} \ell$
wobei $\hat{W} = Q_\Delta(W)$ und $B(W)$ eine gelernte Sensitivitätskarte ist.

Wichtige Designentscheidungen

Parametrisierung: $B(W)$ ist auf diagonal oder blockdiagonal beschränkt, mit einem skalaren Gain pro Quantisierungsgruppe (z. B. 128 Gewichte). Dies hält den Speicheraufwand vernachlässigbar (nur ca. 1,5 % zusätzlich zu den Adam-Optimierer-Zuständen).
Schätzstrategien: Die Arbeit stellt zwei praktische, datengesteuerte Schätzer zur Aktualisierung von $B(W)$ $B (W)$ vor:
- JacQuant-Probe: Verwendet Gaußsche Zufallssonden $\delta_g$ , um die lokale Steigung $\Delta q_g = Q_\Delta(W_g + \delta_g) - Q_\Delta(W_g)$ zu messen. Eine Kleinste-Quadrate-Anpassung schätzt die lokale Sensitivität.
- JacQuant-Dither: Verwendet subtraktives Dithering (Hinzufügen von gleichverteiltem Rauschen), um einen differenzierbaren Proxy für den Quantisierer zu erstellen, und schätzt die Sensitivität basierend auf der erwarteten Reaktion.
Amortisierte Aktualisierungen: Um den Trainingsdurchsatz zu erhalten, wird $B(W)$ nicht bei jedem Schritt aktualisiert. Es wird gelegentlich (z. B. alle 50–200 Schritte) unter Verwendung eines Exponential Moving Average (EMA) aktualisiert. Zwischen den Aktualisierungen wird das aktuelle $B(W)$ wiederverwendet.
Varianzreduktion: Das Framework integriert sich mit Optimierern zur Varianzreduktion (z. B. loopless-SVRG, SAGA), um die durch den Sondierungs-/Dithering-Prozess eingeführte Stochastik zu stabilisieren.

Theoretische Beiträge

Die Arbeit liefert eine Konvergenzanalyse für JacQuant unter nicht-konvexen Zielfunktionen, strukturiert um vier Hauptargumente:

Definition des Zielgradienten: Die Autoren definieren einen „Mean-Field"-Zielgradienten $g^\dagger(W) = J(W)\bar{v}$ , wobei $J(W)$ die wahre Sensitivität des dither-gemittelten Quantisierers ist. Sie beweisen, dass die Verzerrung (Bias) von JacQuant proportional zu $\|B(W) - J(W)\|$ ist, während die Verzerrung des STE proportional zu $\|I - J(W)\|$ ist.
Konsistenz der Schätzer: Sie zeigen, dass sowohl der Probe- als auch der Dither-Schätzer konsistent mit der Ziel-Sensitivität $J(W)$ sind. Konkret ist die Probe-Aktualisierung eine lokale Regression, die gegen die Mean-Field-Sensitivität konvergiert, und die Dither-Aktualisierung hat $J(W)$ als ihren eindeutigen Populations-Fixpunkt.
Stabilität: Die gelernte $B(W)$ wirkt als adaptiver Gain-Regler. Sie bleibt in reaktionsfähigen Bin-Innenbereichen nahe der Identität, dämpft aber Gradienten (Gain $< 1$ ) in gesättigten Bereichen, in denen der Quantisierer unempfindlich ist. Dies verhindert die bei STE-basiertem Training üblichen Oszillationen.
Konvergenzraten: Unter einer Polyak-Łojasiewicz (PŁ)-Bedingung und unter der Annahme der „Code-Erhaltungs"-Eigenschaft (Quantisierungscodes ändern sich innerhalb eines lokalen Fensters nicht häufig) beweist die Arbeit, dass JacQuant lineare Konvergenzraten erreicht. Der Konvergenzfehler ist durch den Jacobi-Tracking-Fehler und den stochastischen Varianzboden begrenzt.

Experimentelle Ergebnisse

Die Autoren bewerten JacQuant an mehreren LLM-Familien (LLaMA3-1B/3B/8B, Qwen3-1.7B) in Ultra-Low-Bit-Bereichen ( $W \le 2$ , $A8/A16$ ).

Leistung: JacQuant übertrifft konsistent STE-basierte Baselines (einschließlich ParetoQ und WinQ) sowie den „Rotation Trick" (eine Methode mit fester geometrischer Transformation).
- Bei LLaMA3-1B (W2A16) verbesserte JacQuant die Perplexität von 13,5 auf 12,3 und die Genauigkeit von 55,7 % auf 56,6 % gegenüber der besten STE-basierten Baseline.
- Bei LLaMA3-8B (W2A16) zeigte JacQuant über 30K, 60K und 90K Trainingsschritte hinweg monotone Verbesserungen gegenüber ParetoQ.
Effizienz: Der Laufzeit-Overhead ist vernachlässigbar. Bei LLaMA3-1B liegt der Overhead je nach Gruppengröße zwischen 1–5 %; bei LLaMA3-8B bleibt er für praktische Gruppengrößen unter 1,6 %.
Mechanismus-Validierung: Finite-Differenzen-Diagnostiken bestätigen, dass JacQuant in beschneidenen/gesättigten Bereichen Gruppengains $< 1$ lernt und damit Gradienten effektiv dämpft, wo der STE sie übermäßig propagieren würde.

Bedeutung und Behauptungen

Die Arbeit positioniert JacQuant als Drop-in-Upgrade für bestehende QAT-Pipelines. Ihre Bedeutung liegt in:

STE-freies Training: Es beseitigt die Abhängigkeit von der heuristischen Identitäts-Jacobi-Matrix, ohne Änderungen am Vorwärts-Quantisierer oder der Modellarchitektur zu erfordern.
Geometrie-Bewusstsein: Durch das Erlernen der lokalen Sensitivität passt es sich natürlich an die Geometrie des Quantisierungslandschafts an und adressiert spezifisch die Sättigungs- und Beschneidungsprobleme, die das Ultra-Low-Bit-Training plagen.
Praktikabilität: Es erzielt diese Verbesserungen mit minimalen Rechenkosten, was es für das Training großer LLMs geeignet macht.

Die Autoren begrenzen ihre empirischen Behauptungen explizit auf Ultra-Low-Bit-LLM-QAT ( $\le 2$ Bit). Sie räumen ein, dass die Methode zwar ein Drop-in-Ersatz für die Rückwärtsregel ist, der Vorwärts-Quantisierer jedoch ein harter, eingesetzter Operator bleibt. Die Arbeit behauptet nicht, alle Quantisierungs-Herausforderungen zu lösen, sondern zielt spezifisch auf die bei STE-basierter Optimierung inhärente Gradienten-Diskrepanz ab. Erwähnte zukünftige Arbeitsrichtungen umfassen die Erweiterung der Methode auf Mixed-Precision-Richtlinien, nicht-uniforme Quantisierer und die Verfeinerung des Schätzungs-Schedulings.