PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Omatharv Bharat Vaidya, Ashwin Vinod, Pedram Akbarian, Aditya Sai Ellendula, Connor T. Jerzak, Nhat Ho

arXiv:2610.04140v1cs.LGcs.CLmath.OC

TL;DR

Muon can give same-shaped expert matrices similarly sized steps even when their updates align poorly with current gradients, a mismatch that can arise as routed data shifts. Compass scales expert steps using alignment while keeping Muon’s direction; it matches or improves on tested optimizers in FineWeb-Edu pretraining, with its clearest advantage under blocked language order.

  • Problem

    With a shared learning rate, Muon gives same-shaped expert matrices roughly equal-sized steps even when an expert’s update aligns weakly with its current gradient as routed data changes.

    Source: §1 · p. 1

  • Method

    Compass scales each expert’s Muon step using whole-matrix alignment relative to other experts and a scalar summary of row alignment, while keeping Muon’s direction and momentum buffer.

    Source: §1 · pp. 1–2

  • Results

    On FineWeb-Edu, Compass performs as well as or better than tested optimizers with matched NorMuon weight decay in longer runs; under blocked language order, it lowers test loss in all eight languages versus Nesterov Muon.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Expert step sizes matter more when the data seen by experts changes during training, as in continual pretraining or time-varying data mixtures.

    Source: §6 · p. 11

    Models trained on these corpora can reproduce biases and harmful content, so deployment needs a separate evaluation.

    Source: Ethics statement · p. 11

Abstract

from arXiv · show

Mixture-of-experts (MoE) language models send each token to a few experts, so each expert is trained on a different part of the data, and this part changes during training. With a shared learning rate, Muon applies updates of roughly the same size to expert matrices of the same shape, even when an expert's update is poorly aligned with its current gradient. We here propose ExpertMuon-Compass (Compass), which multiplies the Muon step of each expert by two factors. A family factor compares the cosine between the expert's orthogonalized update and its gradient with the same cosine for the other experts in its layer. A scalar radius aggregates the alignment between corresponding rows of the update and gradient into one step-length multiplier. Compass keeps the update direction and the momentum buffer of Muon. In pretraining on FineWeb-Edu, Compass with Nesterov momentum on all matrices performs as well as or better than Muon, NorMuon, and other optimizers, with weight decay matched to NorMuon in the longer runs. Adding its factors to NorMuon gives the same or a lower loss than NorMuon. Compass is the most effective when the data seen by each expert varies during training, for example, as when the languages of a multilingual corpus arrive in separate blocks. With Compass, the expert load stays balanced, and the router assigns tokens to experts more decisively. We also prove a perturbation bound for the two factors.

1 Introduction

Muon gives same-shaped expert matrices roughly equal-sized steps even when their updates align poorly with current gradients. Compass scales expert steps using alignment, and its strongest reported advantage appears when experts’ data changes during training.

1 Introduction

  • Muon’s orthogonalized updates have nearly fixed norms, so same-shaped expert matrices move by about the same amount even when an update agrees weakly with its current gradient.This mismatch can occur when routed tokens change while momentum still reflects earlier tokens.Source: §1 · p. 1
  • Compass sets each expert’s step using alignment between its orthogonalized update and gradient, comparing whole-matrix alignment across experts and aggregating row alignment into a scalar radius.Both factors preserve Muon’s update direction and single momentum buffer.Source: §1 · pp. 1–2
  • Under blocked language order, Compass reaches lower test loss than Nesterov Muon in all eight languages; with shuffled batches, their test losses are nearly the same.Figure 1 reports seed-level differences and language-specific mean differences for the blocked-order comparison.Source: §1 · p. 2, Figure 1 · p. 2
  • On FineWeb-Edu, Compass with Nesterov momentum on all matrices performs as well as or better than tested optimizers, with weight decay matched to NorMuon in longer runs.Adding Compass factors to NorMuon gives the same or lower loss than NorMuon.Source: §1 · p. 2
  • The paper also analyzes training behavior and proves a bound on how much the two factors can change an update when gradients change.Compass keeps expert load balanced and makes the router assign tokens more decisively.Source: §1 · p. 2

2 Related work

Related work covers matrix optimization, MoE routing and load balancing, and alignment-based update rules. Compass differs by scaling expert steps from alignment with each expert’s own gradient while retaining Muon’s direction.

Matrix optimizers

  • Muon orthogonalizes matrix momentum, NorMuon adds row-wise second-moment normalization, and SOAP runs Adam in Shampoo’s preconditioner eigenbasis.Compass retains Muon’s single momentum buffer and adds scalar factors for expert matrices.Source: §2 · p. 3

Training mixture-of-experts models

  • MoE training commonly balances load through auxiliary losses or router-bias updates, while cited models use a shared learning rate across experts.Compass instead sets each expert’s step from its own update and gradient.Source: §2 · p. 3

Alignment and expert updates

  • Cautious optimizers mask coordinates by sign alignment, Magma combines block alignment with stochastic masking, and Excitation scales by expert utilization.Compass measures alignment after orthogonalization, compares experts within a layer, and combines row cosines into a scalar that preserves Muon’s direction.Source: §2 · p. 3

3 ExpertMuon-Compass

Compass keeps Muon’s direction while setting expert-matrix step length with family-level and row-level alignment factors. Its update retains Muon’s momentum buffer, and the paper bounds changes in the factors under changes to directions and scores.

3 ExpertMuon-Compass

  • Compass keeps each expert matrix’s Muon direction and sets its length with two alignment-derived factors; the paper bounds how much these factors can change the update.The factors are the family multiplier and scalar radius.Source: §3 · p. 3

3.1 Update direction and parameter groups

  • For expert matrices, Compass forms momentum and a Nesterov input, then applies five Newton–Schulz iterations to produce the orthogonalized direction.Attention and router matrices may use heavy-ball momentum or the Nesterov input; AdamW handles embeddings, biases, and normalization parameters.Source: §3.1 · p. 3

3.2 Expert and row scaling

  • The family factor compares an expert’s clipped cosine alignment score with its family and normalizes relative expert scales; the scalar radius summarizes row alignment while preserving the direction.Row factors move toward one as the learning rate decays, and the radius lies between zero and one.Source: §3.2 · p. 4

3.3 Complete update

  • Experts receiving zero tokens in a step are excluded from family means and keep their weights and momentum.The expert-matrix update also uses the learning rate and decoupled weight decay.Source: §3.3 · pp. 4–5

3.4 Perturbation bound for the two factors

  • The perturbation analysis bounds update changes when directions and scores change, including when gradients come from different batches.The bound allows scores to depend on the same gradient as the direction; the row contribution vanishes at the end of training when row factors equal one.Source: §3.4 · p. 5

4 Experiments

The experiments compare Compass with Muon-based and other optimizers on FineWeb-Edu, then test whether its advantage changes when expert data arrives in language blocks. Compass matches or improves on the tested baselines in short and longer pretraining, with its clearest advantage under blocked data order.

4.1 Experimental setup

  • The experiments use six-layer MoE models with top-2 routing, and tune each method’s learning rate on two seeds before evaluation on eight shared seeds.The study spans FineWeb-Edu budgets from 8.2M to 1.05B tokens and a FineWeb2 language-order comparison.Source: §4.1 · pp. 5–6, §4 · p. 5

4.2 Pretraining at 8.2M tokens

  • At width 1,024, Compass beats the constant expert-scale control on all eight seeds, supporting alignment-dependent expert steps over a fixed multiplier matched to their average scale.Compass also has lower mean loss than expert-only Nesterov at both widths.Source: §4.2 · p. 6, §4.1 · p. 6
  • At 8.2M tokens, Compass matches or improves on each baseline; at width 1,024 it has the lowest mean loss and is below NorMuon on six of eight seeds.At width 2,048 it has lower mean loss than NorMuon and matches Nesterov Muon; Compass with heavy-ball momentum on attention and router matrices beats heavy-ball Muon on all 16 seeds.Source: §4.2 · p. 6, Table 1 · p. 7

4.3 Longer training at context length 1,024

  • With matched NorMuon per-step decay, Compass matches or improves on the baselines at 65.5M tokens; adding its factors to NorMuon lowers loss at 65.5M and gives similar loss at longer budgets.At twice the decay, Compass has the lowest mean loss in Table 2; the extended comparison finds the two methods close at 262M and 1.05B tokens.Source: §4.3 · pp. 7–8, Table 3 · p. 8

4.4 Language-blocked data order

  • Under blocked language order, Compass has lower mean test loss than the tested baselines, while its gap over Nesterov Muon is larger than under shuffled batches.The advantage over Nesterov Muon appears in all eight languages; adding Compass factors to NorMuon also gives lower loss under blocked order and matches NorMuon under shuffled batches.Source: Table 4 · p. 9, §4.4 · p. 9

5 Analysis of training

The training analysis finds that Compass assigns different step scales across experts, with some experts still taking shorter steps late in training. Across the analyzed runs, load remains balanced while routing is more concentrated, and expert scales respond to language switches without observed divergence or defined loss spikes.

Compass factors

  • At 65.5M tokens, Compass’s expert scales remain heterogeneous: late in training some experts take steps about 30% shorter than average.After the first 100 steps, the smallest factor stays around 0.55–0.7 and the largest around 1.05–1.10; the mean rises from about 0.93 toward one.Source: §5 · p. 9

Load balance and routing

  • All four optimizers keep load balanced, while Compass has lower mean router entropy than the same update alone: 0.957 versus 1.020 nats with Muon, and 0.786 versus 0.843 nats with NorMuon.The lower entropy is also observed at longer budgets; with the same load balance, tokens are assigned more clearly to particular experts.Source: §5 · pp. 9–10, Figure 3 · p. 10

Language switches

  • Twenty-five steps after a language switch, Compass’s scale spread is 0.145 versus 0.136 later in the block, and its minimum is 0.26 versus 0.30, across all eight seeds.With shuffled batches, spread at the same steps stays within 0.003 of its later value; Compass shortens some experts’ steps most strongly early in each block.Source: §5 · p. 10

Stability

  • None of the 165 runs diverged, produced non-finite loss, or met the study’s loss-spike definition, including under a 100-step window.A spike is defined as a step’s training loss at least seven standard deviations from the mean of the previous 1,000 steps.Source: §5 · p. 10

6 Conclusion

Compass scales each expert’s step using agreement between its orthogonalized update and current gradient while retaining Muon’s update direction. The experiments find its strongest advantage when expert data changes during training, with balanced load and more decisive routing in the analyzed runs.

6 Conclusion

  • Compass sets each expert matrix’s step from update–gradient agreement while retaining Muon’s update direction.In FineWeb-Edu pretraining, it performs as well as or better than tested optimizers with weight decay matched to NorMuon in longer runs.Source: §6 · p. 10
  • Adding Compass factors to NorMuon lowers loss at 65.5M tokens and yields similar loss at 262M and 1.05B tokens.In these runs, expert load stays balanced and routing is more decisive than with the same update alone.Source: §6 · p. 10
  • The results suggest expert step sizes matter more when each expert’s data changes during training, as in continual pretraining or time-varying data mixtures.Source: §6 · p. 11

Ethics statement

The paper documents data-use and deployment considerations, then provides appendix details on pretraining experiments, model architecture, and the Compass training procedure.

Ethics statement

  • Training on these corpora can reproduce biases and harmful content, so deployment requires a separate evaluation.Usage conditions for the publicly released models and datasets also apply to trained checkpoints and derived data; training, model selection, and held-out evaluation data are kept separate.Source: Ethics statement · p. 11

A Experimental details for the pretraining runs

A.1 Architectures and training

  • Pretraining uses decoder-only Transformers with six pre-norm blocks; each expert has two linear projections with a GELU between them.The embedding and output layers are untied, and the models use the GPT-2 tokenizer; FineWeb-Edu experts have hidden width twice the model width, while the multilingual model has about 71.08M active parameters per token.Source: §A.1 · p. 13, Table 5 · p. 13
  • Compass applies AdamW to embeddings, output heads, biases, and normalization parameters, while its expert-matrix step uses alignment-based factors and retains the Muon update direction.The algorithm computes row factors and α_e, then adds −(η_t c_e)s(α_e O_e) to each included expert matrix after weight decay.Source: §A.1 · p. 14

A.2 Additional runs at 262M and 1.05B tokens

Additional long-budget runs test the optimizer comparisons under different decay rules and larger evaluation sets. The appendix also describes the data splits and learning-rate choices used for these comparisons.

A.2 Additional runs at 262M and 1.05B tokens

  • 0.0003 nats at 262M tokens and 0.0058 nats below NorMuon at 1.05B tokens: pooled seeds leave Compass on the NorMuon update close to NorMuon.The added seeds leave the conclusion unchanged: the two methods are close at these budgets.Source: §A.2 · pp. 14–15
  • 0.0009 nats below NorMuon at 262M tokens and 0.0073 nats below it at 1.05B: larger validation sets agree with the Table 3 conclusions.These differences are small compared with NorMuon’s 0.17 to 0.18 nats of seed-to-seed spread.Source: §A.2 · p. 15
  • 0.0072 and 0.0075 nats at 262M, and 0.039 and 0.050 nats at 1.05B: the fixed-total-decay rule raises final loss for NorMuon and Compass on the NorMuon update, respectively.Under this rule, Compass on the NorMuon update remains within 0.0004 nats of NorMuon at 262M and differs by less than 0.0001 nats at 1.05B.Source: §A.2 · p. 15
  • 3.7088 versus NorMuon’s 3.7154 mean loss: Compass (N/N) has the lowest mean at 262M tokens under the fixed-total-decay rule and is below NorMuon on eight of nine seeds (p = 0.012).Source: §A.2 · p. 15

A.3 Data

  • FineWeb-Edu uses a 95% training and 5% validation split, with larger splits built from the same document stream and training batches drawn as random windows.The short-run split contains 57M training and 3M validation tokens.Source: §A.3 · p. 15
  • The multilingual experiment uses eight FineWeb2 languages, with documents grouped by normalized URL, exact duplicates removed, and training, validation, and test splits of 90%, 5%, and 5%.Blocked-order batches cycle through five rounds of eight 100-step language blocks; shuffled order permutes the same batches.Source: §A.3 · p. 16

A.4 Learning-rate selection and baselines

  • For 65.5M-token tuning, Compass (N/N) and Nesterov Muon both select 0.00175, with mean tuning losses of 4.28936 and 4.28181, respectively; multilingual choices are almost all 0.00125.The exception is Nesterov Muon under blocked order, which selects 0.000625 by a margin of 7 × 10^-6 nats.Source: §A.4 · pp. 17–18

B Paired comparisons and statistical tests

The comparisons use paired per-seed losses, with seeds as independent replicates; in multilingual runs, Compass’s advantage over baselines is clearest under blocked language order. The appendix also specifies Compass’s numerical implementation and scheduling.

B Paired comparisons and statistical tests

B.1 Language-specific results and training dynamics

  • Under blocked order, the replicated Compass–Nesterov Muon gap averaged −0.01995 nats across 16 paired differences; Compass also favored NorMuon in all eight languages.The NorMuon-update comparison pooled to −0.02525 nats across 16 differences, favoring Compass in all eight languages.Source: B Paired comparisons and statistical… · p. 20, Table 12 · p. 21, Table 13 · p. 22
  • The blocked-order runs had 0.47% zero-token experts versus 4.11% with shuffled order, while family factors continued varying through training.Mean family factors over four quarters rose from 0.910 to 0.994 under blocking and from 0.899 to 0.993 under shuffling.Source: §B.1 · p. 21

C Numerical definition of Compass

  • Compass is defined by Equations (1)–(7) and the Equation (9) schedule; implementations add code constants and operation ordering, with all states kept in FP32.The appendix describes properties of the two factors separately.Source: C Numerical definition of Compass · p. 22

C.1 Symbols and constants

  • The symbol definitions specify expert matrices, clipped gradients, momentum, Newton–Schulz direction, alignment scores, family and row factors, and scalar radius.The family is the corresponding projection across experts in one layer.Source: §C.1 · p. 23, Table 14 · p. 23

C.2 Order of operations

  • Compass measures full-matrix and row alignment from the unscaled orthogonalized direction and clipped gradient; row factors interpolate toward one as the learning rate decays.The family factor is normalized across experts, while row scores are normalized within each expert matrix.Source: §C.2 · pp. 23–24

C.3 AdamW update

  • Parameters outside the Muon matrices use AdamW with bias-corrected moments and decoupled weight decay; Compass’s row factors reach one at the final zero-learning-rate step.On steps without a gradient, those AdamW parameters and their update counter remain unchanged.Source: §C.3 · p. 24, §C.4 · p. 24

C.5 Properties of the scalar factors

The scalar factors have bounded and direction-preserving properties under stated conditions. The appendix also specifies a separate heavy-ball Muon comparison at two widths.

C.5 Properties of the scalar factors

  • The family factors are positive and average one when the family-size condition holds; individual factors can exceed one.The authors state that all their model families meet the condition because they have at most 64 experts.Source: §C.5 · p. 25
  • For 0 ≤ λ_t ≤ 1, the row factors lie between 1 − λ_t and 1; above the denominator floor, the radius α lies in the same bounds.Although α is bounded, the product of the family factor and α can exceed one.Source: §C.5 · p. 25
  • The radius rescales the orthogonalized direction to the norm of its row-weighted version, preserving its direction and singular-value ratios; negative alignment can drive an expert’s step near zero.The negative-alignment result applies when other experts in the family have scores well above the score floor.Source: §C.5 · pp. 25–26, Figure 7 · p. 25

D Comparison with heavy-ball Muon at two widths

  • The heavy-ball Muon comparison uses separate learning-rate sweeps at two widths, with shared settings between methods and rates selected on tuning seeds before evaluation.Both methods use auxiliary-loss-free load balancing, and router bias updates are independent of Compass factors.Source: D Comparison with heavy-ball Muon at… · p. 26, Table 15 · p. 26

D.1 Metrics and tests

The section reports paired comparisons and a learning-rate sweep for Compass versus heavy-ball Muon, then describes Compass runs on JetMoE-8B. The JetMoE runs apply Compass to both expert attention and feedforward matrices.

D.1 Metrics and tests

  • At the selected width-2,048 rates, Compass is lower on four of five tuning seeds by an average of 0.00754 nats.The rates were selected using tuning seeds and then used for evaluation.Source: §D.1 · p. 26, Figure 9 · p. 27
  • All 16 evaluation seeds favor Compass over heavy-ball Muon in final validation loss.The comparison uses paired seed-level differences; the section also describes the exact paired permutation test.Source: §D.1 · p. 26, Figure 8 · p. 27

E Continued pretraining of JetMoE

  • Three runs with different seeds lower validation loss on FineWeb-Edu by about 0.043 nats after 8.2M training tokens, showing Compass applies to a released model of this size.Source: E Continued pretraining of JetMoE · p. 27

E.1 Model, data and settings

  • Compass updates JetMoE-8B’s expert attention and feedforward matrices, while router matrices use heavy-ball Muon and other listed parameters use AdamW.The model has about 8B total parameters and 2B active parameters per token.Source: §E.1 · p. 28

E.2 Per-seed results

The per-seed results report that all three standard JetMoE-8B runs reach their lowest validation loss at step 100, while an additional half-rate run reaches a lower loss.

E.2 Per-seed results

  • All three standard runs reach their lowest validation loss at step 100; an additional run at matrix learning rate 0.00005 reaches a lower loss.The additional run is seed 991; Table 18 gives validation losses after 1,000 steps and 8.192M tokens.Source: §E.2 · p. 28, Table 18 · p. 29

F Ablations of the alignment scores

The ablations test family-level and row-level alignment scores against tuned heavy-ball Muon and compare controls. The theory section states the setting for the perturbation bound, including which experts belong to a family.

F Ablations of the alignment scores

  • The family factor alone lowers mean loss at all four widths and on DCLM, with intervals below zero at widths 1,024 and 2,560.The uncapped family factor lowers loss up to width 2,048; at width 2,560 two seeds finish above heavy-ball Muon.Source: F Ablations of the alignment scores · p. 29
  • The combined family × row variant lowers loss in all three tested settings, with 95% confidence intervals below zero in each.The settings are width 1,024, width 2,048, and DCLM; the table reports differences against heavy-ball Muon.Source: F Ablations of the alignment scores · p. 29, Table 19 · p. 30
  • Permuting capped family factors removes most of their benefit, while reciprocal scores or dropping mean-one normalization sharply raise loss.The cosine factor is 0.00417 nats lower than a factor computed from gradient norms.Source: F Ablations of the alignment scores · p. 30
  • Row weighting alone lowers loss at width 1,024 on both corpora, while the combined variant preserves the orthogonalized direction that row weighting changes.Compass keeps capped family factors and converts capped row factors into a scalar radius.Source: F Ablations of the alignment scores · pp. 29–30

G.1 Setting

  • The perturbation-bound setting compares updates for the same expert family; zero-gradient experts belong to it, but experts without a gradient are skipped.The matrices may be rectangular or rank deficient.Source: §G.1 · p. 31, G Theory · p. 31

G.2 Proof of the perturbation bound

The proof bounds changes in Compass updates through changes in alignment factors and update directions. It also connects gradient errors to score errors and establishes properties of the five-step Newton-Schulz map.

G.2 Proof of the perturbation bound

  • The family and row factors are bounded and positive, and the radius map changes by at most the row-weight difference in infinity norm times the input Frobenius norm.The proof’s perturbation comparison also holds when the scores depend on the same gradient as the direction.Source: §G.2 · pp. 32–34

G.3 From gradient errors to score errors

  • The update-error bound can be expressed in terms of clipped-gradient errors for any clipping rule, including each factor’s dependence on its own gradient.The section derives direction and score errors from gradient errors before applying the perturbation bound.Source: §G.3 · p. 35

G.4 The Newton-Schulz map

  • Five Newton-Schulz iterations preserve the input matrix’s rank, and the resulting direction has operator and Frobenius norms at most 1.21.The proposition derives these properties from the map’s action on singular values.Source: §G.4 · p. 35
  • The five-step Newton-Schulz direction is generally not the polar factor, whose nonzero singular values equal one.The proof’s results use the actual transformed singular values.Source: §G.4 · p. 37

Related papers

Loading 2610.04140v1…