Align-DA: Align Score-based Atmospheric Data Assimilation with Multiple Preferences
Authors: Jing-An Sun
Hang Fan
Junchao Gong
Ben Fei
Kun Chen
Fenghua Ling
Wenlong Zhang
Wanghan Xu
Li Yan
Pierre Gentine
LEI BAI
Venue: NeurIPS 2025 poster ()
Keywords: data assimilation
score-based model
preference alignment
reinforcement learning
Primary area: machine_learning_for_sciences

Abstract: Data assimilation (DA) aims to estimate the full state of a dynamical system by combining partial and noisy observations with a prior model forecast, commonly referred to as the background. In atmospheric applications, this problem is fundamentally ill-posed due to the sparsity of observations relative to the high-dimensional state space. Traditional methods address this challenge by simplifying background priors to regularize the solution, which are empirical and require continual tuning for application. Inspired by alignment techniques in text-to-image diffusion models, we propose Align-DA, which formulates DA as a generative process and uses reward signals to guide background priors—replacing manual tuning with data-driven alignment. Specifically, we train a score-based model in the latent space to approximate the background-conditioned prior, and align it using three complementary reward signals for DA: (1) assimilation accuracy, (2) forecast skill initialized from the assimilated state, and (3) physical adherence of the analysis fields. Experiments with multiple reward signals demonstrate consistent improvements in analysis quality across different evaluation metrics and observation-guidance strategies. These results show that preference alignment, implemented as a soft constraint, can automatically adapt complex background priors tailored to DA, offering a promising new direction for advancing the field.

Align-DA: Align Score-based Atmospheric Data Assimilation with Multiple Preferences
Jing-An Sun"
Hang Fan"
Fudan University
Columbia University hf2526@columbia a. edu
Shanghai Artificial Intelligence Laboratory fudaan...uuu.....mmom fudan. edu.
jasun220m.
Junchao Gong
Ben Fei
Shanghai Artificial Intelligence Laboratory Shanghai Jiaotong University gj cchimeel@ssjte edu. cn
The Chinese University of Hong Kong
Shanghai Artificial Intelligence Laboratory benfei@cuhk. edu. hk
Kun Chen
Fenghua Ling
Fudan University
Shanghai Artificial Intelligence Laboratory lingf Seggffengu@@pplal. org.
Shanghai Artificial Intelligence Laboratory kunchen226m. fudan, edu. cn
Wenlong Zhang
Wanghan Xu
Shanghai Artificial Intelligence Laboratory zhangwenl ongepjlab...ppp.bb........ org.cn cn
Shanghai Artificial Intelligence Laboratory Shanghai Jiaotong University @anghannhasssj edu. cn
Li Yan Fudan University cli, www.yannntooo..oom com
Pierre Gentine
Columbia University pg2328@columbia. edu
Lei Bai
Shanghai Artificial Intelligence Laboratory
salsaansh@@maall com
Abstract
Data assimilation (DA) aims to estimate the full state of a dynamical system by combining partial and noisy observations with prior model forecast, com- monly referred to as the background. is fundamentally ill-posed due the sparsity of observations relative to the high- dimensional state space. Traditional methods address this challenge by simplifying background priors regularize the solution, which empirical and require tinual tuning for application. Inspired by alignment techniques in exx-to-image diffusion models, we propose Align-DA, (Dg, which formulates DA generative process and uses reward signals to guide background priors-replacio manual tuning with data-driven alignment. Specifically, train score-based model in the latent space to approximate the background-conditione prior, and align it
atmospheric applications, this problem
"Equal contribution. "Corrsponndig author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

using three complementary reward signals for DA: (1) assimilation accuracy, (2) forecast skill initialized from the assimilated state, and (3) physical adherence of the analysis fields. Experiments with multiple reward signals demonstrate con- sistent improvements in analysis quality across different evaluation metrics and observation-guidee guridance strategies. These results show that preference alignment, implemented as a soft constraint, can automatically adapt complex background priors tailored to DA, offering promising new direction for advancing the field.
1 Introduction
Data assimilation (DA) estimates the posterior state of a dynamical system by integrating prior model forecasts I6 (background) with observations [T 2 BJ. However, in practical applications, the number of observations is often much smaller than the dimensionality of the model state, making DA an underdetermined probabilistic mroblem-multipe of observations. This inherent ambiguity raises a fundamental question: what defines a good analysis for DA, and how can it be reliably obtained in the face of such uncertainty?
plausible states can equally explain the same set
In atmospheric science, data assimilation (DA) provides accurate initial conditions that are essential for reliable weather forecasting. It is also used to generate high-fidelity reanalysis datasets, which serve as a fundamental resource for climate research 2 4]. The quality of atmospheric analysis is commonly evaluated based on three criteria: accuracy, downstream forecast performance, and adherence to physical constraints. Over decades of development, modern atmospheric DA has adopted Bayesian inference framework, typically modeling the background-conditionee prior p(x: (z,) as a Gaussian distribution with empirically constructed covariance [ 3
To ensure the analysis satisfies these criteria, operational systems depend extensive parameter tuning-adjusting correlation length scales in variational methods [9], calibrating localization and inflation in ensemble Kalman filters (10), and modifying ensemble weights in hybrid approaches [1]] While these traditional approaches have proven effective, they are often computationally inefficient and labor- intensive. The remarkable success of machine learning in medium-range weather prediction [K
14
as a core component of end-to-end forecasting systems 223 24 25 However, current ML approaches face a fundamental limitation: they lack explicit mechanisms to incorporate domain-
15
16 17] has spurred a growing interest in ML-based DA [18] 19
20
21
22
knowledge, especially physical constraints 25 26 (7). While some studies have applied reinforcement learning to data assimilation for the above challenge, these approaches typically frame DA sequential pecision-making problem (28 29 30, which is challenging for direct application to large-scale, high- -ineeiinesson systems like global weather [3].
To this end, we propose Align-DA, a novel framework that aligns DA analyses with key preferences in atmospheric applications using reinforcement learning. It is the introduction of a new paradigm shift for DA, replacing the manual, experience-driver tuning process with a preference alignment- based framework, which becomes the core contribution of this paper. Instead of new predictive model architecture, we propose a pluggable alignment module designed based DA method, enabling it to satisfy complex domain preferences that are difficult to model directly. Specifically, we model the background background-conditinn conditiiond prior distribution p(z) ((z))r)) directly using a diffusion model and incorporate observations via guided sampling to approximate the posterior p(z]/r,))) y) [33. Inspired by alignment techniques in diffusion- image generation, we further refine the pretrained diffusion model via direct preference optimization (DPO) (34] using domain-
enhance any score-
specific reward signals tailored for atmospheric DA. As illustrated in Figure [ this alignment process makes the posterior ((x y) more focused, effectively narrowing the solution space and enabling better analysis and forecast performance.
Given the high dimensionality of global atmospheric states, we implement the Align-DA framework in latent space learned via a variational autoencoder (VAE). Here, we consider three complementary reward metrics: assimilation accuracy, forecast skill, and physical adherence. Our experiments reveal two critical insights: 1) Joint optimization across all reward signals leads to consistent improvement in various score- score-based DA implementations, and 2) adjusting the reward composition allows the DA analysis to adapt to different downstream tasks. These results demonstrate the effective integration of prior knowledge through soft-constrain DPO, which opens new avenue for enhancing DA. Moreover, our framework supports flexible reward design, making it adaptable to wide range of DA scenarios. Our contributions are outlined as follows,

---

Novel RL-driven DA framework: We propose Align-DA, a pioneering reinforcement learning-base based DA framework that replaces empirical tuning with preference alignment and consistently improves analysis and forecast performance across score-based methods.
Soft constrain DA formalism: The proposed framework enables the incorporation of implicit or hard-to-model constraints (e.g., forecast performance, geostrophic balance) through soft, reward-based alignment, offering a flexible new avenue for enhancing DA.
2 Related work
Traditional data assimilation. Traditional DA methods are fundamentally grounded in Bayesian inference, aiming to estimate the posterior distribution p(a y). which is proportional to ((z/)))( The most widely used variational and ensemble Kalman filter methods typically as- sume that both the prior p(æ æ5)) and observation errors p x) follow Gaussian distributions B] S 6) Taking variational methods rior probability p(z) xx, (t))))()z)) which is equivalent to minimizing the following cost function in a three-dimensional DA scenario:
example, the analysis state is obtained by maximizing the poste-
1 2 ( xz) J (x 2b) + (y (xx)=>x"""
(zx
Hx)" 'gy Hz))
(I)
where B and R denote the background and observation error covariance matrices, respectively. represents the observation operator. For numerical weather prediction systems, the dimensionality of can reach 10'4, making its estimation and storage computationally intractable. Consequently, B is often simplified and empirically tuned based on the quality of the resulting analysis and the forecasts initialized from it. Recent advances in latent data assimilation (LDA) 35 36 37 38 enable more efficient assimilation in compact latent space. LDA also alleviates the challenge of estimating B, but still relies on empirical tuning.
has enabled new approaches to data assimilation through their ability to approximate complex conditional dis- tributions p(ar/ g). Existing approaches primarily differ in how they incorporate the background state and observations into the diffusion process. Rozet et al. [42] and Manshausen et al. [43 treat observations as guidance signals during the reverse diffusion process but omit the background prior, which may lead to suboptimal performance in scenarios with sparse or noisy observations. Huang et al. [33 train diffusion model conditioned on the background state and incorporate observations via a repainting scheme during sampling; however, this approach struggles to handle complex, nonlinear observation operators such as satellite radiative transfer. Qu et al. 44] encode both the background and observation into a joint guidance signal, but their method is restricted to specific observation distributions, limiting its applicability in general DA settings. Furthermore, current approaches leave two critical challenges unaddressed: (i) they lack explicit mechanisms to enforce assimilation accuracy and forecast skill during optimization, and (ii) the incorporation of physical law remains unexplored.
Score- based data assimilation (SDA). The emergence of diffusion models 39 40
The alignment of the score-based model. Our work presents the first attempt to enhance prior knowledge encoding in score- based data assimilation and embed the physical constraints through preference alignment in soft style, drawing inspiration from recent advances in diffusion model optimization. In the text-to toxt-to--mage domain, following the success of RLHF in language models [45], recent work [46] 47 48 has demonstrated the potential of fine-tuning diffusion models via reward models that capture human preferences
52 53 Seminal works by Black et al. [46) and Fan al. 44 pioneered the formulation of discrete diffusion sampling as reinforcement learning problems, enabling policy gradient optimization, though their approaches faced challenges with com- putational efficiency and training stability. The field advanced significantly with Diffusion-DDPP -PO (54), which adapted Direct Preference Optimization (55) to align diffusion models with human prefer- ences through image pairs. Subsequent innovations like DSPO [56] introduced direct score function correction, while remaining limited to single- eegle-preference alignment. Most recently, breakthroughs such as Capo [57] and BalanceDPO (58] have established multi-preferenee alignment frameworks though their application to data assimilation remains unexplored. In the DA context, evaluations traditionally focus on three key metrics: assimilation accuracy, forecast skill, and physical adherence of analyses. Our ADi framework uniquely bridges these domains by extending and adapting multi-preference alignment techniques to simultaneously optimize this triad of DA objectives through differentiable reward signals.
[49 50 5
3

---

3 Method
3.1 The score-based model and data assimilation
The score-based model employs forward and reverse stochastic differential equations (SDEs) The forward SDE progressively adds noise data distribution po(x) through:
40
dx f(z, t)dt + g(t)
(2)
while the reverse-time SDE removes noise:
pt dx = [f(x,t) v(t)2V= p: (2))dtt + g(t)dw..
Here, f(x, t) and g(t) control the drift and diffusion, with w, w both being the standard Wiener processes. The perturbation kernel follows p(zz|z)
where we adopt variance- preserving SDE with cosine scheduling schedule for putt) [42]. The score function Vz log Pt (x) can be estimated by a neural network $o 5xx t) via minimizing the denoising score matching loss 4Q) Once trained, samples are generated by solving the reverse SDE using the learned score estimate s t) & V Pz).:
((((t).o²(t))
L. = ppzz |0 // (z, 1) V log pe æo)
In SDA, aim to sample the analysis field through variational DA approaches, the posterior score function can be decomposed into two components:
arg max p(z| y). Following the
Vz, log [b, y) Vz, lo vs.) + Vz, logl(y)):)
s0 (z), Vz, logp(y)<x))
The first term represents conditional score function of p(æ: 2t6), leveraging the expressive- ness of diffusion models to learn complex prior distributions of analysis fields directly from data. The observation-aasociatee score function is called guidance, enforcing observational con- straints. In DifIDA 33,, the repainting technique is used to approximate this guidance. While DPS Guidance
34 retains the observation Gaussian likelihood assumption variational DA:
p(y (x)) ~ (2y H (z)), R), where
(x(2,) = 2: + V log log Dt + (0)),
(5)
put)
ju(t)
is the posterior mean given by Tweedie's formula diffusion model
41
3.2 The observation guidance methods latent space
Direct SDA in high-dimensional atmospheric model space poses significant computational challenges. To overcome this, we employ VAE to compress physical fields into low-dimensional manifolds [59] 60) and perform data assimilation within this reduced latent space. Our framework first trains score-based model to learn the background conditional prior p(z 25). For observation integration, adapt two guidance approaches to the latent space.
The score function decomposition forms the foundation for our latent DPS guidance approach:
Vz, log p( (ze y) A 50 2b) + Vz, (og)(()):),
where the observation term assumed Gaussian p(y]zz) ~ / (HH(D(20(z)) (2))))). R). Similar to Equa- tion 5 20 (3, + a² a²(t) Vz, log p(z:
is the posterior mean. The D(-) denotes the decoder of VAE. Assuming the 2; is the denoised latent at reverse time f, latent counterpart of DPS guidance formulated as:
2, = 2. + < log p(y)
=:
- 2
vv. (HH(()) (V-H((((2))) H()(2))),
1
where is the guidance scale. We also implement a latent-space adaptation of Repaint techniques We first sample the observation-informee latent zahe ~ (() o with noise, where æ the ERA5 truth state. Then, we fuse the observation information and background prior through masked composition in the physical space:
žt-1 E (m @ D(z("0) + (1 m) @ D(ž.))
(8)
where is spatial mask indicating observed regions and E(-) is the VAE encoder.

---

reward,
loss
p(Keruch (Kb)
Patign (Frrana (x
rewardd
sempling
tuning
Loss
DA
DA
Patign (Feruth Pbo
P
reward,
loss
win
Without Align
With Align
Figure 1: Schematic Align-DA. Conventional score-based DA learns background-condiiineed diffusion model to approximate broad prior distribution p(Zmuh æb), which, after incorporating observations, yields posterior estimate p p(zuut 2z. that may be overly dispersed (left). Align-DA leverages reward-guit -uidded alignment to adaptively refine and concentrate the prior Palign (I zx). resulting in a posterior Palign (rruth Zo. y) that better reflects observational constraints and more closely aligns with the requirements of DA (right).
3.3 The diffusion preference optimization.
The preliminary stage of SDA establishes conditional priors for potential analysis fields. However, empirical evidence reveals substantial deviations between these prior estimates and the true analysis fields [33 [43, indicating significant room for prior knowledge refinement. Furthermore, traditional data assimilation approaches and existing Al-assisted variants face inherent challenges in enforcing physical adherence constraints through conventional knowledge embedding [18] 19 20 21 Rec- ognizing these dual challenges of unreliable prior knowledge estimation and inadequate physical adherence enforcement, we propose a novel soft-constraint framework on the basis of diffusion pref- erence optimization (Diffusion-DPe) [54). extending its recent success in text-to-image generation atmospheric data assimilation.
The reward finetuning, seeking to maximize the expected reward r(z, c) while maintaining proximity to the reference distribution Pret via KL regularization [61]:
max Be,zorpp [r c)] 3 3D (pe(-cc Il Prer(-lc)),
where 3 modulates the regularization strength. Following the Bradley-Terry preference model one can construct the reward signal through pairwise comparisons:
PBr > <g(c) = a(((=] c) r(z), c)),
(10)
with the reward model trained via negative log-likelihood minimization:
LBT(r) -BCc, = [log (r(c, r(c,,)))]
(11)
where 21" and z/j are the generated "winning" and "losing" samples with condition In this work, the "winning" and "losing" samples are selected across multi-rewards (see details Section 4.2).
The analytical solution of Equation preference optimization objective:
p" (wolc) Pror( exp R(za,;)) e)) yields the direct
Po( (c) Pret (xx) [c) Pent
log (o( (2/lc)
(())= - log log
(12)
making it bypass the reward model training.
To circumvent intractable trajectory-level computations when adopting the DPO for the diffusion model, DiffusionDPPO [54) decomposes the trajectory-level DPO loss into the following step-wise
5

---

form:
LDIt-DPO = -B _ - log
Po po
,
Po a c)
[ Boo
Blog
Pref (x
x8 c)
Tef
(
/
(13)
where xt" and x are the noised zi and z/j respectively. Finally, through noise prediction network reparameterization, the implementable training objective can be derived:
((0) = (B -/)="""/>~>>>>_.) (0, >m " d'~N(o.r) [log a (-3Tw(A) (Ab(a) , A ( }))]] (14)
where (zt,,)) e) = || ea cz t)ll ||e ( t)ll quantifies the deviation of learned noise patterns from reference behaviors. With such refinement, the diffusion model is aligned specific preferences.
Experiments
4.1 Experimental Settings and Evaluations
The background-conditioned model training and sampling We introduce an integrated approach for background ObsFree distributions modeling that combines latent space learning with ObsFree diffusion processes. The framework employs a two-stage architecture beginning with transformer- based VAE [63] that reduces input dimensionality by 16x (processing high-resolution field data 69 x 128 x 256 into compressed latent codes 69 x 32 x 64) (see Appendix [ for details). The second stage implements ObsFree latent diffusion model that learns the mapping p(z|z,) (2))) through stacked transformer blocks [64]. The model utilizes (2,2) patch tokenization and integrates ObsFree information from background latent 2; via cross aooss-attention. We adopt variance preseervin diffusion process [40] with cosine noise scheduling and train the model using fixed learning rate of le (batch size 32), observing stable convergence over 100K training iterations. For ObsFree sampling, we implement hybrid strategy combining 1t8-step ancestral prediction with two rounds of Langevin- based refinement, enhaneing sample fidelity while maintaining computational efficiency
Experimental setting Our data assimilation framework employs the FengWu AI forecasting model (65] (6-hour temporal resolution) generate background fields. The 48-hour background fields are created through 8 auto-regressive 6-hour forecasts initialized from ERA5 data 4 two days prior. Observational data simulated by applying random masking to ERA5 ground truth with 99% sparsity. approximating realistic satellite observation coverage patterns. The spatial resolution 1.40625 (128 256 grid). Here, the FengWWu AI forecasting model serves as a fixed (frozen) background generator for all DA methods, including our proposed methods and the baselines. This ensures fair and controlled comparison.
4.2 The preference dataset construction
In the standard practice of RL-based post-training, preference dataset is typically required [7 54 45). In this paper, our preference dataset is constructed by generating candidate analysis fields latent given background condition from the pre-trained reference model p(> 2z). To address the multi-preference nature of data assimilation-wheee single-step assimilation accuracy, forecast skill, and physical adherence are distinct objectives-wy we adopt a holistic evaluation framework. For a batch of N -egerrated analysis latents, each sample decoded into the analysis field x, and evaluated across three rewards:
Assimilation Accuracy: The latitude-weighte root mean square error (WRMSE) [66 67 between I. and ERA5 ground truth is computed for all 69 variables. The reward for the n-th th sample is derived the average pairwise comparison across variables:
m==69 69
Rassim (n) N
p( > x)"),
(15)
69
jmn
m==1
where the Bradley-Terry preference model applied variables level p(a") > xx')==
a (-wrmssefor [[][m]] [] (emsslj][mm)) and wrmse[n][m the WRMSE of m-th variable in n nt-th sample.

---

Background
40000
50000
52000
S6000
5600D
S2000
48000
50000
$2000
S4000
56000
50000
40300
50000
$2000
54000
56000
50000
Nign-Obsfree-P error reduction
Align-Repain--P eiror reduction
Align-OPS-P error reduction
-100
150
200
100
-50
50
150
-150
-50
100
200
ObsFree- reduction
Align-Repaine M reduction
Align-DPS reduction
150
150
-150
Figure 2: Visualization of Align-DA Assimilation Performance for Z,50 Analysis. (10 averaged analysis with randomly selected timestamp 2019) 'ObsFree' means without observation integration baseline. 'Repaint' and 'DPS' represent repaint and diffusion posterior sampling guidance methods. Note that 'Pp' indicates single physical reward aligmment, while '-M' signifies multi-reward alignment. First row: Reference fields showing ERA5 reanalysis sis (ground truth denoted as GT), background field, and observational data. Second row: Error reduction through single physical reward alignment, quantified as | TGT| - x Align-P - FGT/. Warm hues indicate regions where physical adherence alignment moderately improves accuracy, suggesting potential reward correlations. Third row: Error reduction through multi-reward alignment, quantified |= #cr | M warm hues indicate regions where multi-reward alignment substantially improves accuracy compared the unaligned baseline. The attenuation of color intensity from left to right columns demonstrates diminishing DPO effects as more effective observation guidance methods are applied pattern consistent with theoretical expectations.
Øcr|- More intense
48-hour Forecast Skill: In this metric, Feng acts as both evaluator and reward source. After each DA method produces an analysis field, we use FengWu to run a 48-hour forecast. The forecast RMSEs relative to ERA5 computed, yielding a forecast preference score Rioreaast critical distinction of Align-DA DA that while baseline methods can be evaluated by FengWu, they lack any mechanism internalize this feedback improve their generative process. In contrast, our Align-DA framework is explicitly designed to learn from this feedback, which the core of its innovation.
Physical adherence: To demonstrate the effectiveness and flexibility of our framework in aligning with physical knowledge, we incorporate reward mechanisms based primary examples of atmospheric principles: geostrophic wind balance and hydrostatic balance (see the Appendix FF. The geostrophic wind components [88 are defined as:
100 Oy
104
V
(16)
Ox.
where f denotes the Coriolis parameter and + represents the geopotential. We quantify atmospheric field deviations using the metrie: D = [u-uyl |w-0,] We acknowledge that perfect balance an idealized approximation, We therefore establish the ERA5 truth deviation DeT reference benchmark. The final physical adherence score (Rphys) for

---

Table I: Assimilation accuracy gains of alignment strategies under 1% observation. Percentage values quantify relative accuracy variants compared to the non-aligned baseline for multi-reward experiments.
WRMSE
MSE
MAE
u10
v500
2500
t850
ObsFree
0.0643
0.1382
1.5275
2.7467
104,5531
1.0897
Align-OOssFree-P
0.0645
0.1385
1,5331
2.7501
103.1857
1.0896
Align ObsFree-M M
0.061 (--4440)
0.1348: 0.1348(-2.467)
1.4942: 222235)
2.7202 00.97%)
96.8603c -794%)
1.0866( 0.25%
Repaint
0.0623
0.1368
1.5009
2.7027
103.3036
1.0911
Align-Repaint-P
0,0624
0.1369
1,5063
2.7047
101,9634
1.0904
Alin-Repuint--n
0.0598(-41128
0. 11144(411778)
1. 1.4759(-1668 1.69%)
2.6836 0.71%)
96.8906 (9906c66.620) 66.62)
L08891-0.200
DPS
0.0609
0.1337
1.4441
2.5972
93,4654
1.1159
A
0,0612
0.1342
1,4514
2,6069
92,4301
1,1184
Align-DPS-M DPS-M M
(05933-2.600))
0.1317
1,4187(-
2,5800, 0.06%)
88.4235c 5.70%)
1.1 1.0 0.53%
Table 2: 48-hour forecast skill gains of alignment strategies under 1% observation.
WRMSE
MSE
MAE
u10
v500
2500
1850
ObsFree
0.1146
0,1918
2.1182
4.1571
232.6243
1.4826
Align- ObsFree-P
0,1151
0,1923
2.1203
4.1623
232.5924
1.484
Align-ObsFree-M
0.1080
0.1872(-2447%
2.0965
4.1091(-117%)
225.7758(-3.00)
146666(11300)
Repaint
0.1123
0,1898
2.0899
4.1030
228.2876
1.4672
Align-Repainte
0.1127
0.1902
2.0912
4.1067
1765
1.4677
Aligu- Reeo--Rpaant-MM M
10068
(-5.12%)
061860 2.0%%)
2.0713 0300%)
4.0665(-a.se0
222,9336i
2.40%)
144166(1007%)
DPS
0.0969
0,1757
1.9230
3.7675
193.4167
1.3564
Align: DPS-P
0.0972
0.1761
1.9275
3.7760
193.4700
1.3599
Align-DPS-M
0.0943(-2.748)
0.1734(-1.278)
1.9070
3.7287(-10098
189.6935(-1900)
-.3428(-1011)
given principle (yielding Geo-Score) then computed as:
[D; [Di-DoH DGT
Ryhys = 2-
(17)
Dar
where (-) denotes the sigmoid function. This formulation ensures that generated fields closer to the ERA5 ground truth receive higher scores, with the theoretical maximum score of corresponding perfect agreement with observational data.
To rigorously validate our Align-DA framework, particularly examining the impact of RL-driven physical knowledge embedding and the trial-and-error free formalism, we construct two distinct preference datasets: a physical adherence reward dataset (i.e. focusing on Geo-Score) and multi- reward comprehensive set. Both datasets are constructed through systematic sampling of = 32 instances from the pre-trained diffusion model's 's output space.
For the physical adherence reward dataset, we define winning samples as top-5 performers in Geo-Score, while bottom-5 instances are assigned to the losing group. The comprehensive dataset employs multi-dimensional reward ranking, where winning samples occupy the top decile across all three reward dimensions, with losing counterparts drawn from the bottom decile. To ensure statistical robustness and prevent sampling bias, we implement stratified random pairing with duplicate elimination through hash- hash-based filtering [69]. Our dataset construction leverages 2018 experimental records, yielding a carefully curated 4,000-pair corpus for model fine-tuning. The optimization employs KL regularization (3 = 8000). batch size of 32, and a learning rate of le 5, with model checkpoints saved after 500 fine-tuning steps.
4.3 Results
We evaluate our framework through two distinct experimental paradigms: single physical adher- ence reward optimization on the physics reward dataset and multi-reward optimization on the comprehensive dataset. For observation integration, we implement three established methodolo- 1) ObsFree an observation-free baseline model, 2) Repaint the standard repainting technique, also known as inpainting, and 3) DPS diffusion posterior sampling guidance (detailed in Section 3.2). Model variants are denoted through suffix conventions: -P indicates single phys- ical reward alignment, while -M signifies multi-reward alignment. For instance, Align-DPS-P.
gies 33 4
8

---

Physical adherence reward aligned model with DPS guidance, Align-ObsFred Multi-reward aligned model without observation integration, Repaint: reference model with repaint integration.
Posterior analysis improve- ment. Table [ presents quan- titative comparisons of poste- rior analysis accuracy across different alignment approaches, For each observation integration method, we evaluate three figurations: the reference model with no reward alignment, single- Anysics-reward-aligg model (- P). and multi-reward-alignedd model (-M). We report compre- hensive metrics including over- all MSE, MAE, and variable- specific WRMSE scores, accompanied by percentage improvements relative the reference baseline.
Align offect
Palign (gr
Palign (Ferveh Rg.
Pastgn Ma.
Align effect
Align effect
No observation
Stonger observation constrain
Figure 3: Interaction between observation constraints and alignment. Better incorporation of observation information (i.e., stronger constraints) leads narrower space for improvement by the alignment process.
Multi-reward alignment consistently outperforms the unaligned model, demonstrating the benefit of directly optimizing for assimilation accuracy. For instance, Align-ObsFree-M M reduces overall WRMSE by 5.44%, indicating that multi-objective reward signals lead to tighter posteriors and more accurate reconstructions.
reduces MSE by 4. 4.12%, but its single-reward version (-P) slightly worsens performance (+0.2%). indicating that physics alone insufficient for DA tasks. Nevertheless, Align-DPS-P improves Z500 WRMSE across both Repaint and ObsFree settings, illustrating the localized efficacy of enforcing geostrophic balance. Collectively, these findings underscore the importance of reward formulation: while task-specific objectives are essential for global accuracy, Nomain-specific physical constraints can provide targeted gains. Multi- signal alignment thus offers a more robust and generalizable mechanism for incorporating scientific priors into generative DA.
In contrast, single-reward alignment shows mixed results Aliin--lggaaiiiee
Visualization of analysis improvement. providess a spatial evaluation of reward alignment strategies using Z500 as diagnostic field. The first row illustrates a representative sample of the ERA5 reanalysis ground truth (GT), background field, and the sparse observations utilized for DA. The second row quantifies error reduction achieved through single physics reward alignment while the third row extends this analysis to multi-reward alignment. The results indicate that the single physical reward consistently improves Z500 across different guidance settings, possibly due to its implicit embedding within the geostrophic wind balance equations. Moreover, the multi-reward strategy yields more substantial improvements than the single reward alone, highlighting the enhanced effectiveness of preference alignment in improving assimilation quality.
Forecast skill improvement. Table 2 reports 48-hour forecast performance across different align- ment strategies under 1% observation, showing trend consistent with their effects on analysis. Multi-reward alignment consistently outperforms the baseline across metrics and variables. More importantly, the improvements achieved analysis time are not only preserved but often amplified in the forecast, For instance, Reign-Re improves analysis MSE by 4.12%, while its 48-hour forecast MSE drops by 5.12%, indicating that aligned states provide more reliable initial conditions for downstream prediction. In contrast, on geostrophic balance, continues to yield negligible or even adverse forecast gains. These results highlight the importance of explicitly optimizing for analysis and forecast objectives rather than relying solely on physical priors.
alignment (-P), which focuses solely
Alignment benefits fade with strong observational constrain. The results above demonstrate that, under well-designed reward signals, preference alignment consistently improves both posterior analysis and forecast skill. However, the magnitude of improvement varies across observation incorporation schemes. The observation-fret baseline gains the most from alignment, followed by Repaint, while DPS shows the least benefit.
This can be explained by the overlapping effects of observational constraints and preference alignment (Figure B). The observation-Fi free setting performs sampling without incorporating observational guidance, allowing greater flexibility for reward-basee alignment to shape the posterior distribution.

---

In contrast, DPS imposes the strongest observational constraints as evidenced by its superior
baseline performance in Table leaving limited room for further improvement through alignment.
Physical adherence improvement. In this study, we endeavor to enhance physical ad- herence by employing an alignment technique implemented as a soft constraint, utilizing geostrophic wind balance as a reward signal. Figure 4 presents the Geo-Score, demonstrating that while multi-reward alignment (incorporat- ing physical adherence reward) yields broad im- provements across all guidance methods (e.g., from 93.5 19), optimizing solely with the Geo-Score reward achieves more pronounced enhancements in this specific physical metric (e.g., from 93.85 to 95.00). We acknowledge, however, that slight degradation in overall as- similation accuracy and forecast skill was ob- served when only the physical adherence reward was applied. As demonstrated in Appendix EE] the performance change of its single-reward ver- sion (-P) yields statistically insignificant, sug- gesting that geostrophic balance alone insufficient for DA tasks. Nevertheless, the significant Geo-Score improvement substantiates the efficacy of alignment techniques in enhancing physical consistency, highlighting the promising direction of this approach and underscoring the importance of refining physical reward design in future work.
ObsFree
(i)
91.5
92.0
92.5
93.0
93.5
94.0
94.5
95.0
95.5
Geo-Score (56
Algn-OOssrree-i
(b)
Align Repaint-M M
Nign-OPS-M
91.5
92.0
92.5
93.5
94.0
96.5
95.0
95.5
Geo-Srore
AlignOOssrre..P
Nign-Repeint--P
Align-DPS.P
01.5
92.0
92.5
93.0
93.5
96.0
94.5
95.0
96.5
Gea-Scare (%)
Figure 4: Geo-Score comparisons across align- ment strategies: (a) Reference model, (b) Multi- reward aligned (-M), and (c) Single physical ad- herence reward aligned (-P) results, with colors denoting different guidance methods.
Conclusion
This work presents Align-DA, novel reinforcement learning framework that introduces soft- constraint alignment as new paradigm for advancing DA. Unlike traditional approaches relying on empirical tuning or hard physical constraints, our method reformulates DA as preference optimization problem, enabling the flexible integration of domain knowledge through differentiable reward signals. By refining diffusion- baased prior with DPO, the framework effectively narrows the solution space while maintaining adaptability to diverse atmospheric constraints (analysis accuracy, forecast skill, and physical consistency). The error reduction and physical enhancement after alignment demonstrate the efficacy of our alignment framework. While our results are promising, we acknowledge a key limitation of the current approach, which operates in a learned latent space. This is due the low-dimensional manifold assumption, which makes it inherently difficult for the model to capture rare, outlier events such as extreme weather. Nevertheless, Align-DA establishes soft-constraint DPO where implicit or complex prior knowledge (e.g., forecast skill, geostrophic balance) can now be incorporated without restrictive assumptions.
scalable and generalizable alternative to conventional DA formalisms,
Acknowledgements
This work supported by Shanghai Artificial Intelligence Laboratory and the JC STEM Lab of AI for Science and Engineering, funded by The Hong Kong Jockey Club Charities Trust, the Research Grants Council of Hong Kong (Project No. 4UHK14213224). This work was done during Jing-An An Sun's 's internship at Shanghai Artificial Intelligence
References
[1] Andrew C Lorenc. Analysis methods for numerical weather prediction. Quarterly Journal of the Royal Meteorological Society, 11(447)::177--114 1986.
[2] Nils Gustafsson, Tijana Janjić, Christoph Schraff, Daniel Leuenberger, Martin Weissmann, Hendrik Reich, Pierre Brousscau, Thibaut Montmerle, Eric Wattrelot, Antonín Bučánck, et al.
10

---

Survey of data assimilation methods for sonvective-scall numerical weather prediction at operational centres. Quarterty Journal of the Royal Meteorological Society, 144(713):1218- 1256, 2018.
[3] Mark Asch, Mare Bocquet, and Maëlle Nodet. Data assimilation: methods. algorithms, and applications. SIAM, 2016.
[4] 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, (99(730):199922000 2020.
[5] François-Xavv Xavier Le Dimet and Olivier Talagrand. Variational algorithms for analysis and assimilation of meteorological observations: theoretical aspects. Tellus A: Dynamic Meteorology and Oceanography, 38(2):97-110, 1986.
[6] Florence Rabier and Zhiquan Liu. Variational data assimilation: theory and overview. In Proc. ECMWF Seminar on Recent Developments in Data Assimilation for Atmosphere and Ocean, Reading, UK, pages 29 43, 2003.
[7] Alberto Carrassi, Mare Bocquet, Laurent Bertino, and Geir Evensen. Data assimilation in the geosciences: An overview methods, issues, and perspectives. Wiley Interdisciplinary Reviews: Climate Change, 9(5):e335, 2018.
[8] Alberto Carrassi, Mare Bocquet, Laurent Bertino, and Geir Evensen, Data assimilation in the geosciences: An overview of methods, issues, and perspectives. WIREs Climate Change, 9(5),
September 2018.
[9] G. Descombes, Auligné, F. Vandenberghe, D. M. Barker, and Barré, Generalized back- ground error covariance matrix model (GEN_ BE v2.0). Geoscientific Model Development, 8(3):669-66 696,, March 2015.
[10] Thomas M. Hamill, Jeffrey S. Whitaker, and Chris Snyder. Distance-Dependent Filtering of Background Error Covariance Estimates in an Ensemble Kalman Filter. Monthly Weather Review, 129(1 (0((1))2776-7000 November 2001.
[11] Thomas M. Hamill and Chris Snyder. A Hybrid Ensemble Kalman Filter-3D Variational Analysis Scheme. Monthly Weather Review, 128(8):2905-29) 2919, August 2000.
[12] Ross N Bannister. A review of operational methods of variational and ensemble-variation data assimilation. Quarterly Journal of the Royal Meteorological Society, 143(703):607-633, 2017.
[13] Anna Allen, Stratis Markou, Will Tebbutt, James Requeima, Wessel Bruinsma, Tom Andersson, Michael Herzog, Nicholas D Lane, Matthew Chantry, Scott Hosking, et al. End- to-end data-driven weather prediction. Nature, pages 1-3, 2025.
[14] Lei Chen, Xiaohui Zhong, Feng Zhang, Yuan Cheng, Yinghui Xu, Yuan Qi, and Hao Li. Fuxi: cascade machine learning forecasting system for 15-day global weather forecast. upj climate and atmospheric science, 6(1):190, 2023.
[15] Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium-range global weather forecasting with neural networks. Nature, 6195793) 2023.
[16] Remi Lam, Alvaro Sanzhez-Gonzzlez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Raton-Rosen, Weihua Hu, al. Learning skillful medium-range global weather forecasting. Science, 382(6677):1416-1422) 2023.
[17] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcastnet: global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11 2022.

---

[18] Sibo Cheng, César Quilodrán-Casee Said Ouala, Alban Farchi, Che Liu, Pierre Tandeo, Ronan Fablet. Didier Lucor, Bertrand Iooss, Julien Brajard, et Machine learning with data assimi- lation and uncertainty quantification for dynamical systems: a review. IEEE/CAA Journal of Automatica Sinica, (066):1311113: 2023.
[19] Caterina Buizza, César Quilodrán Casas, Philip Nadler. Julian Mack, Stefano Marrone, Zainab Titus, Clémence Le Cornec, Evelyn Heylen, Tolga Dur, Luis Baca Ruiz, et al. Data learning: In- tegrating data assimilation and machine learning. Journal of Computational Science, 58:101525, 2022.
[20] Alban Farchi, Patrick Laloyaux, Massimo Bonavita, and Mare Bocquet. Using machine learning to correct model error in data assimilation and forecast applications. Quarterly Journal of the Royal Meteorological Sociery, 147(739):3067-3000 2021.
[21] Julien Brajard, Alberto Carrassi, Mare Bocquet, and Laurent Bertino. Combining data assimila- tion and machine learning to infer unresolved scale parametrization. Philosophical Transactions of the Royal Society A, 379(2194)20200088 2021.
[22] Junchao Gong, Jingyi Xu. Ben Fei, Fenghua Ling, Wenlong Zhang, Kun Chen, Wanghan Xu, Weidong Yang, Xiaokang Yang. and Lei Bai. Dawp: framework for global observation forecasting via data assimilation and weather prediction in satellite observation space, 2025.
[23] Xiaoze Xu, Xiuyu Sun, Wei Han, Xiaohui Zhong, Lei Chen, Zhiqiu Gao, and Hao Li. Fuxi-da: A generalized deep learning data assimilation framework for assimilating satellite observations. npj Climate and Atmospheric Science, 8(1):156, 2025.
[24] Yanfei Xiang, Weixin Jin, Haiyu Dong, Mingliang Bai, Zuliang Fang, Pengcheng Zhao, Hongyu Sun, Kit Thambiratnam, Qi Zhang, and Xiaomeng Huang. Adaf: An artificial intelligence data assimilation framework for weather forecasting. arXiv preprint arXiv:241 1. 16807, 2024.
[25] Kun Chen, Lei Bai, Fenghua Ling, Peng Ye, Tao Chen, Jing-Jia Luo, Hao Chen, Yi Xiao, Kang Chen, Tao Han, et al. Towards an end-to-end artificial intelligence driven global weather forecasting system. arXiv preprint arXiv: 2...2211112 12462, 2023.
[26] Yi Xiao, Qilong Jia, Kun Chen, Lei Bai, and Wei Xue. Vae-var: Variational autoencoder- enhanced variational methods for data assimilation in meteorology. In The Thirteenth Interna- lional Conference on Learning Representations, 2025.
[27] Kun Chen, Peng Ye, Hao Chen, kang chen, Tao Han, Wanli Ouyang, Tao Chen, and LEI BAI FNP: Fourier neural processes for arbitrary-resolution data assimilation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[28] Julien Brajard, Alberto Carrassi, Mare Bocquet, and Laurent Bertino. Combining data assimila- tion and machine learning to emulate a dynamical model from sparse and noisy observations: A case study with the lorenz 96 model. Journal of computational science, 44:101171 2020.
[29] Corinna Maier, Niklas Hartung, Charlotte Kloft, Wilhelm Huisinga, and Jana de Wiljes. Re- inforcement learning and bayesian data assimilation for model-informed precision dosing in oncology. CPT: pharmacometrics & systems pharmacology, 10(3):241-2555 2021.
[30] Jinding Zhang, Kai Zhang, Zhongzheng Wang, Wensheng Zhou, Chen Liu, Liming Zhang, Xiaopeng Ma, Piyang Liu, Ziwei Bian, Jinzheng Kang, et al. A latent space method with maximum entropy deep reinforcement learning for data assimilation. Geoenergy Science and Engineering. 243:213275, 2024.
[31] Pouria Behnoudfar and Nan Chen. RI-daunce: Reinforcement learning-driven data assimilation with uncertainty- constrained ensembles, 2025.
[32] Mohamad Abed El Rahman Hammoud, Naila Raboudi, Edriss S Titi, Omar Knio, and Ibrahim Hoteit. Data assimilation in chaotic systems using deep reinforcement learning. Journal of Advances in Modeling Earth Systems, 1603))e2023300044777 2024.
12

---

[33] Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter Dominik Dueben, and Torsten Hoefler. DiftDA: a diffusion model for weather-scale data assimilation. Forty-first International Conference Machine Learning. 2024.
[34] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh Interna- tional Conference on Learning Representations, 2023.
[35] Boštjan Meline and Ziga Zaplotnik. 3d-var data assimilation using variational autoencoder. Quarterly Journal of the Royal Meteorological Society, 150(761);273-2229,, 2024.
[36] Maddalena Amendola, Rossella Areucci, Laetitia Mottet, César Quilodrán Casas, Shiwei Fan, Christopher Pain, Paul Linden, and Yi-Ke Guo. Data assimilation in the latent space a convolutional autoencoder. In International Conference Computational Science, pages 373-386. Springer, 2021.
[37] Mathis Peyron, Anthony Fillion, Selime Gürol, Victor Marchais, Serge Gratton, Pierre Boudier, and Gael Goret. Latent space data assimilation by using deep learning. Quarterly Journal of the Royal Meteorological Society, 147(740):35593377, 2021.
[38] Hang Fan, Ben Fei, Pierre Gentine, Yi Xiao, Kun Chen, Yubao Liu, Yongquan Qu, Fenghua Ling, and Lei Bai. Physically consistent global atmospheric data assimilation with machine learning in a latent space. arXiv preprint arXiv:2502.02884, 2025.
[39] Jascha vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper-
PMLR, 2015.
[40] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[41] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In The fourt Annual Conference on Neural Information Processing Systems, 2020.
[42] François Rozet and Gilles Louppe. Score-based data assimilation. In The Thirty-seventh Annual Conference on Neural Information Processing Systems, 2023.
[43] 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 kilometer scales. arXiv preprint 2441::2700.16944, 2024.
[44] Yongquan Qu, Juan Nathaniel, Shuolin Li, and Pierre Gentine. Deep generative data assimilation multimodal setting. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 449-459, 2024.
[45] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. The Thirty-sixth Annual Conference on Neural Information Processing Systems, 2022.
[46] Kevin Black, Michael Janner, Yilun Du, 1lya Kostrikov, and Sergey Levine. Training diffusion models with reinforeement learning. arXiv preprint 23X1: 2023.
[47] Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. The shirty-seventh Annual Conference on Neural Information Processing Systems, 2023.
[48] Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score: Better aligning text-to-image models with human preference. In Proceedings of the IFFF/CVF International Conference on Computer Vision, pages 2096-2105, 2023.
13

---

[49] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. The Thirty-seventh Annual Conference on Neural Information Processing Systems, 2023.
[50] Fei Deng, Qifei Wang, Wei Wei, Tingbo Hou, and Matthias Grundmann. Prdp: difference prediction for large-scale reward finetuning of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7423-7433, 2024.
Proximal reward
[51] Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to-image synthesis. arXiv preprint arXiv:2306. 09341, 2023.
[52] Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302. 12192, 2023.
[53] Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick-a-pic: An open dataset of user preferences for text-to-image generation. The Thirly-sevnnth Annual Conference on Neural Information Processing Systems, 2023.
[54] Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong. Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 8228-8233 2024.
[55] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. The Thirty-seventh Annual Conference on Neural Information Processing Systems, 2023.
[56] Huaisheng Zhu, Teng Xiao, and Vasant Honavar. DSPO: Direct score preference optimization for diffusion model alignment. In The Thirteenth International Conference Learning Representations, 2025.
[57] Kyungmin Lee, Xiahong Li, Qifei Wang, Junfeng He, Junjie Ke, Ming-Hsuan Yang. Irfan Essa, Jinwoo Shin, Feng Yang, and Yinxiao Li. Calibrated multi-preferent optimization for aligning diffusion models. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 18465-18475, 2025.
[58] Dipesh Tamboli, Souradip Chakraborty, Aditya Malusare, Biplab Banerjee, Amrit Singh Bedi, and Vaneet Aggarwal. Balanceddpo: Adaptive multi-metric alignment. arXiv preprint arXiv:2503. /2555, 2025.
[59] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 1rXiv:1312.6114, 2013.
arXiv preprint
[60] Carl Doersch. Tutorial on variational autoencoders. arXiv preprint
2016
[61] Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pages 10835-10866. PMLR, 2023.
[62] Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: the method of paired comparisons. Biometrika, 39(3/4):324-445 1952.
[63] Tao Han, Zhenghao Chen, Song Guo, Wanghan Xu, and Lei Bai. Cra5: Extreme compression of era5 for portable global climate and weather research via an efficient variational transformer. arXiv preprini urXiv:2405.03377 2024.
[64] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages
2023.
[65] Kang Chen, Tao Han, Junchao Gong, Lei Bai, Fenghua Ling, Jing-Jia Luo, Xi Chen, Leiming Ma, Tianning Zhang, Rui Su, et al. Fengwu: Pushing the skillful global medium-range weather forecast beyond 10 days lead. arXiv preprint arXiv:2304.02948. 2023.
14

---

[66] 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(1 I )=e2020MS00220,, 2020.
[67] Stephan Rasp, Stephan Hoyer, Alexander Merose, Ian Langmore, Peter Battaglia, Tyler Russell, Vivian Yang, Rob Carver, Shreya Agrawal, et al. Weatherbench 2: A benchmark for the next generation of data-driven global weather models. Journal of Advances in Modeling Earth Systems, 166))c2023MS00401,, 2024.
Alvaro
[68] Harold Jeffreys. On the dynamics of geostrophic winds. Quart. J. Roy. Meteor Soc, 52(217):85- 104, 1926.
[69] Zhanhao Liang Yuhui Yuan, Shuyang Gu, Bohan Chen, Tiankai Hang, Mingxi Cheng, Ji Li, and Liang Zheng. Aesthetic post-training diffusion models from generic preferences with step-by-step preference optimization. arXiv preprint arXiv:2406.043114 2024.
[70] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[71] Πly Loshchilov, Frank Hutter, et al. Fixing weight decay regularization adam. arXiv preprint 15500 5:5, 2017.
15

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We clarify our contributions in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have discussed the limitations of our work in the paper.
16

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes].
Justification: We provide the full proof for any theoretical results in the main content or the appendix.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide all experimental settings for reproductivity.
Guidelines:
The answer NA means that the paper does not include experiments.
17

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We train our model on ERA5, an open dataset. The detailed settings are fully clarified in our paper. We provide our code in the supplement material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.cc/publcc/eee.c http://dddde/ooeuuuiiieeeeee ssi ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
18

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g.. data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The complete experimental settings, including the dataset, hyperparameters, type of optimizer, etc, are provided in the Experiments section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The paper performs suitable statistical significance analysis to ensure the improvements are statistically significant.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report I-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The sufficient information on the computer resources is provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
19

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We perform this work following the NeurIPS Code of Ethics exactly.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work discusses RL-driven data assimilation, not involving societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This work does not involve such risks.
20

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The original papers or websites that produced the code dataset are properly cited and we use an open-source dataset for our experiments.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pacerswithcode thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors are encouraged to reach out the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code in the paper is well documented and the documentation is provided alongside the code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
21

---

Answer:
[No]
Justification:
The paper has no crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non stanndaddrrd components.
Please refer to our LLM policy (https://neut) /cC////reeeeccc//222//////////// for what should or should not be described.
22

---

A Dataset and Evaluations
Dataset Our experiments employ the European Centre for Medium-Range Weather Forecasts (ECMWF) atmospheric reanalysis (ERA5) 4 the primary dataset for model development and validation. The inputs consist of a comprehensive set of meteorological variables, including upper-air atmospheric variables (geopotential 2, temperature t, specific humidity q. zonal wind and meridional wind v)) with 13 pressure level (50hPa, 100hPa, 150hPa, 200hPa, 250hPa, 300hPa, 400hPa, 500hPa, 600hPa, 700hPa, 850hPa, 925hPa, and 1000hPa), four surface-level measurements (10-meter zonal component of wind (u10), 10-meter meridional component of wind (y10), 2-meter temperature (12m) and mean sea level pressure (msl)). This configuration yields total of 69 distinct meteorological variables. We follow ECMWF's standard nomenclature (e.g., q500 for specific humidity at 500 hPa). The temporal coverage spans four decades (1979-2018) to ensure robust model training.
Evaluations We implement two standard evaluations for data assimilation quality assessment. (1) Single-step assimilation accuracy: Direct comparison between assimilated states and ERAS ground truth at current timestep through error metries.
(2) Assimilation-basee forecasting: Initializing numerical weather prediction models with assimilated states to compute errors against ERA5 at forecast lead times. Three complementary metries quantifying performance are overall mean square error (MSE), mean absolute error (MAE), and the latitud--eeihhtdd root mean square error (WRMSE), which is a statistical metric widely used in geospatial analysis and atmospherie science. Given the estimate w, and the truth Th,w.e the WRMSE is defined as,
1/2
WRMSE(C)
HW Σn. H.
COS Ofh,a
(18)
hh=1 CDS 0r11,00
Here H and W represent the number of grid points in the longitudinal and latitudinal directions, respectively, and a w is the latitude of point (h, w).
The validation procedure involves 6-hourly assimilation cycles over the entirety of 2019, where each assimilation employs a 48-hour forecast as its background field. The three designated metrics are calculated for each cycle. Final performance scores the annual averages of these metries.
B The VAE training
We employ the established transformer-based architecture (VAFformer) [3 to compress high- dimensional atmospheric fields into a low-dimensional latent space. This implementation uses window attention [70) to capture atmospheric circulation features effectively. This model follows the "vit_ "vit_large" configuration for both encoder and decoder, with patch embedding of size (4,4) and stride (4,4), an embedding dimension of 1024, and 24 stacked transformer blocks equipped with window attention. The encoder-decoder architecture is optimized using Adam [71] with a batch size of 32, trained under a two-phase learning rate policy: a 10,000-step linear warm-up up to 2e-4 followed by cosine annealing. The training of VAE is conducted on the ERAS dataset (1979-2016) with the 2016-2018 period reserved for validation, and runs for a total of 80 epochs. Our trained VAE achieves 0.0067 overall MSE and 0.0486 overall MAE
C Computational Costs
The alignment process is computationally efficient. The main costs are divided as follows:
One-Time Pre-training: This is the most intensive step. The VAE training takes approx- imately 4 days, and the diffusion model pre-training takes about I day, both on 4 A100 GPUs.
DPO Alignment (Fine-tuning): This step is highly efficient. The core alignment process completes in about 30 minutes on 4 A 100 GPUs after 500 update steps. Reward Evaluation: Evaluating the "forecast skill" for single data sample involves running 48-hour forecast, which takes about 4 seconds on one A100 GPU.
23

---

Inference (Sampling): Generating a single analysis field takes approximately 30 seconds on one A100 GPU.
This analysis highlights that the novel alignment stage, which provides significant performance gains, adds a minimal computational overhead compared to the initial pre-training.
D Comparison with Traditional 3D-Var
To position work relative established methods, conducted direct comparison against a traditional 3D- Var baseline, a cornerstone of operational weather forecasting. The results are summarized in Table
Table 3: Comparison of Align-DA (Align-DPS-M) with the traditional 3D-Var baseline. While
3D-Var is marginally better on the metric directly optimizes (DA Accuracy), Align-DA shows significant advantages in the more critical downstream metrics.
Method DA Accuracy (MSE) Forecast Skill (MSE)
Geo-Score
Var Align-DPS-M
0.0581
0.0952
88.34
0.0593
0.0943
94.89
The results are highly revealing. While 3D>-Var achieves marginally better DA Accuracy (MSE), which is unsurprising its cost function is explicitly designed to minimize this value, it falls short in the metrics that are ultimately more critical for weather prediction. Our Align-DA framework demonstrates superior performance in both downstream Forecast Skill and Physical Consistency. The substantial improvement in the Geo-Score (from 88.34 to 94.89) indicates that our method produces analysis fields that are more physically realistic. This superior physical consistency and forecast skill underscore the core value of our contribution: Align-DA's ability to optimize for complex set of preferences simultancously, moving beyond the narrow optimization target of traditional methods.
E
Statistical Significance Testing
To confirm that the performance improvements reported in our paper are statistically significant, we conducted analysis to quantify the inherent stochasticity of our generative model. We ran the diffusion sampling process 10 times for single time step under the most uncertain setting-without any observations alignment (' ("ObsFree'), The relative standard deviation of the resulting metrics, which represents the model's 's baseline random variability, is presented Table 4
Table 4: Relative standard deviation (%) of key metrics for the unaligned 'ObsFree' baseline, calculated over 10 independent runs. This table quantifies the inherent randomness of the generative process.
WRMSE
MSE
MAE
u10
v500
Z500
1850
DA Accuracy (ObsFree)
Relative std. dev. (%)
0.045
0.058
0.115
0.079
0.201
0.083
Forecast Skill (ObsFree)
Relative std. dev. (%)
0.163
0.067
0.145
0.155
0.252
0.106
This analysis confirms that our reported improvements are statistically significant. For instance, the relative standard deviation for the primary DA accuracy metric (MSE) is merely 0.045%. This is orders magnitude smaller than the 2.65% 5.44% performance improvements that our Align-DA framework provides over the baseline (as shown in the main paper). Meanwhile, it is important to note that the performance change of the single-reward version (-P) yields statistically insignificant results, suggesting that geostrophic balance alone is insufficient for complex DA tasks.

---

Table 5: Assimilation accuracy gains with different observation densities. Percentage values quantify relative accuracy variants compared to the non-aligned baseline.
WRMSE
MSE
MAE
ulo
V500
2500
1850
Repaint
0.0623
0.1368
2,7027
103,3036
1
1% observation
Align-Repsint-W M
10098-41188
0,3440(-17560
1,4759 -1,900
2. 2,68.66
-0.766
96.89 96.8906-6..055)
1,0889--6
00.00%
DPS
0.0609
0.1337
1.4441
2,5972
93,4654
11159
Align-DPS-M
0,5992220555
0.1517(-13193)
1.1877-17768)
2.5600
58.4235--5.088
(055)
Repaint
0.0583
0.1327
1.4338
2,6097
100.0290
1,696
3% observation
Align: Repaint-M M
0.0660
14188-117788
2.5929)-006600)
93.88855
L0575
DPS
0,0546
0.1276
1.3698
2,4883
85,5742
1,0586
Aliga-DPS-M M
0,0533/ 2,6850
0.1256 1.90%)
1.3457 L.788)
2.47244
80,5474: 6.243
LO558:
0.20%
Repaint
0,0547
0.1290
1,3891
2.5248
96,9590
LD495
5% observation.
Align-Repaine M
0.5566-59195
1.3656
2,5081
91 41.8533-44 0590
(.477(-02200)
DPS
0,0534
0.1267
1.3597
2,4767
85,9150
1,0514
Aliga-DPS-M
0.0519 27785)
(1246-112200)
1.3337
L.996))
2.4560(-0.808)
N0.5787(-6.00
L0475
Repaint
0,0470
0.1204
1,2721
2,339
9K.17747
1,0057
10% observation
Aligu-Repinte
0,443: 8.77))
0,1186
1.229)
1,2518:
1.52%)
2.32500
0.72%)
84.7320X
6.42%)
1.0012
0.50%
0.0524
0.1
1.3520
2,4673
86,2364
1,468
Aliga-DPS-M M
0.510-22880
0,1239-17700
1.3252
20039)
2,4458
-0.000
80.0644
1,0403
0.02%
F Ablations studies
Observation densities. We further investigate the effect of varying observation densities (1%. 3%, 5%, and 10%) on the efficacy of our alignment strategy, detailed in Table [S] A key finding is that the proposed alignment mechanism significantly enhances assimilation accuracy across all tested densities for both guidance methods. Specifically, Align-Repaint-M and Align-DPS-M consistently outperform their non-aligned counterparts in all scenarios. This benefit is evident even at the highest 10% observation density, where alignment yields substantial error reductions of 6.42% for Aign-Repaint-M and 7.71% for Align-DPS-M on the z500 WRMSE metric. While the absolute errors decrease for all methods with more observations, the relative improvement percentage due to alignment shows nuanced behavior. For Align-Repaint-M, the relative gain slightly diminishes as observation density increases (c.g.. MSE improvement changes from 4.12% 1% observations to 3.67% at 10%). This suggests that while alignment is always beneficial, its proportional contribution might be more pronounced when initial information is sparser. Interestingly, Align-DPS-M does not follow a consistent trend. For some metrics like z500 WRMSE, the relative improvement even increases with higher density (e.g., from -5.70% at 1 1% to -7.11% at 10%). It may be ascribed the aligned conditional model is easier to be guided in DPS method. Nonetheless, the results affirm the significant and persistent positive impact of the alignment strategy across varying data availability.
Observation Error Robustness. 6|s ablation study examines Align-DPS-Mr assimilation ac- curacy gains over the non-aligned DPS baseline under varying observation error with 1% observation. Results consistently demonstrate the alignment effect's 's robustness: Align- DDgg--DPSMM MM outperforms the DPS baseline across all metrics, even as observation error increases. For example, Align-DPS-M shows 1.97% MSE reduction and a 4.24% z500 WRMSE reduction even at the highest error (std=0.05). However, while alignment remains beneficial, its relative improvement magnitude over the baseline decreases with increasing observation error. This is evident as Align-DPS-We percent- age gains, though still negative (indicating improvement), diminish in magnitude as observation error rises; for instance, MSE improvement drops from -2.65% ('Idea") to -1.9%% (std=0.05), and z500 WRMSE improvement from -5.70% ('Idea') to -4.24 4.24% (std=0.05). This suggests that while consistently advantageous, alignment's 's relatíve benefit is somewhat attenuated by higher observation noise.
Alternative Preference Optimization Algorithms
To validate the robustness of our preference alignment framework, we conducted experiments with alternative preference optimization algorithms: Identity Preference Optimisation (IPO) and Direct Preference Optimization with discrete sampling (DSPO). Table 2 presents the comprehensive results for both DA Accuracy and Forecast Skill across our three guidance strategies (ObsFree, Repaint, and DPS). These results confirm that our alignment concept is powerful and generalizable tool, both IPO and DSPO also yield significant performance gains over the unaligned baselines.
Additional Physies Constraints
25

---

Table Assimilation accuracy gains with different observation error under 1% observation. Percent- age values quantify relative accuracy variants compared to the non-aligned baseline.
WRMSE
MSE
MAB
u10
v500
z500
1850
DPS
0,0009
0,1337
1.4441
2.5972
93.4654
1.1159
Iden
AlignDDS-M
0.0593 2.60%)
0.1317(-1.119)
L-4177(-1788)
2,5800: 000000
88.4235 -.700)
1.1100x-0.658)
DPS
0,0607
0.1335
1.4408
2,5959
92.5883
1,1132
std=0.02
Align-DPS-M
0. 13166-1-1088
1.1176(1-6118)
2.5805 0.000%)
18.5310--155000
1.1077-0.415)
1.4388
std-0.05
Align-DPS-M
0.0590(-1995)
0. 0.1312(-1445
1,1446(-1.005
2.5740(-0.685)
87,6790044200
1.1099-0.188)
Table 7: DA Accuracy and Forecast Skill results for DPO (Align-M), IPO, and DSPO.
WRMSE
MSE
MAE
u10
v500
z500
t850
DA Accuracy
ObsFree
0.0643
0.1382
1.5275
2.7467
104.5531
1.0897
Align-ObsFree-h
0.0611
0.1348
1.4942
2.7202
96.8603
1.0866
IPO-ObsFree-M
0.0607
0.1355
1.4831
2.6944
101.7478
1.0807
DSPO-ObsFree-M
0.0604
0.1357
1.4780
2.6900
97.0645
1.0845
Repaint
0.0623
0.1368
1,5009
2.7027
103.3036
1.0911
Jlin-Repaint-NN
0.0598
0.1344
1.4799
2.6836
96.8906
1,0889
IPO-Repaint-M
0.0597
0.1350
1.4682
2.6660
100.7377
1.0868
DSPO- Repaineea
0.0590
0.1338
1.4528
2.6555
96.2449
1.0762
DPS
0.0609
0.1337
1,4441
2.5972
93.4654
1.1159
Align-DPS
0,0593
0.1317
1.4187
2.5800
88,4235
1.1100
IPO-DPS-M
0.0582
0.1315
1.4070
2.5539
89.4790
1.0977
DSPO-DPS-M
0,0585
0.1320
1.4102
2.5576
89.8927
1.1014
Forecast Skill
ObsFree
0.1146
0.1918
2.1182
4.1571
232.6243
1.4826
Align-OsFFree-
0.1080
0.1872
2.0965
4.1091
225.7758
1.4636
IPO-ObsFree-M
0.1093
0.1882
2.0966
4.1162
227.8923
1.4661
DSPO-ObsFree-M
0.1086
0.1879
2.0838
4.1020
226.1687
1.4637
Repaint
0.1123
0.1898
2.0899
4.1030
228.2876
1.4672
Begn--Reaiin-MM M
0.1068
0.1860
2.0713
4.0665
222.9336
1,4516
IPO-Repaint 100-Repant-MM
0.1083
0.1872
2.0712
4.0721
224.3243
1.4555
DSPO-Repaint-M
0.1080
0.1868
2.0663
4.0585
223.3857
1,4504
DPS
0.0969
0.1757
1.9230
3.7675
193.4167
1.3564
Align-DPS-M
0.0943
0.1734
1,9070
3.7287
189.6935
1.3428
IPO-DPS-M
0.0937
0.1732
1.9045
3.7275
189.5445
1.3412
DSPO-DPS-M 0.0932
0.1735
1.9062
3.7269
189.2176
1.3418
To demonstrate the flexibility of the Align-DA framework, we performed preliminary experiment incorporating a second physical constraint: hydrostatic balance. Hydrostatic balance governs the vertical structure of the atmosphere, linking pressure and temperature profiles via the hydrostatic equation:
04
RT,
(19)
0l
where P is pressure, R is the specific gas constant, and T is temperature. We quantify its imbalance using the metric Dnydro = 7 P + RT) Hydro-Score.
RT Similar to the Equation 17 one can also define
This extends the multi-reward alignment from three preferences to four. In the table below, "-M" denotes the original threc-preference alignment, while "-4" denotes the four- foureerefeeenc alignment
26

---

which includes the additional hydrostatic constraint. The results shown in Table validate our framework's ability easily leverage multiple, diverse physical constraints, often leading further performance gains across various metrics.
Table 8: DA Accuracy and Forecast Skill with an additional hydrostatic balance reward (*-4").
WRMSE
MSE
MAE
u10
v500
z500
1850
DA Accuracy
ObsFree
0,0643
0.1382
1.5275
2.7467
104.5531
1.0897
Align-ObsFree-
0.0611
0.1348
1.4942
2.7202
96.8603
1.0866
Align-Obs A
0.0571
0.1335
1.4675
2.6546
94.3716
1.0649
Repaint
0.0623
0.1368
1.5009
2.7027
103.3036
1.0911
Align- Repaint-M
0.0598
0.1344
1.4759
2.6836
96.8906
1.0889
0.0567
0.1333
1.4536
2.6327
95.7225
1.0756
DPS
0.0609
0.1337
1.4441
2.5972
93.4654
1.1159
Align-DPS-M
0.0593
0.1317
1,4187
2.5800
88.4235
1.1100
Align-DPS-4
0.0588
0.1315
1.4155
2.5672
87.3151
1.0996
Forecast Skill
ObsFree
0.1146
0.1918
2.1182
4.1571
232.6243
1.4826
Align-Obsi Align-ObsFree-M
0.10
0.1872
2.0965
4.1091
225.7758
1.4636
-
0.1028
0.1841
2.0456
4.0024
219.8289
1.4423
Repaint
0.1123
0.1898
2.0899
4.1030
228.2876
1.4672
0.1068
0.1860
2.0713
4.0665
222.9336
1.4516
Algn-Repaint-4
0.1017
0.1831
2.0345
3.9911
217.5085
1.4304
DPS
0.0969
0.1757
1,9230
3.7675
193.4167
1.3564
Align-DPS-M 0.0943
0.1734
1.9070
3.7287
189.6935
1.3428
Align-DPS 0.0938
0.1725
1.9059
3.7244
189.3845
1.3423
Realistic observation. We further evaluate the performance of AlignDA using real-world observa- tional data by conducting experiments with the Global Data Assimilation System (GDAS) prepbufr dataset, which incorporates multi-source observational data. To ensure data quality, we apply filtering process that eliminates observations exhibiting deviations exceeding 0.1 standard error when compared to the ERAS reference data. Table 9 shows that when using real-world GDAS observations, Align-DPS-M M still outperforms the standard DPS model. It achieves better accuracy across all measures, with notable improvements like a 3.89% decrease MSE and a 2.82% decrease in WRMSE for z500. These results demonstrate that the alignment strategy is effective and beneficial for practical, real-world applications.
G Limitation and Future work
While our AAign-DA demonstrates promising results, several aspects should be further studied. The current physical constraints, primarily geostrophic balance, may not fully capture the complexities of operational DA, suggesting a need for more comprehensive reward designs. Additionally, the framework relies on offline reinforcement learning, which is inherently sensitive to the quality and diversity of the pre-collected preference dataset. Incorporating online RL strategies could enable more adaptive policy updates and mitigate limitations imposed by static offline data. Moreover, the RI.-based alignment framework introduced in this paper has potential beyond improving initial conditions; it can also generalized to other stages of the weather forecasting workflow, such as forecast processin.....
27

---

Table Assimilation accuracy gains with GDAS observations. Percentage values quantify relative accuracy variants compared to the non-aligned baseline.
WRMSE
MSR
MAE
u10
v500
2500
t850
DPS Align-DPS-M
0.0634
0.1365
1.3342 1.3200 1.08%)
2,5642
99.2263 96,4988 2.82%)
1.1343
0.0610 -3.89%
0.1340X L.87))
2.5498 0.59%)
1.1342(-0.01%))
ERAS
Background
Observation
-40
-20
-40
-20
ObaFreso-P reduction
reduction
reduction
AlignObs error reduction
Align-Repaint-t Maoppann--MMee error reduction
Align-DPS-M ertor reduction
Figure 5: Visulaization of u700 2019-01-26-1888 UTC.
H
More visualization
Here we provide more visualization results. In all figures the appendix, 'ObsFree' means without observation integration baseline. 'Repaint' and 'DPS' represent repaint and diffusion posterior sampling guidance methods. Note that '-p' indicates single physical reward alignment, while '-M' signifies multi-reward alignment. First row: Reference fields showing ERA5 reanalysis (ground truth denoted as GT), background field, and observational data. Second row: Error reduction through single physical reward alignment, quantified as | &GT| Wa reduction through multi-reward alignment, quantified as | Øcrl
Align.P
ZGTl Third row: Error IGT|.
Align-M
28

---

ERAS
Background
Observation
96000
98000
100000
102000
104000
90000
98000
100000
102000
104000
56200
58000
100000
102000
104000
Align Obafree-P errrr reduction
Alinn-Repaint-P error reduction
Nigh-Ur error reduction
-00
-200
100
100
200
300
-300
-200
100
100
200
300
990
-00
-100
100
200
200
Align-
Align- RapaiRRa
Align-DPS-M
100
200
300
100
200
-100
Figure 6: Visulaization of msl at 2019-03-03-00-00 UTC.
ERAS
Background
Observation
300
Olign-OsstreeP error reduction
Align Repai nt enror reduction
Align DPS enror reduction
Align. Obstree error reduction
Align Repaint M enror reduction
Align DPS eIror reduction
Figure 7: Visulaization of t850 2019-07-05-120000 UTC.
29

---

ERAS
Background
Observation
0.002 0.004
0008
6,008
0010
0.012
0.004 0.006 0.005 0.000
0.002
0.004
0.006
0,010
0.12
Alin-Obsfrre--
Align-Repeint-P reduction
Align- DPS-P eiror reduction
-0.003
-0.002
-0.00
0.000
2.001
0007
5.003
-0.002
-0.003
Align Mign-ObsFrre M emor reduction
Align Repaint M eror reduction
Align-DPS M enror reduction
-3.002 -0.001
0.000
-0.002
Figure 8: Visulaization of 9700 at UTC.
30


Reviews:
[{"review_id": "0xNsRg3GvH", "reviewer": "Reviewer_46nx", "summary": "This paper presents a reinforcement learning-based framework for data assimilation (DA), named Align-DA, to estimate the system states of atmospheric applications given their prior states and observations. In this framework, diffusion models are used for DA, and goals of assimilation and domain-specific knowledge can be formulated as reward functions that guide the learning process of the models. As an example, the authors explored regularizing the optimization with three metrics, assimilation accuracy, forecast skill, and physical adherence. Experimental results demonstrate that models trained using the proposed framework with multiple rewards outperform those trained without or with a single reward.", "questions": "Is there any existing reinforcement learning-based DA method? If yes, what are the key aspects that differentiate this work from them?\n\nIs Align-DA specialized in atmospheric applications, and, if yes, how? Do the authors believe that it can generalize to other problem domains, and, if yes, what are the key concerns that make them refrain from experimenting on other tasks?\n\nAnswering these questions can help me better understand the significance of the contribution of this paper, which could increase or decrease my evaluation score.", "limitations": "Yes.", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nThe necessity of the proposed framework is well motivated. The proposed framework using reinforcement learning with reward functions naturally addresses the problem. The method design and the derivation are introduced with reasonable details. The presentation clarity is good. While the method may be inspired by prior work, applying reinforcement learning on DA appears to be novel and original – this can be further clarified by the authors.\n\nWeaknesses:\n\nThe major weaknesses lie in the experimental verification of the proposed framework, which can impact the significance of the contribution of this paper. Specifically:\n\nThe theoretical framework appears general and not limited to the atmospheric data, while the experimental verification only presents one specific case of weather forecasting with one set of data. As a result, although the framework is interesting and offers flexibility, its effectiveness (in atmospheric applications overall) and generalizability (to other tasks in atmospheric science and to other problem domains) are yet to be fully verified. The significance of the contribution is therefore limited.\n\nThe current experiments mainly focus on how the proposed scheme improves certain methods on given metrics. It should be no surprise for the performance gain, as the metrics align with the designed reward function while the baseline methods are trained without the knowledge of the evaluation metrics. It could be more convincing if other frameworks or training schemes are compared against ensuring fairness. If no existing approach can achieve this goal and this work is the first to accomplish this goal, this innovation and significance should be highlighted.\n\nThis paper also presents a lack of comparisons to state-of-the-art methods for the forecast or those methods that can be applied to the task of interest with minor, reasonable modifications.\n\nIn the context that Align-DA is designed for atmospheric applications, the problem of interest regarding atmospheric analysis and weather forecasting should be explained in greater detail. If not, the positioning of this work could be clarified.", "was_revised": true, "final_justification": "While the authors attempt to differentiate their work from my concerns, I still do not see they are positioning this work such that it offers a clear, substantial contribution. The clarified core research problem appears to be narrower than necessary given the other materials presented. Comparing to an existing DA method, the proposed Align-DA has not demonstrated clear superiority. Therefore, I maintain my original recommendation.", "rebuttal": "We sincerely thank the reviewer for their thorough and constructive feedback. We will address each point in detail below.\n\n___\n> **Response to Question 1:** Difference from the existing RL-based DA \n\nThank you for this important question. Align-DA is fundamentally different from other reinforcement learning for data assimilation (RL-DA),  in its problem formulation and optimization paradigm.\n\nExisting RL-DA methods typically frame DA as a sequential decision-making problem within a Markov Decision Process (MDP)[1,2]. In this view, an agent learns a policy through online interaction to maximize a cumulative reward, operating on a small-scale system. It is still challenging to scale this sequential, interactive paradigm to high-dimensional systems like global weather [3,4,5].\n\nIn contrast, Align-DA re-formulates DA as a preference-aligned generative modeling problem. Our framework operates offline, using Direct Preference Optimization (DPO) to directly align a score-based model's entire learned distribution with complex preferences, rather than learning a sequence of actions through trial-and-error.\n\nTo the best of our knowledge, Align-DA is the first work to introduce this preference alignment paradigm into the domain of large-scale DA. To make this crucial claim clearer in the paper, we will add the following sentence and citations to the 'Introduction' section: \"**While some studies have applied reinforcement learning to data assimilation, these approaches typically frame DA as a sequential decision-making problem [1,2,3], which is challenging to apply to large-scale weather systems [4,5].**\"\n\n___\n> **Response for Question 2:** Specialization, Generalization, and experiments for other domains\n\nThank you for the excellent question. While the core alignment mechanism of Align-DA is entirely general, our implementation is specialized for atmospheric science through our specific choice of reward functions—analysis accuracy, forecast skill, and physical adherence. The framework can be readily generalized to any domain where a \"good\" solution can be defined by a set of preference criteria; the main effort is simply designing new, domain-appropriate rewards. We agree that extending our framework to other domains would be highly valuable. However, it requires substantial time and resources for data preparation, model training, and acquiring new domain knowledge. We have limited time to complete it, but we have identified it as a clear and important direction for our future work. \n\n---\n>**Response to Weakness 1:** Generalizability of our framework\n\nWe sincerely thank the reviewer for this insightful comment. In this paper, we choose to validate the preference alignment mechanism in the atmospheric domain for several reasons. First, we are familiar with weather prediction tasks and the underlying physical constraints of the atmosphere, which are important for designing meaningful preferences. Second, atmospheric DA poses greater challenges than many other dynamical systems, owing to its stronger temporal variability, higher‑dimensional state space, and the diversity of its prognostic variables.\n \nWe note, however, that the proposed Align‑DA framework is not limited to atmospheric applications. In practice, operational DA frameworks are highly similar across the atmosphere, ocean, land, and other fluid systems, which shows that DA methods are naturally generalizable. \n\nWe acknowledge that extending Align-DA to other domains would be highly valuable. However, this requires preparing reanalysis datasets, forecasting models, VAE and diffusion models, and domain knowledge to design meaningful preferences, which would take substantial time.\n\nTo better address your concern, we have added a sentence \"... **Our framework demonstrated here are highly transferable, creating a clear pathway for future work in other domains like oceanography and seismology. Adapting the framework primarily requires designing new, domain-specific reward functions.**\" to the 'Conclusion' section. \n\n---\n>**Response to Weakness 2:** \"Teaching to the test\" concern and comparison with other frameworks\n\nTo address your concern about \"teaching to the test,\" we performed a new experiment that demonstrates generalization beyond explicit rewards. In this experiment, the model was aligned solely using preferences for analysis accuracy and forecast skill, without any explicit physics‑based rewards. The resulting forecasts still exhibited substantially improved adherence to physical laws, indicating that Align‑DA learns a fundamentally realistic system state rather than merely overfitting to the specified rewards.\n\n| Geo-Score | Non-Align | Aligned |\n|:---|:---|:---|\n|**ObsFree**|92.25|93.57|\n|**Repaint**|92.16|92.46|\n|**DPS**|93.98|94.89|\n\nOur main contribution, which we will highlight more clearly as you suggest, is that our framework is the first to make preference optimization possible for DA. To further validate the robustness of this new framework, we would like to test alternative preference optimization algorithms (IPO[6] and DSPO[7]). Both methods yielded significant gains, confirming that our alignment concept is a powerful and generalizable tool. We present the DA accuracy and Forecast skill results below. **Bold text** indicates the best. \n\n| DA accuracy | MSE | MAE | u10-WRMSE | v500-WRMSE | z500-WRMSE | t850-WRMSE |\n|:---|:---|:---|:---|:---|:---|:---|\n|**ObsFree**|0.0643|0.1382|1.5275|2.7467|104.5531|1.0897|\n|**Align-ObsFree-M**|0.0611|**0.1348**|1.4942|2.7202|**96.8603**|1.0866|\n|**IPO-ObsFree-M**|0.0607|0.1355|1.4831|2.6944|101.7478|**1.0807**|\n|**DSPO-ObsFree-M**|**0.0604**|0.1357|**1.4780**|**2.6900**|97.0645|1.0845|\n||||||||\n|**Repaint**|0.0623|0.1368|1.5009|2.7027|103.3036|1.0911|\n|**Align-Repaint-M**|0.0598|0.1344|1.4759|2.6836|96.8906|1.0889|\n|**IPO-Repaint-M**|0.0597|0.1350|1.4682|2.6660|100.7377|1.0868|\n|**DSPO-Repaint-M**|**0.0590**|**0.1338**|**1.4528**|**2.6555**|**96.2449**|**1.0762**|\n||||||||\n|**DPS**|0.0609|0.1337|1.4441|2.5972|93.4654|1.1159|\n|**Align-DPS-M**|0.0593|0.1317|1.4187|2.5800|**88.4235**|1.1100|\n|**IPO-DPS-M**|**0.0582**|**0.1315**|**1.4070**|**2.5539**|89.4790|**1.0977**|\n|**DSPO-DPS-M**|0.0585|0.1320|1.4102|2.5576|89.8927|1.1014|\n\n\n|Forecast Skill|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|**ObsFree**|0.1146|0.1918|2.1182|4.1571|232.6243|1.4826|\n|**Align-ObsFree-M**|**0.1080**|**0.1872**|2.0965|4.1091|**225.7758**|**1.4636**|\n|**IPO-ObsFree-M**|0.1093|0.1882|2.0966|4.1162|227.8923|1.4661|\n|**DSPO-ObsFree-M**|0.1086|0.1879|**2.0838**|**4.1020**|226.1687|1.4637|\n||||||||\n|**Repaint**|0.1123|0.1898|2.0899|4.1030|228.2876|1.4672|\n|**Align-Repaint-M**|**0.1068**|**0.1860**|2.0713|4.0665|**222.9336**|1.4516|\n|**IPO-Repaint-M**|0.1083|0.1872|2.0712|4.0721|224.3243|1.4555|\n|**DSPO-Repaint-M**|0.1080|0.1868|**2.0663**|**4.0585**|223.3857|**1.4504**|\n||||||||\n|**DPS**|0.0969|0.1757|1.9230|3.7675|193.4167|1.3564|\n|**Align-DPS-M**|0.0943|0.1734|1.9070|3.7287|189.6935|1.3428|\n|**IPO-DPS-M**|0.0937|**0.1732**|**1.9045**|3.7275|189.5445|**1.3412**|\n|**DSPO-DPS-M**|**0.0932**|0.1735|1.9062|**3.7269**|**189.2176**|1.3418|\n\nWe will revise the manuscript to more explicitly state that this framework is a key contribution and will include these new experimental results to strengthen our claims. \n\n___\n> **Response to Weakness 3:** Comparisons with state-of-the-art methods for the forecast or those methods that can be applied to this task.\n\nWe thank the reviewer for this helpful comment. Our contribution is to utilize Preference Optimization to enhance Score-based DA rather than establishing a new forecast model to compare existing state-of-the-art methods.\n\nConsequently, our experimental design focuses on controlled ablation studies of score‑based DA methods to identify the effect of alignment.  In our experiments,  the frozen forecasting model (FengWu) serves two distinct roles: First, it generates the 48-hour background field that acts as the common input for all DA methods. Second, we use FengWu to assess forecast skill by running 48‑hour forecasts initialized from each analysis field.\n\nWe have tried to use other forecast models, such as GraphCast or Pangu-Weather, but we meet practical constraints like the lack of access to them, at the required 1.4-degree resolution. Nevertheless, we firmly believe our present, focused comparison is enough to state our contributions. We will clarify this positioning more explicitly in the experimental setup section of the revised manuscript.\n\n___\n> **Response to Weakness 4:** Explain the DA background in detail\n\nWe thank the reviewer for this constructive suggestion. In the revised manuscript, we will expand the introduction to include a more detailed explanation of DA in the context of numerical weather prediction, and highlight that traditional approaches often require extensive trial-and-error tuning to obtain satisfactory results. \n\nHere is what we revised in Line 30:\n\"**In atmospheric science, data assimilation (DA) provides accurate initial conditions that are essential for reliable weather forecasting. It is also used to generate high-fidelity reanalysis datasets, which serve as a fundamental resource for climate research.**\n\n___\nThank you for your helpful comments again. We hope this response can solve your concern. \n___\n**Reference:**\n\n[1] https://doi.org/10.1029/2023MS004178\n\n[2] https://doi.org/10.1002/psp4.12588\n\n[3] https://doi.org/10.1016/j.jocs.2020.101171\n\n[4] https://doi.org/10.1016/j.geoen.2024.213275\n\n[5] https://arxiv.org/abs/2505.05452\n\n[6] https://arxiv.org/abs/2310.12036\n\n[7] https://arxiv.org/abs/2504.15176 ICLR2025"}, {"review_id": "7O6pbNQ6WJ", "reviewer": "Reviewer_oCk2", "summary": "This paper proposes a new framework for atmospheric data assimilation (DA) based on preference alignment. Instead of traditional tuning of background priors, the proposed method treats DA as a generative modeling problem and introduces reward-guided alignment through DPO. The authors train a background-conditioned diffusion model in the latent space and fine-tune it using three domain-relevant rewards: assimilation accuracy, forecast skill, and physical consistency. The experiments show that the model consistently improves both analysis and forecast metrics.", "questions": "Could you provide a bit more explanation about the relationship between your method, Fengwu, and the baselines? Does your method's improvement come from the use of fengwu? Can the baselines also leverage fengwu?", "limitations": "Yes.", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n\nThe proposed method replaces manual tuning with preference alignment using a latent-space diffusion model guided by multiple domain-specific rewards. It is the first (to my knowledge) to adapt Direct Preference Optimization (DPO) to atmospheric data assimilation, enabling flexible incorporation of assimilation accuracy, forecast skill, and physical consistency. Experiments show consistent improvements across different guidance methods.\n\n\nI don't see any weakness, but I'm also not an expert on Atmospheric data.", "was_revised": true, "final_justification": "The authors addressed my concerns well.", "rebuttal": "We sincerely thank the reviewer for their positive and encouraging feedback. We are delighted that they recognized the novelty of adapting Direct Preference Optimization (DPO) to atmospheric data assimilation and the significance of our framework.\n\n**Clarifications about Main Contribution**\n\nTo further clarify the motivation for our work, we would like to briefly outline the fundamental challenge we address at first. Since the weather forecasting models are inherently imperfect, initial errors tend to grow rapidly during model integration. To correct these errors and improve forecast accuracy, observations should be continuously assimilated through data assimilation (DA), a process that combines observations with prior forecasts (background states) that already contain accumulated model error. In this study, we use FengWu as the forecasting model, providing background fields and forecasts using analysis fields derived from various score-based DA schemes. Our proposed Align-DA framework adaptively tames the score-based DA methods to meet the accuracy requirements of both analysis and forecast, as dictated by the assimilation objective.\n\n___\n> **Response to Question 1:** Relationship between Align-DA, Fengwu, and Baselines; Source of performance gain; Baselines'  access to FengWu\n\n**The Relationship Between Align-DA, FengWu, and Baselines**\n\nThe core contribution of our work is to introduce the preference alignment for score-based DA. Therefore, the baselines are score-based DA methods without alignment, while the Align-DA is the aligned counterpart. Their shared goal is to refine a coarse background field using sparse observations to produce an optimal initial analysis field. They are the competing methods in our evaluation.\n\nFengWu is a state-of-the-art weather forecasting model. In our work, it is not a competitor but serves two critical roles:\n\n- **Background Generator:** FengWu generates the 48-hour background forecast that serves as the common starting point for all DA methods, ensuring a fair comparison.\n- **Evaluator and Reward Source:** After each DA method produces an analysis field, we use FengWu to run a 48-hour forecast from that field. The accuracy of this forecast serves as a metric for evaluation. For Align-DA specifically, this \"downstream forecast skill\" is also used as a key reward signal during its optimization.\n\n**Does Align-DA's Improvement Come From Using FengWu?**\n\nThank you for this insightful question. The improvement comes not from using FengWu itself, but from our novel preference optimization framework, which allows Align-DA to learn from FengWu's feedback in a way baseline methods cannot. \n\nFengWu acts as a frozen module that provides a common background field and a \"forecast skill\" score for any given analysis. One can replace the Fengwu with any forecast model if available. \n\nBaseline methods can be evaluated by FengWu, but they have no mechanism to internalize this \"forecast skill\" feedback to improve themselves. In contrast, Align-DA, powered by its preference optimization framework, is uniquely designed to do exactly that. It treats the forecast skill score as a reward signal and directly optimizes its own generative process to produce analyses that lead to better future forecasts.\n\n**Can the Baselines Also Leverage FengWu?**\n\nThe Baselines are unable to leverage Fengwu as a reward model in the same way as our Align-DA, highlighting the contribution of our work.\n\nUnlike our approach, baseline methods—whether traditional or standard deep learning approaches—lack an inherent mechanism to integrate complex external reward signals, such as \"48-hour forecast skill,\" into their own optimization process. While FengWu can evaluate these methods, it cannot actively guide their learning.\n\nIn contrast, our Align-DA framework introduces preference alignment that directly optimizes the data assimilation model, especially improving forecast accuracy evaluated by FengWu. This represents a paradigm shift compared to existing methods\n\n___\nWe sincerely appreciate your valuable feedback, which has significantly enhanced the clarity and completeness of our paper."}, {"review_id": "2SdXKGd2L0", "reviewer": "Reviewer_eG1A", "summary": "This paper aims to tackle Data Assimilation which uses realistic observation to update computationally predicted weather model in atmospheric data. Given the current score-based DPO approach on this task, the author proposed a new soft constraint alignment techniques, which uses 3 different metrics (including a physics constraint) as a soft constraint to guide the DPO. The author conducted benchmarking on simulation data and showed comparisons to baseline model.", "questions": "- I'd like to see some baseline comparison to traditional data assimilation methods that is actually used in today's wether forecasting. If that is difficult/not possible, please explain why.\n- I'd like to see the effects of including more physics constraint. In the paper, the author claims that the proposed methods does not require manual fine-tuning, this should means that including more physics constraint is straightforward. If this is not the case, could the author explain why?\n- How well does the proposed forecasting model deal with extreme weather condition?", "limitations": "The author has some discussions about limitation in the paper but there's no dedicated limitation section. There's no discussion on societal impact. I think the possible limitation could be how the model deal with extreme weather (which, if failed, would have significantly high cost on society).", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength:\n- While I am not working directly in the field of atmospheric science, I can see the importance of data assimilation and its relevance as a AI/ML task; this paper has a good motivation\n- It is a novel application effort to use multi-score DPO method on climate science data\n- This paper proposed a suites of rewards relevant to atmospheric data assimilation and aims to align multiple rewards (including physics rewards)\n- The author benchmarked their model on an existing simulation and real data, and the method seems efficient at least on this dataset.\n- The paper is very well written and structure\n\nWeakness:\n- The novelty on model structure is rather limited\n- The author only introduced one physics constraint\n- There is no benchmark against traditional data assimilation methods that is actually used in climate station\n\nIn general, I think this is a well-written paper with good application novelty but limited model novelty.", "was_revised": true, "final_justification": "I think this work makes some novel contributions to the field, and the reviewer's rebuttal addressed some of my concerns. However, I am not completely satisfied with the fact that all training and validation data are AI-generated without comparing to traditional methods. Therefore I would like to retain my original score at 4: Borderline Accept.", "rebuttal": "We sincerely thank the reviewer for their insightful and well-balanced feedback. We are especially pleased that the reviewer recognized the novelty of applying multi-score direct preference optimization (DPO) to climate science, the relevance of our chosen rewards, and the clarity of the paper's writing.  We appreciate the opportunity to elaborate on points raised in the report below.\n\n---\n>**Response to Question 1 & Weakness 3:**  Reasons for Traditional DA Methods\n\nWe completely agree that a direct comparison with operational data assimilation (DA) systems used in today's weather forecasting is important. We sincerely appreciate you raising this point.\n\nModern operational DA systems are highly complex and sophisticated systems, which are mainly written in Fortran and tailored to specific forecast models. Developing a dedicated operational DA system for our forecast models from scratch would be challenging and is infeasible due to the time limit during the rebuttal.\n\nIn addition, this study aims to introduce preference alignment into score‑based DA, which departs from the paradigm of traditional DA. Despite several prior studies, score‑based DA has never been directly compared against traditional approaches. Therefore, we primarily focus on assessing whether preference alignment can enhance score‑based DA.\n\nWe acknowledge that the lack of a direct comparison to an operational system is important for the application. We are grateful for your valuable feedback. We believe our study lays the essential groundwork for this next step, and we plan to explore such comparisons in the future.\n\n___\n>**Response to Question 2 & Weakness 2:** More Physics Constraint\n\nWe sincerely appreciate this insightful comment, and your intuition is indeed correct. A key strength of our Align-DA framework is its flexibility, which makes adding new physics constraints a straightforward process.\n\nInspired by your suggestion, we performed a preliminary experiment incorporating a second physical constraint: hydrostatic balance, across multiple layers (50, 100, 250, 500, 700, 850, 1000 hpa). To complete this within the review period, we used the following setup: we halved both the preference dataset size to 2,000 samples and the finetuning duration to 250 steps, and we retained all win-loss pairs for faster data construction. We show the DA accuracy and Forecast skill below, where  \"**-4**\" denotes the four preferences alignment (one more hydrostatic constraint over three preferences in the manuscript). \n\n|DA accuracy|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|**ObsFree**|0.0643|0.1382|1.5275|2.7467|104.5531|1.0897|\n|**Align-ObsFree-M**|0.0611|0.1348|1.4942|2.7202|96.8603|1.0866|\n|**Align-ObsFree-4**|**0.0571**|**0.1335**|**1.4675**|**2.6546**|**94.3716**|**1.0649**|\n||||||||\n|**Repaint**|0.0623|0.1368|1.5009|2.7027|103.3036|1.0911|\n|**Align-Repaint-M**|0.0598|0.1344|1.4759|2.6836|96.8906|1.0889|\n|**Align-Repaint-4**|**0.0567**|**0.1333**|**1.4536**|**2.6327**|**95.7225**|**1.0756**|\n||||||||\n|**DPS**|0.0609|0.1337|1.4441|2.5972|93.4654|1.1159|\n|**Align-DPS-M**|0.0593|0.1317|1.4187|2.5800|88.4235|1.1100|\n|**Align-DPS-4**|**0.0588**|**0.1315**|**1.4155**|**2.5672**|**87.3151**|**1.0996**|\n\n|Forecast Skill|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|**ObsFree**|0.1146|0.1918|2.1182|4.1571|232.6243|1.4826|\n|**Align-ObsFree-M**|0.1080|0.1872|2.0965|4.1091|225.7758|1.4636|\n|**Align-ObsFree-4**|**0.1028**|**0.1841**|**2.0456**|**4.0024**|**219.8289**|**1.4423**|\n||||||||\n|**Repaint**|0.1123|0.1898|2.0899|4.1030|228.2876|1.4672|\n|**Align-Repaint-M**|0.1068|0.1860|2.0713|4.0665|222.9336|1.4516|\n|**Align-Repaint-4**|**0.1017**|**0.1831**|**2.0345**|**3.9911**|**217.5085**|**1.4304**|\n||||||||\n|**DPS**|0.0969|0.1757|1.9230|3.7675|193.4167|1.3564|\n|**Align-DPS-M**|0.0943|0.1734|1.9070|3.7287|189.6935|1.3428|\n|**Align-DPS-4**|**0.0938**|**0.1725**|**1.9059**|**3.7244**|**189.3845**|**1.3423**|\n\nThe results shown above validate our framework's ability to leverage multiple physical constraints. To improve the quality of our work, we will add a detailed description of this new experiment and its positive results to the Appendix in the revised manuscript.\n\n\n___\n>**Response to Question 3:** Extreme Weather Condition\n\nWe sincerely appreciate your thoughtful comment. This study performs score-based DA in the latent space to mitigate the high dimensionality of atmospheric states. As noted in prior work [1], extreme weather conditions cannot be effectively represented in the latent space. This limitation aligns with the low-dimensional manifold assumption, under which outliers are inherently difficult to capture. Consequently, the performance of our approach under extreme weather conditions cannot be fully assessed.\n\nNevertheless, if score-based DA were conducted in the model space, we believe our framework could enhance the assimilation of extreme events, with performance under such conditions serving as an additional evaluation metric.\n\nWe agree that the capability for handling extreme weather conditions is important. We will carefully discuss this direction in the revised version.\n\n---\n>**Response to Weakness 1:** The Novelty of Model Structure\n\nWe thank the reviewer for raising this point. The reviewer is correct that our work does not introduce a new model structure. Otherwise, we aim at introducing a paradigm shift for the DA task.  We would like to clarify our motivation and solution again. \n\nIn numerical weather prediction, DA is a critical step, as the quality of the initial analysis directly impacts subsequent forecast accuracy. A key challenge is that this analysis must satisfy multiple complex criteria simultaneously, such as assimilation accuracy, physical consistency, and the ability to produce a skillful downstream forecast. The traditional DA relies on extensive and labor-intensive parameter tuning to satisfy these criteria, which is often computationally inefficient and labor-intensive.\n\nOur Align-DA framework tackles this by reformulating DA as a preference optimization problem. Using DPO, we directly and automatically align a generative DA model with multiple preferences—analysis accuracy, forecast skill, and physical consistency. Moreover, our framework is highly flexible for any well-defined rewards incorporation and any score-based DA without dependence on model structure.  We believe this is a critical contribution to the field.\n\n---\nWe appreciate your valuable feedback again, which has enriched our work and improved our manuscript.\n\n**Reference:**\n\n[1] https://doi.org/10.1175/MWR-D-24-0058.1"}, {"review_id": "V8w024dt0R", "reviewer": "Reviewer_r9ir", "summary": "This paper proposes Align-DA, a data assimilation framework that leverages score-based diffusion models to generate posterior estimates of atmospheric states conditioned on forecasts and sparse observations. It introduces a reward-guided alignment mechanism inspired by Direct Preference Optimization (DPO), enabling the diffusion model to prefer samples that improve assimilation accuracy, forecast skill, and physical consistency. The approach combines a pretrained AI forecast model (Fengwu) for generating background estimates, while a latent-space diffusion model learns to adjust plausible estimates via guided sampling and based on the current observations.", "questions": "- The method uses a latent space diffusion model. How well does this approach support point-wise sparse observations? Does conditioning in latent space require observations over entire spatial patches? How does the observation rate affect the performance?\n\n- While the method leverages a pretrained large forecast model, the finetuning of the diffusion models raises the question of whether training a diffusion model from scratch would be better. For example, could you provide a comparison with approaches such as DiffusionPDE [1]?\n\n- How computationally expensive is the alignment process, especially the evaluation of forecast skill?\n\n- Have you tried inpainting methods instead of guidance methods? That is, condition the diffusion model directly on the background and the sparse observations.\n\n[1] Huang et. al. DiffusionPDE: Generative PDE-Solving Under Partial Observation", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strength**\n\n- The paper extends large forecast models to address sparse observation problems by leveraging diffusion generative models and guided sampling.\n\n- The paper thoroughly explores multiple objectives (assimilation accuracy, forecast skill, and physical constraints), effectively combining data-driven priors with physical knowledge as inductive biases.\n\n- The paper provides detailed derivations of the mathematical formulation and losses, including the score-based models, guidance methods, and the DPO loss.\n\n- The experiments provide clear ablations comparing single-reward and multi-reward alignment across different observation guidance strategies. The quantitative analysis is comprehensive, covering multiple error metrics to robustly evaluate performance.\n\n**Weakness**\n\n- The method relies on offline reward alignment rather than integrating forecast skill and physical adherence directly into the primary training objective. There is limited exploration or comparison with fully end-to-end training approaches that could jointly optimize these objectives, particularly in the context of sparse observation scenarios.\n\n- Using both a separate deterministic forecast model and a probabilistic diffusion model introduces duplicated capacity for modeling plausible states, which could potentially be unified in a single generative framework.\n\n- Moreover, relying on a deterministic forecast model while employing a separate diffusion model to capture uncertainty appears somewhat ad hoc and may be a limited remedy for the inherently deterministic nature of the base simulation.", "was_revised": true, "final_justification": "The authors answered all my questions and provided broader contexts of the method to help motivate the work. Therefore, I would like to keep my positive feedback and recommend acceptance.", "rebuttal": "We would like to thank the reviewer for their very detailed and technically insightful feedback. We are especially pleased that the reviewer found our mathematical formulation robust and our experimental ablations comprehensive. We appreciate this opportunity to elaborate on the rationale behind these critical design choices and will address each point in detail below.\n\n___\n>**Response to Question 1:** Handling of sparse, point-wise observations; Observations requirements for conditioning; Impact of observation density.\n\nOur framework handles observations by performing guidance sampling in observation space, not directly in the latent space. As shown in Equation (7), we leverage the VAE decoder $D$ and observation operator $H$ to calculate the distance between the observation and latent state,  $\\pmb y-H(D(\\pmb z_0))$. The gradient then provides guidance towards observation for the latent state. Benefiting from such a guidance mechanism, our framework does not require entire spatial patch observations. \n\nWe have conducted detailed ablation studies on this in the Appendix. We would like to show the DA accuracy again in the following tables. Our results show that our aligned strategies consistently outperform the non-aligned baselines across all tested observation densities (1% to 10%). This confirms the robustness of our method. As expected, the relative improvement from alignment slightly diminishes as observation density increases, since more observations naturally provide a stronger constraint on their own.\n\n|1% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0623|0.1368|1.5009|2.7027|103.3036|1.0911|\n|Align-Repaint-M|0.0598(-4.12%)|0.1340(-1.78%)|1.4759(-1.69%)|2.6836(-0.71%)|96.8906(-6.62%)|1.0889(-0.20%)|\n|DPS|0.0609|0.1337|1.4441|2.5972|93.4654|1.1159|\n|Align-DPS-M|0.0593(-2.65%)|0.1317(-1.51%)|1.4187(-1.78%)|2.5800(-0.66%)|88.4235(-5.70%)|1.1100(-0.53%)|\n\n|3% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0583|0.1327|1.4438|2.6097|100.0290|1.0696|\n|Align-Repaint-M|0.0560(-4.02%)|0.1305(-1.73%)|1.4188(-1.77%)|2.5929(-0.65%)|93.8835(-6.54%)|1.0675(-0.20%)|\n|DPS|0.0546|0.1276|1.3698|2.4883|85.5742|1.0586|\n|Align-DPS-M|0.0533(-2.53%)|0.1256(-1.56%)|1.3457(-1.78%)|2.4724(-0.65%)|80.5474(-6.24%)|1.0558(-0.27%)|\n\n|5% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0547|0.1290|1.3891|2.5248|96.9590|1.0495|\n|Align-Repaint-M|0.0526(-3.91%)|0.1269(-1.68%)|1.3656(-1.72%)|2.5081(-0.67%)|91.0853(-6.45%)|1.0471(-0.23%)|\n|DPS|0.0534|0.1267|1.3597|2.4767|85.9150|1.0514|\n|Align-DPS-M|0.0519(-2.78%)|0.1246(-1.70%)|1.3337(-1.95%)|2.4560(-0.84%)|80.5787(-6.62%)|1.0475(-0.37%)|\n\n|10% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0470|0.1204|1.2721|2.3398|90.1747|1.0037|\n|Align-Repaint-M|0.0453(-3.67%)|0.1186(-1.52%)|1.2518(-1.62%)|2.3230(-0.72%)|84.7320(-6.42%)|1.0012(-0.25%)|\n|DPS|0.0524|0.1261|1.3520|2.4673|86.2364|1.0468|\n|Align-DPS-M|0.0510(-2.83%)|0.1239(-1.76%)|1.3252(-2.03%)|2.4458(-0.88%)|80.0644(-7.71%)|1.0403(-0.62%)|\n\n___\n>**Response to Question 2:** Comparison with DiffusionPDE [1]?\n\nThank you for this excellent question. We would like to clarify that the preference alignment of our AlignDA leads to fine-tuning. The concept of \"alignment\" implies starting from a reference model and then optimizing it towards preferences. \n\nWe appreciate the reference to DiffusionPDE[4]. While both methods use diffusion models, they are designed to solve fundamentally different problems. DiffusionPDE aims to be a general PDE solver, generating an entire spatio-temporal field from partial observations. Align-DA is actually a data assimilation (DA) framework. Our focus is on the core DA problem: at a single time step, how to optimally fuse a background field with sparse observations. Given the highly complex dynamics and dimensionality of weather-scale data, a direct \"inpainting\" approach like DiffusionPDE is not yet available in this community. \n\nWe will add a citation and discussion of DiffusionPDE to clarify this distinction. Thank you for the suggestion again.\n\n___\n> **Response to Question 3:** Computationally Costs of Alignment Process\n\nWe are grateful to the reviewer for highlighting this important practical consideration. The alignment process is computationally efficient.\n\n- **DPO Fine-tuning:** The fine-tuning itself is very fast. It runs on 4 A100 GPUs and completes in about **30 minutes** after 500 update steps.\n- **Reward Evaluation:** Evaluating the \"forecast skill\" for a single data sample involves running a 48-hour forecast, which takes about **4 seconds** on one A100 GPU. \n\nThe alignment training itself takes a very small cost but brings significant benefits, as demonstrated in our experiments. We will add the details to the Experiment setting section accordingly.\n\n___\n> **Response to Question 4:** Inpainting Methods\n\nThat's an excellent point. Indeed, the **Repaint** method, which we evaluate as a core guidance strategy in our paper, is **fundamentally an inpainting technique**.\n\nIts mechanism works by repeatedly enforcing the known observations onto the generated field during the diffusion sampling process. This is precisely the definition of inpainting.\n\n___\n>**Response to Weakness 1:** Offline reward alignment v.s End-to-end training \n\nWe thank the reviewer for this insightful comment. We agree that a fully end-to-end framework is an elegant theoretical goal. Nevertheless, our offline alignment strategy is a practical and efficient choice, critical for making this problem computationally tractable and flexible.\n\nSpecifically, one of our key rewards, \"forecast skill,\" requires running a 48-hour forecast. An end-to-end approach would necessitate backpropagating through this entire forecast model for every sample in every training batch, which is computationally infeasible. In contrast, our two‑stage approach of Align-DA decouples reward evaluation from the training loop, enabling the use of complex signals like forecast skill.\n\nOur design is also highly flexible: new rewards can be incorporated simply, and our trained model can readily adapt to various observation densities without requiring complete retraining, unlike a specialized end-to-end model. We will clarify this critical design choice in the revised manuscript.\n\n___\n>**Response to Weakness 2&3:** Unified framework for DA\n\nWe would like to clarify the reason for the two models' design. Weather forecasting focuses on simulating atmospheric dynamics, while DA is formulated as a statistical optimal state estimation problem. Consequently,  forecast models and DA models have historically been developed separately. Even recent end‑to‑end DA-forecast systems such as Fengwu‑ADAS[1], Fuxi‑Weather[2], and Aardvark[3] treat forecasting and assimilation as separate modules. \n\nMoreover, the combination of a deterministic forecast model with a probabilistic DA model is the most common operational approach. A probabilistic DA model can provide an ensemble of initial states, which a deterministic forecast model naturally evolves into ensemble‑based probabilistic forecasts.\n\nTherefore, the combination is not ad-hoc but a standard practice. In this paper, we follow this framework where the forecast model serves to generate a background field and the score-based model is trained for DA.\n\nWe totally agree that developing a unified framework is a valuable long-term objective. However, this remains a significant open challenge for the field, given the high complexity of atmospheric dynamics. \n\n___\nThank you for the detailed and insightful feedback sincerely again. We hope our response may relieve your concerns.\n\nReferences:\n\n[1] https://arxiv.org/abs/2312.12462\n\n[2] https://www.nature.com/articles/s41467-025-62024-1\n\n[3] https://www.nature.com/articles/s41586-025-08897-0\n\n[4] https://arxiv.org/abs/2406.17763 NeurIPS2024"}, {"review_id": "RruBGmNRm9", "reviewer": "Reviewer_ywbS", "summary": "The authors propose to apply ideas from the alignment of text-to-image diffusion models to data assimilation for atmospheric applications. Their framework uses either one type of reward (physical adherence) or three types (assimilation accuracy, forecast skill, and physical adherence). They then use DPO to preference align a model with parts of the ERA5 data and propose a \"soft constraint\" to incorporate aspects such as forecast performance and geostrophic balance that are hard to model otherwise.", "questions": "1. While the authors claim in the introduction, that their research targets the use of background with observational data, the evaluation shows diminishing return when such data is incorporated into the model. Please elaborate.\n2. The evaluation uses 1% observational data. Did you also experiment with different amounts of the observational data? Any insights?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- the idea to apply DPO alignment with multiple reward signals to weather forecast diffusion models seems novel\n- it shows consistent performance gains across metrics and guidance methods; multi-reward alignment clearly outperforms baselines\n- technical rigor with formal derivations for score-based models, posterior decomposition, and DPO adaptation\n- it includes qualitative and quantitative analysis (e.g., Z500 maps, Geo-Score) supporting claims\nWeaknesses:\n- writing is very dense\n- no discussion of computational overhead (no analysis of training/inference time or scalability beyond latent compression)\n- no external state of the art comparison baseline (e.g., FourCastNet+DA)\n- results lack error bars or statistical significance testing", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to the reviewer for their strong support of our work and for the high rating. We sincerely appreciate the detailed and constructive feedback, which not only validates our core contributions but also provides valuable suggestions for improving the paper's clarity and presentation. The meticulous list of formatting corrections is particularly helpful, and we will ensure every point is addressed in the final version.\n\n___\n> **Response to Question 1:** Elaborations about Evaluations\n\nThank you for this insightful comment. The role of data assimilation (DA) is to correct the background state toward the true atmospheric state, which is fundamentally a probabilistic inference problem. Stronger constraints correspond to narrower posterior distributions with reduced uncertainty, while weaker constraints yield broader distributions and more diverse plausible solutions.\n\nBoth the alignment method and the assimilation of observations act to constrain the solution space. These constraints, however, are not independent—their overlap increases with the number of observations, as illustrated in Fig. 4. As a result, the marginal benefit of Align-DA diminishes when observational coverage becomes dense.\n\n___\n>**Response to Question 2:** Experiments and Insight on Observation densities\n\nThanks for your constructive suggestion. We have conducted additional experiments on different amounts of observational data in the Appendix. We would like to show the DA accuracy again.\n\n|1% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0623|0.1368|1.5009|2.7027|103.3036|1.0911|\n|Align-Repaint-M|0.0598(-4.12%)|0.1340(-1.78%)|1.4759(-1.69%)|2.6836(-0.71%)|96.8906(-6.62%)|1.0889(-0.20%)|\n|DPS|0.0609|0.1337|1.4441|2.5972|93.4654|1.1159|\n|Align-DPS-M|0.0593(-2.65%)|0.1317(-1.51%)|1.4187(-1.78%)|2.5800(-0.66%)|88.4235(-5.70%)|1.1100(-0.53%)|\n\n|3% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0583|0.1327|1.4438|2.6097|100.0290|1.0696|\n|Align-Repaint-M|0.0560(-4.02%)|0.1305(-1.73%)|1.4188(-1.77%)|2.5929(-0.65%)|93.8835(-6.54%)|1.0675(-0.20%)|\n|DPS|0.0546|0.1276|1.3698|2.4883|85.5742|1.0586|\n|Align-DPS-M|0.0533(-2.53%)|0.1256(-1.56%)|1.3457(-1.78%)|2.4724(-0.65%)|80.5474(-6.24%)|1.0558(-0.27%)|\n\n|5% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0547|0.1290|1.3891|2.5248|96.9590|1.0495|\n|Align-Repaint-M|0.0526(-3.91%)|0.1269(-1.68%)|1.3656(-1.72%)|2.5081(-0.67%)|91.0853(-6.45%)|1.0471(-0.23%)|\n|DPS|0.0534|0.1267|1.3597|2.4767|85.9150|1.0514|\n|Align-DPS-M|0.0519(-2.78%)|0.1246(-1.70%)|1.3337(-1.95%)|2.4560(-0.84%)|80.5787(-6.62%)|1.0475(-0.37%)|\n\n|10% observation|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|Repaint|0.0470|0.1204|1.2721|2.3398|90.1747|1.0037|\n|Align-Repaint-M|0.0453(-3.67%)|0.1186(-1.52%)|1.2518(-1.62%)|2.3230(-0.72%)|84.7320(-6.42%)|1.0012(-0.25%)|\n|DPS|0.0524|0.1261|1.3520|2.4673|86.2364|1.0468|\n|Align-DPS-M|0.0510(-2.83%)|0.1239(-1.76%)|1.3252(-2.03%)|2.4458(-0.88%)|80.0644(-7.71%)|1.0403(-0.62%)|\n\nThe results show that our aligned strategies consistently outperform the non-aligned baselines across all tested observation densities (1%, 3%, 5%, and 10%). This further confirms the effectiveness and robustness of our method. As expected, the relative improvement from alignment slightly diminishes as observation density increases, since more observations naturally provide a stronger constraint on their own.\n\n___\n>**Response to Weakness 1:** Dense Writing\n\nWe thank the reviewer for this valuable feedback. We acknowledge that the writing is relatively dense. Following your advice, we will carefully revise the manuscript to improve clarity and readability, breaking down complex sentences and adding clearer transitions where possible.\n\n___\n>**Response to Weakness 2:** Discussion of Computational Overhead\n\nWe appreciate the reviewer for highlighting this point. We will add the following details to the Section Experimental Setting for clarity. \n1. One-Time Pre-training: This is the most computationally intensive part. The VAE training takes approximately 4 days, and the diffusion model pre-training takes about 1 day, both on 4 A100 GPUs.\n2. DPO Alignment (Fine-tuning): This step is highly efficient. The alignment process itself, which is the core of our proposed method, completes in about 30 minutes on 4 A100 GPUs. \n3. Inference (Sampling): Generating a single analysis field takes approximately 30 seconds on one A100 GPU.\n\n\n___\n>**Response to Weakness 3:** More Comparison baseline (e.g., FourCastNet+DA)\n\nThank you for this insightful comment. Our paper's primary contribution is to introduce alignment mechanisms to enhance score-based DA methods. It leads to our main experiment studies are to compare different score-based DA methods with and without alignment, where the forecast model stays unchanged. \n\nWe agree that using different forecast models is valuable to validate the efficiency of our framework; however, we encountered a significant technical barrier to comparing with a model like FourCastNet. The publicly available FourCastNet models operate at a different resolution (0.25 degrees) than our current experimental setup (1.4 degrees). A fair comparison would require us to re-engineer and re-train our entire high-dimensional VAE and diffusion model pipeline from scratch to match this new resolution.\n\nGiven these constraints, performing such an experiment was not feasible. We believe our current study provides a clear and robust validation of our proposed methodology, and we hope to explore such large-scale comparisons in future work.\n\n___\n>**Response to Weakness 4:** Error Bars or Statistical Significance Testing\n\nThank you for this crucial suggestion, which makes our paper more solid. To confirm our results, we ran a statistical significance test under a specific setting—without using any observations or alignment—because this would yield the largest uncertainty. We ran the diffusion sampling 10 times for a single time step and calculated the metrics for each run to see the relative standard deviation. The results are presented in the following table. \n\n|DA accuracy (ObsFree)|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|**Relative standard deviation(%)**|0.045|0.058|0.115|0.079|0.201|0.083|\n\n|Forecast Skill(ObsFree)|MSE|MAE|u10-WRMSE|v500-WRMSE|z500-WRMSE|t850-WRMSE|\n|:---|:---|:---|:---|:---|:---|:---|\n|**Relative standard deviation(%)**|0.163|0.067|0.145|0.155|0.252|0.106|\n\n\nThis test confirmed that our reported improvements are statistically significant. For instance, the relative standard deviation was less than 0.1% for the main accuracy metric. This is much smaller than the 2.65% to 5.44% improvement that Align-DA provides over the baseline, which confirms our gains are real and not just from random chance.\n\nIn our revised paper, we will add a brief statement to the main text explaining that our results are statistically significant. We will also create a new section in the Appendix that shows our main result tables with the full mean and standard deviation values for complete transparency. In that section, we will also add a specific note about the single-preference result, stating that: \"**The performance change of its single-reward version (-P) yields statistically insignificant, suggesting that geostrophic balance alone is insufficient for DA tasks.** \"\n___\n>**Response to Paper Formatting Concerns:**\n\nWe are extremely grateful to the reviewer for taking the time to provide such a detailed list of formatting corrections. This is incredibly helpful for improving the final quality of the paper. We confirm that we will address each point in the revised version, including:\n- Adjusting the page length.\n- Revising the \"conditional\" in the figure 2 title to \"ObsFree\".\n- Ensuring consistent capitalization (FengWu).\n- Revise line 217 and 230 spacing issues. \n- Fixing the bracket in the Table 2 caption.\n- Standardizing the spelling of \"constraint\".\n- Changing the placement of Figure 3 and Figure 4.\n- Thoroughly revising the reference list for consistency in journal/conference capitalization, citation styles (NeurIPS), and inclusion of publication locations.\n\n---\nThank you again for this meticulous and valuable feedback."}]

Author rebuttal:

Meta-review: Paper summary:

This paper presents a reinforcement learning-based framework for data assimilation (DA), named Align-DA, to estimate the system states of atmospheric applications given their prior states and observations. In this framework, diffusion models are used for DA, and goals of assimilation and domain-specific knowledge can be formulated as reward functions that guide the learning process of the models. As an example, the authors explored regularizing the optimization with three metrics, assimilation accuracy, forecast skill, and physical adherence. Experimental results demonstrate that models trained using the proposed framework with multiple rewards outperform those trained without or with a single reward.

Summary of the discussion:

4 reviewers for for acceptance. 2 with confidence and 2 only weakly. 1 reviwer votes weakly for rejection. 

The only reviewer voting for rejection said they cannot see that the authors are positioning the paper such that it offers a clear, substantial contribution. The clarified core research problem appears to be narrower than necessary given the other materials presented. Comparing to an existing DA method, the proposed Align-DA has not demonstrated clear superiority.

None of the other 4 reviewers shared this concern, which reduces the weight of the latter comment.

Recommendation:

This is a borderline paper with 4 reviewers voting for acceptance (two weakly) and another reviwer voting weakly for rejection. Based on the above, I recommend accepting the paper and encourage the authors to use the feedback provided to improve the paper for its camera ready version.
Decision: Accept (poster)