Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

本文提出了名为 LoRA-Pre 的新型低秩优化器，通过将动量矩阵分解为紧凑的低秩子空间来显著降低内存开销，并在 Llama 架构的预训练和微调任务中实现了超越现有高效基线方法的性能表现。

要点

这篇论文提出了一种名为 LoRA-Pre 的新方法，旨在解决训练大型人工智能模型（如 Llama 系列）时“太费内存”的痛点。

为了让你轻松理解，我们可以把训练 AI 模型想象成一位正在学习驾驶新车的司机。

1. 核心问题：司机的“记性”太占地方

在训练 AI 时，优化器（Optimizer，比如常用的 Adam）就像司机的大脑。为了开得稳，司机不能只看眼前的路况（当前的梯度），还需要记住过去的经验（动量 Momentum）。

• 传统做法（Adam/Muon）： 司机的大脑里有一个巨大的“记忆本”，详细记录了每一秒、每一个方向上的过去经验。这个本子非常厚，随着车开得越久（模型越大），本子就越厚，最后把司机的背包（显存/内存）塞爆了，导致他跑不动了，或者只能开很小的车。
• 现状： 现在的 AI 模型越来越大，这个“记忆本”占用的内存甚至超过了模型本身，成了最大的瓶颈。

2. 核心洞察：记忆其实是个“线性回归”

作者发现了一个有趣的数学秘密：“动量”的更新方式，本质上和“在线线性回归”是一回事。

• 通俗比喻：
  • 传统的动量更新像是在说：“记住过去，再结合现在，取个平均值。”
  • 作者发现，这其实等同于：司机在实时训练一个“预测器”。他在问自己：“如果我要预测未来的路况，我该怎么调整我的记忆策略，才能最准确地拟合刚才走过的路？”
  • 这就把“死记硬背”变成了“学习如何预测”。

3. 解决方案：LoRA-Pre（低秩压缩）

既然动量是一个“预测器”，我们就不需要把整个巨大的“记忆本”都存下来。我们可以用两张薄薄的卡片（低秩矩阵）来代表这个记忆本。

• 比喻：
  • 以前，司机要背下整本《世界地图》（全秩矩阵，$p \times q$），这太重了。
  • 现在，LoRA-Pre 告诉司机：“你不需要背整本地图。你只需要记住两个关键点：‘主要方向’（矩阵 A）和‘主要距离’（矩阵 B）。只要把这两个卡片拼起来，就能还原出 99% 的地图信息。”
  • 效果： 内存占用从“一本厚书”瞬间变成了“两张小卡片”。

4. 为什么它比以前的方法更好？

以前的低秩方法（比如 GaLore）有点像**“定期更新地图”**：

• 它们每隔一段时间才重新计算一次“主要方向”，中间这段时间，如果路况变了（梯度变了），司机还在用旧地图，就会走弯路，产生误差积累。

LoRA-Pre 的做法是“实时微调”：

• 它利用刚才提到的“线性回归”原理，每一步都在微调那两张小卡片。
• 比喻： 就像司机手里拿着一个智能导航仪，每走一步，导航仪就自动微调路线，而不是等开了一公里才重新规划。这样，即使卡片很薄（秩很低），也能紧紧跟上路况的变化，不会走偏。

5. 实验结果：小身材，大能量

作者在 Llama 模型（从 6000 万参数到 10 亿参数）上进行了测试：

• 预训练（从头学）： LoRA-Pre 在几乎所有模型大小上都取得了最好的成绩，甚至超过了那些占用巨大内存的传统优化器。
• 微调（学新技能）： 在让大模型学习数学题时，LoRA-Pre 的表现也吊打其他省内存的方法。
• 惊人的效率： 最酷的是，LoRA-Pre 只需要使用其他方法 1/8 甚至更少 的“卡片厚度”（Rank），就能达到同样的效果。

总结

LoRA-Pre 就像给 AI 司机换了一个**“智能压缩记忆法”：
它不再死记硬背庞大的历史数据，而是学会了用两张小卡片实时预测路况**。这不仅让司机（AI）能背起更重的背包（训练更大的模型），还跑得更快、更稳，而且省下的内存空间足以让它在更便宜的电脑上运行。

一句话总结： 把笨重的“死记硬背”变成了灵动的“实时预测”，用极少的内存训练出最强的 AI。

技术摘要

这是一篇发表于 ICLR 2026 的会议论文，题为 《Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation》（驯服动量：通过低秩近似重新思考优化器状态）。该论文提出了一种名为 LoRA-Pre 的新型低秩优化器，旨在解决大语言模型（LLM）预训练和微调过程中优化器状态带来的巨大内存开销问题。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

• 内存瓶颈： 现代大语言模型（如 Llama 系列）的训练依赖于 Adam 或 Muon 等优化器。这些优化器需要维护一阶和二阶动量估计（momentum estimates），这导致优化器状态占用的内存是模型参数本身的 2-3 倍，严重限制了模型的可扩展性和计算效率。
• 现有方法的局限性：
  • 投影类方法（如 GaLore）： 通过 SVD 或随机投影将梯度压缩到低秩子空间。这类方法通常需要周期性地更新子空间，导致子空间与梯度结构之间存在滞后，产生误差累积，影响优化性能。
  • LoRA 类方法： 传统的 LoRA 主要用于微调（Fine-tuning），假设权重更新具有低秩结构。但在从头预训练（Pre-training from scratch）时，权重更新通常是满秩的，直接应用 LoRA 会导致性能显著下降。
• 核心挑战： 如何在大幅减少优化器状态内存占用的同时，保持甚至提升优化器的收敛性能和稳定性，且适用于预训练和微调场景。

2. 核心方法论 (Methodology)

论文的核心创新在于建立了一个新的理论视角，并据此设计了 LoRA-Pre 算法。

2.1 理论洞察：动量即在线线性回归

作者发现，现代优化器中的指数移动平均（EMA）动量更新在数学上等价于通过在线梯度流训练一个线性回归器。

• 传统 EMA 更新： $m_{t+1} = \beta \cdot m_t + (1-\beta) \cdot g_t$
• 等价形式： 这可以重写为对目标函数 $L(m; g) = \frac{1}{2} \|m - g\|_F^2$ 进行梯度下降更新，其中 $m$ 被视为线性回归器的参数，$g$ 是输入数据。
• 意义： 既然动量 $m$ 本质上是一个线性模型的权重，那么就可以像压缩模型参数一样，通过低秩分解来压缩动量矩阵。

2.2 LoRA-Pre 算法设计

基于上述洞察，LoRA-Pre 将完整的动量矩阵 $m \in \mathbb{R}^{p \times q}$ 分解为两个低秩矩阵的乘积：$m = m_B \cdot m_A$，其中 $r \ll \min(p, q)$。

• 一阶动量压缩： 优化目标变为 $\min_{m_B, m_A} \frac{1}{2} \|m_B m_A - g\|_F^2$。
  • 作者利用牛顿法（Newton's Method）推导出了闭式更新规则（Closed-form update rules），无需反向传播即可更新 $m_B$ 和 $m_A$。
  • 更新规则保持了 EMA 的形式：$m_B \leftarrow (1-\gamma_1)m_B + \gamma_1 \cdot g m_A^T (m_A m_A^T)^{-1}$ 等。
• 二阶动量压缩（针对 Adam）： 由于二阶动量 $v$ 需要元素级开方（$\sqrt{v}$），必须保证元素非负。
  • 作者采用重参数化策略：$v = (v_B \circ v_A)^{\circ 2}$（Hadamard 积的平方）。
  • 优化目标变为拟合梯度幅值 $|g|$，确保分解后的矩阵元素非负。
• 兼容性： 该方法不仅适用于 Adam，还扩展到了 Muon 优化器（通过调整目标函数中的线性项），证明了其通用性。

3. 主要贡献 (Key Contributions)

1. 理论突破： 首次建立了 EMA 动量更新与在线线性回归之间的数学等价性，为优化器状态的低秩压缩提供了坚实的理论基础。
2. 算法创新： 提出了 LoRA-Pre，一种专为预训练设计的内存高效低秩优化器。它通过低秩分解压缩动量状态，并推导出了无需反向传播的闭式更新规则。
3. 广泛适用性： 构建了基于 Adam 和 Muon 的 LoRA-Pre 变体，证明了该方法在不同优化器架构下的有效性。
4. 实证验证： 在预训练（从 60M 到 1B 参数的 Llama 模型）和微调（Llama-2-7B, Llama-3.1-8B）任务上进行了全面验证。

4. 实验结果 (Results)

• 预训练性能：
  • 在 C4 数据集上预训练 Llama 系列模型（60M, 130M, 350M, 1B）。
  • LoRA-Pre 在所有模型规模下均取得了最优或次优的验证困惑度（Perplexity）。
  • 相比之前的最佳低秩基线（如 Fira），LoRA-Pre 在 130M、350M 和 1B 模型上分别提升了 0.81、2.45 和 1.6 个困惑度点。
  • 秩效率（Rank Efficiency）： 仅需基线方法 1/8 的秩（Rank）即可达到同等甚至更好的性能（例如，LoRA-Pre 在 Rank=16 时达到了 GaLore 在 Rank=128 时的效果）。
• 微调性能：
  • 在 MetaMath100k 数据集上微调 Llama-2-7B 和 Llama-3.1-8B，并在 GSM8K 和 MATH-500 上评估。
  • 相比标准 LoRA，LoRA-Pre 在 Llama-3.1-8B 上提升了 3.14 分，在 Llama-2-7B 上提升了 6.17 分。
  • 在 Muon 优化器下，LoRA-Pre 同样显著优于其他高效微调基线。
• 消融实验：
  • 证明了 LoRA-Pre 的连续子空间更新机制（Continuous Subspace Adaptation）优于周期性更新机制，能有效避免误差累积。
  • 展示了该方法在不同秩设置下的鲁棒性。

5. 意义与影响 (Significance)

• 打破预训练与微调的界限： 现有的低秩方法（如 LoRA）通常仅适用于微调，而 GaLore 等预训练方法存在性能瓶颈。LoRA-Pre 成功统一了这两个场景，证明了低秩优化器状态压缩同样适用于从零开始的预训练。
• 内存效率的飞跃： 通过减少优化器状态的内存占用（从 $O(p \times q)$ 降至 $O((p+q) \times r)$），使得在消费级显卡或有限显存环境下训练更大规模的模型成为可能。
• 理论指导实践： 将优化器状态视为可学习的线性模型，为未来设计更高效的优化器提供了新的范式（即“优化器即模型”）。
• 开源贡献： 代码已公开，推动了社区对高效 LLM 训练的研究。

总结：
LoRA-Pre 通过深刻的理论洞察（动量即回归），巧妙地将优化器状态压缩问题转化为低秩线性回归问题。它不仅解决了 LLM 训练中的内存瓶颈，还在性能上超越了现有的低秩优化方法，特别是在低秩设置下表现出惊人的效率，是迈向更高效、更可扩展的大模型训练的重要一步。