[論文レビュー] Probabilistic Bayesian optimal experimental design using conditional normalizing flows

Rafael Orozco, Felix J. Herrmann+1 more|arXiv (Cornell University)|Feb 28, 2024

Optimal Experimental Design Methods被引用数 5

ひとこと要約

本論文は、条件付き正規化フローを共同訓練し期待情報ゲインを最大化し、ベイズ実験計画設計のための確率的二値設計マスクを最適化する手法を提案し、高次元MRIデータで実証した。

ABSTRACT

Bayesian optimal experimental design (OED) seeks to conduct the most informative experiment under budget constraints to update the prior knowledge of a system to its posterior from the experimental data in a Bayesian framework. Such problems are computationally challenging because of (1) expensive and repeated evaluation of some optimality criterion that typically involves a double integration with respect to both the system parameters and the experimental data, (2) suffering from the curse-of-dimensionality when the system parameters and design variables are high-dimensional, (3) the optimization is combinatorial and highly non-convex if the design variables are binary, often leading to non-robust designs. To make the solution of the Bayesian OED problem efficient, scalable, and robust for practical applications, we propose a novel joint optimization approach. This approach performs simultaneous (1) training of a scalable conditional normalizing flow (CNF) to efficiently maximize the expected information gain (EIG) of a jointly learned experimental design (2) optimization of a probabilistic formulation of the binary experimental design with a Bernoulli distribution. We demonstrate the performance of our proposed method for a practical MRI data acquisition problem, one of the most challenging Bayesian OED problems that has high-dimensional (320 $\times$ 320) parameters at high image resolution, high-dimensional (640 $\times$ 386) observations, and binary mask designs to select the most informative observations.

研究の動機と目的

• 予算制約の下で事前分布を事後分布へ更新するベイズ最適実験設計（OED）を動機づける。
• 尤度ベースの生成モデルとスケーラブルな学習を用いてEIGの計算上の課題に対処する。
• ロバストでスケーラブルな設計最適化を可能にする確率的で学習可能な二値設計マスクを導入する。
• 高次元のMRIデータ取得問題でアプローチを実証する。
• 最適化された設計が事後不確実性を低減し再構成品質を向上させることを示す。

提案手法

• EIG最適化が、結合分布の下で期待事後対数尤度を最大化することと同値であることを導出する。
• 厳密な尤度を持つ条件付き正規化フロー（CNF）を用いて設計勾配の逆伝搬を可能にし、ネットワークパラメータと設計の同時最適化を行う。
• 二値設計を学習可能な連続マスク w を介してベルヌーイ分布変数としてパラメータ化し、正規化予算 s を用いて二値マスクへ写像する。
• データ-事後分布上の期待対数尤度を最大化することにより、CNFと設計パラメータを同時に訓練する。
• 320×320のパラメータ空間と640×386の観測空間を持つ大規模逆問題を解くことでMRIデータに適用する。
• NMSEとSSIM指標を用いて近似事後サンプリングと設計性能を評価し、最適化設計と手作成マスクを比較する。

(a) (a) Posterior sample $\mathbf{x}\sim p_{\hat{\theta}}$

実験結果

リサーチクエスチョン

• RQ1 条件付き正規化フローは、ベイズ OED のための EIG を最適化するための扱いやすい尤度を提供できるか？
• RQ2 確率的二値設計と CNF の同時訓練は、高次元イメージング問題に対して頑健でスケーラブルな OED をもたらすか？
• RQ3 最適化された確率マスクは MRI における事後不確実性と再構成品質にどのような影響を与えるか？
• RQ4 大規模 MRI OED タスクへ CNF を適用する際の計算的トレードオフは何か？

主な発見

• この手法は EIG 最適化を CNF ベースの事後尤度訓練と結びつけ、設計パラメータへの逆伝搬を可能にする。
• 二値設計はベルヌーイパラメータ化を介して確率的に学習され、予算を守るマスクを再訓練なしに異なる予算で実現できる。
• FAST MRIの膝データで、学習された設計は低周波を滑らかに強調し、MRIの対称性に沿った各向同性で非対称なサンプリングを示す。
• 最適化設計は事後不確実性を低減し再構成品質を向上させ、事後平均の NMSE が低下し SSIM が高くなることで示される。

(b) (b) Posterior sample $\mathbf{x}\sim p_{\hat{\theta}}$