Mixture of Training: Recombining Small-Scale Scaffolded Pretraining Runs into a Larger Language Model

技術要約：Mixture of Training (MoT)

問題提起
大規模言語モデル（LLM）の事前学習は、現在、単一のモノリシックなエンドツーエンドの最適化プロセスとして構成されている。このアーキテクチャは、すべてのレイヤーを同時に学習させる必要があるため、スケーリングのボトルネックを生じさせる。これにより、プロセスは失敗に対して脆弱（単一のエラーがプロセス全体を停止させる）であり、反復的な研究に対しても柔軟性に欠ける（改善を行うには、巨大で結合されたシステムを再開または継続する必要がある）。本論文は、事前学習を、品質を損なうことなく、後に一貫したより大きなモデルへと再構成可能な、より小さく独立して学習可能なジョブへと分解できるか否かを問うものである。

手法：Mixture of Training (MoT)
著者らは、ターゲットとなるTransformerを連続したレイヤーブロックに分割するように設計された、足場（scaffolded）を用いたモジュール式事前学習手順である**Mixture of Training (MoT)**を導入する。この手法は、以下の3つのステージで動作する。

1. ステージ 0（準備）： 学習済みのアライナー（aligner） ($A$) を選択または準備する。ターゲットとなるTransformer ($F$) とアライナーは、共に $K$ 個の連続したレイヤーブロックに分割される。アライナーは、グローバルな幅、アテンション・ヘッドの次元数、フィードフォワード幅、トークン埋め込み、および出力ヘッドを共有し、ターゲットと形状の互換性を持つ。
2. ステージ 1（足場付き学習）： ターゲットモデルを $K$ 個のサブモデル ($f_1, \dots, f_K$) に分解する。各サブモデル $f_i$ に対して、凍結されたアライナー構造の中に $f_i$ を挿入することで、足場付きネットワーク $S_i$ を構築する：
   $$S_i = a_K \circ \dots \circ a_{i+1} \circ f_i \circ a_{i-1} \circ \dots \circ a_1$$
   このセットアップでは、ターゲットブロック $f_i$ のみが学習対象であり、すべてのアライナースライスは凍結されている。各スカフォールド（足場）は、標準的な次トークン予測損失を用いて独立して最適化される。決定的なのは、このフェーズではターゲットブロック間で勾配の交換が行われないことであるが、すべてのブロックは共有されたアライナーによって提供される安定した表現コンテキスト内で学習を行う。
3. ステージ 2（再構成と適応）： アライナーは破棄される。学習済みのターゲットブロックは、フルモデル $\hat{F} = f_K \circ \dots \circ f_1$ へと再構成される。この「コールド・コンポーズド（cold-composed）」モデルは、直接評価することも、あるいは残存する表現の不一致を解消するために、短いエンドツーエンドの適応パスを経由させることもできる。

主な貢献

1. 足場付きモジュール学習： 共通のアライナーを介して共有された表現インターフェースを維持しながら、ターゲットの深さ方向のスライスを独立して学習させ、事後的なレイヤーの継ぎ合わせなしに再構成を可能にする手順。
2. メカニズムの証明： 独立して学習された深さ方向のスライスが、使用可能な1.3Bパラメータの言語モデルへと正常に再構成できることを示す小規模な研究。
3. 計算量とスケジューリングの会計： 計算量（FLOPs）、トークン露出、および理想的なレイヤー等価クリティカルパスに関する詳細な分析。これには、アライナーのコストの明示的な扱いが含まれ、「直接的」な学習コストと、「完全充当（fully charged）」コスト（アライナーの全コストを単一の実行に割り当てる）および複数の実行にわたる「償却」コストを区別している。
4. アブレーション研究： アライナーの必要性、不連続なデータストリームの影響、および分割数 ($K$) の増加に伴う品質と効率のトレードオフを示す実証的証拠。

実験結果
著者らは、C4データセットを用いて、12レイヤー、1.3BパラメータのGemmaスタイルのモデルに対してMoTを評価した。

• 品質の同等性： サブモデルの学習を延長し、30kステップの適応パスを含めた「品質パリティ（quality-parity）」スケジュールは、モノリシックなベースラインと同じパープレキシティ（15.0）を達成した。
• 計算量のトレードオフ：
  • 「品質パリティ」スケジュールは、より多くの総計トークン（47.1B vs. 33.6B）を処理し、より短い理想的なクリティカルパス（1.7倍の高速化の可能性）を持っていたが、アライナーのコストを単一の実行に完全に充当した場合、より多くの総FLOPs（285.0 EFLOPs vs. 268.4 EFLOPs）を必要とした。
  • しかし、アライナーが $R$ 個の独立した実行に対して再利用される場合、実効コストは低下する。論文では、$R \ge 3$ の場合、実効コストがモノリシックなベースラインを下回ることが記されている。
  • より低計算量のスケジュール（MoT + 15k 適応）は、189.4 EFLOPs（完全充当時）において15.9のパープレキシティに達した。これはベースラインの予算を下回っているが、この特定の予算における直接比較のための、同等の計算量を持つモノリシックなベースラインは提供されていない。
• アブレーション：
  • アライナーの必要性： アライナーがない場合、コールド・コンポジションの品質は崩壊した（PPL 38.9）。これは、アライナーがインターフェースの不一致を防いでいることを裏付けている。
  • データストリーム： アライナーが存在する場合、サブモデルに対して不連続なデータストリームを使用する方が、共有データを使用する場合（PPL 20.3）よりもコールド・コンポジションの品質が向上した（PPL 19.3）。
  • 分割数： 分割数を $K=2$ から $K=4$ に増やすと、計算量は減少したが、コールド・コンポジションの品質は低下した（PPL 24.8）。これは、並列性とインターフェースの複雑さの間のトレードオフを浮き彫りにしている。

意義と主張
本論文は、MoTをモノリシックな事前学習の一般的な代替手段としてではなく、**「足場を用いたサブラン（sub-run）が、再利用可能な学習ユニットとして機能するかどうかを研究するための小規模なフレームワーク」**として位置づけている。

著者らは、主に以下の3つの結論を主張している。

1. 共有されたアライナーの足場の中で学習されるならば、独立して学習されたサブモデルは、一貫した言語モデルへと再構成できる。アライナーがない場合、コールド・コンポジションは失敗する。
2. 独立して学習されたブロック間の不一致は、大部分が回復可能である。短い適応パスによってその差の大部分は埋まり、より長いスケジュールを用いることで、モノリシックな学習と同等の水準に達することができる。
3. MoTは、事前学習のエンジニアリングの形態を変化させ、より小さく、独立してスケジューリング可能で、再開可能かつアブレーション可能なジョブを生み出す。これは特に小規模な学習研究に適しており、より安価で、再現性が高く、反復が容易になる可能性がある。

著者らは、現在のエビデンスは、一様な計算量やハードウェアの速度上の優位性を確立するものではなく、**「モジュール的な構成可能性」と「償却再利用スケジューリング・レジーム」**を実証しているものであると明示している。この研究は、具体的な設計空間としての「再利用可能な足場付きサブラン」を検証するための、意図的に限定された範囲（一つのモデルファミリー、一つのデータセット、パープレキシティのみの指標）で行われている。