PaperlayerView exact revision on arXiv

Source-linked AI summary

Swift: An Autoregressive Consistency Model for Efficient Weather Forecasting

Jason Stock, Troy Arcomano, Rao Kotamarthi

arXiv:2509.25631v1cs.LG

TL;DR

Diffusion-based forecasting offers a grounded approach to probabilistic weather prediction, but iterative inference is costly for long-lead S2S forecasts. Swift uses a single-step consistency model finetuned autoregressively with CRPS, producing skillful 6-hourly forecasts stable for up to 75 days while running up to 39× faster than diffusion baselines.

  • Problem

    Diffusion models provide a grounded framework for probabilistic weather forecasting, but their iterative inference is costly for long-lead S2S applications where calibration is important.

    Source: §1 · p. 2

  • Method

    Swift autoregressively finetunes a single-step consistency probability-flow model using a CRPS objective.

    Source: §1 · p. 2

  • Results

    Swift produces skillful 6-hourly forecasts stable for up to 75 days while running up to 39× faster than diffusion baselines.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The results mark a step toward machine-learning ensemble forecasting spanning medium-range to seasonal scales without substantially more compute than deterministic models.

    Source: §5 · p. 9

    Comparisons use fewer initial conditions and ensemble members because of compute and storage constraints, and forecasts remain underdispersive.

    Source: §5 · p. 9

Abstract

from arXiv · show

Diffusion models offer a physically grounded framework for probabilistic weather forecasting, but their typical reliance on slow, iterative solvers during inference makes them impractical for subseasonal-to-seasonal (S2S) applications where long lead-times and domain-driven calibration are essential. To address this, we introduce Swift, a single-step consistency model that, for the first time, enables autoregressive finetuning of a probability flow model with a continuous ranked probability score (CRPS) objective. This eliminates the need for multi-model ensembling or parameter perturbations. Results show that Swift produces skillful 6-hourly forecasts that remain stable for up to 75 days, running $39\times$ faster than state-of-the-art diffusion baselines while achieving forecast skill competitive with the numerical-based, operational IFS ENS. This marks a step toward efficient and reliable ensemble forecasting from medium-range to seasonal-scales.

1 Introduction

The introduction frames reliable, calibrated ensembles at long leads as a gap between efficient deterministic models and costly diffusion forecasts; Swift pairs single-step consistency modeling with autoregressive CRPS finetuning.

1 Introduction

  • Deterministic weather models can match or exceed operational systems efficiently, but spectral biases and weak sensitivity to initial-condition perturbations limit reliable, calibrated ensembles.Diffusion models are presented as a grounded formulation for addressing these issues.Source: §1 · pp. 1–2
  • 20–40 neural function evaluations per forecast step make diffusion forecasts costly at seasonal leads, while 12- or 24-hour steps sacrifice temporal fidelity.Alternative noise-perturbation and multi-model ensembles add maintenance overhead and may remain unstable or artifact-prone by day 15.Source: §1 · p. 2
  • Swift autoregressively finetunes a single-step consistency model with CRPS, preserving diffusion properties while extending forecasts from medium range to seasonal scales with temporal fidelity.The passage attributes accurate, well-calibrated ensembles to the efficiency that makes finetuning possible.Source: §1 · p. 2

2 Background

The background introduces probability-flow models, consistency models, and CRPS as the foundations for Swift’s forecasting and ensemble training.

2.1 Diffusion Models and Flow Matching

  • Diffusion and flow-matching formulations share a probability-flow ODE sampling framework, but sampling requires an iterative numerical solver.TrigFlow unifies EDM and flow matching with a v-prediction parameterization and maps noise schedules onto a common trigonometric trajectory.Source: §2.1 · pp. 2–3

2.2 Continuous-time Consistency Models

  • Consistency models predict clean data from noisy samples along the probability-flow trajectory, enabling direct one-step estimates instead of iterative ODE solving.Their training matches predictions at adjacent trajectory steps; the continuous objective depends on the trajectory tangent, with a warmup used to stabilize training.Source: §2.2 · p. 3

2.3 Continuous Ranked Probability Score (CRPS)

  • CRPS is a strictly proper score measuring the discrepancy between a forecast CDF and an observation, with a fair finite-ensemble estimate.Its first term rewards members close to the observation, while the second accounts for ensemble dispersion by correcting for self-similarity.Source: §2.3 · p. 4

3 Methodology

Swift models global atmospheric transitions from ERA5 and trains a consistency model through pretraining followed by CRPS-based autoregressive finetuning, using a conditional Swin architecture.

3.1 Task and Dataset

  • The task predicts atmospheric states at 6-hour intervals from ERA5, using dynamic 6-, 12-, and 24-hour intervals and a 1979–2018/2019/2020 train-validation-test split.The data are downsampled to 1.40625° resolution; the model predicts four surface variables and five atmospheric variables at 13 pressure levels.Source: §3.1 · p. 4

3.2 Training Objectives

  • Two-stage training first pretrains diffusion and consistency models, then finetunes the consistency model as Swift with an unbiased, fair CRPS objective over autoregressive rollouts.Finetuning uses a curriculum of K = {1, 2}, then K = 3, then K = {4, 8} steps, with loss on the final step.Source: §3.2 · p. 5

3.3 Model Architecture

  • A conditional non-hierarchical Swin transformer takes noisy samples, spatial conditions, noise level, and data time deltas as inputs; shifting windows expand its receptive field.The model has 225M parameters, with 12 attention blocks and 12 heads per block.Source: §3.3 · p. 6

3.4 Implementation Details

  • Training uses 15M-image pretraining and 5M-image finetuning, with Muon for consistency pretraining and AdamW for diffusion and finetuning.A cosine schedule warms up for 2M images, anneals until 15M, then holds η = 1e-5 during finetuning; inference uses EMA parameters and training uses a log-uniform noise schedule.Source: §3.4 · pp. 6–7

4 Experimental Results

The experiments assess medium-range forecast skill, long-term stability for subseasonal-to-seasonal forecasting, and case studies of extreme weather and seasonal trends.

4.1 Global Forecast Skill

  • 30× empirically faster wall-clock time; Swift remains competitive with IFS ENS for most variables, outperforms Swift-B, and marginally underperforms GenCast.It is also competitive with the authors’ diffusion model, despite using 12 ensemble members versus GenCast’s 50; Swift is underdispersive.Source: §4.1 · p. 7

4.2 Long-term Stability

  • At 60 days, Swift reproduces realistic equatorial wave propagation and land–ocean boundaries; q700 fields retain coherent structures and realistic spectra, and 14- and 45-day forecasts remain sharp.Additional Hovmöller diagrams support stability to 75 days.Source: Figure 6 · p. 7, §4.2 · p. 8

4.3 Case Studies

  • For Hurricane Laura, ensemble spread increases when the track turns northward, and several members capture the observed track; individual fields show a tropical cyclone, atmospheric river, and intertropical convergence zone.The 48 forecasts were initialized five days before landfall.Source: §4.3 · p. 8, Figure 7 · p. 8
  • Across seasonal 75-day rollouts, Northern Hemisphere temperatures warm through boreal spring and cool into autumn while the Southern Hemisphere follows the opposite progression.Forecasts remain stable and phase-aligned with ERA5 across initial conditions, with a tendency to slightly understate amplitude.Source: §4.3 · pp. 8–9, Figure 8 · p. 9

5 Conclusion

Swift produces skillful 6-hourly forecasts stable for up to 75 days, with diffusion-baseline speedups of up to 39×. The paper notes evaluation and calibration limitations and identifies further modeling directions.

5 Conclusion

  • Up to 39× faster than diffusion baselines, Swift achieves skillful 6-hourly forecasts stable for up to 75 days.Case studies capture diverse, realistic ensembles during extreme events and model seasonal cycles.Source: §5 · p. 9
  • Forecasts are underdispersive; finetuning improves long-term spread at the expense of short-term SSR on days 0–4.Comparisons also used reduced numbers of initial conditions and ensemble members because of computational and storage constraints.Source: §5 · p. 9
  • Future work includes consistency distillation with a larger, high-resolution model, classifier-free guidance, and benchmarking alternative backbones.The authors identify operational deployment as a motivation for the larger-model exploration.Source: §5 · p. 9

A Full Field Visualizations

The appendix provides visualizations of a seven-day forecast error and forecast power spectra.

A Full Field Visualizations

  • Figure 10 shows forecast error at a seven-day lead time for a forecast initialized on 01-01-2020T0z.Source: Figure 10 · p. 14

B Additional Medium-Range Results

The appendix lists additional medium-range visualizations of power spectra, ensemble mean RMSE, CRPS, and spread-to-skill ratio.

B Additional Medium-Range Results

  • Figure 11 shows power spectra compared with target ERA5 for generated 15-day forecasts.The spectra are averaged over 32 initial conditions and 12 ensembles for each variable between 6 and 360 hours.Source: Figure 11 · p. 15
  • Figure 12 presents latitude-weighted ensemble mean root-mean-squared error (RMSE).Source: Figure 12 · p. 15
  • Figure 13 presents latitude-weighted continuous ranked probability score (CRPS).Source: Figure 13 · p. 16
  • Figure 14 presents latitude-weighted spread-to-skill ratio (SSR).Source: Figure 14 · p. 16

C Additional Hovmöller Diagrams

Additional Hovmöller diagrams show 75-day tropical u850 anomalies across four year-initialized panels.

C Additional Hovmöller Diagrams

  • Figure 15 adds tropical u850 anomaly Hovmöller diagrams averaged between 10°N/S over 75 days, with year initials in panels (a)–(d).Source: Figure 15 · p. 17

Related papers

Loading 2509.25631v1…