Low-rank Momentum Factorization for Memory Efficient Training

📄 arXiv: 2507.08091v1 📥 PDF

作者: Pouria Mahdavinia, Mehrdad Mahdavi

分类: cs.LG

发布日期: 2025-07-10

🔗 代码/项目: GITHUB

💡 一句话要点

提出MoFaSGD以解决大模型微调中的内存效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 内存效率 优化算法 低秩分解 深度学习 模型微调 动态更新 谱归一化 非凸优化

📋 核心要点

  1. 现有方法在微调大型模型时面临内存消耗过大的问题，尤其是使用状态优化器时。
  2. 论文提出MoFaSGD，通过动态维护低秩动量表示，解决了内存效率与优化性能之间的矛盾。
  3. 实验表明，MoFaSGD在大型语言模型对齐任务中，内存减少效果与LoRA相当，同时性能优于其他低秩优化方法。

📝 摘要（中文）

微调大型基础模型面临显著的内存挑战，尤其是使用状态优化器（如AdamW）时，所需GPU内存往往是推理时的数倍。尽管存在如LoRA等内存高效方法，但这些方法在固定子空间或计算成本高的离线重采样方面存在不足。本文提出的Momentum Factorized SGD（MoFaSGD）通过动态更新低秩SVD表示，近似全秩动量，提供了一种内存高效的微调方法。MoFaSGD在每次迭代中自适应更新优化子空间，并在理论上证明其在非凸随机优化中的收敛性。实验结果显示，MoFaSGD在大型语言模型对齐基准上表现出色，内存减少效果与LoRA相当，同时在性能上与现有低秩优化方法竞争。

🔬 方法详解

问题定义：本文旨在解决在微调大型基础模型时，由于使用状态优化器（如AdamW）导致的内存消耗过大的问题。现有方法如GaLore在固定子空间或计算成本高的离线重采样方面存在不足，限制了其应用。

核心思路：论文提出的MoFaSGD通过动态更新低秩SVD表示的动量，近似全秩动量，从而实现内存高效的优化。该方法在每次迭代中自适应更新优化子空间，避免了固定子空间的限制。

技术框架：MoFaSGD的整体架构包括低秩动量的计算、动态更新机制以及谱归一化更新。主要模块包括动量因子的计算和更新、优化子空间的调整等。

关键创新：MoFaSGD的核心创新在于其动态低秩动量表示的维护，与现有方法相比，避免了固定子空间的限制，并且减少了计算成本。

关键设计：在参数设置上，MoFaSGD采用了低秩SVD来表示动量，并设计了谱归一化更新机制。损失函数和网络结构的具体细节未在摘要中明确，但可以推测其与优化过程密切相关。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，MoFaSGD在大型语言模型对齐基准上表现优异，内存减少效果与LoRA相当，同时在性能上与现有低秩优化方法相比具有竞争力，证明了其在实际应用中的有效性。

🎯 应用场景

该研究的潜在应用领域包括大型语言模型的微调、深度学习模型的优化以及资源受限环境下的模型训练。通过提高内存效率，MoFaSGD可以使得更多研究者和开发者能够在有限的硬件条件下进行高效的模型训练，推动AI技术的普及和应用。

📄 摘要（原文）

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.