# COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs

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

## Key facts

- Cited by 30 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 50 references that are in the map.
- First posted on arXiv 24 Feb 2025.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [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)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [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)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [LISA: Layerwise Importance Sampling for Memory-Efficient Large Language Model Fine-Tuning](https://mlsanity.com/paper/2403.17919/) (2024)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [Benchmarking Neural Network Training Algorithms](https://mlsanity.com/paper/2306.07179/) (2023)
- [Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures](https://mlsanity.com/paper/2311.00636/) (2023)
- [Deconstructing What Makes a Good Optimizer for Language Models](https://mlsanity.com/paper/2407.07972/) (2024)
- [Memory-Efficient LLM Training with Online Subspace Descent](https://mlsanity.com/paper/2408.12857/) (2024)

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

- [NorMuon: Making Muon more efficient and scalable](https://mlsanity.com/paper/2510.05491/) (2025)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [ASGO: Adaptive Structured Gradient Optimization](https://mlsanity.com/paper/2503.20762/) (2025)
- [Lions and Muons: Optimization via Stochastic Frank-Wolfe under Heavy-Tailed Noise](https://mlsanity.com/paper/2506.04192/) (2025)
- [Convergence of Muon with Newton-Schulz](https://mlsanity.com/paper/2601.19156/) (2026)
- [On the Convergence of Muon and Beyond](https://mlsanity.com/paper/2509.15816/) (2025)
- [Convergence Bound and Critical Batch Size of Muon Optimizer](https://mlsanity.com/paper/2507.01598/) (2025)
- [AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates](https://mlsanity.com/paper/2509.02981/) (2025)
- [Low-rank Orthogonalization for Large-scale Matrix Optimization with Applications to Foundation Model Training](https://mlsanity.com/paper/2509.11983/) (2025)
- [RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization](https://mlsanity.com/paper/2603.20527/) (2026)
- [HTMuon: Improving Muon via Heavy-Tailed Spectral Correction](https://mlsanity.com/paper/2603.10067/) (2026)
- [FedMuon: Federated Learning with Bias-corrected LMO-based Optimization](https://mlsanity.com/paper/2509.26337/) (2025)
- [Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement](https://mlsanity.com/paper/2602.22681/) (2026)
- [On the Convergence Analysis of Muon](https://mlsanity.com/paper/2505.23737/) (2025)
- [Conda: Column-Normalized Adam for Training Large Language Models Faster](https://mlsanity.com/paper/2509.24218/) (2025)
- [Ky Fan Norms and Beyond: Dual Norms and Combinations for Matrix Optimization](https://mlsanity.com/paper/2512.09678/) (2025)
- [Decoupling Variance and Scale-Invariant Updates in Adaptive Gradient Descent for Unified Vector and Matrix Optimization](https://mlsanity.com/paper/2602.06880/) (2026)
- [A unified convergence theory for adaptive first-order methods in the nonconvex case, including AdaNorm, full and diagonal AdaGrad and Muon](https://mlsanity.com/paper/2604.17423/) (2026)
- [Convergence of Spectral Descent for Non-smooth Optimization](https://mlsanity.com/paper/2605.26977/) (2026)
- [Budget-aware Auto Optimizer Configurator](https://mlsanity.com/paper/2605.04711/) (2026)
- [Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization](https://mlsanity.com/paper/2605.06316/) (2026)
- [MuonQ: Enhancing Low-Bit Muon Quantization via Directional Fidelity Optimization](https://mlsanity.com/paper/2605.11396/) (2026)
- [When Muon Optimizer Meets Adversarial Training: A Theoretical and Empirical Study](https://mlsanity.com/paper/2605.26929/) (2026)
- [OPTScientist: Multi-Agent Discovery of Typed Optimizer Programs for Transformer Pretraining](https://mlsanity.com/paper/2607.20486/) (2026)
- [Lotus: Efficient LLM Training by Randomized Low-Rank Gradient Projection with Adaptive Subspace Switching](https://mlsanity.com/paper/2602.01233/) (2026)
- [Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration](https://mlsanity.com/paper/2605.09034/) (2026)
- [AMO: Operator-level Adaptive Muon Orthogonalization](https://mlsanity.com/paper/2605.17806/) (2026)
- [OmniOpt: Taxonomy, Geometry, and Benchmarking of Modern Optimizers](https://mlsanity.com/paper/2607.04033/) (2026)
- [Blog: Survey of Optimizers](https://mlsanity.com/paper/2608.28557/) (2026)
- [Clean: Second-order LLM Training at Linear Memory Cost via Nystr"om Sketching](https://mlsanity.com/paper/2610.04204/) (2026)

## Abstract (arXiv)

Large Language Models (LLMs) have demonstrated remarkable success across various domains, yet their optimization remains a significant challenge due to the complex and high-dimensional loss landscapes they inhabit. While adaptive optimizers such as AdamW are widely used, they suffer from critical limitations, including an inability to capture interdependencies between coordinates and high memory consumption. Subsequent research, exemplified by SOAP, attempts to better capture coordinate interdependence but incurs greater memory overhead, limiting scalability for massive LLMs. An alternative approach aims to reduce memory consumption through low-dimensional projection, but this leads to substantial approximation errors, resulting in less effective optimization (e.g., in terms of per-token efficiency). In this paper, we propose COSMOS, a novel hybrid optimizer that leverages the varying importance of eigensubspaces in the gradient matrix to achieve memory efficiency without compromising optimization performance. The design of COSMOS is motivated by our empirical insights and practical considerations. Specifically, COSMOS applies SOAP to the leading eigensubspace, which captures the primary optimization dynamics, and MUON to the remaining eigensubspace, which is less critical but computationally expensive to handle with SOAP. This hybrid strategy significantly reduces memory consumption while maintaining robust optimization performance, making it particularly suitable for massive LLMs. Numerical experiments on various datasets and transformer architectures are provided to demonstrate the effectiveness of COSMOS. Our code is available at https://github.com/lliu606/COSMOS.