Explainable Offline‐Online Training of Neural Networks for Parameterizations: A 1D Gravity Wave‐QBO Testbed in the Small‐Data Regime

Self Cite

Neural networks (NNs) are increasingly used for data‐driven subgrid‐scale parameterizations in weather and climate models. While NNs are powerful tools for learning complex non‐linear relationships from data, there are several challenges in using them for parameterizations. Three of these challenges are (a) data imbalance related to learning rare, often large‐amplitude, samples; (b) uncertainty quantification (UQ) of the predictions to provide an accuracy indicator; and (c) generalization to other climates, for example, those with different radiative forcings. Here, we examine the performance of methods for addressing these challenges using NN‐based emulators of the Whole Atmosphere Community Climate Model (WACCM) physics‐based gravity wave (GW) parameterizations as a test case. WACCM has complex, state‐of‐the‐art parameterizations for orography‐, convection‐, and front‐driven GWs. Convection‐ and orography‐driven GWs have significant data imbalance due to the absence of convection or orography in most grid points. We address data imbalance using resampling and/or weighted loss functions, enabling the successful emulation of parameterizations for all three sources. We demonstrate that three UQ methods (Bayesian NNs, variational auto‐encoders, and dropouts) provide ensemble spreads that correspond to accuracy during testing, offering criteria for identifying when an NN gives inaccurate predictions. Finally, we show that the accuracy of these NNs decreases for a warmer climate (4 × CO2). However, their performance is significantly improved by applying transfer learning, for example, re‐training only one layer using ∼1% new data from the warmer climate. The findings of this study offer insights for developing reliable and generalizable data‐driven parameterizations for various processes, including (but not limited to) GWs.

Section: Summary and Discussionmentioning

confidence: 99%

Section: Introductionmentioning

confidence: 99%

Data Imbalance, Uncertainty Quantification, and Transfer Learning in Data‐Driven Parameterizations: Lessons From the Emulation of Gravity Wave Momentum Transport in WACCM

Sun,

Pahlavan,

Chattopadhyay

et al. 2024

Self Cite

“…This has important implications for any data-driven SGS modeling approach, including those using deep neural networks or any other statistical learning method (Fatkullin & Vanden-Eijnden, 2004;Grooms et al, 2021;Guan et al, 2022;Sun et al, 2023;Zanna & Bolton, 2021). In fact, this non-uniqueness and uncertainty of the true SGS term is a major shortcoming of the supervised/offline learning approach to data-driven closure modeling in real-world applications, and is one of the main motivations to pursue online or at least offline-online learning approaches (Pahlavan et al, 2024;Schneider, Stuart, & Wu, 2021;Schneider et al, 2023).…”

Section: Journal Of Advances In Modeling Earth Systemsmentioning

confidence: 99%

“…While some of these studies found the learned data-driven SGS closures to lead to stable and accurate LES (Frezat et al, 2022;Guan et al, 2022Guan et al, , 2023Yuval & O'Gorman, 2020), a number of major challenges remain (Balaji, 2021;Schneider, Jeevanjee, & Socolow, 2021). Perhaps the most important one is interpretability, which is difficult for neural networks, despite some recent advances in explainable ML for climate-related applications (Clare et al, 2022;Mamalakis et al, 2022), including for SGS modeling (Pahlavan et al, 2024;Subel et al, 2023). The black-box nature of neural network-based closures aside, there are also challenges related to generalizability, computational cost, and even implementation (Balaji, 2021;Chattopadhyay et al, 2020;Guan et al, 2022;Kurz & Beck, 2020;Maulik et al, 2019;Subel et al, 2021;Xie et al, 2019;Zhou et al, 2019), limiting the broad application of such closures in operational climate and weather models, at least for now.…”

Section: Introductionmentioning

confidence: 99%

Learning Closed‐Form Equations for Subgrid‐Scale Closures From High‐Fidelity Data: Promises and Challenges

Jakhar,

Guan,

Mojgani

et al. 2024

Self Cite

There is growing interest in discovering interpretable, closed‐form equations for subgrid‐scale (SGS) closures/parameterizations of complex processes in Earth systems. Here, we apply a common equation‐discovery technique with expansive libraries to learn closures from filtered direct numerical simulations of 2D turbulence and Rayleigh‐Bénard convection (RBC). Across common filters (e.g., Gaussian, box), we robustly discover closures of the same form for momentum and heat fluxes. These closures depend on nonlinear combinations of gradients of filtered variables, with constants that are independent of the fluid/flow properties and only depend on filter type/size. We show that these closures are the nonlinear gradient model (NGM), which is derivable analytically using Taylor‐series. Indeed, we suggest that with common (physics‐free) equation‐discovery algorithms, for many common systems/physics, discovered closures are consistent with the leading term of the Taylor‐series (except when cutoff filters are used). Like previous studies, we find that large‐eddy simulations with NGM closures are unstable, despite significant similarities between the true and NGM‐predicted fluxes (correlations >0.95). We identify two shortcomings as reasons for these instabilities: in 2D, NGM produces zero kinetic energy transfer between resolved and subgrid scales, lacking both diffusion and backscattering. In RBC, potential energy backscattering is poorly predicted. Moreover, we show that SGS fluxes diagnosed from data, presumed the “truth” for discovery, depend on filtering procedures and are not unique. Accordingly, to learn accurate, stable closures in future work, we propose several ideas around using physics‐informed libraries, loss functions, and metrics. These findings are relevant to closure modeling of any multi‐scale system.

“…Then, for out-ofsample states (that are from the same distribution as those of the training), the ANN predicts the systematic model tendency correction needed to nudge the state of NWP or climate model, thus improving the trajectory and potentially the simulated statistics. While ANNs are powerfully expressive, they have a number of major shortcomings: (a) they are difficult to interpret, (b) they do not generalize to out-of-distribution, (c) they are datahungry, and (d) their predictions fed into numerical models can cause instabilities and unphysical drifts (Bretherton et al, 2022;Clark et al, 2022;Farchi et al, 2023;Guan et al, 2022;Pahlavan et al, 2024;Slater et al, 2023;Subel et al, 2023). Challenges with interpretability hinder understanding the root cause(s) of the model errors and fixing them.…”

Section: Introductionmentioning

confidence: 99%

Interpretable Structural Model Error Discovery From Sparse Assimilation Increments Using Spectral Bias‐Reduced Neural Networks: A Quasi‐Geostrophic Turbulence Test Case

Mojgani,

Chattopadhyay,

Hassanzadeh

2024

Self Cite

Earth system models suffer from various structural and parametric errors in their representation of nonlinear, multi‐scale processes, leading to uncertainties in their long‐term projections. The effects of many of these errors (particularly those due to fast physics) can be quantified in short‐term simulations, for example, as differences between the predicted and observed states (analysis increments). With the increase in the availability of high‐quality observations and simulations, learning nudging from these increments to correct model errors has become an active research area. However, most studies focus on using neural networks, which while powerful, are hard to interpret, are data‐hungry, and poorly generalize out‐of‐distribution. Here, we show the capabilities of Model Error Discovery with Interpretability and Data Assimilation (MEDIDA), a general, data‐efficient framework that uses sparsity‐promoting equation‐discovery techniques to learn model errors from analysis increments. Using two‐layer quasi‐geostrophic turbulence as the test case, MEDIDA is shown to successfully discover various linear and nonlinear structural/parametric errors when full observations are available. Discovery from spatially sparse observations is found to require highly accurate interpolation schemes. While NNs have shown success as interpolators in recent studies, here, they are found inadequate due to their inability to accurately represent small scales, a phenomenon known as spectral bias. We show that a general remedy, adding a random Fourier feature layer to the NN, resolves this issue enabling MEDIDA to successfully discover model errors from sparse observations. These promising results suggest that with further development, MEDIDA could be scaled up to models of the Earth system and real observations.