Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design

Thank you kindly for your helpful review.

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.

We 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.

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:

Thank 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.

The paper lacks a systematic analysis for balancing additional computational cost against EIG gains.

Thank 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.

The 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.

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)?

While the optimal update schedule will vary between applications, our findings do provide some helpful guidance:

• 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.

• 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.

• Even small amounts of policy refinement (on the order of a couple of minutes) can be beneficial.

We 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.

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)?

Thank 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).

Interestingly 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.

Thank you kindly for your helpful review.

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?

Yes! 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.

How do you evaluate different designs if you just have access to an empirical dataset, and not a DGP (simulator)?

In 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.

Thank you kindly for your helpful review.

I expected to see a plot showing the trade-off between achieved EIG and fine-tuning wall-time

Thank 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.

The 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.

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."

We 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.

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

Thank you kindly for your helpful review.

The Step-DAD setup is reminiscent of Model Predictive Control (MPC)... This approach is sufficiently related that this connection warrants discussion in the paper.

Thank 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.

My main question that remains unanswered is how feasible this re-training of the policy is in practical settings

Thank 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.

To 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.

Why not re-train the policy at each step?

There 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.

To 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).

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.

Thank 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.