Multimodal Spatiotemporal Atmospheric Data Assimilation with Latent Flow-matching

M ULTIMODAL SPATIOTEMPORAL ATMOSPHERIC DATA ASSIMILATION WITH LATENT FLOW- MATCHING

arXiv:2608.05103v1 [cs.LG] 5 Aug 2026

Dibyajyoti Chakraborty Information Sciences and Technology The Pennsylvania State University University Park, Pennsylvania d.chakraborty@psu.edu

Romit Maulik School of Mechanical Engineering Purdue University West Lafayette, Indiana rmaulik@purdue.edu

A BSTRACT Data assimilation (DA) uses Bayesian inference to update the state of a numerical forecast model with observed data. In this study, we propose a fundamentally different, unified approach to atmospheric data assimilation. We use latent video flow-matching to sample temporally consistent trajectories from a prior trained using ERA5 reanalysis (69 variables over an 8-day window). We also use posterior sampling to assimilate real observation sources, such as those from the NOAA Integrated Global Radiosonde Archive and the Integrated Surface Database. Because the prior generates a continuous trajectory, it naturally propagates information between observed and unobserved frames. Therefore, we can perform various DA tasks, such as filtering and smoothing, simply by changing the observed frames. Moreover, we generate full-state ensemble forecasts directly from sparse observations, achieving performance competitive with state-of-the-art observation-to-forecast models. Keywords Generative machine learning · Latent generative models · Data assimilation · Atmospheric super-resolution · Observation-to-forecast

1

Introduction

Earth’s atmospheric dynamics span a wide range of spatial and temporal scales [1, 2, 3]. Forecasting such a system is difficult because it can only be simulated approximately and requires accurate initial conditions. However, atmospheric states can only be partially observed using specialized instruments [4, 5]. This creates a need for data and model fusion, in which first-principles numerical models provide a prior (background) state that is updated given partial observations through a Bayesian formulation [6, 7]. This update, known as data assimilation (DA), is central to modern weather prediction and underpins reanalysis products such as ERA5 [8]. However, these classical variational and ensemble methods [9, 10, 11] are computationally expensive and require several iterations of expensive numerical models. Generative machine learning [12, 13] offers an alternative [14]. Instead of re-running a numerical model, one learns an implicit probability density of the atmospheric state from which new samples can be drawn, and conditions that density in real time on sparse observations. Diffusion models [15, 16, 17] learn such densities by corrupting data with noise through a forward stochastic process and learning the score (the gradient of the log-probability density), so that the reverse process can iteratively denoise pure noise into samples from the data distribution. Flow-matching [18] is an adjacent method that learns a velocity field, which is used to directly interpolate from noise to the data distribution. Recent advances in the computer vision community have shown that flow-matching techniques have an edge over previous methods in high-dimensional generative tasks [19]. A trained generative model also enables zero-shot conditioning: the pretrained unconditional prior can be guided by sparse observations without retraining [16, 20, 21]. This is effectively equivalent to sampling from the posterior distribution. Previous work has demonstrated this for ERA5 snapshots, combining coarse-grid fields, NOAA Integrated Global Radiosonde Archive (IGRA) data, and a climate emulator for super-resolution and multimodal reconstruction [22]. There, temporal consistency came only from conditioning on the output of a separate emulator. In this work, temporal propagation is internal to the prior itself. This has two consequences that are absent in a single snapshot generation

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

setting. First, the prior inherently propagates information from observed frames to unobserved frames, removing the need for a separate emulator or numerical forecast model. Second, different classical DA regimes are accurately recovered simply by choosing which frames (or data) are observed. The frames here refer to the individual timesteps in the video. We summarize our contributions as follows. 1. We train an unconditional latent video flow-matching prior (TrigFlow [23] with a 3D diffusion-transformer backbone [24]) that generates a 32-frame (8-day, 6-hourly) window of the global atmosphere for 69 ERA5 variables. This serves as a background model for a wide range of data-assimilation experiments without any retraining or observation-specific conditioning. 2. On the methods side, we adopt an adaptive noise-level weighting for the TrigFlow loss. We also explore a broad set of guided-sampler designs, utilizing low-noise-level correctors and gradient (and noise) rescaling techniques for both efficiency and accuracy. 3. On the experimental side, we recover the full atmospheric state from both structured coarse observations and real observations from the NOAA Integrated Global Radiosonde Archive (IGRA) [25], the NOAA Integrated Surface Database (ISD) [26], and the International Comprehensive Ocean-Atmosphere Data Set (ICOADS) [27]. We use these to realize the three classical DA regimes (filtering, smoothing, and fixed-interval reconstruction) from the same prior by changing only which frames are observed. We further use the prior for direct observation-to-forecast, producing a six-day forecast competitive with some prior observation-to-forecast models (GraphDOP [28]). Finally, we study its behavior in a particular case, Hurricane Laura (2020). Our results indicate that latent video flow-matching priors can perform accurate multimodal spatiotemporal DA from multiple real data sources, with uncertainty quantification through the spread of posterior draws. An end-to-end comparison with operational assimilation is beyond our scope. The codes used in the work are in https://github.com/ISCLPurdue/era5_videogen. 1.1

Related work

Deep learning has produced strong results in reconstructing the state of high-dimensional multiscale dynamical systems [29, 30, 31, 32, 33, 34, 35]. Specifically, Graph Neural Networks (GNNs) or Voronoi-tessellation-assisted methods have been used for reconstruction from unstructured sensor measurements [36, 37]. Most of these demonstrations, however, rely on coarse grids extracted from the fine state, so inputs and targets come from the same dataset. Additionally, they require retraining when the observation source changes and, being deterministic, lack uncertainty quantification (UQ). Probabilistic networks [38], deep ensembles [39], and stochastic weight averaging [40, 41] provide UQ at the cost of restrictive assumptions or increased computation, which motivates generative priors. Score-based diffusion models are now the dominant generative models for physical systems, beginning with diffusion probabilistic models [42, 43] and score-based generative models [44] (building on score matching [45]), unified through stochastic differential equations [16] and refined by the EDM design space exploration [17]. Flow-matching [18, 46, 47] uses a different construction, prescribing an interpolation path from noise to the data distribution directly rather than reversing a noising process. These priors solve inverse problems in a zero-shot, plug-and-play manner through gradient guidance (DPS [21], ΠGDM [48]), range-null-space methods (DDRM [49], DDNM [50]), variational and variablesplitting solvers (RED-diff [51], DiffPIR [52], PnP-DM [53], DAPS [54]), and latent-space variants [55, 56]. The I NVERSE B ENCH benchmark [57] compares many of these. We also re-implement some of these methods on our prior in App. A.5. In atmospheric applications, they have been applied for downscaling and super-resolution [58, 59, 60], probabilistic forecasting [61, 14, 62], and data assimilation [63, 64, 65, 66, 67, 68, 69]. Similar to us, Qu et al. [66] target a multimodal setting, and Liang et al. [69] and Bao et al. [68] pursue score-based filtering, but on single-time or lower-dimensional states rather than a latent video trajectory of the full ERA5 state. SEEDS [61] generates skillful forecast ensembles cheaply but emulates an ensemble rather than assimilating observations. Manshausen et al. [64] assimilate sparse station data from a single source on a single snapshot, while we assimilate several independently sourced observations over an 8-day trajectory and infer the full vertical column of the atmosphere. Spatiotemporal diffusion priors for scientific video inverse problems [70] and Appa Weather [71] are the closest to our work, and we differ from each as follows. STeP [70] performs spatiotemporal latent posterior sampling but is evaluated on simulated observations. Similarly, Appa uses a latent diffusion prior and synthetic observations for global data assimilation and forecasting. Here, we use real observations to perform various data-assimilation tasks and show over three weeks of competitive forecast performance directly from observations. 2

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Autoencoder

TrigFlow

noising →

← denoise

;

Guided sampling

NOTATION

— ERA5 fields — latent video

×N

— encoder, decoder — flow prior (NN model)

— Gaussian noise — noise angle, level — observations, operator — guidance, via

posterior score = prior score + likelihood score

−

— adaptive guidance

Figure 1: Overview of the proposed method. An autoencoder (E, D) compresses 8-day (32 timesteps) ERA5 windows (69 variables at 128×256 grid) into a latent video z ∈ R128×8×32×64 . The TrigFlow prior F is trained by velocity matching along the spherical noising path zt = cos t z0 + sin t ε. From approximately pure noise, N guided steps sample the posterior: each step decodes the latent Tweedie estimate ẑ0 , scores it against the observations y through the measurement operator A, and injects the guidance term ℓ into the update, with the annealed likelihood variance 2 σlik weakening guidance at high noise. D and F stay frozen while sampling, so a single pretrained prior serves every assimilation task in this paper zero-shot. Note: The latent video looks like Gaussian noise visually; however, it has some small structures that the decoder can decode successfully.

2

Methods

A high-level view of the full methodology is given in the schematic of Fig. 1: an autoencoder compresses spatiotemporal ERA5 data into a latent video, a TrigFlow prior is trained on those latents, and at sampling time the learned prior is combined with an observation likelihood to draw posterior samples. The subsections below describe each ingredient in turn. 2.1

Autoencoder.

We train a 3D-convolutional autoencoder (AE) (E, D) to compress the ERA5 data x to a latent representation z for faster training and sampling. The decoder D maps the latent back to the 69 standardized ERA5 variables at T =32, H=128, and W =256. The autoencoder has a 4× temporal and 4× spatial compression (each direction) for the latent grid. The generative prior is trained on the latents of the autoencoder (similar to a video) of shape (C=128, T =8, H=32, W =64). Because likelihoods are evaluated in the physical field space, guidance passes through the decoder as in latent-diffusion inverse solvers [55, 56]. 2.2

Generative modeling

Flow-matching. Flow-matching [18, 47] prescribes an interpolation path xt = αt x0 + βt ε between data and Gaussian noise. The network is trained to regress the velocity dxt /dt along that path. During sampling, new samples are generated by simulating the ordinary differential equation (ODE) dxt = vθ (xt , t)dt from t = T (pure noise) to t = 0. Various flow models correspond to particular choices of schedule (αt , βt ). TrigFlow. The TrigFlow parameterization [23] is a flow-matching formulation whose interpolation between data and noise is spherical. Our prior is trained not on the ERA5 fields (x) themselves but on the latent z of an autoencoder. The forward interpolation path for the noise is defined as: zt = cos(t) z0 + sin(t) ε,

ε ∼ N (0, σd2 I), t ∈ (0, π2 ). 3

(1)

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

We fix σd =1 throughout. Differentiating Eq. (1) with respect to time gives the target velocity dzt /dt = − sin(t) z0 + cos(t) ε. A neural network F (zt /σd , t), here a 3D diffusion transformer (DiT3D [24]), is trained so that σd F (zt /σd , t) ≈ dzt /dt. At the optimum of the training loss, the learned velocity yields an estimate of the clean latent, which coincides with the Tweedie (minimum mean squared error) estimate [72], ẑ0 = cos(t) zt − sin(t) σd F (zt /σd , t).

(2)

Sampling integrates the ODE dxt = vθ (xt , t)dt from tmax = arctan(σmax /σd ) (approximately pure noise) to t = 0 (data). Adaptive noise-level weighting. The baseline flow model is trained to match the velocity by minimizing the expected squared error: i h L(θ) = Ez0 ,ε,t

t σd F (zt /σd , t; θ) − dz dt

2

.

(3)

The velocity target contains the noise ε, so the regression has a t-dependent irreducible variance. Since no network can predict the specific noise component, the loss floor is disproportionately large at certain noise levels. Under a fixed weighting, high noise levels dominate the gradient, which hampers learning. We therefore utilize an adaptive uncertainty weighting [73], where the network additionally predicts a scalar log-variance u(t) through a small linear head on its noise-level embedding. The loss is formulated as the maximum-likelihood objective for a Gaussian error model with level-dependent variance:   2 t + u(t) . (4) L(θ) = E e−u(t) σd F (zt /σd , t; θ) − dz dt At the optimum, this objective sets eu(t) to the expected squared error at level t, ensuring that every noise level contributes gradients of comparable magnitude regardless of its floor. Training noise levels are drawn from a log-normal distribution centered on the data scale, ln σ ∼ N (0, 1.52 ) with σ = tan t. The u head is discarded at inference, so sampling is unaffected. Posterior sampling and data fusion. To sample the posterior p(x | y) we draw latents from p(z | y) and decode them, utilizing Bayes’ theorem to combine the pretrained unconditional prior with a measurement likelihood. The posterior score is the sum of the prior and likelihood scores, ∇zt log p(zt | y) = ∇zt log p(zt ) + ∇zt log p(y | zt ),

(5)

structurally analogous to a classifier guidance [74]. For a differentiable measurement operator A and Gaussian observation noise with standard deviation σy , DPS [21] evaluates the likelihood at the Tweedie estimate ẑ0 , log p(y | zt ) ≈ − 2 σ21 (σ) y − A(D(ẑ0 ))

2

,

(6)

lik

and score-based data assimilation (SDA) [75] accounts for the variance of ẑ0 through an annealed likelihood variance, 2 σlik (σ) = σy2 + γ (σ/σd )2 ,

(7)

where γ is the SDA annealing coefficient and σd is the latent data scale (fixed to 1 here, so σ/σd =σ is the running noise level while sampling). This is used to automatically weaken guidance at high noise levels. 2.3

Guided TrigFlow sampler

All sampling methods in this article are built on our guided TrigFlow sampler (Algorithm 1), which combines three base techniques. First, it uses the Elucidating Diffusion Model (EDM) noise schedule [17], with each noise level mapped to a TrigFlow angle through t = arctan(σ/σd ). Second, it advances the state with the DPM-Solver++ update [76], which in TrigFlow coordinates is a rotation (the second-order coefficient is derived in Appendix A.1e). Third, at every step it adds DPS-style measurement guidance with SDA annealing. The sampler has an option to follow a predictor-corrector structure [16]. The guided rotation acts as the predictor: it transports the sample from one noise level to the next, with the guidance term bending the trajectory toward measurement consistency. However, since the error of this predictor step is carried to the next level, a Langevin corrector (Algorithm 2) holds the noise level fixed and takes one Langevin step toward the intermediate posterior p(zt | y) ∝ p(zt ) p(y | zt ), using a fresh denoiser evaluation and a fresh likelihood gradient. Since each corrector step costs one extra denoiser c evaluation, we apply it only below a threshold σmax , where the Tweedie estimate is sharp and measurement fidelity is actually established (Sec. 2.3.1). 4

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Algorithm 1: Guided TrigFlow sampler. Tunable options and their defaults are described in Sec. 2.3. input :measurement y, steps N , σd =1, scale; angles t0 > . . . >tN =0 at σ levels, tk = arctan(σk /σd ); clip(·) is the elementwise clip to [−1, 1] (used in all runs) prev 1 zt ← σd · ε, ε ∼ N (0, I); ẑ0 ← ∅; ℓm ← 0; 2 for k = 0, . . . , N − 1 do 3 s ← tk , t ← tk+1 , δ ← s − t, σ ← σd tan s; 4 Fs ← F (zt /σd , s); // denoiser velocity 5 ẑ0 ← cos(s) zt − sin(s) σd Fs ; // Tweedie estimate, Eq. (2) g g 6 if σmin ≤σ≤ then  σmax  7 g ← ∇zt − 2σ12 ∥y − A(D(ẑ0 ))∥2 ; // backprop through D and F lik √ // n=CT HW 8 ℓ ← scale · σ · clip(g) or scale · σ · clip( n g/∥g∥2 ) under DSG; 9 ℓ ← ℓ + m ℓm ; ℓm ← ℓ; // momentum on the scaled score 10 else 11 ℓ ← 0; // guidance band gate z(1) ← cos(δ) zt − sin(δ) σd Fs ; 13 if 0 < k < N − 1 and ẑprev ̸= ∅ then 0 14 r ← (log tan s − log tan tprev )/(log tan s − log tan t); 15 znext ← z(1) + 2 rsinsinδ s (ẑprev − ẑ0 ); 0 16 else 17 znext ← z(1) 18 znext ← znext + sin(δ) σd ℓ; 19 ẑprev ← ẑ0 ; 0 c then 20 if corrector and k < N − 1 and σd tan t ≤ σmax 21 znext ← L ANGEVIN C ORRECTOR(znext , t); 22 ẑprev ← ∅; 0 23 tprev ← s; zt ← znext ; output : zt

// 1st-order rotation

12

// 2nd-order

// inject guidance, F → F − ℓ

// Alg. 2 // 2nd-order unused for next step

Algorithm 2: L ANGEVIN C ORRECTOR(z, s): one SDA-style Langevin step at angle s toward the posterior score. 1 F ← F (z/σd , s); ẑ0 ← cos s z − sin s σd F ; // σ=σd tan s  1  2 2 g ← ∇z − 2 ∥y − A(D(ẑ0 ))∥ ; ℓ ← scale · σ clip(g); // clip/DSG as in Alg. 1; no momentum 2σ lik

spri ← (cos s ẑ0 − z)/(sin2 s σd2 ); 4 slik ← ℓ/(σd tan s); 2 5 η ← (snr · sin s σd ) ; √ 6 z ← z + η(spri +slik ) + λ 2η ε, ε ∼ N (0, I); output :z

// prior score // likelihood score

3

2.3.1

// λ=1:

sample; λ→0:

mean-seeking

Design choices

We have the following tunable options in our sampler. • scale: guidance scale (the strength of the likelihood gradient); g g • [σmin , σmax ]: guidance band, outside which no guidance is applied;

• m: momentum decay on the guidance term; • DSG: normalization of the guidance gradient; c • λ, σmax : noise scale and activation threshold of the Langevin corrector.

We highlight the following designs which are combinations of the above choices. Table 1 mentions them with their design choice settings and Number of Function Evaluations (NFE). Extended reasoning behind these designs is given in Appendix A.1. All the hyperparameters have been tuned using the ablation studies in Appendix A. The DPS mentioned in the rest of the paper is DPS with the annealed likelihood variance of Eq. 7. 5

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Low noise level correction. Every solver step and every corrector step costs one denoiser forward pass, plus one backward pass when guidance is active. At high σ the Tweedie estimate ẑ0 is blurry, and SDA annealing further shrinks the likelihood gradient, so guidance accomplishes little there. Measurement fidelity is only established near t→0. The reduced-step corrector configurations (Table 1) therefore reallocate the NFE budget. The Langevin corrector steps are restricted to the low-σ band (Algorithm 2). The configurations we use further set the step count N and corrector gate c σmax to c • DPS+corr (N25): (N, σmax ) = (25, 3); c • DPS+corr (N30): (N, σmax ) = (30, 5).

Reshaping gradients and noise.

We modify how the sampler uses the gradients and noise in the following ways. √ • λ noise-scaling [77] multiplies the corrector’s injected noise, λ 2η ε in Algorithm 2, by a factor λ ∈ [0, 1]. This interpolates between posterior samples at λ=1, the default, and a deterministic mean-seeking trajectory at λ=0, used by the mean-seeking configurations of Table 1. • Momentum guidance accumulates the scaled likelihood score across steps, ℓt = scale · σt gt + m ℓt−1 , with decay m ∈ [0.3, 0.5] (Algorithm 1). Unlike lazy gradient reuse [78], the gradient is still recomputed at every step; momentum only averages successive directions. The accumulator starts at ℓ−1 =0, so the first step is plain DPS. One variant uses m=0.5; the other momentum variants use m=0.3 (Table 1). • DSG (diffusion with spherical-Gaussian √ constraint) [79] unit-normalizes the guidance gradient per sample and rescales it to magnitude n, where n=CT HW is the number of latent elements (here n=128·8·32·64=2,097,152), before the elementwise clip and the guidance scale are applied as in Algorithm 1. This fixes the guidance step size regardless of the raw gradient magnitude.

Table 1: Configuration of sampler variants used in the results, expressed as the switches of Algs. 1 and 2 hyperparameters: c solver steps N , corrector gate σmax , momentum m, corrector noise scale λ, and DSG normalization. NFE = solver steps + corrector evaluations. Tags are cumulative additions to the DPS baseline; (N25)/(N30) denote reduced-step variants and (all σ) an ungated corrector. Sampler DPS +mom0.5 +corr +corr+λ0 +corr+λ0+mom +corr+DSG +corr+DSG+λ0 +corr(all σ) +corr(all σ)+λ0 +corr(all σ)+λ0+mom +corr (N25) +corr (N30) +corr (N30)+λ0

2.4

steps N

corrector

m

λ

DSG

NFE

50 50 50 50 50 50 50 50 50 50 25 30 30

σ<3 σ<3 σ<3 σ<3 σ<3 all σ all σ all σ σ<3 σ<5 σ<5

0 0.5 0 0 0.3 0 0 0 0 0.3 0 0 0

1 0 0 1 0 1 0 0 1 1 0

✓ ✓ -

50 50 76 76 76 76 76 99 99 99 38 47 47

Measurement operators

For our posterior sampling, we change only the measurement operator A entering the likelihood. We broadly divide the measurement operators used here into these three categories. • Structured (gridded) subsampling. For SR-type tasks with a structured observation, A selects the observed frames of the decoded video and spatially subsamples them by the coarsening factor. • Sparse point interpolation. For station datasets, A bilinearly interpolates the decoded fields at the reported [lat, lon] locations on the observed frames. The observations are fused by taking their union. • Super-obbed point observations. For the dense, strongly clustered observations (ISD and the ICOADS), we use grid-cell super-obbing, the standard assimilation thinning step [80, 81]. All raw observations in a cell are collapsed to a single super-observation weighted by their distance from the cell centers. All likelihoods take the form of Eq. (6)-(7). 6

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

3

Experimental configuration

3.1

Data

We use ERA5 reanalysis [8] regridded to 1.40625◦ (128×256) from WeatherBench [82], organized into 32-frame 6-hourly windows of 69 variables (surface fields plus geopotential, wind, temperature, and specific humidity on 13 pressure levels). The flow prior is trained on autoencoder latents of these windows. The years 1979 − 2018 are used for training, 2019 for validation, and 2020 for testing. For observations we use the NOAA Integrated Global Radiosonde Archive (IGRA) [25] and the Integrated Surface Database (ISD; the ISD-Lite subset) [26]. IGRA provides on average ∼300 active stations per observed frame (ranging from ∼600 at the 00/12 UTC launch times down to a few tens at 06/18 UTC; ∼900 distinct sites), reporting 63 of the 69 variables: the three near-surface fields (2 m temperature, 10 m u, v wind) and, on every one of the thirteen pressure levels, geopotential, zonal and meridional wind, and temperature, together with specific humidity at 300 hPa and below. The soundings carry no mean sea-level pressure, and we discard radiosonde humidity above 300 hPa, where the sensors are unreliable and the reports correlate poorly with the reanalysis. ISD provides four surface variables (2 m temperature, 10 m u/v wind, mean sea-level pressure) from ∼2500 globally grid-thinned stations per step (∼3100 distinct sites, ∼9400 observations per 6 h step): a dense but surface-only observation. For the observations-to-forecast experiments, we additionally ingest the International Comprehensive Ocean-Atmosphere Data Set (ICOADS) [27] ship and buoy reports (marine air temperature, wind, and pressure), which extend the surface coverage over the oceans. The ERA5 data closely match the observations for the temperature fields but less so for the winds, whose local behavior cannot be captured by ERA5 at ≈1.4◦ . Although we note that these observations are closer to the ground truth at those stations (up to measurement accuracy), we use ERA5 data for the evaluation of global reconstructions. 3.2

Autoencoder

The flow-matching prior operates in the latent space of a 3D-convolutional autoencoder trained separately on the same ERA5 windows. The encoder is a ResNet-style 3D network with four resolution levels (hidden widths [96, 192, 384, 768], three residual blocks per level) that maps a 69-variable, 32-timeframe window at 128×256 down to a latent of shape (C=128, T =8, H=32, W =64) via the anisotropic strides [(1, 2, 2), (2, 2, 2), (2, 1, 1)], a 4× temporal and 4× spatial compression; the decoder mirrors p the encoder with one extra residual block per level. The latent is passed through a smooth transformation x 7→ x/ 1 + (x/10)2 that bounds its range without hard clipping. Two static fields, the land-sea mask and the surface geopotential, are area-resampled to each decoder resolution and added to the decoder through zero-initialized 3×3 convolutions. This ensures that the decoder need not expend its capacity to learn the fixed surface structures. The autoencoder is trained deterministically for 155 epochs with Adam [83], a base learning rate 5×10−4 cosineannealed to 10−8 , gradient clipping at 5, and a batch size of 128. A small Gaussian latent noise (std 0.02) is injected before the decoder during training as a regularizer. The reconstruction objective is a latitude-weighted and variancenormalized MSE, so high-variance fields do not dominate. We also have three terms to preserve small-scale details, each weighted 0.05: a spatial gradient, a Laplacian for second-order sharpness, and a temporal-difference penalty. The trained autoencoder is frozen for all posterior sampling, and every likelihood is evaluated in physical field space through its decoder. The complete architecture and training configuration are tabulated in Appendix B. We tune these hyperparameters empirically using a smaller subset of training data. 3.3

Flow-matching training and sampling

The prior F is a DiT3D [24] trained with the modified TrigFlow loss Eq. 4 on (128, 8, 32, 64) latents. It uses a [1, 2, 2] patch so the (8, 32, 64) latent grid becomes a 4096-token sequence. The DiT3D uses hidden size 1536, depth 12, 24 attention heads, and MLP ratio 4 with a size-384 sinusoidal timestep embedding, and is trained with log-noise ln σ ∼ N (0, 1.52 ) and σd =1 using AdamW [84] and learning rate 2×10−4 , weight decay 0.01, batch size 256, with an exponential moving average of the weights with a half-life of 500 kilo-images (kimg) and a 500-kimg learning-rate ramp-up. The prior is never fine-tuned or conditioned on observations; all results use a single EMA checkpoint at 19,901 kimg. The full architecture and training details are in Appendix B. The sampling hyperparameters are given in their respective subsections. 3.4

Metric and cost

The common metric used is the per-variable latitude-weighted RMSE in physical units, following the WeatherBench evaluation convention [82]: for each variable we compute the area-weighted (cosine-latitude) RMSE of the ensemble7

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Pd mean prediction against ERA5, then average over frames. For variable v, with ū = d1 k=1 u(k) the mean over d posterior draws, ERA5 truth u⋆ , latitude ϕi , and frame index f running over the T =32 decoded frames (or a stated frame subset), sP T ⋆ 2 1 X i,j cos ϕi (ūv,f,i,j − uv,f,i,j ) P RMSEv = . (8) T i,j cos ϕi f =1

The reported values in the results are mean ± standard deviation over eight independent initial conditions (distinct, non-overlapping ERA5 test windows), and every prediction is the mean of an 8-member ensemble, each member an independent posterior draw from a distinct seed. For some tasks like forecasts, we have used more initial conditions and ensemble members, which are mentioned in those sections.

4

Results

We first recover the full state from structured coarse observations, which we call super-resolution (SR) in Sec. 4.1, and from real station observations, which we call data fusion in Sec. 4.2. We then evaluate each case for posterior accuracy, ensemble spread, physical consistency, and performance against unseen observations. We perform extensive ablation studies on the different sampling methods (Appendix A) and propose a set of recommended sampling hyperparameters in Table 2. Next, we explore the classical filtering, smoothing, and fixed-interval regimes of traditional DA with the same fixed prior (Sec. 4.3). Finally, we perform forecasting directly from observations (Sec. 4.4), including a case study of Hurricane Laura. Table 2: Recommended sampler configurations per task: the efficiency column gives the best low-NFE configuration and the accuracy column the best overall configuration in the corresponding ablation tables (key switches as in Table 1; NFE in parentheses).

4.1

task

efficiency recommendation

accuracy recommendation

Reference

super-resolution stations (IGRA/ISD/joint) obs → forecast filtering/smoothing study

DPS+corr (N30) (47) DPS+corr (N30) (47) DPS+corr (N30) (47)

DPS+corr(all σ) (99) DPS+corr+DSG+λ0 (76) DPS+corr+DSG+λ0 (76) DPS+corr+DSG+λ0 (76)

Tables 8, 10 Tables 9, 11, 12 Fig. 16 Table 7, Fig. 15

Super-resolution

We first perform a uniform spatio-temporal super-resolution (SR): a 4× temporal and 8×-per-axis spatial down-sampling. From these observations the prior recovers 2 m temperature to 1.02 K RMSE with the most accurate configuration (Table 2; the efficiency-tuned default gives 1.05 K at half the cost; Tables 8, 10). For reference, the autoencoder’s own reconstruction floor is 0.84 K for 2 m temperature. Figure 2 shows the reconstruction of the surface variables on an observed frame. The ensemble mean recovers fine-scale structure absent from the coarsened observations. The ensemble spread (Fig. 3) quantifies uncertainty, which is low over the regularly sampled observation grid and higher over dynamically active regions. The corresponding reconstructions for all 69 output channels (the four surface fields and the five upper-air fields at all 13 pressure levels) are reported in Appendix E.2. Spectral analysis (Fig. 4) confirms that the posterior recovers the correct zonal energy spectrum well beyond the observation cutoff, unlike a naive upsampling of the low-resolution field. The longitudinal traces (Fig. 5) show that the ensemble mean tracks sharp peaks and troughs with an uncertainty band (quantified below). The scatter of reconstructed against true values (Fig. 6) collapses onto the diagonal at observed pixels and remains tight at unobserved pixels, indicating effective full-state recovery. In time, the per-frame RMSE (Fig. 7) has a sawtooth structure that dips at each observed frame (shaded orange). We see the error drop after the first frame as the prior’s internal dynamics (and two-sided information) help reconstruction. We observe higher errors after the last observed frame, as expected. The following paragraphs explore whether the recovered posterior has calibrated spread and physical consistency. Calibration of the posterior spread. An ensemble is useful only if its spread is calibrated, that is, if the mean is wrong by about 1 K (skill), then the members should typically differ by 1 K (spread). We compute a total spread-skill ratio (SSR) that combines all 69 variables in their standardized and latitude-weighted units (1 is ideal). Table 3 reports it for 8-member ensembles of various sampling configurations. As expected, we observe that the setting λ=0 removes the injected noise from the corrector, making the SSR worse. For the sparse radiosonde observations, sampling (using λ=1) is the closest to calibrated (SSR 0.85). On the densely observed uniform SR task, the sampling configurations are more under-dispersed (SSR 0.72 to 0.75), meaning the spread is smaller relative to error. In short, λ=1 remains the 8

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 2: Super-resolution reconstruction on an observed frame for the surface variables: the low-resolution observation (left), ERA5 truth (middle), and posterior ensemble mean (right). The prior recovers fine-scale structure absent from the coarsened observations. All 69 channels are shown in Appendix E.2. choice whenever a calibrated spread is required. In future work, we would need to utilize an inflation-style correction, as operational ensembles routinely apply. Table 3: Ensemble calibration, 8-member ensembles for every configuration. Total Spread-Skill Ratio (SSR) combines p all 69 variables in standardized units (latitude weighted, scaled by (d+1)/d for a d-member ensemble; 1 is ideal, and values below 1 mean the ensemble is overconfident). Empirical ±2σ coverage has Gaussian target 0.954. The unobs. restricts to the unobserved frames. task

configuration

SSR

Empirical coverage

SSR unobs.

Empirical coverage unobs.

SR SR SR SR SR SR

DPS (λ=1) +corr (N30) (λ=1) +corr(all σ) (λ=1) +mom0.5 (λ=1) +corr (N30)+λ0 (λ=0) +corr(all σ)+λ0 (λ=0)

0.72 0.72 0.73 0.75 0.52 0.23

0.772 0.766 0.773 0.778 0.565 0.172

0.74 0.74 0.74 0.76 0.53 0.24

0.777 0.770 0.778 0.784 0.570 0.178

IGRA IGRA

+corr+DSG (λ=1) +corr+DSG+λ0 (λ=0)

0.85 0.83

0.803 0.732

0.87 0.85

0.818 0.763

9

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 3: Super-resolution uncertainty quantification: posterior ensemble mean (left), signed error of the mean from ERA5 (middle), and ensemble standard deviation (right) of an eight-seed ensemble. Spread is lowest over the regular observation grid and highest over dynamically active regions, where the error also concentrates.

Physical consistency. A posterior can score well on RMSE while producing physically inconsistent states, so we test the samples against two atmospheric balances (following the physical-consistency checks in prior work [71]) (Table 4). The first is geostrophic balance: in the extratropics, the wind should blow along the geopotential contours. We assess this by comparing the actual 500 hPa wind with the geostrophic wind computed from the sampled geopotential gradient. The second is hypsometric consistency: the 500-850 hPa thickness dictates the layer-mean temperature, so the thickness-implied temperature should match the directly sampled one. We observe that the posterior samples are at least as geostrophically balanced as the reanalysis in every configuration. In samples conditioned on IGRA measurements, the hypsometric residual is marginally above the reference (0.98-0.99 K versus 0.95 K). In general, the samples are dynamically and physically plausible states. Table 4: Physical consistency of the posterior samples. Geostrophic balance at 500 hPa in the extratropics (25◦ ≤|lat|≤70◦ ): cosine similarity between actual and geostrophic wind, and their speed ratio. Hypsometric check: RMSE between the 500-850 hPa thickness-implied layer-mean temperature and the directly sampled one. task

configuration

geo. cos

speed ratio

hyps. RMSE [K]

SR SR SR IGRA IGRA

ERA5 truth DPS +corr(all σ) +corr (N30)+λ0 +corr+DSG +corr+DSG+λ0

0.982 0.986 0.986 0.987 0.986 0.988

1.045 1.033 1.045 1.040 1.047 1.051

0.95 0.94 0.93 0.94 0.98 0.99

10

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 4: Zonal power spectra. The posterior sampling recovers the correct spectral slope up to high wavenumbers

4.2

Real-observation data fusion

We now condition on real, sparse observations (Fig. 8): IGRA measurements and the ISD surface measurements. From these sparse point observations, the prior reconstructs the full 69-channel field (Fig. 9), including the six unobserved channels (mean-sea-level pressure and specific humidity above 300 hPa). The ensemble spread of the generated samples (Fig. 10) is lower near the observation stations and higher over the sparse regions. The reconstruction scatter (Fig. 11) collapses onto the diagonal at the observed station locations (top row), where it largely reflects the data fit, and, more informatively, stays close to the diagonal at off-station grid points (bottom row), quantifying information gain at the soundings. However, some fields, like the winds, are intrinsically small-scale and highly chaotic, which causes higher pixel-wise RMSE error (Tables 5, 9). Moreover, the truth here is ERA5, which at the 1.4◦ grid differs considerably from the localized observations. Moreover, even a perfect posterior cannot beat the autoencoder reconstruction floor (Table 8; e.g. u10m 0.53 m/s, z500 37.0 m2 /s2 ), so a portion is due to the irreducible decoder error. Multimodal data fusion. Intuitively, the addition of more observation points shows lower errors on joint sampling than individually using IGRA and ISD (Table 5). The surface-only ISD is worst despite its density because the prior has to reconstruct every upper-air variable from surface observations alone. The comparison therefore contrasts two whole observing systems: the soundings are sparser on the map but report 63 channels through the depth of the atmosphere against ISD’s four. Figure 12 visualizes this as the 2 m-temperature error map across modalities, and Fig. 13 shows the truth-vs-reconstruction scatter for the ISD surface observations. Matching held-out observations. Everything above scores the reconstruction against ERA5, which is not the perfect representation of the ground truth. We therefore verify directly in observation space similar to prior work [28]. We withhold a random 20% of stations, assimilate only the remaining 80%, and measure the deviation of the reconstruction from the unseen observations. Because station observations are highly localized, a meaningful reference error is the deviation of ERA5 at the same sites. Table 6 shows the result on the observed frames: on the withheld stations the 11

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 5: Longitudinal traces at 40.4◦ N: ERA5 truth vs. the posterior ensemble mean with a ±2σ band.

Figure 6: Reconstruction scatter: ensemble mean vs. ERA5 truth at observed (red) and unobserved (grey) pixels. Points on the diagonal indicate exact recovery. Table 5: Multimodal data-fusion: per-variable physical RMSE (Units: t2m K, u10m/v10m m/s, q1000 g/kg, z500 m2 /s2 ; bold = best, underline = second-best). modality Joint (IGRA+ISD, 64 var) IGRA (radiosonde, 63 var) ISD (surface, 4 var)

t2m

u10m

v10m

q1000

z500

1.80±0.11 1.90±0.13 2.21±0.11

2.93±0.14 3.01±0.17 3.34±0.18

3.00±0.17 3.08±0.18 3.33±0.20

1.13±0.04 1.15±0.04 1.41±0.09

467.6±48.5 474.9±50.4 701.3±65.8

reconstruction is at or below the ERA5 deviations on three of four ISD variables (10 m winds and pressure) and close for 2 m temperature. For the IGRA radiosondes, the reconstruction deviation from held-out observations closely tracks that of ERA5. 12

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 7: SR per-frame physical RMSE per variable (six panels: t2m, u10m, v10m, q1000, z500, t850). Vertical shading marks the observed frames (every fourth).

Figure 8: The two real observation sources IGRA and ISD (union of reporting sites across the eight observed frames and eight windows). Table 6: Observation-space verification at withheld stations. RMSE of the reconstruction versus that of ERA5 at the same sites and times (lower half: withheld IGRA radiosondes; upper half: withheld ISD surface stations). Bold marks where the reconstruction departure is at or below the ERA5 reference. Observation

Variable

RMSE

ERA5

ISD ISD ISD ISD

2 m temperature [K] 10 m u-wind [m/s] 10 m v-wind [m/s] MSLP [hPa]

3.24 2.27 2.19 2.16

3.06 2.44 2.31 2.50

IGRA IGRA IGRA IGRA IGRA IGRA IGRA IGRA

2 m temperature [K] 10 m u-wind [m/s] z500 [m2 /s2 ] z850 [m2 /s2 ] t500 [K] t850 [K] t300 [K] u200 [m/s]

3.16 2.39 165 124 1.21 1.58 1.42 3.77

2.91 2.40 155 108 0.91 1.23 1.19 2.95

13

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 9: IGRA radiosonde data-fusion surface fields: ERA5 truth, posterior ensemble mean, and signed error.

4.3

Filtering and smoothing

We now experiment with the classical estimation regimes of traditional DA [85] with the same learned prior, changing only which temporal frames are observed. In classical systems, these regimes require distinct frameworks around a numerical forecast model, whereas here, each is just a different observation mask on the same posterior [75, 71]. The smoother observes the middle-8 frames to reconstruct the full 32-frame (8-day) spatio-temporal chunk. The filter observes the leading eight frames, so its unobserved tail is a forecast. We also use a fixed-interval reference which observes one frame per day. A schematic of these tasks is in Fig. 14. We demonstrate these first with the idealized observations (similar to SR), and then with the real IGRA observations. The sampling method is held fixed within each observation source, using the recommended configurations of Table 2. The motivation of this experiment is to show that a single unconditional prior realizes each regime and shows exactly the behavior classical DA associates with it, in Fig. 15. Table 7 has the quantitative analysis. The fixed-interval reference flattens the error across the whole window at a low level, as expected when observations constrain the entire trajectory. The smoother gives the familiar U-shape, with the lowest error within the observation window. The filter shows error that is low over the leading observed frames and grows forward in time as the prior must forecast beyond the last observation. The error grows smoothly forward, suggesting that the prior learns to propagate the atmospheric fields. Therefore, the prior acts as the learned forecast model, which is the function of a numerical model in classical assimilation. We also run the same three regimes on the real IGRA radiosonde observations, the lower half of Table 7. Together, these show that the same unconditional video prior supports filtering, smoothing, and fixed-interval observing, 14

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 10: IGRA radiosonde data fusion uncertainty quantification: ensemble mean (left) and standard deviation (right). Spread is lower over the regions with more observation locations. simply by changing the temporal observation geometry. Classically, each of these would be a separate algorithm wrapped around a forecast model. 4.4

Observations to forecast

The causal filtering geometry of the previous section motivates conditioning on real observations and propagating forward for forecasting, a direct observation-to-forecast [86, 28]. We assimilate a global in-situ observation on the leading frames of the window and let the prior forecast the remaining timesteps. Now, since the likelihood is a bilinear interpolation of the decoded field (Sec. 2.4), it admits only observations that directly measure an ERA5 state variable. This admits three useful in-situ observation sources: IGRA radiosondes, ISD surface stations, and ICOADS ships and 15

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 11: IGRA reconstruction vs ERA5 scatter plots for three observed variables (t2m, u10m, z500). Top: at the radiosonde station locations; bottom: at off-station grid points. Color shows log10 hexbin count; mentioned text gives the RMSE in physical units and Pearson r. Similar plots for all observed variables are in Appendix E.3.

Figure 12: 2 m-temperature signed error by assimilation modality. Among the real-observations data-fusion, the joint ISD+IGRA has the lowest error. The SR panel uses structured dense observations and is shown only for reference.

buoys (Sec. 3.1; compared to [86]). It excludes the satellite radiances and scatterometer backscatter, which require a radiative-transfer or geophysical forward operator. We therefore assimilate the union ISD ∪ IGRA ∪ ICOADS, a land, ocean, and upper-air network of observations per 6-hourly step. We therefore apply the grid-cell super-obbing of Sec. 2.4, to the dense ISD and ICOADS observations. This lets the likelihood gradient not be dominated by the over-sampled regions and equalizes each cell’s contribution. More details on this are in Appendix C. The sampler which performs best in this case is DPS+corr+DSG+λ0 (Table 1: 50 solver steps, one corrector step at σ<3, λ=0, 76 NFE). The DSG guidance scale is re-selected on a set of validation windows and applied to a disjoint test set of thirty-two initial conditions spread evenly across 2020, over which we report the mean. Chaining windows into a three-week forecast. The same posterior sampler supports autoregressive chaining, an analogue of the cycled reinitialization of operational systems. We take the last generated frame of a window and re-assimilate it as the initial condition of the next window, treating it as a dense observation of that window’s first frame. 16

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 13: ISD surface reconstruction vs ERA5 scatter plots for three observed variables (t2m, u10m, MSLP). Top: at the station locations; bottom: at off-station grid points (random subsample). Color shows log10 hexbin count; mentioned text gives the RMSE (physical units) and Pearson r. All four observed ISD variables are in Appendix E.3.

Figure 14: The three temporal observation geometries over the 32-frame window, each observing 8 frames: smoother (two-sided), filter (causal/leading), and the fixed-interval reference (every fourth frame, the geometry of a reanalysis). Only which frames are observed differs; each realizes one classical DA regime with the same frozen prior.

The observation operator for the next window becomes the identity at the first window for every grid cell and all 69 channels. Each additional window extends the forecast by 7 days and 18 hours. Firstly, we observe that the forecast beats persistence of the last observed frame at essentially all leads and stays at or below the annual-climatology error. Figure 16 compares this forecast against GraphDOP [28], the directly comparable observation-only, end-to-end data-driven forecast system, and the operational ECMWF Integrated Forecasting System (IFS), all verified against ERA5. The GraphDOP and IFS skill curves are extracted from the respective figure of [28], every 12 h through day 10, their full range. The comparison is asymmetric in input information: GraphDOP is trained on a broad slice of the operational observing system, ingesting several satellite measurements alongside the surface, marine, aircraft, and radiosonde measurements [28]. However, we still improve on GraphDOP at all leads for 2 m temperature where our dense surface observation is helpful. For other surface fields, our errors start above GraphDOP and converge to it by days 5 to 7 as forecast error saturates. This suggests that the satellite measurements are very 17

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 15: DA experiments: per-frame physical RMSE per variable (six panels: t2m, u10m, v10m, q1000, z500, t850) for the three temporal geometries. The filter experiment’s error grows forward in time, the smoother is U-shaped, and the fixed-interval reference is flat. Table 7: DA experiments: per-variable physical RMSE over all frames. All geometries observe 8 of the 32 frames; only which frames differs. Upper panel: using 8× subsampled observations per axis. Lower panel: using IGRA radiosonde observations; for its fixed-interval row the window starts are aligned to the 00/12 UTC launch hours. Per-frame observed/unobserved breakdown in Fig. 15 and Appendix E.1. The frame here refers to the specific timestep of our video window. geometry

t2m

u10m

v10m

q1000

z500

0.85±0.03 1.05±0.03 0.63±0.03

232.9±12.4 412.7±32.6 62.8±1.3

SR observations (8× subsample, sampling configuration is DPS+corr (N30)) Smoother (frames 12-19) Filter (frames 0-7, causal) Fixed-interval (every 4th frame, ref.)

1.40±0.04 1.85±0.04 1.03±0.04

1.96±0.05 2.58±0.09 1.10±0.02

2.01±0.07 2.73±0.07 1.15±0.02

IGRA radiosondes (real soundings, sampling configuration is DPS+corr+DSG+λ0) Smoother (frames 12-19) Filter (frames 0-7, causal) Fixed-interval (every 4th frame, ref.)

1.90±0.13 2.13±0.09 1.51±0.11

3.01±0.17 3.38±0.11 2.29±0.10

3.08±0.18 3.51±0.10 2.28±0.10

1.15±0.04 1.26±0.06 0.94±0.02

474.9±50.4 621.3±32.1 272.6±35.8

helpful for an accurate forecast, especially at shorter lead times. In longer forecast leads, our methodology shows clear advantages over GraphDOP. Since it is a generative model, we see better stability in the generated forecasts We perform additional analysis for our observation-to-forecast case. First, we observe that the length of the observed analysis window does not affect the forecast performance. Shortening it from two days to one day, or even to a single analysis frame, leaves forecast skill essentially unchanged (Appendix C). Moreover, in Appendix D, we use additional gridded observations from [86], where our forecasts perform close to GraphDOP. However, we could perform limited tests due to the availability of only a single observation from [86]. We leave further tests with operators that can be associated with satellite data to future work. Hurricane tracking. We finally conduct a case study of Hurricane Laura (in 2020), one of the strongest hurricanes on record to make landfall in the U.S. Starting the window at Aug 23 06Z, Laura passes Hispaniola on its way to Cuba. We condition on the real in-situ observations (ISD ∪ IGRA ∪ ICOADS union, exactly the same as the forecast configuration above) over observed windows of increasing length, two, three, and four days of observations, and let the frozen prior forecast the remainder up to when landfall occurs on Aug 27 06Z at Cameron, Louisiana. Figure 17 shows 18

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 16: Observations-to-forecast skill. Latitude-weighted RMSE against ERA5 versus lead over thirty-two initial conditions spread evenly across 2020, 8-member ensembles. The GraphDOP and IFS curves [28] are extracted from their published figures. The persistence and annual-climatology floors are for reference. Persistence is omitted for 2 m temperature because of the diurnal cycle.

sixteen ensemble trajectories for the three cases against ERA5. We observe that the trajectory tightens accordingly, with the forecast with 4 observed days being the closest and showing the least spread. The lower spread here does not directly mean greater accuracy, but that the forecasts become more confident. Firstly, at 1.40625◦ the grid cannot resolve a 19

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 17: Hurricane Laura case study: actual track (ERA5, black; minimum-MSLP) versus the sixteen individual predicted member trajectories for three observed windows starting Aug 23 06Z: two days (blue; landfall at a 54-hour lead), three days (orange), and four days (green). Red dots show the actual observation sites (stations and ships) at the first analysis time; the star marks landfall (Aug 27 06Z, Cameron, Louisiana).

Figure 18: Hurricane Laura case study: the window starts Aug 21 00Z and the first three days are observed as Laura crosses the Lesser Antilles, Puerto Rico, and Hispaniola (last observed frame: open square, Aug 23 18Z); everything afterward is forecast. Thirty-two member trajectories are shown (light blue), and red dots are the observation sites at the first analysis time.

hurricane. Secondly, there are only a few observations over the Gulf. However, some of the ensemble members track the hurricane accurately. The distribution of 10 m wind inside the hurricane region (Fig. 19) shows that the posterior matches ERA5 except the extreme tails. We also repeat the experiment from the storm’s early life: a window starting 20

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 19: Distribution of the 10 m zonal wind u near the Laura region: ERA5 (black) against the posterior ensemble for two, three, and four days of observations.

Aug 21 00Z, observing the three days in which Laura crosses the Lesser Antilles, Puerto Rico, and Hispaniola, followed by a 32-member ensemble forecast in which landfall is a 3.5-day-lead prediction (Fig. 18). The forecast seems to carry the storm into the Caribbean and then spreads, some entering the Gulf. Both these experiments show promising results where our learned prior and posterior sampling methodology can be used for direct obs-to-forecast, given an ideal observation set and the related operator is available. Hindcast. The same framework also runs in reverse. We take the Laura window starting Aug 21 00Z, but now observe the real in-situ observations only over the last four days (Aug 25-28), when the storm runs from maturity to landfall, and hindcast using the prior to infer the first four days (Aug 21-24). This is the initial state that grew into the observed hurricane. We plot individual posterior members, their ensemble mean, and the ensemble spread along with ERA5 (Figs. 20 and 21). The figures show a coherent structure moving westward across the Caribbean. Information from the observed mature storm propagates backwards in time through the prior to constrain its unobserved genesis, with increasing errors backwards in time. The wind distributions (Fig. 22) here reproduce the full range of ERA5’s 10 m wind, as expected when a past state is reconstructed rather than forecasting an extreme event.

5

Discussion

In this work, we train a latent generative video prior of ERA5 and perform spatio-temporal data assimilation. We perform several experiments starting from simplified super-resolution to data fusion using real observations such as IGRA radiosondes, ISD surface stations, and the ICOADS ocean-atmosphere dataset. We also show that our generated samples are physically consistent in the limited tests that we have performed. Moreover, because the prior is a video, the methodology supports the classical data assimilation experiments without requiring an autoregressive forecast model: observing leading frames yields a causal filter whose error grows forward in time, and observing the middle yields a two-sided smoother. The same posterior sampling framework also yields a direct observation-to-forecast pipeline. Driven only by in-situ observations, with no learned encoder for observations and no separate forecast model, it can produce a prediction competitive with other trained obs-to-forecast systems like GraphDOP [28]. We also explore different sampling algorithms, some of which improve efficiency by distributing the sampling budget towards lower noise levels, and others that scale the gradients and noise injection for more accuracy and calibrated spread. Future work includes extending our sequential cycling across consecutive windows for long-term sub-seasonal forecasts. Moreover, an operator that can inject several observations, especially from satellites, and convert them to the ERA5 grid (like Aardvark [86] or HealDA [87]) can be used to further improve our pipeline. We present a preliminary experiment using Aardvark in Appendix D. Another future direction is to directly include these observations as input to the flow model, with the option to exclude them when unavailable. 21

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 20: Hindcast of Hurricane Laura (10 m wind speed): ERA5 (top row), two individual posterior members, the ensemble mean, and the ensemble spread (standard deviation across members, bottom row) on the four inferred days Aug 21-24 00Z. Black contours are each field’s own MSLP; the × marks the storm position.

Acknowledgments This research used resources of the Argonne Leadership Computing Facility, a U.S. Department of Energy Office of Science User Facility operated under contract DE-AC02-06CH11357. DC and RM acknowledge support from the DOE ASCR award “Inertial neural surrogates for stable dynamical prediction.” DC and RM also acknowledge computing resources through an AI for Science allocation at the National Energy Research Scientific Computing Center (NERSC) and from the Penn State Institute for Computational and Data Sciences (ICDS). RM acknowledges the support of a YIP from the ARO Modeling of Complex Systems program (PM: Robert Martin). We also acknowledge Melissa Adrian for the useful discussions related to this work.

22

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 21: Hindcast of Hurricane Laura’s precursor (mean sea-level pressure). Individual members develop a closed low that deepens and tracks with the storm’s position (×) as Aug 24 approaches, recovering the broad structure of ERA5’s developing low.

23

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 22: Distribution of the 10 m zonal wind u inside the box for the hindcast, split into the observed segment (Aug 25-28, green) and the inferred precursor (Aug 21-24, red); ERA5 solid, posterior dashed.

24

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

References [1] A Arakawa, J-H Jung, and C-M Wu. Toward unification of the multiscale modeling of the atmosphere. Atmospheric Chemistry and Physics, 11(8):3731–3742, 2011. [2] Horst J Neugebauer. Dynamics of multiscale earth systems, volume 97. Springer Science & Business Media, 2003. [3] Wei-Kuo Tao and Mitchell W Moncrieff. Multiscale cloud system modeling. Reviews of Geophysics, 47(4), 2009. [4] Jean-Noël Thépaut. Satellite data assimilation in numerical weather prediction: An overview. In Proceedings of ECMWF Seminar on Recent Developments in Data Assimilation for Atmosphere and Ocean, ECMWF, Reading, UK, pages 8–12, 2003. [5] François Bouttier and Graeme Kelly. Observing-system experiments in the ecmwf 4d-var data assimilation system. Quarterly Journal of the Royal Meteorological Society, 127(574):1469–1488, 2001. [6] Christopher K Wikle and L Mark Berliner. A bayesian tutorial for data assimilation. Physica D: Nonlinear Phenomena, 230(1-2):1–16, 2007. [7] Kody JH Law and Andrew M Stuart. Evaluating data assimilation algorithms. Monthly weather review, 140(11):3757–3782, 2012. [8] Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz-Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, et al. The era5 global reanalysis. Quarterly journal of the royal meteorological society, 146(730):1999–2049, 2020. [9] Istvan Szunyogh. Applicable atmospheric dynamics: Techniques for the exploration of atmospheric dynamics. World Scientific, 2014. [10] Bin Wang, Xiaolei Zou, and Jiang Zhu. Data assimilation and its applications. Proceedings of the National Academy of Sciences, 97(21):11143–11144, 2000. [11] Rolf H Reichle. Data assimilation methods in the earth sciences. Advances in water resources, 31(11):1411–1418, 2008. [12] Ruslan Salakhutdinov. Learning deep generative models. Annual Review of Statistics and Its Application, 2(1):361–385, 2015. [13] Han Gao, Sebastian Kaltenbach, and Petros Koumoutsakos. Generative learning for forecasting the dynamics of high-dimensional complex systems. Nature Communications, 15(1):8904, 2024. [14] Ilan Price, Alvaro Sanchez-Gonzalez, Ferran Alet, Tom R Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, et al. Probabilistic weather forecasting with machine learning. Nature, 637(8044):84–90, 2025. [15] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415–1428, 2021. [16] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. [17] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35:26565–26577, 2022. [18] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. [19] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024. [20] Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. Advances in Neural Information Processing Systems, 35:14715–14728, 2022. [21] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.14687, 2022. [22] Dibyajyoti Chakraborty, Haiwen Guan, Jason Stock, Troy Arcomano, Guido Cervone, and Romit Maulik. Multimodal atmospheric super-resolution with deep generative models. Machine Learning: Earth, 2(1):015001, 2026. 25

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

[23] Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-time consistency models. In International Conference on Learning Representations, volume 2025, pages 50611–50649, 2025. [24] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195–4205, 2023. [25] Imke Durre, Russell S Vose, and David B Wuertz. Overview of the integrated global radiosonde archive. Journal of Climate, 19(1):53–68, 2006. [26] Adam Smith, Neal Lott, and Russ Vose. The integrated surface database: Recent developments and partnerships. Bulletin of the American Meteorological Society, 92(6):704–708, 2011. [27] Eric Freeman, Scott D Woodruff, Steven J Worley, Sandra J Lubker, Elizabeth C Kent, William E Angel, David I Berry, Philip Brohan, Ryan Eastman, Lydia Gates, et al. Icoads release 3.0: a major update to the historical marine climate record. International Journal of Climatology, 37(5):2211–2232, 2017. [28] Mihai Alexe, Eulalie Boucher, Peter Lean, Ewan Pinnington, Patrick Laloyaux, Anthony McNally, Simon Lang, Matthew Chantry, Chris Burrows, Marcin Chrust, et al. Graphdop: Towards skilful data-driven medium-range weather forecasts learnt and initialised directly from observations. arXiv preprint arXiv:2412.15687, 2024. [29] Kai Fukami, Koji Fukagata, and Kunihiko Taira. Super-resolution reconstruction of turbulent flows with machine learning. Journal of Fluid Mechanics, 870:106–120, 2019. [30] Mojtaba F Fathi, Isaac Perez-Raya, Ahmadreza Baghaie, Philipp Berg, Gabor Janiga, Amirhossein Arzani, and Roshan M D’Souza. Super-resolution and denoising of 4d-flow mri using physics-informed deep neural nets. Computer Methods and Programs in Biomedicine, 197:105729, 2020. [31] Daniel Kelshaw, Georgios Rigas, and Luca Magri. Physics-informed cnns for super-resolution of sparse observations on dynamical systems. arXiv preprint arXiv:2210.17319, 2022. [32] Kai Fukami, Koji Fukagata, and Kunihiko Taira. Super-resolution analysis via machine learning: a survey for fluid flows. Theoretical and Computational Fluid Dynamics, 37(4):421–444, 2023. [33] Han Gao, Luning Sun, and Jian-Xun Wang. Super-resolution and denoising of fluid flow using physics-informed convolutional neural networks without high-resolution labels. Physics of Fluids, 33(7), 2021. [34] Pu Ren, Chengping Rao, Yang Liu, Zihan Ma, Qi Wang, Jian-Xun Wang, and Hao Sun. Physr: Physics-informed deep super-resolution for spatiotemporal data. Journal of Computational Physics, 492:112438, 2023. [35] Kai Fukami and Kunihiko Taira. Single-snapshot machine learning for super-resolution of turbulence. Journal of Fluid Mechanics, 1001:A32, 2024. [36] Shivam Barwey, Pinaki Pal, Saumil Patel, Riccardo Balin, Bethany Lusch, Venkatram Vishwanath, Romit Maulik, and Ramesh Balakrishnan. Mesh-based super-resolution of fluid flows with multiscale graph neural networks. Computer Methods in Applied Mechanics and Engineering, 443:118072, 2025. [37] Kai Fukami, Romit Maulik, Nesar Ramachandra, Koji Fukagata, and Kunihiko Taira. Global field reconstruction from sparse sensors with voronoi tessellation-assisted deep learning. Nature Machine Intelligence, 3(11):945–951, 2021. [38] Romit Maulik, Kai Fukami, Nesar Ramachandra, Koji Fukagata, and Kunihiko Taira. Probabilistic neural networks for fluid flow surrogate modeling and data recovery. Physical Review Fluids, 5(10):104401, 2020. [39] Romit Maulik, Romain Egele, Krishnan Raghavan, and Prasanna Balaprakash. Quantifying uncertainty for deep learning based forecasting and flow-reconstruction using neural architecture search ensembles. Physica D: Nonlinear Phenomena, 454:133852, 2023. [40] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407, 2018. [41] Masaki Morimoto, Kai Fukami, Romit Maulik, Ricardo Vinuesa, and Koji Fukagata. Assessments of epistemic uncertainty using gaussian stochastic weight averaging for fluid-flow regression. Physica D: Nonlinear Phenomena, 440:133454, 2022. [42] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256–2265. pmlr, 2015. [43] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. [44] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. 26

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

[45] Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005. [46] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022. [47] Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. Journal of Machine Learning Research, 26(209):1–80, 2025. [48] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In International conference on learning representations, 2023. [49] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in neural information processing systems, 35:23593–23606, 2022. [50] Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. arXiv preprint arXiv:2212.00490, 2022. [51] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. In International Conference on Learning Representations, volume 2024, pages 28027–28053, 2024. [52] Yuanzhi Zhu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, and Luc Van Gool. Denoising diffusion models for plug-and-play image restoration. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1219–1229, 2023. [53] Zihui Wu, Yu Sun, Yifan Chen, Bingliang Zhang, Yisong Yue, and Katherine L Bouman. Principled probabilistic imaging using diffusion models as plug-and-play priors. Advances in Neural Information Processing Systems, 37:118389–118427, 2024. [54] Bingliang Zhang, Wenda Chu, Julius Berner, Chenlin Meng, Anima Anandkumar, and Yang Song. Improving diffusion inverse problem solving with decoupled noise annealing. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 20895–20905, 2025. [55] Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alex Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. Advances in Neural Information Processing Systems, 36:49960–49990, 2023. [56] Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen. Solving inverse problems with latent diffusion models via hard data consistency. In International Conference on Learning Representations, volume 2024, pages 7624–7654, 2024. [57] Hongkai Zheng, Wenda Chu, Bingliang Zhang, Zihui Wu, Austin Wang, Berthy Feng, Caifeng Zou, Yu Sun, Nikola Kovachki, Zachary Ross, et al. Inversebench: Benchmarking plug-and-play diffusion priors for inverse problems in physical sciences. In International Conference on Learning Representations, volume 2025, pages 90912–90940, 2025. [58] Morteza Mardani, Noah Brenowitz, Yair Cohen, Jaideep Pathak, Chieh-Yu Chen, Cheng-Chin Liu, Arash Vahdat, Mohammad Amin Nabian, Tao Ge, Akshay Subramaniam, et al. Residual corrective diffusion modeling for km-scale atmospheric downscaling. Communications Earth & Environment, 6(1):124, 2025. [59] Prakhar Srivastava, Ruihan Yang, Gavin Kerrigan, Gideon Dresdner, Jeremy McGibbon, Christopher Bretherton, and Stephan Mandt. Precipitation downscaling with spatiotemporal video diffusion. Advances in Neural Information Processing Systems, 37:56374–56400, 2024. [60] Robbie A Watt and Laura A Mansfield. Generative diffusion-based downscaling for climate. arXiv preprint arXiv:2404.17752, 2024. [61] Lizao Li, Robert Carver, Ignacio Lopez-Gomez, Fei Sha, and John Anderson. Generative emulation of weather forecast ensembles with diffusion models. Science Advances, 10(13):eadk4489, 2024. [62] Martin Andrae, Tomas Landelius, Joel Oskarsson, and Fredrik Lindsten. Continuous ensemble weather forecasting with diffusion models. In International Conference on Learning Representations, volume 2025, pages 26392– 26416, 2025. [63] Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter D Dueben, and Torsten Hoefler. Diffda: a diffusion model for weather-scale data assimilation. arXiv preprint arXiv:2401.05932, 2024. [64] Peter Manshausen, Yair Cohen, Peter Harrington, Jaideep Pathak, Mike Pritchard, Piyush Garg, Morteza Mardani, Karthik Kashinath, Simon Byrne, and Noah Brenowitz. Generative data assimilation of sparse weather station observations at kilometer scales. Journal of Advances in Modeling Earth Systems, 17(10):e2024MS004505, 2025. 27

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

[65] Shangshang Yang, Congyi Nai, Xinyan Liu, Weidong Li, Jie Chao, Jingnan Wang, Leyi Wang, Xichen Li, Xi Chen, Bo Lu, et al. Generative assimilation and prediction for weather and climate. arXiv preprint arXiv:2503.03038, 2025. [66] Yongquan Qu, Juan Nathaniel, Shuolin Li, and Pierre Gentine. Deep generative data assimilation in multimodal setting. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 449–459, 2024. [67] Hao Wang, Jindong Han, Wei Fan, Weijia Zhang, and Hao Liu. Phyda: Physics-guided diffusion models for data assimilation in atmospheric systems. arXiv preprint arXiv:2505.12882, 2025. [68] Feng Bao, Zezhong Zhang, and Guannan Zhang. A score-based filter for nonlinear data assimilation. Journal of Computational Physics, 514:113207, 2024. [69] Siming Liang, Hoang Tran, Feng Bao, Hristo G Chipilski, Peter Jan van Leeuwen, and Guannan Zhang. Ensemble score filter with image inpainting for data assimilation in tracking surface quasi-geostrophic dynamics with partial observations. arXiv preprint arXiv:2501.12419, 2025. [70] Bingliang Zhang, Zihui Wu, Berthy T Feng, Yang Song, Yisong Yue, and Katherine L Bouman. Step: A framework for solving scientific video inverse problems with spatiotemporal diffusion priors. arXiv preprint arXiv:2504.07549, 2025. [71] Gérôme Andry, Sacha Lewin, François Rozet, Omer Rochman, Victor Mangeleer, Matthias Pirlet, Elise Faulx, Marilaure Grégoire, and Gilles Louppe. Appa: Bending weather dynamics with latent diffusion models for global data assimilation. arXiv preprint arXiv:2504.18720, 2025. [72] Bradley Efron. Tweedie’s formula and selection bias. 106(496):1602–1614, 2011.

Journal of the American Statistical Association,

[73] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 24174–24184, 2024. [74] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022. [75] François Rozet and Gilles Louppe. Score-based data assimilation. Advances in Neural Information Processing Systems, 36:40521–40541, 2023. [76] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. Machine Intelligence Research, 22(4):730–751, 2025. [77] Yuhan Wang, Suzhi Bi, Ying-Jun Angela Zhang, and Xiaojun Yuan. Traversing distortion-perception tradeoff using a single score-based generative model. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 2377–2386, 2025. [78] Angela Castillo, Jonas Kohler, Juan C Pérez, Juan Pablo Pérez, Albert Pumarola, Bernard Ghanem, Pablo Arbeláez, and Ali Thabet. Adaptive guidance: Training-free acceleration of conditional diffusion models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 1962–1970, 2025. [79] Lingxiao Yang, Shutong Ding, Yifan Cai, Jingyi Yu, Jingya Wang, and Ye Shi. Guidance with spherical gaussian constraint for conditional diffusion. arXiv preprint arXiv:2402.03201, 2024. [80] AC Lorenc. A global three-dimensional multivariate statistical interpolation scheme. Monthly Weather Review, 109(4):701–721, 1981. [81] Tilo Ochotta, Christoph Gebhardt, Dietmar Saupe, and Werner Wergen. Adaptive thinning of atmospheric observations in data assimilation with vector quantization and filtering methods. Quarterly Journal of the Royal Meteorological Society: A journal of the atmospheric sciences, applied meteorology and physical oceanography, 131(613):3427–3437, 2005. [82] Stephan Rasp, Peter D Dueben, Sebastian Scher, Jonathan A Weyn, Soukayna Mouatadid, and Nils Thuerey. Weatherbench: a benchmark data set for data-driven weather forecasting. Journal of Advances in Modeling Earth Systems, 12(11):e2020MS002203, 2020. [83] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [84] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [85] Stephen E Cohn, NS Sivakumaran, and Ricardo Todling. A fixed-lag kalman smoother for retrospective data assimilation. Monthly Weather Review, 122(12):2838–2867, 1994. 28

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

[86] Anna Allen, Stratis Markou, Will Tebbutt, James Requeima, Wessel P Bruinsma, Tom R Andersson, Michael Herzog, Nicholas D Lane, Matthew Chantry, J Scott Hosking, et al. End-to-end data-driven weather prediction. Nature, 641(8065):1172–1179, 2025. [87] Aayush Gupta, Akshay Subramaniam, Michael S Pritchard, Karthik Kashinath, Sergey Frolov, Kelsey Lieberman, Christopher Miller, Nicholas Silverman, and Noah D Brenowitz. Healda: Highlighting the importance of initial errors in end-to-end ai weather forecasts. arXiv preprint arXiv:2601.17636, 2026. [88] Tuomas Kynkäänniemi, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance in a limited interval improves sample and distribution quality in diffusion models. Advances in Neural Information Processing Systems, 37:122458–122483, 2024. [89] Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. arXiv preprint arXiv:2108.01073, 2021. [90] Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. Advances in Neural Information Processing Systems, 36:76806–76838, 2023. [91] Ross N Hoffman. The effect of thinning and superobservations in a simple one-dimensional data analysis with mischaracterized error. Monthly Weather Review, 146(4):1181–1195, 2018. [92] Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, et al. Learning skillful medium-range global weather forecasting. Science, 382(6677):1416–1421, 2023.

29

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Appendix contents A Exploring different sampling techniques . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30 A.1 Sampler design choices . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30 A.2 Efficiency: distributing the NFE budget . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31 A.3 Reshaping gradients and noise . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31 A.4 Likelihood weighting . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 33 A.5 Comparison with existing samplers . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 33 A.6 Variable-splitting samplers: further ablations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34 B Model and training details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36 C Grid-cell super-obbing and the observed-window length . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36 D Forecasting from the Aardvark latent observations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 37 E Extended results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 37 E.1 Full per-variable physical RMSE . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 37 E.2 Per-channel super-resolution reconstructions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 39 E.3 Per-variable observation scatter . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 41

A

Exploring different sampling techniques

This appendix collects the design rationale and the sampling study behind the recommended configurations of Table 2. Sec. A.1 justifies the central design choices of the guided sampler; the remaining subsections report the experiments. In every experiment the prior is fixed, every run is scored with the same latitude-weighted physical RMSE as the main results, and each method is tested on both observation types: the dense structured SR grid and the sparse, irregular station observations. Unless stated otherwise, the reported rows are 8-member ensembles over the 8 validation windows disjoint from the test set. Conceptually, the sampler has two independent directions to explore: where the NFE budget is spent (Sec. A.2) and how the computed gradient and injected noise are used (Sec. A.3). Sec. A.4 then reports the likelihood-weight sweeps behind the settings used in the paper, and Secs. A.5 and A.6 benchmark existing samplers in literature on the same tasks. A.1

Sampler design choices

We justify the central design choices here. (a) Low-σ corrector helps. At high σ the rotation mostly transports the prior and observation y barely constrains the decoded ẑ0 . The low-σ Langevin corrector targets the posterior score and recovers measurement fidelity only in the low-σ band where it can be determined. (b) Lazy reuse fails. DPS requires ∇zt log p(y | zt ) = (∂ẑ0 /∂zt )⊤ ∇ẑ0 log p(y | ẑ0 ). From (2) the TrigFlow denoiser Jacobian is ∂ẑ0 ∂F = cos(t) I − sin(t) σd . (9) ∂zt ∂zt Through (9) the likelihood gradient depends on cos(t), sin(t), and the changing ẑ0 . Therefore, holding it fixed across steps and reusing injects a stale direction which hampers performance in our test cases. 30

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

(c) Effect of momentum. The per-step likelihood (through a decoder) gradient can be a noisy estimate of the misfit gradient. Momentum averages successive estimates, so the guidance direction has lower variance, similar to several optimizers. We observe empirically in our test cases that it does not help when the Langevin corrector is present: the corrector already re-equilibrates the state at each level, so the accumulated gradient direction adds little. We leave further exploration of the effect of momentum on highly non-linear measurement operators to future work. (d) DSG normalization. The per-timestep coverage of station measurements swings widely (for IGRA it is dense at 00/12Z and sparse at 06/18Z). Therefore, the raw gradient magnitude varies erratically and DSG sphere-normalization is preferred. (e) Second-order update of Algorithm 1. Substituting σd Fs = (cos s zt −ẑ0 )/ sin s (from Eq. (2)) into the first-order rotation gives the exact exponential-integrator form sin t sin δ zt + ẑ0 , (10) sin s sin s the trigonometric analogue of DDIM: the step would be exact if ẑ0 were constant along it. The multistep second-order scheme of DPM-Solver++(2M) [76] replaces ẑ0 by its linear extrapolation in log σ = log tan using the previous 1 estimate, ẑ0 → ẑ0 + 2r (ẑprev − ẑ0 ) with r the log-σ step ratio defined in Algorithm 1 (negative under that sign 0 prev δ convention). Inserting this yields the correction term 2rsin − ẑ0 ) of Algorithm 1. sin s (ẑ0 z(1) =

A.2

Efficiency: distributing the NFE budget

We observe that using NFEs strategically by using the corrector steps at low noise levels helps, as per Table 8 and Fig. 23. On SR, the reduced-step DPS+corr (N25) matches the 50-step DPS baseline with fewer steps and less wall time, and adding more corrector steps improves accuracy consistently. The IGRA radiosondes show a similar same ordering (Table 9), except that the 50-step corrector configuration concentrated at low noise performs best. Restricting guidance to late steps [88] degrades results the more guidance is removed, and reusing gradients across steps hurts accuracy (Appendix A.1b). An SDEdit warm start [89] matches the baseline but does not beat it. The only acceleration that actually helps is the NFE reallocation above. Wall-times per draw are listed alongside NFE in Tables 8 and 9; Fig. 24 shows the SR accuracy against wall time instead of NFE. Table 8: SR sampler ablations: per-variable latitude-weighted physical RMSE. Units: t2m K, u10m/v10m m/s, q1000 g/kg, z500 m2 /s2 . Bicubic = bicubic upsampling of the observed frames; AE floor = Autoencoder reconstruction error. Bold = best, underline = second-best, per variable. method

NFE wall (s)

DPS +corr(all σ) +corr +corr (N30) +corr (N25)

50 99 76 47 38

bicubic (no prior) AE floor

-

123 243 186 116 94

t2m

u10m

v10m

q1000

z500

1.09±0.04 1.02±0.04 1.02±0.04 1.03±0.04 1.05±0.04

1.13±0.02 1.08±0.02 1.09±0.02 1.10±0.02 1.12±0.02

1.17±0.02 1.12±0.02 1.13±0.02 1.15±0.02 1.16±0.02

0.65±0.04 0.61±0.03 0.61±0.03 0.63±0.03 0.64±0.03

66.6±3.1 59.5±1.0 60.4±0.9 62.8±1.3 65.9±2.6

- 5.64±0.30 4.29±0.10 3.97±0.14 2.44±0.09 794.7±30.0 - 0.84±0.03 0.53±0.01 0.48±0.01 0.42±0.02 37.0±1.9

Table 9: IGRA data-fusion sampler ablations: per-variable physical RMSE (mean±std; units as Table 8; bold = best, underline = second-best). method

NFE wall (s)

DPS (50 steps) +corr (50 st, corr σ<3) +corr (N30) (30 st, corr σ<5) +corr (N25) (25 st, corr σ<3)

A.3

50 76 47 38

142 215 134 109

t2m

u10m

v10m

q1000

z500

2.23±0.21 2.16±0.20 2.27±0.25 2.33±0.26

3.32±0.11 3.30±0.11 3.32±0.11 3.35±0.11

3.49±0.14 3.47±0.14 3.49±0.14 3.51±0.14

1.25±0.03 1.23±0.03 1.27±0.03 1.28±0.03

567.5±29.6 560.7±30.5 567.4±22.4 573.1±22.5

Reshaping gradients and noise

Three modifications change how the sampler uses the computed gradient and noise (Sec. 2.3.1). First, on the sparsestation tasks, use λ=1 when the ensemble spread itself matters. On SR task, the λ=0 shows worse performance 31

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 23: SR sampler efficiency: SR physical RMSE per variable (six panels: t2m, u10m, v10m, q1000, z500, t850) vs. NFE; lower is better. Error bars are ±1 standard deviation.

Figure 24: Physical RMSE per variable (six panels: t2m, u10m, v10m, q1000, z500, t850) vs. wall-time per draw; markers are the 8-window mean with ±1 standard-deviation error bars.

(Tables 10, 11). Second, DSG normalization [79] shows significant performance for station data because the misfit gradient varies with the per-frame report schedule and DSG normalization is needed to keep the updates stable. DSG outperforms on the sensor measurements on most of the variables (Table 12). Third, momentum guidance is safe only at small decay; it is roughly neutral up to m=0.3 and degrades beyond that. All headline configurations keep m=0. We hypothesize that it may be helpful for highly non-linear or noisy operators. 32

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Table 10: SR sampler noise, normalization, and momentum ablations: per-variable physical RMSE (8-window mean of 8-member ensembles; units as Table 8; bold = best, underline = second-best per variable). This table is complementary to Table 8, which gives the efficiency ablations. method

NFE

t2m

u10m

v10m

q1000

z500

DPS +corr (N30)+λ0 (low-NFE) +mom0.5 (λ=1) +corr(all σ) (prior best) +corr(all σ)+λ0 +corr(all σ)+λ0+mom

50 47 50 99 99 99

1.09±0.04 1.05±0.03 1.04±0.04 1.02±0.04 1.04±0.04 1.04±0.04

1.13±0.02 1.09±0.02 1.09±0.02 1.08±0.02 1.09±0.02 1.08±0.02

1.17±0.02 1.13±0.02 1.13±0.02 1.12±0.02 1.12±0.02 1.12±0.02

0.65±0.04 0.63±0.03 0.62±0.03 0.61±0.03 0.62±0.03 0.63±0.03

66.6±3.1 63.9±1.7 60.3±0.9 59.5±1.0 64.3±1.9 64.0±2.1

AE floor

- 0.84±0.03 0.53±0.01 0.48±0.01 0.42±0.02 37.0±1.9

Table 11: IGRA sampler noise, normalization, and momentum ablations: per-variable physical RMSE (units as Table 8; bold = best, underline = second-best). This table is complementary to Table 9, which gives the efficiency ablations. method

NFE

t2m

u10m

v10m

q1000

z500

76 76 76 76 76

1.90±0.13 1.91±0.11 2.07±0.19 2.14±0.21 2.16±0.20

3.01±0.17 3.03±0.17 3.23±0.12 3.29±0.11 3.30±0.11

3.08±0.18 3.11±0.20 3.39±0.14 3.45±0.14 3.47±0.14

1.15±0.04 1.15±0.03 1.21±0.02 1.23±0.02 1.23±0.03

474.9±50.4 480.0±50.8 539.4±31.8 561.5±32.3 560.7±30.5

DPS+corr+DSG+λ0 (scale 0.5) DPS+corr+DSG (λ=1, scale 0.5) DPS+corr+λ0+mom DPS+corr+λ0 DPS+corr

Table 12: Sampler ablations on the ISD and joint (IGRA+ISD) observations (per-variable physical RMSE, units as Table 8; bold = best, underline = second-best). Full set in Appendix E.1. Observation

method

t2m

u10m

v10m

q1000

z500

ISD (surface) DPS 2.10±0.08 3.39±0.12 3.55±0.19 1.44±0.06 641.1±34.9 +corr+DSG 2.15±0.10 3.36±0.19 3.35±0.22 1.39±0.09 669.7±59.8 +corr+DSG+λ0 2.21±0.11 3.34±0.18 3.33±0.20 1.41±0.09 701.3±65.8 Joint (both)

A.4

DPS 2.09±0.17 3.27±0.11 3.43±0.13 1.23±0.02 557.3±27.4 +corr+DSG 1.81±0.10 2.95±0.14 3.01±0.17 1.13±0.04 467.2±52.8 +corr+DSG+λ0 1.80±0.11 2.93±0.14 3.00±0.17 1.13±0.04 467.6±48.5

Likelihood weighting

Table 13 sweeps the SR likelihood weight (guidance scale and the annealing coefficient γ) for the base DPS sampler. Table 14 contrasts DPS guidance at three scales against DSG-normalized guidance across six scales on the IGRA network. SR uses scale 4.0, γ=0.1 and all station tasks use DSG normalization at scale 0.5. The structured SR misfit is summed over the observed pixels, whereas the station misfit is averaged over points, channels, and observed frames. Under the averaged reduction, the guidance scale acts as a per-channel weight, so it is re-selected whenever the observed channel set changes, and the variable-splitting baselines require the same retuning. A.5

Comparison with existing samplers

We benchmark against existing inverse-problem samplers. We port them from their respective paper and tune them for our case, on both observation experiments in Table 15. The compared methods include the leading variable-splitting methods of the I NVERSE B ENCH benchmark [57]: DAPS [54] and PnP-DM [53], its most accurate methods on the closed-form operators, together with the efficient DiffPIR [52], alongside RED-diff [51], STeP’s HMC sampler [70], guided Restart [90], and a back-projection/DDNM port [50]. The table also has graded corrector, our own ablation that adds a second batch of Langevin steps in a lowest-σ band on top of the single-band corrector of Alg. 2. The two experiments show different outcomes. On the uniformly observed SR task, several of those methods are genuinely competitive with 8-member ensembles: guided Restart is the most accurate on the winds, and our ablated graded corrector on temperature and geopotential. RED-diff, DAPS, and PnP-DM show poor performance, the latter two at an order of magnitude more cost per draw, because their decoupled annealing runs an inner reverse-diffusion and Langevin chain at every noise level. On the sparse radiosondes, in contrast, every other sampling technique performs inferior to our sampler on every variable, and the variable-splitting methods fail badly (we did not pursue some methods that 33

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Table 13: SR likelihood-weight ablation (base DPS, 50 steps, 4 tuning windows disjoint from the test set): per-variable physical RMSE (units as Table 8; bold = best within each sweep). Guidance scale swept at γ=0.1; γ swept at scale 2.5. The scale trend flattens near 3-3.5; the production SR runs use scale 4.0, γ=0.1. setting

t2m

u10m

v10m

q1000

z500

guidance scale (γ=0.1) scale 1.5 1.78±1.02 scale 2.0 1.39±0.40 scale 2.5 1.14±0.02 scale 3.0 1.13±0.02 scale 3.5 1.14±0.01 scale 4.0 1.15±0.02

1.57±0.09 1.48±0.05 1.41±0.01 1.39±0.02 1.37±0.01 1.36±0.02

1.63±0.08 1.55±0.05 1.48±0.02 1.45±0.02 1.44±0.02 1.43±0.02

0.82±0.09 0.77±0.05 0.74±0.02 0.73±0.01 0.73±0.02 0.74±0.02

97.8±5.3 90.6±5.9 83.3±2.1 80.7±1.8 79.4±2.1 79.5±2.4

annealing γ (scale 2.5) γ=10−2 1.80±0.09 γ=3 × 10−2 1.37±0.07 γ=10−1 1.14±0.02 γ=3 × 10−1 1.76±0.55

1.60±0.03 1.38±0.02 1.41±0.01 1.76±0.09

1.65±0.03 1.44±0.03 1.48±0.02 1.83±0.08

0.93±0.03 110.0±6.2 0.81±0.03 86.2±3.9 0.74±0.02 83.3±2.1 0.87±0.05 129.2±18.9

Table 14: Station likelihood ablation on IGRA (4 validation windows, single seed): DPS guidance scale vs. DSGnormalized guidance, per-variable physical RMSE (units as Table 8; bold = best). setting DPS (scale 3) DPS (scale 5) DPS (scale 7) DSG scale 0.05 DSG scale 0.15 DSG scale 0.3 DSG scale 0.5 DSG scale 0.7 DSG scale 1.0

t2m

u10m

v10m

q1000

z500

3.06±0.45 2.73±0.34 3.31±0.62 3.06±0.41 2.91±0.47 2.44±0.13 2.43±0.09 2.64±0.26 2.90±0.45

4.52±0.20 4.38±0.16 4.44±0.13 4.35±0.22 4.04±0.22 3.94±0.30 3.83±0.25 3.93±0.28 4.09±0.44

4.83±0.21 4.68±0.17 4.63±0.09 4.53±0.15 4.13±0.17 4.10±0.27 3.92±0.29 4.01±0.34 3.83±0.20

1.63±0.04 1.57±0.04 1.63±0.10 1.66±0.22 1.52±0.21 1.42±0.13 1.40±0.11 1.41±0.14 1.48±0.12

834.1±59.6 773.1±59.2 791.8±28.9 793.5±63.1 687.1±66.0 674.2±114.9 626.9±79.3 701.5±135.8 717.7±163.7

showed clear inferiority). We conjecture that the decoupled-annealing robustness that makes these methods accurate on nonlinear closed-form operators offers little advantage for our smooth, differentiable operator, where gradient guidance is already stable. A major issue for the station observations is how the data term is scaled. The strength hyperparameters used for SR case are far too weak for the sub-steps of the variable-splitting samplers (Table 18). Once that strength is re-tuned for the station likelihood (DiffPIR σn =5 × 10−4 ,PnP-DM τ =10−4 , matching the DAPS retune), the gap shrinks considerably and DiffPIR comes close to the guided samplers, but none of them beats our sampler on any variable. The remaining gap can be accounted for by the changes in the number of station locations from frame to frame, which keeps the data sub-step poorly scaled. The performance of these methods reflects our implementations rather than the methods in general. A.6

Variable-splitting samplers: further ablations

Tables 16-17 give further hyperparameter ablations behind the DAPS rows of Table 15, and Table 18 those for DiffPIR and PnP-DM. The likelihood weight τ (the assumed measurement-noise scale in the data term of the inner Langevin) controls data-fitting strength. On SR a small τ performs well, and on the station data a much smaller τ recovers a usable performance. All DAPS runs use 40 annealing levels with 3 inner reverse-diffusion and 30 Langevin steps each, chosen for reasonable compute in comparison to other methods. This is roughly 1/8 of the I NVERSE B ENCH default denoiser budget. Even at this budget DAPS costs an order of magnitude more per draw than our sampler. We also perform ablation studies on other methods. Because the decoder operator is nonlinear, we replace DiffPIR’s closed-form proximal data solve (Zhu et al.) by a single clamped gradient step. Its prior-trust weight w is swept, where a larger w weakens the data step since the step size is ∝ 1/w (data clamp 0.3 at w=100, 0.2 at w=300). DiffPIR runs 50 steps with ξ=0.1 and uses our stabilized re-noising variant, with the effective noise computed from the pre-correction estimate. PnP-DM is the split-Gibbs sampler which uses a likelihood Langevin step anchored on the previous clean 34

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Table 15: Comparison with other existing samplers, on both observation geometries: per-variable physical RMSE (8-window mean of 8-member ensemble means; units as Table 8; bold = best, underline = second-best per panel). NFE counts denoiser evaluations per draw; † marks methods whose inner loops use additional compute. method

NFE

t2m

u10m

v10m

q1000

z500

1.30±0.13 1.32±0.16 1.05±0.04 1.08±0.05 1.02±0.04 1.20±0.08 1.08±0.09 1.56±0.34 1.05±0.03

1.65±0.06 1.17±0.03 1.07±0.02 1.10±0.02 1.09±0.02 1.28±0.04 1.16±0.02 1.50±0.06 1.09±0.02

1.70±0.05 1.20±0.03 1.10±0.02 1.14±0.02 1.13±0.02 1.34±0.04 1.20±0.02 1.53±0.06 1.13±0.02

0.78±0.02 0.75±0.08 0.63±0.03 0.63±0.03 0.62±0.03 0.67±0.03 0.66±0.03 0.77±0.04 0.63±0.03

104.9±4.6 86.9±13.0 66.1±2.5 69.3±2.8 60.7±1.1 78.6±3.0 70.1±3.6 101.2±7.8 63.9±1.7

2.35±0.23 2.25±0.20 2.27±0.21 3.43±0.61 2.53±0.29 4.13±0.58 1.90±0.13

3.43±0.10 3.38±0.10 3.40±0.10 3.86±0.14 3.32±0.21 4.18±0.10 3.01±0.17

3.60±0.14 3.55±0.14 3.57±0.14 3.94±0.13 3.39±0.18 4.06±0.14 3.08±0.18

1.29±0.03 1.27±0.03 1.27±0.03 1.57±0.11 1.28±0.06 1.89±0.18 1.15±0.04

613.9±29.8 590.2±28.4 597.1±29.3 767.9±69.1 600.0±95.5 857.9±37.5 474.9±50.4

Super-resolution (uniform space-time subsampling) RED-diff [51] 150 STeP-HMC [70] 30† guided Restart [90] 85 back-projection / DDNM [50] 76 graded corrector (ours, ablated) 71 DAPS [54] 120† DiffPIR [52] 50 PnP-DM [53] 400† ours (DPS+corr (N30)+λ0) 47 IGRA radiosondes (real soundings) STeP-HMC [70] 50† guided Restart [90] 95 graded corrector (ours, ablated) 110 DAPS [54] 120† DiffPIR [52] 50 PnP-DM [53] 400† ours (DPS+corr+DSG+λ0) 76

Table 16: DAPS ablation on SR (8-window mean; 40 anneal steps, 3 inner reverse-diffusion + 30 Langevin steps each): per-variable physical RMSE (units as Table 8; bold = best). setting τ =0.02 τ =0.01 τ =0.01 + λ = 0 τ =0.01 + DSG (lr 10−2 ) τ =0.01 + DSG + (λ = 0)

t2m

u10m

v10m

q1000

z500

2.16±0.80 1.20±0.08 1.50±0.33 2.52±0.46 2.48±0.32

1.76±0.08 1.28±0.04 1.56±0.06 2.94±0.21 2.92±0.22

1.82±0.07 1.34±0.04 1.64±0.07 3.08±0.15 3.03±0.17

0.87±0.07 0.67±0.03 0.76±0.03 1.31±0.15 1.31±0.13

123.0±16.1 78.6±3.0 96.5±5.5 258.0±34.2 250.6±21.7

Table 17: DAPS ablation on IGRA: per-variable physical RMSE (units as Table 8; bold = best per column). The row carried into the IGRA panel of Table 15 is τ =10−4 . The SR-tuned τ (0.005-0.01) badly underfits (note the large t2m errors); τ ≈10−4 or the scale-invariant DSG step is needed. The τ ∈ {0.01, 0.005} rows use Langevin lr 2×10−4 and the τ ∈ {10−4 , 3×10−5 } rows lr 10−4 . setting τ =0.01 τ =0.005 τ =3 × 10−5 τ =10−4 (8-member) τ =10−4 + (λ = 0) DSG (lr 5 × 10−3 ) DSG (lr 2 × 10−2 ) DSG (lr 10−2 ) DSG (lr 10−2 ) + (λ = 0)

t2m

u10m

v10m

q1000

z500

7.93±3.17 8.84±4.17 10.22±1.03 3.43±0.61 4.24±0.85 5.75±1.64 5.02±0.75 5.03±1.01 4.55±0.81

5.69±0.28 5.62±0.39 5.32±0.31 3.86±0.14 4.97±0.15 5.37±0.43 5.17±0.29 5.25±0.23 5.12±0.26

5.92±0.27 5.77±0.29 4.23±0.15 3.94±0.13 5.13±0.25 5.50±0.30 5.47±0.22 5.45±0.23 5.29±0.14

3.13±0.97 3.39±1.25 2.51±0.11 1.57±0.11 2.12±0.29 2.52±0.59 2.41±0.41 2.36±0.32 2.29±0.33

1530.9±358.4 1586.2±486.5 2118.5±136.9 767.9±69.1 988.9±100.0 1191.6±252.5 1064.7±118.1 1118.8±127.9 1032.7±123.5

sample (coupling ρ set to the annealed σ) alternates with a stochastic reverse-flow prior step, at 40 annealing levels with 10 inner ancestral and 30 Langevin steps each. 35

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Table 18: DiffPIR and PnP-DM ablation, per-variable physical RMSE (units as Table 8). Top block: SR; bottom: IGRA. setting

B

t2m

u10m

v10m

q1000

z500

SR DiffPIR (w=100) DiffPIR (w=300) PnP-DM (τ =0.02) PnP-DM (τ =0.01) PnP-DM (τ =0.005)

1.08±0.09 1.20±0.06 2.32±1.01 1.56±0.34 1.55±0.30

1.16±0.02 1.49±0.04 1.93±0.09 1.50±0.06 1.48±0.04

1.20±0.02 1.56±0.03 1.95±0.08 1.53±0.06 1.54±0.03

0.66±0.03 0.76±0.03 1.03±0.08 0.77±0.04 0.76±0.03

70.1±3.6 94.5±3.3 158.3±12.4 101.2±7.8 95.5±6.5

IGRA DiffPIR (σn =5 × 10−4 , 8-member) DiffPIR (σn =5 × 10−4 , w=300) DiffPIR (σn =5 × 10−3 ) DiffPIR (σn =0.05) PnP-DM (τ =3 × 10−5 ) PnP-DM (τ =10−4 , 8-member) PnP-DM (τ =0.005)

2.53±0.29 3.43±0.68 3.48±0.82 6.73±2.62 14.56±1.87 4.13±0.58 7.76±2.90

3.32±0.21 4.55±0.18 4.72±0.21 5.56±0.28 7.51±0.58 4.18±0.10 6.74±0.59

3.39±0.18 4.75±0.22 4.84±0.27 5.68±0.23 4.69±0.23 4.06±0.14 4.68±0.12

1.28±0.06 1.74±0.12 1.86±0.19 2.84±0.82 3.39±0.29 1.89±0.18 2.95±0.86

600.0±95.5 835.3±52.6 877.2±63.8 1360.5±285.7 2643.2±212.0 857.9±37.5 2210.0±283.1

Model and training details

Tables 19 and 20 collect the complete architecture and training configuration of the two trained components, the autoencoder and the TrigFlow prior; everything else in the paper runs zero-shot on these frozen weights. Table 19: Autoencoder: architecture and training details. Architecture parameters widths / blocks level strides latent saturation surface conditioning dropout

451.1 M (encoder 188.6 M, decoder 262.4 M) [96, 192, 384, 768]; 3 residual blocks per level (+1 per decoder level) [(1, 2, 2), (2, 2, 2), (2, 1, 1)] softclip2, scale 10 land-sea mask + orography, area-resampled per decoder level, injected through zero-initialized 3×3 convolutions 0.05

Optimization objective optimizer learning rate batch latent regularizer epochs

latitude-weighted, variance-normalized MSE (variance floor 10−2 ) + 0.05×(spatial-gradient + Laplacian + temporal-difference) Adam, β=(0.9, 0.95), weight decay 0 5 × 10−4 , cosine-annealed; gradient clip 5 = 128 Gaussian latent noise, std 0.02, before the decoder 155

Data and hardware training data hardware training time

C

32-frame ERA5 windows, 1979-2018 (2019 validation) 16× NVIDIA A100-40GB ≈ 9 days

Grid-cell super-obbing and the observed-window length

Figure 25 shows the effect of the super-obbing step of Sec. 4.4. This process collapses the clustered raw observations to one weighted super-observation per occupied grid cell. We observe that this lowers latitude-weighted forecast RMSE in the first forecast timesteps. Beyond day two the raw and super-obbed forecasts are essentially similar. However, we use super-obbing in our results following previous DA works [91]. The length of the observed window does not drastically change the outcomes either (Fig. 26). Shortening it from two days (eight frames) to one day leaves day-1 to day-3 skill essentially unchanged and even conditioning on a single analysis frame only slightly worsens the first forecast timestep. This result is consistent with previous work on AI based weather forecasting [92]. 36

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Table 20: TrigFlow prior (DiT3D): architecture and training details. Architecture parameters input tokenization width / depth / heads timestep embedding

525.4 M standardized AE latents (128, 8, 32, 64) patch (1, 2, 2) → 4096 tokens 1536 / 12 / 24; MLP ratio 4 sinusoidal, size 384, input scale 1000

Optimization objective optimizer lr schedule batch size EMA precision training length

TrigFlow velocity matching, σd =1, ln σ ∼ N (0, 1.52 ), adaptive log-variance weighting; σ ∈ [0.002, 80] AdamW, lr 2 × 10−4 , β=(0.9, 0.999), weight decay 0.01; gradient clip 1 500-kimg linear ramp-up, floor at 0.05× peak 256 half-life 500 kimg, ramp-up ratio 0.05 bf16 autocast (TF32 matmul) 19,901 kilo images

Data and hardware training data hardware Training time

D

AE latents of the ERA5 training windows (1979-2018) 8× NVIDIA A100-40GB ≈ 5 days

Forecasting from the Aardvark latent observations

Sec. 4.4 uses only the fraction of the global observing system that our interpolation operator ingests directly. Especially fields at higher pressure levels are weak because a few hundred radiosonde measurements sample it sparsely and the satellite observations are excluded. To experiment with richer observations, we use the Aardvark latent observations [86], whose encoder maps several observation modalities (including the satellite radiances, scatterometer, and geostationary imagery) onto a regular grid of 24 ERA5 variables. We use the Aardvark latent observation to forecast a 23.25-day lead (Fig. 27). This experiment is just an illustration, not a rigorous test case benchmark. The public Aardvark release contains only one observation data point, on 2018-08-18, (which falls within our training period). A held-out, multiple initial conditions test would require an extended observation-assembly pipeline, which we leave for future work. In this simplified experiment, we observe that the Aardvark latent lowers the forecast errors, with some variables performing better than GraphDOP.

E

Extended results

This section expands the five-variable tables and the selected-variable figures of the main text across the full channel set: the complete per-variable RMSE tables for every experiment (Sec. E.1), the per-channel super-resolution reconstructions (Sec. E.2), and the reconstruction-versus-observation scatter for every observed variable (Sec. E.3). E.1

Full per-variable physical RMSE

The main-text tables report five representative variables; here we give the latitude-weighted physical RMSE of all 69 output channels. Tables 21-26 are divided by variables and pressure levels and columns have the different tasks: super-resolution and the IGRA, ISD, and joint measurements, each with its recommended sampler configuration from Table 2 (mean±std). Table 21: Full per-variable physical RMSE, Surface (SR: DPS+corr (N30)+λ0; stations: DPS+corr+DSG+λ0). variable

SR

IGRA

ISD

Joint

2m temperature (K) 1.05±0.03 1.90±0.13 2.21±0.11 1.80±0.11 10m u-wind (m/s) 1.09±0.02 3.01±0.17 3.34±0.18 2.93±0.14 10m v-wind (m/s) 1.13±0.02 3.08±0.18 3.33±0.20 3.00±0.17 MSLP (Pa) 77.6±2.1 469.3±44.4 492.8±43.1 447.7±38.1

37

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Table 22: Full per-variable physical RMSE, Geopotential (m2 /s2 ) (SR: DPS+corr (N30)+λ0; stations: DPS+corr+DSG+λ0). variable

SR

IGRA

ISD

Joint

50 hPa 109.6±4.1 857.8±327.3 2461.9±466.1 839.2±319.1 100 hPa 93.6±8.5 709.6±168.8 2202.2±490.6 703.0±164.4 150 hPa 82.7±6.0 678.0±109.0 1969.2±407.9 670.2±109.8 200 hPa 80.3±5.3 684.6±84.0 1721.5±319.8 676.1±83.1 250 hPa 85.1±6.0 694.6±73.7 1495.6±237.1 685.8±71.4 300 hPa 81.8±4.4 668.1±67.4 1287.7±179.8 659.5±64.6 400 hPa 73.9±2.9 566.4±57.4 942.0±108.5 558.7±54.7 500 hPa 63.9±1.7 474.9±50.4 701.3±65.8 467.6±48.5 600 hPa 56.9±1.3 410.1±45.2 540.8±38.3 402.2±43.4 700 hPa 54.7±1.7 370.7±40.5 443.5±25.5 361.2±38.3 850 hPa 54.0±1.7 349.0±34.9 383.0±29.8 336.8±31.4 925 hPa 56.4±1.9 355.6±34.0 383.3±32.8 341.9±29.8 1000 hPa 61.0±1.8 369.2±34.5 392.9±34.7 353.5±29.9

Table 23: Full per-variable physical RMSE, Zonal wind u (m/s) (SR: DPS+corr (N30)+λ0; stations: DPS+corr+DSG+λ0). variable

SR

IGRA

ISD

Joint

50 hPa 100 hPa 150 hPa 200 hPa 250 hPa 300 hPa 400 hPa 500 hPa 600 hPa 700 hPa 850 hPa 925 hPa 1000 hPa

1.78±0.05 2.10±0.06 2.40±0.06 2.57±0.05 2.72±0.05 2.73±0.05 2.50±0.06 2.19±0.04 2.00±0.04 1.87±0.04 1.68±0.03 1.54±0.04 1.20±0.03

4.47±0.92 4.91±0.25 6.10±0.35 7.46±0.47 8.26±0.49 8.19±0.42 6.95±0.33 5.74±0.29 5.01±0.27 4.57±0.24 4.27±0.23 4.28±0.22 3.33±0.19

10.29±1.38 7.59±0.73 9.35±0.64 10.60±0.59 11.06±0.51 10.56±0.46 8.63±0.38 6.96±0.34 5.90±0.30 5.31±0.28 4.86±0.26 4.79±0.25 3.75±0.22

4.43±0.93 4.85±0.26 6.02±0.36 7.35±0.45 8.16±0.45 8.10±0.40 6.88±0.30 5.68±0.27 4.94±0.25 4.50±0.22 4.18±0.20 4.18±0.18 3.25±0.16

Table 24: Full per-variable physical RMSE, Meridional wind v (m/s) (SR: DPS+corr (N30)+λ0; stations: DPS+corr+DSG+λ0). variable

SR

50 hPa 100 hPa 150 hPa 200 hPa 250 hPa 300 hPa 400 hPa 500 hPa 600 hPa 700 hPa 850 hPa 925 hPa 1000 hPa

1.73±0.04 2.07±0.07 2.35±0.07 2.53±0.07 2.70±0.05 2.73±0.04 2.49±0.05 2.17±0.04 1.97±0.03 1.84±0.03 1.67±0.04 1.57±0.04 1.25±0.02

IGRA

ISD

Joint

3.39±0.57 5.10±0.66 3.37±0.58 4.34±0.29 5.76±0.13 4.30±0.26 5.79±0.42 7.74±0.31 5.72±0.39 7.37±0.57 9.57±0.44 7.24±0.55 8.39±0.62 10.42±0.50 8.25±0.61 8.38±0.53 10.07±0.38 8.25±0.52 7.06±0.42 8.21±0.28 6.95±0.40 5.77±0.39 6.60±0.27 5.67±0.37 4.94±0.37 5.55±0.28 4.84±0.35 4.42±0.32 4.89±0.26 4.33±0.30 4.13±0.29 4.45±0.26 4.03±0.27 4.29±0.27 4.62±0.27 4.18±0.26 3.44±0.20 3.75±0.25 3.35±0.20

38

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 25: Effect of grid-cell super-obbing on forecast skill: the raw clustered observations versus one super-observation per cell, both compared against ERA5 and averaged over the eight test initial conditions (8-member ensembles). Gray reference lines mark persistence of the last observed frame (dashed) and the annual climatology (dotted).

E.2

Per-channel super-resolution reconstructions

The qualitative figures in the main text show a small set This appendix shows the reconstruction of every one of the 69 output channels: the four surface fields and the five upper-air fields (geopotential, zonal and meridional wind, 39

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 26: Observed-window ablation for the in-situ obs→forecast task of Sec. 4.4: conditioning on a two-day window, a one-day window, and a single analysis frame. The plots show Latitude-weighted RMSE against ERA5 versus lead.

temperature, and specific humidity), each at all 13 pressure levels (50-1000 hPa; Figs. 28-38). Every panel shows the ERA5 truth, the posterior ensemble mean (the same 8-member protocol as the quantitative tables), and the signed error (mean − truth) for the super-resolution task (8× spatial subsample on every fourth frame, DPS+corr (N30) sampler) at a representative observed frame. 40

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 27: Forecast using a single (train-period) latent observation of Aardvark (2018-08-18). GraphDOP and IFS as reported in [28] through day 10. The plots show Latitude-weighted RMSE against ERA5 versus lead. E.3

Per-variable observation scatter

The main text shows truth-vs-reconstruction scatter for three variables per task (Figs. 11 and 13). Here we give the record of all remaining variables, at the station locations and at off-station grid points, pooled over the eight observed frames and eight windows.

41

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Table 25: Full per-variable physical RMSE, Temperature (K) (SR: DPS+corr (N30)+λ0; DPS+corr+DSG+λ0). variable

SR

IGRA

ISD

Joint

50 hPa 100 hPa 150 hPa 200 hPa 250 hPa 300 hPa 400 hPa 500 hPa 600 hPa 700 hPa 850 hPa 925 hPa 1000 hPa

0.92±0.03 0.86±0.03 0.70±0.02 0.72±0.02 0.73±0.02 0.71±0.02 0.71±0.02 0.72±0.01 0.74±0.01 0.77±0.01 0.97±0.02 0.93±0.03 0.93±0.03

2.47±0.63 2.29±0.41 2.25±0.27 2.68±0.24 2.43±0.11 2.15±0.14 2.27±0.15 2.26±0.14 2.23±0.13 2.19±0.14 2.37±0.15 2.19±0.18 1.96±0.17

5.91±1.05 4.60±0.68 4.57±0.92 5.53±0.98 5.38±1.06 5.03±0.92 4.79±0.72 4.27±0.54 3.84±0.43 3.49±0.37 3.00±0.13 2.62±0.13 2.31±0.11

2.48±0.60 2.26±0.43 2.23±0.26 2.65±0.21 2.42±0.13 2.12±0.15 2.21±0.13 2.20±0.10 2.16±0.09 2.12±0.10 2.29±0.14 2.09±0.15 1.86±0.14

stations:

Table 26: Full per-variable physical RMSE, Specific humidity (g/kg) (SR: DPS+corr (N30)+λ0; stations: DPS+corr+DSG+λ0). variable

SR

IGRA

ISD

Joint

50 hPa 100 hPa 150 hPa 200 hPa 250 hPa 300 hPa 400 hPa 500 hPa 600 hPa 700 hPa 850 hPa 925 hPa 1000 hPa

0.000±0.000 0.000±0.000 0.002±0.000 0.010±0.001 0.034±0.002 0.078±0.003 0.232±0.007 0.447±0.014 0.649±0.012 0.873±0.022 1.064±0.038 0.794±0.027 0.627±0.031

0.000±0.000 0.001±0.000 0.002±0.000 0.014±0.001 0.050±0.004 0.118±0.008 0.361±0.021 0.710±0.022 1.061±0.017 1.419±0.048 1.622±0.057 1.220±0.055 1.151±0.036

0.000±0.000 0.001±0.000 0.003±0.000 0.019±0.002 0.063±0.004 0.148±0.008 0.463±0.022 0.886±0.043 1.248±0.063 1.645±0.087 1.859±0.110 1.497±0.112 1.411±0.094

0.000±0.000 0.001±0.000 0.002±0.000 0.014±0.001 0.050±0.003 0.117±0.008 0.358±0.021 0.703±0.022 1.050±0.015 1.402±0.044 1.603±0.053 1.206±0.054 1.132±0.037

42

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 28: Surface variables: ERA5 truth, 8-member ensemble mean, and signed error (mean − truth; red positive, blue negative) for the super-resolution task. The following per-level figures use the same panel layout.

43

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 29: Geopotential (per pressure level), part 1/2; same panel layout as Fig. 28.

44

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 30: Geopotential (per pressure level), part 2/2; same panel layout as Fig. 28.

45

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 31: Zonal wind u (per pressure level), part 1/2; same panel layout as Fig. 28.

46

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 32: Zonal wind u (per pressure level), part 2/2; same panel layout as Fig. 28.

47

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 33: Meridional wind v (per pressure level), part 1/2; same panel layout as Fig. 28.

48

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 34: Meridional wind v (per pressure level), part 2/2; same panel layout as Fig. 28.

49

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 35: Temperature (per pressure level), part 1/2; same panel layout as Fig. 28.

50

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 36: Temperature (per pressure level), part 2/2; same panel layout as Fig. 28.

51

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 37: Specific humidity (per pressure level), part 1/2; same panel layout as Fig. 28.

52

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 38: Specific humidity (per pressure level), part 2/2; same panel layout as Fig. 28.

53

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 39: IGRA: ERA5 truth vs. reconstruction at the radiosonde station locations, for a representative subset of the 63 observed variables (the surface fields, and each upper-air field sampled at an upper, middle and lower level).

54

Multimodal spatiotemporal atmospheric data assimilation with latent flow-matching

Figure 40: IGRA: ERA5 truth vs. reconstruction at off-station grid points, for the same representative subset of the 63 observed variables. Winds and humidity show the largest off-station scatter.

Figure 41: ISD surface observations: ERA5 truth vs. reconstruction at stations (top) and off-station grid points (bottom), all four observed variables.

55

Related documents

• Distilling deep optical flow stereo methods to retrieve dense three-dimensional wind fields#925465
• Distilling deep optical flow stereo methods to retrieve dense three-dimensional wind fields#921229
• Context-Aware Multimodal Representation Learning for Spatio-Temporally Explicit Environmental Modeling#424855
• Latent-to-Latent Flow for Volumetric Stochastic Segmentation#984381
• SemanticSeg4EO: An open-source unified framework and QGIS plugin for semantic segmentation in Earth Observation#229860