Source-linked AI summary
Variational Bayesian Optimal Experimental Design
Adam Foster, Martin Jankowiak, Eli Bingham, Paul Horsfall, Yee Whye Teh, Tom Rainforth, Noah Goodman
TL;DR
BOED is limited by the difficulty of obtaining fast, accurate EIG estimates. The paper introduces amortized variational estimators, showing theoretical convergence and empirical gains that improve end-to-end performance. The approach also includes an implementation in Pyro.
Problem
BOED is hampered by the difficulty of obtaining fast and high-quality EIG estimates because EIG computation is intractable and nested.
Method
The paper introduces four variational EIG estimators based on amortized distributional approximations, including approaches for implicit likelihoods and imperfect variational families.
Results
The estimators have convergence rates aligned with conventional estimation, provide significant empirical EIG gains over previous methods, and improve end-to-end performance.
Takeaways & Limitations
The estimators are applicable across experimental design problems and are suited to modern deep probabilistic programming through a Pyro implementation.
Takeaways & Limitations
The convergence analysis assumes strong conditions, and stochastic optimization may converge to a local rather than global optimum, causing additional asymptotic bias.
Abstract
from arXiv · showhide
Bayesian optimal experimental design (BOED) is a principled framework for making efficient use of limited experimental resources. Unfortunately, its applicability is hampered by the difficulty of obtaining accurate estimates of the expected information gain (EIG) of an experiment. To address this, we introduce several classes of fast EIG estimators by building on ideas from amortized variational inference. We show theoretically and empirically that these estimators can provide significant gains in speed and accuracy over previous approaches. We further demonstrate the practicality of our approach on a number of end-to-end experiments.
1 Introduction
BOED selects experiments that maximize expected information about parameters, but accurate EIG estimation is computationally difficult. The paper introduces variational estimators that address this difficulty and reports theoretical and empirical gains.
- BOED frames study design, active learning, and feedback surveys as experiments chosen to maximize gathered information.It has been applied in psychology, Bayesian optimization, active learning, bioinformatics, and neuroscience.
- EIG measures the expected reduction in posterior entropy about target parameters after observing an experimental outcome.The design is optimal from an information-theoretic viewpoint when the predictive model is correct.
- Sequential BOED reuses previous experimental results to choose later designs, potentially requiring fewer iterations to uncover the same information.
- Obtaining fast, high-quality EIG estimates is difficult because posterior intractability makes EIG a nested expectation problem.Conventional nested Monte Carlo can achieve at best O(T −1/3), slower than O(T −1/2) for conventional Monte Carlo.
- The paper proposes four variational EIG estimators that sidestep double intractability and support distinct settings, including implicit likelihoods and misspecified variational families.The authors provide a general-purpose Pyro implementation and report improved EIG estimation and end-to-end performance.
2 Background
The BOED background formalizes design selection through expected information gain, interpreted as mutual information between parameters and outcomes. Closed-form computation is generally unavailable, making nested Monte Carlo the standard but slow approach.
- BOED chooses a design d using a prior p(θ) and predictive model p(y|θ, d) to maximize information about parameters θ.
- EIG averages information gain over possible outcomes and equals the mutual information between θ and y given d.
- The Bayesian optimal design is d∗≜arg maxd∈D EIG(d), where D denotes the permissible designs.
- EIG computation is generally intractable because neither p(θ|y, d) nor p(y|d) has a closed form, preventing conventional Monte Carlo estimation.
- Nested Monte Carlo is consistent but has an overall O(T −1/3) convergence rate, making EIG estimation a bottleneck for real-time sequential BOED.
- In sequential BOED, each new design uses the posterior conditioned on earlier designs and outcomes as the updated prior.
3 Variational Estimators
The paper develops four variational estimators that learn amortized approximations to posterior, marginal, or likelihood distributions, reducing EIG estimation costs. The estimators differ in bias, likelihood requirements, and asymptotic behavior.
- Variational estimators learn functional approximations rather than separately estimating the EIG integrand for every outcome.This amortizes computation across outcomes and is intended to alleviate the EIG bottleneck.
- Variational posterior: The variational posterior estimator learns qp(θ|y, d) and uses it in a variational lower bound that is tight when it equals p(θ|y, d).
- Variational marginal: The variational marginal estimator learns qm(y|d), which produces an upper bound tight when qm(y|d) equals p(y|d).
- Variational NMC: VNMC combines fast variational proposal learning with Monte Carlo refinement, trading early biased estimates for asymptotic bias removal.Its bound is tight as L →∞ even for imperfect qv, yielding asymptotically unbiased estimates outside the variational family.
- Implicit likelihood: The implicit-likelihood estimator supports models that can be sampled but whose likelihoods cannot be evaluated, although it is not itself an EIG bound.Its estimation error is bounded, with equality only when both learned densities match their targets almost everywhere.
- Selecting an estimator: Estimator choice depends on output-versus-parameter dimension, explicit likelihood availability, and whether variational families contain target densities.VNMC is the only method guaranteed to converge to the true EIG when the variational families miss their targets, given increasing computational budget.
4 Convergence rates
The variational estimators converge at O(T^-1/2) to variationally determined limits under stated assumptions, while VNMC can remove asymptotic bias through a slower refinement stage.
- The convergence analysis decomposes error into Monte Carlo variance, optimization error, and the variational-family gap to EIG.The first term decreases with more samples, the second with more stochastic-gradient steps, and the third requires a richer variational family or larger L.
- O(T^-1/2) convergence holds for the estimators’ biased variational limits when N is proportional to K and the assumptions are satisfied.The result combines O(N^-1/2) Monte Carlo error with O(K^-1/2) optimization error.
- Local rather than global optimization can introduce an additional asymptotic bias in practice.The theoretical rate relies on strong stochastic-optimization assumptions and convergence to the global optimum.
- VNMC removes asymptotic bias by increasing M with N after variational training, converging to EIG at O((NM)^-1/3).Its total cost is T = O(KL + NM), with fast amortized proposal training followed by slower NMC refinement.
- Standard NMC is a special case of VNMC with the prior as proposal, so it skips amortized proposal learning and typically requires higher cost for equal accuracy.This contrasts the fast training stage of VNMC with standard NMC’s direct use of the prior proposal.
5 Related work
The paper compares its EIG estimators with established alternatives spanning nested Monte Carlo, Laplace approximation, likelihood-free ratio estimation, and the Donsker–Varadhan representation.
- Nested Monte Carlo is an established baseline for estimating EIG in the paper’s empirical comparisons.
- Laplace approximation is fast but limited to continuous variables and can exhibit large bias.
- The comparisons also include LFIRE-based implicit likelihood estimation and a Donsker–Varadhan bound used for mutual-information estimation.The Donsker–Varadhan approach had not previously been considered in BOED according to the passage.
6 Experiments
Experiments show that variational EIG estimators improve accuracy and convergence under fixed computational budgets, while supporting adaptive sequential design. Their practical benefits extend to online psychology and revealed-preference experiments.
- 6.1 EIG estimation accuracy: All variational estimators outperformed NMC across the benchmark scenarios, except that Laplace was best for the Gaussian linear model where its approximation is exact.The methods were evaluated over design grids using fixed computational budgets and analytically or computationally obtained true EIG values.
- 6.2 Convergence rates: ˆµpost substantially outperformed ˆµmarg as the number of MC samples increased after fixed variational training.The RMSE initially decreased with N before plateauing because of estimator bias.
- 6.2 Convergence rates: With N = K increasing together, the estimators followed the predicted convergence rate O(T^-1/2).When T = N + K was fixed, the lowest RMSE appeared for K/T between 0.5 and 0.9.
- 6.2 Convergence rates: ˆµVNMC improved over NMC after variational pre-training and continued improving as M and N increased, at the slower rate O(T^-1/3).The estimator avoided the plateau observed for ˆµpost and ˆµmarg, although increasing N and M remained important.
- 6.3 End-to-end sequential experiments: In the Mechanical Turk experiment, online BOED with ˆµm+ℓ produced a more certain posterior than random design.The system presented stimuli, collected responses, updated the model, and selected the next stimulus online.
- 6.3 End-to-end sequential experiments: The sequential CES experiment compared random design with ˆµmarg and ˆµNMC for learning latent variables ρ, α, and u.The design goal was to select experiments whose agent responses were informative about θ = (ρ, α, u).
7 Discussion
The discussion emphasizes that the estimators combine fast variational learning with the consistency properties of Monte Carlo methods. The appendix results establish monotonicity, limiting behavior, and conditions under which the estimators are exact.
- 7 Discussion: The estimators are broadly applicable and provide improved convergence rates and practical gains across experimental design problems.The implementation is suited to modern deep probabilistic programming languages and is available in Pyro.
- 7 Discussion: The variational and Monte Carlo components offer complementary benefits because NMC is inherently biased while variational methods learn rapidly.ˆµVNMC is presented as a way to trade off variational learning speed against NMC asymptotic consistency.
- 7 Discussion: The variational posterior approximation must have the same support as the target posterior for the stated results.This support requirement is an explicit validity condition on qp(θ|y, d).
- 7 Discussion: ˆµVNMC decreases monotonically with increasing L and converges to EIG(d) as L approaches infinity.It equals EIG(d) for every L ≥ 1 when qv(θ|y, d) equals the true posterior for all y and θ.
A.4 Variational marginal + likelihood ˆµm+ℓ
The variational marginal-plus-likelihood estimator admits an error bound determined by divergences between learned and true predictive distributions. Under a specific coupling of the approximations, it becomes a lower bound on EIG.
- A.4 Variational marginal + likelihood ˆµm+ℓ: The absolute estimation error is bounded by the expected conditional KL divergence plus the marginal KL divergence.The bound decomposes error into mismatch in qℓ(y|θ, d) and mismatch in qm(y|d).
- A.4 Variational marginal + likelihood ˆµm+ℓ: The error bound is tight exactly when qℓ(y|θ, d) and qm(y|d) equal their corresponding true distributions almost everywhere.The bound vanishes when both KL divergences are zero.
- A.4 Variational marginal + likelihood ˆµm+ℓ: When qm(y|d) = Ep(θ)[qℓ(y|θ, d)], Im+ℓ(d) is a lower bound on EIG(d).This condition can hold, for example, when θ takes finitely many values.
- A.4 Variational marginal + likelihood ˆµm+ℓ: When qm and qℓ are learned separately, Im+ℓ(d) is generally neither an upper nor a lower bound on EIG(d).The lower-bound property therefore depends on the stated relationship between the two approximations.
B Details for convergence rates
The convergence analysis establishes O(T^-1/2) error rates under strong regularity conditions, while noting that these assumptions may fail and optimization can settle at a local optimum.
- Assumptions: The theorem assumes convexity, differentiability, Lipschitz gradients and Hessians, strong convexity, moment bounds, and Lipschitz continuity of the objective.These conditions support the stochastic-gradient convergence argument and control the function-value variance.
- Limitations: The assumptions are strong, particularly strong convexity, and may not hold in practice.The authors caution that these conditions limit how directly the theorem applies to practical optimization problems.
- Limitations: When global optimality fails, stochastic gradient descent may converge to a local optimum φ† rather than the global optimum φ∗.The authors state that convergence behavior and rates may remain similar, but the error breakdown changes.
- Convergence results: O(N^-1/2 + K^-1/2) = O(T^-1/2) when N ∝ K combines Monte Carlo estimation and stochastic optimization errors.The result follows from uniform convergence of the sample estimate and O(K^-1/2) convergence of averaged stochastic gradient descent.
- Proof strategy: The proof combines a weak law of large numbers for uniform Monte Carlo convergence with standard stochastic-gradient convergence results.The analysis separately controls estimation error and optimization error before combining them.
- Empirical evaluation: The paper benchmarks variational estimators across experiment-design scenarios, including A/B testing, and compares them with Laplace, LFIRE, DV, and NMC methods.The variational estimators generally outperform the baselines, while Laplace is best for the Gaussian linear model where its approximation is exact.
D.2 End-to-end sequential experiments
The paper applies its estimators to adaptive experiments involving human preferences and stylized-face judgments, using sequential posterior updates to select informative designs. In simulated data, BOED reduces posterior uncertainty more quickly.
- Mechanical Turk experiment: The adaptive face experiment represents images with two three-level feature dimensions and encodes each comparison using a left-right symmetric design matrix.Each image is represented by a six-dimensional one-hot feature vector, and the query compares a left image with a right image.
- Mechanical Turk experiment: EIG is estimated across all 36 candidate designs within a 30-second turnaround for posterior learning, design selection, and presentation.The procedure selects the best design at each step for each independent run.
- Posterior inference: The adaptive model requires a variational family covering fixed effects, random effects, and hyperparameters in the full posterior.A partial mean-field approximation is used, with conventional variational inference for its parameters.
- Simulated experiment: In simulated-data trials, BOED decreases posterior uncertainty more quickly.The entropy curves summarize the mean and ±1 standard error across 10 experimental trials.
- Preference experiment: The preference experiment presents pairs of three-commodity baskets and records preference, weak preference, or indifference on a slider.The response model is based on the difference in utility between the two baskets.
- Preference experiment: Optimal preference queries lie between predictable responses for very similar or very different baskets and noisy comparisons of dissimilar baskets.The informative region is described as a sweet spot in basket differences.
E.1 Death process
The death-process experiment evaluates how variational-family choice affects asymptotic EIG bias for a two-time infection measurement design. The log-normal posterior estimator is nearly exact, while Laplace is less accurate than the truncated-normal posterior estimator.
- Design selection: The exact method identifies the optimal design, while approximate methods generally select designs with large EIG; the log-normal posterior is nearly exact.The optimal design from the exact method is marked with a star, and each method’s selected design with a cross.
- Experiment setup: The experiment designs two infection-count measurements at times (t1, t2) in a fixed-size population with a log-normal prior on infection rate b.The design space satisfies 0 ≤ t1 ≤ t2.
- Estimator comparison: Four estimators are compared: exact quadrature, log-normal posterior, truncated-normal posterior, and Laplace.The comparison isolates the effect of variational-family choice and approximation method.
- Results: Mean absolute errors of about 10^-3 occur for the log-normal posterior estimator, whose variational family best matches the true posterior.The EIG surfaces are scaled so that the exact-method maximum equals 1.0.
- Results: The Laplace estimator has absolute errors about 30% higher than the truncated-normal posterior estimator despite making the same distributional assumption.The figure-level analysis also notes that Laplace produces an upper bound in this case.
F Consistent EIG estimation with control variates
The control-variate estimator decomposes EIG using an approximate marginal density, replacing one expectation with analytic KL terms and retaining nested Monte Carlo for the other. It can reduce variance when the approximation is good and remains consistent as sample sizes grow.
- Control-variate construction: EIG is decomposed into an expected KL divergence to an approximate marginal minus the KL divergence between the true and approximate marginals.This decomposition enables analytic treatment of the first term when the relevant KL divergence is computable.
- Estimator construction: The first expectation can be estimated with conventional Monte Carlo when its KL divergence to the approximate marginal is analytic.The second term is estimated with nested Monte Carlo.
- Variance reduction: Variance is expected to be low when the approximate marginal is close to the true marginal, because the relevant covariance is large.The approximation functions as a control variate for the Monte Carlo estimator.
- Consistency: The control-variate estimator is consistent and converges to EIG as N, M → ∞.Consistency is stated for the estimator when both Monte Carlo sample sizes grow.
G KL ( q || p ) versus KL ( p || q )
The paper compares forward-KL posterior approximation with reverse-KL minimization for EIG estimation. Although reverse KL is common in variational inference, it can produce discontinuities, mode-locking, and design-dependent EIG bias.
- The posterior estimator minimizes an expected forward KL divergence between the true posterior and its variational approximation.
- The forward-KL objective yields an EIG lower bound whose gap equals the expected KL divergence between the true and approximate posteriors.
- Reverse-KL minimization instead compares the variational posterior against the true posterior and is considered as an alternative approximation strategy.
- Reverse KL can exhibit discontinuous behavior in the partial KL, making EIG estimation problematic.
- Mode-locking under reverse KL can drop modes, causing significant EIG misestimates and potentially large bias that changes across neighboring designs.