Probabilistic Bayesian optimal experimental design using conditional normalizing flows (2402.18337v1)
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.
- Dennis V Lindley. On a measure of the information provided by an experiment. The Annals of Mathematical Statistics, 27(4):986–1005, 1956.
- Variational bayesian optimal experimental design. Advances in Neural Information Processing Systems, 32, 2019.
- The frontier of simulation-based inference. Proceedings of the National Academy of Sciences, 117(48):30055–30062, 2020.
- Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516, 2014.
- Minimizing the expected posterior entropy yields optimal summary statistics. arXiv preprint arXiv:2206.02340, 2022.
- A unified stochastic gradient approach to designing bayesian-optimal experiments. In International Conference on Artificial Intelligence and Statistics, pages 2959–2969. PMLR, 2020.
- Bayesflow: Learning complex stochastic models with invertible neural networks. IEEE transactions on neural networks and learning systems, 33(4):1452–1466, 2020.
- Guided image generation with conditional invertible neural networks. arXiv preprint arXiv:1907.02392, 2019.
- On the universality of coupling-based normalizing flows. arXiv preprint arXiv:2402.06578, 2024.
- Learning-based optimization of the under-sampling pattern in mri. In Information Processing in Medical Imaging: 26th International Conference, IPMI 2019, Hong Kong, China, June 2–7, 2019, Proceedings 26, pages 780–792. Springer, 2019.
- Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
- fastmri: An open dataset and benchmarks for accelerated mri. arXiv preprint arXiv:1811.08839, 2018.
- End-to-end sequential sampling and reconstruction for mr imaging. In Proceedings of the Machine Learning for Health Conference, 2021.
- Invertiblenetworks. jl: A julia package for scalable normalizing flows. arXiv preprint arXiv:2312.13480, 2023.
- Extending loupe for k-space under-sampling pattern optimization in multi-coil mri. In Machine Learning for Medical Image Reconstruction: Third International Workshop, MLMIR 2020, Held in Conjunction with MICCAI 2020, Lima, Peru, October 8, 2020, Proceedings 3, pages 91–101. Springer, 2020.
- Image formation in diffusion mri: a review of recent technical developments. Journal of Magnetic Resonance Imaging, 46(3):646–662, 2017.
- A fast and scalable computational framework for large-scale high-dimensional bayesian optimal experimental design. SIAM/ASA Journal on Uncertainty Quantification, 11(1):235–261, 2023.
- Design amortization for bayesian optimal experimental design. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 8220–8227, 2023.
- Sequential experimental design for x-ray ct using deep reinforcement learning. arXiv preprint arXiv:2307.06343, 2023.
- Edge-promoting adaptive bayesian experimental design for x-ray imaging. SIAM Journal on Scientific Computing, 44(3):B506–B530, 2022.