EnergyCal reduces training loss spikes in sparse moes with muon optimizers
EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training
Machine LearningArtificial IntelligenceComputation and Language
Summary
Training large language models with many experts can sometimes cause sudden jumps in error, making the learning unstable. The authors propose EnergyCal, a simple method that adjusts updates during training to avoid these jumps without adding extra memory or changing key calculations. They tested EnergyCal on a 4-billion-parameter model and found it drastically reduces error spikes and improves final accuracy with only a small speed cost. This makes training these complex models more stable and efficient.
What this means in practice
- •For large ai model trainers: Stabilize and improve accuracy when training sparse mixture-of-experts language models using Muon-family optimizers.
- •For deep learning framework developers: Integrate lightweight rules like EnergyCal to reduce instability and loss spikes during sparse expert model training.
Authors
Masato Fujitake
Abstract
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.