技術的サマリー：離散拡散サンプリングとブリッジ

問題定義

本論文は、エネルギー関数 $E$ は既知だが正規化定数 $Z$ が扱いにくい、離散状態空間 $\mathcal{S} = \{1, \dots, C\}^d$ 上の分布 $p(x) \propto e^{-E(x)}$ からのサンプリング課題に取り組む。拡散サンプリングは、確率的最適制御と経路空間マッチングを用いて連続空間へ成功裏に適応されてきたが、離散空間への適用は未発展のままである。既存の離散拡散サンプリング手法は、連続空間で標準的なオフポリシー強化学習（RL）手法を活用できておらず、特に高度に多峰性の分布において、モードカバレッジとサンプリング品質が最適でない傾向にある。さらに、2 つの任意の分布間を繋ぐ問題（シュレディンガー・ブリッジ）は、両方の分布からのサンプルが利用可能な設定（データ対データ）にほぼ限定されており、一方の分布が正規化されていない密度のみで定義される「データ対エネルギー」設定は未探索であった。

手法

著者は、オフポリシー学習戦略を離散拡散サンプリングに統合し、これらを離散シュレディンガー・ブリッジへ一般化するフレームワークを提案する。

1. 離散拡散サンプリングのためのオフポリシー学習

核心的な手法は、前方プロセスの軌道分布と後方（ノイズ付加）プロセスの軌道分布との間の発散を最小化することで、目標分布に一致する前方カーネル $\overrightarrow{p}_\theta$ を学習することである。検討された目的関数は、2 次モーメント発散であり、具体的には以下の通りである：

• 軌道バランス（TB）： 正規化定数を学習可能なパラメータとして扱う。
• 対数分散（LV）： 前方確率と後方確率の対数比の分散を最小化する。

探索とモードカバレッジを改善するため、著者は RL と生成フローネットワーク（GFlowNets）から適応された 3 つのオフポリシー戦略を導入する：

• リプレイバッファ： サンプリングされた状態 $X_N$ を保存し、学習のために後方軌道 $\overleftarrow{p}^{\otimes N}$ をサンプリングして再利用する。
• 重要度重み付きバッファ： 軌道の重要度重みに基づいてバッファ要素にサンプリング確率を割り当て、現在のサンプリング手法が見逃す可能性のある高確率領域を優先する。
• MCMC 探索： 学習に使用される前に、マルコフ連鎖モンテカルロ（MCMC）ステップ（例：Swendsen-Wang または Metropolis-Hastings）を用いてリプレイバッファからのサンプルを洗練させる。これにより、モデルのロールアウトを必要とせず、エネルギー関数を直接利用してサンプリングを高密度領域へ誘導する。

2. 離散データ対エネルギー・シュレディンガー・ブリッジ

本論文は、2 つの任意の分布 $p_0$ と $p_N$ 間のシュレディンガー・ブリッジ（SB）を見つけるサンプリング問題を一般化する。両端点のサンプルを仮定する従来の離散 SB 研究とは異なり、この研究は $p_N$ がエネルギー関数のみで定義されるデータ対エネルギー設定に対処する。

• 本手法は、逐次比例適合（IPF）手順を採用する。
• $p_0$ からのサンプルを用いた最大尤度による後方カーネルの適合と、エネルギー定義の $p_N$ からのサンプルを用いた対数分散損失による前方カーネルの適合を交互に行う。
• 探索を促進するため、オフポリシー手法（バッファと MCMC）を IPF ループに統合し、エネルギー定義分布の探索を支援する。

3. 潜在空間への適用

著者は、これらのサンプリング手法を「外部委託」された事後サンプリングに応用する。事前学習されたベクトル量子化変分オートエンコーダ（VQ-VAE）の離散潜在空間において拡散サンプリングを学習する。目的は、不偏な条件付きサンプルを観測することなく、事後分布 $p(z|y) \propto p_{latent}(z)p(y|f(z))$ からサンプリングすることであり、これにより潜在空間でのサンプリングを学習することで生成モデルを条件付ける。

主要な貢献

1. 離散サンプリングのためのオフポリシーアルゴリズム： 離散拡散サンプリングの学習へのオフポリシー RL 手法（リプレイバッファ、重要度重み付け、MCMC 探索）の導入。本論文は、これらの手法がオンポリシーベースラインに対して性能を大幅に向上させることを実証する。
2. データ対エネルギー・シュレディンガー・ブリッジ： 離散拡散サンプリングをデータ対エネルギー・シュレディンガー・ブリッジ問題へ初めて一般化し、サンプルベースの分布とエネルギー定義の分布間の輸送マップの構築を可能にする。
3. 潜在空間事後サンプリング： 画像生成モデルの離散潜在空間における事後推論のための離散拡散サンプリングの実証。生成モデルの再学習なしにクラスラベル（例：奇数/偶数の数字）を条件として扱える能力を示す。

実験結果

著者は、合成ベンチマークと実世界応用において手法を評価する：

• イジングモデルおよびポッツモデル： $16 \times 16$ 格子において、オフポリシー手法（特にバッファ＋MCMC を用いた TB）は、オンポリシー手法およびベースラインのマスクド拡散ニューラルサンプリング（MDNS）を一貫して上回る。特にモード崩壊が prevalent な低温（高 $\beta$）において、シンクホルン距離および磁化/相関誤差が低くなる。
• 離散化された合成密度： グレーコードを介して離散化された困難な多峰性分布（40GMM および ManyWell）において、オフポリシー手法はすべてのモードを正常に発見するのに対し、オンポリシー手法および MDNS は分布全体をカバートすることに失敗することが多い。
• シュレディンガー・ブリッジ： 単純な分布と多峰性分布間のブリッジタスクにおいて、オフポリシー手法はオンポリシーサンプリングが崩壊する複雑なターゲットへも効果的にスケーリングする。
• 外部委託サンプリング： 潜在空間から生成された MNIST 数字の可視化は、オンポリシーおよびオフポリシーの両方のサンプリング手法が、所望のクラス（例：奇数対偶数）の数字を正常に生成できることを示し、条件付き生成に対するアプローチの有効性を検証する。

意義と主張

本論文は、オフポリシー学習手法の統合が、モード崩壊の問題に対処し、高次元で多峰性の離散空間におけるサンプリング効率を向上させるために、離散拡散サンプリングにとって重要な前進であると主張する。離散拡散とオフポリシー RL の間のギャップを埋めることで、著者はこれらのサンプリング手法が、従来の手法よりも効果的に複雑な統計物理学問題（イジング/ポッツ）および組合せ最適化タスクを処理できることを示す。

さらに、この研究はシュレディンガー・ブリッジの理論的範囲を、離散ドメインにおけるデータ対エネルギー設定へ拡張し、一方がエネルギー関数で定義される分布間の輸送のための新たなツールを提供する。最後に、潜在空間への適用は、完全なデータ分布へのアクセスや生成モデル自体の再学習を必要とせずに、事前学習された生成モデルを条件付けるための実用的な道筋を示す。

著者は、ノイズ付加カーネルを固定する必要があること、および軌道保存の計算コストといった限界を認めているが、連続時間理論の拡張によってサポートされる粗い時間スケールでの学習が、これらのコストの一部を軽減する可能性を指摘している。また、データからのエネルギー関数の学習や、これらの手法をより大規模な画像生成モデルへスケーリングすることなどを将来の方向性として特定している。