Routing Manifold Alignment Improves Generalization of Mixture-of-Experts LLMs

TL;DR

This paper introduces Routing Manifold Alignment (RoMA), a lightweight finetuning method that aligns routing weights with task embeddings to enhance the generalization of Mixture-of-Experts large language models.

Contribution

RoMA is a novel manifold regularization technique that improves MoE LLM routing by encouraging similar task samples to share expert choices, leading to better downstream performance.

Findings

RoMA significantly improves MoE LLM performance across benchmarks.

Lightweight finetuning of routers suffices for effective alignment.

Enhanced task understanding and solution generation integration.

Abstract

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.