Aller au contenu

Stable LatentMoE

98 % des paramètres de Kimi K3 vivent ici.

Le problème d'échelle

Augmenter à la fois le vivier d'experts et le nombre d'experts actifs élargit l'espace des spécialisations possibles. Mais dans un MoE conventionnel, chaque expert sélectionné reçoit le vecteur complet de dimension \(d = 7168\).

Conséquence : la communication et le trafic de poids d'experts croissent avec la multiplicité de routage. Passer de 8 à 16 experts actifs double le volume all-to-all par couche.

LatentMoE : découpler la largeur du modèle et celle des experts

L'idée (Elango et al., 2026) : les experts routés n'ont pas besoin d'opérer à la largeur du modèle.

                     ┌── 2 experts partagés (largeur d = 7168) ────┐
  x ∈ ℝ^7168 ────────┤                                              ├──► y
                     │                                              │
                     └── W↓ ─► z ∈ ℝ^3584 ─► top-16 / 896 ─► u ──► RMSNorm ──► W↑
                          ↑         ↑              ↑                    ↑        ↑
                     7168→3584   espace       experts en          ajout K3   3584→7168
                                 latent      dimension 3584
\[ \begin{aligned} \mathbf{u} &= \sum_{i\in\mathcal{T}_k(\mathbf{x})} p_i\, E_i^{\mathrm{routed}}(\mathbf{W}^{\downarrow}\mathbf{x}) \\[4pt] \mathbf{y} &= \sum_{j=1}^{N_s} E_j^{\mathrm{shared}}(\mathbf{x}) + \mathbf{W}^{\uparrow}\operatorname{RMSNorm}(\mathbf{u}) \end{aligned} \]
Symbole Signification Valeur K3
\(d\) Largeur du modèle 7 168
\(\ell\) Largeur latente des experts routés 3 584 (\(0{,}5\times d\))
\(E_i^{\mathrm{routed}}\) Expert routé, \(\mathbb{R}^{\ell}\to\mathbb{R}^{\ell}\) dim. interne 3 072
\(E_j^{\mathrm{shared}}\) Expert partagé, \(\mathbb{R}^{d}\to\mathbb{R}^{d}\) —
\(N_s\) Nombre d'experts partagés 2, dans chaque couche
\(k\) Experts routés actifs 16
\(n\) Experts routés au total 896

Le gain

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.

L'organisation générale (experts partagés + routés) est celle de DeepSeekMoE ; l'apport de LatentMoE est la séparation des largeurs.

Les deux modes de défaillance à l'échelle extrême

Une rareté de 56 amplifie deux pathologies. Stable LatentMoE existe pour les corriger.

Défaillance 1 : l'explosion des activations

Le diagnostic du rapport

La branche routée compose \(\mathbf{W}^{\downarrow}\), un FFN d'expert à branches multiples, puis \(\mathbf{W}^{\uparrow}\) : une chaîne de quasiment quatre multiplications matricielles consécutives.

Cette structure est mal conditionnée : les amplifications se composent. Combinée à l'échelle de 2,8 billions de paramètres, elle produit des activations internes explosives dans la branche routée.

Deux correctifs :

(a) RMSNorm avant la projection montante.

Le LatentMoE original applique \(\mathbf{W}^{\uparrow}\) directement à l'agrégat \(\mathbf{u}\), dont l'échelle varie selon quels experts ont été sélectionnés et avec quels poids. Kimi K3 insère une RMSNorm entre les deux.

Pourquoi c'est nécessaire ici et pas ailleurs

Dans un MoE classique, la sortie est directement ajoutée au résidu ; une variation d'échelle est absorbée. Ici, l'agrégat traverse encore une matrice (\(\mathbf{W}^{\uparrow}\)) avant d'être combiné à la branche partagée de pleine largeur. Une échelle instable en entrée d'une matrice produit une échelle instable en sortie, qui se propage.

Le rapport indique que cette RMSNorm, au-delà de stabiliser, améliore aussi la perte de validation et les benchmarks. Confirmé par latent_moe_use_norm: true.

(b) SiTU-GLU, pour borner les activations à l'intérieur de chaque expert. Chapitre dédié : SiTU-GLU.

Défaillance 2 : l'équilibrage de charge

Le diagnostic du rapport

Équilibrer la charge de près de \(10^3\) experts dépasse le régime dans lequel les mises à jour de biais sans perte auxiliaire restent bien conduites.

Un routage déséquilibré ralentit l'entraînement à parallélisme d'experts (tous attendent le GPU le plus chargé) et peut laisser certains experts mal entraînés.

Correctif : Quantile Balancing. Chapitre dédié : Quantile Balancing.

Le routage

\[ \mathcal{T}_i = \operatorname{argtop}_{k}(\mathbf{s}_i+\mathbf{b}), \qquad p_{i,j} = \frac{s_{i,j}}{\sum_{r\in\mathcal{T}_i}s_{i,r}}, \quad j\in\mathcal{T}_i \]

avec \(\mathbf{s}_i = \operatorname{Sigmoid}(\mathbf{W}_r\mathbf{x}_i)\).

Le point clé, souvent mal compris

Le biais \(\mathbf{b}\) apparaît dans la sélection (\(\operatorname{argtop}_k\)) mais pas dans les poids de mélange \(p_{i,j}\).

Conséquence : il régule l'aiguillage sans altérer les poids de mélange ni l'optimisation par gradient du routeur. C'est ce qui rend le mécanisme « sans perte auxiliaire » : aucun terme supplémentaire n'entre en compétition avec l'objectif de langage.

Confirmé par la configuration : moe_router_activation_func: "sigmoid", topk_method: "noaux_tc", moe_renormalize: true.

Sigmoïde plutôt que softmax

Le routeur utilise une sigmoïde, donc des scores indépendants entre experts, et non une distribution qui somme à 1. Cela évite la compétition artificielle entre experts au niveau du score, et rend l'ajout d'un biais additif interprétable comme un simple déplacement de seuil.

Les chiffres

Grandeur Valeur
Paramètres d'un expert routé \(3 \times 3584 \times 3072 = 33{,}0\) M
Experts routés par couche 896 → 29,60 G de paramètres
Experts routés actifs par jeton 16 → 528 M
Experts partagés (2, largeur \(d\)) \(3 \times 7168 \times 6144 = 132\) M
\(\mathbf{W}^{\downarrow} + \mathbf{W}^{\uparrow}\) \(2 \times 7168 \times 3584 = 51{,}4\) M
Routeur \(7168 \times 896 = 6{,}4\) M
Total par couche MoE 29,78 G
Total actif par couche MoE 718 M
Sur 92 couches 2 740 G totaux / 66,1 G actifs

Vérification de compréhension

Pourquoi \(\ell = 0{,}5 \times d\) exactement ?

Le rapport ne le justifie pas. C'est un hyperparamètre issu du papier LatentMoE et probablement d'ablations internes. Intuitivement, \(\ell\) trop petit brime la capacité de chaque expert ; \(\ell = d\) annule l'intérêt. Un facteur 0,5 divise le trafic par 2 pour une perte de capacité que les experts partagés compensent.

À retester si vous réimplémentez — ce n'est pas une constante universelle.

Les experts partagés utilisent-ils la même dimension interne que les routés ?

Le rapport ne l'énonce pas directement, mais la reconstruction des paramètres l'impose : pour retrouver 2,78 T, les 2 experts partagés doivent avoir une dimension intermédiaire combinée de \(2 \times 3072 = 6144\) à la largeur \(d = 7168\), soit 132 M par couche. Avec la valeur intermediate_size: 33792 (celle du FFN dense de la couche 1), le total dépasserait 2,80 T. Voir la vérification chiffrée en Retrouver les paramètres.

Que se passe-t-il si un expert ne reçoit jamais de jeton ?

Ses paramètres ne reçoivent aucun gradient et restent à leur initialisation : 33 M de paramètres perdus. À 896 experts par couche et 92 couches, quelques pour cent d'experts morts représenteraient des dizaines de milliards de paramètres gaspillés. C'est l'enjeu économique direct de Quantile Balancing.


Chapitre précédent : Attention Residuals · Chapitre suivant : SiTU-GLU