🤖 machine learning

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

Dit artikel introduceert BCJR-QAT, een differentieerbare relaxatie van trellis-codering voor gewichtsquantisatie die de niet-differentieerbare Viterbi argmax vervangt door een temperatuur-gereguleerd BCJR som-product-algoritme, waardoor end-to-end quantisatie-bewuste training mogelijk wordt die empirisch de state-of-the-art PTQ-grens op LLMs overtreft.

Oorspronkelijke auteurs: Venugopalan Iyengar

Gepubliceerd 2026-05-12

📖 6 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Venugopalan Iyengar

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een enorme bibliotheek met boeken hebt (een Groot Taalmodel) die je wilt verkleinen zodat hij op een standaard laptop of telefoon past. Om dit te doen, moet je de "woorden" (gewichten) binnenin de bibliotheek comprimeren.

Het artikel introduceert een nieuwe methode genaamd BCJR-QAT om deze boeken nog verder te comprimeren—tot slechts 2 bits per woord (wat vergelijkbaar is met het comprimeren van een foto in hoge resolutie tot een klein, korrelig miniatuurbeeld).

Hier is het verhaal van hoe ze dit deden, met gebruikmaking van eenvoudige analogieën:

1. Het Probleem: De "Eenrichtingsstraat"-Valstrik

Voorheen was de beste manier om deze modellen te verkleinen een methode genaamd QTIP. Denk aan QTIP als een zeer slimme bibliothecaris die een pagina tekst bekijkt en zegt: "Oké, ik vervang deze complexe zin door de dichtstbijzijnde simpele zin uit mijn woordenboek."

Er was echter een limiet. Zodra de bibliothecaris die keuze had gemaakt, kon hij niet terug en zijn mening wijzigen om het hele boek beter te laten klinken. Het was als een "eenrichtingsstraat". Als de bibliothecaris vroeg in het proces een iets verkeerde zin koos, leed de rest van het boek hieronder, en hij kon het niet repareren omdat de beslissing "hard" en definitief was.

Om dit op te lossen, gebruiken onderzoekers meestal Training (QAT), waarbij je de bibliothecaris laat oefenen en zijn keuzes laat aanpassen. Maar hier zit de adder onder het gras: het beslissingsproces van de bibliothecaris omvat een complex doolhof (een "trellis"). Om het beste pad door het doolhof te vinden, gebruiken ze een regel genaamd Viterbi, wat vergelijkbaar is met een knop "kies de absoluut beste optie".

Het Probleem: Je kunt een computer niet leren door een "kies de beste" knop te gebruiken, omdat die knop een wiskundige doodlopende weg is. Als je op de knop drukt, weet de computer niet hoe hij zijn keuzes moet bijsturen om beter te worden; hij schiet gewoon naar de dichtstbijzijnde optie. Het is alsof je een auto probeert te sturen door alleen toe te staan dat hij direct naar de volgende rijbaan springt—je kunt geen soepele bochten maken.

2. De Oplossing: De "Zachte" Beslissing (BCJR)

De auteurs, Venugopalan Iyengar, bedachten een nieuwe manier om de bibliothecaris te laten oefenen. In plaats van direct een harde "kies de beste" beslissing af te dwingen, introduceerden ze een concept genaamd Temperatuur ( $T$ ).

Hoge Temperatuur (Heet): Stel je voor dat de bibliothecaris zeer ontspannen en onbeslist is. In plaats van slechts één zin te kiezen, overweegt hij vele zinnen tegelijk, waarbij hij een "waarschijnlijkheid" toekent aan elke zin. Het is als een mistige dag waarop je meerdere paden kunt zien, niet slechts één. Dit "zachte" beeld is wiskundig glad, wat betekent dat de computer gemakkelijk kan berekenen hoe hij de keuzes moet bijsturen om het boek te verbeteren.
Lage Temperatuur (Koud): Naarmate de training vordert, wordt de bibliothecaris "kouder" en beslissender. De mist klart op en hij schiet naar de enige beste zin toe (de harde beslissing).

Deze methode, genaamd BCJR, verandert de "eenrichtingsstraat" in een gladde, berijdbare weg. De computer kan nu precies leren hoe hij de gewichten moet aanpassen om het hele boek beter te laten klinken, en niet alleen individuele zinnen.

3. De "Te Ver Schieten"-Fout

De auteurs ontdekten een verrassende eigenaardigheid in hoe deze "temperatuur" werkt.

In de traditionele fysica, wanneer je probeert de beste oplossing te vinden, begin je meestal zeer heet (zeer onbeslist) om alle mogelijkheden te verkennen, en koel je vervolgens langzaam af.

De Bevinding van het Artikel: Voor dit specifieke type compressie is te "heet" beginnen een ramp. Als de bibliothecaris aan het begin te onbeslist is, dwaalt hij weg naar een slechte buurt (een "slechter Voronoi-bekken") en blijft daar hangen. Tegen de tijd dat hij afkoelt en probeert het beste pad te vinden, is hij al verdwaald in een slecht gebied en kan hij niet terug.
De Oplossing: Ze ontdekten dat beginnen met een matige temperatuur (niet te heet, niet te koud) het beste werkt. Het is alsof je een wandeling begint met een duidelijke kaart in plaats van blind in de mist te dwalen. Door de "super hete" fase over te slaan, voorkwamen ze dat ze verdwaalden en vonden ze een betere oplossing.

4. De Resultaten: Een Overwinning voor Kleine Computers

Het team testte dit op twee verschillende soorten modellen:

De "Proxy"-Test (OLMoE): Ze probeerden de compressie te optimaliseren door alleen te kijken hoe goed elke laag van het model zijn eigen data reconstrueerde (alsof je controleert of een fotokopie eruitziet als het origineel). Resultaat: Het faalde. Het model werd eigenlijk slechter dan de oude methode. Dit leerde hen dat het maken van een "betere fotokopie" niet betekent dat het boek beter leest.
De "Echte" Test (Llama-3.2): Ze optimaliseerden het model om daadwerkelijk tekst beter te lezen en te begrijpen (met behulp van een methode genaamd "distillatie", waarbij een slim leermeestermodel de student begeleidt). Resultaat: Het werkte!
- Op een specifieke laag van het model versloeg hun nieuwe methode de oude beste methode met een klein maar significant verschil.
- Toen ze dit toepasten op meerdere lagen, stapelden de verbeteringen zich op in een "super-additieve" manier (het geheel werd groter dan de som der delen).

5. De Motor Onder de Motorkap

Het doen van deze wiskunde is meestal zeer traag en vereist dure supercomputers. De auteurs bouwden ook een gespecialiseerde "motor" (een Triton-kernel) die deze complexe wiskunde ongelooflijk snel uitvoert op een enkele consumentengrafische kaart (zoals een RTX 4080). Ze maakten het 6,5 keer sneller dan standaardmethoden, wat bewijst dat deze geavanceerde techniek geen datacenter nodig heeft om te draaien.

Samenvatting

Het artikel presenteert een nieuwe manier om AI-modellen te verkleinen zodat ze op consumentenapparaten passen.

De Truc: Ze vervingen een star, onleerbaar beslissingsproces door een "zacht", leerbaar proces dat geleidelijk hard wordt.
Het Inzicht: Je moet het leerproces niet te "heet" beginnen (te willekeurig); een matig begin voorkomt dat het model verdwaalt.
Het Resultaat: Wanneer het correct wordt getraind om tekst te begrijpen (en niet alleen data te kopiëren), levert deze methode betere resultaten op dan eerdere state-of-the-art technieken, en draait het snel genoeg op een reguliere gaming-pc.

De auteurs hebben hun code en de getrainde modellen vrijgegeven zodat anderen het kunnen uitproberen.

Technische Samenvatting: BCJR-QAT

Probleemstelling

Het bereiken van 2 bits per gewicht (bpw) of lager voor Large Language Models (LLM's) op consumentenhardware is een kritieke knelpunt. Hoewel Post-Training Quantization (PTQ)-methoden zoals QTIP een sterke grens hebben vastgesteld voor 2-bit quantisatie met behulp van Trellis-Coded Quantization (TCQ), worden ze beperkt door het "PTQ-plafond". Het overstijgen hiervan vereist Quantization-Aware Training (QAT). Echter, standaard QAT-technieken falen bij trellis-codes omdat het Viterbi-decoderingsalgoritme een niet-differentieerbare argmax-operatie omvat over exponentieel veel paden. Het toepassen van een Straight-Through Estimator (STE) op de trellis vernietigt de koppeling tussen aangrenzende gewichtsbeslissingen, wat het primaire voordeel is van trellis-codering ten opzichte van scalare quantisatie.

Methodologie: BCJR-QAT

Het artikel introduceert BCJR-QAT, een differentieerbare relaxatie van trellis-codering die de harde Viterbi argmax vervangt door een zachte verwachting afgeleid van het BCJR (Bahl-Cocke-Jelinek-Raviv) forward-backward algoritme.

Kernmechanisme

Analogie met Statistische Mechanica: De trellis wordt behandeld als een factorgrafiek die equivalent is aan een 1D Ising-achtige spin-keten. De Viterbi argmax wordt beschouwd als het nultemperatuurlimiet ( $T \to 0$ ) van een Boltzmann-verdeling over paden.
Zacht Codewoord: Bij een eindige temperatuur $T > 0$ berekent de methode de exacte marginale posterior-kansen van elke staat met het BCJR-algoritme. De gequantiseerde output is de Boltzmann-verwachting van het uitgezonden codewoord:
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
Dit zachte codewoord is exact differentieerbaar met betrekking tot het gewicht-pre-image en de trellis-parameters.
Temperatuurverloop (Annealing): Tijdens het trainen wordt de temperatuur $T$ afgekoeld van een initiële waarde $T_0$ naar een bijna-nul $T_{end}$ . Naarmate $T \to 0$ , convergeert het zachte codewoord naar de harde Viterbi-output, waardoor de continue optimalisatie effectief "kristalliseert" tot een discrete oplossing.

Belangrijkste Technische Innovaties

Differentieerbaarheid: De BCJR forward-backward recursie (complexiteit $O(L \cdot S^2)$ ) levert exacte gradiënten door de trellis, waardoor de inter-staatkoppeling behouden blijft die scalare STE-methoden verliezen.
Triton Kernel: De auteurs hebben een gefuseerde Triton-kernel geïmplementeerd die gather-, log-sum-exp- en additie-operaties combineert tot één lancering. Dit reduceert de forward/backward-passage van honderden kernel-lanceringen naar één, wat een 6,57× snelheidswinst oplevert ten opzichte van een referentie-implementation van PyTorch autograd, terwijl numerieke pariteit in FP32 behouden blijft.
Schema-correctie (Skip-High-T): Het artikel diagnosticeert een kritieke faalmodus in standaard gesimuleerde annealing-schema's (startend bij $T_0=1.0$ ). Bij hoge temperaturen is de gradiënt van het zachte codewoord "uitgesmeerd" en oninformatief, waardoor de optimizer in een slechtere Voronoi-bekken afdrijft voordat het afkoelingsfase zich kan herstellen. Het voorgestelde correctiemiddel is om de hoge- $T$ -fase over te slaan, door te starten bij een lagere temperatuur (bijv. $T_0=0.3$ ) waar gradiënten al informatief zijn.

Experimentele Resultaten

Het artikel evalueert BCJR-QAT op twee architecturen: OLMoE-1B-7B (MoE) en Llama-3.2-1B (Dense).

1. OLMoE-1B-7B (Negatief Resultaat: Proxy-gat)

Opzet: Per-laag greedy QAT die optimaliseert voor reconstructie-MSE (een proxy-doel).
Uitkomst: BCJR-QAT slaagde erin de per-laag reconstructie-MSE met 17–20% te verminderen in vergelijking met QTIP-PTQ. Deze verbetering vertaalde zich echter niet naar prestaties op de eindtaak. De modellen degradeerden met $\sim 1.3$ perplexiteit (PPL) op WikiText-2 in vergelijking met de QTIP-PTQ-basislijn.
Conclusie: Dit bevestigt een "proxy-gat" waarbij het minimaliseren van lokale reconstructiefouten geen betere prestaties op de eindtaak garandeert bij 2 bpw.

2. Llama-3.2-1B (Positief Resultaat: End-to-End KL)

Opzet: Enkellaags QAT die optimaliseert voor end-to-end forward-KL-distillatie tegen een FP16-leraar.
Uitkomst:
- Schema-overschrijding: Het gebruik van een naïef hoge- $T$ -schema ( $T_0=1.0$ ) resulteerde in prestaties die licht slechter waren dan PTQ (+0,005 tot +0,022 PPL).
- Skip-High-T Correctie: Het gebruik van het gecorrigeerde schema ( $T_0=0.3$ ) elimineerde de overschrijding. Op Laag 4 bereikte het model 10,1347 PPL, een verbetering van −0,084 PPL ten opzichte van de QTIP-PTQ-basislijn (10,2189).
- Super-additiviteit: Bij het combineren van getrainde lagen (Laag 4 en Laag 8) overtrof de gezamenlijke winst (−0,077 PPL) de som van individuele winsten (−0,062 PPL), wat een "coöperatie-overschot" aantoont waarbij BCJR-getrainde codewoorden elkaar versterken.

3. Theoretische Analyse

Drift-Budget Theorie: De auteurs leiden een voorwaarde af om het PTQ Voronoi-bekken te ontvluchten: de cumulatieve gewichtsdrijving ( $\eta \cdot N \cdot g_{max}$ ) moet de straal van het bekken overschrijden ( $r_{Voronoi} \approx 10^{-3}$ ). Experimenten bevestigen dat schema's die deze voorwaarde niet halen, niet uit het PTQ-oplossing kunnen ontsnappen, terwijl die welke deze overschrijden betere oplossingen kunnen vinden mits het doel en het schema correct zijn.

Betekenis en Beweringen

Het artikel claimt drie primaire bijdragen:

Algoritmisch: BCJR-QAT is de eerste methode die trellis-codering native differentieerbaar maakt, waardoor QAT op TCQ mogelijk wordt zonder de voordelen van padkoppeling te verliezen.
Technisch: De gefuseerde Triton-kernel maakt dit computationeel haalbaar op consumentenhardware (één 12GB GPU), waardoor de implementatiebarrière voor onderzoek op het moment van inferentie wordt weggenomen.
Empirisch Inzicht:
- Het identificeert dat per-laag reconstructie-MSE een slechte proxy is voor 2-bit QAT, terwijl end-to-end KL-distillatie effectief is.
- Het daagt de conventionele wijsheid van gesimuleerde annealing voor discrete relaxaties uit, door aan te tonen dat starten "heet" ( $T=1.0$ ) schadelijk is voor BCJR-QAT omdat de hoge-temperatuurgradiënt door ruis wordt gedomineerd. De optimale strategie is om te starten bij een temperatuur waar de gradiënt al informatief is.

Het artikel concludeert dat hoewel enkellaagsresultaten veelbelovend zijn, het volledige potentieel van BCJR-QAT waarschijnlijk ligt in gezamenlijke end-to-end training over alle lagen, een taak die betaalbaar is gemaakt door de vrijgegeven kernel maar cloud-schaal hardware (H100/H200) vereist voor modellen groter dan 1B parameters. De auteurs stellen expliciet dat zij deze volledige schaal gezamenlijke training in dit werk niet hebben uitgevoerd.