Aller au contenu

Mixture-of-Experts

Le constat de départ

Plus un modèle a de paramètres, mieux il apprend. Mais plus il a de paramètres, plus chaque jeton coûte cher à traiter — dans un Transformer dense, tous les poids participent au calcul de chaque jeton.

Le Mixture-of-Experts (MoE, « mélange d'experts ») casse ce lien :

L'idée en une phrase

Remplacer un gros réseau par beaucoup de petits réseaux, et n'en activer que quelques-uns par jeton.

Chez Kimi K3 : 896 experts par couche, dont 16 sont activés à chaque jeton. Le rapport de rareté (sparsity) vaut \(896/16 = 56\).

Le mécanisme

Dans un Transformer dense, la seconde sous-couche est un FFN unique. Dans un MoE, elle devient :

  1. Routeur : une petite matrice calcule un score pour chaque expert, \(\mathbf{s} = \sigma(\mathbf{W}_r\mathbf{x}) \in (0,1)^{896}\).
  2. Sélection Top-\(k\) : on retient les \(k = 16\) meilleurs scores.
  3. Calcul : les 16 experts sélectionnés traitent le jeton.
  4. Agrégation : leurs sorties sont combinées, pondérées par les scores normalisés.
\[ \mathbf{u} = \sum_{i \in \mathcal{T}_k(\mathbf{x})} p_i\, E_i(\mathbf{x}), \qquad p_i = \frac{s_i}{\sum_{r \in \mathcal{T}_k} s_r} \]
Symbole Signification
\(\mathcal{T}_k(\mathbf{x})\) Ensemble des \(k\) experts sélectionnés pour le jeton \(\mathbf{x}\)
\(E_i\) Le \(i\)-ième réseau expert (un FFN à porte)
\(p_i\) Poids de mélange, normalisé sur les experts retenus

Les experts partagés

Certaines transformations sont utiles à tous les jetons. Les leur faire apprendre par chaque expert serait un gaspillage. La solution (introduite par DeepSeekMoE et reprise par Kimi) : quelques experts partagés, toujours actifs, à côté des experts routés.

Kimi K3 : 2 experts partagés par couche, systématiquement actifs, à pleine largeur (\(d = 7168\)).

Le calcul de la rareté chez Kimi K3

Reconstruisons les chiffres à partir de la configuration publiée.

Un expert routé est un FFN à porte opérant dans l'espace latent de dimension \(\ell = 3584\), avec une dimension intermédiaire de 3 072 :

\[ 3 \times 3584 \times 3072 = 33{,}0 \text{ millions de paramètres par expert} \]

Le facteur 3 correspond aux trois matrices d'un FFN à porte (gate, up, down). Pour 896 experts :

\[ 896 \times 33{,}0\,\text{M} = 29{,}6 \text{ milliards de paramètres par couche} \]

Sur les 92 couches MoE : 2,72 T de paramètres, soit 98 % du modèle. Mais par jeton, seuls 16 experts travaillent :

\[ 16 \times 33{,}0\,\text{M} = 528 \text{ millions de paramètres par couche} \]

À retenir

La reconstruction complète (experts + attention + embeddings) donne 2,779 T de paramètres au total et 104,0 G activés par jeton — à comparer aux 2,78 T et 104,2 G annoncés dans le rapport. L'écart est de 0,2 %. Le calcul détaillé et vérifiable est donné dans Retrouver les 2,8 T de paramètres.

LatentMoE : l'innovation d'échelle de Kimi K3

Doubler le nombre d'experts actifs (8 → 16) double le trafic de communication : dans un MoE classique, chaque expert sélectionné reçoit le vecteur complet de dimension \(d = 7168\).

LatentMoE dissocie la largeur du modèle de la largeur des experts :

                   ┌── experts partagés (largeur pleine, d = 7168) ──┐
   x (7168) ───────┤                                                  ├──► y
                   └── W↓ ──► z (3584) ──► 16 experts routés ──► W↑ ──┘
                       ↑                      (largeur latente)    ↑
                    projection                                RMSNorm
                    descendante                            (ajout de K3)

Ce qui est routé et communiqué entre GPU n'est plus un vecteur de 7 168 dimensions mais de 3 584 : deux fois moins de trafic pour deux fois plus d'experts actifs.

Intuition

Les experts partagés gèrent le générique à pleine résolution ; les experts routés se spécialisent dans un espace comprimé. Le rapport résume : shared experts retain a full-width path for common transformations, whereas specialized routed experts operate in a compact latent space.

Le problème du déséquilibre de charge

C'est le problème pratique du MoE. Rien ne garantit que le routeur répartisse les jetons équitablement. En pratique, il ne le fait pas :

  • quelques experts reçoivent la majorité des jetons (overheated) ;
  • beaucoup en reçoivent très peu, et s'entraînent mal ;
  • certains n'en reçoivent plus du tout (dying experts) — leurs paramètres sont perdus.

Pire, en calcul distribué : si un GPU reçoit deux fois plus de jetons que les autres, tous les autres l'attendent. Le débit d'entraînement est dicté par le GPU le plus chargé.

Trois générations de solutions

Approche Principe Défaut
Perte auxiliaire Ajouter un terme à la fonction de coût qui pénalise le déséquilibre Entre en conflit avec l'objectif principal, dégrade la qualité
Biais sans perte (loss-free) Ajouter un biais \(b_j\) par expert au score de routage, ajusté par pas fixe : \(b_j \leftarrow b_j + \gamma\operatorname{sign}(\bar{\ell}-\ell_j)\) Le pas \(\gamma\) arbitre entre lenteur et oscillation ; ne passe pas l'échelle de ~1 000 experts
Quantile Balancing (Kimi K3) Fixer chaque biais exactement au quantile qui donne la charge cible —

Le biais \(b_j\) n'entre que dans la sélection Top-\(k\), jamais dans les poids de mélange \(p_i\) : il régule l'aiguillage sans perturber l'optimisation du routeur. Le détail complet de Quantile Balancing est en Quantile Balancing.

Intuition sur QB

Le biais sans perte tâtonne : « cet expert est trop chargé, baisse un peu son biais ». QB calcule directement : « pour que cet expert reçoive exactement 2 048 jetons, son biais doit valoir ceci » — en lisant le quantile approprié de la distribution des scores. Aucun taux d'apprentissage à régler, équilibre atteint en quelques pas.

Les conséquences système du MoE

Un MoE n'est pas seulement une architecture : c'est un problème de systèmes distribués. Les 896 experts d'une couche ne tiennent pas sur un GPU ; ils sont répartis. Chaque jeton doit donc voyager vers les GPU qui hébergent ses 16 experts, puis revenir.

C'est le parallélisme d'experts (Expert Parallelism, EP), et il implique deux communications collectives all-to-all par couche MoE — une à l'aller (dispatch), une au retour (combine). Sur 92 couches, cela fait 184 communications globales par passe avant.

Limite importante

Le MoE échange du calcul contre de la mémoire et de la communication. Kimi K3 doit stocker 2,8 T de paramètres (~1,56 To en MXFP4) même s'il n'en calcule que 104 G. C'est pourquoi le modèle ne tient sur aucun GPU unique, et pourquoi le rapport consacre une section entière à MoonEP.

Vérification de compréhension

Pourquoi 2 experts partagés plutôt que 0 ou 10 ?

Le rapport ne justifie pas ce nombre par une ablation publiée ; il indique seulement que \(N_s = 2\) est fixé pour toutes les couches (contre 1 chez Kimi K2). L'intuition est un compromis : trop peu d'experts partagés, et le savoir commun est dupliqué dans les experts routés ; trop, et la part de calcul dense annule l'intérêt de la rareté. C'est un des points où le rapport reste silencieux — voir Ce qui n'est pas publié.

Si 16 experts sur 896 sont actifs, le modèle est-il 56 fois moins cher ?

Non, pour deux raisons. D'abord, le coût d'attention (36 G de paramètres actifs) ne dépend pas du MoE. Ensuite, la mémoire et la communication ne diminuent pas d'un facteur 56. Au total, 104 G actifs sur 2 780 G, soit un facteur ~27 sur le calcul — et aucun sur la mémoire de stockage.


Chapitre précédent : Attention linéaire et récurrence · Chapitre suivant : Entraînement et lois d'échelle