用于高效大语言模型训练的随机梯度子空间
机器学习
2025-10-03 v1
摘要
训练大语言模型（LLMs）常常受限于极高的内存需求，其中优化器状态占据了主要的内存开销。近期的工作通过使用复杂的更新策略将梯度投影到低维子空间来减轻这一开销。在本文中，我们分析了梯度空间及其底层子空间的动力学。我们发现，虽然一个小子空间捕获了大部分梯度能量，但仍有相当一部分存在于残余主体中；此外，核心子空间的影响随时间推移和在更深层中而减弱。我们还观察到梯度空间表现出近乎平坦的曲率，这需要明确考虑这一几何特征的算法。受这些见解的启发，我们引入了一组随机算法 GrassWalk 和 GrassJump，它们利用子空间并实现了最先进的内存节省，同时在 LLaMA-1B 和 LLaMA-7B 预训练上提升了性能。
引用
@article{arxiv.2510.01878,
title = {Randomized Gradient Subspaces for Efficient Large Language Model Training},
author = {Sahar Rajabi and Nayeema Nonta and Samanvay Vajpayee and Sirisha Rambhatla},
journal= {arXiv preprint arXiv:2510.01878},
year = {2025}
}