Depth Registers Unlock W4A4 on SwiGLU: A Reader/Generator Decomposition

Probleemstelling

Post-training kwantisatie naar 4-bit gewichten en 4-bit activaties (W4A4) is een veelbelovende techniek om de inferentie-efficiëntie van Large Language Models (LLMs) te verhogen. Echter, bij SwiGLU-gebaseerde modellen (zoals LLaMA, Mistral, Qwen) leidt naïeve "round-to-nearest" W4A4 kwantisatie vaak tot een catastrofale degradatie van de prestaties.

• Observatie: In een gecontroleerd experiment met een 300M-parameter decoder-only model getraind op 5 miljard tokens, daalde de validatie-perplexiteit (PPL) van 23,6 (FP16) naar 1727 bij naïeve W4A4.
• Oorzaak: Bestaande methoden zoals SmoothQuant, QuaRot en AWQ proberen dit op te lossen door schaling of rotatie toe te passen, maar het is onduidelijk welke specifieke invoer-activatiepunten in het netwerk de fouten veroorzaken en of deze toegankelijk zijn via de "residual axis" (de hoofdstroom van de residual stream).

Methodologie: Reader/Generator Decompositie

De auteurs introduceren een operator-niveau decompositie van de vijf trainbare lineaire lagen binnen een SwiGLU-blok, gebaseerd op de bron van hun invoer:

1. Readers (Lezers):
   • Lagen: qkv, w1, w3.
   • Invoer: De residual stream na RMSNorm.
   • Eigenschap: Hun invoermagnitude wordt direct begrensd door interventies die de residual stream reguleren.
2. Generators (Generatoren):
   • Lagen: o_proj (uit de attention-mechanisme) en w2 (de tweede laag van de SwiGLU feed-forward).
   • Invoer: Interne berekeningen van het blok.
     • o_proj: Een lineaire combinatie van waarden (afhankelijk van softmax-weights).
     • w2: Een bilineair product (silu(w1(x)) ⊙ w3(x)).
   • Eigenschap: De auteurs tonen wiskundig aan dat een interventie op de residual-as de invoer van w2 slechts losjes kan begrenzen via een triviaal product van grenswaarden. Een zeldzame grote waarde in één factor vermenigvuldigd met een typische waarde in de andere kan een extreme outlier creëren die niet aanwezig was in de individuele factoren.

De Interventie: Depth Registers + Sink Loss (DR+sink)
Om dit te testen zonder post-hoc kalibratie of rotatie, ontwikkelen de auteurs een trainingsmethode:

• Ze partitioneren de residual-kanaal-as in een semantisch gebied en een "register" gebied (diepe registers).
• Ze passen een hinge-loss toe op de magnitude van de register-kanalen, waardoor deze een "afvoer" (sink) vormen voor extreme waarden die anders op de semantische kanalen zouden landen.
• Dit is een trainingstijd-probe om te zien of het beperken van de residual-magnitude de fouten vermindert.

Belangrijkste Bijdragen

1. Operator-niveau Decompositie: Een theoretisch en empirisch onderscheid tussen "Readers" (gebonden aan de residual) en "Generators" (interne producten), met een wiskundige onderbouwing waarom residual-interventies de bilineaire tail van w2 niet effectief kunnen onderdrukken.
2. DR+sink Methode: Een trainingstechniek die de maximale magnitude van register-kanalen begrenst, wat leidt tot een drastische reductie van outliers bij de "Readers".
3. Skip-Ablatie Budget: Een nieuwe metriek (∆remove(S)) die de schade van W4A4 kwantificeert door specifieke lagen FP16 te houden terwijl de rest gekwantiseerd is. Dit onthult dat de schade voornamelijk bij de generators ligt zodra de readers gerepareerd zijn.
4. Validatie van Rotatiebeperkingen: Een directe test die aantoont dat orthogonale rotaties (zoals in QuaRot) de bilineaire tail van w2 niet kunnen onderdrukken, in lijn met de theoretische voorspelling.

Resultaten

De experimenten werden uitgevoerd op een 300M-parameter model met 5B tokens.

• Prestatieverbetering:
  • Naïeve W4A4: PPL daalt van 23,6 naar 1727.
  • Met DR+sink: PPL daalt naar 119 (een verbetering van ~14x) bij een gelijkblijvende FP16 PPL (23,7) en gelijke zero-shot capaciteit.
  • Combinatie met SmoothQuant: DR+sink gecombineerd met SmoothQuant bereikt een PPL van 39,9.
• Kurtosis Analyse:
  • DR+sink verlaagt de kurtosis (uitbijters) van de Readers drastisch (bijv. qkv van 35,9 naar 2,85, bijna Gaussiaans).
  • De kurtosis van de Generators blijft onveranderd hoog (bijv. w2 blijft rond 1860-1921).
• Schadeverdeling (Skip-Ablatie):
  • Zodra alle Readers FP16 zijn gehouden, blijft er een "residue" van ~0,28 nats (PPL ~31) over, ongeacht of DR+sink is gebruikt. Dit betekent dat de residual-as geen oplossing biedt voor de resterende schade.
  • De schade wordt gedomineerd door de generator w2 (∆remove(w2) is ~0,69-0,88 nats), terwijl o_proj een verwaarloosbare bijdrage levert zodra de residual schoon is.
• Post-hoc Vergelijking:
  • Per-Linear QuaRot (rotatie zonder residual-fusie) presteert bijna identiek aan DR+sink op de reader-as (PPL ~25,4), maar lost het w2-probleem niet op.
  • Full QuaRot (met online rotatie op w2-invoer) verbetert de situatie niet significant ten opzichte van Per-Linear QuaRot (PPL ~26,0), wat bevestigt dat orthogonale rotatie de bilineaire tail niet effectief onderdrukt.

Significantie en Conclusie

Dit paper levert een fundamenteel inzicht in waarom W4A4 kwantisatie faalt bij SwiGLU-modellen:

1. De "Reader/Generator" Split: De meeste kwantisatiefouten komen voort uit de "Readers", die oplosbaar zijn via residual-axis beheer (zoals DR+sink of Per-Linear rotatie).
2. De w2-Barrière: De resterende fouten (~2 PPL gap met FP16) worden veroorzaakt door de bilineaire aard van de w2-laag. Deze is inherent moeilijk te beheersen met residual-axis interventies of orthogonale rotaties.
3. Toekomstige Richting: Omdat rotatie en residual-beperking de w2-tail niet volledig kunnen oplossen, suggereren de auteurs dat toekomstige oplossingen gericht moeten zijn op mixed-precision voor w2, Quantization-Aware Training (QAT), of het afsnijden van het product (product clipping), in plaats van alleen te vertrouwen op post-training kalibratie of rotatie.

De studie is beperkt tot een enkel seed en een 300M-parameter model, maar de mechanistische inzichten en de decompositie van de schade bieden een nieuwe richting voor het optimaliseren van kwantisatie in moderne LLM-architecturen.