Align-DA: Align Score-based Atmospheric Data Assimilation with Multiple Preferences

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.

Response to Question 1: Elaborations about Evaluations

Thank 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.

Both 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.

Response to Question 2: Experiments and Insight on Observation densities

Thanks 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.

The 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.

Response to Weakness 1: Dense Writing

We 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.

Response to Weakness 2: Discussion of Computational Overhead

We appreciate the reviewer for highlighting this point. We will add the following details to the Section Experimental Setting for clarity.

1. 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.
2. 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.
3. Inference (Sampling): Generating a single analysis field takes approximately 30 seconds on one A100 GPU.

Response to Weakness 3: More Comparison baseline (e.g., FourCastNet+DA)

Thank 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.

We 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.

Given 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.

Response to Weakness 4: Error Bars or Statistical Significance Testing

Thank 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.

This 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.

In 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. "

Response to Paper Formatting Concerns:

We 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:

• Adjusting the page length.
• Revising the "conditional" in the figure 2 title to "ObsFree".
• Ensuring consistent capitalization (FengWu).
• Revise line 217 and 230 spacing issues.
• Fixing the bracket in the Table 2 caption.
• Standardizing the spelling of "constraint".
• Changing the placement of Figure 3 and Figure 4.
• Thoroughly revising the reference list for consistency in journal/conference capitalization, citation styles (NeurIPS), and inclusion of publication locations.

Thank you again for this meticulous and valuable feedback.

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.

Response to Question 1: Handling of sparse, point-wise observations; Observations requirements for conditioning; Impact of observation density.

Our 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.

We 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.

Response to Question 2: Comparison with DiffusionPDE [1]?

Thank 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.

We 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.

We will add a citation and discussion of DiffusionPDE to clarify this distinction. Thank you for the suggestion again.

Response to Question 3: Computationally Costs of Alignment Process

We are grateful to the reviewer for highlighting this important practical consideration. The alignment process is computationally efficient.

• 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.
• 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.

The 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.

Response to Question 4: Inpainting Methods

That's an excellent point. Indeed, the Repaint method, which we evaluate as a core guidance strategy in our paper, is fundamentally an inpainting technique.

Its mechanism works by repeatedly enforcing the known observations onto the generated field during the diffusion sampling process. This is precisely the definition of inpainting.

Response to Weakness 1: Offline reward alignment v.s End-to-end training

We 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.

Specifically, 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.

Our 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.

Response to Weakness 2&3: Unified framework for DA

We 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.

Moreover, 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.

Therefore, 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.

We 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.

Thank you for the detailed and insightful feedback sincerely again. We hope our response may relieve your concerns.

References:

[1] https://arxiv.org/abs/2312.12462

[2] https://www.nature.com/articles/s41467-025-62024-1

[3] https://www.nature.com/articles/s41586-025-08897-0

[4] https://arxiv.org/abs/2406.17763 NeurIPS2024

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.

Response to Question 1 & Weakness 3: Reasons for Traditional DA Methods

We 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.

Modern 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.

In 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.

We 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.

Response to Question 2 & Weakness 2: More Physics Constraint

We 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.

Inspired 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).

The 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.

Response to Question 3: Extreme Weather Condition

We 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.

Nevertheless, 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.

We agree that the capability for handling extreme weather conditions is important. We will carefully discuss this direction in the revised version.

Response to Weakness 1: The Novelty of Model Structure

We 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.

In 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.

Our 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.

We appreciate your valuable feedback again, which has enriched our work and improved our manuscript.

Reference:

[1] https://doi.org/10.1175/MWR-D-24-0058.1

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.

Clarifications about Main Contribution

To 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.

Response to Question 1: Relationship between Align-DA, Fengwu, and Baselines; Source of performance gain; Baselines' access to FengWu

The Relationship Between Align-DA, FengWu, and Baselines

The 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.

FengWu is a state-of-the-art weather forecasting model. In our work, it is not a competitor but serves two critical roles:

• Background Generator: FengWu generates the 48-hour background forecast that serves as the common starting point for all DA methods, ensuring a fair comparison.
• 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.

Does Align-DA's Improvement Come From Using FengWu?

Thank 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.

FengWu 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.

Baseline 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.

Can the Baselines Also Leverage FengWu?

The Baselines are unable to leverage Fengwu as a reward model in the same way as our Align-DA, highlighting the contribution of our work.

Unlike 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.

In 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

We sincerely appreciate your valuable feedback, which has significantly enhanced the clarity and completeness of our paper.

We sincerely thank the reviewer for their thorough and constructive feedback. We will address each point in detail below.

Response to Question 1: Difference from the existing RL-based DA

Thank 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.

Existing 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].

In 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.

To 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]."

Response for Question 2: Specialization, Generalization, and experiments for other domains

Thank 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.

Response to Weakness 1: Generalizability of our framework

We 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.

We 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.

We 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.

To 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.

Response to Weakness 2: "Teaching to the test" concern and comparison with other frameworks

To 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.

Our 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.

We 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.

Response to Weakness 3: Comparisons with state-of-the-art methods for the forecast or those methods that can be applied to this task.

We 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.

Consequently, 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.

We 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.

Response to Weakness 4: Explain the DA background in detail

We 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.

Here is what we revised in Line 30: "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.

Thank you for your helpful comments again. We hope this response can solve your concern.

Reference:

[1] https://doi.org/10.1029/2023MS004178

[2] https://doi.org/10.1002/psp4.12588

[3] https://doi.org/10.1016/j.jocs.2020.101171

[4] https://doi.org/10.1016/j.geoen.2024.213275

[5] https://arxiv.org/abs/2505.05452

[6] https://arxiv.org/abs/2310.12036

[7] https://arxiv.org/abs/2504.15176 ICLR2025

Dear Reviewer 46nx,

Thank you for your thoughtful feedback again.

To respond to your valuable comments, we have clarified how Align-DA's preference alignment paradigm is fundamentally different from sequential RL-DA approaches. We also elaborated on how this framework can be generalized beyond atmospheric science. Crucially, to provide strong empirical evidence for our claims, we performed new experiments. First, we conducted a preference optimization without physics-based rewards explicitly. The physical consistency improvement indicates the model is not simply overfitting on metrics, which addresses the "teaching to the test" concern directly. Second, to further validate that our alignment concept is a powerful and generalizable tool, we also demonstrated its success with other preference optimization methods, including IPO and DSPO.

We have incorporated these new results and other clarifications into the manuscript and hope these additions fully address your concerns. We would be grateful for the chance to discuss this further with you.

Best regards

The authors

Please respond to the authors' rebuttal and indicate if you are happy to reconsider your score based on the rebuttal and the other reviewers' comments.

Thank you for your clarification and the additional experiments. After reading the reply, I now understand the significance of the contribution better.

Ultimately, I consider there two ways to position this work: (1) A general framework to improve existing methods' performance at a specific task by incorporating new, domain-specific reward functions. (2) A new method for weather forecasting providing superior accuracy. At this point I believe the authors are more inclined toward (1); however, if this is the case, multiple small case studies demonstrating the effectiveness and generalizability of the proposed framework must be shown, where the atmospheric application studied in this paper can only be considered as one of them. If (2) is the case, comparison to existing weather forecasting approaches is necessary. In a similar vein, Reviewer eG1A also raised the concern of baseline comparison to traditional data assimilation.

As critical items are still missing for either interpretation, I am refrained from raising my rating. Nonetheless, it appears to me that the proposed framework should be able to transfer to other domains, as the authors suggested, even though I have not seen any evidence. Plus, all other reviewers have given positive ratings. Therefore, although I am more toward keeping my borderline reject recommendation, I acknowledge the paper's potential value to the community and understand it could still be considered for acceptance with appropriate revisions.

I apologize for the narrower time frame, but I am still glad to know the authors' thought about my concern.