[Paper Review] 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 Methods5 citations

TL;DR

The paper proposes jointly training a conditional normalizing flow to maximize expected information gain and optimizing a probabilistic binary design mask for Bayesian experimental design, demonstrated on high-dimensional MRI data.

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.

Motivation & Objective

• Motivate Bayesian optimal experimental design (OED) to update priors to posteriors under budget constraints.
• Address computational challenges of EIG by using likelihood-based generative models and scalable training.
• Introduce a probabilistic, trainable binary design mask to enable robust, scalable design optimization.
• Demonstrate the approach on a high-dimensional MRI data acquisition problem.
• Show that optimized designs reduce posterior uncertainty and improve reconstruction quality.

Proposed method

• Derive an equivalence showing EIG optimization is equivalent to maximizing the expected posterior log-likelihood under a joint distribution.
• Use conditional normalizing flows (CNFs) with exact likelihoods to enable back-propagation of design gradients and joint optimization of network parameters and design.
• Parametrize binary designs as Bernoulli-distributed variables via a learnable continuous mask w, and map to a binary mask with a normalization budget s.
• Train CNF and design parameters jointly by maximizing the expected log-likelihood over data-posteriors.
• Apply to MRI data by solving a large-scale inverse problem with a 320x320 parameter space and a 640x386 observation space.
• Evaluate amortized posterior sampling and design performance via NMSE and SSIM metrics, comparing optimized vs hand-crafted masks.

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

Experimental results

Research questions

• 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?

Key findings

• 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}}$