Uncertainty Quantiﬁcation in Neural Differential Equations
Abstract
Uncertainty quantiﬁcation (UQ) helps to make trustworthy predictions based on collected observations and uncertain domain knowledge.
With increased usage of deep learning in various applications, the need for efﬁcient UQ methods that can make deep models more reliable has increased as well.
Among applications that can beneﬁt from effective handling of uncertainty are the deep learning based differential equation (DE) solvers.
We adapt several state-of-the-art UQ methods to get the predictive uncertainty for DE solutions and show the results on four different DE types
Introduction
Driven by the growing popularity of deep learning, several areas of research have obtained state-ofthe-art performances with deep neural networks (NNs).
Among other applications, deep NNs have been applied for solving differential equations (DEs) [10] — a fundamental tool for mathematical modeling in engineering, ﬁnance, and the natural sciences.
Deep learning based solutions of DEs have recently appeared in, e.g., [20], [18], [19], [22], [9], [11], [17], [5], [21], [8].
Typically, the NN itself approximates the solution of a DE.
Thanks to that, parallelization is natural and, in contrast to classical numerical methods, the solution at any time can be computed without the burden of having to compute all previous time steps.
Furthermore, NNs are continuous and differentiable.
Until recently, the focus of deep learning was on achieving better accuracy in the NN predictions, but now it is increasingly being shifted to measuring the prediction’s uncertainty, especially if the task at hand is safety critical.
Uncertainty quantiﬁcation (UQ) has been considered for deep models in computer vision, medical image analysis, bionformatics, etc [1].
Likewise, UQ is important for deep models that solve DEs.
The uncertainty here stems from the fact that we cannot train a NN on an inﬁnite time and/or space domain.
Therefore, we seek to estimate the solution’s uncertainty in the regions where the model was not trained.
Moreover, another source of uncertainty comes from the model’s limitations such as its architecture.
To the best of our knowledge, this is the ﬁrst work to discuss UQ for deep models that solve DEs.
Contrary to common deep learning setup, we solve DEs without any observed data, relying only on the samples of time and/or space and on the mathematical statement that relates functions and their derivatives.
This makes the application of existing UQ methods not so straightforward.
In this paper, we make the following contributions:
We propose an adaptation of the four state-of-the-art UQ methods in deep learning — Bayes By Backprop [3], Flipout [24], Neural Linear Model [23, 15], and Deep Evidential Regression [2, 13] — to the case of solving DEs.
We test the above-mentioned methods on four different DE types: linear ordinary DE (ODE), non-linear ODE, system of non-linear ODEs, and partial DE (PDE).
Preliminaries 2.1 Solving differential equations with neural networks
A DE can be expressed as L u − f = 0 , where L is the differential operator, u ( x ) is the solution that we wish to ﬁnd on some (possibly multidimensional) domain x , and f is a known forcing function.
We denote the NN approximation of the true solution by u N .
To solve the DE, we minimize the square loss of the residual function R ( u N ) :
= L u N − f , i.e., the optimization objective is [Equation] where w are the NN parameters.
It is also necessary to inform the NN about any initial and/or boundary conditions, u c = u ( x c ) .
One can achieve that in a straightforward way by adding a penalizing term to the loss function. However, the exact satisfaction of initial/boundary conditions is not possible in this case, causing problems in case of high sensitivity to initial conditions, and also yielding unnecessary local uncertainty at x c .
Therefore, we employ an alternative approach and consider a transformation of u N which enforces the initial/boundary conditions and satisﬁes them by construction.
E.g., in one-dimensional case, given an initial condition u 0 = u ( t 0 ) , we consider a transformation ˜
u N ( t ) = u 0 + (1 − e − ( t − t 0 ) ) u N ( t ) . In general, the transformation has the form ˜ u N ( x ) =
A ( x , x c , u c ) + B ( x , x c ) u N ( x ) .
Hereinafter, ˜ u N will denote the enforced solution. Accordingly, we replace u N by ˜ u N in the optimization objective (1).
Besides the advantage of satisfying the initial/boundary conditions exactly, the latter approach can also reduce the effort required during training [12].
2.2 Uncertainty quantiﬁcation under Bayesian framework
While classical learning considers deterministic model parameters θ , the Bayesian framework introduces uncertainty by considering a posterior distribution over the model parameters , p ( θ |D ) , obtained after observing some data D .
The posterior distribution is given by Bayes’ theorem, p ( θ |D )
= p ( D| θ ) · p ( θ ) / p ( D ) , where p ( D| θ ) is the likelihood, p ( θ ) is the prior distribution over the parameters, and p ( D ) is the evidence.
The predictions y at a new test point x are given by the posterior predictive distribution,
[Equation]
For probabilistic deep models, there are two main strategies of estimating (2).
Inference through the posterior distribution of model parameters .
As stated in (2), the posterior predictive is obtained by averaging over the posterior uncertainty in the model parameters.
Thus, we can start with estimating the posterior distribution of the NN weights.
In this case, well-suited is Bayesian NN [14] which places a prior distribution on all the weights (and biases) w .
Since an analytical solution for the posterior is intractable for Bayesian NNs, we have to use numerical approximation methods such as MCMC or variational methods.
Despite the need for sampling in both cases, variational methods are computationally less expensive for high-dimensional parameter spaces and also provide an analytical approximation.
Bayes By Backprop ( BBB ) is a variational, backpropagation-compatible method for training a Bayesian
NN.
Its optimization objective seeks to minimize the Kullback-Leibler divergence between the true posterior and the variational posterior which is re-parametrized as N ( µ, σ = log(1 + exp( ρ ))) to allow for backpropagation.
At each optimization step, weights w = µ + σ ◦ (cid:15) , where (cid:15) ∼ N (0 , I ) and ◦ is pointwise multiplication, are obtained by sampling from the variational posterior.
BBB is followed by Flipout which adds a pseudo-independent perturbation to the weights at each training point x n in the mini-batch, namely, w n = µ + ( σ ◦ (cid:15) ) R n , where R n is the random sign matrix.
Intuitively, the weights get ﬂipped symmetrically around the mean with probability 0 .
Neural Linear Model ( NLM ) is an alternative to Bayesian NN.
It places a prior distribution only on the last layer’s weights, and learns point estimates for the remaining layers.
One can interpret the output of these layers as a basis deﬁned by the feature embedding of the data.
The last layer of NLM performs Bayesian linear regression on this feature basis.
NLM provides tractable inference under the
Gaussian assumption on likelihood; we get analytical solution for the posterior distribution.
Inference through the higher-order evidential distribution .
It is also possible to infer parameters of the posterior predictive directly, using Bayesian hierarchical modeling [6, 7].
In Deep Evidential
Regression ( DER ), the higher-order, evidential prior is placed over the Gaussian likelihood function.
Choosing Normal Inverse-Gamma (NIG) prior yields an analytical solution for the model evidence which is maximized by the optimization objective with respect to the NIG hyperparameters.
DER also proposes an evidence regularizer which minimizes evidence on incorrect predictions.
The posterior predictive mean and variance are computed analytically using the learned hyperparameters.
Uncertainty quantiﬁcation in neural differential equations
Instead of learning a deterministic solution ˜ u N , we now aim to learn a probabilistic solution u θ , characterized by a posterior predictive distribution.
We estimate it using some probabilistic model g θ parametrized by θ .
3.1 Proposed approach
UQ methods described in Section 2.2 rely upon the assumption that the likelihood function is Gaussian, centered at the model’s prediction and evaluated at the observed data points.
Namely,
DER and NLM use it to derive the analytical form of a loss function and a posterior distribution, respectively.
BBB and Flipout can in principle use any likelihood in the loss function, but it has to be of known analytical form.
In case of DEs, a natural way of computing likelihood is to evaluate it at the residual R on the training domain x t , which can be seen as a counterpart to the observed data points in the classical setting.
However, we are left with an open problem of choosing the underlying distribution for the likelihood function.
It makes sense to assume that the probability density is high enough for values close to zero, but no further assumptions immediately follow.
E.g., it may happen that the limitations of NN architecture do not allow for the perfect ﬁt, i.e., the distribution of residuals is not centered around zero.
To circumvent this problem, we propose an alternative way of computing likelihood. In this ﬁrst work on UQ for neural DE solvers, we will focus on comparing predictions outside of the training domain given by different UQ methods, leaving the detailed treatment of the model ﬁt and its associated uncertainty for future work. Although in Bayesian framework this uncertainty also affects the uncertainty outside of the training domain, we hypothesize that even a simpliﬁed treatment, i.e., without using residuals’ distribution, gives reasonable uncertainty estimates. We propose a two-stage training procedure:
1) We ﬁrst train a classical NN on the training domain x t to ﬁnd a deterministic solution ˜ u N , 2)
We use ˜ u N as observed data for our probabilistic model g θ and deﬁne the likelihood using a Gaussian assumption, p (˜ u N | θ ) = (cid:81) x t N (˜ u N ; u θ , ε ) .
Now the optimization objective will be trying to align the probabilistic model with the given reference ˜
u N rather than trying to minimize the residuals at all costs. We note that despite interpreting ˜ u N in stage two as observed data rather than a function that solves DE, its associated variance ε is not of aleatoric nature (i.e. irreducible variance that comes from the noise inherent to the data), as it would be in the classical regression problem.
It can be still interpreted as a source of epistemic (reducible) uncertainty coming from the NN model limitations.
Here, we consider a simpliﬁed treatment of ε .
In case of BBB , Flipout , and NLM , we pre-deﬁne ε with some small number.
In case of DER , we learn ε along with posterior predictive distribution, but the result is not particularly useful since DER does not have direct access to residuals during learning.
Eventually, the probabilistic model allows us to ﬁnd the posterior predictive distribution of u θ .
In case of BBB , Flipout , and NLM , we have ˜ g
θ ( x t , ˜ u N ) =
u θ , i.e., the model outputs a single instance u θ .
For BBB and Flipout , the posterior predictive distribution p ( u θ | x t , ˜ u N ) is computed as an approximation of integral (2) using sampling; for NLM , an analytical form is available.
In case of
DER , we have g θ ( x t , ˜ u N ) =
( γ, ν, α, β ) , where ( γ, ν, α, β ) are the NIG hyperparameters.
The mean of the posterior predictive distribution is equal to ˜ γ and the variance is computed using the remaining hyperparameters.
We note that the predictive uncertainty also requires the initial and/or boundary condition enforcement; this way we are able to eliminate unnecessary uncertainties at x c .
Main drawbacks of the current approach are the double computational burden and the not so useful uncertainty for NN approximation of the true solution in the traning region; both of them are subject to further improvement.
Experiments and discussion
We corroborate our theory with experimental results on four equations: 1.
Linear ODE for squared exponential, dudt = − 2 tu ; 2.
Non-linear ODE for Dufﬁng-type oscillator, ¨ u + ω 2 u
+ (cid:15)u 3 = 0 ; 3.
Lotka-Volterra equations (system of non-linear ODEs), ˙ u
= αu − βuv
∧ ˙ v = − δu + γuv ; 4.
Burgers’ equation (non-linear PDE), ∂u∂t + u ∂u∂x = ν
∂ 2 u ∂x 2 .
Our implementation is based on a DE solver provided by neurodiffeq [4], a Python package built with
PyTorch [16].
Since we are considering relatively simple DEs, we use networks with one to three fully-connected hidden layers.
For the prior distribution of the weights in BBB , Flipout , and NLM , we use ﬂat Gaussian priors with mean zero.
In BBB and Flipout , we estimate the posterior predictive from 1000 samples.
In DER , there is no need for choosing a weight prior, but an appropriate regularization parameter has to be chosen instead.
Here, we tune the regularization parameter manually.
We demonstrate the UQ results for ODEs and for PDE in Figure 1 and Figure 2, accordingly.
For all
ODEs, the deterministic solution is able to approximate the true solution well; we incorporate this fact in our Bayesian inference by choosing small ε which yields that there is almost no uncertainty in the training domain.
We observe that the epistemic uncertainty away from the training domain is high enough for all methods, which is our main desired result in this paper.
For Burgers’ equation, however, we see that the NN is not able to learn the true solution, and our probabilistic model is underestimating the epistemic uncertainty in the training domain and outside of it.
In this case, either a better deterministic model or a better UQ methodology is needed.
Nevertheless, even for a non-perfect ﬁt, the uncertainty starts inﬂating outside of the training domain which proves our initial hypothesis.
We have witnessed comparable performance in sampling-dependent ( BBB , Flipout ) and sampling-free ( NLM , DER ) methods.
Given the computational expense of sampling during Bayesian NN training, the latter two methods could be preferable in the case of complex DEs on a multidimentional domain.
We believe that further enhancement in terms of diversifying experiments (e.g., considering more complex high-dimensional DEs) and developing theory (e.g., calibrating ε with residuals at each optimization step) will help the deep learning based DE solutions to outperform classical ones and lead to their increased presence in applications.