FlowDAS: A Stochastic Interpolant-based Framework for Data Assimilation

We thank the reviewer for their balanced and constructive feedback, and we are encouraged by their recognition of the potential of our work.

A. On the implementation of 3D-Var in high-dimensional settings

We agree that the background error covariance matrix is a critical component of modern data assimilation systems used in numerical weather prediction (NWP). Like in our SEVIR settings, the B matrix measures how an error of the VIL at one location is related to errors at other locations and is quite critical for modern variational data assimilation (VAR) systems [1]. However, it is always hard and complicated to compute this B matrix, especially in high-dimensional systems like the SEVIR and NS. Taking the SEVIR as an example, we only take VIL into consideration and its size is $1\times 128\times 128$, and its B-matrix would have a shape of $16384 \times 16384=268,435,456$. Even though the B matrix is a positive semi-definite, symmetric matrix, it would still be very hard to compute, and we would need to resort to highly specialized scientific software like GEN_BE v2.0 [2].

Given this, and the fact that FlowDAS is intended as a general-purpose, model-agnostic framework, looking into this specific question setup would be an important but orthogonal direction to our current contributions. We would be happy to discuss this point further and add it in the discussion section if the paper is accepted.

Besides, we have added EnKF as a baseline in our Lorenz 63 experiments. In EnKF, the forecast error covariance $B$ is approximated using the sample covariance of the forecast ensemble, which avoids explicit construction of $B$ and instead operates directly on the ensemble. We observe that FlowDAS still outperforms the EnKF in terms of RMSE. Besides, both BPF and EnKF achieve similar RMSE, which is expected as they both have the knowledge of the forward dynamic model.

Note: As suggested by Reviewer Ztob, we also include PDEDiff, the universal amortised extension of SDA proposed in [4], as an additional baseline for completeness. PDEDiff trains a conditional score network to autoregressively predict the next state(s) given a number of past states, enabling more flexible and accurate sequential generation compared to the original all-at-once SDA. This helps address the forecasting limitations of SDA’s joint score formulation.

Finally, as we transition FlowDAS to real-world weather forecasting applications in collaboration with industry, the design and approximation of the background error covariance will definitely be an important consideration for future work.

 

B. On computational cost comparisons:

We agree with the reviewer on the importance of computational efficiency and have added a detailed performance comparison in the appendix. The results, using a single A100 GPU on the SEVIR dataset, highlight FlowDAS’s practical viability:

The slightly longer inference time for FlowDAS stems from its iterative correction process, which is essential for achieving accurate and observation-consistent predictions. While this leads to higher per-sample inference time, it is balanced by significant training efficiency, as FlowDAS reaches convergence much faster than all competing baselines. We believe this reflects a favorable trade-off in practical deployment scenarios. A more detailed discussion will be included in the final version.

To further improve inference speed, one promising direction is to apply post-hoc distillation, where a faster surrogate model is trained to approximate the FlowDAS sampler. Another avenue is to explore hybrid formulations that integrate variational inference techniques into the framework, potentially offering both computational gains and improved flexibility.

 

C. On ensemble sampling and uncertainty modeling in inference:

For uncertainty modeling, yes, FlowDAS does support ensemble-style sampling. At each assimilation step, we can generate multiple samples from the posterior $p( x\_{k+1} | x\_k, y\_{k+1})$, allowing us to form an ensemble of analysis trajectories. These samples can then be propagated forward through the system dynamics, enabling flow-dependent uncertainty evolution similar to EnKF. In our experiments, we used this capability to assess uncertainty. For instance, in the Navier–Stokes (NS) experiments, we conducted inference 25 times independently and used the resulting ensembles to compute standard deviations of RMSE across different settings. We apologize for not making this explicit in the main text. Below we report the mean and standard deviation (std) of RMSE over 25 sampled trajectories (with some additional experiments on more challenging assimilation scenarios suggested by Reviewer 3gVf):

We will clarify this ensemble sampling capability and its usage in uncertainty quantification in the revised paper.

 

Reference

[4] Shysheya, A., Diaconu, C., Bergamin, F., Perdikaris, P., Hernández-Lobato, J. M., Turner, R., & Mathieu, E. (2024). On conditional diffusion models for PDE simulations. Advances in Neural Information Processing Systems, 37, 23246-23300.

We thank the reviewer for their balanced and constructive feedback, and we are encouraged by their recognition of the potential of our work.

Weaknesses:

A. “The writing feels rushed in places”

We thank the reviewer for the careful reading and for pointing out these writing issues. We corrected all the noted misprints. Specifically, we now clarify that “discontinuity situations” refer to assimilation tasks involving abrupt state changes or non-smooth dynamics that often challenge traditional filtering methods.

B. “Comparisons with classical filtering baselines are largely missing.”

We have added EnKF as classical filtering baselines in the Lorenz 63 experiment. The updated RMSE comparison is shown below:

The highly nonlinear system dynamics induce non-Gaussian ensembles, which breaks the assumption of EnKF. Our nonlinear observation model in this case further deteriorates its performance. In contrast, FlowDAS is data-driven and does not rely on linearity or Gaussianity assumptions. As a result, it achieves the lowest RMSE across all baselines, including classical filters BPF and EnKF, and diffusion-based methods, e.g., SDA and PDEDiff. This supports FlowDAS’s suitability for nonlinear and non-Gaussian data assimilation tasks.

C. “It’s unclear what advantages FlowDAS brings over EnSF and DiffDA, etc.”

We have added EnSF to the Lorenz-63 experiment (details below). Its performance is on par with BPF. This is expected since EnSF uses a training-free Monte Carlo estimator for the score function, which faces the same particle collapse issues, i.e., samples concentrate in high-probability regions of the transition dynamics, leading to biased score estimates. This reflects a train-sampling tradeoff: EnSF saves training time but cannot learn accurate scores, especially under highly nonlinear dynamics and observation models like in our noisy Lorenz-63 settings(e.g., our $\text{atan}(x\_1)$ observation). FlowDAS, by learning the drift with collected data, avoids this issue and achieves superior performance.

We do not include a direct comparison to DiffDA for the following reasons:

1). Mismatch in data and task setup: DiffDA is tailored for ERA5 reanalysis, which involves many atmospheric variables. In contrast, SEVIR is a simplified dataset with only vertically integrated liquid fields. Adapting DiffDA to SEVIR is non-trivial and out of scope for this work, but we are happy to include it in future studies.

2). Architectural differences: DiffDA uses an external numerical weather model as its forecasting component and applies a diffusion model purely as a denoiser. FlowDAS handles both forecasting and assimilation jointly with a single learned stochastic model, which is more modular and general-purpose.

We will clarify this distinction in the final version and emphasize that FlowDAS targets general, model-agnostic data assimilation without relying on external solvers.

 

Questions:

1. “What do the acronyms for FlowDAS stand for?”

We apologize for the omission. FlowDAS stands for Flow-based Data Assimilation with Stochastic Interpolants. Here, “Flow” refers to the use of stochastic interpolant flows to model the posterior distribution, and “DAS” is short for Data Assimilation. We will revise the manuscript to introduce and define this acronym explicitly at first mention.

2. “More details on how to roll out FlowDAS”

• To improve clarity in inference, we revised the Inference paragraph in the main text to explicitly explain how FlowDAS proceeds in sequential data assimilation (see below).
• To directly address your questions on the target distribution that FlowDAS samples from: In our current setup, we sample $x\_{k+1} \sim p(x\_{k+1} | x\_k, y\_{k+1})$ with $n=1$, and proceed sequentially. That is, we assimilate one observation at a time and do not marginalize over multiple future samples. FlowDAS targets a local posterior conditioned on the current state and the next observation only. We agree that extending FlowDAS to sample from $p(x\_{k+1} | x\_{k}, y\_{k+1:k+P})$ would provide a more expressive and generalized model and can be better suitable to approximate the full filtering distribution. We mention this as a promising direction for future work.

Inference Given a trained drift model $\hat{b}_s$, FlowDAS performs inference in an autoregressive manner to estimate the latent state trajectory $\{ \hat{x}\_{1:K} \}$. At each time step $k$, we begin from the known state $\hat{x}_k$ and predict the next state $\hat{x}\_{k+1}$ using a discretized stochastic interpolant over a time grid $s_0 = 0 < s_1 < \cdots < s_N = 1$. We set the interpolant state $X\_{s\_0} \gets \hat{x}\_k$ and retrieve the next observation $y\_{k+1}$. For each interpolation step $s\_n$, we simulate a forward transition using the learned drift, given in Eq.7, and noise. To enforce observation consistency, we generate $J$ posterior endpoint samples $\hat{X}\_1^{(j)}$, and compute their likelihoods under the measurement model, i.e., $|| y - \mathcal{A}(\hat{X}\_1^{(j)}) ||_2^2$. A softmax over these values gives importance weights $\{ w\_{1:J} \}$. We then apply a correction to the interpolant using a weighted gradient $- \zeta\_n \nabla\_{X\_{s\_n}} \sum \_{j=1}^J w\_j || y - \mathcal{A}(\hat{X}\_1^{(j)}) ||_2^2$, where $\zeta\_n$ denotes the step size. After reaching $s=1$, we set $\hat{x}\_{k+1} \gets X\_{s\_N}$ and proceed to the next time step. Repeating this procedure over all $k$ yields the full trajectory.

3. Particle collapse issue for FlowDAS

This is a very insightful comment! This is indeed a well-known issue in high-dimensional importance sampling and a valuable observation for our method.

We observed that the likelihood score can first decrease and then unexpectedly increase near the end of the sampling trajectory. When this occurs, we switch to the alternative sampling strategy described in the supplementary material. While it may occasionally produce slightly suboptimal results, it is more robust compared to the importance sampling approach in the main text. We will explicitly mention this issue and our workaround in the revised paper, and plan to investigate more principled solutions in future work.

More specifically, we found that particle collapse tends to occur during the final 30–50 steps of the sampling process. In these cases, the likelihood term $||y - Ax||_2^2$, which is expected to decrease steadily, begins to rise unexpectedly—eventually diverging and causing sampling failures. Upon examining the importance weights, we observed a clear pattern:

• When the likelihood is decreasing as expected, the weights are distributed relatively evenly across particles.
• When the likelihood increases abnormally, the weights become highly concentrated on a single particle (close to 100%), indicating a collapse of the Monte Carlo estimator.

To address this, we use the following two-step strategy:

• Resampling: When the weights are abnormally concentrated, we immediately resample a new set of particles. This is a standard remedy for particle collapse in importance sampling.
• Alternative sampling strategy: If resampling does not resolve the issue, we switch to the alternative sampling method described in the supplementary. Although it is not always as precise, it significantly improves robustness and prevents divergence.

4 & 5. Metrics used when evaluating FlowDAS

We clarify that these metrics were only used in the Lorenz experiments, where our goal was to benchmark FlowDAS against SDA. Since SDA used these evaluation metrics, we adopted the same for direct comparability. Below, we explain the meaning of each metric and we will clarify them in the revised text:

• Expected log-prior $\mathbb E\_{q(x\_{1:K} | y)} [\log p(x\_{2:K} | x\_1)]$: This measures whether the generated trajectories conform to the known prior dynamics. Specifically, given a sampled trajectory $x\_{1:K} \sim q$, we evaluate how likely it is under the true transition model $p(x\_{k+1} \mid x\_k)$ for $k = 1, \ldots, K-1$. A higher value indicates better physical consistency with the governing equations. It penalizes trajectories that deviate from the expected evolution under the dynamics, even if they match observations.
• Expected log-likelihood $\mathbb E\_{q(x\_{1:K} | y)} [\log p(y | x\_{1:K})]$: This evaluates how well the generated trajectories explain the observed data. For each sampled trajectory, we compute the likelihood of the full observation sequence $y$, and average this across samples. This score reflects data fidelity, i.e., higher values indicate that the sampled trajectories are more consistent with the observations under the assumed noisy and nonlinear observation model.
• Wasserstein distance $W\_1(p, q)$: This is computed between the estimated and ground-truth trajectory distributions. In our case, it becomes a trajectory-wise distance between delta measures centered at $\hat{x}\_{1:K}$ and $x\_{1:K}$. While simple, it provides an interpretable measure of deviation in trajectory space.

 

Appendix

Implementation of EnSF

[1] Bao, Feng, Zezhong Zhang, and Guannan Zhang. "An ensemble score filter for tracking high-dimensional nonlinear dynamical systems." Computer Methods in Applied Mechanics and Engineering 432 (2024): 117447.

First of all, we sincerely thank the reviewer for the positive evaluation and encouraging comments, as well as for the insightful suggestions that helped us further improve the paper! :)

Weakness

A. “Stochastic integrators introduce bias; accuracy-speed trade-offs under very stiff dynamics ...”

1). Accuracy-speed trade-offs Indeed, the stochastic integrators in Equations (9) and (10) trade exactness for efficiency and introduce numerical bias. To evaluate this trade-off, we compared three settings: first-order integrator, second-order integrator, and a hybrid method that computes the full stochastic interpolant path to the endpoint $X\_1$ for the final 100 steps (“2nd order × 400 + Flow × 100”). As shown in the table (now included in Appendix D), the full-path method increases computation time by 7× and peak memory usage by 18× compared to the second-order integrator, while improving RMSE by only ~0.1%.

We initially considered the full Flow*500 method, i.e., computing $X\_1$ exactly at every step by fully rolling out the leftward stochastic interpolant, but this would require about 500+(500+499+...+1) $\approx$ 250,000 forward passes and retaining the entire computation graph at each step. This would cause out-of-memory failures even on A100 GPUs. This substantial cost arises because computing the full-path endpoint requires storing the entire interpolation path to backpropagate $\log p(y|X\_1)$. In contrast, the 1st- and 2nd-order methods yield stable, efficient inference while maintaining competitive accuracy.

We believe this experiment validates our design choice: the benefit of observation conditioning significantly outweighs the small bias introduced by the integrator. We will clarify this point in the final version and include the trade-off analysis and table in the appendix.

2). More stiff dynamics In our experiments on the Navier–Stokes system, known for its stiffness and chaotic behavior, we have observed stable rollout behavior and robust error control. While we do not explicitly tailor our integrators for stiffness (e.g., via implicit solvers or adaptive step size control), the short-step nature of the interpolant and the sequential correction mechanism naturally mitigate stiffness-related instabilities.

In addition, we are currently extending FlowDAS to a real-world ocean dynamics application, where the governing equations follow the Quasi-Geostrophic (QG) model. These systems exhibit more complex behavior, such as inverse energy cascades and long-range spatial correlations, making them even more challenging than standard Navier–Stokes in certain regimes. Our preliminary results indicate that FlowDAS still performs well under this setting, though further gains are likely possible through domain-specific strategies tailored to QG-type dynamics. We view this as a promising direction for advancing the applicability of FlowDAS to operational-scale geophysical systems.

B. “Difficult uncertainty quantification...”

While FlowDAS generates a single trajectory per sampling run, its stochastic formulation naturally supports uncertainty quantification via ensemble sampling. In our weather forecasting task, we already report ensemble standard deviation as a measure of uncertainty (see Table S.10). For the Navier–Stokes experiments, we have similarly conducted Monte Carlo sampling to quantify uncertainty and will include the ensemble standard deviation in the updated version of the paper. This is done by running multiple inference trajectories under the same observation setting and computing the spread of predictions. As shown in the table below, the uncertainty increases as the observation becomes sparser or lower in resolution, which aligns with expectations.

This ensemble-based strategy provides a practical and interpretable way to assess uncertainty within the FlowDAS framework.

 

Questions

A. “Could variational inference over the interpolant path yield accurate density spread...”

Thank you for this insightful question. Since FlowDAS is based on MCMC-style sampling, it is natural to ask whether variational inference (VI), a widely used alternative in inverse problems and posterior approximation, could provide comparable or even improved density estimates. We address this in three parts:

• Can VI be integrated into the FlowDAS framework?
  Yes. Following RED-Diff [1], we can construct a RED-FlowDAS variant. All theoretical components generalize naturally to the stochastic interpolant setting (see our brief derivation).

• How does VI compare to EnKF and FlowDAS?

  • Compared to EnKF, VI is more flexible: it accommodates general likelihoods (not limited to Gaussian noise) and models complex state transitions via KL divergence minimization, which goes beyond EnKF’s linear assumptions.
  • Compared to FlowDAS, VI provides an explicit density estimate, while FlowDAS samples from the posterior and thus models the posterior distribution indirectly. However, MCMC may better explore complex multi-modal posteriors, whereas VI can be mode-seeking due to the nature of KL divergence.

• Additional benefits of VI:
  VI avoids the need for conditional score estimation, can be faster than MCMC, and is amenable to stochastic and distributed optimization. These advantages make it a promising direction for extending FlowDAS.

B. “Performance degradation under challenging cases & Localization?”

1). Performance degradation under extremely challenging cases We have conducted additional experiments using super-resolution (SR) and sparse-observation (SO) settings with as little as 0.75% coverage. The results, shown below, indicate a graceful degradation of performance as the amount of observed information decreases. This supports the robustness of FlowDAS even under highly limited measurement scenarios.

These results will be included in the updated appendix.

2). Localization Unlike ensemble-based filters such as EnKF, which rely on covariance estimation and therefore require localization to suppress spurious long-range correlations, FlowDAS does not estimate or use a background covariance matrix. Instead, it performs observation-conditioned updates directly through a likelihood-based correction at each step, guided by the measurement operator A. This naturally confines the influence of observations to their spatial support, eliminating the need for explicit localization.

 

(Optional) Brief Derivation for VI-FlowDAS

For completeness, we include a brief derivation showing how variational inference can be integrated into FlowDAS.

We could format the data assimilation question in a VI settings and that is we want to fit a family of densities $q({X}_1 \mid {X}_0,{Y})$ to a intractable posterior distribution $P({X}_1 \mid {X}_0,{Y})$, where the ${X}_1$ is the state variable we want to predict, and ${X}_0$ is the current state and ${Y}$ is the observation on ${X}_1$. And the variation method is based on KL minimization

$\min_q KL\big(q({X}_1 \mid {X}_0,{Y}) || P({X}_1 \mid {X}_0,{Y}))$

where $q := \mathcal{N}(\mu, \sigma^2 I)$ is a variational distribution. And this objective can be expand:

$KL\big(q({X}\_1 \mid {X}\_0,{Y}) \| P({X}\_1 \mid {X}\_0,{Y})\big) = -{E}\_{q({X}\_1 \mid {X}\_0,{Y})}\big[\log P({Y} \mid {X}\_1,{X}\_0)\big]\+KL\big(q({X}\_1 \mid {X}\_0,{Y})\| p({X}\_1 \mid {X}\_0)\big) + \log p({Y},{X}\_0)$

The last term is the observation likelihood that is constant w.r.t.~$q$. Thus, to minimize the KL divergence, it suffices to minimize the variational bound (omit the last term). This brings us to the next claim.

${Proposition 1}$. The KL minimization is equivalent to minimizing the variational bound, that itself obeys the velocity matching loss in the stochastic interpolants framework:

$\min\_{\{\mu, \sigma\}}-{E}\_{q({X}\_1 \mid {X}\_0,{Y})}\left[\log P({Y} \mid {X}\_1,{X}\_0)\right] +\int_{0}^{1} \frac{1}{2\sigma\_s^{2}} {E}\_{q({X}\_s \mid {X}\_0,{Y})} \Big[\big\|\hat{b}^{q}\_s({X}\_s, {X}\_0) - \hat{b}^{p}\_s({X}\_s, {X}\_0) \big\|^2_2 \Big] ds$

where $q({X}_s | {X}_0,{Y}) = \mathcal{N}( \alpha_s {X}_0 + \beta_s \mu, (\beta_s^2 \sigma^2 + \sigma_s^2 s)I)$ and $\hat{b}^{q}_s({X}_s, {X}_0)$ can be explicitly computed:

$\hat{b}^{q}_s({X}_s, {X}_0) = \dot{\alpha}_s{X}_0 + \dot{\beta}_s{X}_1^{q}+\dot{\sigma}_s{W}_s$

where ${X}_1^{q}$ is produced by sampling from $q({X}_1 \mid {X}_0,{Y}) := \mathcal{N}(\mu, \sigma^2 I)$.

For the $\hat{b}^{p}_s({X}_s, {X}_0)$, this is exactly the stochastic model we trained.

 

Reference

[1]. Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. arXiv preprint arXiv:2305.04391, 2023.

We thank the reviewer for their detailed, balanced and constructive feedback, and we are encouraged by their recognition of the potential of our work.

Weaknesses:

A. “Inappropriate SDA baseline configuration for tasks involving forecasting.”

The SDA results reported in our paper for Navier-Stokes and SEVIR are already based on the universal amortised SDA variant proposed in [1]. The model is conditioned on previous states (P=10 for NS, P=6 for SEVIR) with varying conditioning scenarios during training. We will rename this baseline to PDEDiff for clarity.

For Lorenz-63, we originally used the naive joint-score SDA. We have now added the PDEDiff (P=1, C=1 as FlowDAS) to this setting as well.

We will revise the manuscript to clarify the SDA configurations used and correct any references to joint vs conditional score formulations. Notably, FlowDAS achieves the best performance across most metrics, even compared to BPF, which has access to the exact forward model. While PDEDiff improves upon SDA by modeling conditional scores, it still falls a bit short of FlowDAS.

B. “Lack of computational cost analysis” Response: 1). Regarding the Discretization steps, our stochastic interpolant uses N = 500 time steps during sampling. We will add this to the implementation details section. 2). We have now added quantitative runtime and parameter comparisons (for the weather forecasting task using a single A100 GPU) to the appendix.

The slightly longer inference time for FlowDAS stems from its iterative correction process, which is essential for achieving accurate and observation-consistent predictions. While this leads to higher per-sample inference time, it is balanced by significant training efficiency, as FlowDAS reaches convergence much faster than all competing baselines. We believe this reflects a favorable trade-off in practical deployment scenarios. A more detailed discussion will be included in the final version.

To further improve inference speed, one promising direction is to apply post-hoc distillation, where a faster surrogate model is trained to approximate the FlowDAS sampler. Another avenue is to explore hybrid formulations that integrate variational inference [5] techniques into the framework, potentially offering both computational gains and improved flexibility.

C. “Ablation on conditioning horizon”

To assess the impact of the conditioning horizon, we conducted a controlled ablation using the CRPS metric on both Navier–Stokes (NS) and SEVIR experiments by varying the number of observed previous steps used during inference. The results are summarized below:

• Navier–Stokes CRPS vs. Conditioning Horizon

• SEVIR CRPS vs. Conditioning Horizon

These results clearly show that incorporating a longer observation history leads to consistent and significant improvements in forecasting performance. In particular, we observe a ~30% CRPS reduction in NS when increasing the conditioning length from 1 to 10 steps. Similarly, in SEVIR, although the absolute values are small due to the scale of the task, conditioning on 6 previous frames led to the best predictive accuracy. This supports our claim that incorporating temporal context is important for accurate data assimilation, especially in high-dimensional forecasting tasks.

D. “Missing limitations discussion”

To improve clarity, we will include an explicit Limitations and Future Work section in the revised manuscript.

Limitations and Future Work

FlowDAS has so far been validated on controlled dynamical systems (e.g., Lorenz-63, Navier–Stokes), and its generalizability to more complex, real-world settings such as numerical weather prediction remains to be explored. In particular, we have not yet studied its performance under Sim2Real conditions or applied it to operational-scale, high-dimensional systems like those addressed by GEN_BE.

Although FlowDAS supports probabilistic inference via sampling, the current sampling speed remains a limitation. One possible solution is to apply post-hoc distillation, where a faster surrogate model is trained to mimic the FlowDAS sampler. Exploring hybrid formulations that incorporate variational inference techniques could be another promising direction to improve efficiency.

E. Incomplete coverage of a few related methods

1. (Better coverage of [4]) The reviewer is correct that our Navier-Stokes experiment follows the setup established in [4], and we will certainly add this citation in the relevant section. To further clarify the context of our work, we will also expand our discussion in Related Work section to explicitly describe how FlowDAS builds upon and extends such autoregressive forecasting approaches by incorporating an observation-driven data assimilation component. To reflect this, we plan to revise the paragraph talking about ‘Diffusion model-based methods’ in Section 2 into ‘Generative model-based methods’ and insert the following paragraph discussing [4]:

A closely related line of work is the stochastic interpolant-based forecasting framework recently proposed in [4], which formulates the prediction task as conditional sampling of future system states given the current state, using a fictitious, non-physical stochastic process governed by a learned SDE. While this approach enables probabilistic forecasting, it does not incorporate observation-driven corrections during inference. In contrast, FlowDAS augments the stochastic interpolant dynamics with a measurement-consistent correction term at each step, enabling integration of observational data and thus unifying generative modeling with data assimilation.

2. (non-autoregressive vs. autoregressive diffusion models) We will revise the sentence in the related work section to reflect this more balanced view:

Additionally, non-autoregressive diffusion models may struggle with long sequence forecasting in high-dimensional systems, though autoregressive diffusion approaches, like GenCast [3], have shown promising results in such settings.

F. Typographical and notational errors

We corrected the notation errors. These edits will be included in the final version.

 

Questions

A. “FlowDAS provides clearer physical insight into one-step dynamics?”

We agree that the phrase 'provides clearer physical insight' could imply an unintended comparison. Our intention was simply to state that our method offers insights that are physically clear. To better reflect this, we will revise the phrasing to 'provides clear physical insight' and ensure this meaning is unambiguous.

B. “For the SDA baseline, … a few steps of Langevin sampling …, is this something the authors have also employed in the experiments?”

Yes, we apply this predictor-corrector (PC) sampling settings proposed in the SDA baseline. Specifically, the predictor step is 512 and the corrector step is 3. As we empirically find, these settings show a good balance between the good performance and computation burden. The settings are kept across the Lorenz, NS, SEVIR experiments.

 

Limitations:

A. “Inclusion of several previous time steps -> large memory requirements”

We agree that memory usage can grow with the number of input states, especially in high-dimensional settings such as medium-range weather forecasting. In fact, in our ongoing collaboration with an industry partner, we are applying FlowDAS to real-world meteorological data with shape (71, 181, 360), where 71 is the number of channels. This indeed is significantly higher than in our current Navier–Stokes and SEVIR experiments.

To address this, a promising direction we are actively exploring is lifting FlowDAS into a latent space, inspired by works such as SLAMS [6], which can substantially reduce computational and memory overhead while preserving dynamics. We will add this as a concrete point in the discussion of future work.

B. “Costly conditioning if the likelihood is costly to evaluate”

We agree with the reviewer. In our current experiments, the likelihood evaluation is relatively cheap, so this is not a bottleneck. For settings where the likelihood is expensive to compute, developing approximated likelihood evaluations would be a promising direction for future work.

 

Reference

[5] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. arXiv preprint arXiv:2305.04391, 2023.

[6] ​​Qu, Yongquan, et al. "Deep generative data assimilation in multimodal setting." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

We thank the reviewer for the thoughtful and constructive discussions, these have significantly contributed to improving the quality and clarity of our work.

A. “Baseline configs.” We will make sure the baselines are clearly identified in the revised manuscript to avoid confusion.

B. “Additional details on computational cost.”

1. Empirical results with fewer sampling steps We experimented with reduced sampling steps from 1000 down to 50 on SEVIR. Results are summarized below:

While N = 200–400 achieves similar RMSE, we adopt N = 500 to maintain sampling diversity, which is essential for detecting rare or extreme events in weather assimilation tasks.

2. Why we use an SDE-based sampler instead of an EDM (ODE-based) sampler

SDE and ODE samplers offer different trade-offs in generative modeling. ODE-based samplers (e.g., EDM) are deterministic and typically more efficient, requiring fewer steps for high-quality synthesis. SDE-based samplers introduce stochasticity during inference, which is beneficial for capturing uncertainty and producing diverse samples; this is especially important for data assimilation under fixed observations. Currently, FlowDAS is built upon the stochastic interpolants framework, which naturally defines a reverse-time SDE for sampling. We acknowledge the practical advantages of ODE-based sampling and appreciate the suggestion. We have not explored this direction in FlowDAS yet, but investigating efficient ODE solvers, hybrid methods, or accelerated SDE schemes is a valuable future direction.

C. “FlowDAS provides clear physical insights.”

For the question on physical insights: FlowDAS directly models transitions between two valid physical states $x\_{t-1}$ to $x\_t$, so the learned velocity field can be interpreted as a data-driven approximation of the underlying dynamics over one time step, where both ends stay on the data manifold. The shorter distance between data states (compared to distance between noise and data manifold for diffusion model based frameworks) also leads to faster convergence during training, making it easier to capture underlying physics.

Regarding physical scenarios where diffusion-based frameworks may struggle: in our double-well experiment (Appendix D.3), a multi-modal system with a nonlinear observation model, noise-initialized optimization led to unstable gradients and recovery failure. In contrast, FlowDAS starts from a valid state, improving stability and robustness in such settings.

First of all, we sincerely thank the reviewer for their valuable feedback and constructive suggestions, which have helped us improve the clarity, accuracy, and completeness of the paper!

 

A. “The contribution appears somewhat incremental.”

Response:

We would like to highlight the novelty and significance of our contributions. FlowDAS is not a simple application of existing stochastic interpolants or diffusion models, but a principled and technically novel framework for solving the scientifically important problem of data assimilation (DA)—which differs fundamentally from pure forecasting.

• Data assimilation is harder and more scientifically meaningful than pure forecasting. Unlike forecasting, which predicts future states given the present, data assimilation requires reconstructing a latent trajectory given partial and noisy observations over time. This is critical in real-world systems such as weather prediction and oceanography, etc. It demands algorithms that are measurement-consistent, temporally coherent, and robust under sparse observations, that FlowDAS directly addresses.

• Beyond existing stochastic interpolant (SI) methods. Prior SI methods (e.g., [31]) focus exclusively on probabilistic forecasting and cannot incorporate observations. FlowDAS makes the first extension of SI to the DA setting by introducing:
  (1) A likelihood-based drift correction (Eq. 7) that enables measurement conditioning during rollout;
  (2) An unbiased estimator for the gradient of the observation likelihood (Eq. 8); and
  (3) A Monte Carlo sampling accelerator using stochastic integrators (Eqs. 9–10) for efficient inference. (4) A flexible backbone architecture, where using neural operators (e.g., FNO) instead of UNet yields even better results (see Table S.11, to be moved into main text).

  These advances together enable observation-conditioned generation, which has not been addressed in prior SI literature.

• Advantages over diffusion-based approaches. FlowDAS offers key improvements over diffusion-based methods like SDA and PDEDiff [1]. SDA models the dynamics via a joint score over entire trajectories, which limits temporal resolution and interpretability. PDEDiff improves on this by using conditional scores for state transitions, making it closer in spirit to FlowDAS. However, both methods require sampling from Gaussian noise, which leads to instability (see our toy double-well experiment in Appendix D.3), reduced physical interpretability, and degraded performance as shown in all experiments. In contrast, FlowDAS performs direct state-to-state transitions with an explicit transition density $p(x\_{k+1} \mid x\_k, y\_{k+1})$, enabling more stable, efficient, and physically grounded inference.

Together, these contributions establish FlowDAS as the first generative data assimilation framework based on stochastic interpolants. It represents a substantial and necessary step forward for building physically grounded, observation-aware generative models for scientific applications.

 

B. “The inference procedure is not presented very clearly.”

Thank you for this valuable feedback. We agree that a more detailed explanation of the inference procedure would enhance the clarity of the main paper. While we had placed a full, detailed breakdown in Algorithm 2 (Appendix B) due to initial space limitations, we will incorporate it directly and insert the following description into Section 3.2 in the revision:

Inference Given a trained drift model $\hat{b}_s$, FlowDAS performs inference in an autoregressive manner to estimate the latent state trajectory $\{ \hat{x}\_{1:K} \}$. At each time step $k$, we begin from the known state $\hat{x}_k$ and predict the next state $\hat{x}\_{k+1}$ using a discretized stochastic interpolant over a time grid $s_0 = 0 < s_1 < \cdots < s_N = 1$. We set the interpolant state $X\_{s\_0} \gets \hat{x}\_k$ and retrieve the next observation $y\_{k+1}$. For each interpolation step $s\_n$, we simulate a forward transition using the learned drift, given in Eq.7, and noise. To enforce observation consistency, we generate $J$ posterior endpoint samples $\hat{X}\_1^{(j)}$, and compute their likelihoods under the measurement model, i.e., $|| y - \mathcal{A}(\hat{X}\_1^{(j)}) ||_2^2$. A softmax over these values gives importance weights $\{ w\_{1:J} \}$. We then apply a correction to the interpolant using a weighted gradient $- \zeta\_n \nabla\_{X\_{s\_n}} \sum \_{j=1}^J w\_j || y - \mathcal{A}(\hat{X}\_1^{(j)}) ||_2^2$, where $\zeta\_n$ denotes the step size. After reaching $s=1$, we set $\hat{x}\_{k+1} \gets X\_{s\_N}$ and proceed to the next time step. Repeating this procedure over all $k$ yields the full trajectory.

 

C. “The presentation focuses heavily on these adaptations, rather than a self-contained method.”

We would like to hightlight that data assimilation naturally consists of two essential components: a probabilistic forecasting model and an observation-conditioned correction mechanism. Prior work such as [31] addresses only the forecasting part. To clearly situate our contribution, we chose to introduce that work early on, but this does not mean FlowDAS is not self-contained.

FlowDAS is a complete generative DA method. In the final version, we will revise Section 3 to present FlowDAS more coherently as a standalone algorithm, and move contextual references to prior work into a separate discussion in the Section 2 to improve readability.

 

D. “The inconsistent use of uppercase and lowercase symbols in the notation makes the derivations more difficult to follow.”

Our notation intentionally distinguishes between random variables and their realizations:

• Uppercase letters (e.g., $X, I, R$) denote random variables or stochastic processes.
• Lowercase letters (e.g., $x, y$) represent specific realizations or observations.

For example, $X\_0, X\_1 \sim p(X\_0, X\_1)$ are random variables, while $x\_0, x\_1$ are concrete samples. This convention is common in probabilistic modeling and used consistently throughout. In the final version, we will add an explicit note at the start of Section 2 to make this distinction clear to readers.

 

E. "The paper would benefit from a more thorough discussion of computational efficiency. Runtime, scalability, and comparisons to baseline methods are not sufficiently addressed."

We agree with the reviewer on the importance of computational efficiency and have added a detailed performance comparison in the appendix. The results, using a single A100 GPU on the SEVIR dataset, highlight FlowDAS’s practical viability:

The slightly longer inference time for FlowDAS stems from its iterative correction process, which is essential for achieving accurate and observation-consistent predictions. While this leads to higher per-sample inference time, it is balanced by significant training efficiency, as FlowDAS reaches convergence much faster than all competing baselines. We believe this reflects a favorable trade-off in practical deployment scenarios. A more detailed discussion will be included in the final version.

To further improve inference speed, one promising direction is to apply post-hoc distillation, where a faster surrogate model is trained to approximate the FlowDAS sampler. Another avenue is to explore hybrid formulations that integrate variational inference [2] (as suggested by Reviewer 3gVf) techniques into the framework, potentially offering both computational gains and improved flexibility.

 

Reference

[1] Shysheya, A., Diaconu, C., Bergamin, F., Perdikaris, P., Hernández-Lobato, J. M., Turner, R., & Mathieu, E. (2024). On conditional diffusion models for PDE simulations. Advances in Neural Information Processing Systems, 37, 23246-23300.

[2] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. arXiv preprint arXiv:2305.04391, 2023.

Dear AC,

Thank you for your timely reminder and the efforts for coordinating the review process.

We sincerely appreciate the reviewers’ time and effort in considering our rebuttal and providing their valuable feedback. During the rebuttal period, we did our best to address the raised concerns and questions.

Now, we are working on addressing reviewers' follow-up questions and will respond shortly. Moreover, we are eager to clarify any remaining questions and address any concerns they may have.

Best regards,

The Authors

Dear reviewers,

The reviewers-authors discussion phase will end in less than 24 hours.

If not done already, make sure to submit the "Mandatory Acknowledgement" that confirms that you have read the reviews, participated in the discussion, and provided final feedback in the "Final justification" text box.

Be mindful of the time and efforts the authors have invested in answering your questions and at least acknowledge their responses. Make sure to provide a fair and scientifically grounded review and score. If you have changed your mind about the paper, please update your review and score accordingly. If you have not changed your mind, please provide a clear and sound justification for your final review and score.

Best regards, The AC