LoRA-Muon让低秩微调更稳定,学习率可跨规模迁移。
LoRA-Muon: Spectral Steepest Descent on the Low-Rank Manifold

- 在低秩流形上使用谱最速下降,改进优化器更新机制。
- 秩2时恢复全参数最优学习率,秩32时验证损失低于密集基线。
- 无需QR分解、不存二阶矩,适合硬件加速且内存友好。
低秩微调(LoRA)虽显著降低计算与内存开销,但调参难度高于全参数训练:使用因子级优化器如AdamW时,对初始化敏感,最优学习率难以跨秩迁移,常无法超越密集基线。本文通过将Muon优化器的谱最速下降规则应用于低秩场景,提出LoRA-Muon。结合分拆权重衰减策略,主张其为全秩Muon及Shampoo族优化器的良好低秩代理。其最优学习率可在不同秩、宽度、深度及因子缩放间有效迁移。在计算量匹配的TinyShakespeare实验中,秩2代理成功恢复密集训练的最佳学习率;秩32的LoRA-Muon在种子平均搜索中取得低于密集基线的平均验证损失。进一步表明Spectron依赖任意因子缩放,在初始因子不平衡时表现不佳;而LoRA-RITE的简化QR坐标核心实现相同谱更新。LoRA-Muon无需QR分解且不存储二阶矩,更具加速器友好性与内存效率。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。