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

Das Paper stellt HEAPr vor, einen effizienten Pruning-Algorithmus für Mixture-of-Experts-Modelle, der durch die Zerlegung von Experten in atomare Einheiten und die Nutzung von Hessian-basierter zweiter Ordnung die Speicherkomplexität auf $O(d^2)$ reduziert und dabei bei hohen Pruning-Raten nahezu verlustfreie Kompression ermöglicht.

Das Wesentliche

Stell dir vor, ein großes Sprachmodell (wie ein KI-Assistent) ist wie ein riesiges, hochspezialisiertes Kochteam in einer Super-Küche.

In herkömmlichen Modellen arbeiten alle Köche gleichzeitig an jedem Gericht. Das ist langsam und braucht riesige Küchen.
MoE-Modelle (Mixture-of-Experts) sind cleverer: Es gibt ein "Türsteher-System" (den Router), das nur die wenigen Köche auswählt, die für ein bestimmtes Gericht wirklich gut sind. Wenn du über Mathematik fragst, werden nur die Mathematik-Köche aktiviert. Das macht das Denken schneller.

Aber hier liegt das Problem:
Obwohl nur wenige Köche arbeiten, müssen alle Köche im Gebäude anwesend sein und Platz einnehmen. Die Küche ist vollgestopft mit tausenden von Köchen, von denen die meisten gerade nur herumstehen. Das kostet enorm viel Speicherplatz (Gedächtnis) und macht es schwer, die KI auf normalen Geräten laufen zu lassen.

Bisherige Methoden, um diese Küche zu verkleinern, waren wie ein grober Kettensäge-Einsatz: Man hat ganze Köcheteams (Experten) komplett rausgeworfen oder zusammengefasst. Das war oft zu grob: Man hat vielleicht den einzigen Koch rausgeworfen, der das perfekte Gewürz kannte, und das Gericht schmeckte danach furchtbar.

Die Lösung: HEAPr – Der "Atomare" Küchen-Check

Die Forscher haben eine neue Methode namens HEAPr entwickelt. Hier ist die einfache Erklärung, wie sie funktioniert, mit ein paar kreativen Vergleichen:

1. Die Zerlegung: Vom Kochteam zum einzelnen Gewürz

Statt ganze Köcheteams zu entfernen, schauen sich die Forscher die Köche noch genauer an. Sie stellen fest: Jeder "Experte" (Koch) besteht eigentlich aus vielen kleinen, unteilbaren Einheiten – nennen wir sie "Atomare Experten".

• Vergleich: Stell dir vor, ein Koch ist nicht nur eine Person, sondern ein Set aus verschiedenen Werkzeugen: Ein Messer, ein Schneidebrett und ein Gewürzstreuer.
• HEAPr entfernt nicht den ganzen Koch, sondern nur das unnötige Werkzeug (z. B. ein stumpfes Messer), das niemand braucht. Der Rest des Kochs bleibt erhalten und funktioniert weiter. Das ist viel präziser.

2. Die Intelligenz: Der "Zweite Blick" (Hessische Matrix)

Wie wissen sie, welches Werkzeug sie wegwerfen dürfen, ohne dass das Essen verdirbt?
Frühere Methoden haben nur geschaut: "Wie oft wird dieses Werkzeug benutzt?" (Das ist wie zu sagen: "Wir brauchen das Messer nicht, weil es heute niemand benutzt hat"). Das ist oft falsch, weil es vielleicht nur zufällig heute nicht gebraucht wurde.

HEAPr macht etwas viel Klügeres: Es nutzt eine mathematische Vorhersage, die wie ein "Zweiter Blick" funktioniert.

• Die Analogie: Stell dir vor, du willst testen, ob du das Gewürz aus dem Rezept streichen kannst. Ein einfacher Test wäre, es einfach wegzulassen und zu schmecken.
• HEAPr macht aber einen "Simulationstest": Es berechnet theoretisch, wie sehr der Geschmack (die Leistung der KI) leiden würde, wenn man das Gewürz weglässt. Es schaut nicht nur auf die Oberfläche, sondern berechnet die Auswirkung auf die Struktur des Ganzen.
• Das ist wie ein genialer Küchenchef, der sagt: "Wenn wir das Salz weglassen, wird das Gericht nicht nur etwas salzlos, sondern die ganze Sauce bricht zusammen." Oder: "Wenn wir das Paprika weglassen, merkt das niemand."

3. Der Trick: Die "Flüster-Technik" (Effizienz)

Das Problem bei diesem "Zweiten Blick" ist normalerweise, dass die Berechnungen so komplex sind, als müsste man jede einzelne Kombination von Zutaten in jeder möglichen Küche simulieren. Das würde Jahre dauern und den Computer sprengen.

HEAPr hat einen genialen Trick gefunden, um das zu vereinfachen:

• Statt zu berechnen, wie sich jedes einzelne Werkzeug auf jedes andere auswirkt, schauen sie nur auf das Ergebnis (den Teller mit dem Essen).
• Vergleich: Statt zu analysieren, wie sich das Messer auf das Schneidebrett auswirkt, schauen sie nur auf den fertigen Teller. Wenn der Teller perfekt aussieht, ist alles in Ordnung.
• Durch diesen Trick wird die Berechnung unmengenmal schneller und braucht weniger Speicher. Sie brauchen dafür nur einen kurzen Testlauf mit ein paar Beispielsätzen (eine "Kalibrierungs-Küche"), um zu entscheiden, was weg kann.

Das Ergebnis: Eine schlankere Küche, gleicher Geschmack

Was bringt das alles?

• Platzsparend: Man kann bis zu 20–25 % der "Köche" (bzw. ihrer Werkzeuge) entfernen, ohne dass die KI schlechter wird. Es ist fast wie ein "verlustfreies" Komprimieren.
• Schneller: Da weniger Werkzeuge bewegt werden müssen, ist die KI auch schneller (bis zu 20 % weniger Rechenarbeit).
• Robust: Die Methode funktioniert bei verschiedenen großen KI-Modellen (wie DeepSeek oder Qwen) gleich gut.

Zusammenfassend:
HEAPr ist wie ein hochmoderner Küchen-Optimierer. Statt ganze Teams zu feuern (was oft katastrophal ist), analysiert es präzise, welche kleinen Werkzeuge in jedem Team überflüssig sind, und entfernt nur diese. Und das alles so clever berechnet, dass es schnell geht und die Qualität der KI (den Geschmack des Gerichts) perfekt erhält.

Technische Zusammenfassung

1. Problemstellung

Mixture-of-Experts (MoE)-Architekturen in großen Sprachmodellen (LLMs) bieten zwar eine hervorragende Leistung und reduzieren die Inferenzkosten durch spärliche Aktivierung, stellen jedoch aufgrund der enormen Gesamtzahl an Parametern eine massive Herausforderung für den Speicherbedarf dar. Alle Parameter müssen im GPU-Speicher gehalten werden, auch wenn nur ein Bruchteil aktiviert wird.

Bestehende Pruning-Methoden (Beschneidungsalgorithmen) stoßen hier an Grenzen:

• Feingranulares Pruning (z. B. Gewichts- oder Neuronen-Pruning) erhält oft die Genauigkeit, bietet aber aufgrund hardwarebedingter Ineffizienzen nur geringe Beschleunigung.
• Grobgranulares Pruning auf Ebene ganzer Experten (Expert-Level) ermöglicht direkte Beschleunigung und Speicherreduktion, führt jedoch oft zu signifikanten Genauigkeitsverlusten, da ganze Expertenkomponenten entfernt werden, die möglicherweise wertvolles, komplementäres Wissen enthalten.
• Bestehende Ansätze wie Experten-Merging (Zusammenführung) oder Dekomposition sind oft rechenintensiv, instabil oder führen zu suboptimalen Ergebnissen.

Das Ziel ist es, eine Methode zu finden, die feiner als ganzer Experten, aber effizienter als einzelne Gewichte ist, um Speicher und Rechenlast zu reduzieren, ohne die Modellleistung zu beeinträchtigen.

2. Methodik: HEAPr

Die Autoren stellen HEAPr (Hessian-based Efficient Atomic Expert Pruning in Output Space) vor, einen neuen Algorithmus, der Experten in kleinere, unteilbare Einheiten zerlegt.

A. Konzept der „Atomic Experts"

Anstatt ganze Experten zu entfernen, wird jeder Experte in Atomic Experts zerlegt.

• Ein Experte $E_i$ besteht aus den Matrizen $W_{up}$, $W_{gate}$ und $W_{down}$.
• Ein Atomic Expert $e^{(j)}_i$ wird definiert durch die $j$-te Spalte von $W_{up}$ und $W_{gate}$ sowie die entsprechende $j$-te Zeile von $W_{down}$.
• Der Output eines vollständigen Experten ist die Summe der Outputs seiner Atomic Experts.
• Durch das Entfernen von Atomic Experts wird die Struktur der verbleibenden Komponenten nicht gestört, was eine flexible und präzise Reduktion ermöglicht.

B. Wichtigkeitsmessung basierend auf der Optimal Brain Surgeon (OBS) Theorie

HEAPr nutzt die OBS-Theorie, die den Einfluss des Gewichts-Prunings auf die Verlustfunktion durch eine Taylor-Entwicklung zweiter Ordnung approximiert. Das Ziel ist es, Parameter zu identifizieren, deren Entfernung den geringsten Anstieg des Verlusts verursacht.

Die Herausforderung bei MoE-Modellen ist die Berechnung der Hesse-Matrix (zweite Ableitung der Verlustfunktion), deren Speicherbedarf auf Expertenebene bei $O((3d_{model} \cdot d_{inter})^2)$ liegt – was für große Modelle untragbar ist.

HEAPr löst dies durch zwei Optimierungen:

1. Entkopplung der Parameter: Es wird gezeigt, dass die Parameter verschiedener Atomic Experts innerhalb eines Experten entkoppelt sind (die Kreuz-Hessen sind null). Dies reduziert die Komplexität bereits auf $O((3d_{model})^2 \cdot d_{inter})$.
2. Transformation in den Output-Raum: Statt die Hesse-Matrix im Parameterraum zu berechnen, wird die Analyse in den Output-Raum der Atomic Experts verschoben.
   • Anstatt die Bedingung zu stellen, dass der Output für alle Eingaben null sein muss, wird die Bedingung pro Token betrachtet.
   • Es wird gezeigt, dass die Fisher-Information (die erwartete Hesse-Matrix) effizient berechnet werden kann.
   • Ein entscheidender Befund ist, dass alle Atomic Experts innerhalb eines Experten denselben Gradienten bezüglich des Gesamtverlusts teilen.
   • Ergebnis: Die Speicherkomplexität sinkt drastisch von $O(d^4)$ auf $O(d^2)$ (wobei $d$ die Modell-Dimensionalität ist).

C. Algorithmus

HEAPr benötigt nur zwei Forward-Passes und einen Backward-Pass auf einem kleinen Kalibrierungsdatensatz:

1. Gradient-Covariance-Schätzung: Berechnung einer geteilten Gradient-Kovarianzmatrix für jeden Experten basierend auf den Tokens, die zu diesem Experten routen.
2. Wichtigkeitsberechnung: Berechnung des Scores für jeden Atomic Expert durch Kombination seines Outputs mit der geteilten Kovarianzmatrix.
3. Globales Ranking: Alle Atomic Experts im gesamten Modell werden nach ihrem Wichtigkeits-Score sortiert, und die unwichtigsten $r\%$ werden entfernt.

3. Hauptbeiträge

• Neue Granularität: Einführung des Konzepts der „Atomic Experts" für MoE-Modelle, was eine flexiblere und präzisere Pruning-Strategie als das Entfernen ganzer Experten ermöglicht.
• Effiziente Second-Order Approximation: Entwicklung einer Methode, die die Second-Order-Information von Expert-Parametern auf Atomic-Expert-Parameter und schließlich auf den Output-Raum transformiert. Dies reduziert die Speicherkomplexität von $O(d^4)$ auf $O(d^2)$.
• Hohe Effizienz: Der Algorithmus ist extrem ressourcenschonend und benötigt nur minimale Inferenzschritte auf einem kleinen Kalibrierungsset, ohne Nachtraining (Retraining) zu erfordern.
• Globale Konsistenz: Die Methode liefert einen global konsistenten Wichtigkeits-Score, der es erlaubt, Experten über alle Schichten hinweg fair zu vergleichen und zu beschneiden.

4. Ergebnisse

HEAPr wurde an sieben Zero-Shot-Aufgaben auf verschiedenen MoE-Modellen getestet (DeepSeekMoE-16B-Base, Qwen1.5-MoE-A2.7B-Chat, Qwen2-57B-A14B, Qwen3-30B-A3B).

• Nahezu verlustfreie Kompression: Bei Pruning-Raten von 20% bis 25% erreicht HEAPr in den meisten Modellen eine Leistung, die fast identisch mit dem Originalmodell ist (nahezu verlustfrei).
• Hohe Reduktion: Auf dem Qwen2-57B-A14B-Modell bleibt die Leistung selbst bei 40% Pruning-Rate fast unverändert. Auf dem neuesten Qwen3-30B-A3B sinkt die durchschnittliche Genauigkeit bei 25% Pruning nur um 0,03 Punkte.
• Vergleich mit SOTA: HEAPr übertrifft bestehende State-of-the-Art-Methoden (wie NAEE, MoE-I2, MC-SMoE, D2-MoE, Sub-MoE) deutlich in Bezug auf Genauigkeit bei gleicher Speicher- und Rechenreduktion.
• FLOPs-Reduktion: Der Algorithmus reduziert die FLOPs (Floating Point Operations) um nahezu 20% bei den genannten Pruning-Raten, was eine echte Beschleunigung darstellt (im Gegensatz zu reinem Gewichts-Pruning).
• Robustheit: Die Methode ist robust gegenüber der Wahl des Kalibrierungsdatensatzes (WikiText-2 vs. C4) und skaliert gut mit der Größe des Kalibrierungssets.

5. Bedeutung

HEAPr adressiert eine kritische Lücke in der Effizienzsteigerung von MoE-Modellen. Während frühere Ansätze entweder zu grob (Verlust von Wissen) oder zu ineffizient (kein Hardware-Benefit) waren, bietet HEAPr einen „Sweet Spot":

• Es ermöglicht eine feingranulare Reduktion des Speicherbedarfs und der Rechenlast, ohne die Architektur neu zu erfinden oder aufwendiges Nachtraining zu benötigen.
• Die Reduktion der Komplexität der Hesse-Matrix-Berechnung macht Second-Order-Pruning für große MoE-Modelle erstmals praktikabel.
• Die Ergebnisse zeigen, dass MoE-Modelle eine signifikante Redundanz auf der Ebene der Atomic Experts aufweisen, die entfernt werden kann, ohne die Kernfähigkeiten des Modells zu beeinträchtigen.

Dieser Ansatz ebnet den Weg für den effizienteren Einsatz von MoE-Modellen auf ressourcenbeschränkten Geräten und in Produktionsumgebungen, wo Speicher und Latenz kritische Faktoren sind.