Bayesian Experimental Design via Contrastive Diffusions

TL;DR

CoDiff introduces a pooled-posterior gradient and diffusion-based sampling to scale Bayesian Optimal Experimental Design, enabling efficient, single-loop optimization of the Expected Information Gain. By deriving EIG gradient expressions via a reparameterization trick and employing a pooled posterior as a sampling proxy, it circumvents nested Monte Carlo and leverages diffusion models for data-based priors. The method is demonstrated on sequential source localization and MNIST-based inverse problems, showing substantial gains over state-of-the-art approaches in information gains and posterior quality, and it extends BOED to diffusion-based generative models. While effective, it remains a greedy, likelihood-reliant approach; future work includes nonlinear forward models and simulation-based inference to broaden applicability and non-myopic planning.

Abstract

Bayesian Optimal Experimental Design (BOED) is a powerful tool to reduce the cost of running a sequence of experiments. When based on the Expected Information Gain (EIG), design optimization corresponds to the maximization of some intractable expected contrast between prior and posterior distributions. Scaling this maximization to high dimensional and complex settings has been an issue due to BOED inherent computational complexity. In this work, we introduce a pooled posterior distribution with cost-effective sampling properties and provide a tractable access to the EIG contrast maximization via a new EIG gradient expression. Diffusion-based samplers are used to compute the dynamics of the pooled posterior and ideas from bi-level optimization are leveraged to derive an efficient joint sampling-optimization loop. The resulting efficiency gain allows to extend BOED to the well-tested generative capabilities of diffusion models. By incorporating generative models into the BOED framework, we expand its scope and its use in scenarios that were previously impractical. Numerical experiments and comparison with state-of-the-art methods show the potential of the approach.

Figures (9)

• Figure 1: $28 \times 28$ Image ${\boldsymbol{\theta}}\xspace$ (1st column) reconstruction from seven $7\times 7$ sub-images ${\mathb{y}}\xspace={\mathb{A}}\xspace_{\boldsymbol{\xi}}\xspace{\boldsymbol{\theta}}\xspace + {\boldsymbol{\eta}}\xspace$ centered at seven central pixels ${\boldsymbol{\xi}}\xspace$ (designs) selected sequentially. Optimized vs. random designs: measured outcome ${\mathb{y}}\xspace$ (2nd vs. 3rd column) and parameter ${\boldsymbol{\theta}}\xspace$ estimates (reconstruction) with highest weights (upper vs. lower sub-row).
• Figure 2: Source localisation example. Prior (left) and pooled posterior (right) samples at experiment $k$. Final ${\boldsymbol{\xi}}\xspace_k^*$ (orange cross) at the end of the optimization sequence ${\boldsymbol{\xi}}\xspace_0, \cdot, {\boldsymbol{\xi}}\xspace_{T}$ (blue crosses). This optimization "contrasts" the two distributions by making the pooled posterior "as different as possible" from the prior.
• Figure 3: Source location. Median and standard error over 100 rollouts for SPCE, L$_2$ Wasserstein distance (log-scale), SNMC with respect to number of experiments $k$. Number of samples N+M=400.
• Figure 4: Image reconstruction. First 6 experiments (rows): image ground truth, measurement at experiment $k$, samples from current prior $p({\boldsymbol{\theta}}\xspace|{\mathb{D}}\xspace_{k-1})$, with best (resp. worst) weights in upper (resp. lower) sub-row. The samples incorporate past measurement information as the procedure advances. Each design steps takes $\sim7.3$s
• Figure 5: Source localization example. Experiments 0 (prior samples), 4, 8 and 12. As new design locations are selected (orange crosses), samples concentrate to the true sources (red crosses). Samples with lower weights in blue, higher weights in yellow.

Theorems & Definitions (2)

• Lemma 1: Donsker and Varadhan’s variational formula
• Lemma 2