低秩动量分解实现内存高效训练
机器学习
2025-07-14 v1
摘要
微调大型基础模型因状态优化器（如 AdamW）带来显著内存挑战，往往需要比推理多数倍的 GPU 内存。虽然参数高效微调（如 LoRA）和优化器状态 压缩等内存高效方法已存在，但近期方法如 GaLore 通过低秩梯度投影和子空间动量累积来弥合这两者。然而，这些方法可能在固定子空间或计算密集的离线重抽样（如需要完整矩阵 SVD）方面存在困难。我们提出动量分解 SGD(MoFaSGD)，通过保持对一阶动量的低秩 SVD 表示而实现内存高效，该表示在训练期间贴近其完整秩版本。这种分解使我们能够实现一种自适应更新优化子空间的内存高效微调方法。关键的是，MoFaSGD 利用计算的低秩动量因子执行有效的 spectrally 归一化更新，提供了子空间动量累积的替代方案。我们建立了 MoFaSGD 的理论收敛保证，证明其在标准假设下实现非凸随机优化的最优收敛率。在大语言模型对齐基准测试中，我们证明了 MoFaSGD 的有效性，在内存降低（相当于 LoRA）与性能之间实现了有竞争力的平衡。我们的实现已发布于 https://github.com/pmahdavi/MoFaSGD。
引用
@article{arxiv.2507.08091,
title = {Low-rank Momentum Factorization for Memory Efficient Training},
author = {Pouria Mahdavinia and Mehrdad Mahdavi},
journal= {arXiv preprint arXiv:2507.08091},
year = {2025}
}