EvoESAP: Non-Uniform Expert Pruning for Sparse MoE

Le papier présente EvoESAP, un cadre de recherche évolutionnaire qui optimise l'allocation non uniforme de l'élagage des experts dans les modèles MoE épars en utilisant une métrique de proxy d'acceptation spéculative (ESAP), permettant d'améliorer significativement les performances de génération tout en respectant un budget global fixe.

L'essentiel

Imaginez que vous avez un groupe de 100 experts (des cuisiniers, des mathématiciens, des codeurs) dans une cuisine géante. C'est ce qu'on appelle un modèle de langage "MoE" (Mixture of Experts).

Pour répondre à une question, le chef (le modèle) ne fait pas travailler tout le monde. Il choisit seulement les 3 ou 4 meilleurs experts pour la tâche du moment. C'est très efficace !

Mais il y a un problème :
Même si seuls 3 experts travaillent, vous devez garder les 100 experts dans votre cuisine (sur votre serveur) au cas où ils seraient nécessaires. Cela prend énormément d'espace et coûte cher à faire tourner.

L'objectif de cette recherche est de licencier certains experts pour réduire la taille de la cuisine, tout en s'assurant que le repas (la réponse du modèle) reste aussi délicieux.

Voici comment les auteurs ont résolu le problème, expliqué simplement :

1. Le problème de l'ancienne méthode : "La règle du couloir"

Jusqu'à présent, quand on voulait réduire la taille de l'équipe, on appliquait une règle simple et uniforme : "On licencie 50% des experts dans chaque étage de l'immeuble".
C'est comme si, dans un immeuble de 10 étages, on enlevait 5 personnes par étage, partout pareil.

• Le souci : Certains étages (couches du modèle) sont très importants pour la créativité, d'autres pour la logique. Enlever la même quantité partout, c'est comme couper les ailes d'un oiseau de la même façon sur les deux côtés : ça marche, mais ce n'est pas optimal. Parfois, on coupe trop dans un étage crucial et le modèle devient bête.

2. La nouvelle idée : "L'architecte intelligent" (EvoESAP)

Les auteurs disent : "Attendez, on ne devrait pas couper partout pareil !".
Ils proposent une méthode en deux étapes :

1. Qui licencier ? (On garde l'ordre de départ : on enlève d'abord les moins performants).
2. Où licencier ? (C'est là que la magie opère).

Au lieu de couper 50% partout, ils utilisent un algorithme d'évolution (comme la sélection naturelle) pour trouver le meilleur équilibre.

• Exemple : Peut-être qu'on peut licencier 70% des experts dans les étages 1 à 5 (qui sont moins importants), mais seulement 30% dans les étages 10 à 15 (qui sont cruciaux pour la créativité).
• Résultat : On garde le même nombre total d'experts licenciés, mais on garde les "cerveaux" là où il faut.

3. Le test de réalité : "Le double jeu" (ESAP)

Comment savoir si cette nouvelle répartition est bonne sans attendre des jours pour tester le modèle ?
C'est là qu'intervient ESAP (le nom technique). Imaginez que vous avez un juge invisible.

• Au lieu de faire écrire un roman entier au modèle réduit pour voir si c'est bien (ce qui prendrait des heures), on lui demande de deviner un seul mot à la fois, en regardant ce que le modèle original (le "professeur") aurait écrit.
• Si le modèle réduit devine souvent le même mot que le professeur, c'est qu'il est intelligent.
• Cette méthode est ultra-rapide et permet de tester des milliers de combinaisons de licenciement en quelques minutes.

4. Les résultats : "Plus petit, mais plus malin"

Les chercheurs ont testé ça sur des modèles de 7 à 30 milliards de paramètres.

• Le constat : En utilisant leur méthode "non uniforme" (couper plus ici, moins là), ils ont obtenu des résultats bien meilleurs pour la créativité (écrire des histoires, du code, des maths) que la méthode uniforme, tout en ayant la même taille de fichier.
• L'analogie finale : C'est comme si vous aviez une valise de 20kg. La méthode ancienne consistait à enlever 1kg de vêtements à chaque endroit (chemises, pantalons, chaussettes). La nouvelle méthode consiste à enlever 5kg de chaussettes inutiles et à garder tous les manteaux chauds. Résultat : la valise est plus légère, mais vous êtes toujours au chaud !

En résumé :
Ce papier nous apprend qu'il ne faut pas traiter toutes les parties d'une intelligence artificielle de la même façon. En étant plus stratégique sur où on réduit la taille, on peut rendre les modèles plus petits, moins chers à utiliser, tout en gardant (voire en améliorant) leur capacité à créer et à raisonner.

Résumé technique

1. Problématique

Les modèles de langage à base de Mélange d'Experts (MoE) épars (Sparse Mixture-of-Experts) offrent des performances élevées avec un coût de calcul par token réduit. Cependant, leur déploiement reste limité par la mémoire et le débit, car l'ensemble complet des experts doit être stocké et servi, même si seuls quelques-uns sont activés à chaque étape.

Le pruning (élagage) d'experts après l'entraînement est une solution prometteuse pour réduire ces coûts. Pourtant, les méthodes existantes souffrent de deux limitations majeures :

1. Focalisation sur la sélection : Elles se concentrent principalement sur quels experts supprimer au sein de chaque couche, tout en ignorant combien en supprimer dans chaque couche.
2. Allocation uniforme : La plupart des approches adoptent par défaut une allocation uniforme du budget d'élagage (c'est-à-dire supprimer le même pourcentage d'experts dans chaque couche).
3. Évaluation inadéquate : La plupart des travaux évaluent le pruning uniquement sur des tâches à choix multiples (MCQ), négligeant la qualité de la génération ouverte (code, mathématiques, écriture créative), où les méthodes de fusion d'experts échouent souvent.

L'article postule que l'allocation du budget d'élagage entre les couches est une décision cruciale et sous-étudiée, et qu'une allocation non uniforme pourrait mieux préserver les capacités du modèle.

2. Méthodologie : EvoESAP

Les auteurs proposent EvoESAP, un cadre de recherche évolutionnaire conçu pour optimiser l'allocation non uniforme du budget d'élagage tout en maintenant l'ordre de sélection des experts au sein de chaque couche fixe.

A. Découplage du problème

Le processus de pruning est décomposé en deux étapes :

1. Sélection intra-couche (Fixe) : Un critère d'importance (ex: REAP, Frequency, EAN, SEER) est utilisé pour classer les experts dans chaque couche. L'ordre de suppression est déterminé une fois pour toutes.
2. Allocation inter-couche (Optimisée) : Le nombre d'experts à supprimer dans chaque couche ($r_\ell$) est optimisé sous un budget global fixe ($B$).

B. Fonction de Fitness : ESAP (Expected Speculative Acceptance Proxy)

Pour évaluer efficacement des milliers de candidats d'allocation sans recourir à un décodage autoregressif coûteux, les auteurs introduisent ESAP.

• Inspiration : Basé sur le speculative decoding (décodage spéculatif).
• Principe : ESAP mesure la similarité entre la distribution de probabilité du token suivant du modèle complet (enseignant) et celle du modèle élagué (élève) dans un cadre de teacher forcing.
• Formulation mathématique : Au lieu d'estimer le taux d'acceptation par échantillonnage (bruité et lent), ESAP calcule l'espérance de l'acceptation spéculative. Cela se réduit à la somme des minima des probabilités sur le vocabulaire :
  $$\text{ESAP}(x) = \sum_{v \in V} \min(p(v|x), q(v|x))$$
  où $p$ est la distribution du modèle complet et $q$ celle du modèle élagué.
• Avantage : C'est une métrique bornée, stable, calculable rapidement (sans génération autoregressive) et corrélée à la capacité du modèle élagué à imiter le modèle original.

C. Algorithme de Recherche Évolutionnaire

EvoESAP utilise une recherche évolutionnaire pour trouver la meilleure allocation de budget :

• Espace de recherche : Vecteurs d'entiers représentant le nombre d'experts supprimés par couche, soumis à la contrainte de budget global.
• Opérateur de mutation : Une mutation de type "level-switch" (commutation de niveau) transfère le budget d'élagage d'une couche à une autre tout en respectant les contraintes (au moins $k$ experts restants par couche).
• Processus : Initialisation avec des schémas uniformes et aléatoires, sélection des meilleurs individus via ESAP, mutation, et itération sur plusieurs générations.

3. Contributions Clés

1. Introduction d'ESAP : Une fonction de fitness efficace, inspirée du décodage spéculatif, permettant d'évaluer rapidement des milliers de candidats de pruning pour la préservation de la génération.
2. Identification de l'importance de l'allocation non uniforme : Démonstration que, même avec le même critère de sélection d'experts, une allocation non uniforme (trouvée par recherche) surpasse systématiquement l'allocation uniforme, en particulier pour les tâches de génération ouverte.
3. Cadre EvoESAP : Une méthode "plug-and-play" qui peut être appliquée sur n'importe quel critère de pruning existant (Frequency, REAP, etc.) pour optimiser la répartition du budget.

4. Résultats Expérimentaux

Les expériences ont été menées sur trois modèles MoE de grande taille (7B à 30B paramètres) : OLMoE, ERNIE-4.5, et Qwen3, à des niveaux d'épissure de 25 % et 50 %.

• Performance de Génération Ouverte : EvoESAP améliore significativement les performances sur le code et les mathématiques.
  • Sur ERNIE-4.5 à 50 % d'épissure, l'utilisation de la métrique REAP combinée à l'allocation recherchée par EvoESAP a entraîné une augmentation de +19,6 % sur le benchmark MATH-500 par rapport à une allocation uniforme.
  • Des gains similaires sont observés sur les benchmarks de code (Eval+, LiveCodeBench) et d'écriture créative (WildBench).
• Précision à Choix Multiples (MCQ) : Les performances sur les tâches MCQ (MMLU, ARC, etc.) restent compétitives, avec des variations minimes (parfois légères améliorations, parfois légères baisses négligeables), prouvant que la méthode ne sacrifie pas la compréhension factuelle.
• Efficacité du calcul : L'utilisation d'ESAP réduit le temps de recherche de près de 18 fois par rapport à l'utilisation du taux d'acceptation réel du décodage spéculatif (Tableau 4).
• Robustesse : La méthode fonctionne bien avec différents critères de sélection d'experts et différents modèles, bien que la distribution optimale trouvée varie selon le modèle et le critère utilisé (il n'existe pas de template universel).

5. Signification et Impact

• Optimisation du déploiement : EvoESAP permet de réduire l'empreinte mémoire et les coûts d'inférence des modèles MoE sans nécessiter de réentraînement (finetuning), tout en préservant, voire en améliorant, la qualité de la génération complexe.
• Changement de paradigme : L'article démontre que l'hypothèse d'une épissure uniforme par couche est sous-optimale. La recherche active de schémas non uniformes est une étape nécessaire pour le compression efficace des LLMs.
• Accessibilité : En réduisant les coûts de déploiement tout en maintenant des capacités de pointe, cette méthode facilite l'accès à des modèles performants dans des environnements à ressources limitées.

En résumé, EvoESAP propose une approche rigoureuse et efficace pour le pruning des MoE, transformant le problème d'allocation de ressources en un problème d'optimisation guidé par une métrique de similarité comportementale (ESAP), avec des gains substantiels sur les tâches de génération les plus difficiles.