arXiv:2602.24283cs.LGcs.AI2026-02中稿 · ICLR被引 3

用低秩分解优化器状态,大幅降低内存占用却保持性能

Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

  • 将动量机制视为在线线性回归,通过低秩近似压缩状态空间
  • 在60M到1B参数模型上均表现最佳,仅需1/8基线秩即达同等效果
  • 适用于预训练与微调,显著优于标准LoRA等高效方法

现代优化器如Adam和Muon在训练大语言模型中至关重要,但其对一阶、二阶动量的依赖带来显著内存开销,制约可扩展性与计算效率。本文将这些动量中的指数移动平均(EMA)重构为在线梯度流下的线性回归训练过程。基于此等价性,提出新型低秩优化器LoRA-Pre,用于高效预训练。具体地,通过在在线线性学习器中将完整动量矩阵分解至紧凑低秩子空间,实现内存效率提升,同时保持优化性能。我们通过在Llama系列模型上从60M到1B参数的预训练实证验证了其有效性:所有规模下均达到最高性能。尤为关键的是,LoRA-Pre展现出卓越的秩效率,仅需基线1/8的秩即可获得相当或更优结果。此外,在微调场景中,相同秩下,LoRA-Pre持续优于所有高效微调基线。相比标准LoRA,其在Llama-3.1-8B上提升3.14点,在Llama-2-7B上提升6.17点,验证了该方法在预训练与微调中的普适有效性。代码已公开于https://github.com/mrflogs/LoRA-Pre。

原文摘要 · Abstract (English)

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.

优化器低秩内存效率微调

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。