Résumé Technique : Mélange Hiérarchique d'Experts avec Optimisation en Deux Étapes (Hi-MoE)

1. Énoncé du Problème

Les modèles de Mélange d'Experts (MoE) éparses augmentent leur capacité en acheminant les jetons vers un petit sous-ensemble de réseaux d'experts. Cependant, l'entraînement de ces modèles fait face à un compromis fondamental entre l'équilibrage de charge et la spécialisation des experts :

• Effondrement du Routage : Des mécanismes de diversité agressifs conduisent souvent à un « effondrement du routage », où quelques experts dominent la sélection des jetons tandis que d'autres restent inactifs. Cela entraîne une utilisation inefficace de la capacité et un déséquilibre de charge entre les appareils (phénomène des traînards).
• Le Dilemme du Groupement : Des approches récentes atténuent le déséquilibre en regroupant les experts et en imposant une sélection fixe par groupe (par exemple, sélectionner $K/M$ experts parmi chacun des $M$ groupes). Bien que cela assure un calcul uniforme au niveau de l'appareil, cela restreint la liberté de routage. Cette contrainte conduit souvent à une redondance, où les experts de groupes différents apprennent des fonctions similaires car ils sont optimisés sur des mélanges de jetons comparables, remplaçant ainsi l'effondrement au niveau de l'expert par un « effondrement de la diversité » entre les groupes.

Les méthodes d'équilibrage basées sur la perte existantes (par exemple, GShard) entrent souvent en conflit avec les gradients de la tâche, ralentissant la convergence, tandis que les méthodes sans perte peuvent échouer à prévenir la redondance. Il existe un besoin d'un cadre qui impose simultanément un équilibre aligné sur le matériel et favorise une spécialisation complémentaire et significative.

2. Méthodologie : Hi-MoE

Les auteurs proposent Hi-MoE, un cadre hiérarchique qui décompose le contrôle du routage en deux niveaux couplés : l'équilibrage inter-groupe et la spécialisation intra-groupe.

Architecture

• Structure Groupée : Les experts sont partitionnés en $M$ groupes disjoints (souvent alignés sur des fragments de GPU).
• Contrainte de Routage : Pour chaque jeton, le routeur sélectionne un nombre fixe d'experts ($K_m$) parmi chaque groupe $m$.
• Optimisation à Deux Niveaux : Le cadre introduit des régularisateurs spécifiques pour chaque niveau de la hiérarchie.

Composants Clés

1. Régularisation Inter-Groupe ($R_{inter}$) :

   • Objectif : Imposer une distribution équitable du trafic entre les groupes pour assurer une utilisation équilibrée des appareils.
   • Mécanisme : Minimise la concentration $\ell_2$ des poids de routage post-sélection entre les groupes. En pénalisant la concentration de la masse de routage sur des groupes spécifiques, elle agit comme un proxy pour minimiser le coefficient de variation (CV) des charges de groupe.
   • Fondement Théorique : Le théorème 4.1 prouve que la minimisation de ce terme fournit une borne supérieure certifiée sur le déséquilibre au niveau des groupes.

2. Régularisation Intra-Groupe ($R_{intra}$) :

   • Objectif : Favoriser une spécialisation complémentaire au sein des groupes et empêcher les experts d'apprendre des fonctions redondantes.
   • Mécanisme : Maximise la concentration $\ell_2$ de la distribution de routage pré-sélection (effectivement un terme anti-chevauchement). Cela encourage le routeur à prendre des engagements décisifs et épars envers des experts spécifiques au sein d'un groupe plutôt que de mélanger de nombreux experts avec une faible confiance.
   • Fondement Théorique : Le théorème 4.2 montre que la réduction du chevauchement de routage diminue le couplage des gradients entre les experts, leur permettant de suivre des trajectoires d'optimisation distinctes. Le théorème 4.3 démontre que, sous des marges équilibrées, l'augmentation de la concentration de routage maximise l'information mutuelle entre l'entrée et l'expert sélectionné, un proxy de la spécialisation.

3. Mécanisme de Stabilisation :

   • Pour empêcher le signal de diversité intra-groupe de dériver vers des régimes dégénérés, les auteurs emploient un calcul de routage corrigé par biais utilisant une moyenne mobile exponentielle (EMA) des logits du routeur. Cela décourage la sur-sélection persistante d'experts historiquement préférés sans ajouter de perte d'équilibrage séparée.

Fonction Objectif

L'objectif d'entraînement total combine la perte de tâche, la perte d'équilibrage de charge standard et les termes hiérarchiques proposés :
$$\mathcal{L} = \mathcal{L}_{task} + \mathcal{L}_{load} + R_{intra} + R_{inter}$$
Les auteurs formulent cela comme un problème d'optimisation sous contraintes où les coefficients $\lambda_{intra}$ et $\lambda_{inter}$ agissent comme des multiplicateurs de Lagrange pour les contraintes de diversité et d'équilibre, respectivement.

3. Contributions Clés

• Identification de la Tension : L'article identifie que l'imposition d'un équilibre au niveau de l'appareil via des contraintes de groupe strictes peut stabiliser l'utilisation mais augmenter la redondance entre les groupes.
• Cadre Hiérarchique : Hi-MoE introduit un objectif modulaire, intégrable, qui découple ces préoccupations :
  • Un objectif inter-groupe pour un équilibre aligné sur le matériel.
  • Un objectif intra-groupe pour la diversité fonctionnelle et la spécialisation.
• Garanties Théoriques : L'article fournit des preuves formelles reliant les régularisateurs aux bornes de déséquilibre de groupe, au découplage des gradients et à l'augmentation de l'information de routage (spécialisation).
• Validation Empirique : Des améliorations constantes dans les domaines de la vision et du langage, démontrant que Hi-MoE étend la frontière de Pareto entre la perplexité/précision et l'équilibre de charge (CV).

4. Résultats Expérimentaux

Les auteurs ont évalué Hi-MoE sur trois configurations d'entraînement préliminaire distinctes :

A. Vision par Ordinateur (Swin-MoE sur Tiny ImageNet)

• Configuration : Modèle de 1 milliard de paramètres avec 64 experts (4 groupes de 16).
• Résultats : Hi-MoE a atteint la précision Top-1 la plus élevée (59,13 %) et le meilleur équilibre de charge (CV 0,29) par rapport à GShard Vanilla, l'équilibrage sans perte et MoGE.
• Observation : La visualisation des cartes d'attention a révélé que les groupes Hi-MoE ont développé des spécialisations visuelles distinctes (par exemple, un groupe se concentrant sur les parties du corps, un autre sur les surfaces d'arrière-plan), tandis que les modèles de base montraient plus de redondance.

B. Modélisation du Langage (nanoGPT sur OpenWebText)

• Configuration : Modèle de 1,06 milliard de paramètres avec 8 experts (4 groupes de 2).
• Résultats : Hi-MoE a atteint la perplexité la plus faible (2,947) et le meilleur équilibre (CV 0,23).
• Ablation : La suppression de $R_{inter}$ a dégradé l'équilibre ; la suppression de $R_{intra}$ a dégradé la qualité. La correction de biais a agi comme un stabilisateur.
• Efficacité : Le routeur hiérarchique a ajouté une surcharge minimale (de 1,3 ms à 1,6 ms par couche) tout en réduisant la surcharge de synchronisation d'environ 3 fois par rapport au routage plat.

C. Pré-entraînement à Grande Échelle (OLMoE-7B vs Hi-MoE-7B)

• Configuration : Modèles de 7 milliards de paramètres entraînés sur 58,54 milliards de jetons (mélange Dolma).
• Perplexité : Hi-MoE-7B a réduit la perplexité de 5,6 % en moyenne sur 11 domaines de validation par rapport à la référence standard OLMoE-7B.
• Tâches en Aval : Hi-MoE a montré des gains significatifs sur les tranches MMLU (par exemple, +22,2 % en STEM, +11,9 % en Sciences Humaines) et sur les tâches de raisonnement de bon sens.
• Équilibre de Charge : Hi-MoE a amélioré l'équilibre des experts (CV) de 40 % (0,153 contre 0,257), répondant directement au goulot d'étranglement de l'efficacité des ressources dans l'entraînement à grande échelle.

5. Signification et Revendications

L'article revendique que Hi-MoE résout la tension « équilibre–diversité » inhérente à l'entraînement des MoE éparses. En modélisant explicitement la hiérarchie du routage (groupe vs expert), la méthode permet aux modèles de :

1. Passer à l'échelle Efficacement : Maintenir une utilisation uniforme des GPU sans sacrifier la capacité conditionnelle du modèle.
2. Apprendre des Représentations Spécialisées : Empêcher l'« effondrement de la diversité » où les experts groupés deviennent fonctionnellement redondants.
3. Améliorer la Qualité : Obtenir des gains constants en perplexité et en performance des tâches en aval à travers les tâches de vision et de langage.

Les auteurs positionnent Hi-MoE non pas comme un remplacement des mécanismes de stabilisation existants, mais comme un complément modulaire qui transforme le groupement d'experts en un mécanisme d'optimisation principiel plutôt qu'en une simple contrainte système. Le travail suggère que la coordination hiérarchique est une voie viable pour débloquer des gains de qualité fiables dans les modèles de plusieurs billions de paramètres.