Factorizable Normalizing Flows for parameter-dependent density morphing

技術要約：パラメータ依存型密度モーフィングのための因子分解可能正規化流（Factorizable Normalizing Flows）

問題提起
高エネルギー物理学（HEP）やその他の科学分野において、尤度に基づく推論では、連続的なパラメータの関数として確率密度関数（PDF）がどのように変形するかをモデル化することが求められる。これらのパラメータは、物理的効果、較正定数、または系統誤差（不純パラメータ）を表す。正規化流（Normalizing Flows; NF）は単一の固定された密度のモデリングには優れているが、$K$個のパラメータのあらゆる可能な構成に対して個別のフローを学習することは、結合構成空間が$K$に対して指数関数的に増大するため、計算量的に困難である。

従来のHEPのアプローチは、「テンプレート・バリエーション」に依存している。これは、各不純パラメータに対して個別に$\pm 1\sigma$のシミュレーションを実行し、その結果を低次元のビン化されたデータに対して低次多項式を用いて補間する手法である。この手法は堅牢ではあるが、ビン化への依存と、不純効果が因子分解されるという仮定（相関を無視する）によって制限されている。これを、非ビン化（unbinned）かつ高次元の適合へと拡張するには、以下の条件を満たす手法が必要である：

1. 忠実性（Fidelidty）: 密度の変形を正確に捉えること。
2. プロファイリング能力: パラメータに対する滑らかで微分可能な依存関係を提供すること。
3. スケーラビリティ: 全$K$次元の不純空間をサンプリングすることを回避すること。
4. 解釈可能性: 各パラメータがどのように分布を再形成するかを明示的に示すこと。

手法：因子分解可能正規化流（FNFs）
著者らは、複雑な公称密度（nominal density）のモデリングと、系統的な変動のモデリングを分離する生成フレームワークである、因子分解可能正規化流（FNFs）を導入する。コアとなるアーキテクチャは、以下の2つのコンポーネントで構成される：

1. 固定された公称フロー（Fixed Nominal Flow）: リファレンス・シミュレーション（$\nu = 0$）に対して一度だけ学習される、高忠実度で可逆なフロー（例：Neural Spline Flow）。このコンポーネントは、系統的モデリングフェーズの間は凍結される。

2. 学習可能な因子分解可能変換（$T_\nu$）: 観測されたイベント $y$（系統 $\nu$ の影響を受けたもの）を、公称リファレンスフレーム $y_{\text{nom}}$ へと写像する変換を学習する。尤度は、変数変換の公式を通じて評価される：
   $$p(y | x, \nu) = p_{\text{nom}}(T_\nu(y | x, \nu) | x) \cdot |\det \nabla_y T_\nu(y | x, \nu)|$$

   変換 $T_\nu$ は、スケール（$s$）とシフト（$t$）のパラメータが不純パラメータ $\nu$ の多項式となるオートレグレッシブなアフィン写像として構築される。決定的なのは、これらの多項式が個々の不純に対して因子分解されている点である：
   $$s_j(y_{<j}, x, \nu) = \sum_{k=1}^K (\nu_k \alpha^k_j + \nu_k^2 \beta^k_j) + \sum_{k<\ell} \nu_k \nu_\ell \phi^k_\ell$$
   $$t_j(y_{<j}, x, \nu) = \sum_{k=1}^K (\nu_k \gamma^k_j + \nu_k^2 \delta^k_j) + \sum_{k<\ell} \nu_k \nu_\ell \psi^k_\ell$$

   係数（$\alpha, \beta, \gamma, \delta, \phi, \psi$）は、入力特徴量に条件付けられたマスク付き多層パーセプトロン（MLP）の出力である。

   • 加法性（Additivity）: 主要な構造は、各不純からの加算的な寄与を想定しており、これにより各不純を、その特定のパラメータのみを用いた$\pm 1\sigma$のサンプルのみを使用して個別に学習することが可能になる。
   • 相互作用項（Interaction Terms）: 相関関係を捉えるために、オプションとして交差項（$\phi, \psi$）を含めることができる。これには、結合された変動サンプルを用いた学習が必要となる。

主な貢献

• スケーラブルなアーキテクチャ: FNFは、離散的なテンプレート・バリエーションの手法を、連続的かつ微分可能な高次元の設定へと一般化する。これは、結合空間ではなく、パラメータごとの変動を用いて学習することで、「次元の呪い」を回避する。
• 線形スケーリング: 計算コストおよび学習データの要件は、不純パラメータの数に対して線形にスケールする（交差項を省略した場合）。これにより、数百の不確実性を伴う解析にも適用可能となる。
• 扱いやすい尤度: オートレグレッシブな構造により、ヤコビアン行列の行列式が三角行列となり、計算コストが低く（$O(D)$）、効率的な評価が可能となる。
• 解釈可能性: モデルは、各パラメータの線形、二次、および相互作用の効果を明示的に分離しており、系統がどのように密度を変形させるかについて明確な物理的解釈を提供する。

実験結果
著者らは、2つのクラスと2つの不純パラメータ（$\nu_{\text{shift}}$ および $\nu_{\text{squeeze}}$）を含む制御されたトイ・データセットを用いて、FNFを検証した。これらのパラメータは、運動学的変数および条件付きスコア分布の両方を変形させる。

• 再構成: 学習された変換は、課された変形を正常に再現した。純粋な平行移動の不純に対しては、モデルは単位スケールを持つ線形シフトを学習し、体積保存のスクイーズ（squeeze）に対しては、適切な非線形スケーリングとシフトを学習した。
• クロージャ（Closure）: 因子分解モデル（交差項なし）は、学習軸付近および公称点付近において、最適な尤度に対して良好な一致（closure）を示した（$\sim 10^{-3}$ nats/event以内）。
• 相関: 両方の不純が同時に強く変動する領域では、無視された相関のために因子分解モデルの適合品質に低下が見られた。オプションの交差項を有効にすることで、KLダイバージェンスが4分の1以下に減少した。これは、標準的な因子分解近似と完全な相関モデルとの間を補間する、本アーキテクチャの能力を示している。

意義と主張
本論文は、FNFが連続的な密度のモーフィングを必要とする推論ワークフローのための汎用的なツールを提供することを主張している。FNFは、組み合わせ爆発的な結合サンプリングのコストを負うことなく、高次元の非ビン化空間における系統誤差のモデリングを可能にすることで、非ビン化尤度適合の採用における主要な障壁を下げている。著者らは、本研究はパラメータ密度のモデリングを検証したものであり、パラメータ抽出のためのフルプロファイルされた非ビン化尤度適合へのFNFの展開は、今後の専用の研究課題であるとしている。このアプローチは、確立されたテンプレート・バリエーションのレシピを直接的かつ連続的に一般化したものであり、その堅牢性を維持しつつ、高次元の観測量への適用性を拡張するものである。