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

Omatharv Bharat Vaidya; Ashwin Vinod; Pedram Akbarian; Aditya Sai Ellendula; Connor T. Jerzak; Nhat Ho

cs.LGOct 2, 2026

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

Authors: Omatharv Bharat Vaidya, Ashwin Vinod, Pedram Akbarian, Aditya Sai Ellendula, Connor T. Jerzak, Nhat Ho

Organizations: University of Texas at Austin

Abstract

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.

Figures & tables

Figure 1: Compass under language-blocked data order. Final test loss of Compass minus Nesterov Muon on an eight-language corpus. (a) Each point is one seed; diamonds and bars show the mean and 95% confidence interval over eight seeds. (b) Mean difference and 95% confidence interval for each language under blocked order, in the same units. Compass uses Nesterov momentum for expert matrices and heavy-ball momentum for attention and router matrices, and learning rates are tuned separately for each method and data order (Section 4.4 ).

Panel of Figure 2 — Figure 2: The Compass update for one expert matrix. Alignment is measured on the orthogonalized direction $O_{e}$ . The family factor $c_{e}$ compares whole matrices within a family, and the scalar radius $\alpha_{e}$ summarizes row-level alignment. Both scale the same direction. Router load balancing is handled separately by bias updates.

Method	Momentum	Width 1,024	Width 2,048
Compass	N/H	$5.45397\pm 0.00710$	$5.32207\pm 0.01184$
Compass	N/N	$\mathbf{5.44908}\pm 0.00768$	$5.32340\pm 0.01162$
Nesterov Muon	N/N	$5.45252\pm 0.00861$	$5.32296\pm 0.00973$
NorMuon	N/N	$5.45211\pm 0.00619$	$5.32626\pm 0.00909$
Heavy-ball Muon	H/H	$5.46876\pm 0.00673$	$5.32991\pm 0.01169$
Excitation	H/H	$5.46430\pm 0.00579$	$5.32789\pm 0.01066$

Table 1: Pretraining on FineWeb-Edu at 8.2M tokens (context length 256). Final validation loss (nats), mean

\pm

standard deviation over eight seeds. The momentum column gives the choice for expert matrices and for attention and router matrices (N: Nesterov, H: heavy-ball); Adam marks the AdamW-based methods. The best mean in each column is in bold.

NorMuon’s per-step weight decay
Method	Momentum	Validation loss	Learning rate
Compass	N/H	$4.35013\pm 0.00837$	0.0025
Expert-only Nesterov	N/H	$4.35145\pm 0.00875$	0.0025
Constant expert scale	N/H	$4.34941\pm 0.00849$	0.0025
Compass	N/N	$4.31807\pm 0.00516$	0.0025
Nesterov Muon	N/N	$4.31872\pm 0.00451$	0.00225

Table 2: Pretraining at context length 1,024 (width 1,024, 65.5M tokens). Final validation loss (nats), mean

\pm

standard deviation over eight seeds, and the selected learning rate. The upper block uses weight decay

\omega=0.01

. The lower block gives every method the per-step decay of NorMuon (learning rate 0.016,

\omega=0.01

\eta_{t}\omega=1.6\times 10^{-4}\lambda_{t}

; the rows marked twice the decay use twice this per-step decay. The best mean is in bold.

Tokens	Decay	Seeds	NorMuon	Compass	$\Delta$ [95% CI]	Wins
Compass (N/N)
65.5M	$1\times$	16	$4.30562$	$4.29974$	$-0.00587\;[-0.01024,-0.00151]$	13/16
65.5M	$2\times$	16	$4.30562$	$4.28556$	$-0.02005\;[-0.02384,-0.01627]$	16/16
Compass on the NorMuon update
65.5M	$1\times$	16	$4.30562$	$4.30110$	$-0.00452\;[-0.00727,-0.00176]$	14/16
262M	$1\times$	8	$3.70793$	$3.70774$	$-0.00019\;[-0.00412,0.00375]$	3/8

Table 3: Compass against NorMuon as the training budget grows (width 1,024, context length 1,024). The decay column gives Compass’s per-step weight decay as a multiple of NorMuon’s; NorMuon keeps its own decay throughout. Losses are final validation losses (nats), averaged over the listed number of seeds;

\Delta

is Compass minus NorMuon on the same seeds with its 95% confidence interval, and wins counts the seeds on which Compass has the lower loss.

Method	Shuffled	Blocked
Compass (N/H)	$4.40861\pm 0.00188$	$5.46396\pm 0.04741$
Nesterov Muon	$4.40840\pm 0.00249$	$5.48667\pm 0.04560$
Heavy-ball Muon	$4.40923\pm 0.00207$	$5.47008\pm 0.05157$
Expert-only Nesterov	$4.40961\pm 0.00213$	$5.47661\pm 0.04251$
Constant expert scale	$4.40866\pm 0.00380$	$5.46575\pm 0.04341$
NorMuon	$4.40463\pm 0.00202$	$5.46953\pm 0.04158$

Table 4: Held-out FineWeb2 test loss (nats) after 16.4M training tokens, mean

\pm

standard deviation over eight seeds. The same batches are shuffled or arranged in 100-step language blocks. Compass uses the N/H configuration, and the lower block applies the Compass factors to the expert matrices of the NorMuon update, with the learning-rate grid and weight decay of NorMuon. With shuffled batches all methods lie within 0.005 nats of each other, and the best mean under blocked order is in bold.

Figure 3: Training analysis at 65.5M tokens. All methods use twice NorMuon’s per-step weight decay (33 seeds for Compass (N/N), 32 for Nesterov Muon, and 30 each for NorMuon and Compass on the NorMuon update). (a) Smallest, mean, and largest expert scale

c\alpha

over the 96 expert matrices at each step, averaged over seeds and smoothed over 100 steps. (b) MaxVio; zero means an even load. (c) Mean entropy of the router softmax over all eight experts, computed without the router bias. Panels (b) and (c) use 32,768 fixed training tokens and average over the MoE layers; bands show 95% confidence intervals.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

	FineWeb-Edu
Setting	8.2M tokens	65.5M to 1.05B tokens	Multilingual
Model width	1,024 / 2,048	1,024	512
Attention heads	16	16	8
Experts per layer	8	8	64
Experts per token	2	2	2
Expert hidden width	2,048 / 4,096	2,048	1,024

Appendix

Table 5: Model sizes. In the FineWeb-Edu models, the expert hidden width is twice the model width. The multilingual model has about 71.08M active parameters per token.

Tokens	GPU	Comparison	Pairs	Extra time
65.5M	RTX PRO 6000	Compass (N/N) and Nesterov Muon, twice the decay	4	4.3%
65.5M	RTX PRO 6000	Compass on the NorMuon update and NorMuon	5	3.8%
262M	GH200	Compass on the NorMuon update and NorMuon	8	5.2%
1.05B	GH200	Compass on the NorMuon update and NorMuon	4	6.6%

Appendix

Table 6: Wall-clock cost of the Compass factors. Each pair uses the same seed on the same machine, and the time covers the whole run, including validation. The last column gives the extra time of Compass relative to the method without the factors.

Setting	Method	Learning rate: mean tuning loss	Selected
8.2M, 1,024	Nesterov Muon	0.0005625: 5.58323, 0.001125: 5.46007, 0.00225: 5.61247	0.001125
8.2M, 2,048	Nesterov Muon	0.00034375: 5.43610, 0.0006875: 5.32134, 0.001375: 5.56518	0.0006875
8.2M, 1,024	Compass (N/H)	0.000625: 5.55381, 0.00125: 5.45141, 0.0025: 5.85329	0.00125
8.2M, 2,048	Compass (N/H)	0.000375: 5.41242, 0.00075: 5.31861, 0.0015: 5.66832	0.00075
8.2M, 1,024	NorMuon	0.004: 5.58115, 0.008: 5.46068, 0.016: 5.64487	0.008
8.2M, 2,048	NorMuon	0.003125: 5.46190, 0.00625: 5.32107, 0.0125: 5.46869	0.00625

Appendix

Table 7: Learning-rate sweeps. Each entry gives a learning rate and the mean final validation loss (nats) over the two tuning seeds; the last column gives the selected rate.

Figure 4: Validation loss during training at 65.5M tokens (width 1,024, context length 1,024; seeds 250 to 265, evaluated every 666 updates). (a) Mean over the 16 seeds; the inset enlarges the last three evaluations. (b) Compass minus NorMuon on the same seed, with pointwise 95% paired-

t

confidence bands. NorMuon and Compass on the NorMuon update use NorMuon’s weight decay; Compass (N/N) uses twice NorMuon’s per-step decay, as in Table 3 .

Figure 5: Paired differences in final validation loss between Compass and each baseline. Each point is one seed; diamonds and bars show the mean and 95% confidence interval. Negative values favor Compass. At 8.2M tokens, Compass uses Nesterov momentum on all matrices (N/N), with eight seeds per comparison. At 65.5M tokens, Compass (N/N), Nesterov Muon, and NorMuon all use twice NorMuon’s per-step weight decay, with 46 seeds for Nesterov Muon and 32 for NorMuon (Section 4.3 ).

Width	Ablation	Difference [95% CI]	$t$	$p$
1,024	Expert-only Nesterov	$-0.00603\;[-0.01102,-0.00104]$	$-2.86$	0.024
2,048	Expert-only Nesterov	$-0.00576\;[-0.01212,+0.00059]$	$-2.14$	0.069
1,024	Constant expert scale	$-0.00749\;[-0.00934,-0.00564]$	$-9.56$	0.00003
2,048	Constant expert scale	$-0.00298\;[-0.00872,+0.00276]$	$-1.23$	0.26

Appendix

Table 8: Ablations at 8.2M tokens. Compass (N/H) minus each ablation in final validation loss (nats); negative values favor Compass. Intervals are 95% confidence intervals from paired

t

-tests over eight seeds, and

t

and

p

are the paired

t

statistic and its two-sided

p

-value.

Setting	Compass	Comparison	Difference	95% CI	Wins	$p$
8.2M, 1,024	N/N	NorMuon	$-0.00303$	$[-0.00532,-0.00074]$	6/8	0.016
8.2M, 2,048	N/N	NorMuon	$-0.00286$	$[-0.00913,+0.00340]$	5/8	0.32
8.2M, 1,024	N/N	Nesterov Muon	$-0.00344$	$[-0.00807,+0.00119]$	6/8	0.12
8.2M, 2,048	N/N	Nesterov Muon	$+0.00044$	$[-0.00473,+0.00560]$	3/8	0.85
8.2M, 1,024	N/H	Nesterov Muon	$+0.00144$	$[-0.00415,+0.00704]$	5/8	0.56
8.2M, 2,048	N/H	Nesterov Muon	$-0.00089$	$[-0.00562,+0.00384]$	4/8	0.67

Appendix

Table 9: Paired comparisons of final validation loss: Compass minus the comparison method (nats), 95% confidence interval, seeds favoring Compass, and two-sided paired

t

-test

p

-value. Settings give the token budget and the model width;

1\times

and

2\times

denote NorMuon’s per-step weight decay and twice that value.

Method	Momentum	WikiText-2 test loss
NorMuon	N/N	$5.8705\pm 0.0337$
Nesterov Muon	N/N	$5.8785\pm 0.0265$
Nesterov Muon, NorMuon decay	N/N	$5.8438\pm 0.0220$
Compass, NorMuon decay	N/N	$5.8381\pm 0.0319$
Compass, $2\times$ NorMuon decay	N/N	$\mathbf{5.8031}\pm 0.0350$
Compass on the NorMuon update	N/N	$5.8592\pm 0.0395$

Appendix

Table 10: WikiText-2 test loss (nats) of the models trained for 65.5M tokens at context length 1,024, mean

\pm

standard deviation over eight seeds. The evaluation scores all 286,178 test tokens and resets the context at document boundaries. The best mean is in bold.

Loss	Compass	Comparison	Difference	95% CI	Wins
Validation	N/N, $2\times$	NorMuon, $2\times$	$-0.00642$	$[-0.01027,-0.00256]$	22/32
Validation	N/N, $2\times$	Nesterov Muon, $2\times$	$-0.00242$	$[-0.00485,0.00001]$	30/46
Validation	NorMuon update, $2\times$	NorMuon, $2\times$	$-0.00330$	$[-0.00656,-0.00004]$	20/32
Full validation	N/N, $2\times$	NorMuon, $2\times$	$-0.00744$	$[-0.01075,-0.00412]$	25/32
Full validation	N/N, $2\times$	Nesterov Muon, $2\times$	$-0.00280$	$[-0.00499,-0.00060]$	29/46
Full validation	NorMuon update, $2\times$	NorMuon, $2\times$	$-0.00303$	$[-0.00583,-0.00023]$	20/32

Appendix

Table 11: Additional paired comparisons at 65.5M tokens. The validation rows use all evaluation seeds at twice NorMuon’s per-step decay: seeds 250 to 281 for the comparisons with NorMuon and 250 to 295 for Nesterov Muon. The full-validation rows score the same models on the whole validation split of about 3M tokens, 23 times the usual 131,072, and their paired

t

-test

p

-values are below 0.001, 0.014, and 0.035; for the validation rows they are 0.002, 0.051, and 0.048. The WikiText-2 rows use the models of Table 10 , and their exact two-sided paired permutation

p

-values are 0.023 and 0.031. Differences are Compass minus the comparison method, in nats, with 95% confidence intervals from paired

t

-tests; wins count the seeds on which Compass has the lower loss. “NorMuon update” marks Compass on the NorMuon update.

Order	Compass	Comparison	Difference	95% CI	Wins	$t$	$p$
Blocked	N/H	Nesterov Muon, 0.000625	$-0.01719$	$[-0.03641,+0.00204]$	7/8	$-2.11$	0.070
Blocked	N/H	Nesterov Muon, 0.00125	$-0.00886$	$[-0.02087,+0.00314]$	5/8	$-1.75$	0.133
Blocked	N/N	Nesterov Muon, 0.00125	$-0.00449$	$[-0.01879,+0.00981]$	6/8	$-0.74$	0.500
Blocked	N/H	Constant expert scale	$-0.00223$	$[-0.01341,+0.00894]$	5/8	$-0.47$	0.648
Shuffled	N/H	Nesterov Muon, 0.00125	$+0.00016$	$[-0.00250,+0.00283]$	2/8	$0.14$	0.922

Appendix

Table 12: Multilingual replication, seeds 1250 to 1257. Held-out test loss of Compass minus each comparison method, in nats, with 95% confidence intervals from paired

t

-tests and exact two-sided paired permutation

p

-values. The four blocked-order comparisons were chosen before training,

t

is the paired

t

statistic, and the last row gives the shuffled order. The replication uses the data, schedule, and learning rates of Section 4.4 and differs from the first set of runs only in its seeds.

Order	Comparison	Difference	95% CI	Wins	$t$	$p$
Seeds 250 to 257
Blocked	NorMuon	$-0.02346$	$[-0.04128,-0.00563]$	6/8	$-3.11$	0.031
Blocked	Compass (N/H)	$-0.01788$	$[-0.03087,-0.00489]$	7/8	$-3.25$	0.016
Shuffled	NorMuon	$+0.00106$	$[-0.00079,+0.00292]$	3/8	$1.35$	0.219
Shuffled	Nesterov Muon	$-0.00270$	$[-0.00569,+0.00028]$	6/8	$-2.14$	0.086
Blocked minus shuffled	NorMuon	$-0.02452$	$[-0.04324,-0.00580]$	7/8	$-3.10$	0.023

Appendix

Table 13: Compass on the NorMuon update in the multilingual experiment. Held-out test loss of Compass on the NorMuon update minus each comparison method, in nats, with 95% confidence intervals from paired

t

-tests and exact two-sided paired permutation

p

-values. The upper block uses seeds 250 to 257, and the lower block repeats the same comparisons on seeds 1250 to 1257 with the same learning rates. The first four comparisons were chosen before training, and

t

is the paired

t

statistic. The last row of each block gives the difference between the blocked and shuffled gaps relative to NorMuon.

Figure 6: Test loss by language. Compass minus Nesterov Muon under shuffled and blocked order, with 95% confidence intervals over eight seeds. Results for different languages come from the same models. Under blocked order, all eight languages favor Compass, and under shuffled order, every mean difference is at most 0.0015 nats in magnitude.

Symbol	Definition
$t,T$	Training step, starting at one, and total number of training steps.
$e,E$	Expert index and number of experts in the family.
$r,m,n$	Row index, number of rows and number of columns of an expert matrix.
$W_{e,t}$	Expert weight matrix before step $t$ .
$G_{e,t}$	Expert gradient after global gradient clipping.
$M_{e,t}$	Momentum after its update at step $t$ , with $M_{e,0}=0$ .

Appendix

Table 14: Symbols for one layer and one expert projection. A family contains the corresponding projection from every expert in that layer.

Figure 7: The scalar radius for a matrix with two rows and row factors

h=(0.4,1)

. Row weighting gives

\operatorname{diag}(h)\,O

and changes the direction. Compass applies

\alpha O

, which has the norm of

\operatorname{diag}(h)\,O

(dotted arc) and the direction of

O

. The family factor is omitted.

Setting	Width 1,024	Width 2,048
Layers	6	6
Attention heads	16	16
Experts per layer	8	8
Selected experts per token	2	2
Expert hidden width	2,048	4,096
Sequences per step	16	16

Appendix

Table 15: Settings of the runs at the two widths. Both methods use the same values.

Width	Heavy-ball Muon	Compass	$\Delta$	95% CI	$t$	Wins	$p$
1,024	$5.4661\pm 0.0078$	$5.4551\pm 0.0075$	$-0.0110\pm 0.0065$	$[-0.01638,-0.00555]$	$-4.79$	8/8	0.0078
2,048	$5.3291\pm 0.0114$	$5.3192\pm 0.0094$	$-0.0099\pm 0.0049$	$[-0.01393,-0.00578]$	$-5.72$	8/8	0.0078

Appendix

Table 16: Comparison with heavy-ball Muon at two widths on FineWeb-Edu. Final validation loss over eight evaluation seeds, with 8.192M training tokens per run. Values are mean

\pm

sample standard deviation, and

\Delta

is Compass minus heavy-ball Muon with its 95% confidence interval and statistic from a paired

t

-test. Heavy-ball Muon and Compass use learning rates 0.001125 and 0.00125 at width 1,024 and 0.0006875 and 0.00075 at width 2,048, selected on separate tuning seeds. The last column is the exact two-sided paired permutation

p

-value.

Figure 8: Per-seed differences in the heavy-ball comparison. Final validation loss of Compass minus heavy-ball Muon for each evaluation seed. All 16 seeds favor Compass.

Figure 9: Learning-rate sweep at width 2,048. Points are means over five tuning seeds, and error bars are one sample standard deviation.

Setting	Value
Parameters, gradients and optimizer states	FP32
GPUs per run	$2\times$ H100
Tokens per step	8,192
Sequence length	2,048
Accumulated micro-batches per GPU	2
Sequences per micro-batch per GPU	1

Appendix

Table 17: Settings of the JetMoE runs. Learning rates are peak values.

Seed	Matrix learning rate	Best loss	Final loss
991	0.0001	2.104669	2.115326
997	0.0001	2.102527	2.115145
1009	0.0001	2.104960	2.113939
Mean $\pm$ SD	0.0001	$2.104052\pm 0.001329$	$2.114804\pm 0.000754$
991	0.00005	2.072954	2.075305

Appendix

Table 18: Compass on JetMoE-8B. Validation loss on FineWeb-Edu after 1,000 steps and 8.192M training tokens, with AdamW learning rate 0.00002. The first three rows use matrix learning rate 0.0001, and the last row is an additional run at 0.00005. The initial validation loss is 2.157508.

	FineWeb-Edu, width				DCLM
Variant	1,024	1,536	2,048	2,560	1,024
Family factor	$-0.0057$ (11/16)	$-0.0041$ (5/8)	$-0.0049$ (6/8)	$-0.0070$ (7/8)	$-0.0041$ (6/8)
Family, uncapped	$-0.0039$ (13/16)	$-0.0039$ (7/8)	$-0.0064$ (9/9)	$+0.0175$ (5/8)
Family $\times$ row	$-0.0059$ (14/16)		$-0.0078$ (8/8)		$-0.0075$ (6/8)
Row weighting	$-0.0034$ (11/16)		$-0.0010$ (3/8)		$-0.0069$ (7/8)
Excitation	$+0.0035$ (2/8)		$-0.0080$ (7/8)		$+0.0029$ (2/8)

Appendix

Table 19: Ablations of the alignment scores against heavy-ball Muon. Final validation loss of each variant minus heavy-ball Muon at its selected learning rate; negative values favor the variant, and parentheses give wins over the number of seeds. Top: model widths and corpora; only settings with runs have entries. Bottom: variants and controls at width 1,024 on FineWeb-Edu with seeds 0 to 7, with 95% confidence intervals from paired

t

-tests. The scores

a_{e}

and

b_{er}

are the positive parts of the matrix and row cosines plus

10^{-6}

(Section 3 ), and Cap marks capping at the family or row mean.

Explore similar work

Training Communication-Efficient Mixture-of-Experts Language Models with Layer Re-Configuration
Aug 28, 2026Simeng Sun, Roger WaleffeMixture-of-Experts Language Models Mixture of Experts
NVIDIA
Slicing and Dicing: Configuring Optimal Mixtures of Experts
May 12, 2026Margaret Li, Sneha Kudugunta, Danielle Rothermel +1Mixture-of-Experts Language Models Mixture of Experts
Paul G Allen School of Computer Science University of Washington Seattle, USA · Courant School of Data Science New York University New York, USA
Marco-MoE: Open Multilingual Mixture-of-Expert Language Models with Efficient Upcycling
Apr 28, 2026Fan Jiang, Yu Zhao, Chenyang Lyu +5Multilingual Language Models Mixture-of-Experts Language Models
Alibaba International Digital Commerce

Aug 28, 2026J/K move · Enter open · S save

Training Communication-Efficient Mixture-of-Experts Language Models with Layer Re-Configuration

Simeng Sun, Roger Waleffe

NVIDIA

Mixture-of-Experts Language Models Mixture of Experts Efficient Language Model Training Communication-Efficient Distributed Training

Open paper AskPDF arXiv

When training Mixture-of-Experts (MoE) language models with expert parallelism, all-to-all token dispatch and combine collectives can consume a substantial fraction of end-to-end training time. In this work, we study communication-efficient MoE models (CE-MoE), in which we adopt a heterogeneous layer pattern that decouples token-mixing and channel-mixing depth. Compared to conventional models which interleave MoE layers after each token-mixing layer (e.g., attention, Mamba-2), CE-MoE models concentrate expert capacity in a select few routed MoE layers, while maintaining depth by adding additional token-mixing and dense-FFN layers. Across a scaling ladder from 2B to 31.5B total parameters, under matched total and activated parameters, CE-MoE models consistently reduce training cost while matching validation loss and downstream benchmarks with full-MoE baselines. At the 31.5B scale, CE-MoE uses 33.3% fewer GPU-hours while improving average downstream score and inference throughput.

Abstract

Explore similar work

Training Communication-Efficient Mixture-of-Experts Language Models with Layer Re-Configuration

Slicing and Dicing: Configuring Optimal Mixtures of Experts

Marco-MoE: Open Multilingual Mixture-of-Expert Language Models with Efficient Upcycling