Mixture of Experts (MoE)

Résumé

Le Mixture of Experts (MoE), ou mélange d’experts en français, constitue une architecture d’apprentissage profond révolutionnaire qui repose sur un principe simple mais puissant : au lieu de traiter chaque donnée d’entrée avec un unique réseau de neurones massif, on dispose de plusieurs sous-réseaux spécialisés appelés experts, et un gating network (réseau de routage) détermine dynamiquement quels experts doivent être activés pour chaque échantillon. Cette approche permet de créer des modèles contenant des centaines de milliards, voire des milliers de milliards de paramètres, tout en ne utilisant qu’une faible fraction de ces paramètres pour chaque inférence. Le MoE est aujourd’hui au cœur des plus grands modèles de langage modernes comme Mixtral, Qwen-MoE et GPT-4. Ce guide complet explique le principe mathématique, l’intuition, l’implémentation en PyTorch, les hyperparamètres critiques, ainsi que les cas d’usage pratiques de cette architecture fascinante.

Principe Mathématique

Formulation fondamentale

Le Mixture of Experts repose sur une décomposition élégante du problème d’apprentissage. Considérons K experts notés E₁, E₂, …, E_K, chacun étant un sous-réseau de neurones (généralement un MLP — Multi-Layer Perceptron — avec ses propres poids et biais). Le gating network G est un module distinct qui, pour chaque entrée x, produit une distribution de probabilité sur les K experts.

La sortie globale du modèle MoE s’écrit comme une combinaison pondérée des sorties de chaque expert :

y = Σᵢ G(x)ᵢ · Eᵢ(x)

où la somme parcourt tous les experts i = 1 à K.

Le gating network : fonction softmax

Le gating network calcule pour chaque expert un score linéaire, puis applique une fonction softmax pour obtenir des poids normalisés :

G(x)ᵢ = exp(wᵢ · x + bᵢ) / Σⱼ exp(wⱼ · x + bⱼ)

Ici, wᵢ est le vecteur de poids associé à l’expert i, et bᵢ son biais. Le produit scalaire wᵢ · x mesure la similarité entre l’entrée et le profil de l’expert i. Après softmax, les poids G(x)ᵢ forment une distribution de probabilité : ils sont tous positifs et leur somme vaut 1. L’expert avec le score le plus élevé reçoit le poids le plus fort, ce qui signifie que sa contribution à la sortie finale sera prépondérante.

MoE moderne avec top-k routing

Dans les modèles à grande échelle, activer tous les experts pour chaque entrée serait prohibitif en termes de calcul. L’innovation clé du MoE moderne est le top-k routing : on ne sélectionne que les k experts ayant les poids G(x) les plus élevés. Typiquement, k = 1 (routing expert unique) ou k = 2 (deux experts), ce qui rend le modèle sparse.

La formulation devient alors :

y = Σᵢ∈top-k G(x)ᵢ · Eᵢ(x)

Cette modification est fondamentale : même si le modèle possède des milliers d’experts, seuls k d’entre eux sont activés pour chaque échantillon. La complexité computationnelle par échantillon reste faible et indépendante du nombre total d’experts. C’est cette sparsité qui permet d’entraîner des modèles de taille colossale sur du matériel réaliste.

Auxiliary Load Balancing Loss

Un problème majeur du MoE est le collapsus d’experts : sans régularisation appropriée, le gating network peut apprendre à router systématiquement la majorité des entrées vers un seul expert ou un petit sous-ensemble d’experts. Les autres experts ne reçoivent presque jamais de données, ne s’entraînent pas correctement, et deviennent inutiles. C’est un phénomène bien documenté dans la littérature.

Pour contrer ce problème, on ajoute une perte auxiliaire de répartition de charge (auxiliary load balancing loss) :

L_balance = α · Σᵢ (fᵢ / k)²

où fᵢ représente la fraction des échantillons d’un batch routés vers l’expert i, et k est le nombre d’experts sélectionnés par échantillon (top-k). Le coefficient α pondère l’importance de cette perte auxiliaire par rapport à la perte principale de la tâche.

Cette formule pénalise les distributions déséquilibrées : si un expert reçoit trop d’échantillons par rapport aux autres, le terme (fᵢ/k)² explose quadratiquement. L’idéal est que chaque expert reçoive une fraction égale des données, soit fᵢ ≈ 1/K pour tout i. En pratique, on vise une répartition aussi uniforme que possible tout en conservant la capacité du gating à router intelligemment.

La fonction de coût totale du modèle MoE est donc :

L_total = L_tâche + L_balance

où L_tâche est la perte standard de la tâche (par exemple, l’entropie croisée pour une classification ou la perte de langage pour un LLM).

Intuition : Le Cabinet de Consultants

Imaginez un grand cabinet de consultants international. Ce cabinet dispose de dizaines d’experts dans des domaines variés : un spécialiste en finance, un expert en intelligence artificielle, un conseil en stratégie industrielle, un juriste spécialisé en propriété intellectuelle, et bien d’autres encore.

Maintenant, supposons qu’un client soumette un nouveau dossier — disons, une question sur la régulation financière des algorithmes de trading. Que se passe-t-il ?

Le cabinet n’envoie pas tous ses experts sur ce dossier. Ce serait absurde : pourquoi mobiliser le juriste en propriété intellectuelle ET le spécialiste en supply chain sur un problème de trading algorithmique ? Ce serait non seulement un gaspillage considérable de ressources, mais aussi une source de confusion et de bruit.

À la place, le chef de cabinet — notre gating network — évalue rapidement la nature du problème, identifie les 1 ou 2 experts les plus compétents pour traiter ce dossier spécifique, et leur confie le travail. Le financier pour comprendre les mécanismes de trading, et peut-être un analyste en IA pour comprendre l’algorithme sous-jacent. C’est tout. Deux experts suffisent.

C’est exactement le principe du Mixture of Experts :

• Chaque dossier (chaque donnée d’entrée) est évalué par le gating network
• Seuls les 1-2 experts les plus compétents sont activés pour ce dossier particulier
• Chaque expert se spécialise dans un sous-domaine du problème global, car il ne voit qu’un sous-ensemble des données (celles qui lui sont routées)
• L’efficacité est bien supérieure à celle d’un réseau géant unique qui devrait traiter chaque problème avec l’ensemble de ses paramètres

L’analogie souligne aussi pourquoi le load balancing est crucial : si le chef de cabinet envoyait toujours les mêmes 2 experts sur tous les dossiers, les autres ne gagneraient jamais d’expérience. Il faut veiller à ce que chaque expert reçoive suffisamment de dossiers pertinents pour développer son expertise. C’est précisément le rôle de la perte auxiliaire de répartition de charge.

Implémentation Python avec PyTorch

Voici une implémentation complète et fonctionnelle d’un modèle Mixture of Experts en PyTorch, incluant le gating network, le top-k routing, la passe avant sparse et l’entraînement sur une tâche de classification.

1. Le réseau Expert

Chaque expert est un simple MLP à deux couches avec activation GELU, une architecture courante dans les modèles modernes :

import torch
import torch.nn as nn
import torch.nn.functional as F


class Expert(nn.Module):
    """Un expert individuel : MLP à deux couches avec activation GELU."""

    def __init__(self, dim_in: int, dim_hidden: int, dim_out: int):
        super().__init__()
        self.fc1 = nn.Linear(dim_in, dim_hidden)
        self.fc2 = nn.Linear(dim_hidden, dim_out)
        self.gelu = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(self.gelu(self.fc1(x)))

2. Le Gating Network

Le gating network calcule les scores de chaque expert et applique le top-k routing :

class GatingNetwork(nn.Module):
    """
    Gating network avec top-k routing et load balancing loss.
    Retourne les poids des experts sélectionnés et leur masque.
    """

    def __init__(self, dim_in: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = min(top_k, num_experts)
        self.gate = nn.Linear(dim_in, num_experts, bias=False)

    def forward(self, x: torch.Tensor):
        # Scores bruts pour chaque expert : (batch_size, num_experts)
        scores = self.gate(x)
        probs = F.softmax(scores, dim=-1)

        # Sélection top-k : on garde les k experts avec les scores les plus élevés
        topk_probs, topk_indices = torch.topk(probs, self.top_k, dim=-1)

        # Ré-normaliser les probs après masquage top-k
        topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True)

        # Masque binaire : un expert est 1 s'il est dans le top-k
        mask = torch.zeros_like(probs)
        mask.scatter_(-1, topk_indices, 1.0)

        # Load balancing loss
        # f_i = fraction d'échantillons routés vers l'expert i
        frac_routed = mask.sum(dim=0) / mask.sum()
        # Pénalité quadratique par rapport à la répartition idéale
        ideal_frac = 1.0 / self.num_experts
        load_bal_loss = torch.sum((frac_routed / ideal_frac) ** 2)

        return topk_probs, topk_indices, mask, load_bal_loss

3. Modèle MoE complet

class MixtureOfExperts(nn.Module):
    """
    Modèle Mixture of Experts complet avec :
    - num_experts experts indépendants
    - top-k routing sparse
    - load balancing loss auxiliaire
    """

    def __init__(
        self,
        dim_in: int,
        dim_hidden: int,
        dim_out: int,
        num_experts: int = 8,
        top_k: int = 2,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Créer les experts
        self.experts = nn.ModuleList([
            Expert(dim_in, dim_hidden, dim_out)
            for _ in range(num_experts)
        ])

        # Gating network
        self.gating = GatingNetwork(dim_in, num_experts, top_k)

    def forward(self, x: torch.Tensor):
        """
        Passe avant sparse : on ne calcule que les experts sélectionnés.
        """
        topk_probs, topk_indices, mask, load_bal_loss = self.gating(x)

        batch_size = x.size(0)
        output = torch.zeros(batch_size, x.size(-1), device=x.device)

        # Pour chaque échantillon du batch, on active seulement les top-k experts
        for b in range(batch_size):
            for k_idx in range(self.top_k):
                expert_id = topk_indices[b, k_idx].item()
                weight = topk_probs[b, k_idx]
                # Calcul de l'expert uniquement s'il est sélectionné
                expert_out = self.experts[expert_id](x[b:b+1])
                output[b:b+1] += weight * expert_out

        return output, load_bal_loss

4. Entraînement sur une tâche de classification

def train_moe_classifier(num_samples=2000, dim_in=16, dim_hidden=64,
                         num_classes=4, num_experts=8, top_k=2,
                         epochs=50, lr=1e-3, aux_loss_weight=0.01):
    """
    Entraîne un modèle MoE sur une tâche de classification synthétique.
    Démo de l'architecture avec données aléatoires.
    """
    torch.manual_seed(42)

    # Données synthétiques : clusters dans l'espace d'entrée
    X = torch.randn(num_samples, dim_in)
    # Labels basés sur le signe des coordonnées
    y = ((X[:, :num_classes] > 0).float()).long().argmax(dim=-1)

    # Modèle
    model = MixtureOfExperts(
        dim_in=dim_in,
        dim_hidden=dim_hidden,
        dim_out=num_classes,
        num_experts=num_experts,
        top_k=top_k,
    )

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    print(f"Début de l'entraînement : {num_experts} experts, top-{top_k}")
    print(f"Nombre d'échantillons : {num_samples}")
    print(f"Taux d'apprentissage : {lr}")
    print("-" * 60)

    for epoch in range(epochs):
        model.train()
        optimizer.zero_grad()

        # Passe avant
        logits, bal_loss = model(X)

        # Perte totale : tâche + load balancing
        task_loss = criterion(logits, y)
        total_loss = task_loss + aux_loss_weight * bal_loss

        # Rétropropagation
        total_loss.backward()
        optimizer.step()

        # Suivi
        if (epoch + 1) % 10 == 0:
            with torch.no_grad():
                _, pred = logits.max(dim=-1)
                accuracy = (pred == y).float().mean().item()
            print(
                f"Époque {epoch+1:3d}/{epochs} | "
                f"Perte tâche : {task_loss.item():.4f} | "
                f"Perte balance : {bal_loss.item():.4f} | "
                f"Précision : {accuracy:.4f}"
            )

    # Évaluation finale
    model.eval()
    with torch.no_grad():
        logits, _ = model(X)
        _, pred = logits.max(dim=-1)
        final_acc = (pred == y).float().mean().item()
    print(f"\nPrécision finale : {final_acc:.4f}")

    return model


if __name__ == "__main__":
    model = train_moe_classifier()

5. Version optimisée avec torch.einsum

Pour des performances supérieures sur de gros modèles, on peut vectoriser le calcul des experts :

class MoEEfficient(nn.Module):
    """
    Version optimisée du MoE utilisant einsum pour calculer
    tous les experts en parallèle, puis masquer les non-sélectionnés.
    Plus rapide sur GPU pour un grand nombre d'experts.
    """

    def __init__(self, dim_in, dim_hidden, dim_out, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Poids empilés : (num_experts, dim_out, dim_hidden)
        self.w1 = nn.Parameter(torch.randn(num_experts, dim_in, dim_hidden))
        self.w2 = nn.Parameter(torch.randn(num_experts, dim_hidden, dim_out))

        self.gating = GatingNetwork(dim_in, num_experts, top_k)

    def forward(self, x: torch.Tensor):
        topk_probs, topk_indices, mask, bal_loss = self.gating(x)

        # Calculer tous les experts : (num_experts, batch, dim_out)
        h = torch.einsum('bid,edh->ebih', x, self.w1)
        h = F.gelu(h)
        out = torch.einsum('ebh,eho->ebo', h, self.w2)

        # Appliquer le masque de routing
        masked_mask = mask.transpose(0, 1).unsqueeze(-1)
        out = out * masked_mask

        # Pondérer par les probs du gating et sommer sur les experts
        out = out.sum(dim=0)  # Somme sur les experts

        return out, bal_loss

Hyperparamètres Clés

Le MoE introduit plusieurs hyperparamètres critiques qui n’existent pas dans les architectures denses classiques. Les régler correctement est essentiel pour obtenir de bonnes performances.

num_experts

Le nombre total d’experts dans le modèle. Cette valeur détermine la capacité totale du modèle et son niveau de spécialisation.

• Valeurs typiques : 4 à 256 (dans les LLMs : jusqu’à plusieurs milliers)
• Règle pratique : augmenter progressivement. Commencer avec 8 experts, puis doubler si les performances stagnent. Chaque expert supplémentaire augmente la capacité du modèle mais aussi la complexité de l’entraînement.
• Attention : trop d’experts par rapport à la quantité de données entraîne une sous-utilisation. Chaque expert doit recevoir suffisamment d’échantillons pour apprendre quelque chose d’utile.

top_k

Le nombre d’experts activés pour chaque échantillon d’entrée. C’est le paramètre qui contrôle directement la sparsité du modèle.

• Valeurs typiques : 1 ou 2 (le plus courant dans les modèles modernes)
• k = 1 (expert gating) : chaque échantillon n’est traité que par un seul expert. C’est le maximum de sparsité mais le risque de déséquilibre est plus élevé.
• k = 2 : chaque échantillon est traité par deux experts dont les sorties sont combinées. Cela offre un meilleur compromis entre expressivité et efficacité computationnelle. C’est la valeur recommandée pour la plupart des applications.
• k > 2 : rarement utilisé en pratique, car cela diminue l’avantage de la sparsité.

capacity_factor

Le facteur de capacité détermine la taille du tampon d’experts, c’est-à-dire le nombre maximum de tokens qu’un expert peut traiter dans un batch. C’est un mécanisme de régulation pour éviter la surcharge d’un expert particulièrement populaire.

• Valeurs typiques : 1.0 à 2.0
• 1.0 : capacité exacte. Si un expert reçoit plus de tokens que sa capacité, les tokens excédentaires sont soit ignorés, soit routés vers un expert de secours.
• 1.25 – 1.5 : marge de sécurité de 25 à 50 %. Recommandé pour la plupart des cas d’usage.
• > 2.0 : généralement inutile et gaspille de la mémoire.

aux_loss_weight (α)

Le poids de la perte auxiliaire de répartition de charge. C’est sans doute l’hyperparamètre le plus difficile à régler, car il contrôle directement l’équilibre entre spécialisation des experts et utilisation uniforme.

• Valeurs typiques : 0.001 à 0.1
• Trop faible (< 0.001) : le gating network favorise un petit nombre d’experts, les autres ne sont jamais sélectionnés. L’entraînement se dégrade rapidement.
• Trop fort (> 0.1) : le modèle force une répartition artificiellement uniforme, ce qui empêche les experts de se spécialiser réellement. On perd l’avantage principal du MoE.
• Valeur recommandée : 0.01 est un bon point de départ. Ajuster en observant la distribution effective des routages pendant l’entraînement.

Avantages et Limites

Avantages

1. Efficacité computationnelle extrême : le MoE permet d’entraîner des modèles avec des milliers de milliards de paramètres tout en n’activant qu’une infime fraction de ces paramètres pour chaque token. Un modèle Mixtral 8x7B possède 46,7 milliards de paramètres totaux mais n’en utilise que ~13 milliards par forward pass, soit environ 28 % de la compute d’un modèle dense équivalent. C’est un gain considérable.
2. Spécialisation automatique : les experts apprennent spontanément à se diversifier pendant l’entraînement. Certains deviennent spécialistes de la syntaxe, d’autres de la sémantique, d’autres encore du raisonnement logique. Cette émergence de spécialisation est un phénomène fascinant observé dans les grands modèles de langage.
3. Mise à l’échelle horizontale naturelle : les experts étant des réseaux indépendants, ils peuvent être placés sur différents GPU ou même différents nœuds de calcul. Le gating network détermine quels GPU doivent être activés, ce qui permet une parallélisation efficace à très grande échelle.
4. Meilleure généralisation : en théorie, la composition d’experts spécialisés peut généraliser mieux qu’un modèle dense équivalent. Chaque expert se concentre sur un sous-domaine, ce qui réduit l’interférence catastrophique entre différentes compétences du modèle.
5. Flexibilité architecturale : on peut facilement augmenter le nombre d’experts sans redimensionner l’architecture de base. Ajouter des experts est une opération simple qui augmente la capacité du modèle sans modifier sa structure fondamentale.

Limites

1. Complexité d’entraînement : le MoE est nettement plus difficile à entraîner qu’un modèle dense. Problème de collapsus d’experts, instabilité du gradient, sensibilité aux hyperparamètres (surtout aux_loss_weight). Il faut une expertise significative pour obtenir de bons résultats.
2. Problèmes de routage : si le gating network n’apprend pas un bon routage, les performances chutent drastiquement. Dans les pires cas, un seul expert absorbe la quasi-totalité des données et les autres restent inutilisés. C’est un mode d’échec bien documenté.
3. Coût mémoire : même si la compute est sparse, tous les paramètres de tous les experts doivent être chargés en mémoire. Un modèle 8x7B nécessite la mémoire de 46,7 milliards de paramètres, même si seulement 13 milliards sont activés par token. C’est un goulot d’étranglement majeur pour l’inférence.
4. Communication distribuée : dans un environnement distribué, le routing dynamique implique des communications all-to-all entre les GPU qui transportent les tokens vers les experts appropriés. Ces communications peuvent devenir un goulot d’étranglement sérieux si le réseau n’est pas optimisé.
5. Manque de standardisation : contrairement aux transformers denses où l’architecture est bien standardisée, il existe de nombreuses variantes de MoE (switch routing, hash routing, expert capacity, etc.) et peu de consensus sur les meilleures pratiques.

4 Cas d’Usage Pratiques

1. Grands Modèles de Langage (LLMs)

C’est l’application la plus emblématique du MoE. Des modèles comme Mixtral 8x7B de Mistral AI, Qwen-MoE, et GPT-4 (probablement) utilisent des couches MoE pour atteindre des performances compétitives avec une fraction de la compute d’un modèle dense. Dans un LLM, chaque couche transformer peut contenir une couche MoE à la place du MLP traditionnel. Le gating network route chaque token vers les 1-2 experts les plus appropriés. Cette approche permet des gains de 4x à 8x en efficacité par rapport à un modèle dense de qualité équivalente.

2. Traitement Multi-Domaines

Dans les systèmes qui doivent traiter des données provenant de domaines très différents — par exemple, un modèle unique qui analyse à la fois du texte médical, du code informatique et des documents juridiques — le MoE excelle naturellement. Chaque expert peut se spécialiser dans un domaine particulier : un expert pour le vocabulaire médical, un autre pour la syntaxe du code, un troisième pour le langage juridique. Le gating network apprend automatiquement à router chaque texte vers l’expert compétent. C’est bien plus efficace qu’un modèle unique qui devrait mélanger toutes ces compétences sans jamais se spécialiser.

3. Systèmes de Recommandation

Les plateformes de recommandation (Netflix, Spotify, Amazon) doivent prédire les préférences de millions d’utilisateurs avec des profils extrêmement diversifiés. Le MoE permet de créer des experts qui se spécialisent sur des segments d’utilisateurs ou des catégories de produits. Un expert pour les utilisateurs amateurs de films d’action, un autre pour les passionnés de jazz, un troisième pour les acheteurs de matériel électronique. Le gating network identifie le profil de l’utilisateur courant et active les experts pertinents. Cette approche est déployée à grande échelle chez Google pour son système de recommandation de vidéos YouTube.

4. Vision par Ordinateurs Multi-Tâches

Dans les systèmes de vision qui doivent accomplir plusieurs tâches simultanément — détection d’objets, segmentation sémantique, estimation de profondeur, classification de scènes — le MoE offre une architecture élégante. Chaque expert peut se spécialiser dans une tâche ou un type de représentation visuelle. Par exemple, un MoE intégré dans un backbone ViT (Vision Transformer) permet de partager les couches basses (détection de contours, textures) tout en ayant des experts spécialisés pour les tâches de haut niveau. Cela évite d’entraîner un modèle séparé pour chaque tâche tout en conservant des performances élevées grâce à la spécialisation.

Voir Aussi

• Résoudre le Problème d’Entretien Maximum Subarray avec Python
• Implémenter l’Algorithme de Dekker en Python : Synchronisation des Threads Efficace