LATMiX: Learnable Affine Transformations for Microscaling Quantization of LLMs

1. Probleemstelling

Grote Taalmodellen (LLMs) worden steeds vaker ingezet met behulp van Post-Training Quantization (PTQ) om geheugen- en rekentkosten te verlagen. Een specifiek dataformaat dat aan populariteit wint, is Microscaling (MX), dat tensoren partitioneert in blokken en een unieke schalingsfactor toewijst aan elk blok om dynamische bereiken efficiënter te hanteren dan globale schaling.

Het combineren van MX-quantisatie met bestaande technieken voor het mitigeren van uitbijters (die cruciaal zijn voor quantisatie met weinig bits) vormt echter aanzienlijke uitdagingen:

• Activeringsuitbijters: LLM-activeringen bevatten extreme waarden (uitbijters) die de quantisatiefout domineren.
• Beperkingen van huidige methoden: Bestaande benaderingen gebruiken inverteerbare transformaties (doorgaans rotaties of Hadamard-transformaties) om activeringsenergie te herschikken en uitbijters te onderdrukken.
  • Globale Transformaties: Het toepassen van een globale rotatie op activeringen vóór MX-quantisatie leidt tot ernstige prestatiedegradeatie, omdat de globale transformatie in conflict komt met de blokgewijze schalingsstructuur van MX.
  • Block-diagonale Transformaties: Recent werk suggereert het toepassen van transformaties onafhankelijk binnen elk MX-blok (block-diagonale matrices). Hoewel dit het mismatch-probleem omzeilt, beperkt dit de transformatie tot kleine, geïsoleerde deelruimtes, waardoor een effectieve herschikking van activeringsmassa over de hele tensor wordt verhinderd. Bijgevolg kunnen dominante uitbijters niet voldoende worden gediffuseerd, wat leidt tot suboptimale nauwkeurigheid in regimes met weinig bits (bijv. 4-bit).

2. Methodologie: LATMiX

De auteurs stellen LATMiX (Learnable Affine Transformations for Microscaling) voor, een methode die uitbijterreductie generaliseert naar leerbare, volwaardige inverteerbare affiene transformaties die specifiek zijn geoptimaliseerd voor het MX-dataformaat.

A. Theoretische Analyse

De auteurs leiden een theoretische bovengrens af voor de quantisatiefout onder MX-quantisatie. De analyse toont aan dat de fout afhangt van twee concurrerende factoren:

1. De Spectrale Norm van de Inverse Transformatie ($||A^{-1}||_\sigma$): Een maat voor hoe sterk de transformatie fouten versterkt.
2. De Verwachte Maximale Activering per Blok: Het gemiddelde van de maximale absolute waarden binnen elk MX-blok.

Het artikel betoogt dat block-diagonale transformaties (gebruikt in eerdere werken) een speciaal geval zijn dat faalt in het globaal balanceren van deze factoren. In plaats daarvan is een volledige (niet-block-diagonale) affiene transformatie vereist om expliciet rekening te houden met zowel de MX-blokstructuur als de globale verdeling van kenmerken.

B. LATMiX Architectuur

LATMiX leert een reeks inverteerbare affiene transformaties ($T(x) = Ax + v$) die kunnen worden ingebouwd in de lineaire lagen van het model, wat resulteert in geen extra inferentie- overhead.

• Transformatie Locaties:
  • $T_1$: Globaal toegepast op het residupad (input embeddings).
  • $T_2$: Lokaal toegepast op de waardenvectoren in Multi-Head Attention (MHA) en binnen Feed-Forward Networks (FFN).
  • $T_3$: Een online block-Hadamard-transformatie (behouden voor consistentie met eerdere baselines).
• Parametrisatie: Om inverteerbaarheid te garanderen en standaard deep learning-optimalisatie mogelijk te maken, stellen de auteurs twee parametrisaties voor de matrix $A$ voor:
  1. LU-decompositie: $A = P(L + \text{diag}(s))$, waarbij $L$ een onderunitriangulaire matrix is.
  2. QR-decompositie: $A = Q(R + \text{diag}(s))$, waarbij $Q$ orthogonaal is (via matrixexponentiële van een schuifsymmetrische matrix) en $R$ een boventriangulaire matrix is.
  • Er wordt een volumebehoudende regularisatie toegepast om ervoor te zorgen dat de determinant dicht bij 1 blijft, wat numerieke instabiliteit voorkomt terwijl individuele richtingen wel kunnen uitbreiden of krimpen.
• Optimalisatiedoel:
  • Distillatieverlies: In plaats van strikte computationele invariantie af te dwingen (wat orthogonale matrices vereist), gebruikt LATMiX een Kullback-Leibler (KL-divergentie)-verlies tussen het full-precision teacher-model en het getransformeerde, gekwantiseerde student-model. Dit staat de transformatie toe af te wijken van orthogonaliteit om beter te passen bij de dataverdeling.
  • Regularisatie: Een term om het volume van de transformatie dicht bij 1 te houden.

C. Inbouwmechanisme (Folding)

Een kritieke bijdrage is het vermogen om de geleerde affiene transformaties ($A$ en $v$) direct in de gewichtsmatrices en biases van de LLM te inbouwen (folden).

• De inverse transformatie wordt ingebouwd in de input-projecties (K, Q, V) van het attentiemechanisme.
• De voorwaartse transformatie wordt ingebouwd in de output-projecties van de attentie- en FFN-lagen.
• Dit zorgt ervoor dat het model tijdens inferentie exact hetzelfde gedrag vertoont als een standaard gekwantiseerd model, zonder extra rekentkosten.

3. Belangrijkste Bijdragen

1. Theoretisch Inzicht: Verschaffing van een theoretische grens die aantoont dat quantisatiefout in MX-formaten afhangt van zowel de globale kenmerkverdeling als de blokstructuur, en bewijst dat block-diagonale transformaties ontoereikend zijn voor optimale foutreductie.
2. LATMiX-algoritme: Introductie van een methode om algemene inverteerbare affiene transformaties (niet beperkt tot rotaties) te leren met behulp van LU/QR-parametrisaties en distillatieverlies.
3. Implementatie zonder Overhead: Aantonen dat deze complexe geleerde transformaties wiskundig kunnen worden ingebouwd in bestaande lineaire lagen, waardoor de inferentiesnelheid en het geheugengebruik van standaard gekwantiseerde modellen behouden blijven.
4. Uitgebreide Evaluatie: Validering van de methode over meerdere modelgroottes (1B tot 32B) en architecturen (Llama, Qwen) op diverse benchmarks.

4. Experimentele Resultaten

De auteurs evalueerden LATMiX op zeven zero-shot redeneerbenchmarks (ARC, HellaSwag, WinoGrande, PIQA, BoolQ, OpenBookQA) en de WikiText2-perplexiteitstak.

• Prestatiewinst: LATMiX presteerde consequent beter dan sterke baselines, waaronder:
  • Standaard MX-quantisatie (RTN, GPTQ).
  • Rotatiegebaseerde methoden (QuaRot, SpinQuant).
  • Block-diagonale rotatiemethoden (MR-GPTQ, OSTQuant).
  • Andere transformatiegebaseerde methoden (FlatQuant).
• Nauwkeurigheids Herstel: Bij MXFP4 (4-bit) quantisatie bereikte LATMiX tot 4% verbetering in gemiddeld nauwkeurigheids Herstel ten opzichte van de beste baselines. Bijvoorbeeld, op Llama-3.2-1B bereikte LATMiX-LU 93,88% herstel versus 93,07% voor MR-GPTQ.
• Robuustheid: De methode toonde consistente verbeteringen over verschillende modelgroottes (1B tot 32B) en quantisatiegranulariteiten (MXFP4 en MXINT4).
• Ablatiestudies:
  • Volledig vs. Blok: Volledige (tensorwijze) affiene transformaties presteerden aanzienlijk beter dan block-diagonale.
  • Affijn vs. Orthogonaal: Het leren van algemene affiene transformaties (met bias) was superieur aan het leren van alleen orthogonale (rotatie) matrices.
  • Verliesfunctie: KL-distillatieverlies bleek effectiever voor zero-shot generalisatie dan MSE- of Cross-Entropy-verliezen.

5. Betekenis

• Overbrugging van de Kloof: LATMiX slaagt erin de kloof te overbruggen tussen de behoefte aan agressieve uitbijteronderdrukking (via globale transformaties) en de structurele beperkingen van door moderne hardware ondersteunde MX-formaten.
• Hardware-efficiëntie: Door de transformaties in te bouwen, maakt LATMiX het gebruik mogelijk van krachtige, geleerde wiskundige bewerkingen zonder in te leveren op de inferentie-efficiëntiewinsten die MX-quantisatie biedt.
• Toekomstrichting: Het werk suggereert dat het beperken van transformaties tot rotaties of block-diagonale matrices suboptimaal is. Het opent de deur voor het leren van expressievere, data-afhankelijke transformaties die rekening houden met de specifieke dataformaten van AI-hardware van de volgende generatie.

Kortom, LATMiX toont aan dat door theoretisch de interactie te analyseren tussen activeringsverdelingen en MX-blokstructuren, en door flexibele affiene transformaties te leren, bijna verliesvrije quantisatie van LLM's mogelijk is bij extreem lage bitbreedtes (4-bit) zonder inferentie-overhead.