提出无偏低秩优化方法,让大模型训练更省内存且性能不降。
Unbiased Gradient Low-Rank Projection
- 通过分层采样消除低秩投影偏差,保持优化稳定性。
- 在大模型微调与预训练中优于原版GaLore,甚至超越全参数训练。
- 适合追求高效训练的大模型研究者和工业部署场景。
训练日益庞大的语言模型需要高效的内存优化策略。当前流行的方法是梯度低秩投影,仅存储投影后的优化器状态,例如GaLore。然而,这类方法普遍缺乏收敛性保证,因为低秩投影会引入相对于原始优化算法的固有偏差,导致性能差距。本文研究了分层采样技术以消除低秩投影的偏差,提出基于GaLore机制与Muon算法的新方法GUM(GaLore Unbiased with Muon)。理论证明该方法在收敛性上等价于基础Muon算法,同时保持低秩技术的内存效率。实验证明,在大模型微调与预训练任务中,GUM显著优于GaLore,甚至表现超过全参数训练。进一步分析表明,其优势源于层内知识分布更均匀,提升了参数空间利用效率与记忆能力。
原文摘要 · Abstract (English)
Memory-efficient optimization is critical for training increasingly large language models (LLMs). A popular strategy involves gradient low-rank projection, storing only the projected optimizer states, with GaLore being a representative example. However, a significant drawback of many such methods is their lack of convergence guarantees, as various low-rank projection approaches introduce inherent biases relative to the original optimization algorithms, which contribute to performance gaps compared to full-parameter training. Aiming to tackle this problem, this paper investigates the layerwise sampling technique for debiasing low-rank projection mechanisms. In particular, an instantiation of the paradigm gives rise to a novel and unbiased low-rank optimization method built upon GaLore's mechanism and the Muon algorithm, named GaLore Unbiased with Muon (GUM). We theoretically prove our method matches the convergence guarantees of the base Muon algorithm while preserving the memory efficiency of low-rank techniques. Empirical experiments on LLM fine-tuning and pretraining also demonstrate non-trivial improvements over GaLore and even better performance than full-parameter training. Further investigation shows that the improvement of this technique comes from a more uniform distribution of knowledge inside layers, leading to more efficient utilization of the model parameter space and better memorization.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。