DOT-MoE: MoE화를 위한 미분 가능 최적 수송

초록

대규모 언어 모델(LLM)의 스케일링은 상당한 성능 향상을 가져왔지만 추론 효율성에 있어서는 큰 어려움을 초래했다. 전문가 혼합(MoE) 아키텍처는 모델 크기와 추론 비용을 분리함으로써 이 문제를 해결하지만, MoE를 처음부터 학습시키는 것은 종종 불안정하고 계산 집약적이다. 사전 학습된 밀집 모델을 희소 MoE로 변환하는 것이 대안적인 해결책으로 등장했지만, 기존 방법들은 일반적으로 피드포워드 네트워크(FFN)를 전문가로 분할하기 위해 휴리스틱 뉴런 클러스터링 또는 무작위 분할에 의존한다. 본 연구에서는 DOT-MoE라는 새로운 프레임워크를 제안하며, 이는 밀집 레이어의 분해를 미분 가능한 최적 수송(DOT) 문제로 공식화한다. 정적 휴리스틱 대신, 우리는 뉴런 할당을 균형 잡힌 수송 문제로 모델링하고, 미분 가능한 Sinkhorn-Knopp 반복을 활용하여 엄격한 전문가 용량 제약 조건을 적용한다. 또한, 직선 추정기(STE)를 사용하여 이산적인 뉴런-전문가 할당과 토큰-전문가 라우팅 정책을 종단 간(end-to-end)으로 공동 학습한다. 여러 아키텍처와 벤치마크에 걸친 광범위한 실험을 통해 DOT-MoE가 구조적 가지치기, 휴리스틱 클러스터링 및 무작위 분할 기준선을 크게 능가하며, 활성 파라미터를 50% 줄이면서 원래 밀집 모델 성능의 90%를 유지함을 입증한다.

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