Probabilistic Emulation of a Global Climate Model with Spherical DYffusion

📄 Read the paper:

https://bit.ly/SphericalDY

Motivation

• Climate models are crucial for understanding long-term Earth system evolution.
• But, they are computationally very expensive…

→ Limited exploration of mitigation and adaptation strategies.

→ Large ensembles for reducing uncertainties are prohibitive to run.

Solution: Train a cheap-to-run ML model to emulate (parts of) a climate model!

2

Motivation

• Climate models are crucial for understanding long-term Earth system evolution.
• But, they are computationally very expensive…

→ Limited exploration of mitigation and adaptation strategies.

→ Large ensembles for reducing uncertainties are prohibitive to run.

Solution: Train a cheap-to-run ML model to emulate (parts of) a climate model!

3

Weather forecasting vs. Climate Modeling

• Rapid progress: ML models competitive or better than operational physics-based ones!
• Short-term time-specific focus (+ ML models often long-term unstable).
• Initial-value problem.

• Fairly little work, especially for fully data-driven temporal modeling.
• Stable & accurate reproduction of long-term statistics necessary.
• Boundary-condition problem.

4

Our Approach: Spherical DYffusion

• Build upon the pioneering ACE climate emulator [1]
• Train on coarse-res simulated atmospheric data from FV3GFS, the US primary global forecast model.
• 34 variables: Temperature, humidity, winds, surface pressure.
• Forced by solar radiation and specified sea surface temperatures.
• Evaluate 10-year-long simulation on a 6-hourly interval (rollout of 14600 time steps!).
• Spherical Fourier Neural Operator (SFNO) architecture [2].
• Integrate SFNO into the dynamics-informed diffusion model (DYffusion) framework [3]
• Improved modeling of physical fields on a sphere versus U-Net from vanilla DYffusion.
• Enables ensemble climate simulations naturally.
• Guarantees high efficiency versus other diffusion models and only 3x more than ACE.

​

​

5

[1] Watt-Meyer et al., “ACE: A fast, skillful learned global atmospheric model for climate prediction” (NeurIPS CCAI 2023)

[2] Bonev et al., “Spherical fourier neural operators: Learning stable dynamics on the sphere” (ICML 2023)

[3] Ruhling Cachay et al., “DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting” (NeurIPS 2023)

Forecasting with DYffusion at inference time

6

Forecasting with DYffusion at inference time

7

DYffusion’s Key Idea: Couple diffusion steps with physical time steps

​

8

Inference & Training

(including forcings)

9

SFNO network topology

(single block)

10

SFNO network topology

(single block)

11

The new part for diffusion modeling

Inference including forcings

12

Results: Visualization of two ensemble members

13

Results: 10-yr time-mean RMSEs

14

Results: 10-yr time-mean RMSEs

15

Results: 10-yr time-mean RMSEs (showing more fields)

16

Results: Global bias maps of total water path time-mean

17

Results: 100-year rollouts

18

Computational comparison

19

Results: Encouraging

climate variability

20

Results: Encouraging ensemble weather forecasts

21

Conclusion

• Low climate biases, reaching reference “noise floor” within ~20-60% range
• Ensembling our method effectively reduces climate biases & improves weather forecasts
• Encouraging climate variability of emulated simulations
• Relatively low slowdown at inference time vs. ACE (~2-3x)

→ Using SFNO architecture absolutely key for successful climate emulation!

Outlook/Future work:

• Expand dataset to include climate change scenarios & forcings.
• Emulate or couple with other Earth system’s (ocean, sea-ice, land, …)
• How to correlate climate biases with performance on short rollouts?

22

Thanks for listening!

Feel free to reach out :)

23

📄 Read the paper:

https://bit.ly/SphericalDY

24

More slides

25

Computational comparison

26

Key Idea: Couple diffusion steps with physical time steps

​

27

Forecasting with DYffusion at inference time

28

Forecasting with DYffusion at inference time

29

Results: Encouraging

climate variability

30

Results: Encouraging ensemble weather forecasts

31

Results: SFNO-DYffusion leads to more effective ensembles

32

Methods: Training

​

33

Methods: Sampling

​

34

Motivation

Existing ML models for high-dimensional spatiotemporal forecasting tend to be:

• Deterministic → blurry, unrealistic long-range forecasts

35

Motivation

Existing ML models for high-dimensional spatiotemporal forecasting tend to be:

• Deterministic → blurry, unrealistic long-range forecasts
• Autoregressive → Inference differs from training, errors accumulate, and rollouts become long-term unstable

36

Spherical Fourier Neural Operators: Learning stable dynamics on the sphere, Bonev et al., ICML 23, https://arxiv.org/abs/2306.03838

​

​

Results: Competitive probabilistic rollouts

37

Results: Temporal super-resolution (8x)

38

Summary

• First study on diffusion models for spatiotemporal forecasting
• Novel adaptation of diffusion models to ensemble-based probabilistic forecasting
• Effective training approach for multi-step and long-range forecasting with low memory needs
• Competitive performance on probabilistic evaluations for forecasting complex dynamics in sea surface temperatures, Navier-Stokes flows, and spring mesh systems

+ more details, results and extensive ablations in our paper!

39

Thanks for listening!

Feel free to reach out :)

40

https://github.com/Rose-STL-Lab/dyffusion

Paper: arxiv.org/abs/2306.01984

Blog: https://salvarc.github.io/blog/2023/dyffusion

​