模型訓練與最佳化
sMuon Projects Muon Into the LoRA-Expressible Subspace, Runs About 30% Faster Than Riemannion at the Same Loss
sMuon first projects the orthogonalized full-weight update direction into the low-rank space expressible by LoRA factors, then computes the update using only matrix multiplications. In ReLoRA experiments, it matches Riemannion’s lowest validation loss, although cross-model fine-tuning results depend heavily on the optimizer originally used for the base model.
Muon orthogonalizes the update to each layer’s weight matrix, pushing its nonzero singular values toward uniformity. However, LoRA permits weight changes only through the low-rank matrix `BA`, and the full Muon direction generally lies outside this expressible space. The newly proposed sMuon first linearizes Muon’s spectral-norm steepest-descent objective, then solves a least-squares problem to project the orthogonalized gradient back into the row and column subspaces of the LoRA factors. The authors further exploit the low-rank structure to reformulate the computation, avoiding QR decomposition or SVD on the full weight matrix. The main steps require only matrix multiplications, inverse square roots, and operations on small matrices.
In ReLoRA experiments using a 162-million-parameter Transformer, FineWeb data, and adapter merging every 500 steps, both sMuon and Riemannion achieved a final validation loss of 3.66, outperforming LoRA-Muon at 3.68, per-factor Muon at 3.70, and AdamW at 3.83. The paper also reports that sMuon is about 30% faster. It maintains only one momentum state per parameter, theoretically requiring less optimizer-state memory than AdamW’s first- and second-moment estimates.
The results are not uniformly superior. The authors evaluated Qwen2.5-3B, Llama-3.2-3B, DeepSeek-V2-Lite, and the Muon-pretrained Moonlight-16B-A3B across 11 commonsense-reasoning and coding tasks. sMuon ranked first on six tasks with Moonlight, but AdamW or other low-rank Muon variants often led on models pretrained with AdamW. For example, on HumanEval pass@1 with Qwen2.5-3B, AdamW scored 54.9%, compared with 51.8% for sMuon. This suggests that the implicit biases of pretraining and fine-tuning optimizers continue to interact. Engineering teams should next watch for an official implementation, evidence of stability under distributed and mixed-precision training, and measurements of actual GPU-memory use and wall-clock cost at higher LoRA ranks. The paper currently provides no dedicated public sMuon library.