Data Imbalance, Uncertainty Quantification, and Generalization via Transfer Learning in Data-driven Parameterizations: Lessons from the Emulation of Gravity Wave Momentum Transport in WACCM

Published in Journal of Advances in Modeling Earth Systems, 2024

Sun, Y. Q., H. A. Pahlavan, A. Chattopadhyay, P. Hassanzadeh, S. W. Lubis, M. J. Alexander, E. Gerber, A. Sheshadri, and Y. Guan, 2024: Data Imbalance, Uncertainty Quantification, and Generalization via Transfer Learning in Data-driven Parameterizations: Lessons from the Emulation of Gravity Wave Momentum Transport in WACCM, J. Adv. Model. Earth Syst., 16, e2023MS004145, https://doi.org/10.1029/2023MS004145.

Official Version.

Key points

  • WACCM’s orographic, convective, and frontal gravity wave parameterizations are emulated using neural nets to inform future modeling efforts
  • Out-of-distribution generalization (extrapolation) of the neural nets under 4×CO2 forcing is enabled via transfer learning with 1% new data
  • Data imbalance is addressed via resampling and weighted loss; uncertainty quantifi- cation via Bayesian, dropout, and variational methods

Neural networks (NNs) are increasingly used for data-driven subgrid-scale parameteriza- tion 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 parame- terizations. 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 ra- diative forcing. Here, we examine 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 com- plex, 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 resam- pling 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 accuracy of these NNs decreases for a warmer climate (4×CO2). However, the generaliza- tion 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.