SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 2: Average posterior summary statistics over 64 observations from the low (A) and high (B) frequency observation processes. We observe that, as  $k$  and the number of corrections  $C$  increase, the statistics of the approximate posteriors get closer to the ground-truth, in red, which means they are getting more accurate. However, increasing  $k$  and  $C$  improves the quality of posteriors with decreasing return, such that all posteriors with  $k \geq 3$  and  $C \geq 2$  are almost equivalent. This is visible in C, where we display trajectories inferred ( $C = 2$ ) for an observation of the low frequency observation process. For readability, we allocate a segment of 32 states to each  $k$  instead of overlapping all 192 states. Note that the Wasserstein distance between the ground-truth posterior and itself is not zero as it is estimated with a finite number (1024) of samples.
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' \mid y) = \mathbb{E}_{q(x_{1:L}|y)}[p(y' \mid x_{1:L})]$ .
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 \leq 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  $\mathcal{N}(y \mid \tilde{a}_{1:L:8}, 0.05^2 I)$  and  $\mathcal{N}(y \mid \tilde{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} \mid 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} \mid y)$  defined by each score network. We compare the approximate and ground-truth posteriors with three summary statistics: the expected log-prior  $\mathbb{E}_{q(x_{1:L}|y)}[\log p(x_{2:L} \mid x_1)]$ , the expected log-likelihood  $\mathbb{E}_{q(x_{1:L}|y)}[\log p(y \mid 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 ( $\tau = 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 \mid x_{1:L}) = \mathcal{N}(y \mid \tilde{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
$$
\dot {\boldsymbol {u}} = - \boldsymbol {u} \nabla \boldsymbol {u} + \frac {1}{R e} \nabla^ {2} \boldsymbol {u} - \frac {1}{\rho} \nabla p + \boldsymbol {f} \tag {19}
$$
$$
0 = \nabla \cdot \boldsymbol {u}
$$
where  $\pmb{u}$  is the velocity field,  $Re$  is the Reynolds number,  $\rho$  is the fluid density,  $p$  is the pressure field and  $\pmb{f}$  is the external forcing. Following Kochkov et al. [53], we choose a two-dimensional domain  $[0,2\pi]^2$  with periodic boundary conditions, a large Reynolds number  $Re = 10^3$ , a constant density  $\rho = 1$  and an external forcing  $\pmb{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\times 256$  domain grid. The states  $x_{i}$  are snapshots of the velocity field  $\pmb{u}$ , coarsened to a  $64\times 64$  resolution, and the integration time between two such snapshots is  $\Delta = 0.2$  time units. This corresponds to 82
Figure 4: Example of sampled trajectory from coarse, intermittent and noisy observations. States are visualized by their vorticity field  $\omega = \nabla \times \mathbf{u}$ , that is the curl of 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 one of DPS [41] yields trajectories inconsistent with the observation.
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 \times 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 \times 8$  and perturbed by a moderate Gaussian noise  $(\Sigma_y = 0.1^2 I)$ . Given the observation, we sample a trajectory with SDA  $(C = 1, \tau = 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.
Figure 5: A trajectory consistent with an unlikely observation of the final state  $x_{L}$  is generated with SDA. To verify whether the trajectory is realistic and not hallucinated, we plug its initial state  $x_{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  $x_{L}$  in a circle-shaped sub-domain. Then, we sample a trajectory  $(C = 1, \tau = 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.
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  $\nabla_{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.
François Rozet is a research fellow of the F.R.S.-FNRS (Belgium) and acknowledges its financial support.
[1] Andrew C Lorenc. “Analysis methods for numerical weather prediction”. In Quarterly Journal of the Royal Meteorological Society 112.474 (1986), pp. 1177–1194.
[2] 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.
[3] 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.
[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 an imperfect model in 4D-Var." In ECMWF Technical Memoranda 477 (2005), p. 20. URL: https://www.ecmwf.int/node/12845.
[6] Yannick Trémolet. "Model error estimation in 4D-Var". In ECMWF Technical Memoranda 520 (2007), p. 19. URL: https://www.ecmwf.int/node/12819.
[7] 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.
[8] 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).
[9] ECMWF. "Part II: Data Assimilation". In IFS Documentation CY47R1 2 (2020). URL: https://www.ecmwf.int/node/19746.
[10] 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.
[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.
[70] Jimmy Lei Ba et al. "Layer normalization". 2016. URL: https://arxiv.org/abs/1607.06450.
[71] 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$ ,
$$
\begin{array}{l} \nabla_ {a} \log p (a, b) = \frac {1}{p (a , b)} \nabla_ {a} p (a, b) \\ = \frac {1}{p (a , b)} \nabla_ {a} \int p (a, b, c) \mathrm {d} c \\ = \int \frac {p (c \mid a , b)}{p (a , b , c)} \nabla_ {a} p (a, b, c) d c \\ = \mathbb {E} _ {p (c \mid a, b)} \left[ \nabla_ {a} \log p (a, b, c) \right]. \\ \end{array}
$$
Replacing  $a, b$  and  $c$  by  $x_{i}(t), x_{\bar{b}_{i}}(t)$  and  $x_{\not\in}(t) = \{x_{j}(t): j \neq i \land j \notin \bar{b}_{i}\}$ , respectively, we obtain
$$
\nabla_ {x _ {i} (t)} \log p (x _ {i} (t), x _ {\bar {b} _ {i}} (t)) = \mathbb {E} _ {p (x _ {\notin} (t) | x _ {i} (t), x _ {\bar {b} _ {i}} (t))} \big [ \nabla_ {x _ {i} (t)} \log p (x _ {1: L} (t)) \big ],
$$
meaning that  $\nabla_{x_i(t)}\log p(x_i(t),x_{\bar{b}_i}(t))$  is the expected value of  $\nabla_{x_i(t)}\log p(x_{1:L}(t))$  over  $x\notin (t)$ . In other words, regardless of the elements or the size of  $\bar{b}_i$ ,  $\nabla_{x_i(t)}\log p(x_i(t),x_{\bar{b}_i}(t))$  is an unbiased (exact in expectation) estimate of  $\nabla_{x_i(t)}\log p(x_{1:L}(t))$ .
B On the covariance of  $p(x\mid x(t))$
Assuming a Gaussian prior  $p(x)$  with covariance  $\Sigma_{x}$ , the covariance  $\hat{\Sigma}$  of  $p(x \mid x(t))$  takes the form
$$
\begin{array}{l} \hat {\boldsymbol {\Sigma}} = \boldsymbol {\Sigma} _ {x} - \boldsymbol {\Sigma} _ {x} \left(\boldsymbol {\Sigma} _ {x} + \frac {\sigma (t) ^ {2}}{\mu (t) ^ {2}} I\right) ^ {- 1} \boldsymbol {\Sigma} _ {x} \\ = \frac {\sigma (t) ^ {2}}{\mu (t) ^ {2}} Q \Lambda \left(\Lambda + \frac {\sigma (t) ^ {2}}{\mu (t) ^ {2}} I\right) ^ {- 1} Q ^ {- 1} \\ \end{array}
$$
where  $Q\Lambda Q^{-1}$  is the eigendecomposition of  $\Sigma_{x}$ . We observe that for most of the perturbation time  $t$ , the central diagonal term is close to  $\Lambda (\Lambda +I)^{-1}$ . We therefore propose the covariance approximation
$$
\hat {\Sigma} = \frac {\sigma (t) ^ {2}}{\mu (t) ^ {2}} \Gamma
$$
where  $\Gamma = Q\Lambda (\Lambda +I)^{-1}Q^{-1}$  is a positive semi-definite matrix.
Algorithm 3 Estimating the posterior score  $\nabla_{x(t)}\log p(x(t)\mid y)$
1 function  $s_{\phi}(x(t),t\mid y)$    
2  $s_x\gets s_\phi (x(t),t)$    
3  $\hat{x}\gets \frac{x(t) + \sigma(t)^2s_x}{\mu(t)}$    
4  $s_y\gets \nabla_{x(t)}\log \mathcal{N}\left(y\mid \mathcal{A}(\hat{x}),\Sigma_y + \frac{\sigma(t)^2}{\mu(t)^2}\Gamma\right)$    
5 return  $s_x + s_y$
Algorithm 4 Predictor-Corrector sampling from  $s_{\phi}(x(t), t \mid y)$
1 function  $\mathrm{SAMPLE}(\{t_i\}_{i = 0}^N,C,\tau)$    
2  $x(1)\sim \mathcal{N}(0,\Sigma (1))$    
3 for  $i = N$  to 1 do   
4  $x(t_{i - 1})\gets \frac{\mu(t_{i - 1})}{\mu(t_i)} x(t_i) + \left(\frac{\mu(t_{i - 1})}{\mu(t_i)} -\frac{\sigma(t_{i - 1})}{\sigma(t_i)}\right)\sigma (t_i)^2 s_\phi (x(t_i),t_i\mid y)$    
5 for  $j = 1$  to  $C$  do   
6  $\epsilon \sim \mathcal{N}(0,I)$    
7  $s\gets s_{\phi}(x(t_{i - 1}),t_{i - 1}\mid y)$    
8  $\delta \gets \tau \frac{\dim(s)}{\|s\|_2^2}$    
9  $x(t_{i - 1})\gets x(t_{i - 1}) + \delta s + \sqrt{2\delta}\epsilon$    
10 return  $x(0)$
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 \leq 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.
Table 1: Score network hyperparameters for the Lorenz experiment.
Hyperparameter	k≤4	k>4
Architecture	fully-connected	fully-convolutional
Residual blocks	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-3	10-3
Scheduler	linear	linear
Epochs	1024	1024
Batches per epoch	256	256
Batch size	256	64
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.
Table 2: Score network hyperparameters for the Kolmogorov experiment.
Architecture	U-Net
Residual blocks per level	(3, 3, 3)
Channels per level	(96, 192, 384)
Kernel size	3
Padding	circular
Activation	SiLU
Normalization	LayerNorm
Optimizer	AdamW
Weight decay	10-3
Learning rate	2 × 10-4
Scheduler	linear
Epochs	1024
Batches per epoch	128
Batch size	32
E Assimilation examples
Figure 6: Example of sampled trajectory when the observation is insufficient to identify the original trajectory. The observation is the center of the velocity field every three steps, coarsened to a resolution  $8 \times 8$  and perturbed by a small Gaussian noise  $(\Sigma_y = 0.01^2 I)$ . SDA  $(C = 1, \tau = 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 a trajectory less consistent with the observation.
Figure 7: Example of sampled trajectory when the observation process is non-linear. The observation corresponds to a saturating transformation  $x \mapsto \frac{x}{1 + |x|}$  of the vorticity field  $\omega$ . SDA (512 discretization steps,  $C = 1$ ,  $\tau = 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 a trajectory inconsistent with the observation.
Figure 8: Example of sampled trajectories when the observation is spatially sparse. The observation  $y^{n}$  corresponds to a spatial subsampling of factor  $n$  of the velocity field. SDA ( $C = 1, \tau = 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] leads to trajectories that are progressively less consistent with the observation as  $n$  increases.
Figure 9: Example of long  $(L = 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 to zero, which enforces a looping trajectory  $(x_{1} \approx x_{L})$ . Note that this scenario is not realistic and will therefore lead to physically implausible trajectories.
Figure 10: Example of sampled trajectories when the observation is insufficient to identify the original trajectory. The observation is the center of the velocity field every three steps, coarsened to a resolution  $8 \times 8$  and perturbed by a small Gaussian noise  $(\Sigma_y = 0.01^2 I)$ . SDA  $(C = 1, \tau = 0.5)$  identifies the original state where data is sufficient, while generating physically plausible states elsewhere.
Figure 11: Example of sampled trajectories when the observation is spatially sparse. The observation  $y$  corresponds to a spatial subsampling of factor 16 of the velocity field with small Gaussian noise  $(\Sigma_y = 0.01^2 I)$ . SDA  $(C = 1, \tau = 0.5)$  generates trajectories similar to the original, with physically plausible variations.