LocMoE+ : Routeur amélioré avec sensibilisation aux caractéristiques des tokens pour un pré-entraînement efficace des LLM

Les architectures Mixture-of-Experts (MoE) ont récemment gagné en popularité dans le domaine des grands modèles de langage (LLM) grâce à leur capacité à réduire significativement les coûts d'entraînement et d'inférence. Cependant, ces architectures MoE présentent des défis, tels que d'importantes disparités dans le nombre de tokens assignés à chaque expert et une tendance à l'homogénéisation entre experts, ce qui affecte négativement les capacités de génération sémantique du modèle. Dans cet article, nous introduisons LocMoE+, une version raffinée du LocMoE à faible surcharge, incluant les améliorations suivantes : (1) Quantification et définition de l'affinité entre experts et tokens. (2) Mise en œuvre d'une stratégie de routage adaptatif au niveau global pour réorganiser les tokens en fonction de leurs scores d'affinité. (3) Réestimation de la borne inférieure de la capacité des experts, qui a montré une diminution progressive à mesure que la distribution des caractéristiques des tokens évolue. Les résultats expérimentaux démontrent que, sans compromettre la convergence ou l'efficacité du modèle, le nombre de tokens traités par chaque expert peut être réduit de plus de 60%. Combiné aux optimisations de communication, cela conduit à une amélioration moyenne de l'efficacité d'entraînement allant de 5,4 % à 46,6 %. Après fine-tuning, LocMoE+ affiche une amélioration de performance de 9,7 % à 14,1 % sur les ensembles de données GDAD, C-Eval et TeleQnA.