JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

技術要約：JADAI – 結合型適応設計およびベイズ推論（Jointly Amortizing Adaptive Design and Bayesian Inference）

問題提起
多くの科学・工学分野では、複雑なシミュレータによって生成された観測値 $x$ から、未観測のパラメータ $\theta$ を推論することが求められる。シミュレーションベース推論（SBI）は、明示的な尤度評価なしに事後分布 $p(\theta|x)$ を近似する能力を進展させてきたが、これらの手法は通常、固定された実験設計を前提としている。対照的に、ベイズ実験計画法（BED）および適応的設計（BAD）は、情報利得（例：期待情報利得、EIG）を最大化するために設計変数 $\xi$ を選択することに焦点を当てているが、多くの場合、データ収集後に別途行われる計算負荷の高い推論ステップを独立したものとして扱う。既存のアプローチは、しばしばこれらのタスクを切り離している。BAD手法は、非アモルタイズ（非償却）な低速な推論や単純なパラメトリック族に依存することが多く、一方でSBI手法はデータの取得を最適化しない。この分離は、特に高次元またはマルチモーダルな設定において、情報の獲得を最適化できず、計算効率を低下させる原因となる。

手法
著者らは、適応的設計と事後推論を単一のエンドツーエンドで学習可能なシステムへと統合するフレームワークである JADAI（Jointly Amortizing Adaptive Design and Bayesian Inference）を提案する。その核心的な革新は、以下の3つのコンポーネントを結合的にアモルタイズ（償却）することにある：

1. ポリシーネットワーク ($\pi_\phi$): 現在の実験履歴に基づき、次の設計 $\xi_t$ を提案する。
2. 履歴／要約ネットワーク ($\eta_\omega$): 過去の設計と観測のペア $\{(\xi_k, x_k)\}_{k=1}^t$ のシーケンスを、固定次元の要約統計量 $h_t$ へとエンコードする。
3. 推論ネットワーク ($q_\psi$): 実験シーケンスの各ステップにおいて、事後分布 $p(\theta|h_t)$ を近似する。

最適化目的関数
JADAIは、実験シーケンスに沿った事後誤差の増分減少を集計する汎用的な損失関数を最小化する。期待情報利得（EIG）を直接最大化するのではなく（これは密度評価を必要とするため）、本フレームワークは事後損失の入れ子状の和（telescoping sum）を利用する。

• 訓練目的関数は、全期待情報利得（TEIG）の変分下界から導出される。
• ステップ $t$ における損失 $\ell_t$ は、事後推定器（例：拡散モデルまたはフローマッチングの損失）によって定義される。
• 効用関数 $u_T$ は、連続する損失の差の和として構成される：$u_T = \sum (\ell^*_{t-1} - \ell_t)$。ここで $\ell^*$ は、直前の損失のデタッチ（勾配停止）されたバージョンを表す。
• この定式化により、ポリシーおよび要約ネットワークは、事後密度の明示的な評価を必要とせずに、各ステップにおける事後損失の減少を最大化するように訓練される。これにより、実質的に情報利得を最大化する。

主要なアーキテクチャの選択

• 拡散モデル: 推論ネットワーク $q_\psi$ は拡散モデル（またはフローマッチング）を用いて実装されている。これにより、フレームは各実験ステップにおいて、高次元かつマルチモーダルな事後分布を近似できる。これは、正規化フローに制限されていた従来のBAD手法には欠けていた能力である。
• 訓練の安定性: 時間方向へのネストされたバックプロパゲーション（BPTT）を回避し、安定した訓練を確保するために、以下の手法を採用している：
  • デタッチされた履歴状態: ポリシーと履歴ネットワーク間の循環的な勾配経路を防ぐため、設計はデタッチされた要約状態から生成される。
  • カリキュラム学習: 訓練中にロールアウトの長さを段階的に増加させ、また、ネットワークがポリシーのみに依存する前に多様な設計に触れさせるための「設計事前分布のミックスイン（design prior mix-in）」戦略を使用する。

結果
JADAIは、ポリシーおよび事後構造の複雑さが段階的に増していく3つのベンチマークを用いて評価された：

1. 位置特定（Location Finding, LF）: ノイズを含む強度測定から光源の位置を推論するタスク。JADAIは、長いホライゾン（$T > 10$）において、最先端のベースライン（例：DAD, iDAD, RL-BOED）と比較して優れた、あるいは競争力のあるsPCE（逐次的事前分布対比推定）スコアを達成した。学習されたポリシーは探索と活用のバランスを効果的にとり、拡散ベースの事後分布は、光源のマルチモーダルかつ交換可能な性質を正しく捉えた。
2. 代替の弾力性（Constant Elasticity of Substitution, CES）: 好みの推論を含む困難な設計問題であり、情報量の多い設計は非常に狭い「スイートスポット」に存在する。JAD_AIは、標準的なホライゾン $T=10$ においてすべての既存のベースラインを上回り、ランダムな設計が情報を持たないシナリオにおける堅牢性を実証した。
3. MNIST画像発見: 部分的な測定から画像を再構成することを目的とした高次元タスク。JADAIは、最先端のSSIM（構造類似性指数）およびNRMSE（正規化平方根平均二乗誤差）スコアを達成し、CoDiffのベースラインを上回った。本手法は、数字のクラスを早期に特定し、局所的な詳細を精緻化することに成功し、加法的測定ノイズの下でも性能を維持した。

意義と主張
論文は、JADAIが設計と推論を統合することで、適応プロセス中の個別の高コストな推論ステップを排除し、大きな進歩をもたらしたと主張している。主な貢献は以下の通りである：

• 結合的アモルタイゼーション: 事後誤差の増分減少を集計する汎用的な効用を用いることで、設計と事後推論を結合的にアモルタイズする初のフレームワークであり、密度評価を必要としない。
• 複雑な事後分布へのスケーラビリティ: 拡散モデルを活用することで、JADAIは各実験ステップにおいて高次元かつマルチモーダルな事後分布を扱うことができ、正規化フローやガウス混合モデルに依存していた従来手法の表現力の限界を克服している。
• パフォーマンス: 本フレームワークは、高次元の画像再構成を含む多様なベンチマークにおいて、優れた、あるいは競争力のある性能を達成すると同時に、展開時の計算効率（ロールアウトはミリ秒単位で完了）を維持している。

著者らは、JADAIが現在は微分可能なシミュレータに依存しているものの、そのアプローチは微分可能なシミュレーションフレームワークの拡大する展望に適していると述べている。また、高次元の場合のサブ秒単位の展開において事後サンプリングがボトルネックとなることを認めているが、ポストホックな蒸留（distillation）が将来的な方向性であると示唆している。本研究は、勾配なしのブラックボックス・シミュレータ最適化を解決すると主張するものではなく、古典的なパイプラインが計算上非現実的となる微分可能な設定において、堅牢なソリューションとして位置づけている。