Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

📄 arXiv: 2602.24283v1 📥 PDF

作者: Zhengbo Wang, Jian Liang, Ran He, Zilei Wang, Tieniu Tan

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-02-27

备注: Camera-ready version. Accepted as Oral at ICLR 2026

🔗 代码/项目: GITHUB

💡 一句话要点

LoRA-Pre：通过低秩近似优化器状态，提升大模型预训练和微调效率

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

关键词: 低秩近似 优化器 大语言模型 预训练 微调 内存效率 在线学习

📋 核心要点

  1. 现有优化器（如Adam）在大模型训练中内存开销大，限制了模型规模和训练效率。
  2. LoRA-Pre将动量估计视为在线线性回归，并利用低秩分解降低动量矩阵的存储需求。
  3. 实验表明，LoRA-Pre在预训练和微调任务中均优于现有方法，且秩效率显著提升。

📝 摘要（中文）

现代优化器如Adam和Muon在大语言模型训练中至关重要，但它们对一阶和二阶动量的依赖带来了显著的内存开销，限制了可扩展性和计算效率。本文将这些动量中使用的指数移动平均（EMA）重新定义为通过在线梯度流训练线性回归器。基于这种等价性，我们提出了一种新颖的低秩优化器LoRA-Pre，专为高效预训练而设计。具体来说，LoRA-Pre通过将完整的动量矩阵分解为在线线性学习器中的紧凑低秩子空间，从而减少优化器的内存占用，在提高内存效率的同时保持优化性能。我们通过预训练Llama架构系列的60M到1B参数模型，从经验上验证了LoRA-Pre的有效性。LoRA-Pre在所有模型尺寸上都实现了最高的性能。值得注意的是，LoRA-Pre展示了卓越的秩效率，仅使用基线方法1/8的秩即可实现相当或更好的结果。除了预训练，我们还评估了LoRA-Pre在微调场景中的有效性。在相同秩的情况下，LoRA-Pre始终优于所有高效微调基线。具体而言，与标准LoRA相比，LoRA-Pre在Llama-3.1-8B上实现了3.14个点的显著改进，在Llama-2-7B上实现了6.17个点的显著改进，验证了我们的方法在预训练和微调范例中的有效性。

🔬 方法详解

问题定义：现有优化器，如Adam，在训练大型语言模型时，需要存储一阶和二阶动量，导致显著的内存开销。这限制了模型可以扩展到的规模，并降低了计算效率。因此，如何在保持优化性能的同时，降低优化器的内存占用，是一个关键问题。

核心思路：论文的核心思路是将优化器中的指数移动平均（EMA）过程重新解释为在线梯度流训练线性回归器。基于此，通过对动量矩阵进行低秩分解，将其投影到一个低维子空间，从而显著减少需要存储的参数数量，降低内存占用。

技术框架：LoRA-Pre的核心在于将动量估计过程视为一个在线线性学习问题。具体来说，它维护一个低秩矩阵来近似完整的动量矩阵。在每次迭代中，LoRA-Pre首先使用梯度更新低秩矩阵，然后使用更新后的低秩矩阵来更新模型参数。整体流程包括：1) 将动量矩阵分解为低秩形式；2) 在线更新低秩矩阵；3) 使用低秩矩阵更新模型参数。

关键创新：LoRA-Pre的关键创新在于将动量估计与在线线性学习联系起来，并利用低秩近似来压缩动量矩阵。与传统的优化器相比，LoRA-Pre不需要存储完整的动量矩阵，从而显著降低了内存占用。此外，LoRA-Pre的低秩结构允许它在保持优化性能的同时，实现更高的秩效率。

关键设计：LoRA-Pre的关键设计包括：1) 使用随机梯度下降（SGD）或其变体来更新低秩矩阵；2) 选择合适的秩的大小，以平衡内存占用和优化性能；3) 使用适当的初始化策略来初始化低秩矩阵。损失函数通常是标准的交叉熵损失或均方误差损失，具体取决于任务类型。网络结构与原始模型保持一致，LoRA-Pre主要作用于优化器层面。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LoRA-Pre在Llama架构系列的预训练中表现出色，在所有模型尺寸上均达到最高性能，且秩效率显著提升，仅使用基线方法1/8的秩即可实现相当或更好的结果。在微调任务中，LoRA-Pre也优于现有方法，例如，在Llama-3.1-8B上相比标准LoRA提升了3.14个点，在Llama-2-7B上提升了6.17个点。

🎯 应用场景

LoRA-Pre适用于各种需要高效训练大型模型的场景，例如自然语言处理、计算机视觉等。它可以降低训练成本，加速模型迭代，并支持在资源受限的环境中进行模型训练。未来，LoRA-Pre可以进一步扩展到其他类型的优化器和模型架构，并应用于更广泛的实际问题。

📄 摘要（原文）

Modern optimizers like Adam and Muon are central to training large language models, but their reliance on first- and second-order momenta introduces significant memory overhead, which constrains scalability and computational efficiency. In this work, we reframe the exponential moving average (EMA) used in these momenta as the training of a linear regressor via online gradient flow. Building on this equivalence, we introduce LoRA-Pre, a novel low-rank optimizer designed for efficient pre-training. Specifically, LoRA-Pre reduces the optimizer's memory footprint by decomposing the full momentum matrix into a compact low-rank subspace within the online linear learner, thereby maintaining optimization performance while improving memory efficiency. We empirically validate LoRA-Pre's efficacy by pre-training models from the Llama architecture family, scaling from 60M to 1B parameters. LoRA-Pre achieves the highest performance across all model sizes. Notably, LoRA-Pre demonstrates remarkable rank efficiency, achieving comparable or superior results using only 1/8 the rank of baseline methods. Beyond pre-training, we evaluate LoRA-Pre's effectiveness in fine-tuning scenarios. With the same rank, LoRA-Pre consistently outperforms all efficient fine-tuning baselines. Specifically, compared to standard LoRA, LoRA-Pre achieves substantial improvements of 3.14 points on Llama-3.1-8B and 6.17 points on Llama-2-7B, validating our approach's effectiveness across both pre-training and fine-tuning paradigms. Our code is publicly available at https://github.com/mrflogs/LoRA-Pre.