[ICLR 2025] LoRA-Pre：重新定义动量，开启大模型全参数预训练的“降维打击”

TL;DR

在 LLM 预训练领域，显存瓶颈往往不在于权重本身，而在于优化器的“负重”。本文提出的 LoRA-Pre 巧妙地将优化器的 EMA 动量更新视为一场“在线线性回归”，通过低秩矩阵实时拟合动量，仅需 1/8 的秩（Rank）便能刷新显存效率与性能的双双上限。

背景定位：优化器状态的“脂肪”

现代深层模型依赖于 Adam 这种带有“记忆”（动量）的优化器。这意味着除了存储数千亿个参数，你还得存两份同样大小的梯度历史（一阶和二阶动量）。对于单卡显存而言，这是难以承受之重。虽然之前的 GaLore 或 Flora 采用了梯度投影（Projection）技术来“减肥”，但它们通常是每隔几百步更新一次投影子空间。这种离散型、滞后性的更新窗口会带来严重的轨迹漂移（Error Accumulation）。

核心直觉：你的动量其实是个“潜伏”的回归器

作者揭示了一个令人拍案叫绝的物理直觉： $m\leftarrow \beta \cdot m+(1-\beta )\cdot g$ 这个公式在数学上等价于最小化如下目标函数的梯度下降步骤： ${\min }_m\frac{1}{2}\Vert m-g{\Vert }_F^2$ 这意味着，动量的本质就是通过不断拟合最新的梯度流，来找到一个能够代表历史梯度的最佳“权重”。既然它是线性模型的权重，我们为什么不直接用 LoRA (Low-Rank Adaptation) 的思路去分解它呢？

方法论详解：LoRA-Pre 的手术刀

LoRA-Pre 将全秩动量 $m$ 分构为 $m_B\cdot m_A$。为了避免传统反向传播更新优化器状态带来的额外开销，作者利用**牛顿法（Newton's Method）**推导出了闭式更新规则。

关键技术点：

1. 二阶动量的 Hadamard 平方根处理：为保证二阶动量 $v$ 的元素正定性（因为 Adam 更新需要 $\sqrt{v}$），作者将其参数化为 $(v_B{v}_A{)}^{\circ 2}$，确保了数学计算的严谨性。
2. 连续子空间演化：不同于 GaLore 的周期性 SVD 投影，LoRA-Pre 的子空间在每一步迭代中都在“自我进化”，这极大地减小了动量近似的残差。

实验与结果：小秩办大事

在 Llama 架构的系列模型实验中，LoRA-Pre 表现出了极其恐怖的 秩效率 (Rank Efficiency)。

• 性能反超：在 1B 模型上，LoRA-Pre Adam 的困惑度（13.53）不仅远超 GaLore（15.64），甚至在特定设置下逼近了全参数 Muon 优化器的表现。
• 极低配置也能打：实验显示，LoRA-Pre 在 Rank=16 时的表现竟然比 GaLore 在 Rank=256 时还要好。这意味着你可以用极少的内存开销，换取极高的搜索空间覆盖。

从上图可以看到，即使是极小的 Rank，LoRA-Pre 的收敛轨迹依然非常稳健，这得益于其在线更新机制对梯度流的高保真捕获。

深度洞察与总结

核心价值（Takeaway）： LoRA-Pre 的成功在于它打破了“低秩优化 = 梯度空间投影”的思维定势，转而从“动量即回归”的视角，将静态的投影变成了动态的回归任务。

局限性与未来展望： 虽然预训练表现极其抢眼，但对于极端的超参数（如 $\beta o1$）组合，低秩因子的适应速度可能会跟不上梯度的变化，导致数值不稳定。未来的研究重点可能会放在如何自适应地调整回归器的学习率（$\gamma$），以应对训练后期极其稀疏或剧烈波动的梯度场景。

对于正在苦练大模型、深受显存之苦的开发者来说，LoRA-Pre 提供了一种无需折中性能、更加优雅的“瘦身”方案。