Score-based Data Assimilation
Authors: François Rozet
Gilles Louppe
Venue: NeurIPS 2023 poster (2023)
Keywords: data assimilation
score-based
generative modeling
posterior inference
dynamical systems

Abstract: Data assimilation, in its most comprehensive form, addresses the Bayesian inverse problem of identifying plausible state trajectories that explain noisy or incomplete observations of stochastic dynamical systems. Various approaches have been proposed to solve this problem, including particle-based and variational methods. However, most algorithms depend on the transition dynamics for inference, which becomes intractable for long time horizons or for high-dimensional systems with complex dynamics, such as oceans or atmospheres. In this work, we introduce score-based data assimilation for trajectory inference. We learn a score-based generative model of state trajectories based on the key insight that the score of an arbitrarily long trajectory can be decomposed into a series of scores over short segments. After training, inference is carried out using the score model, in a non-autoregressive manner by generating all states simultaneously. Quite distinctively, we decouple the observation model from the training procedure and use it only at inference to guide the generative process, which enables a wide range of zero-shot observation scenarios. We present theoretical and empirical evidence supporting the effectiveness of our method.

Score-based Data Assimilation
François Rozet University of Liège francois. francois. rozet@ulieeeeeee
Gilles Louppe University of Liège
louppe@uliege be
Abstract
Data assimilation, in its most comprehensive form, addresses the Bayesian inverse problem of identifying plausible state trajectories that explain noisy incomplete observations of stochastic dynamical systems. Various approaches have been proposed to solve this problem, including particle-based and variational methods. However, most algorithms depend on the transition dynamics for inference, which becomes intractable for long time horizons for high-dimensional systems with complex dynamics, such as oceans or atmospheres. In this work, we introduce score-based data assimilation for trajectory inference. We learn score-based generative model of state trajectories based on the key insight that the score of an arbitrarily long trajectory can be decomposed into a series of scores over short segments. After training, inference is carried out using the score model, in a non- autoregressive manner by generating all states simultaneously. Quite distinctively, we decouple the observation model from the training procedure and use it only at inference to guide the generative process, which enables wide range of zero-shot observation scenarios. We present theoretical and empirical evidence supporting the effectiveness of our method.
1 Introduction
Data assimilation (DA) [1-9] is the core of many scientific domains concerned with the study of complex dynamical systems such as atmospheres, oceans or climates. The purpose of DA is to infer the state of a system evolving over time based on various sources of imperfect information, including sparse, intermittent, and noisy observations.
Formally, let Fi.L (x ..,....,,t)) E LLXD denote trajectory of states in discrete-time stochastic dynamical system and p(z+41 I;) be the transition dynamics from state I; to state T++1- An observation RM of the state trajectory Ft.t follows an observation process p(y x1,t), generally formulated as A(t1-L) + 77, where the measurement function A : RR + R& is often non-linear and the observational error 17 E R is stochastic additive term that accounts for instrumental noise and systematic uncertainties. In this framework, the goal of DA is to solve the inverse problem of inferring plausible trajectories F1.L given an observation y, that is, to estimate the trajectory posterior
p(1L. (8) = ply FILL) p(v)
L-1
p(x1) II p( xx)
(I)
1=1
where the initial state prior p(r1) is commonly referred to as background [5-9]. In geosciences, the amount of data available is generally insufficient to recover the full state of the system from the observation alone [8]. For this reason, the physical model underlying the transition dynamics is of paramount importance fill in spatial and temporal gaps in the observation.
State-of-the-arr approaches to data assimilation are based on variational assimilation [1, 2, 5-7J. Many of these approaches formulate the task as maximum-a-posterion (MAP) estimation problem and solve it by maximizing the rog-posterior density log p(11:L y) via gradient ascent. Although this
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

approach only produces a point estimate of the trajectory posterior, its cost can already be substantial for problems of the size and complexity of geophysical systems, since it requires differentiating through the physical model. The amount of data that can be assimilated is therefore restricted because of computational limitations. For example, only small volume of the available satellite data is exploited for operational forecasts and yet, even with these restrictions, data assimilation accounts for a significant fraction of the computational cost for modern numerical weather prediction [10, 11j. Recent work has shown that deep learning can be used in variety of ways to improve the computational efficiency of data assimilation, increase the reconstruction performance by estimating unresolved scales after data assimilation, integrate multiple sources of observations [12-19].
Contributions this work, we propose a novel approach to data assimilation based on score- based generative models. Leveraging the Markovian structure of dynamical systems, we train score network from short segments of trajectories which is then capable of generating physically consistent and arbitrarily-long state trajectories. The observation model is decoupled from the score network and used only during assimilation to guide the generative process, which allows for a wide range of zero-shot observation scenarios. Our approach provides an accurate approximation of the whole trajectory posterior is not limited to point estimates - without simulating or differentiating through the physical model. The code for all experiments is made available https:// //gi thub. com/ wwwaaaannnoooeeeeeeeeeeeee
2 Background
Score-based generative models have recently shown remarkable capabilities, powering many of the latest advances in image, video or audio generation [20-27]. In this section, we review score-based generative models and outline how they can be used for solving inverse problems.
Continuous-timg score-based generative models Adapting the formulation of Song et al. [28], samples x E IR_ from a distribution p(x) are progressively perturbed through continuous-time diffusion process expressed as a linear stochastic differential equation (SDE)
dx (t) f(t) xtt) dt + dw(()
where f(t) E R is the drift coefficient, g(t) E IR is the diffusion coefficient, w(t) E IR_ denotes a Wiener process (standard Brownian motion) and r=(t) E R. is the perturbed sample at time E [o. 1]. Because the SDE is linear with respect to rrf), the perturbation kernel from x to rr(t) is Gaussian and takes the form
p(=(() (x) = (a((() p()) x. Σ())
where u(t) and E(t) = (t) can be derived analytically from f(t) and g()) [29, 30). Denoting p(z(t)) the marginal distribution of z(t), we impose that ((0) 1 and (0)) < 1, such that p(z(0)) ss p(r), and we chose the coefficients f(t) and g(t) such that the influence the initial sample on the final perturbed sample sr(1) (1) is negligible with respect the noise level that is, p(x(1)) s (0, ((1)). The variance exploding (VE) and variance preserving (VP) SDEs [28, 31, 32] widespread examples satisfying these constraints.
Crucially, the time reversal of the forward SDE (2) is given by a reverse SDE [28, 33]
dx(( [5(() ,()² Vato logp(x(())) dt + g(t) dw((t.
(4)
That is, we can draw noise samples æ(1) ~ (O. E(1) and gradually remove the noise therein to obtain data samples a0) p(x(00) by simulating the reverse SDE from t 0. This requires access to the quantity Va(e) log p(x(t)) known as the score of p(=())).
Denoising score matching In practice, the score Va(e) log p(x(t) is approximated by a neural network 8 (x(t), t), named the score network, which trained to solve the denoising score matching objective [28, 35]
arg min E Ep(e)tmemat m) ((t))[|so(a(()) - Va(o) logp(x(() | (x) 1)]
(5)
where p(t) u000 1). The theory of denoising score matching ensures that 8 (1)) R Vu(1) log pot)(z()) for sufficiently expressive score network. After training, the score network

---

is plugged into the reverse SDE (4), which is then simulated using an appropriate discretization scheme [28, 30, 36, 37].
In practice, the high variance Ve(() log p(r(() x) near = 0 makes the optimization (5) unsta-
ble [30]. To mitigate this issue, a slightly different parameterization es (z((),, t) = -o(t) s (zx()), t) of the score network is often used, which leads to the otherwise equivalent objective [30, 32, 36]
[llc@(u(t) + et)) <ll²]
arg min B (p(r)((p(())
(6)
where p(c) (O, I). In the following, we keep the score network notation s (2))), t) for conve- nience, even though we adopt the parameterization es (z((),, t) and its objective for our experiments.
Zero-shot inverse problems With score-based generative models, we can generate samples from the unconditional distribution p(=(0)) p(z). To solve inverse problems, however, need to sample from the posterior distribution p(r y). This could be accomplished by training a conditional score network s (ar()), & 1y) approximate the posterior score Vz(0) log p(x(() y) and plugging
into the reverse SDE (4). However, this would require data pairs (z. y) during training and one would need to retrain a new score network each time the observation process p(y x) changes. Instead, many have observed [28, 38 38-41] that the posterior score be decomposed into two terms thanks to Bayes' rule
Va(sy log p(x(() v) = Vato logp()(()) + Va(1) logp(y x<()).
Since the prior score Vacee p({(()(() can be approximated with single score network, the remaining task is to estimate the likelihood score Vz(e) log (x))). Assuming a differentiable measurement
function A and Gaussian observation process p(y x) = -N A(x), Ey,, Chung et al. [41] propose the approximation
ply xx(t) f ply x) x(t)) dr = N A()(=(t)), E.y)
(8)
where the mean =(x(t)) Ex(ea(t))) [xc] is given by Tweedie's formula [42, 43]
E [p(+)(())=]]= sx(t) + (1)2 Va(e) log p(x(())
(9)
m(t)
xx1) + o(1 (x((,,)
(10)
(t))
As the log-likelihood of a multivariate Gaussian is known analytically and (s(x))), is differentiable, we can compute the likelihood score Va(1) log xt))) with this approximation in zero-shot, that is, without training any other network than (t)
3 Score-based data assimilation
Coming back to our initial inference problem, we want to approximate the trajectory posterior p(Z1.1 y) of a dynamical system. To do so with score-based generative modeling, we need to estimate the posterior score V21 log p( y). which we choose to decompose into prior and likelihood terms, as in (7), to enable a wide range of zero-shot observation scenarios.
In typical data assimilation settings, the high-dimensionalliy of each state x; (e.g. the state of atmo- spheres or occans) combined with potentially long trajectories would require an impractically large score network sy (21.2 (t), t) to estimate the prior score Vz, (t) log 1( (t)) and proportional amount of data for training, which could be prohibitive if data scarce if the physical model is expensive to simulate. To overcome this challenge, we leverage the Markovian structure of dynamical systems to approximate the prior score with series of local scores, which are easier to learn, as explained in Section 3.1. In Section 3.2, we build upon diffusion posterior sampling (DPS) [41] to propose a new approximation for the likelihood score VeL.2()) log ply (1)), which we find more appropriate for posterior inference. Finally, in Section 3.3, we describe our sampling procedure inspired from predictor-corrector sampling [28]. Our main contribution, named score- score-based data assimilation (SDA), is the combination of these three components.
3

---

Forward (data - noise)
Figure 1: Trajectories T1.L of dynamical sys- tem are transformed to noise via a diffusion pro- Reversing this pro- cess generates new tra- jectories, but requires the score of (.).()))), We approximate it by com- bining the outputs of a score network over sub- segments of (1:2(t).
dr(() = /(0))(())t + a(()do())
#(z((1)
21-1001(0) Ve Vennt (opplruztt) te-ene
#1.1(0)
da(()= [[((())) =(1) Vaces (o a(0dw(()
#1.((f)
Reverse SDE fuoise + data)
3.1
How is your blanket?
Given a set of random variables 31.7. {11,, (2,..,,,1,, it is sometimes possible to find small Markov blanket Tb; S Veg such that p(r; | 5z41) = p(a; | vs,) for each element 2; using our knowledge of the set's structure. follows that each element V,, log logp(z).)) of the full score Vanz log p(#1.L) can be determined locally, that is, only using its blanket;
Va, log ((z)()))) = V., log p(z; m4s)) + Vx, log V., log (b,)) V., log logg((>) = Vz; log p(g, arb.)-
(1I)
(12)
This property generally does not hold for the diffusion pnfftusionn-pruurhee set 2112 there is guarantee that ns, (t) is a Markov blanket of the element (t). However, there exists set of indices b; 2 b; such that
Var(e) logp(z)n R Vaa()) (og 27, (1)
(13)
good approximation for all E [0. That is, 27, (t) is a "pseudo" Markov blanket of (t).
In the worst case, b, contains all indices except i, but we argue that, for some structures, there is set b, not much larger than b, that satisfies (13). Our rationale is that, since we impose the initial noise to be negligible, we know that Is, (t) becomes indistinguishable from I6, as t approaches O. Furthermore, as f grows and noise accumulates, the mutual information between elements x; (t) and x, (t) decreases to finally reach 0 when = Hence, even if b; b; the pseudo-blanket approximation (13) already holds near t = 0 and = 1. In between, even though the approximation remains unbiased (see Appendix A), the structure of the set becomes decisive. If it is known and present enough regularities/symeetrr the score network &o (Z (t), t).
(13) could and should be exploited within the architecture of
In the case of dynamical systems, the set Fı.L is by definition first-order Markov chain and the minimal Markov blanket of an element x; Fb. (24-1,7741)) .. For the perturbed element (t). the pscudo-blanket 2z, (t) can take the form of window surrounding 2t). that is b, {i -i-Ke .,,...,................ + k} (i} with & 2 1. The value of is dependent on the problem, but we argue, supported by our experiments, that it is generally much smaller than the chain's length L. Hence, fully convolutional neural network (FCNN) with a narrow receptive field is well suited to the task, and any long-range capabilities would be wasted resources. Importantly, if the receptive field is 2k + 1, the network can be trained on segments x -isks instead of the full chain #1:L. thereby drastically reducing training costs. More generally,
train a local score network (see Algorithm 1)
so( & (t) log (og((xx--4((()
(14)
such that its + 1-th element approximates the score of the i-th state Vay(t) log ((1.. (t)..(t)). We also have that the k first elements of (= (1), approximate the score of the k first states Vay log p( (f)) and the k last elements of (= (t), f) approximate the score of the k last states Vzz-xe(t) (t) log (t). Hence, we can apply the local score network on all sub-segments I -i-kis+k & (t) of T1.1(t), similar to a convolution kernel, and combine the outputs (see Algorithm 2) to get approximation of the full score Vzy.tte) (t) log p(z1.1(t)), (t)). Note that can either condition the score network with or assume the statistical stationarity of the chain, that is p(x+) p(z+11).

---

Algorithm Training Eo (a boi
Algorithm Composing (T (t),
for = 1 to N do
1 function s (1).
F1.L p(11:1)
$1:k+1 f- 86 (11:2++++1 (t), - do
i~ U(fk + I....,, 1:))
3
for = k
t~ t~ U(, 1), 1),
4
8; 5 (z ki (t), 8) +
5
Ti-Rsith(t) + (z + (t)e
5
S1.k:L f- (() 1)(k + +1:]
6
Et |eo(Fi (1),) -
6
return 81-L
Φ GRADIENTIDESCENT (Ø, V
3.2 Stable likelihood score
Due to approximation and numerical emrors in (ar(()). computing the score Vz(0) logp(y x(()) with the likelihood approximation (8) is very unstable, especially in the low signal-to-noise regime, that is when (l) > ((t). This incites Chung et al. [41] to replace the covariance Ey by the identity I and rescale the likelihood score with respect to lVy - A(r(x(t))) to stabilize the sampling process. These modifications introduce significant error in the approximation as they greatly affect the norm of the likelihood score.
We argue that the instability is due (8) being only exact if the variance of p(x xx()) is null or negligible, which is not the case when t > 0. Instead, Adam et al. [40] and Meng et al. [44] approximate the covariance of p(z (t))) with ((0)/(tt), which is valid as long the prior p(æ) is Gaussian with a large diagonal covariance Σ,. We motivate in Appendix B the more general covariance approximation o /u where the matrix I depends the eigendecomposition of Σz,
Then, taking inspiration from the extended Kalman filter, we approximate the perturbed likelihood as
(1 (1IA((((())) M, + a(t)² m(()² AL
ply x(t)) 2 NN(
(15)
where A= 0z A la(a(e)) is the Jacobian of A. In practice, to simplify the approximation, the term AT can often be replaced by constant (diagonal) matrix. We find that computing the likelihood score Ve((t) log ply =z()) with this new approximation (see Algorithm 3) is stable enough that rescaling it or ignoring Mg is unnecessary.
3.3 Predictor-Correccor sampling
To simulate the reverse SDE, we adopt the exponential integrator (EI) discretization scheme introduced by Zhang et al. [30]
x(t At) Mlt - At) x(t)) + p(t)
At) ((())
olt a(t- _ o(t)
A1) (()))
(16)
which coincides with the deterministic DDIM [36] sampling algorithm when the variance preserving SDE [32] is used. However, as we approximate both the prior and likelihood scores, errors accumulate along the simulation and cause it to diverge, leading low-quality samples. To prevent errors from accumulating, we perform (see Algorithm 4) few steps of Langevin Monte Carlo (LMC) [45, 46]
x(t) + xxt) (z)((())) V2
(17)
where c ~ (O. I), between each step of the discretized reverse SDE (16). In the limit of an infinite number of LMC steps with a sufficiently small step size 8 E R4, simulated samples are guaranteed to follow the distribution implicitly defined by our approximation of the posterior score at each time meaning that the errors introduced by the pseudo-blanket (13) and likelihood (15) approximations do not accumulate. In practice, we find that few LMC steps are necessary. Song et al. [28] introduced a similar strategy, named predietor-corrector (PC) sampling, to correct the errors introduced by the discretization of the reverse SDE.
5

---

Results
We demonstrate the effectiveness of score-based data assimilation on two chaotic dynamical systems: the Lorenz 1963 [47] and Kolmogorov flow [48] systems. The former is simplified mathematical model for atmospheric convection. Its low dimensionality enables posterior inference using classical sequential Monte Carlo methods [49, 50] such as the bootstrap particle filter [51]. This allows us to compare objectively our posterior approximations against the ground-truth posterior. The second system considers the state of a two-dimensional turbulent fluid subject to Kolmogorov forcing [48]. The evolution of the fluid is modeled by the Navier-Stokes equations, the same equations that underlie the models of oceans and atmospheres. This task provides good understanding of how SDA would perform typical data assimilation applications, although our analysis is primarily qualitative due to the unavailability of reliable assessment tools for systems of this scale.
For both systems, we employ as diffusion process the variance preserving SDE with a cosine schedule [52], that is p(() = cos(wi)² with w = arccos v10-3 and (1) = / =(()². The score networks are trained once and then evaluated under various observation scenarios. Unless specified otherwise, we estimate the posterior score according to Algorithm with 10 20-21 and simulate the reverse
SDE (4) according to Algorithm 4 in 256 evenly spaced discretization steps.
4.1 Lorenz 1963
The state x (a, b, c) E R$ of the Lorenz system evolves according to a system of ordinary differential equations
à = (b a) b = ap c) - ab Bc
(18)
where a = 10, p= 28 and 3= - are parameters for which the system exhibits a chaotic behavior. We denote ā and the standardized (zero mean and unit variance) versions of and c, respectively. As our approach assumes discrete-timg time stochastic dynamical system, we consider transition process of the form 2++1 M(x))= n. where M R33 + R3 is the integration of the differential equations (18) for A = 0.025 time units and 17 ~ No. AT) represents Brownian noise.
16
Corrections Corrections
a
15
16
DPF
-15
B
200
25
50
33
65
97
129
161
190
(0g2( )7))
log Ima.)
W,
Figure 2: Average posterior summary statistics over 64 observations from the low (A) and high (B) frequency observation processes. We observe that, as l and the number of corrections C increase, the statistics of the approximate posteriors get closer to the ground-trute truth, in red, which means they are getting more accurate. However, increasing and C improves the quality of posteriors with decreasing return, such that all posteriors with & > 3 and C 2 are almost equivalent. This is visible in C, where we display trajectories inferred (C 2) for observation of the low frequency observation process. For readability, we allocate segment of 32 states to each & instead of overlapping all 192 states. Note that the Wasserstein distance between the ground-truth posterior and itself zero as it is estimated with finite number (1024) of samples.

---

45
37
33
25
17
13
-10 0 10
-15
15
20 30
-1
0
(v)
Figure 3: Example of multi-modal posterior inference with SDA. We identify four modes (dashed lines) in the inferred posterior. All modes are consistent with the observation (red crosses), as demonstrated by the posterior predictive distribution q(y/ y) = Egganall (p(p [1:L)]].
We generate 1024 independent trajectories of 1024 states, which split into training (80 %),, validation (10 %) and evaluation (10 %) sets. The initial states are drawn from the statistically stationary regime of the system. We consider two score network architectures: fully-connected local score networks for small l (k I 4) and fully-convolutional score networks for large k. Architecture and training details for each l are provided in Appendix D.
We first study the impact of k (see Section 3.1) and the number LMC corrections (see Section 3.3) on the quality of the inferred posterior. We consider two simple observation processes N( (vr 41.L.,, 0.05²1) and Nw fI.L. 0.25"I). The former observes the state at low frequency (every eighth step) with low noise, while the latter observes the state high frequency (every step) with high noise. For both processes, we generate an observation y for a trajectory of the evaluation set (truncated at L 65) and apply the bootstrap particle filter (BPF) draw 1024 trajectories T1.L from the ground-truth posterior p(æ1. y)- We use a large number of particles (210) to ensure convergence. Then, using SDA, sample 1024 trajectories from the approximate posterior ((11.1 y) defined by each score network. We compare the approximate and ground-truth posteriors with three summary statisties: the expected log-prior Eg(zr.cl)) [log p(12: L x )], the expected log- likelihood E, v) [log p(y æ1:L)]] and the Wasserstein distance W1 WI(P. 9) in trajectory space. We repeat the procedure for 64 observations and different number of corrections (r = 0.25, see Algorithm 4) and present the results in Figure 2. To paraphrase, SDA is able reproduce the ground-truth posterior accurately. Interestingly, accuracy can be traded off for computational efficiency: fewer corrections leads to faster inference at the potential expense of physical consistency.
Another advantage of SDA over variational data assimilation approaches is that it targets the whole posterior distribution instead of point estimates, which allows to identify when several scenarios are plausible. As a demonstration, we generate an observation from the observation process p(y II.L) = Niy Č1.L:4, 0.1²I) and infer plausible trajectories with SDA (k 4, C 2). Several modes are identified in the posterior, which we illustrate in Figure 3.
4.2 Kolmogorov flow
Incompressible fluid dynamics are governed by the Navier-Stokes equations
t = u Vu + v 1 Ap+ff P
Re
(19)
u
where is the velocity field, Re is the Reynolds number, is the fluid density, is the pressure field and ff is the external forcing. Following Kochkov al. [53], we choose -vo-dimensional domain [0. |x with periodic boundary conditions, large Reynolds number Re = 10 a constant density = and an external forcing f corresponding to Kolmogorov forcing with linear damping [48, 54]. We use the jax-cfd library [53] to solve the Navier-Stokes equations (19) on 256 x 256 domain grid. The states 2; are snapshots of the velocity field u. coarsened a 64 x resolution, and the integration time between two such snapshots is A 0.2 time units. This corresponds to 82
7

---

31::144
SDA
DPS
Figure 4: Example of sampled trajectory from coarse, intermittent and noisy observations. States visualized by their vorticity field w V x u, that is the curl the velocity field. Positive values (red) indicate clockwise rotation and negative values (blue) indicate counter-clockwise rotation. SDA closely recovers the original trajectory, despite the limited amount of available data. Replacing SDA's likelihood score approximation with the of DPS [41] yields trajectories inconsistent with the observation.
integration steps of the forward Euler method, which would be expensive to differentiate through repeatedly, required by gradient-based data assimilation approaches.
We generate 1024 independent trajectories of states, which are split into training (80 %), validation (10 %) and evaluation (10 %) sets. The initial states are drawn from the statistically stationary regime of the system. We consider a local score network with k = 2. As states take the form of 64 64 images with two velocity channels, we use U-Net [55] inspired network architecture. Architecture and training details are provided in Appendix D.
We first apply SDA to a classic data assimilation problem. We take trajectory of length L = 32 from the evaluation set and observe the velocity field every four steps, coarsened resolution 8 x 8 and perturbed by moderate Gaussian noise (Eg 1? Given the observation, we sample a trajectory with SDA (C illustrated in Figure 4. A similar experiment where we modify the amount of spatial information is presented in Figure 8. When data is insufficient to identify the original trajectory, SDA extrapolates physically plausible scenario while remaining consistent with the observation, which can also be observed in Figure and
0.5) and find that it closely recovers the original trajectory, as
SDA
A(æl) +
21
Physical model
Figure 5: A trajectory consistent with an unlikely observation of the final state is generated with SDA. To verify whether the trajectory is realistic and not hallucinated, we plug its initial state Φ1 into the physical model and obtain an almost identical trajectory. This indicates that SDA is not simply interpolating between observations, but rather propagates information in a manner consistent with the physical model, even in unlikely scenarios.

---

Finally, we investigate whether SDA generalizes to unlikely scenarios. We design an observation process that probes the vorticity of the final state 2r, in a circle-shaped sub-domain. Then, we sample a trajectory (C'=1) 1, 0.5) consistent with uniform positive vorticity observation in this sub-domain, which is unlikely, but not impossible. The result is discussed in Figure 5.
Conclusion
Impact In addition to its contributions to the field of data assimilation, this work presents new technical contributions to the field of score-based generative modeling.
First, we provide new insights on how to exploit conditional independences (Markov blankets) in sets random variables to build and train score-based generative models. Based on these findings, we are able to generate/infer simultaneously all the states of arbitrarily long Markov chains L.L while only training score models on short segments Ti-ksi thereby reducing the training costs and the amounts of training data required. The decomposition of the global score into local scores additionally allows for better parallelization at inference, which could be significant depending on available hardware. Importantly, the pscudo-blanker approximation (13) is not limited to Markov chains, but could be applicd to any set of variables 21.L, as long as some structure is known.
Second, we motivate and introduce a novel approximation (15) for the perturbed likelihood p(y (x)). when the likelihood p(y|æ)) is assumed (linear or non-linear) Gaussian. We find that computing the likelihood score Vact) log ((y])(())))))))((((())(())))(( with this new approximation leads to accurate posterior inference, without the need for stability tricks [41]. This contribution can be trivially adapted to many tasks such inpainting, deblurring, super-resolution or inverse problems in scientific fields [39-41].
Limitations From a computational perspective, even though SDA does not require simulating or differentiating through the physical model, inference remains limited by the speed of the simulation of the reverse SDE, Accelerating sampling in score-based generative models is an active area of research [30, 36, 37, 56] with promising results which would be worth exploring in the context of our method.
Regarding the quality of our results, we empirically demonstrate that SDA provides accurate ap- proximations the whole posterior, especially i and the number of LMC corrections C increase. However, our approximations (13) and (15) introduce certain degree of error, which precise im- pact on the resulting posterior remains to be theoretically quantified. Furthermore, although the Kolmogorov system is high-dimensional (tens of thousands of dimensions) with respect to what is approachable with classical posterior inference methods, it remains small comparison to the millions of dimensions of some operational DA systems. Whether SDA would scale well to such applications is an open question and will present serious engineering challenges.
Another limitation of our work is the assumption that the dynamics of the system are shared by all trajectories. In particular, a parametric physical model is used, all trajectories are assumed to share the same parameters. For this reason, SDA is not applicable to settings where fitting the model parameters also required, or at least not without further developments. Some approaches [57-60] tackle this task, but they remain limited to owwdiiensional settings. Additionally, if a physical model used to generate synthetic training data, instead of relying on real data, one can only expect SDA to be accurate as the model itself. This is limitation shared by any model-based approach and robust assimilation under model misspecification or distribution shift is left as an avenue for future research.
Finally, posterior inference over entire state trajectories is not always necessary. In forecasting tasks, inferring the current state of the dynamical system is sufficient and likely much less expensive. In this setting, data assimilation reduces to a state estimation problem for which classical methods such as the Kalman filter [61] or its nonlinear extensions [62, 63] provide strong baselines. Many deep learning approaches have also been proposed to bypass the physical model entirely and learn instead a generative model of plausible forecasts from past observations only [64-67].
Related work A number of previous studies have investigated the use of deep learning to improve the quality and efficiency of data assimilation. Mack et al. [12] use convolutional auto-encoders to project the variational data assimilation problem into lower-dimensional space, which simplifies the optimization problem greatly. Frerix et al. [14] use deep neural network to predict the initial state of trajectory given the observations. This prediction is then used as a starting point for traditional

---

(4D- variational data assimilation methods, which proves to be more effective than starting at random. This strategy also possible with SDA (using a trajectory sampled with SDA as a starting point) and could help cover multiple modes of the posterior distribution. Finally, Brajard et al. [17] address the problem of simultaneously learning the transition dynamics and estimating the trajectory, when only the observation process is known.
Beyond data assimilation, SDA closely relates to the broader category of sequence models, which have been studied extensively for various types of data, including text, audio, and video. The latest advances demonstrate that score- based--bbddd-- generative models achieve remarkable results on the most demanding tasks. Kong et al. [26] and Goel et al. [27] use score- based models to generate long audio sequences non-autoregressively Ho al. [23] train a score-based generative model for fixed number of video frames and use it autoregressively to generate videos of arbitrary lengths. Conversely, our approach non-autoregressive which allows to generate and condition all elements (frames) simultaneously. Interestingly, part of their method, Ho et al. [23] introduce "reconstruction guidance" for conditional sampling, which can be seen as special case of our likelihood approximation (15) where the observation is subset of z. Lastly, Ho et al. [25] generate low-frame rate, low-resolution videos which are then up-sampled temporally and spatially with cascade [24] of super-resolution diffusion models. The application of this approach to data assimilation could be worth exploring, although the introduction of arbitrary observation processes seems challenging.
Acknowledgments and Disclosure of Funding
François Rozet is a research fellow of the F.R.S.-ENRS (Belgium) and acknowledges its financial support.
References
[1] Andrew C Lorene. "Analysis methods for numerical weather prediction". In Quarterly Journal of the Royal Meteorological Sociery 1 112.474 (1986), pp. 1177-1194.
[2] François-Xann Xavier Le Dimet and Olivier Talagrand. "Variational algorithms for analysis and assimilation of metcorological observations: theoretical aspects". In Tellus A: Dynamic Meteo- rology and Oceanography 38.2 (1986), pp. 97-110.
[3] Geir Evensen. "Sequential data assimilation with a nonlinear guasi-geostrophic model using Monte Carlo methods forecast error statistics". In Journal of Geophysical Research: Oceans 99.C5 (1994), pp. 10143-10162.
[4] Thomas M. Hamill. "Ensemble- based Atmospheric Data Assimilation".. In Predictability of weather and climate 124 (2006), p. 156.
[5] Yannick Trémolet. "Accounting for imperfect model in 4D- Var." In ECMWF Technical Memoranda 477 (2005), p. 20. URL:: https: //www. ecmwf. httt/node/12844
[6] Yannick Trémolet. "Model error estimation in 4D-Var". In ECMWF Technical Memoranda
520 (2007), p. 19. URL: https://www/ ecmwf int watt//nn//////eeeeee 12819.
[7] Mike Fisher et al. Creetk-Constraintt and Long- Window 4D-Var". In ECMWF Technical Memoranda 655 (201 1), p. 47. URL: https: : /w/ww//.ww///oo ecmwf. int/ 9911/
[8] Alberto Carrassi et al. "Data assimilation the geosciences: An overview of methods, issues, and perspectives". In Wiley Interdisciplinary Reviews: Climate Change 9.5 (2018).
[9] ECMWF. "Part II: Data Assimilation". In IFS Documentation CY47RI 2 (2020). URL: https: / //www. ecmwf. int ht////dde/99//
[10] Peter Bauer et al. "The quiet revolution of numerical weather prediction". In Nature 525.7567
(2015), pp. 47-55. URL: https: h/ttts nature. com/ arti certicless/ee nature14955
[11] Nils Gustafsson et al. "Survey of data assimilation methods for convective-scale numerical weather prediction at operational centres". In Quarterly Journal of the Royal Meteorological Society 144.713 (2018), pp. 1218-1256.
10

---

[12] Julian Mack et al. "Attention-base convolutional autoencoders for 3d-variational data as- similation". Computer Methods in Applied Mechanics and Engineering 372 (2020). URL: https: //axivv./ orgooroaaa a0100101/1101..... 02121..
[13] Rossella Arcucci et al. "Deep data assimilation: integrating deep learning with data assim- ilation". In Applied Sciences 11.3 (2021), p. 1114, URL: https: http//wwwmmpttt com com/ 2076- 3417/11/3/111
[14] Thomas Frerix et al. "Variational data assimilation with a learned inverse observation operator". In International Conference on Machine Learning. 2021, pp. 3449-3458. URL.: https: ://
arxiv. org/ aba/aba//222/ 11192.
[15] Ronan Fablet et al. "Learning variational data assimilation models and solvers". In Journal of Advances in Modeling Earth Systems 13.10 (2021). URL: 12941.
[16] Yu-Hong Yeung al. "Physics-Informed Machine Learning Method for Large-Scale Data Assimilation Problems". Water Resources Reseurch 58.5 (2022). URL: https://arxiv. org/ ab//a/b//11.. 00037.
[17] Julien Brajard et al. "Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: case study with the Lorenz 96 model". Journal of computational science 44 (2020). URL: https // //txxivvvvv/////////////////aavl wor.oo////aaa. aat/a/bs/000000 01520.
[18] Julien Brajard et al. "Combining data assimilation and machine learning to infer unresolved scale parametrization". " Philosophical Transactions of the Royal Society A 379.2194 (2021). URL: https: : / aaxivviiv org/ abs/0099 .03118.
[19] Kai Zhang et al. "Multi-source information fused generative adversarial network model and data assimilation based history matching for reservoir with complex geologies". In Petroleum Science 19.2 (2022), pp. 707-719.
[20] Robin Rombach et al. "High-resolution image synthesis with latent diffusion models". In IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. URL: https
/ //arxiv..t org/ aa//a/bs//1111
[21] Aditya Ramesh al. "Hierarchical text-conditional image generation with clip latents". , In (2022). URL: https: http://axxvvvvooii wttp///////////////oo 06125.
[22] Chitwan Saharia et al. "Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding", In Advances in Neural Information Processing Systems. 2022. URL: https: / //openreview sett/forummid-08/1 n512A1.
[23] Jonathan Ho et al. "Video Diffusion Models". In ICLR Workshop on Deep Generative Mod- els for Highly Structured Data. 2022. URL: https / openreview net / forum rid= BBe1R2NdD25.
[24] Jonathan Ho et al. "Cascaded Diffusion Models for High Fidelity Image Generation". In Journal of Machine Learning Research 23.47 (2022), pp. 1-33. URL: http://jmlr org/
[25] Jonathan Ho et al. "Imagen video: High definition video generation with diffusion models". 2022. URL: https: aariiiiiiivv/aai///aaa//aaaai..m org/ abs/2210 .02303.
[26] Zhifeng Kong al. "DiffWave: A Versatile Diffusion Model for Audio Synthesis" In Inter- national Conference on Learning Representations. 2021. URL: //tpnreve///eevvieee net/ forum?id=a-xPPaa
[27] Karan Goel et al. "It's raw! audio generation with state-space models". In International Conference on Machine Learning. PMLR. 2022. URL: https://proceediii mlr. .lrrsress// v162/goel22
[28] Yang Song al. SScore-Based Generative Modeling through Stochastic Differential Equa- tions". In International Conference on Learning Representations. 2021. URL: https: /// openreview net/forumminn dearrnnidddPXIII122R
[29] Simo Särkkä and Arno Solin. "Applied stochastic differential equations". Vol. 10. Cambridge University Press, 2019.

---

[30] Qinsheng Zhang and Yongxin Chen. "Fast Sampling of Diffusion Models with Exponential Integrator". International Conference on Learning Representations. 2023. URL: https:
// //openreviee...
[31] Yang Song and Stefano Ermon. "Generative Modeling by Estimating Gradients of the Data Distribution". Advances Neural Information Processing Systems. Vol. 32. 2019. URL.: https: //axivv//xxvv wttesii.o.ro///eeooe wrg/oss//ta/eeeoeoooooo ///a//s//1///// .05600.
[32] Jonathan Ho et al. "Denoising Diffusion Probabilistic Models". In Advances in Neural Infor- mation Processing Systems. Vol. 33. 2020. URL:
www./or/////../ooooo 11239. [33] Brian Anderson. "Reverse-time diffusion equation models". Stochastie Processes and their Applications 12.3 (1982), pp. 313-326.
[34] Aapo Hyvärinen. "Estimation of Non-Normalized Statistical Models by Score Matching"... In Journal of Machine Learning Research 6.24 (2005), pp. 695-709, URL: http://jmlr: org/ papers/v6 http://///yrriinn05aaaaa
[35] Pascal Vincent. "A Connection Between Score Matching and Denoising Autoencoders". In Neural computation 23.7 (2011), pp. 1661-1674.
[36] Jiaming Song et "Denoising Diffusion Implicit Models". In International Conference on Learning Representations. 2021. URL: https / openreview net forum ? id=
StIgiarCHLP.
[37] Luping Liu et al. "Pseudo Numerical Methods for Diffusion Models on Manifolds". Inter- national Conference Learning Representations. 2022. URL: https:///ppnrreveew.rt net/ forum?id=P1KWVdPP
[38] Bahjat Kawar et al. "SNIPS: Solving Noisy Inverse Problems Stochastically". Advances in Neural Information Processing Systems. 2021. URL.: htttp:://ppenreiee..nett
id=pBKOx_dxxs dxYAN.
[39] Yang Song et al. "Solving Inverse Problems in Medical Imaging with Score-Basee Generative Models". In International Conference on Learning Representations. 2022. URL: https /openreviee. net/ fortt//frmmmmee iidruaRCHV OuGI.
[40] Alexandre Adam et al. "Posterior samples of source galaxies in strong gravitational lenses with score-based priors". 2022. URL: https:///artiii org/ htab//221.000m 03812.
[41] Hyungjin Chung et al. "Diffusion Posterior Sampling for General Noisy Inverse Problems". In International Conference on Learning Representations. 2023, URL:
net net/forum? httffformmiddddnDip DnDniddDDD.DD
[42] Bradley Efron. "Tweedie's ' Formula and Selection Bias". In Journal of the American Statistical Association 106.496 (2011), pp. 1602-1614.
[43] Kwanyoung Kim and Jong Chul Ye. "Noise2Score: Tweedie's Approach to Self-Supervised Image Denoising without Clean Images". In Advances in Neural Information Processing Systems. 2021. URL: https: h /ttp///ppeeeeeeeee///ee mett:forrumtddd2225ss3tsss
[44] Xiangming Meng and Yoshiyuki Kabashima. "Diffusion Model Based Posterior Sampling for
Noisy Linear Inverse Problems". 2022. URL: https://arxivv org/ abs/ 12343.
[45] Giorgio Parisi. "Correlation functions and computer simulations". In Nuclear Physics 180.3
(1981), pp. 378-384.
[46] Ulf Grenander and Michael I. Miller, "Representations of Knowledge in Complex Systems". In Journal of the Royal Statistical Society 56.4 (1994), pp. 549-603. URL: http: /ww.. jstor. org/ sttssaaale///// sttgg////bb////.......
[47] Edward N. Lorenz. "Deterministic Nonperiodic Flow". In Journal of Ammosperic Sciences 20.2 (1963), pp. 130-141
[48] Gary J. Chandler and Rich R. Kerswell. "Invariant recurrent solutions embedded in turbulent two-dimensional Kolmogorov flow". In Journal of Fluid Mechanics 722 (2013), pp. 554-595
[49] Jun S. Liu and Rong Chen. "Sequential Monte Carlo methods for dynamic systems". Journal of the American statistical association 93.443 (1998), pp. 1032-1044.
12

---

[50]
Arnaud Doucet, Adam M. Johansen, et al. "A Tutorial on Particle Filtering and Smoothing". In Handbook of nonlinear filtering 12 (2011), pp. 656-705. URL: httpps:/wwtt warwick. ac uk/37961.
[51] Neil J Gordon et al. "Novel approach to nonlinear/non-- In Radar and Signal Processing. Vol. 140. IET. 1993, pp. 107-113.
Bayesian state estimation".
[52] Alexander Quinn Nichol and Prafulla Dhariwal. "Improved denoising diffusion probabilis- tic models". In International Conference on Machine Learning. 2021. URL: https proceedings. mlr www.sp/ews/eneeeeeeeeo nicho...a..
[53] Dmitrii Kochkov et al. "Machine Proceedings of the National Academy of Sciences 118.21 (2021). URL: https: //www. http///www.........../../.. pnas.
accelerated computational fluid dynamics". In
/00 10. 1070/ 10w.//ass........ pnas. 2101784118.
[54] Guido Boffetta and Robert Ecke. "Two-dimensional Turbulence". Annual Review of Fluid Mechanics 44 (2012), pp. 427-451.
[55] Olaf Ronneberger et al. "U-Net: Convolutional Networks for Biomedical Image Segmentation". Medical Image Computing and Computer-Assisted Intervention. 2015, pp. 234-241, URL: https:/ //ae:i//axxvv..// wrggaas/////eeeeooooogoooooee abs/15///5 .0499..
[56] Yang Song et al. "Consistency models". 2023. URL: https // arxiv. org/ hes//ab//// 2210. 02303.
[57] Evan Archer et al. "Black box variational inference for state space models". 2015. URL: https://arx:s org/aaga/
[58) Chris Maddison et al. "Filtering variational objectives". In Advances in Neural Information Processing Systems (2017). URL: ttttps://aseee wwwaariwwww..waaa.eeeerreeeeaeee argg/bs//is 1705. 1etps/977009999.
[59] Christian Naesseth et al. "Variational sequential monte carlo".. In International Conference on
Artificial Intelligence and Statistics. 2018, pp. 968-977. URL: http:/ httt:://proceeeingge mlr. Fress/v84//aessethie
[60] Dieterich Lawson et "SIXO: Smoothing Inference with Twisted Objectives". In Advances in Neural Information Processing Systems, 2022. URL: http://pppprrreeeeeeeeeeeeeeeeee id=bDyLgfvZ0QJ.
[61] Rudolf E. Kalman. "A New Approach to Linear Filtering and Prediction Problems". Journal of Basic Engineering 82.1 (1960), pp. 35-45.
[62] Eric A. Wan and Rudolph Van Der Merwe. "The unscented Kalman filter for nonlinear estima- tion". Proceedings of the IEEE Adaptive Systems for Signal Processing, Communications, and Control Symposium. IEEE. 2000, pp. 153-158.
[63] Ienkaran Arasaratnam and Simon Haykin. "Cubature Kalman Filters". In IEEE Transactions on Automatic Control 54.6 (2009), pp. 1254-1269.
[64] Rahul G. Krishnan et al. "Deep Kalman Filters". 2015. URL: htttt::///aaattta org/abst abs/ 1511.05123
[65] Laurent Girin et al. "Dynamical variational autoencoders: A comprehensive review". 2020. URL: / wwwpaaawwwwww////.//.//// org/abs/2008 12595.
[66] Suman Ravuri et al. "Skilful precipitation nowcasting using deep generative models of radar".
In Nature 597.7877 (2021), pp. 672-677. URL: https //www nature. com/articlss/
841586- 0
[67] Remi Lam et al. "GraphCast: Learning skillful medium range global weather forecasting". 2022. URL: https: //te::////xxii .rg/ ab//////.../////////////aaa/2/2a 12794.
[68] Kaiming He et al. "Deep Residual Learning for Image Recognition". In IEEE conference computer vision and pattern recognition. 2016, pp. 770-778, URL: http:::////a/ii..oo org/ abs/1512. 03385.
[69] Stefan Elfwing al. "Sigmoid-Weipe Linear Units for Neural Network Function Ap- proximation in Reinforcement Learning". Neural Networks 107 (2018), pp. 3-11, URL: https://arxiv www.aiiww....aaaaeeeeee wwgg/as/////// 03118.
13

---

[70] Jimmy Lei Ba et al. "Layer normalization". 2016. URL: https://arxx..orgg//////o/ 1607. 06450.
[71] Ilya Loshchilov and Frank Hutter. "Decoupled Weight Decay Regularization". In International Conference on Learning Representations. 2019. URL: https://openreviev "tr.e/rrrrrinnneeeoommn id=Bkg6Ricqy7.
14

---

A
The pseudo-blanket approximation is unbiased
For any continuous random variables a, b and c,
V. log = p(a,b) Va p(a,b)
b(a.6) Vo / c(a,b,c) plc a,
p(a, Va ba
Eptcla.) Va log b, b.c]]
Replacing b and e by n. (t),,, 25, (t) and (t) (z,(t) : + ivj & b,), respectively, we obtain (1) log p (z. (t), 25, (t)) (())) (0) Vz,(0) log log
meaning that (t) log p
(z (1). 25, (t)) is the expected value of (2,()) log over =g((). In
other words, regardless of the elements the size of b;, Va,(1) log (oz,((((()) 26, (t)) is unbiased (exact in expectation) estimate of (1, logp(m::..()).
B On the covariance of p(æ xxt))
Assuming a Gaussian prior p(x) with covariance Ez, the covariance E of p( xx))) takes the form
att)
-1
== Exz
Ex
Σx
m11)²
at) p(1)2
(t)
QA
A
Q-1
m(1)2
where QAQ-1 the eigendecomposition of Eg. We observe that for most of the perturbation time
the central diagonal term is close to A(A [/++1] 7) -)- We therefore propose the covariance approximation
at))
r
m())²
where r = QA (A 11= Q-1 is positive semi-definite matrix.
15

---

C
Algorithms
Algorithm 3 Estimating the posterior score Va(s) (()(()) v)
function (=((t)), t y)
$ + (z())t) t)
+- =()++((0)* m(E)
By f Vz(() log.N A(t), + =
return Sy + Sy
Algorithm 4 Cordector-Correctrr sampling from s,(x(t),t
function SAMPLE((t:) N (i= 1= C,
z(1) ~ (0, (1))
for N to 1 do
a(t:-1) t m(ts-1) x(t;) + (((t-1)) o(ti-i) a(t.)? ss(x(t)), p(t,)
m(L)) o(ti)
for j = t0 do
E 1)
+ s (x0(ti f1-1 5) 8
(im
2
x(t4-1) f x (tz-1) + + /28
10
return z0))
16

---

D Experiment details
Resources Experiments were conducted with the help of cluster of GPUs. In particular, score networks were trained and evaluated concurrently, each on single GPU with at least 11 GB of memory.
Lorenz 1963 We consider two score network architectures: fully- conleceed local score networks for small & S 4) and fully-convolutioonal score networks for large & (A: >> 4). Residual blocks [68], SiLU [69] activation functions and layer normalization [70] are used for both architecture. The number of blocks in the cully-convolutional architecture controls the value of k. We train all score networks for 1024 epochs with the A [71] optimizer and linearly decreasing learning rate. Other hyperparameters are provided in Table
Table 1: Score network hyperparameters for the Lorenz experiment.
Hyperparameter
k S 4
k 4
Architecture Residual blocks
fully-connected
fully-convolutional
5
k - 2
Features/channels
256
64
Kernel size
3
Activation
SiLU
SiLU
Normalization
LayerNorm
LayerNorm
Optimizer
AdamW
AdamW
Weight decay
10-3
10-3
Learning rate
10
10-3
Scheduler
linear
linear
Epochs
1024
1024
Batches per epoch
256
256
Batch size
256
64
Kolmogorov flow The local score network is U-Net [55] with residual blocks [68], SiLU [69] activation functions and layer normalization [70]. This architecture is motivated by the locality of the Navier-Stokes equations, which impose that the evolution of drop of fluid is determined by its immediate environment. This can be seen as application of the pseudo-blanket approximation (13) to grid-strue grid-structured set random variables. train the score network for 1024 epochs with the AdamW [71] optimizer and linearly decreasing learning rate. Other hyperparameters are provided Table 2.
Table 2: Score network hyperpurameters for the Kolmogorov experiment.
Architecture Residual blocks per level
U-Net (3,3, 3) (96, 192, 384)
Channels per level
Kernel size
Padding
circular
Activation
SiLU LayerNorm
Normalization
Optimizer
AdamW 10 10-3 3
Weight decay
Learning rate
2 x 10-4
Scheduler
linear
Epochs
1024
Batches per epoch
128
Batch size
32
17

---

E
Assimilation examples
FEL
SDA
DPS
Figure 6: Example of sampled trajectory when the observation is insufficient to identify the original trajectory. The observation the center of the velocity field every three steps, coarsened resolution 8 x 8 and perturbed by small Gaussian noise (E, = T.. SDA (C= = 1, = 0.5) identifies the original state where data is sufficient, while generating physically plausible states elsewhere. Replacing SDA's likelihood score approximation with the one of DPS [41] yields trajectory less consistent with observation.
SDA
DPS
Figure 7: Example of sampled trajectory when the observation process non-linear. The observation corresponds to a saturating transformation v + I+|2] of the vorticity field w. SDA (512 discretization steps, - 1, = 0.5) identifies the original state where data is sufficient, while generating physically plausible states elsewhere. Replacing SDA's likelihood score approximation with the one of DPS [41] yields trajectory inconsistent with the observation.
18

---

2/1:1.
SDA
DPS
DDE
SDAX
DPS
Figure 8: Example of sampled trajectories when the observation is spatially sparse. The observation y" corresponds to a spatial subsampling of factor n of the velocity field. SDA (C' = 1, = 0.5) closely recovers the original trajectory for all factors n, despite the limited amount of available data. Replacing SDA's likelihood score approximation with the one of DPS [41] l] leads to trajectories that are progressively less consistent with the observation as n increases.
19

---

Figure 9: Example of long (I. 127) sampled trajectory. Odd states are displayed from left to right and from top to bottom. The observation process probes the difference between the initial and final states and the observation is set zero, which enforces looping trajectory (x) xT)). Note that this scenario is not realistic and will therefore lead to physically implausible trajectories.
20

---

SDA
Figure 10: Example of sampled trajectories when the observation is insufficient identify the original trajectory. The observation is the center the velocity field every three steps, coarsened to resolution 8 x 8 and perturbed by small Gaussian noise (E,, 0.01*1). SDA (C 1, 0.5)
identifies the original state where data is sufficient, while generating physically plausible states elsewhere.
SDA
Figure 11: Example of sampled trajectories when the observation is spatially sparse. The observation y corresponds to spatial subsampling of factor 16 of the velocity field with small Gaussian noise (Σ,- 0.. SDA (C' = t 0.5) generates trajectories similar to the original, with physically plausible variations.
21


Reviews:
[{"review_id": "9QTWSNHCe1", "reviewer": "Reviewer_DaY5", "summary": "Under a data assimilation setting —where the hidden state of a dynamical system is only accessed through noisy observations— the authors consider the problem of inferring the latent state as well as learning a faithful generative model.  In order to accomplish this, they take a score-based diffusion model approach by placing a diffusion model prior over latent trajectories.  They reveal many practical problems that arise due to the complexity of training the model with (possibly and practically) very high dimensional latent trajectory samples.\n", "questions": "Can the method only generate trajectories of length L?\n\nComing from a probabilistic inference and signal processing background, I was wondering if the author considers data assimilation models to be a superset of models that contain state-space models?  Prior to this, I have never encountered this interesting field of data assimilation, and I felt like tying back this model to simpler latent variable models such as those would help increase its scope (mentioned a little at the end, but I think moving this up/fleshing it out could help).  \n", "limitations": "As mentioned, the biggest limitation is assumption of the underlying dynamics.  \n\nThere are not many comparisons against other methods made. \n", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written, and I think the flow and order of presentation was great.  Those who may only be somewhat familiar with denoising diffusion models and latent variables should be able to follow along for the most part.  \n\nTo make training a diffusion model with the intent to generate high dimensional latent trajectories scale practical, the authors propose pragmatic solutions to problems that arise: \n\nThey propose to take advantage of the temporal smoothness of a diffusion model prior by approximating the score function with respect to particular time indices of the perturbed trajectories with the score function of an artificially imposed Markov blanket around that index; they are able to show that in practice, only a few neighbors are needed to create a large enough artificial Markov blanket.   \n\nThey propose to approximate the conditional likelihood of an observation given a sample from a point in the chain; unlike other proposed methods, they do not have to rescale the residuals \n\nThe experiments were interesting, especially, the experiment where the observations are taken to be coarse grained versions of the velocity field; the trajectory sampled from SDA is qualitatively similar to a ground truth sample.  Showing a sample generated from SDA  resembles feeding that sample through the true model was a good demonstration. \n", "weaknesses": "Weaknesses:  It is mentioned in the paper, but by far, the biggest weakness would be that the true underlying dynamics are assumed to be known (or at the very least, can be easily sampled from).  This does slightly bring down the experimental results, in my opinion, since learning the underlying dynamics of an unknown dynamical system is very difficult.  \n\nWhile the paper is well written, and the problem considered is both important and challenging, there is not too much in the way of technical advances; therefore, I think there should have also been more focus on the empirical results.  While I am not certain if many classes of models/methods exist for solving the exact same problem — I am wondering if the authors could have compared their method to other denoising diffusion models that operate on time series e.g. audio or video.  \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and the legitimate concerns you have raised. We hope that the following answers will satisfy you.\n\n**Weaknesses & Limitations**\n\n> [...] the biggest weakness would be that the true underlying dynamics are assumed to be known (or at the very least, can be easily sampled from).\n\nWe believe this concern is due to some misleading statements in Section 5 (Conclusion), which will be addressed. We invite you to consult the global rebuttal regarding the assumption of a known physical model.\n\n> [...] there is not too much in the way of technical advances;\n\nWe will address the lack of discussion about the significance of the contributions in Section 5 (Conclusion). We invite you to consult the global rebuttal regarding the novelty of the contributions.\n\n> [...] I think there should have also been more focus on the empirical results.\n\nWe respectfully disagree with this statement. Three pages of the manuscript as well as a few pages in the Appendix are dedicated to empirical results, with both quantitative and qualitative assessments. Besides, we already condensed Sections 2 and 3 to the limit and cannot afford to remove any more content.\n\n> [...] I am wondering if the authors could have compared their method to other denoising diffusion models that operate on time series e.g. audio or video.\n\nIndeed, diffusion models have been extensively used for sequence data, including text, audio and video. Some of these methods could have been applied to the first system, Lorenz 63. However, the purpose of Section 4.1 is to compare our approximations against the ground-truth posterior (which is accessible) and study their convergence with respect to hyper-parameters ($k$ and $C$). The comparison with other approximations would have required a lot of space, without added value.\n\nFor the second system, Kolmogorov flow, the only contenders would have been video diffusion models (VDMs) [19]. However, as briefly mentioned in Section 5, these models generate sequences autoregressively from \"left-to-right\" (past-to-future), which prevents from conditioning and sampling all frames simultaneously. Hence, they cannot be used for trajectory posterior inference. We could have built our own vanilla diffusion model for image sequences, but it would have been very expensive to train and irrelevant as a point of comparison.\n\n**Questions**\n\n> Can the method only generate trajectories of length L?\n\nTraining requires sequences of length $2k + 1$, but the length $L$ of the generated trajectories is chosen arbitrarily (albeit longer than $2k+1$) at inference.\n\n> I was wondering if the author considers data assimilation models to be a superset of models that contain state-space models?\n\nIndeed, state-space models are a special case of dynamical systems. We chose to frame the paper within the context of data assimilation because of its generality.\n\nWe believe that this rebuttal addresses most of your concerns and, therefore, kindly ask you to reconsider your score."}, {"review_id": "jMwIphypE2", "reviewer": "Reviewer_bwmG", "summary": "The authors propose score-based data assimilation framework that relies\non score based generative modeling for trajectory inference/state estimation\nof a dynamical model described by an SDE.\nTo make the procedure efficient, the authors employ three novelties, partly \nadopted from existing literature.\nThey train a score based generative model on short segments of trajectories, and  decompose the posterior score into a prior score and a likelihood score, and approximate the latter with a  Gaussian assuming a Gaussian observation process. \nFor approximating the prior score, they rely on the Markovian nature of the system and employ local (in time)\nprior score approximations.\nThey present their proposed approach on two non trivial model systems.\n\n\n\nI find the contribution interesting, however the novelty the overall ambition\nof the project is slightly limited for a conference like NeurIPS. In particular,\nI would consider more reasonable the option of trying to identify both\nthe dynamics and estimate the state of the model, since in most practical\nscenarios it is the underlying model dynamics that are unknown. In that respect I find the\nuse of word inference throughout the text misleading, since the actual problem that\nis solved here is not model identification but model estimation and prediction.\n\nOverall I find the paper quite well written, and technically sound. However the novelty and the contribution in the \nexisting literature is relatively low compared to existing approaches.", "questions": "- the system in consideration is assumed to be stochastic, however in Fig.2 C the trajectories seem\npretty deterministic to me. I assume that the individual non-opaque trajectories are individual realisations.\nIs the employed noise amplitude 0.025? I consider this negligible noise, especially considering the dynamic range of the system.\n\n- can you quantify either theoretically or empirically the errors incurred by each of the approximations employed in the proposed solution.\n\n- have you tried to predict non-stationary trajectories with the proposed framework? I would think that already the score function estimator\nwould struggle considerably in this scenario.\n\n- doesn't Eq. 12 only hold in expectation?", "limitations": "- the proposed framework assumes an already known dynamical model which is a rather limiting scenario considering both state of the art frameworks and practical applicability.\n\n- according to my understanding, the employed noise in the experiments is weak and the system could be effectively treated as deterministic\n\n- the method is limited only to Gaussian observation model settings.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the approach does not require simulation through the model equations, that can in principle become too computationally demanding for large systems.\n\n- well written manuscript and nice presentation.\n\n- the proposed method provides an approximation to the entire posterior of the path of the system instead of a point estimate, that in principle allows for further Bayesian inference of the underlying dynamical model.", "weaknesses": "- The observation model employed is rather trivial (Gaussian), yet the derivation relies on this assumption. Have you tested the approach with different observation models to test how well this Gaussian approximation holds?\n\n- To my understanding, the trained system is able to generate only\ntrajectories in equilibrium, and it would fail in non-equilibrium scenarios.\n- The authors mention that the introduced approximations introduce errors in the system\nbut I would expect them to at least quantify the resulting errors\nnumerically, if theoretical predictions cannot be formulated.\n\n- The authors assume known physical model (both the parametric form and the parameters) that\nconsiderably limits the applicability of the framework for practical applications.\n\n- To my understanding, the noise employed in the numerical experiments is very small, and thereby one could treat the system as deterministic.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and reading the manuscript in details. We hope that the following answers will satisfy you.\n\n**Weaknesses & Limitations**\n\n> [...] However, the novelty and the contribution in the existing literature is relatively low compared to existing approaches.\n\nWe will address the lack of discussion about the significance of the contributions in Section 5 (Conclusion). We invite you to consult the global rebuttal regarding the novelty of the contributions.\n\n> The authors assume known physical model (both the parametric form and the parameters) that considerably limits the applicability of the framework for practical applications.\n\nWe believe this concern is due to some misleading statements in Section 5 (Conclusion), which will be addressed. We invite you to consult the global rebuttal regarding the assumption of a known physical model.\n\n> The observation model employed is rather trivial (Gaussian), yet the derivation relies on this assumption.\n\nIndeed, the observation process is assumed to be Gaussian. However, we argue that this assumption is by no means trivial. It covers all observations that can be expressed as \"measurement + Gaussian noise\", which corresponds to many real-life scenarios, including some of the most challenging non-linear ones. In fact, an overwhelming amount of literature in the fields of state estimation and Bayesian filtering (Kalman filtering) is based on that Gaussian assumption. In the field of conditional score-based modeling, this is also a very common assumption [19, 24, 34-37] and we are not aware of any literature assuming non-additive measurement noise.\n\n> To my understanding, the trained system is able to generate only trajectories in equilibrium, and it would fail in non-equilibrium scenarios.\n\nIn our experiments, we leverage the statistical stationarity of the dynamics to make the local score network index-agnostic, as mentioned at line 153. If stationarity cannot be assumed, for instance in a periodic environment (e.g. seasons), the local score network would simply need to be conditioned with respect to the indices of the states.\n\nHowever, your comment raises a legitimate point: the score network can only be valid for dynamics it has been trained for, which leads to open questions regarding distribution shifts (e.g. climate change), as mentioned in Section 5.\n\n> The authors mention that the introduced approximations introduce errors in the system but I would expect them to at least quantify the resulting errors numerically, if theoretical predictions cannot be formulated. [...] Can you quantify either theoretically or empirically the errors incurred by each of the approximations employed in the proposed solution?\n\nWe already provide quantitative results for Lorenz 63, for which the true posterior is available via the bootstrap particle filter (BPF). We empirically demonstrate that SDA quantitatively converges to the true posterior as $C$ and $k$ increase (see Figure 2). For the Kolmogorov Flow system, the true posterior is intractable with classic posterior inference or data assimilation techniques, and therefore, cannot be quantitatively compared to.\n\nAs for theoretical estimation, the impact of the approximations (13) and (15) remains to be quantified, which we leave for future work, as stated in Section 5.\n\n> To my understanding, the noise employed in the numerical experiments is very small, and thereby one could treat the system as deterministic.\n\nIndeed, the noise used for both systems is not large. However, because the systems are chaotic, very small changes in the initial state or small stochasticites can lead to vastly different trajectories. Using more noise would only make the system more chaotic and make inference over long time horizons pointless ($x_L$ would be independent from $x_1$). In fact, increasing the noise levels would make the pseudo Markov approximation (13) more accurate as it would reduce the mutual information between far-away states. To paraphrase, the \"low\" amount of noise in our experiments is actually a means to challenge our methods and assumptions. Note that, in practice, physical models considered in data assimilation are often deterministic, or present very low amount of noises.\n\n**Questions**\n\n> The system in consideration is assumed to be stochastic, however in Fig. 2C the trajectories seem pretty deterministic to me. [...]\n\nIn Figure 2C, trajectories are not independent. They are realizations sampled from the true and estimated posteriors. All trajectories are close to each other because the observation $y$ enables to infer a narrow posterior. Conversely, in Figure 3, the observation (red crosses) leads to a multi-modal posterior, and trajectories diverge even when they start at the same initial states. If we were to represent the prior $p(x_{1:L})$ instead of the posterior $p(x_{1:L} | y)$, one could observe a wide variety of trajectories.\n\n> [...] Is the employed noise amplitude 0.025? I consider this negligible noise, especially considering the dynamic range of the system.\n\n$\\Delta = 0.025$ is the variance of the noise at each step. The standard deviation is $0.5$ which is not negligible with respect to the system's domain.\n\n> Have you tried to predict non-stationary trajectories with the proposed framework?\n\nWe are actively working on/looking for collaborations with weather and oceanic research centers to apply SDA on operational data, which typically have periodic (daily/yearly) dynamics.\n\n> Doesn't Eq. 12 only hold in expectation?\n\nThis equation is valid. We can note that $\\nabla_{x_i} \\log p(x_{j}) = 0$ and, by definition of a Markov blanket, $p(x_i | x_{\\neq i}) = p(x_i | x_{b_i})$.\n\nWe believe that this rebuttal addresses most of your concerns and, therefore, kindly ask you to reconsider your score."}, {"review_id": "mutNsmjmV7", "reviewer": "Reviewer_RNbZ", "summary": "This paper presents a technique to tackle data assimilation problems (inverse problems involving partially observed time series with a dynamical model acting as a prior over the reconstructed states) using a score based model to estimate the prior distribution $p(x)$ from data. The method used, at its core, leverages a now well-known capability of unconditional score based models to address inverse problems by incorporating the score of the likelihood into the generating process. This makes the approach appealing in zero-shot problems, where a new inverse problem with the same dynamical prior can be solved without any retraining. The paper is the first (to my knowledge) to adapt these techniques for data assimilation. It further proposes a few technical contributions on score based modeling : i) the capacity to exploit the Markov structure of dynamical systems data to train the score based model by small time chunks rather than a global one for the whole trajectories ii) A new approximation technique to estimate the score of the likelihood term $p(y|x(t))$, which is no trivial task iii) and a predictor corrector sample technique to improve data generation by incorporating a few Langevin Monte Carlo sampling iterations in between the steps of the SDE solver. The method is then tested on various assimilation tasks on two datasets coming from the Lorenz-63 system (in a quantitative way since the posterior can be arbitrarily well approximated by a Bootstrap Particle Filter) and Kolmogorov fluid flow (in a qualitative way).\n", "questions": "- Can the authors highlight the difference/added value between their predictor-corrector sampler and the one in [24] ? The rationale seems very similar even though the equations differ. Please explain.\n- In figure 4, would it be possible to quantify the closeness to the observations quantitatively, since this is the qualitative way of assessing the result on the Kolmogorov flow experiment? And compare it to what a competitor technique would output.\n- Similarly, it would be interesting to show several posterior samples on one of these experiments to have an idea of whether the model exhibits a multimodal posterior in this situation (I would expect this e.g. in the super-resolution and extrapolation applications). An analog of figure 3 for the second case study, basically.\n- Maybe a more detailed explanation of Tweedie’s formula and the computation of equations (9) and (10) would be nice to make the paper more self-contained.", "limitations": "- To me the main limitation is the lack of information about the practical impact of the contributions on the Markov blanket/correction step on the running time/memory in addition to simply their approximation capability. See above for more detailed remarks on this.\n- As mentioned by the authors, the dynamical model is assumed to be perfectly known or at the very least fixed. However, this is not true for the observation model, which makes the proposed method versatile as long as one is interested in various inverse problems with the same dynamical model. An interesting next step would be to see how to enable fitting/fine tuning model parameters, but since here the model is never used but rather emulated, this seems like a difficult endeavor. By the way, going through an emulator of the dynamics rather than having to run/differentiate the model could also be an advantage for complex models that cannot be easily made differentiable (as is often the case for operational models), provided they can be emulated successfully by a score based model. \n- On a related note, another score based model training is required in case of change in dynamical model (but this is the same for variational data assimilation techniques)", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper gives a clear exposition of score based methods and their use to solve inverse problems. It also reviews/improves a number of techniques used for efficient training of the model and incorporation of the inverse problem likelihood into the score function.\n- This seems to be the first application of score based models for posterior sampling in data assimilation, and it shows nice potential for those applications (emulating the physical model without the need to differentiate through the whole dynamical model and a solver, and reusing it for any assimilation problem using the same dynamics). The applications tackled, although relatively simple (in their dynamics), lead to think this could be applied to more ambitious situations in the future.\n- The paper presents three additional technical contributions useful for the generative modeling community as a whole, improving sample generation/training local score based models adapted to Markov chains/improving the approximation of the score of the likelihood. These contributions and the associated claims are globally well supported by the experiments.\n", "weaknesses": "- Section 3.1 should be a bit expanded, I believe. In particular, using the blanket to compute equations (11) and (12) could be a bit more detailed, and the link between the result of appendix A and the unbiasedness of the estimate (13) could be a bit explicited. It would be nice also to have an idea of how the value of k affects training time. However, the idea is sound and the approximation looks okay in the results on the Lorenz system.  So overall, the blanket is good.\n- More generally, there is a lot of emphasis in the paper on making the model easier to train/less expensive, less memory consuming or less prone to technical/computational difficulties (e.g. differentiating through the physical model and its associated solver), but none of this is actually quantified. For example, it would be quite interesting to test the effect of the value of $k$ on the training time compared to a naive global model, to quantify the effect of the number of correction steps $C$ on running time instead of just the quality of the solutions.\n- In addition, the paper would benefit from the comparison to more standard data assimilation techniques. This is done on the Lorenz data since a BPF is used to sample from the true posterior, but it would be nice to have an idea of the quality of the SDA output with respect to say, a classical variational assimilation with differentiation through the model (both in performance and running time). The training time should be reported, as well as the inference time, to be compared with the computation cost of optimizing a variational assimilation cost function.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your in-depth review and feedback. Most of your concerns are sensible and will be addressed in the manuscript.\n\n**Weaknesses & Limitations**\n\n> In particular, using the blanket to compute equations (11) and (12) could be a bit more detailed, and the link between the result of appendix A and the unbiasedness of the estimate (13) could be a bit explicited. [...] Maybe a more detailed explanation of Tweedie’s formula and the computation of equations (9) and (10) would be nice to make the paper more self-contained.\n\nWe will motivate and explain these equations more explicitly in the manuscript.\n\n> There is a lot of emphasis in the paper on making the model easier to train/less expensive, less memory consuming or less prone to technical/computational difficulties, but none of this is actually quantified.\n\nIndeed, the primary goal of Section 3.1 is to make the training of the prior score easier/less expensive. If we were to train a (global) score network on trajectories $x_{1:L}$, we would perform amounts of computation proportional to $L$, which could be huge depending on the dimensionality of the states $x_i$. Conversely, training a local score network only requires segments $x_{i-k:i+k}$ instead of complete trajectories $x_{1:L}$. An immediate consequence is that the cost of training is drastically reduced; roughly by a factor $\\frac{L}{2k + 1}$. A more subtle consequence is that the decomposition of the global score into local scores imposes a form of translational equivariance, which reduces the amounts of training data required.\n\nWe propose to replace (lines 143-144)\n\n\"Still, in some settings where each element $x_i$ is high-dimensional, building and training an FCNN for the full chain $x_{1:L}$ could be technically difficult. As a more tractable alternative, we propose to train a local score network\"\n\nwith\n\n\"Importantly, if the receptive field is $2k + 1$, the network can be trained on segments $x_{i-k:i+k}$ instead of the complete chain $x_{1:L}$, thereby drastically reducing training costs. More generally, we can train a local score network\"\n\nto make the benefits less ambiguous.\n\nAt inference, the decomposition of the global score into local scores additionally allows for better parallelization, which could be significant depending on available hardware. Concerning the number of corrections, $C$ (see Algorithm 4) has a linear impact on the number of network evaluations at inference, and would have the same impact with a global score network. Hence, doubling the number of corrections roughly doubles the inference time, but not the memory consumption, which stays constant.\n\n> In addition, the paper would benefit from the comparison to more standard data assimilation techniques.\n\nWe 100% agree with you. Unfortunately, we were only able to find ad-hoc implementations of 3D-Var and 4D-Var, with barely any documentation. If you point us to a relatively easy-to-use implementation, we'd be happy to oblige. We considered implementing 4D-Var ourselves, but we quickly realized that this would be very hard (which highlights the usefulness of SDA) and beyond the scope of this work.\n\n> As mentioned by the authors, the dynamical model is assumed to be perfectly known or at the very least fixed. [...]\n\nThis is correct, but we invite you to consult the clarifications concerning this limitation in the global rebuttal.\n\n**Questions**\n\n> Can the authors highlight the difference/added value between their predictor-corrector sampler and the one in [24] ?\n\nThe difference between our predictor-corrector sampler and the one from Song et al. [24] (Algorithms 2 and 3) is the predictor. We employ a predictor called exponential integrator introduced by Zhang et al. [26]. Note that we do not claim our predictor-corrector sampling to be a novel contribution.\n\n> Would it be possible to quantify the closeness to the observations quantitatively [...] ?\n\nBecause the observation process $p(y | x_{1:L}) = \\mathcal{N}(y | \\mathcal{A}(x), \\Sigma_y)$ is known, we can measure the likeliness of an observation with respect to an inferred trajectory. However, the likeliness for a single pair $(x_{1:L}, y)$ cannot be interpreted. Instead, as is done for the Lorenz 63 system (see Figure 2), we can compute the expected log-likelihood $E_{q(x_{1:L} | y) p(y)} [ \\log p(y | x_{1:L}) ]$, which should be equal to minus the entropy of $\\mathcal{N}(0, \\Sigma_y)$. In the case of the low and high-frenquency observation processes, the entropies are respectively $\\frac{9}{2} \\log(2 \\pi e 0.05^2) \\approx -14.191$ and $\\frac{65}{2} \\log(2 \\pi e 0.25^2) \\approx 2.122$, which matches what we observe in Figure 2. It should be noted, however, that this equality is a necessary but insufficient condition for exact posterior inference.\n\nWe will add a new section in Appendix computing this expected likelihood for the Kolmogorov system.\n\n> It would be interesting to show several posterior samples on one of these experiments to have an idea of whether the model exhibits a multimodal posterior in this situation [...]\n\nWe will add the suggested examples in the Appendix."}, {"review_id": "d92FppwSY0", "reviewer": "Reviewer_RS1N", "summary": "This paper proposes an approach to solve data-assimilation tasks using techniques from score-based generative modeling. Given observations and trajectories from the data distribution, a posterior distribution over an entire state-trajectory is to be computed. The score function of the posterior is estimated using a learned prior-score network and a Gaussian approximation to the observation model. To make the score-based approach scale to long state trajectories, the authors propose to estimate the prior score over the entire trajectory as a combination of scores computed on local time windows. The theoretical groundwork is backed by a set of experiments in simulated environments.", "questions": "- Is the approach scalable to massive-scale real-world problems, such as weather forecasting?\n  - Is it feasible to generate as many training trajectories as would be needed for such complex dynamics?\n- Was the performance of the method studied under dynamics that develop turbulences or even shocks (such as, e.g., Burgers' equation)?", "limitations": "- Limitations were stated in the main text.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The presentation of the method is clear and understandable\n- This approach to data assimilation is unlike any I have seen before. It combines two established techniques in a natural way without bending one or the other out of shape. The method yielded by said combination appears natural and it seems valuable to test it on real-world applications.\n- The authors provide clean code in their supplementary material\n- The limitation section answered open questions I had while reading the paper. It lists approximations made and work that is left to do in the future.\n- My feeling is that this is the ground work to a fairly interesting and useful approach to inference in dynamical systems", "weaknesses": "- From my understanding, the method has yet to be tested on (settings on the scale of) actual real-world problems. Concretely, the experiments shown -- yet convincing -- are rather small-scale, considering that data assimilation regularly has to tackle problems with several millions of state variables. This is not necessarily to be expected when presenting a novel technique, but the Kolmogorov flow experiment left me wondering if there was something standing in the way of choosing a grid that is larger than 64 x 64.\n(Note: This is not asking for an additional experiment, but rather my opinion that a statement regarding scalability of the approach would be a good addition to the manuscript.)\n- I suspect that -- as is -- the amount of training data that would have to be generated for very large-scale tasks could be computationally quite infeasible", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your enthusiastic review and for acknowledging the code provided in the supplementary material. We hope that the following answers will satisfy you.\n\n**Weaknesses & Limitations**\n\n> From my understanding, the method has yet to be tested on (settings on the scale of) actual real-world problems. [...] a statement regarding scalability of the approach would be a good addition to the manuscript.\n\nThis is a good point. We propose to add the following paragraph in Section 5.\n\n* Concerning our experiments, although the Kolmogorov system is already high-dimensional (tens of thousands of dimensions) with respect to what is approachable with classical posterior inference methods, it remains small in comparison to the millions of dimensions of some operational DA systems. Whether SDA would scale well to such applications is an open question and will obviously present serious engineering challenges.\n\nNevertheless, we would like to mention that operational DA is usually performed on clusters with hundreds to thousands of CPUs/GPUs, while training and inference in our experiments were performed with a single GPU. We are actively working on/looking for collaborations with weather and oceanic research centers to apply SDA on operational data.\n\n> I suspect that -- as is -- the amount of training data that would have to be generated for very large-scale tasks could be computationally quite infeasible.\n\nIn most dynamical systems, many symmetries and inductive biases (e.g. equivariances) can be exploited to reduce the amount of data necessary for training. This is even more true with SDA, as we train on short segments instead of complete trajectories. In fact, for the Kolmogorov system the full training set consists of 820 trajectories of 64 states, which is ridiculously small in comparison to the amounts of data typically used to train image/video generation models.\n\n**Questions**\n\n> Was the performance of the method studied under dynamics that develop turbulences or even shocks (such as, e.g., Burgers' equation)?\n\nThe Kolmogorov system considers the state of a two-dimensional turbulent fluid. We did not apply SDA on systems presenting shocks, but this would be an interesting application.\n"}]

Author rebuttal:
We would like to thank the reviewers for the quality and pertinence of their reviews. All reviewers declare that the methods are sound, interesting and well presented.

The main concerns of reviewers DaY5 and bwmG regard the assumption of a known physical model and the novelty of the contributions. We beleive these concerns are due to some misleading statements in Section 5 (Conclusion) and a lack of discussion about the significance of the contributions, which we propose to address with minor changes in the manuscript. Other comments will be addressed in the per-reviewer rebuttals.

**Assumption of a known physical model**

The introduction starts with the assumption of a known physical model because data assimilation (DA) methods typically and routinely rely on physical models to make up for the lack of data in geophysical systems. In these settings, purely data-driven approaches are simply not applicable. However, SDA is not limited to that assumption and can be applied as long as segments $x_{i-k:i+k}$ ($k \ll L$) are available for training. These segments are not necessarily synthetic and we can apply SDA from data alone. For instance, in cellular biology, the state of some systems can be observed with incredible accuracy, but their dynamics are still relatively unknown.

What we tried to convey at lines 273-280 was that

1. All trajectories are assumed to share the same dynamics, known or unknown.
2. If a physical model is used to generate the training data, one can only expect SDA to be as accurate as the physical model.

We propose to clarify these points in the manuscript by

1. replacing (line 273)

    > Another limitation of our approach is the assumption that the parameters of the physical model are known.

    with

    > Another limitation of our approach is the assumption that the dynamics of the system are shared by all trajectories. In particular, if a parametric physical model is used, all trajectories are assumed to share the same parameters.

2. replacing (line 276)

    > Similarly, SDA relies on the assumption that the physical model is an adequate description of reality. Although we show that it can generalize well to unlikely scenarios (see Figure 5), we can only expect our approach to be as accurate as the physical model itself.`

    with

    > Additionally, if a physical model is used to generate synthetic training data, instead of relying on real data, one can only expect SDA to be as accurate as the model itself.
    
**Novelty of the contributions**

As noted by reviewers RNbZ and RS1N, our work present several novel contributions, in both the fields of data assimilation and score-based generative modeling. As the significance of these contributions are not explicitly discussed in the manuscript, we propose to add the following sub-section in Section 5.

> **Significance** In addition to its contributions to the field of data assimilation, this work presents new technical contributions to the field of score-based generative modeling.
>
> First, we provide new insights on how to exploit the structure of sets of random variables to build and train score-based generative models. Based on these findings, we are able to generate/infer simultaneously all the states of arbitrarily long Markov chains $x_{1:L}$, while only training score models on short segments $x_{i-k:i+k}$, thereby reducing training costs and the amounts of training data required. The decomposition of the global score into local scores additionally allows for better parallelization at inference, which could be significant depending on available hardware. Importantly, the pseudo-blanket approximation (13) is not limited to Markov chains, but could be applied to any set of variables $x_{1:L}$, as long as some structure is known.
>
> Second, we motivate and introduce a novel approximation (15) for the perturbed likelihood $p(y | x(t))$, when the likelihood $p(y | x)$ is assumed (linear or non-linear) Gaussian. We find that computing the likelihood score $\nabla_{\! x(t)} \log p(y | x(t))$ with this new approximation leads to accurate posterior inference, without the need for stability tricks [37]. This contribution can be trivially adapted to many tasks such as inpainting, deblurring, super-resolution or inverse problems in scientific fields [35-37].

We would also like to emphasize that several papers whose sole purpose were to approximate the likelihood score, were published in major venues [34-37].

Meta-review: All four reviewers argue for accepting this paper. They found it methodologically interesting, well written, and as a possible basis for further work down the stream.  

In the rebuttal/discussion, the authors promise to modify the paper (and appendix) with several edits and additional sections. I'm counting on them that the edits will be made as promised. Those edits will help the reader understand the work in better detail and thus ensure a wider impact of the paper.
Decision: Accept (poster)