One-sentence Summary

Researchers from SII, SJTU, and GAIR propose UMoE, a budget-preserving pipeline that realigns mixture-of-experts (MoE) expert pools to a target domain by pruning low-saliency experts, regrowing via perturbation, and then applying supervised fine-tuning, achieving consistent gains over direct SFT across two architectures and five domains—including $+3.4$+3.4 math accuracy and $+6.0$+6.0 on SWE-bench Verified—without extra parameters or inference cost.

Key Contributions

• UMoE is a budget-preserving pipeline that prunes low-saliency experts, regrows them via perturbation-based expansion, and then applies standard supervised fine-tuning, keeping expert count, parameter count, and inference cost unchanged.
• Across two MoE architectures, five domains, and 12 benchmarks, UMoE consistently outperforms direct SFT, with representative gains of 3.4 points in math accuracy and 6.0 points on SWE-bench Verified, using a single frozen recipe without per-domain tuning.
• Analysis reveals that direct SFT wastes capacity on low-saliency experts that can be removed with little loss, while UMoE converts that redundancy into useful domain expertise, achieving lower training loss and improvements across all difficulty levels.

Introduction

Frontier large language models increasingly adopt Mixture-of-Experts (MoE) architectures to scale capacity while keeping inference efficient, and domain-specific post-training of these models is crucial for pushing performance on hard problems or creating specialist models that can later be consolidated into generalists. However, the expert pool inherited from mixed-domain pre-training is not homogeneous: sparse routing encourages functional specialization, leaving some experts with low relevance to a target distribution yet still receiving non-trivial routing during standard supervised fine-tuning (SFT). This misalignment wastes fixed capacity because SFT updates experts in place without reallocating slots. The authors propose UMoE, a domain-aware reorganization method that first prunes the least domain-salient experts using target-domain calibration data, then regrows the pool via perturbation-based expansion to restore the original expert count before applying standard SFT. The resulting model preserves architecture, parameter count, and inference cost while making a larger fraction of the expert pool consequential for downstream performance.

Method

The authors propose UMoE, a method that reorganizes the expert pool of a Mixture-of-Experts model before domain-specific supervised fine-tuning, preserving the original architecture and inference budget. The core idea is to replace a generic pre-trained expert configuration with a domain-aligned initialization, so that subsequent fine-tuning starts from a more relevant set of parameters.

An MoE layer consists of $E$E routed experts $\{E_1, \dots, E_E\}${E1​,…,EE​} and a router. For a token representation $h$h, the router selects a top-$K$K subset $\mathcal{S}(h)$S(h), and the layer output is a gated sum:

$y = \sum_{i \in \mathcal{S}(h)} g_i(h) E_i(h),$y=∑i∈S(h)​gi​(h)Ei​(h),

where $g_i(h)$gi​(h) is the normalized gate weight. Both $E$E and $K$K are fixed by the pre-trained architecture, determining the parameter count and per-token inference cost. UMoE deliberately keeps these numbers unchanged, so the reorganized model has exactly the same deployment budget as the original.

Pre-training endows experts with specialized functions, but under a single target distribution their relevance becomes uneven. Some experts align well with the domain, while others remain poorly aligned despite receiving non-trivial routing. Standard fine-tuning updates the existing pool in place but does not explicitly reallocate expert capacity. UMoE addresses this by first pruning low-saliency experts and then expanding from the survivors, creating a domain-aligned initialization before applying the standard SFT loss.

The method proceeds in three steps. First, a calibration subset $\mathcal{C} \subset \mathcal{D}$C⊂D is used to score every expert per layer, and the least domain-salient half are dropped. Second, the pruned pool is expanded back to the original size $E$E by regenerating experts from the survivors. Third, the reorganized model undergoes domain SFT with a cross-entropy objective. The final model is architecturally identical to the original, differing only in its expert initialization.

For pruning, the authors adopt the REAP saliency measure, which couples an expert’s activation strength with its router gate. For layer $l$l and expert $e$e, let $\mathcal{T}_{l,e}$Tl,e​ be the set of calibration tokens for which the expert is selected, $f_{l,e}(h_t)$fl,e​(ht​) its MLP output, and $g_{l,e}(h_t)$gl,e​(ht​) its renormalized gate weight. The saliency is defined as the conditional mean:

$S(l, e) = \frac{1}{|\mathcal{T}_{l,e}|} \sum_{t \in \mathcal{T}_{l,e}} g_{l,e}(h_t) \left\| f_{l,e}(h_t) \right\|_2.$S(l,e)=∣Tl,e​∣1​∑t∈Tl,e​​gl,e​(ht​)∥fl,e​(ht​)∥2​.

If an expert is never selected, its score is set to zero. This metric separates per-activation saliency from routing frequency, keeping experts that produce large outputs under large gate weights when they are active. Ablations show that REAP and activation frequency are competitive, while unweighted Expert Activation Norm performs worse in this setting.

After pruning, each layer retains $E/2$E/2 domain-salient experts. Expansion regrows the pool to $E$E experts, preserving parameter count, active experts $K$K, and inference cost. Each freed slot is initialized from a survivor assigned in round-robin order, so under a 50% pruning ratio each survivor seeds exactly one new expert. The MLP weights and router row of the survivor are copied into the freed slot, and then both the parent and the regrown expert are perturbed independently. For any parameter tensor $W$W involved in a parent–regrown pair, the update is:

$W_{\text{parent}}' = W + \sigma \, s(W) \, \epsilon_p, \quad W_{\text{regrown}}' = W + \sigma \, s(W) \, \epsilon_r,$Wparent′​=W+σs(W)ϵp​,Wregrown′​=W+σs(W)ϵr​,

where $\epsilon_p, \epsilon_r \overset{\text{i.i.d.}}{\sim} \mathcal{N}(0, \mathbf{I})$ϵp​,ϵr​∼i.i.d.N(0,I), and $s(W)$s(W) is the empirical standard deviation of the corresponding parameter group (e.g., an individual MLP projection, the stored expert tensor, or the full router matrix). The noise scale is set to $\sigma = 0.05$σ=0.05, and noise is sampled independently across groups and between the two sides. For architectures with an expert-score correction bias, the regrown expert inherits the parent’s bias.

Two design principles guide the perturbation strategy. First, both sides of the pair are perturbed, not just the regrown expert. One-sided perturbation creates an initialization asymmetry that causes routing to the perturbed half to decline during SFT and harms downstream performance; symmetric perturbation removes the privileged unperturbed anchor. Second, both the router row and the MLP weights are perturbed. Router perturbations introduce routing asymmetry, while MLP perturbations create functional asymmetry. Either component alone performs well, and their combination yields the best average result.

Experiment

UMoE is evaluated on two MoE architectures across five diverse domains and multiple SFT data scales, consistently outperforming direct fine-tuning. Ablations confirm that REAP-based pruning and symmetric perturbation of both router and MLP weights are crucial, while mechanism analysis reveals that the baseline suffers from redundant expert usage and UMoE’s gains stem from making the regrown experts collectively consequential. Behavioral analysis shows improvements are spread across problem difficulty, with UMoE often correcting critical reasoning errors.

The SFT datasets cover five domains with notable differences in scale and sequence length. Agentic data has the longest average sequences by far, while Science provides the most samples but with moderate length. Total token counts span from 0.19B for Tool-use to 2.6B for Code, reflecting the interplay of sample count and average length. Agentic domain (OpenSWE) has the highest average sequence length at over 50,000 tokens, despite having the smallest sample count. Science domain uses the largest dataset with nearly 700,000 samples, yet its total token count is lower than Code and Math due to shorter average length. Code and Math each contain 300,000 samples, but Code's longer average sequences give it the highest total token count among all domains. Tool-use dataset is relatively compact with around 119,000 samples and an average length of 1,600 tokens, resulting in the smallest total token count.

UMoE consistently improves math benchmark accuracy over Direct SFT across two MoE architectures. Gains are larger when starting from a weaker baseline, with an average lift of +3.40 on Qwen3-30B-A3B, and persist even on a strong in-house trained baseline that already surpasses a comparable thinking model, where UMoE adds +1.36 on average. The expert-pool reorganization is effective regardless of data scale and baseline strength. On a 300k subset of OpenMathReasoning, UMoE improved average accuracy over Direct SFT on all six math benchmarks for both Qwen3-30B-A3B and Qwen3.5-35B-A3B, with the largest single-benchmark gain of +6.56 on AIME 2024. When trained on a high-quality in-house math corpus, the Direct SFT baseline already outperformed Qwen3-30B-A3B-Thinking, yet UMoE still lifted five of six benchmarks and raised the average by +1.36, peaking at +3.03 on HMMT Feb 2026.

UMoE yields consistent improvements over Direct SFT on Qwen3-30B-A3B across all six cross-domain benchmarks, covering code, science, tool-use, and agentic coding. The gains are broad, with an unweighted average increase of 2.55 points, and the largest single boost is a 6.0-point rise in pass@1 on SWE-bench Verified. This demonstrates that expert-pool reorganization transfers effectively beyond math to diverse reasoning and agentic tasks. UMoE outperforms Direct SFT on every evaluated benchmark, including LiveCodeBench, GPQA-Diamond, τ-bench, τ²-bench, and SWE-bench Verified. Agentic coding sees the largest absolute improvement, with SWE-bench Verified pass@1 increasing from 16.2% to 22.2% (+6.0 points). Science reasoning on GPQA-Diamond improves by 3.10 points, while code benchmarks show gains of 2.73 and 0.85 points on two LiveCodeBench versions. Tool-use tasks on τ-bench and τ²-bench rise by 1.20 and 1.40 points, respectively, confirming gains extend to interactive agent settings.

Ablating the two steps of the UMoE method shows that REAP pruning provides the largest boost over direct SFT, with frequency-based pruning close behind. For expansion, symmetric noise applied to both router and MLP weights gives the best results, whereas restricting noise to only the regrown side causes a substantial drop below the baseline. REAP pruning (router-weighted activation norm) improves over direct SFT by 5.16 points, slightly outperforming frequency-based pruning (+4.78). Symmetric noise on both router and MLP yields the highest average score (43.01), while one-sided perturbation drops performance below the direct SFT baseline.

Pruning the least-salient half of experts from Direct SFT barely affects its math benchmark average (a drop of only 0.12 points), showing that those experts were redundant despite receiving a large share of routing. In contrast, pruning UMoE costs 5.00 points and collapses its lead over Direct SFT from +5.16 to +0.28, demonstrating that UMoE's gains depend on making the broader expert pool collectively useful. After removing the bottom-half experts, Direct SFT's average score changes by only -0.12, revealing that nearly 40% of its routed capacity was redundant. Pruning UMoE costs 5.00 points and shrinks its lead over Direct SFT from +5.16 to +0.28, indicating its advantage relies on the additional half of the expert pool.

The evaluation spans five SFT domains with diverse data scales and sequence lengths, and UMoE is tested on math, code, science, tool-use, and agentic coding benchmarks to validate its expert-pool reorganization method. Across both weak and strong baselines, UMoE consistently improves accuracy, with the largest gains on math and agentic tasks, and ablation studies show that REAP pruning and symmetric noise expansion are crucial while the method's advantage depends on the entire expanded expert pool being collectively useful. Overall, UMoE turns previously redundant experts into a broadly beneficial resource, yielding robust improvements across reasoning and interactive agent scenarios.