Long-Context Training: 長文脈学習

概要

Long-Context Training は、Molmo2 の Stage 3 として実施される最終的な訓練段階である。Stage 2 の Supervised Fine-Tuning (SFT) と同じデータミックスを使用しながら、コンテキスト長を大幅に拡張 することで、長尺ビデオや大量の画像を扱う能力を向上させる。

この段階は、オーバーヘッドが大きいため 短期間のみ実施 されるが、長尺ビデオ理解タスクでの性能向上に重要な役割を果たす。

目的

Long-Context Training の主な目的は以下の通りである。

  1. 長尺ビデオの理解: 10分以上の長い動画を処理する能力の向上
  2. 大量フレームの処理: より多くのフレームを同時に扱うことで、時間的な文脈をより正確に把握
  3. 複雑なマルチモーダル入力: 字幕付き動画や、複数の画像セットなど、トークン数が多い入力への対応

Stage 2 との比較

Stage 2 (SFT) と Stage 3 (Long-Context Training) の主な違いは以下の表の通りである。

パラメータ Stage 2 (SFT) Stage 3 (Long-Context) 変化率
シーケンス長 16,384 36,864 +225%
最大フレーム数 (F) 128 384 +300%
訓練ステップ数 30,000 2,000 -93%
バッチサイズ 128 128 変更なし
データミックス SFT データミックス SFT データミックス（同一） 変更なし
並列化手法 標準的な Data Parallelism Context Parallelism (CP) 追加
Noteなぜ短期間のみ実施するか

Long-Context Training は 2,000 ステップ のみ実施される（Stage 2 の30,000ステップと比較して6.7%）。これは、Context Parallelism による 計算オーバーヘッドが大きい ためである。

具体的には以下の通りである:

  • 各例が 8 GPU のグループ で処理される
  • Vision encoder と attentional pooling の分散処理が必要
  • 通常の訓練よりも通信コストが高い

短期間の訓練でも、長尺ビデオベンチマークでの性能が 有意に向上 することが確認されている（Table 11）。

コンテキスト長の拡張

Stage 3 では、最大シーケンス長を 16,384 → 36,864 トークンに拡張する。これにより、以下のような入力が扱えるようになる。

例:

  • 長尺ビデオ: 384 フレーム（2 fps で約3分12秒）+ 字幕
  • 複雑な QA: 長いキャプション（平均924語）を含む動画に対する複数ターンの会話
  • マルチ画像: 大量の高解像度画像（複数のクロップ含む）

シーケンス長の拡張により、packing アルゴリズムは最大 36,864 トークン と 384 画像クロップ を単一のパックシーケンスに詰め込むことができる（Stage 2 では16,384トークン、128クロップが上限）。

フレーム数の拡張

Stage 3 では、動画の最大フレーム数を F = 128 → F = 384 に拡張する。

サンプリング方法:

  • サンプリングレート: 2 fps（変更なし）
  • 最大長: 384 / 2 = 192秒（約3分12秒）
  • フレーム選択: 動画長が F/S を超える場合、F フレームを均等にサンプリング
  • 最終フレーム: 常に含まれる（動画プレイヤーが最終フレームを表示するため）
Stage 最大フレーム数 (F) 最大動画長（2 fps）
Stage 2 (SFT) 128 64秒（約1分）
Stage 3 (Long-Context) 384 192秒（約3分）
Tip推論時のフレーム数拡張

訓練時は F = 384 であるが、推論時には テストタイムスケーリング を使用して、さらに多くのフレームを処理することが可能である。

アブレーション（Table 13）によると、Molmo2-8B（SFT後、Long-Context訓練前）は、推論時に 224 フレーム で最良の性能を示した。Long-Context訓練後は、さらに多くのフレームを処理できる可能性がある。

Context Parallelism (CP)

Long-Context Training では、LLM に対して Context Parallelism (CP) を使用する。これは、長いシーケンスを複数の GPU に分散して処理する並列化手法である。

Ulysses Attention

Molmo2 は、CP の実装として Ulysses attention を採用している。

選択理由:

  • All-gather 操作の柔軟性: Molmo2 の packing と message tree システムで使用される カスタムアテンションマスク に対応可能
  • 通信効率: シーケンス次元を分割し、必要な情報のみを all-gather

動作概要:

  1. 各 GPU が シーケンスの一部 を処理
  2. Attention 計算時に all-gather で他 GPU の情報を集約
  3. カスタムアテンションマスク（packing や message tree 用）を適用
NoteContext Parallelism の設定
  • CP グループサイズ: 8 GPU
  • 各例の処理: 8 GPU のグループで1つの例を処理
  • 適用範囲: LLM のみ（Vision encoder は別途分散）

参考文献: [56] Ulysses attention

Vision Encoder の分散処理

Context Parallelism は LLM に適用されるが、Vision encoder と attentional pooling も CP グループ全体に分散される。

分散方法:

  1. フレーム分割: 384 フレームを 8 GPU に分割（各 GPU が48フレームを処理）
  2. Vision encoder: 各 GPU が担当フレームの Vision Transformer 処理を実行
  3. Attentional pooling: Vision token を LLM 用にプーリング（3x3 pooling for video）
  4. 統合: プーリングされた視覚トークンを LLM の入力として結合

効果:

  • メモリフットプリント削減: Vision encoder のアクティベーションを複数 GPU に分散
  • 計算効率: フレーム処理を並列化
ImportantVision Encoder 分散の重要性

論文では、Vision encoder と attentional pooling の分散処理が 非常に効果的 にモデルのメモリフットプリントを削減したと述べられている。

384 フレームの処理は膨大なメモリを必要とするため（各フレームが複数のパッチに分割され、ViT で処理される）、この分散処理なしでは訓練が困難であった。

訓練の詳細

ハイパーパラメータ

Long-Context Training では、Stage 2 (SFT) と 同じハイパーパラメータ を使用する（Table 12）。

  • Optimizer: AdamW
  • Learning rate: Cosine decay（10% まで減衰）
  • Warmup: ViT と LLM で長めのウォームアップ
  • Weight decay: なし

訓練時間

モデル GPU 数 訓練時間 GPU 時間
Molmo2-4B 128 25.3 時間 3,200 GPU-hr
Molmo2-O-7B 128 25.7 時間 3,300 GPU-hr
Molmo2-8B 128 26.0 時間 3,300 GPU-hr

注: GPU は Nvidia H100 を使用

Stage 2 (SFT) と比較すると、訓練時間は約40%（4B: 7.5k → 3.2k GPU-hr）である。これは、ステップ数が少ない（2k vs 30k）ためである。

アブレーション結果

Long-Context Training の効果を検証したアブレーション（Table 11）の結果は以下の通りである。

設定 短尺ビデオ QA 長尺ビデオ QA Molmo2 Video Cap. 画像 QA
Long-Context SFT あり 69.4 67.4 39.9 80.6
Long-Context SFT なし 69.6 64.4 42.3 80.5

観察:

  • 長尺ビデオ QA: +3.0 ポイント向上（67.4 vs 64.4）
  • 短尺ビデオ QA: ほぼ変化なし（-0.2 ポイント）
  • ビデオキャプション: -2.4 ポイント低下（42.3 → 39.9）
  • 画像 QA: ほぼ変化なし（+0.1 ポイント）
Warningトレードオフ: キャプション性能の低下

Long-Context Training は長尺ビデオ理解を向上させるが、ビデオキャプション の性能がわずかに低下する。

これは、長いコンテキストに適応する過程で、短い出力（キャプション）の生成能力が若干犠牲になる可能性を示唆している。実用上は、タスクに応じて Long-Context Training の有無を選択することが重要である。

長尺ビデオでの課題

Molmo2 は Long-Context Training により長尺ビデオ理解が向上したが、最良のオープンウェイトモデル（Eagle2.5-8B など）には及ばない。

主な原因:

  1. オープンデータの不足: 10分以上の長尺ビデオのアノテーションデータが極めて少ない
  2. 計算制限: 超長尺コンテキストの訓練は計算コストが高く、大規模実施が困難
  3. 訓練期間の制約: わずか2,000ステップの訓練では、長尺ビデオに特化した能力の獲得が限定的
Note今後の展望

論文では、オープンソースの長尺ビデオデータ の不足が主要なボトルネックであると述べられている。

コミュニティが長尺ビデオのアノテーションデータを構築し、より長期間の Long-Context Training を実施できれば、Molmo2 の長尺ビデオ理解能力はさらに向上する可能性がある。

まとめ

Long-Context Training（Stage 3）は、Molmo2 の訓練パイプラインの最終段階として、以下を実現する。

  1. コンテキスト長: 16,384 → 36,864（+125%）
  2. フレーム数: F = 128 → F = 384（+200%）
  3. 並列化: Context Parallelism（Ulysses attention）を使用し、8 GPU で処理
  4. Vision encoder 分散: フレーム処理を CP グループ全体に分散してメモリ削減
  5. 短期間訓練: わずか2,000ステップ（オーバーヘッドのため）

この短期間の訓練でも、長尺ビデオ QA で +3.0 ポイント の性能向上を達成しており、Long-Context Training の有効性が示されている。ただし、ビデオキャプションの性能がわずかに低下するトレードオフも確認されている。

今後、オープンソースの長尺ビデオデータが充実し、より長期間の Long-Context Training が可能になれば、Molmo2 の長尺ビデオ理解能力はさらに向上する可能性がある。