Variable Masking for Climate Emulation in ACE

In the recent AIMIP benchmark, it was shown that the ACE2.1-ERA5 model, which does not have CO$_2$ as an input, produces poor in-sample and out-of-sample tropospheric temperature forecasts. This suggests that the prior ACE2-ERA5 model, which did include CO$_2$ as an input, learned historical warming trends by overfitting to the monotonic increase of CO$_2$ as a forcing variable. In this work, we present M-ACE2S, a modified version of the stochastic ACE2S model, which improves long-term tropospheric temperature prediction. Variable masked training under the AIMIP protocol encourages M-ACE2S to learn proper physical relationships between variables and to avoid spurious correlations. We also present a novel data normalization strategy called ``global mean removal'', which further improves the tropospheric forecasts. We show that M-ACE2S outperforms the ACE2S baseline on in-sample data for ERA5 while improving training stability and generalizability to unseen forcing when paired with global mean removal.