Schéma d'équilibrage de charge basé sur un poids expert adaptatif pour le routage dynamique du MoE

Le déséquilibre de charge est un goulot d’étranglement majeur de performance dans l’entraînement des modèles mixture-of-experts (MoE), car des charges d’experts déséquilibrées peuvent entraîner un effondrement du routage. La plupart des approches existantes traitent ce problème en introduisant des fonctions de perte auxiliaires pour équilibrer la charge ; cependant, les hyperparamètres de ces fonctions de perte doivent souvent être réglés pour différentes tâches. De plus, augmenter le nombre d’experts activés tend à aggraver le déséquilibre de charge, tandis que fixer le nombre d’activations peut réduire la confiance du modèle dans la gestion de tâches difficiles. Pour relever ces défis, cet article propose une stratégie de routage à équilibre dynamique reposant sur un algorithme dynamique de routage à seuil. Après chaque étape de routage, la méthode ajuste les poids des experts pour influencer la distribution de charge lors du routage suivant. Contrairement aux méthodes d’équilibrage basées sur des fonctions de perte, notre approche opère directement au niveau du routage, évitant les perturbations de gradient susceptibles de dégrader la qualité du modèle, tout en routant de manière dynamique pour une utilisation plus efficace des ressources de calcul. Des expériences sur des benchmarks de compréhension du langage naturel (NLU) montrent que la méthode proposée atteint une précision comparable au routage top-2, tout en réduisant significativement l’écart-type de charge (par exemple, de 12,25 à 1,18 sur MNLI). De plus, l’activation dynamique d’experts basée sur un seuil réduit les paramètres du modèle et offre une nouvelle perspective pour atténuer le déséquilibre de charge entre experts.