COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs (2502.17410v2)

Abstract: 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.

• The paper introduces COSMOS, a hybrid adaptive optimizer that uses eigensubspaces of gradient matrices to achieve memory-efficient training for large language models.
• COSMOS significantly reduces memory consumption compared to optimizers like Adam and SOAP while maintaining or improving optimization performance in empirical evaluations.
• This memory efficiency allows for training larger LLMs or using increased batch sizes without the typical prohibitive memory overheads.

COSMOS: A Memory-Efficient Hybrid Optimizer for LLMs

The development of LLMs has reached a pivotal juncture where optimization methodologies must evolve to manage the extensive computational demands of these models. The paper "COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs" addresses this challenge by introducing COSMOS, a sophisticated hybrid optimizer designed to alleviate the memory consumption issues inherent in training LLMs without sacrificing optimization performance.

COSMOS originates from a critical evaluation of existing adaptive optimizers like Adam and its derivatives, which, although effective, are constrained by their memory-intensive operations and limited ability to capture interdependencies among model parameters. The authors identify that contemporary approaches such as SOAP enhance parameter interaction capture but escalate memory usage, hindering scalability in large-scale models.

In response, COSMOS is crafted as a hybrid optimizer that tactically utilizes eigensubspaces of gradient matrices to optimize memory usage. By decomposing gradients into leading and remaining eigensubspaces, COSMOS applies a high-memory efficiency technique similar to SOAP on the leading subspaces, which encapsulate the primary optimization dynamics. A secondary, less memory-demanding method akin to MUON addresses the remaining subspaces. This dual strategy ensures that COSMOS maintains a balance between computational efficiency and optimization effectiveness.

Key to COSMOS's design is the use of exponential moving averages (EMAs) for smoothing gradient estimates and a one-step power iteration to dynamically identify leading eigenvectors. These elements contribute to the optimizer's robust yet efficient handling of the complexities presented by LLMs. In practical evaluations, COSMOS reduced memory consumption notably while maintaining, or even enhancing, optimization effectiveness compared to leading algorithms like SOAP and Adam, particularly as model sizes scale.

Empirical results underscore COSMOS's superiority in per-token efficiency relative to traditional methods. Specifically, COSMOS outperformed Adam and SOAP in extensive numerical experiments conducted on datasets and various transformer architectures, reaffirming its practical utility in large-scale model training. The integration of SOAP-like dynamics for leading eigensubspaces and the efficient memory management capabilities of MUON ensure that COSMOS provides a scalable solution for training LLMs with billions of parameters.

The implications of this research are profound, suggesting potential for increased batch sizes or further model scaling without the prohibitive memory overheads typically associated with such tasks. Future developments could explore the possible enhancements of COSMOS's structures for even greater efficiency or the integration of similar hybrid strategies in other optimization contexts.

In conclusion, COSMOS represents a significant advance in optimizer design for LLMs, demonstrating that sophisticated hybrid approaches can reconcile the demands of memory efficiency and optimization performance. This research contributes not only to the immediate fields of machine learning and AI but also paves the way for future explorations into scalable, efficient neural network optimization methodologies.