Spatiotemporal Pyramid Flow Matching for Climate Emulation

TL;DR

The paper introduces Spatiotemporal Pyramid Flows (SPF), a hierarchical generative model for efficient, multi-scale climate emulation that outperforms existing methods and is scalable to large datasets like ClimateSuite.

Contribution

We propose SPF, a novel flow matching approach that models climate data hierarchically across space and time, enabling fast, stable, and scalable climate emulation.

Findings

SPF outperforms baseline models at yearly and monthly scales.

SPF enables fast sampling at multiple temporal levels.

ClimateSuite is the largest Earth system simulation dataset to date.

Abstract

Generative models have the potential to transform the way we emulate Earth's changing climate. Previous generative approaches rely on weather-scale autoregression for climate emulation, but this is inherently slow for long climate horizons and has yet to demonstrate stable rollouts under nonstationary forcings. Here, we introduce Spatiotemporal Pyramid Flows (SPF), a new class of flow matching approaches that model data hierarchically across spatial and temporal scales. Inspired by cascaded video models, SPF partitions the generative trajectory into a spatiotemporal pyramid, progressively increasing spatial resolution to reduce computation and coupling each stage with an associated timescale to enable direct sampling at any temporal level in the pyramid. This design, together with conditioning each stage on prescribed physical forcings (e.g., greenhouse gases or aerosols), enables efficient, parallel climate emulation at multiple timescales. On ClimateBench, SPF outperforms strong flow matching baselines and pre-trained models at yearly and monthly timescales while offering fast sampling, especially at coarser temporal levels. To scale SPF, we curate ClimateSuite, the largest collection of Earth system simulations to date, comprising over 33,000 simulation-years across ten climate models and the first dataset to include simulations of climate interventions. We find that the scaled SPF model demonstrates good generalization to held-out scenarios across climate models. Together, SPF and ClimateSuite provide a foundation for accurate, efficient, probabilistic climate emulation across temporal scales and realistic future scenarios. Data and code is publicly available at https://github.com/stanfordmlgroup/spf .