技術的サマリー：マージナル条件付きブリッジによるフロー言語モデルからのサンプリング

問題定義

フロー言語モデル（FLM）は、ワンホットエンコードされたトークン系列に対して連続的なフローマッチングまたはノイズ除去を適用する、非自己回帰型言語モデルのクラスを表します。一般的な連続拡散モデルとは異なり、FLM は独特の構造的性質を有しています。すなわち、そのノイズ除去平均 $m_t(x) = \mathbb{E}[X_0 | X_t = x]$ は、各位置におけるクリーンなトークンに関する事後マージナル分布と厳密に対応します。具体的には、出力ベクトルの $(\ell, v)$ 座標は $P(X_{0,\ell} = e_v | X_t = x)$ を表します。

FLM の標準的なサンプリング手順は、連続拡散の慣習（例えば DDPM スタイルのサンプリング）から継承されており、このノイズ除去平均を点推定として扱います。それらは現在のノイズ状態をこの条件付き平均へと橋渡しします。しかし、ワンホット系列の文脈において、この条件付き平均は単体積の内部に位置し、有効なワンホット系列ではありません。事後分布からサンプリングするのではなく、この「不確実性の表現」へと橋渡しすることは、言語モデルに内在する離散構造を崩壊させるモーメント一致近似をもたらします。このアプローチは、FLM の出力において利用可能な完全な事後予測情報を活用できていません。

手法：マージナル条件付きブリッジ（MCB）サンプリング

著者らは、FLM の構造に特化して設計された、トレーニング不要の事後予測サンプリング手法であるマージナル条件付きブリッジ（MCB）サンプリングを提案します。

中核メカニズム

条件付き平均へと橋渡しする代わりに、MCB サンプリングは各離散時間ステップ $k$（ノイズレベル $u_k$ から $u_{k+1}$ へ移動）において、2 段階の逆転移を行います。

1. クリーンな終端点のサンプリング: サンプリング器は、現在の連続状態 $y_{t_k}$ における FLM のトークンごとの事後マージナルを評価します。その後、分解された事後分布から、すべての位置 $\ell$ に対して有効なワンホットトークン $x^I_{0,\ell}$ を独立にサンプリングします。
   $$q^I_{u_k}(x_0 | y) = \prod_{\ell=1}^L q_{u_k, \ell}(x_{0,\ell} | y)$$
   これにより、有効なワンホット系列であるサンプリングされたクリーンな終端点 $x^I_0$ が得られます。

2. 解析的ブリッジによる次の状態のサンプリング: サンプリング器は、現在の状態とサンプリングされた終端点 $x^I_0$ の両方を条件とした解析的 Ornstein-Uhlenbeck（OU）ブリッジから、次の連続状態 $y_{t_{k+1}}$ を引き出します。転移カーネルは以下の通りです。
   $$y_{t_{k+1}} | y_{t_k}, x^I_0 \sim \mathcal{N}\left( \frac{\sinh(\gamma_k)}{\sinh(u_k)}x^I_0 + \frac{\sinh(u_k - \gamma_k)}{\sinh(u_k)}y_{t_k}, \frac{2\sinh(\gamma_k)\sinh(u_k - \gamma_k)}{\sinh(u_k)}I_D \right)$$
   ここで $\gamma_k$ は時間ステップサイズです。

復号制御との統合

終端点が明示的なカテゴリカル分布からサンプリングされるため、MCB 枠組みは再トレーニングなしで標準的な非自己回帰型復号制御を自然にサポートします。

• 温度スケーリング: 分解された終端点事後分布のロジットに直接適用され、鋭さを制御します。
• ヌクレアス剪定: サンプリングされたトークン確率に適用され、低確率の尾部を剪定します。

理論的解析と貢献

本論文は、統計的近似誤差とサンプリング手順誤差を分離し、MCB に対する厳密な理論的正当性を提供します。

1. 事後予測最適性: 著者らは、MCB が FLM によって表現される分解された終端点事後分布によって誘導される、正確な事後予測サンプリングであることを証明します。標準的な DDPM ブリッジと同じ条件付き平均を共有する一方で、MCB は終端点分散を保持し、DDPM はそれを崩壊させるのに対し、MCB はそれを保持します。
2. 分解誤差の定量化: 真の結合クリーン事後分布をトークンマージナルの積に置き換えるコストは、トークン位置間の条件付き多情報量に等しいことが示されます。1 ステップブリッジカーネルの誤差は、この多情報量によって上から抑えられます。
3. パス空間 KL 改善: ガルサノフパス空間比較を用いて、著者らは正確な事後マージナルの下では、MCB サンプリングが真の逆過程に対するパス空間 KL 発散を、標準的な条件付き平均（DDPM）サンプリングのものよりも大きくないことを証明します。
   • 中間座標ごとのブリッジ観測が、固定されたグリッド点推定を超えてクリーンなトークンに関する追加情報を明らかにする限り、この改善は厳密です。
   • これは $L_2$ 射影の議論から導かれます。中間状態 $X_{t,\ell}$（サンプリングされた終端点を介して）を条件付けることは、粗いグリッド点 $X_{t_k}$ のみを条件付けるよりも、$X_{0,\ell}$ に対する厳密に優れた予測子を提供します。

実験結果

著者らは、LM1B および OWT データセットでトレーニングされた事前学習済み FLM チェックポイントにおいて MCB を評価し、サンプリング手順の効果を分離するためにモデルを固定しました。

• 品質と多様性のトレードオフ: LM1B において、MCB は一貫して標準的な ODE サンプリング（条件付き平均を使用するもの）を上回りました。
  • 生成パープレキシティ（PPL）: MCB は、特に離散化ステップ数が増加するにつれて、より低い生成パープレキシティを達成しました。低い温度（$\tau=0.7$）は、MCB 変種の中で最も低い PPL を生み出しました。
  • エントロピー: 標準的な ODE サンプリングは、低ステップ数で退化した低エントロピー生成（反復的なテキストへの崩壊）に苦しんだのに対し、MCB はデータ分布に近い安定したエントロピーを維持しました。
  • トレードオフ: MCB は、温度スケーリングが品質（PPL）と多様性（エントロピー）のバランスを制御し、ベースラインで見られた崩壊を回避する、滑らかなトレードオフ曲線を提供しました。
• 効率性: この手法は、標準的なサンプリングと同じ数のモデル評価を必要とし、トレーニング不要です。

意義と主張

本論文は、この研究の主な意義がフロー言語モデルの出力を正しく解釈することにあると主張しています。著者らは、FLM の出力は単にブリッジのための単体値の条件付き平均として扱うべきではなく、むしろクリーンなワンホットブリッジ終端点に対する直接分布として扱うべきであると論じています。

この離散構造を利用することで、MCB は以下の点を実現します。

1. サンプリング品質の向上: 再トレーニングなしで、標準的なサンプリング手法よりも優れた品質と多様性のトレードオフを達成します。
2. 柔軟な制御の実現: 事後予測枠組みにネイティブな、トークンレベルの復号制御（温度、ヌクレアスサンプリング）のための原理的な非自己回帰型インターフェースを提供します。
3. 理論的保証: 中間状態を条件付けるという情報理論的な利点に基づき、条件付き平均アプローチと比較してパス空間 KL 誤差の削減を保証します。

著者らは、各ステップで離散的な確率性を注入することが、滑らかで決定論的な軌道に依存するフローマップ蒸留のような推論加速技術にとって課題となり得る非滑らかな分岐プロセスを作成するという限界を認識しています。