Joint Parameter and Parameterization Inference with Uncertainty Quantification Through Differentiable Programming

Authors: yongquan  qu
mohamed aziz bhouri
pierre  gentine
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-05 05:13:21

Abstract: Accurate representations of unknown and sub-grid physical processes through parameterizations (or closure) in numerical simulations with quantified uncertainty are critical for resolving the coarse-grained partial differential equations that govern many problems ranging from weather and climate prediction to turbulence simulations. Recent advances have seen machine learning (ML) increasingly applied to model these subgrid processes, resulting in the development of hybrid physics-ML models through the integration with numerical solvers. In this work, we introduce a novel framework for the joint estimation of physical parameters and machine learning parameterizations with uncertainty quantification. Our framework incorporates online training and efficient Bayesian inference within a high-dimensional parameter space, facilitated by differentiable programming. This proof of concept underscores the substantial potential of differentiable programming in synergistically combining machine learning with differential equations, thereby enhancing the capabilities of hybrid physics-ML modeling.

Main paper: 
 introduction: 
 Weather and climate models are critical tools for predicting weather patterns, understanding climate change, and informing future environmental policies and strategies IPCC (2021) . Central to these models is the challenging task of precisely solving time-dependent parametric partial differential equations (PDEs) that encapsulate the intricate dynamics of Earth systems. A key challenge in these models stems from the chaotic and multi-scale nature of atmospheric and oceanic processes. Owing to computational constraints, these models are typically simulated on coarse meshes (O(10)km in the horizontal), leading to miss-representation of crucial sub-grid scale processes (Schneider et al., 2017) . Yet, the mere increase in resolution is insufficient, as the set of PDEs remains unclosed due to the absence of governing equations for certain critical, yet poorly understood or unknown processes, such as the carbon cycle (Trugman et al., 2018) or microphysics. These gaps introduce significant challenges to weather and climate projection Nathaniel et al. (2024) ; Bony et al. (2015) , and underscore the need for a robust methodology to capture and couple all these dynamics that are not directly resolved or described, which is called closure or parameterization (Randall et al., 2003) . Most traditional parameterization schemes contain an empirical functional relationship with tunable physical parameters (Smagorinsky, 1963; Siebesma et al., 2007) . These parameterization schemes contribute to model uncertainty (Draper, 1995) . Estimating the physical parameters of interest, regarded as an inverse problem, can be approached by variational data assimilation (Smith et al., 2009) ensemble methods such as Kalman filter ensemble methods (Evensen, 2009; Cleary et al., 2021) and Monte-Carlo based approaches (Yang et al., 2020) , but the quality of inferred models is challenged by the dynamical systems' strong nonlinearity (Cheng et al., 2023) , and heuristic assumptions behind traditional parameterization schemes (Gentine et al., 2018) . The latter limitation has spurred the use of machine learning for modeling sub-grid scale dynamics from high-resolution simulations to emulate the coarse one (Rasp et al., 2018; Bhouri et al., 2023; Zanna & Bolton, 2020) . Kalman based methods and variational methods can presumably be used to infer physical and machine learning parameters based on uncertain observations (Evensen, 2009) . However, the strong non-linearity of the underlying models presents a challenge to the Gaussian and near-linear assumption underlying Kalman filtering (Van Leeuwen et al., 2015) and the dimensionality of the problem (with millions of degrees of freedom for neural network), limits their applicability and performance. On the other hand, differentiable modeling has been showing great promise in integrating machine learning and physical models (Shen et al., 2023) . It allows taking the derivatives within numerical errors to any model parameters whether neural network based on physically based, thus permitting the use of modern optimization techniques (backpropagation) for model inference. For parameterization, online training of neural networks (NN) with differentiable solver through target trajectories offers numerous benefits. These include enhanced numerical stability and accuracy (Frezat et al., 2022; Qu & Shi, 2023) , flexibility to integrate variational data assimilation (Farchi et al., 2023; Qu & Shi, 2023) without Gaussian assumption, and efficient uncertainty quantification facilitated by gradients of the numerical solver (Yang et al., 2020; Bhouri & Gentine, 2022) . The recent development of differentiable general circulation models (NeuralGCM, Kochkov et al. (2023) ) signals a promising future for scaling up from surrogate models like the Lorenz systems to larger-scale, more realistic systems.In this work, we consider a hybrid model that contains poorly known physical parameter values, and a neural network for sub-grid scale parameterization of turbulence. We approach the joint estimation of physical parameters and machine learning parameters with quantified uncertainty, framed as a Bayesian inverse problem, through a 2-stage approach enabled by differentiable programming. An initial estimate of the set of parameters is obtained using stochastic gradient-based optimization on temporally sparse trajectories. Then, we perform Bayesian inference of the set of parameters using stochastic gradient Hamiltonian Monte Carlo (SG-HMC) (Chen et al., 2014) , also through the set of temporally sparse trajectories. As a proof of concept, the proposed approach is applied to a twolayer quasi-geostrophic model to illustrate the potential of next generation Earth System Models combining Bayesian ML and physics using a differentiable programming framework. 
approach: 
 An abstraction of the coarse-grained dynamical system for climate and weather prediction can be represented by the following differential equation:EQUATIONwith appropriate initial and boundary conditions. Here, X denotes the estimate of true physical states X. The function F encapsulates the resolved dynamics depending on physical parameters θ phy ∈ R d1 . The function G models unknown or sub-grid scale dynamics as a function of X and parameters θ 1 . In this study, we model G as a neural network with parameter θ N N ∈ R d2 . Typically, d 2 is much larger than d 1 , reflecting the higher dimensionality of parameter space in neural networks compared to physical parameters and traditional parameterization schemes. Given numerical solver time step ∆t and initial value X t0 , n step integration using an explicit numerical scheme results inX t0+n∆t = M n (X t0 ; θ phy , θ N N ),where M is a differentiable numerical solver utilized to evolve Equation 1, and M n denotes the n-fold composition of M with itself. Assuming one ground truth observation is available every ∆T = k∆t (i.e. every k model time steps), a temporally sparse trajectory of N + 1 ground truth data points is represented as {X t0+i∆T } N i=0 , and the corresponding forecast trajectory obtained from solving Equation 1 is denoted as{X t0+i∆T (θ phy , θ N N )} N i=0 = {M ik (X t0 ; θ phy , θ N N )} N i=0 .Online deterministic training: An estimate of {θ * phy , θ * N N }, used subsequently to initialize the Markov Chain, is obtained by mini-batch gradient-based optimization of a loss function,EQUATIONwhere I is a random batch of ground truth trajectories' initial time-steps from training dataset D, and L(•, •) evaluates the distance between a ground truth trajectory and the corresponding forecast. In each training iteration, parameters θ phy and θ N N are updated based on ∂J /∂θ phy and ∂J /∂θ N N , respectively. This separation permits the potential deployment of distinct learning rates and stochastic gradient-based algorithms for each parameter type. Moreover, one may choose to cease the update of θ phy upon convergence, focusing solely on fine-tuning θ N N , especially considering the high dimensionality of the latter. The gradients can be obtained conveniently when M is written in programming frameworks that support automatic differentiation, such as JAX (Bradbury et al., 2018) , PyTorch (Paszke et al., 2019) and Julia (Bezanson et al., 2017) .Bayesian inference and uncertainty propagation: In contrast to traditional Markov-Chain Monte Carlo (MCMC) sampling methods, which are computationally intensive for large-scale Bayesian inference (Van Ravenzwaaij et al., 2018) , Hamiltonian Monte Carlo (HMC) methods offer an efficient means to sample high-dimensional parameter spaces (Neal et al., 2011) . To circumvent directly computing the costly gradient of the potential energy over the whole dataset, we adopt the stochastic gradient HMC (SG-HMC) Chen et al. ( 2014), which approximates the gradient by evaluating the likelihood on mini-batches. A Bayesian hierarchical approach is applied to quantify the uncertainty. We combine {θ phy , θ N N } as a set of uncertain model parameters θ ∈ R d1+d2 with a prior parameterized by λ that encodes our prior knowledge about θ. Another random variable γ is introduced to quantify the quality of data. The likelihood is constructed as follows:EQUATIONwhere I denotes a random batch of initial times of ground truth trajectories from training dataset D, and N represents the probability density function for normal distribution. The posterior distribution is then formulated as:EQUATIONwith specifics on the selections of priors detailed in Appendix B. Importantly, the posterior distribution does not depend on initial time t 0 of trajectories as long as the inference is conducted in statistically quasi steady state of the system, taking into account the ergodicity. The likelihood formulation Equation 3 emphasizes the need for a differentiable PDE solver M as SG-HMC necessitates the computation of the gradient of the log-likelihood with respect to θ. The Markov Chain sampling is initialized with θ * = {θ * phy , θ * N N } to favor a short transient phase and improve sampling robustness. The predictive posterior distribution of a forecast at time t, denoted by X * (t), is then given byEQUATIONallowing us to sample X * (t) for a given initial state X t0 , assuming t -t 0 = l∆t for some integer l, as depicted in the following equation:EQUATIONThe posterior mean and variance of X * (t) can be approximated from SG-HMC samples as detailed in Appendix B. 
experiments and results: 
 The proposed framework is applied to the two-layer quasi-geostrophic equations with rigid lid approximation and flat bottom topography, as implemented in PyQG JAX port, which supports automatic differentiation (Otness et al., 2023) . Detailed descriptions of the model's governing equation, the coarse-grained equation, sub-grid scale terms targeted for parameterization, and parameter specifics are detailed in Appendix A. "Ground truth" data is generated from simulations on a 256 × 256 mesh covering a 1, 000km × 1, 000km domain, with convergence testing outlined in Ross et al. (2023) . Simulations span 10 years at a timestep of ∆t = 1hour, using third-order Adams-Bashforth time stepping, and the initial 5 years are excluded as spin-up. 60 simulations are used for deterministic training and Bayesian inference, with additional 10 for testing. Finally, the data is coarse-grained to a 32 × 32 mesh using a sharp spectral cutoff filter (details in Appendix A).The physical parameter set, θ phy = {δ, U 1 }, includes the layer thickness ratio δ and the upper layer background velocity U 1 , with true values of {0.25, 0.025} and initial guess of {0.01, 0.001}. Sub-grid total tendency was modeled using a convolutional neural network (CNN) that takes the upper and bottom layer vorticities as inputs, with an architecture detailed in Appendix C. Temporally sparse trajectories of just 20 daily observations at one observation per day (∆T = 24∆t), were used for online deterministic training. The mean squared error (MSE) served as the loss function.Online deterministic training proceeds with updating both θ phy and θ N N until convergence of To assess the efficacy of the inferred physical parameters and neural network parameterizations, equation 1 is solved using a ground truth initial condition from the statistically steady state (i.e., post-year 5), over a one-year period (8,640 ∆t). Evaluation metrics include the coefficient of determination (R 2 ), Mean Squared Error (MSE), and total kinetic energy. For a comprehensive analysis of long-term behavior, the empirical distribution of the upper layer potential vorticity was examined over the prediction period's final 100 days. Performance comparisons were drawn between our framework's deterministic, maximum a posteriori(MAP) and posterior mean estimates, against traditional Smagorinsky schemes and scenarios without parameterization. The uncertainty was quantified through a 2 posterior standard deviation band, encapsulating 95% of variability, as depicted in Figure 1 . For the prediction window up to 4,000 hours, our framework's deterministic esti-Figure 1 : Metrics evaluating the performance of online predictions over 1 year period. mate outperform both the no-parameterization approach and the Smagorinsky scheme in terms of achieving the highest R 2 and lowest MSE. The MAP estimate's performance closely mirror that of the deterministic estimate. Within the initial 2,500 hours, the posterior mean showcase commendable accuracy with a minimal spread of uncertainty. The Smagorinsky parameterization exhibited over-dissipation, as shown by the continuous decline in total kinetic energy and alterations in the long-term vorticity distribution. Conversely, simulations at low resolution displayed unphysical increases in total kinetic energy, hinting at potential numerical instability for extended simulations. In contrast, our framework's deterministic, MAP and posterior mean predictions manage to better conserve total kinetic energy, albeit with a significant uncertainty spread. These outcomes underscore the viability of the proposed approach in enhancing weather and climate model predictions through efficient parameterization and uncertainty quantification. Additional results are available in Appendix D. 
conclusion and discussion: 
 Our proposed framework demonstrates the potential of integrating Bayesian differentiable programming with physical parameter inference and machine learning parameterization. This integration, complemented by efficient Bayesian inference, paves the way for more accurate and reliable scien-tific simulations and knowledge discovery. Current efforts are directed towards optimizing online training strategies, specifically exploring the optimal selection and configuration of ground truth trajectories, including their length and temporal sparsity. Such design considerations are closely tied to the characteristics of the system under study, aiming to balance the reduction of temporal correlation with the challenges of gradient backpropagation over extensive, sparse trajectories. Given the potential inaccuracies in gradient estimates from SG-HMC in large datasets, especially where high-precision scientific simulation is required, alternatives such as Control Variate Gradient HMC (CVG-HMC) (Zou & Gu, 2021 ) can be considered for gradient estimation improvements. Our future work will extend to accommodate spatially sparse ground truth data and noise, and integrate state inference, as considered in Qu et al. (2024) . For various systems, inferring parameters and parameterizations simultaneously may exhibit the equifinality issue, highlighting the importance of continuous effort to further improve the proposed framework's generalizability. Through these endeavors, we aim to refine and expand the capabilities of our methodology, contributing to the advancement of artificial intelligence for scientific modeling. 
 
 Appendix: We consider the two-layer quasi-geostrophic(QG) equation illustrating flows driven by baroclinic instability of a background velocity shear U 1 -U 2 with rigid lid approximation and flat bottom topography:EQUATIONwhere J(•, •) is the horizontal Jacobian, q i , ψ i is the layer-i potential vorticity and stream function, respectively. They are related throughEQUATIONThe physical meaning of physical parameters and their setting for generating ground truth can be found in Table 1 . The parameters to be inferred is a subset of the physical parameters, denoted by Upper Layer Background x-axis velocity 2.5 × 10 -2 U 2Lower Layer Background x-axis valocity 0 δ Layer Thickness Ratio H 1 /H 2 2.5 × 10 -1 θ phy . Following Ross et al. (2023) , we approximate the effect of grid by a coarse-graining filter (•), and the sub-grid dynamics to be parameterized is the sub-grid total tendencyEQUATIONThe filter applied is a sharp spectral truncation filter as following:EQUATION)where qκ is the Fourier transformation of vorticity at wave number κ, κ c is the cutoff threshold and ∆x LowRes is the spatial resolution of low-resolution simulations. The sub-grid totalt tendency is not available in low-resolution simulations since the first term at the right-hand-side of Equation 9is a filtered ground truth/high-resolution total tendency. Therefore, we use a neural network with parameter θ N N to model it as a function of low-resolution variable.Following Bhouri & Gentine (2022), the choices of priors are given byEQUATIONEQUATIONEQUATIONwhere α 1 , α 2 , β 1 , β 2 are hyperparameters. The the use of logarithm transformation is to ensure λ and γ are always positive. The posterior mean and variance are:EQUATIONEQUATIONHere N s is the number of samples that approximates the posterior distributions obtained through SG-HMC sampling.The sub-grid total tendency was parameterized using a convolutional neural network (CNN), detailed in Table 2 . This CNN incorporates periodic padding in each layer, includes bias terms, omits batch normalization, and employs the ReLU activation function. The parameter space of the CNN is dimensioned at d 2 = 113, 766. Implementation was carried out using the Flax library (Heek et al., 2023) . In the online deterministic training phase, we employed two separate AdaBelief optimizers (Zhuang et al., 2020) for optimizing θ N N and θ phy , each with its distinct learning rate schedule. Specifically, θ phy utilized an exponential decaying learning rate, starting at 0.01 and decreasing to 0.001 with a decay rate of 0.9. Similarly, for θ N N , an exponential decaying learning rate was applied, commencing at 0.0005 and diminishing to 0.0001 with a decay rate of 0.95. Convergence of θ phy was observed near the 50-epoch mark, at which point it was held constant to exclusively continue the training of the CNN for an additional 50 epochs.During the SG-HMC sampling phase, the hyperparameters α 1 , β 1 , α 2 , and β 2 within the prior distributions are all assigned a value of 1. The leapfrog step size for SG-HMC, denoted as ϵ HM C , is set to 5 × 10 -5 , with the number of leapfrog steps per iteration fixed at L = 10. The sampling process is conducted over 2,000 iterations.  


Decision: All reviewers unanimously agree on acceptance. However, author should address the concerns on the camera-ready version

Reviews: review: title: title: Joint probabilistic inference in hybrid partial differential equations for climate modeling: A promising proof of concept for scalable Bayesian inference with differentiable programming.
main_review: main_review: Summary:
The authors propose and explore an approach to do joint inference over a hybrid (semi-parametric) partial differential equation. In a two-step approach, they first fit both the non-parametric and the parametric part jointly with gradient descent and use the result as a starting point for SG-HMC. One key aspect of the pipeline for the integration with numerical solvers is differentiable programming.

Pros:
- The approach is set up for realistic problems of dynamical systems for climate and weather predictions and as such highly relevant.
- The manuscript is well written and the authors do a great job in clearly explaining the different layers of their approach.
- The work shows a proof of concept highlighting the advantage of the approach over traditional and non-parametric methods.
- A discussion on ways forward to improve it or potential issues within semi-parametric modeling, such as equifinality/non-identifiability, are briefly mentioned.

Cons:
- A more thorough explanation of the results for the concrete problem and their consequences would be desirable (though I understand that it goes beyond what fits into a workshop paper of 4 pages). 

Open questions:
- It is interesting to me how this approach would compare to other scalable approaches of deep uncertainty quantification, such as deep ensembles or variational inference, both in performance and posterior distributions.
- As we are doing joint inference over the parameters and the non-parametric part, it would also be nice to show the posterior distributions over the $\delta$ and $U_1$ parameters. 
- Given that the data is synthetic, is there a specific reason why the deterministic estimator does not converge to the exact $\delta$ and $U_1$ value? Does it average out with multiple fitting runs, or is there a structural positive bias? Does it persist in the posterior distribution?
- What is the no-parameterization approach?
- The results show a clear advantage of the approach over traditional and non-parametric approaches up to roughly Time 3500h. After that, the Smagorinsky parametrization seems to take the lead. How can this be explained, and is there a need and/or a way to improve the long-term performance?
title: title: Advancing hybrid modeling parameterizations with uncertainty estimates
main_review: main_review: Summary
------------------
This paper proposes a hybrid modeling approach that allows learning parameters and parametrizations in Earth system models, with the additional advantage of delivering uncertainty estimates. The approach is showcased in a toy example but seems promising and advancing in the right direction.  Authors consider a hybrid model that contains poorly known physical parameters and a neural network for sub-grid scale parameterization. They approach the joint estimation of physical and ML model parameters with quantified uncertainty, framed as a Bayesian inverse problem. The work, even if in its early stages, shows promise, and if thoroughly advanced, it could have a strong impact on the Earth and climate science community. It goes without saying that the approach, in a similar vein to NeuralGCM, would not only learn parametrizations but scale up from surrogate models like the Lorenz systems to larger-scale, more realistic scenarios.

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

**Quality:**
The paper demonstrates a robust research methodology, technical depth, and experimental rigor. It addresses a significant challenge in weather and climate modeling and proposes a novel framework for joint estimation and uncertainty quantification of physical parameters and machine learning parameterizations. The experiments are well-designed, using appropriate simulations and evaluation metrics, and the results are thoroughly analyzed. Overall, a high-quality piece of work.

**Clarity:**
The paper effectively communicates complex concepts in a clear and understandable manner. The introduction provides necessary background information, and the methodology section is logically organized with illustrative equations and figures. The experimental setup and results are presented in a structured format, making it easy to follow the research workflow and understand the findings.

**Originality:**
The paper demonstrates originality by proposing a novel framework for addressing a pressing challenge in weather and climate modeling. Integrating differentiable programming with Bayesian inference represents a unique approach to enhancing hybrid physics-ML modeling, with potential implications for advancing key problems in Earth system science.

**Significance:**
The paper's findings are highly significant for weather and climate modeling, offering potential improvements in prediction accuracy and reliability. The experimental results demonstrate promising performance over traditional parameterization schemes, highlighting practical implications for environmental science and artificial intelligence in scientific modeling.

Clarifications
------------------
- clarify what you mean by "The superscript in M n denotes function composition"

- some comments on computational complexity, a bit of ablation study on all parts (integration, HMC, model runs), and some words about future schemes

- comments about calibration of uncertainties would be welcome.

- what are "Smagorinsky schemes"?

- unlike variational approaches, MC schemes have shown great advantages in dealing with multimodal distributions but are more costly (see Svendsen 2023, https://doi.org/10.1007/s10994-021-05999-4). some discussion about this in the context of learning dynamics?

- when learning params and NN params, which are unevenly distributed and dimensional, some words of caution should be expected, both because I presume the optimization may become quite unstable and also because of the risk of running into identifiability problems, right?

- Why MSE, especially in such a dimensional uneven setting? 

Typos and grammar issues
------------------

- Fix notation around (1); \theta_1 --> \theta_NN (also take the opportunity to harmonize dimensionality notation, e.g. d_phys and d_nn, and eventually give rough numbers of those beyond the d_nn>>d_phys for scientists that are not in the field)

- several typos: "Hamiltonian Monte Calro (HMC)", "which approximate_S_ the gradient", "Peter Jan Van Leeuwen," double author?, "maximum a posteriori_(MAP)", etc. Please give a grammar pass to the text and clean your bib (this is a nice tool, btw: https://flamingtempura.github.io/bibtex-tidy/).
title: title: Parameter inference for weather and climate models.
main_review: main_review: Summary: This paper proposed a novel method for parameter estimation and uncertainty quantification which is then applied to meaningful problems in weather and climate modeling. 

Strengths and Weaknesses:

* Should add (ours) to the legend in Figure 1 to clearly indicate the author's proposed methods.

[+] The paper is well-motivated and clearly presented. It provides a clear description of the problem and the proposed solution. It is very well written, with comprehensive references to relevant material. 

[+] The uncertainty quantification is a nice benefit over the baseline approaches (Smagorinsky & no-parameterization schemes)

[-] It would be nice to have an overview figure or algorithm of the proposed framework to show the flow from data to initialization to model to prediction and what each step entails. 

[-] It would be nice to have a ground truth solution plot, either in the appendix or main text, so readers can better visualize the problem being solved. From my understanding, the author used the problem in Ross et. al. (2023). In that paper, Figure 1 gives a nice visual of the vorticity. This would help connect readers back to the application of the author's method to Earth's weather and climate predictions. 

Conclusion: The paper is concise and to the point. While there could be some improvement in the presentation of the problem and methodology, the discussion is still clear. I recommend accepting the paper and believe it is important to develop better methodologies for this class of problems after reading the paper. While the proposed method didn't excel in all metrics over all times, it shows that it is still a complex and difficult problem to solve, and this work makes a nice step in the right direction.