ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training
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.
Related
- Fast Training of Mixture-of-Experts for Time Series Forecasting via Expert Loss Integration
- Mixture of Censored Partially Linear Experts Models for Symmetric or Asymmetric Data with or without Outlying Observations: A novel approach with two-piece scale-mixture of normal distributions
- The Limits of Speculation: Bounding Speculative Decoding in Mixture-of-Experts
- Non-asymptotic oracle inequalities for the Lasso in high-dimensional mixture of experts
- A Mixture-of-Experts Framework for Multi-Stream Quickest Change Detection
BibTeX
@misc{vaidya2026expertmuoncompass,
title = {{ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training}},
author = {Omatharv Bharat Vaidya and Ashwin Vinod and Pedram Akbarian and Aditya Sai Ellendula and Connor T. Jerzak and Nhat Ho},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.04140},
eprint = {2610.04140},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.04140},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training
AU - Omatharv Bharat Vaidya
AU - Ashwin Vinod
AU - Pedram Akbarian
AU - Aditya Sai Ellendula
AU - Connor T. Jerzak
AU - Nhat Ho
PY - 2026
DA - 2026/10/02
PB - arXiv
DO - 10.48550/arxiv.2610.04140
UR - https://arxiv.org/abs/2610.04140
AB - 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.
ER -