Amortized Bayesian Experimental Design for Decision-Making
Authors: Daolang Huang
Yujia Guo
Luigi Acerbi
Samuel Kaski
Venue: NeurIPS 2024 poster (2024)
Keywords: Bayesian experimental design
amortized inference
Bayesian decision theory
neural processes
Primary area: probabilistic_methods
TL;DR: We introduce a decision-aware amortized Bayesian experimental design framework with a novel Transformer neural decision process architecture to optimize experimental designs for better decision-making.

Abstract: Many critical decisions, such as personalized medical diagnoses and product pricing, are made based on insights gained from designing, observing, and analyzing a series of experiments. This highlights the crucial role of experimental design, which goes beyond merely collecting information on system parameters as in traditional Bayesian experimental design (BED), but also plays a key part in facilitating downstream decision-making. Most recent BED methods use an amortized policy network to rapidly design experiments. However, the information gathered through these methods is suboptimal for down-the-line decision-making, as the experiments are not inherently designed with downstream objectives in mind. In this paper, we present an amortized decision-aware BED framework that prioritizes maximizing downstream decision utility. We introduce a novel architecture, the Transformer Neural Decision Process (TNDP), capable of instantly proposing the next experimental design, whilst inferring the downstream decision, thus effectively amortizing both tasks within a unified workflow. We demonstrate the performance of our method across several tasks, showing that it can deliver informative designs and facilitate accurate decision-making.

Amortized Bayesian Experimental Design for Decision-Making
Daolang Huang
Yujia Guo Aalto University yujia.guo@aale gww.gauoooaaoooooomm fi
Aalto University daolang. www.gggaannn.o.....o
Luigi Acerbi University of Helsinki luigi acerbi@heeehhhll
Samuel Kaski
Aalto University University of Manchester samuel. kaski Caalto fi
Abstract
Many critical decisions, such personalized medical diagnoses and product pric- ing, are made hased on insights gained from designing, observing, and analyzing series of experiments. This highlights the crucial role of experimental design, which goes beyond merely collecting information on system parameters as in traditional Bayesian experimental design (BED), but also plays key part in facili- tating downstream decision-maie maainn.. Most recent BED methods use an amortized policy network rapidly design experiments. However, the information gathered through these methods is suboptimal for down-the worn-thh-linee line decision maainn, as the experiments are not inherently designed with downstream objectives in mind. In this paper, we present amortized decision-aware BED framework that prioritizes maximizing downstream decision utility. We introduce a novel architecture, the Transformer Neural Decision Process (TNDP), capable of instantly proposing the next experimental design, whilst inferring the downstream decision, thus effectively amortizing both tasks within unified workflow. We demonstrate the performance our method across several tasks, showing that can deliver informative designs and facilitate accurate decision-making
Introduction
In wide array of disciplines, from clinical trials (Cheng and Shen, 2005) to medical imaging (Burger al.. 2021), fundamental challenge is the design experiments infer the distribution some unobservable, unknown properties of the systems under study. Bayesian Experimental Design (BED) (Lindley, 1956; Chaloner and Verdinelli, 1995; Ryan et al., 2016; Rainforth al., 2024) is powerful framework in this context, guiding and optimizing experimental design by maximizing the expected amount of information about parameters gained from experiments, see Fig. (a)). However, the ultimate goal in many tasks extends beyond parameter inference to inform a downstream decision- making task by leveraging our understanding of these parameters. For example, personalized medical diagnostics, model is built based on historical data to facilitate an optimal treatment for a new patient (Bica al., 2021). This data typically comprises patient covariates, administered treatments, and observed outcomes. Since the query of such data tends to be expensive due to, e.g., privacy issues, we need to actively design queries to optimize resource utilization. Here, when the goal is to improve decisions, the strategy of experimental designs should not focus solely on inferring the parameters the model, but rather on guiding the final decision decision-making for the new patient, to ensure that each query contributes maximally to the diagnostic decision.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Design
EA, Posterior
OFFLEse
Optinise
Inference
?
Traditional
Ducssdon
(a)
Feetback
OFFILIN trwining stage
ThDP
Update history
Policy
Gemerata
Design
3eserve
Outcome
Our decision-aware anortized BED
Observe
Design
Outcone
(c)
Anortized BED
(b)
Figure 1: Overview of BED workflows. (a) Traditional BED, which iterates between optimizing designs, running experiments, and updating the model via Bayesian inference. (b) Amortized BED, which uses a policy network for rapid experimental design generation. (c) Our decision-aware amortized BED integrates decision utility in training to facilitate downstream decision-making
Traditional BED methods (Rainforth et al., 2018; Foster et al., 2019, 2020; Kleinegesse and Gutmann, 2020) do not take down-the-line decision-making tasks into account during the experimental design phase. As a result, inference and decision-making processes are carried out separately, which is suboptimal for decision-making in scenarios where experiments can be adaptively designed. Loss- calibrated inference. which was originally introduced by Lacoste-Julien et al. (Lacoste-Julier al., 2011) for variational approximations in Bayesian inference, provides a concept that adjusts the inference process to capture posterior regions critical for decision maaiing tasks. Rather than focusing on parameter estimation, the idea is directly maximize the expected downstream utility, recognizing that accurate decision-matee all regions of the posterior contribute equally to the downstream task. Inspired by this concept, we could consider integrating decision -making directly into the experimental design process to align proposed designs more closely with the ultimate decision-making task.
can proceed without exact knowledge of the posterior as
To pick the next optimal design, standard BED methods require estimating and optimizing the expected information gain (EIG) over the design space, which can be extremely consuminn.. This limitation has led to the development of amortized BED (Foster et al., 2021; Ivanova et al., 2021; Blau al., 2022, 2023), a policy-based method which leverages neural network policy trained on simulated experimental trajectories to quickly generate designs, as illustrated in Fig. (bb). Given an experiment history, these policy-based methods determine the next experimental design through single forward pass, significantly speeding up the design process. Another significant advantage of amortized BED methods is their capacity to extract and utilize unstructured domain knowledge embedded in historical data. Unlike traditional methods, which never reuse the information from past data, amortized methods can integrate this knowledge refine and improve design strategies for new experiments. In our problem setting, the benefits of amortization are also valuable where decisions must be made swiftly, such as when determining optimal treatment for patients in urgent settings.
In this paper, we propose amortized decision-mant making- aware BED framework, see Fig. I(c). We identify two key aspects where previous amortized BED methods fall short when applied down- stream dicision-makingg tasks. First, the training objective of the existing methods does not consider downstream decision tasks. Therefore, we introduce the concept of Decision Utility Gain (DUG) to guide experimental design to better align with the downstream objective. DUG is designed to measure the improvement in the maximum expected utility derived from the new experiment. Second, to obtain the optimal decision, we still need to explicitly approximate the predictive distribution of the outcomes to estimate the utility. Current amortized methods learn this distribution only implicitly and therefore require extra computation for the eeciisionmmainng process. To address this, propose

---

novel Transformer neural decision process (TNDP) architecture with dual output heads: one acting as a policy network to propose new designs. and another to approximate the predictive distribution to support the downstream decision-making This setup allows an iterative approach where the system autonomously proposes informative experimental designs and makes optimal final decisions. Finally, since our ultimate goal is to make optimal decisions at the final stage, which may involve multiple experiments, it is crucial that our experimental designs are not myopic or overly greedy by only maximizing next-step decision utility. Therefore, we develop non-myopic objective function that ensures decisions are made with comprehensive consideration of future outcomes.
Contributions. In summary. the contributions of our work include:
We propose the concept of decision utility gain (DUG) for guiding the next experimental design with direct focus on optimizing down-the-line decision-making tasks.
We present novel architecture Transformer neural decision process (TNDP) designed to amortize the experimental designs directly for downstream decision-mant -aaiin..
We empirically show the effectiveness of TNDP across variety of experimental design tasks involving decision-making, where it significantly outperforms other methods that do not consider downstream decisions.
2 Preliminaries
2.1 Bayesian experimental design
BED (Lindley, 1956; Ryan et al., 2016; Rainforth et al., 2024) is a powerful statistical framework that optimizes the experimental design process. The primary goal of BED is to sequentially select a set of experimental designs & E\t and gather outcomes y, to maximize the amount of information obtained about the parameters of interest, denoted as 0. Essentially, BED seeks to minimize the entropy of the posterior distribution of 0 or, equivalently, to maximize the information that the experimental outcomes provide about 6.
At the heart of BED lies the concept of Expected Information Gain (EIG), which quantifies the value of different experimental designs based on how much they are expected reduce uncertainty about 0, measured in terms of expected entropy (H (-) reduction:
EIG(E) = Ep(ste [H [p(0]] H (101E.y)ll.
The optimal design is then defined as <* = arg max EIG(E). In practice, calculating EIG is a computationally challenging task because it involves integrations over p(ylE) and p(0]g, (.fls.y). v). which both intractable. In recent years, various methods have been proposed make the computation of the EIG feasible in practical scenarios, such as nested Monte Carlo estimators (Rainforth et al., 2018) and variational approximations (Foster al., 2019). However, even with these advancements, the computational load remains significant, hindering feasibility in tasks that demand rapid designs. This limitation has pushed forward the development of amortized BED methods, which significantly reduce computational demands during the deployment stage.
2.2 Amortized BED
Amortized BED methods represent significant shift from traditional experimental optimization tech- niques. Instead of optimizing for each experimental design separately, Foster et al. (2021) developed parameterized policy я, which directly maps the experimental history æдя {(E1, // )...... (E., Vo to the next design &++1 x (haz). To train such a policy network, Foster et al. (2021) proposed using sequential Prior Contrastive Estimation (sPCH) optimize the lower bound of the total EIG across the entire T-step experiments trajectory:
p( 16. TT)
s PCF((T, L.) E (0h log
L+1 P(aer 9c,
where 01.1. are contrastive samples drawn from the prior distribution. Although the training of such a policy network is computationally expensive, once trained, the network can act as an oracle to quickly propose the next design through single forward pass, thus amortizing the initial training cost over numerous deployments.
3

---

2.3 Bayesian decision theory
Bayesian decision theory (Berger, 2013) provides an axiomatic framework for decision-making under uncertainty, systematically incorporating probabilistic beliefs about unknown parameters into decision-makee processes. It introduces a sask-specific utility function u(0, a), which quantifies the value of the outcomes from different decisions E AA when the system is in state 0. The optimal decision is then determined by maximizing the expected utility, which integrates the utility function over the unknown system parameters, given the available knowledge h1::
a arg max E aCA
(1..)) (1
[u(0,, a)].
In many scemarios, outcomes are stochastic and it is more typical to make decisions based on their predictive distribution P(ylč, N1:s) By(0()r.e)
( [p(ys. 6 h1:1), such in clinical trials where the optimal treatment is chosen based on predicted patient responses rather than solely on underlying biological mechanisms. A similar setup can be found in (Kušmierczyk et al., 2019; Vadera al., 2021). As we switch the belief about the state of the system to the outcomes and to keep as much information as possible, we need evaluate the effect of 0 on all points of the design space. Thus, instead of the posterior over the latent state 0, we represent our belief directly as p(y= ( = {p(yle, h1aa)) /1a)) ie. a joint predictive (posterior) distribution of outcomes over all possible designs given the current information h1s-"' The utility is then expressed u(v=,a), which relies on the decision and all possible predicted outcomes ye. It is a natural extension of the traditional definition of utility by marginalizing out the posterior distribution of 0. The rule of making the optimal decision is reformulated in terms of the predictive distribution as:
= arr max (u(ys,
(4)
aEA
Traditional methods usually separate the inference and decision-making steps, which optimal when the true posterior or the predictive distribution can be computed exactly. However, in most cases the posteriors are not directly accessible, and we often resort to using approximate distributions. This necessity results in a suboptimal decision-making process as the approximate posteriors often focus on representing the full posterior yet fail to ensure high accuracy in regions crucial for decision- making. Loss-calibrated inference (Lacoste-Julien et al., 201 I) emerges as an effective solution to address this problem. It calibrates the inference by focusing on utility rather than mere accuracy of the approximation, thereby ensuring more targeted posterior estimation. This method has been applied to improving Markov chain Monte Carlo (MCMC) methods (Abbasnejad et al., 2015), Bayesian neural networks (Cobb et al., 201 8) and expectation propagation (Morais and Pillow, 2022).
3 Decision-aware BED
3.1 Problem setup
Our objective is to optimize the experimental design process for down-the-line decision-making, In this paper, we consider scenarios which we design a series of experiments & E and observe corresponding outcomes y to inform final makingg step. We assume we have a fixed experimental budget with T query steps. For decision-making we consider a set of possible decisions, denoted as E A, with the objective of identifying an optimal decision a" that maximizes predefined prediction-basee basse utility function u( u(yz, a).
3.2 Decision Utility Gain
Our method focuses on designing the experiments to directly improve the quality of the final decision- making. To quantify the effectiveness of each experimental design in terms of decision-making we introduce Decision Utility Gain (DUG), which is defined as the difference in the expected utility of the best decision, with the new information obtained from the current experimental design, versus the best decision with the information obtained from previous experiments.
Thii definition assumes conditional independence of the outcomes given the design. More generally, p(y=| (4 defines a joint distribution or a stochastic process indexed by the set E (Parzen, 1999). where familiar example could be Gaussian process posterior defined on 3 C _dd (Rasmussen and Williams, 2006).

---

Definition 3.1. Given a historical experimental trajectory h1:--1, the Decision Utility Gain (DUG) for given design & and its corresponding outcome Y/. at step defined as follows:
(UG yr) max Epívalhie 10((E..2)))) [u(w=, a)] max Ep(valbua 1) [u(ye, a)]-
(5)
DUG measures the improvement in the maximum expected utility from observing new experimental design, differing in this from standard marginal utility gain (see e.g., Garnett, 2023). The optimal design is the one that provides the largest increase in maximal expected utilityy Shifting from parameter-centric to utility-centrie evaluation, we directly evaluate the design's influence on the decision utility, bypassing the need to reduce the uncertainty of unknown latent parameters.
At the time we choose the design ξt, the outcome remains uncertain. Therefore, we should consider the Expected Decision Utility Gain (EDUG) with respect to the marginal distribution of the outcomes to select the next design.
Definition 3.2. The Expected Decision Uiility Gain (EDUG) for design &r, given the historical experimental trajectory 11:--1, is defined as:
(EUG(()) Bp(ulli.Aaaa1)) (UG((E:y)))
(6)
With EDUG, we can guide the experimental design without calculating the posterior distribution. If we knew the true predictive distribution, we could always determine the one-step lookahead optimal design by maximizing EDUG across the design space with 8r = arg max HDUG((). However, in practice, the true predictive distributions are often unknown, making the optimization of EDUG exceptionally challenging. This difficulty arises due to the inherent bi-level optimization problem and the need to evaluate two layers of expectations, both over the unknown predictive distribution.
To avoid the expensive computational cost of optimizing EDUG, we propose leveraging a policy network, inspired by Foster et al. (2021), that directly maps historical data to the next design. This approach sidesteps the continuous need to optimize EDUG by learning design strategy over many simulated experiment trajectories during the training phase. It can dramatically reduce computational demands at deployment, allowing for more efficient real-time decisions.
4 Amortizing decision-aware BED
fully amortized BED framework for dacision-makingg requires not only amortizing the experimental design but also the predictive distribution to approximate the expected utility. Moreover, permutation invariance is often assumed in sequential BED (Foster et al., 2021)², meaning that the sequence of experiments does not influence the cumulative information gained. Conditional neural processes (CNPs) (Garnelo et al., 2018; Nguyen and Grover, 2022; Huang et al., 2023b) provide suitable basis for developing invariance of the inputs by treating them as an unordered set but also amortizes modeling of the predictive distributions. See Appendix A for brief introduction to CNPs and TNPs.
framework due to their design, which not only respects the permutation
4.1 Transformer Neural Decision Processes
The architecture of our model, termed Transformer Neural Decision Process (TNDP), is novel architecture building upon the Transformer neural process (TNP) (Nguyen and Grover, 2022). It aims to amortize both the experimental design and the predictive distributions for the subsequent decision- making process. The data architecture of our system comprises four parts D = (D D(P) D GI]:
A context set DiO)_ = -1at = ((El)) -i { contains all past L-step designs and outcomes.
A prediction set D(P) {(c))) y/ ) consists of np design-outcome pairs used for approximating p(y= (1=:). which is closely related to the training objective of the CNPs. The output from this head can then be used to estimate the expected utility.
A query set D {5 ( 1 consists of ng candidate experimental designs being consid- ered for the next step. In scenarios where the design space E is continuous, we randomly sample a set of query points for each iteration during training. In the deployment phase, optimal experimental designs can be obtained by optimizing the model's output.
'When permutation invariance does not hold in some cases, our model can be easily adapted by adding positional encoding to the input.
5

---

D/d
770
stsp
GT
Prediction Head f
Head f
AW
Transformer Block fite
E7
Data Embedding Block fub
d Context Set pm
d
d
d
Global Info GI
Prediction Set pi
Query Set pia
(a)
Figure 2: Illustration of TNDP. (a) An overview of TNDP architecture with input consisting of observed design-outcome pairs from D(e), 2 designs from D for prediction, and candidate designs from D for query. (b) The corresponding attention mask. The colored squares indicate that the elements on the left can attend to the elements on the top in the self-attention layer of fifm-
Global information GI = [t, 7] where represents the current step in the experimental sequence, and encapsulates task-related information, which could include contextual data relevant to the decision-making process. We will further explain the choice y in Section
TNDP comprises four main components: the data embedder block femb, the Transformer block fimm- the prediction head fp. and the query head fq. Each component plays a distinct role in the overall decision-aware BED process. The full architecture is shown in Fig. 2(a).
At first, each set of D is processed by the data embedder block femb to map to aligned embedding space. These embeddings are then concatenated to form a unified representation E concat ( E(P) E(4) E Please refer to Appendix B for a detailed explanation of how we embed the data. After the initial embedding, the Transformer block fum processes E using self-attention mechanisms produce single attention matrix, which is subsequently processed by an attention mask (see Fig. 2(b)) that allows for selective interactions between different data components, ensuring that each part contributes appropriately to the final output. To explain, each design from the prediction set D(P) is configured to attend to itself, the global information, and the historical data, reflecting the dependence of the predictions on the historical data and the independence from other designs. Similarly, each ((9) in the query set D(9) is also restricted to attend only to itself, the global information, and the historical data. This setup preserves the independence of each candidate design, ensuring that the evaluation of one design neither influences nor is influenced by others. The output of film is then split according to the specific needs of the query and prediction head A = [X() X = fim)
The primary role of the prediction head fp is to approximate (( with family of parameter- ized distributions (yslPtt,, where p. = (A!!)) is the output of fo at the step t. The training of fo is by minimizing the negative log-likelihood of the predicted probabilities:
T
Xip
C(P)
log gl. (p..) =
log N (v
==
where Pi, represents the output of design c(P) at step I. Here choose Gaussian likelihood with ja and a representing the predicted mean and standard deviation split from
The query head f4 processes the transformed embeddings ( from the Transformer block to generate a policy distribution over possible experimental designs. Specifically, it converts the embeddings into probability distribution used to select the next experimental design. The outputs of the query head, ={((()), are mapped to a probability distribution via a Softmax function:
exp (q End exp(q..))
(8)

---

where q3.e represents the t-step's s output for the candidate design 4
To design a reward signal that guides fa in proposing informative designs, we first define a single- step immediate reward based on DUG (Eq. (5)), replacing the true predictive distribution with our approximated distribution:
re r(( = max (ep(sl/l)))) (ps) [u(ye. a)]
max (v(lppp--)) (u(ya,a)) dEA
(9)
aEA
This reward quantifies how the experimental design influences our decision-making by estimating the improvement in expected utility that results from incorporating new experimental outcomes. However, this objective remains myopic, as it does not account for the future or the final Decision-makinn address this, we employ the REINFORCE algorithm (Williams, 1992), which allows us to consider the impact of the current design on future rewards. The final loss of fa can written as the negative expected reward for the complete experimental trajectory:
To
T
(
Rs log ग (11a-1),
(10)
where R² = at represents the non-myopic discounted reward. The discount factor used to decrease the importance of rewards received at later time step. <t obtained through sampling from the policy distribution <1") ~ π (-1//a=-1).
The update of Sq depends critically on the accuracy with which Sp approximates the predictive distribution. Ultimately, the effectiveness of decision-making relies on the informativeness of the designs proposed by fq. ensuring that every step in the experimental trajectory is optimally aligned with the overarching goal of maximizing the decision utility. The full algorithm of our method is shown in Appendix C.
5 Related work
Lindley (1972) proposes the first decision-theoreti BED framework, later reiterated by Chaloner and Verdinelli (1995). However, their utility is defined based on individual designs, while our utility is formulated in terms of a stochastic process and is designed for the final decision-making task after multiple rounds of experimental design. Recently, several other BED frameworks that focus on different downstream properties have been proposed. Bayesian Algorithm Execution (BAX) (Neiswanger et al., 2021) develops a BED framework that optimizes experiments based on downstream properties of interest. BAX introduces a new metric that queries the next experiment by maximizing the mutual information between the property and the outcome. CO-BED (Ivanova et al., 2023) introduces contextual optimization method within the BED framework, where the design phase incorporates information-theorrticc objectives specifically targeted at optimizing contextual rewards. Neiswanger et al. (2022) presents information based acquisition function for Bayesian optimization which explicitly considers the downstream task. Zhong et al. (2024) proposes a goal- oriented BED framework for nonlinear models using MCMC to optimize the EIG predictive quantities of interest. Filstroff et al. (2024) presents a framework for active learning that queries data to reduce the uncertainty the posterior distribution of the optimal downstream decision.
In recent years, various amortized BED methods have emerged. Foster et al. (2021) is the first to introduce this framework; subsequent work extends it to scenarios with unknown likelihood (Ivanova et al., 2021) and improved performance using reinforcement learning (Blau al., 2022; Lim et al., 2022). The latest research proposes a semi-amortized framework that periodically updates the policy during the experiment to improve adaptability (Ivanova et al., 2024). Maraval et al. (2024) proposes fully amortized Bayesian optimization (BO) framework that employs similar TNP architecture, while their work focuses specifically on BO objectives, our approach addresses general downstream decision-makee making tasks. Additionally, our framework introduces a novel coupled training objective between query and prediction heads, providing a more integrated architecture for decision- decision-making.
Our proposed architecture is based pre-trained Transformer models. Transformer-bast neural processes (Müller et al., 2021; Nguyen and Grover, 2022; Chang et al., 2024) serve as the foundational structure for our approach, but they have not considered experimental design. Decision Transformers (Chen et al., 2021: Zheng et al., 2022) can used for sequentially designing experiments. However, we additionally amortize the predictive distribution, making the learning process more challenging.
7

---

0.9
Proportion of decision:
0.3
target function
1.0
0.7
0.5
queried data next query
0.005
0.5
GP-TEI
CR-US
GP-D610
0.4
EP-OUS
TNUP (wen))
0.000
Design steps
(a) ID synthetic regression
(b) Decision-aware active learning
Figure 3: Results of synthetic regression and decision-aware active learning. (a) The top figure represents the true function and the initial known points. The red line indicates the location of x² The blue star marks the next query point, sampled from the policy's predicted distribution shown in the bottom figure. (b) Mean and standard error of the proportion of correct decisions on 100 test points w.r.t the acquisition steps. Our TNDP significantly outperforms other methods.
6 Experiments
In this section, we evaluate our proposed framework on several tasks. Our experimental ap- proach is detailed in Appendix B. In Appendix F33, provide additional ablation studies of TNDP to show the effectiveness of our query head and the non-myopic objective function. The code reproduce experiments is available https: thub. com/huangdaolang/ bed.
amortized-decision--
6.1 Toy example: synthetic regression
We begin with illustrative example to show how TNDP works. We consider ID synthetic regression task where the goal to perform regression at a specific test point x" on an unknown function. To accurately predict this point, need sequentially design some This example can be viewed as a prediction-oriente active learning (AL) task (Smith et al., 2023).
points to query.
The design space E = X' is the domain of and is the corresponding noisy observations of the function. Let (.X' denote the set combinations of distributions that can output by TNDP, we can then define decision space to be A ((X)). The downstream decision is to output predictive distribution for y" given test point x and the utility function w(y=, a) = log g(y" x= J(::)) log probability of under the predicted distribution, given the queried historical data h4
During training, we sample functions from Gaussian Processes (GPs) (Rasmussen and Williams, 2006) with squared exponential kernels of varying output variances and lengthscales and randomly sample point as the test point x². We set the global contextual information À as the test point x For illustration purposes, we consider only the case where T = 1. Additional details for the data generation can be found Appendix
Results. From Fig, 3(a), we can observe that the values of concentrate near x², meaning our query head fg tends query points close to x² maximize the DUG. This is intuitive example demonstrating that our TNDP can adjust its design strategy based on the downstream task.
6.2 Decision-aware active learning
We now show another application of our method in case of decision-aware AL studied by Filstroff et al. (2024). In this experiment, the model will be used for downstream decision- decision-making after performing AL, i.e., will use the learned information to take action towards specific target. practical application of this problem is personalized medical diagnosis introduced Section 1, where doctor needs query historical patient data decide treatment for
patient.
We use the same problem setup as in Filstroff al. (2024). The decision space consists of Na available decisions, A E Na). The design space E = X' x AA is composed of the patient's

---

rünger
rpart
SVm
xeboost
2.0
784
2.20
2.7
270
12
31
Step2
Step
Step:
Stop
01:00: RS
DKCH
81
P1
PINsSBO
-
TNDP
Figure 4: Results on Top-k HPO task. For each meta-dataset, we calculated the average utility across all available test sets. The error bars represent the standard deviation over five runs. TNDP consistently achieved the best performance in terms of utility.
covariates æ and the decisions they receive. The outcome y is the treatment effect after the patient receives the decision, which is influenced by real-world unknown parameters such as medical conditions. For historical data, each patient is associated with only one decision. The utility function u(yz,4) = I[("", aa) is a binary accuracy score that measures whether we can make the correct decision for new patient x= based on the queried history, where â" is the predicted Bayesian optimal decision and a" the true optimal decision. Here, u(ys, a) = 1 if and only if â" = a".
In experiment, use the synthetic dataset from Filstroff et al. (2024), the details of the data generating process can be found in Appendix F. We set N₂ = and use independent GPs to generate different outcomes. Each data point is randomly assigned decision, and the outcome is the corresponding y value from the associated GP. We randomly select a test point x" and determine the optimal decision a" based on the GP that provides the maximum y value at x². We set global contextual information as the covariates of the test point x".
We compare TNDP with other non-amortized AL methods: random sampling (GP-RS), uncertainty sampling (GP-US), decision uncertainty sampling (GP-DUS), targeted information (GP-TEIG) introduced by Sundin et al. (2018), and decision EIG (GP-DEIG) proposed by Filstroff et al. (2024). detailed description each method can be found in Appendix Each method is tested on 100 different x" points, and the average utility, i.e., the proportion of correct decisions, is calculated.
Results. The results are shown in Fig. 3(b), where we can see that TNDP achieves significantly better average accuracy than other methods. Additionally, we conduct an ablation study of TNDP in Appendix F.3 to verify the effectiveness of fq. We further analyze the deployment running time to show the advantage of amortization, see Appendix D.I.
6.3 Top-k hyperparameter optimization
In traditional optimization tasks, we typically only aim to find a single point that maximizes the underlying function f. However, instead of identifying single optimal point, there are scenarios where we wish estimate set of top-k distinct optima. For example, this is the case in robust optimization, where selecting multiple points can safeguard against variations in data or model performance.
In this experiment, we choose hyperparameter optimization (HPO) as our task and conduct experi- ments the HPO-B datasets (Arango al., 2021). The design space E C X' is finite set defined over the hyperparameter space and the outcome is the accuracy of a given configuration on specific dataset. Our decision is to find & hyperparameter sets, denoted
(ai,..., a&) E A x with a, * aj. The utility function is then defined as u(yz, a) = Eh= Vo,, where No, is the accuracy corresponding the hyperparameter configuration a;. In this experiment, the global contextual information = 0.
We compare our methods with five different BO methods: random sampling (RS), Upper Confidence Bound (UCB), Expected Improvement (EI), Probability of Improvement (PI), and an amortized method PFNs4BO (Müller et al., 2023), which is transformer-based model designed for hyperpa- rameter optimization. We set k = and T 50, starting with an initial dataset of points. Our experiments are conducted on four search spaces selected from the HPO-B benchmark. All results

---

are evaluated on a predefined test set, ensuring that TNDP does not encounter these test sets during training. For more details, see Appendix G.
Results. From the experimental results (Fig. 4), our method demonstrates superior performance across all four meta-datasets, particularly during the first 10 queries, achieving clearly better utility gains. This indicates that our TNDP can effectively identify high-performing hyperparameter configurations early in the optimization process.
Finally, we included a real-world experiment on retrosynthesis planning. Specifically, our task is to assist chemists in identifying the top-k synthetic routes for novel molecule, selecting the most practical routes from many random routes generated by the retrosynthesis software can be troublesome. The detailed description of the task and the results are shown in Appendix G.3.
7 Discussion
7.1 Limitations & future work
We recognize that the training of the query head inherently poses a reinforcement learning (RL) (Li, 2017) problem. Currently, we employ basic REINFORCE algorithm, which can result in unstable training, particularly for tasks with sparse reward signals. For more complex problems in the future, we could deploy advanced RI. methods, such Proximal Policy Optimization (PPO) (Schulman et al., 2017); the trade-offs include the introduction of additional hyperparameters and increased computational cost. Like all amortized approaches, our method requires large amount of data and upfront training time to develop a reliable model. Besides, 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. Future work could focus on designing more cample-efficiem methods reduce the data and training requirements. Our TNDP follows the common practice in the neural processes literature (Garnelo et al., 2018) of using independent Gaussian likelihoods. If modeling correlations between points is crucial for the downstream task, we can replace the output with a joint multivariate normal distribution (Markou et al., 2022) or predict the output autoregressively (Bruinsma et al., 2023). Following most BED approaches, our work assumes that the model is well-specified. However, model misspecification or shifts in the utility function during deployment could impact the performance of the amortized model (Rainforth al., 2024). Future work could address the challenge of robust experimental design under model misspecificatiion (Huang et al., 2023a). Another limitation is that our system currently constrained accepting designs of the same dimensionality. Future work could focus on developing ciimension-agnosee methods to expand the scope of amortization. Lastly, our model is trained on fixed-step length, assuming finite horizon for the experimental design process. Future research could explore the design of systems that can handle infinite horizon cases, potentially improving the applicability of TNDP to broader range of real-world problems.
7.2 Conclusions
In this paper, proposed an amortized framework for decision-aware Bayesian experimental design (BED). We introduced the concept of Decision Utility Gain (DUG) to guide experimental design more effectively toward optimizing decision outcomes. Towards amortization, novel Transformer Neural Decision Process (TNDP) architecture with dual output heads: one for proposing new experimental designs and another for approximating the predictive distribution to facilitate optimal cerziiioonmaking. Our experimental results demonstrated that TNDP significantly outperforms traditional BED methods across variety of tasks. By integrating decision-making considerations directly into the experimental design process, TNDP not only accelerates the design of experiments but also improves the quality of the decisions derived from these experiments.
developed
A Acknowledgements
DH, LA and SK were supported by the Research Council of Finland (Flagship programme: Finnish Center for Artificial Intelligence FCAI). YG was supported by Academy of Finland grant 345604. LA was also supported by Research Council of Finland grants 358980 and 356498. SK was also supported by the UKRI Turing AI World-Leading Researcher Fellowship, [EP//W02973/11. The authors wish to thank Aalto Science-IT project, and CSC-IT Center for Science, Finland, for the computational and data storage resources provided.
10

---

References
Abhasnejad, E.,, Domke, J., and Sanner, S. (2015). Loss-calibrated monte carlo action selection, In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29.
Arango, S. P., Jomaa, H. S., Wistuba, M., and Grabocka, J. (2021). Hpo-b: A large-scale reproducible benchmark for black-box hpo based on openml. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2).
Balandat, M., Karrer, B., Jiang, D., Daulton, S., Letham, B., Wilson, A. G., and Bakshy, E. (2020). Botorch: A framework for efficient monte-carlo bayesian optimization. Advances in neural information processing systems, 33.
Berger, J. O. (2013). Statistical decision theory and Bayesian analysis. Springer Science & Business Media.
Bica, I., Alaa, A. M., Lambert, C., and Van Der Schaar, M. (2021). From real-world patient data to individualized treatment effects using machine learning: current and future methods to address underlying challenges. Clinical Pharmacology & Therapeutics, 109(1):87-100.
Bille, P. (2005). A survey on tree edit distance and related problems. Theoretical computer science, 337(1-3):117239.
Blacker, J., Williams, M. T., and Williams, M. (2011). Pharmaceutical process development:
current chemical and engineering challenges, volume 9. Royal Society of Chemistry.
Blau, T., Bonilla, E., Chades, I., and Dezfouli, A. (2023). Cross-entropy estimators for sequential experiment design with reinforcement learning. arXiv preprint arXiv:2305. 18435.
Blau, T., Bonilla, E. V., Chades, I., and Dezfouli, A. (2022). Optimizing sequential experimental design with deep reinforcement learning. In International conference on machine learning, pages 2128. PMLR
Bruinsma, W., Markou, S., Requeima, J., Foong, A. Y., Andersson, T., Vaughan, A., Buonomo, A. Hosking, S., and Turner, R. E. (2023). Autoregressive conditional neural processes. In The Eleventh International Conference Learning Representations.
Burger, M., Hauptmann, A.. Helin, T., Hyvönen, N., and Puska, J.-P. (2021). Sequentially optimized projections in x-ray imaging. Inverse Problems, 37(7):075006.
Chaloner, K. and Verdinelli, (1995), Bayesian experimental design: review, Statistical science, pages 273-304.
Chang. P. E., Loka, N., Huang, D., Remes, U., Kaski, S., and Acerbi, L. (2024). Amortized proba- bilistic conditioning for optimization, simulation and inference. arXiv preprint arXiv:2410. 10..220..
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, (2021). Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34.
Cheng, Y. and Shen, Y. (2005). Bayesian adaptive designs for clinical trials. Biometrika, 92(3):633- 646.
Cobb, D., Roberts, S. J., and Gal. Y. (2018). Loss-calibrated approximate inference in bayesian neural networks. arXiv preprint
Filstroff, L., Sundin, I., Mikkola, P., Tiulpin, A., Kylmuoja, J., and Kaski, S. (2024). Targeted active learning for bayesian decision-making Transactions on Machine Learning Research.
Foster, A., Ivanova, D. R., Malik, L, and Rainforth, T. (2021). Deep adaptive design: Amortizing sequential bayesian experimental design. In International Conference on Machine Learning, pages PMLR.
Foster, A., Jankowiak, M., Bingham, E., Horsfall, P., Teh, Y. W., Rainforth, T., and Goodman, N. (2019). Variational bayesian optimal experimental design. Advances Neural Information Processing Systems, 32.
11

---

Foster, A., Jankowiak, M., O' O'Meara, M., Teh, Y. W., and Rainforth, T. (2020). A unified stochastic gradient upproach designing mayesian-optimall experiments. International Conference on Artificial Intelligence and Statistics, pages 2959-2969. PMLR.
Garnelo, M.. Rosenbaum, D., Maddison, C., Ramalho, T., Saxton, D., Shanahan, M., Teh, Y. W., Rezende, D., and Eslami, S. A. (2018). Conditional neural processes. In International conference machine learning, pages 1704-1713. PMLR.
Gamett, R. (2023). Bayesian optimization. Cambridge University Press.
Huang, D., Bharti, A., Souza, A., Acerbi, L., and Kaski, S. (2023a). Learning robust statistics for simulation-basee baased inference under model misspecification. Advances Neural Information Processing Systems, 36.
Huang, D., Haussmann, M., Remes, U., John, S., Clarté, G., Luck, K., Kaski, S., and Acerbi, L. (2023b). Practical equivariances via relational conditional neural processes. Advances Neural Information Processing Systems, 36.
Ivanova, D. R., Foster, A., Kleinegesse, S., Gutmann, M. U., and Rainforth, T. (2021). Implicit deep adaptive design: Policy-based experimental design without likelihoods. Advances Neural Information Processing Systems, 34.
Ivanova, D. R., Hedman, M., Guan, C., and Rainforth, T. (2024). Step-DAD: Semi-Amortized Policy-Basse Based Bayesian Experimental Design. ICLR 2024 Workshop on Data-centric Machine Learning Research (DMLR).
Ivanova, D. R., Jennings, J., Rainforth, T., Zhang, C., and Foster, A. (2023). Co-bed: information- theoretic contextual optimization via bayesian experimental design. International Conference on Machine Learning. PMLR.
Kleinegesse, S. and Gutmann, M. U. (2020). Bayesian experimental design for implicit models by mutual information neural estimation. International conference on machine learning, pages 5316-5326. PMLR.
Kušmierozyk, T., Sakaya, J., and Klami, A. (2019). Variational bayesian decision-making for continuous utilities. Advances in Neurul Information Processing Systems, 32.
Lacoste-Julien, S., Huszár, F., and Ghahramani, Z. (2011). Approximate inference for the loss- calibrated bayesian. In Proceedings of the Fourteenth International Conference Artificial Intelligence and Statistics, pages 416-424. JMLR Workshop and Conference Proceedings.
Li, Y. (2017). Deep reinforcement learning: An overview. arXiv preprint arXiv:1701.07274
Lim, V., Novoseller, E., Ichnowski, J., Huang, H., and Goldberg, K. (2022). Policy-based bayesian experimental design for non-differentiable implicit models. arXiv preprint arXiv:2203.0427
Lindley, D. V. (1956). On measure of the information provided by an experiment. The Annals of Mathematical Statistics, 2764))98-1
Lindley, D. V. (1972). Bayesian statistics: review. SIAM.
Maraval, A., Zimmer, M., Grosnit, A., and Bou Ammar, H. (2024). End-to-end meta-bayesian optimisation with transformer neural processes. Advances in Neural Information Processing Systems, 36.
Markou, S., Requeima, J., Bruinsma, W., Vaughan, A., and Turner, E. (2022). Practical conditional neural process via tractable dependent predictions. In International Conference Learning Representations.
Mo, Y., Guan, Y., Verma, P., Guo, J., Fortunato, E., Lu, Z.,, Coley, W., and Jensen, K. (2021). Evaluating and clustering retrosynthesis pathways with learned strategy. Chemical science, 44644:1469-1478
Morais, M. J. and Pillow, J. W. (2022). Loss-calibrated expectation propagation for approximate bayesian fecision-making arXiv preprint arXiv:2201.03125
12

---

Müller, S., Feurer, M., Hollmann, N., and Hutter, F. (2023). Pfns4bo:
In-context learning for bayesian
optimization. In International Conference on Machine Learning, pages 25444-24 25470. PMLR.
Müller, S., Hollmann, N., Arango, S. P.. Grabocka, J., and Hutter, F. (2021). Transformers can do bayesian inference. In International Conference on Learning Representations.
Neiswanger, W., Wang. K. A., and Ermon, S. (2021). Bayesian algorithm execution: Estimating com- putable properties of black-box functions using mutual information. In International Conference on Machine Learning, pages 8005-8015, PMLR.
Neiswanger, W., Yu, L., Zhao, S., Meng, C., and Ermon, S. (2022). Generalizing bayesian opti- mization with decision-theoretic entropies. Advances in Neural Information Processing Systems, 35.
Nguyen, and Grover, A. (2022). Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In International Conference on Machine Learning, pages 16569-16594. PMLR.
Parzen, E. (1999). Stochastic processes. SIAM.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. (2019). Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32.
Rainforth, T., Cornish, R., Yang, H., Warrington, A., and Wood, F. (2018). On nesting monte carlo estimators. In International Conference Machine Learning, pages 4267-42777 PMLR.
Rainforth, T., Foster, A., Ivanova, D. R., and Bickford Smith, (2024). Modern bayesian experimen- tal design. Statistical Science, 39(1):10--14
Rasmussen, E. and Williams, C. K. (2006). Gaussian Processes for Machine Learning. MIT Press.
Ryan, E. G., Drovandi, C. C., McGree, J. M., and Pettitt, A. N. (2016). A review of modern computational algorithms for bayesian optimal design. International Statistical Review, 84(1):128- 154.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint orXiv:1707.060
Smith, F. B., Kirsch, A., Farquhar, S., Gal, Y., Foster, A., and Rainforth, T. (2023). Prediction-oriented bayesian active learning. In International Conference on Artificial Intelligence and Statistics, pages 7331-7348. PMLR.
Stevens, S. J. (2011). Progress toward the synthesis of providencin. PhD thesis, Colorado State University.
Sundin, I., Peltola, T., Micallef, L., Afrabandpey, H., Soare, M., Mamun Majumder, M., Dace, P., He. C., Serim, B., Havulinna, A., et al. (2018). Improving genomics-based predictions for precision medicine through active elicitation of expert knowledge. Bioinformatics, 34(13):1395-1443
Szymkué, S., Gajewska, E. P., Klucznik, T.. Molga, K., Dittwald, P.. Startek, M.. Bajezyk. M.. and Grzybowski, B. A. (2016). Cosputer-assisted synthetic planning: the end of the beginning. Angewandte Chemie International Edition, 59940)7590
Vadera, M. P., Ghosh, S., Ng, K., and Marlin, B. M. (2021). Post-hoc loss-calibration for bayesian neural networks. In Uncertainty in Artificial Intelligence, pages 1403-1412. PMLR.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.. Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.
Williams, R. (1992). Simple statistical pradient-followin algorithms for connectionist reinforce- ment learning. Machine learning, 8:229-256.
13

---

Zheng, Q.. Zhang, A., and Grover, A. (2022). Online decision transformer. In international conferncce on machine learning. PMLR.
Zhong, S., Shen, W., Catanach, T., and Huan, X. (2024). Goal-oriented bayesian optimal experimental design for nonlinear models using markov chain monte carlo. arXiv preprini 1.xiv:22433.8072
14

---

Appendix
The appendix is organized as follows:
In Appendix A, we provide brief introduction to conditional neural processes (CNPs) and Transformer neural processes (TNPs).
In Appendix describe the details of model architecture and the training setups.
In Appendix C, present the full algorithm for training TNDP architecture.
In Appendix D. we compare the inference time with other methods and show the overall training time of TNDP.
In Appendix E, describe the details of our toy example.
In Appendix F, we describe the details of the decision-awand active learning example.
Appendix G, we describe the details of the top-k hyperparameter optimization task, along with additional results on the retrosynthesis planning task.
A Conditional neural processes
CNPs (Garnelo et al., 2018) are designed to model complex stochastic processes through a flexible architecture that utilizes a context set and targer set. The context set consists of observed data points that the model uses to form its understanding, while the target set includes the points to be predicted. The traditional CNP architecture includes an encoder and decoder. The encoder architecture to ensure permutation invariance, it transforms each context point individually and then aggregates these transformations into single representation that captures the overall context. The decoder then uses this representation to generate predictions for the target set, typically employing Gaussian likelihood for approximation of the true predictive distributions. Due the analytically tractable likelihood, CNPs can be efficiently trained through maximum likelihood estimation.
DeepSet
A.1 Transformer neural processes
Transformer Neural Processes (TNPs), introduced by Nguyen and Grover (2022), improve the flexibility and expressiveness of CNPs by incorporating the Transformer's attention mechanism (Vaswani et al., 2017). In TNPs, the transformer architecture uses self-attention process the context set, dynamically weighting the importance of each point. This allows the model to create rich representation of the context, which is then used by the decoder to generate predictions for the target set. The attention mechanism in TNPs facilitates the handling of large and variable-sized context sets, improving the model's s performance on tasks with complex input-output relationships. The Transformer architecture is also useful in our setups where certain designs may have a more significant impact on the decision-making process than others. For more details about TNPs, please refer to Nguyen and Grover (2022).
B Implementation details
B.1 Embedders
The embedder fenb responsible for mapping the raw data to a space of the same dimension. For the toy example and the top-k hyperparameter task, we use three embedders: a design embedder
f(e) an outcome embedder flx) and time step embedder f({) Both f(E) and f(v) multi-layer perceptions (MLPs) with the following architecture:
Hidden dimension: the dimension of the hidden layers, set to 32.
Output dimension: the dimension of the output space, set to 32,
Depth: the number of layers in the neural network, set to
Activation function: ReLU is used as the activation function for the hidden layers.
The time step embedder (() is a discrete embedding layer that maps time steps to continuous embedding space of dimension 32.
15

---

For the decision-ave awaree active learning task, since the design space contains both the covariates and the decision, use four embedders: covariate embedder feab a decision embedder fealy an outcome embedder p(y) and time step embedder senbb femb Jemb and fenb are MLPs which use the same settings as described above. The decision embedder siah is another discrete embedding layer.
For context embedding E we first map each g( and v, to the same dimension using their respective embedders, and then sum them to obtain the final embedding. For prediction embedding E and query embedding E we only encode the designs. For E except the embeddings of the time step, we also encode the global contextual information À using flin in the toy example and the decision-aware active learning task. All the embeddings are then concatenated together to form our final embedding E.
B.2 Transformer blocks
We utilize the official TraeEnoorerr layer of PyTorch (Paszke et al., 2019) (https: pytorch. org) for our transformer architecture. For all experiments, we use the same configuration: the model has 6 Transformer layers, with 8 heads per layer, the MLP block has hidden dimension of 128, and the embedding dimension size is set to 32.
B.3 Output heads
The prediction head, fp is an MLP that maps the Transformer's output embeddings of the query the predicted outcomes. It consists of an input layer with 32 hidden units, ReLU activation function, and an output layer. The output layer predicts the mean and variance of Gaussian likelihood, similar to CNPs.
For the query head fy. all candidate experimental designs are first mapped embeddings X by the Transformer, and these embeddings are then passed through fy to obtain individual outputs. We then apply Softmax function these outputs ensure proper probability distribution. fy is MLP consisting of an input layer with 32 hidden units, a ReLU activation function, and an output layer.
B.4 Training details
For all experiments, we use the same configuration to train our model. We set the initial learning rate to 5e-4 and employ the cosine annealing learning rate scheduler. The number of training epochs is set to 50,000 for top-k tasks and 100,000 for other tasks, and the batch size 16. For the REINFORCE, use discount factor of 0.99.
C Full algorithm for training TNDP
Algorithm Transformer Neural Decision Processes (TNDP)
Input: Utility function u(ys, a), prior p(0), likelihood p(yl6, E). query horizon T
2: Output: Trained TNDP
3: while within the training budget do
4: Sample 0 ~ p(0) and initialize D
5: for
do
6: c. ~ (e((1f1aa11)
sample next design from policy
7: Sample yr ~ (,)
> observe outcome
8:
Set /14 h U ((gl0) .w:)]
update history update D
9: Set D = h1r, D D (
Calculate re(E(T) with u(y= a) using Eq. (9)
calculate reward
11: end for
R₂ Σk=r h (Tx)
calculate cumulative reward
Update TNDP using ((P) (Bq. (7)) and ( (Eq. (10)
14: end while
deployment, we can use f( sequentially query T designs. Afterward, based on the queried experiments, we perform one-step final decision decision-making making using the prediction from f(r).
16

---

D Computational cost analysis
D.1 Inference time analysis
We evaluate the inference time of our algorithm during the deployment stage. We select decision- aware active learning as the experiment for our time comparison. All experiments are evaluated on an Intel Core 17-12700K CPU. We measure both the acquisition time and the total time. The acquisition time refers to the time required to compute one next design, while the total time refers to the time required to complete 10 rounds of design. The final results are presented in Table Al, with the mean and standard deviation calculated over 10 runs.
Traditional methods rely on updating the GP and optimizing the acquisition function, which is computationally expensive. D-EIG and T-EIG require many model retraining steps to get the next design, which is not tolerable in applications requiring fast decision- decision-making. making. However, since our model is fully amortized, once trained, it only requires a single forward pass design the experiments, resulting in significantly faster inference times.
Method
Acquisition time (s)
Total time (s)
GP-RS
0.00002 (0. 00.0000)
28(7)
GP-US
0.07(0.01
(9(7)
GP-DUS
0.38 (0.02)
44(5)
T-EIG (Sundin et al., 2018)
1558 (376)
15613
D-EIG (Filstroff al., 2024)
572 105)
5746
TDNP (ours)
0.015 (0.004)
(.31(0.06)
Table A1: Comparison of computational costs across different methods. We report the mean value and (standard deviation) derived from 10 runs with different seeds.
D.2 Overall training time
Throughout this paper, we carried all experiments, including baseline model computations and preliminary experiments not included in the final paper, GPU cluster featuring combination of Tesla P100, Tesla V100, and Tesla A100 GPUs. We estimate the total computational usage be roughly 5000 GPU hours. For each experiment, it takes around 10 GPU hours on a Tesla V100 GPU with 32GB memory to reproduce the result, with an average memory consumption of GB.
E
Details of toy example
E.1
Data generation
In our toy example, we generate data using a GP with the Squared Exponential (SE) kernel, which is defined as:
k(æ, x') = exp (- 2
(x - x²)²
(AI)
where e is the variance, and f is the lengthscale. Specifically, each training iteration, we draw random lengthscale f ~ 0.25 + 0.75 x UO, 1) and the variance 0.1 + U (o, 1), where U(0, 1)
denotes a uniform random variable between and I.
F
Details of decision-aware active learning experiments
F.1
Data generation
For this experiment, we use GP with Squared Exponential (SE) kernel to generate our data. The covariates æ are drawn from a standard normal distribution. For each decision, we use an independent GP to simulate different outcomes. In each training iteration, the lengthscale for each GP is randomly sampled as & 0.25 + 0.75 (0, 1) and the variance as 0.1 + (0, 1), where U(0,) 1) denotes uniform random variable between 0 and 1.
17

---

F.2
Other methods description
We compare method with other non-amortized approaches, all of which GPs as the functional prior. Each model equipped with an SE kernel with automatic relevance determination. GP hyperparameters estimated with maximum marginal likelihood.
Our method is compared with the following methods:
Random sampling (GP-RS): randomly choose the next design & from the query set.
Uncertainty sampling (GP-US): choose the next design & for which the predictive distribu- tion p(yole, 1) has the largest variance.
choose the next design & such that the predictive distribution of the optimal decision corresponding to this design is the most uncertain.
Decision uncertainty sampling (GP-DUS):
Targeted information (GP-TEIG) (Sundin al., 2018): targeted active learning criterion, introduced by (Sundin et al., 2018), selects the next design & that provides the highest EIG on ply" x 1 1 U {Ee, )))
Decision EIG (GP-DEIG) (Filstroff et al., 2024): choose the next design & which directly aims at reducing the uncertainty on the posterior distribution of the optimal decision. See Filstroff et al. (2024) for a detailed explanation.
F33
Ablation study
We also carry out an ablation study to verify the effectiveness of our query head and the non-myopic objective function. We first compare TNDP with TNDP using random sampling (TNDP-RS), and the results are shown in Fig. A (a).). We observe that the designs proposed by the query head significantly improve accuracy, demonstrating that query head can propose informative designs based on downstream decisions.
We also evaluate the impact of the non-myopic objective by comparing TNDP with myopic version that only optimizes for immediate utility rather than long-term gains (a = 0). The results, presented Fig. Al(b), show that TNDP with the non-myopic objective function achieves higher accuracy across iterations compared to using the myopic objective. This indicates that our non-myopic objective effectively captures the long-term benefits of each design choice, leading improved overall performance.
0.9
0.9
Provetto of bonnec dechiiens (55)
Propetion cxmes decisioss
TNDP-RS
TNBP
0.5
Design steps
Design stops
(a) Effeet query head
(b) Impact of non-myopic objective
Figure A1: Comparison of TNDP variants on the decision-aware active learning task. (a) Shows the effect of the query head, where TNDP outperforms TNDP-RS, demonstrating its ability to generate informative designs. (b) Illustrates the impact of the non-myopic objective, with TNDP achieving higher accuracy than the myopic version.
18

---

G
Details of top-k hyperparameter optimization experiments
G.1
Data
In this task, we use HPO-B benchmark datasets (Arango et al., 2021). The HPO-B dataset is large-scale benchmark for HPO tasks, derived from the OpenML repository. It consists of 176 search spaces (algorithms) evaluated on 196 datasets, with a total of 6.4 million hyperparameter evaluations. This dataset designed facilitate reproducible and fair comparisons of HPO methods by providing explicit experimental protocols, splits, and evaluation measures.
We extract four meta-datasets from the HPOB dataset: ranger (id=7609, d,=9), svm (id=5891, d,=8), rpart (id=5859, dg=6),, and xgboost (id=5971, dz=16). In the test stage, the initial context set is chosen based their pre-defined indices. For detailed information the datasets, please refer to https:/ / / thub. cww/releauee http://eeeauniffeebrrrrggg
G.2 Other methods description
In our experiments, we compare our method with several common acquisition functions used in HPO. We use GPs as surrogate models for these acquisition functions. All the implementations are based on BoTorch (Balandat al., 2020) battps://botoree org/). The acquisition functions compared are as follows:
Random Sampling (RS): This method selects hyperparameters randomly from the search space, without using any surrogate model or acquisition function.
Upper Confidence Bound (UCB): This acquisition function balances exploration and exploitation by selecting points that maximize the upper confidence bound. The UCB is defined as:
aUCB (x) μux) + kx((x),
(A2)
where μ(x) is the predicted mean, (xx is the predicted standard deviation, and is parameter that controls the trade-off between exploration and exploitation.
Expected Improvement (EI): This acquisition function selects points that are expected to yield the greatest improvement over the current best observation. The El is defined as:
on (x) = Efmmx((0, f(x) - f(xx)() ))l,
(A3)
where f(x+) is the current best value observed, and the expectation is taken over the predictive distribution of fxx).
Probability of Improvement (PI): This acquisition function selects points that have the highest probability of improving over the current best observation. The PI is defined as:
(Pp( = [1(x) f(xx(((xxx)))))))<)))
(A4)
(xx)
where & is the cumulative distribution function of the standard normal distribution, f(xt) is the current best value observed, and w is parameter that encourages exploration.
In addition to those non-amortized GP-based methods, we also compare our method with amortized surrogate model PFNs4BO (Müller et al., 2023). It is a Transformer-based model designed for hyperparameter optimization which does not consider the downstream task. We use the pre-trained PFNs4BO-BaR BNN model which is trained on HPO-B datasets and choose PI as the acquisition function, the model and the training details can be found in their official implementation (https://githe com/ /attomNPPNaBBB
G.3 Additional experiment on retrosynthesis planning
We now show real-world experiment on retrosynthesis planning (Blacker et al., 2011). Specifically, our task to help chemists identify the top-k synthetic routes a novel molecule (Mo et al., 2021) 1), as it can be challenging to select the most practical routes from many random options generated by the retrosynthesis software (Stevens, 2011; Szymkuć et al., 2016). In this task, the design space for each molecule m is finite set of routes that can synthesize the molecule. The sequential experimental
19

---

design is to select a route for specific molecule to query its score y, which is calculated based on the tree edit distance (Bille, 2005) from the best route. The downstream task to recommend the top-k routes with the highest target-specific scores based on the collected information.
choose k: 3 and T' 10, and set the global information y = m. We train our TNDP dataset, including 1500 molecules with 70 syn- thetic routes for each molecule. The represen- tation dimension of the molecule is 64 and that of the route is 264, both of which through a neural network. Given the high- dimensional nature of the data representation, it is challenging to directly compare TNDP with other GP-based methods, as GPs typically strug- gle with scalability and performance in such high-dimensional settings. Therefore, we only compare TNDP with TNDP using random sam- pling. The final results are evaluated on 50 test molecules that not seen during training, shown in Fig. A2.
In this experiment,
novel non-public meta-
Unilyy
learned
TNDP
Random search
Design steps
Figure A2: Results of retrosynthesis planning experiment. The utility is the sum of the quality scores of top-k routes and is calculated with 50 molecules. Our TNDP outperforms the random search baseline.
20

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In this paper, we propose an amortized decision-aware BED framework (Section and Section 4), addressing the desiderata delineated in the abstract and in the introduction. We explicitly state the contributions at the end of the introduction. The experimental results presented in Section match the claims made.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of our work explicitly in Section 7.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
21

---

Answer:
[NA]
Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We fully describe model architecture in Section 4, and we also provide the implementation details Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
22

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We have included the code repository link the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
somiiisscssccy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ttpp..ccppbblll/ http:cccccppcccccccccecccccccccc htppuullccuuuidssssseepppp guides/ /CodeSummddiubii ssi ssisssoolooooooooo oonooiiyyyyloloooooomoooomiiiilp for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We describe the experimental settings for all experiments Ap- pendix E.Appendix F.Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All the experiment results are accompanied by error bars. We also report the method for caleulating the error bars in each experiment section.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
23

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We report the compute resources in Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't It make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https : htt////eurrrppeeeeee wttpripss.cppppaeee wtw.iiiiiicciceeeeeeeeenn
Answer: [Yes]
Justification: Our work adheres to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The work is foundational research and not tied particular applications. We believe there is no negative societal impact on the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the original paper that produced the dataset, and also provide the URLs for the code packages we use.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
25

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
26

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
27


Reviews:
[{"review_id": "LaytIYhfLY", "reviewer": "Reviewer_Zp5w", "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:\n\n1. They introduce the concept of Decision Utility Gain (DUG) to quantify how much an experimental design improves the expected utility of the downstream decision. \n\n2. 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.\n\n3. 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.\n\n4. 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.\n\nIn 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.", "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)?\n\n- 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?", "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.\n- 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- 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.\n- 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.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive assessment of our work and the points you raised. In the following we address your questions and points raised. \n\n**Weaknesses**\n> 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.\n\nThe 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].\n\nRegarding 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.\n\n> 2. The paper could benefit from including real-world case studies or applications to further validate the practical utility of the proposed framework.\n\nWe conducted a new retrosynthesis planning experiment based on real-world molecule data, please refer to the global response.\n\n> 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.)\n\nThanks, 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.\n\n\n**Questions**\n\n> 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)?\n\nIG 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.\n\n> 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?\n\nWe added an ablation study regarding the choice of the discount factor, please refer to the global response.\n\n**Limitations**\n\n> 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.\n\nThanks 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.\n\n> 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.\n\nThe 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.\n\n**References**\n\n[1] Zhang et al. (2021). Sample efficient reinforcement learning with REINFORCE. *AAAI*.\n\n[2] Schulman et al. (2017). Proximal policy optimization algorithms. *arXiv:1707.06347*.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "VALImQdKjt", "reviewer": "Reviewer_WHLG", "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.", "questions": "- Perhaps section 2 can be organized in a way to avoid singleton nested subsection (i.e., 2.1.1)?\n\n- 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)? \n\n- Are there any practical disadvantages of assuming a diagonal Gaussian predictive distribution? Can complex models induce multimodal or highly correlated predictive distributions that?", "limitations": "The authors openly discuss the current limitations of their approach.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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).\n\n- The proposed architecture and training objectives are novel and seem to unlock both qualitative and quantitative improvements over existing methods. \n\n- 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. \n\n- 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.\n\n- 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.\n\n- 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?\n\n- 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.\n \nI am happy to discuss these points with the authors and increase my score if they are addressed / clarified.\n\n[1] Rainforth, T., Foster, A., Ivanova, D. R., and Bickford Smith, F. (2024). Modern Bayesian\n429 experimental design. Statistical Science, 39(1):100–114.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and the valuable comments. We address your remarks and questions below.\n\n**Weaknesses**\n> 1. Some notational confusion can be avoided…\n\nThanks, 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.\n\n> 2. P4L126 is an abuse of notation with slightly confusing wording, such as “the predictive posterior distribution over all possible designs”. \n\nWe 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. \n\n> 3. Figure 2 is not particularly informative for the data flow ... the right panel comes out of the blue and is not well explained.\n\nWe 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.\n\n> 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. \n\nThank 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.\n\n> 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?\n\nWe 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.\n\nAs 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.\n\n> 6. The topic of model misspecification is never mentioned in the paper.\n\nThat’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.\n\n**Questions**\n> 7. Perhaps section 2 can be organized in a way to avoid singleton nested subsection?\n\nAgreed; will do.\n\n> 8. P4L130: Isn’t there also an assumption that decision are optimal only if there is no model misspecification?\n\nYes, it is a basic common assumption in the field of BED [1, 2]. We will explicitly state this assumption in the revised manuscript.\n\n> 9. Are there any practical disadvantages of assuming a diagonal Gaussian predictive distribution? Can complex models induce multimodal or highly correlated predictive distributions?\n\nThat’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.\n\n**References**\n\n[1] Müller et al. (2023). Pfns4bo: In-context learning for Bayesian optimization. *ICML*.\n\n[2] Foster et al. (2021). Deep adaptive design: Amortizing sequential Bayesian experimental design. *ICML*. \n\n[3] Blau et al. (2022). Optimizing sequential experimental design with deep reinforcement learning. *ICML*.\n\n[4] Rainforth et al. (2024). Modern Bayesian experimental design. *Statistical Science*.\n\n[5] Garnelo et al. (2018). Conditional neural processes. *ICML*.\n\n[6] Markou et al. (2022). Practical conditional neural processes via tractable dependent predictions. *ICLR*.\n\n[7] Bruinsma et al. (2023). Autoregressive conditional neural processes. *ICLR*.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "V0WZjw3Gmy", "reviewer": "Reviewer_Ctfm", "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.", "questions": "In addition to the questions raised in the Weaknesses section:\n\n1.  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?\n2. In the decision-aware AL experiment: why does the random baseline perform as good as all the other ones?\n3. 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.", "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. \n\nThe 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. \n\nI 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.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "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.  \n\n- 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.\n\n- 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]:\n\nMy 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.\n\n- 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. \n\n- 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]).\n\n### Related work: \n\nFor 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:\n- Decision-theoretic BED: as previously discussed, the general framework of utility-based BED was developed by Lindley (1972).\n- 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)\n- Decision-theoretic approaches in related fields such as Bayesian Optimisation, e.g. [7], [8]\n- 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]?\n\n### Other:\n\n- 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. \n- 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].\n- 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?\n\n-----\n#### References\n\n[1] Lindley, D. V. (1972). Bayesian statistics: A review. Society for industrial and applied mathematics.\n\n[2] Chaloner, K., & Verdinelli, I. (1995). Bayesian experimental design: A review. Statistical science, 273-304.\n\n[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.\n\n[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.\n\n[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.\n\n[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.\n\n[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.\n\n[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.\n\n[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.\n\n[10] Zheng, Q., Zhang, A., & Grover, A. (2022, June). Online decision transformer. In international conference on machine learning (pp. 27042-27059). PMLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and the valuable references you provided. We address your questions and points raised below.\n\n**Weaknesses**\n\n> 1. My main issue with the paper is the presentation of DUG and EDUG as novel.\n\nWe 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.\n\n> 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.\n\nYes, $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.\n\n> 3. I am not sure how the utility $u(y_{\\Xi}, a)$ is defined.\n\nOur 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. \n\n> 4. The related works section is quite short and sparse on citations.\n\nThanks 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.\n\n> 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.\n\nAgreed. We will change to “most recent BED methods leverage an amortized policy network to rapidly design experiments”.\n\n> 6. Line 179: \"to ensure the framework satisfied the permutation invariance property of sequential BED\": not all BED problems are permutation invariant.\n\nAgreed, 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.\n\n> 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?\n\nThat 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.)\n\n**Questions**\n\n> 8. Have the authors performed any ablations?\n\nWe 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$.\n\n> 9. In the decision-aware AL experiment: why does the random baseline perform as good as all the other ones?\n\nBecause 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.\n\n> 10. Could you give guidance on choosing utility functions? \n\nOur 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.\n\n**Limitations**\n\n> 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.\n\nThanks. 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.\n\n**References**\n\n[1] Lindley, D. V. (1972). Bayesian statistics: A review. *Society for industrial and applied mathematics*.\n\n[2] Chaloner et al. (1995). Bayesian experimental design: A review. *Statistical science*.\n\n[3] Kuśmierczyk et al. (2019). Variational Bayesian decision-making for continuous utilities. *Neurips*.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "3gJEIEzgWz", "reviewer": "Reviewer_7z19", "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.", "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?\n- 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?\n- 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.\n- What is the scalability of the method in terms of all relevant dimensions, e.g., dimensionality of xi, y, a, etc?\n- 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”?\n- 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?", "limitations": "The paper discusses several limitations. I am missing a discussion on the initial overhead of training, which is usually substantial in amortized methods.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "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.\n- 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.\n- 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.\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive comments and thoughtful questions. We address your remarks and questions below.\n\n> 1. The presentation of p(y_Xi | h_t) between Eq 3 and 4 is partially unclear to me…\n\nThanks 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.\n\n> 2. What if the true predictive distribution is far away from Gaussian?\n\nThat’s a good point. Please refer to the response for question 9 of reviewer WHLG.\n\n> 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.\n\nThe 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. \nTo 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.\n\n> 4. Toy example was hard for me to understand at first. Is this just a standard BO task?\n\nThanks 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).\n\n> 5. How problematic is the fact that we randomly generate some designs from the design space? \n\nThat’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.\n\n> 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?\n\nWe 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. \n\n> 7. Given that (non-decision aware) amortized BED methods exist, why are the benchmarks only comparing against non-amortized methods? \n\nWe 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.\n\n> 8. What is the scalability of the method in terms of all relevant dimensions?\n\nWe 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. \n\n> 9. Figure 4: Can you clarify why you feel that the improvements are indeed “substantial gains”?\n\nWe 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. \n\n> 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?\n\nThe 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.\n\nRegarding 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.\n\n> 11. I am missing a discussion on the initial overhead of training.\n\nThanks 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.\n\n**References**\n\n[1] Garnelo et al. (2018). Conditional neural processes. *ICML*.\n\n[2] Müller et al. (2021). Transformers can do Bayesian inference. *ICLR*.\n\n[3] Settles (2012). Active learning. *Springer*.\n\n[4] Markou et al. (2022). Practical conditional neural processes via tractable dependent predictions. *ICLR*.\n\n[5] Bruinsma et al. (2023). Autoregressive conditional neural processes. *ICLR*.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We thank the reviewers for their thoughtful comments and suggestions. We are glad to see that all reviewers have a positive view of the paper. Specifically, the reviewers agreed on the following strengths of the paper:
* **Relevance**: Zp5w: “tackles an important problem”. Ctfm: “relevant and interesting topic”. 7z19: “This is an innovative paper with high practical relevance”.
* **Novelty**: Reviewers Zp5w, WHLG, and Ctfm agree that “the proposed architecture is novel”.
* **Good presentation**: WHLG: “The paper is clearly written… a pleasure to read”.  Ctfm: “clarity of writing is excellent”. 
* **Experiments**: Zp5w: “The empirical evaluation is comprehensive”. WHLG: “The results indicate superior and stable performance…”.

**New experiments and results**

* As suggested by reviewers WHLG and 7z19, we added an amortized method as a benchmark for top-$k$ optimization experiments. Specifically, we chose PFNs4BO [1], a transformer-based amortized model designed for hyperparameter optimization. The final results are shown in Figure R1 of the rebuttal PDF. Our method outperforms PFNs4BO across all four tasks, as PFNs4BO does not consider the downstream task (i.e., top-$k$ optimization). We will update the results in the revised paper.
* As asked by reviewers Zp5w, Ctfm, and 7z19, we ran an extra ablation study to evaluate the impact of the discount factor $\alpha$. When $\alpha=0$, our objective is purely myopic. We observed that, compared to other non-myopic settings, the policy did not learn to query designs effectively. This could be due to the sparse nature of rewards in this task; when the algorithm only considers immediate rewards, it struggles to learn the value of actions that lead to future rewards. We included the experimental results in the rebuttal PDF (Figure R2).
* As suggested by reviewer Zp5w, we included a real-world experiment on retrosynthesis planning. Specifically, our task is to assist chemists in identifying the top-$k$ synthetic routes for a novel molecule, as selecting the most practical routes from many random routes generated by the retrosynthesis software can be troublesome. We trained our TNDP on a novel meta-dataset, including 1500 molecules and their routes collected by our collaborators. In this task, experimental design refers to selecting a route for a novel molecule to query its score. The downstream task is to recommend the top $k$ routes based on the collected data. Due to limited time, we compared only TNDP and the random search. The results in Figure R3 of the rebuttal PDF show that the utility of TNDP is significantly better than that of random search. We will include more baselines and provide a detailed problem description in the final paper.

**Clarification to Section 2.2**

We thank the reviewers for raising thoughtful questions regarding the definition of $p(y_\Xi | h_t)$ and the utility function. We acknowledge that this section lacks some details, and we would like to provide further explanations here.

Our utility function is defined based on the measured outcomes ($y$) instead of the state of the world ($\theta$), as many downstream tasks directly rely on the predictions of outcomes for decision-making (see P4L124 in our paper as an example). It is a natural extension of the traditional definition of utility by marginalizing out the posterior distribution of $\theta$, a similar decision-theoretical setup can be found in [2]. As we are switching the belief about the state of the world (posterior) to the outcomes (posterior predictive) and to keep as much information as possible about the state of the world, we need to evaluate $\theta$’s effect on all points of the design space, thus, we define the utility based on $p(y_\Xi | h_t)$, which is a stochastic process that defines a joint predictive distribution of outcomes indexed by the elements of the design set $\Xi$, given the current information $h_t$. We formulate the decisions in terms of this stochastic process, which differs from traditional utility based on individual observations, such as those defined in [3, 4]. A familiar example of our framework may be a decision process that depends on the observed values of a Gaussian process simultaneously evaluated at a large number of points. For example, in top-$k$ optimization, the goal is to select $k$ hyperparameter settings from a predefined finite set that maximize the cumulative accuracy. In this task, estimating the predictive distribution of a single hyperparameter setting is not sufficient for making the optimal decision. We need to determine the optimal decision based on the predictive distributions of all candidate hyperparameter settings. We adhere to the standard definitions of decision theory, but the entities now are stochastic processes instead of individual observations. 

Our architecture simultaneously amortizes two tasks. The first task is to amortize the predictive distribution needed for maximizing the utility during inference, which is similar to the goal of neural processes. When we can accurately predict $p(y_\Xi | h_t)$, we can make optimal decisions. For example, if we can accurately predict the outcomes corresponding to all hyperparameter settings, we can directly determine the optimal set of hyperparameters. The second task is to amortize the design of experiments. Our goal is to enable the neural network to propose more informative designs, thereby allowing more accurate prediction of the outcome and facilitating optimal decision-making.

We will include the above explanations in the revised paper.

**References**

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

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

[3] Lindley (1972). Bayesian statistics: A review. *Society for industrial and applied mathematics*.

[4] Chaloner & Verdinelli (1995). Bayesian experimental design: A review. *Statistical science*.

Meta-review: The paper presents an amortized decision-aware Bayesian experimental design framework that explicitly prioritizes maximizing downstream decision utility. The paper also introduces Transformer Neural Decision Process (TNDP), an architecture capable of instantly proposing the next experimental design by inferring the downstream decision. The reviewers were generally positive about the submission and highlighted the novelty of the BED framework and the proposed architecture, and they appreciated the strength of the empirical results.
Decision: Accept (poster)