技術要約：離散拡散言語モデルのための平均場並列デコーディング

1. 問題提起

離散拡散言語モデル（dLLM）は、並列的なトークン生成を可能にすることで、自己回帰（AR）モデルの逐次的な性質とは対照的な、低レイテンシなデコーディングへの道を提供します。しかし、**結合不整合（joint inconsistency）**という決定的なボトルネックが存在します。ほとんどのdLLMは、マスクされた位置のトークン分布を独立して予測する、因数分解された逆ポリシーを採用しています。複数のトークンを、単なる周辺的な確信度（例：最高確率やマージン）のみに基づいて並列に確定させると、結果として生成されるシーケンスはしばしば「結合の不一致（joint mismatch）」に陥ります。個別に見た場合には信頼できそうに見えるトークンであっても、未モデル化の構文的・意味的依存関係（例：主語と動詞の一致や論理的一貫性）により、同時に更新されると互いに矛盾する構成を形成してしまうことがあります。

これを軽減するための既存の解決策には、安全な確定をスケジューリングする依存関係を考慮した推論（DAWNなど）や、ブロック単位の半自己回帰デコーディングがありますが、これらは補助モデル、学習済みの予測器、グラフ構築のヒューリスティック、または追加の学習目的関数を必要とすることが多く、計算オーバーヘッドや複雑さを導入します。そこで、生成品質を損なうことなく、推論中にトークンの相互作用を推定するための、軽量でトレーニングを必要としないメカニタズムが求められています。

2. 手法：平均場並列デコーディング

著者らは、トークンの確定を構造化された推論問題として扱うことで、並列更新を調整する、トレーニングフリーかつシングルパスのフレームメントである**平均場並列デコーディング（Mean-Field Parallel Decoding）**を提案しています。

2.1 エネルギーベースの定式化

確定セット $C_t$（アンマスクされるべきマスクされた位置の集合）の選択は、バイナリインジケータ $S = (s_1, \dots, s_m)$ 上のペアワイズ・マルコフ確率場（MRF）としてモデル化されます（ここで $s_i=1$ は確定を示す）。構成の確率は、ギブス分布を通じて以下のように定義されます：
$$P(S; U_t) \propto \exp\left( \sum_{i} \phi_i(s_i) + \sum_{i<j} \phi_{ij}(s_i, s_j) \right)$$

• 単項ポテンシャル（局所的な確信度）： $\phi_i(s_i) = c_i s_i$。ここで $c_i$ は、位置 $i$ における上位2つの予測トークンのログマージンです。これは、局所的な確信度が高い位置での確定を促進します。
• ペアワイズ・ポテンシャル（予測の重複）： $\phi_{ij}(s_i, s_j) = -D_{ij} s_i s_j$。この項は、強く結合した位置の同時確定を抑制します。相互作用行列 $D$ は、マスクされたトークンの予測分布 $\pi_{t,i}$ と $\pi_{t,j}$ の間の**正規化されたイェンセン・シャノンダイバージェンス（JSD）**を用いて構築されます：
  $$D_{ij} = \frac{1 - \text{JSD}(\pi_{t,i}, \pi_{t,j}) / \ln 2}{\max_{a,b} \tilde{D}_{ab}}$$
  高い $D_{ij}$ は、予測分布が類似していることを示しており、両方が確定された場合に冗長性や衝突のリスクが高いことを意味します。これにより、類似した位置が確定を競い合う「反強磁性的」な結合が生まれます。

2.2 変分平均場推論

最大事後確率（MAP）の割り当てを見つけることはNP困難です。著者らは、結合分布を完全に因数分解された分布 $q(S) = \prod q_i^{s_i}(1-q_i)^{1-s_i}$ で近似する**変分平均場緩和（variational mean-field relaxation）**を採用しています（ここで $q_i$ は確定強度、すなわち確定の周辺確率を表します）。

エビデンス下界（ELBO）を最大化することで、以下の不動点更新式が得られます：
$$q_i = \sigma\left( c_i - \sum_{j \neq i} D_{ij} q_j \right)$$
ここで $\sigma$ はロジスティックシグモイド関数です。この更新により、高確信度の位置が、ペアワイズな相互作用を通じて他の位置を調整できるようになります。アルゴリズムは、安定した確定強度に収束させるために、この更新を $R$ 回（通常は $R=2$）反復します。その後、閾値を適用して最終的な確定セット $C_t$ を決定します。

2.3 実装の詳細

• トレーニングフリー： 補助モデル、再学習、または勾配更新を必要としません。ベースモデルの単一フォワードパスから得られる予測分布のみで動作します。
• 複雑性： 主なオーバーヘッドは、JSD計算のための $O(m^2|V|)$ です（ここで $m$ はマスクされた位置の数、$|V|$ は語彙サイズ）。これはGPU上で高度に並列化可能です。
• 収束： 実験結果によれば、$R=2$ 回の反復で確定強度が安定した不動点に達するのに十分であることが示されています。

3. 主な貢献

1. 定式化： dLLMにおける並列確定選択の問題を、反復的なデノイジングとエネルギーベースモデルを橋渡しする、構造化された推論タスクとして定式化しました。
2. トレーニングフリーの指標： 単一のフォワードパスのみを必要とし、補助コンポーネントも不要な、分布の重複（JSD）に基づく勾配フリーの指標を導入しました。
3. 効率的な推論： 局所的な確信度と相互作用の認識を両立させ、高速に収束する効率的な平均場反復を導出しました。これにより、適応的な並列化が可能になります。
4. 実証的検証： 推論およびコード生成のベンチマークにおいて、本手法が品質とレイテンシのトレードオフを大幅に改善することを実証しました。

4. 実験結果

本手法は、3つのdLLMバックボーン（LLaDA-8B-Instruct, LLaDA-1.5, Dream-v0-Instruct-7B）を用い、4つのベンチマーク（GSM8K, MATH, HumanEval, MBPP）で評価されました。

• スループット vs 正解率： 提案手法は、比較されたすべてのベースライン（Entropy, KLASS, LocalLeap, DAWN）の中で、競争力のあるタスク精度を維持しながら、一貫して**最高のTokens Per Second (TPS)**を達成しました。
  • LLaDA-8B-Instructにおいて、Entropyベースラインに対して4.07倍から6.10倍のスループット向上を実現しました。
  • LLaDA-1.5では、MBPPにおいて最大8.62倍のスループット向上に達しました。
• ベースラインとの比較：
  • Entropyベースライン（保守的、高精度、低スループット）やKLASS（周辺的な確信度のみ）とは異なり、平均場デコーディングは、不釣り合いな精度の低下を招くことなく、トレードオフを高スループット側へとシフトさせます。
  • DAWNやLocalLeapのような依存関係を考慮したベースラインと比較して、提案手法は一貫したスループットの優位性を示しました。これは、平均場の手順が、補助的な検証器や追加のフォワードパスを必要とせずに、より効率的な確定セットを選択できることを示唆しています。
• アブレーション研究：
  • ペアワイズな相互作用を除去すると、精度が大幅に低下し、局所的な確信度だけでは積極的な並列化には不十分であることが確認されました。
  • 一様な相互作用行列を使用すると、JSD由来の行列よりも性能が悪化しました。これは、特定の相互作用構造の重要性を浮き彫りにしています。
  • 反復回数を $R=1$ に減らすとパフォーマンスが低下し、反復的な洗練の必要性が検証されました。
• 感度分析： 本手法は、確定閾値 $\tau$ とブロックサイズによって制御される、滑らかで解釈可能なトレードオフを示します。閾値を下げるとスループットが増加しますが精度が低下し、閾値を上げると精度が維持されますが速度が低下します。

5. 意義と主張

本論文は、平均場並列デコーディングが、離散拡散言語モデルにおける「結合不整合」のボトルネックを緩和するための、原理に基づいたシングルパスのメカニズムを提供すると主張しています。独立した周辺的な更新を、局所的な確信度とグローバルな適合性のバランスを取る構造化されたプロセスに置き換えることで、デコーダは各ステップでより多くのトークンを安全にアンマスクできるようになります。

その意義は、その軽量な性質にあります。補助モデル、再学習、または複雑なスケジューリングのヒューリスティックを必要とせずに、スループットと生成の一貫性を向上させます。これは、モデル自身の予測分布に基づいた協調的な選択が、推論やコード生成のような複雑なテキスト合成タスクにおいて、優れた品質とレイテンシのトレードオフを達成できることを示しています。著者らは、これを既存の拡散デコーディング・パイプラインに対する、実用的でそのまま利用可能なソリューションとして位置付けています。