[論文レビュー] Learning Interpolations between Boltzmann Densities

Bálint Máté, François Fleuret|arXiv (Cornell University)|Jan 18, 2023

Gaussian Processes and Bayesian Inference被引用数 5

ひとこと要約

本稿では、逆KL発散の代わりに、時間に依存するベクトル場とエネルギー関数の補間を用いてボルツマン分布の補間を学習する、連続正規化流れのための新しい学習目的を提案する。この手法は連続性方程式の違反を最小化することで、モード崩壊を軽減し、ガウス混合分布や量子力学的二重井戸ポテンシャルなどのマルチモーダルなターゲットにおいて、訓練効率を向上させる。

ABSTRACT

We introduce a training objective for continuous normalizing flows that can be used in the absence of samples but in the presence of an energy function. Our method relies on either a prescribed or a learnt interpolation $f_t$ of energy functions between the target energy $f_1$ and the energy function of a generalized Gaussian $f_0(x) = ||x/σ||_p^p$. The interpolation of energy functions induces an interpolation of Boltzmann densities $p_t \propto e^{-f_t}$ and we aim to find a time-dependent vector field $V_t$ that transports samples along the family $p_t$ of densities. The condition of transporting samples along the family $p_t$ is equivalent to satisfying the continuity equation with $V_t$ and $p_t = Z_t^{-1}e^{-f_t}$. Consequently, we optimize $V_t$ and $f_t$ to satisfy this partial differential equation. We experimentally compare the proposed training objective to the reverse KL-divergence on Gaussian mixtures and on the Boltzmann density of a quantum mechanical particle in a double-well potential.

研究の動機と目的

• マルチモーダルなターゲット密度の下での期待値推定において、逆KL発散で訓練された正規化流れにおけるモード崩壊問題に対処すること。
• 逆KL発散の制限を克服し、エネルギー関数の補間に沿った連続的輸送に基づく、物理学的インスピレーションを受ける学習目的を導入すること。
• 真のサンプルが入手できない状況でも、エネルギー関数の評価値とその勾配のみに依存して訓練を可能にすること。
• 学習されたエネルギー補間に沿って連続性方程式を満たすように時間に依存するベクトル場を最適化することにより、サンプル品質と訓練速度の両方を向上させること。
• ガウス混合分布や量子力学的ボルツマン密度といった挑戦的なターゲットにおいて、本手法の有効性を示すこと。

提案手法

• 時間に依存するエネルギー関数 $ f_t = (1-t)f_0 + t f_1 + t(1-t)f^\theta(t) $ を定義し、ここで $ f_0 $ は一般化されたガウス分布、$ f_1 $ はターゲットエネルギー関数である。
• エネルギー関数 $ f_t $ に従うボルツマン分布の族 $ p_t \propto e^{-f_t} $ を定義し、この族に沿ってサンプルを輸送する時間に依存するベクトル場 $ V_t $ を学習する。
• 適切な輸送の条件を連続性方程式 $ \partial_t p_t + \nabla \cdot (p_t V_t) = 0 $ として定式化し、これにより $ V_t $ と $ p_t $ を関連付ける。
• 連続性方程式の局所的違反を最小化することで $ V_t $ と $ f_t $ を最適化し、損失関数 $ \mathcal{L} = \int \mathcal{E} + \mathcal{E}^2 \, dt $ を用いる。ここで $ \mathcal{E} $ は残差誤差である。
• 時間軸全体にわたって共通のアーキテクチャを持つ多層パーセプトロン（MLP）を用いて $ f_t $ と $ V_t $ をパラメータ化し、自動微分により勾配を計算する。
• Adam最適化法を用い、コサインスケジュールによる学習率の減少を適用し、4次ルンゲ＝クッタ法を用いて50ステップで流れを統合する。

実験結果

リサーチクエスチョン

• RQ1 エネルギー関数の評価値のみが利用可能な状況において、連続性方程式に基づく学習目的が、正規化流れにおけるモード崩壊を低減できるか。
• RQ2 本手法の性能は、逆KL発散最小化と比較して、KL発散、有効サンプルサイズ、モードカバレッジの観点でどのように異なるか。
• RQ3 基本エネルギー関数とターゲットエネルギー関数の間で非線形補間 $ f_t $ を学習することで、線形補間よりも性能が向上するか。
• RQ4 高次元のマルチモーダルなターゲットにおいて、本手法の訓練速度とサンプル品質のスケーリング特性はどの程度か。
• RQ5 本手法は、二重井戸ポテンシャル内の量子粒子のボルツマン密度のような複雑な物理系にも一般化可能か。

主な発見

• 本手法は、ガウス混合分布および量子力学的二重井戸ターゲットの両方において、逆KL学習よりも顕著に低いKL発散を達成した。
• ガウス混合ベンチマークにおいて、逆KL法と比較して有効サンプルサイズが最大2.5倍向上した。
• モードカバレッジが著しく向上し、2次元ガウス混合分布の4つのモードすべてが、崩壊を伴わず学習された。
• 訓練がより速く収束し、安定性に優れた。10,000イテレーションで学習したモデルが、250,000イテレーションで学習した逆KLモデルを上回る結果を達成した。
• 学習された非線形補間 $ f_t $ は、エネルギー障壁の急勾配領域で線形補間を上回る性能を示した。
• 本手法はPINNフレームワークに自然に適合し、連続性誤差が残差損失として機能するため、PDE制約付き学習への応用可能性が広がると示唆された。