Tractable Optimal Experimental Design using Transport Maps

TL;DR

The paper tackles Bayesian OED for nonlinear inverse problems with non‑Gaussian posteriors by introducing a transport-map surrogate that couples a tractable reference density to the joint design–data–parameters distribution via a Knothe–Rosenblatt rearrangement. Implemented through functional tensor trains (FTTs) and the deep inverse Rosenblatt transport (DIRT), the method yields tractable, sample‑based approximations to A‑ and D‑optimal criteria and extends to sequential designs with posterior preconditioning to reuse information across stages. Key contributions include a flexible KR‑map framework for general priors and optima, error bounds based on the Hellinger distance, TT‑based construction of KR maps, and a comprehensive numerical study (simple nonlinear, SEIR, and permeability inversion) comparing against nested Monte Carlo and illustrating efficiency gains. The approach enables efficient, scalable, and sequential experimental design in challenging Bayesian inverse problems where posteriors are non‑Gaussian and forward maps are nonlinear, with practical impact in areas such as disease modeling and subsurface imaging.

Abstract

We present a flexible method for computing Bayesian optimal experimental designs (BOEDs) for inverse problems with intractable posteriors. The approach is applicable to a wide range of BOED problems and can accommodate various optimality criteria, prior distributions and noise models. The key to our approach is the construction of a transport-map-based surrogate to the joint probability law of the design, observational and inference random variables. This order-preserving transport map is constructed using tensor trains and can be used to efficiently sample from (and evaluate approximate densities of) conditional distributions that are required in the evaluation of many commonly-used optimality criteria. The algorithm is also extended to sequential data acquisition problems, where experiments can be performed in sequence to update the state of knowledge about the unknown parameters. The sequential BOED problem is made computationally feasible by preconditioning the approximation of the joint density at the current stage using transport maps constructed at previous stages. The flexibility of our approach in finding optimal designs is illustrated with some numerical examples inspired by disease modeling and the reconstruction of subsurface structures in aquifers.

Figures (11)

• Figure 4.1: On the left, a contour plot of a two-dimensional density $\pi_{e,d}$ for the illustrative toy problem in remark:toy. A visualization of~$\pi_{d \left.\space\middle|\space\right. e^*}$ for three different design locations $e^*$ is visualized in the right image. The three different designs correspond to the vertical color-coded lines overlaid on the contour plot in the left.
• Figure 5.1: Flowchart for an SOED procedure depicting the feedback loop between finding optimal designs, using them to conduct the experiment and collect data, and updating the state of knowledge and utility function using the newly-collected data.
• Figure 5.2: Flowchart visualization of the greedy preconditioned SOED procedure outlined in algorithm:SOED_EIG.
• Figure 6.1: The expected information gain at different designs for the nonlinear example in subsection:nonlin_comparison (left), and a comparison of the spread of optimal designs obtained using NMC with $N = 255$ and our approach (right). A scatter plot is used to show the optimal designs obtained using our approach, whereas a boxplot is used for the NCM optimizers.
• Figure 6.2: The expected information gain for the two design nonlinear example in subsection:nonlin_comparison. On the left is the EIG as approximated using our approach, on the right is the EIG approximated using NMC with $N = 828$, and the middle is the EIG approximated with NMC using $N = 10000.0$ samples per design point. The optimal designs obtained using both approaches are visualized using the x symbol in the left and right plots.

Theorems & Definitions (17)

• remark 1: On the probability density for the designs
• lemma 1: CuiDolgov:2021:1
• lemma 2
• proof
• lemma 3: CuiDolgovZahm:2023:1
• proposition 1
• proof
• proposition 2
• proof
• theorem 1