[RFC] MuonTR: Muon variant for deterministic Token-Routed architectures

[Complexity-ML]

Hey @KellerJordan, big fan of Muon — we've been using it in our framework.

We're developing MuonTR, a Muon variant specialized for deterministic Token-Routed MoE (no learned router, no top-K, CUDA graph compatible). Architecture is under review at TMLR.

Key additions on top of Muon:

Per-expert Newton-Schulz: orthogonalizes each expert slice [H, I] independently instead of the full [E, H, I] tensor, so each expert maintains its own orthogonal direction
Per-expert LR scaling: experts seeing fewer tokens (Zipf routing) get higher LR
Expert-aware weight decay: lighter decay for routed experts to preserve specialization
Gradient normalization by token count: prevents high-frequency experts from dominating
Draft implementation: https://github.com/Complexity-ML/complexity-framework/blob/main/complexity/training/muon_tr.py

Currently training a 384M Token-Routed MoE vs dense baseline. With standard AdamW the loss gap converges from +0.31 to +0.09 but plateaus. We plan to benchmark MuonTR next.

Would love your thoughts on the per-expert orthogonalization approach — does it make sense from the Newton-Schulz perspective to orthogonalize each expert slice separately?

References:

Architecture: https://openreview.net/forum?id=jZq6EVboC6 (TMLR, under review)

[Whatsonyourmind]

The per-expert Newton-Schulz orthogonalization is mathematically sound, and here is why.

Why per-expert orthogonalization works. The Newton-Schulz iteration computes the polar factor U of a matrix W = U * S by iterating X_{k+1} = X_k * (3I - X_k^T * X_k) / 2. When you orthogonalize the full [E, H, I] tensor, you are finding the polar factor of a much larger matrix where the singular value spectrum mixes all experts together. When you orthogonalize each [H, I] slice independently, you find E separate polar factors, each with its own spectrum.

The key insight: Muon's update rule is W <- W - lr * U where U = polar(G). The polar factor preserves the subspace structure of G — it is the closest orthogonal matrix to G in Frobenius norm. For MoE, each expert's gradient lives in a different subspace (different input token distributions under routing). Orthogonalizing them jointly forces all experts toward a single orthogonal structure, which is wrong when experts are meant to specialize.

Per-expert is actually more correct than joint orthogonalization because:

1. Each expert's gradient G_e has a different singular value distribution (different token frequencies under Zipf routing)
2. The Frobenius-norm initialization X_0 = G / ||G||_F needs to bring all singular values below sqrt(3). For joint orthogonalization, the largest singular value across ALL experts dictates the scaling, shrinking small-expert gradients excessively
3. Per-expert, each expert gets its own optimal scaling

On the +0.09 loss gap plateau with AdamW. This is expected for Token-Routed MoE with deterministic routing. High-frequency experts see ~10x more tokens than tail experts. AdamW's per-parameter adaptive learning rate partially compensates, but it cannot correct the fact that tail experts receive weaker gradient signal overall.

Prediction: Your per-expert LR scaling (inversely proportional to token count) should help close the gap further because it directly addresses the load imbalance that AdamW only partially compensates. I would expect the loss gap to shrink from +0.09 to roughly +0.03-0.05.

One thing to watch: with per-expert orthogonalization, the smallest experts may need 1-2 more Newton-Schulz iterations. Fewer tokens = noisier gradient = wider singular value spread = slower convergence. Monitor ||X_k^T * X_k - I||_F per expert during training.

[Complexity-ML]

Thank you @Whatsonyourmind for this thorough analysis — this is exactly the kind of rigorous feedback we needed.

We've implemented all three of your suggestions:

1.Per-expert Newton-Schulz orthogonalization — each expert's [H, I] slice is now orthogonalized independently. Your point about the Frobenius-norm initialization scaling is particularly important: with joint ortho, the largest σ across all experts dominates X_0, shrinking tail expert gradients. Per-expert, each gets its own optimal scaling.

2.Adaptive NS iterations for tail experts — we added convergence-based early stopping: iterate until ||X_k^T X_k - I||_F < τ or max 10 steps (min 5). Tail experts with fewer tokens consistently need 7-8 iterations vs 5 for high-frequency experts, confirming your prediction about wider SV spread from noisier gradients.

3.Per-expert convergence monitoring — we log ||X_k^T X_k - I||_F per expert every 100 steps. This already revealed that expert load imbalance correlates directly with NS convergence speed.

Your prediction of +0.03–0.05 loss gap with per-expert LR scaling is our target for the next ablation run.

All changes are in complexity-framework — muon_tr.py.

[Whatsonyourmind]

@Complexity-ML Great to hear the analysis was useful. The deterministic routing constraint is what makes MuonTR interesting — it forces the optimizer to work within a fixed communication topology, which is a fundamentally different optimization landscape than standard MoE where routing is learned.

Curious to see how the Nesterov momentum variant performs with the fixed routing. The momentum direction should be more stable when the routing graph doesn't change between steps, which might allow larger effective learning rates.

[Complexity-ML]

Thanks @Whatsonyourmind — spot on about Nesterov + fixed routing.

We just finished a 400M code pre-training run with AdamTR + spectral conditioning (2x B200, 15K steps, 17 languages): loss 10.55 → 0.91.

Following your prediction on the loss gap, we added per-expert LR scaling to MuonTR — tail experts get up to 2x LR based on token count ratio. Combined with adaptive NS (7-8 iters for tail vs 5 for head experts).

You're right that fixed routing stabilizes Nesterov — each expert's gradient subspace is consistent across steps, so momentum accumulates coherently instead of chasing a moving routing target. Next step: benchmark MuonTR vs AdamTR on the same setup.

085acd2

[Whatsonyourmind]

@Complexity-ML That 10.55 → 0.91 loss reduction on 15K steps with 17 languages is a strong result for a 400M run, especially given how much harder code data is to fit cleanly than web text. The fact that AdamTR + spectral conditioning got you most of the way there before MuonTR per-expert LR scaling layered on top is the cleaner version of the experiment — you can attribute the residual gain to MuonTR rather than confounding it with the conditioner.

The per-expert LR scaling decision is the move that matters most for the long-tail experts. When token-count distribution is skewed (which it almost always is in MoE on natural-language code), the tail experts see effective batch sizes 5-20x smaller than the head experts, and the optimal learning rate scales inversely with effective batch size in the small-batch regime. The 2x LR for tail experts is conservative — if you have time, sweep [1.5x, 2x, 3x, 4x] and look at the per-expert grad-norm trajectory. The right multiplier is the one where tail-expert grad-norms match head-expert grad-norms within 10-20% by step 5K. If tail grad-norms are still higher at the chosen multiplier, you are under-scaling.

The adaptive Newton-Schulz iteration count (7-8 for tail, 5 for head) is mathematically sound but worth measuring directly. The Newton-Schulz iteration converges quadratically when the input singular values are in the convergence basin, and the basin is wider for matrices closer to orthogonal. Tail experts that get fewer updates accumulate further from orthogonal between updates, so they need more iterations to recover the polar factor. The diagnostic to track is the Frobenius distance between U_k (the k-th NS iterate) and U_{k-1}: when ||U_k - U_{k-1}||_F drops below 1e-4, you can stop early. This gives you an adaptive NS that uses 5 iterations on most steps and only spikes to 7-8 when needed, without hardcoding the rule.

One measurement that would close the loop on the optimizer choice: log per-expert effective sample count (sum of routing probabilities over the batch) at each step. If the distribution flattens over training (tail experts pick up usage), the per-expert LR scaling becomes less necessary and you can decay it. If it sharpens (rich-get-richer dynamics dominate), the LR scaling should grow. A simple controller: target_lr_ratio[i] = max(1.0, mean_count / count[i]). This gives you 1.0x for above-average experts and >1.0x for below-average, with the multiplier proportional to the imbalance. Smooth it with EMA over ~500 steps to avoid LR oscillation.

The spectral conditioner direction is the part of your stack I would be most interested to see ablated independently. AdamTR + conditioning vs MuonTR + conditioning vs Muon + no conditioning at the same compute budget would tell you whether the conditioner is doing the heavy lifting or whether the optimizer itself is the gain. If the conditioner is the dominant term, that is a result worth its own writeup — most of the field is still asking what the right conditioner shape is for transformer training, and a 400M-scale empirical answer would be cited.

Nice work, the architecture is clearly converging on something interesting.

[Complexity-ML]

Thanks for the detailed feedback — really actionable, and you also helped me find a deeper bug.

First, the AdamTR result needs to be retracted. The 10.55 → 0.91 loss looked great on paper, but inference was incoherent — fluent for a few tokens then drift. Investigation showed that across all 20 layers in both AdamTR and MuonTR runs, the routed expert weights (gate_proj_w, up_proj_w, down_proj_w) had stayed identical to their initialization values (std 0.0200, 0.0200, 0.0032). Bit-exact across optimizers. That's not optimizer dynamics — that's no gradients at all.

Root cause: my CGGR Triton kernels for grouped expert GEMM are written as forward-only Triton ops, no autograd.Function wrapper, no backward. PyTorch had no way to compute dL/dW_expert, so the routed experts received zero gradients for the entire run. The model was learning entirely through the shared experts and attention. Both AdamTR and MuonTR were masking the bug because the shared experts (regular nn.Linear) were doing all the heavy lifting — and Zipf-balanced routing was unaffected because routing is a deterministic lookup, no learnable parameters there.

Fixed by gating CGGR on not torch.is_grad_enabled() — used only at inference, with the masked-loop fallback during training. Re-running with this fix now.

What this means for the optimizer comparison:

The "MuonTR > AdamTR" claim from my hackathon smoke test is invalid until I re-run with working gradients on the routed experts.
The shared experts (which DID train) grew ~10x more under MuonTR than AdamTR (std ~0.18 vs ~0.021). Even on the limited subset of weights that trained correctly, MuonTR was meaningfully different — but this is a partial signal, not a full ablation.
Your suggested ablation (AdamTR + conditioner / MuonTR + conditioner / Muon + no conditioner) is now even more important because I have to redo the whole optimizer comparison from scratch on a working model.
On your specific suggestions:

Iterate-difference NS criterion (||X_k - X_{k-1}||_F < 1e-4): implemented, replacing the orthogonality residual. min_steps=3, max_steps=10.
Per-expert grad-norm logging: added, tracked at each step, dumped every 100 steps.
Adaptive per-expert LR controller (target = max(1, mean_count/count_i), EMA-smoothed over 500 steps): implemented, replacing the hardcoded 2x. Clamped to [1, 4].
All three are in for the GPU MODE hackathon next week. Expect to rerun the AdamTR vs MuonTR ablation properly after that with the gradient bug fixed.

Thank you again — embarrassing but valuable lesson on how a forward-only kernel can produce a low-loss number that means nothing. Posting a writeup on this once the corrected runs are in.

[Complexity-ML]

Update: I have to walk back everything I claimed in this thread, including my reply 21 minutes ago. Your feedback pushed me to instrument the runs more carefully and the actual story is much worse than I described.

The bug: across all 20 layers, in every run I've posted here (AdamTR, MuonTR, AdamTR + spectral conditioning, the 1B-token MuonTR + AdamW retrain — all of them), the routed expert weights gate_proj_w, up_proj_w, down_proj_w are bit-exact identical to their initialization. Same std (0.0200, 0.0200, 0.0032), every layer, every checkpoint. Two different optimizers producing literally the same numbers down to the last digit means no gradients ever reached those parameters.

Root cause: my CGGR Triton grouped-GEMM kernels for routed expert dispatch are forward-only — no autograd.Function wrapper, no backward implementation. PyTorch had no path to compute dL/dW_expert, so the routed experts were frozen at init for every run since I integrated CGGR. The models were learning entirely through the shared experts and attention. The Zipf routing table is a deterministic buffer (no learnable params), which is why nothing complained — and why my MuonTR "clean result" from this morning generates coherent text despite the routed experts being dead. The shared expert was carrying the entire forward.

What this invalidates:

The 10.55 → 0.91 AdamTR + spectral conditioning result on code data: invalid. The down_proj std I called "shrunk to 11%" was actually just the untouched initialization (0.02 × 1/sqrt(800) = 0.0007 → 0.0032 with init scaling). It looked like collapse because I was comparing to a different config's init scale, but it's just the original numbers, never updated.
The "MuonTR per-expert orthogonalization closes the loss gap" claim from the earlier exchanges in this thread: also invalid. Per-expert NS was running on parameters that weren't moving. The convergence diagnostics I posted (5 iters head, 7-8 tail, residual logging) were measuring the orthogonalization of frozen weights.
The shared-expert std comparison (~0.18 MuonTR vs ~0.021 AdamTR) is the only signal that survives. It's a partial result that says the optimizers do behave differently on the parameters that actually got gradients, but it tells us nothing about whether per-expert NS is working.
Fix: gate CGGR on not torch.is_grad_enabled() so it's only used at inference. Training falls back to the masked-loop path that PyTorch knows how to differentiate. Re-running now.

Your three suggestions are in regardless, since they were the right call independently of the routed-expert bug:

Iterate-difference NS criterion (||X_k - X_{k-1}||_F < tol) replacing the orthogonality residual. min_steps=3, max_steps=10, tol=1e-4.
Per-expert grad-norm logging at each step, dumped every 100.
Adaptive per-expert LR controller target = max(1, mean_count / count_i), EMA-smoothed over ~500 steps, clamped to [1, 4]. The hardcoded 2x is gone.
Your suggested 3-way ablation (AdamTR + conditioner / MuonTR + conditioner / Muon + no conditioner) is now what I have to run to actually answer the question about whether MuonTR contributes anything. Will report back with the corrected runs and the per-expert grad-norm trajectories you asked about.

Apologies for the noise on the earlier claims in this thread. This is a "loss curve looked fine, model was empty" failure that I should have caught at the inference step three runs ago. Thank you for your patience and for the detailed feedback — without your push to actually instrument the per-expert dynamics, I would still be tuning hyperparams on a broken model.

[Whatsonyourmind]

@Complexity-ML This is the right way to handle a finding like this. Most people would have buried the bug and quietly fixed it; the public retraction is the move that builds long-term credibility, not the loss curve. The fact that you walked back two claims (the morning AdamTR result and your earlier reply 21 minutes ago) in the same comment is the rigor signal that matters.

A few specific things on the bug and the corrected runs:

On the diagnostic that caught it. The bit-exact identity of the routed expert weights across two different optimizers, every layer, every checkpoint is the cleanest possible signal — there is exactly one explanation (no gradients) and it is unambiguous. Worth turning that diagnostic into a standing assertion in the test suite: at the end of every training step, verify that (W_t - W_init).abs().max() > 0 for every learnable parameter in the routed-expert path. This is a one-line check that runs in microseconds and would have caught the bug at iteration 1. The general lesson is that any custom Triton kernel should ship with an autograd integration test that asserts dL/dW is non-zero on a single forward+backward pass, before it ever sees a real training run.

On what to expect from the corrected runs. With the routed experts actually receiving gradients, the corrected loss curves will almost certainly be worse than the broken numbers, not better. The reason is that you were effectively training a smaller model (just shared experts + attention) which is an easier optimization problem at small scale than the full MoE. The full MoE has higher capacity but also higher gradient noise per parameter, slower per-token convergence, and the load balancing dynamics introduce coupling between experts that the shared-only path did not have. A loss of 1.5-2.5 on the same setup would be a normal corrected number; anything below 1.0 would be suspicious and worth re-checking.

The shared-expert std differential you flagged (~0.18 MuonTR vs ~0.021 AdamTR) is genuinely interesting and survives the bug — those parameters were getting gradients. The question that result raises is whether the std difference reflects genuine optimization geometry (Newton-Schulz orthogonalization producing larger-magnitude updates) or just a learning-rate calibration mismatch between the two optimizers. The clean test is to set the effective LR (= base_LR * estimated update norm) to be approximately equal across the two and see if the std differential persists. If it does, the geometry argument is real. If it does not, it was an LR artifact.

On the 3-way ablation now being the right thing. With the routed expert bug fixed, the AdamTR + conditioner / MuonTR + conditioner / Muon + no conditioner ablation finally answers the question it was supposed to answer. The previous version was confounded because the conditioner was working on parameters that were not learning, so its contribution was indistinguishable from the lack-of-learning baseline. Now you can actually attribute the gain to the conditioner alone vs. the optimizer alone vs. the combination. I would suggest adding a fourth condition — vanilla AdamW with no conditioner and no MuonTR — as the absolute baseline. Without it, you cannot distinguish 'MuonTR helps over Adam' from 'spectral conditioning helps over no conditioning, and Adam was already fine.'

On the per-expert grad-norm trajectories. Once the routed experts are actually training, the diagnostic to watch is the ratio max(grad_norm_per_expert) / mean(grad_norm_per_expert) over the course of training. At init, this ratio is approximately 1.0 (uniform). If load balancing works, it should stay below ~3-4 throughout training. If it grows monotonically (above ~10) by the end, you have rich-get-richer dynamics and the per-expert LR controller should kick in to compensate. The adaptive multiplier you described (max(1, mean_count / count_i) clamped to [1, 4]) is conservative and should handle moderate imbalance — if you see the ratio still growing past 10 with the controller active, the ceiling needs to be raised or the EMA smoothing window needs to be shortened.

Looking forward to the corrected ablation. The story this thread tells now (researcher catches a silent kernel bug via instrumentation pushed by external feedback) is actually a better paper-section than the original loss-curve result would have been. Genuine experimental rigor under public scrutiny is a rarer signal than "my optimizer beats the baseline."