SubZero: Random Subspace Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning

要約

大規模言語モデル (LLM) の微調整は、さまざまな下流タスクに効果的であることが証明されています。
ただし、LLM のサイズが大きくなるにつれて、バックプロパゲーションのメモリ要求はますます法外なものになります。
0 次 (ZO) 最適化手法は、勾配の推定にフォワード パスを使用することでメモリ効率の高い代替手段を提供しますが、勾配推定の分散は通常、モデルのパラメーターの次元に線形に比例するため、LLM にとっては重大な問題となります。
この論文では、LLM の高次元性によってもたらされる課題に対処するために、ランダムな部分空間 0 次 (SubZero) 最適化を提案します。
トレーニングのパフォーマンスを向上させながらメモリ消費を大幅に削減する、LLM 向けに調整された低ランク摂動を導入します。
さらに、勾配推定がバックプロパゲーション勾配に非常に近似し、従来の ZO 法よりも低い分散を示し、SGD と組み合わせた場合に確実に収束することを証明します。
実験結果は、SubZero がさまざまな言語モデリング タスクにわたって、MeZO などの標準的な ZO アプローチと比較して、微調整パフォーマンスを強化し、より高速な収束を達成することを示しています。

要約(オリジナル)

Fine-tuning Large Language Models (LLMs) has proven effective for a variety of downstream tasks. However, as LLMs grow in size, the memory demands for backpropagation become increasingly prohibitive. Zeroth-order (ZO) optimization methods offer a memory-efficient alternative by using forward passes to estimate gradients, but the variance of gradient estimates typically scales linearly with the model’s parameter dimension$\unicode{x2013}$a significant issue for LLMs. In this paper, we propose the random Subspace Zeroth-order (SubZero) optimization to address the challenges posed by LLMs’ high dimensionality. We introduce a low-rank perturbation tailored for LLMs that significantly reduces memory consumption while improving training performance. Additionally, we prove that our gradient estimation closely approximates the backpropagation gradient, exhibits lower variance than traditional ZO methods, and ensures convergence when combined with SGD. Experimental results show that SubZero enhances fine-tuning performance and achieves faster convergence compared to standard ZO approaches like MeZO across various language modeling tasks.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google