Fine-Tuning Language Models with Just Forward Passes

要約

言語モデル (LM) を微調整することで、さまざまなダウンストリーム タスクで成功を収めてきましたが、LM のサイズが大きくなるにつれて、バックプロパゲーションには法外に大量のメモリが必要になります。
ゼロ次 (ZO) 法は、原理的には 2 つの前方パスのみを使用して勾配を推定できますが、大規模なモデルを最適化するには壊滅的に時間がかかることが理論化されています。
この研究では、古典的な ZO-SGD 手法をインプレースで動作するように適応させ、それによって推論と同じメモリ フットプリントで LM を微調整する、メモリ効率の高いゼロ次オプティマイザ (MeZO) を提案します。
たとえば、単一の A100 80 GB GPU を使用すると、MeZO は 300 億のパラメーター モデルをトレーニングできますが、バックプロパゲーションを使用した微調整では、同じ予算で 27 億の LM しかトレーニングできません。
私たちは、モデル タイプ (マスクおよび自己回帰 LM)、モデル スケール (最大 66B)、および下流タスク (分類、多肢選択、生成) にわたる包括的な実験を実施します。
私たちの結果は、(1) MeZO がコンテキスト内学習と線形プローブよりも大幅に優れていることを示しています。
(2) MeZO は、複数のタスクにわたるバックプロパゲーションによる微調整に匹敵するパフォーマンスを実現し、実装では最大 12 倍のメモリ削減と最大 2 倍の GPU 時間削減を実現します。
(3) MeZO は、フルパラメータ調整手法と、LoRA やプレフィックス調整などのパラメータ効率の高い調整手法の両方と互換性があります。
(4) MeZO は微分不可能な目的を効果的に最適化できます (精度や F1 の最大化など)。
古典的な ZO 分析ではそうではないことが示されているにもかかわらず、適切な事前トレーニングとタスク プロンプトによって MeZO がどのように巨大なモデルを微調整できるかを強調し、理論的な洞察によって経験的発見をサポートします。

要約(オリジナル)

Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a prohibitively large amount of memory. Zeroth-order (ZO) methods can in principle estimate gradients using only two forward passes but are theorized to be catastrophically slow for optimizing large models. In this work, we propose a memory-efficient zerothorder optimizer (MeZO), adapting the classical ZO-SGD method to operate in-place, thereby fine-tuning LMs with the same memory footprint as inference. For example, with a single A100 80GB GPU, MeZO can train a 30-billion parameter model, whereas fine-tuning with backpropagation can train only a 2.7B LM with the same budget. We conduct comprehensive experiments across model types (masked and autoregressive LMs), model scales (up to 66B), and downstream tasks (classification, multiple-choice, and generation). Our results demonstrate that (1) MeZO significantly outperforms in-context learning and linear probing; (2) MeZO achieves comparable performance to fine-tuning with backpropagation across multiple tasks, with up to 12x memory reduction and up to 2x GPU-hour reduction in our implementation; (3) MeZO is compatible with both full-parameter and parameter-efficient tuning techniques such as LoRA and prefix tuning; (4) MeZO can effectively optimize non-differentiable objectives (e.g., maximizing accuracy or F1). We support our empirical findings with theoretical insights, highlighting how adequate pre-training and task prompts enable MeZO to fine-tune huge models, despite classical ZO analyses suggesting otherwise.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google