Wasserstein Gradient Flows for Batch Bayesian Optimal Experimental Design

Louis Sharrock

stat.MLMar 12, 2026

Wasserstein Gradient Flows for Batch Bayesian Optimal Experimental Design

Authors: Louis Sharrock

Organizations: Department of Statistical Science, University College London.

Abstract

Bayesian optimal experimental design (BOED) provides a powerful, decision-theoretic framework for selecting experiments so as to maximise the expected utility of the data to be collected. In practice, however, its applicability can be limited by the difficulty of optimising the chosen utility. The expected information gain (EIG), for example, is often high-dimensional and strongly non-convex. This challenge is particularly acute in the batch setting, where multiple experiments are to be designed simultaneously. In this paper, we introduce a new approach to batch EIG-based BOED via a probabilistic lifting of the original optimisation problem to the space of probability measures. In particular, we propose to optimise an entropic regularisation of the expected utility over the space of design measures. Under mild conditions, we show that this objective admits a unique minimiser, which can be explicitly characterised in the form of a Gibbs distribution. The resulting design law can be used directly as a randomised batch-design policy, or as a computational relaxation from which a deterministic batch is extracted. To obtain scalable approximations when the batch size is large, we then consider two tractable restrictions of the full batch distribution: a mean-field family, and an i.i.d. product family. For the i.i.d. objective, and formally for its mean-field extension, we derive the corresponding Wasserstein gradient flow, characterise its long-time behaviour, and obtain particle-based algorithms via space-time discretisations. We also introduce doubly stochastic variants that combine interacting particle updates with Monte Carlo estimators of the EIG gradient. Finally, we illustrate the performance of the proposed methods in several numerical experiments, demonstrating their ability to explore multimodal optimisation landscapes and obtain high-utility batches in challenging examples.

Figures & tables

Panel of Figure 1 — Figure 1: Bayesian optimal experimental design as an optimisation problem over the space of probability measures . We lift the original optimisation problem over a design point $\xi\in\Xi$ (Fig. 1(a) ) to an optimisation problem over a design distribution $\mu\in\mathcal{P}(\Xi)$ (Fig. 1(b) ), before incorporating an entropic regulariser to ensure that this optimisation problem is strictly convex, and thus admits a unique optimum (Fig. 1(c) ).

Figure 2: The interacting particle system (IPS) . We plot the trajectories of

N=100

particles over

T=500

iterations (orange), the kernel density estimate of the final particle distribution (orange), and the target expected information gain (EIG) (blue dashed).

Method	Design-law ansatz	State evolved	Main advantage	Main limitation	Computational scaling
WGF (Joint)	$\nu_{m}\in\mathcal{P}(\Xi^{m})$	full batch $\bm{\xi}\in\Xi^{m}$	most expressive; directly targets joint batch law	scales poorly with batch size $m$	$\mathcal{O}(R\,C_{\nabla G})$ per iteration
WGF (MF)	$\nu_{m}=\mu_{1}\otimes{\cdot}\mkern 1.0mu{\cdot}\mkern 1.0mu{\cdot}\otimes\mu_{m}$	$m$ particle systems $(\xi_{b}^{i,N})^{i\in[N_{\mathrm{mf}}]}_{b\in[m]}$	allows within-batch specialisation via distinct marginals	must learn $m$ coupled marginals	$\mathcal{O}(mN_{\mathrm{mf}}K\,C_{\nabla_{1}G})$ per iteration
WGF (MF-IID)	$\nu_{m}=\mu^{\otimes m}$	one particle system $(\xi^{i,N})^{i\in[N]}$	simplest and most scalable approximation	does not explicitly encourage within-batch diversity	$\mathcal{O}(NK\,C_{\nabla_{1}G})$ per iteration
WGF (MF-IID-REP)	$\nu_{m}=\mu^{\otimes m}$	one particle system $(\xi^{i,N})^{i\in[N]}$	scalable and explicitly promotes diversity	introduces an additional repulsion hyperparameter	$\mathcal{O}(N(K\,C_{\nabla_{1}G}+K_{\mathrm{rep}}\,C_{\nabla r}))$ per iteration

Table 1: Summary of the four optimisation schemes considered in this paper. The methods are ordered from the most expressive but least scalable formulation to the most scalable structured approximation. We write

R

for the number of joint chains,

N_{\mathrm{mf}}

for the number of particles per-coordinate for the mean-field approximation,

N

for the number of particles in the i.i.d. approximation,

K

for the number of sampled partner-tuples used to approximate the interaction term,

K_{\mathrm{rep}}

for the number of repulsion samples,

C_{\nabla G}

for the cost of one batch-gradient evaluation,

C_{\nabla_{1}G}

for the cost of one partial-gradient evaluation, and

C_{\nabla r}

for the cost of one repulsion-gradient evaluation.

Panel of Figure 3 — Figure 3: Comparison of pointwise optimisation and distributional optimisation for a one-dimensional experimental design problem. The top row (Fig. 3(a) - Fig. 3(c) ) shows the results of directly optimising the EIG using GA (purple); the bottom row (Fig. 3(d) - Fig. 3(f) ) shows the results of optimising the entropy-regularised objective using the WGF (blue). To be specific, Fig. 3(a) and Fig. 3(d) show the empirical distribution of the final designs generated by the two approaches, given a uniform initialisation over the interval $[-3.5,3.5]$ . Fig. 3(b) and Fig. 3(e) show the corresponding trajectories; while Fig. 3(c) and Fig. 3(f) show the mapping from initial designs $\xi_{0}$ to final designs $\xi_{T}$ . In this example, gradient ascent converges to the local maximisers associated with its basins of attraction (Fig. 3(a) - Fig. 3(c) ). Conversely, the additional noise allows the WGF to discover the global maximum (Fig. 3(d) - Fig. 3(f) ).

Panel of Figure 4 — Figure 4: Comparison of stochastic pointwise optimisation and stochastic distributional optimisation for a one-dimensional experimental design problem. Fig. 4(a) displays the stochastic estimate of the EIG landscape. Fig. 4(b) reports a histogram of the final EIG values obtained via SGA trajectories (purple) and WGF particles (blue), after initialisation near one of the local maxima. Fig. 4(c) illustrates the posterior entropy associated with the “best” result obtained via stochastic gradient ascent (purple) or via the WGF (blue), as measured by the EIG, after initialisation at this same local maximum.

Panel of Figure 5 — Figure 5: Comparison of the designs obtained using stochastic pointwise optimisation (blue) and stochastic distributional optimisation (orange) for a two-dimensional non-linear sensor placement problem, for three different initialisations. Fig. 5(a) displays the designs obtained using SGA with multiple restarts (blue) or i.i.d. copies of the WGF (orange), given a uniform initialisation around the minor mode. Fig. 5(b) displays the corresponding results given a uniform initialisation over the entire domain $\Xi=[-5,5]^{2}$ . Fig. 5(c) displays the corresponding results given a uniform initialisation far from either mode.

Panel of Figure 6 — Figure 6: A comparison of the EIG (higher is better) achieved by the designs obtained using stochastic pointwise optimisation (blue) and stochastic distributional optimisation (orange) for a two-dimensional non-linear sensor placement problem, for three different initialisations. In all three cases, the WGF outputs designs corresponding to higher values of the EIG.

Panel of Figure 7 — Figure 7: A comparison of the posterior uncertainty (lower is better) of the designs obtained using stochastic pointwise optimisation (blue) and stochastic distributional optimisation (orange) for a two-dimensional non-linear sensor placement problem, for three different initialisations. In all three cases, the WGF outputs designs corresponding to lower values of the posterior uncertainty.

Panel of Figure 8 — Figure 8: A comparison of the EIG obtained using pointwise optimisation and distributional optimisation, including both joint and mean-field approaches, as a function of the batch size. We plot the achieved EIG as a function of the batch size $m\in\{2,3,5,10,20,50,100,200,500,1000\}$ for six design strategies: repeating the best single design, gradient ascent, WGF (Joint) , WGF (MF) , WGF (MF-IID) , and WGF (MF-IID-REP) . The points show the mean EIG over 5 independent random runs, while the error bars denote $\pm$ one standard error. The “distributional” methods (green, red, purple, brown) are superior to the “pointwise” method (orange) across all batch sizes; while the repeat-single baseline is clearly suboptimal due to redundancy in repeated measurements (blue). For larger batch sizes, the single-law i.i.d. approximations (purple, brown) display an increasing advantage over the joint and coordinate-wise mean-field methods (green, red).

Figure 9: EIG summaries for the pharmacokinetic (PK) sampling-time benchmark. Each boxplot summarises the distribution of independently evaluated

\mathrm{EIG}_{m}

values obtained from designs returned over multiple optimisation seeds (higher is better). Boxes show the interquartile range with median (orange line); whiskers extend to

1.5\times

IQR and circles denote outliers; green triangles indicate the mean.

Figure 10: Comparison of designs for the pharmacokinetic (PK) sampling-time benchmark. The grey curves show

100

latent PK mean trajectories under prior draws

\theta\sim\pi

. The coloured tick marks indicate the

m

selected sampling times for each method (one row per method).

Figure 11: EIG summaries for the FitzHugh–Nagumo (FHN) sampling-time benchmark. Each boxplot summarises the distribution of independently evaluated

\mathrm{EIG}_{m}

values obtained from designs returned over multiple optimisation seeds (higher is better). Boxes show the interquartile range with median (orange line); whiskers extend to

1.5\times

IQR and circles denote outliers; green triangles indicate the mean.

Figure 12: Comparison of designs for the FitzHugh–Nagumo (FHN) sampling-time benchmark. The grey curves show

100

latent FHN mean trajectories under prior draws

\theta\sim\pi

. The coloured tick marks indicate the

m

selected sampling times for each method (one row per method).

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 13: EIG summaries for the pharmacokinetic sampling-time benchmark for

m=10

Figure 14: EIG summaries for the pharmacokinetic sampling-time benchmark for

m=20

Figure 15: EIG versus step-size for the pharmacokinetic sampling-time benchmark .

Figure 16: EIG versus number of chains used by WGF (Joint) and WGF (Joint) (FUSE) for the pharmacokinetic sampling-time benchmark .

Figure 17: EIG versus number of particles per coordinate used by WGF (MF) and WGF (MF) (Sub) for the pharmacokinetic sampling-time benchmark .

Figure 18: EIG versus number of particles used by WGF (MF-IID) and WGF (MF-IID-REP) for the pharmacokinetic sampling-time benchmark .

Figure 19: EIG summaries for the Fitzhugh–Nagumo sampling-time benchmark for

m=10

Figure 20: EIG summaries for the Fitzhugh–Nagumo sampling-time benchmark for

m=15

Figure 21: EIG versus step-size for the Fitzhugh–Nagumo sampling-time benchmark . WGF (Joint) (FUSE) does not have a tunable step size parameter, but is included for comparative purposes.

Figure 22: EIG versus number of chains used by WGF (Joint) and WGF (Joint) (FUSE) for the Fitzhugh–Nagumo sampling-time benchmark .

Figure 23: EIG versus number of particles per coordinate used by WGF (MF) and WGF (MF) (Sub) for the Fitzhugh–Nagumo sampling-time benchmark .

Figure 24: EIG versus number of particles used by WGF (MF-IID) and WGF (MF-IID-REP) for the Fitzhugh–Nagumo sampling-time benchmark .

Bayesian Experimental Design

Bayesian experimental design uses probabilistic models to choose measurements that improve learning or support better decisions.

See all 63 papers

Beyond Expected Information Gain: Stable Bayesian Optimal Experimental Design with Integral Probability Metrics and Plug-and-Play Extensions
Apr 23, 2026Di Wu, Ling Liang, Haizhao YangBayesian Experimental Design Integral Probability Metrics
Department of Mathematics The University of Tennessee, Knoxville TN 37916, USA · Department of Mathematics and Department of Computer Science University of Maryland, College Park MD 20742, USA
Action-BED: Task-Driven Bayesian Experimental Design with Singly Intractable Objectives
Jun 22, 2026Tom Rossa, Angus Phillips, Tom RainforthBayesian Decision Theory Bayesian Experimental Design
Department of Statistics University of Oxford
Goal-driven Bayesian Optimal Experimental Design for Robust Decision-Making Under Model Uncertainty
May 25, 2026Jinwoo Go, Xiaoning Qian, Byung-Jun YoonBayesian Experimental Design Decision-Making under Uncertainty
Computing & Data Sciences, Brookhaven National Laboratory, Upton, NY · Department of Electrical & Computer Engineering, Texas A&M University, College Station, TX · Department of Computer Science & Engineering, Texas A&M University, College Station, TX

Apr 23, 2026J/K move · Enter open · S save

Beyond Expected Information Gain: Stable Bayesian Optimal Experimental Design with Integral Probability Metrics and Plug-and-Play Extensions

Di Wu, Ling Liang, Haizhao Yang

Department of Mathematics The University of Tennessee, Knoxville TN 37916, USA · Department of Mathematics and Department of Computer Science University of Maryland, College Park MD 20742, USA

Bayesian Experimental Design Integral Probability Metrics

Open paper AskPDF arXiv

Bayesian Optimal Experimental Design (BOED) provides a rigorous framework for decision-making tasks in which data acquisition is often the critical bottleneck, especially in resource-constrained settings. Traditionally, BOED typically selects designs by maximizing expected information gain (EIG), commonly defined through the Kullback-Leibler (KL) divergence. However, classical evaluation of EIG often involves challenging nested expectations, and even advanced variational methods leave the underlying log-density-ratio objective unchanged. As a result, support mismatch, tail underestimation, and rare-event sensitivity remain intrinsic concerns for KL-based BOED. To address these fundamental bottlenecks, we introduce an IPM-based BOED framework that replaces density-based divergences with integral probability metrics (IPMs), including the Wasserstein distance, Maximum Mean Discrepancy, and Energy Distance, resulting in a highly flexible plug-and-play BOED framework. We establish theoretical guarantees showing that IPM-based utilities provide stronger geometry-aware stability under surrogate-model error and prior misspecification than classical EIG-based utilities. We also validate the proposed framework empirically, demonstrating that IPM-based designs yield highly concentrated credible sets. Furthermore, by extending the same sample-based BOED template in a plug-and-play manner to geometry-aware discrepancies beyond the IPM class, illustrated by a neural optimal transport estimator, we achieve accurate optimal designs in high-dimensional settings where conventional nested Monte Carlo estimators and advanced variational methods fail.

Abstract

Bayesian Experimental Design

Explore similar work

Beyond Expected Information Gain: Stable Bayesian Optimal Experimental Design with Integral Probability Metrics and Plug-and-Play Extensions

Action-BED: Task-Driven Bayesian Experimental Design with Singly Intractable Objectives

Goal-driven Bayesian Optimal Experimental Design for Robust Decision-Making Under Model Uncertainty