Uncertainty Quantification for Fourier Neural Operators

Authors: tobias  weber
emilia  magnani
marvin  pförtner
philipp  hennig
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 00:34:09

Abstract: In medium-term weather forecasting, deep learning techniques have emerged as a strong alternative to classical numerical solvers for partial differential equations that describe the underlying physical system. While well-established deep learning models such as Fourier Neural Operators are effective at predicting future states of the system, extending these methods to provide ensemble predictions still poses a challenge. However, it is known that ensemble predictions are crucial in real-world applications such as weather, where local dynamics are not necessarily accounted for due to the coarse data resolution. In this paper, we explore different methods for generating ensemble predictions with Fourier Neural Operators trained on a simple one-dimensional PDE dataset: input perturbations and training for multiple outputs via a statistical loss function. Moreover, we formulate a new Laplace approximation for Fourier layers and show that it exhibits better uncertainty quantification for short training runs.

Main paper: 
 introduction: 
 Partial differential equations (PDEs) are used to model complex interactions in real-world phenomena like weather and climate. In recent years, several deep learning methods such as FourCast-Net (Pathak et al., 2022) , PanguWeather (Bi et al., 2023) , and GraphCast (Lam et al., 2023) have been developed to provide deterministic medium-term PDE predictions. This marks an important step in moving from simulated PDE datasets towards real-world data, which comes with additional challenges such as local weather dynamics that are not resolved in the input data. To address the uncertainties accompanying reanalysis data, ensemble predictions have a long history in weather forecasting. In existing literature, two primary methodologies have been used for adding ensemble predictions to the architectures mentioned above: perturbing the input (Pathak et al., 2022) and training for multiple trajectories via statistical loss function (Lessig et al., 2023) . An alternative approach arises from the Bayesian perspective on Deep Learning, by sampling new models from a posterior distribution conditioned on the training data (Gal, 2016) . In this context, Laplace approximations have shown to be competitive with more popular alternatives such as variational Bayes or deep ensembles, both in terms of performance and computational cost (Daxberger et al., 2021) . While Laplace approximations have been previously explored for Graph Neural Operators (Magnani et al., 2022) , its application remains unexplored within the context of Fourier Neural Operators (FNOs). In this paper, we hence focus on FNOs and equip the parameters of the last Fourier layer with an uncertainty estimate. Our primary objective is to compare the resulting ensemble predictions with the existing methodologies mentioned above. Since proper comparison studies of such approaches are missing, we start our analysis in this paper with a simple one-dimensional toy problem of producing ensemble predictions for the trajectories of Korteweg-de Vries equations for different initial states. This stands in contrary to new methods like NeuralGCM (Kochkov et al., 2023) and Gen-Cast (Price et al., 2023) that explore new architectures for explicitly modeling ensemble predictions. Such further comparisons and higher-dimensional experiments are left for future work. 
fourier neural operators: 
 Neural Operators (Kovachki et al., 2023; Li et al., 2021a; 2020b; a; 2021b) are a type of neural network f θ -with parameters θ ∈ C d or θ ∈ R d -designed to learn solution operators of PDEs. Their training and test sets consist of input-output functions D = {u n , v n } N n=1 discretized on some spatio-temporal mesh. The input is usually an initial state u n (•, t 0 ), and the output corresponds to the solution of the PDE after a certain time T , i.e. v n := u n (•, T ). Different types of Neural Operator architectures have been developed (Li et al., 2020a; b; 2021a) . One prominent method is the so-called Fourier Neural Operator, which consists of pointwise operating linear mappings and so-called Fourier layers. The original Fourier layer (Li et al., 2021a ) can be written asEQUATIONwhere R is a linear transform acting only on the lower Fourier modes and W is a linear mapping operating pointwise in the spatial domain. The higher Fourier modes are truncated, having a similar effect as the non-linear activation function σ. The restriction to pointwise operations and the transformation into the Fourier domain makes it possible to transfer between different discretizations. 
uncertainty quantification for fnos: 
 We denote for an input u n the ensemble of |K| predictions by either(f θ (u n,k )) k∈K , (f θ (u n ) k ) k∈K or (f θ k (u n ))k∈K , highlighting the different positions where the ensemble is introduced. For evaluating an ensemble, we consider its empirical mean f (u n ) and the standard deviation σn . In the rest of this section, we outline the different approaches and derive the Laplace approximation for the Fourier layer.A direct approach for extending any model to generate ensemble predictions is by forwarding a batch of perturbed versions of a single input u n . One way is to sample noise ϵ x,t ∼ N (0, τ 2 ) for each input value of the discretized function state u n (x, t) (Pathak et al., 2022) .FNO-Ensemble. Recently Lessig et al. (2023) suggested training directly for |K| predictions, by optimizing |K| parallel prediction heads via an extended loss function. This loss is given by summing over the squared error loss for each prediction, a statistical loss, and the ensemble spreadEQUATION)where G is the Gaussian cumulative density function.Laplace approximation for FNOs. The Bayesian approach for modeling weight-space uncertainty is based on the observation, that the supervised learning objective -extended by a weight regularizer r(θ) -can be viewed as summing over the negative log-likelihood and a negative log-prior over the weight vector θ (Bishop, 2006):θ * = arg min θ∈C d N n=1 L(v n , f θ (u n )) + r(θ) = arg min θ∈C d - N n=1 log p (v n |f θ (u n )) -log p(θ)The weight-space uncertainty under a Laplace approximation is then described by the (unnormalized) log densitylog p(θ|D) ≈ log p(θ * |D) - 1 2 (θ -θ * ) T -∇ 2 θθ log p(θ|D) (θ -θ * ) of the Normal distribution N θ * , -∇ 2θθ log p(θ|D) -1 (Daxberger et al., 2021) . In the deep learning setting, the optimized weights θ * are not necessarily a minimizer of the negative logposterior and therefore the Hessian might not be positive definite, i.e. a valid inverse covariance matrix. A common positive definite approximation is given by the generalized Gauss-Newton matrix GGN(θ):= J θ (f θ (u n )) T Λ(f θ (u n ))J θ (f θ (u n )), where J θ (f θ (u n )) := ∇ θ f θ (u n ) and Λ(f θ (u n )) = ∇ 2 f θ (un)f θ (un) L(v n , f θ (u n )) (Botev, 2020).However, for computational efficiency, the GGN is only calculated using a batch of data points and usually only a subset of weights is modeled probabilistically, where the focus is often set on the weights of the last layer. In the case of FNOs, the last linear map is only applied pointwise to the function and therefore does not capture any global structure. As an alternative, we suggest modeling the complex Fourier space parameters θ F L := vec(R) of the last Fourier layer probabilistically. To do so, we consider a Kronecker-factored approximation (K-FAC, cf. Martens & Grosse (2015) ) to further reduce the cost of computing the GGN. This approximation is based on the observation thatJ θ F L (f θ (u n )) := J vec(RF (zn)) f θ (u n )J θ F L RF(z n ) and J θ F L RF(z n ) = F(z n ) T ⊗ I hidden .The K-FAC approximation of the covariance matrix can then be derived similarly to Eschenhagen et al. (2023) and Ritter et al. (2018) as followsEQUATIONwhere we assume a centered Gaussian prior with covariance matrix given by τ 2 I and denote b n := J vec(RF (zn)) f θ (u n ). By sampling new weights θ F L ∼ N (θ * F L , Σ F L ), and passing a single input u n through the ensemble of resulting models (f θ F L,k ∪θ * ) k∈K , we obtain a corresponding ensemble prediction (f θ F L,k ∪θ * (u n )) k∈K . We denote this approach by FNO-FL. 
experimental results: 
 We consider a one-dimensional FNO with 4 Fourier layers, truncation to 16 modes, and 64 hidden dimensions. It was trained to predict 20 time steps of a solution to the Korteweg-de Vries equation (cf. 4) given the solution at 20 previous time steps.EQUATIONThe training, validation, and test set consists of 512 solutions for varying initial conditions, each discretized at 256 spatial and 140 time points (Auzina, 2022) . For comparing the methods outlined above in the context of model convergence, we consider a short and a long training run. The same training is used for the FNO-Perturbed and FNO-FL, while the FNO-Ensemble is trained for |K| = 10 predictions separately. For comparison both short training runs are stopped after the RMSE on the validation set reaches 0.28, and the long training runs after reaching 0.09. The τ parameters in the FNO-Perturbed and FNO-FL are optimized by a simple bisection search and we draw 100 ensemble members.We consider the following evaluation metrics (Chung et al., 2021) :EQUATIONQ = 1 C n,x,t (f θ * (u n )(x, t) -v n (x, t)) 2 σ2 n,x,t, and ( 6)EQUATIONwhere C is the number of summands. The RMSE describes the match of the ensemble prediction mean and should be close to zero. The Q estimate indicates the posterior calibration and should be close to one, while the average Gaussian negative-log-likelihood should be low and represents the trade-off between low standard deviations and error terms.The results of the metrics evaluated on the test set can be found in Table 1 . An ensemble prediction for a single test data sample is shown in Figure 1 . 
conclusion: 
 While FNO-Ensemble is more expensive to train, FNO-Perturbed and FNO-FL need prior precision tuning. Once a suitable τ value is found, FNO-FL is an effective tool for quantifying uncertainty Note that the uncertainty quantification provided by the Laplace approximation in FNO-FL is arguably strictly more powerful than the other two methods since it is given by a full probability measure on weights as opposed to one approximated by samples.These promising findings strongly encourage further exploration of ensemble predictions using weight-space uncertainty with Laplace approximations. However, there are still many aspects to investigate, such as uncertainty in W and scaling these methods to larger networks and highdimensional data. 
 
 Appendix:  


Decision: This paper explores the ensemble predictions with Fourier neural operators. Both reviewers agree unanimously. Therefore, I support the acceptance. However, author needs to address the concerns in the camera ready version

Reviews: review: title: title: Good insights on ensemble prediction for Fouier neural operators
main_review: main_review: Summary
------------------

The paper explores ensemble prediction techniques for Fourier Neural Operators (FNOs) in medium-term weather forecasting. It introduces a novel approach using Laplace approximation for uncertainty quantification in FNOs, addressing a gap in existing research. It evaluates various approaches, including input perturbations and statistical loss functions, and introduces a novel uncertainty quantification technique using Laplace approximation for FNOs. Through experiments on a one-dimensional PDE dataset, the paper demonstrates the effectiveness of the proposed method compared to traditional ensemble methods. The findings suggest that the proposed approach improves uncertainty estimation in FNOs, enhancing the reliability of medium-term weather predictions. Further research in ensemble prediction techniques for neural operators is needed, particularly in scaling the methods to larger networks and high-dimensional data. Overall, the paper contributes to advancing the field of ensemble prediction in weather forecasting and related domains.

Evaluation
------------------

**Quality:** The paper is good quality, offering interesting insights on ensembles of FNOs, which could have some impact on weather forecasting using neural operators and other (related) methods. 

**Clarity:** The paper provides a clear introduction to the problem domain and previous research efforts. It explains the concepts of neural operators and FNOs concisely and, in my opinion, makes it accessible to general readers, which I appreciate. The methods for UQ and ensemble prediction are well-described, with equations and algorithms nicely structured. Experimental results are presented logically, supported by tables and figures - good visuals overall.

**Originality:** The paper contributes to the field by exploring ensemble prediction techniques for FNOs, an area with limited prior research. So, its originality resides in applying a new method to a challenging problem. The introduction of Laplace approximation for uncertainty quantification in FNOs is somewhat novel, as recognized by the authors and properly cited previous related works, e.g. Magnani et al. The comparison with other ensemble methods adds value to the research, offering insights into the relative performance of different approaches. 

**Significance:** The paper's findings have significant implications for weather (and eventually climate?) forecasting and related fields in other domains and disciplines dealing with complex dynamic systems. By improving uncertainty quantification in FNOs, the proposed approach apparently enhances the reliability of medium-term predictions. The evaluation metrics and experimental results demonstrate the practical utility of the method.

Pros
------------------
1. Introduces a novel method using Laplace approximation for UQ in FNOs.
2. Provides a rigorous evaluation of the proposed method through experiments on a one-dimensional PDE dataset.
3. Demonstrates improved uncertainty estimation in FNOs compared to traditional ensemble methods.
4. Enhances the reliability of medium-term weather predictions, addressing a critical need in weather forecasting.
5. Contributes to advancing the field of ensemble prediction techniques for neural operators.
6. Opens avenues for further research in ensemble prediction methods for larger networks and high-dimensional data.

Cons
------------------
1. Requires prior precision tuning for the Laplace approximation, which may add complexity to the implementation. Discussing the limitations would be welcome, yet maybe this could be left for the workshop.
2. Uncertainty quantification in longer training runs using the proposed method appears to be underconfident.
3. Training FNO-Ensemble is computationally expensive compared to other methods. Some comments and ablation study could be added (eventually in the Appendix) and briefly commented (in the main part, just 2 sentences would suffice).
4. The generalization of the proposed method to higher-dimensional datasets needs to be investigated further. Some words of caution about that (substantial) jump to more challenging yet realistic cases would be welcome.
5. The comparison with other SOTA ensemble prediction methods could be more comprehensive. I understand this is not possible in such a short paper format, and consult with the sentence "Such further comparisons and higher-dimensional experiments are left for future work.", some extra lines in Table 1 or appendices wouldn't hurt and would add much value to the analysis.

Clarifications
------------------

- In the intro, "real-world data," I suspect the authors refer to ERA5, which is reanalysis data, not observational data. Please rephrase if so.

- "of input-output functions D = {v_n, u_n}" --> switch to output-input or exchange {u_n,v_n}

- in eq.2 one would expect sigma_n to be constant and not included in the loss, or is it a time-varying process? Please clarify as this would change the solution and ensemble in exponentially varying ways, no?

- clarify in eq.2 to what step prediction K the cdf "G f ˆ (u n ),σ n" refers to.

- is the computational cost an issue for the different approaches? could you include cputime estimates in table 1? i noticed that FNO-Perturbed and FNO-FL perform the same in RMSE but maybe there's other advantage (Q is not a definite score to declare a winner here, esp. in long training runs).

- a discussion on calibration (and eventually some metrics) would be welcome. sentences like "to be underconfident." should be substantiated with numbers - this tool can help: https://github.com/uncertainty-toolbox/uncertainty-toolbox

All these clarifications could be placed in the paper's experimental section or in an appendix/supp.mat section (which I don't find).

Typos & grammar
------------------

- Inconsistent capitalization in "uncertainty quantification" (sometimes capitalized, sometimes not).

- "It exhibits better uncertainty quantification" could be rephrased for clarity (e.g., "it shows improved uncertainty quantification").

- "Moreover, we formulate a new Laplace approximation for Fourier layers" - "formulate" could be replaced with a more precise term like "propose".

- "The training set consists of 520, the validation and test set of 100 solutions" - should include "solutions" after "520" for clarity.

- Inconsistent use of tense throughout the paper (e.g., switching between present and past tense). Choose one and maintain consistency.

- Some sentences are overly long and could be broken down for clarity and readability.

- Punctuation errors, such as missing commas or periods, should be corrected for better readability.

- In the conclusion, "arguably strictly more powerful" could be simplified for clarity.

- clean the bib; missing entries, capital letters. This is a nice tool btw: https://flamingtempura.github.io/bibtex-tidy/.
title: title: This paper explores different ensemble methods with FNOs for PDE data prediction. This paper is well written. The proposed methods are clearly described and supported by empirical studies.
main_review: main_review: This paper explores three methods for generating ensemble predictions with FNOs: 1) input perturbation, 2) training for multiple outputs via a statistical loss function, and 3) using Laplace approximation for Fourier layers for getting ensemble of weights. Empirical studies are conducted on a simple one-dimensional PDE dataset.

Pros
1. This paper is well organized. The outline of this paper is clear and easy to follow.
2. The empirical studies clearly demonstrate the advantages and the disadvantages of different ensemble methods.
3. The proposed FNO-FL is capable of providing a full probability measure on weights for uncertainty quantification, instead of approximation from generations.

Cons
1. "short run" and "long run" is defined by the validation RMSE. However, since there is no information about the convergence, it is hard to judge if the "short run" is too far from convergence and hence less meaningful. Directly comparing the convergence speed can be a good alternative.
2. The validation RMSE are 0.4 and 0.07 for "short run" and "long run" respectively. However, the training RMSE are around 0.3 and 0.11, and the test EMSE are around 0.27 and 0.1, which show different trends compared with the validation RMSE. The dataset split may cause severe distribution mismatch among training, validation and test sets.