DOT-MoE: Transporte Ótimo Diferenciável para MoEficação

Resumo

O escalonamento dos Modelos de Linguagem de Grande Porte (LLMs) impulsionou ganhos significativos de desempenho, mas criou desafios substanciais na eficiência de inferência. Embora as arquiteturas de Mistura de Especialistas (MoEs) abordem isso ao desacoplar o tamanho do modelo do custo de inferência, treinar MoEs do zero é frequentemente instável e computacionalmente intensivo. A conversão de modelos densos pré-treinados em MoEs esparsas surgiu como uma solução alternativa; no entanto, os métodos existentes geralmente dependem de agrupamento heurístico de neurônios ou divisão aleatória para particionar a Rede Feed-Forward (FFN) em especialistas. Neste trabalho, propomos o DOT-MoE, uma nova estrutura que formula a decomposição de camadas densas como um problema de Transporte Ótimo Diferenciável (DOT). Em vez de heurísticas estáticas, modelamos a atribuição de neurônios como um problema de transporte balanceado, utilizando iterações diferenciáveis de Sinkhorn-Knopp para impor restrições estritas de capacidade dos especialistas. Além disso, utilizamos Estimadores Straight-Through (STE) para aprender conjuntamente a atribuição discreta de neurônios a especialistas e a política de roteamento de tokens para especialistas de ponta a ponta. Extensos experimentos em múltiplas arquiteturas e benchmarks demonstram que o DOT-MoE supera significativamente as linhas de base de poda estruturada, agrupamento heurístico e divisão aleatória, mantendo 90% do desempenho do modelo denso original enquanto reduz os parâmetros ativos em 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%.