EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training
Sparse mixture-of-experts (MoE) language-model training can exhibit large loss spikes with Muon-family optimizers. We introduce EnergyCal, a lightweight rule that attenuates routed-expert updates according to the normalized squared Frobenius norm of their finite-step Newton-Schulz directions. EnergyCal leaves the spectral transform unchanged and adds no persistent optimizer state. We evaluate it on a 4B-parameter MoE model. Against a development-selected, lower-learning-rate MuonH recipe, EnergyCal-MuonH reduces mean thresholded loss-spike burden by 99.51% and final geometric-mean perplexity by 4.06%, with a modest throughput cost. With plain Muon at matched nominal learning rate and weight decay, the reduction in mean spike burden averages 95.98% across the two tested exponents, accompanied by small endpoint improvements. These findings support stabilization both with and without Hyperball in the evaluated setting.
Authors
- Masato Fujitake (ORCID: https://orcid.org/0000-0001-7702-499X)
Publication Details
- Journal
- arXiv (Cornell University)
- Published
- 2026-10-27
- DOI
- https://doi.org/10.48550/arxiv.2610.37062
- Primary Topic
- Stochastic Gradient Optimization Techniques
- Type
- preprint