arXiv:2507.08091cs.LG2025-07被引 5

用动态低秩分解优化动量,大幅降低大模型微调显存占用。

Low-rank Momentum Factorization for Memory Efficient Training

  • 通过动态低秩SVD表示一阶动量,实时更新优化子空间。
  • 显存占用与LoRA相当,大语言模型对齐任务性能接近顶尖方法。
  • 理论证明收敛性最优,适合资源受限下的大模型高效微调。

大模型微调因AdamW等状态优化器导致显存需求远超推理,常需数倍于推理的显存。现有方法如参数高效微调(如LoRA)和优化器状态压缩虽有效,但近期方法如GaLore依赖固定子空间或高成本离线重采样(如全矩阵SVD)。本文提出动量因子化SGD(MoFaSGD),以动态更新的低秩SVD表示一阶动量,近似全秩版本并自适应调整优化子空间。该因子化结构支持高效的谱归一化更新,替代子空间动量累积。我们建立了理论收敛性证明,在标准假设下实现非凸随机优化的最优收敛率。实验表明,MoFaSGD在大语言模型对齐基准上表现优异,显存节省程度媲美LoRA,性能优于当前主流低秩优化方法。代码已开源:https://github.com/pmahdavi/MoFaSGD。

原文摘要 · Abstract (English)

Fine-tuning large foundation models presents significant memory challenges due to stateful optimizers like AdamW, often requiring several times more GPU memory than inference. While memory-efficient methods like parameter-efficient fine-tuning (e.g., LoRA) and optimizer state compression exist, recent approaches like GaLore bridge these by using low-rank gradient projections and subspace moment accumulation. However, such methods may struggle with fixed subspaces or computationally costly offline resampling (e.g., requiring full-matrix SVDs). We propose Momentum Factorized SGD (MoFaSGD), which maintains a dynamically updated low-rank SVD representation of the first-order momentum, closely approximating its full-rank counterpart throughout training. This factorization enables a memory-efficient fine-tuning method that adaptively updates the optimization subspace at each iteration. Crucially, MoFaSGD leverages the computed low-rank momentum factors to perform efficient spectrally normalized updates, offering an alternative to subspace moment accumulation. We establish theoretical convergence guarantees for MoFaSGD, proving it achieves an optimal rate for non-convex stochastic optimization under standard assumptions. Empirically, we demonstrate MoFaSGD's effectiveness on large language model alignment benchmarks, achieving a competitive trade-off between memory reduction (comparable to LoRA) and performance compared to state-of-the-art low-rank optimization methods. Our implementation is available at https://github.com/pmahdavi/MoFaSGD.

大模型微调低秩优化显存效率

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。