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

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

cs.LGOct 5, 2026

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

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

Organizations: Vrije Universiteit Amsterdam · University of Amsterdam · Delft University of Technology · Centrum Wiskunde & Informatica

Abstract

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.

Figures & tables

Figure 1: Xaurora produces ensemble forecasts with a nested rollout . Each physical step $\bm{x}^{t}\to\hat{\bm{x}}^{t+\Delta{t}}$ is generated by integrating a stochastic differential equation (SDE) over pseudo-time $\tau\in[0,1]$ (SDE time, grey box). Starting from the current state $\bm{X}_{0}^{t}=\bm{x}^{t}$ , the state $\bm{X}^{t}_{\tau}$ follows an SDE formed by the Xaurora drift term ( $\bm{B}_{\theta}$ ) and a Wiener process ( $\bm{W}_{\tau}$ ). At $\tau=1$ , the terminal SDE state is the forecast, $\hat{\bm{x}}^{t+\Delta{t}}\coloneq\bm{X}^{t}_{1}$ , and the procedure is repeated autoregressively for the physical time rollout. Ensemble members are generated by sampling different Wiener processes $\bm{W}_{\tau}$ .

Figure 2: 2D toy example: only DSI preserves an ideal prior. In each parameterisation ( Table 5 ), we replace the network with an ideal prior $\bm{G}(\bm{x}_{0})=\mathbb{E}[\bm{x}_{1}\mid\bm{x}_{0}]$ and integrate Equation 4 from a common initial state (black square), no training. Only DSI recovers the true mean (gold star) as its network models $\mathbb{E}[\bm{x}_{1}\mid\bm{x}_{0}]$ , exactly what $\bm{G}$ computes. The deterministic prior has no spread, hence the SDE trajectories (thin lines, thick line is the average) land on the expectation. Hence, finetuning ( Section 4.2 ) aims to capture the spread and maintain the skill. We design our algorithm ( Section 4.3 ) for that purpose and apply it to an atmospheric foundation model ( Section 4.4 ).

Figure 3: SDE-rollout training with single-backpropagation CRPS loss. (a) Two members are sampled from the buffer and fed to the network, only member

m

is backpropagated. Both advance one step (physical step if

\tau

reaches 1, SDE step otherwise), and are pushed back. Fresh samples are randomly injected, old (FIFO) or finished rollouts are removed. (b) The buffer tracks pairs of members (yellow and blue circles) throughout nested rollouts (

\tau,k

). Both members start from the same initial condition, but receive different noise samples. A training iteration costs the same at every depth (2 forward, 1 backward, purple box), unlike regular rollout training that require full nested rollouts (shown in grey).

Figure 4: Global medium-range forecast skill in 2022. From left to right, columns represent RMSE (uRMSE for ensemble models), fair CRPS (ensemble only), and SSR (ensemble only). Rows represent different variables; Q850 is a proxy for mesoscale variability, Z500 for large-scale circulation and synoptic-scale variability. Xaurora’s ensemble mean unbiased RMSE outperforms Aurora-S/L on most lead times and is competitive with operational models. All models are well calibrated at 15 days (SSR

\approx

1). Xaurora’s ensemble is generally underdispersive (SSR

<

1). Z500 is overdispersive (SSR

>

1) until approximately 3 days.

Figure 5: Spectral calibration and spatial structure. (a) Average Z500 zonal Power Spectral Density (PSD) at 12 hours (solid) and 5 days (dashed) with white noise (green) and error-fitted noise (purple, see Section B.1.5 ), compared to ERA5 (dotted). The vertical dotted line indicates the cut-off wavenumber ( $k=31$ ; defined as the first $k$ where the power is twice as large as that of ERA5). (b) Cumulative spectral SSR at 12 hours and 5 days, can be interpreted as SSR up to wavenumber k . (c-e) Spatial spread, RMSE, and SSR for Z500 at 12 hours (with white noise). (f-h) Same at 5 days. The noise floor is visible in the spectral plots (k $>$ 31) and the lack of tropical structure in (c) and (f).

		5-day uRMSE		5-day CRPS
		Z500	Q850	Z500	Q850
Xaurora (Replay buf. + CRPS + Aurora-S)		309.9	$1.52\times 10^{-3}$	139.4	$7.55\times 10^{-4}$
Aurora-S		-15.4%	-18.2%	N/A	N/A
Optim.	Single-step + MAE	0 +0.0% †	0 +0.1% †	000 -1.3% †	00 -6.5%
	Curriculum + MAE	0 +3.1%	0 +1.0%	000 -4.1%	0 -15.5%
	Replay buf. + MAE	0 +1.4% †	0 +2.5%	000 -4.7%	0 -14.9%

Table 1: Ablation experiments results. Skill scores against Xaurora,

100\cdot(1-\mathcal{M}_{\text{model}}/\mathcal{M}_{\text{Xaurora}})

: positive values are better than the reference. The best values are underlined. Methods are detailed in Section B.3.1 .

\dagger

marks differences that are not significant at the 95% level (see Section A.4.2 ).

Appendix figures & tables44 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Units	Name	Kind
2T	$\mathrm{K}$	2m-temperature	Surface
MSL	$\mathrm{Pa}$	Mean sea level pressure	Surface
10U	$\mathrm{m}\text{\,}{\mathrm{s}}^{-1}$	10m-zonal wind component (east-west)	Surface
10V	$\mathrm{m}\text{\,}{\mathrm{s}}^{-1}$	10m-meridional wind component (north-south)	Surface
T	$\mathrm{K}$	Temperature	Atmosphere
Z	$\mathrm{m}$	Geopotential	Atmosphere

Appendix

Table 2: List of variables . Atmospheric variables are measured on discrete pressure levels (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000

\mathrm{hPa}

Model	Size	Accuracy	Key characteristics
Aurora-S	113M	Lowest	Fastest and lightest.
Aurora-L	1.3B	High	Pretrained foundation model.
Aurora-O	1.3B	Highest	Aurora-L fine-tuned for operational weather forecasts.

Appendix

Table 3: Aurora model variants in Bodnar et al. (2025) .

Figure 6: Xaurora denoiser architecture . Minimal modifications to Aurora that turn it into the denoiser $\bm{D}_{\theta}(\bm{X}^{t}_{\tau},\bm{h}^{t},\tau)$ . New parameters are initialised such that the model reduces to Aurora before training. Embeddings: $\bm{e}_{\text{surface}}$ , $\bm{e}_{\text{atmosphere}}$ are Aurora’s surface/atmospheric variable embeddings, $\bm{e}_{\text{space+time}}$ the spatiotemporal position embedding, $\bm{e}_{\text{lead time}}$ the lead-time embedding, and $\bm{e}_{\tau}$ the new (zero-initialised) pseudo-time embedding.

Component	Aurora-S	Aurora-L
LoRA	540,672	2,850,816
SDE state $\bm{X}^{t}_{\tau}$	49,664	99,328
Pseudo-time $\tau$	131,584	525,312
Total	721,920	3,475,456
Increase	+0.64%	+0.28%

Appendix

Table 4: Parameter increase in the denoiser compared to Aurora-S and Aurora-L.

Figure 7: Aurora-O validation. RMSE of Aurora-O from the official WeatherBench2 (WB2, dashed black) and from our own run (Ours, red). The two runs are essentially equivalent, with negligible variations coming from different sets of initialisations. The WB2 run only goes up to 10 days and hides the shut down of the rollout training mechanism used in Aurora-O (discontuity at 10 days).

	Network	Target	Drift $(\mathbb{E}[\bm{R}^{t}_{\tau}\mid\bm{X}^{t}_{\tau},\bm{h}^{t}])$
$\bm{B}$ -prediction	$\bm{B}_{\theta}$	$\bm{R}^{t}_{\tau}$	$\bm{B}_{\theta}$
$\bm{W}$ -prediction	$\bm{W}_{\theta}$	$\bm{W}_{\tau}$	$\dot{\alpha}_{\tau}\bm{x}^{t}+\frac{\dot{\beta}_{\tau}}{\beta_{\tau}}\big[\bm{X}^{t}_{\tau}-\alpha_{\tau}\bm{x}^{t}-\gamma_{\tau}\bm{W}_{\theta}(\bm{X}^{t}_{\tau},\bm{h}^{t},\tau)\big]+\dot{\gamma}_{\tau}\bm{W}_{\theta}(\bm{X}^{t}_{\tau},\bm{h}^{t},\tau)$
$\bm{x}_{1}$ -prediction	$\bm{D}_{\theta}$	$\bm{x}^{t+\Delta{t}}$	$\dot{\alpha}_{\tau}\bm{x}^{t}+\dot{\beta}_{\tau}\bm{D}_{\theta}(\bm{X}^{t}_{\tau},\bm{h}^{t},\tau)+\frac{\dot{\gamma}_{\tau}}{\gamma_{\tau}}\big[\bm{X}^{t}_{\tau}-\alpha_{\tau}\bm{x}^{t}-\beta_{\tau}\bm{D}_{\theta}(\bm{X}^{t}_{\tau},\bm{h}^{t},\tau)\big]$

Appendix

Table 5: Neural network, target and drift term for each parameterisation

Figure 8: Replay-buffer rollout training statistics. Top: Time series of the pseudo-time (

\tau

) of the samples Bottom: Time series of the rollout depth of the samples. The values seen by the model at each training steps are in light blue and the replay buffer averages are depicted in dark blue. The x-axis shows the training step

N_{\text{bwd}}

as defined in Table 6 . These plots show that throughout the training, the denoiser is exposed to a variety of pseudo-times (mean=0.5) and rollout steps (mean=4.7).

Figure 9: Replay-buffer training off-policy discrepancy. Top: Buffer mean sample age, i.e. , number of training steps since the trajectory was first pushed to the buffer. Bottom: Average number of training steps since a sample was updated, i.e. , popped, advanced by one SDE step, and pushed back to the buffer. The x-axis shows the training step

N_{\text{bwd}}

as defined in Table 6 . These plots characterise the minor off-policy discrepancy introduced by the replay buffer compared to a naive rollout training strategy (intractable).

Figure 10: Denoiser parameterisation: effect of the interpolation schedule $\beta_{\tau}$ . (a,b) Warm-start SDE trajectories for the denoiser with a linear schedule $\beta_{\tau}=\tau$ and a quadratic schedule $\beta_{\tau}=\tau^{2}$ (same 2D setup and markers as Fig. 2 ). (c) Distance of the mean SDE trajectory endpoint to the unbiased prior, $\lVert\bm{X}_{1}-\bm{G}_{\phi}\rVert$ , versus the number of integration steps (Euler–Maruyama). The linear schedule is exact regardless of the integration discretisation, whereas the quadratic schedule requires more than 20 steps to converge.

Symbol	Meaning	Value
–	Aurora variant	Aurora-S
$p_{\text{dropout}}$	Dropout rate ( Hinton et al., 2012 )	0.0
$p_{\text{droppath}}$	Stochastic depth ( Huang et al., 2016 )	0.0
$\alpha_{\tau}$	Stochastic Interpolant coefficient	$1-\tau$
$\beta_{\tau}$	Stochastic Interpolant coefficient	$\tau$ (according to Section A.6.4 )
$\gamma_{\tau}$	Stochastic Interpolant coefficient	$1-\tau$

Appendix

Table 6: Xaurora hyperparameters

	Aurora-L	Aurora-S	Aurora 1.5	Xaurora	Xaurora init. with Aurora-L
Training samples	3.8M	3.8M	2.1M	500k	N/A
Training time	2 weeks	unknown	2 weeks	42 hours	N/A
Number of training GPUs	32	unknown	24 to 32	16	N/A
$6\text{\,}\mathrm{h}$ forecast time (in seconds)	2.0	0.6	unknown	13.1	54.6
$24\text{\,}\mathrm{h}$ forecast time (in seconds)	8.2	2.4	unknown	52.3	218.2
NFE per $6\text{\,}\mathrm{h}$ -step	1	1	1	20	20

Appendix

Table 7: Computational requirements of training and running Aurora and Xaurora models. Training times only include pretraining for Aurora-S and Aurora-L. GPUs are 80GB A100. Aurora and Aurora 1.5 numbers are taken from ( Bodnar et al., 2025 ; Weyn et al., 2026 ) . Aurora 1.5 numbers indicate additional finetuning over Aurora-O. The last columns shows an untrained Xaurora model initialised with Aurora-L, hence the training numbers are not provided. We also show the number of function evaluation (NFE). Xaurora uses 10 integration steps with a Heun solver (2 NFE per pseudo-step).

Figure 11: Noise fields and PSD used in the spectral calibration analysis. Top three: Fields resulting from scaling the variance of a normal distribution by $\lambda$ (used Section B.1.2 ). Bottom three: fields characterised by a power law PSD, $\mathrm{PSD}(k)\propto k^{-\nu}$ (used in Section B.1.3 ). Right panel: PSD of each fields. The inserts in the field plots show the distribution at a specific location over 500 draws compared to a normal distribution (shown in black). $\lambda=1$ and $\nu=0$ correspond to the same white noise distribution. Varying $\lambda$ only changes the intensity of the pixels, whereas increasing $\nu$ increases the spatial correlation length (large structures).

Figure 12: Spectral calibration of an 8-initialisation

12\text{\,}\mathrm{h}

ensemble rollout with scaled variance (

\bm{\lambda}

). Top row: Zonal PSD of the ensemble spread (solid line) and unbiased ensemble mean error (dashed line). Z500 is unstable on large wavenumbers because of the log scale and

P^{u}_{\text{err}}(k>200)\approx 0

. Middle row: Cumulative PSD of the ensemble spread (solid line) and unbiased ensemble mean error (dashed line). The global spread and uRMSE correspond to the final values (circles). Bottom row: Cumulative spectral SSR; the dotted line is perfect calibration and the global SSR corresponds to the final value (circle). Curve colours correspond to different scaling values (

\lambda

). White noise corresponds to

\lambda=1

(grey). Curves are smoothed with a rolling average of size 7.

Figure 13: Spectral calibration of an 8-initialisation

12\text{\,}\mathrm{h}

ensemble rollout with varying noise colour (

\bm{\nu}

). Top row: Zonal PSD of the ensemble spread (solid line) and unbiased ensemble mean error (dashed line). Z500 is unstable on large wavenumbers because of the log scale and

P^{u}_{\text{err}}(k>200)\approx 0

\nu

exponents. White noise corresponds to

\nu=0

(grey). Curves are smoothed with a rolling average of size 7.

Figure 14: Noise fields and PSD on a flat grid and on the sphere. Maps: Fields resulting from sampling white noise on a flat grid (i.i.d. (grid), equivalent to white noise in other experiments, shown in grey) and on the sphere (isotropic (sphere), equivalent to ( Price et al., 2023 ) , shown in pink). Right panel: Global (solid line) and polar (dashed) PSD of both fields. The grid-based white noise has constant power across wavenumbers and latitudes. Because the earth is a sphere, grid cells in the poles are smaller than on the equator. Therefore, grid-based white noise over-powers the high frequencies in these regions. The spherical version address this problem by explicitly calibrating the variance and spatial correlation to the geometry.

Figure 15: Spectral calibration of an 8-initialisation

12\text{\,}\mathrm{h}

ensemble rollout with i.i.d. grid white noise and spherical isotropic white noise. Top row: Zonal PSD of the ensemble spread (solid line) and unbiased ensemble mean error (dashed line). Z500 is unstable on large wavenumbers because of the log scale and

P^{u}_{\text{err}}(k>200)\approx 0

. Middle row: Cumulative PSD of the ensemble spread (solid line) and unbiased ensemble mean error (dashed line). The global spread and uRMSE correspond to the final values (circles). Bottom row: Cumulative spectral SSR; the dotted line is perfect calibration and the global SSR corresponds to the final value (circle). Isotropic spherical white noise is shown in pink, the flat grid version is shown in grey. Curves are smoothed with a rolling average of size 7.

Figure 16: Random fields fitted to the ensemble error. Top: Samples drawn from the fitted distribution of different variables. Bottom: The PSD of the ensemble mean error (grey) and of the Matérn fit (purple, individual components dashed). Z500 has the most large scale structures (largest $\lambda$ ).

Figure 17: Spectral calibration of an 8-initialisation

12\text{\,}\mathrm{h}

ensemble rollout with flat white noise and the fitted Matérn fields. Top row: Zonal PSD of the ensemble spread (solid line) and unbiased ensemble mean error (dashed line). Z500 is unstable on large wavenumbers because of the log scale and

P^{u}_{\text{err}}(k>200)\approx 0

. Middle row: Cumulative PSD of the ensemble spread (solid line) and unbiased ensemble mean error (dashed line). The global spread and uRMSE correspond to the final values (circles). Bottom row: Cumulative spectral SSR; the dotted line is perfect calibration and the global SSR corresponds to the final value (circle). Matérn fields are shown in purple, the flat grid version is shown in grey. Curves are smoothed with a rolling average of size 7.

Figure 18: SSR and uRMSE against lead time for the white noise and fitted Matérn fields. Top: SSR. Bottom: uRMSE. Metrics are shown for different variables, the lead time of the fit is shown by the dashed vertical line (

12\text{\,}\mathrm{h}

), the white noise curves are in grey and the fitted ones in purple. Generally, the fitted random fields only fix the metrics at the lead time of the fit. The uRMSE is clearly negatively impacted, especially on longer lead times.

Figure 19: Convergence for fixed single-member budget. Unbiased metrics (uRMSE, CRPS, SSR) for Z500 and Q850 at 12 (solid) and 72 (dashed) hours lead time for different generative models and solvers (colours). Xaurora operating 20 NFE per member and physical step is shown in dotted vertical line. The x-axis shows the normalised compute budget (x=10 means 10 integration steps for EM but 5 for Heun).

Figure 20:

12\text{\,}\mathrm{h}

PSD averaged across ensemble members for fixed single-member budget. Ensemble-mean PSD at

12\text{\,}\mathrm{h}

for Z500 and Q850 for different generative models and solvers (columns). The operating Xaurora is DSI+Heun with 20 NFE per member and physical steps. The colour bar shows the normalised compute budget (x=10 means 10 integration steps for EM but 5 for Heun).

			SDE steps $S$
NFE	$M$	$M\times\mathrm{NFE}$	Heun	EM
2	64	128	1	2
4	50	200	2	4
10	20	200	5	10
20	10	200	10	20
50	4	200	25	50

Appendix

Table 8: Hyperparameter grid shared by all three configurations . The number of SDE steps

S

follows from the solver’s drift evaluations per step (Heun: 2, Euler–Maruyama: 1), hence we use NFE as normalised measure of integration steps. Ensemble sizes (

M

) are chosen such that every point costs the same, i.e.

M\times\mathrm{NFE}=200

. Only

\mathrm{NFE}=2

departs from this, capped at

M=64

. The main Xaurora configuration used in the paper is

\mathrm{NFE}=20

M=10

(Heun,

S=10

) and is shown in bold.

Figure 21: Spectral properties of different Xaurora training checkpoints. Top:

12\text{\,}\mathrm{h}

PSD averaged across 8 test set initialisations and 2 ensemble members, colours correspond to different checkpoints from initialisation (step 0) to the end of the training (step 7812). Matching ERA5 average spectra are shown with a dashed line. Bottom: Average power for wavenumbers greater than 100. This measures the flattening of the tail ( i.e. the noise floor) of the PSD for each checkpoint.

Figure 22: PSD of Xaurora and denoiser trained from scratch.

12\text{\,}\mathrm{h}

PSD averaged across 8 test set initialisations and 10 ensemble members. Xaurora is shown in green and the denoiser trained from scratch in orange. Matching ERA5 average spectra are shown with a dashed line.

Figure 23: PSD of Xaurora initialised with different priors.

12\text{\,}\mathrm{h}

PSD averaged across 8 test set initialisations and 10 ensemble members. Xaurora (official, prior is Aurora-S) is shown in green and Xaurora with Aurora-L as the prior is in orange. Matching ERA5 average spectra are shown with a dashed line.

Figure 24: PSD of Xaurora and other optimisation setups.

12\text{\,}\mathrm{h}

PSD averaged across 8 test set initialisations and 10 ensemble members. Xaurora is shown in green and other colours denote the different optimisation approaches discussed in the ablations. Matching ERA5 average spectra are shown with a dashed line.

Figure 25: Effect of the ensemble size on the biased ensemble mean RMSE. 5-day rollout evaluation of the RMSE (no ensemble size correction) for Z500 and Q850 and ensemble sizes of 1, 2, 5 and 10. With

M=1

, the

6\text{\,}\mathrm{h}

RMSE of Z500 is higher than the full ensemble. Q850 behaves as expected, the RMSE is consistently better with larger

M

. At a lead time of 5 day,

M=2

surpasses Aurora-S.

Figure 26: Effect of training from scratch vs finetuning on the ensemble mean uRMSE (unbiased). 5-day rollout evaluation of the uRMSE for Z500 and Q850 after training from scratch or finetuning from Aurora-S weights. The finetuned model performs much better. Both models were trained with "Single-step + MAE" (see Section B.3.1 ) because rollout training requires meaningful predictions at initialisation.

Figure 27: 15-day ablation Xaurora vs Single step + MAE. Top: uRMSE skill score of Single step + MAE compared to Xaurora; positive values mean that Single step + MAE is better than Xaurora, negative means Xaurora is better. Middle: Same for CRPS skill score. Bottom: SSR of Single step + MAE and Xaurora. Columns are different variables. Bands are 95% confidence intervals (see Section A.4.2 ).

Figure 28: Xaurora vs Aurora MC-DropPath CRPS scorecard. Each panel shows the CRPS improvement percentage from Xaurora over a naive probabilistic Aurora-S for one variable across lead times and pressure levels (for atmospheric variables). Both models have 10 ensemble members. Red indicates degraded performance (negative improvement), blue shows improvement over the reference.

Figure 29: Global medium-range weather forecasting RMSE (uRMSE for ensemble models) for all recorded variables in 2022 at

0.25^{\circ}

Figure 30: Global medium-range weather forecasting CRPS for all recorded variables in 2022 at

0.25^{\circ}

Figure 31: Global medium-range weather forecasting SSR for all recorded variables in 2022 at

0.25^{\circ}

Figure 32: 2-day PSD for all recorded variables in 2022 at

0.25^{\circ}

Figure 33: 10-day PSD for all recorded variables in 2022 at

0.25^{\circ}

Figure 34: Calibration maps at $12\text{\,}\mathrm{h}$ for a selection of variables.

Figure 35: Calibration maps at 5 days for a selection of variables.

Figure 36: Spatially-pooled calibration metrics at different lead times. Across lead times (shown in different line styles), pooling generally improves the CRPS score (top row) and tend to decrease the SSR (bottom row).

Figure 37: Rank histograms at different lead times. 2T and Z500 are overdispersive on short lead times (concave) as highlighted by the rest of the analysis. U500 and Q850 are underdispersive (convex).

Figure 38: Xaurora vs Aurora-S RMSE scorecard. Each panel shows the RMSE improvement percentage from Xaurora over Aurora-S for one variable across lead times and pressure levels (for atmospheric variables). We use uRMSE for Xaurora (unbiased) and the plain RMSE for Aurora-S. Red indicates degraded performance (negative improvement), blue shows improvement over the reference.

Figure 39: Xaurora vs Aurora-L RMSE scorecard. Reads as Figure 38

Figure 40: 3-day rollout initialised on January 1st, 2022 at 12:00 UTC. Columns correspond to ERA5, Xaurora first member (member 0), Xaurora ensemble mean (10 members), and Aurora-S. Variables (unnormalised) are shown on each row.

Figure 41: 10-day rollout initialised on January 1st, 2022 at 12:00 UTC. Columns correspond to ERA5, Xaurora first member (member 0), Xaurora ensemble mean (10 members), and Aurora-S. Variables (unnormalised) are shown on each row.

Figure 42: 15-day rollout initialised on January 1st, 2022 at 12:00 UTC. Columns correspond to ERA5, Xaurora first member (member 0), Xaurora ensemble mean (10 members), and Aurora-S. Variables (unnormalised) are shown on each row.

Explore similar work

Sep 29, 2026cs.LG

Proper Scoring Rule-based Diffusion for Probabilistic Weather Forecasting

Recent probabilistic weather forecasters train stochastic predictors with the continuous ranked probability score (CRPS) to generate each ensemble member in a single forward pass. These models learn the predictive distribution from the forecast context alone, which becomes difficult at longer forecast horizons where uncertainty is high. To learn the predictive distribution more effectively, we introduce auxiliary conditional denoising tasks that predict the same future state from the context and its corrupted version, which provides partial future information that can reduce prediction ambiguity. Building on distributional diffusion models, we learn the conditional distributions of these tasks with a single stochastic predictor by minimizing a proper scoring rule across noise levels. At inference, the predictor can still generate each ensemble member in a single forward pass at the fully corrupted endpoint. Standard CRPS training is recovered as the endpoint-only special case of our formulation, so our framework extends existing CRPS-based forecasters with only additional conditioning inputs. Controlled experiments show that the auxiliary tasks improve one-step forecasting across architectures, with larger gains at longer forecast horizons. The gains extend to high-dimensional global weather forecasting under both training from scratch and fine-tuning, along with improved calibration and potential benefits for generalization under distribution shift.

Joonhyeong Park, Giung Nam, Hyungi Lee +3

KAIST · Kookmin University · New York University

Numerical Weather Prediction Iterative Denoising Methods

Ask PDF

Jun 24, 2026cs.LG

Does Aurora Encode Atmospheric Structure? Latent Regime Analysis and Attribution

ML foundation models are able to emulate atmospheric dynamics accurately and efficiently but operate as opaque ``black boxes''. We investigate the internal representations of the Aurora model using spatially pooled PCA and layer-wise relevance propagation (LRP). We find evidence that Aurora's latent space is primarily organized by seasonal cycles, whereas extreme storm events do not form a linearly separable cluster. LRP indicates that the model attends to features consistent with the 3D vertical structure of the Great Storm of 1987. Perturbation tests show masking relevant regions degrades forecasts

3.31\times

more than random masking. These findings suggest that Aurora learns meteorological coherence and vertical structure without explicit instruction.

Emma Kasteleyn, Ana Lucic

University of Amsterdam

Atmospheric Dynamics Latent Space

Ask PDF

Jul 22, 2026cs.LG

Toward Mechanistic Interpretability of an AI Foundation Model Fine-Tuned for Atmospheric Chemistry

Weather forecasting foundation models (FMs) are increasingly fine-tuned to predict air quality, offering fast global pollution forecasts at lower computational cost than conventional chemical transport models. These FMs are typically trained on reanalysis data and generate forecasts through autoregressive rollout. They do not explicitly represent governing physical or chemical processes. Therefore, high forecast skill does not reveal whether a model has learned physical mechanisms or exploits statistical regularities in its training data. Here, we present the first study of what a FM fine-tuned for atmospheric chemistry has learned by examining Microsoft's Aurora model. We impose controlled chemical perturbations on its forecasts and test them against known photochemical relationships. We then examine the internal representations that generate these forecasts. We find that Aurora captures a first-order ozone response to reactive nitrogen but does not enforce the chemical constraints that a process-based model encodes. It generates chemically inconsistent combinations of related species and relaxes localized emission features such as wildfire plumes toward background. Internally, its representations remain largely organized around the meteorology inherited during pretraining, with little structure specific to chemistry. Using sparse autoencoders, we identify internal components that causally control the chemical forecast but do not map cleanly onto individual atmospheric processes. This work provides a framework for testing whether AI forecasting systems learn atmospheric chemistry from reanalysis data. As these models are increasingly positioned to inform environmental policy decisions, we argue that composition forecasts should also be judged by their internal mechanisms rather than by benchmark skill alone.

Jason Y. Hu, Ivan Higuera-Mendieta, Patrick Obin Sturm +1

Stanford University, Department of Computer Science, Stanford, CA, USA. · Stanford University, Department of Earth System Science, Stanford, CA, USA · University of Southern California, Department of Earth Sciences, Los Angeles, CA, USA +1

Atmospheric Dynamics Wireless Foundation Models Mechanistic Interpretability

Ask PDF