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

• URL: http://arxiv.org/abs/2505.24749v1
• Date: Fri, 30 May 2025 16:08:40 GMT
• Title: SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training
• Authors: Yehonathan Refael, Guy Smorodinsky, Tom Tirer, Ofir Lindenbaum,
• Abstract summary: Low-rank gradient-based optimization methods have significantly improved memory efficiency during the training of large language models (LLMs)<n>These methods primarily emphasize memory savings, often overlooking potential acceleration in convergence.<n>In this paper, we propose SUMO (Subspace-Aware Moment-Orthogonalization), an norm that employs exact singular value decomposition.<n>We show that SUMO accelerates convergence, enhances stability, improves performance, and reduces memory requirements by up to 20% compared to state-of-the-art methods.
• Score: 13.180761892449736
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Low-rank Momentum Factorization for Memory Efficient Training [13.464518325870444]
  Momentum Factorized (MoFaSGD) maintains a dynamically updated low-rank SVD representation of the first-order momentum.<n>We demonstrate MoFaSGD's effectiveness on large language model benchmarks, achieving a competitive trade-off between memory reduction (e.g. LoRA) and performance.
  arXiv  Detail & Related papers  (2025-07-10T18:04:52Z)
• Divergence Minimization Preference Optimization for Diffusion Model Alignment [58.651951388346525]
  Divergence Minimization Preference Optimization (DMPO) is a principled method for aligning diffusion models by minimizing reverse KL divergence.<n>Our results show that diffusion models fine-tuned with DMPO can consistently outperform or match existing techniques.<n>DMPO unlocks a robust and elegant pathway for preference alignment, bridging principled theory with practical performance in diffusion models.
  arXiv  Detail & Related papers  (2025-07-10T07:57:30Z)
• Efficient Differentiable Approximation of Generalized Low-rank Regularization [64.73416824444328]
  Low-rank regularization (LRR) has been widely applied in various machine learning tasks.<n>In this paper, we propose an efficient differentiable approximation of LRR.
  arXiv  Detail & Related papers  (2025-05-21T11:49:17Z)
• Architect Your Landscape Approach (AYLA) for Optimizations in Deep Learning [0.0]
  Gradient Descent (DSG) and its variants, such as ADAM, are foundational to deep learning optimization.<n>This paper introduces AYLA, a novel optimization technique that enhances adaptability and efficiency rates.
  arXiv  Detail & Related papers  (2025-04-02T16:31:39Z)
• I3S: Importance Sampling Subspace Selection for Low-Rank Optimization in LLM Pretraining [50.89661053183944]
  Low-rank optimization has emerged as a promising approach to enabling memory-efficient training of large language models (LLMs)<n>Existing low-rank optimization methods typically project gradients onto a low-rank subspace, reducing the memory cost of storing states.<n>We propose importance sampling subspace selection (I3S) for low-rank optimization, which theoretically offers a comparable convergence rate to the dominant subspace approach.
  arXiv  Detail & Related papers  (2025-02-09T06:30:19Z)
• Refining Adaptive Zeroth-Order Optimization at Ease [24.327161891577727]
  This paper introduces Refined Adaptive Zeroth-Order Optimization (R-AdaZO)<n>We first show the untapped variance reduction effect of first moment estimate on ZO gradient estimation.<n>We then refine the second moment estimate based on these variance-reduced gradient estimates to better capture the geometry of the optimization landscape.
  arXiv  Detail & Related papers  (2025-02-03T03:10:44Z)
• Elucidating Subspace Perturbation in Zeroth-Order Optimization: Theory and Practice at Scale [33.38543010618118]
  Zeroth-order (ZO) optimization has emerged as a promising alternative to gradient-based backpropagation methods.<n>We show that high dimensionality is the primary bottleneck and introduce the notion of textitsubspace alignment to explain how the subspace perturbations reduce gradient noise and accelerate convergence.<n>We propose an efficient ZO method using block coordinate descent (MeZO-BCD), which perturbs and updates only a subset of parameters at each step.
  arXiv  Detail & Related papers  (2025-01-31T12:46:04Z)
• Gradient Normalization Provably Benefits Nonconvex SGD under Heavy-Tailed Noise [60.92029979853314]
  We investigate the roles of gradient normalization and clipping in ensuring the convergence of Gradient Descent (SGD) under heavy-tailed noise. Our work provides the first theoretical evidence demonstrating the benefits of gradient normalization in SGD under heavy-tailed noise. We introduce an accelerated SGD variant incorporating gradient normalization and clipping, further enhancing convergence rates under heavy-tailed noise.
  arXiv  Detail & Related papers  (2024-10-21T22:40:42Z)
• Zeroth-Order Fine-Tuning of LLMs in Random Subspaces [66.27334633749734]
  As language models grow in size, memory demands for backpropagation increase. Zeroth-order (ZOZO) optimization methods offer a memory-efficient alternative. We show that SubZero enhances fine-tuning and achieves faster results compared to standard ZOZO approaches.
  arXiv  Detail & Related papers  (2024-10-11T17:01:43Z)
• Enhancing Zeroth-order Fine-tuning for Language Models with Low-rank Structures [21.18741772731095]
  Zeroth-order (ZO) algorithms offer a promising alternative by approximating gradients using finite differences of function values. Existing ZO methods struggle to capture the low-rank gradient structure common in LLM fine-tuning, leading to suboptimal performance. This paper proposes a low-rank ZO algorithm (LOZO) that effectively captures this structure in LLMs.
  arXiv  Detail & Related papers  (2024-10-10T08:10:53Z)
• AdaLomo: Low-memory Optimization with Adaptive Learning Rate [59.64965955386855]
  We introduce low-memory optimization with adaptive learning rate (AdaLomo) for large language models. AdaLomo results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models.
  arXiv  Detail & Related papers  (2023-10-16T09:04:28Z)
• Proximal Subgradient Norm Minimization of ISTA and FISTA [8.261388753972234]
  We show that the squared proximal subgradient norm for the class of iterative shrinkage-thresholding algorithms converges at an inverse square rate. We also show that the squared proximal subgradient norm for the class of faster iterative shrinkage-thresholding algorithms (FISTA) is accelerated to convergence at an inverse cubic rate.
  arXiv  Detail & Related papers  (2022-11-03T06:50:19Z)
• An Accelerated Doubly Stochastic Gradient Method with Faster Explicit Model Identification [97.28167655721766]
  We propose a novel doubly accelerated gradient descent (ADSGD) method for sparsity regularized loss minimization problems. We first prove that ADSGD can achieve a linear convergence rate and lower overall computational complexity.
  arXiv  Detail & Related papers  (2022-08-11T22:27:22Z)
• On Learning Rates and Schr\"odinger Operators [105.32118775014015]
  We present a general theoretical analysis of the effect of the learning rate. We find that the learning rate tends to zero for a broad non- neural class functions.
  arXiv  Detail & Related papers  (2020-04-15T09:52:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.