[论文解读] 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

一句话总结

本文提出将 conditional normalizing flow 联合训练，以最大化 expected information gain，并优化一个 probabilistic binary design mask 用于 Bayesian experimental design，在高维 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 $ imes$ 320) parameters at high image resolution, high-dimensional (640 $ imes$ 386) observations, and binary mask designs to select the most informative observations.

研究动机与目标

• 激励贝叶斯最优实验设计（OED）在预算约束下将先验更新为后验。
• 通过使用基于似然的生成模型和可扩展的训练来解决 EIG 的计算挑战。
• 引入一个概率化、可训练的二进制设计掩码，以实现鲁棒、可扩展的设计优化。
• 在高维 MRI 数据采集问题上演示该方法。
• 表明优化后的设计能降低后验不确定性并提高重建质量。

提出的方法

• 推导一个等价关系，表明 EIG 的优化等价于在联合分布下最大化期望后验对数似然。
• 使用带有精确似然的 conditional normalizing flows (CNFs)，以便对设计梯度进行反向传播并对网络参数和设计进行联合优化。
• 通过一个可学习的连续掩码 w，将二进制设计参数化为 Bernoulli 分布变量，并通过一个归一化预算 s 映射为二进制掩码。
• 通过对数据后验的期望对数似然进行最大化，联合训练 CNF 与设计参数。
• 通过解决一个规模较大的反问题将其应用到 MRI 数据，参数空间为 320×320，观测空间为 640×386。
• 通过 NMSE 和 SSIM 指标评估 amortized 后验采样和设计性能，比较优化掩码与手工设计掩码。

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

实验结果

研究问题

• RQ1 Can a conditional normalizing flow provide tractable likelihoods to optimize EIG for Bayesian OED?
• RQ2 Does joint training of CNFs with probabilistic binary designs yield robust, scalable OED for high-dimensional imaging problems?
• RQ3 How does an optimized probabilistic mask affect posterior uncertainty and reconstruction quality in MRI?
• RQ4 What is the computational trade-off when applying CNFs to large-scale MRI OED tasks?

主要发现

• The approach links EIG optimization to CNF-based posterior likelihood training, enabling back-propagation to design parameters.
• Binary designs are learned probabilistically via a Bernoulli parameterization, allowing budget-respecting masks without retraining for different budgets.
• On FAST MRI knee data, the learned design emphasizes low frequencies smoothly and reveals anisotropic, asymmetric sampling aligned with MRI symmetry.
• Optimized designs reduce posterior uncertainty and improve reconstruction quality, evidenced by lower NMSE and higher SSIM in posterior means.

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