🤖 machine learning

BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization

Dieser Beitrag stellt BCJR-QAT vor, eine differenzierbare Relaxierung der trellis-codierten Gewichtsquantisierung, die den nicht-differenzierbaren Viterbi-Argmax durch einen temperaturgesteuerten BCJR-Summen-Produkt-Algorithmus ersetzt und damit ein durchgängiges quantisierungsbewusstes Training ermöglicht, das empirisch die State-of-the-Art-PTQ-Grenze bei LLMs übertrifft.

Ursprüngliche Autoren: Venugopalan Iyengar

Veröffentlicht 2026-05-12

📖 6 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Venugopalan Iyengar

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

Stellen Sie sich vor, Sie haben eine riesige Bibliothek von Büchern (ein Large Language Model), die Sie verkleinern möchten, damit sie auf einen Standard-Laptop oder ein Smartphone passt. Um dies zu tun, müssen Sie die „Wörter" (Gewichte) innerhalb der Bibliothek komprimieren.

Die Arbeit stellt eine neue Methode namens BCJR-QAT vor, um diese Bücher noch weiter zu komprimieren – bis auf nur noch 2 Bit pro Wort (was so ist, als würde man ein hochauflösendes Foto in einen winzigen, körnigen Thumbnail komprimieren).

Hier ist die Geschichte, wie sie es geschafft haben, unter Verwendung einfacher Analogien:

1. Das Problem: Die „Einbahnstraße"-Falle

Früher war der beste Weg, diese Modelle zu verkleinern, eine Methode namens QTIP. Denken Sie an QTIP als einen sehr klugen Bibliothekar, der eine Textseite betrachtet und sagt: „Okay, ich ersetze diesen komplexen Satz durch den nächsten einfachen Ausdruck aus meinem Wörterbuch."

Es gab jedoch eine Grenze. Sobald der Bibliothekar diese Wahl getroffen hatte, konnte er nicht zurückkehren und seine Meinung ändern, um das gesamte Buch besser klingen zu lassen. Es war wie eine „Einbahnstraße". Wenn der Bibliothekar früh einen leicht falschen Ausdruck wählte, litt der Rest des Buches darunter, und er konnte es nicht beheben, weil die Entscheidung „hart" und endgültig war.

Um dies zu beheben, verwenden Forscher normalerweise Training (QAT), bei dem Sie dem Bibliothekar das Üben und Anpassen seiner Entscheidungen ermöglichen. Aber hier liegt der Haken: Der Entscheidungsprozess des Bibliothekars beinhaltet ein komplexes Labyrinth (ein „Gitter"). Um den besten Weg durch das Labyrinth zu finden, verwenden sie eine Regel namens Viterbi, die wie ein „Wählen Sie die absolut beste Option"-Knopf funktioniert.

Das Problem: Sie können einem Computer nicht beibringen, durch einen „Wählen Sie die Beste"-Knopf zu lernen, weil dieser Knopf eine mathematische Sackgasse ist. Wenn Sie den Knopf drücken, weiß der Computer nicht, wie er seine Entscheidungen anpassen soll, um besser zu werden; er schnappt einfach zur nächsten Option. Es ist, als würde man versuchen, ein Auto zu steuern, indem man es nur erlaubt, sofort in die nächste Spur zu springen – man kann keine sanften Kurven fahren.

2. Die Lösung: Die „weiche" Entscheidung (BCJR)

Die Autoren, Venugopalan Iyengar, erfanden eine neue Möglichkeit, dem Bibliothekar das Üben zu ermöglichen. Anstatt sofort eine harte „Wählen Sie die Beste"-Entscheidung zu erzwingen, führten sie ein Konzept namens Temperatur ( $T$ ) ein.

Hohe Temperatur (Heiß): Stellen Sie sich vor, der Bibliothekar ist sehr entspannt und unentschlossen. Anstatt nur einen Ausdruck zu wählen, betrachtet er viele Ausdrücke gleichzeitig und weist jedem eine „Wahrscheinlichkeit" zu. Es ist wie an einem nebligen Tag, an dem Sie mehrere Pfade sehen können, nicht nur einen. Diese „weiche" Sichtweise ist mathematisch glatt, was bedeutet, dass der Computer leicht berechnen kann, wie er die Entscheidungen anpassen soll, um das Buch zu verbessern.
Niedrige Temperatur (Kalt): Wenn das Training abgeschlossen ist, wird der Bibliothekar „kühler" und entscheidungsfreudiger. Der Nebel lichtet sich, und er schnappt auf den einzigen besten Ausdruck zu (die harte Entscheidung).

Diese Methode, BCJR genannt, verwandelt die „Einbahnstraße" in eine glatte, befahrbare Straße. Der Computer kann nun genau lernen, wie er die Gewichte anpassen soll, um das gesamte Buch besser klingen zu lassen, nicht nur einzelne Sätze.

3. Der „Überziel"-Fehler

Die Autoren entdeckten eine überraschende Eigenart in der Funktionsweise dieser „Temperatur".

In der traditionellen Physik, wenn Sie versuchen, die beste Lösung zu finden, beginnen Sie normalerweise sehr heiß (sehr unentschlossen), um alle Möglichkeiten zu erkunden, und kühlen dann langsam ab.

Die Erkenntnis der Arbeit: Für diese spezifische Art der Komprimierung ist es katastrophal, „zu heiß" zu beginnen. Wenn der Bibliothekar am Anfang zu unentschlossen ist, verirrt er sich in eine schlechte Nachbarschaft (ein „schlechteres Voronoi-Becken") und bleibt dort stecken. Bis er abkühlt und versucht, den besten Weg zu finden, ist er bereits in einem schlechten Bereich verloren und kann nicht zurück.
Die Lösung: Sie fanden heraus, dass ein Start bei einer moderaten Temperatur (nicht zu heiß, nicht zu kalt) am besten funktioniert. Es ist wie ein Wanderweg mit einer klaren Karte zu beginnen, anstatt blind im Nebel herumzuwandern. Indem sie die Phase „super heiß" übersprangen, vermieden sie, sich zu verirren, und fanden eine bessere Lösung.

4. Die Ergebnisse: Ein Sieg für kleine Computer

Das Team testete dies an zwei verschiedenen Modelltypen:

Der „Proxy"-Test (OLMoE): Sie versuchten, die Komprimierung zu optimieren, indem sie nur betrachteten, wie gut jede Schicht des Modells ihre eigenen Daten rekonstruierte (als würde man prüfen, ob eine Fotokopie wie das Original aussieht). Ergebnis: Es scheiterte. Das Modell wurde tatsächlich schlechter als die alte Methode. Dies lehrte sie, dass nur eine „bessere Fotokopie" zu machen nicht bedeutet, dass das Buch besser liest.
Der „echte" Test (Llama-3.2): Sie optimierten das Modell, um tatsächlich Text besser zu lesen und zu verstehen (unter Verwendung einer Methode namens „Distillation", bei der ein intelligenter Lehrer-Modell den Schüler-Modell führt). Ergebnis: Es funktionierte!
- Auf einer bestimmten Schicht des Modells schlug ihre neue Methode die alte beste Methode mit einer kleinen, aber signifikanten Marge.
- Als sie dies auf mehrere Schichten anwendeten, addierten sich die Verbesserungen auf eine „superadditive" Weise (das Ganze wurde größer als die Summe seiner Teile).

5. Der Motor unter der Haube

Das Durchführen dieser Mathematik ist normalerweise sehr langsam und erfordert teure Supercomputer. Die Autoren bauten auch einen spezialisierten „Motor" (ein Triton-Kernel), der diese komplexe Mathematik auf einer einzigen Consumer-Grafikkarte (wie einer RTX 4080) unglaublich schnell ausführt. Sie machten es 6,5-mal schneller als Standardmethoden und bewiesen, dass diese fortschrittliche Technik kein Rechenzentrum zum Ausführen benötigt.

Zusammenfassung

Die Arbeit stellt eine neue Möglichkeit vor, KI-Modelle zu verkleinern, damit sie auf Consumer-Geräten passen.

Der Trick: Sie ersetzten einen starren, nicht lernfähigen Entscheidungsprozess durch einen „weichen", lernfähigen, der sich allmählich verhärtet.
Die Erkenntnis: Man sollte den Lernprozess nicht zu „heiß" (zu zufällig) beginnen; ein moderater Start verhindert, dass das Modell sich verirrt.
Das Ergebnis: Wenn es korrekt trainiert wird, um Text zu verstehen (und nicht nur Daten zu kopieren), liefert diese Methode bessere Ergebnisse als frühere State-of-the-Art-Techniken, und sie läuft schnell genug auf einem normalen Gaming-PC.

Die Autoren haben ihren Code und die trainierten Modelle veröffentlicht, damit andere es ausprobieren können.

Technische Zusammenfassung: BCJR-QAT

Problemstellung

Das Erreichen von 2 Bits pro Gewicht (bpw) oder weniger für Large Language Models (LLMs) auf Consumer-Hardware stellt einen kritischen Engpass dar. Während Post-Training-Quantisierungsmethoden (PTQ) wie QTIP eine starke Grenze für die 2-Bit-Quantisierung unter Verwendung von Trellis-Coded Quantization (TCQ) etabliert haben, sind sie durch die „PTQ-Obergrenze" begrenzt. Ein Vorstoß darüber hinaus erfordert Quantization-Aware Training (QAT). Standard-QAT-Techniken versagen jedoch bei Trellis-Codes, da der Viterbi-Decodierungsalgorithmus eine nicht-differenzierbare argmax-Operation über exponentiell viele Pfade beinhaltet. Die Anwendung eines Straight-Through Estimators (STE) auf das Trellis zerstört die Kopplung zwischen benachbarten Gewichtsentscheidungen, was den Hauptvorteil der Trellis-Codierung gegenüber der skalaren Quantisierung darstellt.

Methodik: BCJR-QAT

Die Arbeit stellt BCJR-QAT vor, eine differenzierbare Relaxation der trellis-codierten Quantisierung, die die harte Viterbi-argmax-Operation durch eine weiche Erwartung ersetzt, die vom BCJR-Algorithmus (Bahl-Cocke-Jelinek-Raviv) Forward-Backward abgeleitet ist.

Kernmechanismus

Analogie zur Statistischen Mechanik: Das Trellis wird als Faktorgraph behandelt, der einer 1D-Ising-ähnlichen Spin-Kette entspricht. Die Viterbi-argmax-Operation wird als Grenzfall der Temperatur Null ( $T \to 0$ ) einer Boltzmann-Verteilung über Pfade betrachtet.
Weiches Codewort: Bei einer endlichen Temperatur $T > 0$ berechnet die Methode die exakten marginalen Posterior-Wahrscheinlichkeiten jedes Zustands unter Verwendung des BCJR-Algorithmus. Das quantisierte Ausgabe ist die Boltzmann-Erwartung des emittierten Codeworts:
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
Dieses weiche Codewort ist exakt differenzierbar bezüglich des Gewichtsvor-Bildes und der Trellis-Parameter.
Ausglühen (Annealing): Während des Trainings wird die Temperatur $T$ von einem Anfangswert $T_0$ auf ein nahe Null liegendes $T_{end}$ ausgeglüht. Wenn $T \to 0$ geht, konvergiert das weiche Codewort zum harten Viterbi-Ausgang und „kristallisiert" die kontinuierliche Optimierung effektiv in eine diskrete Lösung.

Wichtige technische Innovationen

Differenzierbarkeit: Die BCJR-Forward-Backward-Rekursion (Komplexität $O(L \cdot S^2)$ ) liefert exakte Gradienten durch das Trellis und erhält die zwischenzuständige Kopplung, die bei skalaren STE-Methoden verloren geht.
Triton-Kernel: Die Autoren implementierten einen fusionierten Triton-Kernel, der Gather-, Log-Sum-Exp- und Additionsoperationen in einem einzigen Startvorgang kombiniert. Dies reduziert den Forward-/Backward-Durchlauf von hunderten Kernel-Starts auf einen und erzielt eine 6,57-fache Beschleunigung gegenüber einer Referenz-Implementierung mit PyTorch Autograd bei Beibehaltung der FP32-numerischen Parität.
Zeitplan-Korrektur (Skip-High-T): Die Arbeit diagnostiziert einen kritischen Fehlermodus in Standard-Simulated-Annealing-Zeitplänen (Start bei $T_0=1.0$ ). Bei hohen Temperaturen ist der Gradient des weichen Codeworts „verschmiert" und uninformierend, was den Optimierer dazu veranlasst, vor der Abkühlphase in ein schlechteres Voronoi-Becken abzugleiten. Der vorgeschlagene Fix besteht darin, die Phase mit hoher $T$ zu überspringen und bei einer niedrigeren Temperatur zu initialisieren (z. B. $T_0=0,3$ ), wo die Gradienten bereits informativ sind.

Experimentelle Ergebnisse

Die Arbeit bewertet BCJR-QAT an zwei Architekturen: OLMoE-1B-7B (MoE) und Llama-3.2-1B (Dense).

1. OLMoE-1B-7B (Negatives Ergebnis: Proxy-Lücke)

Setup: Greedy-QAT pro Schicht mit Optimierung für Rekonstruktions-MSE (ein Proxy-Ziel).
Ergebnis: BCJR-QAT reduzierte den Rekonstruktions-MSE pro Schicht im Vergleich zu QTIP-PTQ erfolgreich um 17–20 %. Diese Verbesserung übersetzte sich jedoch nicht in die End-Aufgabenleistung. Die Modelle verschlechterten sich um $\sim 1,3$ Perplexity (PPL) auf WikiText-2 im Vergleich zur QTIP-PTQ-Baseline.
Fazit: Dies bestätigt eine „Proxy-Lücke", bei der die Minimierung des lokalen Rekonstruktionsfehlers keine bessere End-Aufgabenleistung bei 2 bpw garantiert.

2. Llama-3.2-1B (Positives Ergebnis: End-to-End KL)

Setup: Ein-Schicht-QAT mit Optimierung für End-to-End Forward-KL-Distillation gegen einen FP16-Lehrer.
Ergebnis:
- Zeitplan-Überschießen: Die Verwendung eines naiven Zeitplans mit hoher $T$ ( $T_0=1,0$ ) führte zu einer Leistung, die leicht schlechter war als PTQ (+0,005 bis +0,022 PPL).
- Skip-High-T-Fix: Die Verwendung des korrigierten Zeitplans ( $T_0=0,3$ ) eliminierte das Überschießen. Auf Schicht 4 erreichte das Modell 10,1347 PPL, eine Verbesserung von −0,084 PPL gegenüber der QTIP-PTQ-Baseline (10,2189).
- Super-Additivität: Bei der Kombination trainierter Schichten (Schicht 4 und Schicht 8) überstieg der gemeinsame Gewinn (−0,077 PPL) die Summe der einzelnen Gewinne (−0,062 PPL) und demonstrierte einen „Kooperationsüberschuss", bei dem BCJR-trainierte Codewörter sich gegenseitig verstärken.

3. Theoretische Analyse

Drift-Budget-Theorie: Die Autoren leiten eine Bedingung für das Entkommen aus dem PTQ-Voronoi-Becken ab: Die kumulative Gewichtsdrift ( $\eta \cdot N \cdot g_{max}$ ) muss den Radius des Beckens ( $r_{Voronoi} \approx 10^{-3}$ ) überschreiten. Experimente bestätigen, dass Zeitpläne, die diese Bedingung nicht erfüllen, nicht aus der PTQ-Lösung entkommen können, während diejenigen, die sie überschreiten, bessere Lösungen finden können, sofern das Ziel und der Zeitplan korrekt sind.

Bedeutung und Behauptungen

Die Arbeit beansprucht drei Hauptbeiträge:

Algorithmisch: BCJR-QAT ist die erste Methode, die trellis-codierte Quantisierung nativ differenzierbar macht und damit QAT auf TCQ ermöglicht, ohne die Vorteile der Pfadkopplung zu verlieren.
Technisch: Der fusionierte Triton-Kernel macht dies auf Consumer-Hardware (eine einzelne 12GB-GPU) rechnerisch machbar und entfernt die Barrieren für den Einsatz in der Inferenzphase für die Forschung.
Empirische Einsicht:
- Es wird identifiziert, dass Rekonstruktions-MSE pro Schicht ein schlechter Proxy für 2-Bit-QAT ist, wohingegen End-to-End-KL-Distillation effektiv ist.
- Es wird die konventionelle Weisheit des Simulated Annealing für diskrete Relaxationen in Frage gestellt, indem gezeigt wird, dass ein „heißer" Start ( $T=1,0$ ) für BCJR-QAT nachteilig ist, da der Gradient bei hohen Temperaturen von Rauschen dominiert wird. Die optimale Strategie besteht darin, bei einer Temperatur zu starten, bei der der Gradient bereits informativ ist.

Die Arbeit schließt, dass, obwohl die Ergebnisse auf einzelner Schicht vielversprechend sind, das volle Potenzial von BCJR-QAT wahrscheinlich im gemeinsamen End-to-End-Training über alle Schichten liegt, eine Aufgabe, die durch den veröffentlichten Kernel erschwinglich gemacht wird, aber Cloud-Scale-Hardware (H100/H200) für Modelle mit mehr als 1 Milliarde Parametern erfordert. Die Autoren stellen ausdrücklich fest, dass sie dieses vollständige gemeinsame Training im Rahmen dieser Arbeit nicht durchgeführt haben.