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

TL;DR

SUMO introduces a subspace-aware optimizer using exact SVD for moment orthogonalization, accelerating large language model training by aligning optimization steps with the loss landscape's spectral properties, leading to faster convergence and memory savings.

Contribution

The paper presents SUMO, a novel optimizer employing exact SVD-based moment orthogonalization within a low-dimensional subspace, improving convergence and memory efficiency in LLM training.

Findings

SUMO accelerates convergence compared to existing methods.

SUMO reduces memory usage by up to 20%.

Theoretical bounds on approximation errors are established.

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.