Spatiotemporal Pyramid Flow Matching for Climate Emulation

AI summary

Overview

Research area: Generative machine learning for climate emulation — specifically a flow-matching method (categorized on arXiv as computer vision, cs.CV) that replaces or accelerates physics-based Earth system models.

Technical level: Advanced. The paper assumes familiarity with diffusion/flow matching, probability paths, and transformer-based generative backbones (DiT), though the high-level idea can be followed without the derivations.

Scope: The paper introduces Spatiotemporal Pyramid Flows (SPF), a flow-matching model that cascades across both space and time to sample climate states at decadal, yearly, and monthly timescales in parallel, and releases ClimateSuite, a dataset of 33,739 simulation-years from ten Earth system models that includes the first simulations of stratospheric aerosol injection (SAI).

What This Paper Is About

Earth system models are the main tools for projecting how climate will respond to emissions and interventions, but they are computationally expensive: a single 10-year trajectory with a state-of-the-art emulator takes nearly three hours, and long rollouts, ensembles, and forcing sweeps are prohibitive even on supercomputers. Prior generative emulators rely on weather-scale autoregression, which compounds small errors into long-term drift and is slow over climate horizons, and they have yet to demonstrate stable rollouts under nonstationary forcings. The goal of this paper is a probabilistic emulator that samples climate states at multiple timescales directly and in parallel, conditioned on physical forcings such as greenhouse gases and aerosols, without step-by-step autoregressive rollout.

Key Contributions

1. A new class of flow matching models, Spatiotemporal Pyramid Flows (SPF). SPF applies a pyramid of resolution-cascaded flows across both space and time, refining noised latent trajectories conditioned on physical forcings so that climate model outputs can be sampled at arbitrary future times without autoregressive simulation.

2. A generalization of prior pyramid flow methods. SPF supports arbitrary resolution resampling between stages (unlike earlier purely spatial pyramids with homogeneous factors) and enables generation of clean samples from any timescale in the pyramid, not just the finest.

3. ClimateSuite, described as the largest dataset for climate model emulation to date. It comprises more than 33,000 simulation-years (reported as 33,739 simulation-years in the data section) spanning 276 non-SAI simulations from 10 ESMs and 39 SAI simulations, and is described as the first ML dataset to include simulations of climate interventions.

4. Empirical results on ClimateBench. SPF is reported to achieve superior accuracy and inference efficiency compared to deterministic baselines, pre-trained models, and flow matching approaches, and training on ClimateSuite is reported to give good generalization to emissions and intervention scenarios across climate models and to improve ClimateBench performance after fine-tuning.

Main Findings

• Yearly timescale accuracy on ClimateBench: SPF achieves the lowest CRPS among all 100M and 200M parameter models at the yearly timescale. With 100M parameters it records CRPS 0.238, RMSE 0.565, and bias 0.004; with 200M parameters it records CRPS 0.222, RMSE 0.511, and bias -0.049. At 200M parameters it has the second lowest RMSE, behind PixelFlow (0.224 CRPS, 0.504 RMSE). With 100M parameters its RMSE is slightly worse than single-year PixelFlow (0.247 CRPS, 0.549 RMSE) and ClimaX (0.347 CRPS, 0.546 RMSE). The authors note ClimaX is pre-trained on a large climate dataset while SPF is trained only on the ClimateBench training set.

• Monthly timescale accuracy: SPF with 100M parameters achieves the best CRPS (0.462), RMSE (1.100), and bias (-0.043) among the compared models; the 200M parameter model achieves the lowest RMSE (1.060) and close to best CRPS (0.453). For comparison, PyramidalFlow at 200M parameters records monthly CRPS 0.473, RMSE 1.083; at 100M parameters it records 0.501 and 1.137.

• Sampling efficiency: The 100M and 200M SPF variants achieve 28× and >15× speedups respectively compared to the autoregressive PyramidalFlow model, reducing sampling time to a few seconds. SPF samples a 10-year trajectory in 3 seconds (yearly) and 6 seconds (monthly) at 100M parameters, and 6 seconds (yearly) and 11 seconds (monthly) at 200M parameters, measured on a single NVIDIA RTX 6000. At the monthly timescale, sampling is >3× faster than PyramidalFlow at 200M parameters and 6× faster at 100M parameters.

• Pyramid design ablation: SPF's three-stage decadal → yearly → monthly design (DYM-Any) achieves CRPS 0.453 and RMSE 1.060 on ClimateBench SSP2-4.5 at the monthly timescale, outperforming a version trained only for monthly samples (DYM-Monthly, 0.453 CRPS, 1.064 RMSE) and two five-stage variants with decoupled spatial and temporal transitions: DYMMM (0.474 CRPS, 1.087 RMSE) and DYYMM (0.463 CRPS, 1.085 RMSE). An alternating pyramid starting with spatial then temporal transitions was attempted but ran out of GPU memory.

• Temporal funneling saves compute: Funneling to a single frame at each stage yields a 10× and 12× reduction in GPU memory and FLOPs per stage, since the temporal windows contain 10 and 12 frames respectively.

• Scaling to ClimateSuite: The paper states that the scaled SPF model generalizes well to held-out scenarios across climate models. The specific ClimateSuite evaluation table is truncated in the provided content, so the full set of fine-tuned ClimateBench and cross-model numbers is not reported here; the visible portion shows a 200M parameter ClimateSuite-pretrained entry with yearly CRPS 0.222 and RMSE 0.51.

Methodology in Plain English

Flow matching trains a model to transform random noise into data samples by learning a velocity field that moves noise toward the target along a path. Earlier "pyramid flow" methods speed up image and video generation by splitting this path into stages: early stages work at low spatial resolution, later stages refine at higher resolution, with one shared network handling all stages.

SPF extends this idea in three ways. First, it adds a temporal dimension to the pyramid, so stages can change temporal resolution as well as spatial resolution. The authors derive a general continuity correction that keeps the probability distributions aligned at stage boundaries for arbitrary resampling factors in height, width, and time — meaning the pyramid can follow natural climate timescale hierarchies such as ×10 from decadal to yearly and ×12 from yearly to monthly. Second, because generating a long, high-resolution sequence at once is memory-prohibitive, they introduce temporal funneling: slicing the latent down to a single time frame before upsampling, which is valid because any subset of a Gaussian latent is still Gaussian and therefore the same rescaling-and-renoising rule applies. Third, they train with randomly sampled "temporal paths" — controlled by Bernoulli indicators (δ₁ sampled with probability 2/3, and δ₂ given δ₁=1 with probability 1/2) — so that a single network can emit clean samples at any timescale in the pyramid, rather than only the finest.

The concrete configuration uses K=3 stages (decadal, yearly, monthly), with spatial resampling factors of 2 and temporal factors of 10 and 12. The backbone is the MM-DiT architecture from SD3, at two scales (100M parameters with 2 layers, 200M with 4 layers), with sinusoidal spatial position encodings and 1D Rotary Position Embeddings for time. Outputs are patchified at 8×8 and input forcings at 16×16, and sequences are pack-packed after patchifying so mixed resolutions fit in a batch. Training and inference fit on a single NVIDIA RTX A4000 with 16GB of VRAM.

Evaluation uses RMSE, bias, and CRPS, all weighted by grid cell area. Probabilistic models are evaluated with 5 samples and 90 inference steps split equally across stages using an Euler solver. On ClimateBench, models train on historical, SSP1-2.6, and SSP5-8.5 scenarios, validate on SSP3-7.0, and test on SSP2-4.5.

Why This Matters

Impact on research. The paper reframes climate emulation from a sequential roll-out problem to a parallel, multi-timescale generation problem. If a single model can produce physically coherent samples at decadal, yearly, and monthly resolutions, it changes how ensembles for uncertainty quantification are constructed: instead of spending compute on millions of sequential steps, researchers can sample many plausible futures at the resolution that a given study actually needs. The release of ClimateSuite (33,739 simulation-years, ten ESMs, first inclusion of SAI experiments) also lowers the barrier to building "superemulators" that span multiple Earth system models.

Real-world applications (drawn from the paper's stated use cases):

• Annual indicators for integrated assessment models and policy assessments, which need coarse timescale outputs rather than high-frequency fields.
• Monthly fields for sectoral impact studies that require seasonal-to-annual structure.
• Exploration of climate intervention scenarios: ClimateSuite's 39 SAI experiments with stratospheric aerosol optical depth as an added input allow emulation of aerosol-injection responses.
• Ensemble-based uncertainty quantification for regional risk assessment, where diverse probabilistic samples under the same forcing are needed.

Industry relevance. The speed claims matter directly for any workflow that currently waits hours for a trajectory. Sampling a 10-year trajectory in seconds rather than nearly three hours opens the door to interactive scenario comparison, large forcing design sweeps, and rapid iteration on emissions pathways. Sectors that consume climate projections but lack supercomputer access — risk analytics, insurance, energy planning, agriculture — are the natural beneficiaries of a cheap probabilistic emulator, and the code and data are publicly available at github.com/stanfordmlgroup/spf.

Future Directions

• Closing the gap with single-timescale models on mean prediction. At 200M parameters SPF trails PixelFlow on RMSE (0.511 vs 0.504) at the yearly timescale; narrowing this while keeping multi-timescale capability is an open engineering question.

• Understanding why multi-timescale training helps. The ablation shows the multi-timescale model (DYM-Any) has slightly better RMSE than the monthly-only variant (1.060 vs 1.064) but the same CRPS, so the source of this benefit is not fully explained.

• Extending the pyramid beyond three stages. Two five-stage variants underperformed the three-stage design, and an alternating spatial-then-temporal pyramid exceeded GPU memory. Whether deeper or differently ordered pyramids can help remains unresolved.

• Broadening generalization guarantees. The held-out UKESM1-0-LL SAI experiment tests transfer to a climate model whose SAI experiments were excluded from training. How well the approach generalizes beyond the ten included ESMs, and how stable rollouts are under stronger nonstationary forcings, are questions the paper frames as ongoing.

Target Audience

This paper is most useful to machine learning researchers working on generative models for scientific simulation, particularly those interested in flow matching, cascaded generation, or efficient long-sequence sampling. It is also aimed at climate scientists and emulator developers who need probabilistic, multi-timescale surrogates for Earth system models, and at applied groups — integrated assessment modelers, impacts researchers, and intervention-studies researchers — who consume annual or monthly climate fields and are constrained by the cost of physics-based simulations. Readers without a background in diffusion or flow matching will find the mathematical sections demanding, but the motivation, dataset description, and benchmark comparisons are accessible.

Authors’ 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 .