State-observation augmented diffusion model for nonlinear assimilation

We start from considering the data distribution we aim to learn, denoted as $\bm{x}_{0}\sim p(\bm{x}_{0})$. By progressively adding Gaussian noise, the distribution undergoes an approximate transformation into a standard normal distribution, and the whole process can be modeled as a forward-time diffusion process. Mathematically, following the notation in [50], we describe the diffusion process as a linear stochastical differential equation (SDE) given by

where $f(t)$ and $g(t)$ represent the drift term and diffusion coefficient, respectively, and $\bm{w}_{t}$ signifies the standard Wiener process. Denoting by $p_{st}(\bm{x}_{t}\mid\bm{x}_{s})$ the transition kernel from $\bm{x}_{s}$ to $\bm{x}_{t}$ for $0\leq s<t\leq 1$, we can derive the conditional distribution

with the mean $\mu_{t}$ and the variance $\sigma_{t}^{2}$ explicitly [51] given by

With appropriate choices of $f(t)$ and $g(t)$, the forward diffusion process (2) transforms the data distribution into standard Gaussian noise $p(\bm{x}_{1})\approx\mathcal{N}(0,\bm{I})$. Under such circumstances, the reverse-time evolution, characterized by

is also a diffusion process [3]. Consequently, the reverse-time evolution can be used to generate samples by transforming noise $\bm{x}_{1}\sim\mathcal{N}(0,\bm{I})$ back into data distribution $\bm{x}_{0}\sim p(\bm{x}_{0})$, provided the score function of each marginal distribution $\nabla_{\bm{x}_{t}}\log p_{t}(\bm{x}_{t})$ is available for all $t$.

To develop an effective score-based generative model, the key lies in minimizing the denoising score matching loss

as discussed in [54]. In practice, $\bm{s}_{\theta}(\bm{x}_{t},t)$ is usually parameterized as a denoising network

and we can rewrite the optimization problem as

Upon successful training of the denoising network, we are abole to generate samples $\bm{x}_{0}\sim p_{0}(\bm{x}_{0})$ by evolving the reverse-time process from an initial random sample $\bm{x}_{1}\sim\mathcal{N}(0,\bm{I})$. Classical numerical solvers for general SDEs such as Euler-Maruyama and stochastic Runge-Kutta methods, and the ancestral sampling [21, 50] are appliable. Recently, to increase the sampling efficiency, the exponential integrator (EI) discretization scheme [58]

has also been proposed. These schemes often work together with the predictor-corrector sampling strategy by performing a few steps of Langevin Monte Carlo (LMC) sampling [50]

for better performances.

While score-based generative models we have introduced above can be directly used for unconditional sampling with well-trained denoising models, assimilation tasks require estimating the conditional distribution $p(\bm{x}\mid\bm{y})$, where we employ the abbreviations $\bm{x}=\bm{x}_{0}=\{\bm{X}_{k}\}_{k=S}^{S+L}$ and $\bm{y}=\{\bm{Y}_{k}\}_{k=S}^{S+L}$. Note that the conditional score function can be decomposed as

by Bayes’ formula, where the additional term is usually called the adversarial gradient

To incorporate conditional information, one may introduce a separate network for the adversarial gradient alongside the original denoising network as suggested in [15, 22]. Alternatively, a conditional denoising network $\bm{\varepsilon}_{\theta}(\bm{x}_{t},t;\bm{y})$ can be trained directly to approximate $-\sigma_{t}\nabla_{\bm{x}_{t}}\log p_{t}(\bm{x}_{t}\mid\bm{y})$, and such an idea has been widely used in many class-conditional generation and text-to-image/video generation tasks [42, 47].

Instead of training a different denoising network, some other approaches focus on estimating the reverse-time transition kernel $p(\bm{x}_{0}\mid\bm{x}_{t})$ in (12) to compute the time-dependent likelihood $p_{t}(\bm{y}\mid\bm{x}_{t})$ explicitly. By Tweedie’s formula, DPS [11] proposes a delta-function approximation

where the estimator $\hat{\bm{x}}_{0}(\bm{x}_{t},t)$ is defined as

to approximate the conditional mean

Based on the uninformative-prior assumption, DMPS [36] replaces the transition kernel with a Gaussian distribution

to take the uncertainty into consideration. In particular, for a linear Gaussian observational model

the DPS and DMPS estimators become

by substituting (14), (16) and (17) into (12). SDA [43, 44] suggests the approximation

to seek a balance between stability and accuracy, where the hyperparameter $\gamma$ related to the variance of the data distribution $p(\bm{x}_{0})$ is tuned empirically. Nonetheless, such an approximation is still based on the linearity assumption for the observational operator, which has shown limitations in our experiments.

To handle the potential nonlinearity of the observational model, we propose to consider an augmented version of (1). Let $\bm{W}_{k}$ denote the concatenation of the state variables $\bm{X}_{k}$ and all possible observation variables $\mathcal{H}(\bm{X}_{k})$, then we have

For the observed states, there exists $\bm{T}_{k}$ linked to $\mathcal{H}_{k}$ such that

In this way, we have essentially established an equivalent augmented system

Clearly, the forward propagation operator $\mathcal{M}_{k}^{W}$ inherits the periodicity of $\mathcal{M}_{k}^{X}$. The observational operator $\bm{T}_{k}$ corresponding to $\bm{W}_{k}$ is now a time-dependent subsampling matrix, which simplifies the assimilation process.

By concatenating $\bm{W}$ and $\bm{Y}$ as $\tilde{\bm{W}}$ and $\tilde{\bm{Y}}$ over consecutive time steps, respectively, the assimilation problem we focus on is transformed into a posterior estimation for $p\left(\tilde{\bm{W}}\mid\tilde{\bm{Y}}\right)$. The prior $p\left(\tilde{\bm{W}}\right)$ implicitly encodes the physical dynamics as well as the model error, and the observational model becomes a linear Gaussian model

with $\tilde{\bm{T}}$ and $\tilde{\bm{\varepsilon}}^{O}$ representing the concatenated versions of the subsampling operator $\bm{T}$ and the observation noise $\bm{\varepsilon}^{O}$, respectively.

In general, our State-Observation Augmented Diffusion (SOAD) model follows the principles of score-based generative models described above. Our SOAD model estimates the joint unconditional distribution for $\tilde{\bm{W}}$, which includes the physical states to be assimilated, the clean observation outputs as well as their interactions induced by the observational operator. For notation convenience, we alias $\tilde{\bm{W}}$ and $\tilde{\bm{Y}}$ with $\tilde{\bm{x}}$ and $\tilde{\bm{y}}$, respectively, and we use $p(\tilde{\bm{x}})$ and $p(\tilde{\bm{y}})$ for the probability densities of the augmented $\tilde{\bm{W}}$ and the noised observation $\tilde{\bm{Y}}$.

To establish the conditional diffusion model, we continue from the Bayes’ rule

adapted from (11)(12) in the current context, where we still use the notation

The prior component $\log p_{t}(\tilde{\bm{x}}_{t})$ can be provided by a trained diffusion model for $p(\tilde{\bm{x}})$ as

according to (6)(7)(8), so it suffices to obtain a good estimation for the time-dependent likelihood component $\nabla_{\tilde{\bm{x}}_{t}}\log p_{t}(\tilde{\bm{y}}\mid\tilde{\bm{x}}_{t})$,

By Applying Bayes’ rule again, we have

up to a constant. If we consider a Gaussian prior for $p(\tilde{\bm{x}}_{0})$ with a covariance matrix $\bm{\Sigma}_{x}$, then $p(\tilde{\bm{x}}_{0}\mid\tilde{\bm{x}}_{t})$ is also Gaussian, whose covariance matrix has a closed form

Similar to [11], the mean of $p(\tilde{\bm{x}}_{0}\mid\tilde{\bm{x}}_{t})$ may be provided by Tweedie’s formula (15). Therefore, we use the following estimation for the reverse-time transition kernel

It follows that $p_{t}(\tilde{\bm{y}}\mid\tilde{\bm{x}}_{t})$ can be computed analytically once the observation likelihood $p(\tilde{\bm{y}}\mid\tilde{\bm{x}}_{0})$ is provided explicitlys. For instance, given the observational model with Gaussian noise

the time-dependent likelihood becomes

To sum up, we have

In practice, we treat $\sigma_{x}$ as a hyperparameter to be tuned as it is hard to evaluate explicitly for a trained diffusion model. Meanwhile, a certain clipping mechanism for the covariance is required to stabilize the generation process when the diffusion step $t$ is near 1. Fortunately, the model performances are not sensitive to the choice of $\sigma_{x}$ as well as the clipping mechanism. See Algorithm 1 for one of the possible implementations.

Since in our augmented model, $\tilde{\bm{Y}}$ is a linear transformation of $\tilde{\bm{W}}$, or more specifically, a subsampling version of $\tilde{\bm{W}}$ up to an observational error, part of the denoised results can be directly replaced by a perturbed version of $\tilde{\bm{Y}}$ to stabilize the reverse-time diffusion process. To be more precise, let $\tilde{\bm{x}}_{t}$ be the diffusion state at time $t$ with the initial state $\tilde{\bm{x}}_{0}=\tilde{\bm{W}}$, then

with a standard Gaussian noise $\bm{\varepsilon}$. We propose to substitute $\tilde{\bm{T}}\tilde{\bm{x}}_{t}$ at diffusion time $t$ with $\mu_{t}(\tilde{\bm{Y}}+\bm{\varepsilon}_{t}^{\prime})$, where $\bm{\varepsilon}_{t}^{\prime}\sim\mathscr{D}_{t}^{\prime}$. To align their means and covariances, we set

whenever $\sigma_{t}^{2}\bm{I}-\mu_{t}^{2}\operatorname{\mathrm{Cov}}\tilde{\mathscr{D}}^{O}$ is positive-semidefinite.

We need to emphasize that, such a condition only fails for small time steps especially when the observation noise is not too large since $\lim_{t\to 1^{-}}r_{t}=+\infty$ under typical choices of diffusion schedulers. In other words, our replacements provide reliable estimation of $\tilde{\bm{T}}\tilde{\bm{x}}_{t}$ most of the time.

We describe the detailed denoising procedure in this subsection, where we set $\tilde{\mathscr{D}}^{O}=\mathcal{N}(0,(\sigma^{o})^{2}\bm{I})$ for simplicity. Similar deduction can be applied to other more complex settings.

Let $r_{\sigma}=\sigma_{t_{-}}/\sigma_{t}$ and $r_{\mu}=\mu_{t_{-}}/\mu_{t}$ be the ratios of the diffusion parameters at adjoint time step $t_{-}$ and $t$, respectively. We choose the EI discretization scheme (9)

for the reverse-time stepper, where we define $\bm{\pi}_{\theta}(\tilde{\bm{x}}_{t},t)=-\sigma_{t}^{2}\nabla_{\tilde{\bm{x}}_{t}}\log p_{t}(\tilde{\bm{x}}_{t}\mid\tilde{\bm{y}})$. By plugging (24), (26) and (33) into the definition of $\bm{\pi}_{\theta}$, we have

The whole denoising process is summarized in Algorithm 1.

In practical applications, assimilation is usually performed sequentially once new observations become available, so we start with the classical assimilation settings that additionally provide the models with a background prior for the first step. We choose the background prior as an ideal Gaussian perturbation with a known covariance $(\sigma^{b})^{2}=0.1^{2}$. Formally, the observational model for the first step is modified as

where $\bm{Y}_{0}$ is the background prior with noise variance $0.1^{2}$. See Figure 2 (3rd and 4th rows) for visualization.

Our experiments start with the element-wise observations $\mathcal{H}^{e}$ and $\mathcal{H}^{h}$. The averaged RMSEs over all the 9 time steps for the assimilated states are summarized in Figure 3.

Our SOAD model performs slightly worse as the time interval $N$ increases and the observational ratio $p$ or $(1/s)^{2}$ decreases, which aligns with our intuition that the assimilation task becomes more difficult with sparser observations. In the most challenging case, where the observations are available only every 8 time steps with a ratio of $1\%$, the RMSEs are around 0.2 and 0.3 for the easy and hard observations, respectively. This indicates that our method is still able to extract the physical features from rare observations to some extent. In contrast, the SDA baseline shows a significant performance drop even when the observations are available everywhere for each time step. The inferior performance is likely due to the linearity assumption for the observation operator in the SDA formulations, which is unsuitable for highly nonlinear observations.

Additionally, the step-wise RMSEs for our SOAD model are exhibited in Figure 4. In both observation cases, The increments of RMSEs at observational ratios of 100% and 25% are mild, and no significant changes in RMSEs are observed when the easier $\mathcal{H}^{e}$ is replaced with the harder $\mathcal{H}^{h}$. Moreover, relatively lower assimilation errors are obtained whenever observational data are available, supporting the idea that more observations enhance the assimilation process. By comparing results with the same observational ratios ($p=(1/s)^{2}$), we may conclude that regular (uniform) observations are more beneficial for the assimilation process than irregular (random) ones.

To explore the potential of our SOAD model in real applications, we also test it with the vorticity-to-velocity mapping $\mathcal{H}^{\mathsf{v2v}}$, and the results are shown in Figure 5. We do not include the performance of the SDA baseline since it diverges in all cases. Performances of the SDA baseline for $\mathcal{H}^{\mathsf{v2v}}$ are not included since it ends with divergence for all the cases. The averaged RMSEs in Figure 5(a) are lower than those in Figure 3(a) and Figure 3(c), which may be attributed to the linearity of the observational operator $\mathcal{H}^{\mathsf{v2v}}$. The results indicate that our SOAD approach is capable of handling the vorticity-to-velocity mapping,a more practical observational model in real applications. Although SOAD shows faster increasing RMSEs with higher uncertainty for $\mathcal{H}^{\mathsf{v2v}}$ (Figure 5(b)) compared with $\mathcal{H}^{e}$ and $\mathcal{H}^{h}$ (Figure 4), the subsequent experiments may offer some evidences that the assimilation process is unlikely to diverge.

In practice, we cannot always expect the availability of the exact probability distribution for background prior. Background errors may result from many aspects, such as unresolved sub-scale features, temporal interpolations and physical parameterizations. Therefore, we investigate whether the discrepancy between our background prior model (46) and the real distribution affects the assimilation performance. We conduct experiments with the same settings as before but feed a background prior with unknown noise

into the model instead of $\bm{Y}_{0}$. Here, $\mathcal{M}$ is the forward propagation model of the QG equations, and $\bm{X}_{-1}$ stands for the snapshot at the previous time step. Comparisons between the step-wise RMSEs are displayed in Figure 6, where we fix $N=1$ and $p=0.25$. Even with incorrect background prior assumptions, our SOAD model has the ability of self-correction with a few time steps of observations. Another interesting observation is that the RMSEs for the hard observation $\mathcal{H}^{h}$ are lower than those for $\mathcal{H}^{e}$ and $\mathcal{H}^{\mathsf{v2v}}$. One possible explanation is that when the observational data is sufficient for the model to capture the dynamics, observations from a harder operator $\mathcal{H}^{h}$ may provide more detailed information since $\mathcal{H}^{h}$ is much more sensitive to the input.

To explore the assimilation performance in the absence of any background prior, we remove the modification (46) and conduct the assimilation process solely from observational data. Even if starting with a background prior, any assimilation method will gradually forget the initial state information as more observations are assimilated. Therefore, the results shown in this section are likely to reveal the long-term behaviors of the assimilation methods.

All assimilation results are recorded as heatmaps of RMSEs in Figure 7. In general, our SOAD model performs well for both the $\mathcal{H}^{e}$ and the $\mathcal{H}^{\mathsf{v2v}}$ observations, with slightly higher accuracy for the latter. The observation $\mathcal{H}^{h}$ appears to be the most challenging across all the testing observational operators.

To visualize the assimilated states, we have plotted the assimilated vorticities with random observational ratios $p\in\{0.25,0.0625,0.01\}$ in Figure 8, Figure 9 and Figure 10 for $\mathcal{H}^{e}$, $\mathcal{H}^{h}$ and $\mathcal{H}^{\mathsf{v2v}}$, respectively. For clarity, we only include the vorticities and the associated observations for the first layer, as they contain more small-scale details. For the $\mathcal{H}^{e}$ observation, our SOAD is capable of recovering most of the physical features until $p$ reaches $1\%$, in which case the assimilated states still share much similarity with the ground truth. With the harder $\mathcal{H}^{h}$, the SOAD model fails to reconstruct the real system states when $p\leq 0.0625=(1/4)^{2}$, consistent with the previous results shown in Figure 7(b). Surprisingly, the assimilated states with rare observations do not collapse and still follow certain dynamics, suggesting that the lack of observational data, rather than incapacity of learning the dynamics, leads to the failure. Finally, when the observation is the the vorticity-to-velocity mapping $\mathcal{H}^{\mathsf{v2v}}$, the assimilated states are closer to the ground truth even when observations are available at only $1\%$ of the grids. We believe that the success may result from the linearity of $\mathcal{H}^{\mathsf{v2v}}$. The results indicate that our SOAD approach is stable when handling highly nonlinear observations in the long term performs well on non-element-wise observations such as the vorticity-to-velocity mapping $\mathcal{H}^{\mathsf{v2v}}$, which is a promising sign for real applications.