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

Résumé

La montée en puissance des grands modèles de langage (LLMs) a généré des gains de performance significatifs, mais a également créé des défis importants en matière d'efficacité d'inférence. Si les architectures de type Mixture of Experts (MoEs) permettent d'y répondre en dissociant la taille du modèle du coût d'inférence, l'entraînement des MoEs à partir de zéro est souvent instable et coûteux en calcul. La conversion de modèles denses pré-entraînés en MoEs éparses est apparue comme une solution alternative ; cependant, les méthodes existantes reposent généralement sur un regroupement heuristique de neurones ou une division aléatoire pour partitionner le réseau feed-forward (FFN) en experts. Dans ce travail, nous proposons DOT-MoE, un nouveau cadre qui formule la décomposition des couches denses comme un problème de transport optimal différentiable (DOT). Au lieu d'heuristiques statiques, nous modélisons l'affectation des neurones comme un problème de transport équilibré, en utilisant des itérations de Sinkhorn-Knopp différentiables pour imposer des contraintes strictes de capacité des experts. De plus, nous utilisons des estimateurs Straight-Through (STE) pour apprendre conjointement l'affectation discrète neurone-expert et la politique de routage jeton-expert de bout en bout. Des expériences approfondies sur plusieurs architectures et bancs d'essai montrent que DOT-MoE surpasse significativement l'élagage structuré, le regroupement heuristique et les lignes de base de division aléatoire, conservant 90 % des performances du modèle dense original tout en réduisant les paramètres actifs de 50 %.

English

The scaling of Large Language Models (LLMs) has driven significant performance gains but created substantial challenges in inference efficiency. While Mixture of Experts (MoEs) architectures address this by decoupling model size from inference cost, training MoEs from scratch is often unstable and compute intensive. Conversion of pre-trained dense models into sparse MoEs has emerged as an alternative solution; however, existing methods typically rely on heuristic neuron clustering or random splitting to partition the Feed-Forward Network (FFN) into experts. In this work, we propose DOT-MoE, a novel framework that formulates the decomposition of dense layers as a Differentiable Optimal Transport (DOT) problem. Instead of static heuristics, we model neuron assignment as a balanced transport problem, utilizing differentiable Sinkhorn-Knopp iterations to enforce strict expert capacity constraints. Furthermore, we utilize Straight-Through Estimators (STE) to jointly learn the discrete neuron-to-expert assignment and the token-to-expert routing policy end-to-end. Extensive experiments across multiple architectures and benchmarks demonstrate that DOT-MoE significantly outperforms structured pruning, heuristic clustering, and random-split baselines, retaining 90% of the original dense model's performance while reducing active parameters by 50%.