PaperlayerView exact revision on arXiv

Source-linked AI summary

Deep Adaptive Design: Amortizing Sequential Bayesian Experimental Design

Adam Foster, Desi R. Ivanova, Ilyas Malik, Tom Rainforth

arXiv:2103.02438v2stat.MLcs.AIcs.LGstat.CO

TL;DR

Designing experiments to be informative is a challenge across fields. DAD learns a policy that selects designs from experiment history, then deploys it rapidly; across experiments, it outperformed baselines with comparable deployment time and was competitive with costly conventional BOED.

  • Problem

    Designing experiments so their outcomes are informative about the underlying process is a challenge across disciplines.

    Source: §1 · p. 1

  • Method

    DAD learns a design network from simulated trajectories, mapping past designs and observations to the next design, and uses contrastive bounds to train its policy.

    Source: §1 · p. 2

  • Results

    Across experiments, DAD performed significantly better than baselines with comparable deployment time and competitively with non-amortized BOED approaches that make costly computations at each stage.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The authors suggest DAD may benefit practitioners in fields including online surveys and clinical trials.

    Source: §7 · p. 9

    DAD requires an explicit likelihood model whose density p(yt|θ, ξt) can be evaluated.

    Source: §7 · p. 9

Abstract

from arXiv · show

We introduce Deep Adaptive Design (DAD), a method for amortizing the cost of adaptive Bayesian experimental design that allows experiments to be run in real-time. Traditional sequential Bayesian optimal experimental design approaches require substantial computation at each stage of the experiment. This makes them unsuitable for most real-world applications, where decisions must typically be made quickly. DAD addresses this restriction by learning an amortized design network upfront and then using this to rapidly run (multiple) adaptive experiments at deployment time. This network represents a design policy which takes as input the data from previous steps, and outputs the next design using a single forward pass; these design decisions can be made in milliseconds during the live experiment. To train the network, we introduce contrastive information bounds that are suitable objectives for the sequential setting, and propose a customized network architecture that exploits key symmetries. We demonstrate that DAD successfully amortizes the process of experimental design, outperforming alternative strategies on a number of problems.

1. Introduction

Sequential BOED can require costly computation between experiment steps, making quick decisions difficult. DAD learns a design policy upfront to choose later designs from past data, enabling rapid adaptive experiments.

1. Introduction

  • Conventional sequential BOED requires substantial computation between steps, while upfront amortization is useful when the same adaptive framework is deployed repeatedly.The introduction gives multiple participants in a survey as an example of repeated deployment.Source: §1 · pp. 1–2
  • DAD learns one network that maps previous designs and observations to the next design, and can be reused for different experiment instances.The network is trained using simulated experimental trajectories.Source: §1 · p. 2
  • A holistic objective and sequential contrastive bounds let DAD train policy parameters without direct posterior or marginal-likelihood estimation.The policy can learn non-myopic strategies that account for future decisions.Source: §1 · p. 2
  • A permutation symmetry of the optimal policy motivates a customized network architecture designed to support amortization across time steps.Source: §1 · p. 2
  • DAD accurately amortizes experiments across application problems, opening the way to real-time adaptive BOED.The introduction names epidemiology, physics, and psychology as relevant application areas.Source: §1 · p. 2

2. Background

BOED selects experiments to maximize expected information about unknown variables. In sequential settings, adaptive designs use prior outcomes, but conventional optimization is costly; contrastive bounds offer a stochastic-gradient approach to EIG optimization.

2. Background

  • BOED optimizes designs for expected information about latent variables; static designs fix all experiments in advance, while adaptive designs use past outcomes to choose later experiments.Source: §2 · p. 2, §2.1 · p. 2

2.1. Conventional adaptive BOED

  • Conventional adaptive BOED fits a posterior after each stage and optimizes the next EIG, requiring costly computations during the experiment.The posterior depends on observed outcomes, and EIG optimization involves a doubly intractable quantity.Source: §2.1 · p. 3

2.2. Contrastive information bounds

  • For continuous designs, contrastive lower bounds such as Prior Contrastive Estimation support joint EIG estimation and optimization by stochastic gradient ascent.The PCE bound becomes tight as the number of contrastive samples L approaches infinity.Source: §2.2 · p. 3

3. Rethinking Sequential BOED

Sequential BOED can be reformulated around a policy that chooses designs from experiment history and optimizes information over the whole sequence. This supports non-myopic policies and direct, low-computation deployment.

3. Rethinking Sequential BOED

  • The design policy maps previous design–observation pairs to the next design, and simulated histories are generated by applying it and sampling outcomes.The process samples a parameter from its prior before generating the sequence.Source: §3 · p. 3
  • The conventional implicit policy fits a posterior and optimizes each stage separately, making myopic choices that do not account for future decisions.Learning a policy to mimic that process would require expensive simulated sequential BOED examples.Source: §3 · p. 3
  • A policy-based objective expresses total expected information gain across the experiment sequence, avoiding posterior and intermediate EIG calculations while allowing non-myopic policies.Source: §3 · p. 3
  • The policy objective strictly generalizes conventional BOED: fixed designs recover static design, while conventional adaptive BOED approximates a policy.Source: §3 · p. 4
  • Once learned, the policy can be evaluated directly during the experiment with negligible deployment computation.Source: §3 · p. 4

4. Deep Adaptive Design

DAD represents sequential experimental design as a policy over past histories and learns that policy upfront, so deployment can select designs with a single fast network pass.

4. Deep Adaptive Design

  • DAD approximates the policy maximizing information in the full experiment history, instead of optimizing each design during the live experiment.After training, the network selects each next design nearly instantaneously from the history.Source: §4 · p. 4

4.1. Contrastive bounds for sequential experiments

  • Because the rollout objective remains intractable, DAD trains with sequential contrastive lower bounds; sPCE is chosen for numerical stability over potentially unbounded ratio estimates.The complementary bounds tighten as the number of contrastive samples increases and become exact as L → ∞ at rate O(1/L).Source: §4.1 · pp. 4–5

4.2. Gradient estimation

  • For reparametrizable continuous observations, gradients can pass through simulated histories; DAD also uses history enumeration or REINFORCE for other observation settings.Enumeration costs O(|Y|^T), limiting its use to relatively small numbers of experiments and possible outcomes.Source: §4.2 · p. 5
  • When histories are impractical to enumerate or continuous observations are non-reparametrizable, DAD uses the score function gradient estimator, also known as REINFORCE.The continuous-observation reparameterization approach assumes observations can be expressed as deterministic functions of the latent parameter, design, and independent noise.Source: §4.2 · p. 5

4.3. Architecture

  • DAD pools encoded design–outcome pairs into an order-invariant history representation, then uses an emitter to choose the next design, sharing encoder weights across pairs and time steps.This architecture improves performance compared with networks that must learn the relevant symmetries themselves.Source: §4 · p. 4, §4.3 · p. 6

5. Related Work

Prior sequential BOED work covers posterior inference, mutual-information estimation, and design optimization, alongside methods that combine estimation and optimization or target particular settings.

5. Related Work

  • Posterior updates in sequential BOED use sequential Monte Carlo (Del Moral et al., 2006; Drovandi et al., 2014), population Monte Carlo, variational inference (Foster et al., 2019; 2020), or Laplace approximation.Source: §5 · p. 6
  • Prior methods estimate mutual information with nested Monte Carlo, variational bounds, Laplace approximation, ratio estimation, or hybrid methods, and optimize designs with Bayesian optimization, particle systems, annealing, regret bounds, or bandits.Examples include Myung et al. (2013), Foster et al. (2019; 2020), Kleinegesse et al. (2020), and Rainforth (2017).Source: §5 · p. 6
  • Other approaches jointly estimate and optimize mutual information with stochastic gradients, while separate work addresses intractable likelihoods or non-myopic strategies in tractable cases.The passages cite perturbation analysis, variational lower bounds, and multi-level Monte Carlo (Goda et al., 2020) among joint approaches.Source: §5 · p. 6

6. Experiments

The experiments compare DAD with random and fixed strategies, variational design, and problem-specific adaptive methods across three design problems. They assess information gained and deployment time.

6.1. Location finding in 2D

  • DAD significantly outperforms all considered location-finding baselines; 30 decisions take 0.0474 ± 0.0003 seconds versus 8963 seconds for the variational method.Table 1 reports upper and lower bounds on total EIG for the location-finding strategies.Source: Table 1 · p. 7, §6.1 · p. 7

6. Experiments

  • The experiments use random, fixed, variational, and task-specific adaptive baselines; evaluation focuses on total EIG and post-training deployment time.The location-finding task chooses observation locations for multiple hidden sources; the other tasks concern human temporal discounting and an epidemiological death process.Source: §6 · p. 7, §6.1 · p. 7, §6.2 · p. 8

6.1. Location finding in 2D

  • A DAD network trained for 30 location-finding experiments generalizes to other sequence lengths and still outperforms baselines; the fixed strategy cannot extend beyond its training length.Figure 2 presents sPCE estimates with error bars.Source: §6.1 · p. 7, Figure 2 · p. 7
  • Across 16 training runs, the lower-bound estimate was 10.91 ± 0.014 and the upper-bound estimate was 12.47 ± 0.046.Performance varied more across rollouts with one network than across the average performance of different networks.Source: §6.1 · p. 7

6.2. Hyperbolic temporal discounting

  • DAD performs best in hyperbolic temporal discounting, surpassing bespoke methods including Badapted, with total deployment time below 0.1 seconds.Its designs, like Badapted's, cluster near the indifference point.Source: §6.2 · p. 8, Table 3 · p. 8, Table 2 · p. 8

6.3. Death process

  • For the death process, DAD exceeds fixed and variational designs in information gain and deploys in under 0.01 seconds, while SeqBED takes hours.The SeqBED comparison used a single rollout, with θ = 1.5; the reported information gains were 1.590 for SeqBED and 1.779 for DAD.Source: §6.3 · pp. 8–9, Table 4 · p. 9

7. Discussion

DAD often outperforms baselines at comparable deployment time and can outperform non-amortized BOED while using a tiny fraction of its deployment resources. Its learned policy can plan across steps, but the method requires an explicit likelihood, conditionally independent experiments, and continuous designs.

Discussion

  • DAD significantly outperforms baselines at comparable deployment time and can outperform costly non-amortized BOED while using a tiny fraction of its deployment resources.The authors suggest that avoiding intermediate posterior approximation and learning a potentially non-myopic policy help explain this performance.Source: §7 · p. 9
  • In the two-step location example, DAD’s first design sacrifices one-step EIG optimality to focus the second experiment on the positive regime, yielding higher total EIG than the exact myopic approach.The learned first design is approximately ξ1 = -0.4; the myopic first design is ξ1 = 0.Source: §7 · p. 9
  • The present method requires an explicit likelihood and experiments that are conditionally independent given θ, a condition that may fail for time-series models.It also requires continuous designs for gradient-based optimization.Source: §7 · p. 9
  • DAD learns a design-policy network offline and deploys it to make adaptive decisions quickly, avoiding intermediate posterior estimates and live design optimization.The authors describe it as the first approach to allow real-time adaptive BOED for general problems.Source: §7 · p. 9

A. Proofs

The proofs establish bounds on total expected information gain, their convergence properties, and invariance of optimal policies to the ordering of past design–observation pairs.

A. Proofs

  • The total objective for a policy is the expected information gain over the full sequence of T experiments.The proof describes it as the expected entropy reduction from the prior to the final posterior, without considering intermediate posteriors.Source: A. Proofs · p. 13
  • The sequential PCE lower bound approaches total expected information gain at rate O(L^-1), under a technical assumption used to establish convergence.The proof also establishes that the lower bound is monotone in the number of contrastive samples L.Source: A. Proofs · pp. 15–16
  • The optimal policy and total expected information gain are invariant to permutations of the observed history’s design–observation pairs.The proof uses the order-invariance of the posterior and induction over the remaining experiment steps.Source: A. Proofs · pp. 16–17
  • The proof shows that UT(pi,L)−IT(pi) tends to zero at rate O(L−1), with higher-order terms o(L−1), and states that monotonicity is established.Source: A. Proofs · pp. 18–19

B. Additional bounds

The appendix presents additional sequential bounds and details the gradient calculations used to train the design policy.

B. Additional bounds

  • The sequential ACE bound lower-bounds total expected information gain using a parametrized proposal q(θ; hT), while sequential VNMC supplies an upper bound.The proposal may be an amortized variational approximation to the posterior, though the bound is not limited to variational posteriors.Source: B. Additional bounds · pp. 19–20

C.1. Score function gradient

  • The appendix differentiates the sPCE objective to provide its score-function gradient.The section states that the objective is differentiated to obtain the gradient.Source: §C.1 · p. 21

C.2. Expanded reparametrized gradient

  • The expanded reparameterized gradient enumerates computational-graph paths by length, summing chain-rule contributions over increasing time-index sequences.The path length is denoted by k, and the derivatives are computed automatically in PyTorch in practice.Source: §C.2 · pp. 21–22

D. Experiment details

The experiments used PyTorch and Pyro, and the authors provide open-source code for reproducing them.

D. Experiment details

  • Experiments used PyTorch and Pyro; open-source DAD code for reproducing each experiment is available.The passage directs readers to the project’s GitHub repository and README for running instructions.Source: D. Experiment details · p. 22

D.1. Location Finding

The appendix describes location-finding and hyperbolic temporal-discounting experiments, including their model assumptions and design procedures.

D.1. Location Finding

  • The location-finding model assumes a known number of sources in d = 1, 2, or 3, with each source emitting an inverse-square signal.The observed total intensity is the superposition of individual source intensities.Source: §D.1 · p. 22
  • Source locations have independent standard normal priors, and the log total intensity is observed with Gaussian noise.The signal model also includes background and maximum-signal constants.Source: §D.1 · p. 23
  • Even in one dimension, the optimal myopic baseline is costly to evaluate, and peaked, multimodal posteriors may require quadrature.The simplified baseline uses numerical integration over a design grid and a parameter grid, with sampling for the inner expectation.Source: §D.1 · pp. 24–25

D.2. Hyperbolic temporal discounting

  • The hyperbolic temporal-discounting experiment models participant latents k and α across 20 iterations, with design choices R and D constrained by 0 < R < 100 and D > 0.The design represents R and D in an unconstrained space before transforming them.Source: §D.2 · p. 25
  • For the temporal-discounting network, the authors used a score-function gradient estimator and Adam, with 500 inner and 500 outer samples.The estimator was used because the experiment had T = 20 iterations.Source: §D.2 · p. 26

D.2.1. ABLATION: TOTAL ENUMERATION

For the T = 10 hyperbolic temporal-discounting experiment, total enumeration and score-function gradient estimation produced statistically similar results and learned designs.

D.2.1. ABLATION: TOTAL ENUMERATION

  • The comparison trained networks for T = 10 experiments, yielding 2^10 = 1024 possible histories.The two gradient estimators compared were total enumeration and the score-function estimator.Source: §D.2.1 · p. 26
  • Figure 7 shows qualitatively similar learned designs for total enumeration and the score-function estimator.The comparison uses T = 10 experiments.Source: §D.2.1 · p. 27
  • The two gradient methods performed statistically equally: Table 5’s mutual-information bound estimates were within 2 standard errors of each other.The table reports final lower and upper bounds for I10(π).Source: Table 5 · p. 27, §D.2.1 · p. 27

D.3. Death process

The Death Process experiment evaluates sequential designs under a truncated-normal prior, including a rollout comparison with SeqBED and a caveat about its timing estimate.

D.3. Death process

  • The infection rate has a truncated-normal prior with mean 1, standard deviation 1, and support from 0 upward.The sequential experiment uses four steps.Source: §D.3 · p. 27
  • Because the outcome space has 51 possible values, the authors used a score-function gradient estimator rather than complete enumeration for T = 4.The four observations are from independent stochastic processes, with each design constrained to be positive.Source: §D.3 · pp. 27–28
  • SeqBED’s reported deployment time is a substantial underestimate because its speed test used a more powerful machine.The test used 40 processors and 189 GiB of memory; SeqBED can also handle a broader class of implicit likelihood models.Source: §D.3 · p. 28
  • In one rollout with θ = 1.5, the methods’ posterior distributions were visually similar and clustered near the true value.Information gains for this comparison were computed by numerical integration because SeqBED was too slow for the usual evaluation.Source: §D.3 · p. 28

Related papers

Loading 2103.02438v2…