SCORENF: Score-based Normalizing Flows for Sampling Unnormalized distributions

技術要約：非正規化分布のサンプリングのための ScoreNF

問題提起
非正規化確率分布（例：$e^{-H(x)}$ の形式のボルツマン分布）からのサンプリングは、統計物理学、生物科学、およびベイズ推論における根本的な課題である。これらの分布は、特に高次元の設定において、解析的に扱いやすい正規化定数（分配関数）を欠いていることが多い。従来のマルコフ連鎖モンテカルロ法（MCMC）は、収束の遅さ、クリティカル・スローイング・ダウン（臨界減速）、モード・ミキシングの不良、および高い自己相関といった問題に直面する。対照的に、ノーマライジング・フロー（NF）、GAN、VAEなどの機械学習アプローチは、より高速なサンプリングを提供するが、重大なトレードオフに直面する。尤度ベースの学習（前方KLまたはFKL）は、「モード・カバリング」を引き起こし、低密度領域に確率質量を割り当てる傾向がある。一方、逆KL（RKL）学習は「モード崩壊」を引き起こしやすく、ターゲット分布の全サポートを捉えきれない傾向がある。さらに、既存のNFベースのソリューションであるボルツマン生成器などは、訓練のために計算コストの高いMCMC生成データに依存しており、目的関数の固有のバイアス・バリアンス・トレードオフに苦慮している。

手法：ScoreNF
著者らは、モード・カバリングとモード崩壊の両方の限界に対処するために、スコアベース学習をノーマライジング・フローのアーキテクチャに統合したフレームワークである ScoreNF を提案する。コアとなる手法は以下の通りである：

1. ハイブリッド目的関数: 単なるFKLやRKLに依存するのではなく、ScoreNFは以下の結合損失関数を最適化する：
   $$L_{net}(\theta) = L_{RKL}(\theta) + \lambda_1 L_{SM}(\theta)$$

   • 逆KL ($L_{RKL}$): モデル分布 $q_\theta$ とターゲット $p$ の間のダイバージェンスを最小化する。この項は本質的に高密度モードを探索するが、モード崩壊のリスクがある。
   • スコアマッチング ($L_{SM}$): モデルのスコア関数 ($\nabla_x \log q_\theta(x)$) とターゲットのスコア関数 ($\nabla_x \log p(x)$) の間のフィッシャー・ダイバージェンスを最小化する。ターゲット密度は定数を除いて既知であるため、正確なスコア $\nabla_x \log p(x) = -\nabla_x H(x)$ を解析的に計算できる。この項は正則化因子として機能し、モデルがすべてのモードにわたって局所的な勾配情報を整合させるよう促し、それによってバランスの取れたカバレッジを促進する。
   • アニーリング: スコアマッチング損失の重み $\lambda_1$ は、訓練中に段階的にアニーリングされる。これにより、スコア項が初期段階でモデルをすべてのモードへと導き、その後、RKL項がモード崩壊を引き起こすことなく密度推定を精緻化することが可能になる。

2. アーキテクチャとサンプリング: モデルは、単純な基底分布をターゲット空間へと写像するために可逆フロー $f_\theta$ を利用する。反復的なランジュバン動力学に依存する従来のスコアベースモデルとは異なり、ScoreNFはフローの可逆性を利用して、シングルパスでサンプルを生成する。

3. バイアス補正 (IMH): フローの潜在的な近似誤差にかかわらず、最終的なサンプルが漸近的に偏りを持たないことを保証するため、著者らは**独立メトロポリス・ヘイスティングス（IMH）**アルゴリズムを採用している。学習された分布 $q_\theta$ は提案分布として機能し、サンプルは非正規化ターゲット密度とモデル密度の比に基づいて受理または拒絶される。

主な貢献

• フレームワークの提案: 非正規化分布を効果的にモデリングするために、スコアベース学習とノーマライジング・フローを組み合わせたScoreNFを導入した。
• サンプリングの病理的挙動の軽減: 本手法は、スコア情報を正則化因子として活用することで、モード崩壊（RKLに共通）とモード・カバリング（FKLに共通）の両方を軽減できることを示している。
• データの効率性: ScoreNFは、大幅に削減された訓練アンサンブルサイズ（250サンプルまで）においても高い性能を維持できることを示しており、高価なMCMC生成訓練データへの依存を軽減する。
• 評価指標: モード崩壊とモード・カバリングの挙動を共同で診断するために、負の対数尤度（NLL）と逆NLL（RNLL）を用いたデュアル・メトリック・アプローチを導入し、検証した。

実験結果
本手法は、合成2D混合ガウス分布（MOG-4およびMOG-8）と、高次元（64次元）のスカラー $\phi^4$ 格子場理論分布を用いて検証された。

• 合成分布 (MOG-4/8): ScoreNFは、FKL、RKL、および純粋なスコアマッチング（SM）のベースラインと比較して優れた性能を示した。可視化により、ScoreNFがFKLのようなカバリング挙動やRKLのような崩壊を起こすことなく、すべてのモードを捉えていることが確認された。定量的に、最も低いNLLおよびRNLLを達成し、高い忠実度とバランスの取れたカバレッジを示した。
• $\phi^4$ 格子場理論: この高次元の物理学的応用において、ScoreNFはFKL、RKL、およびSMを上回った。特筆すべきは、RKLとSMが（モード・カバレッジの限定を示す）低いRNLL値を示した一方で（および高いNLLを示すモード崩壊を示した一方で）、ScoreNFは極端なFKLのRNLLペナルティを伴わずに低いNLLを達成し、より優れたバランスを実現したことである。
• データの希少性に対する堅牢性: 訓練サンプルサイズ（10,000、1,000、250）を変化させた実験では、FKLの性能はサンプル数が減少するにつれて著しく低下した。対照的に、ScoreNFは極めて少ない訓練データであっても、低いNLLと高い有効サンプルサイズ（ESS）を維持し、堅牢であった。
• SOTAとの比較: ScoreNFは、$\phi^4$ 分布においてFABやiDEMといった他の最先端のベースラインをも上回った。

意義と主張
ScoreNFは、特にデータが乏しい、あるいは計算負荷の高い領域において、非正規化分布をサンプリングするためのより効率的で信頼性の高いフレームワークを提供すると主張している。スコアマッチングをNFフレームワークに統合することで、本手法は純粋な尤度ベースの学習（FKL/RKL）に固有の限界を克服しつつ、純粋なスコアベースモデルのような反復的なサンプリングコストを必要としない。著者らは、RKLとスコアマッチングの組み合わせにより、ターゲット分布のすべてのモードを学習しながら密度推定を精緻化できることを強調しており、これは従来のMCMC法が遅く、標準的な生成モデルが全分布サポートを捉えきれない複雑な物理系に対する堅牢なソリューションを提供する。また、本研究は、生成モデルの具体的な失敗モード（崩壊かカバリングか）を診断するためのツールとしてのNLL/RNLLペアの有用性を浮き彫りにしている。