EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training

Masato Fujitake

cs.LGOct 27, 2026

EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training

Authors: Masato Fujitake

Organizations: FA Research, Fast Accounting Co., Ltd.

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.

Figures & tables

Figure 1: EnergyCal overview. (1) Sparse routing creates uneven expert traffic. (2) Each routed matrix’s post-NS energy determines $q_{e}$ ; when $q_{e}>0$ , scaling preserves the individual NS direction, not necessarily the displacement after Hyperball projection or writeback. (3) Continuous calibration reduces observed thresholded loss-spike burden. Curves are schematic, not measurements; coefficients use direction energy, not token counts.

Recipe	Matrix LR	Matrix WD	Fallback WD
MuonH	0.005	0	0
EnergyCal ( $p=1$ )	0.010	0	0
Muon + WD	0.005	0.1	0.1

Table 1: Frozen Cohort H recipes. These are development-selected, best-tested settings rather than a coefficient-only factorial. Fallback peak LR is 0.001; weight decay applies only to eligible parameters.

Recipe	Peak LR	Final CE	PPL	Episodes	$B_{1}$	$L_{\max}$
MuonH	0.005	$3.2760\pm 0.0270$	$26.476\pm 0.721$	$37.33\pm 4.04$	$2046.01\pm 1289.75$	$39.99\pm 12.64$
EnergyCal ( $p=1$ )	0.010	$3.2345\pm 0.0011$	$25.394\pm 0.027$	$3.33\pm 2.89$	$10.11\pm 9.65$	$5.92\pm 1.07$
Muon + WD	0.005	$3.2410\pm 0.0060$	$25.560\pm 0.152$	$8.33\pm 3.51$	$67.41\pm 52.76$	$7.66\pm 1.83$

Table 2: Cohort H confirmation: complete 10k schedule, 1.3107B tokens per run. Values are arithmetic mean

\pm

sample SD over

n=3

unused initialization seeds. PPL is computed per run before aggregation; episodes and burden

B_{1}

are per-run quantities.

L_{\max}

is the maximum post-warmup training loss within each run before aggregation. LRs are selected within bounded development grids.

Figure 2: Fixed-bank validation CE across Cohort H confirmation runs. Lines show means and shaded bands show

\pm 1

sample SD over three initialization seeds. All 20 evaluations are included, including loss excursions. Bands describe run-to-run variation, not confidence intervals. The predefined primary endpoint is step 10,000; development observations are not concatenated with these curves.

Recipe	Tokens/s	Optimizer (s)	$M_{\rm alloc}$ (GiB)
MuonH	$67061.4\pm 570.4$	$0.5460\pm 0.0017$	$39.7224\pm 0.0000$
EnergyCal ( $p=1$ )	$65976.8\pm 353.3$	$0.5487\pm 0.0015$	$39.7224\pm 0.0000$
Muon + WD	$68458.9\pm 487.7$	$0.4824\pm 0.0014$	$39.7245\pm 0.0000$

Table 3: Cohort H production-style benchmark: arithmetic mean

\pm

sample SD over

n=3

execution repeats on one eight-H100 node, not three quality seeds. Throughput is aggregated within each run over steps 501–1000. Optimizer time is summarized first by its within-run median and then across repeats. Peak allocated memory is first maximized over recorded ranks, including initialization, and then summarized across repeats. SD 0.0000 indicates identical recorded peaks, not an absence of hardware variability.

Recipe	Final CE	PPL	Episodes	$B_{1}$
Muon + WD	$3.189786\pm 0.002381$	$24.283\pm 0.058$	$11.67\pm 4.04$	$84.79\pm 33.13$
Muon + EnergyCal ( $p=0.5$ )	$3.185547\pm 0.001877$	$24.181\pm 0.045$	$3.00\pm 2.00$	$4.25\pm 4.06$
Muon + EnergyCal ( $p=1$ )	$3.185975\pm 0.001173$	$24.191\pm 0.028$	$1.67\pm 1.15$	$2.57\pm 2.46$

Table 4: Cohort M: EnergyCal added to plain Muon without Hyperball. All conditions share matrix LR 0.005, eligible matrix and fallback weight decay 0.1, and the same 10k-update token budget. Values are mean

\pm

sample SD over

n=3

new initialization seeds; episodes and

B_{1}

are per-run quantities. The validation bank differs from Cohort H, so absolute CE and PPL values must not be compared directly with Table 2 .

Condition	3k CE	Ep.	$B_{1}$
MuonH	3.791280	4	430.51
Energy ( $p=0.5$ )	3.774685	4	89.03
EnergyCal ( $p=1$ )	3.769136	3	10.59
Stable-rank	3.761237	4	76.17
Amplitude only	9.318367	9	364.67
Near-zero gate	3.814775	6	531.36

Table 5: Single-initialization diagnostic audit at common matrix LR 0.01, over 3k updates. Alternative coefficients are not matched for effective update magnitude. An endpoint occurring during a spike is retained.

Figure 3: Selected diagnostic measurements of direction stable rank and BF16-stored operator displacement. Stored movement is measured directly; the selected observations do not establish a general bound or a causal link to spikes.

Single-update intervention	Event A	Event B
Canonical	6.3165	6.9069
Candidate: freeze	6.3172	6.9110
Control: freeze	6.3160	6.9085
All routed: EnergyCal	6.3024	6.9873

Table 6: Next-step training CE after single-update interventions in two distinct saved-state events. Pre-intervention CE is 4.1928 (A) and 4.2143 (B). No tested branch eliminates the subsequent loss rise.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Recipe	LR	Final CE	Episodes	$B_{1}$
MuonH	0.005000	$3.703740\pm 0.013658$	$8.00\pm 0.00$	$339.40\pm 379.07$
MuonH	0.007071	$3.733119\pm 0.013795$	$7.00\pm 1.41$	$562.76\pm 0.88$
MuonH	0.010000	$3.949324\pm 0.226689$	$11.50\pm 4.95$	$592.87\pm 99.21$
EnergyCal ( $p=1$ )	0.010000	$3.769425\pm 0.000647$	$2.50\pm 0.71$	$8.76\pm 0.24$
Muon + WD	0.005000	$3.703061\pm 0.004366$	$1.50\pm 0.71$	$1.84\pm 1.53$
Muon + WD	0.010000	$3.772394\pm 0.011707$	$5.50\pm 0.71$	$75.87\pm 0.14$

Appendix

Table 7: Completed development screen: mean

\pm

sample SD over

n=2

initialization runs after 3,000 updates. All endpoints, including unfavorable ones, are retained. Episodes and burden are per-run quantities.

EnergyCal minus	CE difference	$B_{1}$ difference	Episode difference
MuonH	$-0.041491\pm 0.026043$	$-2035.91\pm 1297.37$	$-34.00\pm 1.73$
Muon + WD	$-0.006510\pm 0.005709$	$-57.31\pm 55.57$	$-5.00\pm 6.24$

Appendix

Table 8: Paired EnergyCal-minus-baseline differences over

n=3

confirmation initializations (mean

\pm

sample SD). Negative values favor EnergyCal. Pairing precedes aggregation.

Recipe	Observed-range CE	Late-window CE
MuonH	$3.992289\pm 0.339400$	$3.278503\pm 0.007961$
EnergyCal ( $p=1$ )	$3.627002\pm 0.001572$	$3.251435\pm 0.000639$
Muon + WD	$3.611137\pm 0.009433$	$3.256059\pm 0.005045$

Appendix

Table 9: Cohort H curve summaries: mean

\pm

sample SD of run-level results. The late-window comparison is post hoc.

Figure 4: Late-schedule validation CE (mean

\pm

sample SD,

n=3

). EnergyCal’s small endpoint advantage over plain Muon appears late; the full trajectory is shown in Fig. 2 .

Figure 5: Cohort M validation CE: mean

\pm

sample SD over three runs on its own fixed evaluation bank. All observations, including the plain-Muon loss excursion at step 4000, are retained. Bands are not confidence intervals.

Recipe	PPL	$L_{\max}$	Observed-range CE	Late-window CE
Muon + WD	$24.2833\pm 0.0578$	$8.54\pm 0.82$	$3.578701\pm 0.042809$	$3.203759\pm 0.002122$
Muon + EnergyCal ( $p=0.5$ )	$24.1805\pm 0.0454$	$5.61\pm 0.87$	$3.548957\pm 0.002669$	$3.200299\pm 0.001403$
Muon + EnergyCal ( $p=1$ )	$24.1909\pm 0.0284$	$5.49\pm 0.75$	$3.550821\pm 0.000788$	$3.200778\pm 0.001217$

Appendix

Table 10: Additional Cohort M results: mean

\pm

sample SD over

n=3

initialization runs. Observed-range CE integrates the 500–10,000 evaluations; late-window CE is a post hoc check over the final three evaluations.

Coefficient	$\Delta$ CE	$\Delta B_{1}$
$p=0.5$	$-0.004239\pm 0.004194$	$-80.54\pm 36.44$
$p=1$	$-0.003811\pm 0.002849$	$-82.23\pm 30.67$

Appendix

Table 11: Cohort M paired effects: mean

\pm

sample SD over three within-initialization differences (EnergyCal minus plain Muon). Negative values favor calibration.

Intervention	CE A	CE B	$r_{A}$ (%)	$r_{B}$ (%)
Canonical	6.3165	6.9069	0.00	0.00
Candidate ECal	6.3163	6.9068	+0.93	+0.22
Candidate zero	6.3160	6.9109	-1.25	+0.32
Candidate freeze	6.3172	6.9110	+3.09	+0.76
Control freeze	6.3160	6.9085	+1.61	+1.40
All routed ECal	6.3024	6.9873	+2.12	-3.41

Appendix

Table 12: All intervention branches in events A and B (

n=1

per branch and event).

r_{A},r_{B}

are short-window burden improvements, not reductions in the full-schedule

B_{1}

Muon Optimizer

Muon optimizer uses momentum and matrix orthogonalization to shape neural-network updates, especially in large language-model training.

See all 150 papers

EMO: Frustratingly Easy Progressive Training of Extendable MoE
May 13, 2026Linghao Jin, Chufan Shi, Huijuan Wang +4Sparse Mixture-of-Experts Efficient Language Model Training
USC-ISI · MBZUAI-IFM
Reconstruction-Optimized Expert Pruning for Mixture-of-Experts Language Models
Oct 19, 2026Jelena Markovic-Voronov, Petros Prastakos, Kayhan Behdin +7Mixture-of-Experts Pruning Mixture-of-Experts Language Models
LinkedIn · MIT · Google
ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training
Oct 2, 2026Omatharv Bharat Vaidya, Ashwin Vinod, Pedram Akbarian +3Deep Learning Optimization Mixture-of-Experts Language Models
University of Texas at Austin

May 13, 2026J/K move · Enter open · S save

EMO: Frustratingly Easy Progressive Training of Extendable MoE

Linghao Jin, Chufan Shi, Huijuan Wang, Nuan Wen+3

USC-ISI · MBZUAI-IFM

Sparse Mixture-of-Experts Efficient Language Model Training Mixture-of-Experts Models

Open paper AskPDF arXiv

Sparse Mixture-of-Experts (MoE) models offer a powerful way to scale model size without increasing compute, as per-token FLOPs depend only on k active experts rather than the total pool of E experts. Yet, this asymmetry creates an MoE efficiency paradox in practice: adding more experts balloons memory and communication costs, making actual training inefficient. We argue that this bottleneck arises in part because current MoE training allocates too many experts from the beginning, even though early-stage data may not fully utilize such capacity. Motivated by this, we propose EMO, a simple progressive training framework that treats MoE capacity as expandable memory and grows the expert pool over the course of training. EMO explicitly models sparsity in scaling law to derive stage-wise compute-optimal token budgets for progressive expansion. Empirical results show that EMO matches the performance of a fixed-expert setup in large-scale experiments while improving wall-clock efficiency. It offers a surprisingly simple yet effective path to scalable MoE training, preserving the benefits of large expert pools while reducing both training time and GPU cost.

Abstract

Muon Optimizer

Explore similar work

EMO: Frustratingly Easy Progressive Training of Extendable MoE

Reconstruction-Optimized Expert Pruning for Mixture-of-Experts Language Models

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