Technische Samenvatting: CLHA – Cross-Layer Hessian Aggregation voor het kwantiseren van Weight-Shared Mixture-of-Experts Transformers

1. Probleemstelling

Recente efficiënte Mixture-of-Experts (MoE) transformer-architecturen, zoals die in Megrez2, maken gebruik van Cross-Layer Weight Sharing (CLWS). In dit ontwerp hergebruiken aangrenzende transformer-lagen (bijv. $\ell$ en $\ell+1$) dezelfde gerouteerde expert-gewichtstensors, waardoor de parameter-voetafdruk wordt gehalveerd terwijl het aantal geactiveerde parameters per token gelijk blijft.

Echter, bestaande Post-Training Quantization (PTQ) methoden, inclusief GPTQ-achtige benaderingen, behandelen deze gedeelde gewichten alsof ze bij een enkele laag behoren. Ze kalibreren de kwantisatie met behulp van activatiestatistieken van slechts één "primaire" laag. Dit creëert een fundamentele mismatch:

• De representatie van de residual stream evolueert tussen lagen, wat betekent dat de tokenverdelingen die de gedeelde expert activeren in laag $\ell_a$ verschillen van die in laag $\ell_b$.
• Bestaande methoden negeren de activatiestatistieken van de secundaire laag, wat leidt tot een suboptimale minimalisatie van de reconstructiefout, met name in regimes met een lage bitbreedte (bijv. 2-bit).

2. Methodologie

2.1 Theoretische Fundering: Cross-Layer Hessian Aggregation

Het artikel formaliseert het doel voor het kwantiseren van een gewicht $W$ dat gedeeld wordt over lagen $\ell_1, \dots, \ell_n$. Het correcte doel is om de totale reconstructiefout over alle lagen te minimaliseren met behulp van het gewicht:
$$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$$
waarbij $H_{\ell_i}$ de Hessian (activatiecovariantie) van laag $\ell_i$ is, en $\alpha_i$ de routeringsfrequentie (proportie tokens) voor die laag vertegenwoordigt.

Propositie 1 stelt vast dat het minimaliseren van deze gecombineerde fout wiskundig equivalent is aan het toepassen van GPTQ met een gecombineerde Hessian:
$$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$$
Dit staat in contrast met Per-Layer Independent Calibration (PLIC), die de fout minimaliseert met behulp van slechts één enkele laags Hessian ($H_{\ell_a}$).

2.2 CLHA Algoritme

De auteurs stellen CLHA (Cross-Layer Hessian Aggregation) voor, een driefasen-kalibratiealgoritme:

1. Verzameling (Collection): Voor elke kalibratiesequentie worden per-laag Hessians en token-aantallen voor elke actieve expert geaccumuleerd.
2. Aggregatie (Aggregation): Voor elk gedeeld expert-paar $(\ell_a, \ell_b)$ wordt een gecombineerde Hessian berekend. Het artikel evalueert twee weegstrategieën:
   • Routing-Frequency Weighting ($\alpha = n_a / (n_a + n_b)$): Weegt Hessians op basis van het werkelijke aantal tokens dat elke laag ontvangt.
   • Equal Weighting ($\alpha = 0.5$): Fungeert als een regularisator, nuttig wanneer de kalibratiedata beperkt is en de schattingen van de secundaire-laag Hessian ruizig zijn.
3. Kwantisatie (Quantization): Pas standaard GPTQ kolom-sequentiële updates toe met behulp van de geaggregeerde $H_{comb}$.

2.3 CLHA-MP (Mixed Precision)

In het besef dat altijd-actieve gedeelde experts alle $N$ tokens ontvangen (in tegen tegenstelling tot gerouteerde experts die een fractie $K/E$ ontvangen), merkt het artikel op dat hun kwantisatiesensitiviteit $E/K$ keer hoger is. CLHA-MP wijst $b+1$ bits toe aan deze altijd-actieve gedeelde experts, terwijl gerouteerde experts op $b$ bits blijven, wat minimale opslagoverhead veroorzaakt (10–20% van de expert-parameters).

2.4 Implementatiecorrectie: Hook Ambiguity

Het artikel identificeert een kritieke engineering-valkuil: standaard PyTorch forward hooks die zijn gekoppeld aan een gedeelde ExpertFFN module, worden geactiveerd voor elke laag die die module aanroept. Dit veroorzaakt dat activaties van verschillende lagen worden gemengd in een enkele Hessian-accumulator, wat de per-laag attributie corrumpeert.

• Oplossing: De auteurs stellen een laag-bewuste callback-mechanisme voor die via de MoELayer.forward() signature wordt doorgegeven. Dit voorziet de callback expliciet van laag-indices, waardoor aparte Hessians voor $H_a$ en $H_b$ kunnen worden geaccumuleerd, zelfs wanneer de onderliggende gewichtsmodule gedeeld wordt.

3. Belangrijkste Bijdragen

1. Theoretische Propositie: Formalisering dat de optimale Hessian voor gedeelde gewichten de routeringsfrequentie-gewogen som is van per-laag Hessians.
2. CLHA Algoritme: Een praktische implementatie van gecombineerde Hessian-kalibratie binnen het GPTQ-framework.
3. CLHA-MP: Een mixed-precision extensie die een extra bit toewijst aan zeer gevoelige gedeelde experts.
4. Implementatie Fix: Identificatie en oplossing van het "hook ambiguity" probleem dat kalibratie in architecturen met gedeelde modules stilletjes corrumpeert.
5. Empirische Validatie: Demonstratie dat de winsten van CLHA correleren met cross-layer routing divergentie.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op een 34.5M-parameter CLWS-MoE die vanaf nul is getraind op WikiText-2.

• Prestaties bij 2-bit: Dit is het kritieke regime waar miskalibratie zich cumuleert.
  • CLHA-rt (routing-gewogen) verminderde de Perplexity (PPL) met 4,3% vergeleken met PLIC.
  • CLHA-eq (equal-gewogen) verminderde de PPL met 5,4%.
  • CLHA-MP bereikte een verbetering van 18,6% ten opzichte van PLIC.
  • Bij 4-bit en 3-bit lagen de verschillen binnen de ruis-marges (<1%).
• Kalibratie vs. Testfout: Opvallend genoeg behaalde PLIC een lagere Calibration Hessian Error (CHE) op de trainingsset omdat het direct voor die distributie optimaliseert. Echter, CLHA generaliseerde beter naar de testset, wat aangeeft dat de gecombineerde Hessian de ware distributie van de gedeelde gewichten beter vangt dan een enkelvoudige laag-view.
• Routing Sharpness Experiment: Door de router logits kunstmatig te verscherpen (het verhogen van de cross-layer divergentie), observeerden de auteurs:
  • De PLIC-prestaties verslechterden met 5% naarmate de divergentie toenam.
  • CLHA bleef stabiel (<1% degradatie).
  • Het prestatieverschil tussen CLHA en PLIC groeide van +4,3% (natuurlijke routing) naar +8,0% (hoge divergentie), wat de theoretische voorspelling bevestigt dat de winsten van CLHA schalen met cross-layer divergentie.

5. Betekenis en Claims

Het artikel claimt de eerste PTQ-analyse specifief voor cross-layer weight sharing in MoE-modellen te bieden. De significantie ligt in:

• Correctie van het Doel: Het aantonen dat het behandelen van gedeelde gewichten als single-layer entiteiten wiskundig suboptimaal is en dat Hessian aggregatie de noodzakelijke correctie is.
• Schaalbaarheid van Winsten: Het aantonen dat naarmate modellen groter worden en routingpatronen gespecialiseerder worden (toenemende cross-layer divergentie), de voordelen van CLHA waarschijnlijk zullen groeien, wat het essentieel maakt voor toekomstige productie CLWS-MoE modellen (bijv. Megrez2, DeepSeek, Qwen).
• Praktische Noodzaak: Het benadrukken dat de naïeve toepassing van bestaande PTQ-tools op CLWS-architecturen fundamenteel gebrekkig is vanwege hook ambiguity, en het bieden van de noodzakelijke fix voor elke toekomstige kalibratiemethode in dit domein.

De auteurs blijven bescheiden over de absolute prestaties, waarbij zij opmerken dat hun model klein is (34.5M) en voor een beperkt aantal stappen is getraind, maar benadrukken dat de relatieve verbeteringen ten opzichte van standaard kalibratie robuust en theoretisch onderbouwd zijn.