Source-linked AI summary
A Unified Stochastic Gradient Approach to Designing Bayesian-Optimal Experiments
Adam Foster, Martin Jankowiak, Matthew O'Meara, Yee Whye Teh, Tom Rainforth
TL;DR
Existing BOED methods typically estimate expected information gain separately at candidate designs and then use a separate optimizer to select a design. This paper jointly optimizes variational lower bounds and design parameters with stochastic gradient ascent, and reports substantially improved performance, especially for high-dimensional designs.
Problem
Existing BOED approaches typically estimate expected information gain point by point before passing estimates to a separate design optimizer.
Source: §1 · p. 2Method
The paper jointly optimizes variational lower bounds on expected information gain with respect to both variational parameters and experimental designs using stochastic gradient ascent.
Source: §1 · p. 2Results
The unified method yields substantially improved performance, especially on high-dimensional design problems, and ACE generally did as well as the better of BA and PCE in all five experiments.
Source: §5 · p. 8Takeaways & Limitations
The authors recommend ACE as the default among BA, ACE, and PCE; BA performed well when its inference network could closely approximate the true posterior.
Source: §5 · p. 8PCE is effective when the prior and posterior are similar, so the prior is a suitable proposal for estimating p(y|ξ).
Source: §3.3 · p. 5
Abstract
from arXiv · showhide
We introduce a fully stochastic gradient based approach to Bayesian optimal experimental design (BOED). Our approach utilizes variational lower bounds on the expected information gain (EIG) of an experiment that can be simultaneously optimized with respect to both the variational and design parameters. This allows the design process to be carried out through a single unified stochastic gradient ascent procedure, in contrast to existing approaches that typically construct a pointwise EIG estimator, before passing this estimator to a separate optimizer. We provide a number of different variational objectives including the novel adaptive contrastive estimation (ACE) bound. Finally, we show that our gradient-based approaches are able to provide effective design optimization in substantially higher dimensional settings than existing approaches.
1 INTRODUCTION
The paper replaces separate pointwise EIG estimation and design optimization with a unified stochastic-gradient approach using variational lower bounds. It introduces several bounds, including ACE, and reports effective optimization with especially large gains on high-dimensional problems.
1 INTRODUCTION
- Existing BOED approaches estimate EIG separately for each design before an outer optimizer selects the design, making the process computationally difficult.The nested expectation also makes even a single EIG evaluation challenging.Source: §1 · p. 1
- The unified approach jointly optimizes a variational lower bound and the design parameters with stochastic gradient ascent, avoiding a separate outer optimization.Optimizing the variational parameters tightens the bound, while optimizing the design parameters searches for informative experiments.Source: §1 · p. 2
- The adaptive contrastive estimation (ACE) bound can be made arbitrarily tight while remaining suitable for simultaneous stochastic-gradient optimization of variational and design parameters.The paper introduces several classes of variational lower bounds to suit different problem settings.Source: §1 · p. 2
- The unified gradient approaches consistently outperform two-stage baselines, with particularly large gains on high-dimensional problems and more informative resulting experiments.The paper reports these gains across a range of design problems, including a high-dimensional pharmacology example.Source: §1 · p. 2
2 BACKGROUND
The background defines Bayesian optimal experimental design through expected information gain, then reviews its estimation and optimization challenges. Existing methods estimate EIG and search over designs in separate stages, using approaches such as nested Monte Carlo, variational estimators, and Bayesian optimization.
2.1 Bayesian optimal experimental design
- BOED selects the feasible design that maximizes expected information gain, the expected reduction in uncertainty about latent variables after an experiment.The experimental model specifies a prior, a likelihood for outcomes, and posterior updating; iterated designs can adapt future experiments using earlier outcomes.Source: §2.1 · p. 2
2.2 Estimating expected information gain
- Nested Monte Carlo estimates both integrals in EIG, but converges at O(T^-1/3) in total computational budget T.Making a point estimate also requires estimating the marginal outcome distribution or posterior before taking the outer expectation.Source: §2.2 · p. 3
- Amortized variational estimators learn the integrand’s functional form across outer samples; the Barber-Agakov lower bound uses an inference network to approximate the posterior.The BA bound is tight when the network can represent the true posterior for every outcome.Source: §2.2 · p. 3
- VNMC combines an inference network with additional samples to improve the integrand estimate, and its upper bound becomes tight as sample count L tends to infinity.This holds even when the network cannot directly represent the true posterior; the network is trained by minimization.Source: §2.2 · p. 3
2.3 Optimizing the EIG
- Most existing methods use two stages: noisy EIG estimates are made for candidate designs, then a separate optimizer chooses which design to evaluate.Options for the outer search include Bayesian optimization, evaluation of finite candidate sets, adaptive allocation of computation, and MCMC.Source: §2.3 · p. 3
3 GRADIENT-BASED BOED
The section replaces separate EIG estimation and design optimization with joint stochastic-gradient optimization of variational lower bounds. It develops BA, ACE, and PCE objectives and extends ACE to likelihood-free and iterated designs, with gradient estimators for optimization.
3 GRADIENT-BASED BOED
- Jointly maximizing a variational lower bound over (ξ, φ) estimates EIG while optimizing the design, replacing separate pointwise estimation and outer optimization.The approach uses stochastic gradient ascent on bounds expressed as expectations.Source: §3 · pp. 3–4
3.1 Barber-Agakov (BA)
- BA supplies a one-stage instance by jointly optimizing its bound over ξ and φ, using unbiased stochastic-gradient estimators for both.Source: §3.1 · p. 4
3.2 Adaptive contrastive estimation (ACE)
- ACE uses θ1:L as contrasts to θ0, producing a jointly maximizable EIG lower bound that rises with L and converges to EIG.The bound equals EIG for any L when qφ is the true posterior.Source: §3.2 · p. 4
3.3 Prior contrastive estimation (PCE)
- PCE draws contrastive samples from the prior; it can be effective when prior and posterior are similar, and using a fixed distribution can speed training.PCE is less powerful than ACE, but does not require learning φ in this simplified form.Source: §3.3 · pp. 4–5
3.4 Likelihood-free ACE
- For likelihood-free models, ACE remains a valid EIG lower bound with an unnormalized likelihood approximation fψ, trainable jointly with ξ and φ.The approximation, inference network, and design can be optimized in one problem.Source: §3.4 · p. 5
3.5 Iterated experimental design with ACE
- Iterated ACE replaces the prior with the posterior conditioned on past data and designs; an unnormalized density suffices when inference cannot provide a closed-form posterior density.Source: §3.5 · p. 5
3.6 Gradient estimation for ACE
- For ξ, REINFORCE can be high-variance; reparameterizing y and θ typically lowers Monte Carlo gradient variance, while discrete y permits Rao–Blackwellization.Source: §3.6 · p. 5
- For φ, reparameterizable contrastive samples permit Tucker et al.'s double reparameterization.Source: §3.6 · p. 6
4 EXPERIMENTS
The experiments evaluate gradient-based BOED across five settings, from low-dimensional epidemiology to high-dimensional regression and sequential design. Across these settings, gradient methods outperform black-box optimization in several comparisons, while ACE is consistently competitive.
4 EXPERIMENTS
- The experiments span five scenarios: death-process epidemiology, 400-dimensional regression, advertising, 100-dimensional biomolecular docking, and six-dimensional sequential CES design.The paper includes both standard design optimization and an iterated design-and-inference problem.Source: §4 · p. 6, §4.6 · p. 8
4.1 Evaluating experimental designs
- Design quality is assessed primarily by EIG, with design error used when the optimum is known and posterior entropy and RMSE used for iterated inference.When direct EIG evaluation is unsuitable, ACE lower and VNMC upper bounds are used to compare designs.Source: §4.1 · p. 6
4.2 Death process
- On the death-process problem, all gradient methods achieve higher EIG than BO within the allowed time budget.The experiment uses Rao-Blackwellization over the 66 possible outcomes.Source: §4.2 · p. 6, Figure 2 · p. 6
4.3 Regression
- On the 400-dimensional regression problem, gradient methods achieve about double the final EIG of gradient-free baselines.The study compares BA, ACE, and PCE with VNMC optimized by BO and random search.Source: §4.3 · p. 7, Table 1 · p. 7
4.4 Advertising
- In advertising, ACE, PCE, and BA perform best, while BO struggles at D ≥6 despite access to an EIG oracle.PCE degrades as dimension increases; ACE and BA use adaptive proposal distributions, and perform similarly when their variational families include the true posterior.Source: §4.4 · p. 7, Figure 3 · p. 7
4.5 Biomolecular docking
- For biomolecular docking, all gradient methods outperform the expert design in EIG, with ACE appearing strongest among them.The expert design was selected by Lyu et al. (2019); the methods’ final EIGs are reported as bounds.Source: §4.5 · pp. 7–8, Table 2 · p. 8
4.6 Constant elasticity of substitution
- In sequential CES design, ACE and PCE reduce posterior entropy and latent-variable RMSE faster and further than the baseline, whereas BA performs less well.The authors suggest ACE and PCE perform similarly because posterior changes are smaller at middle and late steps, making the preceding posterior an effective proposal for estimating p(y_t|ξ_t).Source: §4.6 · p. 8, Figure 4 · p. 8
5 CONCLUSIONS
The paper replaces separate EIG estimation and design search with stochastic-gradient optimization of a variational lower bound. Its experiments show improved performance especially in high-dimensional design problems, and the authors recommend ACE as a default bound.
5 CONCLUSIONS
- The method uses stochastic gradients to maximize an EIG lower bound and find designs through a single optimization problem.This replaces the separate stages of estimating EIG and optimizing over designs.Source: §5 · p. 8
- The unified approach yields substantially improved performance, especially on high-dimensional design problems.Source: §5 · p. 8
- ACE generally performs at least as well as the better of BA and PCE across the five experiments, so the authors recommend it as the default.BA performs well when the inference network closely approximates the true posterior; PCE performs well when the prior is an adequate proposal and does not require variational-parameter training.Source: §5 · p. 8
A GRADIENT-BASED BOED
The section establishes properties of the ACE variational lower bound and describes gradient estimators for jointly optimizing the design and variational parameters. It also presents an alternative gradient that remains the gradient of a valid EIG lower bound.
A GRADIENT-BASED BOED
- ACE is a lower bound on EIG whose error can be characterized as an expected KL divergence.The bound converges to the true EIG as the number of contrastive samples L grows, increases monotonically with L, and is exact when the inference network equals the true posterior.Source: A GRADIENT-BASED BOED · p. 11
- A theorem gives a condition under which maximizing the ACE objective converges to maximizing EIG as L grows.Source: A GRADIENT-BASED BOED · p. 12
- The section states that ACE's maximum approaches the maximum EIG under a stated condition as L increases.Source: A GRADIENT-BASED BOED · p. 12
A.1 Double reparametrization
- When qφ(θ|y) is reparameterizable, double reparameterization yields a low-variance gradient estimator for ACE.The method adapts the double reparameterization of Tucker et al. (2018) to account for the w0 term.Source: §A.1 · p. 13
A.2 Alternative gradient
- Removing the ∂g/∂ξ term gives a gradient of a valid EIG lower bound, with performance not substantially different from the standard gradient in the experiments.The experiments used the standard gradient for simplicity.Source: §A.2 · p. 13
B.1 Implementation
The experiments were implemented in PyTorch and Pyro, and the paper points readers to supporting code and run instructions.
B.1 Implementation
- Experiments used PyTorch 1.4.0 and Pyro 0.3.4.Source: §B.1 · p. 13
- Supporting code is available in the pyro repository at the sgboed-reproduce branch.Source: §B.1 · p. 13
- The repository README.md provides instructions for running the experiments.Source: §B.1 · p. 13
B.2 Death process
The death-process section reports final EIG values for each method and describes the computational budget and sampling settings used for evaluation.
B.2 Death process
- Table 3 reports final EIG for each method, evaluated with NMC using 200,000 samples.Source: Table 3 · p. 14
- Each method received a 120-second budget, and the experiments used 100 independent runs.Source: §B.2 · p. 14
- ACE, PCE, BA, and NMC used 10 + 660, 10, 10, and 2,000 samples, respectively.Source: §B.2 · p. 14
B.3 Regression
The regression section notes that non-Gaussian priors prevent analytic evaluation of the posterior or true EIG, and describes the evaluation setup for design methods.
B.3 Regression
- With non-Gaussian priors, the posterior and true EIG cannot be computed analytically, so the section imposes a constraint to ensure a finite EIG maximum.Source: §B.3 · p. 14
- Each of five methods had a 15-minute budget and 10 independent runs; gradient methods used Adam with learning rate 10^-3.Source: §B.3 · p. 14
- For BO and random search, pointwise EIG evaluations used VNMC, while random designs were sampled from a standard unit Gaussian.BO used a Matern52 Gaussian process with GP-UCB1 and stopped after 15 minutes.Source: §B.3 · p. 14
- Final ACE/VNMC evaluations used a fixed inference network with L = 2.5 × 10^3 inner samples and N = 10^5 outer samples.Source: §B.3 · p. 15
B.4 Advertising
The advertising example studies information-based budget allocation across correlated regions and derives a tractable route to evaluating and optimizing the design.
B.4 Advertising
- Correlations among the D regions accompany an information-optimal advertising budget that allocates more money to tightly correlated regions.Source: §B.4 · p. 15
- The setup assumes even D, sets σ = 1, and scales the budget as B = D/2.Source: §B.4 · p. 15
- Symmetry makes paired design coordinates equal, leaving ξ1 and ξD to specify the design under the constraint ξ1 + ξD = 1.Source: §B.4 · p. 15
- The exact EIG reduces to a univariate optimization solvable to machine precision, providing ground-truth EIG and an optimal design for evaluation.Source: §B.4 · p. 15
- After normalization in Fig. 3, the uniform-budget design has absolute error equal to 1.Source: §B.4 · p. 15
B.5 Biomolecular docking
The paper reports experimental setups for biomolecular docking and constant-elasticity-of-substitution designs, then identifies discrete design optimization as future work.
B.5 Biomolecular docking
- The docking design uses 100 binary responses, and the expert design's 580 compounds were subsampled to 100 for fair comparison.Source: §B.5 · pp. 15–16
- Each docking method received 5 × 10^5 gradient steps, and all methods converged within that budget.Source: §B.5 · p. 15
B.6 Constant elasticity of substitution
- Insufficient time to effectively train a neural guide led to a prior-matched mean-field family and a linear model using specified response features.Source: §B.6 · p. 16
C FUTURE WORK
- For discrete designs, the lower bounds remain valid, but gradient optimization over ξ is inappropriate; continuous relaxations or MCMC-style updates are suggested.Source: C FUTURE WORK · p. 16