PaperlayerView exact revision on arXiv

Source-linked AI summary

Xaurora: Generative Weather Forecasting with Denoising Stochastic Interpolants from a Foundation Model Prior

Eliot Walt, Miltiadis Kofinas, Nikolaj Mücke, Efstratios Gavves, Dim Coumou

arXiv:2610.06509v1cs.LG

TL;DR

Weather forecasting relies on ensembles for uncertainty assessment, but their computational cost limits size, and training generative models from scratch is expensive. Xaurora finetunes pretrained Aurora using Denoising Stochastic Interpolants and replay-buffer SDE rollout training; its ensemble performance approaches state-of-the-art systems, with 15-day forecasts generated in 13 minutes.

  • Problem

    Ensemble forecasts support uncertainty quantification and extreme-weather risk assessment, but their computational cost limits ensemble size, while training stochastic models from scratch is expensive.

    Source: §1 · p. 2

  • Method

    Xaurora finetunes deterministic Aurora into a generative model using Denoising Stochastic Interpolants and replay-buffer training of SDE rollouts with a CRPS objective.

    Source: §1 · p. 2

  • Results

    Xaurora approaches state-of-the-art ensemble performance, outperforms non-operational Aurora models on most variables and lead times despite 10× fewer parameters, and generates a 15-day forecast in 13 minutes on one A100 GPU.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results show that a deterministic foundation model can be efficiently extended into a stochastic model with skilful 15-day ensemble forecasts.

    Source: §6 · p. 9

    The persistent noise floor in Xaurora’s forecasts is its main limitation, and inference-only interventions can only partially calibrate the noise model’s spectral effects.

    Source: §6 · p. 9

Abstract

from arXiv · show

Deep learning has revolutionised weather forecasting in recent years, especially through atmospheric foundation models, which offer competitive skill for a fraction of the computational costs of classic physics-based models. However, most existing foundation models are deterministic, limiting the generation of large ensembles for accurate uncertainty quantification, extreme weather risk assessment, and long-range weather forecasting. Furthermore, these models incur a large, often prohibitive, computational overhead to train from scratch. To address these shortcomings, we turn a pretrained deterministic prior model, namely the Aurora foundation model, into a generative ensemble-prediction model. To that end, we introduce a novel generative method, Denoising Stochastic Interpolants, combined with a replay buffer for Stochastic Differential Equation (SDE) rollout, enabling probabilistic training of SDE trajectories. Our stochastic foundation model, Xaurora, is finetuned from the small Aurora version, yet it approaches the state-of-the-art on global ensemble metrics and is competitive with the large version of Aurora. Our method is parameter and sample efficient, and generates skilful 15-day forecasts in 13 minutes. Our results demonstrate that deterministic foundation models can be efficiently extended into even stronger stochastic models.

1 INTRODUCTION

Weather ensembles support uncertainty quantification and risk assessment, but their computational cost limits their size, while training generative models from scratch is expensive. Xaurora turns pretrained deterministic Aurora into a generative ensemble model through Denoising Stochastic Interpolants and replay-buffer SDE rollout training.

Motivation

  • Weather agencies rely on ensemble forecasts for uncertainty quantification and extreme-event risk assessment, but their substantial computational cost limits ensemble size.Atmospheric predictability is fundamentally limited, and forecast errors grow non-linearly with lead time.Source: §1 · p. 2

Approach

  • The paper proposes finetuning pretrained deterministic Aurora into a generative model rather than training a stochastic weather model from scratch.Xaurora makes only minimal architectural changes to Aurora.Source: §1 · p. 2

Method

  • Denoising Stochastic Interpolants provide a warm-startable parameterisation, while replay-buffer training combines SDE rollouts with a single-backpropagation CRPS objective.The method is designed to exploit Aurora’s accuracy at initialisation and train probabilistic trajectories efficiently.Source: §1 · p. 2

Performance and computational cost

  • Xaurora’s ensemble mean unbiased RMSE and CRPS are close to state-of-the-art systems, and it outperforms non-operational Aurora models on most variables and lead times despite using 10× fewer parameters.The method adds 0.64% to parameter count and generates a 15-day forecast in 13 minutes on one A100 GPU.Source: §1 · p. 2

2 BACKGROUND

Probabilistic forecasting samples the next atmospheric state conditioned on recent history, then rolls each member forward autoregressively. Stochastic Interpolants connect consecutive atmospheric states through a pseudo-time path whose learned SDE can generate forecast samples.

2.1 PROBLEM STATEMENT

  • Probabilistic weather forecasting samples the next atmospheric state conditioned on recent history, with each ensemble member autoregressively conditioning on its own generated states.The experiments use 69 variables on a 0.25° grid, six-hour steps, and the two most recent states as history; they model atmospheric uncertainty but not uncertainty in initial-condition data assimilation.Source: §2.1 · pp. 2–3

2.2 STOCHASTIC INTERPOLANTS FOR AUTOREGRESSIVE FORECASTING

  • A Stochastic Interpolant bridges consecutive atmospheric states along pseudo-time τ ∈ [0, 1], enabling autoregressive forecasting without requiring a Gaussian initial distribution.The interpolant uses the true endpoints during training, whereas inference generates an SDE state that is fed to the network.Source: §2.2 · p. 3
  • At inference, a neural network approximates the unavailable interpolant to define the SDE drift, which is integrated from its terminal pseudo-time state to draw a forecast sample.The integration scheme and step count depend on the desired sample fidelity and available compute.Source: §2.2 · p. 3

3 RELATED WORK

Prior work developed deterministic, generative, and CRPS-trained deep-learning weather predictors, alongside probabilistic extensions of Aurora. This paper’s approach differs by finetuning a deterministic foundation model into a generative predictor.

Deterministic DLWP

  • Deterministic deep-learning weather models reached global RMSE comparable to or better than NWP, but their predictions became smoother with lead time; Aurora introduced replay-buffer finetuning for cheaper rollout training.This work extends replay-buffer training to SDE-based ensemble models.Source: §3 · p. 3

Generative DLWP

  • Generative weather models have used diffusion for correction, autoregressive forecasting, residual modeling, downscaling, and zero-shot posterior sampling, including at 0.25° resolution.Unlike these approaches, this work finetunes a deterministic foundation model into a generative model at 0.25° resolution with limited training cost.Source: §3 · p. 4

CRPS training

  • CRPS-trained probabilistic models can use small ensembles and cheaper inference, but prior work reports overly smooth predictions and limited expansion to other tasks.Generative models can instead be expanded for posterior sampling or physics-constrained generation.Source: §3 · p. 4

Probabilistic extensions of Aurora

  • Recent probabilistic Aurora extensions use CRPS training and focus on expanding Aurora’s variables and modalities; this work explores generative finetuning as a complementary direction..Source: §3 · p. 4

4 DENOISING STOCHASTIC INTERPOLANTS AND XAURORA

The paper introduces Denoising Stochastic Interpolants (DSI), a warm-startable generative parameterisation, and replay-buffer SDE rollout training, then applies them to Aurora to build Xaurora.

4.1 DENOISING STOCHASTIC INTERPOLANTS

  • DSI reparameterises Stochastic Interpolants as prediction of the clean terminal state, which here is the next atmospheric state.This x1-prediction formulation is called denoising by analogy with diffusion models.Source: §4.1 · p. 4

4.2 DENOISING STOCHASTIC INTERPOLANTS FROM A PRETRAINED PRIOR

  • In a toy test, only DSI preserves the pretrained prior’s correct endpoint without finetuning; finetuning then aims to model spread while maintaining prior skill and potentially surpassing it with ensembles.The paper sets βτ = τ after finding that DSI reaches the correct endpoint in the tested warm-start setting.Source: §4.2 · p. 5, Figure 2 · p. 5
  • DSI’s denoiser and pretrained prior share a conditional-expectation target, unlike parameterisations that target combinations of states and random fields.The paper hypothesises that x-prediction is advantageous for high-dimensional atmospheric states because it remains on the data manifold.Source: §4.2 · p. 5

4.3 SDE-ROLLOUT REPLAY BUFFER FINETUNING WITH PROBABILISTIC LOSS

  • Replay-buffer training blends physical and SDE rollouts, allowing nested rollout training with negligible overhead relative to single-step optimisation.The buffer training extends Aurora’s replay-buffer approach to SDE rollouts.Source: §4.3 · pp. 5–6, Figure 3 · p. 6
  • The method directly optimises SDE rollouts with fair CRPS, whose skill and spread terms measure member errors and ensemble variance, respectively.Only one randomly selected member is supervised, preserving the same local minima while reducing computational cost; finetuning used 10% of Aurora’s pretraining steps and took 42 hours on 16 A100 GPUs.Source: §4.3 · p. 6

4.4 DENOISER MODEL

  • The denoiser is initialised from Aurora-S, with minimal architectural changes that increase parameter count by 0.64%; the Swin Transformer is frozen and finetuned with rank-8 LoRA.The design adds dedicated SDE-state embeddings and conditions on pseudo-time through a Fourier embedding and zero-initialised MLP.Source: §4.4 · p. 6
  • Xaurora’s ensemble mean unbiased RMSE outperforms Aurora-S/L on most lead times, is competitive with operational models, and all models are well calibrated at 15 days.Xaurora is generally underdispersive, while Z500 is overdispersive until approximately three days.Source: Figure 4 · p. 7
  • Parameter-free preconditioning adapts Karras et al.’s approach so model inputs, particularly the SDE state, are approximately zero-mean and unit-variance on average..Source: §4.4 · p. 7

5 EXPERIMENTS

The experiments evaluate Xaurora’s medium-range ensemble forecasts, spectral calibration, and ablations. They compare forecast skill with operational systems and examine calibration behavior and training choices.

5.1 MEDIUM-RANGE WEATHER FORECASTING

  • Xaurora’s ensemble mean tracks operational systems, beats Aurora-S after 1 day and Aurora-L after 3 days, and its CRPS closely follows operational models.Beyond 7 to 12 days, Xaurora has lower error than Aurora-O, though ensemble averaging itself provides a statistical advantage.Source: §5.1 · p. 7

5.2 SPECTRAL CALIBRATION ANALYSIS

  • For Z500, PSD flattens beyond k=31, while large scales remain underdispersive and short-lead small scales can drive global overdispersion.The cutoff is lead-time independent; the authors link short-lead overdispersion to high-frequency power required for calibration falling below the noise floor.Source: Figure 5 · p. 8, §5.2 · p. 8
  • Fitting noise to the error PSD calibrates SSR but costs long-term RMSE; the authors hypothesize the SDE diffusion term drives the noise floor.They suggest similar approaches during training could instead prevent the noise floor.Source: §5.2 · p. 8

5.3 ABLATIONS

  • 3.1% better uRMSE but 15.5% worse CRPS: MAE training trades probabilistic scores for point-error skill, while curriculum and replay buffer further improve uRMSE.At 15 days, Xaurora is significantly better across variables and metrics on longer leads.Source: §5.3 · p. 9
  • B- and W-pred fail to warm-start, while x1-pred + DDPM also underperforms; the replay buffer’s 10-step restriction may disadvantage DDPM, but not DSI.The authors report that insufficient integration steps do not explain the DDPM result.Source: §5.3 · p. 9
  • Aurora-L initialization reproduces the ensemble-over-prior result under single-step + MAE training and incomplete convergence; further probabilistic scaling was not pursued due to cost.The authors expect scaling training and making it probabilistic would improve results.Source: §5.3 · p. 9

6 CONCLUSION

Xaurora efficiently fine-tunes Aurora into an ensemble forecaster whose 15-day forecasts approach state-of-the-art probabilistic systems. The paper identifies a persistent noise floor as its main limitation and notes that the prototype needs independent validation before high-stakes use.

6 CONCLUSION

  • Xaurora generates 15-day ensemble forecasts in 13 minutes and approaches the performance of state-of-the-art probabilistic weather forecasting systems.The conclusion attributes this to Denoising Stochastic Interpolants and probabilistic rollout training for SDE-based autoregressive prediction, enabling efficient fine-tuning from a deterministic prior.Source: §6 · p. 9
  • Low-cost fine-tuning and forecast generation may make skilful probabilistic forecasting more accessible to institutions with limited compute.The authors also expect the method to scale to larger models or output spaces and transfer beyond weather forecasting.Source: §6 · p. 10
  • The persistent noise floor is Xaurora’s main limitation, and the research prototype should not be used for high-stakes decisions without independent validation.Inference-only interventions can only partially calibrate the noise model’s spectral effects; the authors suggest better-suited random fields and other future directions.Source: §6 · pp. 9–10, ETHICS STATEMENT · p. 10

A.1.1 DATASETS; A.1.2 LIST OF VARIABLES

  • The study uses ERA5 reanalysis and HRES-T0 operational initial conditions; Xaurora models Aurora’s set of variables.The ERA5 subset is downsampled to six-hourly resolution with fewer variables and pressure levels.Source: §A.1.1 · p. 18, §A.1.2 · p. 18

A.2 AURORA: THE FOUNDATION MODEL PRIOR

  • Aurora is a deterministic autoregressive atmospheric foundation model that predicts the next state from recent history.Its encoder and decoder use Perceiver-based components around a Swin Transformer backbone.Source: §A.2 · p. 18

A.3.1 ARCHITECTURE DETAILS

  • The Xaurora denoiser makes minimal architectural changes to Aurora, with a total parameter increase below 1% for Aurora-S and Aurora-L.New parameters are initialized so the model reduces to Aurora before training; additions include a pseudo-time embedding.Source: Figure 6 · p. 19, §A.3.1 · p. 20

A.3.2 NORMALISATION AND NETWORK PRECONDITIONING

  • The denoiser reuses Aurora’s variable-normalization statistics and conditions on an approximate inverse standard deviation for the SDE state.The approach ignores correlation between consecutive physical states for simplicity.Source: §A.3.2 · p. 20

A.4.1 METRICS

  • Evaluation metrics are computed per variable and forecast lead time, with latitude weighting to account for the non-uniform grid.The evaluation uses the latest WeatherBench version.Source: §A.4.1 · p. 20

A.4.2 STATISTICAL SIGNIFICANCE FOR ABLATION EXPERIMENTS

This appendix details significance testing and evaluation settings, then derives the DSI parameterisation and describes its training and sampling procedures. Its toy analysis distinguishes parameterisations by whether their mean trajectories reach the correct endpoint.

A.4.2 STATISTICAL SIGNIFICANCE FOR ABLATION EXPERIMENTS

  • Paired bootstrap intervals compare models on the same resampled initialisations, but baseline comparisons have no significance tests because their metrics were downloaded rather than evaluated here.The analysis uses 10,000 replicates; alternative stationary-bootstrap and paired-t-test procedures leave all significance labels unchanged.Source: §A.4.2 · pp. 21–22

A.4.3 FORECAST SETTINGS

  • Forecasts use 56 dates from 2022 and 10-member Xaurora ensembles; comparisons include Aurora versions, operational GenCast, and IFS ENS.Aurora-O metrics extend only to 10 days, so the authors ran Aurora forecasts themselves and report alignment with official evaluation.Source: §A.4.3 · p. 22, §A.4.4 · p. 22

A.4.5 MATCHED BENCHMARK FOR XAURORA MODELS

  • Matched Xaurora evaluations use common initialisations and fixed noise samples to reduce weather and Wiener-process variability, primarily for faster, more accurate benchmarks.The procedure is not used in the main-paper ablation experiments.Source: §A.4.5 · p. 23

A.5 DRIFT PARAMETERISATIONS

  • Table 5 lists each Stochastic Interpolant parameterisation's network, target, and drift, derived by expanding the drift as an expectation and isolating terms.The mean SDE trajectory is approximated by an ODE using the prior in place of the network in each parameterisation's drift.Source: §A.5 · p. 23, §A.5.2 · p. 24

A.5.1 LOSS FUNCTION EQUIVALENCES

  • The loss trained in x1-space can be expressed in B- and W-space using a single τ-dependent reweighting, but the paper leaves their comparison for future work.The derivation is algebraic.Source: §A.5.1 · p. 23

A.5.2 MATHEMATICAL GENERALISATION OF THE 2D TOY EXAMPLE

  • For a direct-prediction prior, the toy analysis finds that x1-prediction reaches the conditional-mean endpoint, whereas B- and W-prediction diverge.With residual prediction, B-prediction instead reaches the correct endpoint; the analysis is extended to arbitrary finite dimension.Source: §A.5.2 · p. 24

A.6.1 GENERAL ALGORITHMS

  • The vanilla DSI training algorithm samples pseudo-time and noise, forms the interpolant, and trains the denoiser; the sampling algorithm iteratively advances SDE states within autoregressive forecast steps.The sampling procedure carries generated states forward as the running history.Source: §A.6.1 · p. 25

A.6.2 SINGLE-BACKPROPAGATION CRPS

  • Single-backpropagation CRPS training supervises one randomly selected ensemble member and stop-gradients the other, reducing trajectory backpropagation while retaining an unbiased one-sided gradient estimate.At τ = 0, the spread term is zero and the loss is equivalent to Aurora's weighted MAE; a replay buffer emulates rollouts with a minor off-policy discrepancy.Source: §A.6.2 · pp. 25–26

A.6.3 PROBABILISTIC SDE-ROLLOUT REPLAY BUFFER TRAINING ALGORITHM

The appendix describes replay-buffer SDE-rollout training and establishes a schedule condition for warm-starting DSI from a pretrained prior. It also specifies the finetuning setup.

A.6.3 PROBABILISTIC SDE-ROLLOUT REPLAY BUFFER TRAINING ALGORITHM

  • Leak-free rollout training evolves SDE states using only the denoiser drift, keeping the target out of the model input and reserving it for the loss.This trains on the rollout distribution rather than the interpolant distribution.Source: §A.6.3 · p. 26
  • The FIFO replay buffer stores SDE states with their histories, pseudo-time coordinates, and rollout information; its off-policy discrepancy is characterized as minor.Training samples average pseudo-time 0.5 and rollout depth 4.7; the buffer averages are also tracked over training.Source: §A.6.3 · p. 26, Figure 9 · p. 26, Figure 8 · p. 26
  • Probabilistic training unrolls pairs of trajectories and supervises them with single-backpropagation CRPS.The method trains the denoiser with a replay buffer to emulate rollouts.Source: §A.6.3 · p. 26

A.6.4 DSI WITH PRETRAINED PRIOR REQUIRES ˙β0̸ = τ

  • A linear interpolation schedule βτ = τ reaches the unbiased prior exactly regardless of integration discretisation, while βτ = τ² needs more than 20 steps to converge.Schedules with zero initial derivative begin with a model-independent drift; a nonzero initial derivative could also work, for example βτ = sin(πτ/2).Source: Figure 10 · p. 27, §A.6.4 · p. 27

A.7 TRAINING DETAILS

  • Aurora-S is finetuned on ERA5 data from 1979–2020, validated on 2021, and tested on 2022, using FP32, AdamW, cosine annealing, and gradient clipping at 1.Aurora-S was pretrained on the same ERA5 splits, which the authors say avoids data leakage.Source: §A.7 · p. 27

A.8 COMPUTATIONAL REQUIREMENTS

The appendix reports the relative training and inference costs of Xaurora, then examines how changing SDE noise affects spectral calibration. Noise-field design alters calibration, but fitted noise parameters calibrated at one lead time do not generalize across lead times.

A.8 COMPUTATIONAL REQUIREMENTS

  • Xaurora finetuning uses substantially fewer samples and GPU hours than Aurora pretraining and about 20 times fewer GPU hours than Aurora 1.5; a single-member 24-hour forecast takes under a minute.Inference uses 10 Heun integration steps; the authors caution that external-model timing comparisons require matched evaluation code.Source: §A.8 · pp. 27–28, Table 7 · p. 29

B.1.1 NO INTERVENTION: WHITE NOISE

  • With unmodified white noise, ensemble spread flattens at higher wavenumbers while error power follows an approximate power law, producing a noise floor.Mesoscale variables are underdispersive across wavenumbers; synoptic-scale variables are underdispersive at large scales but overdispersive across the full range.Source: §B.1.1 · p. 29

B.1.2 VARIANCE SCALING

  • Variance scaling can improve calibration for some variables—λ = 0.5 for overdispersive 2T and Z500, and λ = 2 for underdispersed Q850 and U500—but one λ cannot suit all variables.Scaling changes pixel intensity rather than generating dynamically different members, and large variance changes may be unexpected because the denoiser was trained on unit-variance fields.Source: §B.1.2 · pp. 29–30

B.1.3 NOISE COLOURING

  • Noise colouring changes spread and error spectra across scales more effectively than variance scaling, but tested ν values do not calibrate synoptic-scale 2T or Z500.U500 and Q850 are almost perfectly calibrated at ν = 1 and ν = 2, respectively; the optimal ν depends on the variable.Source: §B.1.3 · p. 31

B.1.4 ISOTROPIC SPHERICAL WHITE NOISE

  • Spherical isotropic noise removes spectral tail flattening, but global SSR changes only slightly because spread and error both rise slightly.This experiment samples white noise on spherical harmonics to respect Earth geometry.Source: §B.1.4 · pp. 31–32, Figure 14 · p. 32

B.1.5 NOISE MODEL FITTING WITH MATÉRN COVARIANCE

  • A variable-dependent Matérn noise fit calibrates ensembles across wavenumbers and variables, but its gains do not generalize beyond the fitted 12-hour lead time.The fit uses 8 validation initialisations from 2021, then evaluates on 8 test initialisations from 2022; no feedback mechanism is implemented.Source: §B.1.5 · pp. 32–33

B.2.1 GENERATIVE MODEL AND SOLVER CONFIGURATION

The configuration tests whether the observed noise floor depends on the generative model or SDE solver, while comparing forecast metrics under a fixed compute budget. DSI with Heun converges near its operating budget, whereas fitted Matérn noise improves calibration mainly at its fit lead time and can impair longer-lead uRMSE.

B.2.1 GENERATIVE MODEL AND SOLVER CONFIGURATION

  • The solver comparison fixes the compute budget at 200 NFE per physical step, with ensemble size and NFE chosen so M × NFE = 200 except at NFE = 2.The main configuration uses NFE = 20, M = 10, Heun, and S = 10.Source: §B.2.1 · p. 33, Table 8 · p. 36
  • Fitted Matérn fields generally correct SSR and uRMSE only at the 12 h fit lead, while uRMSE is negatively affected, especially at longer leads.Figure 18 compares fitted fields with white noise across variables and lead times.Source: Figure 18 · p. 34
  • The noise floor appears across generative-model and solver configurations, indicating that it is a property of the denoiser rather than the solver.The tested configurations also indicate that the main Xaurora setup is properly configured.Source: §B.2.1 · pp. 33, 36
  • DSI with Heun converges at the operating budget of 20 NFE, while DDPM underperforms and degrades beyond 10–20 NFE; DSI with Heun and Euler–Maruyama have no clear winner.Both DSI solvers reach optimal performance around 20 NFE.Source: Figure 19 · p. 35, §B.2.1 · p. 36
  • Increasing NFE reduces prediction power and improves spectral fit, but the noise floor appears in every configuration.With Euler–Maruyama, the DSI noise floor worsens after 20 NFE.Source: Figure 20 · p. 35, §B.2.1 · p. 36

B.2.2 SUFFICIENT DENOISER TRAINING

The noise floor remains essentially constant across training checkpoints, refuting the hypothesis that the denoiser is simply under-trained. It also appears with a denoiser trained from scratch and with Aurora-L as prior, though the scratch model has a slightly lower noise level.

B.2.2 SUFFICIENT DENOISER TRAINING

  • The noise floor remains essentially constant across training checkpoints, refuting the under-training hypothesis.The analysis uses 12 h PSDs from 8 test initialisations and 2 ensemble members; power above wavenumber 100 measures the PSD tail.Source: §B.2.2 · p. 36, Figure 21 · p. 37

B.2.3 CONTRIBUTION OF THE PRETRAINED PRIOR

  • A severely undertrained denoiser trained from scratch also exhibits a noise floor, at slightly lower power, and Aurora-L prior yields very similar spectra to Aurora-S prior.Both comparisons use 12 h PSDs averaged across 8 test initialisations and 10 ensemble members.Source: §B.2.3 · p. 36, Figure 22 · p. 37, Figure 23 · p. 37

B.2.4 CONTRIBUTION OF THE FINETUNING METHOD

The noise floor appears across optimization methods, although its level varies; unlike MAE-based methods, Xaurora reconstructs ERA5 at mid-range wavenumbers for synoptic-scale variables.

B.2.4 CONTRIBUTION OF THE FINETUNING METHOD

  • All tested optimization methods produce a noise floor, but MAE-based methods under-power mid-range wavenumbers for synoptic-scale variables while Xaurora reconstructs ERA5 there.The comparison uses the PSDs of Xaurora and optimization approaches from the ablations.Source: §B.2.4 · p. 38

B.3.1 MAIN PAPER ABLATIONS DETAILS

Xaurora outperforms the single-step MAE baseline especially at longer lead times, while the ablations examine training objectives, parameterizations, ensemble size, prior finetuning, and a stochastic-depth baseline. The findings include improvements with two ensemble members and a finetuning advantage, with a fairness caveat for comparisons against training from scratch.

B.3.1 MAIN PAPER ABLATIONS DETAILS

  • Single-step MAE quickly overfits, curriculum MAE rollout performs poorly probabilistically, and replay-buffer MAE is comparable to curriculum MAE, underscoring the probabilistic loss’s importance.The curriculum uses gradient-less steps based on a conditional-expectation approximation before supervising the final step.Source: §B.3.1 · p. 38
  • Among alternative parameterizations, x1-prediction DDPM falls short of DSI, W-prediction becomes unstable after about 70 steps and is not evaluated, and stable B-prediction has average CRPS 1.8 versus Xaurora’s 0.1 at first validation.The authors hypothesize that the replay buffer’s S = 10 may make DDPM training more challenging.Source: §B.3.1 · p. 39
  • A 1.3B-parameter Aurora-L initialized Xaurora trained with Single-step + MAE outperforms Aurora-L, although the model has not converged and the training-step count is held constant.The authors expect a probabilistic replay-buffer version of the large model to perform even better.Source: §B.3.1 · p. 39

B.3.2 IN-DEPTH COMPARISON: XAURORA VS SINGLE STEP + MAE

  • From 6 to 7 days, Xaurora is significantly better than Single step + MAE; its SSR also stays significantly closer to 1 as lead time increases.Short-term calibration is similar; the figure’s skill-score bands use 95% confidence intervals, and negative values mean Xaurora is better.Source: §B.3.2 · pp. 39–40, Figure 27 · p. 40

B.3.3 ENSEMBLE SIZE

  • With two ensemble members, Xaurora already beats Aurora-S at 5 days, while increasing ensemble size improves Q850 RMSE and the uncorrected RMSE generally improves with more members.The ensemble-mean comparison with deterministic baselines is acknowledged as statistically advantaging the ensemble.Source: Figure 25 · p. 39, §B.3.3 · p. 40

B.3.4 FINETUNING VS TRAINING FROM SCRATCH

  • Finetuning from Aurora-S performs much better than training from scratch across lead times, but a fairer comparison would add Aurora-S pretraining samples to the scratch model’s training.Both models use Single-step + MAE because rollout training requires meaningful initial predictions.Source: Figure 26 · p. 39, §B.3.4 · p. 40

B.3.5 NAIVE PROBABILISTIC AURORA BASELINE: MC-DROPPATH

  • Xaurora is more valuable than the MC-DropPath probabilistic baseline that keeps Aurora’s stochastic-depth mechanism active at inference.Source: §B.3.5 · p. 40

B.4 ADDITIONAL METRICS

The additional-metrics section presents global forecast metrics and calibration diagnostics, while the scorecards compare Xaurora with Aurora-S and Aurora-L. Forecast visualizations show rollouts through 15 days, and Algorithm 3 describes replay-buffer SDE-rollout training.

B.4 ADDITIONAL METRICS

  • Figure 28 compares Xaurora with naive probabilistic Aurora-S using CRPS improvement percentages across variables, lead times, and pressure levels.Both models have 10 ensemble members; red indicates degraded performance and blue indicates improvement over the reference.Source: Figure 28 · p. 41
  • Spatial pooling generally improves CRPS and tends to decrease SSR, while rank histograms show short-lead overdispersion for 2T and Z500 and underdispersion for U500 and Q850.The additional diagnostics also include RMSE or uRMSE, CRPS, SSR, power spectral density at 2 and 10 days, and calibration maps at 12 hours and 5 days.Source: §B.4 · p. 40, Figure 36 · p. 48, Figure 37 · p. 48

B.5.1 AURORA SCORECARDS

  • Xaurora outperforms Aurora-S mostly after 1 day, while Aurora-L is generally better around 3 days before falling short against Xaurora.The comparison uses uRMSE for Xaurora to characterize the maximal gain; an ensemble-size-dependent analysis is provided elsewhere.Source: Figure 38 · p. 49, Figure 39 · p. 49, §B.5.1 · p. 49
  • Figures 38 and 39 present Xaurora-versus-Aurora-S and Xaurora-versus-Aurora-L RMSE scorecards across variables, lead times, and pressure levels.Figure 38 uses unbiased RMSE for Xaurora and plain RMSE for Aurora-S; the scorecards use red for degradation and blue for improvement.Source: Figure 38 · p. 49, Figure 39 · p. 49

B.5.2 FORECAST VISUALISATION

  • Figures 40–42 show 3-, 10-, and 15-day rollouts, with ERA5, Xaurora member 0, the 10-member Xaurora ensemble mean, and Aurora-S in columns.Rows show unnormalised variables; Figures 40 and 41 cover a selection of variables and pressure levels for the stated initialization.Source: Figure 40 · p. 50, §B.5.2 · p. 50, Figure 41 · p. 51
  • Algorithm 3 trains a denoiser with two-member fair CRPS supervision, then advances both members with an SDE solver and stores trajectory states in a replay buffer.It samples training entries from data or the buffer, tracks gradients through only the selected member, and uses stop-gradient for the other member.Source: §B.5.2 · p. 53

Related papers

Loading 2610.06509v1…