Title: Score-based Data Assimilation
Venue: NeurIPS 2023
Accepted: yes
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

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

1 Introduction

Data assimilation (DA) [1-9] is at 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 x 1: L = ( x 1 , x 2 , . . . , x L ) ∈ R L × D denote a trajectory of states in a discrete-time stochastic dynamical system and p ( x i +1 | x i ) be the transition dynamics from state x i to state x i +1 . An observation y ∈ R M of the state trajectory x 1: L follows an observation process p ( y | x 1: L ) , generally formulated as y = A ( x 1: L ) + η , where the measurement function A : R L × D ↦→ R M is often non-linear and the observational error η ∈ R M is a 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 x 1: L given an observation y , that is, to estimate the trajectory posterior

where the initial state prior p ( x 1 ) 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 to fill in spatial and temporal gaps in the observation.

State-of-the-art approaches to data assimilation are based on variational assimilation [1, 2, 5-7]. Many of these approaches formulate the task as a maximum-a-posteriori (MAP) estimation problem and solve it by maximizing the log-posterior density log p ( x 1: L | y ) via gradient ascent. Although this

Gilles Louppe

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 a 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, 11]. Recent work has shown that deep learning can be used in a variety of ways to improve the computational efficiency of data assimilation, increase the reconstruction performance by estimating unresolved scales after data assimilation, or integrate multiple sources of observations [12-19].

Contributions In this work, we propose a novel approach to data assimilation based on scorebased generative models. Leveraging the Markovian structure of dynamical systems, we train a 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 - it is not limited to point estimates - without simulating or differentiating through the physical model. The code for all experiments is made available at https://github.com/francois-rozet/sda .

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-time score-based generative models Adapting the formulation of Song et al. [28], samples x ∈ R D from a distribution p ( x ) are progressively perturbed through a continuous-time diffusion process expressed as a linear stochastic differential equation (SDE)

where f ( t ) ∈ R is the drift coefficient, g ( t ) ∈ R is the diffusion coefficient, w ( t ) ∈ R D denotes a Wiener process (standard Brownian motion) and x ( t ) ∈ R D is the perturbed sample at time t ∈ [0 , 1] . Because the SDE is linear with respect to x ( t ) , the perturbation kernel from x to x ( t ) is Gaussian and takes the form

where µ ( t ) and Σ( t ) = σ ( t ) 2 I can be derived analytically from f ( t ) and g ( t ) [29, 30]. Denoting p ( x ( t )) the marginal distribution of x ( t ) , we impose that µ (0) = 1 and σ (0) ≪ 1 , such that p ( x (0)) ≈ p ( x ) , and we chose the coefficients f ( t ) and g ( t ) such that the influence of the initial sample x on the final perturbed sample x (1) is negligible with respect to the noise level - that is, p ( x (1)) ≈ N (0 , Σ(1)) . The variance exploding (VE) and variance preserving (VP) SDEs [28, 31, 32] are widespread examples satisfying these constraints.

Crucially, the time reversal of the forward SDE (2) is given by a reverse SDE [28, 33]

That is, we can draw noise samples x (1) ∼ N (0 , Σ(1)) and gradually remove the noise therein to obtain data samples x (0) ∼ p ( x (0)) by simulating the reverse SDE from t = 1 to 0 . This requires access to the quantity ∇ x ( t ) log p ( x ( t )) known as the score of p ( x ( t )) .

Denoising score matching In practice, the score ∇ x ( t ) log p ( x ( t )) is approximated by a neural network s ϕ ( x ( t ) , t ) , named the score network, which is trained to solve the denoising score matching objective [28, 34, 35]

where p ( t ) = U (0 , 1) . The theory of denoising score matching ensures that s ϕ ( x ( t ) , t ) ≈ ∇ x ( t ) log p ( x ( t )) for a 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 of ∇ x ( t ) log p ( x ( t ) | x ) near t = 0 makes the optimization of (5) unstable [30]. To mitigate this issue, a slightly different parameterization ϵ ϕ ( x ( t ) , t ) = -σ ( t ) s ϕ ( x ( t ) , t ) of the score network is often used, which leads to the otherwise equivalent objective [30, 32, 36]

where p ( ϵ ) = N (0 , I ) . In the following, we keep the score network notation s ϕ ( x ( t ) , t ) for convenience, even though we adopt the parameterization ϵ ϕ ( x ( t ) , 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 ( x (0)) ≈ p ( x ) . To solve inverse problems, however, we need to sample from the posterior distribution p ( x | y ) . This could be accomplished by training a conditional score network s ϕ ( x ( t ) , t | y ) to approximate the posterior score ∇ x ( t ) log p ( x ( t ) | y ) and plugging it into the reverse SDE (4). However, this would require data pairs ( x, 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-41] that the posterior score can be decomposed into two terms thanks to Bayes' rule

Since the prior score ∇ x ( t ) log p ( x ( t )) can be approximated with a single score network, the remaining task is to estimate the likelihood score ∇ x ( t ) log p ( y | x ( t )) . Assuming a differentiable measurement function A and a Gaussian observation process p ( y | x ) = N ( y | A ( x ) , Σ y ) , Chung et al. [41] propose the approximation

where the mean ˆ x ( x ( t )) = E p ( x | x ( t )) [ x ] is given by Tweedie's formula [42, 43]

As the log-likelihood of a multivariate Gaussian is known analytically and s ϕ ( x ( t ) , t ) is differentiable, we can compute the likelihood score ∇ x ( t ) log p ( y | x ( t )) with this approximation in zero-shot, that is, without training any other network than s ϕ ( x ( t ) , t ) .

3 Score-based data assimilation

Coming back to our initial inference problem, we want to approximate the trajectory posterior p ( x 1: L | y ) of a dynamical system. To do so with score-based generative modeling, we need to estimate the posterior score ∇ x 1: L ( t ) log p ( x 1: L ( t ) | 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-dimensionality of each state x i (e.g. the state of atmospheres or oceans) combined with potentially long trajectories would require an impractically large score network s ϕ ( x 1: L ( t ) , t ) to estimate the prior score ∇ x 1: L ( t ) log p ( x 1: L ( t )) and a proportional amount of data for training, which could be prohibitive if data is scarce or 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 a 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 ∇ x 1: L ( t ) log p ( y | x 1: L ( t )) , 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-based data assimilation (SDA), is the combination of these three components.

L

/greatermuch

k

x

1:

L

(0)

x

1:

L

(

t

)

d

x

(

t

) =

f

(

t

)

x

(

t

)d

t

+

g

(

t

)d

w

(

t

)

Forward SDE (data

→

noise)

x

1:

L

(0)

x

1:

L

(

t

)

d

x

(

t

) =

[

f

(

t

)

x

(

t

)

-

g

(

t

)

2

∇

x

(

t

)

log

p

(

x

(

t

))

]

d

t

+

g

(

t

)d

w

(

t

)

Reverse SDE (noise

data)

x

i

-

1:

i

+1

(0)

x

i

-

k

:

i

+

k

(

t

)

φ

t

∇

x

1:

L

(

t

)

log

p

(

x

1:

L

(

t

))

3.1 How is your blanket?

̸

Given a set of random variables x 1: L = { x 1 , x 2 , . . . , x L } , it is sometimes possible to find a small Markov blanket x b i ⊆ x = i such that p ( x i | x = i ) = p ( x i | x b i ) for each element x i using our knowledge of the set's structure. It follows that each element ∇ x i log p ( x 1: L ) of the full score ∇ x 1: L log p ( x 1: L ) can be determined locally, that is, only using its blanket;

̸

̸

̸

This property generally does not hold for the diffusion-perturbed set x 1: L ( t ) as there is no guarantee that x b i ( t ) is a Markov blanket of the element x i ( t ) . However, there exists a set of indices ¯ b i ⊇ b i such that

is a good approximation for all t ∈ [0 , 1] . That is, x ¯ b i ( t ) is a 'pseudo' Markov blanket of x i ( t ) . In the worst case, ¯ b i contains all indices except i , but we argue that, for some structures, there is a set ¯ b i not much larger than b i that satisfies (13). Our rationale is that, since we impose the initial noise to be negligible, we know that x b i ( t ) becomes indistinguishable from x b i as t approaches 0 . Furthermore, as t grows and noise accumulates, the mutual information between elements x i ( t ) and x j ( t ) decreases to finally reach 0 when t = 1 . Hence, even if ¯ b i = b i , the pseudo-blanket approximation (13) already holds near t = 0 and t = 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/symmetries, (13) could and should be exploited within the architecture of the score network s ϕ ( x 1: L ( t ) , t ) .

In the case of dynamical systems, the set x 1: L is by definition a first-order Markov chain and the minimal Markov blanket of an element x i is x b i = { x i -1 , x i +1 } . For the perturbed element x i ( t ) , the pseudo-blanket x ¯ b i ( t ) can take the form of a window surrounding x i ( t ) , that is ¯ b i = { i -k, . . . , i + k } \ { i } with k ≥ 1 . The value of k 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, a 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 2 k +1 , the network can be trained on segments x i -k : i + k instead of the full chain x 1: L , thereby drastically reducing training costs. More generally, we can train a local score network (see Algorithm 1)

such that its k + 1 -th element approximates the score of the i -th state ∇ x i ( t ) log p ( x 1: L ( t )) . We also have that the k first elements of s ϕ ( x 1:2 k +1 ( t ) , t ) approximate the score of the k first states ∇ x 1: k ( t ) log p ( x 1: L ( t )) and the k last elements of s ϕ ( x L -2 k : L ( t ) , t ) approximate the score of the k last states ∇ x L -k : L ( t ) log p ( x 1: L ( t )) . Hence, we can apply the local score network on all sub-segments x i -k : i + k ( t ) of x 1: L ( t ) , similar to a convolution kernel, and combine the outputs (see Algorithm 2) to get an approximation of the full score ∇ x 1: L ( t ) log p ( x 1: L ( t )) . Note that we can either condition the score network with i or assume the statistical stationarity of the chain, that is p ( x i ) = p ( x i +1 ) .

→

3.2 Stable likelihood score

Due to approximation and numerical errors in ˆ x ( x ( t )) , computing the score ∇ x ( t ) log p ( y | x ( t )) with the likelihood approximation (8) is very unstable, especially in the low signal-to-noise regime, that is when σ ( t ) ≫ µ ( t ) . This incites Chung et al. [41] to replace the covariance Σ y by the identity I and rescale the likelihood score with respect to ∥ y -A (ˆ x ( x ( t ))) ∥ to stabilize the sampling process. These modifications introduce a significant error in the approximation as they greatly affect the norm of the likelihood score.

We argue that the instability is due to (8) being only exact if the variance of p ( x | x ( t )) 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 ( x | x ( t )) with Σ( t ) / µ ( t ) 2 , which is valid as long as the prior p ( x ) is Gaussian with a large diagonal covariance Σ x . We motivate in Appendix B the more general covariance approximation σ ( t ) 2 / µ ( t ) 2 Γ , where the matrix Γ depends on the eigendecomposition of Σ x . Then, taking inspiration from the extended Kalman filter, we approximate the perturbed likelihood as

where A = ∂ x A | ˆ x ( x ( t )) is the Jacobian of A . In practice, to simplify the approximation, the term A Γ A T can often be replaced by a constant (diagonal) matrix. We find that computing the likelihood score ∇ x ( t ) log p ( y | x ( t )) with this new approximation (see Algorithm 3) is stable enough that rescaling it or ignoring Σ y is unnecessary.

3.3 Predictor-Corrector sampling

To simulate the reverse SDE, we adopt the exponential integrator (EI) discretization scheme introduced by Zhang et al. [30]

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 to low-quality samples. To prevent errors from accumulating, we perform (see Algorithm 4) a few steps of Langevin Monte Carlo (LMC) [45, 46]

where ϵ ∼ N (0 , 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 δ ∈ R + , simulated samples are guaranteed to follow the distribution implicitly defined by our approximation of the posterior score at each time t , 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 predictor-corrector (PC) sampling, to correct the errors introduced by the discretization of the reverse SDE.

4 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 a 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 a good understanding of how SDA would perform in 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 µ ( t ) = cos( ωt ) 2 with ω = arccos √ 10 -3 and σ ( t ) = √ 1 -µ ( t ) 2 . The score networks are trained once and then evaluated under various observation scenarios. Unless specified otherwise, we estimate the posterior score according to Algorithm 3 with Γ = 10 -2 I 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 ) ∈ R 3 of the Lorenz system evolves according to a system of ordinary differential equations

where σ = 10 , ρ = 28 and β = 8 3 are parameters for which the system exhibits a chaotic behavior. We denote ˜ a and ˜ c the standardized (zero mean and unit variance) versions of a and c , respectively. As our approach assumes a discrete-time stochastic dynamical system, we consider a transition process of the form x i +1 = M ( x i ) + η , where M : R 3 ↦→ R 3 is the integration of the differential equations (18) for ∆ = 0 . 025 time units and η ∼ N (0 , ∆ I ) represents Brownian noise.

0

1

2

4

8

16

Corrections

A

-

10

0

10

a

-

15

0

15

b

-

200

0

log

p

(

x

2:

L

x

1

)

0

1

2

4

8

16

Corrections

B

-

10

0

10

log

p

(

y

x

1:

L

)

BPF

k

=1

k

=2

k

=3

k

=4

k

=8

0

25

50

W

1

1

33

65

97

129

161

193

i

10

25

40

c

C

|

|

-

10

0

10

a

1

5

9

13

17

21

25

29

33

37

41

45

49

i

-

15

0

15

b

20

30

c

-

1

0

1

q

(

y

′

y

)

|

We generate 1024 independent trajectories of 1024 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 two score network architectures: fully-connected local score networks for small k ( k ≤ 4 ) and fully-convolutional score networks for large k . Architecture and training details for each k are provided in Appendix D.

We first study the impact of k (see Section 3.1) and the number of LMC corrections (see Section 3.3) on the quality of the inferred posterior. We consider two simple observation processes N ( y | ˜ a 1: L :8 , 0 . 05 2 I ) and N ( y | ˜ a 1: L , 0 . 25 2 I ) . The former observes the state at low frequency (every eighth step) with low noise, while the latter observes the state at 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) to draw 1024 trajectories x 1: L from the ground-truth posterior p ( x 1: L | y ) . We use a large number of particles ( 2 16 ) to ensure convergence. Then, using SDA, we sample 1024 trajectories from the approximate posterior q ( x 1: L | y ) defined by each score network. We compare the approximate and ground-truth posteriors with three summary statistics: the expected log-prior E q ( x 1: L | y ) [log p ( x 2: L | x 1 )] , the expected loglikelihood E q ( x 1: L | y ) [log p ( y | x 1: L )] and the Wasserstein distance W 1 ( p, q ) in trajectory space. We repeat the procedure for 64 observations and different number of corrections ( τ = 0 . 25 , see Algorithm 4) and present the results in Figure 2. To paraphrase, SDA is able to 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 | x 1: L ) = N ( y | ˜ c 1: L :4 , 0 . 1 2 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

where u is the velocity field, Re is the Reynolds number, ρ is the fluid density, p is the pressure field and f is the external forcing. Following Kochkov et al. [53], we choose a two-dimensional domain [0 , 2 π ] 2 with periodic boundary conditions, a large Reynolds number Re = 10 3 , a constant density ρ = 1 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 a 256 × 256 domain grid. The states x i are snapshots of the velocity field u , coarsened to a 64 × 64 resolution, and the integration time between two such snapshots is ∆ = 0 . 2 time units. This corresponds to 82

x

1:

L

:4

y

SDA

DPS

SS

integration steps of the forward Euler method, which would be expensive to differentiate through repeatedly, as required by gradient-based data assimilation approaches.

We generate 1024 independent trajectories of 64 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 a 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 a trajectory of length L = 32 from the evaluation set and observe the velocity field every four steps, coarsened to a resolution 8 × 8 and perturbed by a moderate Gaussian noise ( Σ y = 0 . 1 2 I ). Given the observation, we sample a trajectory with SDA ( C = 1 , τ = 0 . 5 ) and find that it closely recovers the original trajectory, as 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 a physically plausible scenario while remaining consistent with the observation, which can also be observed in Figure 6 and 7.

y

=

A

(

x

L

) +

η

SDA

x

1

Physical model

Finally, we investigate whether SDA generalizes to unlikely scenarios. We design an observation process that probes the vorticity of the final state x L in a circle-shaped sub-domain. Then, we sample a trajectory ( C = 1 , τ = 0 . 5 ) consistent with a uniform positive vorticity observation in this sub-domain, which is unlikely, but not impossible. The result is discussed in Figure 5.

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 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 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 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 ∇ x ( t ) log p ( y | x ( t )) 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 as 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 approximations of the whole posterior, especially as k and the number of LMC corrections C increase. However, our approximations (13) and (15) introduce a certain degree of error, which precise impact 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 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 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, if 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 is also required, or at least not without further developments. Some approaches [57-60] tackle this task, but they remain limited to low-dimensional settings. 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. This is a 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 a lower-dimensional space, which simplifies the optimization problem greatly. Frerix et al. [14] use a deep neural network to predict the initial state of a trajectory given the observations. This prediction is then used as a starting point for traditional

(4D-Var) variational data assimilation methods, which proves to be more effective than starting at random. This strategy is 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 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 et al. [23] train a score-based generative model for a fixed number of video frames and use it autoregressively to generate videos of arbitrary lengths. Conversely, our approach is non-autoregressive which allows to generate and condition all elements (frames) simultaneously. Interestingly, as part of their method, Ho et al. [23] introduce 'reconstruction guidance' for conditional sampling, which can be seen as a special case of our likelihood approximation (15) where the observation y is a subset of x . Lastly, Ho et al. [25] generate low-frame rate, low-resolution videos which are then up-sampled temporally and spatially with a 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.-FNRS (Belgium) and acknowledges its financial support.

References

Andrew C Lorenc. 'Analysis methods for numerical weather prediction'. In Quarterly Journal of the Royal Meteorological Society 112.474 (1986), pp. 1177-1194.

François-Xavier Le Dimet and Olivier Talagrand. 'Variational algorithms for analysis and assimilation of meteorological observations: theoretical aspects'. In Tellus A: Dynamic Meteorology and Oceanography 38.2 (1986), pp. 97-110.

Geir Evensen. 'Sequential data assimilation with a nonlinear quasi-geostrophic model using Monte Carlo methods to forecast error statistics'. In Journal of Geophysical Research: Oceans 99.C5 (1994), pp. 10143-10162.

Thomas M. Hamill. 'Ensemble-based Atmospheric Data Assimilation'. In Predictability of weather and climate 124 (2006), p. 156.

Yannick Trémolet. 'Accounting for an imperfect model in 4D-Var.' In ECMWF Technical Memoranda 477 (2005), p. 20. URL: https://www.ecmwf.int/node/12845 .

Yannick Trémolet. 'Model error estimation in 4D-Var'. In ECMWF Technical Memoranda 520 (2007), p. 19. URL: https://www.ecmwf.int/node/12819 .

Mike Fisher et al. 'Weak-Constraint and Long-Window 4D-Var'. In ECMWF Technical Memoranda 655 (2011), p. 47. URL: https://www.ecmwf.int/node/9414 .

Alberto Carrassi et al. 'Data assimilation in the geosciences: An overview of methods, issues, and perspectives'. In Wiley Interdisciplinary Reviews: Climate Change 9.5 (2018).

ECMWF. 'Part II: Data Assimilation'. In IFS Documentation CY47R1 2 (2020). URL: https: //www.ecmwf.int/node/19746 .

Peter Bauer et al. 'The quiet revolution of numerical weather prediction'. In Nature 525.7567 (2015), pp. 47-55. URL: https://www.nature.com/articles/nature14956 .

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.

Julian Mack et al. 'Attention-based convolutional autoencoders for 3d-variational data assimilation'. In Computer Methods in Applied Mechanics and Engineering 372 (2020). URL: https://arxiv.org/abs/2101.02121 .

Rossella Arcucci et al. 'Deep data assimilation: integrating deep learning with data assimilation'. In Applied Sciences 11.3 (2021), p. 1114. URL: https://www.mdpi.com/20763417/11/3/1114 .

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/abs/2102.11192 .

Ronan Fablet et al. 'Learning variational data assimilation models and solvers'. In Journal of Advances in Modeling Earth Systems 13.10 (2021). URL: https://arxiv.org/abs/2007. 12941 .

Yu-Hong Yeung et al. 'Physics-Informed Machine Learning Method for Large-Scale Data Assimilation Problems'. In Water Resources Research 58.5 (2022). URL: https://arxiv. org/abs/2108.00037 .

Julien Brajard et al. 'Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: A case study with the Lorenz 96 model'. In Journal of computational science 44 (2020). URL: https://arxiv.org/abs/2001.01520 .

Julien Brajard et al. 'Combining data assimilation and machine learning to infer unresolved scale parametrization'. In Philosophical Transactions of the Royal Society A 379.2194 (2021). URL: https://arxiv.org/abs/2009.04318 .

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.

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.org/abs/2112.10752 .

Aditya Ramesh et al. 'Hierarchical text-conditional image generation with clip latents'. In (2022). URL: https://arxiv.org/abs/2204.06125 .

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.net/forum?id=08Yk-n5l2Al .

Jonathan Ho et al. 'Video Diffusion Models'. In ICLR Workshop on Deep Generative Models for Highly Structured Data . 2022. URL: https://openreview.net/forum?id= BBelR2NdDZ5 .

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/ papers/v23/21-0635 .

Jonathan Ho et al. 'Imagen video: High definition video generation with diffusion models'. 2022. URL: https://arxiv.org/abs/2210.02303 .

Zhifeng Kong et al. 'DiffWave: A Versatile Diffusion Model for Audio Synthesis'. In International Conference on Learning Representations . 2021. URL: https://openreview.net/ forum?id=a-xFK8Ymz5J .

Karan Goel et al. 'It's raw! audio generation with state-space models'. In International Conference on Machine Learning . PMLR. 2022. URL: https://proceedings.mlr.press/ v162/goel22a .

Yang Song et al. 'Score-Based Generative Modeling through Stochastic Differential Equations'. In International Conference on Learning Representations . 2021. URL: https:// openreview.net/forum?id=PxTIG12RRHS .

Simo Särkkä and Arno Solin. 'Applied stochastic differential equations'. Vol. 10. Cambridge University Press, 2019.

Qinsheng Zhang and Yongxin Chen. 'Fast Sampling of Diffusion Models with Exponential Integrator'. In International Conference on Learning Representations . 2023. URL: https: //openreview.net/forum?id=Loek7hfb46P .

Yang Song and Stefano Ermon. 'Generative Modeling by Estimating Gradients of the Data Distribution'. In Advances in Neural Information Processing Systems . Vol. 32. 2019. URL: https://arxiv.org/abs/1907.05600 .

Jonathan Ho et al. 'Denoising Diffusion Probabilistic Models'. In Advances in Neural Information Processing Systems . Vol. 33. 2020. URL: https://arxiv.org/abs/2006.11239 .

Brian Anderson. 'Reverse-time diffusion equation models'. In Stochastic Processes and their Applications 12.3 (1982), pp. 313-326.

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/hyvarinen05a .

Pascal Vincent. 'A Connection Between Score Matching and Denoising Autoencoders'. In Neural computation 23.7 (2011), pp. 1661-1674.

Jiaming Song et al. 'Denoising Diffusion Implicit Models'. In International Conference on Learning Representations . 2021. URL: https : / / openreview . net / forum ? id = St1giarCHLP .

Luping Liu et al. 'Pseudo Numerical Methods for Diffusion Models on Manifolds'. In International Conference on Learning Representations . 2022. URL: https://openreview.net/ forum?id=PlKWVd2yBkY .

Bahjat Kawar et al. 'SNIPS: Solving Noisy Inverse Problems Stochastically'. In Advances in Neural Information Processing Systems . 2021. URL: https://openreview.net/forum? id=pBKOx_dxYAN .

Yang Song et al. 'Solving Inverse Problems in Medical Imaging with Score-Based Generative Models'. In International Conference on Learning Representations . 2022. URL: https: //openreview.net/forum?id=vaRCHVj0uGI .

Alexandre Adam et al. 'Posterior samples of source galaxies in strong gravitational lenses with score-based priors'. 2022. URL: https://arxiv.org/abs/2211.03812 .

Hyungjin Chung et al. 'Diffusion Posterior Sampling for General Noisy Inverse Problems'. In International Conference on Learning Representations . 2023. URL: https://openreview. net/forum?id=OnD9zGAGT0k .

Bradley Efron. 'Tweedie's Formula and Selection Bias'. In Journal of the American Statistical Association 106.496 (2011), pp. 1602-1614.

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://openreview.net/forum?id=ZqEUs3sTRU0 .

Xiangming Meng and Yoshiyuki Kabashima. 'Diffusion Model Based Posterior Sampling for Noisy Linear Inverse Problems'. 2022. URL: https://arxiv.org/abs/2211.12343 .

Giorgio Parisi. 'Correlation functions and computer simulations'. In Nuclear Physics B 180.3 (1981), pp. 378-384.

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://www. jstor.org/stable/2346184 .

Edward N. Lorenz. 'Deterministic Nonperiodic Flow'. In Journal of Atmospheric Sciences 20.2 (1963), pp. 130-141.

Gary J. Chandler and Rich R. Kerswell. 'Invariant recurrent solutions embedded in a turbulent two-dimensional Kolmogorov flow'. In Journal of Fluid Mechanics 722 (2013), pp. 554-595.

Jun S. Liu and Rong Chen. 'Sequential Monte Carlo methods for dynamic systems'. In Journal of the American statistical association 93.443 (1998), pp. 1032-1044.

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: https://wrap.warwick. ac.uk/37961 .

Neil J Gordon et al. 'Novel approach to nonlinear/non-Gaussian Bayesian state estimation'. In Radar and Signal Processing . Vol. 140. 2. IET. 1993, pp. 107-113.

Alexander Quinn Nichol and Prafulla Dhariwal. 'Improved denoising diffusion probabilistic models'. In International Conference on Machine Learning . 2021. URL: https:// proceedings.mlr.press/v139/nichol21a .

Dmitrii Kochkov et al. 'Machine learning-accelerated computational fluid dynamics'. In Proceedings of the National Academy of Sciences 118.21 (2021). URL: https://www.pnas. org/doi/abs/10.1073/pnas.2101784118 .

Guido Boffetta and Robert E Ecke. 'Two-dimensional Turbulence'. In Annual Review of Fluid Mechanics 44 (2012), pp. 427-451.

Olaf Ronneberger et al. 'U-Net: Convolutional Networks for Biomedical Image Segmentation'. In Medical Image Computing and Computer-Assisted Intervention . 2015, pp. 234-241. URL: https://arxiv.org/abs/1505.04597 .

Yang Song et al. 'Consistency models'. 2023. URL: https://arxiv.org/abs/2210. 02303 .

Evan Archer et al. 'Black box variational inference for state space models'. 2015. URL: https://arxiv.org/abs/1511.07367 .

Chris J Maddison et al. 'Filtering variational objectives'. In Advances in Neural Information Processing Systems 30 (2017). URL: https://arxiv.org/abs/1705.09279 .

Christian Naesseth et al. 'Variational sequential monte carlo'. In International Conference on Artificial Intelligence and Statistics . 2018, pp. 968-977. URL: http://proceedings.mlr. press/v84/naesseth18a .

Dieterich Lawson et al. 'SIXO: Smoothing Inference with Twisted Objectives'. In Advances in Neural Information Processing Systems . 2022. URL: https://openreview.net/forum? id=bDyLgfvZ0qJ .

Rudolf E. Kalman. 'A New Approach to Linear Filtering and Prediction Problems'. In Journal of Basic Engineering 82.1 (1960), pp. 35-45.

Eric A. Wan and Rudolph Van Der Merwe. 'The unscented Kalman filter for nonlinear estimation'. In Proceedings of the IEEE Adaptive Systems for Signal Processing, Communications, and Control Symposium . IEEE. 2000, pp. 153-158.

Ienkaran Arasaratnam and Simon Haykin. 'Cubature Kalman Filters'. In IEEE Transactions on Automatic Control 54.6 (2009), pp. 1254-1269.

Rahul G. Krishnan et al. 'Deep Kalman Filters'. 2015. URL: https://arxiv.org/abs/ 1511.05121 .

Laurent Girin et al. 'Dynamical variational autoencoders: A comprehensive review'. 2020. URL: https://arxiv.org/abs/2008.12595 .

Suman Ravuri et al. 'Skilful precipitation nowcasting using deep generative models of radar'. In Nature 597.7878 (2021), pp. 672-677. URL: https://www.nature.com/articles/ s41586-021-03854-z .

Remi Lam et al. 'GraphCast: Learning skillful medium-range global weather forecasting'. 2022. URL: https://arxiv.org/abs/2212.12794 .

Kaiming He et al. 'Deep Residual Learning for Image Recognition'. In IEEE conference on computer vision and pattern recognition . 2016, pp. 770-778. URL: https://arxiv.org/ abs/1512.03385 .

Stefan Elfwing et al. 'Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning'. In Neural Networks 107 (2018), pp. 3-11. URL: https://arxiv.org/abs/1702.03118 .

Jimmy Lei Ba et al. 'Layer normalization'. 2016. URL: https://arxiv.org/abs/1607. 06450 .

Ilya Loshchilov and Frank Hutter. 'Decoupled Weight Decay Regularization'. In International Conference on Learning Representations . 2019. URL: https://openreview.net/forum? id=Bkg6RiCqY7 .

A The pseudo-blanket approximation is unbiased

For any continuous random variables a , b and c ,

Replacing a , b and c by x i ( t ) , x ¯ b i ( t ) and x ̸∈ ( t ) = { x j ( t ) : j = i ∧ j ̸∈ ¯ b i } , respectively, we obtain

̸

meaning that ∇ x i ( t ) log p ( x i ( t ) , x ¯ b i ( t )) is the expected value of ∇ x i ( t ) log p ( x 1: L ( t )) over x ̸∈ ( t ) . In other words, regardless of the elements or the size of ¯ b i , ∇ x i ( t ) log p ( x i ( t ) , x ¯ b i ( t )) is an unbiased (exact in expectation) estimate of ∇ x i ( t ) log p ( x 1: L ( t )) .

B On the covariance of p ( x | x ( t ))

Assuming a Gaussian prior p ( x ) with covariance Σ x , the covariance ˆ Σ of p ( x | x ( t )) takes the form

where Q Λ Q -1 is the eigendecomposition of Σ x . We observe that for most of the perturbation time t , the central diagonal term is close to Λ(Λ+ I ) -1 . We therefore propose the covariance approximation

where Γ = Q Λ(Λ + I ) -1 Q -1 is a positive semi-definite matrix.

C Algorithms

Algorithm 3 Estimating the posterior score ∇ x ( t ) log p ( x ( t ) | y )

Algorithm 4 Predictor-Corrector sampling from s ϕ ( x ( t ) , t | y )

D Experiment details

Resources Experiments were conducted with the help of a cluster of GPUs. In particular, score networks were trained and evaluated concurrently, each on a single GPU with at least 11 GB of memory.

Lorenz 1963 We consider two score network architectures: fully-connected local score networks for small k ( k ≤ 4 ) and fully-convolutional score networks for large k ( k > 4 ). Residual blocks [68], SiLU [69] activation functions and layer normalization [70] are used for both architecture. The number of blocks in the fully-convolutional architecture controls the value of k . We train all score networks for 1024 epochs with the AdamW [71] optimizer and a linearly decreasing learning rate. Other hyperparameters are provided in Table 1.

Kolmogorov flow The local score network is a 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 a drop of fluid is determined by its immediate environment. This can be seen as an application of the pseudo-blanket approximation (13) to a grid-structured set of random variables. We train the score network for 1024 epochs with the AdamW [71] optimizer and a linearly decreasing learning rate. Other hyperparameters are provided in Table 2.

E Assimilation examples

x

1:

L

y

SDA

DPS

S553

x

1:

L

y

SDA

DPS

x

1:

L

y

2

SDA

DPS

y

4

SDA

DPS

y

8

SDA

DPS

y

16

SDA

DPS

x

1:

L

y

SDA

x

1:

L

y

SDA

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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)?\n\nLIMITATIONS:\n- Limitations were stated in the main text.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n- 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)", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUnder 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nWeaknesses:  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\nQUESTIONS:\nCan 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\nLIMITATIONS:\nAs mentioned, the biggest limitation is assumption of the underlying dynamics.  \n\nThere are not many comparisons against other methods made.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-14"}