Source-linked AI summary
Muon Sublates the Edge of Stability in LLM Pretraining
Yanzhe Chen, Qifang Zhao, Xiaoxiao Xu, Fanghui Liu
TL;DR
Muon’s large-step language-model pretraining dynamics are not captured by the classical GD edge-of-stability picture, which couples loss neutrality and update reversal. The paper analyzes stochastic no-momentum Muon and tests it in controlled models and LLM pretraining. It finds that a stochastic loss-neutral edge survives, while temporal alignment varies and remains far from coherent reversal in the studied LLM settings.
Problem
Muon’s polar update geometry means its large-step LLM-pretraining dynamics cannot be inferred directly from the familiar gradient-descent picture.
Source: §1 · p. 1Method
The paper derives stochastic no-momentum Muon diagnostics and studies them in controlled models and LLM pretraining using exact-polar updates in controlled models and practical NS-5 Muon in language models.
Source: §1 · p. 1, §1.1 · pp. 2–3Results
Takeaways & Limitations
The results support monitoring loss and trajectory signals jointly, rather than treating negative alignment alone as instability.
Source: §4 · p. 10The analysis targets base dynamics from polar normalization; extending the split-edge picture to momentum Muon requires reconsidering diagnostics for coupled parameter–momentum dynamics.
Source: §4 · p. 10
Abstract
from arXiv · showhide
Muon is increasingly used for language-model pretraining, yet its large-step dynamics are not captured by the classical edge-of-stability (EoS) picture of gradient descent (GD). In GD, loss neutrality, equal-magnitude update reversal, and marginal stability meet at a single learning-rate-dependent edge. We show that Muon breaks this coupling. For stochastic no-momentum Muon, we derive a coherence-corrected conditional loss-neutral boundary $2\rho_b/\eta$, while temporal alignment follows a separate geometry. Controlled experiments show that loss balance and temporal alignment respond differently to learning rate and batch size. Across our language model experiments, the 130M Llama-like LLM runs exhibit loss-boundary tracking with weak negative alignment, whereas the studied 1B LLM configuration shows stronger partial cancellation; in both settings, directions remain far from coherent reversal while training continues to improve. These results support a split EoS picture for Muon: a stochastic loss-neutral edge survives, but it is not accompanied by a universal temporal-direction signature. The source code for reproducing the experiments can be found in https://github.com/cyzebra/Muon-Sublates-the-Edge-of-Stability-in-LLM-Pretraining
1 Introduction
Muon’s polar updates change optimization geometry, so the classical GD edge-of-stability signatures need not coincide. The paper derives a stochastic loss-neutral edge and finds that temporal alignment varies separately across controlled and language-model experiments.
1 Introduction
- Muon changes update geometry, separating curvature-dependent loss change from the temporal geometry of how polar directions change between updates.The introduction frames whether any edge-of-stability signatures survive in Muon language-model pretraining as an open question.Source: §1 · p. 1
1.1 Contributions and findings
- For stochastic exact-polar no-momentum Muon, the loss-neutral boundary is 2ρ_k,b/η_k, with minibatch coherence shifting the full-batch value of 2/η_k.Controlled SVD-polar experiments approach the boundary; finite-reference Muon-only probes in 130M runs lie near it while complete runs continue improving over longer horizons.Source: §1.1 · p. 3
- Loss balance does not determine temporal direction: the same effective loss curvature can yield either sign of consecutive-direction alignment, and learning-rate changes can alter loss increments while alignment stays negative.The T1 and T2 crossings can occur in either order.Source: §1.1 · p. 3
- Larger batches make temporal alignment more negative in controlled models, with full-batch trajectories approaching coherent reversal while small-batch trajectories remain closer to orthogonality.The same tendency appears in fixed-checkpoint 22M interventions.Source: §1.1 · p. 3
- LLM runs show loss balance without universal reversal: 130M runs range from near orthogonality, while the 1B run at batch size 512 shows partial cancellation, still far from −1.The paper also reports layerwise c_pol measurements that reveal heterogeneity hidden by the global cosine.Source: §1.1 · p. 3
- The proposed practical implication is to monitor loss and trajectory signals jointly when considering learning-rate and batch-size schedules, rather than treating negative alignment alone as instability.Converting these observations into optimized adaptive schedules is left for future work.Source: §1.1 · p. 3
1.2 Related Work
- The paper’s related-work gap is stochastic Muon: prior geometry-aware Spectral GD analyses were deterministic or controlled, while LLM pretraining uses fresh minibatches and approximate polar updates.The cited literature also covers SGD and adaptive-method edges, LLM EoS studies, and Muon’s matrix-aware optimization geometry.Source: §1 · p. 1, §1.2 · p. 4
2 Analysis of Stochastic Muon with T1 and T2
The section separates Muon’s one-step loss balance (T1) from temporal direction alignment (T2), then uses a quadratic example to show that curvature alone cannot determine alignment.
2.1 T1: stochastic loss balance
- T1 compares effective curvature along the update segment with a coherence-adjusted loss boundary, defining conditional expected loss balance.Coherence measures retained first-order descent; together with effective curvature it determines the one-step loss change.Source: §2.1 · pp. 4–5
- In controlled networks, curvature approaches and oscillates around the coherence-dependent boundary; larger batches increase coherence and reduce fluctuation, while full batch recovers 2/η.The experiments use numerical SVD-polar updates across Transformer, MLP, and CNN settings.Source: §2.1 · p. 5
2.2 T2: the temporal-orthogonality boundary
- Larger batches yield more negative alignment in displayed runs, while small-batch trajectories can remain weakly negative and no displayed global trajectory reaches exact reversal.A negative cosine alone does not imply the two-step return associated with exact reversal.Source: §2.2 · p. 6
- Muon’s T2 temporal-orthogonality boundary depends on gradient geometry through ϑk/ηk, not a universal 1/η threshold; T1 and T2 event order varies with batch size.The diagnostic bars report the first qualifying step for loss-rise, T1, and T2 events.Source: §2.2 · p. 6, Figure 3 · p. 6
- Frozen-state mean cosines are small and positive, but stable-stage temporal means are negative, supporting minibatch dispersion and state changes as distinct sources of alignment behavior.The comparison uses a Tiny Transformer with no-momentum NS-5 Muon, batch size eight, and two seeds.Source: §2.2 · p. 6
2.3 Toy example: A quadratic linear model of separation
- At identical effective curvature 5/(4η), the quadratic example yields consecutive-direction cosines of −1 versus 1/3, separating loss curvature from temporal alignment.The example uses two different spectral configurations under full-batch Muon.Source: §2.3 · p. 7
3 Language model training at the scale
The language-model experiments track loss balance and temporal alignment separately across 22M, 130M, and 1B configurations. The 130M runs show loss-boundary tracking with weak negative alignment, while the 1B run shows stronger partial cancellation.
3.1 130M: loss balance and continued learning
- At 130M, conditional curvature repeatedly lies near the loss-neutral boundary as validation loss improves; sampled directions remain weakly negative, with alignment varying across layers.Larger probe batches generally strengthen reference-gradient coherence and shift the boundary trajectories.Source: §3.1 · p. 8, Figure 5 · p. 8, Figure 6 · p. 8
3.2 Learning-rate and batch-size perturbations
- Learning-rate changes can flip the conditional loss increment while sampled temporal cosines remain negative; a larger peak rate can improve validation loss while alignment stays controlled.After either an increase or decrease, alignment tends to return toward its pre-perturbation value.Source: §3.2 · p. 9, Figure 7 · p. 9
3.3 1B: large-scale pre-training
- At 1B, loss continues to improve despite stronger, fluctuating partial cancellation, while global alignment remains far from coherent reversal.The global cosine summarizes heterogeneous layer directions, and the cross-scale comparison does not isolate batch size.Source: §3.3 · p. 9
4 Conclusion and discussion
Muon separates conditional loss balance from temporal direction: stochastic language-model training can reach the loss edge without coherent reversal. The paper limits its characterization to base dynamics and identifies momentum Muon as future work.
4 Conclusion and discussion
- Across controlled models and language-model runs, the loss edge survives without a universal reversal signature: alignment ranges from weakly negative at 130M to stronger partial cancellation at 1B, still far from reversal.The conclusion also reports that increasing batch size makes alignment more negative in controlled experiments, while small batches stay closer to orthogonality.Source: §4 · p. 9, Figure 8 · p. 10
- A larger peak learning rate can improve validation loss while alignment remains controlled, supporting joint monitoring of loss balance and temporal direction.The paper says rate increases are justified only when the T1 diagnostic leaves sufficient room before loss neutrality.Source: §4 · p. 10
- The experiments characterize base polar-normalization dynamics, not every practical Muon variant; momentum adds temporal filtering and optimizer state that require reconsidering the diagnostics.Source: §4 · p. 10
AI use statement
- AI tools assisted language editing, organization, consistency checks, plotting, and diagnostic summaries; the authors report that measurements were produced by running code on their hardware.Source: AI use statement · p. 10
A.1 Full-batch loss balance and curvature visits
The appendix gives the full-batch specialization of the loss-neutral boundary and a deterministic curvature-visit result under stated assumptions. These claims do not establish persistent boundary tracking or a temporal-alignment condition.
A.1 Full-batch loss balance and curvature visits
- For full-batch exact-polar Muon, ρ=1 and the loss-neutral boundary is 2/η.Source: §A.1 · p. 13
- The η→0 statement is a small-step limit at fixed W_k and P_k, not a continuous-time trajectory.Source: §A.1 · p. 14
- Under the proposition’s full-batch assumptions, curvature cannot eventually remain a fixed distance below 2/η.The result concerns curvature visits, not persistent boundary tracking or temporal alignment.Source: §A.1 · p. 14
A.2 Matrix recursion and separation of T1 and T2
The matrix-quadratic analysis separates loss balance from temporal direction alignment, while the secant analysis shows that Muon’s direction-sign condition depends on update geometry. Stochastic secants also incorporate batch resampling, so they are not simply Hessian averages or conditional loss curvature.
A.2 Matrix recursion and separation of T1 and T2
- For the matrix quadratic, T2 is determined by the sign count n− = d/2, while odd d can permit a cosine sign change without a zero crossing.At cosine −1, equal displacement magnitudes follow for consecutive nonzero canonical polar factors at a fixed learning rate, but not from cosine −1 alone for unequal rates or general implemented directions.Source: §A.2 · p. 15
- Equal loss curvature can accompany opposite alignment signs: the matrix example gives cosines −1 and 1/3 for stretches with the same curvature.This establishes one-step separation, not long-time attraction.Source: §A.2 · p. 15
A.3.1 The GD reference value
- For scalar quadratic GD, direction reversal begins above ηλ = 1, before the loss-neutral and equal-amplitude reversal edge at ηλ = 2.At ηλ = 1, the next gradient vanishes and its cosine is undefined.Source: §A.3.1 · p. 15, §A.3 · p. 15
A.3.2 The full-batch Muon condition
- Muon’s exact full-batch direction-sign relation has a geometry-dependent coefficient ϑk, which need not equal one and does not define a universal threshold.When the next nonzero singular values are equal, the relation reduces to 1/η.Source: §A.3.2 · p. 16
A.3.3 Relation to loss curvature
- For loss curvature, the derived threshold reaches 2/η when ak = 1, but differs from the T2 sign boundary.For an exact two-step return, weighted loss deviations cancel even though neither step must individually be loss-neutral.Source: §A.3.3 · p. 16
A.3.4 Stochastic secants
- The stochastic secant identity requires exact polar factors and a fixed step size, and batch resampling means it is not a pure Hessian average or conditional T1 curvature.Source: §A.3.4 · p. 17
B Experiments on small models
- The small-model experiments section introduces settings and results for MLPs, CNNs, and Transformers, with settings summarized in Table 1.Source: B Experiments on small models · p. 17
B.1 Data and objectives
The MLP and CNN use CIFAR-10 training images, while the Transformer uses fixed-length TinyShakespeare sequences. Their reported objectives are full training losses with model-specific loss definitions.
B.1 Data and objectives
- The MLP and CNN train on all 50,000 CIFAR-10 training images using standardized pixel/channel coordinates and no data augmentation.Their loss is the mean squared error summed over ten one-hot outputs.Source: §B.1 · p. 17
- The Transformer trains on 2,048 fixed-length 64 sequences from the first 90% of TinyShakespeare and uses mean per-token cross entropy.The vocabulary size is 65, and all reported objectives are full training losses.Source: §B.1 · p. 17
B.2 Architectures and initialization
The experiments use deterministic small-model configurations and compare batch sizes under polar updates. Probe-based loss diagnostics estimate conditional increments, while event rules and uncertainty intervals constrain how detected crossings should be interpreted.
B.2 Architectures and initialization
- The MLP, CNN, and Transformer use distinct initialization and architecture settings, with no dropout or BatchNorm and gradients retained in evaluation mode.Source: §B.2 · p. 17
B.3 Optimizer and sampling
- The 17-run suite applies no-momentum SVD polar updates at η = 0.007 across multiple batch sizes, with full batches and finite batches sampled independently across steps.Runs use training seed 0 and share sampled-index prefixes across batch-size comparisons.Source: §B.3 · p. 17
B.4 Conditional probes and numerical normalization
- Conditional loss probes score candidate updates on a fixed full objective, while compact-SVD directions may differ from exact polar factors.The resulting hatted quantities are numerical estimators; the recorded discrepancy |Ak − Nk|/Nk reaches approximately 0.29% in current CIFAR runs.Source: §B.4 · p. 18
B.5 Event rules and uncertainty
- Reported T1 and T2 events are first detected crossings without persistence filtering, and 95% Monte Carlo intervals describe pointwise finite-probe uncertainty rather than across-seed variability or simultaneous confidence.The CNN full-batch loss/T1 event at step 1350 is not a claim about sustained onset.Source: §B.5 · pp. 18–19
B.6 Additional trajectory results
The additional trajectory results describe the batch and model coverage of the diagnostic plots and distinguish global from layerwise alignment views. The plotted sampling grid may omit minima present in complete records.
B.6 Additional trajectory results
- Full-batch global alignment minima in the complete records are −0.946209 for MLP, −0.932072 for CNN, and −0.905996 for Transformer.The displayed sampling grid need not retain every minimum.Source: §B.6 · p. 19
- Figure 9 compares fixed-objective loss, loss balance, and global temporal alignment across four batches within each model.Source: §B.6 · p. 19, Figure 9 · p. 20
- Figure 10 separates layer or block trajectories from global cosines, and its Transformer panels show two saved block aggregates rather than individual Q/K/V matrices.The figure uses full-batch training only.Source: §B.6 · p. 19, Figure 10 · p. 21
- The archive includes configurations and numerical metadata but no executed source hash, raw datasets, environment lock, or checkpoints.Source: §B.6 · p. 19
B.7 Frozen-state and temporal alignment
Frozen-state probes distinguish minibatch dispersion from trajectory alignment: independent-batch cosines are small and positive, while consecutive training directions show a systematic negative bias. The setup and branch comparisons are bounded by single-seed experiments and differing token budgets.
B.7 Frozen-state and temporal alignment
- During training, temporal means are −0.0314 and −0.0324 in two seeds, with more than 99% of consecutive pairs negative; frozen-state means are small and positive.The observed sign difference supports feedback from the evolving parameter state, but joint Muon and AdamW updates prevent attributing it to either optimizer alone.Source: §B.7 · pp. 22–23
C.1 22M setup and probes
- The 22M checkpoint branches vary learning rate or training batch size from a common step-4,000 state, but batch-size branches process different token budgets and use one seed.The branches start from reference loss 4.3578633; batch sizes 8, 16, and 32 reach 393,216,000, 524,288,000, and 786,432,000 tokens after 4,000 updates.Source: §C.1 · pp. 23–24
C.2 130M learning-rate perturbation setup
- The 130M branches change Muon learning rate at step 7,000 while holding the auxiliary AdamW schedule fixed; they report global direction cosines but have no conditional T1 probes.The branches are separate from the 130M runs with 40 tokens per parameter.Source: §C.2 · p. 24
C.3 Loss and direction responses
In 22M checkpoint branches, learning rate and batch size shift temporal alignment, while conditional loss balance and alignment respond as distinct signals. Sparse probes show return trends after perturbations, but do not resolve recovery time or a shared cosine value.
C.3 Loss and direction responses
- Mean consecutive-direction cosines were −0.064, −0.100, and −0.117 at learning rates 0.02, 0.04, and 0.08; at rate 0.04, batches 8, 16, and 32 gave −0.056, −0.100, and −0.167.All five branches had 480 negative observations over steps 4000–6400; these signs describe sampled pairs, not every step.Source: §C.3 · p. 24, Figure 12 · p. 25
- At checkpoint 4,000, doubling batch size raised coherence by 12.7% and curvature by 18.4%, yielding a positive mean increment, while batch eight gave a negative mean increment.A higher loss boundary alone therefore did not imply greater one-step descent; a Muon-only probe also did not determine later loss under complete updates.Source: §C.3 · pp. 24–25
- After learning-rate changes, sparse 22M probes showed loss balance tending to return, while alignment tended toward its pre-perturbation value; 130M alignment showed the same return trend as validation loss improved.In 130M, increasing the rate initially made alignment more negative and decreasing it made alignment less negative; these records did not measure the conditional T1 boundary.Source: §C.3 · p. 25, Figure 13 · p. 25
- Both models entered learning-rate decay during the displayed continuation, so later cosine changes reflect the schedule as well as the earlier intervention, and the return trend does not establish a shared value.The 22M reference-loss samples show transient changes, not stepwise oscillations.Source: §C.3 · p. 25
D.1 130M pre-training protocol
The 130M protocol combines no-momentum NS-5 Muon with auxiliary AdamW and uses conditional probes alongside realized-update diagnostics. Reported results include conditional loss increments, one-seed validation losses, layer and batch analyses, and a 1B run without conditional T1 probes.
D.1 130M pre-training protocol
- The 130M Llama-like model processes 5,226,627,072 tokens, about 40 per parameter, using batch size 32 and sequence length 4,096.Its 39,876 updates use GPT-2-tokenized FineWeb sample-10BT.Source: §D.1 · p. 25, Table 3 · p. 26
D.2 Conditional probes and update scope
- Conditional probes evaluate Muon-only virtual updates with auxiliary parameters fixed, whereas training updates both parameter groups; NS-5 does not assume exact-polar properties such as full-batch coherence one.The reported intervals quantify conditional probe uncertainty, not variation across seeds or reference samples.Source: §D.2 · p. 26
- Final validation losses were 3.35085, 3.32335, 3.31944, and 3.33932 for increasing Muon peak learning rates; these one-seed observations do not establish a statistically resolved ranking.The plotted learning-rate conditions are 0.02, 0.04, 0.08, and 0.12.Source: §D.2 · p. 27, Figure 14 · p. 27
D.3 130M loss increments and complete-update balance
- At peak learning rate 0.02, conditional loss intervals had positive lower endpoints at steps 3,500, 9,000, 9,500, 16,500, and 38,000; these are pointwise observations, not a corrected onset test.The first three approximate lower endpoints are 7.45 × 10^-3, 4.11 × 10^-4, and 8.48 × 10^-4.Source: §D.3 · p. 27
D.4 130M layer structure and batch-dependent coherence
- Figure 19 reports matrixwise heterogeneity, while Figure 5 analyzes fixed-checkpoint probe-batch effects rather than varying the training batch.Source: §D.4 · p. 27
D.5 1B LLM pretraining
- Figure 8 reports loss, global cosine, and layer-level alignment for the 1B run, but there are no conditional T1 probes and the cosines cover Muon directions, not complete Muon-plus-AdamW updates.The model processes approximately 20 tokens per parameter.Source: §D.5 · pp. 27–28