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

Official Version to appear.

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.