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

Le papier présente HEAPr, une méthode de pruning innovante pour les modèles de type Mixture-of-Experts qui, en décomposant les experts en unités atomiques et en exploitant l'information de Hessian dans l'espace de sortie, permet une compression quasi sans perte avec une complexité réduite et des performances supérieures aux méthodes existantes.

L'essentiel

Imaginez que vous avez un groupe de 100 experts dans une entreprise (un modèle d'intelligence artificielle) qui travaillent ensemble pour résoudre des problèmes. C'est ce qu'on appelle un modèle "MoE" (Mélange d'Experts).

Le problème ? Même si seuls quelques experts travaillent à la fois pour chaque tâche, vous devez payer le salaire et louer le bureau pour les 100 experts en permanence, même ceux qui ne font rien. Cela coûte une fortune en mémoire et en énergie.

Les méthodes actuelles pour réduire les coûts sont comme une coupe de cheveux grossière : on renvoie tout un expert (un bureau entier) si on pense qu'il n'est pas utile. Mais le problème, c'est que cet expert avait peut-être des compétences très précieuses qu'on ne peut pas remplacer. Résultat : l'entreprise perd en performance.

Voici comment HEAPr change la donne, expliqué simplement :

1. L'idée géniale : Découper l'expert en "Atomes"

Au lieu de renvoyer un expert entier, HEAPr regarde à l'intérieur de chaque expert. Il découvre que chaque expert est en fait composé de petits blocs indivisibles, qu'ils appellent des "experts atomiques".

• L'analogie : Imaginez que votre expert est un couteau suisse. Les méthodes anciennes jettent tout le couteau s'il est un peu rouillé. HEAPr, lui, retire uniquement la petite lame qui ne sert plus, tout en gardant le tournevis, le couteau et la pince.
• Le résultat : On peut enlever des milliers de ces "petites lames" inutiles sans casser le couteau suisse. On garde la structure, mais on allège considérablement le poids.

2. Comment savoir quoi jeter ? (La "Radiographie" du cerveau)

Comment savoir quelle "petite lame" est inutile sans tout casser ? HEAPr utilise une technique mathématique très puissante appelée OBS (Optimal Brain Surgeon).

• L'analogie : Imaginez un neurochirurgien qui doit retirer une tumeur. Il ne coupe pas au hasard. Il utilise une radiographie de haute précision (l'information du "second ordre" ou Hessian) pour voir exactement quelle partie du cerveau, si on la retire, va causer le moins de dégâts possible.
• Le problème habituel : Faire cette "radiographie" sur un modèle géant est normalement impossible, c'est comme essayer de scanner un océan avec un microscope : ça prend trop de temps et de place.
• La solution de HEAPr : Ils ont trouvé un raccourci génial. Au lieu de scanner les paramètres internes (les neurones), ils regardent simplement la sortie (la réponse donnée par l'expert).
  • C'est comme si, au lieu de disséquer le cerveau du patient, on observait simplement ses réponses à des questions. Si une réponse ne change pas le résultat final, on sait que la partie du cerveau qui l'a produite est redondante.
  • Cela réduit la complexité de calcul de manière spectaculaire (de $O(d^4)$ à $O(d^2)$). C'est passer de l'exploration de l'Univers entier à l'exploration de notre système solaire.

3. Le résultat : Une compression "sans douleur"

Grâce à cette méthode, HEAPr peut :

• Renvoyer 20% à 25% des "petites lames" (les experts atomiques) sans que l'entreprise (le modèle) perde sa capacité à travailler. C'est ce qu'ils appellent une compression "quasi sans perte".
• Économiser du temps de calcul (FLOPs) d'environ 20%.
• Fonctionner très vite : Il ne faut que deux passages en avant et un seul en arrière sur un petit échantillon de données pour prendre toutes les décisions. Pas besoin de réentraîner tout le modèle pendant des semaines.

En résumé

HEAPr est comme un architecte de précision qui rénové un gratte-ciel (le modèle d'IA).

• Les méthodes précédentes : "On démolit tout l'étage 5 !" (Trop brutal, le bâtiment s'effondre un peu).
• HEAPr : "On retire les 20% de briques inutiles dans chaque pièce, on allège la structure, et le bâtiment tient toujours aussi bien, voire mieux, tout en étant plus léger et moins cher à entretenir."

C'est une avancée majeure pour rendre les intelligences artificielles géantes plus légères, plus rapides et plus faciles à installer sur des appareils ordinaires, sans sacrifier leur intelligence.

Résumé technique

1. Problématique

Les modèles de langage de grande taille (LLM) basés sur l'architecture Mixture-of-Experts (MoE) offrent des performances exceptionnelles et réduisent les coûts d'inférence par rapport aux modèles denses, car seuls un sous-ensemble de paramètres est activé lors de l'inférence. Cependant, ces modèles souffrent de deux limitations majeures :

• Coût mémoire prohibitif : Bien que l'activation soit sparse, tous les paramètres (par exemple, 671 milliards pour DeepSeek-V3) doivent être stockés en mémoire GPU, ce qui limite le déploiement sur des dispositifs aux ressources contraintes.
• Limites du pruning existant : Les méthodes de compression actuelles se concentrent principalement sur le pruning au niveau des experts (suppression d'experts entiers). Cette granularité trop grossière entraîne souvent une dégradation significative de la précision. À l'inverse, le pruning fin (au niveau des poids) offre peu d'accélération matérielle. De plus, les méthodes de fusion d'experts souffrent d'instabilité et de conflits de paramètres.

L'objectif est donc de trouver une méthode de compression qui soit à la fois efficace en mémoire, accélère l'inférence et préserve la précision du modèle, en utilisant une granularité plus fine que l'expert complet.

2. Méthodologie : HEAPr

Les auteurs proposent HEAPr, un algorithme de pruning qui décompose chaque expert en unités plus petites et indivisibles appelées "atomic experts" (experts atomiques).

A. Décomposition en Atomic Experts

Dans une couche MoE, chaque expert $E_i$ est un bloc feed-forward. HEAPr décompose cet expert en $d_{inter}$ "atomic experts".

• Un expert atomique $e^{(j)}_i$ correspond à la combinaison d'une colonne spécifique des matrices $W_{up}$ et $W_{gate}$, et de la ligne correspondante de $W_{down}$.
• La sortie de l'expert complet est la somme des sorties de ses experts atomiques.
• Pruner au niveau atomique permet de retirer des unités spécifiques sans perturber la structure des autres composants, offrant une flexibilité accrue.

B. Estimation de l'Importance via l'Information du Second Ordre

Pour déterminer quels experts atomiques supprimer, HEAPr s'inspire de la théorie du Optimal Brain Surgeon (OBS), qui utilise l'expansion de Taylor du second ordre de la fonction de perte pour estimer l'impact de la suppression d'un paramètre.

Le défi principal est le coût computationnel et mémoire du calcul de la matrice Hessienne ($H$), qui est prohibitif ($O(d^4)$). HEAPr résout ce problème grâce à deux optimisations clés :

1. Découplage des paramètres : Les auteurs démontrent que les dérivées secondes croisées entre les paramètres de différents experts atomiques sont nulles. Cela permet de décomposer la Hessienne globale en une somme de Hessiennes individuelles pour chaque expert atomique, réduisant la complexité spatiale.
2. Passage à l'espace de sortie (Output Space) : Au lieu d'analyser les contraintes dans l'espace des paramètres, HEAPr reformule le problème dans l'espace de sortie.
   • Il utilise la Matrice d'Information de Fisher (F) comme approximation efficace de la Hessienne espérée.
   • Une propriété clé est exploitée : les gradients de la perte par rapport aux sorties de tous les experts atomiques d'un même expert sont identiques.
   • Cela permet de calculer une seule matrice de covariance de gradient par expert, partagée par tous ses experts atomiques.

La complexité spatiale est ainsi réduite de $O((3d_{model} \cdot d_{inter})^2)$ à $O(d_{model}^2)$.

C. Algorithme

HEAPr nécessite uniquement deux passages avant (forward) et un passage arrière (backward) sur un petit ensemble de calibration pour calculer l'importance de tous les experts atomiques :

1. Estimation de la covariance de gradient partagée : Calcul d'une matrice de covariance pour chaque expert basée sur les gradients de sortie.
2. Calcul de l'importance : Pour chaque expert atomique, l'importance $s$ est calculée comme le produit scalaire entre sa sortie et la covariance de gradient partagée.
3. Pruning Global : Les experts atomiques sont classés globalement sur l'ensemble du modèle (toutes les couches confondues) et les $r\%$ les moins importants sont supprimés.

3. Contributions Clés

• Nouvelle granularité : Introduction du concept d'atomic expert pour un pruning plus flexible et précis que le niveau expert.
• Efficacité algorithmique : Transformation de l'information du second ordre (Hessienne) des paramètres vers l'espace de sortie, réduisant la complexité de stockage de $O(d^4)$ à $O(d^2)$ et rendant le calcul faisable.
• Faible coût de calibration : La méthode ne nécessite que deux passes avant et une passe arrière sur un petit jeu de données, sans réentraînement (fine-tuning).
• Pruning Global : Capacité à comparer et classer l'importance des experts atomiques à travers toutes les couches du modèle, assurant une cohérence globale.

4. Résultats Expérimentaux

Les auteurs ont évalué HEAPr sur plusieurs modèles MoE modernes (DeepSeekMoE-16B, Qwen1.5-MoE, Qwen2-57B, Qwen3-30B) sur sept tâches zero-shot.

• Compression quasi sans perte :
  • Sur DeepSeekMoE-16B-Base, un pruning de 20% maintient une performance identique au modèle original.
  • Sur Qwen1.5-MoE-A2.7B-Chat, un pruning de 25% est quasi sans perte.
  • Sur Qwen2-57B-A14B, le modèle conserve presque toute sa précision même avec 40% de pruning.
  • Sur le dernier modèle Qwen3-30B-A3B, la précision moyenne ne chute que de 0,03 à un taux de 25%.
• Réduction des FLOPs : Le pruning de 20-25% réduit les FLOPs d'environ 20%, offrant une accélération réelle contrairement au pruning d'experts entiers qui ne réduit pas toujours la charge computationnelle de manière linéaire.
• Comparaison avec l'état de l'art : HEAPr surpasse systématiquement les méthodes existantes (NAEE, MoE-I2, MC-SMoE, D2-MoE, etc.) en termes de précision et d'efficacité, en particulier dans les régimes de pruning agressifs.
• Robustesse : La méthode est robuste au choix des données de calibration (WikiText-2 vs C4) et bénéficie d'une taille d'échantillon plus grande pour une estimation plus fiable.

5. Signification et Impact

HEAPr représente une avancée significative dans le domaine de la compression des modèles MoE. En passant d'une granularité "expert" à une granularité "atomique" et en optimisant le calcul de l'information du second ordre, l'article résout le compromis traditionnel entre précision et efficacité matérielle.

Cette approche permet de déployer des modèles MoE massifs sur des matériels aux ressources limitées (comme les GPU grand public ou les périphériques mobiles) sans sacrifier leurs capacités de raisonnement. De plus, la méthode fournit une perspective plus fine sur la redondance interne des modèles MoE, ouvrant la voie à de futures recherches sur la compensation de paramètres et l'architecture de modèles plus efficaces.