ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Résumé Technique : ProbMoE : Routage Probabiliste Différentiable pour le Mélange d'Experts

Énoncé du Problème
Les modèles de mélange d'experts (Mixture-of-Experts ou MoE) permettent de monter en charge efficacement en n'activant qu'un petit sous-ensemble d'experts par jeton (token). Cependant, l'entraînement de ces modèles reste difficile car le routage standard repose sur un opérateur top-$k$ discret et non différentiable. Les procédures d'entraînement conventionnelles contournent ce problème en ignorant la dépendance de la sélection top-$k$ vis-à-vis des logits du routeur, ne propageant les gradients qu'à travers les probabilités softmax des experts sélectionnés. Cette approximation limite le signal d'apprentissage pour les sous-ensembles d'experts alternatifs, conduisant souvent à des distributions de routage concentrées, une faible utilisation des experts et une instabilité de la dynamique d'entraînement. Les tentatives existantes pour atténuer cela, telles que la rétropropagation dense ou les estimateurs de gradient heuristiques, échouent souvent à modéliser explicitement la distribution sur les sous-ensembles d'experts réalisables ou fournissent des mécanismes limités pour explorer des combinaisons alternatives pendant l'entraînement.

Méthodologie
Les auteurs proposent ProbMoE, un cadre qui reformule le routage d'experts comme un problème d'inférence probabiliste sur des sous-ensembles d'experts à cardinalité contrainte. Au lieu de traiter le routage comme un opérateur déterministe, ProbMoE modélise la sélection d'experts comme une distribution sur des sous-ensembles discrets.

1. Formulation Probabiliste :

   • Routage Exact-$k$ : Le cadre définit une distribution sur des sous-ensembles $S$ où $|S|=k$. Il conditionne des probabilités de sélection de Bernoulli indépendantes (dérivées des logits du routeur via une sigmoïde) sur la contrainte que exactement $k$ experts soient sélectionnés. La constante de normalisation de cette distribution, $Z_k$, est calculée exactement en utilisant la programmation dynamique, rendant la formulation traitable.
   • Routage Dynamique-$k$ : Le cadre se généralise à un contexte de contrainte de plage où la taille du sous-ensemble $|S|$ se situe dans $[k_{min}, k_{max}]$. La distribution est conditionnée sur cette plage, permettant au modèle d'inférer conjointement l'identité des experts et la cardinalité du routage pour chaque jeton.

2. Passage Avant (Forward Pass) :

   • Pendant l'entraînement, ProbMoE échantillonne un sous-ensemble d'experts $S$ à partir de la distribution définie (Exact-$k$ ou contrainte de plage).
   • La couche MoE calcule la sortie en utilisant uniquement les experts échantillonnés, préservant ainsi le schéma d'exécution parcimonieux des modèles MoE standards.
   • À l'inférence, l'ensemble Maximum-A-Posteriori (MAP) est sélectionné, garantissant que les coûts d'inférence correspondent au routage MoE conventionnel.

3. Passage Arrière (Backward Pass) :

   • Pour permettre l'optimisation par gradient, ProbMoE utilise l'estimateur de gradient SIMPLE (Ahmed et al., 2023).
   • Au lieu de différencier à travers l'échantillon discret, le passage arrière utilise les probabilités marginales de chaque expert étant sélectionné sous la distribution contrainte.
   • Les poids de routage sont construits en utilisant un estimateur straight-through : $w = (\text{stopgrad}(z - m) + m) \odot \pi$, où $z$ est le masque échantillonné, $m$ est le vecteur des probabilités marginales exactes et $\pi$ sont les probabilités de routage softmax.
   • Cela permet aux gradients de circuler à travers les probabilités marginales, fournissant des mises à jour informatives au routeur qui reflètent comment les changements dans les probabilités de sélection affectent l'espérance de la perte, tout en maintenant une exécution parcimonieuse des experts.

Contributions Clés

1. Formulation d'Inférence Probabiliste : L'article présente le routage MoE comme une inférence probabiliste sur des sous-ensembles d'experts à cardinalité contrainte, allant au-delà des approximations de gradient heuristiques.
2. Optimisation Traitable Basée sur les Marginales : En exploitant l'estimateur SIMPLE, les auteurs associent l'échantillonnage stochastique de sous-ensembles à des gradients de routeur exacts basés sur les marginales, permettant un entraînement différentiable sans calcul dense du côté des experts.
3. Cadre Unifié Dynamique-$k$ : La méthode s'étend naturellement au routage dynamique, où le modèle infère à la fois l'identité des experts et le nombre d'experts actifs par jeton au sein d'une plage prédéfinie, sans dépendre de règles de partition heuristiques.
4. Améliorations Empiriques :
   • Exact-$k$ : ProbMoE atteint des performances solides sur les benchmarks (GSM8K, Law, Translation, MBPP, MMLU) sur les architectures OLMoE et Qwen, égalant ou surpassant souvent les bases compétitives comme DenseMixer et le Top-$k$ Conventionnel.
   • Dynamique-$k$ : ProbMoE Dynamic-$k$ atteint des performances comparables à Exact-$k$ tout en activant moins d'experts en moyenne, démontrant une allocation adaptative efficace.
5. Diversité du Routage : L'analyse montre que ProbMoE induit des distributions de routage plus larges et une utilisation plus équilibrée des experts par rapport aux bases de référence, atténuant l'effondrement des experts (expert collapse).

Résultats et Signification
L'article démontre que ProbMoE améliore l'utilisation des experts et la diversité du routage. Dans les études d'ablation, il est démontré que la combinaison de l'échantillonnage stochastique et de l'optimisation basée sur les marginales est le principal moteur des gains de performance, surpassant les configurations qui n'utilisent qu'un seul de ces composants.

Les auteurs affirment que ProbMoE fournit une méthode fondée sur des principes pour la sélection d'experts qui améliore la stabilité de l'entraînement et l'efficacité computationnelle. En modélisant le routage de manière probabiliste, le cadre permet l'exploration de combinaisons d'experts alternatives pendant l'entraînement, conduisant à des distributions de routage mieux calibrées. La variante Dynamic-$k$ montre en outre qu'une cardinalité adaptative peut être apprise pour correspondre à la complexité par jeton, réduisant les coûts de routage effectifs sans sacrifier les performances des tâches. Ce travail suggère que les directions futures incluent l'extension de ProbMoE à la pré-entraînement et aux optimisations au niveau système pour exploiter pleinement la parcimonie dans les contextes dynamiques.