今回の論文
今回取り上げるのは、Sadhika Malladiらの「Fine-Tuning Language Models with Just Forward Passes」です。2023年にarXivで公開され、NeurIPS 2023(oral)で発表されました。公開元はarXiv、URLは https://arxiv.org/abs/2305.17333 、研究分野は大規模言語モデルの省メモリ・ファインチューニングです。
選んだ理由は、「学習には逆伝播が必要」という前提を崩し、推論できるGPUなら巨大モデルの調整にも使える可能性を示したからです。提案手法MeZOは、重みを少し正負へ揺らしたときの損失の差だけから更新方向を推定します。モデル、アダプタ、評価指標をどう設計すれば限られたGPUでも改善を回せるかという、実務上のヒントが詰まっています。
どんな技術か
MeZO(Memory-efficient Zeroth-order Optimizer)は、勾配を逆伝播で計算せず、2回の順伝播の損失差から更新方向を見積もるLLMファインチューニング手法です。通常のAdamによる学習では、順伝播の中間活性、勾配、最適化器の状態をGPUに保持します。MeZOはこれらを保持せず、推論とほぼ同じメモリ量で重みを更新します。
一見すると、数十億個の重みをランダムに揺らして本当に学べるのか、という疑問が生じます。論文の要点は、十分に事前学習されたモデルを適切なプロンプト形式で調整するなら、損失地形の実効的な難しさはパラメータ数そのものでは決まらない、という観察です。実験では最大66BのOPTを対象に、分類・多肢選択・生成を扱っています。
課題
LLMを特定ドメインへ合わせるとき、精度より先にGPUメモリが制約になることがあります。逆伝播は、重みだけでなく各層の活性値と勾配を保持し、Adamなら一次・二次モーメントも保存します。そのため、推論できるモデルサイズとフルファインチューニングできるモデルサイズには大きな隔たりがあります。
- 大きなモデルほど活性値・勾配・最適化器状態の合計が増え、バッチを小さくしても単一GPUに載らないことがあります。
- LoRAやprefix-tuningは更新パラメータを減らしますが、一般的な逆伝播では基盤モデルを通る活性値を保存する必要があり、メモリ制約を完全には取り除けません。
- 勾配を使わないゼロ次最適化は昔からありますが、ランダム方向の探索は次元数に弱く、巨大モデルでは極端に遅いと考えられていました。
この問題は、社内データでの専門用語適応、モデルを顧客環境へ持ち込むオンプレミス開発、小規模チームの試作で顕在化します。「高品質な大規模モデルは推論専用」と固定せず、少ない計算資源でも評価・調整できる選択肢を持つことが目的です。
用語解説
- 逆伝播(backpropagation)
- 損失から各重みへの微分(勾配)を連鎖律で求める標準的な学習手法です。精度の高い更新ができる一方、順伝播時の中間活性を保持するため、MeZOが削減するメモリの主因になります。
- ゼロ次最適化
- 微分を直接使わず、入力を変えた前後の目的関数値だけで改善方向を探す最適化です。MeZOは損失を観測できさえすればよいため、AccuracyやF1のような微分不能な指標も目的にできます。
- SPSA
- Simultaneous Perturbation Stochastic Approximationの略で、全パラメータを同時にランダム摂動し、正負2回の評価から勾配らしいベクトルを推定する方法です。MeZOの更新則の核になっています。
- 摂動ベクトル
- 各パラメータをどちらへどれだけ動かすかを表すランダムな方向です。MeZOは保存量を増やさないため、乱数シードから同じベクトルを必要な都度再生成します。
- プロンプトベースのファインチューニング
- 分類ラベルを単なるIDではなく、言語モデルが予測する語句として表す学習形式です。論文では事前学習時に近い目的へ寄せることが、MeZOで安定して学ぶ重要な条件になっています。
技術の仕組み
MeZOは古典的なZO-SGDを、巨大なモデルを複製せずに実行できるように実装し直したものです。ミニバッチの損失を L(theta)、ランダムな摂動方向を z、摂動の大きさを epsilon とします。
1. 正負へ揺らして損失差を測る
まず重みを theta + epsilon z に一時的に動かして順伝播し、損失 L_plus を測ります。次に同じ方向の反対側、theta - epsilon z にして L_minus を測ります。差分 (L_plus - L_minus) / (2 epsilon) が、そのランダム方向に沿って損失が増えるか減るかを表します。
ここで得たスカラーへ z を掛けると、全パラメータに対する勾配推定値になります。厳密な各要素の偏微分ではありませんが、多数ステップで異なる方向を試すことで、平均的には損失を下げる方向を捉えます。更新は概念的に theta <- theta - learning_rate * estimated_gradient です。
2. 重みを複製せず、乱数シードを再利用する
素直な実装では、z 全体を保存するだけでモデルと同じ規模のメモリが必要になります。MeZOは各パラメータテンソルについて、同じ乱数シードから摂動を再生成します。
処理は「正方向を加える→順伝播→負方向へ2倍動かす→順伝播→元へ戻す→同じ乱数を再生成して更新する」という流れです。GPUに残るのは通常の推論と同じ重みと一時的な演算用メモリであり、活性値の保持も勾配テンソルの保持も不要です。これが「in-place」の意味です。
3. フル更新とPEFTを同じ枠組みで扱う
MeZOは全パラメータへ摂動を入れることも、LoRAの低ランク行列やprefixだけへ限定することもできます。対象外の重みは揺らさず、対象テンソルだけで乱数を生成すればよいからです。
ただし、パラメータ数を減らせば必ずMeZOの収束が速くなるわけではありません。論文ではフル更新・LoRA・prefixで似た収束速度も観察されています。PEFTを選ぶ主な理由は、保存・配布する更新差分を小さくすることやタスク切替であり、ゼロ次探索の回数を魔法のように減らすことではありません。
4. プロンプトで事前学習と目的をつなぐ
論文は、MeZOが単独で万能に働くとは主張していません。分類問題でも入力をテンプレートに埋め、クラスを言語モデルの出力語へ対応させています。これにより、事前学習で獲得した次トークン予測の性質と下流目的の隔たりを小さくします。
理論面では、このような条件の下で局所的な損失地形の「有効ランク」が十分小さければ、1ステップ当たりの改善が全パラメータ数 d に直接は依存しないことを示します。これは常に保証される性質ではなく、良い事前学習済みチェックポイントとタスク表現が前提です。
5. 微分不能な目的も直接最適化できる
逆伝播は損失が微分可能である必要があります。MeZOは数値を比較するだけなので、分類AccuracyやQAのF1をそのまま目的関数にできます。論文ではクロスエントロピーを最適化した場合が強いものの、Accuracy/F1を直接最適化してもゼロショットを上回りました。
これは、業務KPIに近いスコアを代理損失へ落とすのが難しいときの選択肢になります。ただし評価指標は離散的でノイズも大きいため、少数データでの過適合やスコアの揺れを監視する必要があります。
実験と結果
論文はRoBERTa-large(350M)と、OPT-13B・30B・66Bを使いました。RoBERTaではSST-2、SST-5、SNLI、TRECなどの感情・自然言語推論・分類を、OPTではSuperGLUE系、SQuAD、DROPなどの多肢選択・生成QAを評価しています。比較対象はゼロショット、in-context learning、linear probing、Adamによる通常のファインチューニングです。
少量・中規模モデルでの精度
RoBERTa-largeでは、クラスごと16例と512例のプロンプトベース設定を検証しました。MeZOは6つのタスクすべてでゼロショットとlinear probingを上回りました。512例/クラスでは、通常のファインチューニングとの差が全タスクで5%以内まで縮みます。
MeZOは100,000ステップ、通常学習は1,000ステップという非対称な比較です。MeZOは1ステップに必要な処理が軽い一方、より多くの方向を試す必要がある、という性質が数字にも現れています。「2回の順伝播だけ」だから短時間で必ず終わるのではなく、総GPU時間で判断すべきです。
13B〜66Bと生成タスクへの拡張
OPT-13Bでは、通常のファインチューニングが約12倍のメモリを使う条件で、11タスク中7タスクにおいてMeZOが通常学習と1%以内またはそれ以上の性能を示しました。SQuADとDROPでは教師強制で学習し、greedy decodingで評価しています。いずれの生成タスクでも、ゼロショット・ICL・linear probingを上回り、通常学習に近い結果でした。
さらに30B・66Bでは各データセットから学習1,000例、検証500例、テスト1,000例を用意し、フルまたはprefixのMeZOを試しています。直接の通常学習が高コストになるスケールでも、MeZOはゼロショット・ICLより強い結果を出しました。巨大モデルを「十分なメモリがないので一切調整できない」とは限らないことを示す実験です。
メモリとGPU時間
平均400トークンのMultiRCを使ったA100 80GBでの測定では、MeZOのピークメモリはゼロショット推論と同等でした。標準ファインチューニングに対して最大12分の1、prefixの通常学習に対しても最大6分の1です。
単一A100では、標準ファインチューニングが収められた最大OPTは2.7B、prefix学習では6.7B、推論相当のMeZOでは30Bでした。30BモデルではMeZOは1ステップが7.74倍速く、必要GPUも8分の1でしたが、収束により多くのステップが必要です。それでも著者らの実装では、フルファインチューニングの約半分のGPU時間で済みました。この値はモデル・実装・ハイパーパラメータに依存するため、一般的な保証ではありません。
微分不能な指標の最適化
512例/クラスのSST-2で、クロスエントロピーをMeZOで最適化した精度は93.3%、Accuracyを直接最適化した場合は92.7%でした。OPT-13BのSQuADではクロスエントロピー最適化のF1が84.7、F1直接最適化は78.5で、どちらもゼロショットの46.2を上回っています。
結果は、微分可能なクロスエントロピーを使えるなら依然として有利であることも示します。MeZOの価値は「どんなKPIでも同じ精度」ではなく、微分を定義しづらい目的にも更新の入口を作れる点です。
何に使える?
GPU一枚でのドメイン適応の試作
専門用語の分類、問い合わせの優先度付け、定型文の生成など、少量のラベル付きデータを使う検証に向きます。標準学習ではVRAMに載らない基盤モデルでも、推論可能ならMeZOで差分を試せる場合があります。まずプロンプトと評価セットを固定し、ゼロショット、ICL、MeZOを同一のモデル・入力で比較すると、導入価値を判断しやすくなります。
微分しにくい業務評価の探索
構造化出力の完全一致、検索結果を含む回答の採点、ルールベース検査の通過率のように、微分可能な損失へ素直に変換できない指標があります。MeZOはスコアが返る限り更新できます。実務では、そのスコアを小さな検証データへ直接最適化すると過適合しやすいため、学習用と独立したホールドアウト評価を必ず分けるべきです。
顧客ごとの差分アダプタ作成
基盤モデルを複数顧客のデータで個別適応する場合、LoRAとMeZOを組み合わせる考え方があります。学習時のVRAMを抑えつつ、保存するのは小さなアダプタだけです。ただし論文ではLoRAがMeZOの収束を常に速くするとは限らないため、主な利点は配布・切替・管理コストの低下だと捉えるのが適切です。
開発や事業へのヒント
学習基盤の入り口を「推論可能か」へ下げる
多くのチームは、モデルを推論できても追加学習用GPUがないため、プロンプトだけで妥協します。MeZOの示唆は、推論環境をそのまま学習実験の足場にできるという点です。最初から大規模な学習基盤を作るのではなく、代表データ数百〜千件、再現可能な採点器、チェックポイント管理を用意する小さなループから始められます。
プロンプトを単なるUI文言と考えない
この論文では、プロンプトは最適化が成功する条件の一部です。分類ラベルの言語化、入出力の区切り、回答形式を設計し、事前学習時に近い次トークン予測へ問題を変換します。ファインチューニングが不安定なときは、学習率だけでなく、タスクをどう文章化したかを先に見直す価値があります。
成功指標とコスト指標を分けて測る
MeZOでは、メモリ使用量、1ステップ時間、収束までのステップ数が逆方向に動くことがあります。「逆伝播なし」を理由に採用せず、(1) GPUメモリ、(2) 総GPU時間、(3) 実タスク品質、(4) 評価の分散を同じ実験表で確認してください。GPU一枚で30Bを動かせることと、サービス品質へ十分に適応できることは別の検証です。
限界
MeZOは各ステップで2回の順伝播を必要とし、勾配推定にはランダム性があります。通常の逆伝播より多くの更新ステップが必要になり、細かい改善や複雑なタスクでは収束が遅い、あるいは不安定になる可能性があります。論文自身も、標準学習より大幅に多いステップ数を使っています。
効果は、事前学習済みモデル、プロンプト形式、タスクとデータ量に依存します。論文の理論は局所的な有効ランクなどの仮定に基づくもので、任意の未学習モデルや分布外タスクに対する保証ではありません。適切なプロンプトが重要だったという実験結果も、実装時の注意点です。
また、推論と同じメモリ量という利点は、モデル重みを載せられることを意味するにとどまります。長いコンテキストのKVキャッシュ、データ前処理、分散通信、チェックポイント保存は別のコストです。AccuracyやF1の直接最適化も、離散値が高分散であるため、評価データを学習に使い切らない設計と複数シードでの検証が必要です。
よくある質問
Q. MeZOはLoRAより省メモリですか?
A. 比較対象によります。LoRAを通常の逆伝播で学習する場合も活性値を保存するため、MeZOの方がさらにメモリを抑えられます。一方でMeZOはLoRAと競合するだけでなく、LoRAのパラメータだけをMeZOで更新する組み合わせも可能です。
Q. MeZOなら巨大モデルをCPUだけで学習できますか?
A. いいえ。逆伝播用メモリは避けられますが、少なくともモデルを推論できる計算資源と、2回の順伝播を繰り返す時間は必要です。CPUへオフロードすれば実行できる構成は増えますが、現実的な所要時間になるかは別途測定してください。
Q. どんな評価指標でも直接最適化できますか?
A. 数値を返せれば原理上は可能です。とはいえ、乱数や外部APIの揺れを含むスコア、段階的なルールしか持たないスコアは推定を不安定にします。固定した評価条件、十分なサンプル数、独立した検証セットを用意してください。
Q. MeZOは通常のファインチューニングを置き換えますか?
A. 一律には置き換えません。逆伝播を使えるメモリがあり、最高精度や速い収束を求めるなら標準手法が有力です。MeZOは、メモリ制約で比較実験すらできないとき、あるいは微分不能な目的を試すときに特に価値があります。
Q. 導入時に最初に確認すべきことは何ですか?
A. まず、ベースモデルのゼロショット品質、実運用に近いプロンプト、独立した評価セットを固定してください。その上で小さなモデルまたは少数データから、MeZOの学習率・摂動幅・ステップ数を探索し、メモリ削減と品質の両方を記録するのが安全です。
今日の学び
この論文は、巨大LLMのファインチューニングで逆伝播が生むメモリ負荷を扱いました。MeZOは、同じランダム摂動の正負で損失を2回測り、乱数シードの再利用で摂動ベクトルも保存せずに更新します。
得られるヒントは、限られたGPUでの学習を「できない」と決める前に、推論相当のメモリで回る最適化とタスク表現を試すことです。プロンプト、評価器、総GPU時間を一緒に設計すれば、モデル適応の実験範囲を小さな基盤から広げられます。