技術要約: ZO-Act

問題提起

大規模言語モデル（LLM）のファインチューニングは、通常、一次最適化（バックプロパゲーション）に依存しているが、これはモデルが数十億のパラメータへとスケールするにつれ、活性化、オプティマイザの状態、および勾配計算のストレージ要件が増大するため、メモリ消費が極めて大きくなる。ゼロ次（ZO）最適化は、損失の評価のみを用いて更新方向を推定するフォワード専用の代替案を提供する。これにより、バックプロパゲーションを回避できる。しかし、既存のZO手法は重大な課題に直面している：

1. 高分散: （MeZOのように）全パラメータ空間を摂動させることは、勾配推定の高分散を招き、収束の遅延と不安定化を引き起こす。
2. 不適切な部分空間: 一部の手法（LOZOやSubZeroなど）は、摂動を低次元の部分空間に制限しているが、これらはランダムに構築された部分空間や、実際の勾配フローと一致しないランダムな摂動に依存していることが多い。
3. オプティマイザとの互換性: 多くのZO手法は、摂動の構造が明示的かつ軽量な学習可能変数を開示しないため、標準的なモーメンタムベースのオプティマイザ（Adamなど）を効果的に利用することに苦慮している。
4. 量子化の制約: 全重みの摂動を実体化する必要がある手法では、量子化されたLLM（例：INT4）のファインチューニングが困難である。

手法: ZO-Act

著者らは、ZO-Act、すなわち、一度の実行で完了する（one-shot）、活性化情報に基づいたZOファインチューニング手法を提案している。その核心となる洞察は、線形層における重み勾配の支配的な成分は、入力活性化行列のトップ右特異ベクトルによって張られる部分空間内に存在するということである。

コアメカニズム

1. ワンショット初期化: 各線形層に対して、ZO-Actはキャリブレーション・バッチを用いた単一のフォワードパスを実行し、入力活性化行列 $X$ を計算する。その後、薄い特異値分解（SVD）を通じて、$X$ の上位 $r$ 個の右特異ベクトル ($V_r$) を計算する。
2. 固定された部分空間のパラメータ化: 重みの更新は、これらの凍結された活性化ベクトルによって定義される低ランクの部分空間に制限される。有効な重み更新は次のようにパラメータ化される：
   $$\Delta W = V_r B$$
   ここで、$V_r \in \mathbb{R}^{m \times r}$ は凍結された活性化基底であり、$B \in \mathbb{R}^{r \times n}$ は軽量で学習可能な係数行列である。元の事前学習済み重み $W$ と基底 $V_r$ は、トレーニング中を通じて凍結されたままとなる。
3. 係数空間における最適化: 全重み行列 $W$ を摂動させる代わりに、ZO-Actは低次元の係数空間においてランダムな摂動 $Z$ をサンプリングする（$Z \sim \mathcal{N}(0, I)$）。摂動後の有効な重みは $W_{eff} = W + V_r(B + \mu Z)$ となる。
4. フォワードのみの推定: $B$ に対する勾配は、前向き差分による損失評価を用いて推定される：
   $$\hat{g}_B = \frac{L(B + \mu Z) - L(B)}{\mu} Z$$
   この推定された勾配は、Adamのような標準的な一次オプティマイザを用いて $B$ を更新するために使用される。

理論的分析

本論文では、ZO-Actを、全重み空間（次元 $d$）ではなく、制限された係数空間 $\beta$（次元 $k$）上のゼロ次最適化として分析している。

• 分散の低減: ZO推定量の分散は、摂動の次元に対して線形にスケールする。次元を $d$（全重み）から $k = \sum r n_\ell$（係数、ここで $r \ll m$）に削減することで、ZO-Actは分散に依存する収束項と有限差分誤差を大幅に減少させる。
• バイアスと分散のトレードオフ: 活性化情報に基づいた部分空間に更新を制限することは、部分空間近似バイアスを導入する。しかし、著者らは、LLMの活性化と勾配は低ランク構造を示すため、支配的な更新方向は入力活性化のトップ特異ベクトルによって十分に捉えられると主張している。したがって、バイアスは制御されており、データに基づいた性質によって軽減されている。

主な貢献

1. 活性化情報に基づいた部分空間: ランダムまたは静的な低ランク行列を使用する従来のサブスペース手法とは異なり、ZO-Actは入力活性化から直接摂動部分空間を導出しており、これにより部分空間がレイヤーの支配的な活性化方向と一致することを保証する。
2. 明示的な係数最適化: 基底を凍結し、係数行列 $B$ のみを最適化することで、ZO-ActはZOファインチューニングを低次元空間上の明示的な最適化問題へと変換する。これにより、モーメンタムベースのオプティマイザ（Adamなど）を直接適用することが可能になり、全重みの摂動の実体化を回避できる。
3. 量子化への適合性: 低ビットの重みは凍結されたままであり、低ビットの係数行列のみが更新されるため、ZO-Actは量子化されたLLM（例：INT4）を自然にサポートする。
4. 理論的保証: 本論文は、ZO-Actが、管理可能な部分空間バイアスの代償として、全重みZOと比較して、勾配推定の分散と有限差分誤差の両方を低減させることを示す収束分析を提供している。

実験結果

著者らは、Llama-3-8B、OPT-13B、および INT4量子化されたLlama-3-8B について、言語理解、質問回答、および常識推論のタスクを用いてZO-Actを評価した。

• パフォーマンス: ZO-Actは、強力なZOベースライン（MeZO、LOZO、SubZero、AGJO、およびZO-Muonを含む）を一貫して上回った。
  • Llama-3-8B において、SST-2、RTE、BoolQ、およびWiCで最高のフル精度ZO結果を達成した。特に、RTEの性能を81.2（ZO-Muon）から87.0へ、CBを69.6から89.3へと向上させた。
  • OPT-13B においても同様に、SST-2、RTE、CB、およびBoolQでトップの結果を達成した。
  • INT4量子化モデル において、ZO-Actはフル精度ZOベースラインと同等の競争力を維持しており、メモリ制約のあるファインチューニングにおける有効性を実証した。
• 効率性: ZO-Actは、比較された手法の中で最も低いメモリ使用量（Llama-3-8Bで約16.1 GB、INT4では5.8 GBまで低下）を実現し、固定されたフォワードクエリ予算の下で最も速い実行時間（他の手法の44.5分以上に対し、42.0分）を達成した。
• ランク感度: 実験により、フル精度モデルでは $r=1$ のランクがしばしば十分であり、摂動次元を最小限に抑えつつ最良の精度をもたらすことが示された。量子化モデルでは、容量の減少を補うために $r=32$ のような高いランクが有益であった。
• 安定性: 診断研究により、活性化情報に基づいた部分空間がファインチューニング中も安定しており、全勾配エネルギーの大部分を捉えていることが確認され、「ワンショット」初期化設計の妥当性が検証された。

意義と主張

本論文は、ZO-Actが、大規模および量子化されたLLMにとって、ゼロ次ファインチューニングを実用的かつ高性能な代替手段にするための重要な一歩であると主張している。その意義は以下の点にある：

• ギャップの解消: データに基づいた部分空間を通じて推定器の分散を低減することで、ZOと一次ファインチューニングの間の性能差を縮めている。
• オプティマイザの有効化: 明示的な低次元の学習可能変数を開示することで、ZO手法へのAdamのような標準的なオプティマイザの適用という困難を解決した。
• スケーラビリティ: リソース制約のあるハードウェアでの展開に不可欠な、量子化モデルのファインチューニングのためのメモリ効率の高いソリューションを提供している。

著者らは、ZO-Actが依然としてZO推定の固有のノイズや部分空間バイアスのために、一次手法との性能差を完全には埋めていないことを認め、限界についても言及している。さらに、ZO-Actは、1回の更新あたりに複数のフォワードパスを必要とするため、ウォールクロックタイムではバックプロパゲーションベースの手法よりも遅くなる。しかし、本手法は、バックプロパゲーションがメモリ的に不可能、あるいは利用できないシナリオにおいて、優れた選択肢として位置付けられている。