HEAPr: Hessian-based Efficient Atomic Expert Pruning in Output Space

Dit artikel introduceert HEAPr, een nieuwe pruning-methode voor Mixture-of-Experts-modellen die door het gebruik van Hessiaanse informatie op atomaire expert-niveau een fijnere en efficiëntere compressie mogelijk maakt met minimaal verlies aan nauwkeurigheid.

De kern

Stel je voor dat een groot taalmodel (zoals een slimme chatbot) een enorm kantoor is met duizenden experts. In de moderne "MoE" (Mixture of Experts) modellen werken deze experts niet allemaal tegelijk. In plaats daarvan is er een slimme manager (de router) die bij elke vraag kijkt: "Wie is hier de beste persoon voor deze taak?" en slechts een paar experts laat werken. Dit maakt het model heel snel en slim.

Maar er is een groot probleem: alle experts moeten op de harde schijf van de computer staan, zelfs als ze niet gebruikt worden. Dit kost enorm veel geheugen en maakt het moeilijk om deze modellen op gewone laptops of telefoons te draaien.

Tot nu toe probeerden mensen dit op te lossen door hele experts weg te gooien. Maar dat is als een bedrijf dat besluit: "We hebben 100 mensen, laten we er 20 volledig ontslaan." Het risico is groot dat je iemand ontslaat die net de perfecte kennis had voor een specifieke, zeldzame vraag. Het resultaat: de chatbot wordt slomer en maakt meer fouten.

HEAPr is een nieuwe, slimme manier om dit op te lossen. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Atomaire Expert": Van Teamleden naar Bouwstenen

Stel je een expert voor als een gereedschapskist. In de oude methode gooi je de hele kist weg als je denkt dat hij niet vaak gebruikt wordt.
HEAPr kijkt echter dieper. Het beseft dat elke "kist" (expert) eigenlijk bestaat uit kleinere, onbreekbare onderdelen: atomaire experts.

• Voorbeeld: Een expert is als een kok die een gerecht maakt. De atomaire experts zijn de specifieke ingrediënten of handelingen: het snijden van de ui, het bakken van het vlees, het kruiden van de saus.
  HEAPr zegt: "We hoeven niet de hele kok te ontslaan. Laten we gewoon de handeling 'snijden van de ui' verwijderen als die niet nodig is voor dit specifieke gerecht." Hierdoor kun je veel meer wegdoen zonder dat het hele team instort.

2. De "Hersenen" van de Operatie: De Tweede Graads Analyse

Hoe weet je nu welke "snij-handeling" je kunt verwijderen zonder dat het gerecht mislukt?
De auteurs gebruiken een wiskundige techniek die lijkt op de "Optimale Hersenchirurg" (een theorie uit de jaren '90).

• De oude manier: Ze keken naar de "eerste graad" (hoe hard iemand werkt).
• De nieuwe manier (HEAPr): Ze kijken naar de "tweede graad". Dit is alsof ze niet alleen kijken naar wie er werkt, maar naar hoe gevoelig het hele systeem is als je iemand weghaalt. Ze vragen zich af: "Als ik deze specifieke handeling stop, hoeveel gaat de kwaliteit van het eindresultaat (de fouten) stijgen?"

3. Het Geniale Trucje: Van Ingewikkeld naar Eenvoudig

Het probleem met deze "tweede graads" berekening is dat het normaal gesproken zo veel rekenkracht en geheugen kost dat het onmogelijk is voor zulke grote modellen. Het zou zijn alsof je een hele bibliotheek moet lezen om één zin te controleren.

HEAPr gebruikt een slimme truc:

1. Decompositie: Omdat de atomaire experts onafhankelijk zijn, hoeven ze niet allemaal tegelijk berekend te worden.
2. De Uitvoer-Space: In plaats van te kijken naar de ingewikkelde interne berekeningen van de experts, kijken ze alleen naar het eindresultaat (de output).

• Analogie: In plaats van te analyseren hoe elke schroef in de motor van een auto is gemaakt (wat enorm veel tijd kost), kijken ze gewoon naar hoe de auto rijdt als je een schroef losdraait. Als de auto nog steeds soepel rijdt, was die schroef niet essentieel.

Dit vermindert de benodigde rekenkracht en geheugen met een factor van duizenden, waardoor het plotseling haalbaar wordt.

4. Het Resultaat: Meer ruimte, zelfde slimheid

De resultaten van de test zijn indrukwekkend:

• Je kunt 20% tot 25% van de "atomaire experts" verwijderen (zoals het weghalen van overbodige ingrediënten uit de keuken).
• De chatbot wordt 20% sneller en neemt 20% minder geheugen in beslag.
• Het slimme is: de chatbot maakt bijna geen enkele fout meer dan voorheen. Het is alsof je een zware jas uittrekt, maar je bent nog steeds even warm en comfortabel.

Kort samengevat:
HEAPr is als een super-efficiënte manager die niet hele afdelingen ontslaat, maar precies weet welke kleine taken overbodig zijn. Door slimme wiskunde toe te passen, kan hij deze taken verwijderen zonder dat het bedrijf (het AI-model) zijn kennis verliest. Hierdoor kunnen we deze slimme chatbots eindelijk op gewone apparaten laten draaien.

Technische samenvatting

Probleemstelling

Mixture-of-Experts (MoE) architecturen in grote taalmodellen (LLM's) bieden uitstekende prestaties en lagere inferentiekosten vergeleken met dichte modellen, omdat ze slechts een fractie van de parameters activeren. Echter, alle parameters moeten toch in het geheugen worden opgeslagen, wat leidt tot prohibitieve opslagvereisten (bijv. DeepSeek-V3 activeert 37B parameters maar vereist 671B geheugen).

Bestaande compressiemethoden hebben te maken met een fundamenteel compromis:

• Fijne granulariteit (bijv. gewichtspruning): Behoudt nauwkeurigheid maar biedt weinig hardware-versnelling.
• Grove granulariteit (Expert-level pruning): Biedt directe versnelling en geheugenreductie, maar leidt vaak tot aanzienlijke nauwkeurigheidsverlies.
  • Expert dropping (weglaten van experts) kan waardevolle kennis verliezen.
  • Expert merging (samenvoegen van experts) is vaak instabiel en leidt tot parameterconflicten.
  • Bestaande decompositiemethoden zijn computatief zwaar en verliezen nog steeds nauwkeurigheid.

Er is behoefte aan een methode die flexibeler is dan expert-level pruning, maar wel direct leidt tot versnelling en geheugenreductie zonder significante kwaliteitsverlies.

Methodologie: HEAPr

De auteurs introduceren HEAPr (Hessian-based Efficient Atomic Expert Pruning in Output Space), een algoritme dat experts decomposeert in kleinere, ondeelbare eenheden genaamd Atomic Experts.

1. Decompositie naar Atomic Experts:
In plaats van een hele expert te verwijderen, wordt elke expert $E_i$ gezien als een lineaire combinatie van $d_{inter}$ atomic experts. Een atomic expert wordt gedefinieerd door het groeperen van de relevante kolommen van de matrices $W_{up}$ en $W_{gate}$, en de corresponderende rij van $W_{down}$. Het verwijderen van een atomic expert verwijdert direct een deel van de berekeningskosten binnen die expert.

2. Importance Scoring via Second-Order Information (OBS):
Om te bepalen welke atomic experts het minst belangrijk zijn, baseert HEAPr zich op de Optimal Brain Surgeon (OBS) theorie. De theorie benadert de impact van het verwijderen van een parameter op de verliesfunctie ($\Delta \mathcal{L}$) via een Taylor-expansie van de tweede orde, waarbij de Hessian-matrix ($H$) wordt gebruikt.

3. Optimalisaties voor Efficiëntie:
Het direct berekenen van de Hessian voor experts is te duur ($O((3d_{model} \cdot d_{inter})^2)$). HEAPr lost dit op met twee cruciale optimalisaties:

• Decoupling van parameters: De parameters van verschillende atomic experts binnen dezelfde expert zijn onafhankelijk. De kruisderivaten (cross-Hessians) zijn nul. Dit reduceert de complexiteit van het sommeren van Hessians aanzienlijk.
• Verschuiving naar Output Space: In plaats van de constraints in de parameter-ruimte te analyseren, verschuift HEAPr de analyse naar de output-ruimte van de atomic expert.
  • Door de constraints te formuleren op de output ($e_P(x) = 0$) en gebruik te maken van de Fisher Information Matrix (die theoretisch equivalent is aan de verwachte Hessian maar efficiënter te berekenen), wordt de complexiteit verder gereduceerd.
  • Een belangrijke observatie is dat alle atomic experts binnen dezelfde expert dezelfde gradienten delen ten opzichte van de verliesfunctie. Hierdoor hoeft er slechts één covariantiematrix per expert te worden opgeslagen.
  • De ruimtelijke complexiteit daalt hierdoor van $O(d^4)$ naar $O(d^2)$.

4. Het Algoritme:
HEAPr vereist slechts twee forward passes en één backward pass op een kleine kalibratieset:

1. Gedeelde Gradient Covariantie: Bereken de gradient van de verliesfunctie ten opzichte van de expert-output en bouw de covariantiematrix $\bar{G}_i$ op.
2. Importance Berekening: Bereken voor elke atomic expert $e_k$ de score $s_k$ door de output van die specifieke atomic expert te projecteren op de gedeelde covariantiematrix:
   $$s_k \approx \mathbb{E}_{x \sim D} \left[ \frac{1}{2} e_k(x)^\top \bar{G}_i e_k(x) \right]$$
3. Global Ranking: Alle atomic experts in het model worden globaal gerangschikt op basis van hun score $s_k$ (kleinere score = minder impact op het verlies = prioriteit voor pruning). De laagst scorende $r\%$ worden verwijderd.

Kernbijdragen

1. Atomic Expert Decompositie: Introductie van een nieuwe eenheid voor pruning die flexibeler is dan hele experts, waardoor directe FLOPs-reductie mogelijk is.
2. Efficiënte Second-Order Schatting: Een innovatieve methode om second-order informatie (Hessian/Fisher) te transformeren van parameter-ruimte naar output-ruimte, waardoor de opslagcomplexiteit drastisch wordt verlaagd ($O(d^2)$) zonder retraining.
3. HEAPr Algoritme: Een schaalbaar algoritme dat de belangrijkheid van alle atomic experts berekent met minimale rekentijd (slechts 2 forward, 1 backward pass).
4. State-of-the-Art Resultaten: Uitgebreide validatie op diverse moderne MoE-modellen (DeepSeek, Qwen-familie) die aantonen dat HEAPr superieur is aan bestaande expert-level pruning en merging technieken.

Resultaten

Experimenten zijn uitgevoerd op modellen zoals DeepSeekMoE-16B-Base, Qwen1.5-MoE-A2.7B-Chat, Qwen2-57B-A14B en Qwen3-30B-A3B op zeven zero-shot taken.

• Bijna verliesloze compressie: HEAPr bereikt bijna geen nauwkeurigheidsverlies bij pruning ratios van 20% tot 25% op de meeste modellen.
  • Bijvoorbeeld: Op Qwen2-57B-A14B behoudt het de prestaties van het originele model zelfs bij 40% pruning.
  • Op Qwen3-30B-A3B daalt de gemiddelde nauwkeurigheid slechts met 0.03 bij 25% pruning.
• Efficiëntie: De methode reduceert de FLOPs met bijna 20% bij de genoemde pruning ratios.
• Vergelijking: HEAPr presteert significant beter dan concurrenten zoals NAEE, MoE-I2, MC-SMoE, HC-SMoE, Sub-MoE en D2-MoE.
• Ablatie Studies:
  • Global vs. Layer-wise: Global pruning (over het hele model) werkt beter dan layer-wise pruning, wat aantoont dat de importance-score consistent is over verschillende lagen.
  • Granulariteit: Pruning op atomic expert niveau levert echte FLOPs-reductie op, terwijl expert-level pruning (alleen hele experts verwijderen) vaak geen versnelling geeft omdat de interne dimensies van de overgebleven experts gelijk blijven.
  • Kalibratie: De methode is robuust en presteert goed met kleine kalibratiesets (128 sequences) en is ongevoelig voor de keuze van de dataset (WikiText-2 vs C4).

Betekenis

HEAPr biedt een doorbraak in het comprimeren van MoE-modellen. Het lost het probleem op dat grove pruning (expert-level) vaak te veel nauwkeurigheid kost, terwijl fijne pruning (gewichten) geen hardware-versnelling biedt. Door over te schakelen naar atomic experts en gebruik te maken van een efficiënte second-order benadering in de output-ruimte, maakt HEAPr het mogelijk om MoE-modellen aanzienlijk lichter en sneller te maken zonder de prestaties te offeren. Dit is cruciaal voor de praktische implementatie van grote MoE-modellen op apparaten met beperkte resources.