Amortized Bayesian Decision Making for simulation-based models
Authors: Mila Gorecki
Jakob H. Macke
Michael Deistler
Venue: Accepted by TMLR ()

Abstract: Simulation-based inference (SBI) provides a powerful framework for inferring posterior distributions of stochastic simulators in a wide range of domains. In many settings, however, the posterior distribution is not the end goal itself --- rather, the derived parameter values and their uncertainties are used as a basis for deciding what actions to take. Unfortunately, because posterior distributions provided by SBI are (potentially crude) approximations of the true posterior, the resulting decisions can be suboptimal. Here, we address the question of how to perform Bayesian decision making on stochastic simulators, and how one can circumvent the need to compute an explicit approximation to the posterior. Our method trains a neural network on simulated data and can predict the expected cost given any data and action, and can, thus, be directly used to infer the action with lowest cost. We apply our method to several benchmark problems and demonstrate that it induces similar cost as the true posterior distribution. We then apply the method to infer optimal actions in a real-world simulator in the medical neurosciences, the Bayesian Virtual Epileptic Patient, and demonstrate that it allows to infer actions associated with low cost after few simulations.

Published in Transactions on Machine Learning Research (05/2024)
Amortized Bayesian Decision Making for simulation-based models
Mila Gorecki
mile.grrehiShhcching....
Social Foundations of Computation, Mox Planck Institute for Intelligent Systems, Tübingen Tübingen AI Center
Jakob H. Macke
jakob. wac.
Machine Learning in Science, Excellence Cluster Machine Learwing, University of Tüihingen Empirical Inference, Max Planck Institute for Intelligent Systems, Tübingen Tühingen Al Center
Michael Deistler
michael dichhct.diisterrrnni--ebbeeeee
Machine Learning Science, Excellence Cluster Machine Learning, University of Tübhingen Tübingen AI Center
Reviewed on OpenReview: https: // openreview. net/ forum? id- BQE4MTAJCE
Abstract
Simulation-based inference (SBI) provides powerful framework for inferring posterior distributions of stochastic simulators in wide range of domains. In many settings, however, the posterior distribution is not the end goal itself - rather, the derived parameter values and their uncertainties are used as basis for deciding what actions to take. Unfortunately, because posterior distributions provided by SBI are (potentially crude) approximations of the true posterior, the resulting decisions can be suboptimal. Here, we address the question stochastic simulators, and how can
of how to perform Bayesian decision making circunvent the need to compute an explicit approxination to the posterion Our method trains a neural network on simulated data and can predict the expected cost given any data and action, and can, thus, be directly used to infer the action with lowest cost. We apply our method to several benchmark problems and demonstrate that it induces similar cost as the true posterior distribution. We then apply the method to infer optimal actions a real-world simulator in the medical neurosciences, the Bayesian Virtual Epileptic Patient, and demonstrate that it allows to infer actions associated with low cost after few simulations.
1 Introduction
Many quantities in the sciences and engineering cannot be directly observed, but knowledge these quantities (often called 'parameters' 0) crucial make well-informed decisions on how to design experiments, claim discoveries, or assign medical treatments. Take, for example, a mechanistic model in a medical context that describes specific brain function, but whose parameters observations from patients, various methods have been proposed to determine likely parameter values and their uncertainty. Ultimately, however, the goal in this scenario is to assign medical diagnoses or to decide whether an intervention is needed. Therefore, we have to turn our beliefs about parameter values into concrete and accurate actions. To do so, Bayesian decision theory provides a framework to select optimal docisions (or "actions") under noisy and polentially non-linear measurements the parameters (Berger 2013). Given the measurement X, and cost function that assigns a cost each action (for example corresponding to the
not directly casily measurable. Given
'Coneurrently work, another work on Bayesian decision making for simulators has been published Alaing et 2023 Discussion similarities and differences of works.

---

Published in Transactions on Machine Learning Research (05/2024)
training
prior
simulated
BAM:
feedforward NN
costs c(e,
action
inference
observation
BAM: feedforward NN
optimal action
expected costs
action
for every action
neural density
posterior
estimator
cost estimation
action
Figure 1: Amortized Bayesian Decision Making, Given (stochastic) simulator and prior over parameters, we propose to learn neural network which directly estimates the expected cost associated with any possible observation and action, i.c. a 'Bayesian Amortized Decision Making' (BAM) network. (a) Analogous to simulation-based methods like Neural Posterior Estimation (NPE), BAM starts by generating a database of (parameter, data) pairs by drawing prior samples and simulating them. BAM then uses a fcodforward neural network trained to directly prodict the estimated costs of taking a particular action given observations. This way, BAM circumvents the computation of an explicit posterior approxination. (b) At inference, BAM requires only a single pass through the network estimate the expected posterior costs of each data-action- pair (Zo,a). In contrast, NPE builds an explicit parameter estimate of the posterior. This can be used to estimate the expected cost of given action through Monte-Care estimation (NPE-MC). For both algorithms, we select the action which leads to the lowest expected costs.
associated risk of each intervention), Bayesian decision making chooses the action which induces the lowest cost, averaged over the distribution of parameters given the measurements p(0)x.).
a² arg min / c(0.a)) p(0)x.)d0.
a
The main difficulty associated with Bayesian decision making is to infer the posterior distribution (())).)
over (potentially many) parameters. The posterior distribution
(0)x,)) x p(x.|0)p(0
is intractable for all but the simplest models, and requires methods such as Markov-Chai Monte Carlo (MCMC) or variational inference to approximate. However, these methods require that the likelihood can be evaluated, but many models in the sciences can only be simulated (Gonççaves et al.) 2020 Cranmer et al. 2020), i.e. their likelihood can only be sampled from. Tn addition, such methods require to re-run inference for new data, which is computationally expensive. Contrary that, recent methods from the field of "simulation-based inference' (SBI) (Cranmer al. 2020) require large upfront cost to generate database of simulations (but likelihood evaluations), and they then infer the posterior distribution given any data without further simulator runs (Hermans et al. 2020 Gonçalves et al. 2020 Radev et al. 2020).
Here, we first. investigate how faithful the actions induced by a particular amortized SBI algorithm, *Neural Posterior Estimation' (NPE) Papamakarios and Murray 2016 Greenberg et al.) 2019), are. We demonstrate that, in cases where the posterior is inaccurate ςLueckmann et al.) 2021 Hermans et al.) 2022), the induced opttmal action will also be (largely) inaccurate, and that this failure can occur even for simulators with few free parameters [Lacoste Julien al.] 2011). Tn addition, while NPE amortizes simulation and
2

---

Published in Transactions on Machine Learning Research (05/2024)
training, it still has to compute the (potentially expensive) cost of several (parameter, action) pairs for every observed data, which incur substantial additional cost at decision time.
To address these limitations, we propose a method which we term Bayesian Amortized decision Making (BAM) (Fig. D. Like NPE, BAM uses an (upfront generated) database of simulations to train a neural network. Unlike NPE, however, BAM circumvents the need to learn the (potentially high-dimensional) posterior distribution explicitly and instead only requires to train a feedforward neural network which is trained to directly predict the expected costs for any data and action. At inference time, BAM requires only one forward pass through the neural network, without the requirement of computing the cost by averaging over the posterior. Because of this, BAM can perform amortized Bayesian decision making in milliscconds even in cases where the cost function itself is expensive to compute.
We define decisions on four benchmark tasks (Lueckmaan et al., 2021 and evaluate both algorithms on all of those tasks. We demonstrate that, on tasks where NPE is known to perform poorly (and returns a poor estimate of the posterior), our alternative method, BAM, can lead to significantly better decisions, sometimes leading to reduction of cost by almost order of magnitude. We then apply both methods to a real-world problem from medical neuroscience: given potentially epileptic measured brain activity, we infer which treatment option would be most suitable. We demonstrate that, given a specified loss function, both algorithms can identify optimal decisions with only a few simulations.
2
Background
2.1
Problem setting
We consider a stochastic simulator with free parameters 0 and assume access to prior distribution p(0)- Running the (potentially computationally expensive) simulator allows generating samples corresponding to the likelihood x ~ p(x(0). We assume that the likelihood is only defined implicitly through the simulator and cannot be evaluated. Given observed data X, we aim to obtain the Bayes-optimal action
a = arg / c(O,a) p(0)x.)d0
where we used the posterior p(0)x,), defined as
x00)).) x p(x, ())(O),
and we defined a cost function c(O,a) a) (Schervish 2012). This cost function quantifies the consequences of taking a particular action a if the true parameters of the system were known. The cost function can be an (casy to compute) analytical expression of parameters and actions, but it can also require (computationally expensive) computer simulations (Alsing et al. 2023). The Bayes-optimal action found by minimizing over actions a, sampled from a distribution p(a). We choose p(a) such that its support covers all possible actions, i.e. it also contains the true optimal action. This is reasonable, as the range of available measures is usually known to the decision maker.
To evaluate the performance of different decision making algorithms, we use the 'incurred cost' of an action. Given a synthetically generated observation x, and its ground truth parameters 0st, the incurred cost of an action a defined as
C c(0st a).
2.2 Neural Posterior Estimation
The posterior distribution p(0) is intractable for all but the simplest models. Neural Posterior Estimation (NPE) a method that overcomes this by using samples from the joint distribution p(0,x), obtained by sampling 0 from the prior and simulating to learn parametric density for the posterior distribution Paaaaaaaaioss and Murray 2016). Concretely, NPE trains deep neural density estimator by minimizing its
negative log-likeihood:
N log q (0,(x)).
L(Φ) = N
Σ
3

---

We discuss advantages and Published in Transactions on Machine Learning Research (05/2024)
where Φ are learnable parameters. After training, the deep neural density estimator can be evaluated for any observation X, and will return an estimate for the posterior distribution, which can be sampled and whose log-probs proobabbbyy can be computed for any 0. Crucially, once the deep neural density estimator has been trained, it can be applied to any observation without running further simulations and without retraining, thereby amorlizing the cost of inference.
3 Methods
We propose two methods obtain amortized Bayesian decisions (Fig. disadvantages of both methods in Sec.
3.1 Neural Posterior Estimation and Monte-Carlo Sampling (NPE-MC)
The first method is a direct extension of Neural Posterior Estimation (NPE) and serves as a baseline. We use two properties of NPE, namely that it is (1) amortized across observations and (2) returns posterior estimate that can be sampled quickly (1000s of samples in milliseconds in typical applications). The method proceeds in three steps: (1) train NPE obtain an amortized neural density estimator 96 (p(0)x) x) and (2) draw M samples from the approximate density qs xx) and (3) compute a Monte-Carlo estimate for the expected cost for actions
M
e x.) ((0,aa)
R
Σ c(O,,a) a)
M
If the action a is low-dimensional then we evaluate actions experiments use a 1D action and we therefore obtain the optimal action via grid search. If the action is high-dimensional (or if the cost function c(O, a) is expensive to compute) then one can use generic optimization method to identify the optimal action a" (c.g., genetic algorithms, or gradient descont if the cost function is differentiable). We term this algorithm NPE-Monte- Carlo (NPE-MC).
dense grid. All low-dimensional
For inference step (2) and (3) have be repeated for every observation zo and action, resulting in M #Zo #actions cost evaluations. If evaluating the cost function c(0,a) a) is expensive, this comes with considerable computational costs. At the same time, decreasing M or the number of actions to compute the cost for can significantly impact selection quality of the optimal action.
We note that every SBI approach which returns (approximate) posterior samples could be used in order to estimate the expected cost with a Monte Carlo estimate. This includes approaches like Approximate Bayesian Computation (ABC) (Csilléry et al., 2010) and neural likelihood (Papamakarios et al. 2019) or likelihood-ratio estimation (Hermans al.) 2020 Miller et al. 2022). Here, we chose NPE because fully amortizes the cost of inference, whereas Neural Likelihood Estimation (NLE) and Neural Ratio Estimation (NRF) require additional MCMC steps and ABC requires re-computing distances for every observation. We additionally report results using ABC and NLE for posterior estimation, referred to as ABC-MC and NLE-MC respectively.
3.2 Bayesian Amortized Decision Making (BAM)
The approach outlined above returns the correct action the neural density estimator indeed captures the true posterior distribution and if the Monte-Carlo estimate provides a sufficiently accurate estimate of the expected costs. However, it may seem overly complicated to infer the full posterior distribution if one is only interested in the expected cost induced by the posterior. In addition, the above method still requires computing the expected cost from Monte-Carlo simulations for every observation, which can bocome prohibitively expensive if the cost. function c(O,a) a) is expensive to evaluate (e.g., if one has to run a simulator to evaluate c(0, a), Alsing et al. (2023)) ). To overcome these two limitations, we introduce a method which (1) does not require estimating the full posterior and (2) directly returns an estimate for the expected cost, thereby circumventing the need to recompute the cost for every observation (Fig. I We term this method Bayesian Amortized decision Making (BAM).
4

---

Published in Transactions on Machine Learning Research (05/2024)
Algorithm 1: Bayesian Amortized decision Making (BAM)
Inputs: prior p(0); simulator with implicit likelihood p(x(0), cost function c(x, a), distribution of actions p(a), number of simulations N feedforward NN fs with parameters Φ, NN learning rate 77, observations (not i.i.d.) x (k.I...K)
Outputs: action a* with minimal expected cost given an observation.
Generate dataset:
sample prior and simulate: 0,x ~00), x, p(x(0:))aa...
Training:
while not converged do
C =0
for (0,x) in {(); x: do
a ~ p
C = N ( a) c(0, a))²
1 - Φ ll Adamm VoL)
Inference:
for k E [1, ..., K] do a" min fo (xb (M)
// find optimal action, e.g., with grid search
Like NPE, BAM is trained on samples from the joint distribution p(0, Unlike NPE, however, it employs a feedforward neural network that takes data and independently sampled actions a ~ p(a) as inputs and aims to predict the cost of the associated underlying parameter set c(0,a). We train this leedforward neural network with Mean Sraan-Squrred- Error (MSE) loss and thereby can recover the expected cost Ep(0)x.)(c((,, a)) under the true posterior for every data-action-pair:
Proposition 1. Let p(0,x) be the joint distribution over porameters and data and let p(a) be distribution of actions whose support covers the range of permissible actions. Let c(f, a)) ) be cost function and fo(:) parameterized function. Then, the loss function L(Φ) Ep( (ff a) c(0 a))² | is minimized if and only if, for all E supp(p(x)) and all a E supp(p(a)) we have J (x, a) = xx(0 (((Oa)))
Proof in Appendix Sec. A2 The algorithm is summarized in algorithm for p(a) to cover the full support of such that all possible action can be observed during training. We further note that, like NPE-MC, BAM requires final phase to search for the optimal action that minimizes the expected cost. Importantly, BAM requires only a single forward pass for every action and no further evaluations of the cost function, which renders this step computationally cheap.
We note that it important
4 Results
4.1
Illustrative example
Inspired by the nuclear power plant example by Lacoste Julien al. (2011), we illustrate NPE-MC and BAM simple example with one-dimensional parameter space and observations, but which induces bimodal posterior (Fig. b).. To illustrate potential failure modes of NPE-MC, we parameterize the posterior using neural spline flow with two splines which fails to accurately capture the true, bimodal posterior and assigns too little probability mass lo the left mode (Fig. We then define a cost function ((0,a) which assigns zero cost the action is optimal and whose cost increases monotonically with distance from the true parameter.
Crucially, the cost function asymmetric: for large 6, the cost is high even if the estimation error is small, whereas for low 6, the cost remains low for larger region of 9 (Fig. 22). See Appendix Fig. A3 for more parameterizations of the NPE posterior.
5

---

Published in Transactions on Machine Learning Research (05/2024)
b
incurred costs
expected costs
ground truth
%
p(0)x,))
NPE-MC
small large 0
BAM
B
8
action a
action
Figure 2: Illustrative Example. (a) Joint distribution of parameter and data (gray) and an observation (green). (b) True posterior distribution p p())x,) and an estimate obtained with NPE with NSF with two
splines. (c) Cost function for two values of 0 (dotted and full). (d) Expected costs function of the action obtained by the ground truth posterior (green) and as estimated with BAM (orange) and NPE-MC (with
NSF, blue). Both NPE and BAM were trained with a simulation budget of 5k. See Appendix Fig. A3 for more parameterizations NPE.
Due this asymmetry of costs, the approximation to the posterior obtained with NPE systematically underestimates the optimal action, and the optimal action chosen by NPE-MC incurs much larger cost than the true posterior (Fig. 21).
4.2 Benchmark decisions
Next, we systematically evaluated the performance of BAM, NPE-MC and ABC-MC on tasks where the ground truth posterior is available. We used the toy example introduced above and three previously published simulators with ground truth posteriors [Lueckmann et al. 2021). For all tasks, we defined cost functions on top of which to perform Bayesian decision making (details Appendix Sec. A4 All cost functions used for the evaluation are based on flipped bell-shaped curve whose width depends on parameter 0:
(0 (0-a)² a)²
c(0.a) = 1 exp
where parameters and actions are normalized to take values in the same range, = (0 mmnnm) //min//(0 maz - 0 minn), à = (a- amin) ) (amaz amin) and the width of the bell curve is determined by w = (ʃ )
Toy example: A toy simulator with ID O and ID x (Fig, 22). We use the cost function described above with f = 2, m = 2. We note that, for the benchmark results, we used a neural spline flow as approximate density Durkka et al. 2019
Linear Gaussian: Simulator with 10D Gaussian prior and linear Gaussian likelihood, leading to a Gaussian posterior. We use the a similar cost function as for the toy example, but with f = 0.5 and m 2 and offset manl to increase sensibility of the cost function towards small and large parameter values.
model is epidemiological model that describes the basic spreading behaviour of an infectious disease that is transmitted from human to human over time. The simulator is parameterized by two parameters, the contact rate 3 and the recovery rate 7. The simulator returns the number of infected individuals I at 10 evenly-spaced points in time. For the decision task we aim to estimate the reproduction rate as the ratio Rg = L. We assume that the value of around and, therefore, define a cost function which is more sensitive deviations from the optimal action for Ro (ʃ 2, m 2, o 1, details in Appendix Sec. A4).
SIR: The SIR
particularly crucial to estimate
The Lotka- Volterra model describes the oscillating dynamics between two species in a predator-prey relationship. These dynamics are determined based on four parameters: the prey population grows at rate a and shrinks at rate 3. In the presence of prey, predators are born at rate and die at rate 7- We evaluated four different cost functions with f ==3, 3, m = 2, where each individual cost function assesses only the accuracy of the estimated marginal distribution (i.c., the cost function ignores three out of four parameters). Analogous the toy example, the costs function is asymmetric and depends on the underlying true parameter value. The higher the parameter value, the more accurate the prediction of the optimal action
Lotka- Volterra:

---

Published in Transactions on Machine Learning Research (05/2024)
Example Gaussian Linear
SIR
Lotka-Volterra Lotka-Volterra B
Lotka-Votterra Lotka-Votterra 6
Latka-Volterra all
0.4
BAM
ditt incurred cost
ABC-MC
NPE-MC
0.2
10
10"
10
10"
10
4
10
10
10"
10"
10"
10"
10"
10"
10"
10"
10"
10
10
simulations
Figure 3: Benchmarking decision tasks. Gap between the incurred cost of optimal actions computed from the ground truth posterior and the predicted optimal actions of BAM (orange), NPE-MC (blue) and ABC-MC (red) as a function of the simulation budget. Results are averaged across ten observations and five seeds. Shaded areas show the standard error of the mean across seeds. Columns: four different simulators (with four different decision tasks for Lotka- Volterra).
Table 1: Runtimes for NPE-MC and BAM. Exemplary runtimes for the ceati-ddmensioonll Lotka-Von Volterra task measured on a NVIDIA GeForce RTX 2080 Ti GPU. Training dataset size is = 10k, inference runtime is computed based on 10k (data, action) pairs. For NPE, M = 1000 samples were drawn from the posterior. Mean and standard deviation are based on five seeds. method
training
inference per (x, a) pair
NPE-MC
1041.0525 + 230.2177s
81. $1.4303 + 5.1675
BAM 17.5475 + 1.8723
0.6816 + 0.0408 us
has to be incur small costs. Intuitively, the action entails various measures to intervene on the rate of shrinkage or growth, this means that overestimation of the rate should be prevented.
Lotka-Vone Volterra (4d): In addition to the low-dimensional tasks described above, we add multi-dimensional setup based on Lotka- Lotka-Volters Volterra. Using all four parameters 0 {a, 8, 8), also define four-dinensional action space and define the costs to be the mean incurred costs over all parameters, where individual cost functions are characterized in the same way as described above.
For all lasks, we trained NPE-MC, ABC-MC and BAM with 6 different simulation budgets: 500, 1k, 5k, 10k, 50k, and 100k and averaged results over five random seeds. We then computed the Bayes optimal action for ten different obscrvations with all algorithms outlined in Sec. 3 computed the incurred cost associated with these actions (Sec. 2.1 and compared this to the incurred cost of the optimal action obtained from the ground truth posterior Fig. For NPF-MC, we computed the expected costs based on M = 1000 samples from the deep neural density estimator (but using M = 10 and M = 100 impacts the results only weakly, see Appendix Fig. A1]. For ABC-MC, we used Rejection ABC with only accepting the top 100 samples for the posterior approximation. For BAM, we resampled one action per (0,x) x) pair at every epoch (keeping the same action every epoch deteriorates performance, see Appendix Fig. A1 See Appendix Fig. A2 for additional results using NLEMMC in place of to NPE-MC.
We first applied NCE-MC, ABC-MC and BAM to the statistical task of the 1D toy example and the linear Gaussian simulator. On the simple 1D toy task all algorithms perform well and lead to essentially optimal actions after around 1k simulations. Similarly, on the linear Gaussian task, the algorithms perform very well and lead to actions with little associated cost. NPE-MC slightly outperforms BAM on this task. We expect that this is because the neural density estimator employed by NPE-MC and ABC-MC has learn an affine mapping from data to parameters (because the task linear Gaussian task), whereas BAM has to learn a non-linear function of the actions (which used during training of BAM).

---

Published in Transactions on Machine Learning Research (05/2024)
b
distance to ground truth BAM NPE-MC
ground truth
NPE-MC
BAM
actions on actions on actions on acãons on
20
40
60
80
100
NPE-MC better
20
1.0
0.6
40
60
expected costs
80
0.2
0.6
0.0
0.4
20
-0.2
40
80
0.0
-0.6
BAM better
0
20
40
60
80
2
4
6
90
2
6
9 10
6 8 9
2
90
observation #
observations
Figure 4: Detailed comparison of NPE-MC and BAM Lotka-Volteree Volterra task. (a) Expected costs induced by the ground truth posterior (left), NPE-MC (middle), and BAM (right), for four different cost functions (rows) based on the Lotka-Volteria simulator and a simulation budget of 10k simulations. The x-axis shows the index of the observation, taken from [Lueckmann et al. 2021 (b) Difference in expected costs between NPEMMC and BAM.
We then applied all algorithms the more complex simulators of the SIR model and the Lotka- Volteraa simulator. On both of these simulators and for both algorithms, the induced costs are significantly higher compared to the previous tasks. For the SIR model, the actions obtained by BAM induce significantly less cost compared to both NPE-MC and ABC-MC. In order to achieve a similar level of induced cost, NPE-MC often requires one order of magnitude more simulations than BAM. Similarly, for Lotka- Volterra, BAM outperforms both NPEMMC and ABC on all decisions and for all simulation budgets. For fixed simulation budget, the cost induced by BAM can be almost one order magnitude less than for NPE-MC (c.g. on the SIR lask for 10d simulations).
Finally, we applied NPE-MC and BAM to the four-dimensional variant of Lotka Volterra. For this task, we found that BAM consistently outperforms NPEMMC for lower simulation budgets. For very large simulation budgets, we found NPE-MC to improve over BAM. Runtimes for training and inference show that BAM outperforms NPE-MC in terms of computational costs (Table
In order to gain insights into the strong performance of BAM compared to NPE-MC, we evaluated the expected costs for individual observations and for all possible actions for BAM and NPF-MC for the Lotka- Volterra simulator (Fig. 4. Across all four docision tasks, BAM closoly matches the expected costs induced by the ground truth posterior for all actions, whereas NPE-MC often fails capture the expected costs. In particular, for almost every action and observation, BAM provides closer approximation to the ground truth expected cost than NPE-MC (Fig. Ap). This failure of NPF-MC occurs because NPE does not accurately capture the true posterior, and is not due to the Monte-Carlo estimation procedure (Appendix Fig. A1
Overall, our results demonstrate that BAM performs similarly to NPE-MC on simple statistical tasks. On more complex and nonlincar simulators, BAM outperformed NPE-MC, sometimes by large margin. This indicates that, on tasks where estimating the posterior difficult in itself, it is possible save a large fraction of simulations if only the Bayes-optimal action is desired (as compared to the full posterior).
8

---

Published in Transactions on Machine Learning Research (05/2024)
Observations and NPE predictives Healthy Zone (HZ)
b
NPE posterior estimate
c
Cost function
C HZ)
ground
0.00
NPE post. predictives
-5
C(e, PZ)
-2.00
0.1
50
C(0,
Propagation Zone (PZ)
0
-5
npz nEZ
fast-variable
0.00
Xiii
8
-2.00
Epileptogenic Zone (EZ)
d
Comparison of incurred costs
MVMMMMM
Incurred costs
0.1
0.00
NPE-MC
BAM
-2.00
0
20
40
60
80
100
10
10
10"
10
1
10
10
10"
10"
10"
10
10
10"
Time (sec)
# simulations
Figure 5: Decision making for the Bayesian Virtual Epileptic Patient. (a) Three observations (green, different rows) and posterior predictives obtained with NPE (different shades of blue) for each of the observations. Each row corresponds to one of the zones of the data when classifying by excitability value of the brain region (healthy, propagation, epileptogenic). (b) Marginals and pairwise marginals of the posterior distribution obtained with NPE given an observation from the epileptogenic zone. True parameter values are shown in red. (c) The cost function for predicting 0 to be in one of the three zones (zones shown in gray scale, one decision per row). (d) Incurred costs of NPE-MC and BAM, averaged over 1k prior predictives (left), as well split by the decision made. Costs are shown for simulation budgets between 500 and 100k simulations. Shaded areas show the standard error of the mean across sceds.
4.3 Bayesian Virtual Epileptic Patient
Finally, we performed Bayesian decision making on a real-world task from medical neuroscience. The framework, termed the Bayesian Virtual Epileptic Patient (BVEP), allows to simulate neural activity in a connected network brain regions to model opilepsy spread (Hashemi et al., 2020; Jirsa et al. 2017). Previous work has demonsttated that NPE close to an approximate ground truth (obtained with Hamiltonian Monte- Monte-Carlo) in this model (Hashemi al.] 2023). The dynamics are controlled by a simulator called 'Epileptor' and in this work, focus single isolated brain region. Based on the level of excitability n) (one of the free parameters) simulations either lead to simulated data that resembles healthy brains to activity of brains during epileptic seizure (Fig. 5)).
The simulator has four free parameters and generates time series which is reduced into ten summary statistics (details in Appendix Sec. A5 NPE posterior Fig. 5)). The goal to predict, given observations, which of three actions is most appropriate (Hashemi et al. 2023). The first. action corresponds intervention (i.e., a healthy brain region, healthy zone, HZ), the second action corresponds to further tests (unclear if the brain region will be recruited during an epileptic seizure, propagation zone, PZ), and the third action corresponds to intervention such as surgery (if the brain is clearly unhealthy and triggers seizures autonomously, epileptogenic zonc, EZ). We assume that, if the excitability of the brain tissue were known exactly, then one could exactly infer the optimal action and that the optimal action would induce zero 'cost!. For simplicity, we also assume that the cost a miscclassification (i.e., predict an action which does not match the excitability region the true paramcters) induces a constant cost (Fig 5)). We use this simple cost function to demonstrate that both NPEMMC and BAM can generally solve this task and to avoid making (potentially controversial) assumptions.
no

---

Published in Transactions on Machine Learning Research (05/2024)
We train NPE-MC and BAM on this simulator with simulation budgets ranging from 500 to 100k. For training BAM, we sample actions from a discrete distribution over the three possible actions. After training, we evaluate the expected costs under the three possible actions for both algorithms (we use 1k Mlnte-Carlo samples for NPE). BAM and NPE-MC both converge with only a few thousand simulations and lead to low expected cost (Fig. 1 l.ft). We then evaluated both algorithus for cases where the ground truth parameter lies in either of three zones (HZ, PZ, EZ). Both algorithms perform similarly well in each zone, and have the highest cost in cases where the ground truth parameter lies in the propagation zone.
Overall, these results demonstrate that NPE-MC and BAM can both be used to perform Bayesian decision making on the BVEP real-world simulator and that the resulting actions can be obtained with few simulations and induce low cost.
5 Discussion
We presented methods to perform amortized Bayesian decision making for simulation baaset models. The first method, NPF-MC, uses the parametric posterior returned by neural posterior estimation to infer the optimal action with a Monte-Carlo average. However, this method requires to compute the cost for many parameter sets and actions for every observation, and its performance can suffer if the parametrie posterior approximation is poor. Using other methods in place of NPE to estimate the posterior distribution comes with the additional disadvantage that they not fully amortize the cost of inference. To overcome the limitations of NPE-MC, we developed a second method, BAM (Bayesian Amortized decision Making), which circunvents the need to compute an explicit posterior approximation and directly returns the expected cost. We demonstrated that this method can largely improve accuracy over NPE-MC on challenging benchmark tasks and we demonstrated that both mothods can infer decisions for the Bayesian Virtual Epileptic Patient.
Related Work Most applications of Bayesian decision making rely using Markov-chain Monte Carlo variational inference methods to obtain estimate of the posterior distribution. Like NPE, variational inference requires a parametric estimate of the posterior distribution, and Lacoste-Julien et al. (2011) have studied how such an approximation can impact the quality of downstream actions and proposed to take the decision task into account for training the variational posterior. Indeed, if the decision task and the associated loss- funecion are already known at inference time, then the 'quality' of a posterior approximation can be assessed by the accuracy with which it allows computation of expected costs and optimal actions- some approximation errors will not impact the expected cost, whereas others might have substantial impact. This makes it possible to quantify the quality of different approxinations, and potentially also directly optimize the posterior approximations for this, an approach referred to as 'loss-calibration' (Lacoste-Julic) et al., 2011 b. Several works applied this principle to classification tasks with Bayesian Neural Networks (Cobb et al. 2018
regression tasks for variational inference (Kušmicrzyyk et al. 2019) and in the context of importance sampling (Abbasnejad et al. 2015). Instead of adapting the training process of variational inference, which requires lo recompute the posterior approximation whenever specifics of the decision task change, Kušmierczyyk et al. (2020) and Vadera et al. (2021) propose post-hoc correction procodures that calibrate the docision making process itself (given an approximate posterior). We expect that, in cases where the parametric posterior approximation of NPF-MC is poor, similar methods could be used to calibrate the NPF posterior for the decision task at hand. BAM, on the other hand, circumvents the need for an explicit posterior approximation and thus does not make any parametric assumptions about the posterior distribution.
Finally, another work on the topic of Bayesian decision making for simulators has been developed concurrently to our work (Alsing et al., 2023). Alsing al. (2023) focus on simulators which incorporate a mechanistic model for the actions, thereby requiring to simulate in order to estimate the costs. For estimating the cost, our method BAM parallels their method, and they introduce distribution of cost (instcad of directly estimating its expected valuc). Alsing et al. (2023) also introduce a method to focus on inference on a particular observation with sequential inference (Papamakarios and Murray 2016 Lueckmann et al. 2017 Greenberg et al., 2019 Deistler et al. 2022). whereas we always sample parameters from the prior to avoid correction steps in the loss function. We provide empirical results across eight benchmark decision tasks and for several simulation budgets (Sec. 4.2 and we empirically demonstrate
additional algorithm to estimate the
the ability to amortize across observations, whereas Alsing et. al. 2023 provide empirical results for single
10

---

Published in Transactions on Machine Learning Research (05/2024)
observations on two benchmark problems. Finally, we provide a proof of convergence for BAM (Appendix Sec. A2 and we apply the methods to Bayesian Virtual Epileptic Patient (Sec. 4.3)
Limitations For any real-world problem, it is difficult to assess the accuracy of the action proposed by either NPE-MC BAM. In particular, if the real-world observation differs from the simulated data (i.e., if the observation is misspecified)), then the neural networks employed by NPE-MC and BAM might give poor results, and any estimated performance on the simulated data might no longer be correct (Cannon et al. Gao et al. 2023
2022
In the following paragraphs we briefly outline further considerations regarding the space of actions and the cost. function more detail.
Action Space In Bayesian decision theory the space of actions is assumed to be known to the decision maker at inference time. Unlike NPE-MC, BAM requires the space of avaílable actions to be fixed already at training time as BAM is specifically tailored to the goal of decision making. We thus introduced a distribution of actions p(a) which defines the range of actions across which to amortize, and require the support of p(a) to cover all possible actions. We provided proof that BAM converges to the expected costs for any action in the support of p(a) and any data in the support of p(r), regardless of the exact form of p(a). Nonetheless, we note that, different distributions p(a) might influence the results in practice. Being amortized over the actions, BAM can can provide insights into the overall cost landscape. This can provide a helpful source of information to guide the final choice of the action, especially if there are additional trade-offs or constraints associated with different actions that are not considered in the simulator or the cost function.
Cost Function The cost function determines which inaccuracies hurt the decision process take the decision task into account, motivating the call for loss-calibration. Another aspect to consider is the computational cost of evaluating the cost function as the proposed methods differ on when to compute the costs of (data, action) pairs: Training of NPE-MC is cost-agnostic, but at inference, the costs have to be computed for every data Toy action a and the number of samples drawn from the posterior estimate M. When evaluating the cost function is expensive, e.g.; in scenarios where evaluating the cost function requires simulate (Alsing) ct al., 2023), this comes with considerable computational costs. In contrast, BAM requires evaluations of the cost function only during training, in total #epochs (Nerain + Neal) evaluations. At inference, BAM requires no further evaluation of the cost function, but #actions many passes through the network, which makes this step is computationally very cheap. For high-dimensional actions in particular, the amortization over the decision process makes BAM largely more efficient than NPE-MC.
Whether to use NPE or BAM We presented two approaches to perform amortized Bayesian decision making, NPE-MC and BAM. Their key difference is that NPE-MC builds a full parametric estimate of the posterior distribution (which is independent of the cost function or the specific action), whereas BAM circumvents the requirement to build a posterior approximation and immediately regresses onto the expected cost for every action the support of p(a) and for every observation in the support of p(r). The key advantages of NPEMMC are that it provides a full posterior that is interpretable and can be used for multiple downstream applications (one of which might be Bayesian decision making), and that it does not require retraining when the cost function or the space of actions changes. Contrary to that, we expect BAM to shine especially in applications where the posterior is complex and hard to estimate (as evidenced by the results on and SIR). In addition, BAM immediately regresses onlo the expected cost. for every action in the support of p(a) and every observation in the support of p(x), whereas NPE-MC requires additional Monte-Carlo simulations and additional evaluations of the cost function for every observation, action, and parameter set.
Conclusions We developed BAM, an algorithm to perform amortized Bayesian decision making. We demonstrated that the approach recovers the Bayes-optimal action for a diverse set of tasks and that it can be used to perform optimal decision making in an example from the medical neurosciences. In addition, we provided a systematic evaluation of commonly used methods from the field of simulation- based inference, namely NPF, ABC and NLE, in the context of decision making on a diverse set of tasks.
11

---

Published in Transactions on Machine Learning Research (05/2024)
Acknowledgements
MG and MD are supported by the International Max Planck Rescarch School for Intelligent Systems (IMPRS- IS). MD and JHM are funded by the Machine Learning Cluster of Excellence, EXC number 2064/1-3907211 This work was supported by the Tübingen AT Center (Agile Research Funds), the German Federal Ministry of Education and Research (BMBF)): Tübingen Al Center, FKZ: 011S18039A. We thank Julius Vetter, Guy Moss, and Manuel Glocckler for discussions and comments on the manuscript.
References
Ehsan Abbasnejad, Justin Domke, and Scott Sanner. Loss-calibrated Monte Carlo action selection. In Twethyy-Nnnth AAAT Conference on Artificial Intelligence, 2015. 10
Justin Alsing, Thomas DP Edwards, and Benjamin Wandelt. Optimal iimmaaaion-baseed Bayesian decisions.
Xrxiv preprint @Xix:2811.057// 2023.
James O Berger. Statistical decision theory and Bayesion analysis. Springer Science & Business Media, 2013.
Patrick Cannon, Daniel Ward, and Sebastian M Schmon. Investigating the impact model misspecification
in neural simulation-basee inference. arXiv preprint arXiv:2209.018. 2022. 11
Adam D Cobb, Stephen J Roberts, and Yarin Gal. Loss-calibrated approximate inference in Bayesian Neural Networks. arXie preprint arXiv:1805.03900 2018. 10
Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-based inference. Proceedings of the National Academy of Sciences, 117(48)330055 30062, 2020. 2
Katalin Csilléry, Michael GB Blum, Oscar Gaggiotti, and Olivier François. Approximate Bayesian computation (ABC) in practice. Trends in ecology & evolution, 25(7):410 418, 2010.
Michael Deistler, Pedro J. Goncalves, and Jakob H. Macke. Truncated proposals for scalable and hassle-free free simulation-based inference. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022, I0
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows Advances in Neural Information Processing Systems, 32, 2019. 6
Richard Gao, Michael Deistler, and Jakob H Macke. Generalized Bayesian inference for scientific simulators via amortized cost estimation. arXiv preprint 1XXx:2305..1208 2023.
Pedro J Gonçalves, Jan-Matthis Lueckmann, Michael Deistler, Marcel Nonnenmacher, Kaan Öcal, Giacomo Bassetto, Chaitanya Chintaluri, William F Podlaski, Sara A Haddad, Tim P Vogels, et al. Training deep neural density estimators to identify mechanistic models of neural dynamics. Elife, 9re56261, 2020.
David Greenberg, Marcel Nonnenmacher, and Jakob Macke. Automatic posterior transformation for likelihood- free inference. In International Conference on Machine Learning, pages 2404-2414, PMLR, 2019. 2 10
Meysam Hashemi, AN Vattikonda, Viktor Sip, Maxime Guye, Fabrice Bartolomei, Michael Marmaduke Woodman, and Viktor K Jirsa. The Bayesian Virtual Epileptic Patient: A probabilistic framework designed to infer the spatial map of epileptogenicity in a personalized large-scale brain model of epilepsy spread. Neurolmage, 217:116839, 2020. 17
Meysam Hashemi, Anirudh N Vattikonda, Jayant Jha, Viktor Sip, Marmaduke M Woodman, Fabrice Bartolomei, and Viktor K Jirsa. Amortized Bayesian inference generative dynamical network models of epilepsy using deep neural density estimators. Neural Networks, 163:178-194. 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Proceedings of the IEFE conference on computer vision and paltern recognition, pages 770- (700-78, 2016. 16
12

---

Published in Transactions on Machine Learning Research (05/2024)
Joeri Hermans, Volodimir Begy, and Gilles Louppe. Likelihood-free MCMC with amortized approximate ratio estimators. In International Conference on Machine Learning, pages 4239 4248. PMLR, 2020. 2
Jocri Hermans, Arnaud Delaunoy, François Rozet, Antoine Wehenkel, Volodimir Begy, and Gilles Louppe. A crisis in simulation-based inference? Beware, your posterior approximations can be unfaithful. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. 2
Viktor K Jirsa, William C Stacey, Pascale P Quilichini, Anton I Ivanov, and Christophe Bernard. On the nature of seizure dynamics. Brain, 137(8):2210 2230, 2014. 17
Viktor K Jirsa, Timothée Proix, Dionysios Perdikis, Michael Marmaduke Woodman, Huifang Wang. Jorge Gonzalez-Martinnz, Martinez, Christophe Bernard, Christian Bénar, Maxime Guye, Patrick Chauvel, et al. The virtual epileptic patient: individualized wraie-brain models of epilepsy spread. Neuroimage, 145:377-388. -88, 2017.
17
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. 16
Tomasz Kušmierczyk,, Joseph Sakaya, and Arto Klami. Variational Bayesian decision maaiiin for continuous utilitics. Advances in Neural Information Processing Systems, 32, 2019. 10
Tomasz Kušmierczyk, Joseph Sakaya, and Arto Klami. Correcting predictions for approximate Bayesian inference. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 4511-4518,
2020. 10
Simon Lacoste-Julien, Ferenc Huszár, and Zoubin Ghahramani. Approximate inference for the loss-calibrated Bayesian. In Geoffrey Gordon, David Dunson, and Miroslav Dudík, editors, Proceedings of the Fourleenth International Conference on Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Research, pages 416 424, Fort Lauderdale, FL, USA, 11 13 Apr 2011. PMLR. 2
5 10
Jan-Matthis Lueckmann, Pedro J Goncalves, Giacomo Bassetto, Kaan Öcal, Marcel Nonnenmacher, and Jakob H Macke. Flexible statistical inference for mechanistic models of neural dynamics. Advunces in Neural Information Processing Systems, 30, 2017. I0
Jan-Matthis Lueckmann, Jan Boelts, David Greenberg, Pedro Goncalves, and Jakob Macke. Benchmarking simulation- inference. In International Conference on Artificial Intelligence and Statistics, pages 343-351. PMLR, 2021.
16
Benjamin K Miller, Christoph Weniger, and Patrick Forré. Contrastive neural ratio estimation. Advances in Neural Information Processing Systems, 35:3262-322 32788, 2022. 4
George Papamakarios and lain Murray. Fast e-free inference of simulation models with Bayesian conditional
density estimation. Advances in Neural Information Processing Systems, 29, 2016. 2 10
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. Advances in Newral Information Processing Systems, 30, 2017. 15
George Papamakarios, David Sterratt, and Iain Murray. Sequential neural likelihood: Fast likelihood inference with autoregressive flows. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 837-848. PMLR, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019. 15
Stefan T Radev, Ulf K Mertens, Andreas Voss, Lynton Ardizzone, and Ullrich Köthe. Bayesflow: Learning complex stochastic models with invertible neural networks. IEEE brunsactions on newral nelworks and learning systems, 33(4):1452-15 1466, 2020. 2
Mark J Schervish. Theory of statistics. Springer Science & Business Media, 2012. 3
13

---

Published in Transactions on Machine Learning Research (05/2024)
Alvaro Pejero-Canters Jan Bocits, Michael Deistler, Jan-Matthis Lueckmann, Conor Durkan, Pedro J. Gonçalves, David S. Greenberg, and Jakob H. Macke. shi: A toolkit for simulation-based inference. Journal of Open Source Software, 5(52):2505, 2020. doi: 10.21105/joss.000 02505. 15
Meet P. Vadera, Soumya Ghosh, Kenney Ng, and Benjamin M. Marlin. Post-hoc loss-calibration for Bayesian Neural Networks. In Cassio de Campos and Marloes H. Maathuis, editors, Proceedings of the 37th Conference Uncerlainly in Artificial Intelligence, volume 161 of Proceedings of Machine Learning Research, pages PMLR, 27-30 30 Jul 2021. 10
Omry Yadan. Hydra- framework for elegantly configuring complex applications. Github, 2019. 15
14

---

Published in Transactions on Machine Learning Research (05/2024)
Appendix
A1 Code and Reproducibility
Code to reproduce all experiments, including the full gil commit history, is available al https://githh: hte///ccckellamoottizeeeeeedd decision- All neural networks and optimization loops are written in pytorch (Paszke et al. 2019). We tracked all experiments with hydra (Yadan 2019 For NPE, we used the implementation in the shi toolbox Tejero-Canters et al. 2020
A2 Convergence proof for proposition
The proof closely follows standard proofs that regression converges to the conditional expectation.
Proof. We aim to prove that
pa [(c(O, a) o(x,a))) 2
Exxx-p(0xx))a [(c(O, a) po0p a)])
for every function g(x,a) a). We begin by rearranging expectations:
Paxx-p(0x))a (cc0,, g(x, a)x]
pp [xx (c) g(x, a))
Below, we prove that, for a given a and for any x, the optimal g(x,a) is the conditional expectation Eg 1c000 a)]]:
Eexxp((xx (c(0,a) g(x, a)
Ex (c(0,a) a) Exr a)] Egr Eer=piox))((00'a g(x, a))
xxxx-x(0x (c(0,a) a) Fg F
(Fg x000)(0))((000 a)] xx 2)) |] +
with
X = [xxxx~(00x - Ep /c(0', a) (00 a)] - g(xx a))],
By the law of itcrated expectations, one can show that
X Exx 1all a) Ee
a)] a))]
The first term in the product reads difference of the same term, and, thus, is zero. Thus, since (Eor (p00 a)] g(x, a)) 2 O, we have:
Ex (c((0. g(x, a)) ((xx= a) -or a)))
Because this inequality holds for any a, the average over p(a) will also be minimized and only if g(x,a) matches the conditional expectation Epr a)] for any within the support of p(a).
A3 Experimental setup
We used mainly the same hyperparameters for all benchmark tasks and for the Bayesian Virtual Epileptic Patient task. Only the learning rate was adjusted to 0.005 for Lotka-Volteree while it was set to 0.001 for all other tasks.
For NPE, we used the default parameters of the sbi package (Tejero- Cantero et al. 2020), apart from the density estimator, for which we used a neural spline (Durkan et al. 2019). Following Hashemi al. (2023). we used masked autoregressive flow (Papamakarios et al. 2017) for the Bayesian Virtual Epileptic Patient with default parameters as set in the sbi package.
15

---

Published in Transactions on Machine Learning Research (05/2024)
For BAM, used feedforward residual network (He et al 2016 with 3 hidden layers and 50 hidden units each. We use ReLU activation functions and, for cases where we know that the expected cost is going to be positive and bounded by 1, squash the output through a sigmoid. We use the Adam optimizer (Kingma and Ba. 2015). As described in Sec. 4 the size of the training dataset varied between 500 and 100k. In all cases, the training dalaset was split. 90:10 inlo training and validation and with a batchsize of 500.
As distribution over all permissible actions, p(a), use uniform distribution over continuous actions between 0 and 100 for all benchmark tasks, (100, 100), and discrete uniform distribution over three classes for the Bayesian Virtual Epileptic Patient.
A4 Benchmark tasks
Below, we provide details on the simulators and cost functions used for benchmarking NPE-MC and BAM. Toy example: The prior follows uniform distribution
p(0) (1,55)
and the likelihood is given by:
p(x(0) = 50 + 0.50 - +
where v NO, 10). We obtained (an approximation to) the ground truth posterior via quadrature.
As cost function, we used flipped (along the y-axis) bell-shaped curve whose width decreases with increasing parameter 0:
(0 - a)
c0,aa) = exp
0 +
with c = 0.1. As the paramcter space and the action space differ, and a were rescaled to both lic within the range (0, 10].
Linear Gaussian: We used exactly the same simulator as in Lueckmann al, (2021) As cost function, used flipped (along the y-axis) bell-shaped curve whose width decreases towards the extremes of the parameter range:
- aa)
c00,a) a) = 1 -xpp 10000.5 0min)] [l
0.5
with 6 0.1. As the parameter space and the action space differ, and a were rescaled both lie within the range of [0, 10). Additionally, a shift was introduced such that the variance is widest at 5 and decreases towards both ends of the value range.
Lotka-Volterra Volterra: We used exactly the same simulator as in Lueckmann et al. (022) ). We evaluated four different cost functions, where cach individual cost function assesses only the accuracy of the estimated marginal distribution (i.e., the cost function ignores three out of four parameters). The cost. function for individual parameters 0, with i = 0, 1, is given by:
(0 a)²
(0 exp
( )a
with e 0.1. As the parameter space and the action space differ, and a were rescaled both lie within the range of [0, 10).
We used exactly the same simulator in Lueckmann et al. 2021 The cost function is based on the ration between parameters, r B and is given by: y.
SIR:
(r a)
c(r,a) = 1- exp
10- r
16

---

Published in Transactions on Machine Learning Research (05/2024)
with 0.1. As the parameter space and the action space differ, and were rescaled to both lie within the range of [0, 10). Additionally shift by 1 in parameter space was introduced such that the cost function is very sensitive around r = and variance increases with increasing distance to 1.
A5 Details on Bayesian Virtual Epileptic Patient (BVEP) simulator
The Bayesian Virtual Epileptic Patient provides a framework to simulate neural activity in connected network of brain regions to model epilepsy spread (Hashemi et al. 2020 Jirsa et al. 2017). The dynamics of individual brain regions are controlled by a simulator called 'Epileptor' that allows to reproduce dynamics of electrical activity during seizuro-like events. Hashemi et al. (2020) introduce two versions of the Epileptor, the full 6D model and a reduced variant, the 2D Epileptor. The full Epileptor model (Jirsa et al., 2014 Hashemi et al. 2020) is composed of five state variables that couple two oscillatory dynamical systems on three different time scales. Variables and Y1 associated with fast electrical discharges during seizures at the fastest time scale. Variables 22 and y/2 act on an intermediate timescale and the permittivity state variable that governs the transition between ictal, i.e. during a seizure, and interictal, i.e. in between seizures, states on a slow time scale. See Jirsa et al. (2014) and Hashemi et al. (2020) for a detailed description of the Epileptor equations. For this work, we focus on a single isolated brain region and used the reduced '2D Epileptor' model which results from applying averaging methods to make the effects of ø2 and y/2 negligible and further assuming time scale separation that 2 and y1 collapse the slow manifold (Hashemi et al. 2020
The 2D Epileptor takes four parameters to generate a time series of neural activity: describes the excitability of the tissue, controls the time scale separation and Zinst and Zinit are initial values for the state variables. See Hashemi et al. (2020) for a full description. Based on the level of excitability n,, brain regions can be categorized into three lypes (Hashemi al. 2020
Epileptogenic Zone (EZ): if 17 > 7ez, the Epileptor can trigger seizures autonomously
Propagation Zone (PZ): if 1pz 1Ez An 7 yez; the Epileptor cannol trigger seizures autonomously, but (in connected network of brain regions) be recruited during the evolution of an epileptic seizure
Healthy Zone (HZ): if y < nrz, the Epileptor does not trigger scizures
Following Hashemi et al. 2023), we set the critical value Ø = 2.05 and An 1.0, i.e. 1P2 = 1Ez An -3.05. These types of brain regions form the basis for the docision task described in Sec. 4.3
When observations are high-dimensional as it the case for the time series generated by the Epileptor, low-dimensional summary statistics are commonly used for inference. For both, NPE-MC and BAM, we used 10 timo-independent summary statistics to describe the observed neural activity. These the mean, median, standard deviation, skew, and kurtosis of the time series well higher moments (up to the 10th moment of the observation), the power envelope, scizure onset, the amplitude phase, and spectral power.
17

---

Published in Transactions on Machine Learning Research (05/2024)
A6
Supplementary figures
Toy Example Gaussian Linear
SIR
Lotka-Volterra Lotka-Volterra Lolka-Volterra Lotka-Volterra
BAM
N=1
0.5
NPE-MC
(M=10)
NPE-MC
(M=100)
0.0
NPE-MC
(M=1000)
NPE-MC
(M=10000)
N=5
0.5
0.0
difference incurred
N=10
0.5
0.0
N=100
0.5
0.0
in loop
0.0
10
10
10
10
10
10
10
10
10
10
10"
10
10"
10"
10
10
10
10 10
# simulations
Figure A1: Expected costs for sampling diverse numbers of actions per data point and for a different number of Monte-Carde (MC) samples. NPE-MC is the same for every row (blue). Different shades of blue indicate NPE-MC with different number of MC-samples. For BAM, we sampled (1,5, 5, 10, 100) actions per (0, x) pair the training datasel (rows). In the last row we sampled one new action at every epoch. This is shown in the main paper in Fig. 3
18

---

Published in Transactions on Machine Learning Research (05/2024)
Toy Example
Lolka-Volterra Lolka-Volterra Lolka-Volterra Lolka-Volteraa 6
0.4
ditt incurred cost
BAM
0.3
NPE-MC
ABC-MC
0.2
NLE-MC
0.1
0.0
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10
simulations
Figure A2: Benchmarking decision tasks. Gap between the incurred cost of optimal actions computed from the ground truth posterior and the predicted optimal actions of BAM (orange), NPE-MC (blue), ABC-MC (red) and NLE-MC (purple), averaged across ten observations and five seeds. Columns: four different simulators (with four different decision tasks for Lotka- all benchmark tasks. In particular, we found that NLE-MC outperformed NPE-MC on all benchmark lasks and also BAM for most simulation budgets. We note that it has previously been reported that NLE performs particularly well on these tasks (Lueckmann et al. 2021). These results indicate that, if the cost inference (e.g., with MCMC) and of estimating the cost (via simulation) not critical the user, then NLE-MC can provide strong alternative.
NLE-MC performed well across
NSF splines
NSF with splines
NSF with 100 splines
ground truth
p(e
BAM
NPE-MC
8
8
expected costs
action a
action 8
action
action a
MAF
MDN (2 components)
MDN (10 components)
MDN (100 components)
p(8
expected costs
action
action
action
action
Figure A3: Illustrative Example. (a) Simulator with Gaussian noise and (b) Simulator with uniform noise. First row: Truc posterior distribution p(0)x.) and an estimate obtained with NPE. Sccond row: Expected costs as function of the action obtained by the ground truth posterior (green) and as estimated with BAM (orange) and NPEMMC (blue). Columns correspond different parametrizations of the posterior estimate. Both NPE-MC and BAM were trained with simulation budget of 5k.
19


Reviews:
[{"review_id": "iu4FirTvzV", "reviewer": "Reviewer_s4Bh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- The paper is well organized and written. Introduction to subject matter is clear and concise. Contributions are presented in a logical and intuitive order. Saving a broader discussion of related works until section 5 makes the paper accessible to non-experts of the subject matter.\n- I find the empirical results to support the main claims of the paper. Below I mention some weaknesses in the set up of the experiments, but I want to emphasize that the main claims are supported. Specifically, when the posterior is difficult to estimate or if it is misspecified, BAM can output actions with noticeably lower cost than NPE. \n- To the best of my knowledge, the proof of proposition 1 is correct.\n\n**Weaknesses**\n\n- Actual run times are not reported. It would be helpful to know how the run times of NPE vs BAM, training and inference, compare. Each iteration in training BAM involves evaluating the cost, and so an iteration should be more expensive. If one of the benefits is amortizing the computation, it’s necessary to understand this effect. \n- The experiments only consider scenarios where the cost function is computationally inexpensive to evaluate. This is especially beneficial for BAM, which relies on calling this cost function repeatedly during training, while existing methods do not. When the cost function isn’t cheap to evaluate or must be approximated, it’s not clear that BAM does not break down (i.e. is training prohibitively long). This is a concern since one of the main benefits of BAM is to avoid evaluating the cost function repeatedly for many samples of the posterior at inference time. Note that even if BAM breaks down in this setting, I still maintain that there is a useful contribution from this paper. However, this is critical information to report to the community alongside being able to use BAM in the settings already considered.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work considers the problem of Bayesian Decision Making, i.e. selecting an action with the minimum expected cost, and proposes a solution that avoids explicitly approximating a posterior by fitting a neural network that predicts the expected cost of an action. In situations where the end goal is solely to find the minimal expected-cost action, the proposed method, named BAM, is computationally faster at test time since the neural network avoids having to sample the posterior to approximate the expected cost. It is instead trained to output the expected cost of an action in a single forward pass. Empirical results demonstrate that when the posterior is misspecified, an existing method’s performance in the Bayesian Decision Making task significantly deteriorates, while the same deteriorating is not observed in BAM. When a posterior is correctly specified, the performance of BAM is comparable to existing methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Critical**\n\n- The toy example presented in figure 2 articulates that severe misspecification is very detrimental to NPE but not to BAM. The toy example is great for gaining intuition, but as it stands now, this is an extreme example. Does more subtle misspecification, perhaps for this example skewness or heavy-tails, still lead to the same degradation in performance? I believe that similar figures to figure 2 in the appendix could be useful. Does the posterior need to be extremely misspecified for NPE to fail while BAM is not affected?  This is an important question for someone considering using BAM since BAM forces the decision-making process to be (and stay) risk-neutral. If only a small performance degradation in NPE is observed, it may be desirable to accept the (unknown) drop in expected performance to maintain having access to an approximate posterior and tailor the decision-making to a different risk profile.\n\n- What happens to the performance/training time of BAM when the cost function is expensive to evaluate? What if the cost function must be approximated? The contribution of BAM is clear in the situations currently simulated, but if cost is more expensive to evaluate then it’s not clear if BAM offers performance benefits. Empirical results for a scenario like this would be desired.\n\n**Not Critical**\n\n- At the end of Section 2 “…it can be applied to any observation without running further simulations and without retraining…”. To me this sentence could be interpreted as suggesting that the neural network can generalize to data not necessarily in the training set. I suggest revising the sentence if this is not implied.\n\n- Adding in the explicit run times of NPE vs BAM would be helpful.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8NnTlprinE", "reviewer": "Reviewer_z9cw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strength\n- The paper is well-written, easy to follow, I would even say pleasant to read.\n- I believe the paper may be of interest to a community of non-experts in SBI, who encounter such decision making and have a way to simulate the observation process and cost associated with decisions. \n- I have no major concerns on the experiments presented in the paper and find they support appropriately the paper's message.\n- I also like the argument that, sometimes, maybe it is reasonable to say often, learning the map from state and action directly to expected cost is an easier ML task than learning the posterior instead. It is even more salient if the cost function is expensive to evaluate (the question is, when is it?)\n\n# Weaknesses\nAlthough I enjoy reading the paper, I have several concerns regarding the relevance of publishing the paper in its current form. I am now expressing my main concerns and will suggest some potential modifications to address these in the next box.\n1. I think the introduction does not a very good job at motivating why using the first equation to make a choice is reasonable (and applicable) in real-world scenario. \n2. They are a bunch of vague statements made about the \"inaccuracy\" of NPE. Although I agree in principle with these statements, they should be clarified to explain what kind of inaccuracy may hurt the decision process (e.g., over vs under confidence) and when such inaccuracies are hard to avoid.\n3. The loss of interpretability behind deciding with BAM vs deciding with NPE-MC should be discussed. It is unclear to me also, why, if the problem targeted has high stakes we couldn't simply train NPE with more simulations. If I were to use such a strategy to make decision I would rather chose a more costly tools which provides me with some interpretability than the black box BAM.\n4. It seems the method would not scale well to higher decision space for two reasons: 1. p(a) in algo 1 might be an important design choice; 2. the grid search might be ineffective, or even dangerous, if the decision+observation spaces lead to a complicated relationship between the cost and the pair action/observation.\n5. Figure 2 looks like a caricature and seems dis-honest to NPE-MC in my opinion. It is unlikely anyone having some tiny experience in SBI would obtain such a green NPE. It is thus unclear in which extreme scenario BAM is a significantly wiser choice than NPE-MC\n6. The claim about how to solve the problem if it is high dimensional are vert hand-wavy. I do not think it is correct to say that using a generic optimisation method to search the space would be a good choice, especially given the high chance to end up in a place where the surrogate models are innacurate.\n7. I am not sure how relevant is this paper for any person already aware of SBI as, in the one hand, NPE-MC is a straightforward extension and BAM is just fitting a deterministic surrogate to simulated data.\n8. I would expect to see error bars on figure 3 and 5d.\n9. It is unclear how this problem setup applies to any real-world problem as it is unlikely modelling assumptions are as correct in the real-world as they are in simulations. The lack of real-world experiments is what causes me this worry.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Rooted in a simulation-based inference (SBI) setting, this paper presents two methods to perform Bayesian decision making when the model relating parameters $\\theta$ to observations $x$ is a black-box, maybe stochastic, simulator. The goal is to find the bayes optimal action $a^{\\star}(x)$ that minimise the expected cost of this action given an observation $x$, under the assumed model $p(\\theta \\mid x)$ implicitly defined by a prior over parameters, \n$$a^{\\star}(x) = \\arg\\min_{a} \\mathbb{E}_{p(\\theta \\mid x)}[c(\\theta, a)].$$\nThis decision strategy should enforce the expected cost over the marginal distribution of possible observation is minimal under correct modelling assumption, that is the simulator and prior represent exactly the true generative process.\n\nAuthors suggest two methods to solve this problem. The first method, named NPE-MC, trains a neural posterior estimator of $p(\\theta \\mid x)$, then, given an observation, it approximates the expected cost with Monte Carlo and use grid search to determine an estimate of the optimal action. The second method, named BAM, observes that the estimation of $p(\\theta \\mid x)$ can often be strictly more complicated than directly estimating the expected cost  of an action given an observation. Thus, BAM directly trains an estimator $f(a, x)$ that maps a pair of observation and action to the associated expected cost. After training, $f(a, x) \\approx \\mathbb{E}_{p(\\theta \\mid x)}[c(\\theta, a)]$, and obtaining the optimal action is performed by a grid search on the estimator of the expected cost of performing action $a$ given the observation $x$. Clearly, the two methods are consistent in the sense that given perfect training (and expressive enough class of surrogates) and sufficiently fine grid search, they return the Bayes optimal action.\n\nAuthors compare the two methods' performance over 4 synthetic tasks derived from standard simulation-based inference benchmarks and decision making problem on virtual epileptic patients. All tasks have low-dimensional action space.\nExperiments demonstrate that none of the two methods is strictly superior than the other on all the benchmark. Nevertheless, BAM is overall slightly less data-hungry than MC-NPE on the SBI-inspired benchmarks. Authors conclude BAM is an effective way to perform Bayesian decision making under implicit model (forward black box simulators).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would suggest the following things to improve the manuscript and convince me this a work of interest to the TMLR audience:\n1. Improve intro, motivate with real-world problems.\n2. Discuss how and when NPE fails, say why it can be difficult to address, and how BAM can better perform in such settings maybe?\nAn even better point would be to show that NPE can fail silently (marginal calibration ok, C2ST ok, etc...) under some metrics but give very bad decision process?!\n3. Discuss.\n4. Tone down your contribution and be more clear about why doing it in a higher dimension-space would add new challenges. Maybe, also explain why there are already important decision problems that are hard to solve even though the action space is low-dimensional. This goes back to remark 1.\n5. Please chance that figure, or do not cite NPE-MC there. This would be a very wrong signal send to people not familiar with NPE.\n6. Improve the discussion.\n7. This one is bit harsh, but I think improving the intro and trying to provide better insights on why applying NPE-MC may be hard could be nice? I simply believe having a real-world experiment where you can show you achieve non-obvious performance would be much stronger. I also understand this is difficult to show... But at the same time, if you cannot show that, I am not sure the paper has any interest to anyone. Indeed, I did not learn much when reading your paper and so I am skeptical another researcher would really gain anything by building on your paper to develop a method that apply to real-world scenario. \n8. Please discuss some variance of your estimators and how this can also guide a practitioner to trust or not the decision process. \n9. Really unclear how this can really be helpful in practice. In contrast to standard SBI, which is either applied on real-world data that are very well modelled as in some scientific fields, or just on simulation to understand better the modelling assumptions made by the forward model, here I do not totally see these applications. Again, it would be awesome to see a real-world experiment and I would directly remove my concerns.", "broader_impact_concerns": "I think applying this methodology to make eal-world decision would be dangerous as the method could miserably fail under model misspecification. This should probably be discussed.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XKXo0ayh3n", "reviewer": "Reviewer_NspT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**: \n- Interesting paper that proposes a novel way to do decision-making for simulation-based models. The idea of amortizing the inference of the action without explicitly approximating the posterior  distribution over the parameters of the model is interesting and novel (to the best of my knowledge, without considering the concurrent work highlighted in the paper).\n- The proposed method is simple and elegant, and it seems to work well in the experiments. \n\n\n**Weaknesses**:\n- The writing can be improved, some content is duplicated and it lacks some rigor in the mathematical definitions. \n- Contributions are unclear (or better, confusing): from reading the abstract, I would say that the contribution of the paper is _only_ BAM, but Section 3 (and Figure 1) puts forward also NPE-MC as contribution (\"Our first method is a direct extension\"). Personally, I would consider NPE-MC more of a baseline and I would clearly present it as such.\n- The experimental evaluation is a bit weak. The authors compare BAM with NPE-MC, but they do not compare with other methods. I do understand the motivation for focusing on NPE-MC, but I think it would be nice to have a comparison with at least one of the other methods mentioned (e.g. an ABC method).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a novel way to do decision-making for simulation-based models. The proposed method, called Bayesian Amortized decision Making (BAM), is based on the idea of amortizing the inference of the action without explicitly approximating the posterior distribution over the parameters of the model. The main motivation is that the posterior over the parameters is often multimodal, and therefore it is difficult to approximate it well. Instead, the authors propose to learn a mapping from the observations and actions, and then use this mapping to select the best action. The mapping is learned using a neural network, and the parameters of the model are sampled from a prior distribution. The proposed method is then compared with a baseline method on a set of experiments (from toy examples to classic benchmarks), showing that it is able to perform well in all the considered scenarios.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Armonize \"Introduction\" and \"Background\" sections. There is some duplicated content between the two sections, and it would be best if you could find a way to introduce the problem and the background content in a coherent way\n- Also related to the point above, please be a bit more rigorous in the mathematical definitions. What is $\\boldsymbol x_o$? $\\boldsymbol \\theta$? Likelihood $p(\\boldsymbol x_o|\\boldsymbol \\theta)$ and prior $p(\\boldsymbol \\theta)$ are never explicitely introducted. \n- Position of Algorithm and Figures closer to the discussion or after most of the content is discussed in the main text (e.g. Fig 1 is in page 2 and discussed in page 5)\n- The message of Figure 1 is unclear to me. It seems that the entire diagram is your _unique_ contribution, while actually it shows two different methods. What about using two different colors to indentify the part of the figure related to BAM and the one related to NPE-MC?\n- One key element of the proposed algorithm that it's never mentioned in the text (unless I'm wrong) is the prior on the actions $p(\\boldsymbol a)$ which only appears in the Algorithm. To me, this is critical and it requires a careful analysis, or at least a discussion on the effect it has on the solution.\n- Regarding the experimental setup, I think it would be best if you could be explicit on the cost functions used (since you already have them in the Appendix---a table might be enough). To be such a critical part of the evaluation, the current descriptions are a bit too general. \n- Regarding the experiments, are they averaged across multiple seeds? Is it robust to different initializations? \n\n**Minor**\n\n- \"NLE\" is not defined in the text\n- The \"Conclusions\" section is very short and it does not add much to the discussion. I would suggest to remove it and move the content to the \"Discussion\" section.\n- At the beginning of the \"Introduction\" section, when you say that \"The posterior distribution is intractable for all but the simplest models, and requires methods such as Markov-Chain Monte Carlo (MCMC) to approximate.\" which it reads as if only MCMC can be used to approximate the posterior. I would suggest to rephrase it to include also variational inference methods (which are mentioned anyways in later sections the paper).\n \n\n**Open questions**\n- I wonder if BAM can be seen as NPE on a different _joint_ model $\\{\\boldsymbol \\theta,\\boldsymbol a\\}$, with a factorized prior $p(\\boldsymbol \\theta,\\boldsymbol a)=p(\\boldsymbol \\theta)p(\\boldsymbol a)$ and likelihood $p(\\boldsymbol x_o|\\boldsymbol \\theta,\\boldsymbol a)$. Can you maybe have a quick comment on this?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All three reviewers are in agreement that this paper meets TMLR's acceptance criteria. The review process and converged in a manuscript where all reviewers are aligned in vetting the conclusions. I kindly ask the authors to fill out the runtimes as delineated in the manuscript whenever convenient.
Decision: Accept (poster)