COSMOS：面向 LLM 内存高效训练的混合自适应优化器
机器学习
2025-10-13 v3
摘要
大型语言模型（LLMs）在各个领域展现出了卓越的成功，但由于其处于复杂且高维的损失景观中，其优化仍然是一个重大挑战。虽然诸如 AdamW 之类的自适应优化器被广泛使用，但它们存在关键局限性，包括无法捕捉坐标之间的相互依赖性以及高内存消耗。以 SOAP 为代表的后续研究试图更好地捕捉坐标相互依赖性，但产生了更大的内存开销，限制了对大规模 LLMs 的可扩展性。另一种方法旨在通过低维投影来减少内存消耗，但这会导致显著的近似误差，从而导致优化效果不佳（例如，在每 token 效率方面）。在本文中，我们提出了 COSMOS，一种新颖的混合优化器，它利用梯度矩阵中特征子空间的不同重要性来实现内存效率，同时不妥协优化性能。COSMOS 的设计受到我们经验性洞察和实际考量的启发。具体而言，COSMOS 将 SOAP 应用于主导特征子空间（捕捉主要优化动态），并将 MUON 应用于剩余特征子空间（该部分不太关键但使用 SOAP 处理计算成本高）。这种混合策略显著降低了内存消耗，同时保持了稳健的优化性能，使其特别适合大规模 LLMs。我们提供了在各种数据集和 Transformer 架构上的数值实验，以证明 COSMOS 的有效性。我们的代码可在 https://github.com/lliu606/COSMOS 获取。
引用
@article{arxiv.2502.17410,
title = {COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs},
author = {Liming Liu and Zhenghao Xu and Zixuan Zhang and Hao Kang and Zichong Li and Chen Liang and Weizhu Chen and Tuo Zhao},
journal= {arXiv preprint arXiv:2502.17410},
year = {2025}
}
备注
23 pages, 9 figures, 6 tables