Deep Operator Networks
for Bayesian Parameter Estimation in PDEs

We start by considering a parametrized partial differential equation (PDE), defined over time span $t\in[t_{0},t_{f}]$ and spatial domain $x\in\Omega$:

${F}$ may contain partial derivatives of the solution of the forms $\partialderivative[\alpha]{y}{x}$ and $\partialderivative[\beta]{y}{t}$. $y(t,x;\lambda)$ is the forward solution of the PDE and is assumed to be well defined with the necessary initial and boundary conditions. Given a dataset of noisy observations of $y$ at different times and positions:

we are interested in training a probabilistic machine learning model that simultaneously approximates $y(t,x)$ and $\lambda$ based on the data eq. 2 and the PDE eq. 1. We will start from a Bayesian formulation of this problem and derive a corresponding loss function for training the model.

The posterior distribution over the latent parameter $\lambda$ given data $\mathcal{D}$ is expressed as

where $p_{0}(\lambda)$ represents the prior distribution over $\lambda$, $p(\mathcal{D}\mid\lambda)$ is the likelihood of the observed data given parameter $\lambda$, and $p(\mathcal{D})$ is the marginal likelihood or evidence.

The likelihood term $p(\mathcal{D}\mid\lambda)$ is constructed from two components: the data likelihood $p(y_{i}\mid t_{i},x_{i},\lambda)$, which captures the relationship between the inputs $t_{i},x_{i}$ and outputs $y_{i}$, and the likelihood of the predicted solution satisfying the governing equations. The latter is achieved via the residuals of eq. 1.

Assuming independence across training data points, the joint likelihood is written as

The data likelihood $p(y_{i}\mid x_{i},\lambda)$ is modeled with normal distribution:

where $\mathcal{M}_{\theta}(t_{i},x_{i},\lambda)$ is the model output for input $(t_{i},x_{i})$ given latent variable $\lambda$, and $\sigma_{y}^{2}$ is the observation noise variance. The log-likelihood is therefore

We are interested in models that, in addition to adherence to observation data, also produce likely predictions that satisfy the governing equations eq. 1. The likelihood of the residuals from the PDE is similarly modeled as:

where $\sigma_{R}^{2}$ is the variance of the residual noise. $\mathcal{R}_{\text{PDE}}$ can either be simplified as $|F(t_{i},x_{i},y_{i},\lambda)|$ or modified to include other constraints such as the boundary and initial conditions. The corresponding log-likelihood of the residuals is

We sample the posterior $p(\lambda\mid\mathcal{D})$ in equation eq. 3 using variational inference (VI). Considering an approximate posterior distribution $q(\lambda)$ we minimize the Kullback-Leibler (KL) divergence:

Substituting eq. 3 for $p(\lambda\mid\mathcal{D})$ and rearranging terms, we obtain the well-known Evidence Lower Bound (ELBO) objective function:

Monte Carlo approximation of the integral in eq. 10 yields Expanding $p(\mathcal{D}\mid\lambda)$ using the joint likelihood, eq. 10 becomes

Finally, substituting the log-likelihoods eqs. 6 and 8, the loss function becomes

Equation 12 can be evaluated using Monte Carlo sampling of the latent variable $\lambda$ from the posterior $q(\lambda)$. The observation and residual noise variances $\sigma_{y}^{2}$ and $\sigma_{R}^{2}$ can be considered constants or learned during training. This enables a flexible uncertainty quantification method that can be utilized to estimate model or data uncertainty. We experiment with different variance estimation strategies in the experiments section 3.

The KL divergence term in eq. 12 can be approximated empirically or analytically: If both the prior $p_{0}(\lambda)$ and the variational distribution $q(\lambda)$ are normally distributed, with means $(\mu_{0}=0,\mu_{q})$ and variances $(\sigma_{0}^{2}=1,\sigma_{q}^{2})$, the KL divergence can be written as:

$\mathcal{M}_{\theta}(t,x;\lambda)$ is a deep operator network (DeepONet) that consists of two components: the branch network and the trunk network [30, 31]. The branch network approximates the forward solution $y(t,x)$ of the PDE. The trunk network processes independent prior samples to generate samples of the approximate posterior $q(\lambda)$ that also influence the branch network’s predictions. See fig. 1 for a schematic representation of the network architecture. Our proposed architectue, hereafter referred to as DeepBayONet uses the DeepONet structure to estimate the forward solution and parameters of the PDE at the same time:

where $W_{0}$ and $b_{0}$ are the output layer weights and bias, $\sigma(\cdot)$ is the activation function, and $\odot$ denotes the element-wise product. The branch network $\mathsf{b}(t,x)$ and trunk network $\mathsf{\tau}(\lambda)$ are general parametrized functions that can take the form of any neural network architectures including convolutional, recurrent, or transformer networks. The entire set of trainable parameters in the model is denoted by $\theta$.

Uncertainty estimation is critical in deep neural networks, particularly in scientific modeling and prediction tasks, to assess the reliability of predictions and the robustness of the model. There are two sources of uncertainty that can be considered: Aleatoric uncertainty arises from inherent noise in the observations, reflecting variability that cannot be reduced even with additional data. This type of uncertainty is often captured by modeling the variance of the output distribution, as demonstrated in probabilistic models such as heteroscedastic neural networks [32]. In contrast, epistemic uncertainty reflects the model’s lack of knowledge or suffiecient training about the underlying data distribution and can be reduced by incorporating additional data or improving the model architecture. Bayesian neural networks (BNNs) and dropout are among various methods that cab address epistemic uncertainty [32] by learning a distribution over the model parameters, allowing the model to express uncertainty in regions of the input space the model’s predictions are ambiguous. These two sources of uncertainty provide complementary information: aleatoric uncertainty quantifies noise in the data, while epistemic uncertainty assesses the model’s confidence in its predictions.

In the proposed model, both aleatoric and epistemic uncertainties are quantified explicitly through different components of the loss function and the model’s probabilistic formulation. Aleatoric uncertainty is captured through the observation noise variance $\sigma_{y}^{2}$, which governs the data likelihood term in the evidence lower bound (ELBO) formulation eq. 12 where $\sigma_{y}^{2}$ represents the variance of the observation noise. This term accounts for the inherent variability in the measurements, and by treating $\sigma_{y}^{2}$ as a learnable parameter, the model can adaptively estimate the aleatoric uncertainty directly from the data. High values of $\sigma_{y}^{2}$ indicate regions with significant observation noise, leading to reduced confidence in the predictions.

Epistemic uncertainty, on the other hand, is inferred from the posterior distribution $q(\lambda)$ over the latent parameters $\lambda$, as derived through variational inference. The posterior distribution is approximated by minimizing the Kullback-Leibler (KL) divergence in eq. 12. The variability in $q(\lambda)$, often characterized by its variance $\sigma_{q}^{2}$, quantifies the model’s uncertainty about $\lambda$. Furthermore, the residual noise variance $\sigma_{R}^{2}$, associated with the residuals of the governing equation eq. 8 also contributes to epistemic uncertainty by capturing how well the predicted solutions satisfy the physical constraints imposed by the partial differential equations (PDEs).

Through Monte Carlo sampling of eq. 12, the model evaluates both aleatoric and epistemic uncertainties simultaneously, enabling a robust and interpretable quantification of uncertainty in predictions.

A distinctive feature of the DeepBayONet architecture is the interaction between the branch and trunk networks, which not only enables the model to represent the forward solution but also plays a role in uncertainty quantification. Specifically, the branch network $\mathsf{b}(t,x)$ approximates the forward operator by mapping the inputs $t$ and $x$ to a latent feature space. Simultaneously, the trunk network $\mathsf{\tau}(\lambda)$ processes samples of the latent parameter $\lambda$ drawn from the approximate posterior distribution $q(\lambda)$. These two networks interact through an element-wise product in eq. 14. The stochastic sampling of $\lambda\sim q(\lambda)$ introduces variability in $\mathsf{\tau}(\lambda)$, thereby perturbing the forward operator represented by $\mathsf{b}(t,x)$ during training.

This mechanism of perturbation bears a conceptual resemblance to dropout [33], widely-used regularization as well as epistemic (model) uncertainty quantification technique in neural networks. Dropout introduces noise during training by randomly setting activations to zero with a probability sampled from a Bernoulli distribution [33]. In DeepBayONet, instead of discrete binary perturbations, the trunk network introduces continuous perturbations drawn from the posterior distribution $q(\lambda)$. This continuous modulation of the branch network activations improves robustness while allowing the model to incorporate epistemic uncertainty in its predictions. Sampling from $q(\lambda)$ enables the model to account for uncertainty in the latent parameters $\lambda$. This uncertainty is propagated through the interaction between the branch and trunk networks. This continuous perturbation mechanism has two significant implications. First, it ensures that the model’s predictions generalize well by preventing the branch network from overfitting to deterministic solutions. Second, it provides a means to quantify epistemic uncertainty explicitly, as the variability in $\mathsf{\tau}(\lambda)$ represents the model’s confidence in its parameter estimates. Thus, the DeepONet architecture naturally integrates regularization and uncertainty quantification in a unified framework, leveraging the interaction between the branch and trunk networks to achieve both objectives.

The loss function defined in eq. 12 enforces adherence to the governing equations of the dynamical system while capturing uncertainties in both the observations and the parameters. To ensure practical implementation, the total loss is decomposed into distinct components, each addressing specific aspects of the problem. Assuming $\mathcal{M}_{\theta}(t_{i},x_{i};\lambda)=\tilde{y}_{i}$ is the predicted PDE solution, the total loss is expressed as:

where the individual loss components are defined to correspond with the probabilistic formulation in eq. 12.

To align these loss terms with the evidence lower bound (ELBO), variance-based scaling is introduced to reflect the uncertainties in the observations and residuals. The revised definitions of the individual components are:

where $\sigma_{R}^{2}$ is the variance of the residual noise and $\sigma_{y}^{2}$ is the variance of the observation noise. These variances play a critical role in balancing the contributions of the individual loss terms, ensuring that the model accounts for uncertainties appropriately.

The weights $w_{\text{Interior}},w_{\text{IC}},w_{\text{BC}},$ and $w_{\text{STD}}$ are hyperparameters used to balance the influence of the respective loss components. These weights can be selected through hyperparameter optimization or by employing adaptive weighting techniques, as explored in [34, 35]. The inclusion of variance scaling in eqs. 15a, 15b and 15c ensures consistency with the ELBO formulation, connecting the practical implementation to the probabilistic foundations of the model. This alignment also allows the model to accurately capture both aleatoric and epistemic uncertainties during training and inference.

The common configurations for all the experiments are as follows: each experiment utilizes a neural network architecture comprising three layers on both the branch and the trunk. The hyperbolic tangent ($\tanh$) activation function is employed universally. Training and optimization are conducted using the Adam optimizer. All implementations are performed in PyTorch, leveraging its advanced features for computational graphs and gradients. Specific details, such as the number of neurons per layer, epochs, batch size, and loss weights, are tuned for each experiment based on its requirements and complexity.

This experiment demonstrates DeepBayONet’s capability to estimate aleatoric and epistemic uncertainty in a synthetic regression problem. The regression task is defined as follows:

where $y_{i}$ are noisy observations with state-dependent inhomogeneous noise. The goal is to learn the function $f(x)$ while estimating two sources of uncertainty: Aleatoric uncertainty arising from state-dependent input noise in the data, and the epistemic uncertainty due to gaps in the data and testing outside the training range eq. 16.

The results are compared against various deep learning uncertainty quantification methods, including a Simple Neural Network with aleatoric uncertainty estimate (SNN), a Bayesian Neural Network (BNN), a Monte Carlo Dropout Network (MCDO), and a Deep Ensemble Network (DENN). Each architecture contains approximately 3,800–4,000 parameters to ensure a fair comparison. The data are split into a training set, an in-training-distribution (IDD) test set, and an out-of-training-distribution (OOD) test set.

The test problem and implementations are adopted from [36]. For a comprehensive review of uncertainty quantification methods in deep learning, refer to [37]. All networks are trained for 150 epochs with the Adam optimizer (batch size 16). Each model is run ten times with random initializations, and accuracies are averaged to mitigate random effects. Common hyperparameters are kept consistent across all models, while method-specific hyperparameters use default library values. To limit computational resource requirements, we have not done method-specific hyperparameter tuning.

We report training and testing mean squared errors (MSE) and the coverage of 95% confidence intervals (CI) predicted by the networks. This design facilitates an evaluation of both predictive accuracy and the quality of uncertainty estimates within and outside the training range.

The results demonstrate that DeepBayONet outperforms other architectures by achieving consistently lower mean squared errors (MSE) and significantly higher confidence-interval coverage. Notably, DeepBayONet excels in out-of-training-distribution scenarios, highlighting its robustness and reliability in quantifying uncertainty across diverse testing conditions.

This experiment evaluates the ability of DeepBayONet to approximate a one-dimensional function while simultaneously estimating an unknown parameter and quantifying associated uncertainties. The function to be approximated is defined as:

where $x\in[-1,1]$ and $\omega$ is an unknown parameter. For this experiment, the true value of $\omega=6$. The dataset consists of sampled values of $x$ from the domain paired with noisy observations:

Two noise scenarios are considered: a low-noise case with $\sigma=0.01$ and a high-noise case with $\sigma=0.1$.

Figure 2 shows the mean predicted function as well as the two standard deviation uncertainty bands for both high- and low-noise cases. The model demonstrates good agreement with the data within the training range, while predicting elevated uncertainty outside the training range. Higher observation noise ($\sigma=0.1$) results in more pronounced uncertainty, reflecting the model’s ability to adapt to noisy data.

The posterior distribution of the parameter $\omega$ is shown in fig. 3, with the mode close to the true value of $\omega=6.0$ in both cases. The narrower posterior distribution in the low-noise scenario ($\sigma=0.01$) indicates greater confidence in the parameter estimate compared to the high-noise scenario.

This experiment explores DeepBayONet’s capability to solve a one-dimensional unsteady heat equation while simultaneously estimating the diffusion parameter $D$ and the decay rate $\alpha$. The heat equation is given as:

where $D$ and $\alpha$ are unknown parameters. When $D=1$ and $\alpha=1$, the equation has the exact solution:

A DeepBayONet model is trained to approximate this solution while learning the parameters $D$ and $\alpha$. The training dataset consists of 100 spatio-temporal samples with the corresponding exact solution $y(t_{i},x_{i})$. A standard normal prior is used for both parameters. Training is performed for 15,000 epochs with a batch size of 100, using the Adam optimizer and an initial learning rate of 0.01. The loss component weights are configured as follows:

The model contains 2,103 trainable parameters.

During training, several metrics were recorded to monitor the model’s performance and convergence, including the total loss and individual loss components: PDE (Interior) loss, boundary condition (BC) loss, initial condition (IC) loss, data loss, and the standard deviation loss for predicted parameters (STD loss). These metrics are shown in fig. 4, highlighting the stability and convergence of the model.

Figure 5 compares a sample of the posterior predicted solution $y(t,x)$ with the exact solution. The grainy appearance of the predicted solution reflects the uncertainty in the parameters $D$ and $\alpha$, which affects the pointwise forward solution. However, obtaining multiple samples of the output and computing statistical measures, such as the mean and mode, is computationally efficient once the model is trained.

The histograms in fig. 6 illustrate the posterior distributions of the parameters $D$ and $\alpha$. The model successfully predicts these parameters, with the posterior mode aligning closely with the true values ($D=1.0$ and $\alpha=1.0$). Additionally, the secondary mode observed in the posterior distributions corresponds to larger values of $D$ and $\alpha$, indicating a highly diffusive and fast-decaying mode that still satisfies the PDE relatively well.

Figure 7 further examines the errors and residuals for the DeepBayONet learned solution. Figure 7a shows the absolute error compared to the exact solution when using the mean of the posterior distribution of the parameters. Small errors are observed across the domain, except near the initial condition boundaries ($t=0,x=\pm 1$).

The PDE residuals, when evaluated using the primary mode (fig. 7b) and the secondary mode (fig. 7c) of the posterior distribution, provide additional insights. The secondary mode results in significantly larger residuals, reflecting suboptimal parameter values that deviate from the true solution.

This experiment evaluates DeepBayONet’s capability to solve a two-dimensional reaction-diffusion equation while estimating the reaction rate $k$. The equation is defined as:

For $k=1$ and $\lambda=0.01$, the PDE has a known solution as given in eq. 19c. The training dataset consists of 100 randomly placed sensors within the domain that observe $u$ and $f$, with added observation noise of variance $0.01^{2}$. Additionally, 25 sensors are uniformly placed along each boundary to enforce Dirichlet boundary conditions.

Figure 8 shows the sensor positions along with contours of $u$ and $f$. The model is configured with 300 neurons per hidden layer, totaling 272,703 trainable parameters. It is trained for 35,000 epochs using the Adam optimizer with a batch size of 500. Loss component weights are calibrated as:

This setup ensures the optimizer appropriately balances the various constraints.

The training history, as shown in fig. 9, demonstrates consistent reduction in the PDE (Interior) and data losses during early epochs, followed by a decrease in the boundary condition loss. The variance of the posterior parameter distribution remains relatively unchanged due to the KL divergence term in the loss function.

Table 3 compares the estimated mean and standard deviation of the reaction rate $k$ using DeepBayONet and alternative Bayesian methods. DeepBayONet achieves results closest to the true value $k=1$ with the lowest standard deviation, demonstrating both accuracy and certainty in parameter estimation.

We conducted a second experiment by relaxing boundary condition enforcement ($w_{\text{BC}}=0$). Figure 10 shows the residuals of the PDE when using the mean posterior parameter value. Higher residuals are observed at the boundaries when boundary conditions are excluded from training.

The posterior distributions of $k$ are shown in fig. 11. With boundary conditions, the posterior is narrow and centered around $k=1.0$. Without boundary conditions, the posterior is wider, skewed, and exhibits a long tail, indicating increased uncertainty towards positive reaction rates.

Figure 12 shows the mean predicted solution and its error along with the predicted uncertainty for boundary-enforced training. In fig. 13, the corresponding plots for boundary-free training reveal reduced accuracy and elevated uncertainty, highlighting the ability of model to adapt to missing boundary conditions.