Probabilistic Bayesian optimal experimental design using conditional normalizing flows

TL;DR

This work tackles Bayesian optimal experimental design (OED), a computationally challenging problem due to double intractability and high dimensionality, by proposing a joint optimization that trains a conditional normalizing flow (CNF) to maximize the expected information gain $EIG$ and learns a probabilistic binary design for sensor placement. The key idea is that maximizing $EIG$ is equivalent to maximizing the posterior likelihood of a conditional generative model, enabling end-to-end training of the CNF parameters $\theta$ alongside design parameters, with CNFs offering exact likelihoods and memory efficiency through invertible mappings. A probabilistic mask design is introduced to handle binary designs, parameterized by weights $\mathbf{w}$ and a budget constraint, enabling gradient-based optimization and post-hoc budget adjustments. The method is demonstrated on a high-dimensional MRI imaging task using the FASTMRI knee dataset, where an amortized posterior sampler and an optimized sampling mask are learned; results show reduced posterior uncertainty and improved image reconstruction metrics (e.g., lower NMSE and higher SSIM) compared to a hand-crafted baseline, illustrating a scalable pathway for Bayesian OED in large-scale imaging. Overall, the paper contributes a principled, scalable framework that unifies likelihood-based CNFs with probabilistic design to enable efficient, robust OED in high-dimensional, real-world problems.

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.

Figures (5)

• Figure 1: Optimized design compared to hand-crafted design.
• Figure 2: Posterior samples from the baseline method with hand-crafted design.
• Figure 3: Posterior samples from our method with optimized design
• Figure 4: Pointwise statistics from the baseline compared to our method.
• Figure 5: Comparing reduction in uncertainty achieved by our method.