라우팅 매니폴드 정렬을 통한 Mixture-of-Experts LLM의 일반화 성능 향상

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

November 10, 2025

저자: Zhongyang Li, Ziyue Li, Tianyi Zhou

cs.AI

초록

희소 전문가 혼합 모델(Sparse Mixture-of-Experts, MoE)은 추론 비용을 증가시키지 않으면서 모델 성능을 효율적으로 확장할 수 있어 최근 대규모 언어 모델에 널리 채택되고 있습니다. 그러나 다양한 다운스트림 작업에 대한 평가 결과, 기존 MoE LLM의 라우터에서 일관된 최적성 부족이 확인되었으며, 이는 최적 라우팅 대비 심각한 성능 격차(예: 정확도 기준 10-20%)를 초래합니다. 본 논문에서는 라우팅 가중치의 매니폴드를 작업 임베딩의 매니폴드와 정렬시키면 이러한 격차를 효과적으로 줄이고 MoE LLM의 일반화 성능을 개선할 수 있음을 보여줍니다. 우리의 방법인 "라우팅 매니폴드 정렬(Routing Manifold Alignment, RoMA)"은 사후 훈련 목적 함수에 추가적인 매니폴드 정규화 항을 도입하며, 라우터의 경량 파인튜닝만을 요구합니다(다른 매개변수는 고정). 구체적으로, 이 정규화는 각 샘플의 라우팅 가중치가 작업 임베딩 공간 내에서 해당 샘플의 성공적인 이웃(라우팅 가중치가 정답을 도출한 샘플)들의 라우팅 가중치에 가까워지도록 유도합니다. 그 결과, 유사한 작업을 대상으로 하는 샘플들은 여러 계층에 걸쳐 유사한 전문가 선택을 공유하게 됩니다. 서로 다른 샘플들에 대해 작업과 전문가 간의 이러한 연결을 구축하는 것은 더 나은 일반화를 달성하는 데 필수적입니다. 더욱이, RoMA는 작업 이해(임베딩 모델 담당)와 해결책 생성(MoE LLM 담당)을 통합하는 이점을 보여줍니다. 실험에서는 RoMA를 사용하여 OLMoE, DeepSeekMoE 및 Qwen3-MoE의 라우터를 파인튜닝했습니다. 다양한 벤치마크에 대한 평가와 베이스라인과의 광범위한 비교를 통해 RoMA가 가져오는 상당한 성능 향상을 확인하였습니다.

English

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.