[Paper] 驯服动量：通过低秩近似重新思考优化器状态

Source: arXiv - 2602.24283v1

概述

训练当今规模庞大的语言模型在很大程度上依赖于 Adam 及其变体等自适应优化器。虽然这些方法能够加速收敛，但它们也需要为每个参数存储一阶和二阶的“动量”向量，从而导致内存占用增加，限制了在单个 GPU 上能够训练的模型规模。论文 “Taming Momentum: Rethinking Optimizer States Through Low‑Rank Approximation” 提出了一种全新的视角，将这些动量缓冲区视为在线线性回归器的参数，并利用该视角通过低秩矩阵分解对其进行压缩。由此产生的优化器 LoRA‑Pre 在显著降低内存开销的同时，仍能在预训练和微调任务上实现业界领先的性能。

关键贡献

• EMA 动量的重新解释：表明 Adam 系列优化器中使用的指数移动平均在数学上等价于通过在线梯度流训练线性回归器。
• 低秩优化器设计（LoRA‑Pre）：将完整的动量矩阵分解为紧凑的低秩子空间，显著降低优化器的状态大小。
• 在 Llama 系列上的实证验证：在参数规模从 60 M 到 1 B 的模型上均展示出一致的性能提升，取得所有基线中最佳的结果。
• 秩效率：仅使用 1/8 的秩（即存储的参数更少）即可匹配或超越基线性能。
• 微调优势：在相同秩预算下，优于流行的高效微调方法（如标准 LoRA），在 Llama‑3.1‑8B 上提升 3.14 分，在 Llama‑2‑7B 上提升 6.17 分。
• 开源发布：完整实现和训练脚本已在 GitHub 上公开。

方法论

1. EMA 作为在线线性回归

   • 论文首先将梯度的指数移动平均（EMA）——Adam 的一阶动量的核心——表述为在线线性回归问题的解。在此视角下，每一次优化器步骤都会更新一个微小的线性模型，该模型尝试从过去的梯度预测当前梯度。

2. 动量矩阵的低秩近似

   • LoRA‑Pre 将完整的动量矩阵（尺寸 = 参数数量 × 隐藏维度）分解为两个更小的矩阵 U 和 V（秩‑r 近似）。乘积 U·Vᵀ 在运行时重建动量估计。
   • U 和 V 的更新遵循从 EMA 解释推导出的相同在线梯度流动力学，确保低秩表示能够忠实于原始优化器的动态。

3. 与现有训练流水线的集成

   • LoRA‑Pre 是 Adam/Muon 的即插即用替代品：优化器 API 保持不变，只有内部状态的处理方式不同。
   • 在预训练阶段，秩 r 被视为超参数；作者发现即使是非常低的秩（例如 1 B 模型的 r = 8）也足够。

4. 评估协议

   • 使用 Llama 系列模型（60 M、160 M、410 M、1 B）在标准语言建模语料上进行预训练。
   • 微调实验采用下游基准（如 Alpaca、MMLU），与 LoRA、QLoRA 以及其他参数高效方法进行比较。

结果与发现

• 预训练：LoRA‑Pre 在使用更少优化器状态的情况下，始终优于全内存 Adam 基线。
• 微调：在相同的低秩下，LoRA‑Pre 超过标准 LoRA 3.14 分（Llama‑3.1‑8B）和 6.17 分（Llama‑2‑7B）。
• 秩效率：实验表明，当秩超过约全维度的 1/8 时收益递减，证实大多数有用的动量信息位于低维子空间。

实际影响

• 在通用硬件上进行规模扩展：通过将优化器内存削减至原始大小的约 12 %，开发者可以在相同的 GPU 内存预算下训练更大的模型，或在单个节点上同时运行多个实验。
• 更快的迭代周期：更小的优化器状态意味着 GPU 与主机内存之间的数据传输更少，这在多 GPU 流水线中（优化器同步是瓶颈时）可以带来适度的加速。
• 参数高效的微调：LoRA‑Pre 可以与现有的 PEFT（参数高效微调）框架结合使用，提供“双重效率”提升——低秩适配器权重 加上 低秩优化器状态。
• 简化的基础设施：无需自定义检查点技巧（例如对优化器状态进行分片），因为状态大小已经非常小。这简化了在按 GPU 内存 GB 收费的云平台上的部署。
• 在其他领域的潜力：EMA‑作为在线回归器的洞见并不限于语言模型；任何使用 Adam 类优化器的大规模训练（视觉、语音、强化学习）都可以采用 LoRA‑Pre 来节省内存。

局限性与未来工作

• 秩选择仍然是启发式的：虽然论文提供了经验性的指导，但如果能够自动为每个模型或每层选择最优秩，将使该方法更加即插即用。
• 与二阶矩的兼容性：LoRA‑Pre 侧重于一阶动量（EMA）。将低秩思想扩展到 Adam 的方差项（二阶矩）可能带来进一步的节省，但这留待未来研究。
• 基准仅限于 Llama 系列：作者在参数量最高达 1 B 的 Llama 模型上进行评估；在真正大规模模型（数百亿参数）以及其他架构（如 GPT‑NeoX、T5）上进行测试，将更有力地证明其通用适用性。
• 潜在的数值稳定性问题：低秩分解可能会为非常深或更新稀疏的层引入条件数问题；论文指出当秩设得过低时会出现偶发的发散，暗示需要更稳健的防护措施。

如果你有兴趣亲自尝试 LoRA‑Pre，作者已经在 https://github.com/mrflogs/LoRA-Pre 开源了代码。将其接入你现有的 PyTorch 训练脚本，设置期望的秩，即可在不牺牲模型质量的前提下立刻看到内存节省。

作者

• Zhengbo Wang
• Jian Liang
• Ran He
• Zilei Wang
• Tieniu Tan

论文信息

• arXiv ID: 2602.24283v1
• 分类: cs.LG, cs.AI, cs.CL
• 发布日期: 2026年2月27日
• PDF: 下载 PDF