技術的サマリー：カテゴリカルフローマップの拡張

問題定義

自己回帰（AR）モデルが言語モデリングを支配する一方、連続拡散モデルやフローマッチングモデルは、加速されたサンプリングや分布の傾き付けなど、連続モダリティにおいて顕著な利点を示しています。これらの恩恵を離散データに拡張することが、最近の研究の主要な目標となっています。従来のアプローチ、例えば離散拡散モデルは離散グリッド上で動作するため、確立された連続技術の応用が制限されていました。一方、連続状態プロセス（例えばカテゴリカルフローマップ、CFM）は、連続単体または $\mathbb{R}^d$ 上の確率分布を進化させることで離散データを生成します。

CFM が少数ステップの領域で競争力のあるサンプル品質を達成することが最近実証されたにもかかわらず、そのスケーラビリティは未証明のままでした。CFM が現代の大規模言語モデル（LLM）に必要な数十億パラメータ、数兆トークンの領域に拡張できるかどうかは不明確でした。主な技術的障壁は計算およびメモリオーバーヘッドです：語彙サイズ $|V|$ とコンテキスト長 $L$ に対して、CFM は各ステップで $L \times |V|$ のロジット行列を具体化（materializing）する必要があり、$|V| \approx 200,000$ で $L$ が大きい場合、これは許容できないほど高くなります。

手法

著者は、2.1 兆トークンでトレーニングされた 17 億パラメータモデルに CFM を拡張するための 2 段階トレーニングパイプラインを提案します。

1. フローモデルの事前学習

第一段階では、確率フローのドリフトを近似する標準的なフローマッチングモデル（デノイザー）のトレーニングを行います。

• アーキテクチャ: リプシッツ制約に従って安定性を確保した、AdaLN-Zero 変調、ロータリー位置エンベディング、SwiGLU MLP を備えた 17 億パラメータのトランスフォーマ。
• 目的: 予測された終点と真のデータ間のクロスエントロピー損失（式 2.10）を最小化するようにモデルをトレーニングし、実質的にベクトル場 $v_\theta$ を学習します。
• 設計空間の探索: スケーリング前に、著者は 4 億パラメータの代理モデル上で広範なハイパーパラメータスイープ（350 以上の構成）を行いました。主な発見は以下の通りです：
  • 時間スケジューリング: 線形と「誤り復号（error decoding）」スケジュールの混合（$\alpha_t = (1-\lambda)t + \lambda\gamma_t$）で $\lambda \approx 0.75$ とした場合が最適なパフォーマンスを示しました。純粋な誤り復号（$\lambda=1$）はパフォーマンスの崩壊を引き起こしました。
  • 適応的損失重み付け: 予測誤差に基づいて損失を再重み付けすること（$w(x) = \text{sg}((\Delta(x) + c)^{-r})$）は生成パープレキシティ（Gen PPL）を改善しましたが、分布忠実度（MAUVE*）を維持するために慎重な調整が必要でした。
  • ランダムアンマスク: ランダムなプレフィックスを確率 $p > 0$ でアンマスクしてトレーニングすることは、一貫性のある条件付きテキストを生成するために不可欠でした。これがないと、モデルは分布外出力を生成しました。

2. カテゴリカルフローマップ（CFM）への自己蒸留

第二段階では、事前学習されたフローモデルを、高品質かつ少数ステップでの生成が可能な CFM に変換します。

• メカニズム: モデルは、ODE を多数ステップにわたって積分するのではなく、確率質量を時間 $s$ から $t$ へ単一ステップで輸送するフローマップ $X_{s,t}$ を学習します。
• 損失関数: 著者は、プログレッシブ自己蒸留（PSD）と安定化オイラーロジット自己蒸留（ESD）を評価しました。
  • PSD: 連続時間ショートカットモデルアプローチを使用します。
  • 発散の選択: 著者は Forward KL、Reverse KL、およびジェンセン・シャノン発散（JSD）を比較しました。Forward KL と JSD は少数ステップ領域でトークンエントロピーと MAUVE* を改善しましたが、Reverse KL は不安定さとモード崩壊を引き起こしました。
  • 安定性: 17 億パラメータ規模での ESD トレーニングは、有限差分で計算された高分散のヤコビアン推定値により不安定であることが判明したため、著者は PSD の結果に焦点を当てました。
• サンプリング: 最終的な CFM は、事前分布からサンプリングし、学習されたフローマップ $X_{0,1}$ を単一パス（または少数パス）で適用することでテキストを生成し、推論時間を大幅に削減します。

3. 尤度推定

標準的な尤度ベースのベンチマーク（例：MCQA）で CFM を評価するために、著者は半離散 ELBOを導出しました。この境界は、高次元の離散設定で特にコストがかかる ODE の逆積分を必要とせずに、回答選択肢にスコアを付けることを可能にします。

主要な貢献

1. スケーラビリティの実証: 著者は 2.1 兆トークンで 17 億パラメータのフローモデルのトレーニングに成功し、それを CFM に自己蒸留することで、CFM が数十億パラメータの領域に拡張できることを証明しました。
2. ハイパーパラメータの発見: 4 億パラメータの代理モデル上での体系的なスイープを通じて、時間スケジューリング、損失重み付け、入力マスクに関する堅牢な設定を特定し、これらがより大規模なスケールへ転移することを確認しました。
3. 少数ステップ生成: 自己蒸留された CFM は、データレベルに近いトークンエントロピーを維持しながら、わずか 4 推論ステップ（NFE）で多様性が高く高品質なテキストを生成し、低 NFE 領域において非蒸留フローモデルを上回ります。
4. 尤度境界: 半離散 ELBO の導入により、CFM を標準的な言語モデリングベンチマークで評価することが可能となり、そのパフォーマンスを他の非自己回帰手法と同様の範囲に位置付けることができました。
5. 限界の分析: 本論文は、$L \times |V|$ 行列の具体化に伴うメモリボトルネックと、時間条件付きアーキテクチャにおけるゼロショットハイパーパラメータ転送の課題を明示的に扱っています。

結果

• 品質とステップ数の関係: 自己蒸留された CFM は、1024 ステップで約 58.66 の Gen PPL を達成し（AR ベースラインの 45.33 と競争力あり）、MAUVE* スコアは約 0.94 です。重要なのは、わずか 4 ステップで CFM が約 0.88 の MAUVE* を達成し、少数ステップ領域において AR ベースライン（0.7）を大幅に上回っていることです。
• トークンエントロピー: 自己蒸留は、低 NFE でも高いトークンエントロピー（多様性）を維持しますが、非蒸留モデルは 1〜2 ステップの領域で崩壊（低エントロピー）する傾向があります。
• ベンチマークパフォーマンス: MCQA ベンチマーク（ARC-E、BoolQ、OBQA など）において、CFM は他の非自己回帰手法（例：Eso-LM）と同等の精度を達成しますが、ELBO 最大化に直接トレーニングされた離散拡散モデルには及びません。
• 条件付き生成: 自己蒸留モデルは、非蒸留モデルと比較して、少数ステップ領域で $\Delta$PPL（条件付きパープレキシティ）が大幅に改善されました。

意義と主張

本論文は、カテゴリカルフローマップが現代の LLM の規模における離散データ生成の viable な道であることを主張しています。17 億パラメータへの拡張に成功したことで、著者はガイダンス、分布の傾き付け、報酬微調整を含む「フローマップツールボックス」が言語モデリングで利用可能になったことを実証しました。

この研究は、連続状態生成が計算リソースをサンプル品質と交換でき、従来の拡散モデルや AR モデルよりも著しく少ない関数評価で高品質なテキスト生成を可能にすることを示唆しています。著者は尤度推定が（導出された ELBO が境界であるため）依然として改善の余地がある分野であると指摘していますが、結果は CFM が離散拡散ベースラインと競争力のあるパフォーマンスを達成しつつ、連続フローマッチングの恩恵を解放できることを示しています。本論文は、CFM のスケーラビリティが単なる妥当な出力生成を超えて言語モデルの制御や微調整のための新たな道を開くと結論付けています。