SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training

📄 arXiv: 2505.24749v2 📥 PDF

作者: Yehonathan Refael, Guy Smorodinsky, Tom Tirer, Ofir Lindenbaum

分类: cs.LG, cs.CL, math.OC

发布日期: 2025-05-30 (更新: 2025-10-25)

期刊: The Thirty-Ninth Annual Conference on Neural Information Processing Systems (NeurIPS 2025)

💡 一句话要点

SUMO：子空间感知矩正交化加速内存高效的大语言模型训练

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 优化器 内存效率 奇异值分解 矩正交化

📋 核心要点

  1. 现有基于低秩梯度的优化方法虽然节省了LLM训练的内存，但忽略了在高度各向异性的损失地形中加速收敛的潜力。
  2. SUMO通过在动态调整的低维子空间内使用精确SVD进行矩正交化，实现了与损失地形谱特征对齐的最速下降优化。
  3. 实验结果表明，SUMO加速了收敛，提高了稳定性和性能，并降低了高达20%的内存需求，优于现有方法。

📝 摘要（中文）

本文提出了一种名为SUMO（Subspace-Aware Moment-Orthogonalization，子空间感知矩正交化）的优化器，旨在加速内存高效的大语言模型（LLM）训练。SUMO采用精确的奇异值分解（SVD）在动态调整的低维子空间内进行矩正交化，从而实现范数诱导的最速下降优化步骤。通过将优化步骤与损失地形的谱特征显式对齐，SUMO有效缓解了常用方法（如Newton-Schulz正交化近似）相关的近似误差。论文在理论上建立了这些近似误差的上界，证明了其对矩的条件数的依赖性，并分析证明了LLM训练过程中会遇到这些条件。此外，论文在理论和实验上都表明，通过SVD进行精确正交化可以显著提高收敛速度，同时降低整体复杂度。实验评估证实，与最先进的方法相比，SUMO加速了收敛，增强了稳定性，提高了性能，并将内存需求降低了高达20%。

🔬 方法详解

问题定义：现有的大语言模型训练方法，特别是基于低秩梯度的方法，虽然在内存效率方面有所提升，但往往忽略了优化过程中的收敛速度。这些方法通常采用各向同性的最速下降法，在深度网络（尤其是LLM）高度各向异性的损失地形中表现不佳。因此，如何既保证内存效率，又能加速LLM的训练收敛是一个关键问题。

核心思路：SUMO的核心思路是在一个动态调整的低维子空间内，通过精确的奇异值分解（SVD）进行矩正交化。这种方法能够更准确地估计损失地形的曲率信息，从而实现范数诱导的最速下降优化步骤。通过将优化步骤与损失地形的谱特征对齐，可以有效地缓解传统近似方法（如Newton-Schulz正交化）带来的误差，从而加速收敛。

技术框架：SUMO优化器的整体框架包括以下几个主要步骤：1) 计算梯度和动量；2) 动态选择一个低维子空间；3) 在该子空间内进行矩正交化，具体是通过SVD分解动量矩阵；4) 使用正交化后的动量更新模型参数。关键在于子空间的动态选择和SVD分解，前者保证了计算效率，后者保证了正交化的精度。

关键创新：SUMO最重要的技术创新在于使用精确的SVD进行矩正交化，而不是采用近似方法。与Newton-Schulz等近似方法相比，SVD能够更准确地估计损失地形的曲率信息，从而避免了因近似带来的误差。此外，SUMO还提出了一种动态调整子空间的方法，以在计算效率和正交化精度之间取得平衡。

关键设计：SUMO的关键设计包括：1) 子空间的选择策略，需要根据梯度和动量的变化动态调整子空间的大小和方向；2) SVD分解的实现，需要考虑计算效率和数值稳定性；3) 学习率的调整策略，需要根据正交化后的动量进行调整，以保证优化过程的稳定性和收敛速度。论文中还分析了矩的条件数对近似误差的影响，并给出了相应的理论保证。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SUMO在加速收敛、提高稳定性和性能方面优于现有方法。具体来说，SUMO能够将内存需求降低高达20%，同时在各种LLM训练任务中实现了更快的收敛速度和更高的模型精度。这些结果验证了SUMO在实际应用中的有效性。

🎯 应用场景

SUMO优化器可以广泛应用于各种需要高效训练的大语言模型场景，尤其是在资源受限的硬件环境下。其潜在应用领域包括自然语言处理、机器翻译、文本生成等。通过加速模型训练和降低内存需求，SUMO能够帮助研究人员和工程师更高效地开发和部署LLM，从而推动人工智能技术的发展。

📄 摘要（原文）

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.