Source-linked AI summary
EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training
Masato Fujitake
TL;DR
Sparse MoE training can show intermittent loss spikes with Muon-family optimizers. EnergyCal attenuates routed-expert updates using finite-step Newton–Schulz direction energy, and in a 4B-parameter MoE evaluation it reduced reported spike burden with both MuonH and plain Muon, alongside endpoint improvements and a throughput cost.
Problem
Muon-family optimizers can exhibit intermittent loss spikes in sparse MoE training, where routing creates uneven, time-varying optimization conditions.
Source: §1 · p. 1Method
Results
99.51% lower mean thresholded loss-spike burden and 4.06% lower final geometric-mean perplexity were reported for EnergyCal–MuonH versus a development-selected, lower-learning-rate MuonH recipe; plain Muon showed a 95.98% average burden reduction across two tested exponents.
Source: §1 · p. 1, §7 · p. 9Takeaways & Limitations
Stabilization was observed with and without Hyperball in the evaluated MoE setting, with a modest throughput cost in the MuonH study.
Source: §1 · p. 1, §8 · p. 9Both cohorts use one custom 4B-parameter MoE model and three initialization seeds, so the results do not establish generalization across architectures, scales, data distributions, or training budgets.
Source: §7 · p. 9
Abstract
from arXiv · showhide
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.
1. INTRODUCTION
Uneven routing and momentum in inactive experts create time-varying update conditions for Muon-family optimizers, where intermittent loss spikes are observed. EnergyCal attenuates routed-expert updates using post-NS direction energy, and the reported full-schedule studies show lower spike burden with both MuonH and plain Muon.
1. INTRODUCTION
- Experts can retain nonzero momentum when they receive no tokens in the current step, motivating attention to expert-wise update magnitudes.The motivation is especially pertinent to MuonH's fixed-norm normalization.Source: §1 · p. 1
- EnergyCal attenuates routed-expert updates using the normalized squared Frobenius norm of their finite-step NS directions, without changing the NS transform or adding persistent optimizer state.Its coefficient is applied continuously; the overview curves are schematic, not measurements, and the coefficient uses direction energy rather than token counts.Source: §1 · p. 1, Figure 1 · p. 1
- 99.51% lower mean thresholded loss-spike burden and 4.06% lower final geometric-mean perplexity were reported for EnergyCal–MuonH versus a development-selected, lower-learning-rate MuonH recipe.Mean throughput was 1.62% lower; the plain-Muon study reported a 95.98% average burden reduction across two tested exponents.Source: §1 · p. 1
- The studies support stabilization with and without Hyperball in the evaluated MoE setting, but do not establish general optimizer superiority or identify the cause of individual spikes.The two examined single-update interventions failed to rescue their events, and the plain-Muon study lacks a routed-only uniform-scaling control.Source: §1 · pp. 1–2
2. RELATED WORK
EnergyCal applies a matrix-level scalar to routed-expert updates after the unchanged Newton–Schulz transform. It differs from approaches using utilization, cross-expert alignment, or other normalization signals, and the paper does not report native-method head-to-head comparisons with several related alternatives.
Expert-wise update scaling
- EnergyCal computes an independent coefficient from each routed matrix's normalized squared post-NS Frobenius norm, without utilization inputs or cross-expert normalization.For positive coefficients, scaling preserves the individual NS direction, but not necessarily the displacement after Hyperball projection or low-precision writeback.Source: §2 · p. 2
- Excitation uses expert utilization [6], while ExpertMuon-Compass uses post-orthogonalization gradient alignment across experts and rows [7].EnergyCal instead uses each routed matrix's post-NS direction energy.Source: §2 · p. 2
Magnitude and spectral corrections
Training stability in sparse models
Evaluation and optimizer comparisons
- Optimizer rankings depend on tuning and training budget, while fixed-token quality and equal-wall-time quality are distinct comparisons.Matched nominal learning rate does not imply matched parameter movement.Source: §2 · p. 2
3. ENERGY-CALIBRATED UPDATES
EnergyCal scales routed-expert updates using the normalized squared Frobenius norm of finite-step Newton–Schulz directions, while leaving other optimizer components unchanged. Its geometric motivation does not establish a causal explanation for loss spikes.
3.1. Logical matrices and the base update
- Muon uses a five-step Newton–Schulz map on a momentum-based update direction, and the finite-step result is an approximate orthogonalized direction rather than an exact polar factor.Gate and up projections may share fused storage but are optimized as separate logical matrices.Source: §3.1 · p. 2
3.2. Expert-wise energy calibration
- Each routed logical matrix receives an independent coefficient based on normalized squared post-NS Frobenius energy, with q_t bounded between zero and one.For the evaluated q_min = 0 choice, directions below the reference energy are attenuated, while those at or above it receive no additional attenuation.Source: §3.2 · p. 3
- EnergyCal replaces the learning-rate factor only for routed matrices; the NS transform, momentum, non-routed updates, and Adam fallback remain unchanged.The coefficient is computed per routed matrix at each update without token-utilization statistics or persistent optimizer state.Source: §3.2 · p. 3
3.3. Transfer to plain Muon
- The same coefficient is applied to plain Muon without Hyperball normalization or projection, but it acts on a different base update and does not imply the same parameter displacement.For plain Muon, q_t = 0 removes the NS update term but does not disable weight decay.Source: §3.3 · p. 3
3.4. Geometric motivation, not a causal theorem
- The energy statistic is not stable rank: it changes with direction rescaling and can reflect spectral concentration, direction amplitude, or both.The geometric relations motivate attenuation but do not establish a causal mechanism for loss spikes; the stated trial-update identity does not describe plain Muon's unnormalized update.Source: §3.4 · p. 4
4. EXPERIMENTAL PROTOCOL
The evaluation uses two three-seed full-schedule studies of a custom 4B-parameter MoE model, with different comparisons for MuonH and plain Muon. Spike burden, validation quality, and throughput are reported separately, with explicit limits on cross-cohort ranking and statistical interpretation.
4.1. Model, data, and training setup
- The experiments train a custom 4B-parameter DeepSeekMoE-style model with 128 routed experts, top-2 routing, and 1.31B tokens per run.Each run uses 10,000 updates; the experiments do not claim compute-optimal or fully converged pretraining.Source: §4.1 · p. 4
4.2. Cohort H: development and confirmation
- Cohort H selects MuonH and plain-Muon baselines by screening-endpoint validation cross-entropy, with both selected learning rates at the lower boundaries of their bounded search grids.Confirmation compares the frozen recipes over three new seeds; because learning rates differ, it estimates recipe-level outcomes rather than the isolated coefficient effect.Source: §4.2 · p. 4
4.3. Cohort M: plain-Muon extension
- Cohort M compares plain Muon with EnergyCal exponents p ∈ {0.5, 1} over three new seeds, holding nominal learning rate and weight decay fixed.The routed coefficient is the sole optimizer change, with no learning-rate retuning or retrospective exponent selection; its evaluation bank and implementation revision differ from Cohort H.Source: §4.3 · p. 4
4.4. Stability and statistical reporting
- The adaptive detector flags losses above a rolling threshold after update 500 and merges flags separated by at most 16 steps into an episode.Thresholded burden sums loss above the rolling median at flagged updates, not excess above the threshold; it is not a reduction in divergence probability.Source: §4.4 · p. 5
4.5. Timing and integrity
- Reported standard deviations are descriptive rather than confidence intervals, and three initializations per condition do not support general significance inference.Throughput uses short benchmarks on one eight-GPU node and does not measure equal-wall-time training quality; infrastructure failures and retries are not counted as optimizer divergence.Source: §4.4 · p. 5, §4.5 · p. 5
5. FULL-SCHEDULE RESULTS
Across the two studies, EnergyCal reduced loss-spike burden with MuonH and plain Muon; the MuonH recipe also improved final validation quality, while intermediate quality and throughput results qualify the comparison.
5.1. Cohort H: recipe-level quality and stability
- EnergyCal–MuonH reduced mean thresholded burden by 99.51% and final geometric-mean PPL by 4.06% versus development-selected MuonH; all three paired endpoints favored EnergyCal.Mean final CE fell from 3.2760 ± 0.0270 to 3.2345 ± 0.0011; the recipes differed in learning rate, so this is a recipe-level comparison.Source: §5.1 · p. 5, Table 2 · p. 6
5.2. Intermediate quality and the late crossover
- EnergyCal’s endpoint advantage over Muon + WD emerged late: its mean CE was higher at steps 3000 and 5000, while normalized observed-range CE favored Muon + WD.The late-window ordering reversed between the observed 8000- and 8500-step evaluations; normalized trapezoidal CE depends on evaluation spacing and spike timing.Source: §5.2 · pp. 5–6
5.3. Throughput, memory, and equal-time interpretation
- EnergyCal throughput was 1.62% below MuonH and 3.63% below Muon + WD, with peak allocated memory effectively unchanged relative to MuonH.No full wall-clock-matched accuracy experiment was conducted; throughput measures the full execution rather than the coefficient alone.Source: §5.3 · p. 6
5.4. Cohort M: stabilization without Hyperball
- Without Hyperball, EnergyCal reduced mean thresholded burden by 95.98% on average across two tested exponents; endpoint PPL was 0.423% and 0.380% lower for p = 0.5 and p = 1.Both exponents reduced burden in all paired runs, but neither consistently dominated; the endpoint differences are descriptive with three seeds.Source: §5.4 · p. 6
- Plain-Muon telemetry showed mean coefficients of 0.9679 ± 0.0021 for p = 0.5 and 0.9598 ± 0.0017 for p = 1, but no routed-only uniform-scaling control isolates energy-dependent attenuation.These coefficient summaries describe applied attenuation, not its causal contribution to stabilization.Source: §5.4 · pp. 6–7
6. DIAGNOSTICS AND LOCAL INTERVENTIONS
The diagnostics distinguish post-NS energy from stable rank and test whether one-step interventions can avert two selected loss excursions. The interventions reduced nominated parameter movement but did not prevent the subsequent loss rise.
6.1. Energy, rank, and amplitude ablations
- EnergyCal with p = 1 had the lowest thresholded spike burden, while a stable-rank coefficient had the lowest endpoint CE in the single-initialization audit.Alternative coefficients were not matched for effective update magnitude, and the audit was exploratory rather than an extra confirmation seed.Source: §6.1 · p. 7, Table 5 · p. 8
- A zero-token expert can retain momentum and exhibit a direction with stable rank 1.0409, illustrating that low stable rank does not imply small operator amplitude.The sampled matrices were biased toward low-energy directions rather than forming an all-expert census; selected displacement measurements establish neither a general bound nor a causal link to spikes.Source: §6.1 · p. 7, Figure 3 · p. 8
6.2. Single-update interventions do not rescue two events
- The measurements motivate post-NS energy attenuation but do not establish energy as a calibrated reliability estimate, a universally optimal exponent, or Hyperball as the sole source of spikes.A fixed-input 512-token probe showed no large output-KL reduction; it differed from the next high-loss training batch and omitted the subsequent router-bias update.Source: §6.1 · p. 7, §6.2 · p. 8
- No tested single-update branch eliminated the subsequent loss rise in either examined event, despite substantial reductions in nominated parameter movement.Candidate-only EnergyCal reduced stored operator displacement by 86.3% in event A and 76.3% in event B; the events were selected for diagnosis, not as representative spike samples.Source: Table 6 · p. 8, §6.2 · p. 8
7. DISCUSSION AND LIMITATIONS
The discussion reports reduced spike burden with both MuonH and plain Muon, alongside improved validation quality in the tested comparisons. It limits the interpretation to empirical stabilization in this setting, not general optimizer superiority or an isolated energy-calibration effect.
Empirical contribution and practical trade-offs
- EnergyCal–MuonH reduced mean burden by 99.51% versus the development-selected lower-rate MuonH recipe, with improved endpoint quality and 1.62% lower mean throughput.In Cohort M, burden fell by 95.98% on average across two tested exponents at matched nominal learning rate and weight decay; its plain-Muon comparator was faster and had better intermediate quality.Source: §7 · p. 9
The unresolved uniform-scaling alternative
- Cohort M lacks a routed-only uniform-scaling or update-budget-matched control, so it cannot distinguish energy-dependent attenuation from generally reducing routed-update magnitudes.No native-method head-to-head comparisons establish superiority over Compass or Excitation.Source: §7 · p. 9
Geometry and causal interpretation
- Single-update interventions reduced nominated movements but failed to prevent two examined spikes, so immediate removal of those events does not explain the full-schedule benefits.Accumulated changes in expert parameters or routing dynamics remain possible explanations, but are not established.Source: §7 · p. 9
Scope, measurement, and generalizability
- Both cohorts use one custom 4B-parameter, 128-expert, top-2 model, three seeds, and a shared data order, so results do not establish generalization across architectures, scales, data, or budgets.The spike detector may miss persistent loss elevation, and periodic validation may miss transient quality changes.Source: §7 · p. 9
8. CONCLUSION
EnergyCal attenuates routed-expert updates using finite-step Newton–Schulz direction energy, with observed stabilization in the evaluated MoE setting. The specific contribution of energy-dependent attenuation remains unresolved without a routed-only uniform-scaling control.
8. CONCLUSION
- EnergyCal reduced mean thresholded spike burden by 99.51% and final geometric-mean perplexity by 4.06% versus the development-selected lower-rate MuonH recipe, with 1.62% lower mean throughput.For plain Muon at matched nominal learning rate and weight decay, mean burden reductions averaged 95.98% across two tested exponents, with small endpoint gains.Source: §8 · p. 9
A. IMPLEMENTATION AND EVALUATION DETAILS
The appendix records implementation, geometry, validation, screening, and confirmation details for interpreting EnergyCal. Its results are bounded by the tested recipes, validation-bank differences, and the reported comparison design.
A. IMPLEMENTATION AND EVALUATION DETAILS
- The appendix provides numerical and parameter-partition details for interpreting EnergyCal, with the model, training budget, and comparison design described elsewhere.Source: A. IMPLEMENTATION AND EVALUATION DETAILS · p. 11
A.1. Finite-step Newton–Schulz
- Finite-step BF16 Newton–Schulz uses five quintic iterations, and its output need not have exact unit singular values, leaving post-NS direction energy informative.The norm is accumulated in FP32 and the output returned in FP32 with its original orientation.Source: §A.1 · p. 11
A.2. Logical matrices and writeback
- EnergyCal acts on routed gate, up, and down logical matrices, while attention and shared-expert matrices retain their matrix-optimizer updates.With q = 0, projection/writeback may still change MuonH weights, and plain Muon can still apply weight decay.Source: §A.2 · p. 11
A.3. Geometry and coefficient interpretation
- The stated geometric identities characterize the trial update, not projected or stored displacement or a bound on network loss change.The appendix also notes stable rank is undefined at the zero matrix and an equation condition requires the denominator guard to be inactive.Source: §A.3 · p. 11
A.4. Validation boundaries
- Validation uses fixed banks within comparisons, but different banks and implementation revisions prevent direct cross-cohort ranking of absolute CE or PPL.The observations are not certified as previously unused external test examples, and near-duplicate contamination was not independently audited.Source: §A.4 · p. 11
B.1. Development screening
- Baseline learning rates were selected by screening-endpoint validation CE after 3,000 updates, and both selected values were at the lower boundaries of their tested grids.EnergyCal–MuonH settings were fixed before confirmation, rather than reselected using confirmation outcomes.Source: §B.1 · p. 11
B.2. Confirmation effects
- The reported 99.51% burden reduction versus MuonH uses condition-level mean burdens, not run-wise percentage reductions.All confirmation runs contribute, including the zero-burden EnergyCal run and the unfavorable burden comparison against Muon + WD.Source: §B.2 · p. 11
B.3. Intermediate-quality summaries
The intermediate-quality summaries favor different optimizers depending on the time window: observed-range CE favors plain Muon, while a post hoc late-window check favors EnergyCal. The separate plain-Muon transfer study uses matched nominal optimizer settings and a complete schedule, but its evaluation bank prevents direct ranking against Cohort H.
B.3. Intermediate-quality summaries
- Observed-range CE favors plain Muon, whereas the post hoc late-window summary favors EnergyCal.Neither summary replaces the predefined final endpoint; the late-window advantage appears late in validation CE.Source: §B.3 · p. 12, Figure 4 · p. 12
- The adaptive spike detector can miss persistent high-loss periods, while narrow excursions can produce large maxima but modest burden.Final validation CE, maximum training loss, episodes, and thresholded burden are complementary outcomes; all final runs complete.Source: §B.3 · p. 12
C. PLAIN-MUON TRANSFER
- Cohort M compares plain Muon with EnergyCal exponents p = 0.5 and p = 1 across three new initializations and the full 10,000-update schedule.The study uses the same architecture, data order, and token budget as Cohort H, but a separate implementation revision and evaluation bank preclude direct cross-cohort ranking.Source: C. PLAIN-MUON TRANSFER · p. 12
- Cohort M holds matrix learning rate, weight decay, and fallback learning rate fixed; only the routed coefficient changes.EnergyCal uses ρref = 0.8, qmin = 0, and the two stated exponents.Source: C. PLAIN-MUON TRANSFER · p. 12
C.1. Quality and stability
Cohort M reduces thresholded burden with both tested EnergyCal exponents, while endpoint differences between exponents are small. Coefficient measurements, unperformed controls, and single-update interventions qualify what these results establish.
C.1. Quality and stability
- 95.98% is the equal-weight average reduction in thresholded burden across the two EnergyCal exponents, not a pooled six-run estimate.Both exponents reduce burden in all three paired runs; endpoint CE improves in two of three runs for p = 0.5 and all three for p = 1, with a paired difference of −0.0004 ± 0.0016 nats/token between calibrated variants.Source: §C.1 · p. 12, Table 11 · p. 14
- The observed-range advantage is partly influenced by a plain-Muon spike at one evaluation, and calibrated variants do not dominate at every point.Final CE at 10,000 updates is the primary endpoint; the late-window check is post hoc.Source: §C.1 · p. 12, Figure 5 · p. 13
C.2. Coefficient measurements
- Mean recorded coefficients are 0.9679 ± 0.0021 for p = 0.5 and 0.9598 ± 0.0017 for p = 1.Some coefficients equal zero; positive scaling preserves a nonzero NS direction, while q = 0 suppresses that direction without necessarily freezing the stored parameter.Source: §C.2 · p. 13
C.3. Unperformed routed-only scalar controls
- The matched nominal learning-rate and decay comparison does not establish whether energy-dependent assignment outperforms uniform routed-update reduction.Neither fixed-scalar nor norm-budget-matched controls were evaluated; the latter matches only summed squared direction-step norms at a common state, not stored changes, decay, or future trajectories.Source: §C.3 · p. 13
D.1. Alternative coefficients and spectral measurements
- The alternative spectral quantities are estimated by power iteration, and the alternatives were not tuned to equal update budgets.Selected routed-down matrices were also analyzed using CPU float64 SVD, but selection deliberately included low-energy directions, so the sampled distribution is not an unbiased census.Source: §D.1 · p. 14
D.2. Single-update interventions
- Candidate-only EnergyCal reduces stored operator displacement by 86.3% in event A and 76.3% in event B, but the next-step loss rise persists.Single-update effects on short-window burden are small and inconsistent; they do not establish substantial immediate rescue or a zero causal effect.Source: §D.2 · p. 14, Table 12 · p. 14
D.3. Fixed-input probes and causal limitations
Fixed-input probes measure CE and output KL immediately before and after treated updates, but they do not identify what produced the examined loss events.
D.3. Fixed-input probes and causal limitations
- The fixed 512-token probe differs from the next high-loss training batch and omits the subsequent router-bias-controller update.Thus, these local observations do not identify the combination of history, routing, and inputs responsible for the events, and do not negate full-schedule improvement under continuous EnergyCal.Source: §D.3 · p. 14