Un réseau de routage, également connu sous le nom de fonction de portage ou fonction de pondération, est un composant des modèles d'apprentissage automatique à mélange d'experts (MoE). Il prend une entrée et produit un ensemble de poids qui déterminent la contribution de chaque réseau expert à la sortie. Dans les grands modèles de langage modernes, le réseau de routage est souvent utilisé pour ne sélectionner qu'un petit sous-ensemble d'experts pour chaque jeton d'entrée, une technique appelée activation éparse, qui réduit le coût de calcul tout en maintenant la capacité du modèle.
Le concept de routage dans les MoE remonte aux premiers travaux sur les machines à comité et les mélanges adaptatifs d'experts locaux. Dans une architecture MoE typique, il existe plusieurs réseaux experts, chacun pouvant être spécialisé dans différentes régions de l'espace d'entrée, et un réseau de routage qui apprend à assigner les entrées aux experts appropriés. La sortie du routeur peut être une distribution de probabilité sur les experts, ou un ensemble de poids non négatifs qui sont combinés avec les sorties des experts via une somme pondérée.
Architecture de base
Dans une configuration MoE standard, le réseau de routage est une fonction \( w(x) \) qui mappe une entrée \( x \) vers un vecteur de poids \( (w(x)_1, ..., w(x)_n) \), où \( n \) est le nombre d'experts. La sortie finale est calculée comme \( f(x) = \sum_{i=1}^n w(x)_i f_i(x) \), où \( f_i \) est le \( i \)-ème expert. Le routeur et les experts sont entraînés conjointement, généralement par descente de gradient sur une fonction de perte telle que l'erreur quadratique moyenne ou l'entropie croisée.
Le réseau de routage lui-même est souvent un petit réseau neuronal, comme une couche linéaire suivie d'une activation softmax. Dans les modèles MoE épars, le routeur peut utiliser un mécanisme de sélection top-k, où seuls les experts avec les poids les plus élevés sont activés, et les autres sont ignorés. Cette éparsité est essentielle pour réduire le coût de calcul, car elle permet au modèle d'utiliser seulement une fraction de ses paramètres totaux pour chaque entrée.
Développement historique
L'idée d'utiliser un réseau de portage pour combiner plusieurs experts a été explorée dans les années 1990. Un exemple précoce est le réseau méta-pi, rapporté par Hampshire et Waibel, qui utilisait une somme pondérée des sorties des experts et était entraîné sur une tâche de classification de phonèmes. Dans leurs expériences, ils ont entraîné six experts, chacun un réseau neuronal à délais temporels, pour classer la parole de six locuteurs japonais. Ils ont observé que le réseau de routage apprenait à dédier cinq experts à cinq locuteurs individuels, tandis que la voix du sixième locuteur était gérée par une combinaison d'experts pour d'autres locuteurs masculins.
Un autre modèle précoce influent était les mélanges adaptatifs d'experts locaux, qui utilisaient un modèle de mélange gaussien. Dans cette approche, chaque expert prédisait une distribution gaussienne avec une covariance fixe, et le routeur était une fonction linéaire-softmax qui assignait des poids en fonction de l'entrée. Ce modèle a démontré que le routage pouvait être utilisé pour partitionner l'espace d'entrée en régions homogènes, chaque expert se spécialisant dans une région locale.
Rôle dans les grands modèles de langage modernes
Dans les années 2020, les réseaux de routage sont devenus un composant central des grands modèles de langage (LLM) qui utilisent des couches MoE. Des modèles tels que ceux développés par des entreprises comme Google DeepMind et OpenAI ont adopté des architectures MoE éparses pour augmenter le nombre de paramètres sans augmenter proportionnellement le coût d'inférence. Dans ces modèles, chaque couche de transformeur peut contenir plusieurs experts, et un réseau de routage sélectionne quelques experts par jeton.
Par exemple, dans un LLM basé sur un transformeur, les plongements de jetons d'entrée sont passés à travers un réseau de routage qui calcule une distribution de probabilité sur les experts. Le routeur sélectionne ensuite les top-k experts (par exemple, k=2 ou k=4) et calcule la sortie comme une somme pondérée des sorties de ces experts. Cela permet au modèle d'avoir des milliards de paramètres tout en n'activant qu'un petit sous-ensemble pour chaque jeton, rendant l'entraînement et l'inférence plus efficaces.
Activation éparse et équilibrage de charge
L'activation éparse est un avantage clé des réseaux de routage dans les modèles MoE. En n'activant que quelques experts par entrée, le modèle peut avoir un nombre total de paramètres important mais un coût de calcul effectif beaucoup plus faible. Cependant, cela introduit des défis en matière d'équilibrage de charge : si le routeur sélectionne systématiquement les mêmes quelques experts, ceux-ci deviennent surchargés tandis que d'autres restent sous-utilisés. Pour remédier à cela, les modèles MoE modernes incluent souvent des fonctions de perte auxiliaires qui encouragent le routeur à distribuer les jetons plus uniformément entre les experts.
Diverses techniques ont été proposées pour améliorer les performances du routeur, comme l'ajout de bruit aux logits du routeur pendant l'entraînement pour favoriser l'exploration, ou l'utilisation d'une sélection top-k différentiable. Certains modèles utilisent également un schéma de routage hiérarchique, où un routeur de premier niveau sélectionne un groupe d'experts, et un routeur de second niveau sélectionne au sein de ce groupe.
Relation avec d'autres techniques
Les réseaux de routage sont étroitement liés à d'autres concepts en apprentissage automatique, tels que attention multi-têtes et attention croisée, qui impliquent également la pondération de différents composants en fonction de l'entrée. Cependant, alors que les mécanismes d'attention pondèrent différentes parties de la séquence d'entrée, un réseau de routage pondère différents réseaux experts, qui sont généralement des approximateurs de fonctions indépendants.
Les réseaux de routage partagent également des similitudes avec élagage de modèle dans la mesure où les deux visent à réduire le coût de calcul, mais l'élagage supprime les paramètres de manière permanente, tandis que le routage sélectionne dynamiquement les paramètres à utiliser pour chaque entrée. De plus, les réseaux de routage peuvent être considérés comme une forme d'apprentissage d'ensemble, car ils combinent les sorties de plusieurs modèles, mais contrairement aux ensembles traditionnels, les experts sont entraînés conjointement avec le routeur.
Implémentation et entraînement
En pratique, les réseaux de routage sont implémentés comme une couche linéaire qui prend la représentation d'entrée et produit des logits pour chaque expert. Les logits sont ensuite passés à travers une fonction softmax pour produire des poids. Pendant l'entraînement, le routeur et les experts sont mis à jour conjointement par rétropropagation. La fonction de perte inclut généralement à la fois la perte de tâche (par exemple, l'entropie croisée pour la modélisation du langage) et une perte auxiliaire d'équilibrage de charge.
Un défi dans l'entraînement des réseaux de routage est que la sélection discrète des experts (par exemple, top-k) n'est pas différentiable. Pour surmonter cela, de nombreuses implémentations utilisent directement les poids softmax pendant l'entraînement, ou utilisent un estimateur de passage direct pour la sélection top-k. Certains modèles utilisent également un réseau de routage séparé pour chaque couche, permettant à différentes couches de se spécialiser dans différents types de décisions de routage.
Applications au-delà des modèles de langage
Bien que les réseaux de routage soient plus proéminents dans les LLM, ils sont également utilisés dans d'autres domaines. Par exemple, en vision par ordinateur, des couches MoE avec routeurs ont été appliquées à des tâches de classification et de génération d'images. En reconnaissance vocale, les premiers modèles MoE utilisaient des routeurs pour partitionner l'espace acoustique. Les réseaux de routage sont également utilisés en apprentissage par renforcement, où différents experts peuvent représenter différentes politiques, et le routeur sélectionne la politique appropriée en fonction de l'état.
Dans le contexte de la recherche en intelligence artificielle, les réseaux de routage sont un domaine d'étude actif, avec des travaux en cours sur l'amélioration de l'efficacité du routage, de l'interprétabilité et de l'évolutivité. Des chercheurs d'institutions telles que Stanford AI Lab et Berkeley AI Research ont contribué à la compréhension du comportement des routeurs dans les modèles à grande échelle.
Défis et orientations futures
Malgré leur succès, les réseaux de routage font face à plusieurs défis. Un problème est que le routeur peut devenir un goulot d'étranglement, car il doit traiter chaque entrée et prendre des décisions de routage rapidement. Un autre défi est que le routeur peut apprendre à router les entrées de manières qui ne sont pas optimales pour la tâche globale, entraînant une dégradation des performances. De plus, la perte auxiliaire d'équilibrage de charge peut interférer avec la perte de tâche principale, nécessitant un réglage minutieux.
Les recherches futures pourraient explorer des mécanismes de routage plus sophistiqués, tels que des politiques de routage apprises qui considèrent le contenu de l'entrée de manière plus nuancée, ou un routage hiérarchique qui réduit le nombre d'experts considérés à chaque niveau. Il y a également un intérêt à rendre les routeurs plus interprétables, afin de pouvoir comprendre pourquoi une entrée particulière est routée vers un certain expert.
En résumé, le réseau de routage est un composant fondamental des modèles MoE, permettant une mise à l'échelle et une spécialisation efficaces. Son développement, des premières fonctions de portage aux routeurs épars modernes dans les LLM, illustre l'évolution des techniques d'apprentissage automatique sur plusieurs décennies.