Muon Makes the First Mixture-of-Experts Layer Harder to Balance

Marcos Hernanz

cs.LGOct 27, 2026

Muon Makes the First Mixture-of-Experts Layer Harder to Balance

Authors: Marcos Hernanz

Organizations: Escuela Politécnica Superior, Universidad Autónoma de Madrid

Abstract

Many open mixture-of-experts (MoE) language models train with Muon instead of AdamW. Most balance the load with a bias on each expert's routing score, adjusted after every step. Their recipes disagree on whether Muon should also train the router. In models with 55M and 108M active parameters, we measure how the optimizer changes the span of these biases, the largest minus the smallest in a layer. The span grows when all tokens favor the same experts, so it shows how hard a layer is to balance. At learning rates tuned on shorter runs, Muon raised the span of the first MoE layer at every seed, roughly doubling it on average, while the other layers changed much less. The rise came from Muon on the rest of the network, whichever optimizer trained the router. In those runs, the inputs to the first router shared a direction, and soon after the warmup its busiest expert averaged four to seven times its fair share. Over the first 3,000 steps, subtracting the batch mean from these inputs removed this overload while barely changing the loss. On the smaller model, halving Muon's learning rate removed most of the rise and slightly improved the loss. Moving only the router to AdamW made the loads shift more during the warmup and slightly raised the loss, so at this scale the results favor keeping the router on Muon.

Figures & tables

Figure 1: The span, the largest minus the smallest bias of a layer, under AdamW and Muon. Left, the mean span of each MoE layer in the main comparison (small model, 1B tokens, seeds 0 to 2), with its range over the seeds. Center, the span of the first MoE layer and the mean span of the other six layers in the same runs, colored as in the left panel, in bins of 100 steps, with the range over the seeds shaded for the first MoE layer. Right, the span averaged over the layers in the tuning runs (0.5B tokens, seed 0) against the peak learning rate, with filled marks at the rates selected by the tuning. Where marks overlap, at

2.5\times 10^{-4}

and

2\times 10^{-3}

, the spans differ by 0.011 or less. The control is Muon without the orthogonalization ( Section 4.2 ).

Model	Router	Balancer	about the router or the balance
			What the report says
DeepSeek-V4	Muon	bias, balance loss	loss spikes tied to outliers in the MoE layers, which routing appears to exacerbate
Kimi K3	Muon	bias set from the scores	a fixed bias step adapts slowly or oscillates
GLM-5.3	Muon	bias, balance loss	cites GLM-5, which follows the training setting of GLM-4.5, whose Muon excludes only embeddings, bias vectors, and normalization gains
Qwen3.8-Flash-Next	AdamW	balance loss	Muon made the router unstable early, with no significant gain later
MiMo-V2.6	AdamW	bias	frozen in reinforcement learning after it unbalanced the loads

Table 1: The optimizer of the router and the balancer in recent open MoE models. In the router column, Muon means that the report, or the earlier report it cites, does not exclude the router from Muon. The balancer is stated in the report, in the earlier report it cites, or in the released configuration [ 72 , 49 , 69 ] .

Partition	Groups trained with Muon	Experiments
AdamW	none	all
Muon	all five	all
Muon, router on AdamW	all but the router	stress test, ablation, balance loss, large model
AdamW, router on Muon	the router	ablation
Experts on Muon	the routed experts	ablation
Attention on Muon	the attention	ablation

Table 2: The partitions and the control. The five groups are the routers, the routed experts, the shared experts, the attention, and the dense FFN of block 0. A partition trains with AdamW every group that it does not put on Muon, and the embedding table, the output head, and the normalization gains always with AdamW.

Metric	AdamW	Muon	Difference
Validation loss	3.255	3.177	$-0.078$ [ $-0.086$ , $-0.062$ ]
Span, all layers	0.105	0.137	$+0.033$ [ $+0.030$ , $+0.036$ ]
Span, first MoE layer	0.084	0.194	$+0.110$ [ $+0.085$ , $+0.151$ ]
Update–load correlation	0.379	0.024	$-0.355$ [ $-0.362$ , $-0.345$ ]
Update–load correlation, spectral norm	0.102	0.035	$-0.068$ [ $-0.073$ , $-0.061$ ]
Load churn	0.032	0.038	$+0.006$ [ $+0.004$ , $+0.008$ ]

Table 3: The main comparison, on the small model with 1B tokens and seeds 0 to 2. Every metric is the integrated value averaged over the MoE layers unless the row names another layer or window. The validation loss is that of the last step and the specialization is measured at the last checkpoint. Each difference is Muon minus AdamW, followed by the range of the three paired differences.

			Input	Score	Early	Common	Span,	Warmup
Arm	Loss	Span	cosine	offset	overload	share	all layers	churn
AdamW	3.253	0.095	0.071	0.095	1.14	0.129	0.103	0.102
AdamW, router on Muon	3.248	0.099	0.074	0.100	1.16	0.151	0.117	0.092
Muon, router on AdamW	3.177	0.175	0.391	0.137	6.63	0.497	0.123	0.122
Muon	3.175	0.198	0.408	0.140	4.06	0.160	0.140	0.084
Muon, half rate	3.166	0.119	0.401	0.100	1.73	0.161	0.115	0.074

Table 4: The arms of the ablation with the bias, on the small model with 1B tokens, as means over seeds 3 to 5. The span, the input cosine, the score offset, the early overload (steps 501 to 1,262), and the common share are measured in the first MoE layer, and the last two columns over all MoE layers. The loss is the validation loss at the last step, and the other metrics are integrated values unless the column names a window. The arm “Muon, half rate” trains with Muon at half its tuned rate, which is the rate of AdamW.

Main comparison, small model, seeds 0 to 2
Arm $-$ AdamW	Span	Early overload	Input cosine	Loss
Muon	$+0.110$ [ $+0.085$ , $+0.151$ ]	$+3.15$ [ $+1.47$ , $+4.58$ ]	–	$-0.078$ [ $-0.086$ , $-0.062$ ]
Ablation, small model, seeds 3 to 5
Muon	$+0.103$ [ $+0.045$ , $+0.174$ ]	$+2.93$ [ $+2.17$ , $+3.52$ ]	$+0.338$ [ $+0.324$ , $+0.345$ ]	$-0.077$ [ $-0.083$ , $-0.067$ ]
Muon, router on AdamW	$+0.080$ [ $+0.065$ , $+0.091$ ]	$+5.50$ [ $+5.06$ , $+5.75$ ]	$+0.321$ [ $+0.313$ , $+0.335$ ]	$-0.075$ [ $-0.081$ , $-0.066$ ]
Muon, half rate	$+0.024$ [ $+0.010$ , $+0.037$ ]	$+0.59$ [ $+0.48$ , $+0.78$ ]	$+0.330$ [ $+0.321$ , $+0.342$ ]	$-0.087$ [ $-0.095$ , $-0.075$ ]

Table 5: The excess over AdamW across experiments. Each entry is the paired difference of an arm from AdamW with the same balancer, followed by its range over the seeds. The span, the early overload, and the input cosine are measured in the first MoE layer, the span and the input cosine as integrated values and the early overload over steps 501 to 1,262 (501 to 2,788 on the large model). The loss is the validation loss at the last step. The main comparison did not record the input cosine. The half rate is half the tuned rate of Muon, and the fast bias moves by

4u

per step against

u

for the usual bias.

Figure 2: The input cosine, the score offset, the overload, and the span of the first MoE layer over the first 3,000 steps of the ablation on the small model, in bins of 30 steps, under AdamW, Muon, and Muon with the router on AdamW, and in the probe that repeats the runs of Muon with the batch mean subtracted from the input of the first router. The input cosine of the probe is measured before the subtraction. Lines are the means over seeds 3 to 5 and bands their range. The thin gray vertical line marks the end of the warmup.

Metric	Small model	balance loss	Large model
		Small model,
Warmup churn	$+0.038$ [ $+0.038$ , $+0.039$ ]	$+0.045$ [ $+0.043$ , $+0.046$ ]	$+0.038$ [ $+0.037$ , $+0.039$ ]
Early overload	$+2.57$ [ $+1.96$ , $+3.52$ ]	$+3.44$ [ $+2.96$ , $+3.93$ ]	$+0.91$ [ $-0.78$ , $+1.94$ ]
Common share	$+0.337$ [ $+0.328$ , $+0.342$ ]	$+0.339$ [ $+0.334$ , $+0.346$ ]	$+0.377$ [ $+0.339$ , $+0.429$ ]
Span	$-0.023$ [ $-0.092$ , $+0.020$ ]	$-0.025$ [ $-0.048$ , $-0.005$ ]	$-0.027$ [ $-0.119$ , $+0.058$ ]
Loss	$+0.002$ [ $+0.001$ , $+0.004$ ]	$+0.001$ [ $+0.001$ , $+0.002$ ]	$+0.004$ [ $+0.003$ , $+0.004$ ]

Table 6: The router on AdamW minus the router on Muon, both with the backbone on Muon, as paired differences with their range over the seeds, which are 3 to 5 on the small model with 1B tokens and 0 to 2 on the large model with 3B tokens. The early overload, the common share, and the span are measured in the first MoE layer, the early overload over steps 501 to 1,262 (501 to 2,788 on the large model) and the other two as integrated values. The warmup churn is averaged over all MoE layers, and the loss is the validation loss at the last step.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

	Small model	Large model
Blocks (dense, MoE)	8 (1, 7)	12 (1, 11)
Width $d$ , attention heads, head width	512, 8, 64	768, 12, 64
Hidden width of the dense FFN	2,048	3,072
Routed experts $E$ , selected $k$ , shared experts	32, 4, 1	32, 4, 1
Hidden width of an expert	768	512
Vocabulary	16,384	16,384

Appendix

Table 7: The two models. An expert, routed or shared, is a SwiGLU FFN with three matrices. The active parameters are those one token uses, which in every MoE layer are the router, the shared expert, and

k

routed experts, and they exclude the embedding table and the output head.

Setting	Value
Muon	Nesterov momentum $\mu=0.95$ , five Newton–Schulz steps, factor $0.2\sqrt{\max(p,q)}$ , $\lambda=0.1$
AdamW	$\beta_{1}=0.9$ , $\beta_{2}=0.95$ , $\epsilon=10^{-8}$ , $\lambda=0.1$ on every weight matrix, including the embedding table and the output head, and none on the gains
Control	the Nesterov momentum of Muon rescaled to the Frobenius norm $\sqrt{\min(p,q)}$ , then the factor and $\lambda$ of Muon
Tuning grids	AdamW from $2.5\times 10^{-4}$ to $4\times 10^{-3}$ (five rates), Muon from $2\times 10^{-3}$ to $3.2\times 10^{-2}$ (five), the control from $2.5\times 10^{-4}$ to $8\times 10^{-3}$ (six), each rate twice the one before
Rates, small model	AdamW $2\times 10^{-3}$ , Muon $4\times 10^{-3}$ , the control $5\times 10^{-4}$ (tuning only)
Rates, large model	tuned with 1B tokens, AdamW $10^{-3}$ and Muon $2\times 10^{-3}$ , and used with 3B tokens, AdamW $7\times 10^{-4}$ and Muon $1.4\times 10^{-3}$

Appendix

Table 8: The settings of the optimizers, the learning rates, and the balancers. The rates are peak rates.

Experiment	Model	Tokens	Arms and seeds	Runs
Tuning	small	0.5B	AdamW and Muon at five rates, the control at six, AdamW with $\beta_{2}=0.98$ , seed 0	17
Main comparison	small	1B	AdamW and Muon, seeds 0 to 2	6
Stress test	small	0.5B	three recipe partitions at five multiples and Muon with the fast bias at multiple 4, seeds 0 to 2, and three calibration runs	51
Ablation	small	1B	eight partitions, Muon at half rate, and AdamW and Muon without the bias, seeds 3 to 5	33
Bias speed	small	1B	AdamW and Muon with the usual and the fast bias, seeds 6 to 8	12
Balance loss	small	1B	AdamW, Muon, and Muon with the router on AdamW, all with the balance loss, seeds 3 to 5	9

Appendix

Table 9: The runs of every experiment. The tuning trains one run per rate and a run of AdamW with

\beta_{2}=0.98

, and the stress test adds a calibration run of each partition. The 145 runs above the rule do not include the 30 exploratory runs.

Figure 3: The stress test, on the small model with 0.5B tokens, against the multiple of the tuned rates. The panels show the final validation loss, the warmup churn, and the largest absolute entries of the residual stream and of the router logits over the run. Lines join the means over seeds 0 to 2, and in the last three panels the small marks are the seeds. The losses of the two Muon partitions coincide.

Figure 4: The span, the input cosine, and the score offset of the first MoE layer in the arms of the ablation with the bias, on the small model with 1B tokens. The bars are the means over seeds 3 to 5 and the marks are the seeds. Dashed lines separate the four combinations of the optimizers of the router and the backbone, Muon at half rate, and the four partitions with part of the backbone on Muon. The arms with the whole backbone on Muon are in blue.

Figure 5: The span of the first MoE layer against its score offset (left, with the least-squares line of the logarithms) and against its input cosine (right), with one mark per run of the ablation with the bias. The span and the score offset are on logarithmic axes. Squares are the runs with the whole backbone on Muon, triangles those with only the routed experts on Muon, diamonds those with only the attention and the dense FFN on Muon, and circles the other four arms.

Mixture-of-Experts Language Models

Mixture-of-Experts language models use input-dependent routing to activate a small subset of expert networks, increasing capacity without proportional compute.

See all 229 papers

Oct 2, 2026

ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training

Mixture-of-experts (MoE) language models send each token to a few experts, so each expert is trained on a different part of the data, and this part changes during training. With a shared learning rate, Muon applies updates of roughly the same size to expert matrices of the same shape, even when an expert's update is poorly aligned with its current gradient. We here propose ExpertMuon-Compass (Compass), which multiplies the Muon step of each expert by two factors. A family factor compares the cosine between the expert's orthogonalized update and its gradient with the same cosine for the other experts in its layer. A scalar radius aggregates the alignment between corresponding rows of the update and gradient into one step-length multiplier. Compass keeps the update direction and the momentum buffer of Muon. In pretraining on FineWeb-Edu, Compass with Nesterov momentum on all matrices performs as well as or better than Muon, NorMuon, and other optimizers, with weight decay matched to NorMuon in the longer runs. Adding its factors to NorMuon gives the same or a lower loss than NorMuon. Compass is the most effective when the data seen by each expert varies during training, for example, as when the languages of a multilingual corpus arrive in separate blocks. With Compass, the expert load stays balanced, and the router assigns tokens to experts more decisively. We also prove a perturbation bound for the two factors.

Omatharv Bharat Vaidya, Ashwin Vinod, Pedram Akbarian +3

University of Texas at Austin

Deep Learning Optimization Mixture-of-Experts Language Models

Ask PDF

Apr 5, 2026

Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training

We model Mixture-of-Experts (MoE) token routing as a congestion game with a single effective parameter, the congestion coefficient gamma_eff, that quantifies the balance-quality tradeoff. Tracking gamma_eff across training checkpoints of two open-source MoE models, OLMoE-1B-7B (20 checkpoints, with dense sampling in the surge region) and OpenMoE-8B (6 checkpoints), reveals a three-phase trajectory: a surge phase where the router learns to balance load (gamma_eff: 14 to 36-39, peaking in the step 30K-40K region), a stabilization phase where experts specialize under steady balance (B_0: 2.4 to 2.3, steps 100K-400K), and a relaxation phase where the router trades balance for quality as experts differentiate (gamma_eff: 27 to 9, steps 400K-1.2M). This non-monotone trajectory, invisible to post-hoc analysis of converged models, reveals that early MoE training prioritizes balance while late training prioritizes quality. The theoretical framework is honest about its limits: the single-type equilibrium reduces to temperature-scaled softmax (held-out L1: MFG = 0.199 vs. softmax = 0.200). The game is not a better predictor; it reveals what the temperature means and, critically, how that temperature evolves. We complement the dynamics with an effective congestion decomposition, a multi-type extension that improves load prediction via token clustering on all 16 layers (mean: 30%), scope diagnostics (K/M, epsilon_l), and robustness verification across four independent quality estimators (r >= 0.89). All confidence intervals are from bootstrap resampling over 50 independent text batches.

Charafeddine Mouzouni

OPIT – Open Institute of Technology, and Cohorte AI, Paris, France.

Congestion Games Expert Load Balancing Mixture-of-Experts Language Models Mixture of Experts

Ask PDF

May 12, 2026

Slicing and Dicing: Configuring Optimal Mixtures of Experts

Mixture-of-Experts (MoE) architectures have become standard in large language models, yet many of their core design choices - expert count, granularity, shared experts, load balancing, token dropping - have only been studied one or two at a time over narrow configuration ranges. It remains an open question whether these choices can be optimized independently, without considering interactions. We present the first systematic study of over 2,000 pretraining runs spanning models up to 6.6B total parameters, in which we exhaustively vary total experts, expert dimension, heterogeneous expert sizing within a single layer, shared expert size and load-balancing mechanisms. We find that at every active-parameter scale that we study, performance consistently improves with total MoE parameters even at extreme active expert parameter ratios like 128.Further, the optimal expert size is nearly invariant to total parameter count and depends only on active parameter count. Third, we see that other choices like shared experts, heterogeneous experts and load-balancing settings have small effects relative to expert count and granularity, although dropless routing yields a consistent gain. Overall, our results suggest a simpler recipe: focus on expert count and granularity, other choices have minimal effect on final quality.

Margaret Li, Sneha Kudugunta, Danielle Rothermel +1

Paul G Allen School of Computer Science University of Washington Seattle, USA · Courant School of Data Science New York University New York, USA

Mixture-of-Experts Language Models Mixture of Experts Sparse Mixture-of-Experts AI Scaling Laws

Ask PDF

Abstract

Mixture-of-Experts Language Models

Explore similar work

ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training

Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training

Slicing and Dicing: Configuring Optimal Mixtures of Experts