技術的概要：マスク付き離散系列モデルにおけるペアワイズ相互情報のニューラル推定

1. 問題定義

マスクド拡散モデル（MDM）は、自己回帰（AR）モデルの固定された回帰順序を回避する、離散系列（テキスト、タンパク質、数独など）のための強力な生成モデルである。しかし、標準的な MDM は主に周辺条件付き分布（$p(x_i | x_{\text{context}})$）を露出させるに留まり、変数間の依存関係を明示的に表現していない。

この明示的な依存関係のモデリングの欠如は、2 つの主要な課題を生み出している：

1. 解釈性：変数がどのように互いに関連しているかに関するモデルの内部信念構造を理解することが困難である。
2. 並列デコーディングの効率性：現在の並列デコーディング戦略（Mask-Predict、EB-Sampler など）は、通常、同時にマスクを解除するトークンを決定するために周辺自信度（エントロピー）に依存している。このアプローチはペアワイズ依存関係を考慮していない。相互情報量が高い（強く相関する）トークンを、互いに条件付けずに同時にマスクを解除すると、グローバルな矛盾（数独のルール違反やタンパク質の構造制約の違反など）を引き起こし、しばしば逐次デコーディングへのフォールバックを強いるか、低品質な生成結果をもたらす。

相互情報量（MI）の従来の計算は、密度推定を必要とするため、高次元設定において計算的に非現実的である。

2. 手法

著者は、事前学習済みの MDM の隠れ状態から直接ペアワイズ条件付き相互情報（$I(X_i; X_j | C)$）を推定するためのニューラルフレームワークを提案する。このアプローチは、3 つの主要なコンポーネントから構成される：

A. 真の MI 計算（教師信号）

軽量な推定器を訓練するために、著者はまず、事前学習済み MDM 自身の条件付き分布に基づいて「真の」MI を計算する正確だが高価な方法を定義する。

• 定義：コンテキスト $C$（マスクされていないトークン）に対して、2 つのマスク位置 $i$ と $j$ 間の MI は、結合分布 $P(X_i, X_j | C)$ と周辺分布の積との間の KL 発散として定義される。
• 計算戦略：MDM は周辺分布を出力するため、著者は摂動ベースのブルートフォース・プロービング戦略を使用する：
  1. ベースパス：マスク付き系列に対してモデルを実行し、周辺分布 $P(X_i | C)$ を取得して個別のエントロピー $H(X_i | C)$ を計算する。
  2. 条件付きパス：各位置 $i$ と可能なすべてのトークン $v$ について、$X_i = v$ を固定し、条件付き分布 $P(X_j | X_i=v, C)$ を取得するためにフォワードパスを実行する。
  3. 計算：条件付きエントロピー $H(X_j | X_i, C)$ を計算し、エントロピーの減少として MI を導出する：$I(X_i; X_j | C) = H(X_j | C) - H(X_j | X_i, C)$。
• コスト：これには $1 + N \cdot |V|$ 回のフォワードパスが必要であり、推論には非現実的であるが、訓練データの生成には適している。

B. ニューラル MI 推定器

軽量なニューラルネットワーク（$f_\phi$）が、凍結された MDM の隠れ状態（$h$）から直接 MI 行列を近似するように訓練される。

• アーキテクチャ：推定器は隠れ状態 $h \in \mathbb{R}^{N \times D}$ を入力とし、すべての位置に対する推定されたペアワイズ MI を表す対称行列 $\hat{I} \in \mathbb{R}^{N \times N}$ を出力する。
• 訓練目的：モデルは、マスクされたインデックスにおいて、予測された行列 $\hat{I}$ と真の行列 $M_{GT}$ 間の平均二乗誤差（MSE）を最小化するように訓練される。

C. MI 誘導型並列サンプリング

著者は、予測された MI 行列を利用して、マスク解除されたトークン間の条件付き独立性を確保する、並列デコーディングのための貪欲な選択アルゴリズムを導入する。

• 戦略：単にエントロピーが最も低い（自信度が最も高い）トークンを選択するのではなく、このアルゴリズムはコンテキストを与えられたときに相互に独立しているトークンのバッチ $S$ を選択する。
• アルゴリズム：
  1. マスクされたインデックスをエントロピーの昇順（自信度が最も高い順）にソートする。
  2. 候補を反復処理し、依存コストを計算する：$d(i | U) = \sum_{j \in U} \hat{I}_{i,j}$。ここで $U$ はすでに選択されたトークンの集合である。
  3. トークン $i$ の総コスト（エントロピー + $\lambda \times$ 依存コスト）が残余予算 $\gamma$ 以内の場合にのみ、トークン $i$ を選択する。
  4. コストが高すぎる場合（すでに選択されたトークンとの高い MI を示す）、そのトークンは逐次ステップに延期される。
• 結果：これにより、強く相関する変数は逐次処理され、条件付き独立な部分集合は並列処理される。

3. 主要な貢献

1. ニューラル MI 推定フレームワーク：推論中に高価な密度推定や真値の計算を必要とせず、MDM の隠れ状態から直接ペアワイズ条件付き MI を推定する方法。
2. MI 誘導型並列デコーディング：推定された MI を使用して条件付き独立な変数の部分集合を特定する新しいサンプリング戦略。これにより、グローバルな整合性を保ちながら安全な並列化が可能になる。
3. 解釈性ツール：MI マップは、明示的なプログラミングなしに（数独のルールやタンパク質のフォールディング依存関係など）学習された制約を明らかにするモデルの内部信念構造の可視化として機能する。

4. 実験結果

このアプローチは、2 つのドメインで評価された：数独（構造化された論理）とタンパク質系列生成（ESM-C を使用）。

数独

• 設定：10 万パズルで訓練され、1,000 の未見の難易度の高いパズルで評価された。
• 性能：
  • 逐次ベースライン：平均 53.9 回のフォワードパス、61.6% の精度。
  • 単純な並列（k=7）：9.0 回のパスだが、精度は 36.8% に低下。
  • MI 誘導型（γ=0.3）：15.2 回のパスで63.6% の精度（逐次ベースラインを上回る）。
  • MI 誘導型（γ=0.6）：9.7 回のパスで 56.2% の精度。
• 観察：MI 誘導型サンプリングは、逐次デコーディングと比較してフォワードパスを 3〜5 倍削減し、単純な並列手法と比較して精度を維持または向上させた。

タンパク質系列（ESM-C）

• 設定：500 個のランダムなタンパク質（長さ 50〜100）を生成し、Jensen-Shannon 発散（JSD）を使用して UniRef50 から 500 個の参照サンプルと比較した。
• 性能：
  • 逐次：74.8 回のパス、JSD 0.093。
  • 単純な並列（k=12）：6.2 回のパス、JSD 0.218（著しい品質の低下）。
  • MI 誘導型（γ=4）：10.0 回のパス、JSD 0.174。
• 観察：MI 誘導型サンプリングは、単純な並列ベースラインよりも優れた速度 - 精度のトレードオフを達成し、生成品質をエントロピーベースの手法よりもよく維持しながら、パス回数を大幅に削減した（逐次と比較してほぼ 1 桁減少）。

5. 意義と主張

本論文は、変数依存関係を明示的にモデル化することが、離散拡散モデルの潜在能力を最大限に引き出すために不可欠であると主張している。

• ギャップの橋渡し：この研究は、高品質な逐次サンプリングと効率的な並列デコーディングの間のギャップを埋める。
• 内部表現：MI マップは、MDM が明示的なプログラミングなしに（数独のルールやタンパク質の依存関係など）厳格な構造的制約を自然に獲得しており、これらが推定器を通じて抽出可能であることを示している。
• 効率性：この手法は、条件付き独立な部分集合を特定する MI 誘導型並列デコーディングを可能にし、逐次デコーディングと比較して推論時のフォワードパスを 3〜5 倍削減する。

認められた限界：
著者は、予測器が完璧ではなく、訓練データに対して真の MI をその場で計算するなど、相当なセットアップと訓練を必要とすることを指摘している。将来の研究では、訓練段階における真値計算の計算コストを回避するために、最適な予測器アーキテクチャと改善されたカリキュラム学習戦略を調査することが提案されている。