CompeteSMoE -- Entraînement de Mélange d'Experts à Garantie Statistique via Compétition

CompeteSMoE -- Statistically Guaranteed Mixture of Experts Training via Competition

May 19, 2025

Auteurs: Nam V. Nguyen, Huy Nguyen, Quang Pham, Van Nguyen, Savitha Ramasamy, Nhat Ho

cs.AI

Résumé

Le mélange parcimonieux d'experts (Sparse Mixture of Experts, SMoE) propose une solution attrayante pour augmenter la complexité du modèle au-delà des méthodes traditionnelles consistant à accroître la profondeur ou la largeur du réseau. Cependant, nous soutenons que l'entraînement efficace des SMoE reste un défi en raison d'un processus de routage sous-optimal, où les experts effectuant les calculs ne contribuent pas directement au processus de routage. Dans ce travail, nous proposons la compétition, un nouveau mécanisme pour router les tokens vers les experts présentant la réponse neuronale la plus élevée. Théoriquement, nous montrons que le mécanisme de compétition offre une meilleure efficacité d'échantillonnage que le routage traditionnel par softmax. De plus, nous développons CompeteSMoE, un algorithme simple mais efficace pour entraîner des modèles de langage de grande taille en déployant un routeur pour apprendre la politique de compétition, permettant ainsi d'obtenir de solides performances avec une faible surcharge d'entraînement. Nos évaluations empiriques approfondies sur les tâches de réglage d'instructions visuelles et de pré-entraînement linguistique démontrent l'efficacité, la robustesse et l'évolutivité de CompeteSMoE par rapport aux stratégies SMoE de pointe. Nous avons rendu l'implémentation disponible à l'adresse suivante : https://github.com/Fsoft-AIC/CompeteSMoE. Ce travail est une version améliorée de l'étude précédente disponible sur arXiv:2402.02526.

English

Sparse mixture of experts (SMoE) offers an appealing solution to scale up the model complexity beyond the mean of increasing the network's depth or width. However, we argue that effective SMoE training remains challenging because of the suboptimal routing process where experts that perform computation do not directly contribute to the routing process. In this work, we propose competition, a novel mechanism to route tokens to experts with the highest neural response. Theoretically, we show that the competition mechanism enjoys a better sample efficiency than the traditional softmax routing. Furthermore, we develop CompeteSMoE, a simple yet effective algorithm to train large language models by deploying a router to learn the competition policy, thus enjoying strong performances at a low training overhead. Our extensive empirical evaluations on both the visual instruction tuning and language pre-training tasks demonstrate the efficacy, robustness, and scalability of CompeteSMoE compared to state-of-the-art SMoE strategies. We have made the implementation available at: https://github.com/Fsoft-AIC/CompeteSMoE. This work is an improved version of the previous study at arXiv:2402.02526