🤖 AI Summary
This study addresses the incompatibility of the Muon optimizer with LoRA fine-tuning, which arises from the inability to orthogonalize low-rank parameters. To overcome this limitation, we propose sMuon, a novel method that achieves the first approximate solution to the Muon objective in low-rank settings. By leveraging linearization and least-squares formulations, sMuon transforms computationally intensive decomposition operations into efficient matrix multiplications. Extensive experiments demonstrate that sMuon successfully adapts Muon to parameter-efficient fine-tuning paradigms, yielding consistent performance improvements in both supervised fine-tuning and ReLoRA pretraining tasks. Consequently, this work presents an effective optimization strategy for low-rank adaptation, significantly enhancing computational efficiency while maintaining robust model performance across diverse training scenarios.
📝 Abstract
The Muon optimizer shows clear benefits versus alternatives when pretraining neural networks. However, it is used less frequently for parameter-efficient fine-tuning (PEFT). One potential reason is that the most common PEFT method, LoRA, does not naturally combine with Muon since it is not mathematically possible to orthogonalize the weight update given by a low-rank parameterization. In this paper, we address this issue by approximating the solution to a relaxed Muon objective in the low-rank setting via linearization and then least-squares. We provide an efficient implementation that uses matmul operations only, as opposed to more complex linear algebra decomposition routines. Our method, sMuon (small Muon), performs favourably across SFT and a ReLoRA pretraining experiment. While results are model- and eval-dependent, we find overall that using Muon for low-rank fine-tuning provides moderate performance improvements.