Le routage dans les Mixtures d'Experts (MoE) est un mécanisme central des architectures modernes de apprentissage automatique qui assigne dynamiquement chaque jeton d'entrée à un sous-ensemble de réseaux experts spécialisés. Au lieu d'activer tous les paramètres pour chaque entrée, un routeur, ou réseau de gating, apprend à diriger les jetons vers les experts les plus pertinents, équilibrant la capacité du modèle avec l'efficacité computationnelle. Cette approche permet aux grands modèles de langage de passer à l'échelle avec des billions de paramètres tout en maintenant des coûts d'inférence gérables, car seule une fraction du réseau est active par jeton.
Le concept provient des travaux du début des années 1990 sur les mixtures adaptatives d'experts locaux, mais il a gagné en importance dans les années 2010 avec l'introduction des couches MoE à gating épars dans l'apprentissage profond. Les implémentations modernes, comme celles utilisées dans les modèles de Google DeepMind et OpenAI, reposent sur des fonctions de routage apprises qui produisent des distributions de probabilité sur les experts, incorporant souvent des pertes auxiliaires pour garantir une utilisation équilibrée. Le routage est devenu un choix de conception critique dans les systèmes d'IA générative de pointe, influençant à la fois l'efficacité de l'entraînement et la qualité finale du modèle.
Développement Historique
L'idée de router les entrées vers des composants spécialisés précède l'apprentissage profond, avec des racines dans les méthodes d'ensemble et les réseaux neuronaux modulaires. En 1991, Michael I. Jordan et Robert Jacobs ont publié des travaux fondateurs sur les mixtures hiérarchiques d'experts, où un réseau de gating pondérait les sorties de plusieurs réseaux feedforward. Ce cadre précoce a établi le principe central : apprendre à partitionner l'espace d'entrée pour que différents experts traitent différentes régions.
En 2017, des chercheurs de Google DeepMind (alors Google Brain) ont introduit la couche MoE à gating épars dans un article dirigé par Noam Shazeer. Cette conception utilisait une fonction de gating softmax pour sélectionner les top-k experts pour chaque jeton, avec k généralement fixé à 1 ou 2. Les auteurs ont démontré que les couches MoE pouvaient passer à l'échelle avec des centaines d'experts, atteignant des résultats de pointe sur des benchmarks de modélisation du langage et de traduction automatique tout en réduisant le coût computationnel par exemple. Ces travaux ont directement influencé des architectures ultérieures comme le Switch Transformer (2021), qui a simplifié le routage pour sélectionner un seul expert par jeton, et le framework GShard pour l'entraînement massivement parallèle.
Mécanismes de Routage
Les fonctions de routage peuvent être catégorisées en plusieurs types. Le plus courant est le routage par choix de jeton, où chaque jeton sélectionne indépendamment ses top-k experts en fonction d'un score appris. Cela est implémenté comme une projection linéaire de l'état caché du jeton suivie d'un softmax sur les indices des experts. Les experts sélectionnés traitent ensuite le jeton, et leurs sorties sont pondérées par les probabilités de routage et sommées.
Une alternative est le routage par choix d'expert, où chaque expert sélectionne les top-k jetons d'un lot, garantissant que tous les experts reçoivent une charge minimale. Cette approche, popularisée dans le Switch Transformer et affinée plus tard dans des modèles comme Mixtral, résout les problèmes de déséquilibre de charge mais nécessite une gestion attentive des assignations jeton-expert.
D'autres variantes incluent le routage basé sur le hachage, qui utilise des fonctions de hachage déterministes pour assigner les jetons sans paramètres appris, et le routage hiérarchique, où un système de gating à deux niveaux sélectionne d'abord un groupe d'experts puis un expert spécifique au sein de ce groupe. Chaque méthode fait un compromis entre flexibilité, surcoût computationnel et stabilité de l'entraînement.
Équilibrage de Charge et Pertes Auxiliaires
Un défi majeur dans le routage MoE est le déséquilibre de charge : le routeur peut apprendre à toujours sélectionner quelques experts populaires, laissant les autres sous-entraînés. Pour atténuer cela, la plupart des implémentations ajoutent une perte auxiliaire d'équilibrage de charge qui pénalise la distribution inégale des jetons entre les experts. La formulation standard, introduite dans l'article de Shazeer de 2017, calcule la fraction de jetons routés vers chaque expert et encourage l'uniformité via un terme d'entropie croisée.
Des approches plus récentes, comme la perte d'équilibrage de charge du Switch Transformer, utilisent un coefficient plus simple qui multiplie le produit scalaire entre la probabilité moyenne du routeur et le nombre réel de jetons par expert. Certains systèmes, comme GShard, emploient des facteurs de capacité qui limitent le nombre de jetons que chaque expert peut traiter, forçant le routeur à distribuer le travail. Ces techniques sont essentielles pour un entraînement stable et pour prévenir l'effondrement des experts, où certains experts deviennent des poids morts.
Intégration Architecturale
Le routage est généralement appliqué dans les blocs transformer, remplaçant le réseau feedforward (FFN) par une couche MoE. Dans un Transformer (architecture) standard, chaque jeton passe par une sous-couche d'attention multi-têtes suivie d'un FFN positionnel. Dans un transformer MoE, le FFN est remplacé par un ensemble de FFN experts, chacun avec ses propres paramètres, et un routeur qui sélectionne les experts à activer.
Cette intégration permet aux modèles d'avoir un nombre total de paramètres important tout en maintenant constant le nombre de paramètres actifs par jeton. Par exemple, un modèle avec 64 experts, chacun avec 1 milliard de paramètres, a 64 milliards de paramètres totaux mais n'active que 2 milliards par jeton si k=2. Cette propriété est cruciale pour passer à l'échelle avec des modèles à des billions de paramètres, comme on le voit dans les systèmes de Mixture d'Experts de Anthropic et d'autres laboratoires.
Le routeur lui-même est un petit réseau neuronal, souvent une seule couche linéaire suivie d'un softmax, avec ses propres paramètres entraînés conjointement avec les experts. Certaines architectures utilisent des routeurs séparés pour différentes couches ou têtes d'attention, et des travaux récents explorent des politiques de routage apprises qui s'adaptent en fonction du type de jeton ou de sa position.
Dynamique d'Entraînement
L'entraînement de modèles MoE avec routage introduit des défis uniques. Les décisions discrètes du routeur (sélection des top-k experts) ne sont pas différentiables, donc les gradients ne circulent que à travers les sorties des experts sélectionnés pondérées par les probabilités de routage. Cela crée un problème de cible mobile : à mesure que les experts s'améliorent, les préférences du routeur changent, ce qui peut provoquer une instabilité.
Pour y remédier, les chercheurs utilisent des techniques comme le gating top-k bruité, où un bruit gaussien est ajouté aux logits de routage pendant l'entraînement pour encourager l'exploration. Une autre approche consiste à utiliser un estimateur à passage direct pour la décision de routage, traitant la sélection comme une assignation dure mais faisant passer les gradients à travers les poids de probabilité. De plus, certaines méthodes ajustent la température du softmax pendant l'entraînement pour affiner progressivement les décisions de routage.
Les pertes d'équilibrage de charge sont généralement pondérées avec un petit coefficient (par exemple, 0,01) pour éviter de dominer la perte de tâche principale. En pratique, l'entraînement de modèles MoE nécessite un réglage minutieux des hyperparamètres, et de nombreux systèmes emploient des pertes auxiliaires pour l'équilibrage de charge et la confiance du routeur.
Applications dans les Grands Modèles de Langage
Le routage dans les MoE est devenu une technique standard dans les grands modèles de langage. Des exemples notables incluent :
- Le Switch Transformer (2021) de Google, qui a atteint 1,6 billion de paramètres avec une architecture MoE éparse.
- GShard (2020), qui a démontré un entraînement efficace de modèles MoE sur des milliers de cœurs TPU.
- Mixtral 8x7B (2023) de Mistral AI, qui utilise 8 experts avec routage top-2, atteignant des performances comparables à des modèles denses plus grands.
- DeepSeek-V3 (2024), qui emploie un MoE à grain fin avec 256 experts et routage top-8.
- Plusieurs modèles de OpenAI et Anthropic sont rapportés comme utilisant des couches MoE, bien que les détails exacts soient souvent propriétaires.
Ces modèles démontrent que le routage permet une mise à l'échelle rentable, car le calcul par jeton reste constant même lorsque les paramètres totaux augmentent. Cela a fait des MoE un choix privilégié pour servir de grands modèles en environnement de production, y compris sur des plateformes cloud comme Amazon Web Services, Microsoft Azure et Google Cloud.
Considérations d'Efficacité et Matériel
Le routage introduit un surcoût de communication dans l'entraînement et l'inférence distribués, car les jetons doivent être envoyés à l'expert approprié, qui peut résider sur un dispositif différent. Ce schéma de communication all-to-all peut devenir un goulot d'étranglement, surtout avec de nombreux experts. Les fabricants de matériel comme NVIDIA (bien que non listé, implicite) et AMD ont développé des noyaux optimisés pour les opérations MoE, et des accélérateurs IA spécialisés de Groq et SambaNova sont conçus pour gérer efficacement les schémas d'activation épars.
Pour réduire la communication, certains systèmes utilisent le parallélisme d'experts, où les experts sont répliqués sur plusieurs dispositifs, ou des approches hybrides combinant parallélisme de données et d'experts. Le choix de la granularité du routage (au niveau du jeton vs. au niveau du bloc) affecte également l'efficacité. Le routage au niveau du jeton offre un contrôle plus fin mais plus de communication, tandis que le routage au niveau du bloc regroupe les jetons pour réduire le surcoût.
Défis et Directions Futures
Malgré son succès, le routage dans les MoE fait face à plusieurs problèmes ouverts. L'un est le compromis entre spécialisation des experts et généralisation : des experts trop spécialisés peuvent ne pas bien se transférer à de nouvelles tâches. Un autre est la difficulté de l'ajustement fin des modèles MoE, car le routeur peut avoir besoin de s'adapter à de nouvelles distributions de données. La recherche sur le routage adaptatif, où le nombre d'experts actifs varie par jeton, et sur les politiques de routage apprises qui considèrent le contexte au-delà du jeton courant, est en cours.
Les directions futures incluent la combinaison du routage avec d'autres techniques d'efficacité comme la élagage de modèle et la quantification, et le développement de routeurs capables de gérer des entrées multimodales. À mesure que les modèles continuent de passer à l'échelle, le routage restera un mécanisme clé pour équilibrer capacité et calcul, avec des applications potentielles au-delà du langage, dans la vision et l'apprentissage par renforcement.
Voir Aussi
- Transformer (architecture)
- Large language model
- Multi-Head Attention
- Model Pruning
- Google DeepMind
- OpenAI
- Anthropic
Références
- Shazeer, N., et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.
- Fedus, W., et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.
- Lepikhin, D., et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.
- Jiang, A. Q., et al. (2023). Mixtral of Experts.
- DeepSeek-AI (2024). DeepSeek-V3 Technical Report.