Résumé Technique : Routage par Index Inversé Adaptatif pour les Mélanges d'Experts Granulaires (AIR-MoE)

1. Énoncé du Problème

Les modèles de type Mélange d'Experts (MoE) sont devenus une pierre angulaire des Modèles de Langage à Grande Échelle (LLM) évolutifs en activant uniquement un sous-ensemble d'experts par token. De récentes lois d'échelle empiriques suggèrent que les performances s'améliorent avec des experts granulaires — en utilisant de nombreux petits experts plutôt que quelques grands. Cependant, ce régime introduit un goulot d'étranglement computationnel significatif : le coût de routage. À mesure que le nombre d'experts ($E$) augmente, le routeur doit calculer les produits scalaires entre chaque représentation de token et chaque centroïde d'expert pour identifier les $K$ meilleurs experts. Cette surcharge peut dominer le calcul total, rendant l'entraînement de modèles hautement granulaires prohibitif.

Les solutions existantes pour réduire les coûts de routage imposent généralement des contraintes structurelles qui limitent l'expressivité du modèle :

• Les MoE Hiérarchiques partitionnent les experts en groupes fixes et non chevauchants, restreignant le routage à l'intérieur de ces groupes.
• Les Contraintes de Clé Produit/Structurelles (par exemple, He, 2024) divisent les centroïdes d'experts en composants pour permettre une récupération exacte, mais restreignent l'espace des représentations d'experts admissibles (par exemple, en imposant une structure de produit cartésien).
• Les approches de Hachage Fixe manquent de capacité d'apprentissage et se sont révélées fondamentalement limitées.

Le défi consiste à approximer le routage top-$K$ efficacement sans imposer de contraintes structurelles sur les paramètres des experts ni exiger un score complet de tous les experts.

2. Méthodologie : AIR-MoE

Les auteurs proposent le Routage par Index Inversé Adaptatif pour MoE (AIR-MoE), une architecture de routage inspirée des index de fichiers inversés utilisés en recherche d'information et en recherche de voisins les plus proches approximatifs. La méthode approxime les vrais experts top-$K$ via un processus de récupération en deux étapes, du grossier au fin, évitant ainsi le score complet de tous les experts.

2.1 Récupération en Deux Étapes

Pour chaque token d'entrée $h$, le routage procède en parallèle :

1. Sélection Grossière : Le token est mappé au mot de code $c_g$ le plus proche dans un codebook appris $\mathcal{C} = \{c_1, \dots, c_G\}$ en utilisant la similarité cosinus. Chaque mot de code $c_g$ est associé à une liste de sélection pré-calculée $L_g$ de $M$ experts candidats. Ces listes de sélection sont essentiellement des "listes de publication" dans un index inversé, stockant les indices des experts ayant les produits scalaires les plus élevés par rapport au centroïde du mot de code.
2. Score Fin : Le routeur calcule les produits scalaires exacts uniquement entre le token $h$ et les experts au sein de la liste de sélection récupérée $L_g$. Les $K$ meilleurs experts sont ensuite sélectionnés à partir de cet ensemble de candidats réduit.

Cette hiérarchie garantit que l'ensemble des experts actifs $T$ est un sous-ensemble de la liste de sélection $L$, elle-même un sous-ensemble de l'ensemble complet des experts $[E]$.

2.2 Entraînement et Optimisation

Un défi central est que l'affectation des tokens aux mots de code est non différentiable. AIR-MoE y répond via une stratégie d'optimisation bi-niveau :

• Entraînement Basé sur le Gradient : Les représentations de tokens et les centroïdes d'experts sont mis à jour par rétropropagation standard à partir de la perte de modélisation du langage en aval.
• Entraînement du Codebook Sans Gradient : Le codebook $\mathcal{C}$ et ses listes de sélection associées sont mis à jour indépendamment en utilisant un algorithme de k-moyennes sphériques adaptatif. Cet algorithme :
  • Utilise une Moyenne Mobile Exponentielle (EMA) pour suivre la distribution des représentations de tokens.
  • Opère sur la sphère unité pour gérer efficacement la géométrie cosinus de haute dimension.
  • Inclut un mécanisme de réinitialisation pour les mots de code "morts" (ceux avec un faible nombre d'affectations) pour éviter l'effondrement.
  • Met à jour les listes de sélection $L_g$ uniquement après les étapes de l'optimiseur, garantissant que l'index reste cohérent avec l'état actuel des centroïdes d'experts.

Crucialement, AIR-MoE n'utilise pas l'astuce d'estimation directe (straight-through estimation) ; les gradients ne circulent pas à travers l'étape de quantification vers le codebook.

2.3 Garanties Théoriques

L'article dérive une borne inférieure sur le Rappel de Masse (la masse de probabilité totale des vrais experts top-$K$ conservée dans la liste de sélection). La borne dépend de :

1. L'erreur de quantification $\epsilon(h) = \|h - c(h)\|$ : Une erreur plus faible conduit à une meilleure rétention.
2. La masse de la liste de sélection $\rho_M(c)$ : La masse de probabilité capturée par les $M$ meilleurs experts par rapport au mot de code.
   Théoriquement, AIR-MoE conserve la majeure partie de la masse de routage lorsque les tokens sont bien quantifiés et que la taille de la liste de sélection $M$ est suffisamment grande.

3. Contributions Clés

1. Architecture AIR-MoE : Un mécanisme de routage inspiré de l'index inversé qui approxime le routage top-$K$ pour les MoE granulaires sans imposer de contraintes structurelles (par exemple, regroupement ou restrictions paramétriques) sur les représentations d'experts.
2. Optimisation Bi-niveau : Une stratégie d'entraînement qui découple l'adaptation sans gradient du codebook (via k-moyennes sphériques adaptatives) de l'entraînement basé sur le gradient des paramètres du modèle, évitant ainsi le besoin d'estimation directe.
3. Analyse Théorique : Une borne inférieure dérivée sur le rappel de masse, fournissant une intuition sur les conditions dans lesquelles l'approximation préserve la qualité du routage.
4. Performance Empirique : Démonstration que AIR-MoE atteint un compromis perplexité–FLOPs supérieur par rapport aux méthodes de routage granulaire efficaces existantes.

4. Résultats Expérimentaux

Les auteurs ont évalué AIR-MoE sur WikiText-103, C5 et OpenWebText2 à travers des échelles de modèles petite, moyenne et grande (jusqu'à environ 0,45 milliard de paramètres) avec 65 536 experts.

• Performance vs. Baselines : AIR-MoE a constamment surpassé les baselines granulaires efficaces existantes, notamment PEER (He, 2024) et les MoE Hiérarchiques.
  • Il a réalisé des améliorations de perplexité allant jusqu'à 10 % par rapport à la baseline PEER à des coûts FLOPs similaires.
  • Bien que le routage granulaire standard (score de tous les experts) ait atteint la perplexité la plus faible, il a engendré des coûts computationnels prohibitifs (par exemple, ~467,5 P FLOPs contre 324,1 P pour AIR-MoE sur WikiText-103 Small).
• Efficacité : AIR-MoE a réduit les FLOPs d'entraînement de manière significative par rapport au routage granulaire exact tout en maintenant une perplexité compétitive.
• Études d'Ablation :
  • Supprimer la normalisation des centroïdes ou utiliser la distance euclidienne au lieu des k-moyennes sphériques a légèrement dégradé les performances.
  • L'utilisation d'un codebook statique (non adaptatif) a donné les pires performances, confirmant la nécessité d'un codebook adaptatif.
  • Le gating par choix d'expert a conduit à un fort chevauchement mais à une famine sévère des experts (78,6 % d'experts morts), indiquant que le chevauchement seul n'est pas une métrique suffisante pour la qualité du routage.
• Utilisation des Experts : Contrairement aux routeurs hiérarchiques ou PEER qui imposent des structures rigides menant à des artefacts de routage, AIR-MoE a permis une utilisation flexible des experts sans contraintes structurelles.

5. Signification et Revendications

L'article revendique qu'AIR-MoE fournit un mécanisme efficace pour le routage dans les architectures de mélanges d'experts clairsemés, permettant spécifiquement l'utilisation d'experts granulaires sans la pénalité computationnelle du routage complet.

• Remplacement Direct : La méthode ne nécessite aucune modification de l'architecture du modèle sous-jacent ou de la fonction de perte, servant de remplacement direct pour les routeurs standards.
• Flexibilité : En évitant les contraintes structurelles sur les paramètres des experts, AIR-MoE permet au modèle d'apprendre les dépendances token-expert plus flexiblement que les méthodes nécessitant des regroupements fixes ou des structures de clés-produit.
• Évolutivité : L'approche évolue favorablement dans les régimes de grands lots effectifs car le coût de la sélection grossière est amorti et indépendant du nombre de tokens routés.

Les auteurs reconnaissent des limites, notant que bien que les réductions de FLOPs soient significatives, les accélérations réelles en temps réel dépendent du support matériel et d'implémentations optimisées pour atténuer le trafic mémoire et les surcharges d'indexation inhérentes aux MoE granulaires. Ils notent également que l'implémentation actuelle n'emploie pas de réduction de dimensionnalité, laissant cela comme une direction orthogonale pour les travaux futurs.