💬 NLP

UMoE:Unlocking Every Expert in Domain-Specific Training

UMoE is a budget-preserving pipeline that enhances domain-specific training of Mixture-of-Experts models by pruning low-saliency experts and regrowing the pool via perturbation before fine-tuning, thereby consistently outperforming standard supervised fine-tuning across various architectures and domains without increasing inference costs.

Original authors: Xuefeng Li, Pengfei Liu

Published 2026-07-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xuefeng Li, Pengfei Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, super-smart robot brain called a "Mixture-of-Experts" (MoE). Think of this brain not as a single giant muscle, but as a team of 128 or 256 tiny specialists, each an expert in something different. When the robot gets a question, a smart "router" decides which few specialists (say, 8 of them) should chime in to help solve it.

Here's the problem: This team was trained on everything—math, coding, science, cooking, you name it. But now, you want to train this robot to be a math wizard. The paper argues that if you just start teaching the whole team math (a process called "Direct SFT"), you're wasting time. Why? Because some of those 256 specialists are terrible at math. They might be great at poetry or history, but they're just "noise" in a math class. Yet, the router keeps accidentally calling on them anyway, and the robot tries to force them to learn math, which is inefficient.

The Big Idea: UMoE (Unlocking MoE Capacity)
The authors, Xuefeng Li and Pengfei Liu, propose a clever trick called UMoE. Instead of just teaching the whole team, they do a quick "audition" first.

The Prune (The Cut): They take a small sample of math problems and ask, "Who is actually good at this?" They find the bottom 50% of the specialists who are the least helpful for math. They kick them out of the room.
The Regrow (The Clone): Now the team is too small! The robot needs its full budget of 256 experts to keep its speed and cost the same. So, they take the remaining math-savvy experts and make clones of them. But here's the twist: they don't just copy-paste. They give the original and the clone a tiny, random "shake" (a mathematical perturbation). This makes the clone slightly different from the original, so they can both learn and specialize in math without stepping on each other's toes.
The SFT (The Training): Now, with a room full of math-ready (or math-ready-to-be) experts, they teach the whole team math.

The Results: Does it work?
The paper shows that this "audition and regrow" method works surprisingly well, and the authors are pretty sure about it because they tested it on real benchmarks, not just simulations.

Math: On a set of hard math competitions, UMoE improved the average score by 3.4 points on one model and 1.67 points on a newer, bigger model. Even when they used a super-strong, high-quality math dataset where the standard method was already beating other famous models, UMoE still squeezed out an extra 1.36 points.
Coding & Science: It wasn't just for math. When they tried it on coding, science, and even "agentic" tasks (where the AI uses tools), it consistently won. For example, on a tough coding benchmark called SWE-bench Verified, the score jumped by 6.0 points (from 16.2% to 22.2%).
Data Size: The paper checked if this only works with a little data. They tested with data sizes ranging from 0.5 billion to 4.1 billion tokens. In every single case, UMoE was better.

What the Paper Says is NOT the Answer
The authors are very clear about what doesn't work or isn't the main point:

Just adding more data isn't enough: They showed that even with huge amounts of data, the standard method (Direct SFT) still leaves a lot of "useless" experts in the mix.
One-sided shaking doesn't work: When they tried to only "shake" the new clones and leave the original experts untouched, the performance actually got worse than the standard method. The paper suggests that you must shake both the original and the clone to keep them balanced.
Simple counting isn't best: They tested different ways to decide who to kick out. Just counting how often an expert is used (frequency) was okay, but a more complex score called REAP (which looks at how strong the expert's answer is and how often it's used) was the best.

Why It Works (The "Aha!" Moment)
The paper dug deep to see why this works. They found that in the standard method, the robot wastes about 42% of its computing power on experts that are actually quite useless for the task. If you took a standard trained model and manually cut out the bottom half of experts after training, the score barely dropped (only 0.12 points). This proves the standard method was carrying dead weight.

But with UMoE? If you cut out the bottom half of their trained team, the score crashed by 5.0 points. This suggests that UMoE successfully turned that "dead weight" into useful, math-solving power.

A Tiny Example
The paper gives a specific math problem (AIME 2026, Problem 25) to show the difference.

Standard Model: It tried to simplify a fraction (200/225) and got it wrong (saying it was 4/5), which led to a wrong final answer of 405.
UMoE: It simplified the fraction correctly (to 8/9) and got the right answer of 850.

The Bottom Line
The paper suggests that by reorganizing the team before the heavy training starts—kicking out the wrong specialists and cloning the right ones—you can make a model smarter without spending more money or time. It's like realizing you don't need to teach your whole soccer team how to play chess; you just need to swap in the chess players and let them practice together. The authors measured this across 12 different benchmarks and found it consistently helpful, suggesting this is a solid way to make AI experts even better at their specific jobs.

Technical Summary: UMoE (Unlocking Every Expert in Domain-Specific Training)

1. Problem Statement

Mixture-of-Experts (MoE) models have become a standard architecture for frontier Large Language Models (LLMs) due to their ability to scale capacity without proportional increases in inference cost. However, a critical inefficiency arises during domain-specific post-training (e.g., Supervised Fine-Tuning or SFT).

Pre-trained MoE models inherit an expert pool shaped by mixed-domain pre-training. Under a specific target domain, this pool is heterogeneous: a substantial subset of experts exhibits low "domain-aligned saliency," meaning they contribute little to the target task despite receiving non-trivial routing. Standard SFT updates these inherited experts in place but does not reallocate expert slots. Consequently, a portion of the fixed model capacity remains poorly aligned with the target domain, limiting the effectiveness of the fine-tuning process.

The core question addressed is: Can the expert pool be reorganized into a more domain-aligned initialization before SFT, such that more expert slots become useful for learning target-domain capabilities, without altering the model's architecture or inference budget?

2. Methodology: UMoE

The authors propose UMoE (Unlocking MoE Capacity), a budget-preserving pipeline that reorganizes the expert pool prior to standard SFT. The method operates in three distinct steps and preserves the original expert count ( $E$ ), active expert count ( $K$ ), parameter count, and inference cost.

Step 1: Domain-Aligned Expert Pruning

The model first identifies and removes the least relevant experts for the target domain.

Calibration: A subset of the target domain data ( $C \subset D$ ) is used to score experts.
Saliency Metric: The authors utilize REAP (Lasby et al., 2026) saliency, which couples an expert's activation strength with its router gate weight. For an expert $e$ at layer $l$ , the saliency $S(l, e)$ is the conditional mean of the gate-weighted output norm over tokens where the expert is selected.
Action: The bottom 50% of experts (per layer) based on this saliency score are pruned.

Step 2: Budget-Preserving Expansion

To restore the original model size and maintain the fixed inference budget, the pruned slots are regrown from the remaining "survivor" experts.

Initialization: Each freed slot is initialized by copying the weights (MLP and router row) of a retained "parent" expert in a round-robin fashion.
Perturbation: To break symmetry and encourage specialization, both the parent and the regrown expert are independently perturbed. For a parameter tensor $W$ , the perturbed weights $W'$ are calculated as:
$W' = W + \sigma s(W) \epsilon$
where $\epsilon \sim \mathcal{N}(0, I)$ , $s(W)$ is the empirical standard deviation of the parameter group, and $\sigma = 0.05$ .
Design Principle: Crucially, the perturbation is applied symmetrically to both the parent and the regrown expert. Ablation studies show that one-sided perturbation (only on the regrown expert) leads to routing asymmetry and performance degradation. Both the router row and MLP weights are perturbed to provide complementary routing and functional asymmetry.

Step 3: Domain SFT

The reorganized model, now initialized with a domain-aligned expert pool, undergoes standard full-parameter Supervised Fine-Tuning on the target domain data. No per-domain hyperparameter tuning is required; the same recipe is applied across all domains.

3. Key Contributions

Expert Pool Reorganization: A novel method to realign the expert composition of a pre-trained MoE to a specific target domain before fine-tuning, converting "redundant" capacity into "domain" capacity.
Budget Preservation: The method strictly maintains the original architecture, parameter count, and inference cost, making it a drop-in replacement for standard SFT pipelines.
Symmetric Perturbation Strategy: The identification that symmetric perturbation of both parent and regrown experts is critical for maintaining balanced routing dynamics and preventing the collapse of the regrown half's utility.
Empirical Validation: Comprehensive evaluation across two MoE architectures (Qwen3-30B-A3B and Qwen3.5-35B-A3B), five diverse domains (math, code, science, tool-use, agentic coding), and 12 benchmarks.

4. Experimental Results

UMoE consistently outperforms Direct SFT (fine-tuning the pre-trained model without reorganization) across all evaluated settings.

Math Domain:
- On Qwen3-30B-A3B (OpenMathReasoning data), UMoE improved the average accuracy across six competition benchmarks by +3.40 points, with a peak improvement of +6.56 on AIME 2024.
- On a strong in-house math corpus where the Direct SFT baseline already surpassed the Qwen3-30B-A3B-Thinking model (82.81 vs. 81.06), UMoE further raised the average to 84.17 (+1.36 points), demonstrating robustness even against strong baselines.
- On Qwen3.5-35B-A3B, UMoE improved the average by +1.67 points.
Cross-Domain Performance:
- Code: +2.73 / +0.85 points on LiveCodeBench v5/v6.
- Science: +3.10 points on GPQA-Diamond.
- Tool-use: +1.20 / +1.40 points on $\tau$ -bench / $\tau^2$ -bench.
- Agentic Coding: A significant +6.0 points improvement on SWE-bench Verified (pass@1 rising from 16.2% to 22.2%).
Data Scaling: Experiments varying SFT data size (from 0.5B to 4.1B tokens) showed that the performance gains of UMoE persist and remain positive across all data scales.

5. Mechanism Analysis

The paper provides three key insights into why UMoE works:

Redundancy in Direct SFT: Analysis of the Direct SFT model reveals that it allocates a substantial portion of routed-expert compute (approx. 42%) to a low-saliency subset of experts. Post-hoc pruning of this subset in Direct SFT models results in negligible performance degradation (−0.12 points), confirming this capacity is functionally redundant.
Consequential Capacity in UMoE: In contrast, pruning the trained UMoE model to 64 experts causes a significant performance drop (−5.00 points). This indicates that UMoE successfully converts the previously redundant half of the expert pool into capacity that is consequential to downstream performance.
Optimization Dynamics: UMoE starts with a slightly higher training loss than Direct SFT due to the reorganization but converges to a lower final training loss (0.396 vs. 0.404), suggesting a more efficient optimization trajectory.
Error Correction: Case studies (e.g., AIME 2026 Problem 25) show UMoE correcting specific reasoning errors (e.g., fraction reduction in the Angle Bisector Theorem) that Direct SFT propagates, leading to correct final answers.

6. Significance and Claims

The paper claims that UMoE offers a simple, budget-preserving, and universally applicable approach to domain adaptation for MoE models. Its significance lies in:

Efficiency: It unlocks latent capacity within existing models without requiring additional parameters or inference costs.
Generalizability: The method works across different model architectures and diverse domains (reasoning, coding, agentic tasks) using a single frozen recipe.
Robustness: It provides gains even when the baseline SFT is already strong and when training data scales up.

The authors conclude that reorganizing inherited expert capacity allows models to better align their fixed compute budget with the specific requirements of a target domain, transforming "dead" experts into active contributors to domain-specific performance.