Deep Bayesian Filter for Bayes-Faithful Data Assimilation

A physical system is defined by variables $z_{t}$, with its evolution described by the dynamics model $p(z_{t+1}|z_{t})=\mathcal{N}(z_{t+1};f(z_{t}),Q)$, where $\mathcal{N}(x|\mu,\Sigma)$ denotes a Gaussian whose mean and covariance are $\mu$ and $\Sigma$. The nonlinear function $f$ is the dynamics operator and $Q$ is the system covariance. The Markov property holds, as $z_{t+1}$ depends only on $z_{t}$. An observation model $p(o_{t}|z_{t})=\mathcal{N}(o_{t};h(z_{t}),R)$ relates observations to physical variables via the observation operator $h$ and covariance $R$. The panel (a) of Fig. 1 shows the system’s graphical model. The objective of sequential DA is to compute the posterior of $z_{t}$ given $o_{1:t}$.

In the KF, the dynamics and observation models are both linear Gaussian, with the mean being linear with respect to the conditional variable and a constant covariance matrix. Given that the dynamics and observation operators $f,h$ are linear, we can represent them using matrices $A$ and $C$, respectively. All matrices ($A,C,Q$, and $R$) are constant. The filter distribution $p(z_{t}|o_{1:t})$ remains Gaussian, provided that the initial distribution $p(z_{1})$ is Gaussian. We can recursively compute the posterior parameters (means $\mu_{t}$ and covariance matrices $\Sigma_{t}$) using the following equations:

where $K_{t}=(A\Sigma_{t-1}A^{T}+Q)H^{T}(H(A\Sigma_{t-1}A^{T}+Q)H^{T}+R^{-1})^{-1}$ is the Kalman Gain.

In this scenario, Gaussianity of the test distribution is lost during the KF update step. We introduce an inverse observation operator (IOO) $r(z_{t}|o_{t})$ (see also Frerix et al. 2021):

where $r(z_{t}|o_{t})=\frac{p(o_{t}|z_{t})\rho(z_{t})}{\int p(o_{t}|z_{t})\rho(z_{t})dz_{t}}$ and $\rho(z_{t})$ is a virtual prior. By approximating both the IOO and the virtual prior as Gaussians, $r(z_{t}|o_{t})=\mathcal{N}(f_{\theta}(o_{t}),G_{\theta}(o_{t}))$ and $\rho(z_{t})=\mathcal{N}(m,V)$, respectively, the posterior $q(z_{t}|o_{1:t})$ can be analytically computed as a Gaussian, where the mean $\mu_{t}$ and covariance $\Sigma_{t}$ are given as:

where $f_{\theta}(o_{t})$ and $G_{\theta}(o_{t})$ are NNs with parameters $\theta$, and $m$ and $V$ are constants set to $m=0$ and $V=10^{8}I$. These values bias the NNs’ outputs without affecting performance.

The recursive formula for the exact posterior (Equation 3) requires no approximation. Thus, DBF computes the exact posterior when the true IOO $r_{\rm true}(h_{t}|o_{t})$ is Gaussian, i.e., the SSM is a LGSS. In that case, the posterior update formula agrees with the KF (see Equations 1, 2 and 4, 5). The key difference is that nonlinear functions are applied to both the mean, $f_{\theta}(o_{t})$, and the covariance, $G_{\theta}(o_{t})$. In the KF, $f_{\theta}(o_{t})$ is linear, and $G_{\theta}(o_{t})$ is a constant covariance matrix (see Equations 1 and 2). The dependence of $G_{\theta}(o_{t})$ on observations provides flexibility in adjusting the impact of the new observation on the state estimation. The importance of adjusting the internal state updates based on observations has also been discussed in recent SSM-based approaches (Gu & Dao, 2023).

In this scenario, the Gaussianity of the test distribution is lost during the predict step, making it impossible to apply the original dynamics over the physical variables $z_{t}$. Therefore, we introduce a new set of latent variables $h_{t}$ and assume a dynamics model over $h_{t}$: $p(h_{t+1}|h_{t})=\mathcal{N}(h_{t+1}|Ah_{t},Q)$ (see panel (b) in Fig. 1). The IOO maps observations into the latent variables $h_{t}$: $r(h_{t}|o_{t})$. The recursive formula follows Equations 4 and 5. To retrieve the distribution of the original physical variables $z_{t}$, we introduce an emission model $p(z_{t}|h_{t})=\mathcal{N}(z_{t};\phi(h_{t}),R)$, where $\phi$ is represented by a NN. By marginalizing over $h_{t}$ with this emission model, a trained DBF can generate samples of $z_{t}$ that follow the test distribution $q(z_{t}|o_{1:t})$ given observations $o_{1:t}$.

Although the dynamics operator $A$ for the latent variables $h_{t}$ is linear, it is able to express any nonlinear dynamics if the latent space is sufficiently high-dimensional. The Koopman operator (Koopman, 1931) provides a framework for representing nonlinear systems by mapping observables—functions of the system’s state—into a higher-dimensional space where the dynamics become linear. Mathematically, for a dynamical system $z_{t+1}=f(z_{t})$, the Koopman operator $\mathcal{K}$ is a linear operator acting on a set of observables $g(z)$, such that $\mathcal{K}g(z_{t})=g(f(z_{t}))$. This allows the system to be represented as $h_{t+1}=Ah_{t}$ in the latent space $h_{t}$, where $A$ is the dynamics matrix learned by DBF. While the physical dynamics $f(z)$ are nonlinear, the Koopman operator ensures the existence of an embedding that linearizes the dynamics, enabling recursive computation of test distributions. Discovering such embeddings in finite dimensions has been widely studied (Takeishi et al., 2017; Lusch et al., 2018; Azencot et al., 2020). In high-dimensional simulations, the true degrees of freedom are often far fewer than the simulated variables, making surrogate modeling with the Koopman operator a promising approach to reducing computational costs.

When assimilating in the physical space (i.e., when the dynamics are linear), we train the IOO (i.e., $f_{\theta}$ and $G_{\theta}$) by optimizing the evidence lower bound (ELBO) without using the teacher signal $z_{t}$:

where $KL[p||q]$ denotes the Kullback-Leibler divergence between distributions $p$ and $q$. If the SSM contains any unknown parameters, we can train these parameters as well.

For SSMs with nonlinear or unknown dynamics, we have two approaches:

We discuss related works in the order of relevance.

This experiment demonstrates DBF’s ability to handle linear dynamics where key parameters of the observation operator are unknown. The dataset consists of 2D figures containing two embedded images, each moving at a constant speed and bouncing off frame edges. The system’s physical state is described by eight variables: the positions $(x,y)$ and velocities $(v_{x},v_{y})$ of the two embedded images. The dynamics matrix is block-diagonal, composed of four (two-body times two dimensions) translation matrices, $A_{tr}$: $A_{tr}=\begin{pmatrix}1&1\\ 0&1\\ \end{pmatrix},z_{t}=\begin{pmatrix}x_{t}\\ v_{x_{t}}\\ \end{pmatrix}$. Observations are corrupted by additive Gaussian noise with a standard deviation of $\sigma=50$ per pixel, where the original pixel values range from $0$ to $255$ (see panel (a) of Fig. 2 for an example of the data provided).

The aim is to show that DBF can track the linear dynamics while estimating unknown system parameters. DBF learns the pixel values of the embedded images from noisy observations, while maintaining consistency with physical motion. The observation model contains $1{,}568$ unknown parameters, corresponding to the number of pixels in the images. In classical DA algorithms, it is not possible to train unknown system parameters. However, it may be possible to infer these parameters by incorporating them as new physical dimensions. We have adopted this strategy for classical DA algorithms (EnKF, ETKF, and PF). We tested at three different noise levels ($\sigma_{sys}$ of $1$, $0.1$, and $0.01$). DVAEs are not comparable as they need to undergo supervised training. We were unable to compare with KalmanNet, as its input dimensionality, $x_{\rm dim}^{2}=(44\times 44)^{2}$, exceeded the available GPU memory.

Fig. 2 summarizes the experiment. Panel (a) shows an example from the test set, illustrating the challenges posed by strong noise and overlapping images. Panel (b) presents the DBF learning process. In the rightmost table, we compare the success rates of DBF against model-based approaches (EnKF, ETKF, PF). We define success as achieving a root-mean-square error (RMSE) of less than 1.0 for both position ($x_{1},y_{1},x_{2},y_{2}$) and velocity ($v_{x_{1}},v_{y_{1}},v_{x_{2}},v_{y_{2}}$) of the two digits over the final ten steps. DBF successfully performs assimilation without explicit knowledge of the images, while all the other model-based approaches fail. The KF-inspired approaches (EnKF, ETKF) failed because of very strong non-Gaussianity in the observation process and the high system dimension. Similarly, PF underperformed because the number of particles ($10{,}000$) was insufficient for the problem dimension ($z_{\rm dim}=8$ and two digits images $2\times 28\times 28=1{,}568$). Figures for visualizing the assimilation results for all the algorithms are given in the appendix (Fig. 7).

Panel (b) of Fig. 2 illustrates the evolution of the estimated figures. Initially, DBF assumes two random shapes. As training progresses, it first identifies one of the numbers (“9”) and subsequently detects the second shape (“5”). By the end of the training process, DBF nearly perfectly estimates the parameters of the observation model, including the positions of the figures, which is crucial for adjusting their reflective behavior. The ability to learn system parameters through gradient descent sets DBF apart from traditional model-based approaches.

This section presents our experiments with a double pendulum system, selected for its nonlinear and chaotic behavior. The pendulum consists of two 1 kg masses, P1 and P2, connected by two 1 meter bars, B1 and B2. One end of the bar B1 is fixed at the origin (“O”), with the other end attached to P1. Mass P2 is connected to P1 via bar B2. A schematic of the setup is shown in panel (a) of Fig. 3.

We use the angles $\theta_{1}$ and $\theta_{2}$, and the two angular velocities, $\omega_{1}$ and $\omega_{2}$, as target physical variables. The latent dimension for DBF, VRNN, SRNN, and DKF is set to 50. Observation data consists of the two-dimensional spatial positions of masses $P_{1}$ and $P_{2}$, corrupted by Gaussian noise. The observation operator combines trigonometric functions for $\theta_{1}$ and $\theta_{2}$, creating a highly nonlinear relationship. Experiments are conducted with noise levels of $\sigma=0.1,0.3$, and $0.5$ meters, with a time step of 0.03 seconds between observations. In the emission model $p(z_{t}|h_{t})$, we assume von Mises distributions for $\theta_{1}$ and $\theta_{2}$, while $\omega_{1}$ and $\omega_{2}$ follow Gaussian distributions.

Table 1 presents the RMSE between the physical variables and the mean of the filtered distribution. Training for KalmanNet was unsuccessful under all conditions. For the DVAEs, we exclude failed initial conditions when calculating the RMSE. DBF outperforms both model-based and latent assimilation methods across all settings, showing significant improvements in estimating $\omega$, which cannot be inferred from a single observation. Fig. 3 (b) illustrates an example of RMSE evolution during assimilation, where DBF consistently outperforms the other methods. The assimilation of $\omega$ occurs within the first $\sim 20$ steps, maintaining an excellent estimation accuracy throughout the experiment.

A key feature of DBF is its ability to generate samples of $z_{t}$ and assess the uncertainty in state estimates. To evaluate this capability, we analyze the distributions of normalized errors defined as $\epsilon_{norm,t,i}=(z_{t,sample,i}-z_{t,i})/\delta_{i}$, where $z_{t,i}$ represents the true value of dimension $i$ at time $t$, and $\delta_{i}$ is the standard deviation of $z_{t,sample,i}$. We collect $\epsilon_{norm,t,i}$ across all time steps, focusing on $i=\omega_{1}$ and $i=\omega_{2}$, since $\theta_{1}$ and $\theta_{2}$ follow von Mises distributions. If the uncertainty estimates are accurate, $\epsilon_{norm,t,i}$ should approximate a Gaussian distribution with a standard deviation of one. To quantify the accuracy, we compute the symmetric KL divergence (Jeffreys divergence) $KL_{sym}[p,q]=(KL[p||q]+KL[q||p])/2$ between the histogram of $\epsilon_{norm,t,i}$ and a unit Gaussian. DBF exhibits very low $KL_{sym}$ values, indicating accurate error estimation. Panels (c) and (d) display example histograms of $\epsilon_{norm,t,i}$ for DBF and ETKF.

In the final experiment, we focus on state estimation in the Lorenz96 model (Lorenz, 1995), a benchmark for testing data assimilation algorithms on noisy, nonlinear observations. The Lorenz96 model describes the evolution of a one-dimensional array of variables, each representing a physical quantity over a spatial domain, like an equilatitude circle. The dynamics are governed by the following coupled ordinary differential equations:

where $z_{i}$ is the value at grid point $i$, $N$ is the number of grid points, and $F$ is external forcing. For our experiments, we take $F=8$ and $N=40$.

We consider two observation operators. The first adds Gaussian noise to direct observations: $o_{t,j}=z_{t,j}+\epsilon$, with noise levels $\sigma=1,3,5$. The second uses a nonlinear operator: $o_{t,j}=\min(z_{t,j}^{4},10)+\epsilon$, with the same noise levels. The dynamic range of $z_{t,j}$ is around $\pm 10$, and observations are capped at 10 when $z_{t,j}$ exceeds 1.8. This makes it highly challenging for classical DA methods, as each observation offers limited information. The filter must integrate data over long timesteps, where nonlinear dynamics distort the probability distribution. Fig. 4 illustrates observations and target values. All models use 80 observation steps with a 0.03 time interval. The latent dimension for DBF, VRNN, SRNN, and DKF is set to 800. For further details for the experiment, see Sec. A.3.

Table 2 presents the assimilation performance across different noise levels and observation settings. DBF outperforms existing methods in direct observations with $\sigma=3,5$, and across all noise levels for nonlinear observation cases. In the $\sigma=1$ setting with direct observation, traditional algorithms like EnKF and ETKF outperform DBF.

The superior performance of EnKF and ETKF with direct observations at the lowest noise level can be attributed to the minimal non-Gaussianity in the posteriors within physical space. Non-Gaussianity can originate from both the dynamics model (predict step) and the observation model (update step). In this setting, the linearity of the observation operator prevents non-Gaussianity from being introduced during the update step, provided that the prior $q(z_{t}|o_{1:t-1})$ is Gaussian. Additionally, state estimation from each observation is highly accurate due to small noise. As a result, the prior $p(z_{t+1}|o_{1:t})$ remains close to a Gaussian distribution, as the locally linear approximation of the dynamics adequately captures the time evolution of probability distributions. The poorer performance of EnKF and ETKF in the $\sigma=5$ experiment is attributed to the increased non-Gaussianity introduced during each predict step. Similarly, when the observation operator is nonlinear, each update step introduces substantial non-Gaussianity. This results in a significant drop in performance for traditional filtering methods across all noise levels. In these scenarios, DBF consistently maintains an advantage over classical DA algorithms.

We observe that training DVAE-based methods is highly unstable, while that for DBF exhibits stability. Dynamics in DVAEs are modeled by RNNs, which often suffer from unstable training due to exploding or vanishing gradients. In contrast, DBF employs matrix multiplication for dynamics. If the eigenvalues of the matrix exceed one by a large margin, the model predictions, and consequently the loss function, would explode irrespective of inputs. Fig. 6 shows the histogram of the absolute values of eigenvalues at the end of training, which are distributed around or below one, indicating stable training.