🤖 machine learning

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

JacQuant introduceert een nieuw Quantization-Aware Training-framework dat de breekbare Straight-Through Estimator vervangt door een lichtgewicht, geleerde Jacobiaan-surrogaat om het trainen te stabiliseren en hogere nauwkeurigheid te bereiken bij ultra-low-bit LLM-quantisatie, zonder de forward-quantizers te wijzigen of aanzienlijke runtime-kosten te veroorzaken.

Oorspronkelijke auteurs: Kai Yi, Vignesh Vivekraja, Harshit Khaitan, Steven Li

Gepubliceerd 2026-05-26

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Kai Yi, Vignesh Vivekraja, Harshit Khaitan, Steven Li

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

Het Grote Probleem: Het "Gebroken Kompas" bij AI-training

Stel je voor dat je probeert een robot te leren een doolhof te navigeren. Om ruimte en energie te besparen, besluit je de robot een kaart met zeer lage resolutie te geven (dit heet kwantisatie). In plaats van de exacte afstand tot een muur te weten, weet de robot alleen of hij zich in "Zone A", "Zone B" of "Zone C" bevindt.

Het probleem doet zich voor wanneer de robot precies op de lijn tussen Zone A en Zone B staat.

De Realiteit: Als de robot een heel klein beetje beweegt, kan hij in Zone A blijven. Beweegt hij een heel klein beetje de andere kant op, dan kan hij naar Zone B springen. De kaart is "blokkerig" en verandert niet soepel.
De Oude Methode (STE): Jarenlang gebruikte het trainingsalgoritme, wanneer de robot probeerde te leren van zijn fouten, een snelweg genaamd de Straight-Through Estimator (STE). Het deed alsof de kaart glad en continu was. Het vertelde de robot: "Je hebt 1 inch bewogen, dus de kaart is 1 inch veranderd."
Het Resultaat: Dit is alsof je de robot een gebroken kompas geeft. Wanneer de robot zich in de buurt van een grens bevindt, draait het kompas wild rond of wijst het de verkeerde kant op, omdat de kaart eigenlijk niet glad is. Dit maakt training onstabiel, vooral wanneer de kaart een zeer lage resolutie heeft (zoals 1 of 2 bits).

De Oplossing: JacQuant (De "Slimme Sensor")

De auteurs van dit paper, werkzaam bij Meta AI, introduceerden JacQuant. In plaats van te doen alsof de kaart glad is, leren ze de robot een lichtgewicht sensor die precies begrijpt hoe de blokkerige kaart reageert op beweging.

Zo werkt het, stap voor stap:

1. De "Surrogaat" Sensor

Stel je voor dat de robot een kleine, goedkope sensor aan zijn voet heeft bevestigd. Deze sensor verandert de kaart niet; de kaart blijft blokkerig. Maar deze sensor meet: "Als ik mijn voet een heel klein beetje bewig, verandert de zone dan echt?"

Als de robot zich in het midden van een zone bevindt, zegt de sensor: "Ja, een kleine beweging verandert dingen." (Gevoeligheid = 1).
Als de robot tegen een muur zit (verzadiging) of precies op een grens waar bewegen de zone niet verandert, zegt de sensor: "Nee, een kleine beweging doet niets." (Gevoeligheid = 0).

Deze sensor heet een Learned Jacobian Surrogate. Het is een eenvoudige wiskundige kaart die het trainingsalgoritme vertelt hoe "gevoelig" de huidige locatie is.

2. Het Kompas Repareren

Wanneer de robot een fout maakt, kijkt het trainingsalgoritme naar de sensorlezing voordat het een correctiesignaal stuurt.

Oude Weg (STE): "Je hebt een fout gemaakt! Duw hard in deze richting!" (Zelfs als de robot tegen een muur zit, duwt hij hard, waardoor hij nutteloos heen en weer stuitert).
JacQuant Weg: De sensor zegt: "Hé, je zit tegen een muur; bewegen helpt niet." Dus dempt het algoritme de duw. Het zegt: "Oké, duw hier niet zo hard; laten we een andere richting proberen."

Dit voorkomt dat de robot wild heen en weer stuitert in de buurt van de grenzen en helpt hem veel sneller de beste route te vinden.

3. De "Geamortiseerde" Kosten (Waarom het goedkoop is)

Je zou kunnen denken: "Het meten van deze gevoeligheid klinkt duur!"
De auteurs losten dit op door lui te zijn (op een slimme manier). Ze controleren de sensor niet elke seconde.

Ze controleren de sensor af en toe (zoals eens per 100 stappen).
Ze gebruiken die lezing langere tijd totdat de robot ver genoeg beweegt om een nieuwe controle nodig te hebben.
Dit heet amortisatie. Het is alsof je 's ochtends een keer het weerbericht controleert en die informatie de hele dag gebruikt, in plaats van elke 5 minuten naar buiten te stappen. De kosten zijn zo laag (minder dan 2% extra tijd) dat het gratis voelt.

De Resultaten: Soepeler, Sneller, Beter

Het paper testte dit op Grote Taalmodellen (LLMs) zoals LLaMA en Qwen, specifiek wanneer ze worden samengeperst naar ultra-lage precisie (1 of 2 bits).

De Analogie: Stel je voor dat je probeert een gigantisch, hoog-resolutie schilderij in een klein, gepixelde lijst te passen. De oude methode (STE) zorgde ervoor dat het beeld wazig en onstabiel leek. JacQuant werkt als een slim filter dat precies weet welke pixels scherp moeten blijven en welke glad moeten worden gemaakt, wat resulteert in een veel duidelijker beeld.
Het Bewijs: In hun tests leerden modellen die met JacQuant waren getraind:
- Sneller (convergeerden sneller).
- Minder fouten (lagere "perplexity", een maatstaf voor hoe verward de AI is).
- Beter op vragen (hogere nauwkeurigheid bij redeneertaken).
- Stabiel (crashten niet of oscilleerden niet in de buurt van de "muren" van de kwantisatiezones).

Samenvatting

JacQuant is een nieuwe manier om AI-modellen te trainen die zijn gecomprimeerd om zeer klein te zijn. In plaats van blind te gokken hoe het gecomprimeerde model leert (met behulp van de oude "Straight-Through Estimator"), leert het een eenvoudige, goedkope "gevoeligheidskaart". Deze kaart vertelt de AI precies hoeveel het zijn eigen bewegingen kan vertrouwen in de buurt van de randen van zijn gecomprimeerde wereld. Het resultaat is een stabielere, nauwkeurigere en efficiëntere trainingsprocess voor de kleinste, meest geheugenvriendelijke AI-modellen.

Technische Samenvatting: JacQuant

Probleemstelling

Kwantiseringsbewuste training (QAT) is een standaardbenadering voor het implementeren van Large Language Models (LLM's) onder strikte geheugen- en latentiebeperkingen, met name in regimes met ultra-lage bitdiepte (bijv. gewichten van $\le 2$ bits). Moderne QAT-pijplijnen vertrouwen echter bijna universeel op de Straight-Through Estimator (STE). De STE benadert de niet-differentieerbare kwantisator door gradiënten door te geven alsof de kwantisator de identiteitsfunctie is.

Het artikel identificeert een kritieke gebrek in deze aanpak: harde kwantisatoren zijn stuksgewijs constant met binnengrenzen en verzadigings-/clippingregio's. In ultra-lage bit-instellingen liggen veel parameters dicht bij deze grenzen of raken ze de clippinglimieten. De aanname van de STE van een identiteits-Jacobaan creëert een aanhoudende mismatch, waardoor gradiënten worden overgepropageerd in verzadigde regio's waar de kwantisator lokaal ongevoelig is. Dit leidt tot breekbare training, instabiliteit nabij binnengrenzen en oplossingen die zwak zijn afgestemd op het daadwerkelijke gedrag van het model met lage precisie.

Methodologie: JacQuant

JacQuant stelt een kader voor dat de vaste identiteits-Jacobaan van de STE vervangt door een geleerde, lichtgewicht surrogaat-Jacobaan, aangeduid als $B(W)$ . De kernfilosofie is om QAT te behandelen als een "ontbrekende-Jacobaan"-probleem: terwijl de forward-pass de ingebouwde harde kwantisator $Q_\Delta(\cdot)$ behoudt, leert de backward-pass hoe verstoringen in full-precision gewichten hun gekwantiseerde proxies beïnvloeden.

Kernmechanisme

In plaats van de standaard STE-update $\nabla_W \ell \approx \nabla_{\hat{W}} \ell$ , gebruikt JacQuant:
$\nabla_W \ell \approx B(W) \nabla_{\hat{W}} \ell$
waarbij $\hat{W} = Q_\Delta(W)$ en $B(W)$ een geleerde gevoeligheidskaart is.

Belangrijkste Ontwerpkiezen

Parametrisatie: $B(W)$ is beperkt tot diagonaal of blok-diagonaal, met één scalair versterkingsfactor per kwantisatiegroep (bijv. 128 gewichten). Dit houdt de geheugenoverhead verwaarloosbaar (slechts ~1,5% toevoeging aan Adam-optimizer-toestanden).
Schattingstrategieën: Het artikel introduceert twee praktische, datagedreven schatters om $B(W)$ $B (W)$ bij te werken:
- JacQuant-Probe: Gebruikt Gaussische willekeurige probes $\delta_g$ om de lokale helling te meten $\Delta q_g = Q_\Delta(W_g + \delta_g) - Q_\Delta(W_g)$ . Een kleinste-kwadratenfit schat de lokale gevoeligheid.
- JacQuant-Dither: Gebruikt aftrekkende dithering (het toevoegen van uniforme ruis) om een differentieerbare proxy voor de kwantisator te creëren, waarbij de gevoeligheid wordt geschat op basis van het verwachte antwoord.
Geamortiseerde Updates: Om de trainingsdoorvoer te behouden, wordt $B(W)$ niet bij elke stap bijgewerkt. Het wordt af en toe ververst (bijv. elke 50–200 stappen) met een Exponentiële Glijdende Gemiddelde (EMA). Tussen de verversingen door wordt de huidige $B(W)$ hergebruikt.
Variance-Reductie: Het kader integreert met variance-gereduceerde optimalisatoren (bijv. loopless-SVRG, SAGA) om de stochasticiteit die wordt geïntroduceerd door het probeer-/dither-proces te stabiliseren.

Theoretische Bijdragen

Het artikel biedt een convergentieanalyse voor JacQuant onder niet-convexe doelfuncties, gestructureerd rond vier kernargumenten:

Doelgradiëntdefinitie: De auteurs definiëren een "mean-field" doelgradiënt $g^\dagger(W) = J(W)\bar{v}$ , waarbij $J(W)$ de ware gevoeligheid is van de door dithering-gemiddelde kwantisator. Ze bewijzen dat de bias van JacQuant evenredig is met $\|B(W) - J(W)\|$ , terwijl de bias van de STE evenredig is met $\|I - J(W)\|$ .
Consistentie van Schatters: Ze tonen aan dat zowel de Probe- als de Dither-schattingen consistent zijn met de doelgevoeligheid $J(W)$ . Specifiek is de Probe-update een lokale regressie die convergeert naar de mean-field gevoeligheid, en heeft de Dither-update $J(W)$ als zijn unieke populatie-vaste punt.
Stabiliteit: De geleerde $B(W)$ fungeert als een adaptieve versterkingsregelaar. Het blijft dicht bij de identiteit in responsieve binnengebieden, maar dempt gradiënten (versterking $< 1$ ) in verzadigde regio's waar de kwantisator ongevoelig is. Dit voorkomt de oscillaties die gebruikelijk zijn bij STE-gebaseerde training.
Convergentiesnelheden: Onder een Polyak-Łojasiewicz (PŁ) conditie en aannemende de "code-bewarende" eigenschap (kwantisatiecodes veranderen niet vaak binnen een lokaal venster), bewijst het artikel dat JacQuant lineaire convergentiesnelheden bereikt. De convergentiefout wordt begrensd door de Jacobaan-trackingfout en de stochasticiteitsvloer.

Experimentele Resultaten

De auteurs evalueren JacQuant op verschillende LLM-families (LLaMA3-1B/3B/8B, Qwen3-1.7B) in ultra-lage bit-regimes ( $W \le 2$ , $A8/A16$ ).

Prestaties: JacQuant presteert consequent beter dan STE-gebaseerde baselines (inclusief ParetoQ en WinQ) en de "Rotation Trick" (een vaste geometrische transformatiemethode).
- Op LLaMA3-1B (W2A16) verbeterde JacQuant de perplexiteit van 13,5 naar 12,3 en de nauwkeurigheid van 55,7% naar 56,6% ten opzichte van de beste STE-gebaseerde baseline.
- Op LLaMA3-8B (W2A16) toonde JacQuant monotoon verbeteringen ten opzichte van ParetoQ over 30K, 60K en 90K trainingsstappen.
Efficiëntie: De runtime-overhead is verwaarloosbaar. Op LLaMA3-1B varieert de overhead van 1–5% afhankelijk van de groepsgrootte; op LLaMA3-8B blijft deze onder de 1,6% voor praktische groepsgroottes.
Validatie van het Mechanisme: Finite-difference-diagnostiek bevestigt dat JacQuant groepsgain $< 1$ leert in afgekapte/verzadigde regio's, waardoor gradiënten effectief worden gedempt waar STE ze zou overgepropageerd.

Betekenis en Claims

Het artikel positioneert JacQuant als een drop-in upgrade voor bestaande QAT-pijplijnen. De betekenis ligt in:

Ste-Vrije Training: Het verwijdert de afhankelijkheid van de heuristische identiteits-Jacobaan zonder wijzigingen aan de forward-kwantisator of de modelarchitectuur te vereisen.
Geometrie-bewustzijn: Door de lokale gevoeligheid te leren, past het zich van nature aan aan de geometrie van het kwantisatielandschap, specifiek de verzadigings- en clippingproblemen die ultra-lage bit-training teisteren.
Praktijktoepasbaarheid: Het bereikt deze winsten met minimale rekenkosten, waardoor het levensvatbaar is voor grootschalige LLM-training.

De auteurs beperken hun empirische claims expliciet tot ultra-lage bit LLM QAT ( $\le 2$ bits). Ze erkennen dat hoewel de methode een drop-in vervanging is voor de backward-regel, de forward-kwantisator een harde, ingebouwde operator blijft. Het artikel claimt niet alle kwantisatie-uitdagingen op te lossen, maar richt zich specifiek op de gradiëntmismatch die inherent is aan STE-gebaseerde optimalisatie. Toekomstige werkrichtingen die worden genoemd, omvatten het uitbreiden van de methode naar mixed-precision-beleid, niet-uniforme kwantisatoren en het verfijnen van de schattingsplanning.