Observation Assimilation Transformer: Probabilistic Weather Analysis from Sparse Station Observations

Christian Donner; Michele Cattaneo; Maxim Samarin; Francesco Zanetta; Icíar Lloréns Jover; Gabriela Aznar Siguan; Michele Volpi

cs.LGOct 20, 2026

Observation Assimilation Transformer: Probabilistic Weather Analysis from Sparse Station Observations

Authors: Christian Donner, Michele Cattaneo, Maxim Samarin, Francesco Zanetta, Icíar Lloréns Jover, Gabriela Aznar Siguan, Michele Volpi

Organizations: Swiss Data Science Center, Eidgenössische Technische Hochschule Zürich (ETHZ) · Federal Office of Meteorology and Climatology MeteoSwiss, Switzerland

Abstract

We present the Observation Assimilation Transformer (OAT), a geometry-aware neural process that corrects and refines a prior numerical weather prediction (NWP) field into a high-resolution posterior field by assimilating sparse observations. OAT is an encoder-decoder transformer that considers a set of observed locations and their NWP prior to condition predictions at a set of arbitrary unobserved query locations. Its sparse attention mechanism makes it scalable to thousands of station observations and depends on relative positions between locations, never on the absolute coordinates. The model allows for partially missing observations, predictions at arbitrary continuous locations without grid projection, and outputs are parametric distributions. A station-randomized training algorithm is developed with continuous ranked probability score (CRPS) loss. We test OAT on two datasets - one at regional scale with hundreds of stations, and one at continental scale with thousands of stations. On these datasets OAT outperforms several baselines and shows well calibrated predictions. In an illustrative case study, OAT's downscaling capabilities are demonstrated by downscaling ERA5 (~31km) to COSMO (1km), where OAT achieves lower mean absolute error (MAE) and CRPS for temperature and precipitation on test stations compared to COSMO.

Figures & tables

	NWP prior	Geo-attn	Station SA	Pred. anywhere	Scalable
(Deep Kernel) GP ( Zanetta et al., 2026 )	✓	$\sim$	✓	✓	$\sim$
ConvCNP ( Gordon et al., 2019 )	✓	✗	✗	✓	✓
TE-TNP ( Ashman et al., 2024a ; Ashman et al., 2024b )	✗	✓	✓	✓	✗
Gridded TNP ( Ashman et al., 2025 )	✗	✗	$\sim$	✓	✓
SMACNP ( Bao et al., 2024 )	✗	$\sim$	✗	✓	✓
OAT (ours)	✓	✓	✓	✓	✓

Table 1: Structural comparison of OAT with closely related models. ✓ = present; ✗ = absent;

\sim

= partial.

Figure 1: Schematic of data and observation assimilation transformer (OAT). Left: Data snapshot with high-resolution topography, coarse NWP field, station locations (orange circles) and illustrative query locations (blue diamonds). Query locations can lie on a grid or at arbitrary locations. Right: Schematic of the OAT architecture. Topography and NWP fields are first embedded at observed and query locations in the Pre-OAT stage . Then observed embeddings are updated with observation values resulting in the informed station tokens . The OAT block comprises a self-attention encoder on the observed embeddings, and a cross-attention decoder between the observed embeddings and query embeddings. The query embeddings, informed by the observational context, are then producing a probabilistic prediction via an output head to produce the final output distributions . The same shared output head can provide predictions on query tokens before going through the OAT block.

Figure 2: Dataset extent and station distribution The dashed line shows the extent for which the NWP product is available. The square, triangle, and circle markers are train, validation, and test stations, respectively. Left: The Switzerland dataset. Center: Europe dataset. Right: Same data as the Europe dataset, but with a regional hold-out split, which coincides with small dashed box on the right representing the extent Switzerland dataset.

Figure 3: Example of predicted fields From left to right: The interpolated NWP, the pre-OAT predicted, and the OAT field with assimilated and test stations. The color code of the stations are the measurement values. Rightmost is the difference between the pre-OAT and OAT field, where the station color code is between pre-OAT and measurement at that station. Top row: Temperature fields. Bottom row: Precipitation fields. The prediction fields are sampled on a regular $250\text{\,}\mathrm{m}$ grid.

		Temperature			Precipitation
Dataset	Model	MAE $\downarrow$	CRPS $\downarrow$	Skill $\uparrow$	MAE $\downarrow$	CRPS $\downarrow$	Skill $\uparrow$
Switzerland	COSMO Interp.	$1.287$	$1.287$	$-0.637$	$0.146$	$0.146$	$-0.587$
	Uninf. baseline	$1.089$	$0.786$	$0.000$	$0.122$	$0.092$	$0.000$
	GP	$1.109$	$0.815$	$-0.037$	$0.130$	$0.111$	$-0.206$
	DK-GP	$1.125$	$0.827$	$-0.052$	$0.119$	$0.101$	$-0.094$
	Gridded TNP	$1.429$	$1.115$	$-0.418$	$0.097$	$0.073$	$0.207$

Table 2: Model performance on the test set. Skill is computed with CRPS and the comparison is to the uninformed baseline.

Panel of Figure 4 — Figure 4: Calibration for the OAT on the test set. In each panel, from left to right: the PIT histogram for temperature, the reliability diagram for the probability of zero precipitation, and the precipitation PIT histogram for wet events.

Panel of Figure 5 — Figure 5: Stratification analysis of CRPS From left to right, the CRPS on the test set as function of season, elevation, and distance to closest observed station. The white bar denotes the median, and the boxes are the first and third quartile. Since the uninformed baseline is not conditioned on any observation, it is not displayed in the distance panels.

Figure 6: Region hold-out and downscaling temperature from the ERA5 grid to the COSMO grid ( $\sim$ 1 million points) on 2024-06-21 12:00 UTC for Swiss domain. In the first row, the input ERA5 field, the uninformed mean prediction, and the OAT mean prediction. In the second row, the corresponding COSMO analysis field as well as the differences between the informed and uninformed models against that analysis. The third row shows the test stations and the corresponding prediction error compared to the observed temperature for COSMO, the uninformed baseline, and OAT.

Figure 7: OAT’s predictive performance for exemplary extreme events. Top: Extreme precipitation event causing flooding in the Maggia valley end of June 2024. Left: Cumulated precipitation map over 24h for NWP, OAT and their difference, respectively. Right: Time-series of precipitation at two test-stations. Bottom: Heat wave in July 2024 in south eastern Europe. Left: Temperature maps of NWP (COSMO/ERA5) and OAT, with differences. Right: Two test-stations for one week.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Hyperparameter	Value
NWP encoder	Type	pointwise MLP, GELU
	Input channels	4 (CH) / 3 (EU)
	Hidden layers	1, width 16
	Output channels	16
Topographic encoder	Type	pointwise MLP, GELU
	Input channels	5 (CH) / 2 (EU)

Appendix

Table A.1: OAT architecture hyperparameters. CH: Switzerland dataset, EU: Europe dataset. Values are identical for both datasets except where noted. The uninformed baseline has only the encoders, the NWP–topography fusion and the output head. It uses widths 32 in the encoders and 64 in the fusion layer and output head, and has

\approx 20

k parameters.

Setting	OAT / Uninf.	GP / DK-GP	Gridded TNP
Optimizer	AdamW	AdamW	Adam
Learning rate	$10^{-3}$	$10^{-2}$ (CH) / $5\times 10^{-3}$ (EU)	$5\times 10^{-4}$
Weight decay	$0.01$	$0.01$	—
LR schedule	cosine annealing, $T_{\max}=500$ epochs c		none
Batch size (time steps) a	$1024$ (CH) / $256$ (EU) c		$1$ b
Gradient clipping (norm)	$1.0$	$1.0$	$0.5$

Appendix

Table A.2: Training configuration. CH: Switzerland dataset, EU: Europe dataset. a Per GPU; effective batch size

\times 4

. One sample is one hourly time step with all its stations. b The Gridded TNP generator emits one sample per step with

10\,000

steps per epoch. c The region hold-out models use a batch size of

32

T_{\max}=100

and at most

100

epochs.

	Temperature			Precipitation
Method	MAE $\downarrow$	CRPS $\downarrow$	Skill $\uparrow$	MAE $\downarrow$	CRPS $\downarrow$	Skill $\uparrow$
OAT (additive)	$1.046$	$0.751$	$0.045$	$0.086$	$0.064$	$0.306$
OAT (vanilla)	$1.076$	$0.771$	$0.019$	$0.088$	$0.066$	$0.287$
OAT (TE-TNP)	$1.035$	$0.743$	$0.054$	$0.085$	$0.064$	$0.310$

Appendix

Table A.3 : Ablation of different attention mechanisms. Evaluated on the Switzerland dataset.

Component	Hyperparameter	Value
Feature trunk	Architecture	as uninformed baseline, no output head
Feature trunk	Embedding dim.	64
Mean function	MLP on embedding	$64\to 64\to 1$ , ReLU
Stationary kernel	Type	RBF with ARD
	Inputs (CH)	easting, northing, elevation
	Inputs (EU)	longitude, latitude, elevation

Appendix

Table A.5: GP and DK-GP hyperparameters. Values are identical for both GP variants and for the Switzerland (CH) and Europe (EU) datasets except where noted. Kernel inputs are standardized. For the training configuration see Table A.2 .

Component	Hyperparameter	Value
Latent grid	NWP stride	10 (CH) / 4 (EU)
	Pseudo-token grid	$80\times 120$ (CH) / $40\times 72$ (EU)
	NWP input channels	4 (CH) / 3 (EU)
Position encoding	Type	spherical harmonics (absolute lon./lat.)
Position encoding	Legendre polynomials / embedding dim.	10 / 128
Token encoders	Observation / NWP MLP	2 layers, width 128

Appendix

Table A.6: Gridded TNP hyperparameters. CH: Switzerland dataset, EU: Europe dataset. For the training configuration see Table A.2 .

Panel of Figure 15 — Figure A.1 : Scalability analysis of OAT and GP. We empirically measure the time per step under two scenarios (the rows in the figure) as the number of observed stations grows up to $2^{16}\approx 65^{\prime}000$ . In the first scenario, there is a fixed number of 1024 query points, while the second scenario has a number of query points proportional to the observed stations to maintain a 80%-20% ratio. The columns show the results for a training step, which includes forward and backward pass, and an evaluation step. A major bottleneck of the OAT’s forward pass is the nearest neighbour computation, which we include in the plot as a dashed line. For the GP, we compare two ways of computing its predictive distribution: GP (exact) relies on a Cholesky decomposition, whose cost grows cubically with the number of points, while GP (iterative) uses GPyTorch’s default, which switches from Cholesky to iterative conjugate-gradient and Lanczos methods above 800 points.

Panel of Figure 16 — Figure A.2 : Comparison between a full-attention scheme and the default sparse attention scheme used in the OAT model to perform training and inference steps. The OAT setting is the same as in Figure A.1 , with the only difference that the full-attention scheme is turned on.

Data Assimilation

Data assimilation combines observations with dynamical models to estimate hidden system states and their uncertainty over time.

See all 58 papers

FuXiWeather2: Learning accurate atmospheric state estimation for operational global weather forecasting
Mar 16, 2026Xiaoze Xu, Xiuyu Sun, Songling Zhu +5Weather Forecasting Data Assimilation
Shanghai Academy of Artificial Intelligence for Science, Shanghai, & 200232, China. · Artificial Intelligence Innovation and Incubation Institute, Fudan University, Shanghai & 200433, China. · FuXi Intelligent Computing Technology Co., Ltd., Shanghai & 200233, China.
Generative 3D Gaussian Splatting for Arbitrary-ResolutionAtmospheric Downscaling and Forecasting
Apr 9, 2026Tao Han, Zhibin Wen, Zhenghao Chen +4Weather Forecasting 3D Gaussian Splatting
Department of Computer Science and Engineering, The Hong Kong University of Science and Technology, 999077, Hong Kong, China · Shanghai Artificial Intelligence Laboratory, Shanghai, 200232, China · Department of Computer Science and Engineering, Southern University of Science and Technology, Shenzhen, 518055, Guangdong, China +2
UtoMe: Observation-Uncertainty-Guided Token Merging for Weather Foundation Models
Oct 18, 2026Yunbei Zhang, Janet Wang, Xi Xiao +2Token Merging Weather Forecasting
Tulane University · Oak Ridge National Laboratory

Mar 16, 2026J/K move · Enter open · S save

FuXiWeather2: Learning accurate atmospheric state estimation for operational global weather forecasting

Xiaoze Xu, Xiuyu Sun, Songling Zhu, Xiaohui Zhong+4

Shanghai Academy of Artificial Intelligence for Science, Shanghai, & 200232, China. · Artificial Intelligence Innovation and Incubation Institute, Fudan University, Shanghai & 200433, China. · FuXi Intelligent Computing Technology Co., Ltd., Shanghai & 200233, China.

Weather Forecasting Data Assimilation

Open paper AskPDF arXiv

Numerical weather prediction has long been constrained by the computational bottlenecks inherent in data assimilation and numerical modeling. While machine learning has accelerated forecasting, existing models largely serve as "emulators of reanalysis products," thereby retaining their systematic biases and operational latencies. Here, we present FuXiWeather2, a unified end-to-end neural framework for assimilation and forecasting. We align training objectives directly with a combination of real-world observations and reanalysis data, enabling the framework to effectively rectify inherent errors within reanalysis products. To address the distribution shift between NWP-derived background inputs during training and self-generated backgrounds during deployment, we introduce a recursive unrolling training method to enhance the precision and stability of analysis generation. Furthermore, our model is trained on a hybrid dataset of raw and simulated observations to mitigate the impact of observational distribution inconsistency. FuXiWeather2 generates high-resolution (

0.25^{\circ}

) global analysis fields and 10-day forecasts within minutes. The analysis fields surpass the NCEP-GFS across most variables and demonstrate superior accuracy over both ERA5 and the ECMWF-HRES system in lower-tropospheric and surface variables. These high-quality analysis fields drive deterministic forecasts that exceed the skill of the HRES system in 91% of evaluated metrics. Additionally, its outstanding performance in typhoon track prediction underscores its practical value for rapid response to extreme weather events. The FuXiWeather2 analysis dataset is available at https://doi.org/10.5281/zenodo.18872728.

Abstract

Data Assimilation

Explore similar work

FuXiWeather2: Learning accurate atmospheric state estimation for operational global weather forecasting

Generative 3D Gaussian Splatting for Arbitrary-ResolutionAtmospheric Downscaling and Forecasting

UtoMe: Observation-Uncertainty-Guided Token Merging for Weather Foundation Models