提出两种新优化技术，大幅降低大模型训练内存占用并提速。
Lean and Mean Adaptive Optimization via Subset-Norm and Subspace-Momentum with Convergence Guarantees
- 通过子集归一化共享步长，将内存从O(d)降至O(√d)
- 结合两项技术可减少80%内存开销，仅用一半训练数据达相同效果
- 适合资源受限下大模型预训练与微调，尤其关注内存效率的场景
我们提出两种互补的高效优化技术，可在降低内存消耗的同时加速大规模神经网络的训练。第一项技术——子集归一化步长（Subset-Norm, SN），通过步长共享推广了AdaGrad-Norm和AdaGrad(-Coordinate)，将AdaGrad的内存开销从O(d)降至O(√d)，其中d为模型规模。在坐标独立子高斯噪声下的非凸光滑目标中，证明了具有改进维度依赖性的高概率收敛性。第二项技术——子空间动量（Subspace-Momentum），将动量状态限制在低维子空间，同时在正交补空间中执行SGD，显著降低动量状态内存。在标准假设下证明了其高概率收敛性。在大语言模型预训练与微调中的实验表明，结合两项技术可在约68亿训练标记内达到Adam的验证困惑度（原需131亿），且优化器状态内存减少超过80%，几乎无需额外调参。
原文摘要 · Abstract (English)
We introduce two complementary techniques for efficient optimization that reduce memory requirements while accelerating training of large-scale neural networks. The first technique, Subset-Norm step size, generalizes AdaGrad-Norm and AdaGrad(-Coordinate) through step-size sharing. Subset-Norm (SN) reduces AdaGrad's memory footprint from $O(d)$ to $O(\sqrt{d})$, where $d$ is the model size. For non-convex smooth objectives under coordinate-wise sub-gaussian noise, we show a noise-adapted high-probability convergence guarantee with improved dimensional dependence of SN over existing methods. Our second technique, Subspace-Momentum, reduces the momentum state's memory footprint by restricting momentum to a low-dimensional subspace while performing SGD in the orthogonal complement. We prove a high-probability convergence result for Subspace-Momentum under standard assumptions. Empirical evaluation on pre-training and fine-tuning LLMs demonstrates the effectiveness of our methods. For instance, combining Subset-Norm with Subspace-Momentum achieves Adam's validation perplexity for LLaMA 1B in approximately half the training tokens (6.8B vs 13.1B) while reducing Adam's optimizer-states memory footprint by more than 80\% with minimal additional hyperparameter tuning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。