Probablistic Emulation of a Global Climate Model with Spherical DYffusion

Authors: Salva Rühling Cachay, Brian Henn, Oliver Watt-Meyer, Christopher S. Bretherton, Rose Yu

NeurIPS 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental Here, we present the first conditional generative model that produces accurate and physically consistent global climate ensemble simulations by emulating a coarse version of the United States primary operational global forecast model, FV3GFS. Our model integrates the dynamics-informed diffusion framework (DYffusion) with the Spherical Fourier Neural Operator (SFNO) architecture, enabling stable 100-year simulations at 6-hourly timesteps while maintaining low computational overhead compared to single-step deterministic baselines. The model achieves near gold-standard performance for climate model emulation, outperforming existing approaches and demonstrating promising ensemble skill.
Researcher Affiliation Collaboration Salva R uhling Cachay UC San Diego Brian Henn Allen Institute for AI Oliver Watt-Meyer Allen Institute for AI Christopher S. Bretherton Allen Institute for AI Rose Yu UC San Diego
Pseudocode Yes C.1 Training and inference pseudocode. In Algorithms 1 and 2 we provide the procedures used to train and sample from our proposed method, respectively.
Open Source Code Yes Code is available at https://github.com/Rose-STL-Lab/spherical-dyffusion
Open Datasets Yes The final training and validation data can be downloaded from Google Cloud Storage following the instructions of the ACE paper at https://zenodo.org/records/10791087. The data are licensed under Creative Commons Attribution 4.0 International.
Dataset Splits Yes We train on 100 years of simulated data from FV3GFS, and evaluate the models on how well they can emulate a distinct 10-year-long validation simulation (i.e. H = 14600 = 10 365 4).
Hardware Specification Yes All models were trained on A6000 GPUs using distributed training on 2 up to 8 GPUs... For a fair inference runtime comparison measuring the wall clock time needed to simulate 10 years (i.e. one full validation rollout), we run all deep-learning baselines on one A100 GPU. We also include the runtime for the physics-based FV3GFS climate model which was run on 96 cores (24 cores for the 2 coarser version) of AMD EPYC 7H12 processors.
Software Dependencies No We use Py Torch Lightning [16] and Weights & Biases [6] as part of our software stack.
Experiment Setup Yes To fairly compare against the deterministic SFNO model from [67], we use exactly the same hyperparameters for training the interpolator and forecasting networks for our method, as described in Table 7. For the stochastic version of ACE, ACE-STO, we re-train ACE from scratch with the only difference being that we use a dropout rate of 10% for the MLP in the SFNO architecture. We train the stochastic interpolator model, SFNOϕ, in our method using the same dropout rate. Both of these stochastic models are run using MC dropout (i.e. enabling the dropout layers at inference time). For our interpolator network, we also use a 10% rate for stochastic depth [28], which is also enabled at inference time.