技術要約: FMLM+ による事後的な洗練（Posterior Refinement）

問題提起

現代の言語モデルは主に自己回帰型（AR）であり、長さ $L$ のシーケンスを生成するために $L$ 回の逐次的なフォワードパスを必要とする。この固有の逐次的ボトルネックは推論を遅くし、メモリ制限を生じさせ、推論や強化学習のようなダウンストリームアプリケーションにおけるスケーラビリティを阻害している。

非自己回帰（NAR）の代替案が、主に以下の2つのパラダイムを通じてこの問題に対処するために登場した：

1. マスク拡散モデル (MDMs): 任意のマスキングパターンで学習することで、柔軟な任意の順序での推論を提供する。これらは適応的な「アンマスク・アンド・リマスク（unmask-and-remask）」戦略を可能にする。しかし、**因子化誤差（factorization error）**に苦しみ、複数のトークンを同時に生成しようとするとサンプル品質が崩壊する。これにより、高品質な出力のための数ステップでの生成が不可能となっている。
2. フローマップ言語モデル (FMLMs): ユークリッド空間におけるガウスノイズからワンホットベクトルへの結合輸送を学習し、MDMの因子化誤差を解決する。これらは優れた数ステップ生成（単一の関数評価ですら）を可能にする。しかし、すべてのトークンを同時に進化させるため、MDMのような推論時の柔軟性に欠け、生成中の高信頼度トークンの適応的な保持ができない。

既存のNARパラダイムは、反復的な洗練（iterative refinement）——すなわち、出力を改善するために任意のトークン集合をグローバルに批判、消去、および再生成する能力——の可能性を十分に実現できていない。MDMはマルチトークンのドラフト作成には効率性に欠け、FMLMは適応的なトークン保持の柔軟性に欠けている。

手法: FMLM+ と事後的な洗練（Posterior Refinement）

著者らは、MDMとFMLMの間の溝を埋める統一フレームワークである FMLM+ と、**事後的な洗練（Posterior Refinement, PR）**と呼ばれる新しい推論戦略を導入する。

1. FMLM+ のアーキテクチャと学習

FMLM+ は、マスキング形式のノイズスケジュールを備えることで、フローマップ言語モデルを一般化する。

• ノイジング・プロセス: シーケンス全体を一様にノイズ化する代わりに、FMLM+ はコンテキスト位置 $C$ におけるクリーンなトークン集合 $y_C$ を保持しつつ、残りをノイズ化する。これは、Diffusion Transformer (DiT) アーキテクチャにおける時間条件付けを、位置ごとに（クリーンな位置に対して $t=1$ と設定）変更することで実装される。
• 学習目的: モデルは、対角線上（$s=t$）でクリーンなデータを再構成し、より小さなジャンプを教師として用いて、より大きなオフダイアゴナル・ジャップ（$s \to t$）を漸進的に蒸留する、漸進的自己蒸留目的関数を用いて学習される。
• 学習済みMDMの活用: 著者らは、$s=t=0$ において FMLM+ の目的関数が MDM の目的関数に帰着するという理論的な接続を示した。これにより、2つの加速戦略が可能になる：
  • 蒸留: 学習済みのMDMを教師として用いてトークン確率を一致させ、勾配分散を低減する。
  • 初期化: 学習済みのMDMのチェックポイントから直接 FMLM+ の重みを初期化する。著者らは、最初の適応型レイヤー正規化（adaLN）層が、ガウスノイズ入力をMDMが [MASK] トークンに対して生成するのと同じ表現にマッピングすることを確認しており、これにより学習済みの知識が効果的に保持される。

2. 事後的な洗練 (Posterior Refinement, PR)

PR は、FMLM+ が提供する 事後的な（a posteriori） 信頼度スコアを利用する推論時の洗練戦略である。

• メカニズム: 独立したトークンごとの周辺分布に基づいて 事前的な（a priori） 信頼度を推定する（そのためマルチモーダルな分布で失敗する）MDMとは異なり、FMLM+ はシーケンス全体を結合して生成する。これにより、モデルは生成されたドラフト全体を条件とした各トークンの条件付き確率を評価できる。
• ループ:
  1. 単一のフローマップ・ジャンプ（例：$\delta_{0,1}$）を介して、候補シーケンス $\hat{x}_1$ を生成する。
  2. 各トークンに対する 事後的な信頼度 を計算する。これは、予測された分布の最大確率 $p_{max}$ によって定義される。高い $p_{max}$ は、ワンホットベクトルへの投影における丸め誤差が小さいことを意味する。
  3. $p_{max}$ が閾値を超えるトークン（またはトップ-$k$ トークン）を、コンテキスト集合 $C$ に コミット する。
  4. 残りの未コミットのトークンを 再ノイズ化 し、プロセスを繰り返す。
• 理論的洞察: 著者らは、MDMが 偶然的（aleatoric） 不確実性（データのランダム性）と 認識論的（epistemic） 不確実性（モデルの確信度）を混同していると主張する。FMLM+ の軌跡は、結合輸送を学習することで、真の認識論的不確実性を表面化させ、MDMが見逃すエラーをPRが確実に特定・修正することを可能にする。

主な貢献

1. 統一フレームワーク (FMLM+): FMLMの数ステップの効率性とMDMの任意の順序での柔軟性を組み合わせた、マスキング形式のノイズスケジュールを持つFMLMの一般化。これは、大規模な学習済みマスキング拡散モデルの進展を利用してFMLM+をスケールさせるための直接的な経路を確立する。
2. 事後的な洗練 (Posterior Refinement): 出力を再帰的に自己修正することを可能にする、事後的な信頼度に基づく新しい戦略。これは、誤ったトークンが一貫して低信頼度の生成物の中に含まれることを特定し、信頼性の高いエラーフィルタリングと改訂を可能にする。
3. 広範な評価: TinyStories、OpenWebText、GSM8K（推論）、Sudoku（構造的制約）の4つのベンチマークにおいて、FMLM+ と PR が、MDM および FMLM ファミリーよりも効果的に速度と品質のトレードオフをナビゲートできることを実証した。

結果

本論文は、目標性能を達成するために必要な関数評価回数（NFE）によって測定される、速度と品質のトレードオフの大幅な改善を報告している：

• Sudoku (構造的制約): 無条件のSudoku生成において、MDMおよび 事前的な 信頼度を用いたFMLM+は、複数のトークンをサンプリングする際に致命的に失敗する。PRを用いたFMLM+は、中程度の難易度のパズルにおいてわずか 9 NFE で 87.1% の精度 を達成し、難しいパズルでは 16 NFE で 71.2% を達成した。これは、同等の精度を得るために1024 NFEを必要とするベースラインと比較して 32倍の高速化 である。
• GSM8K (推論): PRを用いたFMLM+は、GSM8Kテストセットにおいて 32 NFE のみで 19.0% の精度 を達成し、すべての非自己回帰型ベースラインを上回った。
• 無条件生成 (TinyStories/OpenWebText): FMLM+は、最大 8倍少ない NFE で、最強のベースライン（ARモデルを含む）と同等またはそれ以上の性能を発揮する。例えば、TinyStoriesにおいて、FMLM+は32 NFEで17.53の生成パープレキシティを達成したが、これは128 NFEを用いたMDLMの18.74と比較して優れている。
• 学習効率: 学習済みのMDMを蒸留または初期化によって活用することで、ゼロからの学習と比較して、学習の収束と最終的なパフォーマンスが大幅に向上した。

意義と主張

本論文は、FMLM+ と事後的な洗練（Posterior Refinement） が、非自己回帰型言語モデリングにおける新たな最先端（SOTA）を確立すると主張している。その意義は以下の点にある：

• 反復的な洗練の実現: FMLMによる効率的なマルチトークン・ドラフト作成と、PRによる適応的なトークン保持を組み合わせることで、NARモデルの理論的利点である「反復的な洗練」を成功裏に実装した。
• スケーラビリティ: MDMとFMLMの間の溝を埋めることにより、大規模な学習済みマスキング拡散モデル（LLaDA-8B、Dreamなど）を活用できる、高忠実度な言語モデリングのためのスケーラブルな基盤を提供する。
• 推論時のスケーリング: 再学習なしで、追加の洗練ラウンドによってサンプル品質が向上するという、フローベースモデル特有の推論時スケーリングを明らかにした。

著者らは、実験が小規模かつ主に合成データセットで行われたこと、およびゼロからの学習がMDMの重みによる初期化に比べて最適ではないことを認め、限界に対して謙虚な姿勢を保っている。彼らは、今後の研究として、数十億パラメータ規模へのスケーリングや、事後的な信頼度信号を用いた新しいテスト時戦略の設計を提案している。