Learning What to Predict: Downstream-Guided Task Design for Continued Pretraining

技術要約：何を予測すべきかを学習する：継続的事前学習のためのダウンストリーム誘導型タスク設計

1. 問題提起

継続的事前学習（Continued pretraining）は、基盤モデルを新しいドメイン（数学的推論や高密度なビジョン予測など）に適応させるための標準的な手法である。しかし、現在の慣習は粗いフィードバックループに依存している。最適化の単位は単一の自己教師あり更新（例：次トークン予測やマスクされた再構成）であるが、選択の単位はダウンストリームのベンチマークによって評価される事前学習のプロセス全体である。

実務者は通常、チェックポイントを評価し、データ混合や目的関数を手動で修正して、実行を再開する。このプロセスにより、個々の事前学習の更新が、特定のダウンストリーム能力に寄与しているかどうかに「盲目」になってしまう。ファインチューニング（SFT、RL）はダウンストリームのラベルを使用して学習者を直接更新するが、これは自己教師ありの更新を置き換えるものであり、ガイドするものではない。本論文は次のように問いかける：検証可能な少数のダウンストリーム・サンプルが、学習者を直接教師なしにすることなく、また汎化性能を損なうことなく、継続的事前学習中のステップレベルのフィードバックを提供できるか？

2. 手法：V-pretraining

著者らは、学習プロセスを2つのコンポーネントに分離するV-pretraining（Value-based pretraining）というフレームワークを導入している。

1. 学習者（The Learner）: ラベルのないストリーム上の自己教師あり損失を通じてのみ訓練される。そのパラメータ（$\theta$）は、ダウンストリームの勾配によって更新されることは決してない。
2. タスク設計者（The Task Designer）: ラベルのないバッチに対して、自己教師ありタスク（例：ターゲットやビューの選択）を動的に構築する、軽量なモジュール（パラメータ $\phi$ を持つ）。

コアメカニズム:
設計者は、ダウンストリームのラベルを含む少数のフィードバックセット（$D_{fb}$）を使用して、候補となる自己教師あり更新の「価値（Value）」を推定する。

• 勾配のアライメント（Gradient Alignment）: 提案された自己教師あり更新の勾配を $g_{pre}$、フィードバックセットにおけるダウンストリーム損失の勾配を $g_{down}$ とする。
• 価値の目的関数（Value Objective）: 本論文は、ダウンストリーム損失の一階の変化が、内積 $V = g_{down}^\top g_{pre}$ によって近似されると仮定している。
• 学習ループ:
  1. 設計者が、現在のラベルなしバッチに対してタスク（例：ソフトターゲット分布や特定のビュー）を構築する。
  2. 設計者は、$V$ を最大化するように更新される（すなわち、自己教師ありの勾配をダウンストリーム損失を減少させる方向へとアライメントさせる）。
  3. 構築されたターゲットまたはビューには、デタッチ（detach）（ストップグラディエントの適用）が行われる。
  4. 学習者は、これらのデタッチされたターゲットを用いて、標準的な自己教師あり損失を用いて更新される。

具体例（Instantiations）:

• 言語（Language）: 設計者は、次トークン予測のための適応的なTop-Kソフトターゲットを学習する。ハードなOne-hotターゲットの代わりに、設計者は、真のトークンとTop-$K-1$ のモデル予測を含む小さな候補セットに対する確率分布を、ゲーティング定数によって重み付けして出力する。
• ビジョン（Vision）: 設計者は、DINOスタイルの自己教師あり学習のためのインスタンスごとのビューまたはマスクを学習し、入力のクロップやオーグメンテーションを修正して、ダウンストリームの密な予測タスク（セグメンテーション、深度推定）により良く適合させる。

3. 主な貢献

• 概念的貢献: 本論文は、ダウンストリーム誘導型の継続的事前学習を学習可能なタスク構築として再定義している。これは、自己教師ありの更新を教師ありの更新に置き換えるのではなく、間接的（学習者に対しては間接的だが、設計者に対しては直接的）なガイダンスの経路を導入している。
• アルゴリズム的貢献: ダウンストリームの改善に対する一階のサロゲート（代理）として機能する、計算効率の高いステップレベルの目的関数（$g_{down}^\top g_{pre}$）を提案している。これにより、フルな事前学習の軌跡を微分するという困難なバイレベル最適化を回避している。
• 実証的貢献: 計算量（Wall-clock time）を一致させた広範な実験により、V-pretrainingが汎用的な表現を崩壊させることなく、ターゲットとなる能力を向上させることを示している。

4. 実験結果

著者らは、言語およびビジョンのモダリティにおいて、計算量（Wall-clock time、ハードウェア、初期条件が同一）を一致させた実験を行った。

言語（数学的推論）:

• 設定: 1,024個のGSM8Kサンプルをフィードバックとして使用し、NuminaMath-CoT上でQwenモデルの継続的事前学習を行う。
• 結果: V-pretrainingはGSM8KのPass@1を大幅に向上させた。
  • Qwen2.5-0.5B: 単一実行での利得は +7.4 ポイント（22.20 $\to$ 29.60）であり、33%の相対的改善を示した。
  • Qwen1.5-4B: 56.48から58.98へと向上した。
• 汎化: この手法は広範な能力を損なわなかった。MMLUスコアは4Bおよび7Bモデルで安定していたが、0.5Bモデルではわずかな低下が見られ、小さな学習者は狭いフィードバックに対してより敏感であることを示唆している。

ビジョン（密な予測）:

• 設定: ADE20K（セグメンテーション）およびNYUv2（深度）をフィードバックとして、ImageNet-1K上でDINOv3バックボーンの継続的自己教師あり学習を行う。
• 結果:
  • ADE20K: ViT-Lにおいて、mIoUが51.33から52.40に向上した。
  • NYUv2: ViT-Lにおいて、RMSEが0.5752から0.5522に減少した。
  • 汎化: ImageNetの線形分類精度は維持またはわずかに向上しており、モデルがグローバルな認識を犠牲にして密なタスクに過剰特化していないことを示している。
  • 転移: Oxford/Parisデータセットにおけるインスタンス検索において、ほとんどのプロトコルで改善が見られ、表現が依然として堅牢であることを裏付けている。

アブレーションとコントロール:

• 汚染（Decontamination）: 学習ストリームからGSM8K/MATHの近接重複を削除しても利得は消失せず、単純な記憶（Memorization）ではないことが示された。
• ショートカットの排除: 利得は、一般的なスムージング、自己蒸留、またはランダムノイズによるものではない。ランダムなフィードバックや一様スムージングを用いたバリアントは、アライメントされたV-pretrainingよりも性能が低かった。
• オーバーヘッド: タスク設計者は、トークンスループットの約9.4%の減少と、ピークメモリの約16.4%の増加をもたらすが、これらのコストはウォーククロック一致比較の中で考慮されている。

5. 意義と主張

本論文は、V-pretrainingが直接的なポストトレーニング手法（SFTやRLなど）に対する計算効率の高い補完策を提供すると主張している。

• 効率性: 継続的事前学習フェーズの最中にダウンストリームの能力信号を注入することを可能にし、後の高価なダウンストリーム・ファインチューニングの必要性を減らす可能性がある。
• メカニズム: ダウンストリームのフィードバックが、学習者そのもの（どのようにモデルを更新するか）ではなく、自己教師ありタスクの構築（何を予測するか）を制御できることを示している。
• 謙虚な姿勢: 著者らは、価値目的関数が長期的な収束定理ではなく、局所的な一歩の保証（即時的なダウンストリーム改善の下限を最大化すること）を提供するものであると明示している。彼らは、手法が小さなモデルを狭いタスクへと過剰に誘導する可能性があり、アライメント信号がノイズを含み得ることを認めている。

要約すると、V-pretrainingは、少数のダウンストリーム・サンプルが、基礎モデルの汎化特性を維持しながら、継続的事前学習の「何を予測すべきか」という側面を効果的に導くことができることを確立している。