Towards understanding of orthogonalization in Muon

Muon is a recent optimizer that relies on matrix orthogonalization of updates and has been shown to improve large language model (LLM) training. It does so by introducing additional momentum and Newton-Schulz iteration to the stochastic spectral descent method (SSD). However, it incurs higher communication cost if tensor parallelism is enabled, and its hyperparameter transfer propertiesare not yet fully explored. We first introduce block-wise orthogonalization, splitting weight matrices into independent tiles that are orthogonalized separatelyand recombined and empirically analyze its influence on training. This retains the validationloss while allowing up to $16$x tensor parallel splits of weight matrices. Second, we show that under spectral regularization a single learning rate transfers whendepth, width of the model, and token count are co-scaled under Chinchilla guidelines. Finally, we show that a higher weight decay value of $0.1$ underperforms during the first 80\% of the training but outperforms lower values after that, which can be attributed to the tighter spectral norm constraint. Based on this, we propose weight decay clipping and scheduling to capture both regimes. Overall, we demonstrate experimentally for nanoGPT models from 124M to 1.4B parameters that spectral regularization, both with block-wise and full-matrix orthogonalization, allows for learning rate transfer across multiple scaling dimensions and better generalization with weight decay due to the tighter spectral norm constraint.