DOT-MoE: Differentieerbaar Optimaal Transport voor MoE-ficatie

Samenvatting

De schaalvergroting van Grote Taalmodellen (Large Language Models, LLM's) heeft aanzienlijke prestatieverbeteringen opgeleverd, maar heeft ook substantiële uitdagingen gecreëerd op het gebied van inferentie-efficiëntie. Hoewel Mengsel van Experts (Mixture of Experts, MoE)-architecturen dit aanpakken door de modelgrootte los te koppelen van de inferentiekosten, is het trainen van MoE's vanaf nul vaak instabiel en rekenintensief. Conversie van voorgetrainde dichte modellen naar schaarse MoE's is naar voren gekomen als een alternatieve oplossing; echter, bestaande methoden vertrouwen doorgaans op heuristische neuronclustering of willekeurig splitsen om het Feed-Forward Netwerk (FFN) in experts te verdelen. In dit werk stellen wij DOT-MoE voor, een nieuw raamwerk dat de decompositie van dichte lagen formuleert als een Differentieerbaar Optimaal Transport (Differentiable Optimal Transport, DOT)-probleem. In plaats van statische heuristieken modelleren we neuron-toewijzing als een gebalanceerd transportprobleem, waarbij we differentieerbare Sinkhorn-Knopp-iteraties gebruiken om strikte expertcapaciteitsbeperkingen af te dwingen. Verder maken we gebruik van Straight-Through Schatters (Straight-Through Estimators, STE) om zowel de discrete neuron-naar-expert toewijzing als het token-naar-expert routeringsbeleid end-to-end gezamenlijk te leren. Uitgebreide experimenten met meerdere architecturen en benchmarks tonen aan dat DOT-MoE aanzienlijk beter presteert dan gestructureerde pruning, heuristische clustering en willekeurige splitsing als basislijnen, waarbij 90% van de oorspronkelijke dichte modelprestaties behouden blijft terwijl het aantal actieve parameters met 50% wordt verminderd.

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%.