Résumé Technique : Routage Sinkhorn Sélectif pour un Mélange d'Experts Creux Amélioré

Énoncé du Problème

Les modèles de Mélange d'Experts Creux (SMoE - Sparse Mixture-of-Experts) offrent une approche évolutive pour augmenter la capacité du modèle tout en maintenant l'efficacité computationnelle en n'activant qu'un sous-ensemble d'experts par jeton (token). Cependant, les modèles SMoE conventionnels reposant sur un routage basé sur Softmax souffrent souvent de « l'effondrement du routage » (routing collapse), où un petit sous-ensemble d'experts est surutilisé tandis que d'autres restent sous-utilisés.

Les stratégies de mitigation existantes reposent généralement sur :

1. Objectifs Auxiliaires : Des pertes d'équilibrage de charge ou des pertes z qui pénalisent une utilisation inégale des experts. Ces méthodes peuvent introduire un désalignement de l'objectif, une instabilité de l'entraînement et une complexité accrue.
2. Bruit de Gating : L'injection de bruit pour encourager l'exploration, ce qui ajoute des paramètres ou des hyperparamètres entraînables.
3. Routage basé sur Sinkhorn : Les méthodes antérieures utilisant l'algorithme de Sinkhorn pour le transport optimal (OT) utilisent souvent la carte de transport uniquement pour la sélection des experts (top-$k$) plutôt que pour l'assignation des poids. De plus, elles peuvent réduire la flexibilité du routage car la matrice de gating n'est pas directement optimisée par apprentissage par gradient de la même manière que les méthodes Softmax, ou elles engendrent un surcoût d'entraînement substantiel dû à la nature itérative des mises à jour de Sinkhorn.

Méthodologie

Le papier propose le Routage Sinkhorn Sélectif (SSR - Selective Sinkhorn Routing), un mécanisme léger qui reformule l'assignation jeton-expert comme un problème de transport optimal régularisé par l'entropie.

Formulation Centrale

Les auteurs modélisent l'assignation de $m$ jetons à $n$ experts comme un problème de transport optimal à coût maximal avec régularisation entropique. Le but est de trouver un plan de transport $\hat{\Pi}$ qui maximise la compatibilité tout en satisfaisant des contraintes spécifiques :

• Positivité : $\Pi > 0$.
• Somme des lignes : Chaque jeton route sa masse entière vers les experts ($\Pi \mathbf{1}_n = \mathbf{1}_m$).
• Somme des colonnes : Chaque expert reçoit une charge totale attendue égale ($\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$).

La matrice de coût $C$ est dérivée des scores de gating $S$. Les auteurs proposent deux variantes :

• Coût Linéaire : $C = S$.
• Coût Softmax : $C_{i,:} = \text{softmax}(S_{i,:})$, utilisé pour éviter le dépassement numérique lors des itérations de Sinkhorn.

Contrairement aux méthodes antérieures basées sur l'OT qui utilisent le plan de transport uniquement pour la sélection, le SSR utilise directement les valeurs de $\hat{\Pi}$ pour calculer les poids de routage. Pour chaque jeton, les experts top-$k$ sont sélectionnés sur la base des entrées les plus élevées dans la ligne correspondante de $\hat{\Pi}$, et les poids sont normalisés à partir de ces entrées.

Stratégie Sélective et Injection de Bruit

Un défi clé identifié est que le routage Sinkhorn ne met pas à jour la matrice de poids de gating $W_g$ car $W_g$ est découplée de l'objectif OT. Pour y remédier, le SSR emploie une stratégie d'entraînement hybride :

• Application Sélective : Pendant l'entraînement, chaque bloc MoE utilise le routage Sinkhorn avec une faible probabilité $p$ (par exemple, 0,001) et un gating Softmax standard par ailleurs. Cela permet au gating Softmax de mettre à jour $W_g$ efficacement tout en utilisant le routage Sinkhorn pour fournir des signaux d'utilisation équilibrée des experts.
• Injection de Bruit : Pour prévenir l'effondrement des experts et encourager l'exploration, un bruit gaussien est ajouté à la matrice de coût ($\tilde{C} = C + \alpha_{noise} \cdot \epsilon$). L'analyse théorique (Proposition 4.3) montre que cela garantit que chaque expert a une probabilité non nulle d'être sélectionné.

Comportement à l'Inférence

Le papier soutient que l'imposition de contraintes d'équilibre (comme des sommes de colonnes égales) lors de l'inférence sur une seule entrée est contre-productive, car cela force un routage uniforme indépendamment des scores de compatibilité spécifiques de l'entrée. Par conséquent, le routage Sinkhorn et l'injection de bruit sont tous deux désactivés lors de l'inférence, revenant à un routage Softmax déterministe pour garantir des prédictions cohérentes et dépendantes de l'entrée.

Principales Contributions

1. Nouveau Cadre de Routage : Une méthode intégrant le transport optimal régularisé par l'entropie avec une injection de bruit stochastique pour favoriser une utilisation équilibrée des experts sans pertes d'équilibrage auxiliaires.
2. Perspectives Théoriques : Des preuves démontrant que le routage basé sur Sinkhorn et l'injection de bruit aident l'entraînement en encourageant l'exploration et l'équilibre, mais doivent être désactivés à l'inférence pour éviter des assignations déformées.
3. Efficacité et Performance : Des évaluations approfondies montrant que le SSR améliore l'efficacité de l'entraînement, la précision et la robustesse par rapport aux lignes de base de l'état de l'art.

Résultats Expérimentaux

Les auteurs ont évalué le SSR sur la modélisation du langage (WikiText-103, Enwik-8) et la classification d'images (ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R).

• Modélisation du Langage :
  • Sur WikiText-103, les variantes du SSR surpassent systématiquement le Vanilla SMoE et les bases utilisant des pertes d'équilibrage de charge, des pertes z ou du bruit.
  • SSR-L avec bruit a atteint une perplexité de test (PPL) de 34,367, soit une réduction de 1,183 par rapport au Vanilla SMoE (35,550), surpassant de manière significative la meilleure ligne de base.
  • Crucialement, le SSR a engendré un surcoût d'entraînement minimal (0,33 % – 0,65 %), alors qu'un SMoE complet basé sur Sinkhorn engendrait un surcoût de 72,47 %.
  • Sous les réglages de Momentum, le SSR a préservé l'évolution caractéristique de la norme du Vanilla SMoE tout en améliorant les performances, là où d'autres méthodes d'équilibrage ont dégradé les performances.
• Modélisation au niveau des octets (Enwik-8) :
  • SSR-S avec bruit a obtenu le meilleur Bits-per-Character (BPC) de 1,128, améliorant le Vanilla SMoE de 0,010 BPC.
  • Il était 2,28x plus rapide qu'un SMoE complet basé sur Sinkhorn.
• Tâches de Vision :
  • Sur ImageNet-1K, SSR-L avec bruit a atteint une précision Top-1 de 77,420, un gain de +2,368 par rapport au Vanilla SMoE.
  • Il a également montré une forte robustesse sur ImageNet-O et ImageNet-R.
• Études d'Ablation :
  • La probabilité $p$ d'appliquer le routage Sinkhorn est critique ; des valeurs très faibles (par exemple, $10^{-3}$ ou $10^{-4}$) étaient suffisantes pour fournir des signaux d'équilibrage sans trop contraindre le routeur.
  • L'injection de bruit a systématiquement amélioré les performances dans tous les réglages.
  • Désactiver les mécanismes d'équilibrage à l'inférence a donné les meilleurs résultats, confirmant les affirmations théoriques.

Signification et Revendications

Le papier affirme que le Routage Sinkhorn Sélectif offre une solution pratique et efficace pour la conception de SMoE. En remplaçant les pertes auxiliaires complexes par une optimisation intermittente et efficace basée sur Sinkhorn, le SSR atteint :

• Une Meilleure Stabilité de l'Entraînement : Grâce à un équilibrage intrinsèque des experts sans désalignement de l'objectif.
• Une Précision Supérieure : Démontrée sur les benchmarks de texte et de vision.
• Une Robustesse : Performance améliorée face à la corruption des entrées et aux exemples adverses.
• Une Efficacité : Surcoût de calcul minimal par rapport aux méthodes de perte auxiliaire et au routage Sinkhorn complet.

Les auteurs soulignent que la méthode est légère, ne nécessite aucun paramètre entraînable supplémentaire (dans la variante avec bruit) et maintient la flexibilité de la sélection des experts tout en garantissant une utilisation équilibrée durant la phase d'entraînement.