SIMANF: Sample Free Learning of Unnormalized Distributions via Simulated Annealing in Normalizing Flows

技術要約: SIMANF

問題提起
高次元かつマルチモーダルな未正規化分布からの学習およびサンプリングは、統計物理学、化学、ベイズ推論などの科学分野において依然として大きな課題となっている。これらの分布は通常、エネルギーまたはハミルトニアン関数 $H(x)$ によって定義され、確率密度は $p(x) \propto e^{-H(x)}$ となるが、正規化定数 $Z$ は計算上、手に負えない（intractable）ものである。

従来のサンプリング手法であるマルコフ連鎖モンテカルロ法（MCMC）は、混合の遅さ、クリティカル・スローイング・ダウン、および十分に離れたモード間の遷移の困難さにより、これらの分布に対して苦戦することが多い。一方、深層生成モデル、特にノーマライジング・フロー（NF）は効率的なサンプリングのための代替案となるが、ターゲットのサンプルなしで学習する場合、以下の決定的なトレードオフに直面する：

• 前方KL（FKL） 学習はモードの被覆（coverage）を促進するが、ターゲット分布からのサンプルを必要とするため、サンプリングのボトルネックが再導入される。
• 逆KL（RKL） 学習は、未正規化密度のみを用いて最適化できる（サンプルフリー）が、その「モード探索的（mode-seeking）」な挙動により、モデルが少数の支配的なモードに確率質量を集中させて他のモードを無視してしまう**モード崩壊（mode collapse）**を引き起こしやすい。

手法: SIMANF
著者は、SIMANF（Simulated Annealing in Normalizing Flowsによる未正規化分布のサンプルフリー学習）を提案する。これは、シミュレーテッド・アニーリングをノーマライジング・フローと統合することで、ターゲットサンプルを必要とせずにモード崩壊を軽減するフレームワークである。

1. アニーリング・パスの構築: 直接 $p(x)$ に対するRKLを最適化する代わりに、SIMANFは一連のアニーリングされた分布を構築する：
   $$p_j(x) \propto p(x)^{\beta_j}, \quad 0 < \beta_1 < \beta_2 < \dots < \beta_n = 1$$
   ここで、小さな $\beta$ 値は高温で平滑化された分布に対応し、エネルギー地形が平坦化されることで、状態空間全体への探索を容易にする。$\beta$ が増加するにつれ、分布は段階的にターゲットへと鋭くなっていく。

2. 重み付きRKLによる逐次学習: ノーマライジング・フローは、これらのステージを通じて逐次的に学習される。目的関数は、重み付きのRKLダイバージェンスの総和である：
   $$L(\theta) = \sum_{j=1}^n \lambda_j KL(q_\theta \| p_j)$$
   これは、有効なアニーリングパラメータ $\beta_{eff} = \sum \beta_j \lambda_j$ を持つRKL目的関数と数学的に等価である。学習は $\beta_{eff}$ が0に近い状態（平滑な分布）から始まり、徐々に1へと増加していく。極めて重要な点は、ステージ $j$ で学習されたフローのパラメータがステージ $j+1$ の初期化に転送されることで、モデルが複数のモードの被覆を維持しながら、表現を段階的に洗練できる点である。

3. アニーリング後の精緻化（Post-Annealing Refinement）: アニーリング・スケジュールがターゲット分布（$\beta_{eff}=1$）に達した後、精緻化ステージが導入される。このステージでは、標準的なRKLと、重要度重み付きの前方KL目的関数を組み合わせる：
   $$L_R(\theta) = \alpha_1 KL(p \| q_\theta) + \alpha_2 KL(q_\theta \| p)$$
   ターゲットサンプルは利用できないため、前方KL項はフロー自身によって生成されたサンプルを用い、重要度重み $w(x) = p(x)/q_\theta(x)$ によって推定される。このステップは、FKLのモード被覆の傾向とRKLのモード探索的な性質の両方を活用することで、サンプリング効率と密度推定を向上させることを目的としている。

主な貢献

1. サンプルフリー・フレームワーク: 著者は、学習中にターゲット分布からのサンプルを一切必要とせずに、高次元かつマルチモーダルな未正規化分布を学習するフレームワークであるSIMANFを提案している。
2. 重み付きアニーリング目的関数: 一連のアニーリングされた分布に対する重み付きRKL目的関数を定式化した。これにより、スムーズな初期形態からターゲットへの緩やかな遷移が提供され、学習された表現を各ステージ間で効果的に転送することで、時期尚早なモード崩壊を防いでいる。
3. ハイブリッド精緻化ステージ: アニーリング後の精緻化ステージを導入し、フロー生成サンプルを用いた重要度重み付きFKL目的関数とRKLを組み合わせることで、真のターゲットサンプルなしにサンプリング効率を向上させている。
4. 実証的検証: 本手法は、Many-Well分布および高次元スカラー $\phi^4$ 格子場理論分布を用いて評価され、複雑なマルチモーダル構造の学習における有効性が示されている。

結果
SIMANFを、Boltzmann Generators（FKLおよびRKLで学習）、Flow Annealed Importance Sampling Bootstrap (FAB)、Iterated Denoising Energy Matching (iDEM)、およびFUNDを含むベースラインと比較評価した。

• Many-Well (MW) 分布 (8Dおよび16D): SIMANFは両方の次元において最も低い負の対数尤度（NLL）を達成しており、これはターゲット分布との優れた一致を示している。また、高い有効サンプルサイズ（ESS）も達成しており、FUND (FKL) に匹敵し、FABやiDEMよりも有意に優れている。対照的に、BG (RKL) はモード崩壊の兆候（高いNLL、低いRNLL）を示し、BG (FKL) は過剰なモード被覆（低いESS）を示した。
• スカラー $\phi^4$ 格子場理論 (8x8および10x10): 8x8格子において、SIMANFはすべての手法の中で最も低いNLLと最高のESSを達成した。より困難な10x10格子においても、最も低いNLLを維持し、ESSにおいてFABおよびiDEMを大幅に上回り、FUND (FKL) に匹敵する性能を示した。
• 計算効率: アニーリングは学習のオーバーヘッドを導入するものの、SIMANFはFABやFUNDよりも学習時の計算効率が高い。推論時には、ノーマライジング・フローの効率性を保持しており、iDEMのような反復的なSDEシミュレーションを必要とせず、単一のフォワードパス（10,000サンプルで約1.2秒）でサンプルを生成できる。
• アブレーション研究: 10x10 $\phi^4$ 分布において、アニーリング後の精緻化ステージはNLLに対してはわずかな改善しか示さなかったが、ESSを大幅に向上させた。これは、サンプリング効率を高める上での同ステージの役割を裏付けている。

意義と主張
本論文は、SIMANFが複雑な未正規化分布のサンプルフリー学習において、学習コストとサンプリング効率の間に有利なバランスを提供すると主張している。シミュレーテッド・アニーリングを利用して、フローをスムーズな初期形態からターゲットへと導くことで、直接的なRKL学習に固有のモード崩壊を軽減しつつ、FKL学習に伴うサンプリングコストを回避している。著者は、アニーリングベースの学習を、ターゲットサンプルが入手できない、あるいは生成にコストがかかる高次元マルチモーダル分布を学習するための有望なアプローチとして位置づけている。結果は、SIMANFが合成（Many-Well）および物理ベース（$\phi^4$）の両方の分布において、次元が増加しても複雑な構造を効果的に学習できることを示唆している。