## α -VI D eep ON et : A prior robust variational -B ayesian approach for enhancing D eep ON ets with uncertainty quantification

## Soban Nasir Lone

## Subhayan De

Department of Applied Mechanics Indian Institute of Technology Delhi

Department of Mechanical Engineering Northern Arizona University

## Rajdip Nayek

Department of Applied Mechanics Indian Institute of Technology Delhi rajdipn@am.iitd.ac.in

## A bstract

We introduce a novel deep operator network (DeepONet) framework that incorporates generalised variational inference (GVI) using Rényi's α -divergence to learn complex operators while quantifying uncertainty. By incorporating Bayesian neural networks as the building blocks for the branch and trunk networks, our framework endows DeepONet with uncertainty quantification. The use of Rényi's α -divergence, instead of the Kullback-Leibler divergence (KLD), commonly used in standard variational inference, mitigates issues related to prior misspecification that are prevalent in Variational Bayesian DeepONets. This approach o ff ers enhanced flexibility and robustness. We demonstrate that modifying the variational objective function yields superior results in terms of minimising the mean squared error and improving the negative log-likelihood on the test set. Our framework's e ffi cacy is validated across various mechanical systems, where it outperforms both deterministic and standard KLD-based VI DeepONets in predictive accuracy and uncertainty quantification. The hyperparameter α , which controls the degree of robustness, can be tuned to optimise performance for specific problems. We apply this approach to a range of mechanics problems, including gravity pendulum, advection-di ff usion, and di ff usion-reaction systems. Our findings underscore the potential of α -VI DeepONet to advance the field of data-driven operator learning and its applications in engineering and scientific domains.

K eywords Prior-Robust Bayesian Inference · Alpha Divergence · Variational Bayes DeepONet · Bayesian Neural Networks · Uncertainty quantification · PDE Surrogate

## 1 Introduction

As scientific machine learning methodologies take centre stage across diverse disciplines in science and engineering, there is an increased interest in adopting data-driven methods to analyse, emulate, and optimise complex physical systems. The behaviour of such systems is often described by laws expressed as systems of ordinary di ff erential equations (ODEs) and partial di ff erential equations (PDEs) [1]. A classical task then involves the use of analytical or computational tools to solve such equations across a range of scenarios, e.g., di ff erent domain geometries, input parameters, and initial and boundary conditions (IBCs). Solving these so-called parametric PDE problems requires learning the solution operator that maps variable input entities to the corresponding solution of the underlying PDE system. Tackling this task using traditional tools (e.g., finite element methods [2]) bears a formidable cost, as independent simulations need to be performed for every di ff erent domain geometry, input parameters, or IBCs. Driven by this challenge, a growing field of neural operators for solving parametric PDEs has come forth recently.

Learning to solve PDEs is closely related to operator learning. Instead of learning to solve a specific PDE, it can be advantageous to learn the operator that maps a functional parameter of the PDE (such as initial values, boundary conditions, force fields, or material parameters) to the solution associated with the given parameter. Neural operators

exemplify this approach by learning to solve entire classes of PDEs simultaneously. These specially designed deep neural networks can represent the solution map of parametric PDEs in a discretisation-invariant manner, meaning the model can be queried at any arbitrary output location. Unlike neural networks that map between two finite-dimensional vector spaces, neural operators can represent mappings between spaces with infinite dimensions [3]. Popular operator learning frameworks include deep operator networks (DeepONets) [4] and Fourier neural operator (FNO) [5]. This article focuses on DeepONet, which has been applied to a wide range of problems, including solid and fluid mechanics [6], reliability analysis [7], heat transfer [8], and fracture mechanics [9]. De et al. [10] utilised DeepONets for modelling uncertain and partially unknown systems. Cai et al. [11] applied DeepONets to model field variables across multiple scales in multi-physics problems while He et al. [12] employed DeepONets to predict full-field, nonlinear elastic-plastic stress responses in complex geometries. Other notable applications of DeepONets include weather forecasting [13], reduced-order modelling [14], finance [15], and others [16, 17].

While DeepONets excel at solving di ff erential equations (DEs), a crucial aspect often overlooked is the uncertainty quantification of the predictions. Traditional DeepONet architectures have been found to produce over-confident predictions, implying they are poorly calibrated [18]. This means they might underestimate the true range of possible outcomes, potentially leading to unreliable results. Additionally, they do not quantify the inherent uncertainties associated with their predictions. In engineering applications like aerodynamics or structural analysis, uncertainty quantification is paramount. By quantifying uncertainties in predictions related to phenomena such as stress distribution, engineers can make informed decisions about the risk and reliability of their designs.

Several attempts have been made to incorporate uncertainty quantification within the DeepONet architecture. However, existing approaches face some limitations. Lin et al. [19] proposed a Bayesian DeepONet based on replica-exchange stochastic gradient Langevin di ff usion [20] that considers the standard deviation of the output to be known, which may not always be realistic. Additionally, training the replicas in this approach is computationally expensive. Yang et al. [21] used randomised priors and trained an ensemble of deterministic models, which can also be computationally demanding depending upon the size of the ensemble. While the method of ensembles o ff ers some degree of uncertainty estimation through multiple model predictions, it does not adhere to the more principled framework that is provided by Bayes' rule and hence cannot be considered truly Bayesian.

Variational inference (VI)-based Bayesian methods o ff er a promising framework for incorporating uncertainty quantification into deep learning models. This approach has given rise to Bayesian neural networks (BNNs) [22, 23]. In BNNs, the parameters of the neural network are treated as random variables, which allow the uncertainty in the parameters to propagate through the network, ultimately reflecting the uncertainty in the model predictions. The variational inference approach is particularly advantageous for neural networks with a large number of parameters, as traditional Markov chain Monte Carlo (MCMC)-based methods become computationally expensive to use. Variational inference, on the other hand, provides a scalable alternative for approximating the posterior distribution of the parameters in BNN.

It seems natural to extend the idea of uncertainty quantification in BNNs to that of a Bayesian DeepONet. Building on this idea, Garg et al. [24] proposed VB-DeepONet that uses BNNs as building blocks for the DeepONet architecture. However, VB-DeepONet inherits certain challenges associated with BNNs. Their approach employs a fully factorised standard Gaussian distribution as the prior over the DeepONet model parameters. Such a prior might not accurately reflect the true distribution of these parameters, leading to a sub-optimal approximation of the posterior distribution. As indicated in Li et al. [25], Knoblauch et al. [26], and Wenzel et al. [27], such a misspecified prior can hinder the overall e ff ectiveness of variational inference, where the approximated posterior concentrates around a single point leading to overconfident predictions and underestimated uncertainty [28]. This hinders the model's ability to e ff ectively quantify uncertainties.

We propose to address these limitations in VB-DeepONet by introducing prior-robust variational inference for DeepONets. Standard VI relies on the Kullback-Leibler divergence (KLD), which is sensitive to prior selection, especially when the number of parameters is large. Alternative divergence measures exist that o ff er robustness to misspecified priors. Knoblauch et al. [26] introduced the Generalised Variational Inference (GVI) framework that allows for defining divergence metrics beyond KLD in the context of BNNs. GVI incorporates a hyperparameter that controls the degree of robustness to prior misspecification. This leads to posteriors that are less influenced by priors which deviate significantly from the observed data [26].

In this paper, we extend the GVI framework to DeepONets to achieve greater robustness to prior misspecification. In particular, we propose the use of Rényi's α -divergence [29] as a robust alternative to KLD within the VI framework. This novel operator learning approach using GVI leads to an improved approximated posterior distribution, demonstrably enhancing both prediction accuracy and the quality of uncertainty estimates as evidenced by log-likelihood values on unseen test data.

The rest of the article will provide background on deterministic DeepONets, followed by details of the proposed variational Bayesian framework and four numerical examples highlighting the e ffi cacy of our approach. We conclude by discussing the results of di ff erent hyperparameter settings and recovering standard KLD-VI results as a special case.

## 2 Background on DeepONets

Operator networks were first introduced by Chen and Chen [30], where they considered shallow networks with a single hidden layer. Subsequently, this concept was significantly developed with deep neural network architectures, leading to the creation of DeepONets [4]. To define operators, it is useful to consider two di ff erent classes of functions residing in two separate Banach spaces, A and S , respectively. One class of functions, a ( y ) ∈ A , with domain in Ω Y ∈ R D , while the other class of functions, s ( y ) ∈ S , also having domain in Ω Y . An operator G can be defined as a mapping between these two function spaces, G : A→S such that s ( y ) = G ( a ( y )).

DeepONet is a specialised deep neural network designed to approximate the operator G . In the following sections, we review the DeepONet and its variational Bayesian extension (VB-DeepONet), and discuss their limitations in the context of uncertainty quantification.

## 2.1 Deterministic DeepONets

Consider a governing equation for a physical phenomenon within a spatial domain Ω over a time period (0 , T ]:

$$\mathcal { F } ( s ) ( x, t ) = u ( x, t ), \ ( x, t ) \in \Omega \times ( 0, T ],$$

$$\mathcal { B } ( s ) ( x, t ) = s _ { b } ( x, t ), \ ( x, t ) \in \partial \Omega \times ( 0, T ],$$

$$\mathcal { I } ( s ) ( x, 0 ) = s _ { 0 } ( x ), \ x \in \Omega.$$

Here, F is a nonlinear di ff erential operator, x is the location in space, t is the time; s denotes the solution of the di ff erential equation, and u denotes the external source function. Furthermore, B ( s )( x t , ) = sb ( x ) and I ( s )( x , 0) = s 0 ( x ) define the boundary and initial conditions, respectively. The true solution operator is denoted by G ( u )( x t , ). To keep the notations compact, we denote the tuple ( x t , ) by y .

DeepONets [4] are designed to learn an approximate operator G θ parameterised by a deep neural network with parameter vector θ . The architecture for the DeepONet consists of two main sub-networks: a branch neural network and a trunk neural network (see Fig. 1).

- · Branch Network : This is a neural network, with trainable parameters θ B , that maps from space R M → R P . It takes as input M number of sensor measurements of the function a ( y ) ∈ A , represented by the M -dimensional vector a = GLYPH&lt;2&gt; a GLYPH&lt;0&gt; y 1 GLYPH&lt;1&gt; , . . . , a GLYPH&lt;0&gt; y M GLYPH&lt;1&gt;GLYPH&lt;3&gt; (where y k B { xr , t v } denotes some location xr and some time instant t v ), and it outputs a vector of weights [ b 1 ( a ) , . . . , bP ( a )].
- · Trunk Network : This is a neural network, with trainable parameters, θ T , that takes as input y ∈ Ω Y and outputs a P -dimensional basis vector GLYPH&lt;2&gt; ψ 1 ( y ) , . . . , ψ P ( y ) . GLYPH&lt;3&gt;

The final output of the DeepONet G θ : R M × R D → R is obtained by taking a dot product of the outputs from the branch and trunk networks, which approximates the value of s ( y ):

$$s ( y ) \approx s ( y ; \theta ) = \mathcal { G } _ { \theta } ( a ) ( y ) \approx b _ { 0 } + \sum _ { p = 1 } ^ { p } b _ { p } \left ( a ; \theta _ { B } \right ) \, \psi _ { p } \left ( y ; \theta _ { T } \right ).$$

The trainable parameters of the DeepONet θ consist of the bias b 0 and the combined parameters of the branch and trunk networks, i.e., θ = [ b 0 , θ B , θ T ].

Training a DeepONet involves supervised learning using pairs of the representative function a ( y ) and the corresponding solution s ( y ). A representative training set looks like:

$$\mathcal { D } = \left \{ \left ( a ^ { ( i ) }, y _ { k }, s ^ { ( i ) } \left ( y _ { k } \right ) \right ) \colon 1 \leq i \leq N _ { 1 }, 1 \leq k \leq N _ { 2 } \right \}$$

where a ( ) i is the i th training example of the M -dimensional vector a obtained at M locations of y , and s ( ) i is the corresponding solution function obtained at N 2 locations of y from solving the PDE (1). Note that the N 2 locations may di ff er from M sensor locations where the representative input functions a are measured. The M sensor locations can be either random or uniformly spaced and remain fixed during training.

Figure 1: DeepONet architecture (stacked) used in the present study. Figure 1: Architecture of the Deterministic DeepONet Model . The figure illustrates the deep neural network architecture of the deterministic DeepONet, showcasing the branch network that processes the input function and the trunk network that processes the coordinates, culminating in a point estimate of the network output.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-f47e6fc7-a63a-4b1b-b89b-cb5a64115b6e.jpg)

The parameter vector θ is estimated by solving the optimisation problem over the training dataset D :

$$\hat { \theta } & = \underset { \theta } { \arg \min } \frac { 1 } { N _ { 1 } N _ { 2 } } \sum _ { i = 1 } ^ { N _ { 1 } } \sum _ { k = 1 } ^ { N _ { 2 } } \left ( G _ { \theta } \left ( a ^ { ( i ) } \right ) ( y _ { k } ) - s ^ { ( i ) } ( y _ { k } ) \right ) ^ { 2 } \\ & = \frac { 1 } { N _ { 1 } N _ { 2 } } \sum _ { i = 1 } ^ { N _ { 1 } } \sum _ { k = 1 } ^ { N _ { 2 } } \left ( b _ { 0 } + \sum _ { p = 1 } ^ { P } b _ { p } \left ( a ^ { ( i ) } ; \theta _ { B } \right ) \ \psi _ { p } \left ( y _ { k } ; \theta _ { T } \right ) - s ^ { ( i ) } ( y _ { k } ) \right ) ^ { 2 }. \\. &. &. &. &. &. &. &. &. &. &. &. &. &. &. &. &. &. &. &. &. &.$$

. . . The optimisation problem can be solved using stochastic gradient descent [31] to obtain the optimal parameter values for the DeepONet. This approach is referred to as D-DeepONet in the rest of the article, where 'D' stands for deterministic.

## . 2.2 Variational Bayesian DeepONets (VB-DeepONet)

. . . . . . . . . . . A key limitation of D-DeepONet is its inability to quantify uncertainty in the predictions. As it produces a point estimate, it fails to account for potential measurement errors or uncertainties in the model parameters.

. . References VB-DeepONet [24] addresses the limitations of D-DeepONet by incorporating layers of Bayesian neural networks within both the branch and trunk networks of the DeepONet architecture. This integration allows the model to capture prediction uncertainties, enhancing its robustness against overfitting. However, VB-DeepONet faces challenges related to posterior inference using standard variational inference. A common practice in VB-DeepONet is to employ a fully factorised standard Gaussian distribution as the prior for the DeepONet model parameter vector θ . This prior assumes pairwise independence among model parameters, implying that each parameter distribution is unimodal . Although intended to be non-informative, this choice of prior can still exert significant influence on the posterior distribution given the number of parameters of deep neural networks can be quite large. Hence, this prior might not accurately reflect the true distribution of the parameters, potentially leading to sub-optimal uncertainty quantification. Moreover, standard KLD-VI with mean-field Gaussian variational families often exhibit 'mode-seeking' behaviour, where the approximated posterior tends to concentrate around a single point. Additionally, if the model parameters exhibit high correlations, these VI approximations can produce overly confident predictions, failing to capture the true range of uncertainty. These issues can significantly hinder the VB-DeepONet's ability to e ff ectively quantify uncertainties, a crucial aspect for reliable engineering applications.

.

2

The next section discusses the proposed α -VI DeepONet approach that uses an alternative divergence measure, the Rényi'sα divergence, to address the limitations of VI in VB-DeepONet and achieve robust uncertainty quantification.

## 3 Proposed α -VI DeepONet

We propose using Bayesian neural networks in both the trunk and branch networks (shown in Fig. 2) instead of their deterministic counterparts, similar to VB-DeepONet. However, a fully factorised normal distribution is hardly ideal as a prior for BNNs. Constructing alternative prior beliefs that accurately reflect our judgments could be computationally prohibitive too. Nonetheless, by employing an alternative divergence metric ( D ), more robust posterior beliefs can be produced with an imperfect prior. Therefore, in the proposed approach, we seek to use a divergence metric, di ff erent from KLD, that results in posteriors more robust to poorly specified priors and provides reliable uncertainty quantification.

In the deterministic DeepONet, the parameter vector θ of the deep neural network (from Eq. (4)) is a point estimate. However, as noted in [32], multiple di ff erent parameter settings can perform equally well, implying that the set of values that each parameter can take may be captured by a distribution over those plausible values. A Bayesian paradigm allows us to place distributions on the model parameters, representing uncertainty about their exact values. Using Bayes' rule, we update these parameter distributions with the data ( D from Eq. (3)) as follows:

$$p ( \theta | \mathcal { D } ) = \frac { p ( \mathcal { D } | \theta ) p ( \theta ) } { p ( \mathcal { D } ) },$$

where p ( D| θ ) is the likelihood and p ( θ ) is the prior distribution over the model parameter vector. The denominator, p ( D ) = R p ( D| θ ) p ( θ ) d θ , is often an intractable constant known as marginal likelihood. The posterior distribution, p ( θ |D ), allows us to perform inference on unseen data ( D ∗ ). For prediction, we use this posterior to compute the predictive distribution:

$$p ( \mathcal { D } ^ { * } | \mathcal { D }, \theta ) = \, \int p ( \mathcal { D } ^ { * } | \theta ) p ( \theta | \mathcal { D } ) d \theta.$$

We now examine these individual components representing our modelling choices that influence the learned posterior distribution.

## 3.1 Fully Factorised Prior

In BNNs, we often default to a fully factorised standard normal distribution as a prior over the model parameters. This implies that each parameter component is assumed to be independent and drawn from a Gaussian distribution (denoted by N ) with zero mean and unit variance. Mathematically, this can be expressed as:

$$p ( \theta ) = \prod _ { l = 1 } ^ { L } p ( \theta _ { l } ) = \prod _ { l = 1 } ^ { L } \mathcal { N } ( \theta _ { l } ; 0, 1 ) = \mathcal { N } ( \theta ; 0, \mathbf I _ { L } ),$$

where L denotes the dimensionality of the model parameter vector θ . It is important to acknowledge that this fully factorised prior is a 'misspecified prior' as it does not capture the interdependencies that exist between the individual components θ l of the parameter vector θ . To alleviate the negative influence of this 'misspecified prior,' we use a di ff erent divergence metric in Section 3.4.

## 3.2 Likelihood Function

The likelihood function describes the probability of observing the solutions given the model parameters (and deterministic input functions). We write the likelihood function considering that each example i is independent and identically distributed, given the input functions and parameters, as follows:

$$p ( \mathcal { D } | \theta ) = \prod _ { i = 1 } ^ { N _ { 1 } } p \left ( s ^ { ( i ) } ( y _ { 1 } ), \dots, s ^ { ( i ) } ( y _ { N _ { 2 } } ) \, | \, a ^ { ( i ) }, \theta \right ),$$

where N 1 represents the total number of examples of pairs a and s . We further assume that the outputs s ( ) i ( y k ) at di ff erent locations are also statistically independent of each other k = 1 , . . . , N 2 . Consequently, the joint distribution of the solution for the i th example across all N 2 locations factorises into a product of independent marginal distributions for each location:

$$p ( \mathcal { D } | \theta ) = \prod _ { i = 1 } ^ { N _ { 1 } } \prod _ { k = 1 } ^ { N _ { 2 } } p \left ( s ^ { ( i ) } ( y _ { k } ) \, | \, a ^ { ( i ) }, \theta \right ).$$

Figure 1: DeepONet architecture (stacked) used in the present study. Figure 2: Architecture of the α -VI DeepONet Model . The figure depicts the deep neural network architecture of the α -VI DeepONet, where both the branch and trunk networks are replaced with Bayesian neural networks. The output is a random variable characterised by a mean and a standard deviation, providing a probabilistic representation of the response.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-582abd51-aacd-4c2a-b7be-bfafda1af65d.jpg)

Furthermore, we assume that the individual output distributions p GLYPH&lt;16&gt; s ( ) i ( y k ) | a ( ) i , θ GLYPH&lt;17&gt; at each location follow a Gaussian distribution. The mean and standard deviation of these distributions depend on the input function a , location y , and model parameters θ . This dependence reflects the influence of the input function and location on the solution value, mediated by the model parameters. Mathematically, this is expressed as:

.

$$\text{re} \, \text{mouvel parameters.  ivaueuemabulary,   units is expressedu as.} \\ p ( \mathcal { D } | \theta ) & = \prod _ { i = 1 } ^ { N _ { i } } \prod _ { k = 1 } ^ { N _ { 2 } } \mathcal { N } \left ( s ^ { ( i ) } ( y _ { k } ) ; \mu _ { k } \left ( a ^ { ( i ) }, y _ { k }, \theta \right ), \sigma _ { k } \left ( a ^ { ( i ) }, y _ { k }, \theta \right ) \right ) \\ & = \prod _ { i = 1 } ^ { N _ { 1 } } \prod _ { k = 1 } ^ { N _ { 2 } } \mathcal { N } \left ( s ^ { ( i ) } _ { k } ; \mu _ { k } ( \theta ), \sigma _ { k } ( \theta ) \right ), \\ \text{sents the mean and } \sigma _ { k } \, \text{represents the standard deviation.   For notational convenience, we use the shorthand } \, \text{and} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$} \, \text{$\lambda$.}$$

. . A sample prediction from this network represents a stochastic pass through the random BNN parameters, followed by a sample from the Gaussian likelihood. A number of these samples are taken to quantify the total uncertainty and construct the confidence intervals.

. . . . . . . . . . . . . where µ k represents the mean and σ k represents the standard deviation. For notational convenience, we use the shorthand notation s ( ) i k ≡ s ( ) i ( y k ) to represent the solution value at location y k for the i th solution data. Additionally, we only highlight the dependence of the mean and standard deviation on the model parameters θ for brevity. The other variables, a ( ) i and y k , are treated as deterministic (and known) quantities in this context.

## 3.3 Posterior Approximation via Variational Inference

Within a fully Bayesian framework, our goal is to determine the posterior distribution of the model parameters given the observed data. This posterior distribution, denoted as p ( θ |D ), represents the probability of various parameter configurations ( θ ) after observing the data D . Unfortunately, due to the intractable normalising constant as explained in Section 3, directly computing the posterior distribution analytically is not feasible.

References Generally, there are two approaches to solve this class of problems sampling-based approaches , e.g., Markov chain Monte Carlo and approximation-based optimisation approaches , e.g., variational inference, expectation propagation (EP). Sampling-based approaches, exemplified by MCMC-techniques, are often regarded as the gold standard for Bayesian inference [33]. They involve constructing a Markov chain that gradually explores the parameter space,

2

eventually converging to a distribution proportional to the posterior. However, for problems with high-dimensional parameter spaces (often exceeding 20,000 parameters in our case), these approaches can become computationally impractical [34]. We note here that the proposed approach puts no limit on the complexity of the model, that is the number of parameters. Hence, we will focus on the approximation-based approaches, where a simpler variational distribution is introduced to approximate the posterior, and the key lies in optimising the parameters of the variational distribution to minimise the Kullback-Liebler divergence between itself and the intractable posterior. This optimisation process allows us to e ffi ciently train complex models despite the high dimensionality of the parameter space.

## 3.3.1 Variational Inference with KLD

VI [35] introduces a surrogate posterior family of distributions, denoted by Q , that is tractable to sample from. The goal is to find a member of Q distribution that closely resembles the true posterior in terms of the shape. VI accomplishes this by minimising the KLD between the approximate posterior q ( θ η ; ) ∈ Q and the true posterior p ( θ |D ). The KLD, denoted as DKL q [ || p ], quantifies the di ff erence between the q distribution and the reference p distribution weighted by the q distribution, and is mathematically defined as:

$$D _ { K L } \left [ q ( \theta ) \left \| \ p ( \theta ) \right ] = \int q ( \theta ) \log \frac { q ( \theta ) } { p ( \theta ) } d \theta = \mathbb { E } _ { q ( \theta ) } \left [ \log ( q ( \theta ) ) - \log ( p ( \theta ) ) \right ].$$

Using the KLD, the minimisation problem in VI is formulated as finding the best parameters η ∗ of the variational distribution that brings it closer to the true posterior:

$$q ( \theta ; \eta ^ { * } ) & = \underset { q ( \theta ; \eta ) \in Q } { \arg \min } D _ { K L } \left [ q ( \theta ; \eta ) \left \| \, p ( \theta | \mathcal { D } ) \right ],$$

where η are the parameters of the variational distribution, q ( θ η ; ). Here, the KLD measures the information lost by approximating the true posterior with the variational distribution q . However, directly computing this integral is impractical because the true posterior is unknown. VI addresses this by leveraging the connection between the KLD and the evidence lower bound ( ELBO ). VI rewrites the KLD using the following identity:

$$D _ { K L } \left [ q ( \theta ; \eta ) \left \| p ( \theta | \mathcal { D } ) \right ] & = \mathbb { E } _ { q ( \theta, \eta ) } \left [ \log ( q ( \theta ; \eta ) ) - \log \left ( \frac { p ( \mathcal { D } | \theta ) p ( \theta ) } { p ( \mathcal { D } ) } \right ) \right ] \\ & = \log p ( \mathcal { D } ) - \underbrace { \mathbb { E } _ { q ( \theta, \eta ) } \left [ \log ( p ( \mathcal { D } | \theta ) \right ) \right ] - D _ { K L } \left ( q ( \theta ; \eta ) \left \| p ( \theta ) \right ) } _ { \mathbb { E L B O } } ( 1 3 ) \\ & = \log p ( \mathcal { D } ) + \mathcal { L } ( q ), \\ \dots \, a. \, \dots \dots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots \cdots.$$

where the expectation in the ELBO (in Eq. (13)) is taken with respect to the approximate posterior q ( θ η ; ). Note that p ( D ) is the model evidence, which is a constant term for a given dataset and also constant with respect to q ( θ η ; ) and can therefore be ignored during the minimisation process Eq. (12). From Eq. (13), we see that maximising the ELBO is equivalent to minimising the KLD (as KLD is non-negative). So, VI aims to maximise the ELBO which turns out to be the minimisation of the variational free energy L ( q ) (see Eq. (14). This translates to finding an approximate posterior q ( θ η ; ) that balances two key factors (common in Bayesian inference):

- · Data fit : The first term, E q ( θ η ; ) [log( p ( D| θ ))], represents the expected log-likelihood of the data under the approximate posterior. A high value for this term indicates a good fit between the model and the data.
- · Prior regularisation : The second term, DKL q [ ( θ η ; ) || p ( θ )], is the KL divergence between the approximate posterior and the prior distribution, p ( θ ). It acts as a regulariser, penalising overly complex posterior distributions that deviate significantly from the prior.

The expressions of these terms highlight the impact of our prior modelling choices on the optimisation problem. A poorly specified prior - independent standard Gaussian distribution - over parameters, as is the case with BNNs, can degrade the quality of the posterior approximation. However, it remains the most expedient choice for BNNs from the perspective of computational expense. As suggested by Knoblauch et al. [26], we can opt for a di ff erence divergence metric instead of KLD to alleviate the negative influence of misspecified prior. For example, KLD is known to exhibit a mode-seeking behaviour [25, 36]. By selecting a di ff erent divergence metric one can improve the quality of the posterior distribution, both in terms of fitting the data and providing better uncertainty estimates [25, 26, 37].

In this context, we examine the e ff ect of using the flexible Rényi'sα divergence on the posterior predictive performance. We modify the variational free energy objective by replacing the KLD with Rényi'sα divergence:

$$\mathcal { L } ( q ) = - \mathbb { E } _ { q ( \theta ; \eta ) } [ \log p ( \mathcal { D } | \theta ) ] + D _ { A R } ^ { ( \alpha ) } [ q ( \theta ; \eta ) | | p ( \theta ) ].$$

To compute the expected value of the log-likelihood term, we use Monte-Carlo integration by drawing samples from the variational distribution q given the parameters η . The likelihood function is computed using Eq. (10) and we approximate the first term in Eq. (15), as follows:

$$\text{ approximate the first term in Eq. (15), as follows:} \\ \mathbb { E } _ { \varrho ( \theta, \eta ) } [ \log ( p ( \mathcal { D } ) \theta ) ] \approx \frac { 1 } { N _ { q } } \sum _ { c = 1 } ^ { N _ { q } } \log p ( \mathcal { D } ) \theta ^ { ( c ) }, \quad \theta ^ { ( c ) } \sim q ( \theta ; \eta ) \\ = \frac { 1 } { N _ { q } } \sum _ { c = 1 } ^ { N _ { q } } \log \left ( \prod _ { i = 1 } ^ { N _ { i } } \prod _ { k = 1 } ^ { N _ { 2 } } \mathcal { N } \left ( s _ { i } ^ { ( i ) } ; \mu _ { k } \left ( \theta ^ { ( c ) } \right ), \sigma _ { k } \left ( \theta ^ { ( c ) } \right ) \right ) \right ), \quad \theta ^ { ( c ) } \sim q ( \theta ; \eta ) \\ = \frac { 1 } { N _ { q } } \sum _ { c = 1 } ^ { N _ { q } } \sum _ { i = 1 } ^ { N _ { 2 } } \sum _ { k = 1 } ^ { N _ { 2 } } \log \left ( \mathcal { N } \left ( s _ { k } ^ { ( i ) } ; \mu _ { k } \left ( \theta ^ { ( c ) } \right ) \sigma _ { k } \left ( \theta ^ { ( c ) } \right ) \right ) \right ), \quad \theta ^ { ( c ) } \sim q ( \theta ; \eta ). \\ \text{where} \, N _ { \varrho } \text{ is the number of Monte Carlo samples, and } \theta ^ { ( c ) } \text{ is the c ^ { \text{th} } \text{ sample drawn from } q ( \theta ; \eta ). \text{ In practice, we found that}$$

where Nq is the number of Monte Carlo samples, and θ ( c ) is the c th sample drawn from q ( θ η ; ). In practice, we found that using 25 Monte Carlo samples ( Nq = 25) is su ffi cient for accurate approximation, aligning with findings in [38] which suggest that even a single sample can yield satisfactory results. In the next section, we detail Rényi's α -divergence as an alternative to KLD, which represents the second term of our objective function.

## 3.4 Rényi'sα divergence

Rényi's α -divergence, introduced by Rényi [29], o ff ers a robust alternative to KLD. It is characterised by a hyperparameter α , which controls the degree of robustness, and it converges to KLD as a limiting case when α approaches 1. We denote this divergence by D AR ( α ) and follow the parameterisation provided in [39]:

$$D _ { A R } ^ { ( \alpha ) } [ q ( \theta ) | | p ( \theta ) ] = \frac { 1 } { \alpha ( \alpha - 1 ) } \log \left ( \mathbb { E } _ { q ( \theta ) } \left [ \left ( \frac { p ( \theta ) } { q ( \theta ) } \right ) ^ { 1 - \alpha } \right ] \right ).$$

This divergence can be computed in closed form when both q and p are normal distributions. However, in general, we approximate this divergence across various distributions by Monte-Carlo sampling. In the limit as α approaches 1, D AR ( α ) simplifies to DKL :

$$\lim _ { \alpha \to 1 } D _ { A R } ^ { ( \alpha ) } [ \cdot ] = D _ { K L } [ \cdot ].$$

In the context of our problem, D AR ( α ) represents the second term of the objective function in Eq. (15), which is approximated as follows:

$$D _ { A R } ^ { ( \alpha ) } [ q ( \theta ; \eta ) | | p ( \theta ) ] \approx \frac { 1 } { \alpha ( \alpha - 1 ) } \log \left ( \frac { 1 } { N _ { q } } \sum _ { c = 1 } ^ { N _ { q } } \left [ \left ( \frac { p ( \theta ^ { ( c ) } ) ) } { q ( \theta ^ { ( c ) } ; \eta ) } \right ) ^ { 1 - \alpha } \right ] \right ),$$

where θ ( c ) denotes the c th Monte Carlo sample from q ( θ η ; ).

In Fig. 3, we plot divergence values D AR ( α ) between q ( θ ) and p ( θ ), assuming both are normal distributions with p ( θ ) as a standard normal distribution. The divergence value decreases as α increases, with KLD as the case when α = 1. Note that the objective function in Eq. (15) contains a data-loss term and a complexity term. The behaviour of the optimisation problem di ff ers significantly depending on whether α is greater than or less than 1. For α &gt; 1, D AR ( α ) exhibits a mode-seeking behaviour, causing the posteriors to become more concentrated while being penalised less for deviating from the prior distribution. In this regime, the objective function places more emphasis on the data-loss term, and the complexity term - which measures the divergence between the prior and the variational posterior - receives less weight as α increases. In contrast, for α &lt; 1, D AR ( α ) encourages a mass-covering behaviour, akin to expectation propagation. As α decreases below 1, the values of D AR ( α ) increase, resulting in more weight being assigned to the complexity term and less to the data-loss term. It is important to note, however, that a larger regularisation term does not necessarily imply that a misspecified prior will dominate the inference, as demonstrated in the experiments by [26].

Given that the true posterior is unknown, it is uncertain whether the mode-seeking or mass-covering behaviour is preferable, making the choice problem-dependent. For scenarios where the true posterior is multimodal, a divergence that produces wider parameter variances (i.e., mass-covering behavior) may be beneficial for achieving more robust predictive estimates. In particular, we aim to set D AR ( α ) such that the posterior is robust to priors that strongly contradict observed data, and provides reliable uncertainty quantification. We rely on validation sets and adjust α as a hyperparameter to minimise metrics such as the normalised mean squared error or negative log-likelihood.

Figure 3: The figure depicts the variation of Rényi's α -divergence D AR ( α ) between two normal distributions for di ff erent values of α . We take p ( θ ) to be a standard normal distribution and q ( θ ) to be a normal distribution with a randomly selected mean and standard deviation.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-4539c9f9-82e5-403d-8807-14e074ca47fc.jpg)

## 3.5 GVI Posterior Approximation

We choose the approximate posterior to be the mean-field normal (MFN) variational family given by:

$$Q _ { M F N } = \prod _ { l = 1 } ^ { L } \mathcal { N } ( \theta _ { l } ; \mu _ { l } ^ { q }, \sigma _ { l } ^ { q } ).$$

We find the q ( θ η ; ) within the QMFN family that minimises the objective function. This class of approximate variational family is commonly used for BNNs. Once the variational family is decided, the optimisation reduces to finding the parameters, η = [ µ q , σ q ], of the variational family that minimises the loss, L ( q ) (Eq. (15)). Thus the optimisation problem tries to find:

$$\eta ^ { * } = \underset { \eta } { \arg \min } \{ \mathcal { L } ( q ) \}.$$

## 3.6 Training

To optimise the variational free energy objective function (acting as the loss function), as defined in Eq. (21), we employ the Bayes-by-Backprop algorithm [40]. This algorithm approximates the expectation within the objective function through Monte Carlo sampling. Specifically, we use Nq = 25 samples for this approximation.

A key technique for e ffi cient training is the reparameterisation trick [41]. This allows us to automatic di ff erentiation to compute the unbiased gradients of the loss function with respect to the parameters and hence enables the use of gradient-based optimisers commonly employed in deep learning. In our implementation, we utilise the Adam optimiser [42] with its default settings. The entire training procedure is carried out in TensorFlow Probability [43].

For Gaussian distributions, the reparameterisation trick involves sampling a standard Gaussian random variable ϵ , and then scaling it by the standard deviation parameter, σ q , and shifting it by the mean, µ q . To ensure the positivity of the standard deviation, we employ a soft-plus activation function. The resulting parameterisation can be expressed mathematically as follows:

$$\theta = \mu ^ { q } + \log ( 1 + \exp ( \sigma ^ { q } ) ) \odot \epsilon$$

where ⊙ denotes element-wise multiplication. Combining the expected negative log-likelihood (Eq. (16)) and the Rényi's α divergence (Eq. (19)), we obtain the final variational free energy objective:

$$\mathcal { L } ( q ) = - \frac { 1 } { N _ { q } } \sum _ { c = 1 } ^ { N _ { q } } \sum _ { i = 1 } ^ { N _ { 1 } } \sum _ { k = 1 } ^ { N _ { 2 } } \log \left ( \mathcal { N } \left ( s _ { k } ^ { ( i ) } ; \mu _ { k } \left ( \theta ^ { ( c ) } \right ), \sigma _ { k } \left ( \theta ^ { ( c ) } \right ) \right ) \right ) + D _ { A R } ^ { ( a ) } \left [ q \left ( \theta ^ { ( c ) } ; \eta \right ) \| p \left ( \theta ^ { ( c ) } \right ) \right ],$$

where θ ( c ) ∼ q GLYPH&lt;16&gt; θ ( c ) ; η GLYPH&lt;17&gt; .

To optimise the variational parameters, η = [ µ q , σ q ]. we compute the gradients using the Bayes-by-Backprop equations:

$$\Delta _ { \mu ^ { q } } = \frac { \partial \mathcal { L } ( q ) } { \partial \theta } + \frac { \partial \mathcal { L } ( q ) } { \partial \mu ^ { q } }$$

$$\Delta _ { \sigma ^ { q } } = \frac { \partial \mathcal { L } ( q ) } { \partial \theta } \frac { \ep$$

where the ∂ L ( q ) ∂ θ term is obtained through standard backpropagation. The overall training procedure is summarised in Algorithm 1.

## Algorithm 1 Training algorithm for α -VI DeepONet

1: procedure T raining A lgorithm (algo) 2: Given training dataset: Arrange raw dataset, D , according to Eq. (3). 3: Initialise: Initialise the α -VI DeepONet parameters, θ ∼ N ( 0 I , L ). 4: for β = 1 to epochs do 5: for c = 1 to Nq do 6: Generate samples from ϵ ( c ) ∼ N ( 0 I , L ). 7: Reparameterisation trick: θ ( c ) = µ q + log(1 + exp( σ q )) ⊙ ϵ ( c ) . 8: Input training data to branch and trunk nets (Eq. (3)). 9: Obtain the predicted distribution (Eq. (10)). 10: end for 11: Compute the loss as per Eq. (23). 12: Compute the gradients ∆ µ q and ∆ σ q using Eq. (24) and Eq. (25). 13: Update the variational parameters, with learning rate λ : 14: µ q ← µ q - λ ∆ µ q 15: σ q ← σ q - λ ∆ σ q 16: end for 17: Output: Trained model.

18: end procedure

## 4 Numerical studies

To assess the performance of our proposed approach, we conduct a numerical study involving four prototypical problems: two ODEs - the antiderivative operator and the gravity pendulum - and two PDEs - di ff usion-reaction and advection-di ff usion equations. These problems represent a diverse spectrum of complexities, allowing for a rigorous evaluation of our model's capabilities.

We investigate the influence of the hyperparameter α on posterior predictive performance while maintaining consistent DeepONet architecture and optimisation parameter settings across all experiments. Detailed architectural specifications for each problem are tabulated in Table 1. To optimise variational parameters, we employ a full-batch Adam optimiser with standard hyperparameter settings: maximum 10000 epochs, learning rate = 0.001, β 1 = 0 9, and . β 2 = 0 999. We . execute ten independent optimisation runs with varying initialisations for each problem, discarding non-convergent cases identified by oscillatory loss behaviour.

The average values and standard deviations of the converged runs are reported for eleven di ff erent values of α in Table 2 and Table 3. We evaluate the model test performance using two complementary metrics: the normalised mean squared error (NMSE) and the negative log-likelihood (NLL). NMSE provides a straightforward measure of prediction error by averaging the squared di ff erences between observed and predicted mean values of responses, normalised by the norm of the observed values. NLL, on the other hand, assesses how well the predicted probability distribution aligns with actual outcomes, accounting for both the prediction mean and its uncertainty.

Table 1: DeepONet architectural details: number of layers (depth) and number of neurons per layer (width).

| Problem               |   Branch width |   Branch depth |   Trunk width |   Trunk depth |
|-----------------------|----------------|----------------|---------------|---------------|
| Anti-derivative       |             25 |              3 |            25 |             3 |
| Gravity pendulum      |             25 |              3 |            25 |             3 |
| Di ff usion-reaction  |             25 |              4 |            25 |             4 |
| Advection-di ff usion |             35 |              4 |            35 |             4 |

## 4.1 Problem 1: Antiderivative Operator

We commence our numerical investigation with the fundamental problem of the antiderivative operator over the domain x ∈ (0 , 1]:

$$\frac { d s } { d x } = u ( x ), \quad \begin{array} {$$

subject to the initial condition, s (0) = 0. In this case, the independent variable is exclusively x , rendering y = x . Our objective is to approximate the solution s x ( ), driven by the input function u x ( ):

Z

x

s x

(

)

=

s

(0)

+

u

(

τ

)d

τ,

x

∈

[0

,

1]

,

0

Here, the input function a y ( ) from the general formulation is equivalent to the source term u x ( ).

To construct the training dataset, we sample input functions u x ( ) from a zero-mean Gaussian random field (GRF) characterised by:

$$u ( x ) & \sim \mathcal { G P } ( 0, \kappa ( x _ { 1 }, x _ { 2 } ) )$$

$$\kappa ( x _ { 1 }, x _ { 2 } ) & = \exp \left ( - \frac { \| x _ { 1 } - x _ { 2 }$$

where κ is the radial basis function (RBF) covariance kernel with length-scale ℓ . As outlined in Eq. (3), we generate a training set comprising N 1 = 3000 input functions u x ( ) discretised at M = 100 equidistant points within the domain (implying a ( ) i ∈ R 100 , i = 1 , . . . , N 1 ). For each input vector a ( ) i , we compute the corresponding solution at N 2 = 20 randomly selected points, i.e., s ( ) i ( xk ) , k = 1 , . . . , N 2 . Notably, once the operator is learned, solutions can be evaluated at any arbitrary location within the domain. To assess generalisation performance, we employ a test set of 10000 input functions, evaluating predicted solutions s x ( ) for each of these test input functions at 100 equidistant locations in the domain.

Fig. 4 showcases the model's predictive capabilities by visualising solutions for two representative test cases chosen to highlight its performance. This ensemble of points, representing the solution queried at the 100 equidistant locations, provides a visual approximation of the mean solution and the associated uncertainty. A comparative analysis with predictions from D-DeepONet and ground truth solutions for these representative examples further underscores the e ffi cacy of using a di ff erent value of α than 1. For a comprehensive assessment, we present a comparison of performance using NMSE and NLL metrics averaged over all 10000 test cases for eleven di ff erent α values ranging from 0.25 to 3.00. Tables 2 and 3 present these results. In both metrics, α = 1 25 achieved the lowest NMSE and NLL. Notably, . the NMSE is reduced by more than 50% compared to the standard KLD-VI (with α = 1 00), indicating a substantial . improvement in mean predictions.

## 4.2 Problem 2: Gravity Pendulum under External Forcing

We next consider the nonlinear dynamics of a gravity pendulum subjected to an external force. The system is governed by the following nonlinear ODE:

$$\frac { d ^ { 2 } \phi } { d t ^ { 2 } } = - k \sin \phi + u ( t ), \quad \dots \cdots$$

where ϕ represents the angular displacement, k is a constant determined by gravitational acceleration and pendulum length, and u t ( ) is the time-dependent external forcing function. The time domain for simulation is set to t ∈ (0 1]. , Transforming this second-order ODE into a state-space form yields:

$$\frac { d s } { d t } = \left [ \frac { \frac { d s _ { 1 } } { d s _ { 2 } } } { \frac { d t } { d t } } \right ] = \begin{bmatrix} s _ { 2 } \\ - k \sin s _ { 1 } + u ( t ) \end{bmatrix},$$

Figure 4: Predictive performance comparison for the Antiderivative Operator . The figure presents a comparative visualisation of the output function predictions generated by α -VI DeepONet and D-DeepONet for two representative test cases (columns (a) and (b)). Each column displays results for a distinct test input function sample. The top row illustrates mean predictions and corresponding 95% confidence intervals (CIs) from α -VI DeepONet with α = 1 25, . which demonstrates optimal performance for this problem. The second row presents predictions from standard-VI DeepONet by setting α = 1 00. The bottom row showcases predictions from the deterministic D-DeepONet model. .

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-d4bc7f4f-5d15-4a80-bac3-fd2df4e94a61.jpg)

with initial conditions s (0) = [0 0] T , where s 1 = ϕ and s 2 = d ϕ d t . Given the purely time-dependent nature of the problem. the domain reduces to a scalar temporal variable, i.e., y = t . For this problem, we consider the generalised input function a y ( ) as equal to the time-dependent forcing function u t ( ).

To construct the training dataset, we set k = 1 and generate N 1 = 3500 training input functions, u t ( ), sampled from the GRF defined in Eq. (27) at M = 100 time points. For each training input function, we compute the corresponding reference solution, s 1 , using a fourth-order Runge-Kutta integrator at N 2 = 20 randomly selected time points. This data is used to train the operator. A test set of 10000 input functions is employed to evaluate the predicted solution for each test input function at 100 equidistant points within the time domain t ∈ (0 1]. Fig. 5 presents the model's predictions , for two representative test inputs.

A visual comparison of the test prediction cases reveals that the deterministic D-DeepONet struggles to accurately capture the reference solution. The standard KLD-VI solution, obtained at α = 1 00, yields an improvement in the . mean predictions, while the α -VI DeepONet with α = 2 00 demonstrates superior performance over both in terms of . mean prediction accuracy, as evident in Table 2. In terms of distributional fit, as measured by negative log-likelihood,

the standard KLD-VI has the lowest NLL values and exhibits a better performance (Table 3). However, the NLL for α = 2 00 is only marginally higher. Considering both the metrics, . α = 2 00 o . ff ers a favourable balance between accurate mean predictions and reasonable distributional fit.

Figure 5: Predictive performance comparison for the Gravity Pendulum . The figure presents a comparative visualisation of the output function predictions generated by α -VI DeepONet and D-DeepONet for two representative test cases (columns (a) and (b)) . Each column displays results for a distinct test input function sample. The top row illustrates mean predictions and corresponding 95% confidence intervals (CIs) from α -VI DeepONet with α = 2 00, . which demonstrates optimal mean prediction performance for this problem. The second row presents predictions from standard-VI DeepONet by setting α = 1 00. The bottom row showcases predictions from the deterministic D-DeepONet . model.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-6b51d009-a950-4e82-8b66-89aa1b7c494a.jpg)

## 4.3 Problem 3: Di ff usion-Reaction System

Weextend our analysis to the di ff usion-reaction PDE, which involves derivatives in both spatial and temporal coordinates. The two coordinates are denoted by the tuple y = { x t , } , where x represents the spatial location and t represents time. Di ff usion-reaction equations model the combined e ff ects of di ff usion and chemical reactions within a system, finding applications in diverse fields where heat transfer, mass transfer, and chemical kinetics occur simultaneously.

Given an external source term u x ( ), we consider the following di ff usion-reaction PDE:

$$\frac { \partial s } { \partial t } = D _ { c } \frac { \partial ^ { 2 } s } { \partial x ^ { 2 } } + k s ^ { 2 } + u ( x ), x \in ( 0, 1 ), \ t \in ( 0, 1 ].$$

This equation represents a di ff usion-reaction system influenced by a source term and characterised by a quadratic dependence on the solution. We set the di ff usion coe ffi cient Dc = 0 01 and the reaction rate . k = 0 01. In this case, the . source term u x ( ) serves as the input function, equivalent to a x ( ) without the dependence on the temporal variable t . We consider the PDE with zero initial and boundary conditions i.e., s x ( , 0) = s (0 , t ) = s (1 , t ) = 0.

To generate the training data, we numerically solve the PDE using a finite di ff erence method on a ( x × t ) ≡ (100 × 100) grid, following a similar approach as in [4]. The training dataset comprises N 1 = 500 distinct source terms, u x ( ), generated from a Gaussian random field with ℓ = 0 5 (Eq. (27)), evaluated at . M = 100 equidistant points in space. For each training input function, the solution s ( y ) is obtained at N 2 = 100 random points y sampled from the 100 × 100 grid; these sampled solution points are denoted by filled white circles in Fig. 6. Subsequently, we employ this training dataset to learn the operator mapping from u x ( ) to the PDE solution s x ( , t ).

Figure 6: Data Generation Process for Di ff usion-Reaction PDE . The figure illustrates the steps involved in creating training data for the di ff usion-reaction PDE. The top left panel depicts a sampled source term, u x ( ), discretised at 100 equidistant points. Using this source term, the PDE is solved numerically on a 100 × 100 spatiotemporal grid, with the resulting solution visualised in the top right panel. The bottom right panel shows the random selection of N 2 = 100 points from this grid, represented by white circles, where the solution is sampled to complete the training data. The final bottom left panel summarises the overall training dataset structure, where s ( y k ) denotes the solution at the corresponding spatial-temporal point y k .

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-879aea5c-3d66-4ee0-9041-7b6d1f29c6aa.jpg)

To assess generalisation performance, we employ a separate test set comprising 10000 diverse source terms. Fig. 7 presents a visual comparison of the model's predicted solution against the reference solution for a representative test source term across the entire 100 × 100 spatiotemporal grid. The figure also includes the associated uncertainty (one standard deviation) and absolute error. It is seen that regions of low response values (visualised as bluish areas) exhibit reduced smoothness in the predicted mean solution compared to the reference. As such, these regions demonstrate higher uncertainty, as evidenced by increased standard deviation.

A comprehensive performance assessment of predictive performance across di ff erent values of α is conducted using NMSE and NLL metrics, the averages of which are summarised in Tables 2 and 3, respectively. Consistent with

previous findings, α values greater than 1 generally outperform the standard KLD-VI case ( α = 1 00). For this specific . problem, α = 3 00 achieved the lowest NMSE, while . α = 2 50 yielded the lowest NLL. Notably, the performance . metrics at α = 2 00 are only marginally inferior to the optimal values, suggesting a degree of robustness in the model's . performance within this parameter range.

Figure 7: Predictive Performance of α -VI DeepONet for the Di ff usion-Reaction Equation . The figure presents a visual comparison of the α -VI DeepONet prediction with the corresponding reference solution for a representative case of a test source term. The top left panel displays the reference solution, while the top right panel shows the predicted mean from α -VI DeepONet at α = 3 00. The bottom left panel illustrates the predicted standard deviation, representing . the uncertainty in the prediction. The bottom right panel depicts the absolute error between the predicted mean and the true solution. All plots share the same colour bar, with values indicated as a function of spatial dimension, x , and time, t , represented on the horizontal and vertical axes, respectively.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-110b0e26-0d2e-42c8-b4fc-14ce1bc4283c.jpg)

## 4.4 Problem 4: Advection-Di ff usion Equation

We now extend our analysis to the advection-di ff usion PDE, which combines both advection (transport due to fluid flow) and di ff usion (random movement of particles). This equation is fundamental to various physical processes, including solute transport in fluids. We consider the following advection-di ff usion equation:

$$\frac { \partial s } { \partial x } + \frac { \partial s } { \partial t } - D _ { c } \frac { \partial s ^ { 2 } } { \partial x ^ { 2 } } = 0, x \in ( 0, 1 ), \ t \in ( 0, 1 ],$$

subject to a parametric initial condition s ( x , 0) = u GLYPH&lt;16&gt; sin 2 (2 π x ) GLYPH&lt;17&gt; and periodic boundary conditions i.e., s (0 , t ) = s (1 , t ). The di ff usion coe ffi cient is set to Dc = 0 1. In this case, the operator maps the initial condition . s x ( , 0) to the solution s x ( , t ) at the final time. As such, the input function is the initial condition itself, defined as a x ( ) = s x ( , 0) = u (sin 2 (2 π x )).

To construct the training dataset, we generate N 1 = 1000 unique initial conditions by sampling u from a GRF (with input domain defined by sin (2 2 π x )) with length-scale ℓ = 0 5 (Eq. (27)), discretised at . M = 100 spatial points. The PDE is then solved numerically on a spatiotemporal grid ( x × t ) ≡ (100 × 100) using a finite di ff erence method. For each initial condition, the solution s x ( , t ) is evaluated at N 2 = 100 randomly sampled points within the grid, following similarly as the data generation process outlined in Fig. 6. This dataset serves to train the operator.

For evaluation, a separate test set of 10000 initial conditions is employed. Fig. 8 presents a comparison of the model's predicted solution against the reference solution for a sample test initial condition across the entire spatiotemporal domain. A quantitative evaluation of the model's predictive performance across di ff erent values of α was conducted using the NMSE and NLL metrics, summarised in Tables 2 and 3. In contrast to the previous examples where larger α values (i.e., values greater than 1) typically yielded superior performance, this problem exhibits optimal results at α = 0 5 in terms of both NMSE and NLL. While values of . α greater than 1 typically led to subpar performance in this case, the specific value of α = 1 75 can be considered as a reasonable choice for . α values greater than 1.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-37a56e8c-d0ea-483f-ad73-e3d24ba7fc45.jpg)

x

x

Figure 8: Predictive Performance of α -VI DeepONet for the Advection-Di ff usion Equation . The figure presents a comparison of the α -VI DeepONet prediction with the corresponding reference solution for a representative case of a test initial condition. The top left panel displays the reference solution, while the top right panel shows the predicted mean from α -VI DeepONet at α = 0 5. The bottom left panel illustrates the predicted standard deviation, representing . the uncertainty in the prediction. The bottom right panel depicts the absolute error between the predicted mean and the true solution. All plots share the same colour bar, with values indicated as a function of spatial dimension, x , and time, t , represented on the horizontal and vertical axes, respectively.

Table 2: Mean and standard deviation of NMSE values of α -VI DeepONet for di ff erent α values across four problems: antiderivative operator, gravity pendulum, di ff usion-reaction, and advection-di ff usion. The predicted values from deterministic DeepONet have also been included for comparison. The best-performing model (having the lowest NMSE) for each problem is highlighted in bold.

| α          | Antiderivative 10 - 4 ± 10 - 4   | Gravity pendulum 10 - 4 ± 10 - 4   | Di ff usion-reaction 10 - 3 ± 10 - 3   | Advection-di ff usion 10 - 3 ± 10 - 3   |
|------------|----------------------------------|------------------------------------|----------------------------------------|-----------------------------------------|
| 0.25       | 1.072 ± 0.236                    | 1.578 ± 0.489                      | 5.185 ± 1.242                          | 2.131 ± 1.017                           |
| 0.50       | 0.976 ± 0.333                    | 1.467 ± 0.777                      | 7.259 ± 1.510                          | 1.921 ± 0.285                           |
| 0.75       | 1.112 ± 0.488                    | 1.857 ± 0.683                      | 7.133 ± 3.476                          | 3.659 ± 0.432                           |
| 1.00 (KLD) | 1.482 ± 0.670                    | 1.695 ± 0.302                      | 6.436 ± 1.628                          | 3.270 ± 0.530                           |
| 1.25       | 0.659 ± 0.188                    | 2.173 ± 0.190                      | 7.205 ± 0.189                          | 3.090 ± 1.187                           |
| 1.50       | 1.167 ± 0.358                    | 2.056 ± 0.613                      | 6.352 ± 1.782                          | 2.939 ± 0.399                           |
| 1.75       | 1.064 ± 0.314                    | 1.715 ± 0.260                      | 8.279 ± 2.609                          | 2.288 ± 0.747                           |
| 2.00       | 1.137 ± 0.760                    | 1.111 ± 0.366                      | 4.776 ± 1.045                          | 2.591 ± 0.271                           |
| 2.50       | 0.977 ± 0.224                    | 2.306 ± 0.411                      | 5.824 ± 0.567                          | 2.582 ± 0.431                           |
| 3.00       | 1.765 ± 0.152                    | 1.575 ± 0.305                      | 4.463 ± 0.851                          | 2.676 ± 0.617                           |
| 3.50       | 0.950 ± 0.262                    | 1.580 ± 0.706                      | 9.707 ± 1.986                          | 4.734 ± 0.554                           |
| D-DeepONet | 2.930 ± 0.181                    | 2.922 ± 0.321                      | 8.343 ± 1.452                          | 6.950 ± 1.581                           |

Table 3: Mean and standard deviation of NLL values of α -VI DeepONet for di ff erent α values across four problems: antiderivative operator, gravity pendulum, di ff usion-reaction, and advection-di ff usion. The best-performing model for each problem (having the lowest NLL) is highlighted in bold.

| α          | Antiderivative   | Gravity pendulum   | Di ff usion-reaction   | Advection-di ff usion   |
|------------|------------------|--------------------|------------------------|-------------------------|
| 0.25       | -5.156 ± 0.039   | -4.789 ± 0.218     | -3.474 ± 0.265         | -3.824 ± 0.136          |
| 0.50       | -5.230 ± 0.061   | -4.500 ± 0.976     | -3.637 ± 0.275         | -3.897 ± 0.066          |
| 0.75       | -5.185 ± 0.063   | -4.634 ± 0.247     | -3.345 ± 0.422         | -3.624 ± 0.303          |
| 1.00 (KLD) | -5.250 ± 0.061   | -5.015 ± 0.047     | -3.510 ± 0.302         | -3.755 ± 0.309          |
| 1.25       | -5.304 ± 0.121   | -4.668 ± 0.313     | -3.650 ± 0.329         | -3.614 ± 0.249          |
| 1.50       | -5.177 ± 0.080   | -4.594 ± 0.397     | -3.444 ± 0.284         | -3.740 ± 0.220          |
| 1.75       | -5.199 ± 0.090   | -4.829 ± 0.397     | -3.656 ± 0.390         | -3.805 ± 0.111          |
| 2.00       | -5.149 ± 0.218   | -4.928 ± 0.163     | -3.676 ± 0.282         | -3.739 ± 0.248          |
| 2.50       | -5.243 ± 0.062   | -4.438 ± 0.803     | -3.725 ± 0.178         | -3.766 ± 0.194          |
| 3.00       | -5.238 ± 0.074   | -4.550 ± 0.469     | -3.658 ± 0.289         | -3.740 ± 0.214          |
| 3.50       | -5.201 ± 0.046   | -4.517 ± 0.922     | -3.563 ± 0.324         | -3.450 ± 0.348          |

To assess the model's robustness and generalisation capabilities beyond training data distribution, we conducted out-of-distribution (OOD) testing. Two distinct OOD datasets were generated, each comprising 100 examples.

The first OOD dataset was constructed by altering the kernel length-scale of the GRF (Eq. (27)) used to generate initial conditions. While the training data utilised a length-scale of ℓ = 0 5, the OOD dataset employed a reduced length-scale . of ℓ = 0 2 within the RBF kernel. This modification introduces increased fluctuations in the generated initial conditions . compared to the training distribution.

The second OOD dataset was generated using a fundamentally di ff erent kernel, the rational quadratic kernel, which is defined as:

$$\kappa ( x _ { 1 }, x _ { 2 } ) = \exp \left ( 1 + \frac { \| x _ { 1 } - x _ { 2 } \| ^ { 2 } } { 2 \rho \ell ^ { 2 } } \right ) ^ { - \rho }, \Big. _ { - }$$

This kernel is parameterised by an additional scale mixture parameter, ρ . For our experiments, we set ρ = 1 0 and . ℓ = 0 5 to generate 100 OOD test cases. This represents 100 di . ff erent initial conditions for the advection-di ff usion example.

The predictive performance of models trained with three di ff erent α values: 0.5, 1.0 (standard KLD-VI), and 1.75, was evaluated on these OOD test datasets. Similar to the in-distribution analysis, the NMSE and NLL metrics were computed. The average results, summarised in Table 4, indicate that the models exhibit reasonable generalisation capabilities on OOD data. Consistent with the in-distribution findings, α = 0 5 outperformed both the standard KLD-VI .

( α = 1 00) and the . α = 1 75 model in terms of both NMSE and NLL for both OOD datasets. This improvement over . the standard KLD-VI model translated to upto a 7% reduction in NMSE and an 8.7% reduction in NLL for the tested OOD scenarios. Interestingly, the performance on the rational quadratic kernel samples was superior to that on the RBF kernel samples.

Table 4: Out-of-distribution average test NMSE and NLL values for the advection-di ff usion example. The bestperforming model for each metric is highlighted in bold.

|           | RBF with ℓ = 0 . 2   | RBF with ℓ = 0 . 2   | Rational quadratic   | Rational quadratic   |
|-----------|----------------------|----------------------|----------------------|----------------------|
| α         | NMSE 10 - 3 ± 10 - 3 | NLL 10 0 ± 10 0      | NMSE 10 - 3 ± 10 - 3 | NLL 10 0 ± 10 0      |
| 0.50      | 2.062 ± 0.269        | -3.856 ± 0.069       | 0.197 ± 0.032        | -4.352 ± 0.068       |
| 1.00 (VI) | 2.207 ± 0.640        | -3.547 ± 0.414       | 0.212 ± 0.040        | -4.093 ± 0.358       |
| 1.75      | 2.585 ± 0.807        | -3.752 ± 0.095       | 0.210 ± 0.056        | -4.266 ± 0.154       |

## 5 Discussion

The results from the numerical investigation demonstrate the superior performance of the proposed α -VI DeepONet compared to both the deterministic D-DeepONet and the standard KLD-VI DeepONet, as evidenced by the consistently lower NMSE and NLL values across all four problems (Figures 9 and 10). This underscores the e ff ectiveness of the α -VI DeepONet framework in capturing complex input-output relationships and quantifying associated uncertainties.

Figure 9: Comparison of NMSE values for Di ff erent DeepONet variants . The figure presents a comparative analysis of NMSE values for D-DeepONet, standard KLD-VI DeepONet at α = 1 (KLD-VI), and α -VI DeepONet (at the corresponding optimal values of α ) across four numerical problems: antiderivative operator, gravity pendulum, di ff usion-reaction, and advection-di ff usion. Each bar represents the mean NMSE computed over ten independent runs, with error bars indicating the corresponding standard deviation. Lower NMSE values signify better mean predictive accuracy.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-4ee90aa3-dc95-4cd2-8c0d-ace9ee58c9c7.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-d8c9154c-278f-4809-876f-3c2c12269474.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-b617730d-6e0f-4b89-8f96-1b2649c0e091.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-6f46087a-7339-4817-9f8c-179578874517.jpg)

Figure 10: Comparison of NLL values for Di ff erent DeepONet variants . The figure presents a comparative analysis of NLL values for D-DeepONet, standard KLD-VI DeepONet at α = 1 (KLD-VI), and α -VI DeepONet (at the corresponding optimal values of α ) across four numerical problems: antiderivative operator, gravity pendulum, di ff usion-reaction, and advection-di ff usion. Each bar represents the mean NLL computed over ten independent runs, with error bars indicating the corresponding standard deviation. Lower NLL values signify better distributional fit.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-5b79a7c3-7d86-456b-8f68-ae758e0699bb.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-381fc294-ccac-4705-869c-6640419bf9bf.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-f1a47bd4-6eda-4d1c-93f1-791c7e9e9873.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00681v1-figure-ac0c9efb-367d-4fc5-9deb-5ccf5fc613d4.jpg)

A key finding is the influence of the hyperparameter α on model performance, which varies across di ff erent problems. While a universal optimal α value is desirable, our results indicate that cross-validation is necessary to determine the most suitable setting for each specific problem. Notably, the optimal α values for all four problems were mostly di ff erent from the standard KLD-VI case ( α = 1), highlighting the limitations of the KL divergence under prior misspecification.

The choice of α is greatly influenced by the underlying structure of the true posterior distribution. Mode-seeking methods may struggle in scenarios with multi-modal posteriors, as they are susceptible to local optima. Conversely, mass-covering methods may be less e ff ective when the true posterior exhibits widely separated modes. Therefore, the optimal α value is inherently problem-dependent. For problems like the antiderivative, gravity pendulum, and di ff usion-reaction problems, larger α values (greater than 1) yielded superior performance, suggesting a preference for mode-seeking behaviour. Conversely, the advection-di ff usion problem favoured smaller α values (less than 1), indicating a need for mass-covering behaviour. The selection of α should consider both NMSE and NLL, as these metrics provide complementary insights - NMSE looks only at the mean prediction, whereas NLL considers distributional fit - into model performance.

While the α -VI DeepONet o ff ers significant flexibility, it is important to acknowledge the associated computational overhead. Compared to the deterministic DeepONet, the α -VI DeepONet incurs notably higher training times. This increase arises from two primary factors: the approximation of the Rényi α -divergence, which necessitates Monte Carlo sampling, and the doubling of trainable parameters due to the estimation of both mean and variance for each weight. However, compared to a standard KLD-VI DeepONet, the training times are approximately 1.1x longer, which is reasonable for the gain in flexibility.

## 6 Conclusions

In this work, we have introduced a novel, prior-robust, uncertainty-aware DeepONet framework grounded in generalised variational inference (GVI). This framework enables the learning of complex nonlinear operators while providing robust uncertainty quantification. Unlike previous approaches, our framework adopts an optimisation-centric perspective on Bayesian modeling by minimising the GVI objective. By replacing the Kullback-Liebler divergence with Rényi's α divergence, we enhance the model's robustness to prior misspecification and achieve superior predictive performance compared to the standard KLD-VI approach. The hyperparameter α o ff ers flexibility in controlling the degree of robustness, allowing adaptation to di ff erent problem characteristics.

Our numerical investigation across four benchmark problems consistently demonstrates the advantages of the α -VI DeepONet over both deterministic and standard variational methods in terms of NMSE and NLL. The optimal α value varies across problems, highlighting the importance of tuning the hyperparameter. While the α -VI DeepONet o ff ers significant improvements, it is not without limitations. The introduction of the α hyperparameter increases model complexity and requires additional computational resources due to the approximation of Rényi's α divergence. Furthermore, the mean-field assumption underlying the variational posterior may not be adequate in all cases. Future research directions include exploring alternative variational families (e.g., normalising flows) to address the limitations

of the mean-field approximation and developing more e ffi cient approximation techniques for Rényi's α -divergence. Additionally, investigating the application of the proposed framework to a wider range of complex systems is a promising avenue for future work.

## 7 Acknowledgements

R. Nayek acknowledges the financial support received from the Science and Engineering Research Board (SERB) via Grant number SRG 2022 001410, and the Ministry of Ports, Shipping and Waterways. / /

## References

- [1] R. Courant, D. Hilbert, Methods of mathematical physics: partial di ff erential equations, John Wiley &amp; Sons, 2008.
- [2] T. J. Hughes, The finite element method: linear static and dynamic finite element analysis, Courier Corporation, 2012.
- [3] N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, A. Anandkumar, Neural operator: Learning maps between function spaces with applications to PDEs, Journal of Machine Learning Research 24 (2023) 1-97.
- [4] L. Lu, P. Jin, G. Pang, Z. Zhang, G. E. Karniadakis, Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators, Nature Machine Intelligence 3 (2021) 218-229.
- [5] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, A. Anandkumar, Fourier neural operator for parametric partial di ff erential equations, arXiv preprint arXiv:2010.08895 (2020).
- [6] S. Goswami, A. Bora, Y. Yu, G. E. Karniadakis, Physics-informed deep neural operator networks, in: Machine Learning in Modeling and Simulation: Methods and Applications, Springer, 2023, pp. 219-254.
- [7] S. Garg, H. Gupta, S. Chakraborty, Assessment of DeepONet for time dependent reliability analysis of dynamical systems subjected to stochastic loading, Engineering Structures 270 (2022) 114811.
- [8] L. Lu, R. Pestourie, S. G. Johnson, G. Romano, Multifidelity deep neural operators for e ffi cient learning of partial di ff erential equations with application to fast inverse design of nanoscale heat transport, Physical Review Research 4 (2022) 023210.
- [9] S. Goswami, M. Yin, Y. Yu, G. E. Karniadakis, A physics-informed variational DeepONet for predicting crack path in quasi-brittle materials, Computer Methods in Applied Mechanics and Engineering 391 (2022) 114587.
- [10] S. De, M. Reynolds, M. Hassanaly, R. N. King, A. Doostan, Bi-fidelity modeling of uncertain and partially unknown systems using DeepONets, arXiv preprint arXiv:2204.00997 (2022).
- [11] S. Cai, Z. Wang, L. Lu, T. A. Zaki, G. E. Karniadakis, DeepM&amp;Mnet: Inferring the electroconvection multiphysics fields based on operator approximation by neural networks, Journal of Computational Physics 436 (2021) 110296.
- [12] J. He, S. Koric, S. Kushwaha, J. Park, D. Abueidda, I. Jasiuk, Novel DeepONet architecture to predict stresses in elastoplastic structures with variable complex geometries and loads, Computer Methods in Applied Mechanics and Engineering 415 (2023) 116277.
- [13] J. Pathak, S. Subramanian, P. Harrington, S. Raja, A. Chattopadhyay, M. Mardani, T. Kurth, D. Hall, Z. Li, K. Azizzadenesheli, et al., Fourcastnet: A global data-driven high-resolution weather model using adaptive Fourier neural operators, arXiv preprint arXiv:2202.11214 (2022).
- [14] N. Demo, M. Tezzele, G. Rozza, A DeepONet multi-fidelity approach for residual learning in reduced order modeling, Advanced Modeling and Simulation in Engineering Sciences 10 (2023) 12.
- [15] I. M. S. Leite, J. D. M. Yamim, L. G. da Fonseca, The DeepONet for finance: an approach to calibrate the Heston model, in: EPIA Conference on Artificial Intelligence, Springer, 2021, pp. 351-362.
- [16] R. Ranade, K. Gitushi, T. Echekki, Generalized joint probability density function formulation inturbulent combustion using DeepONet, arXiv preprint arXiv:2104.01996 (2021).
- [17] C. Xu, B. T. Cao, Y. Yuan, G. Meschke, A multi-fidelity deep operator network (DeepONet) for fusing simulation and monitoring data: Application to real-time settlement prediction during tunnel construction, Engineering Applications of Artificial Intelligence 133 (2024) 108156.
- [18] L. Lu, X. Meng, S. Cai, Z. Mao, S. Goswami, Z. Zhang, G. E. Karniadakis, A comprehensive and fair comparison of two neural operators (with practical extensions) based on FAIR data, Computer Methods in Applied Mechanics and Engineering 393 (2022) 114778.

- [19] G. Lin, C. Moya, Z. Zhang, Accelerated replica exchange stochastic gradient Langevin di ff usion enhanced Bayesian DeepONet for solving noisy parametric PDEs, arXiv preprint arXiv:2111.02484 (2021).
- [20] M. Welling, Y. W. Teh, Bayesian learning via stochastic gradient Langevin dynamics, in: Proceedings of the 28th international conference on machine learning (ICML-11), Citeseer, 2011, pp. 681-688.
- [21] Y. Yang, G. Kissas, P. Perdikaris, Scalable uncertainty quantification for deep operator networks using randomized priors, Computer Methods in Applied Mechanics and Engineering 399 (2022) 115399.
- [22] D. J. MacKay, A practical Bayesian framework for backpropagation networks, Neural computation 4 (1992) 448-472.
- [23] R. M. Neal, Bayesian learning for neural networks, volume 118, Springer Science &amp; Business Media, 2012.
- [24] S. Garg, S. Chakraborty, VB-DeepONet: A Bayesian operator learning framework for uncertainty quantification, Engineering Applications of Artificial Intelligence 118 (2023) 105685.
- [25] Y. Li, R. E. Turner, Rényi divergence variational inference, Advances in Neural Information Processing Systems 29 (2016).
- [26] J. Knoblauch, J. Jewson, T. Damoulas, An Optimization-centric view on Bayes' rule: Reviewing and Generalizing Variational Inference, Journal of Machine Learning Research 23 (2022) 1-109. URL: http://jmlr.org/ papers/v23/19-1047.html .
- [27] F. Wenzel, K. Roth, B. S. Veeling, J. Swi˛tkowski, L. Tran, S. Mandt, J. Snoek, T. Salimans, R. Jenatton, ´ a S. Nowozin, How good is the Bayes posterior in deep neural networks really?, arXiv preprint arXiv:2002.02405 (2020).
- [28] R. E. Turner, M. Sahani, Two problems with variational expectation maximisation for time series models, Cambridge University Press, 2011, p. 104-124.
- [29] A. Rényi, On measures of entropy and information, in: Proceedings of the fourth Berkeley symposium on Mathematical Statistics and Probability, Volume 1: Contributions to the theory of Statistics, volume 4, University of California Press, 1961, pp. 547-562.
- [30] T. Chen, H. Chen, Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems, IEEE transactions on neural networks 6 (1995) 911-917.
- [31] S. Ruder, An overview of gradient descent optimization algorithms, arXiv preprint arXiv:1609.04747 (2016).
- [32] A. Choromanska, M. Hena ff , M. Mathieu, G. B. Arous, Y. LeCun, The loss surfaces of multilayer networks, in: Artificial Intelligence and Statistics, PMLR, 2015, pp. 192-204.
- [33] W. R. Gilks, S. Richardson, D. Spiegelhalter, Markov Chain Monte Carlo in practice, CRC press, 1995.
- [34] P. Izmailov, S. Vikram, M. D. Ho ff man, A. G. G. Wilson, What are Bayesian neural network posteriors really like?, in: International conference on machine learning, PMLR, 2021, pp. 4629-4640.
- [35] D. M. Blei, A. Kucukelbir, J. D. McAuli ff e, Variational inference: A review for statisticians, Journal of the American Statistical Association 112 (2017) 859-877.
- [36] J. M. Hernández-Lobato, R. Adams, Probabilistic backpropagation for scalable learning of Bayesian neural networks, in: International Conference on Machine Learning, PMLR, 2015, pp. 1861-1869.
- [37] J. Knoblauch, J. E. Jewson, T. Damoulas, Doubly robust Bayesian inference for non-stationary streaming data with β -divergences, Advances in Neural Information Processing Systems 31 (2018).
- [38] A. Kucukelbir, D. Tran, R. Ranganath, A. Gelman, D. M. Blei, Automatic di ff erentiation variational inference, Journal of Machine Learning Research 18 (2017) 1-45.
- [39] A. Cichocki, S.-i. Amari, Families of alpha-beta-and gamma-divergences: Flexible and robust measures of similarities, Entropy 12 (2010) 1532-1568.
- [40] C. Blundell, J. Cornebise, K. Kavukcuoglu, D. Wierstra, Weight uncertainty in neural network, in: International Conference on Machine Learning, PMLR, 2015, pp. 1613-1622.
- [41] D. P. Kingma, T. Salimans, M. Welling, Variational dropout and the local reparameterization trick, Advances in Neural Information Processing Systems 28 (2015).
- [42] D. P. Kingma, J. Ba, Adam: A method for stochastic optimization, arXiv preprint arXiv:1412.6980 (2014).
- [43] J. V. Dillon, I. Langmore, D. Tran, E. Brevdo, S. Vasudevan, D. Moore, B. Patton, A. Alemi, M. Ho ff man, R. A. Saurous, Tensorflow distributions, arXiv preprint arXiv:1711.10604 (2017).
Soban Nasir Lone Subhayan De Rajdip Nayek
Alpha-VI DeepONet: A prior-robust variational Bayesian approach for enhancing DeepONets with uncertainty quantification
We introduce a novel deep operator network (DeepONet) framework that
incorporates generalised variational inference (GVI) using R\'enyi's
$\alpha$-divergence to learn complex operators while quantifying uncertainty.
By incorporating Bayesian neural networks as the building blocks for the branch
and trunk networks, our framework endows DeepONet with uncertainty
quantification. The use of R\'enyi's $\alpha$-divergence, instead of the
Kullback-Leibler divergence (KLD), commonly used in standard variational
inference, mitigates issues related to prior misspecification that are
prevalent in Variational Bayesian DeepONets. This approach offers enhanced
flexibility and robustness. We demonstrate that modifying the variational
objective function yields superior results in terms of minimising the mean
squared error and improving the negative log-likelihood on the test set. Our
framework's efficacy is validated across various mechanical systems, where it
outperforms both deterministic and standard KLD-based VI DeepONets in
predictive accuracy and uncertainty quantification. The hyperparameter
$\alpha$, which controls the degree of robustness, can be tuned to optimise
performance for specific problems. We apply this approach to a range of
mechanics problems, including gravity pendulum, advection-diffusion, and
diffusion-reaction systems. Our findings underscore the potential of
$\alpha$-VI DeepONet to advance the field of data-driven operator learning and
its applications in engineering and scientific domains.