SDECast: Probabilistic Weather Forecasting in Continuous Time with Neural SDEs

AI summary

Overview

Research area: Machine learning for weather forecasting, specifically continuous-time probabilistic forecasting with Neural Stochastic Differential Equations (Neural SDEs).

Technical level: Advanced. The paper assumes familiarity with stochastic differential equations, variational inference, score functions, and neural weather prediction benchmarks.

Scope: The paper introduces SDECast, a simulation-free Neural SDE framework that learns stochastic atmospheric dynamics directly in physical space and generates probabilistic forecasts at arbitrary lead times, evaluated on a simulated geophysical flow (SQG) and on global ERA5 weather forecasting at hourly resolution.

What This Paper Is About

Most machine learning weather models produce forecasts by repeatedly applying a learned one-step predictor at a fixed time resolution. This is efficient for long-range prediction but accumulates severe errors when applied at shorter time steps, and it does not encode the locality and temporal continuity that differential-equation models of the atmosphere naturally express. The paper's goal is to learn the continuous-time stochastic dynamics of the atmosphere directly, so that a single trained model can generate forecast ensembles at any lead time.

Key Contributions

1. SDECast, a physical-space Neural SDE for weather. The authors adapt SDE Matching (a simulation-free objective originally formulated for latent SDEs) so that the learned SDE state is identified with the physical atmospheric state, making the drift and diffusion functions directly interpretable as dynamics models.

2. Continuous-time ensemble forecasting. Because forecasting amounts to simulating one learned SDE forward from the latest observed state, ensembles at arbitrary lead times are produced from a single dynamical model, with independent Wiener paths supplying the ensemble spread.

3. Simulation-free training at weather scales. The posterior marginal is parameterized directly as a transformation of Gaussian noise, so the training objective can be estimated by sampling a time and a noise vector instead of repeatedly integrating the SDE — removing a key obstacle to applying Neural SDEs in high dimensions.

4. Demonstration on two settings. SDECast is validated on the Surface Quasi-Geostrophic (SQG) system, where the true simulator is available for drift comparison, and on global ERA5 forecasting at hourly resolution with lead times up to five days.

Main Findings

• SQG drift recovery: Access to the numerical SQG simulator allowed the authors to compare the learned prior drift against the true drift and a six-hour finite-difference estimate. The learned drift closely resembles the true drift and gives a more faithful description of the instantaneous dynamics than the finite-difference estimate, showing that SDECast recovers continuous-time behavior from coarsely sampled observations.

• High frequencies are under-represented: The learned drift under-represents the highest spatial frequencies. The authors attribute this to coarse temporal sampling, under which the fastest dynamics are only weakly resolved, and note that SDECast may attribute part of that high-frequency variability to stochasticity rather than to the deterministic drift.

• Diffusion is essential: Removing the diffusion at inference time ("SDECast (deterministic)") substantially degrades forecast performance and produces overly smooth, mean-like trajectories, even though the underlying SQG simulator is deterministic.

• SQG skill versus the flow-matching baseline: Compared with an autoregressive conditional flow-matching (CFM) baseline, SDECast obtains higher RMSE and CRPS but maintains an almost ideal spread–skill ratio (SSR). SDECast generates forecasts at a substantially finer temporal resolution than it was trained on, whereas CFM is restricted to its trained resolution. In the reported SQG tables (means over 24 test trajectories), SDECast at 6 h lead reaches RMSE 1.280 against CFM's 0.832, and CRPS 0.717 against CFM's 0.441; SDECast's SSR stays between 0.552 and 0.981 across lead times, while CFM's ranges from 0.810 to 0.935.

• The learnable posterior matters: Replacing the learned posterior with a fixed linear posterior (removing the corrections μ_φ and σ_φ) substantially degrades SQG performance. For instance, SDECast (linear) has SSR values as low as 0.179 at 3 h and reaching 1.419 at 96 h, versus 0.552 and 0.966 for the full model.

• ERA5 skill: On global ERA5 forecasting at hourly resolution, SDECast produces stable and skillful probabilistic forecasts throughout the evaluation horizon. Despite being trained on the same hourly data, it substantially outperforms the one-hour autoregressive baseline (AR-1h) at longer lead times, demonstrating that its continuous-time formulation supports high temporal resolution without the severe error accumulation of repeated one-hour autoregressive rollouts.

• Position among baselines: SDECast improves on ClimODE and is comparable to DYffusion, but remains behind the specialized ARCI-24-1h model. The authors suggest this may be partly explained by model inputs: SDECast and DYffusion assume Markovian dynamics and propagate a single state, while ARCI-24-1h conditions on the two previous states.

• Representative ERA5 numbers: In the latitude-weighted RMSE table, z500 at 6 h lead is 66.7 for SDECast versus 61.7 for ARCI-24-1h, 68.5 for AR-1h, 93.9 for DYffusion, 102.9 for ClimODE, and 235.3 for persistence. At 24 h, z500 RMSE is 172.0 for SDECast versus 142.9 for ARCI-24-1h, 206.7 for AR-1h, 186.5 for DYffusion, 193.4 for ClimODE, and 603.1 for persistence. In the latitude-weighted CRPS table, z500 at 24 h is 86.1 for SDECast versus 70.6 for ARCI-24-1h, 121.7 for AR-1h, 96.7 for DYffusion, and 352.6 for persistence. ClimODE is deterministic and therefore absent from the CRPS table.

• Spatially varying uncertainty: The diffusion coefficient is allowed to vary spatially on ERA5, and the learned volatility magnitude exhibits meaningful spatial variation — for example, representing differences in forecast uncertainty between land and ocean.

• Qualitative trajectory behavior: Visualizing SDECast trajectories for a single spatial location, the generated trajectories follow the overall evolution of the ground truth while exhibiting meaningful ensemble variation at hourly resolution; the figure's mean and standard deviation are computed from a 24-member rollout.

Methodology in Plain English

The authors start from a prior SDE — a drift network that describes the systematic tendency of the atmosphere, plus a diffusion network that injects randomness: dz = h_θ(z,t) dt + g_θ(z,t) dW. This is the model used to generate forecasts. To train it without simulating it, they introduce a second, "posterior" SDE that shares the same diffusion coefficient but whose drift also depends on the observed trajectory; the difference between the two drifts defines a control signal.

Training minimizes two things jointly: the size of that control signal (a dynamics loss) and the error in reconstructing the observations (a reconstruction loss). The trick that makes this simulation-free is that the posterior distribution at any time is written directly as a differentiable, invertible function of Gaussian noise, rather than as the solution of an SDE. Sampling a time, an interval index, and a noise vector is then enough to evaluate the loss; automatic differentiation supplies all needed quantities, including the score of the posterior.

To make the SDE's state mean the physical state, the authors use a very concentrated observation model (a Gaussian with a fixed, very small standard deviation δ) and restrict the posterior so that, between two consecutive observations, it depends only on the two observations bracketing that interval. The resulting posterior interpolates linearly between the two endpoints, plus two learnable corrections — one for the mean and one for the scale — both multiplied by α(1−α), where α is normalized time within the interval. That factor forces the corrections to vanish at the observation times, anchoring the posterior to the observations. Because both corrections vanish there, the reconstruction loss reduces to a constant that depends only on δ and the state dimension, so it can be dropped during optimization.

At inference, the prior SDE is initialized from the most recent observed physical state and simulated forward; different Wiener paths give different ensemble members.

Experimental setup: On SQG (two surface-potential-temperature channels on a 64×64 doubly periodic grid, no conditioning, diffusion parameterized as a single scalar), models were trained on six-hour snapshots but evaluated at one-hour resolution, using 200 Euler–Maruyama steps at inference. On ERA5 (WeatherBench-1 copy at 5.625°, a 32×64 latitude–longitude grid, variables z500, t850, t2m, u10, v10), drift and diffusion are conditioned on time of day and time of year, longitude wraps circularly, and inference uses 64 Euler–Maruyama steps. Both models use AdamW with linear warmup for 500 steps followed by

Authors’ abstract

Existing machine learning weather forecasting models typically generate forecasts through autoregressive rollouts at a fixed temporal resolution. While highly efficient for long-range prediction, this formulation can suffer from severe error accumulation when used with shorter time steps and does not explicitly encode the locality and temporal continuity of atmospheric dynamics. To address these limitations, we introduce **SDECast**, a Neural Stochastic Differential Equation (SDE) framework for continuous-time probabilistic weather forecasting. SDECast extends SDE Matching to learn stochastic dynamics directly in physical space, without requiring repeated SDE simulation during training. On a simulated geophysical flow, we show that SDECast recovers meaningful drift dynamics and faithfully reproduces the underlying continuous-time behavior. We then demonstrate its scalability to global weather forecasting at hourly resolution, where SDECast produces skillful probabilistic forecasts for lead times of up to five days.