Dust: バックプロパゲーションを凌駕するスカラー階級ゼロ最適化法

要約（TL;DR）

• 新手法の提示: 変換器型言語モデル（Transformer LLM）の事前学習において、バックプロパゲーションと競合しうる初めてのスカラー階級ゼロ最適化法「Dust」を提案した。
  • Dust はアクティベーション（活性化値）に対し、すべてのトークンごとに独立して擾乱（ノイズ加算）を行うことで、各トークンを仮想的な集団メンバーとし、1 回の順伝播でこれら全員を並列評価する。
• 計算資源が豊富な領域での優位性: 人口規模が大きくなるほどバックプロパゲーションに近づき、複数の設定においてはそれを凌駕する性能を示す。これは、計算資源が豊富な領域においてバックプロパゲーションを超える可能性を示唆している。
• 進化戦略（ES）との比較: Dust は重量空間ベースのエバリュエーション戦略（Weight-space ES）より桁違いに効率的である。
  • 推計結果から、100 万トークン以降では、最先端の ES 手法である EGGROLL の変換器実装に対して $10^3$ 倍乃至 $10^4$ 倍の効率性を示すと見られる。
• 大規模モデルでの効率性の向上: 通説とは逆に、モデルが大きくなるほど集団に対する効率が低下するのではなく、むしろ向上することが多く、大規模な集団を有効に活用できる。
  • 具体的には、2.43 億パラメータのモデルが、120 倍も小さいモデルよりも多数の人口規模において優れた性能を示した。
• 勾配推定とバックプロパゲーションの整合性: Dust の勾配推定は人口が増加するにつれてバックプロパゲーションとの整合性が向上し、検証したすべてのスケール（最大 10 億トークンまで）で良好な整合性を維持している。

---

はじめに

背景と動機

• 深層学習は主にバックプロパゲーション（Backpropagation）に基づいて構築されてきたが、これは微分可能性を必要とする誘導バイアスである。
• 地球上の利用可能な計算資源が増加するに伴い、汎用で brute-force（蛮力）的な学習アルゴリズム、すなわち探索に基づくアプローチへの転換が期待される（Sutton, 2019）。
  • AlphaGo Zero は、計算資源に比例してスケールする汎用な手法として人間のデータに基づいた Bootstrapping を凌駕した好例である。
• 高計算資源領域において、微分可能性は動作可能なアーキテクチャの探索範囲を制限する要因となり得る（Liu et al., 2020）。より柔軟な探索に基づくクレジット割り当てアルゴリズムが、大幅に優れた一般化性能への一歩となる。

Dust の概要

• 目的: バックプロパゲーションを置き換え、解析的な構造よりも計算蛮力に基づいた学習アルゴリズムを採用すること。
• 基本原理:
  • スカラー階級ゼロ最適化法であり、アクティベーションの擾乱を行い、各擾乱による損失低下量に応じて報酬を割り当てる。
  • 集団全体で重み付けされた擾乱を平均することで勾配を推定する。
• 重量空間 ES との違い:
  • 従来の重量空間ベースのエバリュエーション戦略（例：EGGROLL）は、各メンバーを「materialize（実装・物質化）」し評価する必要がありコストがかかる。
  • Dust はアクティベーション空間を直接操作する「仮想的な集団（virtual population）」の概念を用い、重み空間を全く経由せず、ノード擾乱のようにアクティベーションそのものを擾乱する。
• 主要な貢献:
  1. Transformer の事前学習において、バックプロパゲーションと競合しうる初めてのスカラー階級ゼロ法を提示した。大規模な集団において複数の設定でバックプロパゲーションを上回る性能を示している。
  2. 重量空間ベースの ES より桁違いに効率的であり（$10^3 \sim 10^4$ 倍）、過剰パラメータ化を「より大きな探索空間」と捉える新たな視点を提供した。
  3. トークン予算や集団サイズに対して広く一般化するチューニング可能であることを実証した。

---

手法

アクティベーション空間における擾乱（2.1）

• 概念: 重みではなくアクティベーションそのものを擾乱する「仮想的な集団」を採用する。
  • すべてのトークンに対して独立にガウスノイズを加え、順伝播を行う。
  • これにより、1 回の順伝播で数千個のトークン（メンバー）を並列評価できる。
• 次元論的根拠:
  • 層の出力次元は $d_{\mathrm{out}}$、重み行列の次元は $d_{\mathrm{out}} \times d_{\mathrm{in}}$ であるため、アクティベーションノイズははるかに小さな空間に存在する。
  • しかし、シーケンスごとのノイズは $T \times d_{\mathrm{out}}$ のテンソルとなり、トークン数 $T$ が入力次元 $d_{\mathrm{in}}$ を超えると重み行列以上の成分数を持つことになる。
• 効率的な集団: トークンごとの独立した擾乱と報酬を用いることで、バッチ軸とは直交する新たな効率的な集団を提供可能となる。

クレジット割り当て（2.2）

• 報酬の計算: 各トークン $s$ について、中心化された損失削減量 $c_s = \tilde{\ell}_s - \ell_s$ を計算し、これに減衰係数 $\gamma$ を用いてその後続のトークンの影響を加算する。
  • $$r_t = \sum_{s \ge t} \gamma^{,s-t} c_s$$
• 勾配推定: 報酬と重み付けされたノイズをサンプル間で平均したものを層出力における推定誤差とし、入力との外積をとることで重量勾配を得る。
  • $$\hat g_t = -\frac{1}{K\sigma}\sum_{i=1}^{K} r_t^{(i)} a_t^{(i)}$$
  • $$\widehat{G}_W = \sum_t \hat g_t, x_t^\top$$
• 大規模集団での挙動: 無制限の集団サイズにおいて、勾配は線形に積み上がり、誤差は平方根の法則に従うため比率が低下し、極限で干渉が消失する。

相互干渉とチューニング（2.3）

• 干渉の軽減策:
  1. 層タイプの独立化: 異なる層タイプを独立した順伝播で微調整し、クリーンな順伝播をキャッシュする。
  2. アテンション内部の独立化: クエリ、キー、バリューなどを独立して微調整する。
  3. 言語モデリングヘッドの最適化: キャッシュされたログジット上直接に微調整し、交叉エントロピーと語彙の一部のみを再評価する。
• アテンションスコア: 推定されたアテンション出力勾配との整合性をスコアとして用い、損失削減量の代わりに報酬を計算する。
  • $$c_s = -\langle \hat g_s, \Delta o_s \rangle$$
• チューニング方法:
  1. グリッドサーチ（学習あり）: 小さなトークン予算で学習を行い、損失を低下させる設定を保持する。
  2. グリッドサーチ（学習なし）: 推定量とバックプロパゲーション勾配のコーサイン値を最大化する設定を選ぶ（訓練不要）。

---

バックワードパスなしでの事前学習

セットアップ（3.1）

• データセット: FineWeb を使用し、GPT 様変換器を訓練。
• 設定: 4096 トークン BPE、バッチサイズ 16k トークン（2048 シーケンス 8 個）、学習エポック数 1。
• 比較手法: Dust, EGGROLL, バックプロパゲーション（Adam/SGD）を同一のプロトコルで比較。
  • Dust: draw を集団サイズとしてカウント（クリーンな順伝播がキャッシュされているため安価）。
  • EGGROLL: バッチの順伝播回数をカウント。

主な結果（3.2）

• バックプロパゲーションとの比較:
  • 小規模なトークン予算では Dust は劣るが、集団サイズが増えるにつれて差が縮まり、大規模な予算（1000 万〜2000 万）では上回る傾向にある。
  • 重量空間ベースの ES（EGGROLL）は非効率的で、Dust に比べて数千から約 $10^4$ 倍の集団サイズが必要となる場合がある。
• 訓練中の乖離: Dust はバックプロパゲーションの曲線に沿って推移する一方、EGGROLL は早期に飽和する傾向がある。

Dust と Adam の組み合わせ（3.3）

• Adam への遷移: Adam を用いた再検証では、EGGROLL は性能向上が限定的であったが、Dust とバックプロパゲーションは改善が見られた。
  • チューニングされた Adam ラダーと SGD ラダーの間には 0.01 以内の差しかない。
• 共進化の可能性: Dust の推定量は現代の最適化アルゴリズム（Adam）でも機能し、さらに性能向上が見込める。

---

高次元空間における探索

過剰パラメータ化（4.1）

• 通説への挑戦: スカラー階級ゼロ法は大規模ネットワークを訓練できないという通説に反する結果を得た。
  • 4 つのモデルサイズ（2M, 7M, 38M, 243M パラメータ）を固定トークン予算で訓練。
• 驚くべき発見:
  1. より大きなモデルは集団効率が良い: 大規模なモデルほど、同じ集団サイズでも損失が低く、特に大規模な集団において顕著に優れた性能を示した。
  2. 大規模モデルは集団サイズとともに改善し続ける: 微小なモデルは早期に飽和するのに対し、大規模なモデルは大規模な集団サイズでも改善を続ける（約 30% の性能向上）。
• 解釈: より大きなモデルはより大きい探索空間を持ち、良好な幾何構造を持つ損失ランドスケープを提供するため、大規模な集団を活用できる。

バックプロパゲーション様の勾配の出現（4.2）

• コーサイン値の向上: トークン数が 10M から 1B に跨る際でも、Dust の推定勾配とバックプロパゲーション勾配の間のコーシン値は集団サイズとともに維持・向上する。
  • $$\cos(K) = \frac{c_{\max}}{\sqrt{1 + c/K}}$$
• EGGROLL との違い: EGGROLL は 128k の集団サイズでもコーサイン値 0.05 未満で頭打ちとなる可能性がある。
• 拡張性の保証: トークン数が増加してもコーシン値が維持されるため、十分に大きな集団においては、勾配推定の大規模化の必要性が頭打ちになる可能性を示唆している。
  • Dust の勾配はバックプロパゲーションとは類似な方向を指すため、異なる最適化軌跡を生じさせ、パフォーマンスが向上する可能性がある。

---

おわりにと関連研究

結論

• 計算資源が豊富になるにつれ、バックプロパゲーションを超える代替案が可能である。
• Dust は既存の ES アルゴリズムを大幅に改善し、大規模な計算資源を用いることでバックプロパゲーションを凌駕する。
• 未解決の問題:
  1. 損失ランドスケープの暗黙的な探索による高次勾配の利用可能性。
  2. エンドツーエンド微分性を必要としないため、バックプロパゲーションで困難な領域（時間伝搬勾配など）での応用可能性。
  3. 計算効率性のさらなる向上と実用的な代替案としての確立。

関連研究との位置付け

• 初期 ES: 重量摂動を用いたゼロからの変換器訓練の困難さを検討した全体的な研究に対し、Dust はアクティベーション空間を探索し、順伝播からより多くのクレジットを引き出す。
• ファインチューニング: MeZO や Scale of Evolution Strategies などが示すように、ES は適応的であるが、Dust はゼロからの事前学習を通じた表現そのものの学習を扱う。
• ノード擾乱: GEMINI や Zoop などの過去のノイズ注入手法とは異なり、各トークンに対する報酬とアテンション出力の局所的ターゲットを組み合わせる点で独自性を持つ。