August 6, 2024
Journal Article

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

Abstract

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 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 accuracy of these NNs decreases for a warmer climate (4×CO2). 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.

Published: August 6, 2024

Citation

Sun Y., H.A. Pahlavan, A. Chattopadhyay, P. Hassanzadeh, S. Lubis, M. Alexander, and E. Gerber, et al. 2024. Data Imbalance, Uncertainty Quantification, and Transfer Learning in Data-driven Parameterizations: Lessons from the Emulation of Gravity Wave Momentum Transport in WACCM. Journal of Advances in Modeling Earth Systems 16, no. 7:Art. No. e2023MS004145. PNNL-SA-195382. doi:10.1029/2023MS004145

Research topics