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

要約

大規模な言語モデル（LLM）は、さまざまなドメインで顕著な成功を示していますが、彼らが住んでいる複雑で高次元の損失の状況により、それらの最適化は依然として重要な課題です。
ADAMWなどの適応型オプティマイザーは広く使用されていますが、座標と高いメモリ消費との間の相互依存関係を捉えることができないなど、重大な制限に苦しんでいます。
SOAPによって例示されるその後の研究は、相互依存をよりよく捕捉しようとする試みを試みますが、メモリオーバーヘッドが大きくなり、大規模なLLMのスケーラビリティが制限されます。
別のアプローチは、低次元投影を通じてメモリの消費を減らすことを目的としていますが、これは実質的な近似誤差につながり、効果的な最適化（たとえば、トークンあたりの効率の点で）になります。
この論文では、最適化パフォーマンスを損なうことなくメモリ効率を実現するために、グラデーションマトリックス内のEigensubspacesのさまざまな重要性を活用する新しいハイブリッドオプティマイザーであるCosmosを提案します。
コスモスのデザインは、私たちの経験的洞察と実際的な考慮事項によって動機付けられています。
具体的には、COSMOSは、主要な最適化ダイナミクスをキャプチャする主要なEigensubspaceに石鹸を適用し、Muonは残りのEigensubspaceに適用します。
このハイブリッド戦略は、堅牢な最適化パフォーマンスを維持しながら、メモリ消費を大幅に削減し、大規模なLLMに特に適しています。
宇宙の有効性を実証するために、さまざまなデータセットと変圧器アーキテクチャでの数値実験が提供されています。
私たちのコードは、https：//github.com/lliu606/cosmosで入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google