Technische Samenvatting: WUSH: Nabij-optimale Adaptieve Transformaties voor LLM Kwantisatie

1. Probleemstelling

Het kwantiseren van gewichten en activaties van Large Language Models (LLM's) is een standaardmethode voor efficiënte implementatie. Een centrale moeilijkheid ontstaat echter door "outliers"—enkele extreme waarden in gewichten en activaties die het dynamische bereik oprekken. Deze outliers dwingen de kwantisatieschalen tot grote waarden, wat de effectieve resolutie van low-bit representaties verslechtert en kwantisatiefouten versterkt.

Eerdere mitigatiestrategieën maken vaak gebruik van lineaire transformaties (bijv. rotaties) vóór kwantisatie om variantie gelijkmatiger over kanalen te verspreiden. Hoewel Hadamard-rotaties empirisch effectief zijn gebleken, zijn ze doorgaans vast en data-agnostisch. Ze passen zich niet aan de specifieke statistieken van de onderliggende gewichten of activaties aan. Dit roept een fundamentele vraag op: als de Hadamard-transformatie niet data-bewust is, in welke zin kan deze dan als optimaal voor kwantisatie worden beschouwd? Bestaande methoden die transformaties leren (bijv. SpinQuant, FlatQuant) vertrouwen op iteratieve optimalisatie (gradient descent) op kalibratiedata, wat kostbaar en heuristisch kan zijn.

2. Methodologie

De auteurs stellen WUSH voor, een closed-form, data-bewuste lineaire blockwise transformatie ontworpen voor gezamenlijke gewicht-activatie kwantisatie. De methode is afgeleid om de $L_2$ outputverlies te minimaliseren onder standaard Round-to-Nearest (RTN) AbsMax-geschaalde block kwantiseerders, waarbij zowel floating-point (FP) als integer (INT) formaten worden gedekt.

Theoretische Afleiding

Het probleem wordt geformuleerd als het minimaliseren van de verwachte kwantisatiefout voor een blok gewichten $W$ en activaties $X$. De auteurs modelleren de kwantisatiefout probabilistisch:

• Voor Floating-Point (FP): De fout wordt gemodelleerd als proportioneel aan de absolute waarde van de input ($\varepsilon(\alpha) = \text{diag}(\eta)\alpha$).
• Voor Integer (INT): De fout wordt gemodelleerd als proportioneel aan de maximale absolute waarde binnen een groep ($\varepsilon(\alpha) = \|\alpha\|_\infty \eta$).

Door de tweede-orde statistieken (covariantie) van de gewichts- en activatieblokken te analyseren, leiden de auteurs een closed-form optimale transformatie af. Laat $W'$ en $X'$ matrices zijn afgeleid van de Cholesky-decompositie van de genormaliseerde tweede momenten van de gewichten en activaties, respectievelijk. Laat $U, S, V$ de Singular Value Decomposition (SVD) zijn van $W'^\top X'$. De optimale transformaties $T_W$ en $T_X$ worden geconstrueerd als:
$$T_X = H S^{-1/2} V^\top X'^\top$$
$$T_W = H S^{-1/2} U^\top W'^\top$$
waarbij $H$ een genormaliseerde Hadamard-matrix is.

Kerneigenschappen van deze constructie:

1. Data-bewust: In tegen talijke vaste Hadamard-transformaties incorporeert WUSH $S^{-1/2}$ en de SVD-componenten ($U, V$), die afhankelijk zijn van de specifieke tweede-orde statistieken van de data.
2. Niet-orthogonaal: De resulterende transformatie is over het algemeen niet-orthogonaal, waardoor het in staat is om schalen tussen gewichten en activaties te balanceren op een manier die orthogonale transformaties niet kunnen.
3. Optimaliteit: Het paper bewijst dat deze constructie bewezen optimaal is voor FP block kwantiseerders en asymptotisch optimaal voor INT block kwantiseerders (binnen een factor $d$ of $d^{o(1)}$ afhankelijk van de distributie).

Integratie met GPTQ

De methode integreert natuurlijk met GPTQ (Greedy Per-Block Quantization). Omdat GPTQ de Hessiaan (tweede-orde informatie) gebruikt voor foutpropagatie, kan WUSH dezelfde statistieken gebruiken om de transformatie te construeren. De auteurs stellen een geïnterleefd computationeel schema voor waarbij gewichtsupdates en transformatieconstructie gekoppeld zijn, waardoor WUSH toegepast kan worden naast GPTQ met verwaarloosbare extra overhead.

Implementatie

Om de uitdaging aan te pakken van het toepassen van een afzonderlijke transformatie per blok tijdens inferentie, implementeren de auteurs een fused GPU kernel. De WUSH-transformatiematrices worden in een getransponeerde lay-out $(G, G, C)$ opgeslagen om het geheugenbandbreedtegebruik te maximaliseren. De kernel fuseert de blockwise transformatie met de activatiekwantisatie, waarmee een throughput wordt bereikt die vergelijkbaar is met geoptimaliseerde Hadamard-kernels.

3. Belangrijkste Bijdragen

1. Closed-Form Optimale Transformaties: De afleiding van de eerste closed-form, data-bewuste lineaire transformaties voor gezamenlijke gewicht-activatie kwantisatie die bewezen nabij-optimaal zijn voor zowel FP als INT formaten.
2. Theoretisch Inzicht: Het paper verheldert de rol van de Hadamard-rotatie. Het laat zien dat hoewel Hadamard een effectieve data-agnostische orthogonale transformatie is, de optimale oplossing een niet-orthogonale component ($S^{-1/2}$) vereist die afgeleid is van de datastatistieken om de kwantisatiefout werkelijk te minimaliseren.
3. Efficiënte Implementatie: Een gespecialiseerde fused GPU kernel die per-block adaptieve transformaties afhandelt zonder de inferentiesnelheid op te offeren, waarbij de prestaties van fixed-transform baselines worden geëvenaard.
4. Empirische Validatie: Uitgebreide experimenten die significante nauwkeurigheidsverbeteringen laten zien ten opzichte van state-of-the-art baselines (Hadamard-gebaseerde MR-GPTQ, SmoothQuant, QuaRot, SpinQuant) over meerdere modelfamilies (Llama, Qwen) en kwantisatieformaten (MXFP4, NVFP4, INT4).

4. Experimentele Resultaten

Het paper evalueert WUSH op Llama-3.1-8B-Instruct, Llama-3.2-3B-Instruct, en Qwen3-modellen (8B, 14B, 32B) met behulp van de LM Evaluation Harness en Platinum Benchmarks.

• Layerwise Loss: WUSH bereikt consistent de laagste layerwise RTN kwantisatiefout over MXFP4, NVFP4, en INT4 formaten vergeleken met Identity, Random Rotation, Hadamard, en niet-adaptieve varianten.
• End-to-End Nauwkeurigheid:
  • Op Llama-3.1-8B-Instruct met MXFP4 (W4A4), verbetert WUSH de nauwkeurigheid met +2,8 gemiddelde punten ten opzichte van de Hadamard-gebaseerde baseline (RTN) en +0,7 punten (GPTQ).
  • Het verkleint de prestatiekloof tussen NVFP en MXFP formaten aanzienlijk. Bijvoorbeeld, op Qwen3-14B vermindert WUSH de NVFP4-MXFP4 gemiddelde kloof tot 0,36 punten.
  • WUSH gecombineerd met GPTQ behaalt de beste resultaten over de meeste benchmarks, en herstelt vaak meer dan 96% van de BF16 prestaties in W4A4 settings.
• Throughput: De fused kernel implementatie levert tot wel 5,8× per-layer throughput over BF16 op RTX 5090 GPU's, met slechts een 1,3% gemiddeld throughput verschil vergeleken met de geoptimaliseerde Hadamard baseline.
• Overhead: De offline preprocessing kosten zijn vergelijkbaar met standaard GPTQ. De opslagoverhead voor de activatie-zijde transformaties is verwaarloosbaar (bijv. ~1,4% voor MXFP4 op Qwen3-8B).

5. Betekenis en Claims

Het paper claimt dat WUSH het empirische succes van Hadamard-rotaties verheldert door te tonen dat zij een specifieke, data-agnostische component zijn van een bredere optimale oplossing. Door een vaste Hadamard-backbone te combineren met een data-afhankelijke tweede-moment component, bereikt WUSH een bewijsbaar nabij-optimale niet-orthogonale transformatie.

De auteurs benadrukken dat hun aanpak geprincipieerd en data-bewust is, wat verder gaat dan heuristische of iteratieve leermethoden. Ze stellen dat WUSH het MXFP formaat competitief maakt met NVFP, waarmee een belangrijke beperking in huidige low-precision inferentie wordt aangepakt. Het werk demonstreert dat geprincipeerde, data-bewuste block transformaties significante verbeteringen in nauwkeurigheid kunnen opleveren zonder de efficiëntie in gevaar te brengen die vereist is voor grootschalige LLM-implementatie.