Epidemiological model calibration via graybox Bayesian optimization

Abstract

In this study, we focus on developing efficient calibration methods via Bayesian decision-making for the family of compartmental epidemiological models. The existing calibration methods usually assume that the compartmental model is cheap in terms of its output and gradient evaluation, which may not hold in practice when extending them to more general settings. Therefore, we introduce model calibration methods based on a “graybox” Bayesian optimization (BO) scheme, to enable more efficient calibration for general epidemiological models. This approach uses Gaussian processes as a surrogate to the expensive model, and leverages the functional structure of the compartmental model to enhance calibration performance. Additionally, we develop model calibration methods via a decoupled decision-making strategy for BO, which further exploits the decomposable nature of the functional structure. The calibration efficiencies of the multiple proposed schemes are evaluated based on various data generated by a compartmental model mimicking real-world epidemic processes and COVID-19 datasets. Experimental results demonstrate that our proposed graybox variants of BO schemes can efficiently calibrate computationally expensive models and further improve the calibration performance measured by the logarithm of mean squared errors and achieve faster performance convergence in terms of BO iterations. We anticipate that the proposed calibration methods can be extended to enable fast calibration of more complex epidemiological models, such as the agent-based models.

1. Introduction

Pandemics like the recent coronavirus disease (COVID-19) have shown enormous impacts on public health worldwide. The development of computational epidemiological models is crucial for gaining a better quantitative understanding of disease spread and enabling swift decision-making to design effective mitigation strategies. In fact, well-calibrated epidemiological models can serve as a useful guide for forecasting and quantifying epidemic risk and implementing effective public health measures to control the spread of epidemic diseases (Anastassopoulou et al., 2020; Bentout et al., 2020).

Compartmental models, such as the commonly adopted SIR (Susceptible-Infectious-Recovered) model, constitute an important family of population-based epidemiological models. Models in this family take the form of Ordinary Differential Equation (ODE) systems that capture the general trends and dynamics in the subpopulations represented by different compartments. In this work, we focus on this family of models while considering a new “Quarantined” state represented by an additional compartment. This new state aims at better mimicking the dynamic trajectories of the epidemic spread, typical in airborne diseases such as seasonal flu and COVID-19, resulting in a four-compartment SIQR model (Odagaki, 2021).

To accurately model real-world physical processes, it is crucial for computer models to fine-tune their model parameters based on observed data to capture inherent attributes of the physical processes, which can not be directly or easily measured by means of physical experiments. For instance, material properties, such as porosity and permeability, are important computer inputs in computational material simulations, which cannot be measured directly in physical experiments. In the applied mathematics and computational science literature, the methods used to identify those parameters are called model calibration techniques and the resulting parameters are called calibration parameters (Kennedy & O'Hagan, 2001; Santner et al., 2003). The basic idea of calibration is to find the combination of the calibration parameters, under which the simulated computer outputs align with the observed physical data. A well-calibrated SIQR model may allow people to make more accurate predictions of pandemic risk and enable them to make better-informed decisions on how to mitigate transmission and disease impact.

The existing methods often assume that the computer model is cheap (Tuo & Jeff Wu, 2016). This means that the explicit form of the optimization objective f(x) for calibration parameters x is known and we can evaluate the first- or second-order derivatives of f(x) with respect to (w.r.t.) x. For an ODE system, deriving the explicit form of f(x) and ∇_xf(x) can be infeasible when the ODE is non-linear. Recent advances show that sensitivity analysis allows us to evaluate the gradient ∇_xf(x) and perform the gradient-descent optimization method for the estimation of x (Chen et al., 2018; Ghosh et al., 2021). Correspondingly, physics-informed machine learning (Karniadakis et al., 2021) is drawing increasing attention as it integrates data and mathematical models into deep neural networks or other regression models by enforcing fundamental physical laws.

In real-world scenarios when calibrating a computer simulation software, we are usually blind to the explicit forms of both f(x) and ∇_xf(x) and we may be able to access only the output of f(x). We classify this type of computer model as expensive. Besides, even if we know the explicit form of f(x) and ∇_xf(x), there are usually some key hyper-parameters x′, which are fixed during gradient-based optimization. Given x, the explicit form of f_x′ is intractable in most cases. Therefore, f_x′ can also be classified as expensive. Hence, it is valuable to pay attention to the calibration problem for such expensive computer models. Bayesian Optimization (BO) is one conventional approach for such scenarios when evaluating f(x) is expensive, where we use a surrogate model to approximate f(x) and the optimization process is sequential. In each BO round, we make a Bayesian decision so that some expected utility is maximized and then the corresponding output is queried from the computer model based on this decision to refine the surrogate model to guide the optimization of the objective function f(x) (Frazier, 2018a; Garnett, 2023).

This paper focuses on SIQR model calibration via BO. The remainder of this paper is organized as follows. First, we explain the configurations of the SIQR model and the model calibration setup. Then we introduce our calibration methods based on a “graybox BO” approach, instead of the usual “blackbox BO”, where the experts’ prior knowledge about the computer models is integrated into the BO formulation to improve its optimization performance. Finally, we explain how the graybox BO works for decoupled decision-making, accompanied by a new acquisition function. In our experiments, we evaluate the performance of our proposed graybox BO-based model calibration methods and compare them with traditional blackbox BO-based implementations to demonstrate the benefits of integrating prior model knowledge into the calibration process.

2. Preliminaries

2.1. SIQR model

In this model, people in a given population are classified into four compartments: Susceptible (S), Infectious (I), Recovered (R), and Quarantined (Q). Those who have not been infected are considered susceptible. Transmission occurs between susceptible and infected individuals. The number of symptomatic patients, equivalent to the number of infected individuals, decreases through treatment and/or quarantine. Recovered individuals are assumed to remain immune to further infection. Therefore, all people will be in the recovered group after an adequate period. The dynamics of the SIQR model are given by the following set of ordinary differential equations (ODEs) (Odagaki, 2021):

$$\left\{\begin{array}{ll}{\nabla }_{t}S(t)=& -\beta (t)S(t)\\ {\nabla }_{t}I(t)=& \beta (t)S(t)-\lambda (t)I(t)-\gamma I(t)\\ {\nabla }_{t}R(t)=& \gamma I(t)+\delta Q(t)\\ {\nabla }_{t}Q(t)=& \lambda (t)I(t)-\delta Q(t)\end{array}\right)$$ (1)

where S(t), I(t), Q(t) and R(t) ∈ [0, N] are the populations of compartments ‘S’, ‘I’, ‘Q’ and ‘R’ respectively. It can be verified that the ODEs guarantee the conservation of the population, S(t) + I(t) + Q(t) + R(t) = N. In this formulation, β(t) denotes the transmission rate at which susceptible individuals become infected after contact with infectious individuals; λ(t) is the quarantine/isolation rate at which infectious individuals are identified and moved into quarantine; γ is the rate for infectious individuals who are not quarantined, recover naturally; δ is the recovery rate for quarantined individuals. In the specific model configurations used in the experimental section, λ(t) is not only treated as a constant but also modeled as dependent on epidemic conditions reflects the fact that public-health responsiveness typically increases when the number of infectious individuals rises (e.g., more frequent testing, isolation compliance, and policy enforcement), and relaxes when case numbers decline. In our setting, β(t) is closely dependent on the number of infectious individuals, since higher prevalence typically leads to more frequent contact with infected individuals and a greater probability of exposure, so setting β(t) to explicitly depend on I(t) captures this adaptive mechanism. In contrast, γ is primarily determined by biological and clinical factors and changes very little with social behavior, making it reasonable to treat it as constant. Likewise, δ is largely governed by the clinical course and treatment duration during quarantine (i.e., a relatively fixed isolation period), rather than by epidemic intensity, so it is also appropriate to model it as constant. A schematic illustration of such an SIQR model is shown in Fig. 1. We note that there are many compartmental epidemiological models based on different extensions of the original SIR model (Brauer et al., 2019; Piovella, 2020). While we focus on SIQR model calibration in this paper, the presented model calibration methods can be applied to all these variants in a straightforward manner. More importantly, calibration of more complicated epidemiology models based on agent-based models (ABMs) can also leverage our proposed methods if their critical parameters can be identified, where calibration efficiency is even more important due to their significantly higher computational complexity compared to the ODE-based compartmental models (Aleta et al., 2020; Anirudh et al., 2022; Beyeler et al., 2020).

Fig. 1.

Diagram depicting an SIQR compartmental epidemiological model. Each arrow indicates that the population rise of the ending compartment is caused by the population decrease of the starting compartment.

2.2. Model calibration

Most of the existing calibration methods (Brown & Atamturktur, 2018; Cox et al., 2001; Sargsyan et al., 2015; Sung & Hung, 2024; Tuo et al., 2023; Tuo & Wu, 2015) follow or adapt from the Bayesian approach pioneered by Kennedy and O'Hagan (KOH) (Kennedy & O'Hagan, 2001). In practice, we can only collect epidemiological data in the occurrence of an epidemic. In other words, we can observe the population's dynamic trajectory only once, making model calibration more challenging than in other statistical inference or machine learning tasks. Supposing one such trajectory with time length T is observed, we denote the observed epidemiological data as ${\left\{{\mathbf{d}}^t=(d_1^t,\dots ,d_4^t)\right\}}_{t=1}^T$ and make the following assumption:

$${\mathbf{d}}^t={\zeta }^t+ϵ\approx {\eta }^t(\mathbf{x})+ϵ,$$ (2)

where ζ is the unknown ground-truth underlying epidemic process, η is a computer model (e.g., SIQR model) that we would like to calibrate to simulate the behavior of ζ and fit the observed epidemiological data, $\mathbf{x}\in {\mathbb{R}}^D$ is the D-dimensional calibration parameter vector of η, and $ϵ\sim \mathcal{N}(0,\mathrm{D}\mathrm{i}\mathrm{a}\mathrm{g}(\mathbf{1}))$ represents the observation noise. In this work, ${\eta }^t(\mathbf{x})=({\eta }_1^t(\mathbf{x}),{\eta }_2^t(\mathbf{x}),{\eta }_3^t(\mathbf{x}),{\eta }_4^t(\mathbf{x}))=(S(t;\mathbf{x}),I(t;\mathbf{x}),Q(t;\mathbf{x}),R(t;\mathbf{x}))$ corresponds to the output of the SIQR model with x representing parameters of rate functions {β(t; x), λ(t; x), δ(x), γ(x)}. By our assumption, $d^t|{\eta }^t(\mathbf{x})\sim \mathcal{N}({\eta }^t(\mathbf{x}),1)$, and the Maximum Log-likelihood Estimation (MLE) of x corresponds to (Sung & Hung, 2024; Tuo & Jeff Wu, 2016):

$$\begin{array}{c}{\mathbf{x}}^{*}={\mathrm{argmax}}_{\mathbf{x}}\sum _t^{T}g\left({\eta }^t\left(\mathbf{x}\right)\right),\\ g\left({\eta }^t\left(\mathbf{x}\right)\right)≔-\frac{1}{T}\sum _{i=1}^4{\left(d_i^t-{\mathit{\eta }}_i^t\left(\mathbf{x}\right)\right)}^2,\end{array}$$ (3)

where g denotes the likelihood function as the negative Square Error (SE) function based on the Gaussian assumption. For notation compactness, we denote g(η^t(x)) as f^t(x). The goal of model calibration is the maximization of the corresponding negative Mean Square Error (MSE), ∑_t f^t(x).

3. Methodology

Instead of following the traditional KOH Bayesian model calibration scheme (Kennedy & O'Hagan, 2001), we resort to a new model calibration framework leveraging Bayesian Optimization (BO) (Garnett, 2023). Under this framework, model calibration is formulated into the optimization of the correspondingly designed utility function, called acquisition function. In BO, Gaussian Processes (GPs) have been typically used as cheap surrogates of f(x)1 In this way, we circumvent the difficulties in optimizing f(x) directly. To be more specific, the standard BO approach treats f(x) as a blackbox function and starts modeling it by a set of GP prior distributions, $\mathcal{N}(\mu (\mathbf{x}),\mathrm{\Sigma }(\mathbf{x}))$. Then, it iteratively chooses the next point at which to evaluate f with the procedure described as follows.

Given n observations of the objective function customized for model calibration, f(x¹), …, f(xⁿ), we first infer the posterior distribution, $\mathcal{N}({\mu }_n(\mathbf{x}),{\mathrm{\Sigma }}_n(\mathbf{x}))$, which conditions on the ${\{f({\mathbf{x}}^l),{\mathbf{x}}^l\}}_{l=1}^n$. Then, this posterior distribution is used to compute an acquisition function. Finally, we make a Bayesian decision that chooses the next point to be evaluated, xⁿ⁺¹, as the point that maximizes this acquisition function. After N iteration, the candidates x^best = argmax_x μ_N(x) is the calibrated parameter choice.

3.1. Graybox BO

It is clear that given the simulated dynamics η(x), evaluation of the model calibration performance metric g is cheap and does not need to be approximated. Instead of directly modeling f(x) as a blackbox objective function as typically done in BO, for SIQR model calibration, we here first use a set of GPs, denoted by y(x), as the surrogate to η(x). With that, we then model f(x) as a graybox objective, by a Composite Function (CF) (Astudillo & Frazier, 2019):

$$g(\mathbf{y}(\mathbf{x}))=\frac{1}{T}\sum _i{(d_i-y_i(\mathbf{x}))}^2.$$ (4)

Spared from the modeling of g, the model calibration performance can be expected to be further improved. Furthermore, as shown in Fig. 1, the workflow of the SIQR model encompasses a set of input-output structures between functions of each compartment. For example, infectious population ‘I’ is exclusively transferred from the susceptible population ‘S’, and individuals in the recovery compartment ‘R’ only come from compartments ‘I’ and ‘Q’. Therefore, one may expect that leveraging these functional structures may lead to improved calibration performance. Integrating these relationships into the BO method gives rise to a Function Network (FN), $\mathcal{G}(\mathcal{V},\mathcal{E})$, as shown in Fig. 2. In the graph, the set of nodes $\mathcal{V}=\{v|y_1,\dots ,y_4,x,g\}$ denote four GP sets associated with the four compartments ‘S’, ‘I’, ‘Q’ and ‘R’ respectively, calibration parameters x, and performance metric g. The graph also contains a set of edges $\mathcal{E}=\{(v\to v^{\prime })|v,v^{\prime }\in \mathcal{V}\}$, where (v → v′) means node v′ takes the output of nodes v as its input (Astudillo & Frazier, 2021). Suppose we have already made queries from the computer model η by simulating trajectories of subpopulations in each compartment at different settings and obtained ${\mathcal{D}}_n={\left\{\left({\hat{\mathbf{y}}}^l=\eta ({\mathbf{x}}^l),{\mathbf{x}}^l\right)\right\}}_{l=1}^n$ at n calibration parameter vectors. We have:

$$\begin{array}{l}{P}_n(\mathbf{y}(\mathbf{x}))=\prod _{i=1}^4{P}_n(y_i(\mathbf{x},P{a}_i)),\\ y_i\sim \mathcal{N}\left({\mu }_n^i(\mathbf{x},P{a}_i),{\sigma }_n^i(\mathbf{x},P{a}_i)\right),\end{array}$$ (5)

where $P{a}_i=\{y_j:(y_j\to y_i)\in \mathcal{E}\}$ denotes the GPs that are the parent nodes of y_i in graph $\mathcal{G}$, $P_n(y_i(\mathbf{x},P{a}_i))=P\left(y_i(\mathbf{x},P{a}_i)|{\hat{\mathbf{y}}}^{1:n}({\mathbf{x}}^{1:n})\right)$ denotes the posterior distribution of GP y_i with mean $m_n^i$ and variance ${\sigma }_n^i$. Given ${\mathcal{D}}_n$, we choose the GP prior $P\left(y_i^{1:n}({\mathbf{x}}^{1:n},P{a}_i^{1:n})\right)$ to be multivariate normal with a Matérn covariance kernel and a constant mean, whose hyper-parameters are obtained by MLE over $P({\hat{y}}_i^{1:n}({\mathbf{x}}^{1:n},{\widehat{Pa}}_i^{1:n}))$. Then we infer the posterior from $P({\hat{y}}_i^{1:n}({\mathbf{x}}^{1:n},{\widehat{Pa}}_i^{1:n}),y_i(\mathbf{x},P{a}_i))$ (Frazier, 2018b).

Fig. 2.

Function network structure of the SIQR model calibration, where the two bold arrows indicate that x and g are the parent node and child node of y_1:4 respectively.

Taking the Knowledge Gradient (KG) (Frazier, 2018b; Frazier et al., 2008), a typical formulation of acquisition functions as an example, we define the graybox acquisition function according to the FN as:

$${\alpha }_n(\mathbf{x})={\mathbb{E}}_{P_n(\mathbf{y}(\mathbf{x}))}\left[u_{n+1}^{*}(\mathbf{y},\mathbf{x})\right]-u_n^{*}$$ (6)

where $u_n^{*}={\max }_{\mathbf{x}}{u}_n(\mathbf{x})$, $u_n(\mathbf{x})={\mathbb{E}}_{P_n(\mathbf{y}(\mathbf{x}))}\left[g(\mathbf{y})\right]$, $u_{n+1}^{*}(\mathbf{y}(\mathbf{x}))={\max }_{{\mathbf{x}}^{\prime }}{u}_{n+1}({\mathbf{x}}^{\prime };\mathbf{y}(\mathbf{x}))$, and $u_{n+1}({\mathbf{x}}^{\prime };\mathbf{y}(\mathbf{x}))={\mathbb{E}}_{P_{n+1}\left({\mathbf{y}}^{\prime }({\mathbf{x}}^{\prime });\mathbf{y}(\mathbf{x})\right)}\left[g({\mathbf{y}}^{\prime })\right]$. On the right-hand side of the formula above, the second term is the largest expected metric value based on the current GPs, and the first term is the largest expected metric value if the GPs are further conditioned on predicted data (y(x), x). In comparison, the expectations ${\mathbb{E}}_{P_n}[\mathbf{y}]$ and ${\mathbb{E}}_{P_{n+1}}[\mathbf{y}]$ are computed in the original KG.

Due to the maximization operator within the expectation, neither the original nor graybox versions of KG can be computed explicitly. Thus, we resort to Monte Carlo estimation (Andrieu et al., 2003) of the expectations and the reparametrization trick (Kingma & Welling, 2013), so that $\hat{\alpha }$ is an unbiased estimation of α and ${\nabla }_{\mathbf{x}}\hat{\alpha }\left(\mathbf{x}\right)$ is feasible. To be more specific, we estimate the acquisition function via:

$$\begin{array}{l}\hat{\alpha }(\mathbf{x})=\frac{1}{K}\sum _{k=1}^K{\hat{u}}_{n+1}^{*}\left(\mathbf{x},{\hat{\mathbf{y}}}_{(k)}(\mathbf{x})\right)-{\hat{u}}_n^{*},\\ {\hat{u}}_{n+1}^{*}\left(\mathbf{x},{\hat{\mathbf{y}}}_{(k)}(\mathbf{x})\right)={\max }_{{\mathbf{x}}^{\prime }}\frac{1}{L}\sum _{l=1}^{L}g\left({\hat{\mathbf{y}}}_{(l,k)}^{\prime }({\mathbf{x}}^{\prime };\mathbf{x})\right),\\ {\hat{u}}_n^{*}={\max }_{\mathbf{x}\prime \prime }\frac{1}{L}\sum _{l=1}^{L}g\left(\hat{\mathbf{y}}{\prime \prime }_{(l)}(\mathbf{x}\prime \prime )\right),\end{array}$$

$$\begin{array}{l}{\hat{y}}_{i,(k)}(\mathbf{x})={\mu }_n^i\left(\mathbf{x},{\widehat{Pa}}_{i,(k)}(\mathbf{x})\right)+{\sigma }_n^i\left(\mathbf{x},{\widehat{Pa}}_{i,(k)}(\mathbf{x})\right){\widehat{ϵ}}_i^{(k)}\\ {\hat{y}}_{i,(l,k)}^{\prime }({\mathbf{x}}^{\prime };\mathbf{x})={\mu }_{n+1}^i\left({\mathbf{x}}^{\prime },{\widehat{Pa}}_{i,(l,k)}^{\prime }({\mathbf{x}}^{\prime });{\hat{\mathbf{y}}}_{(k)}(\mathbf{x})\right)\\ +{\sigma }_{n+1}^i\left({\mathbf{x}}^{\prime },{\widehat{Pa}}_{i,(l,k)}^{\prime }({\mathbf{x}}^{\prime });{\hat{\mathbf{y}}}_{(k)}(\mathbf{x})\right){\widehat{ϵ}}_i^{(l,k)}\\ \hat{y}{\prime \prime }_{i,(l)}(\mathbf{x}\prime \prime )={\mu }_n^i\left(\mathbf{x}\prime \prime ,Pa{\prime \prime }_{i,(l)}(\mathbf{x}\prime \prime )\right)+{\sigma }_n^i\left(\mathbf{x}\prime \prime ,Pa{\prime \prime }_{i,(l)}(\mathbf{x}\prime \prime )\right){\widehat{ϵ}}_i^{(l)},\end{array}$$ (7)

where ${\widehat{ϵ}}_i$ denotes a sample of ${ϵ}_i\sim \mathcal{N}(0,1)$, K = 8, and L = 128. The computational cost for inferring P_n is about $\mathcal{O}(n^3)$ and does not scale up with the dimension of GPs' input (i.e., x or [x, Pa_i]) (Seeger, 2004). Thus, the computational complexity of the original and graybox versions of KG are similar. The workflow of our calibration method based on graybox BO is illustrated in Fig. 3 with the pseudo-code summarized in Algorithm 1.

Fig. 3.

Schematic illustration of the model calibration workflow based on graybox BO, where we leverage expert knowledge about functional dependency and metric function.

3.2. Decoupled decision-making

In practice, epidemiological data d from the unknown epidemic process ζ may not be observed at synchronized time points. In other words, some observations at the corresponding time point may be missing in practice. For example, when d = {∅, d₂, d₃, d₄}, where d₁ is missing in the observed data, we may only be able to compute $g(\mathbf{y}(\mathbf{x}))=\frac{1}{T}{\sum }_{i\ne 1}{(y_i(\mathbf{x})-d_i)}^2$. In such a scenario of decoupled observations, the function network organization of GPs allows us to infer the function form y₁(x) by leveraging the ground-truth functional dependency represented by $\mathcal{G}$. To be more specific,

$$g\perp /v,\forall v\in {\mathrm{A}\mathrm{n}}_g,$$ (8)

where v⊥̸v′ indicates that v and v′ are statistically dependent so that P(v|v′) ≠ P(v) for any two random variables v and v′ (Koller & Friedman, 2009). This mimics the functional dependency h(⋅)⊥̸h′(⋅), which means that a perturbation in function h(⋅) will not affect the solution space of function h′(⋅); ${\mathrm{A}\mathrm{n}}_g=\{v|(v\to \dots \to g)\in \mathcal{G}\}$, denoting the ancestor nodes of g with a directed path starting at the ancestor node and ending at g. The corresponding graph with d = {∅, d₂, d₃, d₄} will have the edge (y_i → g) removed, while g⊥̸y₁ still holds.

Moreover, it is reasonable to assume that the approximation of the computer model's output {η₁(x), …, η₄(x)} by the four GPs varies in computational complexity. The reason is that the complexities of the chosen function forms of {η₁, …, η₄} are different. Therefore, some GPs may need more queried data from the computer model than others. Taking into account of the decomposable structure of P_n(y(x)), we extend the graybox acquisition function by decoupled decision-making and define the acquisition function as:

$$\begin{array}{l}{\alpha }_n(\mathbf{x},\mathbf{z})≔\frac{1}{{\mathbf{1}}^{\top }\mathbf{z}}{\mathbb{E}}_{P_n(\mathbf{y}(\mathbf{x}))}\left[u_{n+1}^{*}(\mathbf{z},\mathbf{y}(\mathbf{x}))\right]-u_n^{*}\\ u_{n+1}^{*}(\mathbf{z},\mathbf{y}(\mathbf{x}))≔{\max }_{{\mathbf{x}}^{\prime }}{\mathbb{E}}_{P_{n+1}({\mathbf{y}}^{\prime }({\mathbf{x}}^{\prime });\mathbf{z},\mathbf{y}(\mathbf{x}))}[g({\mathbf{y}}^{\prime })]\end{array}$$ (9)

where z ∈ {0,1}⁴ and P_n+1(y′(x′); z, y(x)) is defined as:

$$\prod _{i=1}^4{P}_{n+1}{(y_i^{\prime }({\mathbf{x}}^{\prime },P{a}_i^{\prime });y_i(\mathbf{x},P{a}_i))}^{z_i}{P}_n{(y_i^{\prime }({\mathbf{x}}^{\prime },P{a}_i^{\prime }))}^{1-z_i}.$$ (10)

New candidates xⁿ⁺¹ is found by argmax_x,zα_n(x, z). When z = 1, P_n+1(y′(x′); z, y(x)) = P_n+1(y′(x′); y(x)). The proposed objective can be seen as a generalization of the KG objective, which measures expected improvement when GPs are conditioned on predicted data (y(x), x). Here we further accommodate the case when a subset of GPs is conditioned on predicted data. Thus, candidates xⁿ⁺¹ with higher expected improvement can be expected by the proposed objective. Convergence analysis of graybox BO using our proposed acquisition function is provided in the appendix.

Due to the combinatorial nature of z, it is hard to optimize w.r.t. z directly. The most naive approach is to optimize α_n w.r.t. x under every possible choice of z, and the resulting computational cost is 2^|z| times that of the non-decoupled version. For the SIQR model in this paper, |z| = 4 and the increased cost is acceptable. Besides, the non-decoupled version is just a special case of the decoupled version, where z = 1. Therefore, we can accelerate the computation by selecting a subset. In this paper, we choose $\left\{(\mathbf{z}:z_1=1,z_{\ne 1}=0),\dots ,(\mathbf{z}:z_{|\mathbf{z}|}=1,z_{\ne |\mathbf{z}|}=0)\right\}\cup \left\{\mathbf{z}=\mathbf{1}\right\}$, resulting |z| + 1 times the computational cost of the non-decoupled version.

Algorithm 1

Model Calibration Workflow

Require: Number of iteration N, Computer model η(x) $(X,Y)\leftarrow \left(X_{\mathit{init}},{\{\eta (\mathbf{x})\}}_{\mathbf{x}\in X_{\mathit{init}}}\right)$//Xinit is a randomly initialized set of calibration parameter vectors.

forn = {0, …, N − 1} do

${\mathcal{D}}^n\leftarrow (X,Y)$

$P_n(y_i(\mathbf{x}))\leftarrow \mathcal{N}({\mu }_n^i(\mathbf{x}),{\sigma }_n^i(\mathbf{x}))$

${\mathbf{x}}^{n+1}\leftarrow {\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{a}\mathrm{x}}_{\mathbf{x}}{\hat{\alpha }}_n(\mathbf{x})$.

(X, Y) ← (X ∪ {xn+1}, Y ∪ {η(xn+1)})

endfor

return${\mathbf{x}}^{\mathit{best}}\leftarrow {\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{a}\mathrm{x}}_{\mathbf{x}}{\hat{u}}_N(\mathbf{x})$

4. Experimental results of simulated datasets

In this section, we conduct experiments to evaluate and compare the performances of model calibration using various Bayesian Optimization (BO) methods. We consider different setups of observed epidemiological data generated from simulated “ground-truth” models. In our experiments, both the ground-truth model ζ and the computer model to calibrate η are SIQR models. Once a trial of epidemiological data ${\{{\mathbf{d}}^t\}}_{t=1}^{30}$ is simulated at 30 time points, the ground-truth model is no longer accessible. It is also assumed that we are only accessible to computer model outputs. The goal is to validate the effectiveness of our graybox BO methods for model calibration in expensive scenarios and investigate how our proposed acquisition functions can improve the performance of graybox BO-based model calibration.

To generate data based on the ground-truth SIQR model ζ, we set β∗(t) = 0.9I(t)/N, λ∗(t) = 0.2, δ∗ = 0.2, γ∗ = 0.1. The ground-truth parameters are chosen to ensure realistic 30-day epidemic dynamics, identifiability across BO methods, and numerical stability for repeated ODE-based simulation. The model to calibrate, η, takes the same form with undetermined parameters x, that is, λ(t, x) = x₁, β(t, x) = x₂I(t)/N, δ = x₃, and γ = x₄. The initial conditions for both ζ and η are: S(0) = 0.99∗N, I(0) = 0.01∗N, R(0) = 0 and Q(0) = 0. We consider two scenarios representing the gap between the ground-truth model and the computer model to calibrate. The first way is to add noise to the ground-truth observations d. The second is to have different setups of η from those of ζ. In our experiments, we set λ∗(t) in ζ as a non-linear function, that is, λ∗(t) = log([I(t),S(t),R(t)]^⊤[0.3, 0.06, 0.12]/N + 1), while λ(t) in η still takes the linear form mentioned above. The corresponding plots for the 30-day simulations are shown in Fig. 4.

Fig. 4.

The simulated ground-truth population fraction trajectories of each compartment for 30 days. (a) Trajectories from a SIQR model with linear derivative functions. (b) Trajectories corresponding to the case when function λ∗(t) is non-linear. (c) The trajectory of λ∗(t) when it is set to be non-linear.

We compare calibration performances based on the following BO methods: (i) the standard Expected Improvement (EI) in the blackbox BO framework without integrating functional structures, (ii) the standard Knowledge Gradient (KG), (iii) the Knowledge Gradient with composite function (KG-CF), (iv) Knowledge Gradient with function network (KG-FN), and (v) the Decoupled Knowledge Gradient with composite function (DG-CF). All algorithms have been implemented based on the BOTorch package (Balandat et al., 2020) while the latter three are all graybox BO variants. In all our experiments, we have run each calibration method five times from the different random seeds, where each run starts with an initial dataset from the computer model, ${\mathcal{D}}_0={\left\{\left({\hat{\mathbf{y}}}^l=\eta ({\mathbf{x}}^l),{\mathbf{x}}^l\right)\right\}}_{l=1}^{2D+1}$, where x^l is randomly selected in [0,1]^D.

Experiments under all three ground-truth models are performed in complete observation and incomplete observation setups. In the latter setup, observations of compartment ‘S’ are assumed to be unknown, so d₁ is not involved with the calibration procedure. This helps to investigate whether the function network organization of GPs allows more reliable model calibration under incomplete (decoupled) observation data d by leveraging the ground-truth functional dependency.

The experimental results under the complete-observation setting are shown in Fig. 5. In all these experiments, it can be observed that our graybox BO methods, including KG-CF, perform significantly better than the blackbox BO methods, EI and KG. This demonstrates that integrating knowledge of g does help improve the calibration performance. The performance of KG-FN is slightly worse than KG-CF. This phenomenon can be ascribed to the following reasons. With complete observations, all the updating of GPs, {y₁, …, y₄} that approximate {η₁, …, η₄} are well-informed, making the information propagating via introducing statistical dependency less important. Besides, the function network organization of GPs make the acquisition estimator $\hat{\alpha }$ more challenging to calculate reliably. As shown in (7), inputs of each GP include not only x but also random samples of Pa, which makes the estimated $\hat{\alpha }$ unstable due to the stochastic nature of inputs, and consequently increases the difficulty of finding the optimal point of $\hat{\alpha }(\mathbf{x})$ in each BO iteration. Last but not least, our proposed DG-CF performs better than KG-CF. This validates that the decoupled acquisition can better inform decision-making by conditioning on the predicted data of a subset of GPs, rather than the whole set of GPs.

Fig. 5.

Calibration performance. The logarithm of the MSE is shown w.r.t. the number of BO iterations/epochs. For the n_th iteration, the MSE is computed as $-{\sum }_t{f}^t({\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{a}\mathrm{x}}_{\mathbf{x}}{\hat{u}}_n(\mathbf{x}))$. Solid lines show the mean values of five runs and the shaded regions correspond to the standard deviations around the means. (a) Performance for a linear SIQR model ζ. (b) Performance for a linear ζ in the presence of noise.(c) Performance for a non-linear ζ.

The experimental results under the incomplete-observation setting are shown in Fig. 6. Again, our graybox BO methods perform significantly better than the blackbox BO methods. It should be noted that KG-FN can outperform KG-CF when the observations are noisy. This shows that integrating functional dependency does help to inform the calibration process, while the performance gain is not significant due to the accompanied optimization difficulty that we mentioned above. Finally, our proposed DG-CF performs better than KG-CF, validating again the advantages of the decoupled acquisition function.

Fig. 6.

Calibration performance runs under incomplete observations. (a) Performance for a linear SIQR model ζ. (b) Performance for a linear ζ in the presence of noise. (c) Performance for a non-linear ζ. In all cases, we assume that the ground-truth observations of the susceptible population (compartment ‘S’) are missing.

In Fig. 7, we also visualize population fraction trajectories generated from the computer models under the calibrated parameters derived by all methods in the setting with noisy and incomplete observations from a linear SIQR model. It can be observed that all methods can provide non-trivial calibration results, showing that our BO-based calibration methods have a promising potential for epidemiological dynamic model calibration, when the computed model is expensive. Among all the BO variants, graybox BO variants KG-FN and DG-CF render the mean trajectories with smaller bias to the ground-truth ones as well as smaller variance, compared to KG-CF, EI, and KG. This again confirms that utilizing the functional structure and decoupled acquisition functions can help better inform decision-making during BO iterations.

Fig. 7.

Population fraction trajectories from the computer model η under the calibrated parameters x^best, calibrated by (a) KG-CF, (b) KG-FN, (c) DG-CF, (d) EI, and (e) KG, with noisy observation from a linear ζ. The solid lines represent the mean trajectories simulated from the corresponding calibrated computer models of five runs and the shaded regions illustrate the standard deviation intervals around the means. The dotted lines are the ground-truth population fraction trajectories.

To further study the converged calibration results, we investigate the difference between the calibrated parameter x of the computer model η and the ground-truth parameter x∗ of the SIQR model ζ. When fewer than three population compartments can be observed or when only a limited number of discrete time points within the considered interval [0, 30] are available, parameter identifiability cannot be guaranteed, even under a linear SIQR parameterization. However, the proposed graybox BO method enable reducing the uncertainty and accelerates the calibration by incorporating the expert prior knowledge through the composite function and function network. To illustrate this effect, we still analyze both the previous complete-observation and incomplete-observation settings considered previously. In these settings, x∗ is, in principle, identifiable since at most one compartment is missing and there are observations at 30 time points available for estimating 4 scalar parameters. The L₂ norm distance between x and x∗ w.r.t. BO epochs are shown in Fig. 8. We also provide the calibrated parameters x^best in Table 1, Table 2. For both settings, graybox methods, KGCF, KGFN and DGCF successfully identify parameters that are very close to x∗, as indicated by distances approaching zero. In the incomplete-observation setting, the missing ‘S’ compartment does increase the uncertainty of KG-CF during the early calibration stage, whereas KG-FN and DG-CF remain more stable by leveraging additional structural knowledge. By contrast, the blackbox methods (EI and KG) fail to recover parameters close to x∗ within 50 iterations.

Fig. 8.

The L₂-norm distances between the parameter x∗ of the a linear SIQR model ζ and the parameter x of the calibrated computer model η. The distances are shown w.r.t. the number of BO iterations. (a) Distances under complete observations. (b) Distances under incomplete observations, where observations of the susceptible population (compartment ‘S’) are missing.

Table 1.

Mean and standard deviation of the parameters x^best calibrated by different BO methods for complete observations and linear ζ.

Method	x1	x2	x3	x4
EI	0.139 ± 0.070	0.913 ± 0.077	0.234 ± 0.095	0.084 ± 0.062
KG	0.196 ± 0.018	0.914 ± 0.075	0.197 ± 0.106	0.081 ± 0.065
KGCF	0.200 ± 0.000	0.900 ± 0.000	0.200 ± 0.000	0.100 ± 0.000
KGFN	0.200 ± 0.000	0.900 ± 0.000	0.200 ± 0.000	0.100 ± 0.000
DGCF	0.200 ± 0.000	0.900 ± 0.000	0.200 ± 0.000	0.100 ± 0.000

Table 2.

Mean and standard deviation of the parameters x^best calibrated by different BO methods for incomplete observations and linear ζ.

Method	x1	x2	x3	x4
EI	0.109 ± 0.100	0.921 ± 0.071	0.242 ± 0.080	0.084 ± 0.066
KG	0.276 ± 0.055	0.892 ± 0.014	0.204 ± 0.026	0.119 ± 0.041
KGCF	0.200 ± 0.000	0.901 ± 0.001	0.199 ± 0.001	0.100 ± 0.000
KGFN	0.200 ± 0.001	0.900 ± 0.000	0.199 ± 0.001	0.100 ± 0.001
DGCF	0.200 ± 0.000	0.900 ± 0.000	0.200 ± 0.000	0.100 ± 0.000

5. Experimental results of real-world datasets

To further confirm the effectiveness of our proposed methods in real-world scenarios, we apply them to calibrate the SIQR model η w.r.t. observed population trends during the COVID-19 pandemic. We use the data collected by the Johns Hopkins University Center for Systems Science and Engineering, via an R data package covid19.analytics (Ponce & Sandhel, 2021). We focus on the infectious populations of U.S. and U.K. For each country, there are 365 temporal observations of the infectious populations collected from June 1st, 2020, to June 1st, 2021.

We consider three variants of the SIQR models, denoted by η_S, η_W, and η_NN. Calibrated Model η_S takes the same form as η, described in the previous experiment section. For model η_W, λ(t) adopts a non-linear form, λ(t) = log([I(t),S(t),R(t)]^⊤[x_1,1, x_1,2, x_1,3]/N + 1), while other components remain the same as those of η. Last but not least, λ(t) of model η_NN is set to be f_NN([I(t), S(t), R(t)]/N; x_NN). Here, f_NN is a neural network with parameters x_NN and three latent layers. The detailed information about the dimension and activation function in each layer of f_NN is summarized as follows: Input Layer–(3 × 20, Tanh); Output Layer–(20 × 1, Sigmoid); Latent Layers–${\{{(20\times 20,\text{Tanh})}_i\}}_{i=1}^3$; The other components of η_NN follow those of η. This calibrated model setup is motivated by the common adoption of neural ODE for modeling the dynamics of complex systems, making it meaningful to validate proposed methods w.r.t. more flexible neural ODE models.

Algorithm 2

Two-Stage Model Calibration Workflow

Require:N, ηNN(x), ηS(x)

x1st ←Algorithm-1(N, ηS(x))

${\{S(t),Q(t),R(t)\}}_{t=0}^{T-1}\leftarrow {\eta }_S({\mathbf{x}}^{\mathit{1st}})$

$\mathcal{D}\leftarrow {\{\hat{I}(t),S(t),Q(t),R(t)\}}_{t=0}^{T-1}$

$\mathbf{x}\leftarrow [{\mathbf{x}}_{NN},{\mathbf{x}}_{2:4}^{\mathit{1st}},{\mathbf{x}}_{5:6}^{\mathit{1st}}]$

Optimize ηNN(x) w.r.t. xNN and ${\mathbf{x}}_{2:4}^{\mathit{1st}}$ by stochastic gradient descent, based on $\mathcal{D}$, and obtain x2nd

returnx2nd

Unlike the simulated experiments, we know the initial state of each population and the total population. In the real-world scenarios, these values should be carefully selected. Since we only have the observations of the infectious population, the simplest choice is to set I(0) equal to the first-day's infectious observation, denoted by $\hat{I}(0)$; and set the total population N to the population of each country. However, during the COVID-19 pandemic, most countries issued policies to discourage the spread of the disease. For example, people were suggested to work at home. The true total population that should be taken into account is much smaller than the population of each country. In light of the above reason and calibration quality, we set I(0) = x₅ and N = x₆ as two additional parameters in the calibrated model, and consider the ranges $10^{-1}\hat{I}(0)\le x_5\le 10\hat{I}(0)$ and $10\hat{I}(0)\le x_6\le 10^3\hat{I}(0)$. In addition, the calibration of η_NN is divided into two stages. In the first stage, we calibrate η_S to the real-world observations. In the second stage, we replace λ(t) = x₁ by f_NN([S(t), I(t), R(t)]/N; x_NN), obtaining an uncalibrated η_NN. Then, we adopt the adjoint gradient method (Chen et al., 2018) to compute gradients efficiently and stochastic gradient descent to optimize parameters of f_NN(⋅; x_NN) and [β(t, x₂), δ(t, x₃), γ(t, x₄)]. Values of x_NN are randomly initialized, and values of x_2:6 found in the first stage are used as the initial values for the second stage. The batch size, the number of total iterations, the learning rate are set to 30, 2000 and 5 × 10⁻⁴. The reason for this two-stage calibration is as follows. First, it should be noted that the goal of model calibration is to efficiently identify the value of key parameters rather than the optimization of all parameters. Compared to the parameters of f_NN, which collectively govern the λ(t), x_2:6 act as hyper-parameters, each of which controls an individual component of the SIQR model. Therefore, properly setting the values or ranges x_2:6 is crucial for the calibration quality. Second, while neural ODEs are typically optimized by stochastic gradient descent, the optimization procedure requires complete observations of all population compartments, or the marginal formulation that only depends on the observed population compartments. To be more specific, we are required to randomly select a window [t₀, t₂₉] ⊂ [0, T] during optimization. Then η_NN is used to simulate sub-trajectories for this window, starting from the ground-truth observation $(\hat{I}(t_0),\hat{S}(t_0),\hat{Q}(t_0),\hat{R}(t_0))$. However, in our experiment, only $\hat{I}(t_0)$ is known, making the stochastic gradient descent infeasible. Performing normal gradient descent usually leads to poor optimization performance, where we always simulate trajectories from $(\hat{I}(0),N-\hat{I}(0),0,0)$. Since we only care about the calibration w.r.t. the infectious observations, a preferable choice is to use $(\hat{I}(t_0),S(t_0),Q(t_0),R(t_0))$, where S(t₀), Q(t₀), R(t₀) are simulated from η_S. The corresponding calibration workflow is presented in Algorithm 2.

In Fig. 9, Fig. 10, we visualize the comparison of predicted trajectories from the three types of computer models, calibrated by KG-CF and KG-FN based on infectious observations of the U.S. and U.K. respectively. We note that, since only the infectious population compartment is observed, the DG-CF implementation is the same as KG-CF. From these figures, all calibrated models by either KG-CF or KG-FN successfully capture the trends of the true observations. Compared to KG-CF, the calibration performances of KG-FN for all three model setups are generally better with lower variance as shown in Table 3. With nonlinear λ(t), η_W fits better to the observations than η_S. Enhancing λ(t) by a neural network and applying the two-stage calibration process, the calibrated models emulate infection population curves that fit the observations better as expected. While the calibrated curves under all model setups exhibit single modes (Odagaki, 2021) due to the limited model expressiveness of the SIQR model family, the capability of our proposed calibration methods is well validated by multiple setups. It can be expected that calibrating more expressive models like ABMs (Aleta et al., 2020) will render better fitting to the true observations.

Fig. 9.

Trajectories of the U.S. infectious population from the computer models (a) η_S, (b) η_W, and (c) η_NN under the calibrated parameters, calibrated by KG-CF and KG-FN. The solid lines represent the mean trajectories simulated from the corresponding calibrated computer models of five runs and the shaded regions illustrate the standard deviation intervals around the means. The scattered points are the ground-truth observations of the infectious population.

Fig. 10.

Trajectories of the U.K. infectious population from the computer models (a) η_S, (b) η_W, and (c) η_NN under the calibrated parameters, calibrated by KG-CF and KG-FN. The solid lines represent the mean trajectories simulated from the corresponding calibrated computer models of five runs and the shaded regions illustrate the standard deviation intervals around the means. The scattered points are the ground-truth observations of infectious population.

Table 3.

The mean and standard deviation of the smallest MSEs achieved by KG-CF and KG-FN on the U.S. and U.K. COVID-19 datasets, from five model calibration runs with randomly initialized ${\mathcal{D}}_0$.

Method	US ( × 106↓)	UK ( × 104↓)
KGCF	11.3087 ± 2.6014	44.6875 ± 4.1894
KGFN	8.0117 ± 0.5123	45.0617 ± 3.6110

6. Conclusion & future work

In this work, we proposed epidemiological model calibration methods based on the BO framework. To improve calibration performance, we formulated the calibration problem as a customized graybox BO task, where expert knowledge about functional dependency and calibration performance metric function is integrated into the acquisition function. Furthermore, we proposed a decoupled acquisition function, which further exploits the decomposable nature of the functional structure. We conducted experiments using simulated data under three types of ground-truth models and two observation types to validate the model calibration performance of our proposed BO-based methods. Within a small number of BO iterations $(\le 50)$, the proposed graybox BO methods can achieve good performance, leading to lower MSEs and faster convergence (in terms of BO iterations) compared to methods based on standard BO. To further confirm the capability of our methods to real-world problems, we calibrated three types of models to COVID-19 data. The experimental results demonstrate the efficacy of our BO-based methods for enabling the calibration of expensive computer models. While the experimental results show that utilizing ground-truth functional dependency can help the calibration process, the resulting formulation of the acquisition function can lead to optimization difficulties in practice, impairing the overall calibration performance. Thus, future work will focus on further investigation of the function network organization of GPs to achieve better performance. Besides, in the current work, the function network is limited to a single system, where each node corresponds to an element within that system. In a more general case, there can be multiple sub-systems interacting with one another, with each node in the network representing an entire sub-system. Our future investigation will consider extending the proposed calibration methods to accommodate this more general case. Finally, agent-based epidemiological models have a much larger number of critical parameters than compartmental models. Along with their significant computational complexity, the underlying functional dependency is also more complex, making their calibration tasks more challenging. Efficient calibration of ABMs for epidemiological modeling may require developing and validating different approximate solutions for BO acquisition functions, which remains an open problem for future research.

CRediT authorship contribution statement

Puhua Niu: Writing – review & editing, Writing – original draft, Software, Methodology, Formal analysis, Data curation. Byung-Jun Yoon: Writing – review & editing, Validation, Supervision, Resources, Project administration, Funding acquisition, Conceptualization. Xiaoning Qian: Writing – review & editing, Validation, Supervision, Resources, Project administration, Funding acquisition, Conceptualization.

Code availability

Code is available at github.com/niupuhua1234/Epi-KG.

Funding

This work was supported in part by the United States Department of Energy (DOE) Office of Science under Award B&R# KJ0403010 / FWP# CC138.

Conflict of interest

The authors declare that there is no known conflict of interest.

Handling Editor: Dr Yijun Lou

Appendix A.

We simplify the notations of ${\mathbb{E}}_{P_n(y_i(\mathbf{x},P{a}_i))}[\cdot ]$ and ${\mathbb{E}}_{P_n(P{a}_i(\mathbf{x}))}{\mathbb{E}}_{P_n(y_i(\mathbf{x},P{a}_i))}[\cdot ]$ as ${\mathbb{E}}_n[\cdot ]$ and ${\overline{\mathbb{E}}}_n[\cdot ]$, respectively (or ${\mathbb{E}}_{n,(\mathbf{x},P{a}_i)}[\cdot ]$ and ${\overline{\mathbb{E}}}_{n,x}[\cdot ]$ to emphasize the dependency relationships). Accordingly, ${\mathbb{E}}_{P_n(y_i(\mathbf{x},P{a}_i))}\left[{\mathbb{E}}_{P_{n+1}\left(y_i^{\prime }({\mathbf{x}}^{\prime },P{a}_i^{\prime });y_i(\mathbf{x},P{a}_i)\right)}[\cdot ]\right]$ and ${\mathbb{E}}_{P_n(P{a}_i(\mathbf{x}))}{\mathbb{E}}_{n,(\mathbf{x},P{a}_i)}\left[{\mathbb{E}}_{P_{n+1}(P{a}_i^{\prime }({\mathbf{x}}^{\prime }))}{\mathbb{E}}_{n+1,({\mathbf{x}}^{\prime },P{a}_i^{\prime })}[\cdot ]\right]$ can be denoted as ${\mathbb{E}}_n[{\mathbb{E}}_{n+1}[\cdot ]]$ and ${\overline{\mathbb{E}}}_n[{\overline{\mathbb{E}}}_{n+1}[\cdot ]]$.

Both ${\mathbb{E}}_n[g(y_i(\mathbf{x},P{a}_i))]$ and ${\mathbb{E}}_n[\max g(y_i)](={\mathbb{E}}_{n,x}[{\max }_{{\mathbf{x}}^{\prime }}g(y_i({\mathbf{x}}^{\prime },P{a}_i))])$ are martingales in that ${\mathbb{E}}_n[{\mathbb{E}}_{n+1}[g(y_i^{\prime })]]={\mathbb{E}}_{P_n(y_i^{\prime })}[g(y_i^{\prime })]={\mathbb{E}}_n[g(y_i)]$, and ${\mathbb{E}}_n[{\mathbb{E}}_{n+1}[\max g(y_i)]]={\mathbb{E}}_n[\max g(y_i)]$ by the law of total expectations (Bect et al., 2019; Kallenberg & Kallenberg, 1997). Furthermore, $\overline{\mathbb{E}}[g(y_i)]$ and $\overline{\mathbb{E}}[\max g(y_i)]$ are also martingales. The reason is as follows: Conditioning on Pa_i ∼ P_n, $P{a}_i^{\prime }\sim P_{n+1}$ is independent of y_i ∼ P_n as Pa_i is the common parent node within the dependency structure $y_i\leftarrow P{a}_i\to P{a}_i^{\prime }$ (Koller & Friedman, 2009). Therefore, ${\overline{\mathbb{E}}}_n\left[{\overline{\mathbb{E}}}_{n+1}[\cdot ]\right]={\mathbb{E}}_{P_n(P{a}_i(\mathbf{x}))}{\mathbb{E}}_{P_{n+1}(P{a}_i^{\prime }({\mathbf{x}}^{\prime }))}\left[{\mathbb{E}}_n{\mathbb{E}}_{n+1}[\cdot ]\right]={\mathbb{E}}_{P_n(P{a}_i(\mathbf{x}))}{\mathbb{E}}_{P_{n+1}(P{a}_i^{\prime }({\mathbf{x}}^{\prime }))}{\mathbb{E}}_n[\cdot ]={\overline{\mathbb{E}}}_n[\cdot ]$.

Lemma 1

Let $G_n^i\left(\mathbf{x}\right)={\overline{\mathbb{E}}}_{n,\mathbf{x}}\left[{\overline{\mathbb{E}}}_{n+1}\left[g\left(y_i^{\text{'}}({\mathbf{x}}_{*}^{\text{'}},P{a}_i^{\text{'}})\right)\right]-{\overline{\mathbb{E}}}_n\left[g\left(y_i({\mathbf{x}}_{*},P{a}_i)\right)\right]\right]$, where ${\mathbf{x}}_{*}={\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{a}\mathrm{x}}_{\mathbf{x}}{\overline{\mathbb{E}}}_n[g(y_i(\mathbf{x},P{a}_i))]$, and ${\mathbf{x}}_{*}^{\prime }={\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{a}\mathrm{x}}_{{\mathbf{x}}^{\prime }}{\overline{\mathbb{E}}}_{n+1}[g(y_i^{\prime }({\mathbf{x}}^{\prime },P{a}_i^{\prime }))]$. Then $G_n^i(\mathbf{x})\ge 0$.

Proof. By definitions, ${\overline{\mathbb{E}}}_{n+1}[g(y_i^{\prime }({\mathbf{x}}_{*},P{a}_i^{\prime }))]\le {\overline{\mathbb{E}}}_{n+1}[g(y_i^{\prime }({\mathbf{x}}_{*}^{\prime },P{a}_i^{\prime }))]$. Thus, $0={\overline{\mathbb{E}}}_n[{\overline{\mathbb{E}}}_{n+1}[g(y_i^{\prime }({\mathbf{x}}_{*},P{a}_i^{\prime }))]]-{\overline{\mathbb{E}}}_n[g(y_i({\mathbf{x}}_{*},P{a}_i))]\le G_n^i(\mathbf{x})$.

According to the definition of $G_n^i(\mathbf{x})$, we further define $H_n^i={\mathbb{E}}_{n,{\mathbf{x}}_{*}}[\max g(y_i)]-{\mathbb{E}}_n[g(y_i({\mathbf{x}}_{*},P{a}_i))]$, where ${\mathbf{x}}_{*}={\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{a}\mathrm{x}}_{\mathbf{x}}{\overline{\mathbb{E}}}_n[g(y_i(\mathbf{x},P{a}_i))]$. We also denote $J_n(\mathbf{x})={\mathbb{E}}_{n,x}[H_{n+1}]$. It can be observed that $H_n^i\ge 0$ and $J_n^i(\mathbf{x})\ge 0$. Since ${\mathbb{E}}_n[\max g(y_i)]$ is a martingale, $H_n^i-J_n^i(\mathbf{x})={\mathbb{E}}_{n,\mathbf{x}}[\max {\mathbb{E}}_{n+1}[g(y_i^{\prime })]]-\max {\mathbb{E}}_n[y_i]=G_n^i(\mathbf{x})$.

Assumption 1

We make the following assumptions:

• •$\mathbf{x}\in \mathbb{X}$ and $\mathbb{X}$ is a compact metric space (e.g. $({\mathbb{R}}^n,\Vert \cdot {\Vert }_2)$). and the sample path of y_{i∈{1, …,4}} is continuous.
• •η is in the state space of y ∼ P_n for n ≥ 1.

Lemma 2

(Lemma B.3 in (Astudillo & Frazier, 2021)). Supposing Assumption 1 holds and taking n → ∞, then $m_n^i$ converge to some $m_{\infty }^i$ almost surely (a.s.) and uniformly on $\mathbb{X}$, and ${\sigma }_n^i$ converges to some ${\sigma }_{\infty }^i$ a.s., and uniformly on $\mathbb{X}\times \mathbb{X}$.

Theorem 1

Supposing Assumption 1 holds and taking n → ∞, the decoupled acquisition function α_n(x) converge to 0 a.s..

Proof. First of all, it is clear that $G_n^i$ and $H_n^i$ are both continuous w.r.t. P_n and P_n+1. Thus, by Lemma 2, G_n → G_∞ and H_n → H_∞ when P_n → P_∞.

Denoting $K^i=H_n^i-H_{n+1}^i$, we have ${\overline{\mathbb{E}}}_{n,x}[K_n^i]=H_n^i-J_n^i(\mathbf{x})=G_n^i(\mathbf{x})$. By the assumption of η, we construct G_n based on a sample trajectory ${\{{\eta }^l({\mathbf{x}}^l),{\mathbf{x}}^l\}}_l^n$ and have:

$$\begin{array}{l}\sum _{n=0}^{N-1}{\overline{\mathbb{E}}}_{1:n,{\mathbf{x}}^{1:n}}\left[G_n^i\left({\mathbf{x}}^{n+1}\right)\right]=\sum _{n=0}^{N-1}{\overline{\mathbb{E}}}_{1:n+1,{\mathbf{x}}^{1:n+1}}\left[K_n^i\right]\\ ={\overline{\mathbb{E}}}_{1:N,{\mathbf{x}}^{1:N}}\left[\sum _{n=0}^{N-1}{K}_n^i\right]=H_0^i-{\overline{\mathbb{E}}}_{1:N,{\mathbf{x}}^{1:N}}\left[H_N^i\right]\le H_0^i<\infty ,\end{array}$$

where P₀ is generated from a given ${\mathcal{D}}_0$, P_n(≥1) is generated from ${\mathcal{D}}_0\bigcup {\{(y_i^l(\mathbf{x}),{\mathbf{x}}^l)\}}_{l=1}^n$ with $y_i^l(\mathbf{x})\sim P_l$. Taking N → ∞, we have ${\sum }_{n=0}^{\infty }\mathbb{E}[G_x^i({\mathbf{x}}^{n+1})]\le \infty$. Since $G_n^i(\mathbf{x})\ge 0$ by Lemma 1, the inequality implies $G_{\infty }^i(\mathbf{x})=0$ a.s.

By the definition of α_n, we have ${\alpha }_n\left(\mathbf{x}\right)={\max }_{\mathbf{z},\mathbf{x}}\frac{1}{1^{\top }\mathbf{z}}{\sum }_i(H_i-z_i{J}_n^i\left(\mathbf{x}\right)-(1-z_i)H_i)={\max }_{\mathbf{z},\mathbf{x}}\frac{1}{1^{\top }\mathbf{z}}{\sum }_i{z}_i{G}_n^i\left(\mathbf{x}\right)$, and $0\le {\alpha }_n({\mathbf{x}}^n)\le {\sum }_i{G}_n^i({\mathbf{x}}^n)$. Therefore, lim_n→∞α_n(xⁿ) = 0 a.s.