Probabilistic neural operators for functional uncertainty quantification
Authors: Christopher Bülte
Philipp Scholl
Gitta Kutyniok
Venue: Accepted by TMLR ()

Abstract: Neural operators aim to approximate the solution operator of a system of differential equations purely from data. They have shown immense success in modeling complex dynamical systems across various domains. However, the occurrence of uncertainties inherent in both model and data has so far rarely been taken into account\textemdash{}a critical limitation in complex, chaotic systems such as weather forecasting. In this paper, we introduce the probabilistic neural operator (PNO), a framework for learning probability distributions over the output function space of neural operators. PNO extends neural operators with generative modeling based on strictly proper scoring rules, integrating uncertainty information directly into the training process. We provide a theoretical justification for the approach and demonstrate improved performance in quantifying uncertainty across different domains and with respect to different baselines. Furthermore, PNO requires minimal adjustment to existing architectures, shows improved performance for most probabilistic prediction tasks, and leads to well-calibrated predictive distributions and adequate uncertainty representations even for long dynamical trajectories. Implementing our approach into large-scale models for physical applications can lead to improvements in corresponding uncertainty quantification and extreme event identification, ultimately leading to a deeper understanding of the prediction of such surrogate models.

Published in Transactions on Machine Learning Research (03/2025)
Probabilistic neural operators for functional uncertainty quan- tification
Christopher Bülte
buelte @matth.mmam
Ludwigs-Masse Mozzimitions---- Universität München Munich Center for Machine Learning (MCML) Munich, Germany
Philipp Scholl
Scholl@mmath.mmee
Muddigg-Mozzimiteee Universität München Munich Center for Machine Learning (MCML) Munich, Germany
Gitta Kutyniok
kutyniok Almath Imai
Masiiigg-Maimmmlianss---- Manaaiimmiian---nnnsssstttt Universität München
University of Tromse
Aerospace Center
Munich Center for Machine Learving (MCML)
Munich, Germany
Reviewed on OpenReview: https: // openreview. net/ forum? id= gangoPXSRu
Abstract
Neural operators aim to approximate the solution operator of system of differential equations purely from data. They have shown immense success in modeling complex dynamical systems across various domains. However, the occurrence of uncertainties inherent in both model and data has so far rarely been taken into account-a a critical limitation in complex, chaotic systems such as weather forecasting. In this paper, we introduce the probabilistic neural operator (PNO), framework for learning probability distributions over the output function space of neural operators. PNO extends neural operators with generative modeling based on strictly proper scoring rules, integrating uncertainty information directly into the training process. We provide thcoretical justification for the approach and demonstrate improved performance in quantifying uncertainty across different domains and with respect to different baselines. Furthermore, PNO requires minimal adjustment to existing architectures, shows improved performance for most probabilistic prediction tasks, and leads to well-ca well-calibrated predictive distributions and adequate uncertainty representations even for long dynamical tratocroees. Implementing our approach into large-scale models for physical applicat ions can lead to improvements in corresponding uncertainty quantification and extreme event identification, ultimately lcading to deeper understanding of the prediction of such surrogate models.
1 Introduction
Many complex phenomena in the sciences are described via parametric partial differential equations (PDEs), making their study a crucial research topic for systems in biology, physics, or engineering sciences. Since traditional solvers, such as finite element methods are based on a space discretization there exists a tradeoff between solution accuracy and the computational complexity imposed by the discrctization resolution. A recent breakthrough has been achieved by neural operators (Kovachki al., 2023; Lu al., 2021), which are class of neural networks that have the capacity to approximate any continuous operator, including any

---

Published in Transactions on Machine Learning Research (03/2025)
solution operator of given family of parametric PDEs. Neural operators, and especially Fourier neural operators (Ti et al., 2021), are now commonly used in practice and have been applied to problems such as weather forecasting (Kurth et al., 2023), fluid dynamics (Renn et al., 2023), seismic imaging (Li et al., 2023), carbon capture (Wen et al., 2023), or plasma modeling in muclear fusion (Gopakumar et al., 2024). However, in application scenarios, the entire operalor learning process is usually corrupled by different sources of noise. In addition, the model itself can be misspecified or-as in the case of dynamical systems long--s-onnntenm terg- predictions might be too chaotic to be predictable in reliable manner. For these cases, it becomes crucial to have a precise understanding of the uncertainty of the corresponding model predictions.
A common approach to uncertainty quantification for standard neural networks is the Laplace approximation (Daxberger et al., 2021), which has been extended to Fourier neural operators by Weber et al. (2024) and Magnani et al. (2024). While this approach is efficient to implement and does not require additional training, it has notable limitations. Specifically, it does not incorporate uncertainty during network training, and its assumption of a Gaussian predictive distribution may be overly restrictive. Guo al. (2023) introduce an information bottleneck structure that utilizes decoder with predictive Gaussian distribution. While innovative, this method shares the limitations of the Gaussian assumption.
Our work addresses these challenges by proposing a theoretically grounded approach to incorporate predictive uncertainty during training in the infiniteddimensional setting of neural operators. By leveraging strictly proper scoring rules, our method enables better-calibrated uncertainty estimates while supporting more flexible predictive distributions.
1.1 Our contributions
In this article, we introduce the probabilistic neural operator (PNO), which extends the neural operator framework to a probabilistic setting, where the output is in the form of (functional) probability distribution. Our novel approach is based on recent advances in training generative neural networks via proper scoring rules (Pacchiardi et al., 2024). PNO is designed to generate samples from an (empirical) probability distribution and is trained on an objective function that fits the predictive distribution to the corresponding observations, thereby integrating uncertainty directly into the training process. We build upon prior research by introducing a theoretical framework for proper scoring rules in separable Hilbert spaces. Specifically, we prove that the energy score is a strictly proper scoring rule in these spaces and, therefore, leads to desirable properties as a loss function. The resulting predictions the network are well-calibrated and can be-due to their sample- based character-applies for different downstream tasks, including extreme event analysis and estimation of dependencies. We evaluate PNO on PDE datasets and real-world data lask, namely the prediction of surface temperature across Europe. PNO requires minimal adjustment to existing architectures and leads to well-calibrated predictive distributions and adequate uncertainty representations even for long dynamical trajectories. To summarize, we make the following key contributions:
We introduce the probabilistic neural operator (PNO) as an extension of the neural operator to the probabilistic setting based on proper scoring rules.
We prove that. the eneryy score a strictly proper scoring rule in separable Hilbert spaces and, therefore, a suitable loss function for the PNO.
In extensive experiments including 2- and 3-dimensional PDEs, real-world weather data, and different model architectures we show that the PNO setup improves the calibration error while maintaining the accuracy of competing approaches.
1.2 Related work
In order to assess uncertainties in machine learning models, a variely of approaches are available; for a detailed overview compare Abdar et al. (2021). The most common way to perform uncertainty quantification is under the Bayesian paradigm. By utilizing Monte Carlo dropout during the inference phase of a neural network, Gal & Ghahramani (2016) show that the corresponding predictive distribution Bayesian approximation of deep Gaussian process. While the approach requires no additional training, the approximation does not.
2

---

Published in Transactions on Machine Learning Research (03/2025)
hold for all architectures and no uncertainty information is used in the training process. Another method that has been popularized in deep learning recently is the Laplace approximation (Daxberger et al., 2021). Tt. builds on the idea that, in a Bayesian setting, the posterior distribution over the network weights can be approximated by a Gaussian distribution, centered around a local maximum. The advantage is that this local maximum is available by regular training, while the covariance matrix of the Gaussian can be recovered by efficient approximations of the Hessian of the loss function (Botev et al., 2017).
A different class of approaches to uncertainty quantification is given by ensemble methods, where the main idea is to obtain an ensemble of predictions that differ slightly from each other and act empirical predictive distribution. A common way to utilize this for neural networks is to use so-called deep ensembles (Lakshminarayanan et al., 2017); where the ensemble is created by training several networks separately, with each optimization leading to a different local minimum. Although this has shown to perform well in practice, it is hindered by the computational demands of training multiple neural networks. Another method to obtain ensembles is by perturbing the input data before the forward pass through the model, which leads to an ensemble-based distribution in the output. This approach is used mainly in dynamical systems and weather forecasting (Bi et al., 2023; Bülte et al., 2025). However, this requires additional fine-tuning of the input perturbations and can lead to instabilities in the model predictions.
Finally, another way of uncertainty quantification is provided by scoring rule-based approaches. A scoring rule (Gneiting & Raftery, 2007) assigns a numerical score to a cetrribution-obsevattinn pair and can be minimized as a loss function, in order to fit a predictive distribution. Scoring rule approaches have been combined with nonparametric distributional regression (Walz et al., 2024), or neural networks (Rasp & Lerch, 2018), including multivariate cases (Pacchiardi et al., 2024; Chen et al., 2024). However, to the best of our knowledge, scoring rule approaches have not yet been transferred to the infniitedinmensional setting.
Since neural operators map to function spaces, the previously mentioned approaches need to be adjusted for uncertainty quantification in infniite-dimensional spaces. Weber et al. (2024) and Magnani et al. (2024) extend the Laplace approximation to Fourier neural operators, by performing the approximation in the last layers of the Fourier domain. While this approach efficient to implement and requires no additional training, the assumption of Gaussian can be restrictive, especially when it comes to calibration. Furthermore, the method is only applied post-hoc and no uncortainty information is used in the training process. Guo et al. (2023) propose an uncertainly quantification method that utilizes an information bottleneck and a Gaussian distribution to recover mean and standard deviations in the output space. However, this approach only leads to an estimate of the mean and standard deviation and does not allow for other measures of uncertainty. Ma et al. (2024) propose a quantile neural operator that is based on conformal prediction and provides a risk-controlling prediction set that includes the true solution with a user-specified probability. While this approach leads to a risk control for the error of the neural operator, it does not lead to a predictive distribution that can be used for assessing the full uncertainty in the prediction. In different line of work, Ma et al. (2024) utilize the deep ensemble and Monte Carlo dropout methods as baselines, as they can be extended to the infiniteddicnsioon setting in straightforward manner, leading to a functional ensemble prediction. However, the usability of the deep ensemble approach is limited by the computational overhead of training multiple models independently, especially a large number of ensembles are required.
To the best of our knowledge, there are no other methods than those mentioned above that are specifically designed to combine uncertainty quantification with neural operators. Our work overcomes the previously discussed limitations, it provides a predictive distribution that is directly incorporated into the training process by combining neural operators with scoring rule training.
1.3 Outline
This article is structured as follows: Section 2 provides preliminaries regarding the problem setting and neural operators, while Section introduces neural network training via scoring rule minimization in scparable Hilbert spaces. Section 4 is dedicated to the numerical experiments for several benchmark tasks, as well as a real-world application. The article finishes with the conclusion in Section 5.
3

---

Published in Transactions on Machine Learning Research (03/2025)
2 Operator learning
Let A= A(D, and 4 = U(D;R** denote separable Banach spaces of functions over a bounded domain D C Rd with values in Rda and IRd", respectively. In this work, we consider paramelric partial differential equations of the form:
La(u](>) f(x); u(x) ub(x),
x EED, I E OD,
(1)
for some u euf U and boundary condition us U. We assume that the solution U exists and is bounded for every Wp and that La E L( is mapping from the parameter space A to the (possibly unbounded) space of linear operators mapping U to its dual. This gives rise to the nonlinear solution operator G (a, ub) (d) + u mapping the parameters to the corresponding solution. Note that this includes initial condition as well as time-dependent problems, as D can be chosen as a spatio-temporal domain. We typically assume that we have observations {(a,, u,)} where aj ~ I generated from some probability mcasure μ supported on A and uj = (az).
2.1 Neural operators
The aim of neural operators (Kovachki et al., 2023) is to generalize deep neural networks to the infinite- dimensional operator setting in order to learn an approximation of G via a parametric map
go A u- U, E
(2)
with parameters from a finttedimmensional space C Rp. By minimizing a suitable loss function one wants lo
choose d* O such that Gg- z Gt. Following the framework of Li al. (2020), neural operators are constructed using a sequence of layers consisting of learnable integral operators combined with pointwise non-lincarities. For input function E U(D; IRd)) at the 7th layer, the map GB UU(D; Rd)) U(D; Rd. ) es + 0441 is computed as
Ggv((1)= = o (K(a ()vs (x) (o)) Vo
Here, Wio Rdd - Rdv is a linear transformation, a R - denotes a pointwise nonlinearity, and K: Ax Ox + C (U(;; Rd.) (.d.), (, Rd)) ) parameterized by o Ox S R.. Typically, K(a; is chosen kernel integral transformation that is parameterized by a neural network.
is map to the bounded linear operators on U(D; which is
More specifically, in this article we mainly focus on the Fourier neural operator (FNO) (Li et al., 2021), which specifies the integral kernel as a convolution operator defined in Fourier space as
Govi(x) = oo F(E))(()) + Wioer(a)). VI D,
=
where F the Fourier transform of function f : - RR,,, F-1 inverse, and Rss = F(RØ), where Kg D + gd. is a periodic function parameterized by neural network parameters @ E Ox. This means that in each layer i, Ris is directly parameterized in the Fourier domain by neural network with parameters & E Øg. Typically, several of these layers are combined together with additional pointwise lifting and projection operators (Li et al., 2021). Several other variants and extensions of Fourier neural operators have been proposed, such as U-shaped neural operators (Rahman et al., 2023), spherical Fourier neural operators (Boney et al., 2023) or Wavelet neural operators (Tripura & Chakraborty, 2023).
3 A probabilistic neural operator
In the neural operator setting introduced above, the data pairs can be corrupted by noise and in time- dependent PDEs long prediction horizons might lead to unstable predictions. Moving to a probabilistic setting, instead of mapping directly to the solution, we require the neural operator to output probability distribution over the function space 4. Assume that we have observational data {(a u;)) with prior
4

---

Published in Transactions on Machine Learning Research (03/2025)
distribution 03 PA and conditional distribution U, P (a) and let Pu denote set of probability measures over U. We consider a probabilistic neural operator to be map
Go : A x u 0 E
that aims to learn distribution Pa(-Iaj = (a(;))))))) E Pu parameterized by 0 E θ that fulfills Po(. (aj)))))))
P a, ). Samples from the conditional distribution are generated via Go(a, where 2 ~~ Az and μæ is a probability distribution over some space 2, usually chosen as a simple tractable distribution such as a Gaussian. For a fixed initial condition a, we denote the approximate (empirical) conditional distribution parameterived by the neural nelwork as P& = [G)(a, }m with Am ~ μiz, m = I,.., M, where we drop the dependence on the initial condition a from the notation. In this article, we restrict ourselves to data from separable Hilbert spaces, which include most solution spaces of PDEs, such as the Sobolev spaces H k N.
3.1 Proper scoring rules for neural operators
Our methodology builds around the theory of proper scoring rules and its extensions to separable Hilbert
spaces (Steinwart & Ziegel, 2021; Ziegel et al., 2024). Let (X, be a measurable space and let. Mi (X) denote the class of all probability measures on X. For P S M1 (X²), scoring rule is function S P X + [- oo]
such that the integral f SS(P,))dQ(C exists for all P,Q E P. Define the so-called expected score P) = S(Q, x(Q = Ex=r|S(Q X)]. Then S is called proper relative to the class P if
P) S S)Q,P), for all P.QEP.
(3)
and it is called strictly proper if equality in (3) implies P P=Q (Gneiting & Raftery, 2007). In other words, for strictly proper scoring rule, P and Q coincide if and only if Q is the unique minimizer of the expected score. This allows for efficient and simple training of neural networks via strictly proper scoring rules, as shown later.
In this article, we focus on the so-called eneryy score (Gneiting & Raftery, 2007), which is a commonly used scoring rule for the multivariate finite-dimensional setting that has been successfully used in combination with neural networks (Pacchiardi et al., 2024; Chen et al., 2024). In the following theorem we prove that the energy score can be extended to a strictly proper scoring rule over separable Hilbert spaces.
Theorem 3.1 (Energy score). Let H denote seporable Hilbert space and E H. The eneryy score ES : (H) x H. + R defined as
1 2
ES(P,z) Epll|X
x||x]
E X'1l
(4)
with X,X' lidd (F(
~ PEME = (PEM IN [=(l)d((()) < oo) is strictly proper relative to the class
Proof. As a first step, we show that a separable Hilbert space induces a positive definite kernel that is characteristie. measurable and bounded kernel called characteristic the kernel embedding defined by ((P) f (w)))))(((), with P E (H) { PP (2(H)) fx Vk(zzx)dP(() <o0. is injective (Steinwart & Ziegel, 2021).
Let (X.d) d) be metric space and define Mi(x)= = M)()) 3o E X' such that f fd(o, x) dP(x) < oo). It is said to be of negative type if d is negative definite, i.c. for all n E N, In E X and a1...., On E R with Et Ou; = 0 it holds that Σ a,a,d(a,, x,) S 0 (compare Berg et al., 1984, Definition 3.1.1). Furthermore, let Pi, E Mf(x'. Then, following Lyons (2013), if X' has negative type, it holds that
f [((1) - P,) S.
(5)
Furthermore, the space (X, is said to be of strong negative type if it is of negative type and equality in (5) holds if and only if P P2, which is the case for separable Hilbert spaces (Tyons, 2013, Theorem 3.25).
5

---

Published in Transactions on Machine Learning Research (03/2025)
Furthermore, metric space (X, of negative type induces positive definite kernel k X' x X - R given k(r, v) = d(x, #0) + d(y, =o) - d(z,1) v) for some fixed E X' Finally, Sejdinovic et al. (2013, Proposition 29) state that for a metric space of strong negative type the corresponding induced kernel k is characteristic.
Next, we cite Steinwart & Ziegel (2021, Theorem 1.1), who prove that the kernel score Sk. :M x X R, defined as (P, ] x')
Ep[k(X, 2>] + {k(2, æ)), where a X and X i.dd P M is a strictly proper scoring rule if ls is a characteristic kernel. Therefore, only remains to show that the energy score corresponds to the kernel score Sk, given by the induced kernel ka == d(x, 20) + d(y. 20) - d(z, v) in the Hilbert space H with metric d(x,, x') = 1| xlllx and for some fixed E H. We obtain
Ska (P.x) = -Ep[ka(X x)] + Ep (p(k
ka x))
2
-Epld((X, 20) + d( 20) d(x, x]] +]]]+]]]]]]] + Ep (5p[d(X,aa)) +0) +((', d(x,
1 (d((1, (0) + d(ar, - d(r,z))
2
-Ep|d(X,))) Epld(( (0,()) x1) - E =n)1 2 "Epdd(X'' =0)] | Ep|d(x, 20)) + d(x, 2o)
=Er(d(X,x)) 2 Br(d(x, x²) FF..P.))
Since H is of strong negative type, ka characteristic and, thus, by Steinwart Ziegel (2021), the energy score is strictly proper relative the class M (H).
3.2 Training neural operators via the energy score
After introducing the underlying concepts and showing that the energy score strictly proper separable Hilbert spaces, we now introduce the corresponding training regine for neural networks. This follows the general framework of Pacchiardi et al. (2024), adapted to operator learning. Given a strictly proper scoring rule, the minimization objective is given as
66)
0
and leads to P)(- objective can be approximated with a Monte-Carlo estimator. While closed-form expressions of S are not always available, the energy score has a representation that admits an unbiased estimator, which requires the output from our neural network to consist of multiple samples of the predictive distribution, i..., M > 1." The minimization objective for the neural network with the energy score is then given by
almost everywhere (Pacchiardi et al., 2024). In the finite data setting, the
M
M
argmin
llu| - uell 2M(M - 1)
Σ
1
u 1|
(7)
N
M
j=1
j,h=l 3/th
Different mothods are available to specify a neural network to output a conditional (ompirical) distribution. Here, we focus on two main approaches. One way to obtain samples is to utilize the reparametrization trick (Kingma et al., 2015), which learns pointwise normal distribution in the last layer of the network, by including two projection layers corresponding the mean and standard deviation of normal distribution. Although this has a computational advantage as only one forward pass is required, the predictive distribution is lixed to Gaussian, which might not be expressive enough for every task. A different way to obtain samples is via stochastic dropout (Gal & Ghahramani, 2016). By adding Bernoulli or Normal random samples to spocific weights of the network, a forward pass becomes stochastic and, therefore, lcads to a prodictive distribution. While this approach has higher computational demands due to the need to perform multiple
Note that this setting still includes the classical neural operator as the special case, where the predictive distribution Dirac measure. The objective function then becomes the || loss.

---

Published in Transactions on Machine Learning Research (03/2025)
forward passes, the partial stochasticity of the network leads to higher expressivity with regard to the conditional distribution (Sharma et al., 2023). We add additional dropout in the Fourier domain, setting random frequency parts as zero with the pre-specified dropout probability, making use of the global structure of the Fourier domain. Since both approaches are quite different and come with different strengths and weaknesses, we included them both in the experiments. For the remainder of this article, we will refer to the methods as PNOd and PNOR for the dropout variant and the parametrization variant, respectively.
4 Numerical experiments
This section describes our experimental sctup. In Subsection 4.1 we introduce the baseline uncertainty quantification methods for comparisons, while Subsection 4.2 includes the evaluation metrics. In the remaining subsections we present the different experiments. All experiments are evaluated on a previously unseen test set with predictions aggregated over ten training runs with different random seeds. The different. datasets as well as the corresponding grids are normalized before they are used as input to the network. As architectures, for our experiments, we focus on the Fourier Neural Operator (FNO) (Li et al., 2021) and provide additional results for the U-shaped Neural Operator (UNO) (Rahman al., 2023) architecture in Appendix B. For the spherical shallow water equation employ the spherical FNO (Bonev al., 2023). The training process and the corresponding hyperparameters are explained in detail in Appendix A. Although all uncertainty quantification methods use the same underlying architecture in the experiments, the dropout rate has been tuned separately, as it has a significant impact on the shape of the predictive distribution for the different methods. While we keep the dropout rate in the Fourier layers and the standard layers equal during the experiments in section 4, we assess the effect of both changes independently in Appendix C. In order to evaluate our approach beyond the pure performance, we include an additional runtime analysis of the different methods in Appendix D. In Appendix F we assess the infiuence of the number of samples M during training for PNO. Finally, in Appendix G we provide more detailed comparison of PNOp and PNOR, highlighting differences and providing guidelines on which approach to choose.
4.1 Baseline models
As a basic comparison and a sanity check, we use a deterministic neural operator (DET), where uncertainty is not directly accounted for. Here, the predictive distribution is a degenerate Dirac distribution with
Ppd (G (a))
(8)
i.e. the predictive are all constant value independent of m. In addition, compare our method to different uncertainty quantification methods, adapted to the task of operator learning, Gal & Ghahramani (2016)
show that a noural network with dropout before each layer is mathematically equivalent to a variational approximation of a Gaussian process. This leads to a simple and efficient way of creating a predictive distribution, known as MCDropout (MCD), which has been extended to neural operators by Ma et al. (2024). In our setting, the predictive distribution for a fixed input a is given as
PM ( (a, zmm) M
where 2m is the random dropout mask and Gj denotes the trained neural operator. As opposed to the PNO, the MCDropout model trained on the L loss and no uncertainty information is used in the training process.
Weber et al. (2024); Magnani et al. (2024) propose to utilize the Laplace approximation (LA) for FNOs, which is based on building a second-order approximation of the weights around the maximum a posteriori (MAP) estimate fMAP- By assuming a Gaussian weight prior, the weight-space uncertainty of the LA is given by
C(0,C) R (0; IMAP, E), Σ 0))xar)
(10)
where C = (an, N By sampling from this Gaussian we obtain the predictive distribution PM We
utilize a last-layer Laplace approximation, which leads to an analytically available Gaussian predictive distribution in function space. Since the Laplace approximation is built around the MAP estimate, also does not utilize uncertainty information during the training process.
7

---

Published in Transactions on Machine Learning Research (03/2025)
4.2 Evaluation
To Uhoroughly assess the performance of the prediction methods, we use a range of evaluation metrics, each focusing on a different aspect of the probabilistic prediction. Denote by D the spatio-temporal domain of the governing equation, by Pgd = {u"m]m Im= the prodictive distribution, and by u the corresponding true observation. Furthermore, let uM == M -m=1 un denote the mean prediction, 6² = -ーM Em= (u -
denote the estimate of the (pointwise) variance and Ø denote the empirical (pointwise) quantiles of Pj at the level a. We utilize the following evaluation metrics:
CL u) == juu ull 2
(11)
M
M
ES(PM ()) = M
Σ
a1ah
20lL> 2M - 1)
Σ
|10t0 athl 111,
(12)
m=1
m. m, migh
(L/
CRPS(PMM u) = D
Σ
QSa (6) (), da)
(13)
(x,1)ED
NLL(PM u) 2
1
Σ
log (2TO2 t)) + (2ns(2r t(2
(14)
(2(2,))
(x,
C. (PM ,u)= = D
.u)=
Σ
uu(,,t) E (x,, 1/-0/2222 (xr, (x,
().00 00
(15)
(x,1)ED
CPMP ,u)] = [D] (a.()ED
1
d0
- 10 (x,,)
1 ala/22
(16)
The LLa loss evaluates the match between the mean prediction äM and the observation, while the energy score (FS) evaluates the match for the predictive distribution as a whole. The contimuous ranked probability score (CRPS) (Gneiting & Raftery, 2007) evaluates the predictive distribution at a pointwise level and can be represented as an integral of the quantile score QSa v) = 2(a 1/ < qa )y qa) over all quantile levels a E (0, 1). The CRPS therefore assesses whether the predicted uncertainty fits the observations at all quantile levels at each gridpoint (Ir f), i.e., whether the predictive distribution is well-calibrated. Furthermore, we analyze the negative log-likeldhood (NLL) of a predictive pointwise Gaussian distribution M (z, f), 62, f)). Although not all methods produce Ganssian output, this criterion is commonly used and describes the fit of the predictive distribution in terms of the first two estimated moments. The NLL is also a strictly proper scoring rule (Gneiting& & Raftery, 2007). Finally, we report the coverage C. of the predictive distribution for selected a-quantile levels, as well as the width of the corresponding prediction interval [C,,]- Following the notions of calibration and sharpness (Gneiting et al., 2007), in an ideal setting, estimator should obtain C. ~ 1 - a, while simultancously predicting the interval |Cal as small as possible. To make the basclines comparable, M = 100 predictive samples are drawn from the respective predictive distributions at test time, and all metrics are averaged over the test. dataset. For the deterministic model, the energy score reduces to the L22 loss and the CRPS reduces to the pointwise mean absolute error, while the negative log- ligllioel and the coverage are not applicable.
4.3 Darcy flow
First, we consider a steady-state solution of the two-dimensional Darcy flow over the unit square, which is defined by a second-order linear elliptic equation of the form
Vu(z)) f(x), u(x) = 0,
x E 10, x E 0 1)2.
Here, a E I ((0,1))) is the diffusion coefficient, u E ((, 1)2, R) is the velocity function, and f [20 R)): the forcing function. The aim is learn the solution operator G* L ((2, + (O, R)) The PDE is usually used for modeling fluid flow in a porous medium. We use data
8

---

Published in Transactions on Machine Learning Research (03/2025)
Model
Method
L1
ES
CRPS
NLL
Co.o5
Co.o5
0.0972
0.0972
0.0783
DET
(+ 0.0076)
(= 0.0076)
(# 0.0062)
0.0772
0.0569
0.0462
-1.3444
0.9813
0.3682
PNOp
(+:00015))
(+00.009))
(=0.0007)
(=0.3355)
(+0.0032)
(+000185)
0.1031
0.0817
0.0707
685.15
0.4702
0.1249
FNO
PNOR
(+0.0069)
(+0.0056)
(=.00040)
(=913.21)
(+0.0559)
(=000019))
0.1002
0.0830
0.0672
1.5973
0.5187
0.1203
MCD
(+00.088)
(+000080)
(=0.0060)
(=0.8223)
(+0.0608)
(+000031)
0.0985
0.0767
0.0621
-0.8851
0.9762
0.5713
LA
(=0.0074)
(=0.0040)
(=0.0033)
(=0.0531)
(=0.0085)
(=0.0076)
Table 1: Results for the Darcy Flow equation using the FNO architecture with the different baselines. The best model highlighted in bold for all metrics, except for the interval width, as it only comparable for models with the same coverage. The standard deviation across the different experiment runs is given in the brackets.
from Takamoto et al. (2022), with the forcing function fixed as f(x) = 1, which corresponds to the setting in Li et al. (2021). The provided data consists of 10000 samples with a resolution of 128 x 128. We train the networks on a spatial resolution of 64 64 and evaluate on 128 128. The results are presented in Table 1. The PNOd obtains the best scores across all metrics when using the FNO architecture. Specifically, it is worth noting that it obtains lower 2 loss, although the deterministic model is dircctly trained on that metric, indicating that learning the whole predictive distribution can also improve the mean prediction. While both the PNOd and the Laplace approach achieve the coverage goal, the prediction interval of the PNOp is significantly narrower despite the coverage of PNOp being better as well. visualization of the different predictions can be found in Appendix
4.4 Kuramoto-Sivashinst equation
Model
Method
L1²
ES
CRPS
NLL
Co.05
Co.05
0.9028
0.9028
0.7100
DET
(+ 0.0086)
(#: 0.0086)
(# 0.0106)
0.8793
0.6195
0.5496
3.0389
0.7640
3.0852
PNOp
(#0.0072)
(#0.0049)
(#0.0097)
-..8772)
(+0.0191)
(#:00422)
0.8640
0.6081
0.4743
1.2268
0.9401
3.2594
FNO
PNOn
(100.038)
(+0.0027)
(10.0037)
(=.00883)
(+0.0030)
(100.099)
0.8635
0.7541
0.5974
107.10
0.3600
0.6046
MCD
(#0.0048)
(+0.0049)
(+0.0058)
(=22.941)
(=0.0082)
(#00.014)
0.8772
0.7332
0.6081
43.9302
0.3955
0.8175
LA
(+0.0055)
(+0.0103)
(10..007)
(1:0.8764)
(+0.0189)
(+.0.669))
Table 2: Results for the Kuramoto-Sivaslinsky equation. The best model is highlighted in bold and the standard deviation is given in the brackets.
The Kuramoto-Sivashinsky (KS-)equation is a fourth-order parabolic PDE with the following form in one spatial dimension:
aul(r,t) vd.u(x,t) +0 ou((,t) + Ju t) 0, u(I, 0) uo(x))
x x (0,T)
EED
where D C R, u E C((0 (D; (D;R), R)), and ug E 2 (D; R).. Furthermore, un is L-periodic, where L is the size the specific domain. The PDE exhibits strong temporal chaotic behavior and is used for modeling
2 HHe (D; RR), Tee (D; R) denote the periodic Sobolev and 1 spaces, respectively.

---

Published in Transactions on Machine Learning Research (03/2025)
MCD
LA
PNOp
PNOR
30
Ground truth
10
Mean prediction
20
10
2
30
Absoluzto errer
20
2.0
Standard deviation
30
20
Coverage (a
10
40
80
100
20
Figure 1: The figure shows the ground truth, mean prediction, absolute error, standard deviation, and coverage for the different methods on sample of the Kuramoto-Sivashinste equation.
instabilities in laminar flame front, cacctiondiiffsion systems, other chemical reactions. We simulate the KS-equation from random uniform noise 1) periodic domain D = [0, 100) using the py-pde package (Zwicker, 2020). We generate 10000 samples with resolution of 256 x 300 and train a 2D (spatio-temporal) FNO, which takes the first twenty timesteps as input and solves the equation for the next twenty timesteps with step size At = 2, The results for the KS KS-equation can be found in Table 2. Here, the MCDropout approach lcads the best I loss, while the PNOR leads the best distributional metrics. Although no method achieves coverage of 95%, the PNOR is closest. A corresponding visualization can be found in Figure 1.
4.5 Spherical shallow water equation
Boney et al. (2023) have extended the FNO to convolutions on the sphere using spherical harmonic basis functions, leading to the so-called Spherical Fourier Neural Operator (SFNO). Following their experiments, we test our method on the shallow water equation on the rotating sphere. The shallow water equations (SWE) set of hyperbolic partial differential equations, derived from the general Navier-Stokes equations
10

---

Published in Transactions on Machine Learning Research (03/2025)
and present framework for modeling free-surface flow problems. In the spherical setting, these given
(,h(2,1) Vh(a, 1000 =0 =0, 0.h(a, 1) + V ( unh( ()) = S,
a x E &ESL L (0, * #ES E (0,77,
with the vector of directional velocitics u (u, u) Here, E C((O, T); (5?; IB)) the geopotential of the water and S contains flux terms, such the Coriolis force. Similarly to Bonev al. (2023), we simulate 5000 training samples with resolution 128 256 with PDE parameters that model the Farth and additional 500 evaluation samples. Gaussian random fields are used as initial conditions with a burn-in period of three hours for the numerical solver. The model is trained and evaluated by using the corresponding losses in a weighted version, adjusted to the latitude weighting of the sphere. Furthermore, Bonev et al. (2023) propose a fine-tuning step, where the model is trained with an additional autoregressive step, which leads to more stable results for long rollout times, which we additionally include in our experiments. For evaluation, we report both training versions with one-hour and ten-hour prediction horizons each. The results for the two-step autoregressive training can be found in Table 3 and the results for the single-step training in Table 12 in Appendix I. For the one hour prediction, the PNOp obtains the best metrics across all models, and especially good coverage. For the ten-hour prediction lask, the deterministic model performs best, although the PNOR and PNOd perform quite similar. It is notable that the PNOR leads to very large values for the negative log- oog-likelhhod,, due to the predicted variance collapsing zero, which could indicate an optimization lailure. However, the method is still compelitive the remaining metrics. This might. be due overconfident. pointwise predictions, which lead to a small variance and, therefore, a large NLL. This behavior is also visible in the visualization in Figure 2. The MCD approach does not lead stable predictions for the ten-hour prediction task, as is reflected in all metrics. However, for the single-step training, the performance is much better (compare Table 12), indicating that the method cannot be combined with the autoregressive training approach.
Method
L1
ES
CRPS
NLL
Co.05
Co.o5
0.3150
0.3150
0.1244
DET
(+: 0.0129)
(#: 0.0129)
(# 0.0046)
0.2970
0.2216
0.0903
-0.4561
0.9828
0.8959
PNOp
(+0.0232)
(+0.0142)
(+0.0056)
(#0.0447)
(+0.0031)
(+0.0040)
0.3871
0.2740
0.1561
10
0.0274
0.0844
1h
PNOR
(+00,0126
(+.0.0888)
(+0.0050)
(+1.6x109)
(+00.206)
(1:0.5511)
0.3022
0.2368
0.0944
0.8878
0.6023
0.2550
MCD
(++00204)
(+0.0196)
(+000066)
(+0.2608)
(+0.0250)
(+000072))
0.4262
0.3159
0.1400
0.2374
0.9096
1.0866
LA
(+-0.5522)
(+000299)
(+-001200)
(+0.2229)
(+000762)
(+.0.3472)
0.9215
0.9215
0.3933
DET
(+ 0.0678)
(#: 0.0678)
(#: 0.0254)
1.2421
0.9915
0.4330
5.4385
0.4935
0.9210
PNOp
(+.00650)
(+.00535)
(1000215)
(+.0.448)
(+.0.000)
(1.0.7993)
1.1803
0.9535
0.4995
2.1
0.0212
0.0850
10h
PNO
(+:0.0345)
(+0.0413)
(+0.0130)
(+1.4x10))
(+0.0200)
(+00.544)
10"
7.5
2.8x
22.292
0.3530
6.7x10"
MCD
(+.27x100%)
(+2.2x
(+8.3x105)
(+19.623)
(+001164)
(.2.0x109)
1.5405
1.2612
0.5299
6.4958
0.5759
1.1925
LA
(+.033697)
(+.0.3338)
(+30.336)
(15.5154)
(+.0.1480)
(1:0.4394)
Table 3: Results for the spherical shallow water equation with two-step auloregressive training. The best. model is highlighted in bold and the standard deviation is given in the brackets.
11

---

Published in Transactions on Machine Learning Research (03/2025)
MCD
LA
PNOp
PNOR
Ground truth
-2
2
Mcan prediction
0.6
Absolute error
0.0
0.6
0.5
Baaddrd deviation
0.2
0.1
0.0
Coverage = 0.05)
The figure shows the ground truth, mean prediction, absolute error, standard deviation, and coverage for the different methods on a sample of the spherical shallow water equation. The shown variable is the water geopotential and the prediction horizon one hour with predictions obtained via two-step autoregressive training.
Figure 2:
4.6 Forecasting of 2m surface temperature
The experiments so far have focused on systems in which the underlying equations are known and compu- fationally tractable. We want to demonstrate method complex large-scale dynamical
12

---

Published in Transactions on Machine Learning Research (03/2025)
system with real data. More concretely, of the Farth's surface temperature. Recently, neural networks have shown huge success in medium-range weather prediction tasks, with models such as FourCastNet (Kurth et al., 2023), or PanguWeather (Bi et al., 2023). However, while some approaches are available that directly learn probabilistic prodictions (Price et al., 2025), most models are either deterministic or require an additional post-processing step (Bülte al., 2025).
approach to create probabilistic medium-range forecasts
Here, we want demonstrate the performance of our model in predicting the 2-meter surface temperature across Europe. For computational reasons restrict the prediction task to only one imput variable. We use the ERA5 dataset (Hersbach et al., 2020), provided via the benchmark dataset WeatherBench (Rasp et al., 2024). The data has a spatial resolution of 0.25% x 0.25 and a time resolution of 6/. For computational reasons, we restrict ourselves to spatial domain over Europe. We use a 3D spatio-temporal neural operator that takes the last 10 timesteps (60 hours) as the imput and predicts the next 10 timesteps. Therefore, we directly predict several future timesteps, instead of using an autoregressive model. The results of the temperature prediction task can be found in Table All models obtain fairly similar and consistent 1 loss, whereas PNOp exceeds the alternative methods in the remaining metrics. However, none of the methods obtain optimal coverage. It is noticeable that the PNOR leads to very poor NLL score, while the PNOp outperforms the other methods significantly with regard to that metric. Figure 3 shows the CRPS and coverage across the spatial domain, aggregated over the test data. The CRPS shows a clear performance distinction between land and sea, which is supported by the findings in Biile et al. (2025). While for MCDropout and the PNOR, the coverage is clearly different between land and sea, for Laplace and PNOp this difference negligible. The latter is preferable, as the models should provide adequate uncertainty even if the underlying spatial domain makes predictions more difficult. Figure 4 shows additional visualization of the corresponding model predictions.
MCD
LA
PNOp
PNO
2.5
2.0
-55
1.0
-0.
0.0
100
Coverage 1961
Figure 3: The figure shows the spatial distribution of the CRPS and coverage, aggregated over the test samples for the different methods on the two-meter surface temperature prediction task with a prediction horizon of 24h.
13

---

Published in Transactions on Machine Learning Research (03/2025)
Model
Method
ES
CRPS
NLL
Co.o5
Co.o5
0.0250
0.0250
0.0267
DET
(# 0.0001)
(# 0.0001)
(= 0.0001)
0.0242
0.0174
0.0187
-1.9587
0.8623
0.0975
PNOD
(=0.0000)
(=0.0000)
(#000001)
(=0.0111)
(+0.0033)
(#0.0012)
0.0241
0.0187
0.0204
616.98
0.6536
0.0576
FNO
PNOR
(#0.0001)
(#0.0001)
(#0.0004)
(+1.2x10*)
(#0.0222)
(#0.0022)
0.0240
0.0202
0.0214
2.3623
0.4456
0.0351
MCD
(+0.0001)
(+0.0000)
(+00.0001)
(=0.0237)
(#00.0015)
(+00.001))
0.0251
0.0195
0.0210
-0.1502
0.6956
0.0802
LA
(=0.0001)
(=000006)
(=0.0005)
(=1.0099)
(#0.1169)
(#0.03444)
Table 4: Results for the ERA5 surface temperature prediction task. The best model is highlighted in bold and the standard deviation is given in the brackets.
MCD
PNOp
PNOR
20
Ground truthh (C')
10
Prodiction (C-)
D0
Absolute errer
Standard deviation
Coverage (a -
Figure 4: The figure shows the ground truth, mean prediction, absolute error, standard deviation, and coverage for the different methods sample of ERA5 dataset. The prediction horizon is 60 hours.
14

---

Published in Transactions on Machine Learning Research (03/2025)
5 Discussion
In this work, we introduce PNO, framework for learning probability distributions over the output function space of neural operators. Our approach combines neural operators with generative modeling utilizing strictly proper scoring rules. We provide a theoretical justification of the approach and analyze how scoring rules can be extended to separable Hilbert spaces, More specifically, we demonstrate how the energy score can be utilized create a predictive distribution in the output function space of a neural operator. We present and examine difffernt ways of generating a probability distribution using neural operators that are agnostic to the underlying architecture. Our framework enables consistent performance in the respective mean prediction task. These advantages are especially noticeable when the underlying system is highly uncertain or chaotic.
accurate uncertainty quantification while maintaining
We evaluate our approach across several experiments, including stationary PDE solutions, dynamical systems, and weather forecasting. The experiments demonstrate that PNO is transferrable to different network architectures, as well as different domains. Furthermore, PNO performs well in ppatio-eepoors as well autoregressive settings. Compared to the baseline methods, PNO shows the best performance across variety of metrics, especially with regard to the fit between the predictive distribution and the observations. Surprisingly, the mean prediction of the PNO is an improvement against the deterministic model, indicating that modeling the predictive distribution can remove bias in the mean prediction. Finally, the predictions are better calibrated in terms of the coverage of corresponding - a prediction intervals.
5.1 Limitations
The main limitation of the PNO is the increased computational complexity and memory requirements. These are mainly induced by the energy score, as it requires multiple samples to evaluate. In the case of the PNOp this also requires multiple forward passes the underlying model. runtime analysis shows that the training time per epoch of the PNO is around 1.5 3 times higher than for the comparing methods. Another limiting factor of PNO is that inherits the expressivity and (lack of) understanding of the underlying generative model. While PNOn, corresponds to a predictive Gaussian in output space, for PNOp the exact distribution and its properties are not available. A better understanding of the predictive distribution could allow for convergence or coverage guarantecs. Finally, the results show that PNO leads to unstable predictions for autoregressive inference with long rollout times.
5.2
Future work
One of the most important tasks for future research revolves around reducing the computational complexity of our method. As a first step, analyzing the capacities of training a deterministic model and fine-tuning the PNO afterwards appears to be a promising direction. In Appendix E we provide first results in that direction, showing that utilizing fine-tuning can significantly reduce computational costs. In a similar manner, the amount of generated samples could be adapted throughout the training process, incrcasing performance while reducing training time simullaneously. In addition, other parameters of the training regime could be adjusted to further reduce training and inference time. Another promising research direction is the application of alternative suitable scoring rules such as the Gaussian kernel score. Alternative kernels might induce different characteristics in the predictive distribution and the optimal kernel might be dependent on the underlying data scenario. Future work could also revolve around extending PNO to more complex architectures and larger models. Especially in the context of weather forecasting, models such as the adaptive Fourier neural operator (Kurth et al., 2023) have shown great success and a combination with proper scoring rules might provide a practical method for uncertainty quantification for large-scale systems. In addition, implementing PNO into large-scale models for physical applications such weather forecasting lead to improvements in the corresponding uncertainty quantification and extreme event identification, ultimately leading to a deeper understanding of the prediction of such surrogate models.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Acknowledgements
C. Bülte and G. Kulyniok acknowledge support by the DAAD programime Konrad Zuse Schools of Excellence in Artificial Intelligence, sponsored by the Federal Ministry of Education and Research.
P. Scholl and G. Kutyniok acknowledges support by the project 'Genius Robot' (01IS24083), funded by the Federal Ministry of Education and Research (BMBF), as well as the ONE Munich Stralegy Forum (LMU Munich, TU Munich, and the Bavarian Ministery for Science and Art).
G. Kutyniok acknowledges partial support by the Munich Center for Machine Learning (BMBF), well as the German Research Foundation under Grants DFG-SPP-2298, KU 1446/31-1 and KU 1446/32-1. Furthermore, G. Kutyniok is supported by the GAIn project, which is funded by the Bavarian Ministry of Science and the Arts (SLMWK Bayern) and the Saxon Ministry for Science, Culture and Tourism (SMWK Sachsen). G. Kutyniok also acknowledges support from LMUexcellent, funded by the Federal Ministry of Education and Research (BMBF) and the Free State of Bavaria under the Excellence Strategy of the Federal Government and the Länder as well by the Hightech Agenda Bavaria.
References
Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavanzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U. Rajendra Acharya, Vladimir Makarenkov, and Sacid Nahavandi. A review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information Fusion, 76:243 297, December 2021. ISSN 1566-2535, doi: 10.1016/..inffus.2000
Christian Berg, Jens Peter Reus Christensen, and Paul Ressel. Hormonic Analysis on Semigroups, volume 100 of Graduate Texts in Mathematics. Springer, New York, NY, 1984. ISBN 978-1-4612-7017-1 978-1- 4612-1128-0. doi: 10.1007 191-0079978-1-412-111
Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium-range global weather forecasting with 3D neural networks. Nature, 619(7970):533-556) 538, July 2023. ISSN 1476-4687. doi: 10.1038/s41586- 0001:///445560022066---53 023-06185-3. Publisher: Nature Publishing Group.
Boris Bonev, Thorsten Kurth, Christian Hundt, Jaidcep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkunar. Spherical Fourier neural operators: learning stable dynamics on the sphere. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of ICML '29, pp- 2806 2823, Honolulu, Hawaii, USA, July 2023. JMLR.org.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss-Newton Optimisation for Deep Learning. In Proceedings of the 34th International Conference on Machine Learning, pp. 557 565. PMLR, July 2017. ISSN: 2640-3498.
Christopher Bülte, Nina Horat, Julian Quinting, and Sebastian Lerch. Uncertainty quantification for data-driven weather models. Artificial Intelligence for the Earth Systems, 2025. doi: 10.1175/ AIES-D-24-0044 URL https:/ //j //pp:://///rall/// ametsoc. wtgg:v/is///ourmaaaaaaais/////// /ttp://///AII------0000-0- 1/AIES-DA D-05--D244 24-0249 1. xml.
Jieyu Chen, Tim Janke, Florian Steinke, and Sebastian Lerch. Generative machine learning methods for multivariate ensemble postproccssing. The Annals of Applied Statistics, 18(1):199-183, March 2024. ISSN 1932-6157, 1941-7330. doi: 00.1214/23AAASTTE Publisher: Institute of Mathernatical Statistics.
Erik Daxberger, Agustimus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hemnig. Laplace redux- - effortless bayesian deep learning. In M. Rauzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in neural information processing systems, volume 34, pp. 20089-20103 Curran Associates, Inc., 2021.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty deep learning. In Maria Florina Balcan and Kilian Q. Weinberger (ods.), Proceedings of the 33rd international conference on machine learning, volume 48 of Proceedings of machine learning research, pp. 1050-1058 New York, New York, USA, June 2016. PMI.R
16

---

Published in Transactions on Machine Learning Research (03/2025)
Tilmann Gneiting and Adrian E. Raftery. Strietly proper scoring rules, prediction, and estimation. Journal of the American Statisticol Associalion, 102(477):359-333 378, 2007. ISSN 0162-1459. doi: 10.1198/00 014198/00622445000000111333
Tilmann Gneiting, Fadoua Balabdaoui, and Adrian E. Raftery. Probabilistic forecasts, calibration and sharpness. Journal of the Royal Statislical Sociely Series B: Statistical Melhodology, 69(2):243 268, 03 2007. ISSN 1369-7412. doi: N0.1111j.11467-10 90....8....2003... URL https://dos www....gggg.... 1111/j11jjjjj... 1467-9868. 2007.00587.. xx
Vignesh Gopakumar, Stanislas Pamela, Lorenzo Zanisi, Zongyi Li, Ander Gray, Daniel Brennand, Nitesh Bhatia, Gregory Stathopoulos, Matt Kusner, Mare Peter Deisenroth, Anima Anandkumar, the JOREK Team, and MAST Team. Plasma surrogate modelling using Fourier neural operators. Nuclear Fusion, 64 (5):056025, April 2024. ISSN 0029-5515. doi: (0008/1741-43222 ///////........
Publisher: IOP Publishing.
Ling Guo, Hao Wu, Wenwen Zhou, Yan Wang, and Tao Zhou. IB-UQ: Information bottleneck based uncertainty quantification for neural function regression and neural operator learning. orXiv preprint, 2023. doi: 10.48550/ar arv. wtt.iii2302003227.. URL http://arxiie http://aarii..orrrab/22222. 03271.
Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz-Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, Adrian Simmons, Cornel Soci, Saleh Abdalla, Xavier Abellan, Gianpaolo Balsamo, Peter Bechtold, Gionata Biavati, Jean Bidlot, Massimo Bonavita, Giovanna De Chiara, Per Dahlgren, Dick Dee, Michail Diamantakis, Rossana Dragani, Johannes Flemming, Richard Forbes, Manuel Fuentes, Alan Geer, Leo Haimberger, Sean Healy, Robin J. Hogan, Flías Hólm, Marta Janisková, Sarah Kecley, Patrick Laloyaux, Philippe Lopez, Cristina Lupu, Gabor Radnoti, Patricia De Rosnay, Iryna Rozum, Freja Vamborg, Sebastien Villaume, and Jean-Noël Thépaut. The ERA5 global reanalysis. Quarterly Journal of the Royal Meteorological Society, 146(730):199)) 2049, July 2020. ISSN 0035-9009, 1477-870X. doi: 10.1002/g 2.11002/j..800 qj.3803.
Jiri Hron, Alex Matthews, and Zoubin Ghahramani. Variational Bayesian dropout: pitfalls and fixes. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2019-20200 PMLR, 10-15 Jul 2018. URL attps://proceediing mlr. www.eprresssseeeeeeeeeeee hww./////////...///////.. html.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. arXiv preprinl, 2017.
doi: 10.48550/ 10.4455///Xii.1112353 arw.a...ii...1.66 URL
Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in neural information processing systems, volume 28. Curran Associates, Inc., 2015.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications pdes. Journal of Machine Learning Research, 24(89):1 97, 2023. URL http://jont org/ http:/ppprs//////e v244/ 21-1524. htmll
Thorsten Kurth, Shashank Subramanian, Peter Harrington, Jaideep Pathak, Morteza Mardani, David Hall, Andrea Miele, Karthik Kashinath, and Anima Anandkumar. Fourcastnet: Accelerating global high-resolution weather forecasting using adaptive fourier neural operators. In Proceedings of the Platform for Advanced Scientific Computing Conference, PASC '23, New York, NY, USA, 2023. Association for Computing Machinory. ISBN 97984400701900. doi: 1145/3592979, 3593412.
015599999.35339111 URL
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in neural information processing systems, volume 30. Curran Associates, Inc., 2017.
17

---

Published in Transactions on Machine Learning Research (03/2025)
Bian Li, Hanchen Wang, Shihang Feng, Xiu Yang, and Youzuo Lin. Solving Scismic Wave Equations Variable Velocity Models With Fourier Neural Operator. IREE Transactions on Geoscience and Remote Sensing, 61:1 18, 2023. ISSN 1558-0644. doi: 10.1109/TGRRS.022.... 202..3.022333333666 Conference Name: LEEE Transactions on Geoscience and Remote Sensing.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kanshik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural Operator: Graph Kernel Network for Partial Differential Equations. arXiv preprint, 2020. doi: 10.48550/ 10.48550/arXiv:200000 (atoa/arrii:.:.3355533 URL http://xxivv.rggbb/////e/// /00/005522003004885...
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadeneselli, Burigede liu, Kaushik Bhaltacharya, Andrew
Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. International Conference on Learning Representations, 2021. URL https://openrevier http:://pppeereeeeeeeeeerorrrrr net/forum?id= c8P9NQVtmnO.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218 229, March 2021. ISSN 2522-5839.
doi: 10.108//$4255-021-000022
Russell Lyons. Distance covariance in metric spaces. The Annals of Probability, 41 -3054, September 2013. ISSN 0091-1798, 2168-894X. Statisties.
doi: 10.1214/12-AOSS0. Publisher: Institute of Mathematical
Ziqi Ma, David Pitt, Kamyar Aviwadenesheli, and Anima Anandkumar. Calibrated uncertainty quantification for operator learning via conformal prediction. Transactions on Mochine Learning Research, 2024. ISSN 2835-8856. URL httpss///ppntt attps://openrrevew net pet/forrumiid=cpegyy11"
Emilia Magnani, Marvin Pförtner, Tobias Weber, and Philipp Hennig. Linearization Turns Neural Operators into Function-Valued Gaussian Processes. arXiv preprint, 2024. doi:
URL
http://arxiv.orgg 2b000220000072 05072.
Lorenzo Pacchiardi, Rilwan A. Adewoyin, Peter Dueben, and Ritabrata Dutta. Probabilistic forecasting with generative networks via scoring rule minimization. Journal of Machine Learning Research, 25(45):1-64 14,, 2024.
Ilan Price, Alvaro Gannhez-Gonzzlez, Ferran Alet, Tom R. Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, Remi Lam, and Matthew Willson. Probabilistic weather forecasting with machine learning. Nature, (87(8444):84499, January 2025. ISSN
1476-4687. doi: 10. 10.1038/s4155 022-55588---2---2--9.
Md Ashiqur Rahman, Zachary F, Ross, and Kamyar Aviwwadenesheli. U-NO: U-shaped neural operators. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https://openrevier. forum?id==3oQfocoodd
Stephan Rasp and Sebastian Lerch. Neural Networks for Postprocessing Ensemble Weather Forecasts. Monthly Weather Review, 146(11):3885 3900, November 2018. ISSN 0027-0644, 1520-0493. doi: 10.1175/
MWR-D-18-0117..1
Stephan Rasp, Stephan Hoyer, Alexander Merose, Ian Langmore, Peter Battaglia, Tyler Russell, Alvaro Sanchez-Gonzalez, Vivian Yang. Rob Carver, Shreya Agrawal, Matthew Chantry, Zied Ben Bouallegue, Peter Dueben, Carla Bromberg, Jared Sisk, Luke Barrington, Aaron Bell, and Fei Sha. Weatherbench 2: A benchmark for the next generation of data-driven global weather models. Journal of Advonces in Modeling Earth Systems, 16(6), 2024. doi: https: https://doi.org// /ttp://do.oor10.1100/////// http
Peter Renn, Cong Wang, Sahin Lale, Zongyi Li, Anima Anandkumar, and Morteza Gharib. Forccasting subcritical cylinder wakes with Fourier Neural Operators. arXiv preprint, 2023. doi: 10.48550/ar URL http://arxiv. http://aaavvvorrba/a/a/a/aaaa org//b///30...oe0/ 08290.
18

---

Published in Transactions on Machine Learning Research (03/2025)
Dino Sejdinovie, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equivalence of distance-based and RKHS-based statistics in hypothesis testing. The Annals of Stalistics, (155):2263-22)) October 2013. ISSN 0090-5364, 2168-8966. doi: 10.1214/13-ASSIA
Publisher: Institute of Mathematical Statistics.
Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. Do Bayesian Neural Networks Need To Be Fully Stochastic? In Proceedings of The 26th International Conference Artificial Intelligence and Stalistics, pp. 7694 7722. PMLR, April 2023. ISSN: 2640-3498.
Ingo Steinwart and Johanna F. Ziegel. Strictly proper kernel scores and characteristic kernels on compact spaces. Applied and Compulational Harmonic Analysis, 51:510 542, March 2021. ISSN 1063-5203. doi:
10.1016/j.acha.20199 201.9.....0011110000...
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Piliger, and Mathias Niepert. PDEBench: An extensive benchmark for scientific machine learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in neural informalion processing systems, volume 35, pp. 1596-1611
Curran Associates, Inc., 2022.
Tapas Tripura and Souvik Chakraborty. Wavelet Neural Operator for solving parametric partial differential equations computational mechanics problems. Computer Methods Applied Mechanics and Engineering, 404:115783, February 2023. ISSN 0045-7825. doi: /t.00//j.ccaa.222..111.
Eva-Maria Walz, Alexander Henzi, Johanna Ziegel, and Tilmann Gneiting. Easy Uncertainty Quantification (EasyUQ): Generating Predictive Distributions from Single-Valued Valued Model Output. SIAM Review, 66(1): 91-122, February 2024. ISSN 0036-1445. doi: N0.1137/2MM154199 Publisher: Society for Industrial and Applied Mathematics.
Tobias Weber, Emilia Magnani, Marvin Pförtner, and Philipp Hemnig. Uncertainty quantification for fourier neural operators. In ICLR 2024 workshop on ATJDifferentaal@uuttoos in science, 2024. URL
https://openreviier mett/forunndddkksyooolllll
Gege Wen, Zongyi Li, Qirui Long, Kamyar Azizzadenssheli, Anima Anandkumar, and Sally M. Benson. Real-time resouionn CO2 geological storage prediction using nested Fourier neural operators. Energy & Environmental Science, Publisher: The Royal Society of Chemistry.
April 2023. ISSN 1754-5706. doi: (0.1039/D2EE04000
Johanna Ziegel, David Ginsbourger, and Lutz Dümbgen. Characteristic kernels on Hilbert spaces, Banach spaces, and on sets of measures. Bernoulli, 30(2):1441-14) May 2024. ISSN 1350-7265. doi: 10.3150/ 23- 33-BEJ1639.
Publisher: Bernoulli Society for Mat hematicaal Statistics and Probability.
David Zwicker. py-pde: A Python package for solving partial differential equations. Journal of Open Source Software, 5(48):2158, April 2020. ISSN 2475-9066. doi: 10.21105/joss.022
19

---

Published in Transactions on Machine Learning Research (03/2025)
A Model training and hyperparameters
Across the experiments, all uncertainty quantification methods share the same underlying neural operator architecture. The architectures for the Darcy-Flow experiments are laken from Li et al. (2021) and Rahman et al. (2023) for the FNO and UNO, respectively. An overview relevant model parameters can be found in table Table 5, As the underlying function space, we use the L space and its corresponding norm. The models are trained for a maximum of 1000 epochs and we use carly stopping, to stop the training process, if the validation loss does not improve for ten epochs in order to avoid overfitting. For optimization, we employ the popular Adam optimizer (Kingma & Ba, 2017) with gradient clipping. Due to its size, additional learning rate scheduler for the ERA5 dataset, which halves the learning rate if the validation loss does not improve for five epochs. The PNO methods are trained with generated samples, while we use an ensemble size of M = 100 across all methods for evaluation. All experiments were implemented PyTorch and run on an NVIDIA RTX A6000 with 48 GB of RAM. The accompanying code can be found tatps://githhn com/ cbuelt/pfno.
Batch
Learning rate
Hidden
Projection channels
Lifting channels
Model
Modes
size
channols
FNO
64
0.001
32
256
256
(12,12)
Darcy
(64,128,128,
(8..
Flow
UNO
64
0.001
32
32
64,32]
(188))(18188)
FNO
0.001
128
(10,12)
KS
[16,32,64,128,
(4,4),,(
UNO
32
0.001
32
32
64,32,16]
(10)))(1))))))))1)
SSWE
SFNO
64
0.005
32
256
256
(32,32)
ERA5
SNO
16
0.005
20
128
128
(10,12,12)
Table 5: Overview of relevant model hyperparameters.
20

---

Published in Transactions on Machine Learning Research (03/2025)
B Results for the U-shaped neural operator
This section provides additional results for our method using the U-shaped neural operator (Rahman et al., 2023), highlighting the transferability across different architectures. While the model generally performs a bit worse than the FNO on our experiments, the performance of the uncertainty quantification methods shows a similar ordering as before. The PNOD leads to the best performance on most metrics for Darey flow and the Kuramoto-Sivashinsky equation. Only once, for the L loss for Darey flow the MCD approach is preferable and PNOr for the CRPS and the coverage on the KS equation.
Experiment
Method
L
ES
CRPS
NLL
Co.o5
Co.05|
0.1949
0.1949
0.1565
DET
(+ 0.0073)
(+ 0.0073)
(+ 0.0060)
0.1894
0.1354
0.1097
-0.2856
0.9069
0.6998
PNOp
(+0.0066)
(+0.0050)
(+0.0041)
(+0.0614)
(+0.0243)
(=0.0477)
Darcy Flow
0.2005
0.1473
0.1275
20.018
0.5765
0.4218
PNOR
(+0.0060)
(+0.0055)
(+0.0044)
(+11.926)
(+0.0698)
(=0.055))
0.1875
0.1491
0.1233
2.2782
0.5468
0.2822
MCD
(+0.0058)
(+0.0076)
(+0.0069)
(#1.4980)
(#0.0600)
(=002599
0.1954
0.1530
0.1229
0.9996
0.7477
0.4374
LA
(#0.0074)
(#0.0052)
(#0.0040)
(+0028335
(+0.0411)
(=006433)
0.9492
0.9492
0.7652
DET
(+: 0.0250)
(+: 0.0250)
(+ 0.0241)
0.9315
0.6560
0.6014
2.0875
0.7828
3.3892
PNOp
(+0.0046)
(+0.0032)
(+0.0088)
(+002564)
(+0.0315)
(=0.1214)
Kuramoto-
0.9371
0.6595
0.5326
5.4528
0.9110
3.4039
PNOR
Sivashinsky
(#0.0022)
(#0.0015)
(#0.0023)
(+6.2975)
(#0.0150)
(=0.0442)
0.9444
0.8500
0.7163
449.08
0.1965
0.4052
MCD
(+0.0014)
(+0.0034)
(=0.0026)
(+.2.839)
(+0.0050)
(=0.0109)
0.9494
0.8544
0.6928
54.338
0.2560
0.5770
LA
(#0.0250)
(#0.0265)
(=0.0249)
(#37.019)
(#0.0473)
(=0.1251)
Table 6: Results for the Darcy Flow and Kuramoto-Sivashinsk. equation using the UNO architecture, The best model is highlighted in bold for all metrics, except for the interval width, as it is only comparable for models with the same coverage. The standard deviation across the different experiment runs is given in the brackots.
21

---

Published in Transactions on Machine Learning Research (03/2025)
C Effect of Fourier dropout
For the PNOD stochastic forward passes are required to generate the predictive distribution as the output of the neural operator. As mentioned in Subsection 3.2, we use an additional dropout by masking random components in the Fourier domain (Fourier dropout) of each Fourier block addition to the standard dropout (weight dropout). In the main paper, we always set the same value for the weight dropout as for the Fourier dropout. Here, we provide an additional study examine the dependence between the different dropout rates and the performance of the PNOp. An FNO architecture is trained for the Kuramolo-Sivashinkk, equation with different combinations of the dropout. The corresponding results be found in Table 7, while Figure shows corresponding visualization. The dropout rates seem to have structured effect on the predictions, as the L loss and the energy score show better performance with higher dropout rates, However, the NLI. seems to be smaller with lower dropout rates. Generally, several of the metrics are minimized by a weight dropout of 0.05 and a higher (L², ES) or lower (CRPS) Fourier dropout. However, detecting a siguificant effect is difficult due to the stochastic nature of the approach and the large computational demand for performing detailed experiments. Still, the experiment suggests that different types of dropout seem beneficial, as they induce more stochasticity in the model that could lead to higher expressivity with regard to the prediclive distribution.
622
ES
0.91
0.64
10 1.
0.64
0.90
Weight dropout
0.90
0.63
0.63
0.89
10-2.
0.89
10-2
0.62
10-2
10-1
10-2
10-1
CRPS
NLL
0.56
6.00
10
10
Weight dropout
5.00
0.55
0.53
10-7-
10-7
2.00
0.52
10'22
10-1
10-²
10-1
Fourier dropout
Fourier dropout
Figure 5: The figure shows the different metrics in dependence of the Fourier and weight dropout for the Saramoto-Sivashinske equation. The axes are shown on a log-scale.
22

---

Published in Transactions on Machine Learning Research (03/2025)
Weight dropout
Fouricr dropout
1
ES
CRPS
NLL
Cu.os
|Co.us
0.01
0.8998
0.6337
0.5666
1.4866
0.8505
3. 3.4649
0.02
0.8942
0.6296
0.5469
1.4281
0.8746
3. 3.4825
0.01
0.05
0.9032
0.6361
0.5315
1.5714
0.8366
3.2629
0.1
0.8899
0.6266
0.5236
1.5835
0.8336
3.2151
0.2
0.8858
0.6237
0.5263
1.6129
0.8294
3.2427
0.01
0.9042
0.6367
0.5462
1.4577
0.8983
3.6815
0.02
0.9054
0.6376
0.5410
1.4299
0.9051
3.6425
0.02
0.05
0.8887
0.6259
0.5223
1.5018
0.8667
3.3519
0.1
0.8942
0.6298
0.5296
1.6132
0.8294
3.2363
0.2
0.8848
0.6231
0.5276
1.6191
0.8279
3.2578
0.01
0.8907
0.6274
0.5181
1.4798
0.8712
3.2865
0.02
0.8922
0.6284
0.5221
1.4968
0.8660
3.3190
0.05
0.05
0.8908
0.6273
0.5342
1.6776
0.8198
3.2360
0.1
0.8900
0.6268
0.5466
1.9742
0.7766
3. 1370
0.2
0.8806
0.6201
0.5327
1.7949
0.8078
3.2122
0.01
0.9072
0.6390
0.5357
1.5733
0.8418
3.2286
0.02
0.8914
0.6280
0.5344
1.6708
0.8151
3. 3.1668
0.1
0.05
0.8910
0.6276
0.5425
1.8126
0.7929
3. 1333
0.1
0.8877
0.6255
0.5643
3.1770
0.7391
3.0874
0.2
0.8852
0.6234
0.5537
2.7249
0.7638
3.0821
0.01
0.8891
0.6264
0.5520
2.2083
0.7598
2.9729
0.02
0.8870
0.6249
0.5563
2.5560
0.7431
2.9557
0.2
0.05
0.8895
0.6266
0.5561
2.4287
0.7527
3.0111
0.1
0.8869
0.6247
0.5640
3.4221
0.7352
2.9904
0.2
0.8828
0.6217
0.5675
6.7351
0.7319
3.0469
Table 7: Metrics for different combinations of dropout for training the Saaamoto-Sivsshiskyy equation.
23

---

Published in Transactions on Machine Learning Research (03/2025)
D Runtime analysis
Since using a scoring rule as the loss function requires predicting multiple samples for cach optimization step, a runtime analysis of PNO is of high interest. In Table 8, we show the estimated runtime of MCD, LA, PNOD, and PNOR- As expected, both PNO methods require significantly more time per epoch, with PNOD taking about three times as long as the baseline methods and PNOR about 50% longer. The main bottleneck is the evaluation of the energy score, which requires to calculate distance matrix for all predictive samples. Furthermore, PNOp requires several forward passes per optimization step, increasing the runtime even further. PNOR the other hand requires size (16 instead of 32), which also increases the computing time per epoch compared to MCD and LA, which are both trained without uncertainty.
memory and had to be trained with smaller batch
Method MCD LA
Avg. Number of Epochs
Time per Epoch (s)
Total Time (s)
112.89
22.36
22718
140.00
21.56
27161
PNOd
132.33
58.71
69920
PNOR
209.78
31.92
60272
Table 8: Runtime comparison on SSWE with one-step autoregressive training. The number of epochs is averaged across the ten training runs using early stopping with 10 epochs of patience. The timings are averaged over 25 training epochs on an RTXA6000 and the best models are highlighted bold.
24

---

Published in Transactions on Machine Learning Research (03/2025)
E Improving computational efficiency
As previously mentioned, the PNO requires substantially more memory and computing time than the deterministic neural operator. In this section, we brielly want to present a strategy to mitigate this behavior. The most straightforward way to improve computing time is to consider a given pre-trained (deterministic) model, which is a realistic scenario a practical setting. that case, the PNOD can used to fine-tune that model, as only requires stochastic dropout with nonzero probability. Table 9 compares the runtime and metrics for the deterministic model, the PNOp and the fine-tuned PNOp. The results show that the pre-trained model requires substantially fewer epochs and training time, while at the same time leading to improved performance. Figure 6 shows additional visualization of the training and validation loss against the training time of the model. While training deterministic model and PNOp together, does not yield significant improvements in the computing time, the pre-trained PNOp requires significantly fewer epochs to converge and therefore can be efficiently used if one has an already trained deterministic model.
PNOD
DET
PNOd (pre-trained)
7043.38
4689.33
2404.10
(+: 895.79)
(#: 408.54)
(+ 700.00)
150.78
112.33
46.33
Ø epochs
(+ 19.18)
(# 9.79)
(# 13.49)
0.8793
0.8635
0.8702
CL²
(+ 0.0072)
(+= 0.0048)
(+: 0.0046)
0.6195
0.7541
0.6130
ES
(#: 0.0049)
(#: 0.0049)
(#: 0.0033)
0.5496
0.5974
0.5343
CRPS
(+: 0.0097)
(# 0.0058)
(= 0.0080)
3.0389
107.1038
6.0401
NLL
(+ 0.5872)
(+ 22.9412)
(+ 1.9573)
0.7640
0.3600
0.7950
Co.o5
(#: 0.0191)
(= 0.0082)
( 0.0154)
3.0852
0.6046
3.1327
|Co.05
( 0.0422)
(= 0.0164)
(=: 0.0288)
Table 9: Performance comparison of the pre-trained PNOp against the deterministic model and the regular PNOp using ten different model runs. The standard deviation across the runs is given in the brackets and the best model is highlighted in bold.
0.72
PNOp: Training
PNOp: Validation
Energy score
PNOp (pre-trained): Training
0.68
PNOp (pre-trained): Validation
0.64
0.62
0.60
0
1000
2000
3000
4000
5000
6000
t [(]
Figure The figure shows the training time against the energy score for the regular PNOp and the pre- trained PNOp for specific experiment run. The solid and dashed lines show the training and validation loss, respectively.
25

---

Published in Transactions on Machine Learning Research (03/2025)
F Influence of the number of predictive samples on the performance of PNO
Our PNO approach depends on utilizing an unbiased estimator of the energy score in Equation 4, which requires the model to generate M > samples. While in practice we found that even for only M = 3 generated samples, the model shows improved performance, increasing the sample size is mainly hindered by computational reasons. In order to further analyze the model training in dependence of the sample size equation with different training sample ensemble size of one hundred members across the test data.
M, we trained both PNO approaches on the sizes M. For evaluation, all models still Table 10 shows the corresponding results, which again aggregated over ten runs. For the PNOn, the L loss and energy score already perform well for M = samples, while the other metrics are improved by increasing the sample size. While the training time per epoch increases, the differences seem negligible for less than twenty samples. On the contrary, the PNOp shows the best performance, M 2 10 samples are already used during the training process. Especially the negative log-likeli lieeliol improves significantly. However, for the PNOD, the training time per epoch increases drastically, as the model also requires M stochastic forward passes. On the other hand, for the PNOD, the number of training epochs decreases with the sample size. A solution might be using different training regimes, starting with low sample size and incrcasing it during the training process, to reduce the computational demand.
Model
M
t
epochs
ES
CRPS
NLL
Co.o5
0.8793
0.6195
0.5496
3.0389
0.7640
3
37.72s
151
(+000072)
(+000049)
(+00.0997)
(+:0.5872)
(+.0.1911)
0.8790
0.6191
0.5577
5.2872
0.7541
5
44.58s
154
(+0.0059)
(+0.0038)
(+000080)
(+1.8802)
(+0.0189)
0.8799
0.6196
0.5348
2.3471
0.8001
PNOp
10
59.94s
126
(=0.0021)
(#0.0015)
(+0.0038)
(+0.2235)
(=0.0095)
0.8784
0.6188
0.5566
5.3147
0.7592
20
90.64s
135
(=0.0091)
(+0.0061)
(#0.0088)
(+1.5163)
(=000170)
0.8857
0.6239
0.5356
1.9174
0.8036
50
181.45s
98
(=0.0050)
(+0.0036)
(+000033)
(#001199))
(=0.0092)
0.8640
0.6081
0.4743
1.2268
0.9401
3
31.08s
163
(=0.0038)
(+0.0027)
(+0.0037)
(+0.0083)
(=0.0030)
0.8661
0.6096
0.4751
1.3768
0.9421
5
29.75s
146
(+-0.0036)
(++0.0025)
(+0000332)
(+0.4410)
(+-0.0027)
0.8649
0.6088
0.4737
1,2180
0.9437
PNOR
10
31.35s
152
(=0.0036)
(+0.0025)
(+000038)
(+0.0158)
(=0.0020)
0.8680
0.6110
0.4756
1,2198
0.9453
20
35.06s
135
(=0.0053)
(+0.0037)
(+000051)
(+0.0181)
(=0.0027)
0.8649
0.6088
0.4733
1.2138
0.9449
50
43.43s
181
(+0.0035)
(+0.0025)
(#0.0022)
(#0.0085)
(=0.0028)
Table 10: Comparison of the performance of PNOp and PNOR in dependence of the training sample size M for the Sanaamoo-SSvashinskky equation. The best results are highlighted in bold, where denotes the average training time per epoch.
26

---

Published in Transactions on Machine Learning Research (03/2025)
G Comparison of PNOd and PNOR
In this section we compare the two presented approaches PNOp and PNOR in more detail and provide some guidance as to which method might be more useful for a specific problem set. While section 4 evaluates both approaches with respect to different metrics and experiments, these absolute values make it difficult to see which method is generally preferable. For that purpose we here also report the skill scores across different metries, which are a common tool for asscssing relative improvements over a reference model. The skill score for performance metric over lest set is given by
Stet SF
S
Sopt Sref
where SF denotes the average performance of the method of interest, Seet denotes the corresponding average metric of reference method and Sopt denotes the optimal achievable value. For the L?, ES and CRPS metrics, Sopt corresponds to zero. For the other metrics, the skill score is not applicable (for the negalive log-likelihood Supt would be minus infinity). The skill score is positively oriented with maximum of one, negative values indicating worse performance than the reference, and zero indicating no improvement. Table 11 shows the skill scores of all methods with respect to the deterministic baseline. The results show that in the cases where PNOp outperforms PNOn, the relative improvement is quite high, while in the opposite direction, PNOn only achieves a relatively small improvement against PNOp. Averaged over all experiments, it is evident that PNOp is preferable to all metrics, it generally leads improvement against PNOR-
Experiment E5²
ES
CRPS
Darey Flow
0.2509
0.303
0.3471
KS
-0.0178
-0.0188
-0.1588
SSWE
0.2328
0.191
0.4216
ERA
-0.0059
0.0704
0.0836
Ø
0.115
0.1364
0.1734
Skill scores for the different experiments and PNOp with respect to PNOR as baseline. The last row shows the average over all experiments.
Table 11:
While from a performance perspective, PNOp is preferable, there might be other considcrations determining which model lo choose. If one has fixed architecture or already trained model, PNO, might be better choice, as it is easier to incorporate into a given model, requiring only nonzero dropout probability, while PNOR also requires additional weights. On the other hand, PNOp comes with higher computational cost, thus if the computational resources are limited, PNOR might be preferable. Finally, one might choose the method based on knowledge the underlying data distríbution and the corresponding expressivity of the method. Both methods, PNOR and PNOp only offer a crude approximation to the posterior predictive distribution (Hron al., 2018). While for PNO,, the expressivity and approximation quality of the resulting distribution, for the PNOR case, the resulting distribution is standard Gaussian. Therefore, if the data distribution is known to be similar to Gaussian, PNOr might be suitable method. On the other hand, if the data distribution is known to be complex multimodal, PNOp might perform better, although no guarantees can be made that it matches the underlying distribution.
in general very difficult make any statements about
27

---

Published in Transactions on Machine Learning Research (03/2025)
H
Additional figures
MCD
PNOp
PNO
0.0
Ground truth
0.6
0.5
0.4
0.0
0.2
-0.5
0.0
1.0
1.0
2.5
2.0
0.0
1.5
Mean prediction
0.6
0.4
0.2
-0.5
-1.0
1.0
0.5
0.8
0.4
Absolute error
0.6
0.2
0.1
1.0
0.5
0.8
Standard deviatio
0.6
0.3
0.2
1.0
0.8
Coverage lar 0.05)
0.6
0.00
0.25
0.30
0.75
1.00
0.25
0.50
0.75
1.000.00
0:25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Figure The figure shows the ground truth, mean prediction, absolute error, standard deviation, and coverage for the different methods on a sample of the Darcy-Flow equation.
28

---

Published in Transactions on Machine Learning Research (03/2025)
Additional results for SSWE
At
Method
ES
CRPS
NLL
Co.05
Co.05
0.3215
0.3215
0.1260
DET
(+ 0.0106)
(#: 0.0106)
(+ 0.0035)
0.2955
0.2202
0.0896
-0.4609
0.9840
0.8923
PNOD
(+0.0207)
(+0.0132)
(+0.0051)
(+0.0463)
(+000032)
(+0.0467)
0.3132
0.2208
0.1048
205
0.6912
0.4971
1h
PNOR
(+0.0244)
(+0.0170)
(#0.0069)
(+3.6x107)
(#0.0897)
(#0.0842)
0.3845
0.2972
0.1193
0.4516
0.7063
0.3819
MCD
(+0.0206)
(+0.0194)
(+0.0102)
(+0.5340)
(+0.0496)
(+0.0123)
0.4222
0.3063
0.1346
0.1821
0.9024
0.9606
LA
(#0.0511)
(#0.0345)
(+0.0135)
(+0.2472)
(+0.0694)
(+0.2178)
1.1346
1.1346
0.4771
DET
(+ 0.1861)
(+ 0.1861)
(+ 0.0665)
2.8x10
2.7x10
22.3235
6.0553
0.4695
43.754
PNOD
(+8.3x109)
(+8.0x109))
(+65.581)
(+0.7769)
(+0.0207)
(+-18.49))
9.4x10
7.6x
3.6x
102x1020
0.3056
1.3x 10
10h
PNOR
(+2.0010*))
(+1.6x10*)
(=7.7x10))
(#4.6x 1020)
(+0.1644)
(+3.2x10*)
1.3324
1,2066
0.5043
25.394
0.2511
0.4104
MCD
(+0.2282))
(+0.2177)
(+0.0739)
(+333598)
(+0.0047)
(+0.0304)
1.5445
1.2835
0.5412
7.2994
0.5254
1.0291
LA
(=..2696)
(+0.2493)
(#0.1041)
(#6.0597)
(+0.1033)
(#0.2616)
Table 12: Results for the spherical shallow water equation with single-step training. The best model is highlighted in bold and the standard deviation is given in the brackets.
29


Reviews:
[{"review_id": "U8gfrz6jtv", "reviewer": "Reviewer_irfP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n-\tThe paper is mostly well written and the technical details are clearly described. \n\n-\tThe paper nicely applies pre-existing techniques (variational posterior approximations and kernel scoring rules) into a very useful application. \n\n-\tExperiments look like they’re evaluated on problems of real world interest.\n\nWeaknesses:\n-\tMissing comparisons to cited relevant work such as Guo et al, ’23, and Magnani et al, ’24, Ma et al, ’24. \n\n-\tPersonally, I’d also strongly appreciate the “basic” baselines of deep ensembles (simply retraining 3x models), a single model (CRPS score becomes MAE) and a heteroscedastic predictor (predict mean and variance for the predictor).\n\n-   see also the comments below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a probabilistic neural operator, which extend the neural operator frameworks into probabilistic (mostly forecasting based) settings, producing distributions over operators (or functions). The main application is in forecasting partial differential equations or physical systems, where the function is over the entirety of the domain of the system (e.g. the region that we wish to predict the weather). The authors give several ways of training the probabilistic neural operator by constructing an energy scoring rule and finding a way to approximate its gradient using samples from a generative base model.\n\nThe paper builds on Pacchiardi et al, ’24 who come up with a method for differentiating through a proper scoring rule (not necessarily a log likelihood) by using a generative model. In the Bayesian setting, the generative model is the posterior predictive distribution. Then, the generative model is compared to either the real observation or an adversary. \n\nThe proposed energy score is a specific type of kernel, and then Thm 3.1 follows standard kernel scoring rule analyses (Steiwart & Ziegel, ’21). For finite posterior samples, Thm 3.1 reduces to Eq. 7, which is the sample based approximation used to train the network. Eq. 7 requires posterior predictive samples to evaluate the scoring rule and the kernel aspect, which is achieved by either dropout or (what is essentially) a variational approximation.\n\nExperiments are performed on a couple of simulated differential equations, where the proposed method tends to do reasonably well. It seems like neither version strongly outperforms the other, which is a shame. Perhaps a bit more work on the dropout method or on the parameterization of the variational distribution would help.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "-\tPNO_R and PNO_D are not defined. It took me quite a while to realize that PNO_D is the dropout version of the score (eq 7), while PNO_R is another variant without posterior samples by using reparameterization. What does Eq 7 simplify to in the PNO_R setting?\n\n-\tSec 4.2 – I would probably state (again?) that `u` is the given heldout observation, while P_\\theta^M is the probability distribution over functions that we are predicting. It would be good practice to go through and define each of these metrics explicitly in terms of each variable used in them (and then these equations can go in the appendix to save space). As of now, it’s a bit unclear what \\bar u_M – is it the posterior predictive mean or is it averaged over all `u`?\n\n-\tAdditionally, use \\bar \\sigma^2 instead of \\hat \\sigma^2 to keep the naming the same for variables that use the same thing.\n\n-\tTable 1: a bit more editorializing here and in the following tables would be good for understanding if you’re going to leave these as a table (and not say a bar plot). For example, “PNO_D tends to perform best across both architectues, with the FNO architecture tending to have better performance [due to the inductive bias of the problem].” Otherwise, as a reader, I’m not sure what the bold in the lower table is supposed to show / represent and and why PNO_D is not bolded for coverage in the UNO side of the table. \n\n-\t“if the correct distribution is learned, then C_\\alpha(x) = 1-\\alpha… “ The following statement here is a bit imprecise. Ideally, we would want the coverage to be exactly $1 - \\alpha$, with under-calibration occurring (e.g. overfitting) if the estimated region covers a smaller region than 1-\\alpha, and over-calibration occurring if the estimated region covers a larger region than 1 - \\alpha. Thus, we’d probably want the C_0.05 band to be “closest” to 0.95, and the width of the prediction interval (the next column) to be as small as possible.\n\n-\tSec 4.5: “PNO_R leads to extremely large values for the negative log likelihood …” Indeed, this could be quite problematic for decision making and forecasting. Often, this is due to improper parameterization of \\sigma in the reparameterization trick. How was \\sigma parameterized here?\n\n-\tSec 4.5: “MCD approach does not lead to stable predictions for the ten hour…” This is probably because for forecasting with dropout, you need to modify how the dropout is done by dropping out the same units and weights at each time step in the rollout to get consistently dropped out paths like in Gal & Ghahramani, neurips, 16. We can also see this visually in the MCD version of Figure  2 where the coverage failure is uniform, as well as the standard deviations for both dropout models there as well. By comparison, LA is at least somewhat varying between regions (as is PNO_R).\n\n-\tSec 3.2: Dropout is an extremely crude posterior predictive approximation that isn’t really theoretically grounded (see Hron, Mathews, & Ghahramani, “variational Bayesian dropout: pitfalls and fixes, ICML, ’18), and really neither is the explicit gaussian reparameterization form that you’re using from Kingma et al, ’15 due to the lack of the likelihood term. Of course, many Laplace approximations are also similarly quite crude for neural networks due to the overparameterization of the model. \n\nHowever, as the authors correctly point out, using dropout does introduce noise into the training process, which might be reasonable for producing semi reasonable uncertainty estimates. It could be better to explicitly parameterize a Gaussian posterior over the weights and toss the scoring rule in a stochastic variational inference objective (Hoffman et al, JMLR, ’13) with corrections for the generalization from the likelihood itself into the scoring rule (Lacoste-Julien et al, AISTATS, ’11 and Maus et al, Neurips, ’24 come to mind).\n\nThe takeaway here is that the claim of “we wish to take into account uncertainty during training” is not sufficient a priori to justify PNO_D over PNO_R or the other Laplace approximations. Thus, I would suggest toning down those claims.\n\n-\tSection 4.2: “if the correct distribution is learned, then C_\\alpha(x) = 1-\\alpha..”, this isn’t strictly true as we could have extremely underconfident predictive distributions that have C_\\alpha >> 1- \\alpha (as is seen in some of the experimental results). I’d suggest rephrasing this sentence to something like “if the correct distribution is learned, then C_\\alpha(x) is close to 1-\\alpha”.\n\n-\tSection 4.3: “PNO_R leads to extremely large values for the negative log likelihood..” This probably suggests a failure in optimization (e.g. the posterior variance collapsed to zero), which is interesting here because in many situations with Bayesian NNs there’s severe underfitting.\n\n\nReferences:\n\nPacchiardi et al, Probabilistic forecasting with generative networks via scoring rule minimization, JMLR, 2024.\n\nLacoste-Julien et al, Approximate inference for the loss calibrated Bayesian, AISTATS, 2011.\n\nMaus et al, Approximation aware Bayesian optimization, Neurips, 2024.\n\nHron, Mathews, & Ghahramani, “variational Bayesian dropout: pitfalls and fixes, ICML, ’18\n\nHoffman et al, Stochastic variational inference, JMLR, 2013.\n\nGal & Ghahramani, A Theoretically Grounded Application of Dropout in Recurrent Neural Networks, Neurips, 2016.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "KHgjcSd78j", "reviewer": "Reviewer_fAH7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- Proves that the energy score is strictly proper in separable Hilbert spaces.\n- Transforms neural operators from deterministic solvers into models that learn long-term probability distributions over function spaces by incorporating uncertainty directly during training.\n- Supports multiple sampling strategies and various neural operator architectures.\n- Shows better results on several PDE problems and real world scenario- weather forecasting.\n\n**Weaknesses**\n- Requires multiple forward passes or samples per training step (increased memory and runtime).\n- The increased computations can limit applicability to very large-scale or high-resolution problems.\n- Can lead to instability or growing error in long-term autoregressive forecasting scenarios.\n- The reliance on specific dropout or Gaussian assumptions may potentially still restrict flexibility in capturing complex uncertainty patterns.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a novel framework that extends neural operators by learning the probability distributions over function spaces instead of single deterministic predictions. \nIt further provides uncertainty quantification for PDEs by incorporating this uncertainty directly into the training process. \nThey prove that the energy score (scoring rule) remains valid in separable Hilbert spaces, which ensures the theoretical side of their approach. \nThe framework has flexible sampling strategies using dropout or reparameterization tricks to make it compatible with several neural operator architectures in a plug-and-play manner. \nExtensive experiments on several PDE problems and weather forecasting have provided robust results of PNO compared to traditional uncertainty quantification methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please consider below suggestions, which are not critical but can help to improve the work.\n- Discuss more regarding the *excessive computational overhead* due to the multiple forward passes along with *potential strategies/attempts* to address this issue. For example, considering the framework’s plug‐and‐play nature, can you provide results on how integrating architectures like the Adaptive Fourier Neural Operator help to reduce some of this overhead burden or solve the scalability issue?\n- Please clarify/elaborate the *reliance on specific dropout or Gaussian assumptions* and discuss how these can affect the flexibility in learning  *more complex uncertainty patterns*.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "tvE6szHtIB", "reviewer": "Reviewer_AXTt", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper introduces energy scores as a principled approach to learn solution distributions, supported by the proof that the energy score is a strictly proper scoring rule in a separable Hilbert space.\n- The authors validate their approach across multiple PDE problems and real-world data, which show both its capabilities and limitations quantitatively.\n\nWeaknesses:\n- The method introduces significant computational costs. The loss function of PNO-D is especially slow to compute.\n- Neither variant demonstrates consistently superior performance across all test cases, and there's no clear guide on which variant to choose for a specific problem.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Their key contribution is to propose the use of energy score, a widely used scoring rule, as the target to have neural operators learn the distribution of possible solutions. Their approach maps both the coefficient field and a hidden noise to a possible solution. They present two variants - PNO-D using stochastic dropout and PNO-R using the reparameterization trick. The work is thoroughly validated across multiple datasets. While one of the variants will perform well on a given dataset, neither variant has consistently superior performance across all test cases, leaving room for improvement in developing more robust approaches.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The authors should provide a more detailed analysis of when and why each variant performs better, as well as rule-of-thumb guidelines for choosing between variants based on problem type.", "broader_impact_concerns": "None that I think need to be specifically highlighted here.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies an important problem of probabilistic neural operators. The reviewers kindly provided insightful suggestions that the authors incorporate in the current draft. 

Throughout the discussion, it came out that despite MCDropout was first extended to neural operators in Ma et al. (2024) "Calibrated Uncertainty Quantification for Operator Learning via Conformal Prediction" it was not clear in the current writing of this paper. The authors are encouraged to integrate this edit.
Decision: Accept (poster)