Do MoE Experts Need Different Learning Rates?

Do routed experts need a larger learning rate simply because sparse routing shrinks their gradients? For Moonlet bf16 with AdamW, the scoped answer is no. Across two seeds and a 200-step sweep, lr_expert = lr_dense is the best tested baseline; the inherited 15x setting produces worse loss or reaches the collapse criterion earlier.

Routing does shrink the raw gradient. At 1x, the expert gradient-to-parameter ratio is roughly 580x to 1450x smaller than the dense ratio. But AdamW's moments shrink with that gradient too, so the applied adaptive update is still 0.36x to 0.62x the dense update. At 15x, the raw-gradient ratio remains tiny while the update rises to 5.63x, 6.12x, and 12.12x dense before routing collapses.

The old rule was plausible; it just answered the gradient question one boundary too early. If Moonlet stepped raw gradients directly, a larger expert LR would be the natural correction. Moonlet steps AdamW-preconditioned gradients, so the quantity that matters is how much attenuation remains after the optimizer.

Why 15x looked defensible

The experiment uses the actual configs/moonlet.toml architecture: dimension 2048, 12 layers, dense FFN width 11264, routed expert width 1408, 64 routed experts with top-6 selection, and 2 shared experts. Dense and router peak learning rates are both 3e-3. I swept expert LR over {0.5, 1, 2, 4, 15} × 3e-3 on FineWeb Edu under the same Moonlet WSD schedule. In bf16, expert_opt = auto selects AdamW; the nvfp4 diagnostic later in this post selects ExpertAdamW.

For one token, write the MoE output as

where $R_i$ is the selected expert set, $g_{i,e}$ is the routed weight, and $s(x_i)$ is the always-on shared branch. With mean-reduced loss

one expert parameter block receives

The routing indicator is absent from a comparable dense block. Under balanced routing and normalized top-K weights, its expected coefficient is approximately

Sparse routing really does attenuate the raw expert gradient. Moonlet's router also leaves the normalization intact: it multiplies logits by route_scale = 2.446 before sigmoid, selects top-K on scores + bias, gathers the un-biased scores, and renormalizes the selected weights to sum to one. The only post-normalization multiplier is routed_scaling_factor, which Moonlet leaves at 1.0. The value 2.446 changes routing pressure; it is not a hidden 2.446x multiplier on the expert output or its gradient.

If the optimizer were SGD and the expert gradient were a scaled signal $g_t^e = c h_t$ with $0 < c < 1$, then

Matching a dense step $-\eta_d h_t$ would suggest $\eta_e \approx \eta_d/c$. The counting argument does not derive exactly 15x, but it gives a larger expert LR the right sign. This is why the rule survived: it was attached to a mechanism that would be correct one optimizer earlier.

AdamW changes what needs equalizing

For the same scaled signal, AdamW's bias-corrected moments scale as

The expert update is therefore

When $|c|\sqrt{\hat v_t^h} \gg \epsilon$, the amplitude $c$ largely cancels between the first and second moments. Weight decay is decoupled, so it does not restore the raw-gradient argument. In the opposite regime (when epsilon, intermittency, or quantization noise matters), the cancellation can be incomplete and a different expert LR can still be justified.

This derivation is deliberately not a proof that every expert should share the dense LR. Expert and dense gradients are not scalar copies of one another; their noise, parameter norms, and moment histories differ. It shows the narrower point: token counts and raw-gradient scale cannot determine the multiplier after adaptive preconditioning. The applied step has to be measured.

For the update-proof runs I compared 1x and 15x at seed 42, logging every 10 steps. grad_to_param is each group's global gradient norm divided by its parameter norm. optimizer_update_to_pre_param is the norm of the adaptive optimizer component, separated from decoupled weight decay, divided by the pre-step parameter norm. The expert path derives it from a before/after parameter snapshot; the dense ZeRO-2 path accumulates the corresponding update directly. That statistic is the quantity the 15x story was supposed to fix.

At 1x, the raw expert ratios are between 6.89e-4 and 1.72e-3 of dense, about 1450x to 580x smaller, while the applied adaptive update remains within the same order of magnitude as dense. At 15x, the raw ratios remain between 7.41e-4 and 1.75e-3, but the applied update is already more than five times dense at step 10 and reaches 12.12x by step 100. That separately instrumented run collapses at step 130.

The direct proof is only one seed, but it tests the mechanism at the quantity the derivation predicts. The two-seed sweep asks whether the training outcome follows it.

The sweep followed the update, not the gradient

The main bf16 sweep kept the architecture, data, schedule, optimizer family, and dense/router LRs fixed. Only expert LR and seed changed.

Across these two seeds and this 200-step window, 1x is the best tested operating point. I do not mean universally stable: seed 42 reaches the collapse criterion at step 200. I mean the alternatives fail to improve on it in the ways the table can distinguish. At 0.5x, seed 42 collapses earlier and seed 43 finishes with higher loss. At 2x, one seed finishes with higher loss and the other collapses at step 80. Both 4x runs finish worse than their 1x counterparts. 15x has the highest reported loss in each seed, with seed 43 collapsing at step 110.

I also tested 1x with adam_beta2_expert = 0.95 against the 0.99 sweep setting. The ordering flipped between seeds 42 and 43, so these runs do not support a beta-two rule. That ambiguity is another reason not to replace one inherited constant with a new universal constant.

NVFP4 moves the boundary, but not the sign

The primary result above is bf16 with AdamW. The available nvfp4 experiment uses ExpertAdamW, has one seed (42), and compares only 1x with 15x; I treat it as a diagnostic, not a second settled result. Both runs complete 200 steps. 1x ends at loss 7.9618, while 15x ends at 8.1885.

More attenuation survives here than in the bf16 proof: the 1x optimizer-update ratio spans about 0.22x to 0.53x across the logged checkpoints. It still does not make 15x corrective. The full logged 15x range is 2.38x to 6.10x, and its final loss is worse. That establishes directional over-update for this pair; it does not locate the best nvfp4 multiplier.

The separate seed-42, 50-step grad-health canary is consistent with caution rather than a causal story. At step 50, 1x has loss 9.5620 and exact-zero gradient fractions of 0.71% on each of W1, W2, and W3. 15x has the slightly lower early loss, 9.4193, but exact-zero fractions of 13.35% on all three matrices. I cannot infer from this canary that the zero fraction causes the later ordering. I can say that the early loss edge coexists with a much sparser exact-zero gradient field and is gone in the separate 200-step proof pair.

For Moonlet bf16 with AdamW, I would use lr_expert = lr_dense as the baseline. It is the best tested setting in the two-seed window, and the update instrumentation explains why the old 15x setting overshoots. I would not transfer that conclusion unchanged to SGD, where the raw-gradient derivation predicts a different pressure, or to ExpertMuon, whose update dynamics were not tested here. I would also keep the nvfp4 claim at its actual size: single-seed evidence against 15x, not a completed search for its optimum.

References

Key references are Shazeer et al., "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer"; Fedus, Zoph, and Shazeer, "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity"; Kingma and Ba, "Adam: A Method for Stochastic Optimization"; and Loshchilov and Hutter, "Decoupled Weight Decay Regularization".