驯服动量：通过低秩近似重新思考优化器状态
机器学习
2026-03-02 v1 人工智能
计算与语言
摘要
像 Adam 和 Muon 这样的现代优化器是训练大型语言模型的核心，但它们对一阶和二阶动量的依赖引入了显著的内存开销，这限制了可扩展性和计算效率。在这项工作中，我们将这些动量中使用的指数移动平均 (EMA) 重新解释为通过在线梯度流训练一个线性回归器。基于这种等价性，我们引入了 LoRA-Pre，一种专为高效预训练设计的新型低秩优化器。具体来说，LoRA-Pre 通过将完整的动量矩阵分解为在线线性学习器内的一个紧凑低秩子空间来减少优化器的内存占用，从而在提高内存效率的同时保持优化性能。我们通过预训练从 60M 到 1B 参数的 Llama 架构系列模型来实证验证 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 个点的提升，验证了我们的方法在预训练和微调范式中的有效性。我们的代码可在 https://github.com/mrflogs/LoRA-Pre 公开获取。
引用
@article{arxiv.2602.24283,
title = {Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation},
author = {Zhengbo Wang and Jian Liang and Ran He and Zilei Wang and Tieniu Tan},
journal= {arXiv preprint arXiv:2602.24283},
year = {2026}
}
备注
Camera-ready version. Accepted as Oral at ICLR 2026