DOT-MoE: Transporte Óptimo Diferenciable para MoEficación

Resumen

El escalado de los Modelos de Lenguaje de Gran Tamaño (LLMs) ha impulsado mejoras significativas en el rendimiento, pero ha creado desafíos sustanciales en la eficiencia de inferencia. Si bien las arquitecturas de Mezcla de Expertos (MoEs) abordan este problema al desacoplar el tamaño del modelo del costo de inferencia, entrenar MoEs desde cero suele ser inestable y computacionalmente intensivo. La conversión de modelos densos preentrenados en MoEs dispersos ha surgido como una solución alternativa; sin embargo, los métodos existentes típicamente dependen de agrupamiento heurístico de neuronas o división aleatoria para particionar la Red de Avance Alimentado (FFN) en expertos. En este trabajo, proponemos DOT-MoE, un marco novedoso que formula la descomposición de capas densas como un problema de Transporte Óptimo Diferenciable (DOT). En lugar de heurísticas estáticas, modelamos la asignación de neuronas como un problema de transporte balanceado, utilizando iteraciones diferenciables de Sinkhorn-Knopp para imponer restricciones estrictas de capacidad de experto. Además, utilizamos Estimadores de Paso Directo (STE) para aprender conjuntamente la asignación discreta neurona-a-experto y la política de enrutamiento token-a-experto de extremo a extremo. Experimentos exhaustivos en múltiples arquitecturas y benchmarks demuestran que DOT-MoE supera significativamente a las líneas base de poda estructurada, agrupamiento heurístico y división aleatoria, reteniendo el 90% del rendimiento del modelo denso original mientras reduce los parámetros activos en un 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%.