Technische Samenvatting: ADMM-Q voor Post-Training Quantisatie van LLM's

Probleemstelling

Grote Taalmodellen (LLM's) staan voor aanzienlijke uitdagingen bij implementatie vanwege hun reken- en geheugeneisen. Quantisatie, en specifiek Post-Training Quantisatie (PTQ), is een primaire strategie voor het comprimeren van deze modellen door de precisie van de gewichten te verlagen (bijvoorbeeld van 16-bit naar 4-bit of lager) zonder hertraining.

Huidige standaardmethoden voor gewichtsquantisatie, zoals Round-to-Nearest (RTN) en Greedy Per-Column Quantization (GPTQ), vertonen beperkingen, met name bij agressieve bit-breedtes (sub-4-bit):

• RTN rondt gewichten onafhankelijk af, waarbij het dependencies tussen gewichten en de structuur van het landschap van reconstructiefouten negeert.
• GPTQ integreert Hessian-informatie om fouten te compenseren, maar verwerkt gewichten op een hebberige manier, kolom voor kolom. Deze sequentiële aanpak is gevoelig voor de verwerkingsvolgorde en cumuleert fouten over kolommen heen.

De kernuitdaging bestaat uit het formuleren van quantisatie op laagniveau als een geconstrueerd optimalisatieprobleem dat reconstructiefout minimaliseert terwijl het discrete quantisatiebeperkingen respecteert; een probleem dat steeds moeilijker wordt naarmate de bit-breedte afneemt.

Methodologie: ADMM-Q

De auteurs stellen ADMM-Q voor, een nieuw algoritme voor gewichtsquantisatie gebaseerd op de Alternating Direction Method of Multipliers (ADMM). In tegenstelling tot de hebberige kolomsgewijze aanpak van GPTQ, behandelt ADMM-Q quantisatie op laagniveau als een gezamenlijk optimalisatieprobleem.

1. Probleemformulering

Het algoritme formuleert de quantisatie van een laag met voorgeöefende gewichten $\hat{W}$ en input-activaties $X$ als een Hessian-gewogen geconstrueerd optimalisatieprobleem:
$$\min_{W} \frac{1}{2} \text{Tr}\left((W - \hat{W})^\top H (W - \hat{W})\right) \quad \text{s.t.} \quad W \in \mathcal{Q}$$
waarbij $H = X^\top X + \lambda I$ de Hessiaan-matrix is, en $\mathcal{Q}$ de verzameling van gekwantiseerde gewichtsmatrices is.

2. Belangrijke Algoritmische Componenten

Om dit niet-convexe probleem efficiënt op te lossen op de schaal van LLM's, integreert ADMM-Q enkele verbeteringen:

• Diagonale Schaling voor Numerieke Stabiliteit: LLM-activaties bevatten vaak "kanaaluitbijters" (dimensies met onevenredig grote magnitudes), wat leidt tot een slecht geconditioneerde Hessiaan. ADMM-Q past een diagonale schalingsmatrix $\Sigma = \text{Diag}(H)^{-1/2}$ toe om het probleem te normaliseren, zodat alle kanalen vergelijkbaar bijdragen aan de updates. Dit is cruciaal; ablatiestudies tonen aan dat het verwijderen van deze schaling de prestaties aanzienlijk verslechtert (bijvoorbeeld een toename van de perplexiteit van 10,06 naar 14,68 bij W3A16 op Qwen3-8B).
• ADMM-splitsing: Het algoritme introduceert een hulpvariabele $D$ om het continue Hessian-gewogen doel te ontkoppelen van de discrete quantisatiebeperking. Het wisselt af tussen:
  1. W-update: Een continue update die de Hessian-gewogen afstand minimaliseert, efficiënt opgelost via eigendecompositie van $H$.
  2. D-update: Een discrete projectie op het quantisatierooster (equivalent aan RTN, maar rekening houdend met schaling).
  3. Dual Update: Aanpassen van de Lagrange-multiplicator om consistentie af te dwingen.
• Grid Refresh: In tegenstelling tot GPTQ, dat het quantisatierooster tijdens de optimalisatie vastzet, past ADMM-Q de roosterparameters (schaal en nulpunt) dynamisch opnieuw aan op basis van de huidige continue iteratie. Dit zorgt ervoor dat de projectie blijft uitgelijnd met het bewegend doel, waardoor de veroudering van een vast rooster wordt vermeden.
• Lokaal Zoeken voor Verfijning: Om suboptimale discrete oplossingen veroorzaakt door gecorreleerd afronden aan te pakken, voert een nabewerkingsstap een efficiënte pair-swap lokale zoektocht uit. Deze stap evalueert discrete bewegingen (het wisselen van waarden tussen rijen) die de reconstructiefout direct verlagen, waarbij GPU-parallelisatie wordt benut voor snelheid.
• Convergentiegaranties: De auteurs breiden bestaande ADMM-convergentiebewijzen (voorheen vastgesteld voor sparsiteitsbeperkingen) uit naar de quantisatiesetting, en bewijzen dat de continue en discrete iteraties convergeren naar hetzelfde limietpunt onder een specifieke straffschema ($\rho_t = \rho_0 \cdot \gamma^t$).

3. Modulariteit en Integratie

ADMM-Q is ontworpen als een drop-in vervanging voor RTN of GPTQ binnen bestaande PTQ-pijplijnen. Het is volledig composable met:

• Schaaltransformaties (bijvoorbeeld SmoothQuant, AWQ).
• Rotatietransformaties (bijvoorbeeld SpinQuant, QuaRot).
• Clipping-heuristieken (bijvoorbeeld MSE-gebaseerde grid search).

Experimentele Resultaten

De auteurs evalueerden ADMM-Q in drie scenario's: alleen-gewicht quantisatie, gewicht-en-activatie quantisatie met SpinQuant, en gewicht-en-activatie quantisatie met SmoothQuant.

1. Alleen-Gewicht Quantisatie (Qwen3-familie)

Over Qwen3-modellen (1,7B tot 32B) heen presteerde ADMM-Q consistent beter dan GPTQ en AWQ in perplexiteit en zero-shot nauwkeurigheid, met winsten die toenamen bij lagere bit-breedtes:

• W3A16 (Qwen3-8B): WikiText-2 perplexiteit verbeterde van 12,85 (GPTQ) naar 10,06 (ADMM-Q).
• W2A16 (Qwen3-8B): WikiText-2 perplexiteit verbeterde van 32,88 (GPTQ) naar 17,69 (ADMM-Q).
• Laagsgewijze Analyse: ADMM-Q verlaagde de reconstructiefout over alle decoderlagen, met tot 10% van de fout van GPTQ in vroege lagen.

2. Gewicht-en-Activatie Quantisatie

• SpinQuant (WbA4KV4): Op LLaMA-3 en Qwen3-modellen leverde het vervangen van GPTQ door ADMM-Q consistente verbeteringen op. Bijvoorbeeld, op LLaMA-3 8B bij W2A4KV4 daalde de perplexiteit van 50,27 (GPTQ) naar 15,65 (ADMM-Q).
• SmoothQuant (WbA8): Op Qwen3-8B verbeterde ADMM-Q de W3A8 perplexiteit van 16,00 naar 12,21.

3. Inferentieprestaties

Wanneer geïntegreerd in de vLLM-inferentiestapel, produceerde ADMM-Q gewichtsmatrices in hetzelfde formaat als GPTQ. Bijgevolg bereikte het identieke inferentiethroughput en geheugenwinst. Bijvoorbeeld, op Qwen3-32B met SpinQuant W8A8 bereikten beide methoden een snelheidswinst van ongeveer 1,52× ten opzichte van de dichte baseline.

4. Rekenkosten

Het offline quantisatieproces voor ADMM-Q is rekenkundig intensiever dan GPTQ vanwege de iteratieve optimalisatie. De auteurs merken echter op dat quantisatie een eenmalige stap is, en dat de runtime praktisch blijft (bijvoorbeeld minder dan twee uur voor een 32B-model op een enkele GPU). Het piekgebruik van GPU-geheugen is iets hoger (ongeveer 6%) dan bij GPTQ.

Betekenis en Claims

Het artikel beweert dat ADMM-Q een principiële verschuiving vertegenwoordigt van hebberige, sequentiële quantisatie naar gezamenlijke, geconstrueerde optimalisatie voor LLM's. De belangrijkste bijdragen zijn:

1. Algoritmische Innovatie: Het slaagt erin ADMM toe te passen op het harde combinatorische probleem van gewichtsquantisatie, en introduceert noodzakelijke stabilisaties (diagonale schaling) en verfijningen (grid refresh, lokaal zoeken) voor de schaal van LLM's.
2. Prestatiewinst: Het toont aanzienlijke verbeteringen in modelnut (perplexiteit en nauwkeurigheid), met name in agressieve low-bit regimes (W3 en W2) waar eerdere methoden moeite hebben.
3. Veelzijdigheid: Het dient als een modulaire, drop-in component die bestaande pijplijnen (inclusief die met rotaties en schaling) verbetert zonder het uiteindelijke inferentieformaat of de snelheid te veranderen.

De auteurs erkennen beperkingen, en merken op dat ADMM-Q hogere offline quantisatiekosten met zich meebrengt en afhankelijk is van de representativiteit van de kalibratieset. Bovendien is de huidige evaluatie beperkt tot dichte autoregressieve decoder-only LLM's, en zou aanpassing aan mixture-of-experts of diffusiemodellen verder werk vereisen.