EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training

概要

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.

参考リンク

• Paper ↗

リンク先のライセンス・再現性・データ条件をご確認ください。自動分類や実装の目安は保証ではありません。

実装を始めるプロンプト

参考リンクを確認したうえで、次の内容をコピーして利用できます。

# 論文・記事
EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training

# 概要
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…

# 原典
https://arxiv.org/abs/2610.37062

上記の手法を実装してください。実行前に依存関係とライセンスを確認してください。