Mélange d'experts

Réseau méta-pi

Le réseau méta-pi, présenté par Hampshire et Waibel^[4], utilise

${\displaystyle f(x)=\sum _{i}w(x)_{i}f_{i}(x)}$

comme sortie. Le modèle est entraîné par descente de gradient sur l'erreur quadratique moyenne

${\displaystyle L:={\frac {1}{N}}\sum _{k}\|y_{k}-f(x_{k})\|^{2}}$.

Les experts peuvent être des fonctions arbitraires.

Dans leur publication initiale, les auteurs cherchaient à classer des phonèmes dans des signaux de parole produits par six locuteurs japonais, deux femmes et quatre hommes. Ils entraînèrent six experts, chacun constitué d'un « réseau neuronal à délais temporels »^[5], c'est-à-dire, pour l'essentiel, d'un réseau de neurones convolutif multicouche appliqué à un spectrogramme de Mel. Le mélange obtenu attribuait cinq experts à cinq des locuteurs. Le sixième locuteur, un homme, ne disposait pas d'expert propre : sa voix était classée au moyen d'une combinaison linéaire des experts correspondant aux trois autres locuteurs masculins.

Mélanges adaptatifs d'experts locaux

Les mélanges adaptatifs d'experts locaux^[6],[7] utilisent un modèle de mélange gaussien. Chaque expert prédit simplement une loi gaussienne et ignore entièrement l'entrée. Plus précisément, le ${\displaystyle i}$-ème expert prédit

${\displaystyle y\sim N(\mu _{i},I)}$,

où ${\displaystyle \mu _{i}}$ est un paramètre appris. La fonction de pondération est une fonction linéaire suivie d'un softmax :

$${\displaystyle w(x)_{i}={\frac {e^{k_{i}^{T}x+b_{i}}}{\sum _{j}e^{k_{j}^{T}x+b_{j}}}}}$$

Le mélange d'experts prédit que la sortie est distribuée selon la densité logarithmique :

$${\displaystyle \ln f_{\theta }(y|x)=\ln \left[\sum _{i}{\frac {e^{k_{i}^{T}x+b_{i}}}{\sum _{j}e^{k_{j}^{T}x+b_{j}}}}N(y|\mu _{i},I)\right]=\ln \left[(2\pi )^{-d/2}\sum _{i}{\frac {e^{k_{i}^{T}x+b_{i}}}{\sum _{j}e^{k_{j}^{T}x+b_{j}}}}e^{-{\frac {1}{2}}\|y-\mu _{i}\|^{2}}\right]}$$

Le modèle est entraîné par maximum de vraisemblance, c'est-à-dire par montée de gradient sur ${\displaystyle f(y|x)}$. Le gradient relatif au ${\displaystyle i}$-ème expert est :

$${\displaystyle \nabla _{\mu _{i}}\ln f_{\theta }(y|x)={\frac {w(x)_{i}N(y|\mu _{i},I)}{\sum _{j}w(x)_{j}N(y|\mu _{j},I)}}\;(y-\mu _{i})}$$

et le gradient relatif à la fonction de pondération est :

$${\displaystyle \nabla _{[k_{i},b_{i}]}\ln f_{\theta }(y|x)={\begin{bmatrix}x\\1\end{bmatrix}}{\frac {w(x)_{i}}{\sum _{j}w(x)_{j}N(y|\mu _{j},I)}}(f_{i}(x)-f_{\theta }(y|x))}$$

Pour chaque couple entrée-sortie ${\displaystyle (x,y)}$, la fonction de pondération est modifiée de façon à accroître le poids des experts ayant obtenu une performance supérieure à la moyenne et à réduire celui des experts moins performants. Elle apprend ainsi à sélectionner les experts qui prédisent correctement chaque type d'entrée.

Le ${\displaystyle i}$-ème expert est modifié pour rapprocher sa prédiction de ${\displaystyle y}$, mais l'amplitude de cette modification est proportionnelle à ${\displaystyle w(x)_{i}N(y|\mu _{i},I)}$. Cette quantité admet une interprétation bayésienne. Pour une entrée ${\displaystyle x}$, la probabilité a priori que l'expert ${\displaystyle i}$ soit le bon vaut ${\displaystyle w(x)_{i}}$, tandis que ${\displaystyle N(y|\mu _{i},I)}$ est la vraisemblance de l'observation ${\displaystyle y}$. La quantité

${\displaystyle {\frac {w(x)_{i}N(y|\mu _{i},I)}{\sum _{j}w(x)_{j}N(y|\mu _{j},I)}}}$

est donc la probabilité a posteriori de l'expert ${\displaystyle i}$. La vitesse d'apprentissage de cet expert est proportionnelle à cette probabilité.

Autrement dit, les experts qui apparaissent rétrospectivement comme les plus pertinents pour un exemple donné sont entraînés sur cet exemple, tandis que les autres sont peu ou pas modifiés.

L'effet combiné de ce mécanisme conduit à la spécialisation des experts. Si deux experts prédisent correctement un certain type d'entrée, mais que l'un est légèrement meilleur, la fonction de pondération finit par favoriser ce dernier. L'expert moins performant reçoit alors un signal de gradient plus faible et se spécialise progressivement dans d'autres catégories d'entrées. Cette rétroaction positive éloigne les experts les uns des autres et amène chacun d'eux à prendre en charge une région locale de l'espace des données, d'où l'expression « experts locaux ».

Mélange d'experts hiérarchique

Les mélanges d'experts hiérarchiques^[8],[9] utilisent plusieurs niveaux de fonctions de routage organisés en arbre. Chaque fonction de routage définit une distribution de probabilité sur les fonctions du niveau suivant, tandis que les experts se trouvent aux feuilles de l'arbre. Leur structure est proche de celle des arbres de décision.

Par exemple, un mélange hiérarchique à deux niveaux possède une fonction de routage de premier niveau ${\displaystyle w_{i}}$, des fonctions de second niveau ${\displaystyle w_{j|i}}$ et des experts ${\displaystyle f_{j|i}}$. La prédiction totale est alors :

${\displaystyle \sum _{i}w_{i}(x)\sum _{j}w_{j|i}(x)f_{j|i}(x)}$.

Variantes

Comme les modèles de mélange gaussien, les mélanges d'experts peuvent être entraînés au moyen de l'algorithme espérance-maximisation. Lors de l'étape d'espérance, la responsabilité d'expliquer chaque donnée est répartie entre les experts. Lors de l'étape de maximisation, chaque expert améliore l'explication des données dont il porte une responsabilité élevée, tandis que la fonction de routage améliore l'attribution de ces responsabilités. Cette méthode peut converger plus rapidement qu'une montée de gradient directe sur la log-vraisemblance^[9],[10].

La fonction de routage est souvent un softmax. Elle peut aussi reposer sur des lois gaussiennes^[11] ou, plus généralement, sur des familles exponentielles^[10].

Dans un mélange d'experts « dur » (hard MoE)^[12], seule la sortie de l'expert le mieux classé est retenue, au lieu d'effectuer une somme pondérée de tous les experts :

${\displaystyle f(x)=f_{\arg \max _{i}w_{i}(x)}(x)}$.

Cette méthode peut réduire le coût de l'entraînement et de l'inférence^[13].

Les experts peuvent utiliser des formes plus générales de lois gaussiennes multivariées. Jordan et Jacobs^[8] proposent par exemple

${\displaystyle f_{i}(y|x)=N(y|A_{i}x+b_{i},\Sigma _{i})}$,

où ${\displaystyle A_{i}}$, ${\displaystyle b_{i}}$ et ${\displaystyle \Sigma _{i}}$ sont appris. Chaque expert réalise ainsi une régression linéaire accompagnée d'une estimation apprise de l'incertitude.

D'autres distributions que la loi gaussienne peuvent être utilisées, notamment la loi de Laplace^[14], ou la loi de Student^[15]. Pour la classification binaire, il est également possible d'employer des experts de régression logistique :

$${\displaystyle f_{i}(y|x)={\begin{cases}{\frac {1}{1+e^{\beta _{i}^{T}x+\beta _{i,0}}}},&y=0\\1-{\frac {1}{1+e^{\beta _{i}^{T}x+\beta _{i,0}}}},&y=1\end{cases}}}$$

où ${\displaystyle \beta _{i}}$ et ${\displaystyle \beta _{i,0}}$ sont appris. Cette approche a ensuite été généralisée à la classification multiclasse au moyen d'experts de régression logistique multinomiale^[16].

Une autre variante, appelée « mélange de softmax », a été proposée pour la modélisation autorégressive du langage^[17]. Dans un modèle de langue, un texte antérieur ${\displaystyle c}$ est encodé sous la forme d'un vecteur ${\displaystyle v_{c}}$, puis la probabilité du mot suivant est prédite par ${\displaystyle \mathrm {Softmax} (v_{c}W)}$, où ${\displaystyle W}$ est une matrice de plongement. Dans un mélange de softmax, le modèle produit plusieurs vecteurs ${\displaystyle v_{c,1},\dots ,v_{c,n}}$ et prédit le mot suivant selon

${\displaystyle \sum _{i=1}^{n}p_{i}\,\mathrm {Softmax} (v_{c,i}W_{i})}$,

où ${\displaystyle p_{i}}$ est une distribution de probabilité obtenue par une opération linéaire suivie d'un softmax sur les activations cachées du modèle. L'article d'origine montre l'efficacité de cette méthode pour les réseaux de neurones récurrents. Elle s'est ensuite révélée applicable aux transformeurs^[18].

Couche de mélange d'experts à routage parcimonieux

La couche de mélange d'experts à routage parcimonieux^[21], publiée par des chercheurs de Google Brain, emploie des réseaux à propagation avant comme experts et un routage linéaire suivi d'un softmax. Comme dans les mélanges d'experts « durs », la parcimonie est obtenue en effectuant une somme pondérée des seuls ${\displaystyle k}$ experts les mieux classés.

Une couche MoE contient des experts ${\displaystyle f_{1},\dots ,f_{n}}$ et un réseau de routage ${\displaystyle w}$. Celui-ci est défini par

${\displaystyle w(x)=\mathrm {softmax} (\mathrm {top} _{k}(Wx+{\text{bruit}}))}$,

où ${\displaystyle \mathrm {top} _{k}}$ conserve les ${\displaystyle k}$ plus grandes composantes d'un vecteur et remplace toutes les autres par ${\displaystyle -\infty }$. L'ajout de bruit contribue à équilibrer la charge entre les experts.

Le choix de ${\displaystyle k}$ est un hyperparamètre dépendant de l'application. Les valeurs usuelles sont ${\displaystyle k=1}$ ou ${\displaystyle k=2}$. La variante ${\displaystyle k=1}$ est également appelée Switch Transformer. Le premier Switch Transformer a été appliqué à un modèle de langue T5^[22].

À titre de démonstration, les auteurs entraînèrent plusieurs modèles de traduction automatique alternant couches MoE et couches LSTM, puis les comparèrent à des modèles LSTM profonds^[23]. Les modèles MoE demandaient moins de calcul lors de l'inférence malgré un nombre de paramètres environ trente fois supérieur.

Ce module architectural a été publié en janvier 2017, quelques mois avant l'architecture Transformeur en juin 2017. Les deux furent combinés dans une architecture multimodale appelée MultiModel, publiée quatre jours après l'article fondateur sur les transformeurs^[24].

Équilibrage de charge

Les mélanges d'experts simples rencontrent souvent un problème d'équilibrage : certains experts sont très fréquemment sollicités, tandis que d'autres le sont rarement, voire jamais. Pour encourager le routeur à sélectionner chaque expert à fréquence comparable dans un lot, les couches MoE peuvent inclure des fonctions de coût auxiliaires. Le Switch Transformer^[22] regroupe ce mécanisme en une seule fonction auxiliaire.

Soit ${\displaystyle n}$ le nombre d'experts et ${\displaystyle \{x_{1},\dots ,x_{T}\}}$ un lot de requêtes. La perte auxiliaire vaut :

$${\displaystyle n\sum _{i=1}^{n}f_{i}P_{i}}$$

où

${\displaystyle f_{i}={\frac {1}{T}}\#({\text{requêtes envoyées à l'expert }}i)}$

est la proportion de jetons ayant choisi l'expert ${\displaystyle i}$, et

${\displaystyle P_{i}={\frac {1}{T}}\sum _{j=1}^{T}{\frac {w_{i}(x_{j})}{\sum _{i'\in {\text{experts}}}w_{i'}(x_{j})}}}$

est la proportion de poids attribuée à cet expert. Cette fonction atteint son minimum, égal à ${\displaystyle 1}$, lorsque tous les experts reçoivent le même poids ${\displaystyle 1/n}$ dans toutes les situations.

Les chercheurs de DeepSeek ont conçu une variante comprenant des « experts partagés », toujours sollicités, et des « experts routés », qui ne le sont que lorsque le routeur les sélectionne. Selon eux, l'équilibrage standard pousse les experts à reproduire des capacités fondamentales identiques, par exemple la grammaire anglaise. Les experts partagés sont donc chargés d'apprendre les capacités générales fréquemment utilisées, tandis que les experts routés apprennent des capacités plus périphériques ou spécialisées^[26].

DeepSeek a également proposé une stratégie d'équilibrage sans perte auxiliaire. Chaque expert ${\displaystyle i}$ reçoit un biais supplémentaire ${\displaystyle b_{i}}$. Si un expert est trop peu utilisé, ce biais augmente ; dans le cas contraire, il diminue. Pour l'affectation des jetons, chaque jeton choisit les ${\displaystyle k}$ experts obtenant les plus grands scores après ajout de ce biais :

$${\displaystyle f(x)=\sum _{i{\text{ appartient au top-}}k{\text{ de }}\{w(x)_{j}+b_{j}\}_{j}}w(x)_{i}f_{i}(x)}$$

Le biais intervient dans le choix des experts, mais pas dans la pondération de leurs sorties^[27].

Facteur de capacité

Supposons qu'une couche comporte ${\displaystyle n}$ experts et qu'un lot contienne ${\displaystyle T}$ requêtes. Si chaque requête est envoyée vers un expert, comme dans le Switch Transformer, et si la charge est parfaitement équilibrée, chaque expert reçoit en moyenne ${\displaystyle T/n}$ requêtes. En pratique, l'équilibrage n'est jamais parfait : un expert peut être peu sollicité dans certains lots et surchargé dans d'autres.

Comme les entrées ne peuvent quitter la couche qu'une fois tous les experts terminés, l'équilibrage est important. Un facteur de capacité peut imposer une contrainte stricte : chaque expert ne peut traiter que

${\displaystyle c\cdot T/n}$

requêtes au maximum par lot. Le rapport ST-MoE indique que des valeurs de ${\displaystyle c}$ comprises entre ${\displaystyle 1{,}25}$ et ${\displaystyle 2}$ fonctionnent bien en pratique^[28].

Routage

Dans le mélange d'experts parcimonieux d'origine, seuls les ${\displaystyle k}$ experts les mieux classés sont interrogés et leurs sorties sont additionnées avec pondération. D'autres méthodes existent^[28]. Le routage est un problème d'affectation : il s'agit d'associer les jetons aux experts tout en respectant différentes contraintes, notamment le débit de calcul et l'équilibrage de charge.

On distingue généralement trois familles :

• les experts choisissent les jetons (expert choice) ;^[29]
• les jetons choisissent les experts, comme dans le mélange parcimonieux d'origine ;
• un algorithme global apparie experts et jetons^[30].

Lors de l'inférence, un mélange d'experts traite simultanément un grand lot de jetons. Si les jetons choisissent les experts, certains experts peuvent recevoir très peu de jetons tandis que d'autres dépassent leur capacité maximale, ce qui oblige à ignorer une partie des jetons. Inversement, lorsque les experts choisissent les jetons, certains jetons peuvent n'être sélectionnés par aucun expert. Ce phénomène est appelé « abandon de jetons » (token drop). Dans un transformeur, il n'est pas nécessairement catastrophique : en raison des connexions résiduelles, un jeton abandonné ne disparaît pas ; sa représentation traverse simplement la couche à propagation avant sans modification^[30].

D'autres méthodes traitent le routage comme un problème de programmation linéaire sous contraintes^[31], ou utilisent l'apprentissage par renforcement pour apprendre l'algorithme de routage, le choix d'un expert étant une action discrète^[32]. L'appariement peut aussi ne comporter aucun apprentissage : ce « routage statique » repose alors sur une fonction de hachage déterministe^[33] ou sur un générateur aléatoire^[34].

Applications aux transformeurs

Les couches MoE sont employées dans les plus grands modèles transformeurs, pour lesquels entraîner et exécuter l'ensemble des paramètres serait trop coûteux. Elles sont généralement parcimonieuses, avec un ou deux experts activés par jeton. Dans les transformeurs, les couches MoE servent souvent à remplacer ou sélectionner les couches à propagation avant, généralement de type linéaire–ReLU–linéaire, situées après l'attention multi-têtes dans chaque bloc. Ces couches représentent une part croissante du coût de calcul lorsque la taille des modèles augmente. Dans PaLM-540B, par exemple, environ 90 % des paramètres appartiennent aux couches à propagation avant^[35].

Un transformeur dense déjà entraîné peut être converti en mélange d'experts en dupliquant ses couches à propagation avant, en initialisant aléatoirement le routeur, puis en poursuivant l'entraînement. Cette technique est appelée sparse upcycling^[36].

Les transformeurs MoE présentent de nombreux choix de conception influant sur la stabilité de l'entraînement et les performances finales. Le rapport OLMoE en propose une analyse détaillée^[37].

À partir de 2023, les modèles suffisamment grands pour recourir à des mélanges d'experts sont le plus souvent de grands modèles de langage, chaque expert pouvant comporter de l'ordre de dix milliards de paramètres. La méthode a aussi été appliquée à la vision : Vision MoE^[38] est un transformeur doté de couches MoE, dont une version à 15 milliards de paramètres a été entraînée. Des transformeurs MoE ont également été employés dans des modèles de diffusion^[39].

Plusieurs grands modèles de langue de Google reposent sur cette architecture. GShard^[40] active jusqu'à deux experts par couche. Le premier expert est toujours sélectionné, tandis que le second l'est avec une probabilité proportionnelle au poids que lui attribue le routeur. GLaM^[41] a ensuite atteint 1,2 billion de paramètres, chaque couche MoE activant deux experts parmi 64. Les Switch Transformers n'en activent qu'un seul.

NLLB-200, développé par Meta AI, est un modèle de traduction automatique couvrant 200 langues^[42]. Chaque couche MoE utilise une hiérarchie à deux niveaux. Au premier niveau, le routeur choisit entre une couche à propagation avant partagée et un ensemble d'experts. Dans ce second cas, un autre routeur calcule les poids et sélectionne les deux meilleurs experts^[43].

Les grands modèles de langue MoE peuvent être adaptés à des tâches en aval au moyen de l'ajustement par instructions^[44].

En décembre 2023, Mistral AI publie Mixtral 8x7B sous licence Apache 2.0. Ce modèle de langue MoE comporte 46,7 milliards de paramètres, huit experts et une parcimonie de deux experts actifs par jeton. Une version ajustée pour le suivi d'instructions est également publiée^[45],[46].

En mars 2024, Databricks publie DBRX, un modèle MoE de 132 milliards de paramètres, comportant seize experts et une parcimonie de quatre experts actifs. Une version ajustée pour le suivi d'instructions est également proposée^[47],[48].