PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhongyang Li, Ziyue Li, Tianyi Zhou

arXiv:2511.07419v2cs.LG

TL;DR

Existing MoE routers show a substantial accuracy gap from oracle routing on downstream tasks. RoMA fine-tunes routers with manifold alignment regularization, encouraging similar tasks to share routing patterns. Across diverse benchmarks, it improves accuracy by 7–15%.

  • Problem

    Existing MoE routers have suboptimal routing weights, leaving a 10–20% accuracy gap to oracle routing across downstream tasks.

    Source: §1 · p. 1

  • Method

    RoMA adds manifold regularization to router fine-tuning, encouraging samples to move toward the routing weights of successful neighbors in task-embedding space.

    Source: §1 · p. 2, §4.1 · p. 4

  • Results

    Across diverse benchmarks, RoMA improves accuracy by 7–15% and consistently improves MoE performance.

    Source: §6 · p. 10

  • Takeaways & Limitations

    RoMA enables small active-parameter MoEs to rival or surpass much larger dense models.

    Source: §5.2 · p. 7, §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Existing MoE routers can leave a 10–20% accuracy gap to oracle routing, partly because semantically similar tasks receive different routing weights. RoMA adds manifold alignment regularization to router fine-tuning, encouraging successful neighbors in task-embedding space to share routing patterns.

1 INTRODUCTION

  • Existing MoE routers show a 10–20% accuracy gap to oracle routing across broad downstream tasks.Routers account for a small fraction of parameters, but their routing decisions are central to expert usage.Source: §1 · p. 1
  • Pretrained routers assign semantically similar samples dramatically different routing weights, hindering knowledge sharing across tasks.This misalignment underuses experts’ collective expertise and undermines generalization.Source: §1 · p. 1
  • Existing router fine-tuning approaches such as Dense BP do not address manifold misalignment between targeted tasks and routing weights across samples.RoMA is motivated by incorporating manifold alignment into the fine-tuning objective.Source: §1 · p. 2
  • RoMA adds manifold regularization to router fine-tuning, moving each sample’s intermediate-layer routing weights toward successful, similarity-weighted neighbors in task-embedding space.The method aligns routing weights with task understanding while retaining the output task loss.Source: §1 · p. 2, Figure 2 · p. 2
  • RoMA improves accuracy by 7–15% across eight benchmarks, outperforming all models shown in Figure 1.The figure compares RoMA on OLMoE-7B-A1B with dense LLMs from 7B to 34B.Source: Figure 1 · p. 2

2 RELATED WORK

Related work covers MoE routing, manifold regularization for LLMs, and routing optimization. The paper identifies embedding-space geometry as an aspect those routing approaches do not explicitly use.

MoE LLMs

Manifold Regularization of LLMs

  • I-STAR (Rudman & Eickhoff, 2023) controls embedding isotropy, while CROW (Min et al., 2024) enforces consistency across layers; neither explicitly uses manifold structure to improve MoE routing.The paper motivates RoMA from geometric insights about embedding strata and expert alignment.Source: §2 · p. 3

Routing Optimization

3 TASK-EXPERT ROUTING MANIFOLD MISALIGNMENT

Existing routers have a 10–20% accuracy gap from oracle routing, and their routing weights fail to reflect task-embedding clusters. Oracle routing instead shows a similar cluster structure to task embeddings.

3 TASK-EXPERT ROUTING MANIFOLD MISALIGNMENT

  • 10–20% accuracy separates existing routers from oracle routing on broad downstream tasks.The oracle routing weights are defined per sample using cross-entropy loss on the model output and ground-truth label.Source: §3 · pp. 3–4
  • Pretrained routing weights lack the clear clusters found in task embeddings, scattering samples from the same semantic group.Figure 3 visualizes task embeddings, pretrained routing weights, and oracle routing weights for ARC-C.Source: Figure 3 · p. 3, §3 · p. 4
  • Oracle routing weights form clusters aligned with task-embedding structure, with similar routing patterns for samples in the same semantic group.The paper identifies task-expert routing manifold misalignment as a bottleneck limiting router generalization.Source: §3 · p. 4

4 ROUTING MANIFOLD REGULARIZATION (ROMA)

RoMA aligns routing weights with task embeddings through lightweight router fine-tuning. It imitates successful, semantically similar samples and adds a manifold regularizer to the output loss.

4 ROUTING MANIFOLD REGULARIZATION (ROMA)

  • RoMA aligns routing-weight and task-embedding manifolds by fine-tuning routers so samples with similar task embeddings share similar routing patterns.The method adds manifold regularization in post-training and concatenates routing weights across multiple layers.Source: §4 · p. 4

4.1 SUCCESSFUL NEIGHBORHOOD TO IMITATE

  • The method selects correctly answered training samples as successful examples and finds neighbors using task similarity in a pretrained embedding space.Neighborhoods can use k-nearest neighbors or an epsilon-ball; filtering avoids imitating routing patterns that produce suboptimal outputs.Source: §4.1 · p. 4

4.2 TRAINING OBJECTIVE WITH MANIFOLD REGULARIZATION

  • The manifold regularizer penalizes routing discrepancies between semantically similar samples, weighted by their task-embedding similarity.It encourages each sample to follow successful neighbors’ routing patterns, while the output cross-entropy loss remains part of the objective.Source: §4.2 · pp. 4–5
  • Training updates router parameters while freezing experts; tuning only the last five layers achieved superior performance and significantly reduced training cost.Router parameters account for 0.0095% of total model parameters.Source: §4.2 · p. 5

5 EXPERIMENTS

Across three MoE families, RoMA leads the adaptation comparisons and its post-trained models compare favorably with larger dense models; OLMoE ablations identify effective routing choices.

5.1 EXPERIMENTAL SETTINGS

  • The evaluation covers OLMoE-7B-A1B, DeepSeekMoE-16B-A3B, and Qwen3-30B-A3B across eight benchmarks, using a 49,000-sample training set spanning five task categories.Source: §5.1 · pp. 5–6

5.2 MAIN RESULTS

  • RoMA leads overall across eight benchmarks; on MMLU, gains are DeepSeekMoE 46.2%→56.8% (+10.6%), OLMoE 57.8%→69.0% (+11.2%), Qwen3 74.2%→78.8% (+4.6%), beating C3PO on all three.Source: §5.2 · p. 6, Table 1 · p. 7
  • 1B-active OLMoE+RoMA scores 69.0% on MMLU and 86.7% on HellaSwag, surpassing several 7–8B and 13B dense models; 3B-active Qwen3 scores 78.8% on MMLU, surpassing 27–34B dense models.Source: §5.2 · p. 7, Table 2 · p. 8

5.3 ABLATION STUDY

  • 76.2% accuracy comes from regularizing the last five layers, exceeding the All-Layer configuration’s 75.1%; single-layer application yields 69.1–69.7%.The reported ablation finds selective regularization of the final layers more effective and efficient than applying RoMA uniformly across layers.Source: §5.3 · p. 8, Figure 6 · p. 9
  • k=3 nearest neighbors achieves 76.2%; random neighbors yield 67.8% versus the 67.6% baseline, while ε-neighborhoods peak at 74.1% for ε=0.5.Accuracy with ε-neighborhoods rises from 68.9% at ε=0.3 to 74.1% at ε=0.5, then falls slightly at ε=0.7.Source: Figure 8 · p. 9, §5.3 · p. 9
  • Last1 achieves 76.2%, above First1 at 71.4% and Middle1 at 69.2%; Last3 reaches 74.5%.The results indicate that final tokens contain richer task-relevant information for routing guidance, while one well-chosen token can outperform aggregating multiple tokens.Source: §5.3 · pp. 8–9, Figure 7 · p. 9
  • Using 30% of the training data yields 70.8% accuracy versus the 67.6% baseline; the full dataset achieves the best result, 76.2%.RoMA’s accuracy rises with additional data, while the reported 30% setting already provides a gain over baseline.Source: Figure 9 · p. 9, §5.3 · p. 9
  • Manifold regularization achieves 76.2%, ahead of L1 at 68.2%, entropy regularization at 70.7%, and L2 at 71.5%.The comparison favors RoMA’s task-embedding-space alignment over the other tested regularization strategies.Source: Figure 10 · p. 9, §5.3 · p. 10

6 CONCLUSIONS

RoMA aligns routing weights with task embeddings through manifold regularization, improving accuracy across diverse benchmarks while updating only routers and keeping experts frozen.

  • RoMA improves accuracy across diverse benchmarks by 7–15% without increasing inference cost.Source: §6 · p. 10
  • The method updates only router parameters while keeping experts frozen.Source: §6 · p. 10
  • RoMA aligns routing weights with the task embedding manifold through manifold regularization.Source: §6 · p. 10
  • RoMA enables small active-parameter MoEs to rival or surpass much larger dense models.Source: §6 · p. 10
  • The findings underscore geometric alignment between task representation and expert selection as important to routing strategies.Source: §6 · p. 10

A.1 COMPARISON BETWEEN ROMA AND BASELINES ON DEEPSEEKMOE-16B-A3B AND QWEN3-30B-A3B

The appendix compares RoMA with baselines on DeepSeekMoE-16B-A3B and Qwen3-30B-A3B; radar figures report consistent performance improvements across multiple benchmarks.

  • RoMA consistently improves performance on multiple benchmarks in the DeepSeekMoE-16B-A3B comparison.Source: Figure 11 · p. 13
  • The Qwen3-30B-A3B radar figure likewise reports RoMA improvements across multiple benchmarks.Source: Figure 11 · p. 13
  • The comparison presents radar figures for DeepSeekMoE-16B-A3B and Qwen3-30B-A3B.Source: Figure 11 · p. 13

A.2 ACCURACY AND COST OF DEEPSEEKMOE-16B-A3B AND QWEN3-30B-A3B

For DeepSeekMoE-16B-A3B, RoMA achieves accuracy improvement similar to C3PO while retaining nearly the base model’s inference efficiency.

  • RoMA achieves accuracy improvement similar to C3PO on DeepSeekMoE-16B-A3B.Source: Figure 12 · p. 13
  • RoMA maintains nearly the base model’s inference efficiency, while C3PO requires test-time optimization and induces 6–7× more FLOPs.Source: Figure 12 · p. 13

A.3 DETAILS OF TRAINING SET

The training set contains 49,000 samples across five task categories, and the evaluation benchmarks are reserved strictly for evaluation.

  • The training set comprises 49,000 samples distributed across five task categories.Source: §A.3 · p. 13
  • The evaluation benchmarks are reserved strictly for evaluation, as shown in the benchmark overview.Source: Figure 15 · p. 14
  • The listed benchmark descriptions cover reasoning and language tasks, including commonsense, science, question answering, and reading comprehension.Source: §A.3 · p. 14
  • The listed benchmarks also include multi-sentence reading comprehension and contextual coreference resolution.Source: §A.3 · pp. 14–15

A.4 DETAILS OF BENCHMARKS

The appendix describes benchmark and model coverage, then tests RoMA settings on DeepSeekMoE. Among the reported layer configurations, applying RoMA to the last five layers achieves the highest average accuracy.

A.4 DETAILS OF BENCHMARKS

  • Eight benchmarks cover general knowledge, commonsense reasoning, science QA, and mathematical problem-solving; GSM8k is treated as OOD because math is absent from training.The benchmark descriptions include MMLU, HellaSwag, PIQA, ARC-Challenge, ARC-Easy, WinoGrande, BoolQ, and GSM8k.Source: §A.4 · p. 15

A.5 MODEL DESCRIPTIONS

  • The model set spans 1B to 34B active parameters and includes dense LLMs and sparse MoE variants.The appendix directs readers to Table 2 for the reported model results.Source: §A.5 · p. 15

A.6 DETAILS OF ABLATION STUDY ON DEEPSEEKMOE

  • 74.7% average accuracy is achieved by applying RoMA to the last five layers, surpassing the 73.7% all-layer setting.The comparison highlights the benefit of selecting critical layers rather than regularizing every layer.Source: Table 3 · p. 16, §A.6 · p. 16
  • 74.7% average accuracy is achieved with Last1 token positions; first and middle positions yield limited improvements.The results indicate stronger performance from applying RoMA to the last token positions.Source: §A.6 · p. 16, Table 4 · p. 17
  • 74.7% average accuracy is achieved using the full training dataset, compared with 73.6% using 70% and 67.5% using 10%.Accuracy increases steadily across the reported training-set proportions.Source: Table 6 · p. 17, §A.6 · p. 17
  • 74.7% average accuracy is reached with k = 3 neighbors, while k = 5 yields 73.4% and performance saturates beyond k = 3.The ε-neighbor strategy reaches 72.8% at ε = 0.5; ε = 0.3 gives 68.0%.Source: §A.6 · p. 16, Table 5 · p. 17
  • 74.7% average accuracy with manifold regularization exceeds L1 at 67.2%, L2 at 70.3%, and entropy regularization at 69.8%.The comparison evaluates alternative regularization objectives on DeepSeekMoE.Source: Table 7 · p. 17, §A.6 · p. 17

Related papers

Loading 2511.07419v2…