Routing Manifold Alignment Improves Generalization of Mixture-of-Experts LLMs
Authors: Zhongyang Li
Ziyue Li
Tianyi Zhou
Venue: ICLR 2026 Poster (2026)
Keywords: Large Language Models ;Mixture of Experts; Manifold Regularization;
Primary area: foundation or frontier models, including LLMs
TL;DR: Aligning the routing weight manifold with the manifold of task embedding can significantly improve existing MoE LLMs' downstream task performance by 6-16% in accuracy with lightweight post-training of routers.

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 via post-training 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 two recent MoE LLMs using RoMA. Evaluations on diverse benchmarks and extensive comparisons with baselines show the substantial improvement brought by RoMA.

ROUTING MANIFOLD ALIGNMENT IMPROVES GENER- ALIZATION OF MIXTURE-OF-EXPERT LLMs
Zhongyang Li Johns Hopkins University Baltimore, MD, USA z1i3000 jh.edu
Ziyue Li
Tianyi Zhou MBZUAI Abu Dhabi, UAE t tianyi www.wmouuu.....c.ooooo ac
University of Maryland College Park, MD, USA 1itzy 96umd. edu
ABSTRACT
Sparse Mixture-of-Experts (MoE) have been widely adopted in recent large lan- guage models since it can efficiently scale up the model capability without increas- ing 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, show that aligning the manifold of routing weights with that of task em- bedding via post-training 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 three recent MoE LLMs using RoMA. Evaluations on di- verse benchmarks and extensive comparisons with baselines show the substantial improvement brought by RoMA. Our code can be accessed here.
INTRODUCTION
Sparse Mixture-of-Experr (MoE) have emerged as cornerstone architecture in scaling large lan- guage models (LLMs), enabling significant capacity increases without proportional computational overhead during inference (Fedus al., 2022: Lepikhin et al., 2020). At the core of this mechanism lies the router, which assigns input tokens to a small subset of experts through routing weights in each layer. Despite the small portion of router parameters in MoE L.MSs (e.g.. 0.03% in 7B model), they are the key to the success of expert usage in MoE (Shazeer et al., 201 7). However, evaluations across broad downstream tasks reveal that routers in existing MoE LLMs cause major failures. As shown in Table 1, their suboptimal routing weights lead to a performance gap of 10-20% in accuracy when compared to the optimal routing weights (oracle). This gap underscores a major untapped bottleneck in MoE LLMs, suggesting that improving routing is critical to boosting MoE LLMs' generalization performance on downstream tasks.
Recent findings on "grokking" in LLM pretraining reveal that routing pathways naturally evolve to become more structured and consistent across layers during the transition from memorization to generalization (Li et al., 2025c). Building on this insight, our analysis further investigates the reasons behind the performance gap and the poor generalization capabilities of pretrained routers. As illustrated in Figures 3(a) and (b), pretrained routers assign semantically similar samples in the task embedding space to distinct experts with dramatically different routing weights. Such misalignment between the task embedding manifold and routing weight manifold hinders effective knowledge sharing across tasks and underutilize the collective expertise of the experts. This misalignment between the targeted tasks and the assigned experts undermines the generalization of MoE and its

---

core principle, which is to leverage specialized experts, share skills, and transfer knowledge for related inputs.
A natural solution is to finetune the routers. Existing approaches such Dense BP (Panda et al., 2025) developed more effective pretraining objectives for routers but do not address the manifold misalignment between the targeted tasks and the routing weights across samples. This limitation motivates our exploration incorporat- ing manifold alignment into the fine-tuning objective. Specifically, our manifold alignment aims to enforce the consistency between task understanding (encoded by an embedding model) and task solving in an MoE LLM (encoded by the routing weights). As illustrated in Figure 2, for each training sample, in addition to min- imizing loss defined on the output, we encourage its intermediate layers' routing weights to move to those of its "successful neighbors" (samples with correct MoE predictions) in the task embedding space. These neigh- bors are weighted by their similarity to the sample. This training objective regularization (Belkin et al., 2006), a well-established technique in machine learning that aims preserve the local neighborhood structure of nigh-ddiensional inputs on the manifold of low-dimensional representa- tions or outputs. Unlike its original setting, we apply such regularization the routing weights across MoE layers rather than the final outputs, and establish coher- ent bindings between the expert choices (weights) and the task embedding instead of the raw inputs.
Masual-T
Mult-7B-AAR ROMA
RoMA Amprovenact
ARC-C
HelaSwag
90.7
MMLU
be formulated as manifold
75.4
60.3
WinoGrandz
Figure I: RoMA on OLMoE-7B-AIB vs. 7- 34B dense LLMs across eight benchmarks. RoMA leads to 7-15% accuracy improve- ment, consistently outperforming mod- els over eight benchmarks, demonstrating the effectiveness of post-training by RoMA.
To this end, we propose "Routing Manifold Alignment (RoMA)", a router post-training method that aligns the manifold of routing weights with task embeddings through lightweight fine-tuning of a few routers in MoE LLMs. RoMA introduces manifold regularization term to the training objective that encourages routing weights of each sample to approximate those of its successful neighbors with similar task embedding, thereby promoting consistent expert selection for semantically related inputs. Extensive periments on three recent MoE LLMs (OLMoE, DeepSeekMoE, Qwen3-MoE) demonstrate that RoMA brings substantial im- provements (7-15% in accuracy) across diverse benchmarks and outperforms SOTA routing methods, as shown Figure 1, by merely finetuning O. 0.0095% parameters of the base model, without affecting inference cost. Notably, RoMA-finetuned MoE LLMs with only 1-3B active parameters achieve competitive or su- perior performance over much larger dense models with 34B parameters. We conduct comprehensive ablation studies that further investigate the effects of key designs in RoMA, including layer selection,
Training Sample
Succeseful Neighber
Successful Neighborhood
Manilol Aligrmere
Wight +! Jol
Task Embedding Manifold
Objective: min Cunl C
Rosting Weights Manitild
186g
Frozen
Maddi
Router
Trainable
Figure 2: Overview of RoMA. RoMA finetunes routers in MoE LLM (bottom, yellow) with a training objective defined on each sample (I,, y.). which is composed of (1) the task loss Lms (i) defined on the model output f(x,, rz); and (2) the manifold align- ment regularization Lmanifold (i), which aligns the manifolds of routing weights (right, green) and the task embedding (left, blue). improves MoE's generalization by unifying solution generation in MoE with task understanding.
2

---

neighborhood configuration, and regularization strategies, validating the effectiveness of RoMA in bridging the performance gap between pretrained routers and optimal routing for MoE LLMs.
RELATED WORK
MoE LLMs Mixture of Experts (MoE) architectures have been extensively incorporated into large language models (LLMs) to enhance computational efficiency and task-specific specializa- tion (Shazeer al., 2017). Recent work such OLMoE (Muennighoff et al., 2024) and DeepSeek- MoE (Dai et al., 2024a) demonstrate the effectiveness of sparse MoE layers in reducing active parameters while maintaining model capacity. These MoE. models fundamentally rely routers to determine expert selection, typically employing token-choice routing that selectively activates subsets of experts for each input token (Fedus et al., 2022: Lepikhin al., 2020). Beyond training MoE models from scratch, MoEfication (Zhang et al.. 2022) proposes converting pretrained dense models into MoE architectures by splitting cerw-forward network parameters into functional partitions as experts. However, the quality of these routing decisions remains a critical bottleneck. Our study shows current routers often produce suboptimal routing weights that fail to fully leverage expert specialization, resulting in load imbalance and expert underutilization.
Manifold Regularization of LLMs Recent work reveals that LLM embeddings exhibit stratified manifold structures with varying dimensions across semantic domains (Li & Sarwate, 2025; Robinson et al., 2025). Furthermore, recent studies show that the routing weights MoE LLMs inherently capture rich semantic representations and can serve effective off-the-shelff embedding models (Li Zhou, 2025). While traditional manifold regularization assumes smooth global structures (Belkin et al., 2006), LLMs require more sophisticated approaches. Methods like I-STAR (Rudman & Eickhoff, 2023) control isotropy in embedding spaces, while CROW (Min et al., 2024) enforces consistency across layers. However, these techniques do not explicitly leverage manifold structures to improve MoE routing. The geometric insights from stratified manifolds suggest that different experts naturally align with different embedding strata, yet current routing mechanisms fail exploit this alignment. This gap motivates our routing manifold alignment approach, which guides routing decisions based on the data's 's inherent geometric structure.
Routing Optimization in MoE architectures has emerged as a critical component for achieving effi- cient expert utilization and balanced computation. Routing optimization methods have evolved from simple load balancing (Fedus et al., 2022; Lepikhin et al., 2020) to sophisticated strategies including differentiable top-k selection (Zhou al., 2022) and test-time optimization such as C3PO (Li et al., 2025a) and R2-T2 (Li et al., 2025b) that dynamically re-weights expert pathways. However, these ap- proaches optimize routing without considering the embedding space's s geometric structure. Moreover, C3PO introduce additional computational overhead for task embedding and nearest neighbor search, requiring 6-7x the cost of standard inference by the base model.
TASK-EXPERT ROUTING MANIFOLD MISALIGNMENT
Task Cluse:
Task Chater
Task Clastrr
Tesk
A4
(4) Task Emhisding
(b) Routing Wrights (Before RaMA)
(c) Routing Wrights (After ReMA)
(4) Optinall Routing Weights (Orade)
Figure 3: UMAP visualization of task embedding and routing weights manifolds for samples in ARC-C. (a) Their task embedding shows cluster structures. (b) Routing weights by pretrained MoE scattered and misaligned with the task embedding clusters. (c) RoMA aligns routing weights with the task embedding manifold's cluster structure. (d) RoMA also achieves a similar manifold structure that of the optimal routing weights (oracle), which explains the improvement in generalization.
3

---

MoE LLMs employ routers to assign input tokens to small subset of experts through routing weights in each layer. For sample (I,, 3s), let h, denote the hidden representation at layer l. The router at each layer & produces routing weight vector:
= R(M (9) (
(runte) E IR
where K is the number of experts, (( denotes the router function, and O represents the router parameters at layer f. The routing weight matrix ry is computed the concatenation of routing weights across L MoE layers:
(2) (2). P m= 1=-r::n0 (T)
(1) T R
(2)
Throughout this paper, r, denotes the routing weights (the softmax output used to aggregate expert outputs for sample 2;)), while fouter denotes the router parameters (the learnable weights in the router network). Our method optimizes Orowter to align the manifold of Ty with task embeddings.
As shown in Table 1, evaluations across broad downstream tasks reveal that routers in existing MoE LLMs produce suboptimal routing weights r;, which lead to a performance gap 10-20% accuracy when compared to the optimal routing weights (oracle) rf:
r A arg min CCe (f v.))..
where -(-.-) represents the MoE model that takes input x; and routing weight vector r to produce output, y is the ground truth label for input T₁, and Lce is the cross-entropy loss. We obtain rद by initializing r with the pretrained routing weights and performing gradient descent with access to the ground truth label until convergence. This oracle serves as an empirical upper bound to quantify the untapped potential of existing routers.
To investigate the root causes behind the observed performance gap MoE LLMs, we conduct comprehensive analysis of the relationship between task embeddings and routing weights in Figure 3.
The comparison between task emheddings (Figure 3()) and pretrained routing weights (Figure 3(b)) reveals severe misalignment. While the task embedding space presents clear cluster structures where semantically similar samples are grouped together, the pretrained routing weights show no corresponding clustering patterns. Instead, samples from the same semantic cluster are scattered across the routing weights space. This manifold misalignment indicates that the pretrained routers fail to capture the underlying task structure, leading to inconsistent expert selection for semantically related inputs. To further substantiate this visual intuition, we provide quantitative alignment metrics (e.g., CKA similarity, Trustworthiness) in Appendix A.3, which consistently confirm the misalignment in baselines and the improvement brought by RoMA.
In contrast, the oracle routing weights (Figure 3(d) demonstrate clear cluster structure the task embedding structure, with samples from the same semantic group receiving similar routing patterns. This alignment between task understanding and expert assignment is precisely what enables the oracle to achieve superior performance, highlighting that the task-expert routing manifold misalignment is the key bottleneck limiting router generalization MoE LLMs.
4 ROUTING MANIFOLD REGULARIZZTION (ROMA)
To address this limitation, we propose "Routing Manifold Alignment (RoMA)". a post-training method that aligns the manifold of routing weights with task embeddings through lightweight router fine-tuning. Our key insight is that samples with similar task embeddings should share similar routing patterns to leverage specialized expertise effectively. To achieve this, we introduce manifold regularization term that encourages alignment between the routing weight manifold and the task embedding manifold. Given training set D {(x,, })) and their associated routing weights {r:)i, (where r; denotes the concatenated routing weights across multiple layers), our goal is to optimize the routers such that samples with similar task embeddings share similar routing patterns.
4

---

4.1 SUCCESSFUL NEIGHBORHOOD TO IMITATE
We first identify the subset of training samples where the MoE produces correct predictions:
S = {j
[n] f(z,. - ツ
This filtering ensures that our finetuning only imitates from routing patterns for samples in S that lead to successful outputs, preventing the propagation of suboptimal routing strategies.
Given the set of successful samples S, we construct neighborhood. ((x;) for each sample based on the task similarity in an embedding space. Let E(-) denote pre-trained embedding model that maps input task descriptions/instrrettt can be defined via K-Nearest Neighbors or e-ball:
to semantic representation space. The neighborhood of x;
k-NN: M(x;) arg
max ASS.AAERR
Sim(E(z,),
(5)
jEA
e-ball: M(x;)) = (j S: sim (x.,,, ((x;)) 2 }
(6)
where sim(:,) is similarity metric, for example, the Gaussian similarity is defined
[[(,,,)) Exxx))
(im(E(z),EEzz))) exp
20
4.2 TRAINING OBJECTIVE WITH MANIFOLD REGULARIZATION
Having identified the successful neighborhood for each sample, our next step is to incorporate this structure into the training objective align routing behaviors with the task embedding geometry. The key idea is that semantically similar samples should not only cluster in the embedding space but also share consistent routing patterns. To achieve this, we introduce manifold regularization term that plicitly aligns the routing weights manifold with the task embedding manifold by encouraging samples to follow the routing patterns of their successful neighbors, weighted by their semantic similarity.
The (normalized) adjacency Wi, between sample æ; and xz is defined as sim(E(x)) F(x;)))
W Σ (x,)( sim{ (E(z,) E(x;))) Vj EVx)),
(8)
where higher weights indicate stronger semantic similarity in the task embedding space. Given Way- the manifold regularization applied to the routing weight r, of sample x; is defined as
Cmanifold (i)
Σ
Wiellr: 13.
(9)
JEV(s:
By penalizing routing discrepancies ||r; -r,- (manitold( enforces the routing weights manifold to be aligned with the task embedding manifold. Moreover, it moves each sample's 's routing weights to those of its "successful neighbors" in the task embedding space.
ry ll2 between semantically similar samples with large Wa
As consequence, the manifold regularization consolidates the bindings between tasks and their expert choices, and thus improves the generalization.
To ensure that aligned routing patterns also lead to correct predictions, the training objective in RoMA applies the manifold regularization the cross-entropy loss LaE defined on the outputs.
Lusk (i) Lce(f(z,rrt) y.).
(10)
With a regularization coefficient > O, the final objective on sample x; is
CROMA (i) Lusk (i) (i)
(11)
During training, we only update router parameters while keeping all expert parameters frozen. The gradient update is performed via backpropagation on PROMA with respect router parameters:
0 (t++) ((t)
V Proms LROMA,
(12)
where fouter represents the parameters of routers and 7 is the learning rate. While router parameters represent only small fraction of the total model parameters (0.0095%), we empirically find that only finetuning routers in the last five layers achieves superior performance while significantly saves the training cost, demonstrated in Figure 6.
5

---

5 EXPERIMENTS
5.1 EXPERIMENTAL SETTINGS
Models We evaluate three recent MoE LLMs: OLMoE-7B-AIB. DeepSeckMoE-IGE 16B-ABBB and Swen3-30B-A3B. OLMoE features a 16-layer transformer with 64 experts per layer, activating 8 per token, totaling 6.9B parameters with 1.3B active per token. DeepSeeKMoE uses 28-layer transformer with shared and 64 routed experts per layer, activating all shared plus 6 routed experts per token, totaling 16.4B parameters with 2.8B active per forward pass. Qwen3-30B-A3B employs 48-layer transformer with 128 experts per layer, activating 8 per token, totaling 30.5B parameters with 3.3B active per token. These models exemplify distinct MoE designs and scales, enabling a comprehensive evaluation of routing dynamics and generalization behavior.
Baselines We evaluate RoMA against both different adaptation methods (See Table 1) and other mod- els (See Table 2) across eight benchmarks. For adaptation methods, we compare with: (1) In-Context Learning (ICL) (Brown al., 2020) with embedding- baasd retrieval for few-shot demonstrations; (2) Router Tuning that directly updates the routers; (3) Oracle Tuning that fine-tunes routers with access to optimal routing weights (oracle); (4) Prefix Tuning (Li & Liang, 2021) and Soft Prompt Tuning (Lester et al., 2021) that introduce lightweight trainable parameters while keeping the base model frozen; (5) Dense Backpropagation (Dense BP) (Panda et al., 2025) that enables gradient flow through the full model while updating few parameters; (6) C3PO (Li et al., 2025a), ofate-of-the-ee art test-time routing weights optimization method. For model comparison, we evaluate against models grouped by active parameters (IB, 3B, 7-8B, 13-14B, 27-34B) including recent models like Llama3.2 2, Gemma2, Qwen2, and Mistral to assess the efficiency of MoE architectures enhanced with RoMA.
Training Set comprises 49,000 samples distributed across five task categories, as shown in Figure 4. The dataset in- cludes General Knowledge tasks (BIG-Bench and Super- GLUE), Commonsense reasoning (CommonsenseQA and SocialIQA), Science QA (OpenBookQA and SciQ), Read- ing comprehension (MultiRC), and Coreference resolution (KnowRef). This diverse composition ensures comprehen- sive coverage across different reasoning capabilities for ef- fective training.
Krowledgs PipeCLUT
Training Set
I55)
Figure 4: Training set statistics.
Benchmarks We evaluate RoMA on eight diverse benchmarks. The evaluation suite includes MMLU, HellaSwag, PIQA, ARC- ARC-Challenge, ARC-Easy WinoGrande, BoolQ and GSM8K. Notably, GSM8K serves Diittribbbtti (OOD) benchmark since training set doesn' 't contain math-related data. Details about training set and benchmarks in Appendix A.8 and A.9.
5.2 MAIN RESULTS
Advantage of RoMA over different adaptation methods. Table compares adaptation methods on OLMoE, DeepSeekMoE, and Qwen3-MoE across eight benchmarks. Lightweight methods (ICL, Router/Prefix/Promp Tuning) yield only modest gains, while Oracle tuning and Dense BP achieve stronger, though still limited, improvements relative to the Oracle upper bound. C3PO performs better than these baselines, yet RoMA achieves the highest overall accuracy. On MMLU, RoMA boosts DeePSeekMoE from 46.2% to 56.8% (+10.6%) and OLMoE from 57.8% to 69.0% (+11.2%), surpassing C3PO by +1.4% and +3.5%, respectively, Although C3PO achieves comparable accuracy as RoMA, its inference cost is 6-7x higher than both RoMA and the base model (See Figure 5), highlighting RoMA superior efficiency-effeeeeee effectiveness trade-off. In addition, RoMA shows more advantages over C3PO on larger models such as Mee and Qwen3-MoF. The accuracy and cost of the other two models are reported in Appendix A. and A.2. We further compare RoMA with parfeeter-efficieet fine-tuning (PEFT) methods, including LoRA, DoRA, and MoLE, applied to the router parameters. As detailed in Appendix A.4, RoMA outperforms these methods by significant margins (+7.5% ~ +8.6% on average) while introducing zero new parameters, highlighting that manifold alignment is more effective than merely increasing parameter capacity for routing optimization.
Comparison of routing weights manifold before and after RoMA. Figure illustrates the effect of RoMA on routing weights. After applying RoMA, routing weights form clear clusters (Figure 3(a))

---

OLM Bose
OLMoE C3PO
OLMoË RoMA
ARC-C
ARC-C
ARC-B
672
HellaSwag
ARCHE
HellaSwag
66.3
12.9
86.7
13.2
87
85.3
2.1
85.3
51.3
77.9
19
PIQA
38.0
85.8
80.7
57.8
69.0
MMLU
PIQA
12:4
19
1.6
2.3
15.8
MMLU
72.2
2.4
1
26
1.9
81
49.4
16.3
17.9
82.7
79.6
50
GSM-8K
WinoGrande
81.7
13.7
BoolQ
Bool(]
(a) Accuracy
(b) Inference Cost
Figure 5: Performance and inference cost of OLMoE (base model), OLMoE + C3PO and OLMoE + RoMA across eight benchmarks. (a) Accuracy: RoMA consistently improves the base model's performance to be comparable or better than C3PO. (b) Inference cost (in FLOPs x 1002); RoMA maintains nearly the same efficiency as base model, while C3PO requires test-time optimization and induces 6-7x more FL.OPs. These results highlight the effectiveness and efficiency of RoMA.
that closely align with the task embedding structure (Figure 3(c)). In contrast, the pretrained routing weights show little alignment with task clusters in Figure 3(b), highlighting that RoMA effectively resolves the manifold misalignment problem. Furthermore, the post-RoMA routing patterns closely resemble the oracle routing weights as shown in Figure 3(d), suggesting that our optimization moves the model toward thcoretically optimal expert assignments. As a result, samples within the same task cluster receive similar routing patterns, enabling more consistent and efficient use of specialized expertise and bridging the performance gap between suboptimal pretrained routing and ideal oracle routing.
Advantage of RoMA over State-of-the-Art models. Table 2 reports LLM performance across eight benehmarks with varying active parameter counts. Notably, OL-oFE7B-AA B+RoMA, with only 1B active parameters, achieves 69.0% on MML.U and 86.7% on HellaSwag, surpassing several 7-8B and even 13B dense models. Similarly, DeepSeckMoE- Bee-----------BBB AOBARRMA (3B active) delivers substantial gains, matching exceeding the performance of dense LLMs up to 34B parameters. These results demonstrate that RoMA consistently improves routing quality, enabling small active active-paramen parameter MoEs to rival outperform much larger dense counterparts. Details of models in Appendix A.10.
5.3 ABLATION STUDY
We perform series ablation studies to systematically analyze the design choices behind RoMA. Specifically, we investigate: (i) which layers to regularize, (ii) which token positions to use for routing guidance, (iii) how select neighbors for manifold alignment, (iv) the effect of training set size, and (v) the choice of regularization method. These ablations help identify the most effective and efficient configuration, revealing which factors are critical for performance. All experiments are conducted on OLMoF, and experiment results on DeepSeekMoF are provided in the Appendix A.
Layer Selection Figure 6 examines how applying routing manifold regularization to different subsets of layers affects model performance. Applying RoMA to single layer yields only modest gains (69.1-69 (9),,,)) while extending it to two layers improves accuracy above 71%. Performance continues to increase as more layers regularized, with the last five layers (L5) achieving the highest accuracy of 76.2%, even surpassing the All-Layer configuration (75.7 1).. These results highlight that the final layers are particularly critical for routing quality, and that selectively regularizing small set of strategically important layers is both more effective and more efficient than uniformly applying RoMA across all layers.
7

---

Table 1: Comparison of RoMA with the Base model, Oracle, test-time adaptation methods (ICL, C3PO), training-basee methods RRouterOOnclee/Peiii/Proooo DeepSeekMoF, OLMoE, and Swen3-330-A33- Details of the baselines and benchmarks are provided in Section 5.1. Bold numbers denote the best performance (excluding Oracle), and underlined numbers denote the second best. RoMA improves DeepSeekMoE from 46.2% to 56.8% (+10.6%), improves OLMoE from 57.8% to 69.0% (+1 1.2%), and improves Qwen3-30B-A3I from 74.2% to 78.8% (+4.6%) on MMLU, outperforming C3PO on all three models.
Tuning), across eight benchmarks on
Hella- Swag
Wino- Grande
Method
MMLU
ARC-C ARC-E
PIQA
BoolQ
GSM8K
Avg
DeepSeckMoE-16B--
Base model
46.2
78.0
50.3
79.9
70.1
72.3
62.2
66.6
Oracle
63.8
92.5
70.8
85.2
90.3
82.1
83.2
71.8
80.0
ICI.
49.0
81.6
56.3
76.2
81.4
72.3
75.8
65.7
69.8
С3PO
55.4
85.7
80.7
85.8
77.5
78.2
68.5
74.2
Router Tuning
49.3
81.5
57.2
76.6
82.0
73.8
74.5
64.8
70.0
Oracle Tuning
54.2
84.3
60.1
79.5
84.0
76.0
77.5
66.2
72.7
Prefix Tuning
73.8
79.2
70.3
73.1
64.8
67.4
Prompt Tuning
49.3
78.6
74.7
80.5
72.0
74.2
65.5
68.7
Dense BP
50.1
80.2
54.8
77.3
81.7
74.2
76.1
63.9
69.8
RoMA (Ours)
56.8
87.9
61.4
81.5
85.2
76.8
80.6
67.4
74.7
OLMoE-7B-AII
Base model
$7.8
77.9
51.3
79.8
80.7
72.2
75.4
45.5
67.6
Oracle
72.2
91.5
74.8
91.4
93.6
87.7
84.5
53.2
81.1
60.3
80.6
82.5
83.6
76.8
78.9
48.5
71.2
C3PO
65.5
85.3
66.3
87.4
88.0
82.7
79.6
50.8
75.7
Router Tuning
63.2
81.7
62.5
83.8
80.9
75.3
77.8
47.2
71.6
Oracle Tuning
66.8
84.2
65.4
86.1
86.2
80.5
79.9
49.0
74.8
Prefix Tuning
59.3
78.2
54.5
80.4
82.1
73.5
76.8
46.7
68.9
Prompt Tuning
59.7
79.5
55.9
81.3
82.4
74.1
77.2
47.3
69.7
Dense BP
61.8
82.4
57.3
84.1
83.9
76.9
75.2
48.1
71.2
RoMA (Ours)
69.0
86.7
67.2
88.0
85.8
81.8
81.7
49.4
76.2
Qwen3-30B-A.B9
Base model
74.2
68.5
56.8
84.3
78.5
65.2
81.3
83.4
74.0
Oracle
82.5
80.3
69.2
92.6
87.4
77.3
90.5
90.9
83.8
ICL C3PO
75.8
70.7
59,3
86.1
80.2
67.8
83.5
84.7
76.0
77.9
74.1
63.4
88.1
81.7
71.9
85.4
86.0
78.6
Router Tuning
75.3
70.3
60.1
85.7
79.8
68.5
82.8
84.2
75.8
Oracle Tuning
77.2
73.5
62.8
87.6
81.3
71.2
84.9
85.5
78.0
Prefix Tuning
74.5
68.9
$7.9
84.8
79.1
66.3
82.1
83.8
74.7
Prompt Tuning
75.0
69.6
58.6
85.2
79.6
67.0
82.7
84.0
75.2
Dense BP
76.1
71.4
59.8
86.5
80.5
69.2
83.8
84.9
76.5
RoMA (Ours)
78.8
74.8
65.5
88.6
83.1
73.8
85.1
86.3
79.5
Token Selection Figure 7 presents the effect different token selection strategies when applying RoMA. Using multiple tokens (e.g., the first three or middle three) provides moderate improvements over the baseline, with last 3 tokens (Last3) reaching 74.5%. Among single-token choices, the last token (Lastl) performs best (76.2%), outperforming both the first one token (Firstl) (71.4%) and the middle one token (Middle1) 1) (69.2%). These results indicate that the final tokens contain richer task-relevant information for guiding expert routing than earlier middle tokens. Moreover, the superiority of Lastl over Last3 highlights that a single, well-chosen token can be more effective and efficient than aggregating multiple tokens.
8

---

Table 2: Comparison of LLMs with varying active parameters (IB, 3B, 7-8B, 13-14B, 27-34B) evaluated on eight benchmarks. MoE models post-trained by RoMA achieve strong performance, surpassing or matching the performance of much larger dense models. For example, OLMoE-7B-AN (IB active) achieves 69.00 69.0% on MMLU and 86.7% on HellaSwag, outperforming several 7-8B and even 13B dense counterparts. Qwen3-30B-A3B (3B active) achieves 78.8% on MMLU, surpassing even 27-34B dense models, highlighting the effectiveness of MoE+RoMA.
Hella- Swag
Wino- Grande
MMLU
ARC-C ARC-E
PIQA
BoolQ
GSM8K
1B active parameters
Llama3.2 2-18
27.4
57.9
32.1
53.9
72.4
57.4
63.7
39,4
OLMo-1B
24.1
61.8
29.6
55.7
75.6
56.8
64.2
28.5
Ol--F-BB-AIB
57.8
77.9
51.3
79.8
80.7
72.2
75.4
45.5
3B active parameters
Gemma2-3B
43.7
66.3
58.4
75.2
71.8
64.5
73.1
41.4
DeepSeeKMoE. 16B-A3B
46.2
78.0
50.3
73.8
79.9
72.3
62.2
Qwen3-30B-ASS
74.2
68.5
56.8
84.3
78.5
65.2
81.3
83.4
77-BB active parameters
Qwen2-7B
53.4
74.9
45.8
69.7
77.2
68.1
84.8
79.9
Mistral-7B
59.6
81.0
53.8
79.6
82.2
74.0
68.1
Llama3.1-8B 1-8BB
57.7
77.9
48.7
80.8
81.4
73.5
81.9
49.6
1-3-14B active parameters
Llama2-131
53.8
78.6
50.1
74.5
75.7
35.2
Vicuna-11 13B
51.3
76.2
47.4
72.8
78.0
68.2
71.5
32.2
Qwen1 .5-14B
66.7
81.5
58.0
85.3
82.1
76.9
81.3
58,4
~77334 active parameters
Gemma2-27B
75.2
86.4
71.4
88.9
83.2
79.0
84.5
Yi-34B
73.5
83.1
58.2
82.6
82.6
78.9
83.1
63.5
Llama2-34B
62.6
79.4
54.5
77.5
81.9
76.0
78.1
42.2
RoMA (Ours)
DeepSeekMoE- 16B-A3B
56.8
87.9
61.4
81.5
85.2
76.8
80.6
4
OLMoE-7B -BB
69.0
86.7
67.2
88.0
85.8
81.8
81.7
49.4
Qwen3-30B-A3B
78.8
74.8
65.5
88.6
83.1
73.8
85.1
86.3
76.2
Ang Accuricy (%6
Layers Layens
Layens
75.
75
Layor
71.8
72.1
72.6
72.3
73.8
72.9
70.6
69
71
67.6
69.1
63.4
69.7
69.9
Bas
L1
FIMI
FILI
MILI
12
F2MS
P2L1
MIL1
E3
M5
L5 (Qars)
Allié
Figure 6: Applying RoMA at different layers (F: early layers, M: middle layers, L: late layers).
Fine-tuning the routers in the last five layers (L5, RoMA) achieves the best performance.
Neighborhood Selection Figure 8 compares different strategies for selecting neighbors in RoMA. Random neighbor selection yields almost no improvement over the baseline (67.8% vs. 67.6%). Using --eighborrhoods shows sensitivity to the choice of radius: performance improves steadily from 68.9% (c=0.3) to a peak of 74.1% 1% at c=0.5, but drops slightly when the radius grows larger (c=0.7). In contrast, k-nearest neighbor selection provides more stable gains, with k=3 achieving the best overall accuracy of 76.2%. Notably, this surpasses both smaller (k=1) and larger (k=5) settings, suggesting that moderate number of neighbors balances robustness and noise. These results

---

Avg Accuracy (%g
Baseiie
Tokens
Tokzn
74.5
71.4
36.2
25
sndighbur
k-aelahbor
Baeliie
72.5
72.3
76.2
70.1
68.7
69.2
Ouze
67.6
Avg
67.6
67,8
Rane
Fine?
Midd:
Last
Fieel
Middlel
Lael
Booc
Rand
0.0.
0.0
2-1
k-1
k-5
Figure 7: Applying RoMA to routing weights of tokens at different positions. Regularizing the Lastl token's s routing weights performs the best.
Figure 8: Comparing neighbor selection strategies in RoMA. Rand-randon neighbors. k-NN with achieves the best performance.
Arg Arg.Ascumo Ascunay t50
75.0
76.2
Accuncc (Pg)
Barelize Ragulirisainn Unes
76.2
Set Size
73.6
71.5
70.7
70
67.6
67.6
68.2
85
Bas:
106
305
30%
Busc
Eniruyy
Manilell
Figure 9: Comparing different training set sizes for RoMA on OLMoE. While the full training set (100%) yields the best performance, 30% suffices to achieve substantial gains over the baselines.
Figure 10: Comparing different regularization methods with RoMA. RoMA's manifold regular- ization achieves the best performance.
highlight that careful neighborhood design is crucial for effective manifold alignment, and that chosen 3=3 strategy offers the most reliable improvement.
Training Set Size Figure 9 examines how the size of the training used for RoMA affects performance. Starting from the baseline accuracy of 67.6%, even using only 10% of the training data yields noticeable gain (68.5%). Performance improves steadily more data is available, reaching 70.8% at 30% and 73.6% at 50%. With 70% of the data, accuracy rises further to 75.0%, and using the full dataset achieves the best performance of 76.2%. These results demonstrate that RoMA benefits consistently from additional training data, but also that substantial improvements can already be obtained with a relatively small fraction of the dataset, highlighting data efficiency.
Regularization Methods Figure 10 compares different regularization strategies applied to the router. Standard techniques such as L1 and L2 penalties yield only modest improvements over the baseline (68.2% and 71.5%, respectively), while entropy regularization reaches similar level (70.7%). In contrast, our proposed manifold regularization achieves best result 76.2%, substantially outperforming all alternatives. This demonstrates that aligning routing weights in the task embedding space provides a more effective inductive bias than generic sparsity or entropy- based constraints, highlighting the unique advantage of RoMA's manifold perspective.
Embedding Models We further investigate the sensitivity of RoMA the choice of embedding models. We evaluate RoMA using diverse embedding models ranging from 22M to 7.8B parameters (including all-MiniLM, BGE, and Qwen-embedding). As detailed in Appendix A.7, RoMA achieves consistent accuracy improvements (+3.6% ~ +8.6%) across all tested models.
6 CONCLUSIONS
Our work introduces RoMA, lightweight router post-training method for sparse Mixture-of-Eaae of-Exxerts LLMs. By aligning routing weights with the underlying task embedding manifold through manifold regularization, RoMA addresses the fundamental misalignment between task understanding and expert utilization in MoE models. Our approach requires updating only router parameters while keeping experts frozen, yet consistently improves accuracy across diverse benchmarks by 7-15% without increasing inference cost. Extensive experiments demonstrate that RoMA enables small parameee MoEs to rival or even surpass much larger dense models, highlighting both the efficiency and effectiveness of RoMA. Beyond performance gains, findings underscore the importance of geometric alignment between task representation and expert selection, offering new insights for advancing routing strategies in future MoE architectures.
10

---

REFERENCES
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023.
Mikhail Belkin, Partha Niyogi, and Vikas Sindhwani. Manifold regularization: geometric frame- work for learning from labeled and unlabeled examples. Jourmal of machine learning research, 7 (11), 2006.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: physical commonsense natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
Reasoning about
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Wei-Lin Chiang, Zhuohan Li, 7i Lin, et al. Vicuna: An open-source chatbot impressing gpt-4 with 90% chatgpt quality. Blog post, 2023.
Christopher Clark, Kenton Lee, Ming-Wor Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL), 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing (EMNLP), 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. In arXiv preprint arXiv:2 14168, 2021.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Yu Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint arXiv:2401. 06066, 2024a.
Wenxuan Dai et al. Deepscekmoe: Towards ultimate expert specialization mixture-of-cxperts language models. arXiv preprint 1.xiv:2401.06066 2024b.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad AI-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024.
Ali Emami, Paul Trichelair, Adam Trischler, Kaheer Suleman, Hannes Schulz, and Jackie Chi Kit. Cheung. The knowref coreference corpus: Removing gender and number cues for difficult pronominal anaphora resolution. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL), 2019.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
Dirk Groeneveld, Kyle Lo, et al. Olmo: Accelerating the science of language models. arXiv preprint arXiv:2402.00833 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations (ICLR), 2021.
Albert Jiang al. Mistral 7b. arXiv preprint arXiv:2310.0066 2023.

---

Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Proceed- ings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL), 2018.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter- efficient prompt tuning. arXiv preprint arXiv: (et::2204408991, 2021.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2 10:::2101.001000 000090,00000000 2021.
Xin Li and Anand Sarwate. Unraveling the localized latents: Learning stratified manifold structures in Ilm embedding space with sparse mixture-of of-eere arXiv preprint arXiv:2502.13577 2025.
Zhongyang Li, Ziyue Li, and Tianyi Zhou, C3po: Critical-layer, core-expert, collaborative pathway optimization for test-time expert re- mixing. arXiv preprint arXiv:2504.07964, 2025a.
Zhongyang Li et al. R2-T2: Re-routing in test-time for multimodal mixture-of-experts. In Interna- tional Conference on Machine Learning, 2025b.
Ziyue Li and Tianyi Zhou, YOUR MIXTURE-OF-EXPERT LLM SECRETLY AN EMBED-
DING MODEL FOR FREE. In The Thirteenth International Conference on Learning Representa- tions, 2025.
Ziyue Li et al. Grokking in LLM pretraining? monitor ceamoorzzaionntt-geeerraaiiliiii without test. arXiv preprint arXiv:2506.2155/ 2025c.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of 2018 Conference on Empirical Methods in Natural Language Processing (EMNLP), 2018.
Nay Myat Min, Long H Pham, Yige Li, and Jun Sun. Crow: Eliminating backdoors from large language models via internal consistency regularization. arXiv preprint arXiv:2411.12766. 2024.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. Olmoe: Open mixture-of-experts language models. arXiv preprint arXiv:2409.0200 020000000000...... 2024.
Ashwinee Panda, Vatsal Baherwani, Zain Sarwar, Benjamin Therien, Supriyo Chakraborty, and Tom Goldstein. Dense backpropagation improves training for sparse mixture-of-expper arXiv preprint arXiv:2504. /2463, 2025.
Michael Robinson, Sourya Dey, and Tony Chiang. Token embeddings violate the manifold hypothesis. arXiv preprint 20X002250..01002, 2025.
William Rudman and Carsten Eickhoff. Stable anisotropic regularization. arXiv preprint arXiv:2305.19358, 2023.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi Winogrande: An adversarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing (EMNLP), 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-exppen layer. arXiv preprint arXiv: /00::/700000000 065388 2017.
12

---

Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseqa: A question answering challenge targeting commonsense knowledge. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL), 2019.
Gemma Team. Gemma 2: Improving open language models practical size. arXiv preprint arXiv:2408.0011 2024a.
Qwen Team. Qwen2 technical report. arXiv preprint arXiv:2407. 2024b.
Hugo Touvron al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXivv2337.09228, 2023.
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: stickier benchmark for general-purpose language understanding systems. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Jason Wei et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint urXiv:2206.04611 2461200060011, 2022.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsoureing multiple choice science questions. In Proceedings of the 3rd Workshop Noisy 'ser-generated Text (W-NUT), 2017.
Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Guoyin Wang, Heng Li, Jiangcheng Zhu, Jianqun Chen, et al. Yi: Open foundation models by 01. ai. arXiv preprint arXiv:2403.04652 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL), 2019.
Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Moefication: Transformer feed-forward layers mixtures of experts. In Findings of the Association for Computational Linguistics: ACL 2022, pp. 877-890, 2022.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. ofxtere-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
13

---

A APPENDIX
A.1 COMPARISON BETWEEN ROMA AND BASELINES ON DEEPSEEKMOE--OB-AB AND
QWEN3-30B-A3B
Gemma2-18
Misral-7
58
Llama2
OLMoE 7B-A1B
OLMoE 7B-AIB RoMA
RoMA Improvement
ARC-C
ARC-C
61.4
65.5
HellaSwag
HellaSwag
50
F%-Q)
$8.5
MIQA
MMLU
PIQA
83.1
28.5
74.2
16.
MMLUU
20:1
699
62.2
76.8
$1.3
WinoGrande
GIM--KK
WanGGrande
19.5
GSM-8K
85.1
BoolQ
BoolQ
(a) DeepSeck
(b) Qwen3-30B-A3B
Figure 11: (a): Radar figure of DoepSeckMoE-16B-A.. (b) Radar figure Qwen3-30B-A33
RoMA consistently improves model's performance on multiple benchmarks.
A.2 ACCURACY AND COST OF DUESSEEKMOE-16B-A33 AND QWEN3-30B-AM
DoegSeckkMo Hose
DespStciMat C3PO
DaggSecMMe RoMA
61.6
HellaSwag
ARC-E E
47.5
HelliSwag
7.8
73.8
50.3
78.0
6,3
5.0
5.5
PIQA
55.2
55.4
56.8
MMLU
PIQA
62
7.0
44
MMLU
4.7
72.3
6.7
49.4
WinsCardee
8066
GSM-SK
Winoirande
GSM 8K
33.8
BuolQ
(a) Accuracy
(b) Inference Cost
Figure 12: (a) Accuracy: RoMA achieves similar accuracy improvement C3PO DeepSeekMoF. 16B-A3B, (b) Inference cost (in FLOPs x x10"1): RoMA maintains nearly the same efficiency as the base model, while C3PO requires test-time optimization and induces 6-7x more FLOPs.
A.3 QUANTITATIVE ALIGNMENT ANALYSIS
To supplement the UMAP visualizations in Figure 3, we conduct quantitative alignment analyses ARC-C and MMLU benchmarks. We report three metries: (1) Subspace Similarity (CKA) to
14

---

QwesMoFF
QwenMoPF+
QwenMeE
ARC-C
ARC-C
63.5
ARC-E
HellaSSag
HellsSwap
36.6
[8.1
74.1
58.8
8,5
10,]
6.2
$4.3
PIQA
83.1
$1.7
74.2
76.M
MMLU
PIQA
6.8
6.0
3.7
7.7
6.2
7.3
71.9
3.9
53.0
418
WinoClande
85.1
GSM-AK
WanGGrande
GSM
BoolQ
Accuracy
(b) Inference Cost
Figure 13: (a) Accuracy: RoMA achieves similar accuracy improvement as C3PO on Qwen3-30B- A3B. (b) Inference cost (in FLOPs x10² RoMA maintains nearly the same efficiency as the base model, while C3PO requires test-time optimization and induces 6-7x more FLOPs.
measure the similarity between routing-weight subspaces and task embeddings; (2) k-NN Neighbor Consistency (k: = 5) to measure how well routing decisions preserve semantic neighborhoods; and (3) Trustworthiness Score (k- 10) evaluate local structure preservation.
Table 3: Quantitative alignment metries before and after RoMA. All metrics show significant improvement, quantitatively confirming the manifold alignment.
CKA Similarity (f)
k-NN Consistency (f)
Trustworthiness (1)
Dataset
Base
RoMA
Base
RoMA
Base
RoMA
ARC-C MMLU 0.21
0.18
0.47
24.3%
51.7%
0.53
0.76
0.52
26.8%
54.2%
0.56
0.79
A.4 COMPARISON WITH PEFT METHODS
We compare RoMA with representative PEFT baselines (LoRA, DoRA, MoLE) applied to the router parameters. The results averaged over 8 benchmarks. As shown in Table 4, RoMA achieves superior performance without introducing any new parameters during inference.
Table 4: Comparison with PEFT methods on OLMoE- (-oE-B--AIB (Last 5 Layers). RoMA outperforms PEFT variants with fewer trainable parameters and new parameters.
Method
Rank
Trainable Params
New Params
Avg. Ace (%)
A Acc
Baseline
67.6
LoRA
16
331K
331K
71.2
43.6
DoRA
16
334K
334K
71.6
+4.0
MoLE (3xLoRA)
16
1.00M
1.00M
72.5
44.9
RoMA (Ours)
658K
0
76.2
+8.6
A.5 TRAINING COST ANALYSIS
We provide a breakdown of the training cost for RoMA on DOMMME-BB-AIL (49K samples). The training process consists of task embedding computation, k-NN search, and router fine-tuning.
15

---

Table FLOPs breakdown for RoMA training. The k-NN search cost is negligible.
Component
FLOPss (x10'))
Share
Task Embedding (Pre-computed)
3.8
3,0%
k-NN Search (FAISS)
0.002
0.01%
Router Fine-tuning
122
97.0%
Total
125.8
100%
As shown in Table 5, the overhead from k-NN search negligible (< 0.01%) thanks efficient approximate nearest neighbor search (FAISS). The task embeddings are computed once and reused.
A.6 BIAS MITIGATION WITH CURRICULUM LEARNING
To address the concern that imitating only "successful neighbors" might introduce confirmation bias, explored curriculum learning strategy. We tested 3-stage curriculum: (1) Early stage: include neighbors from both correct and incorrect predictions (top-30% similarity); (2) Mid stage: top-50% similarity mixed; (3) Late stage: strict filter (only correct predictions).
The curriculum strategy yields slight improvement (76.4%) over the strict filter (76.2%), while an always-soft filter degrades performance (7). This suggests that while strict filter is robust, relaxing it early in training can provide marginal benefits. For simplicity and robustness, the main results in the paper use the strict filter.
A.7 ROBUSTNESS ACROSS EMBEDDING MODELS
We evaluated RoMA on OLMoE-7B-AI AB using various embedding models to assess its sensitivity. (22M), all-mpnet-base-v2 (110M), Qwen3-0.6B-emmeeddee
The models include (0.6B), multilingual-gemmieeeee (2.6B), Qwene--w--insteeeeeeeeeeee (7B), and NV-Embed-v2 (7.8B).
Table 6: Performance of RoMA with different embedding models on OLMoF-7B-AIB. RoMA brings consistent improvements across varying embedding model sizes.
Embedding Model
Size
Avg. Acc (%)
A Acc
Baseline (OLMoE)
67.6
all-MiniLM--L6-v2
22M
71.2
+3.6
all-mpnet-base-v2
110M
72.5
+4.9
(owe33-0.6Beembeddee
0.6B
75.9
+8.3
bge-multilnggualgeent
2.6B
73.4
45.8
gte- gte-Qwen2-7B. -B-istrrutttt instruct
7B
75.1
+7.5
NY-Embed-v2 (Ours)
7.8B
+8.6
As shown in Table 6, RoMA consistently improves performance regardless the embedding model size. Larger embedding models generally bring better alignment (higher accuracy gain), but compact models provide meaningful improvements.
A.8 DETAILS OF TRAINING SET
Our training set comprises 49,000 samples distributed across five task categories, ensuring compre- hensive coverage across diverse reasoning skills:
BIG-Bench (Wei et al., 2022): A large-scale collaborative benchmark covering a wide range of tasks such logical reasoning, linguistic phenomena, and commonsense knowledge. It is designed to probe broad generalization and emergent capabilities in large language models.
SuperGLUE (Wang et al., 2019): benchmark suite for general natural language under- standing, consisting of challenging tasks such natural language inference, word sense disambiguation, and question answering. extends the original GLUE benchmark to push models toward higher-level reasoning.
16

---

Benchmarks
Size
General Know Knowledge Jedge (40.8%)
MMLL
14,042
BiGi-Bnch (10,000)
Hellaawag
10,042
SuperGLUE (10,000)
PIQA
1,838
Comminnscnue (32.7%)
CommonsauseQA (8,000)
ARC-C ARC-E
1.172
SociallQA (8,000)
2,376
Training Set
Science QA (12.2%)
WinoGrande
1,267
Ope (3.000)
(3,000)
BoniQ
3,227
Coreference (6))
(3,000)
Reading (8. (%4%
Figure 15: Overview evaluation bench- marks we use for router post-training. Benchmarks are reserved strictly for eval- uation.
MuliRC (4,000)
Figure 14: Training Set by Task
Commonsense (Talmor et al., 2019): A multiple-choice dataset focusing on common- sense reasoning, requiring models to connect concepts and apply everyday knowledge.
SocialIQA (Sap et al., 2019): benchmark for social commonsense reasoning, where models must infer likely intents, reactions, and motivations of people in everyday situations.
OpenBookQA (Mihaylov et al., 2018): A science-oriented QA dataset requiring both retrieval from small "open book" of facts and additional commonsense reasoning.
SciQ (Welbl et al., 2017): A dataset science exam-style questions covering physics, biology, and chemistry, testing factual recall and reasoning in scientific contexts.
MultiRC (Khashabi et al., 2018): A reading comprehension benchmark with multi-sentence passages and multi-answer questions, requiring deeper reasoning across long contexts.
KnowRef (Emami al., 2019): A coreference resolution dataset where multiple entities mentioned, and models must resolve ambiguous pronouns using contextual cues.
A.9 DETAILS OF BENCHMARKS
We evaluate method eight widely used benchmarks covering general knowledge, commonsense reasoning, science QA, and mathematical problem-solving
MMLU (Hendrycks et al., 2021): A comprehensive benchmark of subjects spanning STEM, humanities, social sciences, and professional domains. It measures models' multitask accuracy and general world knowledge.
HellaSwag (Zellers et al., 2019): A commonsense benchmark requiring models select the most plausible continuation of a given context. It emphasizes grounded reasoning about everyday scenarios.
PIQA (Bisk et al., 2020): A physical commonsense reasoning dataset, where models must infer the correct solution to physical problems from everyday settings.
ARC-Challenge (Clark al., 2018): 9): benchmark of challenging grade-school science questions requiring reasoning, knowledge retrieval, and integration across multiple facts.
ARC-Easy (Clark et al., 2018): The easier subset of ARC focusing on factual recall and simpler reasoning in grade-school science.
WinoGrande (Sakaguchi et al., 2020): large-scale benchmark for pronoun resolution and commonsense reasoning, designed to be adversarially filtered and less susceptible dataset artifacts.
BoolQ (Clark al., 2019): A reading comprehension dataset of yes/no questions paired with passages from Wikipedia, requiring models to integrate text understanding with factual reasoning.
GSM8k (Cobbe et al., 2021): A dataset of grade-school math word problems requiring multi-step numerical reasoning. We treat GSM8k as an ontto--distriuttion (OOD) evaluation since math is not included in our training set.
17

---

A.10 MODEL DESCRIPTIONS
In this section, we provide additional details of the models reported in Table 2. These models cover broad range of active parameters (IB, 3B. 7-8B, 13-14B, 27-34B), including dense LLMs and sparse cesture-of-Expeeee (MoE) variants. All results are reported in the main text (see Table 2).
Models with 1B Active Parameters
Llama3.2-1B (Dubey al., 2024): A IB-parameter dense model Llama 3.2 family,
OLMo-IB (Groeneveld et al., 2024): Dense model from the AllenAI OLMo family.
OLMoE-7B-AIE Muennighoff al. (2024): A sparse MoE variant of OLMoE with 7B total parameters and ~B active per token.
Models with ~3B Active Parameters
Gemma2-3B (Team, 2024a): Google's Gemma2 family dense model.
Qwenl.5-14B-A3B (Bai et al., 2023): Sparse MoE variant Qwenl.5 with 14B total parameters and ~3B active per token.
DeepSeekMoEE-66-AAA (Dai al., 2024b): Sparse MoE model with 16B total parameters and ~3B active per token.
Models with ~7-BB Active Parameters
Qwen2-7B (Team, 2024b): Dense model from the Qwen2 family.
Mistral-7B (Jiang al., 2023): Dense model emphasizing efficient training and inference.
Llama3.1-8B (Dubey al., 2024): A dense model from the Llama 3.1 family.
Models with 13-114 Active Parameters
Llama2-13B (Touvron et al., 2023): Dense model from the Llama 2 release.
Vicuna-13B (Chiang et al., 2023): Instruction-tuner tundd LLM based on Llama2- -3BB.
Qwen1.5-14 1een1.5-14p 55114B (Bai et al., 2023): Dense version of Qwen1.5.
Models with ~-27-34B Active Parameters
Gemma2-27B (Team, 2024a): Largest Gemma2 dense variant.
Yi-34B (Young al., 2024): Dense model with 34B parameters.
Llama2-34B (Touvron et al., 2023): Large dense model from the Llama 2 family.
A.11 DETAILS OF ABLATION STUDY ON DEEPSHEKMOE
In this section, we provide a detailed ablation study of RoMA the DeepSeekMoE model.
Layer configurations (Table 7). When applying RoMA to different numbers of layers, find that using a single layer or two layers provides modest gains (e.g., Fl at 68.0%, L. at 70.6%) Increasing to five layers yields further improvements (up to 72.4%), while applying RoMA all layers gives 73.7%. Interestingly, restricting the method to the last five layers achieves the best result (74.7%), surpassing even the all-layer setting, which highlights the importance of critical-layer selection.
Token positions (Table 8). We next evaluate applying RoMA to the routing weights of different tokens. Regularizing the first or middle tokens shows only limited improvements (69.0% and 67. 67.6%, respectively), while the last token positions provide stronger performance. In particular, Lastl achieves the best result (74.7%), indicating that the most informative supervision signal for routing weights lies in the final tokens.
18

---

Table 7: DeepSeekMoE average accuracy (%) across different layer configurations
Group
Configuration
Avg. Acc. (%)
Base
Base
66.6
Fl
68.0
Layer
MI
67.4
LI
68.6
FIMI
69.0
FILI
70.2
MILI
69.6
Layers
F2
69.4
M2
68.8
L2
70.6
F2M3
70,8
F2L3
71.3
Layers
M2L3
71.0
F5
72.4
M5
71.6
All Layers
All16
73.7
1.5 (Ours)
74.7
Table 8: DeepSeekMoE average accuracy (%) when applying RoMA different token posi- tions.
Table 9;
DeepSeekMoE average accuracy (%)
with different neighbor selection strategies.
Group
Configuration
Avg. Acc. (%)
Group
Configuration
Avg. Acc. (%)
Baseline
Base Rand
66.6
Baseline
Base
66.6
Baseline
66.8
First3
69.0
E 0.3
68.0
3 Tokens
Middle3
67.6
e-neighbor
E 0.5
72.8
Last3
0.7
71.2
1 Token
Firstl Middlel
71.0
68.1
k-neighbor
k= (Ours)
74.7
Ours
Lastl (Ours)
74.7
73.4
Neighbor selection (Table 9). We compare e-neighbor and k-neighbor strategies. Small (0.3) gives minor improvement (68.0) (8). while moderate (0.5) achieves a strong 72.8%. For k-neighbors, increasing & improves accuracy up to k = 3 (74.7%), after which performance saturates (k 5 at 73.4%). This suggests that balanced neighbor selection (neither too sparse nor too dense) is crucial for generalization.
Training set size (Table 10). We investigate different proportions of training data used for regu- larization. Performance grows steadily with larger set sizes, from 67.5% at 10% to 73.6% at 70%. Using the full dataset (100%) achieves the best result (74.7%), confirming that more data consistently strengthens the alignment of routing weights.
Regularization methods (Table 11). We compare different regularization objectives. L1 and L2 losses improve the baseline to 67.2% and 70.3%, respectively. while entropy regularization achieves 69.8%. Our proposed manifold regularization significantly outperforms all alternatives, reaching 74.7%, which demonstrates the effectiveness of aligning routing weights with the manifold structure of successful neighbors.
19

---

Table 10: DeepScekMoE average accuracy (%) with different training set sizes.
Table 11: DeepSeekMoF average accuracy (%)
with different regularization methods.
Configuration
Avg. Acc. (%)
Configuration
Avg. Acc. (%)
Base
Base
66.6
10%
67.5
L1
67.2
30%
69.8
L2
70.3
50%
72.2
Entropy
69.8
70%
73.6
Manifold (Ours)
74.7
100% (Ours)
74.7
THE USE OF LARGE LANGUAGE MODELS (LLMS)
We used large language models (LLMs) solely for grammar checking and language polishing. No part of the research design, methodology, experimental results, core writing was generated by L.Mss. All scientific ideas, analyses, and conclusions are entirely the work of the authors.
20


Reviews:
[{"review_id": "6lgaY6Cweo", "reviewer": "Reviewer_suA3", "summary": "The authors propose a lightweight finetuning method for MoE models built upon the insight that there is typically misalignment between the clustering structure of the task embedding space and the routing weights, which undermines expert specialization as the semantic structure of the inputs is not preserved across the experts. The authors thereby propose a regularizer termed 'Routing Manifold Alignment', which trains the router to align the task embedding manifold, obtained by embedding task samples with a pretrained embedding model, with the routing weights manifold.", "questions": "1. See the note about unclear methodology. Clarification on preliminaries and notation would be helpful. \n\n2. The authors' key insight is that samples with similar task embeddings should share similar routing patterns. There's a conceptual distinction here though between a sample-level task embedding and the token-level routing patterns. Could the authors clarify or add some intuition here why these conceptually distinct objects can be unified?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The high level idea of manifold alignment is intuitive and original. \n\n2. The result that a substantive performance gain is attainable purely through better routing is compelling.\n\n3. The authors present their work in two frontier backbones, OLMoE and DeepSeekMoE, and extensively validate their work across a host of model sizes. The authors also do a good job of validating their method against a range of downstream tasks.", "weaknesses": "**Missing baselines**. Perhaps the largest issue is that the authors essentially present a parameter-efficient finetuning method for MoE but do no compare with existing MoE-PEFT / PEFT works. The authors do a good job of comparing with a wide range of lightweight baselines including prompt and prefix tuning up to Dense BP and C3PO, but, in my view, comparison with LoRA [1] or one of its new sota variants (perhaps DoRA [2]), and a PEFT-MoE method, such as MoLE [3] is vital to assessing the validity of the performance gains and efficiency. \n\n**Unclear methodology**. The authors refer extensively to 'routing weights' as $\\\\{r_i\\\\}_{i=1}^n$ where $r_i$ is the concatenated routing weights across multiple layers. But this sounds more like the authors are describing the (concatenated) *projections* of the $n$ samples through the router, not the weights themselves. Indeed, there are not $n$ different routing weights for each sample, but shared routing weights that project the $n$ samples in each layer. Furthermore, without any preliminaries on what kind of router architecture the authors are referencing, it becomes harder to work this out. To me, the method still makes sense if we are referring to routed projections, i.e hidden representations of the tokens after applying routing weights $h_i = W_r x_i$, where $W_r \\in \\mathbb R^{K \\times d}, x_i \\in \\mathbb R^d$ for $K$ experts, but it seems this is a misreference to refer to $n$ routing weights. Alternatively, it could be the authors mean the hidden representation after the entire MoE layer (i.e after both routing projection and expert feed forward), but, again, without preliminaries this is unknown. If the authors could clarify this that would be much appreciated.\n\n[1] LoRA: Low-Rank Adaptation of Large Language Models (Hu et al, 2021)\n\n[2] DoRA: Weight-Decomposed Low-Rank Adaptation (Liu et al, 2024)\n\n[3] Mixture of LoRA Experts (Wu, 2024)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g8bmlpl1L7", "reviewer": "Reviewer_kpLh", "summary": "This paper proposes Routing Manifold Alignment (RoMA), a lightweight post-training method for Mixture-of-Experts (MoE) large language models. RoMA addresses the suboptimal routing problem in MoE LLMs by aligning the routing weight manifold with the task embedding manifold via a novel manifold regularization term. The method fine-tunes only the router parameters (a tiny fraction of the total model) while keeping experts frozen, leading to significant accuracy improvements (7–15%) across diverse benchmarks without increasing inference cost.", "questions": "1. How did you compute the Inference cost\n2. It is good to show the routing weights in Table 1. \n3. For Figure 6-10, did you use the average accuracy for the 8 tasks in Table 1? If the answer is yes, how did you select these hyperparameters (k, layers, token position) in zero-shot tasks? \n4. For Figure 7, why can the last token obtain good performance?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. RoMA introduces a novel manifold alignment perspective to MoE routing, unifying task understanding (via embeddings) with expert selection.\n2. The paper extensively evaluates RoMA on two recent MoE architectures (OLMoE and DeepSeekMoE) across eight diverse benchmarks. It outperforms strong baselines, including C3PO, Dense BP, and tuning methods, and shows competitive or superior performance to dense models with up to 34B parameters.\n3. The authors systematically analyze key design choices—layer selection, token positions, neighborhood strategies, training set size, and regularization methods.", "weaknesses": "1. While inference cost is unchanged, the training cost of RoMA—especially the nearest-neighbor search over the training set—is not discussed. \n2. The paper uses a pre-trained embedding model \nE(⋅) to compute task embeddings but does not justify the choice or explore its sensitivity. The impact of different embedding models on RoMA’s performance remains unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D7GGiO8hhj", "reviewer": "Reviewer_yq6K", "summary": "The paper proposes RoMA (Routing Manifold Alignment), a post-training method for Mixture-of-Experts (MoE) LLMs that updates only the routers. The idea is to build a neighbor graph in a task-embedding space over a subset of “successful” training examples (those the current model gets right), then add a *manifold regularizer* that encourages a sample’s routing pattern ($r_i$) to resemble those of its successful neighbors ($r_j$), with similarity-normalized weights ($W_{i,j}$). The method adds no inference-time cost and is reported to improve accuracy across benchmarks on MoE backbones. The paper’s motivation leans on a claimed gap to an “oracle routing” upper bound and on UMAP visualizations of routing vs task-embedding manifolds.", "questions": "1. **Oracle:** How were the Oracle numbers in Table 1 produced? Please detail the optimization/relaxation, steps, etc.\n2. **Router gradients:** What is the exact differentiability strategy for top-(k) routing during FT (STE vs Gumbel vs soft routing)?\n3. **Embedding ($E(\\cdot)$):** Which model, inputs (instruction vs raw text), metric/σ, and sensitivity? Is the neighbor graph static or periodically recomputed?\n4. **Stats:** Please report means error bars over $\\geq3$ seeds and the tuning protocol for $\\lambda$, $k/\\epsilon$, $\\sigma$, LR.\n5. Providing the following would really benefit the paper's claims: quantitative evidence that routing space aligns with task-embedding space e.g. through subspace similarity (CCA/CKA) and neighbor consistency (k-NN overlap, trustworthiness) reported pre/post RoMA on at least two datasets to substantiate Figure 3.\n6. **Successful-neighbor filter:** Have you tried variants that include a small fraction of incorrect but semantically close neighbors or a curriculum that relaxes the filter over time?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "- Clear, practical objective: The manifold regularizer is standard, well-posed, and easy to implement for router-only fine-tuning.\n- **Empirical breadth:** The paper reports consistent gains over a wide suite (MMLU, ARC-C/E, HellaSwag, etc.) and shows competitive performance relative to a strong test-time routing method (C3PO).\n- **Focused ablations:** The paper varies the number of routed layers, token choice, and neighborhood size, which are some of the design levers that plausibly matter for this approach (though see “Weaknesses” on explanations/variance).\n- Writing is generally clear at the high level; figures are visually appealing. The method section is readable.", "weaknesses": "**Related works section is really sparse.** \nA huge body of related works are present in the literature that are missing and many works such as Moefication etc. are relevant. \n\n**Oracle” routing is defined but not computed, yet used as an empirical anchor.**\nThe paper defines a per-example oracle  $r_i^*=\\arg\\min_r \\mathcal L_{\\mathrm{CE}}(f(x_i,r),y_i)$ to claim a 10–20% “oracle gap,” and it reports Oracle rows in Table 1. However, no procedure is provided for obtaining these oracle numbers (search/relaxation/heuristic/optimization), nor is the compute or fairness (e.g., use of labels at evaluation) specified. As written, the “oracle” is kind of like a tautological bound, not really a reproducible baseline. Using it in tables materially affects the paper’s narrative but is methodologically unclear.\n\n**Heavy reliance on UMAP visuals without quantitative alignment metrics**\nFigure 3 is central to the paper’s motivation (misalignment before RoMA, alignment after, similarity to “oracle” etc.), but UMAP is a non-linear, hyperparameter-sensitive projection, and the paper provides *no quantitative* alignment measures nor any robustness checks across UMAP settings/datasets. As such, the figure is suggestive but insufficient as primary evidence.\n\n**Incomplete reproducibility for key components**\nImportant knobs/levers are underspecified in the main text:\n- **Embedding model ($E(\\cdot$)):** name/version, what text is embedded (task descriptions vs inputs), metric/σ and sensitivity.\n- **Neighbor graph schedule:** fixed vs recomputed during training; frequency/criteria.\n- **Router differentiability:** how gradients traverse the top-(k) gating (STE, Gumbel-Softmax, soft relaxation), and whether training uses hard or soft routing during FT.\n- **Hyper-parameters & tuning protocol:** $\\lambda$ for the regularizer, $\\sigma$, ($k/\\epsilon$ seems to be provided in the appendix)..\nWithout these, independent reproduction will be difficult and also makes the method incredibly opaque and the results hard to interpret. \n\n*Successful neighbors* may introduce *confirmation bias*\nThe method imitates routing only from samples the current model already solves correctly ((\\mathcal S)). This is a plausible heuristic, but it risks **locking in** existing modes and under-serving hard/rare cases; the paper does not probe this trade-off. Can the authors comment on this?\n\n**Statistical rigor and significance**\nMain tables lack error bars and multiple seeds. Router-only fine-tuning and MoE routing can be high-variance and single-run improvements, especially single-digit deltas, need **CIs**/**p-values** to rule out noise.\n\n**In general** the main proposal of this paper is a manifold regularizer for routing. My main issue with this paper is that, although the paper proposes this regularizer, it does not study it well. It adds it and gets \"good results\". But whether it is because of that is not well studied nor are the experimental details clearly expressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3Mr23s13fB", "reviewer": "Reviewer_8yaH", "summary": "This paper proposes a novel post-training method, termed Routing Manifold Alignment (RoMA), that fine-tunes routing parameters to address the misalignment between routing weights and task semantics in MoE models. The authors first analyze the mismatch between routing selection and tasks in existing MoE models. Based on this, RoMA aligns the manifold of task and routing. Sufficient experiments have demonstrated the effectiveness of the proposed method.", "questions": "1. Although it achieves good performance on a range of downstream tasks, does this adjustment to the expert distribution compromise the model's expressive power?\n\n2. The authors need to provide some theoretical analysis of the ablation study. The current analysis is limited to experimental phenomena, and the authors should provide some theoretical analysis of the causes.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed post-routing training idea of ​​explicitly aligning the \"task embedding manifold\" with the \"routing weight manifold\" is very novel.\n\n2. Experiments have fully demonstrated the effectiveness of the proposed method and achieved significant improvements on different datasets.", "weaknesses": "1. The proposed method requires an additional embedding model to provide annotations, and its performance may be limited by the additional model.\n\n2. Further quantitative analysis is needed to illustrate the impact of the proposed method. (detailed in questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)