Tyche: One Step Flow for Efficient Probabilistic Weather Forecasting

2026-05-07

AI Summary

This work addresses the high computational cost of existing diffusion-based probabilistic weather forecasting, which relies on multi-step sampling and struggles to balance long lead times with large ensemble sizes. The authors propose Tyche, a single-step conditional flow model that maps Gaussian noise directly to future weather states via a target-aware mean velocity field, enabling well-calibrated probabilistic forecasts with only one function evaluation. To facilitate learnable single-step transport in high-dimensional flow fields, they introduce a Jacobian-vector product (JVP) regularization loss and enhance ensemble reliability through rolling fine-tuning and a curriculum CRPS calibration strategy. Parameterized by an isotropic Swin Transformer, Tyche achieves or surpasses the forecast skill and calibration performance of both state-of-the-art diffusion models and the operational ECMWF IFS ensemble system on ERA5 data, using just a single forward pass.

Abstract

Probabilistic weather forecasting requires not only accurate trajectories, but calibrated distributions over plausible atmospheric futures. Recent data-driven systems have achieved remarkable deterministic skill, and diffusion-based ensemble forecasters have substantially improved sample realism and uncertainty quantification. However, their inference cost scales with forecast horizon, ensemble size, and the number of denoising steps required for each transition, making large operational ensembles expensive. To address this, we present Tyche, a one-step conditional flow model for efficient probabilistic weather forecasting. Tyche models the conditional forecast distribution with a destination-aware average-velocity flow that maps Gaussian noise directly to future weather states in a single function evaluation (1-NFE). To make this one-step transport learnable in high-dimensional geophysical fields, we derive a JVP-regularized rectification objective that enforces temporal self-consistency across source and destination flow timesteps without explicitly forming Jacobians. The transport field is parameterized by an isotropic Swin-style transformer that preserves fine-scale spatial structure while remaining scalable on global grids. To improve ensemble reliability under autoregressive forecasting, we further introduce a rollout-based finetuning stage with curriculum CRPS calibration supervision. Experiments on ERA5 at 1.5$^\circ$ and 6-hour resolution show that our Tyche, using merely a single NFE, matches or exceeds the forecast skill and calibration of state-of-the-art multi-step generative baselines and the operational ECMWF IFS ensemble.

Problem

• probabilistic weather forecasting
• diffusion models
• inference cost
• ensemble forecasting
• computational efficiency

Innovation

• one-step flow
• probabilistic weather forecasting
• JVP-regularized rectification
• conditional flow model
• ensemble calibration