Amortized Decision-Aware Bayesian Experimental Design

Authors: daolang  huang
yujia  guo
luigi  acerbi
samuel  kaski
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-08-28 14:43:14

Abstract: Many critical decisions 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 two tasks, showing that it can deliver informative designs and facilitate accurate decision-making.

Main paper: 
 introduction: 
 A fundamental challenge in a wide array of disciplines is the design of experiments to infer unknown properties of the systems under study [9, 7] . Bayesian Experimental Design (BED) [17, 8, 24, 22] is a powerful framework to guide and optimize experiments by maximizing the expected amount of information about parameters gained from experiments, see Fig. 1(a) . 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 time-consuming. This limitation has led to the development of amortized BED [10, 14, 5, 6] , a policy-based method which leverages a neural network policy trained on simulated experimental trajectories to quickly generate designs, as illustrated in Fig. 1(b) .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, such as in personalized medical diagnostics [4] . Previous amortized BED methods do not take down-the-line decision-making tasks into account, 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. [16] for variational approximations in Bayesian inference, provides a concept that adjusts the inference process to capture posterior regions critical for decision-making tasks. Inspired by this concept, we consider integrating decision-making directly into the experimental design process to align the proposed experimental designs more closely with the ultimate decision-making task.In this paper, we propose an amortized decision-making-aware BED framework, see Fig. 1(c ). We identify two key aspects where previous amortized BED methods fall short when applied to downstream decision-making 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. 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 do not fully amortize the decision-making process. To address this, we propose a novel Transformer neural decision process (TNDP) architecture, where the system can instantly propose informative experimental designs and make final decisions. Finally, we train under a non-myopic objective function that ensures decisions are made with consideration of future outcomes. We empirically show the effectiveness of our method through two tasks.2 Decision-aware BED 
preliminaries and problem setup: 
 In this paper, we consider scenarios in which we design a series of experiments ξ ∈ Ξ and observe corresponding outcomes y to inform a final decision-making step. The experimental history is denoted as h 1:t = {(ξ 1 , y 1 ), ..., (ξ t , y t )} and we assume a fixed experimental budget with T query steps. Our objective is to identify an optimal decision a * from a set of possible decisions A at time T .To make decisions under uncertainty, Bayesian decision theory [3] provides an axiomatic framework by incorporating probabilistic beliefs about unknown parameters into decision-making. Given a taskspecific utility function u(θ, a), which quantifies the value of the outcomes from different decisions a ∈ A when the system is in state θ, the optimal decision is then determined by maximizing the expected utility under the posterior distribution of the parameters p(θ|h 1:t ).In many real-world scenarios, outcomes are stochastic and it is more typical to make decisions based on their predictive distribution p(y|ξ, h 1:t ) = E p(θ|h1:t) [p(y|ξ, θ, h 1:t )], such as in clinical trials where the optimal treatment is chosen based on predicted patient responses. A similar setup can be found in [15, 28] . Thus, we can represent our belief directly as p(y Ξ |h 1:t ) ≡ {p(y|ξ, h 1:t )} ξ∈Ξ , which is a stochastic process that defines a joint predictive distribution of outcomes indexed by the elements of the design set Ξ, given the current information h 1:t . Our utility function is then expressed as u(y Ξ , a), which is a natural extension of the traditional definition of utility by marginalizing out the posterior distribution of θ. The rule for making the optimal decision is then reformulated as:a * = arg max a∈A E p(yΞ|h1:t) [u(y Ξ , a)].(1)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.Definition 2.1. Given a historical experimental trajectory h 1:t-1 , the Decision Utility Gain (DUG) for a given design ξ t and its corresponding outcome y t at step t is defined as follows:EQUATIONDUG measures the improvement in the maximum expected utility from observing a new experimental design, differing in this from standard marginal utility gain (see e.g., [12] ). The optimal design is the one that provides the largest increase in maximal expected utility. At the time we choose the design ξ t , the outcome remains uncertain. Therefore, we should consider the Expected Decision Utility Gain (EDUG) to select the next design, which is defined asEDUG(ξ t ) = E p(yt|ξt,h1:t-1) [DUG(ξ t , y t )].The one-step lookahead optimal design can be determined by maximizing EDUG with ξ * = arg max ξ∈Ξ EDUG(ξ). 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.To avoid the expensive cost of optimizing EDUG, we propose using a policy network that directly maps historical data to the next design. This approach sidesteps the need to iteratively optimize EDUG by learning a design strategy over many simulated experiment trajectories beforehand.Our architecture, termed Transformer Neural Decision Process (TNDP), is a novel architecture building upon the Transformer neural process (TNP) [20] . It aims to amortize both the experimental design and the subsequent decision-making. A general introduction to TNP can be found in Appendix A.The data architecture of our system comprises four parts: A context setD (c) = {(ξ (c) i , y (c) i )} t i=1contains all past t-step designs and outcomes; A prediction setD (p) = {(ξ (p) i , y (p) i )} np i=1consists of n p design-outcome pairs used for approximating p(y Ξ |h 1:t ). The output from this head can then be used to estimate the expected utility; A query set D (q) = {ξ (q) i } nq i=1 consists of n q candidate experimental designs being considered for the next step; Global information GI = [t, γ] where t represents the current step, and γ encapsulates task-related information, which could include contextual data relevant to the decision-making process.TNDP comprises four main components, the full architecture is shown in Fig. 2(a) . At first, the data embedder block f emb maps each set of D to an aligned embedding space. The embeddings are then concatenated to form a unified representation E = concat(E (c) , E (p) , E (q) , E GI ). After the initial embedding, the Transformer block f tfm processes E using attention mechanisms that allow for selective interactions between different data components, ensuring that each part contributes appropriately to the final output. Fig. 2 (b) shows an example attention mask. The output of f tfm is then split according to the specific needs of the query and prediction head λ = [λ (p) , λ (q) ] = f tfm (E).The primary role of the prediction head f p is to approximate p(y Ξ |h 1:t ) at any step t with a family of parameterized distributions q(y Ξ |p t ), where p t = f p (λ (p) t ) is the output of f p . We choose a Gaussian likelihood and train f p by minimizing the negative log-likelihood of the predicted probabilities:EQUATIONwhere p i,t represents the output of design ξ (p) i at step t, µ and σ are the predicted mean and standard deviation split from p.Lastly, the query head f q processes the embeddings λ (q) from the Transformer block to generate a policy distribution over possible experimental designs. The outputs of the query head, q = f q (λ (q) ), are mapped to a probability distribution π(ξ (q) t |h 1:t-1 ) via a Softmax function. To design a reward signal that guides the query head f q in proposing informative designs, we first define a singlestep immediate reward based on DUG (Eq. ( 2)), replacing the true predictive distribution with our approximated distribution:EQUATIONThis 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-making. To address this, we employ the REINFORCE algorithm [30] . The final loss of f q can be written as:EQUATIONwhereR t = T k=t α k-t r k (ξ (q)k ) represents the non-myopic discounted reward. The discount factor α is used to decrease the importance of rewards received at later time step. ξ (q) t is obtained through sampling from the policy distribution ξ (q) t ∼ π(•|h 1:t-1 ). The details of implementing and training TNDP are shown in Appendix B. 
toy example: targeted active learning: 
 We begin with an illustrative example to show how our TNDP works. We consider a synthetic regression task where the goal is to perform regression at a specific test point x * on an unknown function. To accurately predict this point, we need to actively collect some new points to query. The design space Ξ = X is the domain of x, and y is the corresponding noisy observations of the function.Let Q(X ) denote the set of combinations of distributions that can be output by TNDP, we can then define decision space to be A = Q(X ). The downstream decision is to output a predictive distribution for y * given a test point x * , and the utility function u(y Ξ , a) = log q(y * |x * , h 1:t ) is the log probability of y under the predicted distribution.During training, we sample functions from Gaussian Processes (GPs) [23] with squared exponential kernels of varying output variances and lengthscales and randomly sample a 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 in Appendix C.Results. From Fig. 3 , we can observe that the values of π concentrate near x * , meaning our query head f q tends to query points close to x * to maximize the DUG. This is an intuitive example demonstrating that our TNDP can adjust its design strategy based on the downstream task. Step t In traditional optimization tasks, we typically only aim to find a single point that maximizes the underlying function f . However, instead of identifying a single optimal point, there are scenarios where we wish to estimate a set of top-k distinct optima, such as in materials discovery [18, 27] .In this experiment, we choose hyperparameter optimization (HPO) tasks and conduct experiments on the HPO-B datasets [1] . The design space Ξ ⊆ X is a finite set defined over the hyperparameter space and the outcome y is the accuracy of a given configuration. Our decision is to find k hyperparameter sets, denoted as a = (a 1 , ..., a k ) ∈ A ⊆ X k , with a i ̸ = a j . The utility function is then defined as u(y Ξ , a) = k i=1 y ai , where y ai is the accuracy corresponding to the configuration a i . We compare our methods with five different BO acquisition functions: random sampling (RS), Upper Confidence Bound (UCB), Expected Improvement (EI), Probability of Improvement (PI), and an amortized method PFNs4BO [19] . We set k = 3 and T = 50. Our experiments are conducted on four search spaces. All results are evaluated on a predefined test set. For more details, see Appendix D.Results. From the results (Fig. 4 ), our method demonstrates superior performance across all four meta-datasets, particularly during the first 10 queries. 
discussion and conclusion: 
 In this paper, we introduced a decision-aware amortized BED framework with a novel TNDP architecture to optimize experimental designs for better decision-making. Future work includes conducting more extensive empirical tests and ablation studies, deploying more advanced RL algorithms [26] to enhance training stability, addressing robust experimental design under model misspecification [22, 13, 25] , and developing dimension-agnostic methods to expand the scope of amortization. 
 
 Appendix: CNPs [11] are designed to model complex stochastic processes through a flexible architecture that utilizes a context set and a target 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 a decoder. The encoder is a DeepSet architecture to ensure permutation invariance, it transforms each context point individually and then aggregates these transformations into a single representation that captures the overall context. The decoder then uses this representation to generate predictions for the target set, typically employing a Gaussian likelihood for approximation of the true predictive distributions. Due to the analytically tractable likelihood, CNPs can be efficiently trained through maximum likelihood estimation.Transformer Neural Processes (TNPs), introduced by [20] , enhance the flexibility and expressiveness of CNPs by incorporating the Transformer's attention mechanism [29] . In TNPs, the transformer architecture uses self-attention to process the context set, dynamically weighting the importance of each point. This allows the model to create a 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 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 [20] . for t = 1 to T do 6: Setξ (q) t ∼ π t (•|hD (c) = h 1:t , D (q) = D (q) \ {ξ (q) t } ▷ update D 10:Calculate r t (ξ (q) t ) with u(y Ξ , a) using Eq. ( 4) ▷ calculate reward 11:end for 12:R t = T k=t α k-t r k (ξ (q) k ) ▷ calculate cumulative reward 13:Update TNDP using L (p) (Eq. ( 3)) and L (q) (Eq. ( 5)) 14: end while 15: At deployment, we can use f (q) to sequentially query T designs. Afterward, based on the queried experiments, we perform one-step final decision-making using the prediction from f (p) .The embedder f emb is 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 embedderf (ξ)emb , an outcome embedder f emb are 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 4.• Activation function: ReLU is used as the activation function for the hidden layers.The time step embedder f (t) emb is a discrete embedding layer that maps time steps to a continuous embedding space of dimension 32.For the decision-aware active learning task, since the design space contains both the covariates and the decision, we use four embedders: a covariate embedder f emb is another discrete embedding layer. For context embedding E (c) , we first map each ξ (c) i and y (c) i to the same dimension using their respective embedders, and then sum them to obtain the final embedding. For prediction embedding E (p) and query embedding E (q) , we only encode the designs. For E GI , except the embeddings of the time step, we also encode the global contextual information λ using f (x) emb in the toy example and the decision-aware active learning task. All the embeddings are then concatenated together to form our final embedding E.We utilize the official TransformerEncoder layer of PyTorch [21] (https://pytorch.org) for our transformer architecture. For all experiments, we use the same configuration, which is as follows: The prediction head, f p is an MLP that maps the Transformer's output embeddings of the query set to the predicted outcomes. It consists of an input layer with 32 hidden units, a ReLU activation function, and an output layer. The output layer predicts the mean and variance of a Gaussian likelihood, similar to CNPs.For the query head f q , all candidate experimental designs are first mapped to embeddings λ (q) by the Transformer, and these embeddings are then passed through f q to obtain individual outputs. We then apply a Softmax function to these outputs to ensure a proper probability distribution. f q is an MLP consisting of an input layer with 32 hidden units, a ReLU activation function, and an output layer.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 the REINFORCE algorithm, we select a discount factor of α = 0.99.In our toy example, we generate data using a GP with the Squared Exponential (SE) kernel, which is defined as:The full version of this work can be found at: https://arxiv.org/abs/2411.02064.Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (NeurIPS2024). 



Reviews: review: title: title: An Interesting Concept of Decision-Aware Bayesian Experimental Design, but Lacks Strong Empirical Justification
main_review: main_review: This paper introduces a novel approach to Bayesian Experimental Design (BED) by directly incorporating downstream decision-making utility into the design process. The approach is well-motivated by Bayesian decision theory, and the authors propose an amortized architecture to predict experimental designs that maximize a non-myopic, discounted cumulative Decision Utility Gain (DUG), which is a nice to the literature. However, while the framing is appealing, the paper has several shortcomings in its empirical validation and comparisons with relevant baselines.

Strengths:

- The paper presents a clean and novel setup that directly considers downstream decision-making utility in the experimental design process, which is an innovative contribution.
- The decision-aware framing is a compelling extension of Bayesian decision theory, and the concept of Decision Utility Gain (DUG) is well justified within this framework.
- The authors successfully motivate the challenge of optimizing DUG and propose a Transformer Neural Decision Process (TNDP) to bypass the complexity of traditional lookahead optimization methods, which is a creative architectural choice.

Weaknesses:

- Marginal improvement: The proposed method shows no statistically significant improvement compared to baseline methods, with only marginal gains observed. The lack of substantial empirical improvement undermines the paper's claim of its practical relevance.
- Inadequate experimental comparisons: The experimentation feels incomplete and lacks an apple-to-apple comparison with similarly motivated lookahead acquisition functions, such as the Knowledge Gradient (KG). Since TNDP aims to bypass the complexity of optimizing EDUG by directly predicting the next design, it should have been compared to lookahead acquisition functions, which also perform a form of lookahead optimization.
- Pretraining and transfer learning: The proposed TNDP architecture requires pretraining, but the paper does not adequately discuss how the pretraining data could be leveraged by GP-based baselines, for example, through transfer learning or task-specific priors.
- Overly simple decision problem: Despite the paper’s emphasis on decision-making, the decision problem it tackles seems overly simple, which raises the question of whether the proposed method offers significant advantages over standard Bayesian Optimization (BO) methods. It is unclear how the decision-aware approach diverges significantly from a regular optimization problem that current BO methods could handle adequately.
- Limited diversity in experiments: While the paper includes four experiments, they are all variations of Hyperparameter Optimization (HPO) problems, offering little variety in experimental tasks. This lack of diversity in experimentation weakens the generalizability of the proposed method.
- No reported training time: The training process for pretrained models like TNDP is arguably non-trivial in terms of computational resources and time, yet the paper does not report any details on the training time. This omission leaves an important factor in the practical applicability of the method unaddressed.
title: title: Review
main_review: main_review: The paper introduces an innovative approach to Bayesian Experimental Design (BED), termed Amortized Decision-Aware BED. The key novelty lies in the integration of decision-making objectives directly into the experiment design process, ensuring that experimental designs maximize downstream decision utility.

**Strengths**
The method seem to be novel, as most prior BED methods focus purely on parameter inference without explicitly considering the decision-making tasks. TNDP architecture design is sensible by utilizing transformer and special masking designs. Experimental results on two tasks: a toy example of targeted active learning and a more complex Top-k Hyperparameter Optimization task. The TNDP demonstrates superior performance in terms of utility maximization compared to several baseline methods, including Random Sampling, UCB, EI, PI, and PFNs4BO.

**Weakness**
Although the idea is interesting, architectural novelty of TNDP is very limited, as well as the theoritical results on the algorithm (soundness and does it actually maximize EDUG). Emperical results are good for lower steps, but it would be a stronger paper if more emperical results are provided since this is an emperical paper.