技術要約: ハイパーエポック・プリトレーニングのためのプリミティブ (q0)

問題提起

言語モデルの分野は、計算量とデータの両方をスケーリングする体制から、データ制約のある体制へと移行している。計算能力は成長し続けている一方で、高品質なテキストの供給は有限であり、最先端のモデルはすでに利用可能なデータの相当部分を消費している。このような文脈において、標準的なマルチエポック学習（固定されたコーパスに対する繰り返しのパス）は収穫逓減を示します。モデルは通常、数回のパス内で飽和し、多大な計算予算が未利用のまま残されてしまいます。

本稿が扱う中心的な問いは、「固定されたデータセットと $N$ 回の学習エポックの予算がある場合、計算量をどのように割り当てるべきか？」という点である。従来のアプローチは単一のモデルの洗練に焦点を当ててきたが、著者らは、追加のエポックが単一のインスタンスを改善しなくなる場合、これは最適ではないと主張している。代わりに、彼らは単一の洗練されたモデルを訓練することから、**多様なモデルの集団（ポピュレーション）**を探索し、その予測を集約することへとシフトすることを提案している。

手法: ハイパーエポック・プリトレーニング (q0)

著者らは、マルチエポックの予算を多様なモデルの集団へと変換するフレームワークである q0 を導入する。この手法は、ナイーブなアンサンブル（高コストであり、同等の品質のモデルしか生成できず、一様な重み付けを用いる）の限界に対処する3つのコア・プリミティブに基づいている。

1. 重み空間の高速探索 (スナップショット・アンサンブル)

すべてのモデルをゼロから訓練するという膨大なコストをかけずに仮説空間を探索するために、q0 は少数の並列な軌跡にわたるサイクリック・スケジュールを採用している。

• メカニズム: 学習は $C$ 個のサイクルに分割される。各サイクルでは学習率 (LR) を高い値から再開し、低いフロアへと減衰させる。各サイクルの終了時に、パラメータをスナップショット ($\theta_{n,c}$) として保存する。
• 反相関ウェイトディケイ (Anti-correlated Weight Decay): 著者らは、ウェイトディケイ (WD) を学習率と反相関させている。高 LR/低 WD のフェーズは探索を促進し、低 LR/高 WD のフェーズは、スナップショットが取られる直前にパラメータを低ノルムの盆地（basin）へと引き込む。
• 結果: これにより、$M = N \times C$ 個のスナップショットが生成される。これらは個別に十分に収束しており（局所解に落ち着いている）、かつアンサンブルに価値をもたらすほど多様である。

2. チェイン蒸留による能力の累積 (Capability Compounding via Chain Distillation)

同一のデータに対する標準的な勾配降下法は、同程度の品質のモデルを生成する。集団の能力が停滞するのではなく累積的に向上することを確実にするため、q0 はチェイン蒸留を導入している。

• メカニズム: ウォームアップ期間の後、各新しいスナップショット $\theta_{n,c}$ は、グラウンドトゥルースのラベルだけでなく、現在のサイクル中に凍結された直前のモデル $\theta_{n,c-1}$ を教師として用いて訓練される。
• 損失関数: 損失は、標準的なクロスエントロピー ($L_{CE}$) と、教師のソフトターゲットに対するカルバック・ライブラー (KL) ダイバージェンス項を組み合わせたものである：
  $$L_c = (1 - \alpha) L_{CE}(\theta_{n,c}) + \alpha T^2 \cdot KL(\sigma(f(x; \theta_{n,c-1})/T) \parallel \sigma(f(x; \theta_{n,c})/T))$$
• 根拠: これは「ダークナレッジ」（教師の誤ったクラスにおけるクラス間類似性構造）と暗黙的な重要度重み付けを活用しており、サイクルを通じて集団の品質が累積的に向上することを可能にする。

3. 学習された汎化事前分布 (Learned Generalization Prior)

一様な平均化ではなく、q0 は保持された適合セット (fitness set) に基づいて、アンサンブルメンバーに対する最適な重み付けスキームを学習する。

• メカニズム: $M$ 個のスナップショットに対する分布は、制約のないロジット $\beta$ によってパラメータ化される。重み $w = \text{softmax}(\beta)$ は、保持された適合セットにおけるアンサンブルの平均負の対数尤度を最小化するように最適化される。
• 推論: いかなる推論予算 $K \le M$ に対しても、適合された重みが最も高い上位 $K$ 個のスナップショットが選択され、再正規化される。この単一の適合はすべての予算サイズに対して再利用され、補完的な仮説を経験的に選択することで、ソロモノフ的な複雑性の事前分布を近似する。

主な結果

実験は、1.8B パラメータのデコーダーのみのトランスフォーマーを用いて、100M トークンの FineWeb データセットで実施され、保持された検証セットおよびダウンストリーム・ベンチマーク (ARC-Easy, PIQA, SciQ) で評価された。

• データ効率:
  • q0 は、強力なベースライン・アンサンブル（32エポックずつ訓練された8つの独立したモデル、合計256エポック）に対し、わずか ~56 エポックを使用して同等の性能を実現した（4.6倍の削減）。
  • ベースラインのアンサンブルサイズ（固定 $k=8$）に合わせた場合でも、q0 は ~67 エポックで同等の性能に達した（3.8倍の削減）。
  • 「Slowrun」設定において、q0 は 960 エポックにおいてベースラインに対して累積 ~12.9倍のデータ効率を達成した。
• ダウンストリーム性能: プリトレーニングでの利得はゼロショット・ベンチマークに転移する。q0 は ARC-Easy, PIQA, SciQ のすべてにおいてベースラインを上回り、256エポックの予算で ~14.2倍、スケール版では ~16.0倍のデータ効率を達成した。
• 最適な割り当て: 著者らは、計算量を割り当てるための処方的なルールを提供している：
  • 小規模予算 ($E \lesssim 120$): 単一のベースモデルに多くのサイクル（または学習終了時のスナップショット）を適用するのが最適である。
  • 中規模予算: 並列な軌跡の数 ($N$) を増やすべきである。経験的に、$N=2$ は ~240 エポックまで有益であり、$N$ は予算に対して対数的に増加する ($N \approx \log_2(E/64)$)。
  • 大規模予算: $N \approx 3$ を超えると、ウォームアップのコストと多様性のトレードオフが飽和するため、限界的な利得は減少する。

意義と主張

本論文は、q0 が、データ制約のあるレジームにおける計算量の活用方法に関する概念的な転換を提示していると主張している。単一のモデルを限界点を超えて洗練させるのではなく、最適戦略は、多様で互いに情報を与え合うモデルの集団を重ね合わせることである。

• 効率性: この手法は、単一の洗練されたモデルやナイーブなアンサンブルよりも、はるかに少ないエポック数で、より低い検証損失と高い汎化性能に到達できることを示している。
• スケーラビリティ: 得られる利得は、広範な予算（1 から 960 エポックまで）にわたって持続し、最適な構成は体系的に変化する。
• 新規性: サイクリックなスナップショット、軌跡内のチェイン蒸留、および学習された汎化事前分布の組み合わせは、マルチエポック・プリトレーニングのための新しいプリミティブのセットを提供する。著者らは、集団の上での探索は、勾配降下法による単一モデルの訓練とは異なる問題であることを指摘しており、q0 が、限られたデータの時代における、より効率的な仮説空間の探索への第一歩であることを示唆している。

著者らは、アンサンブルには $K$ 回のフォワードパスが必要であるため、推論オーバーヘッドが主な制限事項であることを認めている。しかし、これは根本的な問題ではないと述べている。なぜなら、アンサンブルは単一の生徒モデルに蒸留可能であり、その方向性は彼らのレシピと互換性があるが、本研究では追求されていないためである。