loading page

Data Imbalance, Uncertainty Quantification, and Generalization via Transfer Learning in Data-driven Parameterizations: Lessons from the Emulation of Gravity Wave Momentum Transport in WACCM
  • +6
  • Y. Qiang Sun,
  • Hamid Pahlavan,
  • Ashesh Chattopadhyay,
  • Pedram Hassanzadeh,
  • Sandro W. Lubis,
  • M. Joan Alexander,
  • Edwin P Gerber,
  • Aditi Sheshadri,
  • Yifei Guan
Y. Qiang Sun
Rice University

Corresponding Author:[email protected]

Author Profile
Hamid Pahlavan
Author Profile
Ashesh Chattopadhyay
UC Santa Cruz
Author Profile
Pedram Hassanzadeh
Rice University
Author Profile
Sandro W. Lubis
Pacific Northwest National Laboratory (DOE)
Author Profile
M. Joan Alexander
NorthWest Research Associates, CoRA Office
Author Profile
Edwin P Gerber
New York University
Author Profile
Aditi Sheshadri
Stanford University
Author Profile
Yifei Guan
Rice University
Author Profile


Neural networks (NNs) are increasingly used for data-driven subgrid-scale parameterization in weather and climate models. While NNs are powerful tools for learning complex nonlinear relationships from data, there are several challenges in using them for parameterizations. Three of these challenges are 1) data imbalance related to learning rare (often large-amplitude) samples; 2) uncertainty quantification (UQ) of the predictions to provide an accuracy indicator; and 3) generalization to other climates, e.g., those with higher radiative forcing. 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 the test case. WACCM has complex, state-of-the-art parameterizations for orography-, convection- and frontal-driven GWs. Convection- and orography-driven GWs have significant data imbalance due to the absence of convection or orography in many 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 on when a NN gives inaccurate predictions. Finally, we show that the accuracy of these NNs decreases for a warmer climate (4XCO2). However, the generalization accuracy is significantly improved by applying transfer learning, e.g., 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.
22 Dec 2023Submitted to ESS Open Archive
27 Dec 2023Published in ESS Open Archive