Investigating Inductive Biases for Machine Learning Emulation of Sudden Stratospheric Warmings in Idealised Isca Simulations

Abstract: Machine-learning emulators are increasingly used for weather prediction and have the potential to extend skill on subseasonal-to-seasonal timescales by learning dynamically important sources of predictability. A key challenge is whether the models can exploit predictability anchors, such as stratospheric variability, that influence tropospheric circulation beyond short lead times. We test how architectural inductive bias affects emulation of sudden stratospheric warming (SSW) dynamics using paired idealised Isca simulations that differ only in an imposed wave-2 heating perturbation. Across convolutional, transformer, and graph-based architectures trained for one-step prediction, model differences are modest when the stratosphere is dynamically quiet but widen substantially when SSW-like variability is active. Our results identify explicit three-dimensional vertical coupling as a key inductive bias for machine-learning emulation of stratospheric dynamics. However, Eliassen-Palm flux diagnostics show that low forecast error does not guarantee physically faithful wave-mean-flow interaction, with coherent errors remaining in stratospheric wave-driving structure.

• The paper shows that incorporating explicit vertical coupling in neural architectures significantly reduces errors in emulating SSW dynamics compared to horizontal-only models.
• The comparative analysis of 2D/3D CNNs, transformers, and GNNs demonstrates that 3D models better capture multi-scale stratospheric-tropospheric interactions under SSW forcing.
• Mechanism-aware diagnostics using EP flux analysis reveal that pointwise error minimization can miss key nonlinear wave interactions, highlighting the need for physically consistent objectives.

Inductive Biases and Machine Learning Emulation of Sudden Stratospheric Warmings in Idealised GCM Simulations

Introduction

The paper "Investigating Inductive Biases for Machine Learning Emulation of Sudden Stratospheric Warmings in Idealised Isca Simulations" (2606.18857) systematically assesses how architectural inductive biases in ML-based emulators affect the accurate emulation of sudden stratospheric warming (SSW) dynamics in a dry, idealised general circulation model (GCM). The critical motivation is to determine which neural network design choices are most conducive to preserving physically meaningful vertical and wave activity features, especially under regimes characterized by strong stratospheric variability, such as SSWs.

Isca Design and SSW Simulation Protocol

A central methodological innovation is the use of paired Isca GCM datasets—one baseline configuration with suppressed stratospheric variability, and one forced configuration with a persistent zonal wave-2 heating perturbation to enable SSW-like disturbances. Both setups employ a dry dynamical core with 60 vertical pressure levels and no topography, moisture, or radiation scheme, thereby isolating the effect of planetary wave forcing and upward propagation in SSW formation.

Figure 1: Zonal wind at $60^\circ$N and 10 hPa in no-SSW and SSW-enabled configurations, with vortex morphology illustrated for both compact (no-SSW) and disturbed (split/displacement SSW) events.

The SSW-enabled regime demonstrates recurrent vortex disruptions as opposed to the stable polar vortex in the control case, providing a challenging prediction task for emulators that tests their ability to capture multi-scale vertical coupling.

Emulation Architectures and Comparative Framework

The study benchmarks a variety of neural architectures: 2D and 3D CNNs (including U-Net variants), 2D/3D transformer models (global and Swin-style windowed self-attention), and GNNs based on both regular grid and spherical mesh graphs. The architectural differentiation explicitly targets how vertical coupling and spatial interaction are embedded, with 3D models allowing direct convolution/attention across pressure levels, while others restrict explicit coupling to the horizontal plane or via message passing.

Figure 2: Overview illustrating how 2D models stack pressure levels in channels, 3D models operate jointly on latitude-longitude-pressure tensors, and GNNs encode spatial structure as graph connectivity with node/edge features.

Each model is trained for one-step prediction on the state transition operator $\mathcal{M}(\mathbf{x}_t)$, using temperature, zonal wind, and meridional wind fields as target variables. Hyperparameter tuning is strictly controlled and cross-validated.

Error Structure and Vertical Coupling Analysis

Error diagnostics reveal minimal performance separation among architectures in the no-SSW (dynamically quiescent) regime. All trained networks drastically outperform persistence baselines but yield similar MAEs for both horizontal and vertical wind/temperature across the stratified atmospheric column.

By contrast, in the SSW-enabled regime, models with explicit 3D vertical coupling (3D global transformers, 3D U-Nets, and simple 3D CNNs) sustain lower errors in the upper troposphere and stratosphere, particularly for the zonal wind and temperature. Models with only horizontal context (2D CNNs/transformers and GNNs) display significantly amplified error, especially in zones of active vortex disruption.

Figure 3: Mean absolute error as a function of pressure for temperature and zonal wind; 3D-structured models maintain superior performance where stratospheric dynamics dominate.

This result formally establishes that vertical inductive bias is crucial when emulating atmospheric phenomena characterized by strong vertical coupling and wave propagation, whereas the regression problem is less sensitive to architectural details in feature-suppressed regimes.

Mechanism-Aware Diagnostics: Eliassen–Palm (EP) Flux Evaluation

Pointwise accuracy does not guarantee physical consistency. The analysis of Eliassen–Palm (EP) fluxes, which encapsulate eddy momentum and heat transport, exposes a fundamental limitation: even the most accurate networks can fail to reproduce the correct wave–mean-flow interaction, especially in the high-latitude, upper stratosphere during SSW onset.

Figure 4: EP-flux vectors and divergence for true simulator trajectory and model predictions; divergence error is pronounced in high-latitude stratosphere and largely unaffected by the presence of SSWs in the dataset.

Decomposing the EP-flux divergence into zonal wavenumber components reveals that model errors are concentrated in the wave-1 structure of the stratospheric EP-flux, despite accurate reproduction of the wave-2 component (which is directly forced in the Isca configuration).

Figure 5: Zonal-wavenumber decomposition of EP-flux divergence; the dominant error modes appear in the stratospheric wave-1 component.

This diagnostic demonstrates that models can nominally fit the data, yet miss essential nonlinear stratospheric interactions—mainly wave–wave interactions generating wave-1 features from imposed wave-2 forcing. The result calls into question the sufficiency of pointwise error minimization for the mechanistic fidelity required in scientific emulation.

Spatial Error Structure and Additional Analysis

Supplemental figures provide granular diagnostic views, including latitude–pressure structure of errors and persistence benchmarks, and architecture-specific EP-flux diagnostics. These reinforce that vertical coupling is primarily beneficial in the high-latitude, mid-to-upper stratosphere under SSW forcing. The largest MAEs—and the most consequential divergence errors—are co-located with the regions critical for SSW dynamics.

Implications and Outlook

The empirical results demonstrate that explicit vertical coupling must be embedded in ML-based weather/climate emulators when the target phenomena entail multiscale wave activity and stratospheric-tropospheric interaction. This holds for both convolutional and self-attention-based networks, and GNNs can only partially ameliorate this via geometry-aware connectivity.

At a broader level, the study reveals the limitations of static, local MAE-based training and evaluation: reproducing the correct phase and amplitude of wave-mean-flow coupling may require mechanism-aware objectives, regularization, or even hybrid models that inject physical constraints directly. The separation in performance between pointwise and physically consistent error metrics is a critical insight for the future of ML-based atmospheric modeling.

Potential future directions include multi-step and event-conditioned training protocols, integrating additional mechanistic diagnostics (e.g., vortex geometry and wave activity measures), and scalable data augmentation leveraging controlled GCM experiments.

Conclusion

The paper provides a rigorous framework for dissecting the effect of architectural inductive biases on the emulation of SSW-like events in atmospheric GCMs. It establishes the necessity of explicit vertical coupling for capturing stratospheric dynamics under active wave forcing, identifies critical diagnostic gaps with pointwise metrics, and highlights methodological considerations indispensable for advancing scientific machine learning in geophysical applications. These findings have direct implications for both future AI system design in climate/atmospheric modeling and benchmark construction for mechanism-rich physical domains.