Can a Deterministic Weather Foundation Model Be Turned Into a Generative Ensemble Forecaster?

Xaurora converts a deterministic prior into a stochastic ensemble

Most atmospheric foundation models are deterministic: given the same initial conditions, they return the same forecast. That limits their use for uncertainty quantification, extreme weather risk assessment, and long-range forecasting, all of which need large ensembles. Training a new stochastic foundation model from scratch also carries a large, often prohibitive, computational cost.

The paper addresses both problems at once. It takes a pretrained deterministic prior, the Aurora foundation model, and turns it into a generative ensemble-prediction model. The new model is called Xaurora. The authors describe the conversion as parameter and sample efficient, and report that deterministic foundation models can be efficiently extended into stronger stochastic models.

The conversion is not a new model trained from zero. Xaurora is finetuned from the small Aurora version, Aurora-S, which has 113M parameters. The paper reports that Xaurora approaches state-of-the-art on global ensemble metrics and is competitive with the large Aurora version, Aurora-L, at 1.3B parameters.

Denoising Stochastic Interpolants replace drift matching with SDE rollout

The generative method is Denoising Stochastic Interpolants (DSI). A Stochastic Interpolant learns a continuous-time map between two endpoint distributions. Unlike flow matching or DDPMs, it does not assume a Gaussian initial distribution. Setting the endpoints to two consecutive atmospheric states turns the interpolant into a stochastic bridge, which enables autoregressive forecasting.

The interpolant connects state xt and state xt+Δt along a path indexed by pseudo-time τ in [0,1], combining the two endpoints with a Wiener process term. During training the drift of this path depends on the target state xt+Δt, so it is unavailable at inference. A neural network Bθ approximates that drift, and after training, samples are drawn by integrating an SDE from Xt0 = xt to the terminal state Xt1, which becomes the forecast. Repeating this autoregressively gives the physical-time rollout.

The paper distinguishes the interpolant Itτ, built from true endpoints and usable only in training, from the SDE state Xtτ, which the model generates and feeds back to the network at inference. The rollout training algorithm never uses the interpolant and trains exclusively on generated SDE states. A replay buffer for SDE rollout makes probabilistic training of SDE trajectories possible.

The setup models atmospheric uncertainty only, at 0.25° and 6-hour steps

Xaurora follows Aurora's configuration. It uses 69 variables on a 0.25° grid, with 721 latitude points and 1440 longitude points. The timestep is 6 hours, and the history ht is the two most recent states, xt−Δt and xt. Probabilistic forecasting here means sampling from p(xt+Δt | ht) and applying it autoregressively, so each ensemble member conditions on its own generated states.

The scope is deliberately narrow in one respect: the paper strictly focuses on modelling atmospheric uncertainty and deliberately ignores other sources of uncertainty, such as the data assimilation process that generates the initial conditions. So the ensemble spread reflects atmospheric uncertainty as represented by this model, not the full uncertainty of an operational forecast.

Training data comes from ERA5 via the WeatherBench2 project, downsampled to 6-hourly resolution with a reduced variable and pressure-level set. Initial conditions for operational-style evaluation use HRES-T0, also in its WeatherBench2 version.

Ensemble mean RMSE beats Aurora-S and Aurora-L at most lead times

Global medium-range skill is evaluated for 2022. The reported metrics are RMSE (unbiased RMSE for ensemble models), fair CRPS (ensemble only), and SSR (ensemble only). Rows cover different variables: Q850 as a proxy for mesoscale variability, and Z500 for large-scale circulation and synoptic-scale variability.

Xaurora's ensemble mean unbiased RMSE outperforms Aurora-S and Aurora-L at most lead times and is competitive with operational models. The comparison set includes GenCast-O, Aurora-O, Aurora-O, Aurora-L, Aurora-S, IFS ENS, and Xaurora.

Calibration is reported separately from error. All models are well calibrated at 15 days, where SSR is approximately 1. Xaurora's ensemble is generally underdispersive, with SSR below 1. Z500 is overdispersive, with SSR above 1, until approximately 3 days. These are different diagnostics: RMSE measures error, fair CRPS measures the ensemble's probabilistic skill, and SSR measures the ratio of spread to unbiased error.

Finetuning uses 500k samples and 42 hours on 16 GPUs

The efficiency claim rests on concrete training and inference numbers. Xaurora was trained on 500k samples for 42 hours on 16 A100 80GB GPUs. For comparison, Aurora-L used 3.8M training samples over 2 weeks on 32 GPUs, and Aurora-S used 3.8M samples with an unknown training time.

Inference cost is reported per forecast step. Xaurora takes 13.1 seconds for a 6-hour forecast and 52.3 seconds for a 24-hour forecast, using 10 integration steps with a Heun solver, which is 20 function evaluations per 6-hour step. Aurora-L takes 2.0 seconds for 6 hours and 8.2 seconds for 24 hours with 1 function evaluation per step; Aurora-S takes 0.6 and 2.4 seconds respectively.

The paper also reports an untrained Xaurora model initialised with Aurora-L, which takes 54.6 seconds for a 6-hour forecast and 218.2 seconds for 24 hours. The headline result is that skilful 15-day forecasts are generated in 13 minutes. The parameter increase over Aurora is small: 721,920 added parameters for Aurora-S, a 0.64% increase, and 3,475,456 for Aurora-L, a 0.28% increase. Most of that comes from LoRA layers in the Swin Transformer, with rank 8, while the Swin Transformer itself is entirely frozen.

A persistent noise floor leaves mesoscale variables underdispersive

The paper's own stated main limitation is a persistent noise floor. The analysis shows that the power spectral density of the ensemble spread is flat after a certain wavenumber, much lower for mesoscale variables such as Q850 and U500. That flat tail is the noise floor. Mesoscale variables are underdispersive across all wavenumbers, whereas synoptic-scale variables are overdispersive at smaller scales, which dominate when integrating SSR.

The authors test whether this is an artifact of under-training by evaluating different Xaurora checkpoints. The noise floor remains essentially constant across training, with a slight increase for Q850. They conclude the noise floor is an intrinsic property of the denoiser rather than a learned or forgotten trait, and note that whether it comes from the Aurora prior itself cannot be settled by the current analysis.

A denoiser trained from scratch is severely undertrained yet also yields a noise floor, at a slightly lower noise level. Using the large Aurora version as prior gives very similar results. The paper identifies the noise model as the major influence on the spectral properties of the forecasts.

Inference-only noise interventions calibrate some variables but not all

The paper runs four inference-only experiments that manipulate the random fields used in the diffusion term of the SDE, with no retraining and the same Xaurora checkpoint, producing a 10-member ensemble on 8 initial conditions from 2022 and evaluating at 12-hour lead time.

Scaling the noise variance by a scalar λ shows that higher variance generally means higher spread and lower or unchanged error. Overdispersive variables such as 2T and Z500 are better calibrated and have better error at λ = 0.5, while underdispersed variables such as Q850 and U500 do better at λ = 2. But a single λ cannot fix all variables, and the denoiser was trained on unit-variance random fields, so changing variance too much could have unexpected effects.

The authors add a stronger caveat: the spread change from variance scaling is unlikely to be meaningful because it only changes pixel intensity rather than creating dynamically different members, which is what ensemble forecasts are expected to provide. Changing noise colour, by sampling random fields with a power-law PSD, has a consistent effect on error and spread spectra, but the optimal exponent ν is again variable-dependent. The paper concludes that inference-only interventions can only partially calibrate the forecasts.

The practical gain is cheap ensembles, not a solved calibration problem

The paper's broader-impact argument is about access: finetuning Xaurora and generating forecasts has a relatively low computational cost, making skilful probabilistic weather forecasting more accessible to institutes with limited compute, such as academia or national meteorological offices. The method is also described as scalable to larger models or output spaces and expected to transfer to problems beyond weather forecasting.

For a reader deciding whether this changes their workflow, the relevant split is by variable and scale rather than by headline skill. If the use case depends on large-scale circulation and synoptic variability, Xaurora's ensemble mean unbiased RMSE outperforms Aurora-S and Aurora-L at most lead times and is competitive with operational models, and Z500 is overdispersive only until roughly 3 days. If the use case depends on mesoscale variability, the persistent noise floor and underdispersion across all wavenumbers for Q850 and U500 are the binding constraint, and inference-only noise tuning cannot fully remove it.

The paper's own forward-looking priority follows the same split: build random fields better suited to the spectral properties of the task, or learn the diffusion term together with the denoiser. The second option would mean optimising the SDE alongside model parameters, a non-trivial task for which the paper's probabilistic SDE-rollout training algorithm is presented as particularly well suited. Distillation for sampling efficiency and physics-constrained modelling are listed as other directions.

Sources used in this answer