Terry O’Kane: Novel neural operator and transformer architectures for climate prediction
Novel neural operator and multiscale wavelet transformer climate emulators are described and their performance in training and prediction of the global atmosphere and tropical convection are evaluated. Inference and ablation studies at varying resolutions are performed using a novel approach to ensemble generation based on early stopping and regularization of the objective function.
Of note is a first development of multi-scale wavelet transformer architectures to learn system dynamics in a tokenized wavelet domain using cross attention to address spectral bias, that is the attenuation of high-frequency components encoding small-scale structure which can further induce long-horizon instability The wavelet transform explicitly separates low- and high-frequency content across scales leveraging a wavelet-preserving down-sampling scheme that retains high-frequency features and employ wavelet-based attention to capture dependencies across scales and frequency bands.
Trained on daily ERA5 data, the developed operator and transformer models demonstrate long term stability and physical consistency with reliable emulation of key internal climate modes relevant to tropical variability exhibiting low biases over 40 year autoregressive in-sample inference for the teleconnections considered. Examination of neural operator out-of-sample inference and ablation studies indicate that long term forecast skill of the key tropical climate processes and teleconnections is dependent on accurate prediction of large-scale variations in the incident short wave radiation and to a lesser degree outward longwave radiation. Dependencies on resolution and the impact of inclusion of additional boundary constraints are described.
This work is a collaboration spanning CSIRO, NCI and international partners contributing to the ACCESS ML community.