Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
Authors: Marcel Hedman
Desi R. Ivanova
Cong Guan
Tom Rainforth
Venue: ICML 2025 poster (2025)
Keywords: Bayesian experimental design
Bayesian optimal design
Bayesian adaptive design
adaptive design optimization
information maximization
Primary area: probabilistic_methods->bayesian_models_and_methods
TL;DR: We develop a semi-amortized, policy-based, approach to Bayesian experimental design (BED) called Stepwise Deep Adaptive Design (Step-DAD), demonstrating superior decision-making and robustness compared with current state-of-the-art BED methods.

Abstract: We develop a semi-amortized, policy-based, approach to Bayesian experimental design (BED) called Stepwise Deep Adaptive Design (Step-DAD). Like existing, fully amortized, policy-based BED approaches, Step-DAD trains a design policy upfront before the experiment. However, rather than keeping this policy fixed, Step-DAD periodically updates it as data is gathered, refining it to the particular experimental instance. This test-time adaptation improves both the flexibility and the robustness of the design strategy compared with existing approaches. Empirically, Step-DAD consistently demonstrates superior decision-making and robustness compared with current state-of-the-art BED methods.

Step-DAD: Semi-Amortizee
Policy-Based Bayesian Experimental Design
Marcel Hedman"
Desi R. Ivanova"
Cong Guan
Tom Rainforth
Abstract
signs. The goal is then choose designs that are maximally informative about 0. Namely, we maximize the Expected Information Gain (EIG) (Lindley] 1956 1972). which mea- sures the expected reduction in our uncertainty about 0 from running an experiment with a given design.
We develop a semi-amortized, policy-based, ap- proach to Bayesian experimental design (BED) called Stepwise Deep Adaptive Design (Step- DAD). Like existing, fully amortized, policy- based BED approaches, Step-DAD trains a design policy upfront before the experiment. However, rather than keeping this policy fixed, Step-DAD periodically updates it as data is gathered, refin- ing it the particular experimental instance. This test-time adaptation improves both the flexibility and the robustness of the design strategy com- pared with existing approaches. Empirically, Step- DAD consistently demonstrates superior decision- making and robustness compared with current state-of-the-arr BED methods.
The traditional adaptive BED approach, illustrated in Fig la involves iterating between making design decisions by op- timizing the EIG of the next experiment step, and updat- ing the underlying model through Bayesian updates that condition on the data obtained so far. Unfortunately, this approach leads to sub-optimal design decisions, as it is a greedy, myopic, strategy that fails to consider future ex- periment steps (Foster [021]] Huan and Marzouk 2016). Furthermore, it requires substantial computation be un- dertaken at each experiment iteration, making it impractical for real time applications (Rainforth al. 2024).
Foster et al. 2 showed that this traditional framework can be significantly improved upon by taking a policy-based approach (PB-BED) As shown in Fig lb] their Deep Adap- tive Design (DAD) framework, and its extensions (Blau ct al. 2022 Ivanova et al., 202T Lim et al., 2022). are based on learning design policy network that maps from experimental histories to new designs. This policy is trained before the experiment, then deployed to make design de- cisions automatically at test time. This provides a fully amortized approach that eliminates the need for significant computation during the experiment itself, thereby enabling real-time, adaptive, and non-myopic design strategies that represent the current state-of-the-ar in adaptive BED.
1. Introduction
Adaptive experimentation plays crucial role in science and engineering: it enables targeted and efficient data ac- quisition by sequentially integrating information gathered from past experiment iterations into subsequent design deci- sions (Atkinson et al., 2007: MacKay| 1992 Myung et al.] 2013). For example, consider an online survey that aims to infer individual preferences through personalized questions. By strategically tailoring future questions based on insights from past responses, the survey can rapidly hone in rele- vant questions for each specific individual, enabling precise preference inference with fewer, more targeted questions.
In principle, these fully amortized approaches can learn the- oretically optimal design strategies (in terms of total EIG), In practice, learning a policy that remains optimal for all possible experiment realizations is rarely realistic. In par- ticular, the dimensionality of experimental history expands as the experiment progresses, making it increasingly diffi- cult account for all possible eventualities through upfront training alone. Moreover, deficiencies in our model can mean that data observed in practice can be highly distinct from the simulated data used to train the policy.
Bayesian experimental design (BED) offers a princi- pled framework for solving such optimal design prob-
lems (Chaloner and Verdinelli (1955] RRainforhh al. 2024
Ryan al., 2016). In BED, the quantity of interest (e.g. individual preferences), is represented as an unknown pa- rameter 0 and modelled probabilistically through a joint generative model on 0 and experiment outcomes given de-
'Equal contribution 'Department of Statistics. Uni- versity of Oxford. Correspondence to: Marcel Hed- man ceaacee. @ stats,tx....... axk, Desi Ivanova <desi.ivanova @ ax....o...c.kk..
To address these limitations, and allow utilisation of any available computation during the experiment, we introduce hybrid, semi-amortized, PB-BED approach, called Step- wise Deep Adaptive Design (Step-DAD). As illustrated in
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Step-DAD: Semi-Amortized Policy-Based Based Bayesian Experimental Design
Step-DAD periodically updates the policy during the experiment. This allows the policy to be adapted using pre- viously gathered data, refining it to maximize performance for the particular realization of the data that we observe. In turn, this allows Step-DAD to make more accurate design decisions and provides significant improvements in robust- ness to observing data that is dissimilar to that generated in the original policy training. Empirical evaluations reveal Step-DAD is able to provide significant improvements in ofte-of-theeart design performance, while using substan- tially less computation than the traditional BED approach.
Fig Ic.
significant computational expense incurred from the itera- tive posterior inference and EIG optimization. In general, the posterior computation is intractable and the EIG ( esti- mation is doubly intractable (Foster et al.) 2019 Rainforth et al. 2018). Since both of these steps must be conducted at each step of the experiment, the traditional adaptive BED approach is often impractical for real-time applications.
2.2. Amortized Policy-Based BED
In response to the limitations of traditional adaptive BED, Foster et al. (2021) introduce the idea of amortizing the adaptive design process through learnt policies. This amor- tized policy-based BED (PB-BED) approach represents a significant advancement over the traditional framework, de- livering state-of-the- non-myopic design optimization whilst enabling real-time deployment.
2. Background
Guided by the principle of information maximization, Bayesian experimental design (BED, [Lidley] 1956) is model-based framework for designing optimal experiments. Given a model p(y)p(y 0, E), describing the relationship between experimental outcomes y, controllable designs & and unknown parameters of interest 0, the goal is to select the experiment& & that maximizes the expected information gain (EIG) about 0. The EIG, which is equivalent to the mu- tual information between 0 and y. is the expected reduction in Shannon entropy from the prior to the posterior of #:
BE-BED reformulates the design problem using a policy т, which maps experimental histories to subsequent design choices, π : h + &t- The optimal policy is now the one that maximizes the total EIG across the entire sequence of
T experiments (Foster et al. 2021 Shen and Huan| 2021)
I (TT) = Ep((rrn)) (m) H (0)
H Hp(( hr]]]
= =Ep(0)((r 10,z) pl 0, X) - log p(rr [x)] (3)
T(C. = Fl H7 [0)]
H (p(( I [.,)
where p(hr 10, = - P(u. | 0, h1-1), p(hr (T) (p(T) 0, T).. and & = π (ht-1) 1) are all evaluated au- toregressively. This policy- based formulation strictly gen- eralizes the traditional adaptive BED approach, which can be viewed as using a specific policy that maximizes the incremental ane-step-ahead EIG (I at each iteration, that is Tuad Tiaa (h+-1) = arg maxç, plu (Ee)
where p(y {) p0) [p(y 0, [) is the prior predictive distribution of our model.
2.1. Traditional Adaptive BED
BED becomes particularly powerful in adaptive contexts, where we allow the future design decision at time t, Et. to be informed by the data acquired up to that point, h4-1 == (E1. 3/1) .n.....,. ((- ((--), which we refer to as the history. In the traditional adaptive BED framework (Ryan) et al. 2016), this is done by assimilating the data into the model by fitting the posterior p(0) hx-1), followed by the maximization of the one-step ahead, or incremental, EIG
Whilst the total EIG formulation ( provides a unified train- ing objective for the policy, it remains doubly intractable like the standard EIG. The original Deep Adaptive Design (DAD) method (Foster et al. 2021) addressed this by using tractable variational lower bounds of the EIG (Foster et al., 2019 2020 Kleinegesse and Gutmann 2020) coupled with stochastic gradient ascent (SGA) schemes to directly train policy network taking the form of deep neural network directly mapping from histories design decisions. It thus provided a foundation for conducting PB-BED in practice.
~~ = E [p(0] | h(-1) - H [p(0 ho]]],
(1)
where the expectation is taken with respect to the marginal distribution p(y &t, h1-1) = Epcon._,)ll( 0, &,, hr- )]-
We use the superscript h to emphasize conditioning on
A number of extensions to the DAD approach have since been developed (Blau et al. 2022 Ivanova et al. 2021 Lim et al. 2022), broadening its applicability to a wider class of models by proposing alternative policy training schemes. All share core methodology, where the policy is trained only once, offline, with experimental histories simulated from the model p(O)p(hrr 1).)). Once trained, it remains unchanged during the live experiment and across multiple experimental instances (e.g. different survey participants), as illustrated in Fig Lb This fully amortized approach elimi- nates the need for posterior inference and EIG optimization
the history currently available, setting ho = Ø. This is a closed-loop approach (Foster) 2022] Huan and Marzouk] 2016), explicitly integrating all of the acquired data to refine beliefs about and inform subsequent design decisions.
Whilst this traditional framework offers a principled and systematic way to optimize experimental designs, it comes with some limitations. One drawback is its myopic nature that greedily maximizes for the next best design and over- looks the impact of future experiments, ultimately leading to sub-optimal design decisions. Another limitation is the
2

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
Live experiment
Offline
Live experiment
Offfine
Live experiment
training
training
Policy
optimize
Design
observe
Design
Initial Policy
refine
T iterations
Policy
T Fiterations
Posterior
Design
Posterior
infor
Outcome
Outcome
infer
K steps
Outcome
iterations
(a) Traditional adaptive BED
(b) Fully amortized PB-BED
(c) aemi-amottized PB-BED
Figure 1: Overview of adaptive BED approaches. The traditional BED approach fits a posterior after each experiment iteration and optimizes for the next step best designs (i.e. greedily). Fully amortized policy-based BED approaches like DAD train a policy once offline, before the live experiment, then deploy this as a fixed policy to make adaptive design decisions during the experiment. Our semi-amortized approach enables periodic policy refinement during test-time.
at each step of the experiment, thereby allowing design decisions to be made almost instantly at deployment.
finite data. Together these lead to a discrepancy between the learned policy т and the true optimal design strategy n², producing an approximation gap for the learned policies.
3. Semi-Amortized PB-BED
Double reliance on the generative model Fully amortized PB-BED relies on the generative model to both simulate experimental histories for policy training and to evaluate the success of our design decisions via the total resulting information gained. In other words, we use the model in both the expectation and information gain elements of the EIG in 2 This dual reliance magnifies the consequences of model misspecifications (Go and Isaac 2022 Overstall and McGree 2022). Moreover, even if the model is well- specified from Bayesian inference perspective, there might still be significant discrepancies between the prior-pr prioo-predictive distribution, p(hr /r), used to simulate data in the policy training and the true underlying data generating distribution.
Fully amortized PB-BED methods enable real-time deploy- ment and provide design decisions that typically superior to those of the traditional framework. However, there are many problems where we can afford to perform some test- time training during the experiment itself. It is therefore natural to ask whether we can usefully exploit such com- putational availability to further improve the quality of our design decisions? In particular, the fact that the current state-of-the-art approaches for design quality are all fully amortized suggests that improvements should be possible when this is not a computational necessity.
To address this, we note that the computational gains of fully amortized PB-BED methods comes at the cost of their inability adapt the policy itself in response acquired experimental data. We argue that this rigidity leads to sub- optimal designs decisions, partieularly in scenarios where the real-world experimental data significantly deviates from the simulated histories used during training of the policy. Two primary factors contribute to this issue:
The upshot of this is that we may see data at deployment that highly distinct from any simulated during the policy training. The lack of mechanisms for integrating real exper- imental data means that fully amortized approaches have no mechanism to overcome this issue. This can be charac- terized as form of generalization gap-the learned policy fails to generalize to the real-world experimental conditions, due to its inability to integrate and respond to the actual experimental data gathered so far (Hastie et al. 2009).
Imperfect training In fully amortized PB-BED we simu- late experimental histories to try and leam policy that will generalize across the entire experimental space- effeccilly learning a regressor from all possible histories to design decisions. However, the effectiveness of any learner with finite data/training is inevitably limited, especially in re- gions of the input space where training data is sparse. In short, we are learning a policy to cover all possible histo- ries we might see, but at deployment we are dealing only with specific history that may be similar few, if any, of the histories we simulated during training. This challenge is particularly exacerbated in experiments with extended horizons, due to the high dimensionality of the resulting histories. Additionally, the finite representational capacity of the policy hinders perfect approximation even with in-
3.1. Online policy updating
To address these limitations, we propose a semi-amortized PB-BED BED framework. which introduces dynamic adapt- ability by allowing periodic updates to the policy during deployment in response to acquired experimental data. In short, it will update the original policy at one or more points during the experiment, refining it to maximize the EIG of the remaining steps, conditioned on the data gathered so far.
The motivation behind this semi-amortized framework is the intuition that while a fully amortized policy is strong starting point, it can be significantly enhanced through tar- geted refinements leveraging gathered data. Focusing for
3

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
now on the case of single policy update, the following proposition formalizes this intuition and lays the theoretical foundation for semi-amortized PB-BED.
Ty, which is trained to maximize the total remaining EIG, I (11-1((T))) as defined in (5)
This gives us an infer-refine process for semi- semi-amortization in PB-BED that mirrors the two stage procedure characteris- tic of traditional adaptive BED (cf Fig. [a Fig. [c). infer stage entails fitting the posterior distribution p(O htr) with the data up to r. The subsequent refine stage learns customized policy Ty for the remaining steps of the ex- periment by maximizing (5) It therefore allows for more effective design decisions than the fully amortized approach, However, unlike the traditional BED approach, which is greedy and requires updates at every experimental step, our semi-ameon amot-izort method offers a superior non con-myoiic design strategy and allows for selective updates.
Proposition 3.1 (Decomposition of total EIG). For any design policy n, the total EIG of a Ttsepp experiment can be decomposed as
[c). The
(T) = In II-+((T) (T) )+ + Ep(hr 1n) 1 (T),,
(4)
for any intermediate step S S T, where
The (-))
p(h++1:r | (ir 0,
(5)
E pp Yolh, nar (,,. log p(1 |hr, π
Proof. We can write the likelihood and marginal as
It is important to acknowledge that this approach requires some computation during the experiment, which can pose challenges in applications where design decisions must be made very quickly. However, in many cases there is some computation time available, and our semi-amortized ap- proach can exploit this, even if the available time is limited. In particular, as we will show in subsequent sections, im- provements the policy can often be achieved with minimal additional training, such that substantial gains often pos- sible without drastically compromising deployment speed. As such, semi-amortiee PB- FB-BED still maintains large com- putational benefits over traditional adaptive BED.
p(hr 0, = p(h, 0, x ())((((h+ 0,,0, p(hr x) = p(h, x (h (hr n).
Substituting in 2 and rearranging now yields
log p(h, 00 TT) p(h, |T)
I1- (TT) (p(0)r(h, (0,z)
log p(hrter hy, 0)
1000, 11th.
p(hr+1e My, x)
= (T) +
+ Ep(h- m) (+))))(())) (T) as required.
This decomposition of the total EIG into two distinct tomponents-the EIG accumulated up to an intermediate step т, and the expected EIG for subsequent steps condi- tional on the history at that point h, demonstattss that the optimality of a policy for the latter phases of the experiment, from + 1 to T, is solely determined by the model h,.
Multi-step policy updates We can naturally extend our approach to include a multi-step update mechanism, not- ing that Proposition 3.1 can be applied recursively break down the total EIG into more segments. To this end, we define refinement schedule, T = To, Ti, Tk-an an in- creasing sequence defining the points at which the policy is refined. We adopt the convention To = 0 and ho = Ø, mark- ing the offline optimization of the fully-amortized policy no- For Tı >> we follow our two-stage refer-refine procedure. In general, the more often we refine the policy, the better it will be (albeit with diminishing returns), at the cost of increasing the required deployment-time computation.
To this, first note that, without loss of generality, we can break down the definition of our policy into how it behaves when given histories of length less than t and when it is given longer histories, such that (he) = (h,) if < T and T (ht) = тz (h+, if 2 r. We can thus rewrite 4 as
(m) (mo) + (p,,, | 10 (++)]
4. Stepwise Deep Adaptive Design
Here the first term is independent of n,, while wo affects the second one only through its influence on the distribution of
We introduce Stepwise Deep Adaptive Design (Step-DAD) to implement our semi-amortized PB-BED framework in practice. Building on DAD and the infer-refine procedure outlined in the last section, Step-DAD employs stochas- tic gradient ascent schemes optimize variational lower bounds on the remaining EIG (5) to sequentially train the step policy Λ" in scalable manner. An overview of Step- DAD is presented in Algorithm I in Appendix A
h,. At deployment time, we will have a specific h, once we reach step r of the experiment, and so the conditional opti- mal policy for the remaining steps given the data gathered so far is arg max,, 1-;1-r((7,) (TT), which is independent of mp-
Our amdrtizor framework is now based around exploit- ing this independence to refine the policy midway through the experiment by introducing a step design policy x². Ini- tially. n" uses the fully amortized policy To for the first 7 steps of the experiment. Here Wo trained as if would be used for the full experiment, such that we maximize its total FIG, (TT). After step T, T switches to a new policy
The two key components of Step-DAD's aforementioned infer-refine procedure are an inference method for approxi- mating (O),,,, and a refinement strategy for using this to update our policy. Standard inference techniques (such as
4

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
variational inference or Monte Carlo methods) can be used for the former as discussed in our experiments. Our focus here will therefore instead be our specialized procedure for policy refinement and the policy architecture itself.
In principle, one could train an entirely new policy Tr, at each refinement step Tıe potentially even varying the specific architecture between these. Though such a strategy may occasionally be advantageous, we instead, propose more pragmatic and lightweight approach: leveraging the already established fully amortized policy To as a baseline and fine- tuning it for subsequent steps. In our experiments we do this using full fine-tuning of all policy parameters, but one could instead implement more parameter-efficiet methods if needed, for example, only adjusting the last few layers.
4.1. Policy refinement
Due to its doubly intractable nature, the task of optimizing the remaining EIG, -ー+1--(T)) presents notable challenge. In selecting appropriate scalable and efficient estimator for it, we wish to ensure compatibility with wide range of inference schemes for p(0 h..). Namely, as this serves as an updated 'prior' during the policy refinement, it is important that we use an EIG estimator that does not require evaluations of the prior density, to ensure compatibility with sample-based inference schemes.
5. Related Work
The idea of using a design policy in the context of adap- tive BED was first proposed by Huan and Marzouk (2016). Leveraging dynamic programming principles, the policy they learn aims to establish mapping from explicit poste- rior representations-sevvvii as the state in reinforcement learning (RL) terminology- subsequent design choices. As a result, each iteration of the experiment necessitates sub- stantial computational resources for updating the posterior. The concept of fully amortized policy-basee BED, which directly maps data collected to design decisions, has only re- cently been introduced (Foster et al. 2021) and subsequently extended to differentiable implicit models (Ivanova et al.) 2021) and downstream tasks (Huang et al., 2024). While Step-DAD uses the policy training approach of DAD and iDAD based on direct SGA of variational bounds (e.g. ( our semi-amortized PB-BED framework also compatible with more RL-based design policy training approaches. like those of Blau et al. (2022) and Lim et al. (2022), which are more suited to discrete design spaces. We emphasize that none of these previous approaches have looked to refine the policy during the experiment itself.
Lower bound estimators such as the applicit-ikkelhhodd based sequential Prior Contrastive Estimator (sPCE, Fos- ter et al. (021), as well as the implicit likelihood In-
foNCE (Ivanova et al. [02]] van den Oord et al.] 2018)
and NWJ (Kleinegesse and Gutmann 2020 Nguyen et al.
2010) bounds, align with this requirement. For generative models with explicit likelihoods (implicit models are dis- cussed in Appendix
B)
we therefore use the sPCE bound:
Nog C
p(hr 11: 10, RTy; (T) =E log
- T
L+I Eizo pl t 10 hops mT)
Step-DAD parameterizes by neural network and opti- mizes an appropriate objective, such as (6). with respect to the network parameters using stochastic gradient ascent (SGA) schemes (Kingma and Ba 2014 Robbins and Monro
1951). Following Foster et al. (2021) ), we use path-wise gradients in the case of reparametrizable distributions (Mo- hamed et al. 2020; Rezende et al.) 2014), and score function (REINFORCE) otherwise (Williams 1992).
As discussed in 2. adaptive BED has traditionally em- ployed a two-step greedy strategy, involving posterior infer- ence followed by an EIG optimization (Foster al. 2019
2020
Huan and Marzouk 2014 Kleinegesse and Gutmann
4.2. Policy architecture
2019
2020
2021
Kleinegesse et al.
2021 Myung et al.
Similar to Foster et al. (2021), our policy architecture is based on individually embeding each design-outcome pair
2013
Overstall and McGree 2020 Price et al. 2018 Ryan
et al.
2016 Vincent and Rainforth 2017). While Step-DAD
(E). E he into fixed-dimensional representation, be- fore aggregating them across t to produce a summary vector. This allows for condensing varied-length experimental histo- ries into a consistent dimensionality, to handle variable his- tory sizes. The summary vector is then mapped to the next experimental design fr+1. For the aggregation mechanism, the choice between permutation invariant and autoregressive architectures depends the nature of the data. When the data h4 is exchangeable, permutation invariant architectures like DeepSets (Zaheer et al. (017)) or SerTransformer (Lee et al., 2019) are suitable. In contrast, sequential or time- series data would benefit from autoregressive models like transformers /Vaswani et al. 2017).
diverges from these in its EIG optimization by training poli- cies instead of designs, it does share their need to perform posterior inference. The inference scheme used by previous work has varied between problems and the needs of the un- derlying Bayesian model being used, with sequential Monte Carlo (Del Moral et al.) 2006 Drovandi et al. 2014 Rain- Forth 2017] Vincent and Rainforth 2017) and likel likelihood-
free (Huan and Marzouk 2013 Lintusaari et al.) 2017]
Sisson et al. 2018 Thomas ct al. 2016) inference schemes
proving popular. There is also growing recent interest in approaches that utilize inference itself as part of the EIG
optimization. (Amzal et al.; 2006 Iollo et al. 2024 2025
5

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
Iqbal et al. 2024a b). Important considerations in choosing this scheme include the availability of an explicit likelihood, the ability to take derivatives, and computational budget.
Table I: Source Location Finding, Upper and lower bound estimates of total EIG. We report T = 6, rest in Table [T in the Appendix. Errors show + +ss.e., com- puted over 16 (2048) histories for step methods (rest). DAD was trained for 50K steps, Step-DAD for 2.5K.
The challenge of model misspecification in BED remains a critical, but relatively underexplored, problem ( (Farquhar
iet al. 2021 Feng et al. 2015 Go and Isaac 2022 Overstall
Method
Lower hound (1)
Upper bound (4)
and McGree 2022] Rainforth et al.) 20244 Sloman et al.
Random
3.612 + 0.012
3.613 +: 0.012
2022). Fully amortized PB-BED is particularly vulnerable
Static
3,945 + 0.026
3,946 +: 0.026
to model misspecification due to its reliance on a singular learning phase without the capacity to integrate real-world experimental feedback. As we will see in the experiments, our -motiiz PB-BED methodology. whilst not di- rectly tackling the issue of misspecificction, typically does enhance robustness to misspecification, due to enabling iter- ative data integration and policy refinement.
Step-Static (T 66
3,974 + 0.008
3,975 + 0.008
DAD
7,040 + 0.012
7,089 + 0.013
Step-DAD (7 6)
7.759 + 0.114
7.765 + 0.114
icy parameters (i.e. it learns non-adaptive policy whose design choices are fixed). The Step-Static baseline is a two- stage approach that first trains a set of static designs by optimizing a PCE bound on ( ... ( before approximating the posterior and training new conditional set of static designs for the last T - T steps by optimizing a PCE bound on -5 Er). When possible, we include Problem-Speciffe baselines used for the relevant experiment before, instead considering Random design strategy when not. In all cases, the number of contrastive samples used during training was L = 1023.
Finally, BED shares important connections to reinforcement learning (Sutton and Barto] 2018). Most notably, it has been shown that the adaptive BED problem can be formulated as various forms of Markov Design Processes (MDPs (Doshi Velez and Konidaris 20166; Guez ct al. 2012] Ross et al. 2007)), using the incremental EIG as the reward and either the posterior (Huan and Marzouk 2016) or, more 2022 Foster
practically, the history as the state (Blau et al.
(021). Here PB-BED approaches are most c closely linked
Our main metric for assessing the quality of various design strategies is the total EIG, (x), as given in (2). For the baselines, we approximate it via a version of the sPCE lower bound (6) with L = 10th to ensure tight bound, along with its upper bound counterpart- the sequential Nested Monte Carlo estimator (sNMC, Foster et al. 2021 (see Ap- pendix B For Step-DAD, we instead use conservative lower bound estimate on its difference in performance com- pared to the original DAD network, before adding this to the corresponding DAD estimate (see Appendix sures any biases from using bounds instead of the true EIG lead underestimation of the gains Step-DAD gives. Full details experimental details are provided in Appendix C
with offline model-based RL (Kidambi et al.
2020 Levine
let al. 2020 Moerland et al.) 2023 Ross and Bagnell 2012
Yu et al. 2020). in that they learn a policy upfront which can then be deployed. However, PB-BED varies in many significant ways from typical RL settings. For example, we do not have access to any data train our policy, but are instead focused on optimal sequential decision-making under a given model. We can also typically directly use SGA to train the policy as we have access to end-to-end differentiable objectives and our rewards are typically not sparse. We also note that our extension of fully-amortize PB-BED to -mortized PB-BED is quite distinct to the typical generalization of offline RL approaches to hybrid RL approaches (Ball et al.) 2023 Song et al.) 2023 Zheng et al. (022), as we are making refinements to the local policy during single rollout using the same objective as original amortized policy. Our approach of refining a learned policy as new data becomes available also shares some similarities with Model Predictive Control (Qn and Badgwell 2003).
C.3) This en-
6.1. Source Location Finding
We first consider the source location finding experiment from Foster al. (202T). which draws upon the acoustic energy attenuation model, detailed in [Sheng and Hu (2005) ) The objective of the experiment is to infer the locations of some hidden sources using noisy measurements, y, of their combined signal intensity. Each source emits a signal that decreases intensity according to the inverse-square law. A full description of the model given in Appendix C.5
6. Experiments
We empirically evaluate Step-DAD on a range of design problems, comparing its performance against DAD to de- termine the additional EIG achieved by the step policy x² over the fully amortized policy To. We further consider several other baselines for comparison. Static design learns a fixed set of designs prior to the experiment by optimising a PCE bound (Foster et al. 2020) that is equivalent to (6). but which is defined in terms of the designs rather than pol-
We begin by learning fully amortized DAD policy to per- form T 10 experiment steps to locate single source. We chose training budget of 50K gradient steps this policy, as we found that further training did not significantly improve performance with our chosen architecture.
Single policy update We systematically evaluate our Step-
6

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
8.0
Table 2: Location finding multiple sources. Reported
for r = 7, Step-DAD (10K, 2.5K). Errors show Ils.e.
Totol EIG
2.5
0 dim
EIG difference
DAD, total EIG (upper)
7.0
4
0.701 + 0.023
6,483 + 0.055
8
0.426 + 0.014
7.111 = 0.067
6.5
12
0.423 + 0.012
6.956 + 0.056
Tuming step.
DAD (50K), upper
Step-DAD (50K (2.5K). lowes
0.0
Sleg-DD (10K)2.5K) lower
0 DAD (1NK), upper
. Step-DAD (10K/2.5K), Jower
Total EIG fiven
DAD (16K), upper
2.0-
Figure 2: Sensitivity to training budget for location finding experiment. DAD policies are trained for 50K or 10K steps, Step-DAD policies are refined for 2.5K. Errors show ls.e.
1.0
0.0
Priot Mean Shift
DAD approach by exploring all possible fine-tuning steps, = 1,...,3 9. We use importance sampling for posterior inference and fine-tune the policy for 2.5K steps. Results for r = 6 are presented in [able whilst Table i the Appendix shows performance for all values r.
Figure 3: Sensitivity to prior perturbations in the loca- tion finding experiment. Total EIG for Step-DAD remains more robust compared to DAD, which drops to zero.
Fully amortized PB-BED approaches can be particularly prone to pathologies from imperfect prior choices, as the prior dictates the generated data the policy is trained on. Namely, the prior predictive poorly matches the true data generating process (DGP), we may observe data at deploy- ment that is highly distinct anything seen in the policy training. To evaluate if Step-DAD can improve robustness to prior imperfections, we now consider case where the prior, p(0), used for the offline policy training leads to a DGP that is significantly different to the true DGP ob- served at deployment. To this end, we introduce a test-time prior p(0), and evaluate the policy performance under the DGP p(yn:rlEur)) = Eptr) III p(ye (0, Ex, h4-1)], using the EIG under this alternative model as an evaluation metric
Sensitivity to training budget We investigate the overall resource efficiency of Step-DAD by comparing to DAD under two training budgets. The full budget is as before at 50K gradient steps, whilst the reduced budget is limited to 10K steps, both then followed by 2.5K finetuning steps for the Step-DAD networks. Figure presents a conservative comparison, showing upper bound estimates for DAD and lower bound estimates for Step-DAD. The results reveal that Step-DAD consistently outperforms its respective DAD baseline for all r > both budget levels (the apparent slight drop for is likely due the conservative estimation scheme used). Interestingly, Step-DAD with the reduced budget matches or exceeds the DAD with the full budget for all 7 > 3, thereby achieving better results with nearly 5 times fewer total training steps.
[ (T) = (p(0) [log p(hr 0, T) - log where p(hrl (p m)].
We note that the performance advantage of Step-DAD over DAD appears to be most pronounced when fine-tuning oc- curs just past the midpoint of the experiment, that is for 7,7 or 8. At this stage, our method can effectively leverage the accumulated data to refine the policy, while ensuring there are enough experiment steps remaining to benefit from the improved, customized policy.
The results for the source location finding design problem are shown in Figure 3 It reveals that Step-DAD consistently outperforms the DAD baseline across all degrees of prior shift we consider, with the EIG for DAD decreasing to essen- tially zero with the increased prior shift, whilst Step-DAD able deliver positive information gains. This robustness is anticipated due to Step-DAD's 's ability to assimilate data gathered and adjust policies in light of new evidence.
Multiple sources We next consider a more complex setting of locating 2,4 and 6 sources, which correspond to a 4-, 8- and 12-dimensional unknown parameter, respectively. Table 2 shows the results, indicating a consistent positive EIG difference for Step-DAD over DAD.
6.3. Test- Dist-Time& Compute Ablations
We now perform ablations better understand how test- time performance is impacted by restrictions on computa- tional budget. Firstly, for the location finding experiment we consider three different per-step budgets for the inference and then vary the amount of fine-tuning steps performed for updating the StepDAD network. As shown in Figure 4 total EIG generally improves with higher computational
6.2. Robustness to prior perturbations
In BED, selecting an appropriate prior is critical as it both influences our final posterior, and the data we gather in the
first place (Go and Isaac 2022 Simchowitz et al. 2021)
7

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
(PCE)
6.5
Total EIG
6.0
5.5
DAD (LOOS),
Stap-TAD (100 koor
5.0
600
Wall Time (6)
Timing steg,
Figure 6: Hyperbolic temporal discounting. EIG im- provement of Step-DAD over DAD after fine-tuning the policy at step The fully amortized DAD policy is trained for 100K steps, step-policy refined for 1K steps.
100% inference
20%
200% Inforence
Figure 4: EIG as a function of wall time in the location finding experiment (T=10, r = 6). varying both the infer- ence budget and the number of fine-tuning steps (between 250 and 10000 steps). Here 100% inference budget corre- sponds to taking 20000 importance samples.
Table 3: Hyperbolic temporal discounting, Estimates of total EIG.II-22 (2048) histories for the step methods (rest). Baselines as reported in Foster et al. (2021). except DAD and Random.
Errors indicate = s.e., over 16
Method
Lower bound (1)
Upper bound (L)
Random
2.249 + 0.010
2.249 0.010
Total EIG (POS)
8.5
Kirby 2009
1.861 + 0.008
1.864 0.009
Statie
2.518 + 0.007
2.524 + 0.007
Frye 2016
3.500 = 0.029
3,513 0.029
Greedy (RADapted)
4.454 0.016
4.536 0.018
7.5-
DAD
4.778 + 0.013
4.808 0.014
Step-DAD (r=10)
6,711 + 0.040
6.721 0.040
7.0
Nurober Interwentions
6.4. Hyperbolic Temporal Discounting
100% Inference
20% Inforenico
Temporal discounting describes the tendency for individuals to prefer smaller immediate rewards over larger delayed ones. This phenomenon is a key concept in psychology and economics and has been used to study important social and individual behaviors (Critchfield and Kollins 2001).
Figure 5: Sensitivity to number of interventions for loca- tion finding experiment (T=10) for different inference bud- gets (100% 20000 importance samples). EIG increases with more interventions, then plateaus. Each update corre- sponds to posterior inference 2.5k fine-tuning steps.
including dietary choices (Bickel al.] 2021), exercise habits (Tate al. 2015), patterns of substance abuse and addictive behaviours (Story et al.) 2014). An individual's time delay preference is typically measured by asking them series of questions, such as "Would you prefer $R now or $100 in D days time?" Here the tuple ç = (R, D) defines our experimental design, and the experiment outcome y is the participant's 's decision either accept or rejecr the delay.
budgets, with diminishing returns for large budgets. The performance for the 100% and 200% inference budgets are quite similar, indicating our inference has been successful, but a performance drop is seen when only using 20% of our previous inference budget. Importantly though, significant gains relative DAD (corresponding Os wall time in the plot) are still achieved with small budgets corresponding to around a minute of wall time.
Single update Using the hyperbolic discounting model introduced in Mazur 1987) and as implemented by Vincent (2016), we train DAD policy for 100K gradient steps, aimed at designing = 20 experiments. We select grid of tuning steps 7 in the range from 2 to 18 in increments 2. For posterior inference, we use simple importance sampling to draw samples from the posterior p(0 h,) and 1% of the original training budget (i.e. 1K gradient steps).
Secondly, we evaluate how performance evolves as a func- tion of the number interventions. Each intervention con- sists of updating the posterior and applying 2.5k fine-tuning steps. Figure5]showe shows that Total EIG consistently increases with more interventions across all inference budgets, before eventually plateauing. We note that even for small number of interventions, there is increased performance over DAD (corresponding to 0 interventions in the plot), demonstrat- ing the utility of Step-DAD even in settings of constrained test-time compute budget.
Figure 6 reports the results and illustrates that Step-DAD yields an improvement in total EIG for all choices of 7 when compared to the baseline DAD policy. The largest increase occurs around, and shortly after, the middle the exper- iment, aligning with our previous intuition: here sufficient
8

---

Step-DAD: Semi-Amortized Policy-Based Based Bayesian Experimental Design
data has been accumulated inform meaningful posterior update, whilst sufficient number of experiments remain to effectively deploy the refined policy. Table 3 demonstrates the superiority of Step-DAD over conventional baselines, in-
Table 5: Constant Elasticity of Substitution. Estimates on total EIG, 11-11 DAD and Static trained for 50K steps; step variants finetuned for 10K steps. Errors denote + s.e.
cluding those derived from psychology research (Frye et al.
Method
Lower bound (t)
Upper bound (L)
2016; Kirby 2009; Vincent and Rainforth, 2017) and tradi- tional BED approaches such as the BADapted (Vincent and Rainforth 2017) approach which was specifically designed for this problem. It also outperforms the static BED strategy, highlighting the effectiveness of adaptive design strategies in extracting more valuable information from experiments.
Random
2.487 + 0.007
2.487 = 0.007
Greedy (vPCF)
13.333 + 0.975
13.343 + 0.975
Statie
9.279 + 0.020
11.183 + 0.0453
Step-Static (T = 5)
13.010 + 0.185
13.682 + 0.189
DAD Step-DAD (T (T= 13.879 1 0.352
10.181 + 0.021
11.478 + 0.042
14.623 0.363
Multiple updates and design extrapolation We extend the deployment of DAD and Step-DAD policies for this problem to T = 40 experiment steps, doubling the scope at which they were originally trained, i.e. without retraining the DAD network. Step-DAD is fine-tuned two steps, T and 27, with 7 E (5, 6,7, 8). As @ Step-DAD demonstrates significantly improved capacity to extract in- formation in later stages, beyond its initial training. This highlights the robustness and flexibility of our method in extending experimental horizons compared to DAD.
convergence to suboptimal local maxima.
7. Conclusions
In this work, we introduced the idea of semi-amortized approach to PB-BED that enhances the flexibility, robust- ness and effectiveness of fully amortized design policies. Our method, Stepwise Deep Adaptive Design (Step-DAD), dynamically updates its step policy in response to new data through a systematic 'infer-refine' procedure that refines the design policy for the remaining experiments in light of the experimental data gathered so far. This iterative refinement enables the step policy to evolve as the experiment progresses, ensuring more robust and tailored design decisions, as demonstrated in our empirical evaluation. Step-DAD thus improves our ability to conduct more efficient, informed, and robust experiments, opening new avenues for exploration in various scientific domains.
6.5. Constant Elasticity of Substitution (CES)
We conclude our evaluation with the Constant Elasticity of Substitution (CES) model, a framework from behavioral economies analyse the relative utility of two baskets of goods (Arrow al. [961). This model emulates how eco- nomic actors specify their relative preference y between these baskets on a sliding scale. We follow the experimen- tal setup of Foster et al. (2019), with full details given in Appendix C.7] The CES model faces challenges due to y being sampled from a censored normal distribution, concen- trating probability at observation boundaries and creating local maxima (Blau et al. 2022 Foster al. 2019).
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Table shows that Step-DAD outperforms all basclines. Step-Static also achieves on on-par results with Greedy (vPCE), highlighting the benefits of semi-amortized design strategies in this model. We note the performance of DAD is worse than Step-Static. This can be attributed to the discontinu- ities in the censored likelihood (Eq. (9) in the Appendix) that complicate the training of the policy, often resulting in
Acknowledgements
MH is supported by funding provided by Novo Nordisk and by the EPSRC Centre for Doctoral Training in Modern Statistics and Statistical Machine Learning. TR is supported by the UK EPSRC grant EP/Y037200/1.
Table 4: Hyperbolic temporal discounting: extrapolating designs. Comparison of EIG upper bound for DAD and lower bound for Step-DAD across tuning steps 7 and 7
References
Billy Amzal, Frédéric Y Bois, Eric Parent, and Christian P Robert. Bayesian optimal design via interacting particle systems. Journal of the American Statistical association, 101(474):773-88, 2006.
40. Errors indicate ls.e., computed over 16 histories.
EIG from (t)
EIG from 27 (t)
DAD
Step-DAD
DAD
Step-DAD
5
3.9 + 0.24
4.8 + 0.14
2.1 + 0.36
4.6 0.18
Kenneth J Arrow, Hollis B Chenery, Bagicha S Minhas, and Robert M Solow. Capital-labor substitution and economic efficiency. The review of Economics and Statistics, pages
6
3.4 0.30
5.1 - 0.07
0.30
0.12
7
3.0 - 0.35
0.30
0.23
4.11
8
2.6 - 0.36
0.13
1.0
0.13
9

---

Step-DAD: Semi-Amortized Policy-Based Based Bayesian Experimental Design
225-250, 1961. 19
Chi Feng et al. Optimal Bayesian experimental design in
the presence of model error. PhD thesis, Massachusetts Institute of Technology, 2015. 6
Anthony Atkinson, Alexander Donev, and Randall Tobias. Optimum Experimental Designs, with SAS. Oxford Uni- versity Press, 2007. [
Adam Foster. Variational, Monte Carlo and Policy-Based Approaches to Bayesian Experimental Design. PhD the- sis, University of Oxford, 2022. 2
Philip J. Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with of- fline data, 2023. 0
Adam Foster, Martin Jankowiak, Elias Bingham, Paul Hors- fall, Yee Whye Teh, Thomas Rainforth, and Noah Good- man. Variational Bayesian Optimal Experimental Design. In Advances in Neural Information Processing Systems 32, pages 14036-14047. Curran Associates, Inc., 2019.
Warren K Bickel, Roberta Lentos, Devin C Tomlin- son, William H Craft, Diana R Keith, Liqa N Athamneh, Julia C Basso, and Leonard H Epstein. Temporal dis- counting as a candidate behavioral marker of obesity. Neuroscience & Biobehavioral Reviews, 129:307-329,
3 9
2021. 3
Adam Foster, Martin Jankowiak, Matthew O' O'Meara, Yee Whye Teh, and Tom Rainforth. A unified stochastic gradient approach to designing bayesian-optimal experi- ments. In International Conference on Artificial Intelli- gence and Statistics, pages 2959-29966 PMLR, 2020. 2
Eli Bingham, Jonathan P Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah Goodman. Pyro: Deep universal probabilistic programming- Journal of Machine Learning Research, 2018. [4
5
Adam Foster, Desi R Ivanova, Ilyas Malik, and Tom Rain- forth. Deep adaptive design: Amortizing sequential bayesian experimental design. Proceedings of the 38th International Conference on Machine Learning (ICML., PMLR 139, 2021.
Tom Blau, Edwin V Bonilla, ladine Chades, and Amir Dez- fouli. Optimizing sequential experimental design with deep reinforcement learning. In International conference on machine learning. pages 2107-2128. PMLR, 2022. [
15
9
15
Adam Evan Foster. Variational, Monte Carlo and Policy- Based Approaches to Bayesian Experimental Design. PhD thesis, University of Oxford, 2021.
Kathryn Chaloner and Isabella Verdinelli. Bayesian ex- perimental design: A review. Statistical Science pages
6
273-304, 1995. [
Charles CJ Frye, Ann Galizio, Jonathan E Friedel, W Brady
Thomas S Critchfield and Scott H Kollins. Temporal dis- counting: Basic research and the analysis of socially important behavior. Journal of applied behavior analysis,
DeHart, and Amy L Odum. Measuring delay discounting in humans using an adjusting amount task. JOVE (Journal of Visualized Experiments), (107):e5354, 2016.
8
34(1):101-12) 2001.
Jinwoo Go and Tobin Isaac. Robust expected information gain for optimal bayesian experimental design using am- biguity sets. In Uncertainty in Artificial Intelligence. pages 728-737. PMLR, 2022.
Pierre Del Moral, Arnad Doucet, and Ajay Jasra. Sequen- tial monte carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68(3):411- 436, 2006. 5
Arthur Guez, David Silver, and Peter Dayan. Efficient Bayes-adaptive reinforcement learning using sample- based search. In Advances in neural information pro- cessing systems, pages 1025-1033, 2012. 6
Finale Doshi- Velez and George Konidaris. Hidden parame- ter markov decision processes: A semiparametric regres- sion approach for discovering latent task parametrizations. In IJCAI: proceedings of the conference, volume 2016, page 1432. NIH Public Access, 2016. 6
T. Hastie, R. Tibshirani, and J.H. Friedman. The Elements of Statistical Learning: Data Mining, Inference, and Predic- tion. Springer series in statistics. Springer, 2009. ISBN 9780387848846. URL
Christopher C Drovandi, James M McGree, and Anthony N Pettitt. A sequential monte carlo algorithm to incorporate model uncertainty in bayesian sequential design. Journal of Computational and Graphical Statistics, 23(1):3-24 2014. 5
http:/oooooo..oll..... 1
co.uk wtw.boooss?...eeeS......... Chapter 7.4, p. 228.1
Xun Huan and Youssef Marzouk. Gradient-based stochastic optimization methods in bayesian experimental design. International Journal for Uncertainty Quantification, 4
Sebastian Farquhar, Yarin Gal, and Tom Rainforth. On statistical bias in active learning: How and when to fix it. arXiv preprint 2ix:::2:11.... 166,, 2021. 6
(6). 2014.
10

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
Xun Huan and Youssef M Marzouk. Simulation-based opti- mal bayesian experimental design for nonlinear systems. Journal of Computational Physics, Physics
d811408 1411003110 014444666669ccc33110111881111111
pdf
2 5
(32(1)2288317,200
5
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline
Xun Huan and Youssef M Marzouk. Sequential bayesian optimal experimental design via approximate dynamic programming. arXiv preprint arXiv:1604.08320 2016.
reinforcement learning. Advances in neural information
processing systems, 33:21810-21825, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXivv1412.6998,
Daolang Huang, Yujia Guo, Luigi Acerbi, and Samuel Kaski. Amortized bayesian experimental design for decision- making. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 109460-109486. Curran Associates, Inc., 2024. [
2014. 5
Kris N Kirby. One-year temporal stability of delay rates. Psychonomic bulletin & review, 16(3):457-46) 2009.
S. Kleinegesse and M.U. Gutmann. Efficient Bayesian experimental design for implicit models. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of the International Conference on Artificial Intelligence and Statistics (AISTATS), volume 89 of Proceedings of Machine Learning Research, pages 1584-1592. PMLR, 2019. 5
Jacopo Iollo, Christophe Heinkelé, Pierre Alliez, and Flo- rence Forbes. PASOA- PArticle baSed Bayesian opti- mal adaptive design. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Pro- ceedings of the 41 st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learn- ing Research, pages 2024. URI. /ttR.//htttss///roooce v235/io11o24a.htm 5
PMLR, 21-27 Jul ngs. ml r. rrrsss//
Steven Kleinegesse and Michael Gutmann. Bayesian exper- imental design for implicit models by mutual information neural estimation. In Proceedings of the 37th Interna- tional Conference on Machine Learning, Proceedings of Machine Learning Research, pages 5316-53163 5326.. PMLR,
Jacopo lollo, Christophe Heinkelé, Pierre Alliez, and Flo- rence Forbes. Bayesian experimental design via con- trastive diffusions. 2025.
2020. 2 5
URL https://arxii
arg/abs/ 2210 11826
5
Steven Kleinegesse and Michael U. Gutmann. Gradient- based bayesian experimental design for implicit models using mutual information lower bounds. arXiv preprint DXXiv:2105.0437, 2021.
Sahel Iqbal, Hany Abdulsamad, Sara Pérez-Vieites, Simo Särkkä, and Adrien Corenflos. Recursive nested fil- tering for efficient amortized bayesian experimental design, 2024a. URL https://arxiv.org/ abs/ 2409. 05354] 6
Steven Kleinegesse, Christopher Drovandi, and Michael U. Gutmann. Sequential Bayesian Experimental Design for Implicit Models via Mutual Information. Bayesian Analysis, pages 1 30, 2021. doi: 10.1214/00-BA1225
Sahel Iqbal, Adrien Corenflos, Simo Särkkä, and Hany Ab- dulsamad. Nesting particle filters for experimental design in dynamical systems. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Pro- ceedings of the 41st International Conference on Ma- chine Learning, volume 235 of Proceedings of Machine Learning Research, pages 21047-21068. PMLR, 21- 27 Jul 2024b. URL /tttps://proceeeeee .mmr Tetess/v233/iqqaa11 html
5
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention- based permutation cermutation-invariant neural networks. In International conference machine learning, pages 3744-3753, PMLR, 2019. 5
6
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.01643. 2020.
Desi R Ivanova, Adam Foster, Steven Kleinegesse, Michael Gutmann, and Tom Rainforth. Implicit Deep Adaptive Design: Poseey-Based Experimental Design without Like- lihoods. In Advances in Neural Information Processing Systems, volume 34, pages 25785 25795-25798.
Vincent Lim, Ellen Novoseller, Jeffrey Ichnowski, Huang Huang, and Ken Goldberg. Policy-based bayesian ex- perimental design for non-differentiabbe implicit models. arXiv preprint arXiv:2203.04222. 2022.
Curran As-
sociates, Inc., 2021. URL hntps: htt.ttpps///roceeddngg.. wttpssccpppppe//eeee/ htarripps.ccppaper/2////1111

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
Dennis V Lindley. On measure of the information provided by an experiment. The Annals of Mathematical Statistics, pages 986-1005, 1956.
deep learning library. In Advances in Neural Informa- tion Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019.
14
Dennis V Lindley. Bayesian statistics, a review, volume 2. SIAM, 1972. [
David J Price, Nigel G Bean, Joshua V Ross, and Jonathan Tuke. An induced natural selection heuristic for finding optimal bayesian experimental designs. Computational Statistics & Data Analysis, 126: 1 12-124, 2018. 5
Lintusaari, M.U. Gutmann, R. Dutta, S. Kaski, and Corander. Fundamentals and recent developments in approximate Bayesian computation. Systematic Biology, 66( 1):6666-8222 January 2017. 5
S.Joe Qin and Thomas A. Badgwell. A survey of industrial model predictive control technology. Control Engineer- ing Practice, 11(7):733-764, 2003, ISSN 0967-0661. doi: http://bbboog///oo00666fff666660 URL httpps:/www.scieeceddlccee
David JC MacKay. Information-based objective functions for active data selection. Neural computation, 4(4):590~ 604, 1992.
James Mazur. An adjusting procedure for studying de- layed reinforcement. Commons, MI.; Mazur, JE.; Nevin, JA. pages 55-73, 1987.
Tom Rainforth. Automating Inference, Learning, and De- sign using Probabilistic Programming. PhD thesis, Uni- versity of Oxford, 2017.
5
Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based reinforcement learning: A survey. Foundations and Trends@ in Machine Learning, 16(1):1-118, 2023. 6
Tom Rainforth, Rob Cornish, Hongseok Yang. Andrew War- rington, and Frank Wood, On nesting monte carlo estima- tors. In International Conference on Machine Learning, pages 4267-4276, PMLR, 2018. 2
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. Journul of Machine Learning Research, 21
Tom Rainforth, Adam Foster, Desi Ivanova, and Freddie Bickford Smith. Modern bayesian experimental design. Statistical Science, 39(1):100-111 2024.
6
(132):1-63 62, 2020.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wier- stra. Stochastic backpropagation and approximate in- ference in deep generative models. In Proceedings of the 31st International Conference on Machine Learning, volume 32, pages 1278 1286, 2014. 5
Jay| Myung, Daniel R Cavagnaro, and Mark A Pitt. A tutorial on adaptive design optimization. Journal of math- ematical psychology, 57(3-4):53-67, 2013. 15
Xuanlong Nguyen, Martin J. Wainwright, and Michael Jor- dan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11). 2010. ISSN 00189448. doi: 111100/TT.201002066833
Herbert Robbins and Sutton Monro. A stochastic approx- imation method. The annals of mathematical statistics, pages 400-407, 1951. 5
14
Stephane Ross and J Andrew Bagnell. Agnostic system identification for model- reinforcement learning. In International Conference on Machine Learning, 2012.
Antony Overstall and James McGree Bayesian Design of Experiments for Intractable Likelihood Models Using Coupled Auxiliary Models and Multivariate Emulation. Bayesian Analysis, 15(1):103 131, 2020. doi: 10.1214/
Stephane Ross, Brahim Chaib-draa, and Joelle Pineau. Bayes-adaptive pomdps. In NIPS, pages 1225-1232, 2007. 6
19-BA1144, 5
Antony Overstall and James McGree. Bayesian decision- theoretic design of experiments under an alternative model. Bayesian Analysis, 1044 2022. 3
Elizabeth G Ryan, Christopher C Drovandi, James M Mc- Gree, and Anthony N Pettitt. A review of modern com- putational algorithms for bayesian optimal design. Inter- national Statistical Review, 84(1):
2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zem- ing Lin, Natalia Gimelshein, Luca Antiga, Alban Des- maison, Andreas Kopf, Edward Yang, Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chin- tala. Pytorch: An imperative style, high-performance
5
Wanggang Shen and Xun Huan. Bayesian sequential optimal experimental design for nonlinear models us- ing policy gradient reinforcement learning. CoRR,
Abs/2110.153355 2021. URL /tttts://aaatt /aaaii.v.oor/////./......... abs/2110.15335
2
12

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
Xiaohong Sheng and Yu Hen Hu. Maximum likelihood multiple-source localization using acoustic energy mea- surements with wireless sensor networks. IEEE Transac- tions on Signal Processing, 2005. ISSN 1053587X. doi: 10.1 109/TSSP2204:8889230
Benjamin T Vincent. Hierarchical bayesian estimation and
hypothesis testing for delay discounting tasks. Behavior
research methods, 48(4): 1608-1620, 2016.
8 17
Benjamin T Vincent and Tom Rainforth. The DARC
toolbox: automated, flexible, and efficient delayed and risky choice experiments using bayesian adaptive design.
Max Simchowitz, Christopher Tosh, Akshay Krishnamurthy, Daniel Hsu, Thodoris Lykouris, Miro Dudik, and Robert E Schapire. Bayesian decision-me decision-making misspecified priors with applications to meta-learning. Advances in Neural Information Processing Systems, 34: 2021.11
PsyArXiv preprint, 2017.
making under
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Ma- chine learning, 8(3-4):229-256, 1992.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33:
S.A. Sisson, Y. Fan, and M. Beaumont. Handbook of Ap- proximate Bayesian Computation. Chapman & Hall/CRC Handbooks of Modern Statistical Methods. CRC Press,
2018. ISBN 9781351643467.1 S
14142, 2020.
Sabina J Sloman, Daniel M Oppenheimer, Stephen B Broomell, and Cosma Rohilla Shalizi. Characterizing the robustness of bayesian adaptive experimental designs to active learning bias. arXiv preprint arXiv:2205.13668, 2022. 6
M. Zaharia, Andrew Chen, A. Davidson, A. Ghodsi, S. Hong, A. Konwinski, Siddharth Murching, Tomas Nykodym, Paul Ogilvie, Mani Parkhe, Fen Xie, and Corey Zumar. Accelerating the machine learning life- cycle with MI.flo. IEFE Data Eng. Bull., 41:39-45, 2018. 14
Yuda Song, Yifei Zhou, Ayush Sekhari, Andrew Bagnell, Akshay Krishnamurthy, and Wen Sun. Hybrid rl: Using both offline and online data can make rl efficient, 2023. 6
Manzil Zaheer, Satwik Kottur, Siamak Ravanbhakhsh, Barn- abás Póczos, Ruslan Salakhutdinov, and Alexander J Smola. Deep sets. In Proceedings of the 31 Inter- national Conference on Neural Information Processing Systems, NIPS' 17, 2017. 5
Giles W Story, Ivo Vlaev, Ben Seymour, Ara Darzi, and Raymond J Dolan. Does temporal discounting explain unhealthy behavior? systematic review and reinforce- ment learning perspective. Frontiers in behavioral neuro- science, 8:76, 2014. g
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In International Conference Machine Learning, pages 27042-2705 PMLR, 2022.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018. 6
Linda M Tate, Pao-Feng Tsai, Reid D Landes, Mallikarjuna Rettiganti, and Leanne Lefler. Temporal discounting rates and their relation to exercise behavior in older adults.
Physiooogy & behavior, 152:295-299,
2015. g
Owen Thomas, Ritabrata Dutta, Jukka Corander, Samuel Kaski, and Michael U Gutmann. Likelihood- free inference by ratio estimation. arXiv preprint arXiv:1611.10242. 2016.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Rep- resentation learning with contrastive predictive coding. arXiv preprint 107.0:4807, 2018. 5 14
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszko- reit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008,
2017. 5
13

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
A. Algorithm
Algorithm Overview Step-DAD
Input: Generative model py 0) experimental budget T, refinement schedule T [= . Tk+t), T with 7o=0, TK+1=1 - training budgets {Na, )=
Output: Dataset br = {St,, h=)r
OFFLINE STAGE: BEFORE THE LIVE EXPERIMENT Set hn Ø.
while Computational budget does not exceed No do D Train fully-an aully-amortized my in Foster et (2021
end
ONLINE STAGE: DURING THE LIVE EXPERIMENT
for k = 1,..., K 1 do
for Tk-1 <ts do
Compute design Ee = TT- (h4-1)
Run experiment ε., observe outcome V/ Update the dataset h2 14-1 (Et, ve)
end
If k = K + 1 then return hy end
while Computational budget does not exceed N₂ do
Fit posterior p(0 h Fine-tune policy Tr. by optimizing 6
end end
B. Further EIG bounds
The sequential Nested Monte Carlo (sNMC) (Foster al. 2021 upper bound given by
t1h hre Bu, T))
1 log
(7)
I E/ plhrutier Dc,
which we use to evaluate different design strategies. For implicit models we can utilize the InfoNCE bound (van den Oord et al. 2018), which given by
exp(U 0o))
(IT U) = (((70)(( Epeen.)
log T.+1 En ( 0.))
(8)
(hT,
or the NWJ bound [guuyen et al. 2010), given by:
(T, U) := (p(0)p(hrr (.x)) [U -E [exp(U(kT,)) 0))]]
(9)
where in both bounds U is a learnt critie function, hr 0 IR
C. Experiment details
C.1. Computational resources
The experiments were conducted using Python and open-source tools. PyTorch (Paszke et al. 2019) and Pyro (Bingham eet al. 2018) were employed to implement all estimators and models. Additionally, MIFlow (Zaharia et al. (0188) was utilized for experiment tracking and management. Experiments were performed on two separate GPU servers, one with 4xGeForce RTX 3090 cards and 40 epu cores; the other one with 10xA40 and 52 cpu cores. Every experiment was ran on a single GPU.
14

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
C.2. Policy architecture
In the same vein as Foster et (202T) we leverage the permutation invariance of EIG in our BED problem settings allow
for more efficient network training and weight sharing. That is say represent histories of varying lengths by single fixed dimensional representations R(he) parameterized by encoder network Es, (inheriting notation from Foster et al.
(201
R(hx)= = Σ
Es, (5k: yk).
(10)
Our policy then becomes: (T = (02g2,((R)) (R(h,) where Føz is a decoder network. The specific architectures for the encoder and decoder networks are given in Sections C.5 and C.6
C.3. Evaluation details
When evaluating fully amortized policies, we employ the sPCE (6) lower bound and sNMC 7 upper bound using a large number of contrastive samples, L 100K, drawn from the prior to approximate the inner expectation. The outer expectation is approximated using N 2048 draws from the model p(0)p(hr 6, .). To approximate the total EIG quantity for Step-DAD efficiently, we
AI(n" ko) In (TT" In ( (mu)
(11)
-Ephh, (5o) (TT) (T²) I (T)]
(12)
2 (T))
che (-(1-(((7)) - 1 (T))
(13)
-
and add that difference to our corresponding lower/upper bound estimates of I1 (m)) (which can be directly estimated using the sPCE and sNMC bounds respectively). Here we are using a combination of an upper and lower bound is done to ensure any reported gains from StepDAD are conservatively underestimated, as well as reducing the variance in our estimates.
Note, found in Blau et al. (2022), the lower and upper bounds are less tight compared to other experiments. In this experiment, total EIG was simply calculated with the higher variance alternative as follows: (zo)+++ (pp [EIG)-
C.4. Baselines
Static The Static (fixed) baseline pre-selects a fixed EI ahead of the experiment before any observations. As in all cases, designs are optimized to maximise the EIG. This non-adaptive approach used PCE bound to optimize the design set Ei,...,Er Er and can be thought of treating the entire sequence of experiments as a single experiment (Foster et al. 2021).
Step-Static Step-Static computes a set of designs for Er before a posterior update and subsequent computation of designs ....., Êr. Each set of designs are selected following the static methodology outlined above.
Random As the name implies, this baseline selects random sample of designs Ei,..., Er. Thus the most non-informed naive approach.
C.5. Location Finding
The objective of the experiment is to ascertain the location, 8, of K sources. K is presumed to be predetermined. The intensity at each selected design choice, &, represents noisy observation log 0,g centered around the logarithm of the underlying model, p)
K
OK
((05E) = b+
(14)
(m + 1|0₂ - E|I)
A=
In the given context, may be either predetermined constants or random variables, b > 0 represents fixed background signal, and m is a constant representing maximum signal
logly 0, E] ~ N (log o²).
(15)
We assumed a normal standard prior at training: 0g Lid.
A0a,, Ta).
15

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
C.5.1. TRAINING DETAILS
Tables 6 and outline the architecture of the DAD policy network. The model hyperparameters used are outlined in Tables
and
10
Table 6: Source location finding, Encoder network E architecture as in Foster et al. 2002
Layer
Overview
Dimension
Activation
Design-outcome
E
3
H1
Fully connected
64
RELU
H2
Fully connected
256
RELU
Output
Fully Connected
16
Table 7: Source location finding, Decoder network F,., architecture as in Foster et al. 2021
Layer
Overview
Dimension
Activation
Input
E(he)
16
Fully connected
128
RELU RELU
Fully connected
16
Output
5
Table 8: Source location finding. Parameter Values
Parameter
Value
Org
for all
Max signal, m Base signal,
10 10-4
10-1
Observation noise scale,
0.5
Table 9: Source location finding, Parameters for training pre-training DAD/Step-DAD
Parameter
Value
Batch size
1024
Number of negative samples
1023
Number of gradient steps (default)
50K
Learning rate (LR)
0.0001
Table 10: Source location finding, Parameters for Step-DAD finetuning
Parameter
Value
Number of theta rollouts
16
Number of posterior samples
20K
Finetuning learning rate (LR))
0.0001
16

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
7.5
Total EIG
6.5
Thming step,
DAD (50K), upper Step-DAD (50K/0.5K), lower
DAD (10K), upper
Step DAD (0.50..K). lower
Figure 7: Sensitivity to training budget for location finding experiment. DAD policies are trained for 50K or 10K steps, Step-DAD policies are refined for 0.5K steps. Step-DAD outperforms its respective DAD baseline for all т. Errors show =ls.c.
Table 11: Source location finding. Total EIG for Step-DAD for various tuning steps Base DAD network pretrained for 50k steps before subsequent 2.5k finetuning steps at T.
T
Lower bound (1)
Upper bound (4)
(Worst)
7.050 (#: 0.103)
7.065 (= 0.103)
2
7.295 (# 0.119)
7.305 (= 0.120)
3
7.476 (#: 0.102)
7.482 (= 0.103)
4
7.676 (# 0.118)
7.680 (= 0.118)
5
7.536 (# 0.182)
7.540 (= 0.182)
6
7.759 (# 0.114)
7.765 + 0.114)
7
7.748 (# 0.091)
7.757 (= 0.091)
(Best)
7.877 (= 0.082)
7.892 (= 0.082)
9
7.623 (= 0.035)
7.652 (= 0.035)
DAD
7.040
(+ 0.012)
7.089 (= 0.013)
C.5.2. OPTIMAL
The optimal value for EIG occurs at range around [6, 7, (Table I1p. Figure deguonnnmmsssstttts a further ablation from Section 6. where now finetune the Step-DAD policies for 0.5K steps instead of 2.5K. We observe the same behavior for this reduced finetuning as before. Additionally, we note Step-DAD under the reduced 10K pretraining budget still outperforms DAD under the full 50K budget for r 6 6.
C.5.3. MULTIPLE SOURCES
As further ablation, we test the robustness of semi-amortizee approach to the more complex task of location finding with multiple sources of signal. We find positive EIG difference in all cases, once again demonstrating the benefits of using the semi-amortized Step-DAD network compared to the baseline fully amortized DAD network. Increasing the number of sources leads to reduction in the EIG difference. However, this is expected given the increasing complexity of the task compared to the fixed number of steps post to adjust the decision making policy in the semi amortized setting. All experiments were run with 7. Refer to Table 2 in main paper for results.
C.6. Hyperbolic Temporal Discounting Model
Building on Foster et al. (2021) Mazur (1987) and Vincent (2016), we consider hyperbolic temporal discounting model. A participant's 's behaviour is characterized by the latent variables (k: a) with prior distributions as follows:
log k NN-4.25, 1.5)
HalfNormal 2).
(16)
17

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
HalfNormal distribution denotes Normal distribution truncated O. For given k, a, the value of the propositions "ER today" and "E100 in D days" with design ç = -R, D) are given by:
100
Vo = R.
V = 1 1 kD
(17)
Participants select Vi in place of Vo with probability modelled as:
Ví - Vo
ply l|,; a, a, R, D) - -1-
22)
(18)
a
We fix 0.01 and ø is the c.d.f of the standard Normal Distribution
((2) = S
1 exp
= 22
(19)
-oc
V2T
2
As in Foster et al. (2021). the design parameters R, D have the constraints D > 0 and 0 < R < 100. R, D are represented in an unconstrained space Ed. &r and transformed using the below maps
D = exp(Ea)
R = 100 sigmoid(Ex).
(20)
Tables 12 and 13 outline the architecture of the DAD policy network. Tables 20 and 15 give the hyper-parameters for training the DAD DAD/Step-DAD policies for the hyperbolic temporal discounting model.
Table 12: Hyperbolic Temporal Discounting model. DAD encoder network.
Layer
Overview
Dimension
Activation
Design input
Er Fully connected
2
H1
256
Softplus
H2
Fully connected
256
Softplus
H3
Fully connected
16
H3'
Fully connected
16
Output
y C H3 + (1 y) o H3'
16
Table 13: Hyperbolic Temporal Discounting model. DAD decoder (emission) network
Layer Input HI
Overview
Dimension
Activation
R(h,)
16
Fully connected
256
Softplus Softplus
H2 Output
Fully connected
256
Ea, &
Table 14: Hyperbolic Temporal Discounting model. Parameters for training of the DAD network.
Parameter
Value
Batch size
1024
Number of negative samples
1023
Number of gradient steps (default)
100K
Learning rate (LR)
5 x 10
Annealing frequency
1K
Annealing factor
0.95
18

---

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design
Table 15: Hyperbolic Temporal Discounting model. Parameters for Step-DAD policy finetuning.
Parameter
Value
Number of theta rollouts
Number of posterior samples
20K
Finetuning learning rate (LR) 5 10-5
C.7. Constant Elasticity of Substitution (CES)
This experiment builds upon behavioral economic model in which participants compare the differing utility, (xx), of two presented baskets of goods x (Arrow et al. 1961). In this experiment E [0, 1000s represents non-negative quantities of three goods which together form the basket for evaluation.
The agent compares the two baskets æ and s' by evaluating their individual utility U(x) and subsequently indicating their preference using sliding scale ranging from 0 following the probabilistic model defined below. The latent variables governing this framework would in practice vary across individuals; representing their unique preferences. Thus the experimental design objective is to infer each of these latent parameters and therefore understand the individual preference model of the decision maker in question. The priors are defined as follows:
~ Beta 1,
(21)
a Dirichlet( [1, 1]) log (V,
(22)
(23)
(24)
The probabilistic model is expressed as:
1/0
U(x)= = Σ zfa;
(25)
14 = (U(x) U(r)))u
(26)
0n = (1+ + ||- ||)T
(27)
(Mn, o
(28)
У (simmsig (n),
(29)
The values of the soper-parametert used in the model are detailed in Table [6
Table 16: CES model. Hyper-parameter values.
Parameter Value
0.005
2 2-22 22
Table 17: CES model. DAD embedding layers. Outcomes and Designs are embedded and then concatenated before being passed into the encoder.
Layer
Overview
Dimension
Activation
Input
Eu or y
6 or 1
Output
Fully connected (Layer Norm)
32
ReLU
19

---

Step-DAD: Semi-Amortized Policy-Basce Baicy-Based Bayesian Experimental Design
14.0
Total ElG
12.0
10.0
2
3
6
7
8
TTuning step,
DAD (50K), upper
Step-DAD (50K/10 /10K), lower
DAD (50K), lower
Figure 8: Constant Elasticity of Substitution. Improvement in EIG by Step-DAD over DAD after fine-tuning, DAD is trained for 50K steps, with Step-DAD undergoing an additional 10K finetuning steps. Step-DAD (lower) outperforms the DAD baseline. DAD was initialized with designs from the Static methodology overcome challenges introduced by the accumulated probability mass at the boundaries following censoring.
Table 18: CES model. DAD encoder network.
Layer
Overview
Dimension
Activation
Input
Embedding
64
HI
Fully connected (Layer Norm)
128
ReLU
Output
Fully connected (Layer Norm)
32
ReLU
Table 19: CES model. DAD decoder (emission) network
Layer
Overview
Dimension
Activation
Input
Embedding
32
H1
Fully connected
128
ReLU
Output
Fully connected
Table 20: CES model. Parameters for training of the DAD network.
Parameter
Value
Batch size
100
Number of negative samples
1024
Number gradient steps (default)
50K
Learning rate (LR)
1 x 10
1
x 1
Annealing frequency
1K
Table 21: CES model. Parameters for Step-DAD policy finetuning.
Parameter
Value
Number of theta rollouts
16
Number of posterior samples
20K
Finetuning learning rate (LR)
1 x 10-5
20


Reviews:
[{"review_id": "QFdWAKVXmv", "reviewer": "Reviewer_W51C", "summary": "This paper introduces Step-DAD, a hybrid between traditional and fully amortized policy-based approaches to Bayesian Experimental Design (DAD) that retrains its policy on-line as it gathers additional observations. The authors argue that this allows to retain the benefits of policy-based BED while overcoming its key limitation of not being able to adapt the policy in response to data collected. They discuss training procedures (mostly based on existing technical contributions). Finally, Step-DAD is evaluated in a number of experiments against DAD and other baselines, demonstrating the expected performance improvements in terms of EIG improvements and robustness to model misspecification.\n\n## Update after rebuttal \nMy assessment remains unchanged - this is a good and well-executed paper and deserves to be accepted, but for me to champion the paper with a strong accept I would want to see some more novel ideas or some harder technical challenges solved.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you kindly for your helpful review. \n\n> The Step-DAD setup is reminiscent of Model Predictive Control (MPC)... This approach is sufficiently related that this connection warrants discussion in the paper.\n\nThank you for drawing this perceptive connection and we agree there are nice analogues between the approaches examined here and the Model Predictive Control literature. We will happily add references and discussion on this to the paper.\n\n\n>My main question that remains unanswered is how feasible this re-training of the policy is in practical settings\n\nThank you for raising this important and practical point. We will add results on precise wall-clock times and additional discussion on when we expect the retraining to be feasible, including some real-world example applications and practical issues with, e.g., running on embedded devices. In short, we find that helpful retuning can usually be achieved in as little as a minute, so the retraining can be done whenever such time can be justified between one or more experiment iterations.\n\nTo characterize things more precisely, we have conducted a new ablation of performance vs amount of time spent to refine the policy, the results of which can be found here: https://tinyurl.com/EIG-wall-time (anonymous). The charts demonstrate all computation is on the order of minutes and there is an initial increase in EIG with increasing wall time before a plateauing. Note that we implemented importance sampling and the varying percentage of inference corresponds to the number of samples drawn from the proposal (100% = 20,000 samples). For each row, the initial inference cost is a fixed cost and increasing finetuning steps is what increases the wall time.\n\n>Why not re-train the policy at each step?\n\nThere is absolutely nothing wrong with this if sufficient computational budget is available and one may well often do so in practice. The main reasons we did not do this in our experiments were a) to keep evaluation costs down (noting that we are running many seeds over lots of different possible true theta, which won’t be needed when deploying StepDAD in practice) and b) we tended to see quite quickly diminishing returns for conducting multiple retrainings in practice.  \n\nTo more explicitly demonstrate these diminishing returns in the paper itself, we have run a new ablation for the location finding experiment where we increase the number of times that the policy is refined. The results, which can be found here: https://tinyurl.com/eig-interventions (anonymous), demonstrate the diminishing returns behaviour, with Total EIG roughly flat beyond 2 intervention steps. Once again note that the varying percentage of inference corresponds to the number of samples generated from the importance sampling proposal (100% = 20,000 samples).\n\n> It would be interesting to further explore the notion of a \"compute arbitrage\" between investing into training a high quality policy offline vs. using a simpler policy network and investing less compute upfront and instead recover the loss in performance from that by re-training the policy online. Some of this is contained in the results of Fig 2, though that only considers the number of training steps, not the complexity of the policy network.\n\nThank you for the insightful suggestion. We hope that the new plots above provide further insight into this compute arbitrage, beyond what was already provided in Figure 2. Further investigating this trade-off with simpler policy networks instead of fewer training steps would certainly also be interesting and we will look into adding this. One important thing to note though, is that a noticeable part of the time to refine the network is in the inference rather than the network updates itself, so there will be limits to the gains which can be achieved from just simplifying the policy network itself.", "claims_and_evidence": "Claims are clear and supported by convincing evidence.", "methods_and_evaluation_criteria": "* The methods is a very natural evolution of DAD, and largely relies on existing methods and results used in a novel way.\n* The evaluation criteria make sense, and the empirical evaluation is relatively comprehensive and includes relevant and insightful ablations.", "theoretical_claims": "Yes, Prop 3.1 (the only claim) is straightforward.", "experimental_designs_or_analyses": "* The analyses all appear proper and valid.\n* The authors use a conservative estimate of the improvement in performance from Step-DAD over DAD, which lends additional credibility to the results.", "relation_to_broader_scientific_literature": "* Combining the \"best of both worlds\" from traditional and policy-based BED where \"online\" computation during the experiment is permissible is a very natural thing to do and makes a lot of sense.\n* The idea isn't exactly groundbreaking given the parallels with Step-Static policies or MPC (see below), but is executed solidly in the paper.\n* As this extension didn't require a lot of new technical innovations, the primary contribution of the work is therefore to set up the problem and demonstrate the performance in the empirical evaluation, which was executed well.\n* The Step-DAD setup is reminiscent of Model Predictive Control (MPC), which in a similar fashion re-optimizes a sequence of control inputs for the remaining experimental horizon based on a model of the process and the observations obtained so far. The main difference is that wile MPC optimizes input trajectories, Step-DAD optimizes a full policy - in that sense, MPC is the control-theory analogue of the Step-Static policy. This approach is sufficiently related that this connection warrants discussion in the paper.", "essential_references_not_discussed": "N/A (I have limited overview of the relevant literature)", "other_strengths_and_weaknesses": "* The proposed Step-DAD method is a natural evolution from traditional and fully amortized approaches. The empirical results convincingly demonstrate the benefits Step-DAD - while those are intuitive and expected, there is still substantial value in confirming this in a well-designed evaluation.\n* My main question that remains unanswered is how exactly how feasible this re-training of the policy is in practical settings.\n  * The examples provided don't have a notion of cost or duration of obtaining a measurement, and while the studies do evaluate different and multiple steps at which to re-train the policy, this seems somewhat artificial. Why not re-train the policy at each step? It would be highly illustrative and helpful to have some real-world examples with specific budget and time considerations to understand how Step-DAD would be used in practice.\n  * From a practical perspective, it appears substantially simpler to deploy a trained (and possibly inference-optimize) policy network to perform BED - e.g. this could easily be done on edge or embedded devices. It seems less straightforward to run a full training setup for re-training the policy network in such a setting. This doesn't make it infeasible or devalues the methodology more generally, but I feel these practical aspects should be discussed as well.", "other_comments_or_suggestions": "* It would be interesting to further explore the notion of a \"compute arbitrage\" between investing into training a high quality policy offline vs. using a simpler policy network and investing less compute upfront and instead recover the loss in performance from that by re-training the policy online. Some of this is contained in the results of Fig 2, though that only considers the number of training steps, not the complexity of the policy network.\n* In Section 6.1, you note that \"the performance advantage of Step-DAD over DAD appears to be most pronounced when fine-tuning occurs just past the midpoint of the experiment, that is for\nτ = 6,7 or 8.\". At this point this is a bit of reading the tea leaves - this is much clearer from the Hyperbolic Temporal Discounting example in Sec 6.3.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LuPH51Ch8P", "reviewer": "Reviewer_2XqQ", "summary": "The paper deals with Bayesian adaptive experimental design for identifying model parameters. Fully adaptive strategies are costly and myopic. Recent work has proposed amortized experimental design in which a neural net maps from observed data directly to the experimental design policy. That work, however, is not sufficiently adaptive as it cannot learn directly from the data in the current experiment. The paper introduces Step-DAD which attempts to blend the two. Essentially it is adopts the amortized neural network approach, but adds in fine tuning based on the results of the current experiment. This improves performance relative to the amortized approach and various static approaches.\n\n## update after rebuttal\nincreased score", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you kindly for your helpful review.\n\n\n>I expected to see a plot showing the trade-off between achieved EIG and fine-tuning wall-time\n\nThank you for the great suggestion. We have now implemented this analysis, and the chart illustrating this trade-off can be found here: https://tinyurl.com/EIG-wall-time (anonymous). We will incorporate these results into the final revised version of the paper. \n\nThe results show that significant improvements can be achieved with only a few minutes of fine-tuning time, with further tuning providing additional benefits but with diminishing returns. By contrast, traditional BED approaches often require test-time computation on the order of hours (Foster et al., 2021) and similarly training the original DAD network (50k steps) also was on the order of hours.\n\n \n\n> There was one claim that I did not see well-supported. The paper claims that the per-iteration cost of the fully adaptive methods is too time-consuming, and that the proposed Step-DAD is ideal for the \"many problems where we can afford to perform some test-time training during the experiment itself.\"\n\nWe would like to emphasize that the main benefit we expect from StepDAD over the traditional greedy approach is in the quality of the designs, rather than simply in terms of cost, as is demonstrated in our numerical results. While the above results show that there are also computational cost benefits as well, our primary motivation is still to give the best possible design performance when there is computational time available during the experiment. We will make edits to the paper to ensure this is clear.\n\n\n*Adam Foster, Desi R Ivanova, Ilyas Malik, and Tom Rainforth. Deep adaptive design: Amortizing sequential bayesian experimental design. Proceedings of the 38th International Conference on Machine Learning (ICML), PMLR 139, 2021*", "claims_and_evidence": "Generally the paper is well written and the claims are supported by clear and convincing evidence. There was one claim that I did not see well-supported. The paper claims that the per-iteration cost of the fully adaptive methods is too time-consuming, and that the proposed Step-DAD is ideal for the \"many problems where we can afford to perform some test-time training during the experiment itself.\"\n\nHowever, nowhere does the paper actually give results for how much test-time training is required by the method. This is given only in terms of number-of-fine-tuning iterations, as opposed to wall time (what actually matters for the framing of the problem). I expected to see a plot showing the trade-off between achieved EIG and fine-tuning wall-time, where with 0 fine-tuning wall time we would match the EIG of DAD, and then we could see how much wall time is required to improve significantly over that, and if indeed there are many problems where we can afford that much test-time training.\n\nIf this plot were to be added, I'd be more supportive of the paper; I think its absence is significant.", "methods_and_evaluation_criteria": "I have no concerns with the selection of baselines or problems.", "theoretical_claims": "Yes, Prop 3.1", "experimental_designs_or_analyses": "The evaluation was all reasonable to me.", "relation_to_broader_scientific_literature": "The framing with respect to past work was well described.", "essential_references_not_discussed": "Not aware of any.", "other_strengths_and_weaknesses": "The paper is well-written and everything is well motivated.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "What does the EIG vs. wall time trade-off look like as we increase the amount of fine tuning?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C4UJuSUkNe", "reviewer": "Reviewer_U4V6", "summary": "They propose a semi-amortized approach to Bayesian experimental design, in which a policy is learned offline (as in the standard fully amoritzed approach) and then is adapted online. This increases computational cost but results in more adaptive designs, which are more robust to model misspecification. The paper is extremely well written and the method is elegant and effective.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you kindly for your helpful review.\n\n> For the eval metric, have you considered using synthetic data where the ground truth theta* is available, and then assessing distance between E[theta|hT) and theta* where hT is from a particular design policy?\n\nYes! Assessing the distance between the posterior mean and the true parameter value is indeed a reasonable metric, and one we did consider. We ultimately chose to focus on using the EIG to assess performance as this is the metric most commonly used in the BED literature and it tends to be more robust than this distance (in particular, it is possible for the posterior mean to very close to theta* while still having significant posterior variance). Preliminary results with this metric were very similar to that of the EIG and we can add such comparisons if you think they are important. \n\n\n\n>How do you evaluate different designs if you just have access to an empirical dataset, and not a DGP (simulator)?\n\nIn general, this is very difficult as we are trying to evaluate the quality of a data gathering process, which cannot be directly done simply by having access to an existing dataset. In the scenario where we have some existing data and then want to gather more data, the most natural approach is to train a model to the data we already have, then use this model within an experimental design framework. Depending on context, problems like this can also sometimes be tackled using active learning or reinforcement learning methods, instead of experimental design ones.", "claims_and_evidence": "They derive a new algorithm and claim it results in improved information gain for a given number of experiments/ samples.\nThey empirically demonstrate this is true on two different datasets, Source Location Finding and Constant Elasticity of Substitution (CES).\nThey do careful ablation studies to show where the gains come from.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "This is an algorithms paper, so does not have new theory.", "experimental_designs_or_analyses": "Very solid.", "relation_to_broader_scientific_literature": "Authors propose a novel and useful variant of amortized BED in which they allow the agent to learn an amortized policy offline in the usual way, and then adapt it at run time in light of observed data, by sampling potential outcomes from the posterior predictive rather than the prior predictive. This is a very elegant and useful idea.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "- For the eval metric, have you considered using synthetic data where\nthe ground truth theta* is available, and then asseessing distance between\nE[theta|hT) and theta* where hT is from a particular design policy?\n\n- sec 6.3. How do you evaluate diffrent designs if you just have access\nto an empirical dataset, and not a DGP (simulator)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0t0vX3Gaun", "reviewer": "Reviewer_EbFb", "summary": "This paper proposes Step-DAD, a semi-amortized approach to Bayesian experimental design (BED) that extends the existing Deep Adaptive Design (DAD) framework. While DAD pre-trains a fixed policy network before experimentation, Step-DAD allows for test-time adaptation of this policy during deployment, periodically refining it based on accumulated experimental data.\n\nThe authors evaluate Step-DAD on three experimental design problems: source location finding, hyperbolic temporal discounting, and constant elasticity of substitution. Results consistently show that Step-DAD outperforms DAD and other BED baselines, including when dealing with prior misspecification and extended experimental horizons.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you kindly for your helpful review. \n\n>It would be nice to see a comparison with RL-based BED approaches such as the work by Blau et al. (2022). Given that the semi-amortized framework is architecture-agnostic, it would be valuable to see how Step-DAD principles could be applied to RL-BED methods.\n\nWe agree that this could be a nice addition if time and space allow. As you note, the StepDAD approach can be directly applied within RL-BED policy training frameworks, so these simply provide an alternative strategy for the policy refinement we are proposing, rather than a directly competing approach. The main reason we focused on direct policy training strategies in our experiments was that they generally allow the policy to be updated far more cheaply than using an RL-based approach, thus making them more suitable for deployment-time policy refinements where time is more critical than the original policy training. We will add further discussion on this to the paper.\n \n>It would be nice to also mention some other latest BED work in the related work section, where some of these approaches can benefit from the proposed semi-amortized framework:\n\nThank you for this excellent suggestion. We will update the related work section to include discussion of the additional papers you mention, several of which could benefit from or be extended by the proposed semi-amortized framework.\n\n> The paper lacks a systematic analysis for balancing additional computational cost against EIG gains.\n\nThank you for highlighting this. We agree that this analysis would benefit the paper and have added new results on the trade-off between cost and performance. Specifically, we have done two new ablations for the location finding experiment. The first is to vary the amount of compute spent on inference and fine-tuning. The results, which can be found here: https://tinyurl.com/EIG-wall-time (anonymous), show that meaningful gains over DAD can be achieved with only a couple of minutes of computation. The benefits improve as the fine-tuning budget is increased, before plateauing.\n\nThe second new ablation is to increase the number of times that the policy is refined through the experiment. The results can be found here: https://tinyurl.com/eig-interventions (anonymous). They show that there is an initial benefit to increasing the number of intervention steps before a plateauing once again.\n\n>For real-world applications where online policy updates might be computationally constrained, do you have recommendations for determining the optimal update schedule (when and how often to update)?\n\nWhile the optimal update schedule will vary between applications, our findings do provide some helpful guidance:\n- Our new results above show that there are often diminishing returns from conducting many policy refinements, so it will not usually be necessary to refine the policy at every step, in fact there is a general plateau beyond two interventions for the T=10 budget case. \n\n- Figures 2 and 4 suggest that refining the policy in the region of halfway to three quarters of the way through the experiment may have the biggest impact when only doing a single refinement.\n\n- Even small amounts of policy refinement (on the order of a couple of minutes) can be beneficial.\n\nWe also note that in many real applications, the update schedule may be directly dictated by the computate constraints, with little flexibility in the update schedule.  For example, if each experiment itself takes, ~10 minutes to run, this gives us a clear per-iteration budget that we can use while we wait for the next experiment to complete. We will add further discussion on these practical considerations.\n\n>The paper shows Step-DAD is more robust to prior misspecification. Have you investigated how it performs when other aspects of the model are misspecified (e.g., likelihood functions)?\n\nThank you for the great suggestions. We have not explicitly investigated robustness to likelihood misspecification. Like other amortized BED methods, Step-DAD remains sensitive in such cases, and understanding this susceptibility is an important direction for future work in BED literature. We believe that Step-DAD in its current form is less likely to help guard against likelihood misspecification than prior misspecification as the same likelihood is still used in the policy refinement, whereas the prior is replaced by the intermediary posterior (which may have corrected some of the issues of the original prior, such as if it is not sufficiently informative). \n \nInterestingly though, Step-DAD does open up avenues for future work in this direction.  For example, one could consider doing model checking before policy refinement, then training under a new likelihood if the data collected indicates the original one should be rejected. We feel it is beyond the scope of the current paper to fully investigate this, but we will add further discussion as it is certainly an interesting future direction.", "claims_and_evidence": "The paper's main claim that Step-DAD outperforms existing BED methods is well-supported by empirical evidence across multiple experimental settings, which are commonly used in BED literature. The authors present thorough comparisons with appropriate baselines using consistent evaluation metrics (EIG bounds).", "methods_and_evaluation_criteria": "The proposed methods are appropriate for the BED problem. Using EIG bounds as the evaluation metric aligns with the standard objective in BED literature. \n\nThe problems selected in the paper are appropriate, which are commonly used benchmarks in BED literature. The inclusion of additional evaluation scenarios such as prior misspecification strengthens the assessment of the method's practical utility.\n\nHowever, one minor concern is that all evaluations are on synthetic problems rather than real-world datasets, it would be nice to see one such case which can support the need of doing test-time training.", "theoretical_claims": "I've checked the theoretical claims, particularly Proposition 3.1 regarding the decomposition of total EIG. The proof appears sound, correctly leveraging the factorization of joint likelihoods and probabilities to establish the independence of policy optimality for later steps given the history at the intermediate point.", "experimental_designs_or_analyses": "The experimental designs are generally sound. However, it would be nice to see a comparison with RL-based BED approaches such as the work by Blau et al. (2022). Given that the semi-amortized framework is architecture-agnostic, it would be valuable to see how Step-DAD principles could be applied to RL-BED methods. This would provide a more comprehensive understanding of the benefits of semi-amortization across different policy learning paradigms.", "relation_to_broader_scientific_literature": "The paper properly situates itself within the BED literature, acknowledging connections to:\n\n* Traditional adaptive BED frameworks\n* Policy-based BED approaches\n* Reinforcement learning (RL) literature, particularly offline model-based RL", "essential_references_not_discussed": "It would be nice to also mention some other latest BED work in the related work section, where some of these approaches can benefit from the proposed semi-amortized framework:\n\n[1] Iollo, Jacopo, et al. \"Bayesian Experimental Design Via Contrastive Diffusions.\" ICLR.\n\n[2] Iollo, Jacopo, et al. \"PASOA-PArticle baSed Bayesian optimal adaptive design.\" ICML.\n\n[3] Huang, Daolang, et al. \"Amortized Bayesian Experimental Design for Decision-Making.\" Neurips.\n\n[4] Iqbal, Sahel, et al. \"Nesting Particle Filters for Experimental Design in Dynamical Systems.\" ICML.\n\n[5] Iqbal, Sahel, et al. \"Recursive Nested Filtering for Efficient Amortized Bayesian Experimental Design.\" Neurips BDU workshop.", "other_strengths_and_weaknesses": "Strengths:\n* The semi-amortized framework is a natural extension of existing BED approaches.\n* The empirical improvements are significant and consistent across problems.\n* The paper is well-written and easy to understand.\n\nWeaknesses:\n* The contribution is somewhat incremental, building directly on DAD without fundamental architectural innovations. \n* While the paper demonstrates robustness to prior misspecification, it lacks exploration of other types of model misspecification. Given that the framework can be used to address model misspecification, investigating additional forms of misspecification (e.g., likelihood function misspecification) would strengthen the paper's claims about robustness.\n* The paper lacks a systematic analysis for balancing additional computational cost against EIG gains. A metric or guideline for determining the optimal update schedule (number and timing of updates) would provide practical value for deployment.", "other_comments_or_suggestions": "The paper is very well written, I can't spot any typos or apparent mistakes.", "questions_for_authors": "* For real-world applications where online policy updates might be computationally constrained, do you have recommendations for determining the optimal update schedule (when and how often to update)?\n* The paper shows Step-DAD is more robust to prior misspecification. Have you investigated how it performs when other aspects of the model are misspecified (e.g., likelihood functions)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes Step-DAD, a semi-amortized approach to Bayesian experimental design (BED) that extends the existing Deep Adaptive Design (DAD) framework. While DAD pre-trains a fixed policy network before experimentation, Step-DAD allows for test-time adaptation of this policy during deployment, periodically refining it based on accumulated experimental data.  The reviewers largely agree that the paper's main claim, that Step-DAD outperforms existing BED methods, is well-supported by empirical evidence across multiple experimental settings, which are commonly used in BED literature.  Reviewer 2XqQ raised a concern that the paper proposes a more efficient method, but does not present wall-clock time for its experiments.  The authors have since provided the requested analysis, satisfying that reviewer's primary concern.  While the paper lacks an overt champion, all reviewers agree that the ideas are novel and worthy of publication.  I therefore recommend acceptance based on the clear consensus among reviewers.
Decision: Accept (poster)