Leveraging Coordinate Momentum in SignSGD and Muon: Memory-Optimized Zero-Order

要約

大型言語モデル（LLMS）の微調整は、事前に訓練されたモデルをダウンストリームタスクに適応させるために不可欠です。
しかし、確率的勾配降下（SGD）やAdamなどの従来の1次オプティマイザーは、モデルサイズで縮小する法外なメモリと計算コストを負担します。
このホワイトペーパーでは、特にLORAのようなパラメーター効率の高い微調整技術のコンテキストで、メモリおよび計算効率の高い代替品としてゼロオーダー（ZO）最適化方法を調査します。
ZO SignsGDを拡張するZOモメンタムベースのアルゴリズムである$ \ texttt {jaguar signsgd} $を提案します。
私たちの知る限り、これは確率的ZOケースでSignsgDの厳密な収束保証を確立する最初の研究です。
さらに、モデルパラメーターのマトリックス構造を活用するMuonオプティマイザーの新しいZO拡張である$ \ texttt {Jaguar Muon} $を提案し、任意の確率ノイズの下で収束速度を提供します。
挑戦的なLLM微調整ベンチマークに関する広範な実験を通じて、提案されたアルゴリズムが標準の1次方法の収束品質を満たしているか、それを超えて、大幅なメモリ削減を達成することを実証します。
当社の理論的および経験的結果は、リソースに制約のあるLLM適応のための実用的および理論的に根拠のあるアプローチとして、新しいZO最適化方法を確立しています。
私たちのコードはhttps://github.com/brain-mmo-lab/zo_llmで入手できます

要約(オリジナル)

Fine-tuning Large Language Models (LLMs) is essential for adapting pre-trained models to downstream tasks. Yet traditional first-order optimizers such as Stochastic Gradient Descent (SGD) and Adam incur prohibitive memory and computational costs that scale poorly with model size. In this paper, we investigate zero-order (ZO) optimization methods as a memory- and compute-efficient alternative, particularly in the context of parameter-efficient fine-tuning techniques like LoRA. We propose $\texttt{JAGUAR SignSGD}$, a ZO momentum-based algorithm that extends ZO SignSGD, requiring the same number of parameters as the standard ZO SGD and only $\mathcal{O}(1)$ function evaluations per iteration. To the best of our knowledge, this is the first study to establish rigorous convergence guarantees for SignSGD in the stochastic ZO case. We further propose $\texttt{JAGUAR Muon}$, a novel ZO extension of the Muon optimizer that leverages the matrix structure of model parameters, and we provide its convergence rate under arbitrary stochastic noise. Through extensive experiments on challenging LLM fine-tuning benchmarks, we demonstrate that the proposed algorithms meet or exceed the convergence quality of standard first-order methods, achieving significant memory reduction. Our theoretical and empirical results establish new ZO optimization methods as a practical and theoretically grounded approach for resource-constrained LLM adaptation. Our code is available at https://github.com/brain-mmo-lab/ZO_LLM

arxiv情報

提供元, 利用サービス

arxiv.jp, Google