Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels

テクニカル・サマリー：サンプラーを止めろ！分類器ベースの適応的停止を用いたサンプリング・カーネル

問題提起
$\pi(x) = r(x)/Z$ と定義される、非正規化された複雑な確率密度からのサンプリングは、ベイズ推論および確率的モデリングにおける根本的な課題である。マルコフ連鎖モンテカルロ（MCMC）法は漸近的な保証を提供する一方で、特に高次元空間において、混合の遅さや高い計算コストに頻繁に悩まされる。標準的なMCMCの決定的な制限は、固定または手動で調整された軌跡長への依存である。既存の停止時間の決定手法は、生成されたサンプルから計算されるオンライン診断（例：Gelman–Rubin、Geweke）に依存することが多く、これらは必ずしも肯定的な停止基準を提供せず、早期終了や推定のバイアスを招く可能性がある。対照的に、アモルタイズ（償却）されたアプローチ（例：正規化流、拡散サンプラー、GFlowNet）は、学習に計算負荷を転嫁するが、MCMCの厳密なサンプリング保証を回避してしまうことが多い。

手法
著者らは、軌跡の終了を学習可能なコンポーネントとして扱うフレームワークを提案し、MCMCとアモルタイズされたサンプリングの架け橋となる手法を提示している。コアとなる手法は以下の通りである：

1. 非非巡回型GFlowNet定式化： 著者らは、適応的停止を伴うMCMCを、非非巡回型Generative Flow Networks（GFlowNet）の理論の枠組みで捉えている。彼らは、構造化された状態空間 $\bar{S} = \{s_0\} \cup S \cup \{s_f\}$ を定義する。ここで、$S = \mathbb{R}^d$ は連続状態空間を、$s_0$ はソースを、$s_f$ は吸収されるシンク状態を表す。標準的な非巡回型GFlowNetとは異なり、この設定では $S$ 内でのサイクルを許容している。
2. 分類器ベースの停止： 固定されたステップ数のマルコフ連鎖を実行する代わりに、本手法は状態依存のニューラル分類器 $d_F(s)$ および $d_B(s)$ を用いて、軌跡がいつ終了すべきかを決定する。
   • 前方分類器 $d_F(s)$ は、直接シンク状態 $s_f$ へ遷移する（停止する）確率を決定する。
   • 後方分類器 $d_B(s)$ は、ソース状態 $s_0$ へ移動する確率を制御する。
   • 前方遷移密度は、$p_F(s'|s) = (1 - d_F(s))q_F(s'|s)$ （連続遷移の場合）および $p_F(s_f|s) = d_F(s)$ （終了の場合）として定義される。
3. 理論的接続： 本論文は、学習された停止ポリシーとターゲット密度 $\pi$ との間の厳密な関係を、詳細釣り合い条件を通じて確立している。理論的には、最適な前方分類器は $d_F(s) = r(s)/f(s)$ を満たすことが示されている（ここで $f(s)$ はフロー測度である）。これは、停止の決定が、注入された質量とターゲット分布との間の累積的な不一致に関連していることを示している。
4. 学習アルゴリズム： 著者らは、非非巡回環境において計算効率の高い、プレフィックス・トラジェクタリー・バランス（TB）目的関数を利用している。学習プロセスには以下が含まれる：
   • 探索を促進するために、分類器を無効にした状態で最大長 $N_{max}$ の部分的な軌跡をサンプリングする。
   • プレフィックスを終端遷移で拡張し、完全な軌跡を形成する。
   • プレフィックス・トラジェクタリー・バランス損失とフロー正則化項を組み合わせた損失関数を最小化する。この正則化は、総フローを罰することでより短い軌跡を推奨し、期待軌跡長を実質的に最小化する。
5. マルチレベル汎化： 複雑な幾何学的構造を扱うため、フレームワークはマルチレベル・スキームへと拡張されている。状態空間はレベル変数 $\ell \in \{1, \dots, L\}$ によって拡張され、モデルがレジームを切り替えることを可能にする。中間報酬 $r(s, \ell)$ は、初期分布とターゲットとの間を補間するように定義され、アニーリングされた重要度サンプリングと同様に、学習と探索を安定させる。

主要な貢献

1. 理論的フレームワーク： 著者らは、任意の時点で停止の可能性がある $\mathbb{R}^d$ におけるサンプリング力学が、非非巡回型GFlowNet環境と見なせることを証明した。彼らは、フローおよび詳細釣り合い条件を用いて、最適な分類器とターゲット密度の関係を導出した。
2. 実践的アルゴリズム： 学習された停止と、適応的な前方および後方カーネルを組み合わせた実践的なアルゴリズムを導入した。これには、複雑な幾何学的構造における訓練の安定性と探索を改善するために設計されたマルチレベル拡張が含まれる。
3. 実証的評価： 本手法は、様々なベンチマーク密度（GMM9、Funnel、ManyWell、および高次元のDigitsタスク）で評価されている。結果は、本アプローチが標準的なMCMCベースライン（ULAなど）と比較して、平均軌跡長を大幅に短縮しつつ、モード被覆率と混合を向上させていることを示している。

実験結果

• 軌跡長： 分類器ベースのアプローチは、固定ステップのULAおよび標準的な収束診断（Gelman–Rubus、Geweke）を用いたULAと比較して、平均軌跡長（$E[n_\tau]$）を大幅に減少させた。例えば、GMM9 ($d=2$) タスクにおいて、学習された分類器は、ULAの約104ステップから約15ステップへと軌跡長を削減した。
• サンプル品質： 分類器単体でも軌跡長を短縮するが、フルモデル（学習された前方カーネルの補正を含む）は近似品質を大幅に向上させる。これは、より低いSinkhorn距離（SD）、より低い最大平均偏差（MMD）、およびよりタイトなEvidence Lower/Upper Bounds（ELBO/EUBO）によって裏付けられている。
• モード被覆： 可視化と定量的指標は、提案手法が遠く離れた高密度領域間を効果的に遷移できる一方で、標準的なULAはしばしば局所領域に閉じ込められることを示している。
• マルチレベルの性能： マルチレベル拡張（$L=5$）は、一貫してシングルレベルモデルよりも優れた性能を示し、高次元タスクであるDigits ($d=196$) において、確立された拡散サンプリング手法（PIS、DDS）に匹敵する性能を発揮した。

意義と主張
本論文は、サンプリングの軌跡が、高密度領域では動的に短縮され、その他の場所では探索的にあり続けるような、統一されたフレームワークを提供すると主張している。停止を学習可能なコンポーネントとして扱うことで、フローベースのサンプリングに関連する理論的保証を損なうことなく、固定長MCMC軌跡の非効率性の問題に対処している。著者らは、このアプローチが追加の最適化の複雑さ（探索と正則化のバランス）を導入するものの、固定カーネル（ULAなど）をデバイアスし、複雑なターゲット分布の幾何学に適応するための原理的な方法を提供することを指摘している。本研究は、連続的な非非巡回型GFlowNetのためのより優れた訓練目的関数の開発や、ハミルトニアンモンテカルロのようなより表現力豊かな力学への拡張といった、将来の方向性を示唆している。