ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training
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
| Method | Momentum | Width 1,024 | Width 2,048 |
|---|---|---|---|
| Compass | N/H | ||
| Compass | N/N | ||
| Nesterov Muon | N/N | ||
| NorMuon | N/N | ||
| Heavy-ball Muon | H/H | ||
| Excitation | H/H |
| Method | Momentum | Validation loss | Learning rate |
|---|---|---|---|
| Compass | N/H | 0.0025 | |
| Expert-only Nesterov | N/H | 0.0025 | |
| Constant expert scale | N/H | 0.0025 | |
| Compass | N/N | 0.0025 | |
| Nesterov Muon | N/N | 0.00225 | |
| NorMuon’s per-step weight decay | |||
| Tokens | Decay | Seeds | NorMuon | Compass | [95% CI] | Wins |
| Compass (N/N) | ||||||
| 65.5M | 16 | 13/16 | ||||
| 65.5M | 16 | 16/16 | ||||
| Compass on the NorMuon update | ||||||
| 65.5M | 16 | 14/16 | ||||
| 262M | 8 | 3/8 | ||||
| Method | Shuffled | Blocked |
|---|---|---|
| Compass (N/H) | ||
| Nesterov Muon | ||
| Heavy-ball Muon | ||
| Expert-only Nesterov | ||
| Constant expert scale | ||
| NorMuon |
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 |
| 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% |
| 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 |
| Width | Ablation | Difference [95% CI] | ||
|---|---|---|---|---|
| 1,024 | Expert-only Nesterov | 0.024 | ||
| 2,048 | Expert-only Nesterov | 0.069 | ||
| 1,024 | Constant expert scale | 0.00003 | ||
| 2,048 | Constant expert scale | 0.26 |
| Setting | Compass | Comparison | Difference | 95% CI | Wins | |
|---|---|---|---|---|---|---|
| 8.2M, 1,024 | N/N | NorMuon | 6/8 | 0.016 | ||
| 8.2M, 2,048 | N/N | NorMuon | 5/8 | 0.32 | ||
| 8.2M, 1,024 | N/N | Nesterov Muon | 6/8 | 0.12 | ||
| 8.2M, 2,048 | N/N | Nesterov Muon | 3/8 | 0.85 | ||
| 8.2M, 1,024 | N/H | Nesterov Muon | 5/8 | 0.56 | ||
| 8.2M, 2,048 | N/H | Nesterov Muon | 4/8 | 0.67 |
| Method | Momentum | WikiText-2 test loss |
|---|---|---|
| NorMuon | N/N | |
| Nesterov Muon | N/N | |
| Nesterov Muon, NorMuon decay | N/N | |
| Compass, NorMuon decay | N/N | |
| Compass, NorMuon decay | N/N | |
| Compass on the NorMuon update | N/N |
| Loss | Compass | Comparison | Difference | 95% CI | Wins |
|---|---|---|---|---|---|
| Validation | N/N, | NorMuon, | 22/32 | ||
| Validation | N/N, | Nesterov Muon, | 30/46 | ||
| Validation | NorMuon update, | NorMuon, | 20/32 | ||
| Full validation | N/N, | NorMuon, | 25/32 | ||
| Full validation | N/N, | Nesterov Muon, | 29/46 | ||
| Full validation | NorMuon update, | NorMuon, | 20/32 |
| Order | Compass | Comparison | Difference | 95% CI | Wins | ||
|---|---|---|---|---|---|---|---|
| Blocked | N/H | Nesterov Muon, 0.000625 | 7/8 | 0.070 | |||
| Blocked | N/H | Nesterov Muon, 0.00125 | 5/8 | 0.133 | |||
| Blocked | N/N | Nesterov Muon, 0.00125 | 6/8 | 0.500 | |||
| Blocked | N/H | Constant expert scale | 5/8 | 0.648 | |||
| Shuffled | N/H | Nesterov Muon, 0.00125 | 2/8 | 0.922 |
| Order | Comparison | Difference | 95% CI | Wins | ||
| Seeds 250 to 257 | ||||||
| Blocked | NorMuon | 6/8 | 0.031 | |||
| Blocked | Compass (N/H) | 7/8 | 0.016 | |||
| Shuffled | NorMuon | 3/8 | 0.219 | |||
| Shuffled | Nesterov Muon | 6/8 | 0.086 | |||
| Blocked minus shuffled | NorMuon | 7/8 | 0.023 | |||
| Symbol | Definition |
|---|---|
| Training step, starting at one, and total number of training steps. | |
| Expert index and number of experts in the family. | |
| Row index, number of rows and number of columns of an expert matrix. | |
| Expert weight matrix before step . | |
| Expert gradient after global gradient clipping. | |
| Momentum after its update at step , with . |
| 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 |
| Width | Heavy-ball Muon | Compass | 95% CI | Wins | |||
|---|---|---|---|---|---|---|---|
| 1,024 | 8/8 | 0.0078 | |||||
| 2,048 | 8/8 | 0.0078 |
| Setting | Value |
|---|---|
| Parameters, gradients and optimizer states | FP32 |
| GPUs per run | H100 |
| Tokens per step | 8,192 |
| Sequence length | 2,048 |
| Accumulated micro-batches per GPU | 2 |
| Sequences per micro-batch per GPU | 1 |
| 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 SD | 0.0001 | ||
| 991 | 0.00005 | 2.072954 | 2.075305 |
| FineWeb-Edu, width | DCLM | ||||
|---|---|---|---|---|---|
| Variant | 1,024 | 1,536 | 2,048 | 2,560 | 1,024 |
| Family factor | (11/16) | (5/8) | (6/8) | (7/8) | (6/8) |
| Family, uncapped | (13/16) | (7/8) | (9/9) | (5/8) | |
| Family row | (14/16) | (8/8) | (6/8) | ||
| Row weighting | (11/16) | (3/8) | (7/8) | ||
| Excitation | (2/8) | (7/8) | (2/8) | ||