提出最优低秩梯度估计，显著降低大模型训练内存占用。
Optimal low-rank stochastic gradient estimation for LLM training
- 将梯度投影至随机低维子空间再还原，保持无偏且方差最小。
- 罗伯特大模型微调中峰值显存仅3.83GB，远低于全反向传播的16.7GB。
- 适用于主流训练范式，适合资源受限的大模型训练场景。
大规模语言模型（LLM）训练常受高维参数空间中内存限制和随机梯度噪声的制约。基于实证发现训练过程中许多LLM梯度矩阵实质上是低秩的，本文提出一种无偏、内存高效、方差最低的通用低秩梯度估计器。核心思想是将高维随机梯度估计器投影到随机低维子空间并回映，通过最优设计的投影分布（包括Haar–Stiefel投影）控制均方误差，同时减少内存开销。该投影分布由约束泛函优化问题推导得出，指导算法设计。实验表明，所提低秩梯度估计器在实际应用中兼具内存节省与更优训练表现：在RoBERTa-large微调中，峰值GPU内存仅为3.83GB（相比全反向传播的16.7GB），且精度相当；在自回归LLM预训练（LLaMA-20M/60M/100M）中优于传统方法，验证了最优投影策略的有效性。
原文摘要 · Abstract (English)
Large language model (LLM) training is often bottlenecked by memory constraints and stochastic gradient noise in extremely high-dimensional parameter spaces. Motivated by empirical evidence that many LLM gradient matrices are effectively low-rank during training, we present an unbiased, memory-efficient, low-rank matrix estimator with the lowest variance that is applicable across common stochastic gradient estimation paradigms. The core idea is to project a high-dimensional stochastic gradient estimator onto a random low-dimensional subspace and lift it back, reducing memory while keeping the estimator unbiased and controlling mean-squared error via an optimally designed projection distribution, including Haar--Stiefel projections. The projection distribution is derived by solving a constrained functional optimization problem, yielding an optimal random projector that guides algorithm design. Empirically, the resulting low-rank gradient estimators deliver both practical memory savings and improved training behavior. In RoBERTa-large fine-tuning, our method attains the lowest peak GPU memory among compared methods (e.g., 3.83GB versus 16.7GB for full BP) while remaining competitive in accuracy; in autoregressive LLM pretraining (LLaMA-20M/60M/100M), our method outperforms the traditional methods, supporting the benefit of the proposed optimal projection strategy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。