Routing Manifold Alignment Improves Generalization of Mixture-of-Experts
  LLMs
Authors: name: Zhongyang Li, name: Ziyue Li, name: Tianyi Zhou

Sparse Mixture-of-Experts (MoE) have been widely adopted in recent large
language models since it can efficiently scale up the model capability without
increasing the inference cost. However, evaluations on broad downstream tasks
reveal a consistent suboptimality of the routers in existing MoE LLMs, which
results in a severe performance gap (e.g., 10-20% in accuracy) to the optimal
routing. In this paper, we show that aligning the manifold of routing weights
with that of task embedding can effectively reduce the gap and improve MoE
LLMs' generalization performance. Our method, "Routing Manifold Alignment
(RoMA)", introduces an additional manifold regularization term in the
post-training objective and only requires lightweight finetuning of routers
(with other parameters frozen). Specifically, the regularization encourages the
routing weights of each sample to be close to those of its successful neighbors
(whose routing weights lead to correct answers) in a task embedding space.
Consequently, samples targeting similar tasks will share similar expert choices
across layers. Building such bindings between tasks and experts over different
samples is essential to achieve better generalization. Moreover, RoMA
demonstrates the advantage of unifying the task understanding (by embedding
models) with solution generation (by MoE LLMs). In experiments, we finetune
routers in OLMoE, DeepSeekMoE, and Qwen3-MoE using RoMA. Evaluations on diverse
benchmarks and extensive comparisons with baselines show the substantial
improvement brought by RoMA.