Amortized Bayesian Decision-Making for Simulation-Based Models

Bayesian inference is a popular tool for parameter estimation. However, the posterior distribution might not be sufficient for decision-making. Bayesian Amortized Decision-Making is a method that learns the cost of data and action pairs to make Bayes-optimal decisions.

Simulation-Based Inference (SBI) is a powerful tool for estimating the posterior distribution $p(\theta \mid x)$ over the parameters $\theta$ of a simulator, given observed data $x$. However, the goal is sometimes not the posterior itself but but making a decision in a downstream task based on the inferred posterior distribution. Often, these decisions are associated with a cost, which one wishes to minimize. This is where Bayesian decision-making comes into play, aiming to choose actions that minimize the expected cost under uncertain conditions.

Approximate Bayesian Decision Making

Given an observation $x_o$ and a posterior $p(\theta \mid x)$, Bayesian decision making provides the action with the lowest cost, averaged over the distribution of parameters.

$$ a = \underset{a \in \mathcal{A}}{\operatorname{arg\ min}} \int c(\theta, a) p(\theta \mid x_o) d\theta $$

The function $c(\theta, a)$ quantifies the cost of taking an action $a$ if the true parameters $\theta$ of the system were known but is flexible enough to allow for different cost structures as well. The true posterior $p(\theta \mid x)$ is usually not known and is approximated using SBI, usually by a conditional density estimator $q_{\phi}(\theta \mid x)$. The quality of the decision hinges then on the accuracy of this posterior approximation.

Figure 1. [Gor23A], Figure 1. Illustration of the difference between the proposed method for amortized Bayesian decision making and numerical approximation of the integral over the cost function, weighted by the estimated posterior. Each marker describes the averaged cost difference for ten observations.

To address this challenge, [Gor23A] introduce Bayesian Amortized Decision Making (BAM) in the context of SBI. Within the same setting as neural posterior estimation (NPE), BAM learns the cost of data and action pairs. Instead of averaging the cost over the posterior, the proposed method requires only one forward-pass through the network. Therefore, BAM performs amortized Bayesian decision-making.

BAM aims to estimate the expected cost $\mathbb{E}_{p(\theta | x_o)}[c(\theta, a)]$ under the true posterior. This is achieved by sampling from the joint distribution $(\theta, x) \sim p(\theta, x)$ and an action distribution $a \sim p(a)$. A feedforward neural network $f_{\omega}(x,a)$ is then used to regress the cost of parameter sets, utilizing a Mean-Squared-Error loss.

$$ \mathcal{L}(\omega) = \mathbb{E}_{p(\theta, a)p(a)} \left[ \left(f_{\omega}(x,a) - c(\theta, a)\right)^2 \right] $$

To define the cost function, the authors assign zero cost where $a = \theta_{\text{true}}$ and increase the cost the more $a$ deviates from $\theta_{\text{true}}$. The exact manner in which the deviation is penalized depends on the task at hand. In real-world scenarios, the cost function could also include the economic cost of the action. Taking epidemiology as an example, the cost function could include a cost for vaccination, quarantine, or even a lockdown.

The authors prove that BAM accurately yields the expected cost and frames Bayesian decision making as a regression task, offering an efficient alternative to weighted parameter selection. In contrast to a Monte-Carlo approximation of the integral (NPE-MC) with samples from the (approximate) posterior

$$ \mathbb{E}_{p(\theta | x_o)}[c(\theta, a)] \approx \frac{1}{N} \sum_{i=1}^{N} c(\theta_i, a). $$

BAM directly learns the expected cost $f_{\omega}(x,a)$. It thereby circumvents the need to learn the full posterior distribution and repeatedly evaluating the cost function $c(\theta, a)$.

Numerical Experiments

Figure 2. [Gor23A], Figure 3. Comparison of the proposed BAM method and the Monte-Carlo based approximation of the expected cost against the incurred cost using the true posterior of each task.

To illustrate the effectiveness and limitations of BAM, the authors compare it with a Monte-Carlo-based approach (NPE-MC) across various benchmark tasks typical to SBI, such as the Lotka-Volterra and SIR models. Additionally, they present an application to a real-world scenario in medical neuroscience. The comparison on the SBI tasks (Figure 2), based on six different simulation budgets, reveals that the Monte-Carlo variant necessitates a larger sample size to achieve a quality of solution comparable to BAM, suggesting significant savings in simulation resources when only Bayes-optimal solutions are sought.

References

In this series