技術要約：FUND – 非正規化分布のサンプリングのための密度フロー

問題提起

非正規化確率分布、特に $p(x) \propto e^{-E(x)}$ の形式を持つボルツマン分布からのサンプリングは、統計物理学、量子力学、計算化学における複雑な物理系のモデリングにおいて極めて重要である。従来のマルコフ連鎖モンテカルロ（MCMC）や分子動力学（MD）は、高次元のエネルギーランドスケープ、特に高い障壁によって隔てられたメタステーブルな状態をナビゲートする際に、深刻なスローイングダウン、高い自己相関、および不十分なモード混合に直面する。

ノーマライジングフロー（NF）のような深層生成モデルは効率的なサンプリングを提供するが、この領域においては以下のような重大な限界がある：

1. データ依存性： 標準的なNFの学習（例：最大尤度推定/前方KL）は、ターゲット分布からの真のサンプルに関する大規模なデータセットを必要とするが、これらはMCMCやMDを用いて生成するには非常に高価であるか、入手困難な場合が多い。
2. モード崩壊： 真のサンプルを必要としない逆KLダイバージェンス（$KL(q_\theta \| p)$）による学習は、モデルを低エネルギー領域へと偏らせ、ターゲット分布の全サポートをカバーできなくなるモード崩壊を引き起こす。
3. 既存の代替手法： 最近のFAB（Flow Annealed Importance Sampling Bootstrap）やiDEM（Iterated Denoising Energy Matching）といった手法は、真のサンプルの要求を回避しようと試みているが、多大な計算オーバーヘッド（例：繰り返しのAISや逆SDEの解決）を導入するか、あるいは追加の近似段階なしには正確な尤度を提供できない。

手法：FUND (Density Flow for Unnormalised Distributions)

著者らは、ターゲット分布からの真のサンプルを一切必要とせずに、非正規化分布に対する扱いやすいサンプラー $q(x)$ を学習するフレームワークである FUND を提案する。その核心的な革新は、学習のパラダイムを「サンプル・トラジェトリ・マッチング（フロー・マッチングのように）」から「分布・トラジェトリ・マッチング」へと転換した点にある。

コアフレームワーク

1. アニールド分布・トラジェトリ： 個々のサンプルの経路をモデル化するのではなく、FUNDは単純な単峰性事前分布 $p_0(x)$ から複雑なターゲット $p_1(x) = p(x)$ へと変換する連続的な確率密度 $\{p_t(x)\}_{t \in (0,1]}$ の軌跡をモデル化する。中間分布は次のように定義される：
   $$p_t(x) \propto p(x)^t$$
   $t$ が0から1へと増加するにつれ、分布は徐々に鋭くなり、モードが分離していく。$t \to 0$ では分布は滑らかで単峰性であり、$t=1$ ではターゲットの全複雑性が回復する。

2. 重要度重み付けを用いた反復学習： モデルは $p_t$ に対応する一連の中間分布 $q_t$ を学習する。

   • 初期化： 学習は $p_{t_1}$ が単峰性である小さな $t_1$ から開始される。$p_{t_1}$ からのサンプルは存在しないため、モデルはエネルギー評価のみに依存する**逆KL（RKL）**目的関数を用いて学習される。
   • 進行： 後続のステップ $t > t_1$ では、以前に学習された $q_{t'}$（ここで $t' < t$）からサンプルを生成する。これらのサンプルは、$q_t$ の学習を監督するために、重要度重み $w(x) = p_t(x) / q_{t'}(x)$ を用いて再重み付けされる。
   • 損失関数： 学習目的関数は、監督付き項（前方KLまたはスコアマッチング）と非監督付きのRKL項を組み合わせたものである：
     $$L_{final}(\theta) = \lambda_1 L(\theta) + \lambda_2 L_{RKL}(\theta)$$
     • $L(\theta)$: 重み付けされたサンプルを用いて、モードの被覆（探索）を促進する。
     • $L_{RKL}(\theta)$: 主要なモードへの集中（活用）を促し、密度の不一致を修正する。
   • ハイパーパラメータのアニーリング： 重み $\lambda_1$ と $\lambda_2$ はアニーリングされる。初期段階では $\lambda_1 \approx 0$（RKLに依存）であり、学習が進むにつれて、モード被覆のために $\lambda_1$ が増加し、モード分離を精緻化するために $\lambda_2$ が後に増加する。

3. アーキテクチャと実装：

   • シングルブロック・アーキテクチャ： 各ステップで結合ブロックを追加する増分的なアプローチとは異なり、FUNDはすべてのタイムステップにわたって精緻化される単一のノーマライジングフローモデルを使用する。これにより、メモリ効率が向上し、破滅的忘却が防止される。
   • リプレイバッファ： 永続的なバッファが、後続のステップのための重要度重みを計算するために、以前の中間分布からのサンプルと対数尤度を保持する。これにより、学習の安定性とモードの保持が保証される。
   • 離散化： 連続的な時間 $t$ は $N$ ステップに離散化される。ステップ数と初期 $t_1$ は、ターゲットの複雑さと多峰性に基づいて調整される。

主な貢献

1. サンプルフリー学習： FUNDは、ターゲットからの真のサンプルへのアクセスなしに、非正規化分布に対するノーマライジングフローの学習を可能にする。これにより、高価なMCMC/MDによる事前サンプリングの必要性を排除する。
2. 分布・トラジェトリ・マッチング： 本論文は、個々のサンプル・トラジェトリではなく、密度自体の進化をモデル化する新しい定式化を導入し、推論時の直接的かつ効率的な密度評価を可能にする。
3. バランスの取れた目的関数： 前方KL/スコアマッチング（被覆のため）と逆KL（集中のため）を逐次的なアニーリング・スケジュール内で組み合わせることで、FUNDはモード崩壊とモード被覆の間のトレードオフを緩和し、複雑で高次元なランドスケープの包括的な被覆を保証する。
4. スケーラビリティ： FUNDは、従来のベースラインが著しく低下する高次元の問題（例：144次元のスカラー $\phi^4$ 理論）において、スケーラビリティを示す。

実験結果

著者らは、以下の3つのベンチマークファミリーでFUNDを評価している：

1. MOG-40： 広く離れたモードを持つ2次元の40個のガウス混合分布。
2. Many-Well (MW-8, MW-16)： モードの数が指数関数的に増加する（$2^{d/2}$）合成分布。
3. スカラー $\phi^4$ 理論： $8\times8, 10\times10, 12\times12$ グリッド（最大144次元）上の格子場理論モデル。

性能指標：

• 負の対数尤度 (NLL)： ターゲットへの適合度を測定する。FUNDのバリアントは、FAB、iDEM、および標準的なRKL/FKLベースラインを一貫して上回り、最も低いNLLを達成している。
• 逆NLL (RNLL) および 有効サンプルサイズ (ESS)： これらの指標はモード被覆を評価する。FUNDは高いESSと低いRNLLを達成しており、これは複数のモードを効果的に捉えていることを示している。対照的に、RKLベースラインはNLLは低いがESSが低い（モード崩壊）、FABとiDEMは高次元で性能が低下する。
• 計算コスト：
  • FUND(FKL)： 良好なコスト・品質のトレードオフを提供する。学習時間はFABと同等であるが、推論は大幅に高速である（AISやMCMCと比較して、単一のフォワードパス）。
  • FUND(SCORE)： 最高レベルのサンプリング精度と密度推定を実現するが、スコアマッチングの計算により学習コストは高くなる。
  • ベースライン： MCMC法は高い推論コストを持つ。iDEMは逆SDEシミュレーションにより高い推論コストを持つ。FABは高い学習コストを持つ。

意義と主張

本論文は、FUNDが複雑な物理分布のサンプルフリー学習のための新しいパラダイムを確立すると主張している。中間的なアニールド状態を通じてターゲット分布を段階的に形成することにより、FUNDはエネルギーベースの学習に固有の「モード崩壊」と、尤度ベースの学習に固有の「データ依存性」の両方を克服する。

著者らは、FUNDが高次元のボルツマン分布（特に格子場理論や統計物理学において、スローイングダウンやトポロジカルなトラッピングが持続的な課題となる分野）に対する、堅牢でスケーラブルな代替手段を提供することを強調している。このフレームワークは、伝統的なサンプリング手法が苦戦するベイズ推論やゲージ理論のシミュレーションにおいて、高い忠実度（ターゲットの多峰構造への適合）を維持しながら、償還サンプリング（amortized sampling）（単一のフォワードパスによる高速な並列生成）を可能にする。