SUMO优化器通过精准正交化加速低秩LLM训练,兼顾速度与内存效率。
SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training
- 在动态低维子空间中用SVD实现精确梯度矩正交化
- 收敛速度提升,内存减少最多20%,性能更稳定
- 适合追求高效训练的LLM研究者和工程部署人员
低秩梯度优化方法显著提升了大语言模型(LLMs)训练的内存效率,可在资源受限硬件上运行而不损失性能。然而,这些方法多关注内存节省,常忽略因依赖标准各向同性最速下降法导致的收敛缓慢问题,这类方法在深度网络高度各向异性的损失曲面中表现不佳,尤其在LLMs中更为明显。本文提出SUMO(Subspace-Aware Moment-Orthogonalization),一种在动态自适应低维子空间中采用精确奇异值分解(SVD)进行矩正交化的优化器,实现范数诱导的最速下降步长。通过显式对齐优化步骤与损失曲面的谱特性,SUMO有效缓解了如Newton-Schulz近似正交化等常用方法带来的近似误差。我们理论上建立了这些近似误差的上界,证明其依赖于矩的条件数,并在理论上分析表明该条件在LLM训练中普遍存在。进一步,我们从理论和实证两方面证明,通过SVD实现的精确正交化可显著提升收敛速率并降低整体复杂度。实验验证显示,SUMO在加速收敛、提升稳定性、改善性能的同时,相比现有最优方法,内存消耗最多降低20%。
原文摘要 · Abstract (English)
Low-rank gradient-based optimization methods have significantly improved memory efficiency during the training of large language models (LLMs), enabling operations within constrained hardware without sacrificing performance. However, these methods primarily emphasize memory savings, often overlooking potential acceleration in convergence due to their reliance on standard isotropic steepest descent techniques, which can perform suboptimally in the highly anisotropic landscapes typical of deep networks, particularly LLMs. In this paper, we propose SUMO (Subspace-Aware Moment-Orthogonalization), an optimizer that employs exact singular value decomposition (SVD) for moment orthogonalization within a dynamically adapted low-dimensional subspace, enabling norm-inducing steepest descent optimization steps. By explicitly aligning optimization steps with the spectral characteristics of the loss landscape, SUMO effectively mitigates approximation errors associated with commonly used methods like Newton-Schulz orthogonalization approximation. We theoretically establish an upper bound on these approximation errors, proving their dependence on the condition numbers of moments, conditions we analytically demonstrate are encountered during LLM training. Furthermore, we both theoretically and empirically illustrate that exact orthogonalization via SVD substantially improves convergence rates while reducing overall complexity. Empirical evaluations confirm that SUMO accelerates convergence, enhances stability, improves performance, and reduces memory requirements by up to 20% compared to state-of-the-art methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。