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

arXiv: https://arxiv.org/abs/2505.24749 · HTML: https://mlsanity.com/paper/2505.24749/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 61 references that are in the map.
- First posted on arXiv 30 May 2025.

## Builds on (most cited first, 30 of 44)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms](https://mlsanity.com/paper/1905.13319/) (2019)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Deep Learning without Poor Local Minima](https://mlsanity.com/paper/1605.07110/) (2016)
- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Momentum Improves Normalized SGD](https://mlsanity.com/paper/2002.03305/) (2020)
- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be](https://mlsanity.com/paper/2304.13960/) (2023)
- [A Spectral Condition for Feature Learning](https://mlsanity.com/paper/2310.17813/) (2023)
- [Flora: Low-Rank Adapters Are Secretly Gradient Compressors](https://mlsanity.com/paper/2402.03293/) (2024)
- [Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization](https://mlsanity.com/paper/2503.12645/) (2025)
- [A Note on the Convergence of Muon](https://mlsanity.com/paper/2502.02900/) (2025)
- [Robustness to Unbounded Smoothness of Generalized SignSGD](https://mlsanity.com/paper/2208.11195/) (2022)

## Built on by (most cited first, 19 of 19)

- [LiMuon: Light and Fast Muon Optimizer for Large Models](https://mlsanity.com/paper/2509.14562/) (2025)
- [Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning](https://mlsanity.com/paper/2603.09697/) (2026)
- [Powering Up Zeroth-Order Training via Subspace Gradient Orthogonalization](https://mlsanity.com/paper/2602.17155/) (2026)
- [ICR-RL: Deep Reinforcement Learning via In-Context Regression](https://mlsanity.com/paper/2509.11259/) (2025)
- [FedMuon: Federated Learning with Bias-corrected LMO-based Optimization](https://mlsanity.com/paper/2509.26337/) (2025)
- [Hybrid Autoencoders for Tabular Data: Leveraging Model-Based Augmentation in Low-Label Settings](https://mlsanity.com/paper/2511.06961/) (2025)
- [PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training](https://mlsanity.com/paper/2601.22137/) (2026)
- [Evolution of Optimization Methods: Algorithms, Scenarios, and Evaluations](https://mlsanity.com/paper/2604.12968/) (2026)
- [MiMuon: Mixed Muon Optimizer with Improved Generalization for Large Models](https://mlsanity.com/paper/2605.19619/) (2026)
- [Train Less, Infer Faster: Efficient Model Finetuning and Compression via Structured Sparsity](https://mlsanity.com/paper/2602.09169/) (2026)
- [Low-rank surrogate modeling and stochastic zero-order optimization for training of neural networks with black-box layers](https://mlsanity.com/paper/2509.15113/) (2025)
- [No Prior, No Leakage: Revisiting Reconstruction Attacks in Trained Neural Networks](https://mlsanity.com/paper/2509.21296/) (2025)
- [MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning](https://mlsanity.com/paper/2608.05088/) (2026)
- [Spectral Allocation: Why Muon Outperforms Adam, and How to Improve Muon](https://mlsanity.com/paper/2608.25990/) (2026)
- [Uncovering a Winning Lottery Ticket with Continuously Relaxed Bernoulli Gates](https://mlsanity.com/paper/2603.08914/) (2026)
- [No Subspace to Track: Non-Identifiability and Optimizer State in Low-Rank Training](https://mlsanity.com/paper/2607.05872/) (2026)
- [MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training](https://mlsanity.com/paper/2609.15037/) (2026)
- [Which Optimizer Leaks Most? Optimizer Implicit Bias as a Training-Data Privacy Controller](https://mlsanity.com/paper/2610.18294/) (2026)
- [Memory-Efficient Subspace Optimization with Residual Correction: From Momentum SGD to Muon](https://mlsanity.com/paper/2610.19844/) (2026)

## Abstract (arXiv)

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.