DOT-MoE: Дифференцируемый оптимальный транспорт для MoEфикации

Аннотация

Масштабирование больших языковых моделей (LLM) привело к значительному повышению производительности, но создало существенные проблемы с эффективностью инференса. Хотя архитектуры Смеси экспертов (MoE) решают эту проблему, разделяя размер модели и стоимость инференса, обучение MoE с нуля часто оказывается нестабильным и вычислительно затратным. Преобразование предварительно обученных плотных моделей в разреженные MoE стало альтернативным решением; однако существующие методы обычно полагаются на эвристическую кластеризацию нейронов или случайное разбиение для разделения сети прямого распространения (FFN) на эксперты. В данной работе мы предлагаем DOT-MoE — новую структуру, которая формулирует декомпозицию плотных слоев как задачу дифференцируемого оптимального транспорта (DOT). Вместо статических эвристик мы моделируем назначение нейронов как сбалансированную транспортную задачу, используя дифференцируемые итерации Синкхорна-Кноппа для соблюдения строгих ограничений на вместимость экспертов. Кроме того, мы используем Straight-Through Estimators (STE) для совместного обучения дискретного назначения нейронов на эксперты и политики маршрутизации токенов к экспертам в рамках сквозного подхода. Обширные эксперименты на множестве архитектур и эталонных тестов показывают, что DOT-MoE значительно превосходит базовые методы структурной обрезки, эвристической кластеризации и случайного разбиения, сохраняя 90% производительности исходной плотной модели при сокращении активных параметров на 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%.