Amortised and provably-robust simulation-based inference

1. 問題定義と背景

• シミュレーションに基づく推論（SBI）の課題:
  • 複雑な科学・工学モデル（粒子物理学、疫学など）では、尤度関数が解析的に扱えない（intractable）場合が多く、シミュレーションを用いた推論が一般的です。
  • 近年、ニューラルネットワークを用いた条件付き密度推定（NLE, NPE など）により、一度学習すれば複数の観測データに対して推論を高速に行える「償却的（amortised）」推論が可能になりました。
• 頑健性の欠如:
  • 既存の SBI 手法の多くは、測定誤差や欠陥による**外れ値（outliers）**やモデルの誤指定に対して非常に脆弱です。外れ値が含まれると、事後分布が真のパラメータから大きく偏り、過信された推論結果（over-confident posteriors）を生み出します。
  • 既存の頑健な手法（MMD ベースの手法など）は、推論のたびに MCMC を実行する必要があり、償却的ではありません。また、理論的な頑健性の保証がない手法も多数あります。
• 研究の目的:
  • 外れ値に対して理論的に頑健であり、かつ一度学習すれば任意の新しい観測データに対して高速に推論可能な（償却的な）SBI 手法の開発。

2. 提案手法：Neural Score-Matching Bayes (NSM-Bayes)

提案手法は、**一般化ベイズ推論（Generalised Bayesian Inference: GBI）とスコアマッチング（Score Matching）**の枠組みを組み合わせ、ニューラルネットワークで近似した重み付きスコアマッチング損失を用います。

2.1 基本的なアプローチ

1. ステップ 1（学習）: 事前分布からサンプリングしたパラメータ $\theta$ と、それに対応するシミュレーションデータ $x$ を用いて、条件付き密度モデル $q_\phi(x|\theta)$（ニューラルネットワーク、例：Normalizing Flow や Mixture Density Network）を学習します。
   • この段階では、外れ値が含まれていないシミュレーションデータのみを使用するため、標準的なスコアマッチング損失（重みなし）で学習します。
2. ステップ 2（推論）: 観測データ $x^o_{1:n}$ が与えられたとき、一般化ベイズの枠組みを用いて事後分布を更新します。
   • 通常のベイズ更新（対数尤度の和）の代わりに、**重み付きスコアマッチング損失（Weighted Score-Matching Loss）**を最小化するように事後分布を定義します。
   • 事後分布は以下の形で得られます：
     $$\pi_{\text{NSM}}(\theta | x^o_{1:n}) \propto \exp\left( -\beta n L_{\text{NSM}}(\theta; x^o_{1:n}, \hat{\phi}_m) \right) \pi(\theta)$$
   • ここで、$L_{\text{NSM}}$ は観測データに対するスコアマッチングの誤差を重み関数 $w(x)$ で調整したものです。

2.2 頑健性のメカニズム

• 重み関数 $w(x)$ の導入:
  • 外れ値の影響を抑制するため、データ点 $x$ に応じて重み $w(x)$ を動的に調整します。
  • 具体的には、**逆マルチ二次関数（Inverse Multi-Quadratic: IMQ）**を重み関数として採用します：
    $$w_\zeta(x) = \left( 1 + \|x - \hat{\nu}_n\|^2_{\hat{\Xi}_n^{-1}} \right)^{-1/\zeta}$$
  • $\hat{\nu}_n, \hat{\Xi}_n$ は観測データからの頑健な平均・共分散推定量（例：MCD 推定量）です。
  • 外れ値（$\hat{\nu}_n$ から遠い点）は重み $w(x)$ が小さくなるように設計されており、これにより外れ値が事後分布に与える影響が抑制されます。
• パラメータ $\zeta$ の制御:
  • $\zeta$ の値を適切に設定することで（特に $0 < \zeta \le 1$ の場合）、理論的に外れ値に対する頑健性が保証されます。

2.3 2 つの変種

1. NSM-Bayes（一般形）:
   • 任意の条件付き密度モデルと事前分布を使用可能。
   • 推論時には MCMC が必要ですが、密度モデルの再学習は不要（部分的に償却的）。
2. NSM-Bayes-conj（共役形）:
   • 事前分布をガウス分布とし、密度モデルを特定のエネルギーベースモデル（指数族の形式）に制限します。
   • この制限により、事後分布が解析的にガウス分布として計算可能になります（共役更新）。
   • MCMC を不要とし、推論を完全に償却的（fully amortised）かつ高速に行うことができます。

3. 理論的貢献

論文の最大の強みは、償却的 SBI 手法として初めて、外れ値に対する頑健性の数学的証明を提供している点です。

• NLE の非頑健性の証明:
  • 従来のニューラル尤度推定（NLE）は、外れ値に対して「大域的バイアス頑健（globally-bias robust）」ではないことを証明しました（影響関数が有界でない）。
• NSM-Bayes の頑健性の証明:
  • 提案手法 NSM-Bayes および NSM-Bayes-conj が、適切な重み関数（IMQ など）と事前分布の下で、大域的バイアス頑健であることを証明しました。
  • 具体的には、外れ値が任意に大きくても、事後分布への影響（Posterior Influence Function）が有界に保たれることを示しています。
• 実用的な条件:
  • 混合密度ネットワーク（MDN）やマスクド自己回帰フロー（MAF）といった一般的なニューラルアーキテクチャにおいても、特定の条件下で頑健性が保証されることを示しました。

4. 実験結果

著者らは、疫学（SIR モデル）、無線通信（ラジオ伝播モデル）、および g-and-k 分布（SBI のベンチマーク）を用いた実験を行い、既存の頑健な手法（RSNLE, NPL-MMD, NPE-RS, ACE, GBI-SR）と比較しました。

• g-and-k 分布（外れ値あり）:
  • 10% の外れ値を含むデータに対して、NLE は事後分布が真のパラメータから大きく偏り、カバレッジ（信頼区間内の真値の割合）が 0% になりました。
  • 一方、NSM-Bayes と NSM-Bayes-conj は100% のカバレッジを達成し、真のパラメータの周りに正しく集中しました。
  • 推論時間においても、NSM-Bayes-conj は数秒で完了し、他の手法（特に MCMC を必要とするもの）よりも大幅に高速でした。
• SIR モデル（感染症の過少報告）:
  • 感染症データの過少報告（一方向の外れ値）をシミュレートした実験でも、NSM-Bayes-conj が最も高い精度と最短の推論時間を示しました。
  • 既存の手法（GBI-SR など）は、非微分可能なシミュレータでは MCMC の収束が悪く、計算コストが膨大になることが示されました。
• ラジオ伝播モデル（受信機故障）:
  • 故障したアンテナによるノイズ（外れ値）を含むデータに対しても、NSM-Bayes は頑健な推論を行い、NLE が大きくバイアスされるのに対し、真のパラメータを正確に推定しました。

5. 結論と意義

• 主要な成果:
  • 償却性と頑健性の両立: 既存の手法では「頑健だが遅い」か「速いが非頑健」かのトレードオフが存在しましたが、NSM-Bayes は両方を達成しました。
  • 理論的保証: 単なる経験的な頑健性ではなく、数学的に証明された頑健性を提供しています。
  • 計算効率: NSM-Bayes-conj により、MCMC を不要とした高速な推論が可能になり、実用的な応用が容易になりました。
• 意義:
  • 科学シミュレーションにおいて、測定誤差や不完全なデータが一般的である現実の問題に対し、信頼性の高い推論を提供する新しい基盤を築きました。
  • 疫学、無線通信、天体物理学など、高コストなシミュレーションを伴う分野において、外れ値に強い推論システムを構築する道を開きます。

この論文は、SBI の分野において、理論的厳密性と実用的な効率性を両立させた重要な進展をもたらすものです。