🤖 machine learning

DOT-MoE: Differentiable Optimal Transport for MoEfication

DOT-MoE est un nouveau cadre qui convertit les LLM denses pré-entraînés en modèles de mélange d'experts efficaces en formulant la décomposition des neurones comme un problème de transport optimal différentiable, permettant un apprentissage de bout en bout des assignations qui surpasse de manière significative les méthodes heuristiques existantes tout en conservant 90 % de la performance originale avec 50 % de paramètres actifs en moins.

Auteurs originaux : Udbhav Bamba, Arnav Chavan, Aryamaan Thakur, Steve Teig, Deepak Gupta

Publié 2026-06-02

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Udbhav Bamba, Arnav Chavan, Aryamaan Thakur, Steve Teig, Deepak Gupta

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous possédez une bibliothèque massive et incroyablement intelligente (un Grand Modèle de Langage) où chaque livre est ouvert et lu simultanément pour chaque question que vous posez. Bien que cela rende la bibliothèque très intelligente, c'est aussi extrêmement lent et coûteux à exploiter car vous dépensez de l'énergie à lire des livres dont vous n'avez pas réellement besoin.

Pour corriger cela, des scientifiques ont créé un système de « Mélange d'Experts » (MoE - Mixture of Experts). Considérez cela comme le recrutement d'une équipe de spécialistes. Au lieu que toute la bibliothèque lise votre question, vous ne sollicitez que quelques experts spécifiques qui détiennent la réponse. Cela permet d'économiser de l'énergie et d'accélérer le processus.

Le Problème : Comment choisir les experts ?
L'article explique que transformer une bibliothèque standard de type « tout-lit-tout » en une équipe de « spécialistes » est délicat.

Les anciennes méthodes consistaient à mélanger aléatoirement les livres dans différentes pièces en espérant que les bonnes personnes se retrouvent au bon endroit. Ou bien, elles examinaient les couvertures des livres (les poids) et devinaient quels livres allaient ensemble.
Le problème : Ces méthodes séparent souvent le processus de « tri des livres » de celui de « l'apprentissage du bibliothécaire ». Elles trient les livres d'abord, puis tentent d'apprendre au bibliothécaire comment les trouver. Cela conduit souvent à une équipe désordonnée où les experts ne sont pas réellement spécialisés, ou à un bibliothécaire confus.

La Solution : DOT-MoE (Le Trieur Intelligent)
Les auteurs proposent une nouvelle méthode appelée DOT-MoE. Ils traitent le problème du tri des livres dans des salles de spécialistes comme un puzzle logistique appelé « Transport Optimal ».

Voici l'analogie :
Imaginez que vous avez un entrepôt rempli de milliers de travailleurs (neurones) et un ensemble de chantiers de construction (experts). Chaque chantier a besoin d'exactement le même nombre de travailleurs, et chaque travailleur doit aller dans un seul site.

L'Objectif : Vous voulez envoyer les bons travailleurs aux bons sites afin que les bâtiments soient construits parfaitement, tout comme l'entrepôt d'origine l'aurait fait.
L'Innovation : Au lieu de deviner ou d'utiliser une liste aléatoire, DOT-MoE utilise un « contrôleur de trafic » mathématique (Transport Optimal Différentiable). Il calcule la manière la plus efficace de déplacer chaque travailleur vers un site, garantissant qu'aucun site n'est surchargé et qu'aucun travailleur n'est laissé de côté.
La « Recette Secrète » : Le système apprend pendant qu'il trie. Il ne se contente pas de trier les livres pour ensuite embaucher un bibliothécaire. Il détermine quels livres vont dans quelle pièce en même temps qu'il apprend au bibliothécaire comment choisir la bonne pièce pour chaque nouvelle question. Ils apprennent ensemble, ajustant leurs positions jusqu'à ce que tout s'emboîte parfaitement.

Pourquoi est-ce meilleur ?
L'article affirme qu'en utilisant cette approche de « logistique intelligente » :

Il conserve l'intelligence : La nouvelle équipe de spécialistes conserve environ 90 % de l'intelligence de la bibliothèque originale.
Il réduit les coûts : Il n'utilise que 50 % de la « puissance cérébrale » active (paramètres) nécessaire pour répondre à une question.
C'est stable : Contrairement à d'autres méthodes qui pourraient planter ou mal performer lorsqu'elles tentent de diviser l'équipe, cette méthode garantit une équipe équilibrée et fonctionnelle dès le départ.

Le Résultat
Dans leurs tests, cette nouvelle méthode a mieux fonctionné que les précédentes manières de diviser la bibliothèque ou de choisir les experts de manière aléatoire. Elle a prouvé que si l'on traite l'organisation d'un réseau neuronal comme un problème logistique précis et soluble, on peut rendre les modèles d'IA géants beaucoup plus rapides et moins coûteux à exploiter sans perdre leur capacité à comprendre et à générer du langage.

En bref : Ils ont trouvé un moyen de réorganiser un cerveau d'IA géant et lent en une équipe de spécialistes rapide et efficace, en utilisant un « contrôleur de trafic » mathématique qui trie parfaitement les parties du cerveau tout en apprenant au système comment les utiliser, le tout dans un processus fluide et unique.

Résumé technique : DOT-MoE : Optimal Transport Différentiable pour la MoEfication

1. Énoncé du problème

La montée en charge rapide des grands modèles de langage (LLM) a créé un goulot d'étranglement significatif dans l'efficacité de l'inférence. Bien que les architectures Mixture-of-Experts (MoE) décomposent la taille du modèle du coût d'inférence en n'activant qu'un sous-ensemble de paramètres par jeton, l'entraînement des MoE à partir de zéro est notoirement instable, gourmand en données et coûteux en calculs.

Par conséquent, la MoEfication — la conversion de modèles denses pré-entraînés en MoE parcimonieux — est apparue comme une alternative viable. Cependant, les méthodes existantes pour convertir les réseaux Feed-Forward (FFN) denses en experts parcimonieux reposent largement sur des stratégies heuristiques (par exemple, la division aléatoire, le regroupement basé sur les poids ou le regroupement basé sur les activations). Ces approches souffrent de trois limitations principales :

Séparation des préoccupations : Elles traitent généralement l'assignation des neurones (le partitionnement du FFN) et l'entraînement du routeur comme deux étapes distinctes et séquentielles, empêchant leur co-adaptation.
Optimisation par procuration : Elles optimisent sur la base de statistiques intermédiaires (poids ou activations) plutôt que sur la contribution réelle des neurones à la sortie, ce qui conduit souvent à une reconstruction sous-optimale du comportement du modèle dense.
Absence de différentiabilité : La nature discrète de l'assignation des neurones empêche le flux de gradient de bout en bout, rendant difficile l'optimisation conjointe de la décomposition structurelle et de la politique de routage.

2. Méthodologie : DOT-MoE

Les auteurs proposent DOT-MoE, un cadre qui formule la décomposition de couches denses en experts comme un problème de Transport Optimal Différentiable (DOT). Cette approche permet l'apprentissage de bout en bout de la décomposition des experts et du mécanisme de routage.

2.1. Formulation du Transport Optimal

L'idée centrale est de considérer l'assignation des neurones comme un problème de transport équilibré où la « masse » (les neurones) doit être transportée d'une distribution source (les neurones du FFN dense) vers une distribution cible (les experts) sous des contraintes de capacité strictes.

Source : $dffn$ neurones, chacun avec une unité de masse.
Cible : $E$ experts, chacun nécessitant exactement $s$ neurones ( $s \cdot E = dffn$ ).
Fonction de coût : Au lieu d'une matrice de coût fixe, la méthode utilise une matrice d'affinité apprenable $A$ , où $A_{i,e}$ représente l'affinité d'assigner le neurone $i$ à l'expert $e$ . L'objectif est de maximiser l'affinité totale tout en respectant les contraintes marginales.

2.2. Relaxation Différentiable via Sinkhorn-Knopp

Pour rendre le problème d'assignation discret différentiable, DOT-MoE emploie une régularisation entropique.

Le problème de transport optimal est relaxé en utilisant l'algorithme de Sinkhorn-Knopp, qui normalise de manière itérative les lignes et les colonnes de la matrice de plan de transport.
Cela produit une matrice d'assignation souple $M_{soft} \in [0, 1]^{dffn \times E}$ qui satisfait les contraintes de capacité en moyenne.
À mesure que le paramètre de température $\tau \to 0$ , la solution se rapproche de l'optimum discret.

2.3. Optimisation Conjointe avec des Estimateurs à Passage Direct (STE)

Un défi critique est que le déploiement nécessite des assignations discrètes, et le routage implique une sélection top-k discrète, tous deux non différentiables. DOT-MoE aborde cela en utilisant des Estimateurs à Passage Direct (Straight-Through Estimators - STE) :

Passage direct (Forward Pass) : Utilise des décisions dures (arrondi glouton de $M_{soft}$ pour l'assignation ; sélection top-k pour le routage).
Passage arrière (Backward Pass) : Dirige les gradients à travers les contreparties souples ( $M_{soft}$ et les logits du routeur $P$ ).
Cela permet aux poids du routeur ( $W_r$ ) et à la matrice d'affinité d'assignation ( $A$ ) de co-s'adapter pendant l'entraînement.

2.4. Objectif d'Entraînement

Le cadre est entraîné via une Phase d'Alignement où les poids du modèle dense sont gelés. Le modèle minimise une combinaison de pertes :

Reconstruction de la Sortie : Mesurée via la divergence KL entre les sorties du modèle enseignant dense et de l'élève MoE parcimonieux.
Modélisation du Langage : Perte d'entropie croisée sur l'objectif de modélisation du langage.
Auxiliaires de Stabilité : Perte $z$ du routeur (pour éviter l'explosion des logits) et perte d'équilibrage de charge (pour éviter l'effondrement des experts).

3. Contributions Clés

Cadre de Transport Optimal : L'article introduit une nouvelle formulation de la conversion dense-vers-MoE comme un problème de transport équilibré, résolu via des itérations de Sinkhorn différentiables pour imposer des contraintes strictes de capacité d'expert.
Co-adaptation à Double Niveau : Contra-rentionnellement aux méthodes précédentes qui fixent les assignations avant d'entraîner les routeurs, DOT-MoE optimise conjointement la décomposition neurone-vers-expert et le routage jeton-vers-expert en utilisant des STE complémentaires.
Construction Sensible à la Sortie : En modélisant directement la contribution du neurone à la sortie finale (plutôt qu'en optimisant des variables de substitution comme les poids ou les activations intermédiaires), la méthode atteint une fidélité de reconstruction plus élevée.

4. Résultats Expérimentaux

Les auteurs ont évalué DOT-MoE sur trois familles de modèles (LLaMA-2, LLaMA-3, Qwen2.5) et six benchmarks (ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ).

Performance vs Élagage (Pruning) : Sur LLaMA-2 7B avec un budget de paramètres de 50 %, DOT-MoE a obtenu une perplexité de 7,99 sur WikiText-2, surpassant nettement la méthode d'élagage structuré de pointe DISP-LLM (9,84).
Performance vs MoEfication Existante : En réglages zero-shot, DOT-MoE a substantiellement surpassé des méthodes comme CMoE et LLaMA-MoE-v2. Par exemple, sur LLaMA-2 7B, DOT-MoE a atteint 61,5 % de précision moyenne contre 44,5 % pour CMoE.
Efficacité : DOT-MoE conserve environ 90 % des performances du modèle dense original tout en réduisant les paramètres actifs de 50 %.
Ajustement Fin (Fine-Tuning) : Avec seulement 1,2 milliard de jetons de fine-tuning continu, les modèles DOT-MoE ont récupéré des performances proches de leurs homologues denses (ex: 66,6 % contre 76,6 % sur LLaMA-2 7B), surpassant les bases de référence qui nécessitaient des volumes de données similaires ou plus importants.
Études d'Ablation :
- Granularité : La performance s'améliore avec une granularité d'expert accrue jusqu'à un point de saturation, après quoi les gains diminuent.
- Vitesse d'Inférence : Malgré un nombre d'experts plus élevé, le débit d'inférence reste stable lors de l'utilisation de noyaux MoE fusionnés (ex: vLLM), car la dimension intermédiaire totale fusionnée reste constante.
- Parcimonie : Les modèles entraînés à des niveaux de parcimonie plus élevés (75 %) se généralisent de manière robuste à des régimes de parcimonie encore plus extrêmes lors de l'inférence.

5. Signification et Revendications

L'article affirme que DOT-MoE représente une avancée significative dans l'efficacité de l'inférence des LLM en fournissant une méthode différentiable et fondée sur des principes pour convertir des modèles denses en MoE parcimonieux.

Changement de Paradigme : Il s'éloigne du partitionnement statique et heuristique vers une décomposition dynamique et apprenable qui respecte la cohérence fonctionnelle des experts.
Praticité : La méthode permet une sparsification post-hoc efficace de modèles pré-entraînés sans nécessiter un réentraînement complet à partir de zéro, conservant une haute capacité tout en réduisant drastiquement les paramètres actifs.
Robustesse : Le cadre démontre que la préservation de la fidélité des FFN nécessite de modéliser explicitement la contribution du neurone à la sortie, un principe que les méthodes basées sur des variables de substitution échouent à capturer.

Les auteurs concluent que DOT-MoE offre une initialisation supérieure pour les modèles parcimonieux, comblant l'écart de performance entre les architectures denses et parcimonieuses avec un minimum de données d'entraînement supplémentaires.