Stratégie de répartition de charge sans perte auxiliaire pour Mixture-of-Experts

Pour les modèles Mixture-of-Experts (MoE), une charge d'expert déséquilibrée entraînera un effondrement de routage ou une augmentation de la charge computationnelle. Les méthodes existantes utilisent couramment une perte auxiliaire pour encourager l'équilibre de charge, mais une perte auxiliaire importante introduit des gradients d'interférence non négligeables dans l'entraînement, ce qui nuit à la performance du modèle. Afin de contrôler l'équilibre de charge sans produire de gradients indésirables pendant l'entraînement, nous proposons Loss-Free Balancing, caractérisé par une stratégie de répartition de charge sans perte auxiliaire. Plus précisément, avant la décision de routage top-K, Loss-Free Balancing applique d'abord un biais par expert aux scores de routage de chaque expert. En mettant à jour dynamiquement le biais de chaque expert selon sa charge récente, Loss-Free Balancing peut maintenir de manière cohérente une distribution équilibrée de la charge des experts. De plus, puisque Loss-Free Balancing ne produit aucun gradient d'interférence, il élève également la borne supérieure de la performance du modèle obtenue par l'entraînement MoE. Nous validons la performance de Loss-Free Balancing sur des modèles MoE allant jusqu'à 3 milliards de paramètres entraînés sur jusqu'à 200 milliards de tokens. Les résultats expérimentaux montrent que Loss-Free Balancing atteint à la fois une meilleure performance et un meilleur équilibre de charge comparé aux stratégies traditionnelles de répartition de charge contrôlées par perte auxiliaire.