Amortized Bayesian Experimental Design for Decision-Making

TL;DR

The paper addresses designing experiments to optimize downstream decisions, noting that traditional BED maximizes the expected information gain $EIG(\xi)=\mathbb{E}_{p(y|\xi)}[H[p(\theta)]-H[p(\theta|\xi,y)]]$, which neglects decision utilities. The authors propose Decision Utility Gain (DUG) and the Transformer Neural Decision Process (TNDP), a fully amortized architecture with dual heads for proposing designs and approximating the predictive distribution, enabling non-myopic decision-centric optimization. Through toy regression, decision-aware active learning, and top-$k$ hyperparameter optimization, TNDP outperforms baselines by delivering more informative designs and better final decisions, while enabling faster deployment. This framework offers rapid, decision-aligned experimentation with broad applicability to domains such as personalized medicine and robust optimization under resource constraints.

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.

Figures (6)

• 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.
• Figure 2: Illustration of TNDP. (a) An overview of TNDP architecture with input consisting of 2 observed design-outcome pairs from $D^\text{(c)}$, 2 designs from $D^\text{(p)}$ for prediction, and 2 candidate designs from $D^\text{(q)}$ 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 $f_{\text{tfm}}$.
• 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.
• 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.
• 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.

Theorems & Definitions (2)

• Definition 3.1
• Definition 3.2