Low-Rank Adaptation (LoRA) significantly reduces compute and memory costs for finetuning Deep Learning models but is often harder to tune than dense training: when using factor-wise optimizers such as AdamW, it is sensitive to initialization choices, its optimal learning rates transfer poorly across ranks, and it often fails to beat dense baselines. We derive LoRA-Muon by applying the Muon optimizer's spectral steepest-descent rule to the low-rank setting. Along with our split weight-decay rule, our main claim is that LoRA-Muon is a good low-rank proxy for full-rank Muon and Shampoo-family optimizers. Its optimal learning rates transfer across rank, width, depth, and factor-rescaling. In our compute-matched TinyShakespeare study, a rank-$2$ proxy recovers the dense best tested learning rate, and a rank-$32$ LoRA-Muon run attains lower mean validation loss than the dense baseline in the seed-averaged sweep. We further show that the Spectron optimizer depends on arbitrary factor scaling, so it would likely be a poor fit when finetuning starts from badly imbalanced factors, and that LoRA-RITE's simplified QR-coordinate core implements the same spectral update. LoRA-Muon computes that update without QR-decomposition and avoids storing second moments, making it more accelerator-friendly and memory-efficient.
翻译:摘要:低秩适配(LoRA)显著降低了深度学习模型微调的计算与内存成本,但其调优难度通常高于全秩训练:当使用逐因子优化器(如AdamW)时,该方法对初始化选择敏感,最优学习率难以在不同秩之间迁移,且常无法超越全秩基线。我们通过将Muon优化器的谱最速下降法则应用于低秩场景,推导出LoRA-Muon方法。结合我们提出的拆分权重衰减规则,核心观点是:LoRA-Muon可作为全秩Muon及Shampoo族优化器有效的低秩代理。其最优学习率可跨秩、宽度、深度及因子重缩放迁移。在计算量匹配的TinyShakespeare实验中,秩为2的代理模型恢复了全秩训练的最优已验证学习率,而秩为32的LoRA-Muon运行在种子平均扫描中实现了低于全秩基线的平均验证损失。我们进一步证明,Spectron优化器依赖于任意因子缩放,因此在从严重不平衡因子状态开始微调时可能不适用;同时,LoRA-RITE简化后的QR坐标核心实现了相同的谱更新。LoRA-Muon无需QR分解即可完成该更新,且避免存储二阶矩,从而对加速器更友好且内存效率更高。