Amortized Bayesian Experimental Design for Decision-Making

Summary

This paper proposes a method for decision-aware Bayesian experimental design, where the design is not optimized with respect to the most accurate posterior distribution of the latent parameters but rather with respect to the expected utility gain of the actual (down-stream) decision task.

Strengths

This is an innovative paper with high practical relevance. The proposed method appears sound and the corresponding neural networks well designed to suit the goal. Despite my questions and concerns (see below), I am positive about this paper overall and eager to increase my score should my points be addressed.

Weaknesses

- The presentation of p(y_Xi | h_t) between Eq 3 and 4 is partially unclear to me. From the definition, it seems this is not actually a distribution but a set of distributions. To me, then notation p(y_Xi | h_t) appears to be quite the abuse of notation because we cannot readily read this it as a single distribution. Can you perhaps think about a different notation that makes this easier to parse and understand? Relatedly, in Equation 4, it appears that we compute an expectation over p(y_Xi | h_t). But how do we compute an expectation over a set of distributions? I think I get what the authors do and want to imply but to me this notation doesn’t help in understanding it. - Equation 7: It seems we approximate the predictive distribution always by a Gaussian. I mean this of course works if the true underlying function is some kind of GP, but what if the true predictive distribution is far away from Gaussian? I don’t see this choice to be discussed properly so I consider it a weakness of this paper for now. - The discussion of training and inference time can only be found in the appendix. Specifically, training speed seems to be substantial, which of course makes sense for an amortized method. However, I don’t see any discussion for when the training actually amortizes. That is, how many BED tasks do we need to run at minimum before the total (training + “inference”) time of the new method becomes better than those of the competing methods. More generally, I think a discussion of speed should be more prominent in the paper. - 6.1 toy example was hard for me to understand at first. Is this just a standard BO task to find the point where the unknown function is maximal?

Questions

- In 4.1 Query set: How problematic is the fact that we randomly generate some designs from the design space. Doesn’t this mean we need a distribution over the design space? How can we obtain (or define) such a distribution in general? - In 4.1 Query set: 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? - Given that (non-decision aware) amortized BED methods exist, why are the benchmarks only comparing against non-amortized methods? I suggest to also add amortized methods to the benchmarks unless you can convince me that this is not sensible for some reason. - What is the scalability of the method in terms of all relevant dimensions, e.g., dimensionality of xi, y, a, etc? - Figure 4: you say that your method provides substantial gains, but at least on the scale in the figure, gains seem to be small. Can you clarify why you feel that the improvements are indeed “substantial gains”? - The method has quite a lot of components, I wonder which of the components is responsible for the improved results? For example, how relevant is it to consider non-myopic designs, i.e., how does the method perform when only trained in a myopic setup? Relatedly, are the alternative methods myopic or non-myopic?

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

4

Limitations

The paper discusses several limitations. I am missing a discussion on the initial overhead of training, which is usually substantial in amortized methods.

Summary

The paper looks at the problem of designing Bayesian optimal experiments taking into account the downstream decision making. At the core is a Transformer Neural Decision Process (TNDP) architecture that is trained to amortise the experimental design process whilst simultaneously inferring the optimal downstream decision.

Strengths

- Relevant and interesting topic: Downstream decision making is what ultimately matters, so taking this into account when designing experiments to collect data can result in more cost- and sample-efficient learning. - Motivation for the paper as well as clarity of writing are excellent. Contextualisation relative to prior work can be improved as outlined in the next section. - The proposed Transformer Neural Decision Process (TNDP) architecture is tailored to the BED problem, is well-explained and adds some novelty to the architectures typically used in the field.

Weaknesses

### Sections 2.2 & 3.2 and Lindley's decision-theoretic BED [1]: My main issue with the paper is the presentation of DUG and EDUG as novel. This framework was first formulated in [1], and is very well summarised in Section 1.3 of [2]. I strongly recommend the authors read that section, and present their Section 3.2 accordingly, acknowledging they follow Lindley, 1972. The questions/comments in the next 2 bullets are a consequence of this omission of literature. - Second paragraph of Sec 2.2: I am not sure how the predictive distribution $p(y | \xi, h_t)$ is defined. I would think it is $p(y | \xi, h_t) = \mathbb{E}_{p(\theta |h_t)} [p(y | \xi, \theta)]$. Whether or not you compute/approximate the posterior $p(\theta |h_t)$, or seek to directly approximate $p(y | \xi, h_t)$ (eg variationally), I think you should explicitly define what this quantity is. - I am not sure how the utility $u(y_\Xi, a)$ is defined. From a Bayesian decision-theoretic approach, the utility has to depend on the state of the world $\theta$, as well as the experiments $\xi$ you are going to perform (which I guess is implicit in $y_\Xi$). So shouldn't the "lowest level" utility be a function $u(y, \theta, \xi, a)$, which you then integrate over $p(\theta|h_t)$, to obtain $u(y, \xi, a) = \mathbb{E}_{p(\theta|h_t)} [u(y, \theta, \xi, a)]$, then take $\max$ wrt $a$, and finally integrate over the predictive $p(y |\xi, h_t)$ to obtain an expected utility, which can then act a design ranking criterion, as you do in Eq 4 and (cf Eq 2 in [2]). ### Related work: For a field that has such rich history and renewed interest from the ML community recently, the related works section is quite short and sparse on citations. Some areas that are missing include: - Decision-theoretic BED: as previously discussed, the general framework of utility-based BED was developed by Lindley (1972). - BED + RL: this work touches on some aspects of RL; It might be good to discuss relations recent works in the intersection such as [5] and [6] (in addition to those mentioned) - Decision-theoretic approaches in related fields such as Bayesian Optimisation, e.g. [7], [8] - Finally, I'm not too familiar with this line of literature, but more recent work around decision transformers---is there any relation between TNDP with works like [9] and [10]? ### Other: - Line 6: "most recent BED methods use amortised inference with a policy network" is not quite correct in the sense that no "real inference" (posterior updates on the parameters $\theta$) are performed. - Line 179: "to ensure the framework satisfied the permutation invariance property of sequential BED": not all BED problems are permutation invariant. For example, designing experiments for time series models (e.g SIR in [3] and [4]), permutation invariance does not hold. This aspect has been discussed in e.g. Section 3.3 of [3]. - Assuming you do want a permutation invariant architecture (most design problems fall in that category): by conditioning on $t$ as part of the global information (GI) set, I think you actually break that invariance. This is because encoding $(\xi, y)$ at time $t$ or at time $s$ will give you different outputs. As far as I can tell from Fig2b), $D_c$ does attend to GI. Could you please explain if that's the case or I have misunderstood something? ----- #### References [1] Lindley, D. V. (1972). Bayesian statistics: A review. Society for industrial and applied mathematics. [2] Chaloner, K., & Verdinelli, I. (1995). Bayesian experimental design: A review. Statistical science, 273-304. [3] Ivanova, D. R., Foster, A., Kleinegesse, S., Gutmann, M. U., & Rainforth, T. (2021). Implicit deep adaptive design: Policy-based experimental design without likelihoods. Advances in neural information processing systems, 34, 25785-25798. [4] Kleinegesse, S., & Gutmann, M. U. (2019, April). Efficient Bayesian experimental design for implicit models. In The 22nd International Conference on Artificial Intelligence and Statistics (pp. 476-485). PMLR. [5] Mehta, V., Paria, B., Schneider, J., Ermon, S., & Neiswanger, W. (2021). An experimental design perspective on model-based reinforcement learning. arXiv preprint arXiv:2112.05244. [6] Mehta, V., Char, I., Abbate, J., Conlin, R., Boyer, M., Ermon, S., ... & Neiswanger, W. (2022). Exploration via planning for information about the optimal trajectory. Advances in Neural Information Processing Systems, 35, 28761-28775. [7] Neiswanger, W., Yu, L., Zhao, S., Meng, C., & Ermon, S. (2022). Generalizing Bayesian optimization with decision-theoretic entropies. Advances in Neural Information Processing Systems, 35, 21016-21029. [8] Ivanova, D. R., Jennings, J., Rainforth, T., Zhang, C., & Foster, A. (2023, July). CO-BED: information-theoretic contextual optimization via Bayesian experimental design. In International Conference on Machine Learning (pp. 14445-14464). PMLR. [9] Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., ... & Mordatch, I. (2021). Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 15084-15097. [10] Zheng, Q., Zhang, A., & Grover, A. (2022, June). Online decision transformer. In international conference on machine learning (pp. 27042-27059). PMLR.

Questions

In addition to the questions raised in the Weaknesses section: 1. I think the main contribution of the paper is the TNDP architecture. Have the authors performed any ablations, e.g. not sharing the same embedding block? Not including $t$ in the GI? 2. In the decision-aware AL experiment: why does the random baseline perform as good as all the other ones? 3. Could you give guidance on choosing utility functions? For the experiments in the paper it is quite straightforward to define them, but in real-world practical application that might not be the case. This is the reason why the mutual information has become the de facto standard utility in BED.

Rating

6

Confidence

5

Soundness

2

Presentation

3

Contribution

3

Limitations

Some limitations of the work were outlined in the Discussion section of the paper. Regarding negative societal impact, the field of experimental design (which boils down to efficient data collection), generally warrants some discussion. The experiments presented in this paper mostly use synthetic data and do not have negative impact; the HPO experiment, which uses real data does not (directly) represent an application with negative impact. However, applying these methods in real-world applications, particularly if decisions directly affect humans, as in e.g. personalised medicine, could raise concerns around bias, fairness, explainability and privacy. I would suggest to the authors to add 1-2 sentences in their limitations section 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.

Summary

The paper proposes a transformer-based architecture for jointly sampling designs and decisions in Bayesian Experiment Design (BED) using a forward-looking criterion. The latter considers the improvement in maximum expected utility brought about by a new design-outcome pair, where the expectation is taken with respect to the predictive distribution of the model. The main innovation of the paper lies in the coupling between information gain and utility maximization in an amortized, transformer-based framework in the spirit of attentive neural processes. The performance of the new architecture is evaluated on a toy regression task and two more representative models, exhibiting stable performance gains over contender methods.

Strengths

- The paper is clearly written, the ideas and formulations are stringent and well-justified, overall making it easy to follow and a pleasure to read (with the exception of Section 4.1, see below). - The proposed architecture and training objectives are novel and seem to unlock both qualitative and quantitative improvements over existing methods. - The results indicate superior and stable performance of the proposed architecture on two interesting tasks, along a toy 1D GP model which seems to be a standard proof-of-concept task in the neural process (NP) literature.

Weaknesses

- Some notational confusion can be avoided by consistently using the notation $a_{1:t}$ to denote a sequence of $t$ elements and $a_t$ to denote the $t$-th element in the sequence. Currently, $h_t$ denotes a sequence, but, e.g., $y_t$ denotes an element, and then again $\theta_{1:L}$ also represents a sequence. Also, P4L126 is an abuse of notation with slightly confusing wording, such as “the predictive posterior distribution over all possible designs”, whereas the predictive distribution(s) are over future \textit{outcomes}. This is in no way different than the posterior predictive in Bayesian (non-linear or linear) regression, where the posterior predictive is conditioned on the training data set and the set of (unlabeled) predictors available at test time. Hence, I struggle to understand the need for the convoluted abuse of notation, but I may be missing something. Also section 4.1 suddenly starts using bold font for vectors, which was not the case in the preceding sections. - Figure 2 is not particularly informative for the data flow, as it does not clearly communicate weight sharing, input-output operations and dependencies (left panel); the right panel comes out of the blue and is not well explained (i.e., what are the elements on the “left” and on the “top”); the description below on P6 does indeed disambiguate the idea behind the construction of the masks, but I believe it is best when figures support and enhance the text and not vice versa. - 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. For instance, there is no need to reiterate established transformer-based computations, but it could be helpful to explicate the construction of the masks, the representation types (e.g., vectors, sequences of vectors,...?), and the precise partitioning of the components into keys, queries, and values. - 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 (e.g., as highlighted in the Related Work) make for suitable benchmarks, despite not optimizing for future decisions? - The topic of model misspecification is never mentioned in the paper, even though the comprehensive review paper [1] states that it remains a major unsolved issue in BED and in amortized Bayesian inference more generally [2]. I believe this should also be acknowledged in the current paper and the authors can potentially think about quantifying the impact of model misspecification in a small ablation study in the final version of the manuscript. I am happy to discuss these points with the authors and increase my score if they are addressed / clarified. [1] Rainforth, T., Foster, A., Ivanova, D. R., and Bickford Smith, F. (2024). Modern Bayesian 429 experimental design. Statistical Science, 39(1):100–114. [2] Schmitt, M., Bürkner, P. C., Köthe, U., & Radev, S. T. (2024). Detecting Model Misspecification in Amortized Bayesian Inference with Neural Networks: An Extended Investigation. arXiv preprint arXiv:2406.03154.

Questions

- Perhaps section 2 can be organized in a way to avoid singleton nested subsection (i.e., 2.1.1)? - P4L130: Isn’t there also an assumption that decision are optimal only if there is no model misspecification (i.e., that we are working with the posterior of the “true” model)? - Are there any practical disadvantages of assuming a diagonal Gaussian predictive distribution? Can complex models induce multimodal or highly correlated predictive distributions that?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors openly discuss the current limitations of their approach.

Summary

This paper tackles an important problem of designing experiments in a way that directly optimizes downstream decision-making tasks, going beyond just inferring parameters of interest. The authors make several valuable contributions: 1. They introduce the concept of Decision Utility Gain (DUG) to quantify how much an experimental design improves the expected utility of the downstream decision. 2. They propose a novel neural architecture called the Transformer Neural Decision Process (TNDP) that amortizes both the experimental design selection and the approximation of the predictive distribution needed for decision-making. This unified amortized framework is a key innovation. 3. The authors develop a non-myopic training objective that looks beyond just the immediate decision utility to account for effects of the current design on future rewards. 4. Empirically, they demonstrate TNDP's effectiveness over traditional methods on various tasks like active learning, hyperparameter optimization, showing it can find informative designs and make accurate downstream decisions. In summary, this work makes valuable conceptual and technical contributions to the area of Bayesian experimental design by pioneering decision-aware amortized methods. It opens up new research directions for further enhancing real-world decision-making via optimized experimental data acquisition.

Strengths

- The paper presents a novel problem formulation by introducing the concept of Decision Utility Gain (DUG), which shifts the focus of experimental design from reducing parameter uncertainty to directly optimizing downstream decision utility. This new perspective is a creative departure from traditional Bayesian experimental design (BED) approaches. - The application of amortized inference techniques to decision-aware experimental design can be considered an original contribution, as it represents a new domain for these methods beyond traditional BED. - The empirical evaluation is comprehensive, spanning diverse tasks such as active learning, hyperparameter optimization, and synthetic regression problems. The results demonstrate the consistent superiority of TNDP over traditional methods.

Weaknesses

- 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 experimental evaluation, while comprehensive, focuses primarily on synthetic and benchmark datasets. While these serve as important proof-of-concept demonstrations, the paper could benefit from including real-world case studies or applications to further validate the practical utility of the proposed framework. - While the amortized nature of TNDP is highlighted as a key advantage, 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.

Questions

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

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

- The authors mention the use of a basic REINFORCE algorithm for training the query head, which can lead to unstable training, especially in tasks with sparse reward signals. While they suggest the use of more advanced reinforcement learning methods as a potential solution, a more detailed discussion on the specific challenges faced during training and the trade-offs involved in selecting different RL algorithms would be beneficial. - The authors mention that their model is trained on a fixed-step length, assuming a finite horizon for the experimental design process. A discussion on the limitations of this assumption and the potential difficulties in extending their approach to infinite horizon or open-ended experimental scenarios would be valuable.