Amortized Bayesian Experimental Design for Decision-Making

Thank you for your positive comments and thoughtful questions. We address your remarks and questions below.

1. The presentation of p(y_Xi | h_t) between Eq 3 and 4 is partially unclear to me…

Thanks for the question. $p(y_\Xi | h_t)$ is a joint distribution and is well-defined as a stochastic process. Please refer to the global response for further clarification.

2. What if the true predictive distribution is far away from Gaussian?

That’s a good point. Please refer to the response for question 9 of reviewer WHLG.

3. The discussion of training and inference time can only be found in the appendix…I think a discussion of speed should be more prominent in the paper.

The key advantage of amortized methods lies in their efficiency during the inference phase. While the upfront training time can be substantial, this investment pays off when the model is used repeatedly for multiple BED tasks. To answer your specific question about when the training amortizes. For example, in our active learning experiment, we observed that traditional methods like DEIG require approximately 100 minutes for a complete BED task. In contrast, our model's total training time is around 10 hours, with the inference time being negligible. Therefore, after conducting more than 6 experiments during the deployment phase, the cumulative cost of our method becomes more efficient than the traditional method.

4. Toy example was hard for me to understand at first. Is this just a standard BO task?

Thanks for the question. The toy experiment is not a standard BO task but a special active learning task. Our goal is to estimate the value of an unknown function at a specific target point $x^*$ by actively querying points. Traditional active learning methods are not target-oriented and typically query points based on the overall uncertainty of the function, which is not optimal for estimating the value at $x^*$. Our method, however, considers the downstream task, meaning our ultimate goal is to perform regression at the target location $x^*$. Thus, our decision-making-aware policy will strategically query points, specifically around $x^*$, as shown in Figure 3(a).

5. How problematic is the fact that we randomly generate some designs from the design space?

That’s a good question. For most problems, we usually either have some prior knowledge about the design space or the candidate designs come from a predefined finite set, as is commonly practiced in the pool-based active learning literature [3]. The more prior knowledge we have about the problem, the better the candidate designs we can generate.

6. You say that in the deployment phase we can obtain the optimal design by optimizing the models (which model’s?) output. How do you optimize this exactly?

We apologize for the unclear description. For continuous design spaces, we can optimize $\xi^{(q)}$ to maximize the query head's output $\mathbf{q}$, thereby obtaining the design with the highest probability. This can be achieved using optimization techniques such as gradient ascent.

7. Given that (non-decision aware) amortized BED methods exist, why are the benchmarks only comparing against non-amortized methods?

We have included a new amortized method baseline in the top-$k$ experiments and also explained why we didn’t compare with any amortized BED method previously. Please refer to the response for question 5 of reviewer WHLG.

8. What is the scalability of the method in terms of all relevant dimensions?

We have validated our method's effectiveness on high-dimensional design spaces, e.g., in the top-$k$ optimization task, $d_x$ for ranger is 9, and 16 for xgboost. Regarding the output dims, our paper focuses on single-output cases, the same as most BED works. However, our architecture can be readily generalized to multidimensional outputs by adding separate prediction heads for different outputs, as demonstrated in CNP literature [4, 5]. Lastly, for the action dims, we believe this does not significantly affect the performance since the action is made after the model's predictions.

9. Figure 4: Can you clarify why you feel that the improvements are indeed “substantial gains”?

We acknowledge that the term "substantial" may have been an overstatement, we will revise the wording. However, it is clear that our method outperforms the other baselines, particularly within the first 10 queries.

10. I wonder which of the components is responsible for the improved results? How does the method perform when only trained in a myopic setup? Relatedly, are the alternative methods myopic or non-myopic?

The improved results of our method are primarily due to our query head. This module amortizes the experimental design process, which learns common features from different tasks, allowing it to propose more valuable experiments. We conducted an ablation study in Appendix F.3 to verify the performance improvement brought by the query head.

Regarding the difference in performance between non-myopic and myopic objectives, we conducted an additional ablation study, as detailed in our global response. All the alternative methods we considered in our paper are myopic.

11. I am missing a discussion on the initial overhead of training.

Thanks for your suggestion. Indeed, like all amortized approaches, our method requires a large amount of data and upfront training time to develop a reliable model. However, once the model is trained, it offers long-term benefits of faster inference. We will discuss this trade-off in the revised manuscript.

References

[1] Garnelo et al. (2018). Conditional neural processes. ICML.

[2] Müller et al. (2021). Transformers can do Bayesian inference. ICLR.

[3] Settles (2012). Active learning. Springer.

[4] Markou et al. (2022). Practical conditional neural processes via tractable dependent predictions. ICLR.

[5] Bruinsma et al. (2023). Autoregressive conditional neural processes. ICLR.

Thank you for your detailed review and the valuable references you provided. We address your questions and points raised below.

Weaknesses

1. My main issue with the paper is the presentation of DUG and EDUG as novel.

We greatly appreciate your provided references and insights. We will include a discussion in Sections 2.2 and 3.2 about the connections and distinctions between our work and the decision-theoretic BED frameworks mentioned in [1, 2]. For further details, please refer to the global response and our answers to the following two questions.

2. I am not sure how the predictive distribution $p(y|\xi, h_t)$ is defined… I think you should explicitly define what this quantity is.

Yes, $p(y|\xi, h_t) = \mathbb{E}_{p(\theta|h_t)}[p(y|\xi, \theta, h_t)]$. It is obtained by marginalizing the posterior distribution of the parameters. In our architecture, we directly approximate the predictive distribution with neural processes, thereby bypassing the need to define and approximate the underlying parameters. We will add a description of the predictive distribution in P4L126 to clarify this point.

3. I am not sure how the utility $u(y_{\Xi}, a)$ is defined.

Our utility function is defined in terms of outcomes, which slightly differs from the traditional definition. The distribution of outcomes is obtained by (implicitly) marginalizing over $\theta$. A similar decision-theoretical setup can be found in [3]. As mentioned in the global response, the major difference between EDUG and Equation 2 in [2] is that we include an additional expectation over $p(y_\Xi | h_t)$, as the optimal Bayes action in our setup is based on the predictions of outcomes.

4. The related works section is quite short and sparse on citations.

Thanks for your references, we will add them. Regarding the differences with decision transformers, while both our architecture and theirs are based on transformers, their "decision" is similar to the step of experimental design in our context. However, we additionally amortize the downstream decision-making process, making the learning process more challenging.

5. Line 6: "most recent BED methods use amortised inference with a policy network" is not quite correct in the sense that no "real inference" are performed.

Agreed. We will change to “most recent BED methods leverage an amortized policy network to rapidly design experiments”.

6. Line 179: "to ensure the framework satisfied the permutation invariance property of sequential BED": not all BED problems are permutation invariant.

Agreed, we acknowledge that permutation invariance was meant more as an assumption than an inherent property of BED. We will fix it and add that for tasks where permutation invariance does not hold, our model can be easily adapted by incorporating positional encoding to add sequential/temporal information to the design.

7. Assuming you do want a permutation invariant architecture… I think you actually break that invariance… Could you please explain if that's the case or I have misunderstood something?

That is a misunderstanding. Permutation invariance here only means that changing the order of the historical data does not affect the design of the next experiment. Our architecture satisfies that. (Note that even without the timestamp information, the encoding as such would give different outputs at different points of time.)

Questions

8. Have the authors performed any ablations?

We conducted an experiment in Appendix F.3 to verify the effectiveness of the query head. Additionally, we have added a new set of ablation experiments on the discount factor; please refer to the global response for details. Regarding the inclusion of time $t$, we performed ablation studies across different dimensions. The results were inconsistent. Since we only had time to run our algorithm once for each dimension, we cannot draw definitive conclusions. We will further investigate the role of $t$.

9. In the decision-aware AL experiment: why does the random baseline perform as good as all the other ones?

Because methods like US and DUS are traditional active learning methods which do not take the decision-making problem into account. Hence there is no guarantee that the queries made would improve the quality of the decision-making. However, we can see that DEIG actually performs better than the random baseline, as it considers the decision-making problem by taking into account the posterior distribution of the optimal decision for the target point.

10. Could you give guidance on choosing utility functions?

Our utility function comes directly from the downstream decision-making task. We suspect the popularity of mutual information arises from an unwillingness to commit to a specific task, which makes sense if the task is not known. However, if we know our downstream task well, such as in drug design where we aim to maximize the efficacy of the drug while minimizing its risks, we can leverage input from domain experts to define our objectives. We will comment on this in the Discussion.

Limitations

11. I would suggest to the authors to add 1-2 sentences to acknowledge 1) the synthetic or semi-synthetic nature of the experiments, and 2) potential concerns that might arise when applying their method in real-world applications.

Thanks. We conducted a new retrosynthesis planning experiment with real-world data; please refer to the global response. We will also discuss the potential negative societal impact, especially when experiments or decisions directly affect humans, in the limitations section.

References

[1] Lindley, D. V. (1972). Bayesian statistics: A review. Society for industrial and applied mathematics.

[2] Chaloner et al. (1995). Bayesian experimental design: A review. Statistical science.

[3] Kuśmierczyk et al. (2019). Variational Bayesian decision-making for continuous utilities. Neurips.

Thank you for your detailed review and the valuable comments. We address your remarks and questions below.

Weaknesses

1. Some notational confusion can be avoided…

Thanks, we will improve the notations according to your suggestions in the revised paper. Specifically, we will replace $h_t$ with $h_{1:t}$ to better denote a sequence of elements. And we will consistently use bold font for vectors across the whole paper.

2. P4L126 is an abuse of notation with slightly confusing wording, such as “the predictive posterior distribution over all possible designs”.

We apologize for the confusing wording and the lack of a clear explanation of the predictive distribution. Please refer to the global response for a detailed answer.

3. Figure 2 is not particularly informative for the data flow ... the right panel comes out of the blue and is not well explained.

We will improve Figure 2 to better help readers understand our work. Specifically, we will use distinct modules to more clearly visualize the data embedding block (e.g., designs in different sets share the same embedder). We will also refine the Transformer block to highlight the attention mechanisms between different sets. For the causal mask on the right panel, we will remove the elements with complex notation on the left and top, replacing them with set names for better clarity.

4. Overall, I feel that Section 4.1 is the weakest link in the paper, and I believe the authors can think about optimizing the ratio of details dispersed between the main text and the appendix.

Thank you, we will revise. To address your specific points, we will provide a more detailed explanation of the construction of the masks, along with the improvements to the masks in Figure 2 as mentioned above. As stated in our response to the first question, we will improve the notation of some elements to enhance reader comprehension. Finally, we will clarify the elements involved in the attention mechanism. Specifically, we concatenate all sets together and use self-attention to obtain a single attention matrix, then use masks to determine the dependencies between different sets.

5. According to my understanding, none of the contender methods in the experiments is an amortized method. Wouldn’t some of the existing amortized BED methods make for suitable benchmarks?

We agree that it is important to compare with an amortized method. We added PFNs4BO [1], an amortized optimization framework, as a new baseline for our top-$k$ experiments. Please refer to the global response for the experiment results.

As for why we did not choose to compare with any amortized BED method before: Amortized BED mainly refers to DAD [2] and its subsequent works. These methods cannot be directly extended to our tasks for two reasons. First, DAD is a deterministic policy suitable only for continuous design spaces, whereas our tasks such as decision-aware active learning involve discrete data with covariates and their decisions. Second, even though there are subsequent works [3] extending DAD to discrete design spaces, these methods require an additional model or post-processing to make the final decision, making a direct comparison with our method inappropriate without significant alterations to DAD-based frameworks.

6. The topic of model misspecification is never mentioned in the paper.

That’s a very good point. We fully agree on the importance of model misspecification which has received significant attention recently in some areas, such as Simulation-Based Inference, but it has not been sufficiently studied in BED (beyond the handful of papers cited in [4]). Most BED works assume that the models are well-specified. In our work, model misspecification can indeed occur during the inference stage. Additionally, the utility might also shift during the deployment phase, potentially influencing the model's performance. In addition to detecting model misspecification, tackling the problem of robust experimental design under model misspecification would also be very interesting. We will add this discussion and outline it as an important area for future work.

Questions

7. Perhaps section 2 can be organized in a way to avoid singleton nested subsection?

Agreed; will do.

8. P4L130: Isn’t there also an assumption that decision are optimal only if there is no model misspecification?

Yes, it is a basic common assumption in the field of BED [1, 2]. We will explicitly state this assumption in the revised manuscript.

9. Are there any practical disadvantages of assuming a diagonal Gaussian predictive distribution? Can complex models induce multimodal or highly correlated predictive distributions?

That’s a good question. Our TNDP follows the common practice in the neural processes literature [5] of using independent (diagonal) Gaussian likelihoods. If modeling correlations between points is crucial for the downstream task, we can replace the output with a joint multivariate normal distribution (similar to GNP [6]) or predict the output autoregressively (similar to AR-CNP [7]). For modelling multimodal predictive distributions, we could replace the Gaussian head with a mixture-of-Gaussians head. These modifications can be easily implemented in the TNDP architecture. We will mention this in the Discussion.

References

[1] Müller et al. (2023). Pfns4bo: In-context learning for Bayesian optimization. ICML.

[2] Foster et al. (2021). Deep adaptive design: Amortizing sequential Bayesian experimental design. ICML.

[3] Blau et al. (2022). Optimizing sequential experimental design with deep reinforcement learning. ICML.

[4] Rainforth et al. (2024). Modern Bayesian experimental design. Statistical Science.

[5] Garnelo et al. (2018). Conditional neural processes. ICML.

[6] Markou et al. (2022). Practical conditional neural processes via tractable dependent predictions. ICLR.

[7] Bruinsma et al. (2023). Autoregressive conditional neural processes. ICLR.

Thank you for your positive assessment of our work and the points you raised. In the following we address your questions and points raised.

Weaknesses

1. The authors could provide a more rigorous analysis of the properties and characteristics of the TNDP architecture, such as its convergence behavior, sample complexity, and theoretical guarantees (if any) regarding the quality of the proposed designs and decisions.

The primary purpose of this paper is to propose a new amortized BED framework targeted at downstream decision-making. We leverage the established REINFORCE algorithm to aid in model training. Since we did not introduce a new RL algorithm, convergence analysis can be referenced from existing works, such as the analysis provided in [1].

Regarding sample complexity, for the toy example and active learning case, we generate synthetic data online. Specifically, we spent 50,000 steps to train our model and generated 16 synthetic datasets for each step. We will include these additional details in the appendix. For top-$k$ optimization, our model is trained using the publicly available HPO-B dataset, and detailed information about this dataset is already provided in the appendix.

2. The paper could benefit from including real-world case studies or applications to further validate the practical utility of the proposed framework.

We conducted a new retrosynthesis planning experiment based on real-world molecule data, please refer to the global response.

3. The paper could provide a more detailed analysis of the computational complexity and scalability of the proposed approach. (This analysis could include factors such as the training time required for different problem sizes, the memory footprint, and the scalability of the attention mechanisms used in the Transformer architecture.)

Thanks, we will add the following discussion in the revised manuscript. We have already included the rough total training time in the appendix. We will further provide descriptions of the training time and memory footprint required for each task. For example, for active learning experiments, a Tesla V100 GPU with 32GB memory was used, with an average memory consumption of 8 GB. The training time took about 10 GPU hours. Regarding scalability, our architecture is based on the Transformer, which suffers from quadratic complexity with respect to the input sequence length. This can become a bottleneck when the query set is very large. We will also include a further discussion of this issue and potential optimizations in the revised manuscript.

Questions

4. Can the authors provide a more in-depth theoretical analysis of the Decision Utility Gain (DUG) concept, including its relationship with existing concepts like Value of Information (VoI) or Information Gain (IG)?

IG can be regarded as a special case of DUG. First, we define $\mathcal{P}(\theta)$ as a set of distributions that we assume contains the true posterior distribution $p(\theta|h_t)$. We then define the decision space as $\mathcal{A} = \mathcal{P}(\theta)$, and the utility function as $\log a$, where $a \in \mathcal{P}(\theta)$. The optimal action in this case will be $a^* = p(\theta|h_t)$ based on the definition of entropy. DUG can be reinterpreted as the entropy reduction when we observe a new design pair $\\{\xi, y\\}$, which corresponds to the definition of IG.

5. Have the authors explored the sensitivity of TNDP's performance to different hyperparameter choices, such as the discount factor α used in the non-myopic objective? If so, can they share insights into this analysis?

We added an ablation study regarding the choice of the discount factor, please refer to the global response.

Limitations

6. A more detailed discussion on the specific challenges faced during training and the trade-offs involved in selecting different RL algorithms would be beneficial.

Thanks for your suggestion. In our experiments, the REINFORCE algorithm has proven sufficient for training an effective model. However, for more complex problems in the future, REINFORCE may be prone to high variance in policy gradient estimates. If needed, RL training can be improved by using more advanced algorithms like PPO [2], but the trade-offs include introducing more hyperparameters, such as the clip ratio, and increased computational cost, which requires more tuning of the system. Importantly, our work shows that our method can be trained effectively without needing complex, ad hoc RL techniques. We will expand the discussion on these aspects in the revised manuscript.

7. The authors mention that their model is trained on a fixed-step length… A discussion on the limitations of this assumption and the potential difficulties in extending their approach to the infinite horizon or open-ended experimental scenarios would be valuable.

The finite horizon assumption is sufficient for most BED problems, as we usually operate with a limited budget for experimental designs. However, for more complex BED problems, such as long-term medical trials, extending our approach to an infinite horizon setting could be valuable. Potential challenges include increased instability during training and higher computational costs. We will expand the discussion on these aspects in the revised manuscript.

References

[1] Zhang et al. (2021). Sample efficient reinforcement learning with REINFORCE. AAAI.

[2] Schulman et al. (2017). Proximal policy optimization algorithms. arXiv:1707.06347.

Thank you for bringing this paper to our attention. While both works use a similar Transformer neural process-based architecture, the main distinction is that our work focuses on general downstream utility functions and decision-making tasks, while the meta-BO paper focuses specifically on BO objectives. Moreover, our work introduces a coupling between query and prediction heads in the training objective, where the prediction head helps compute rewards to guide the query head, providing a more integrated architecture for decision-making tasks. Thanks again -- we will cite this paper and discuss the connections in the Related Work section of the final version of our paper.