Learning Optimal Filters Using Variational Inference

TL;DR

This paper introduces a variational inference framework to learn parameterized analysis maps for data assimilation, aiming to improve the filtering distribution in high-dimensional nonlinear systems. By formulating an objective that minimizes the KL divergence between approximate and forecast posteriors and incorporating a likelihood term, the authors derive offline, online, and sample-based formulations for learning $A_\theta$ (or an equivalent transport map $T_\theta$) that update forecasts with observations. The framework is validated across linear and nonlinear models, including a stable linear system, Lorenz '96, and Kuramoto–Sivashinsky, demonstrating learning of a fixed Kalman gain, inflation, and localization parameters in EnKF, and illustrating the potential to design new learned filtering algorithms. Overall, the method provides a probabilistically principled route to tailor data assimilation operators to data, reducing bias and enabling flexible, learned filters with practical impact for weather, climate, and related applications.

Abstract

Filtering - the task of estimating the conditional distribution for states of a dynamical system given partial and noisy observations - is important in many areas of science and engineering, including weather and climate prediction. However, the filtering distribution is generally intractable to obtain for high-dimensional, nonlinear systems. Filters used in practice, such as the ensemble Kalman filter (EnKF), provide biased probabilistic estimates for nonlinear systems and have numerous tuning parameters. Here, we present a framework for learning a parameterized analysis map - the transformation that takes samples from a forecast distribution, and combines with an observation, to update the approximate filtering distribution - using variational inference. In principle this can lead to a better approximation of the filtering distribution, and hence smaller bias. We show that this methodology can be used to learn the gain matrix, in an affine analysis map, for filtering linear and nonlinear dynamical systems; we also study the learning of inflation and localization parameters for an EnKF. The framework developed here can also be used to learn new filtering algorithms with more general forms for the analysis map.

Figures (6)

• Figure 1: Learning the gain matrix for a linear dynamical system with full state observations.
• Figure 2: A comparison between the steady-state Kalman gain (left) and the learned gain (right) with partial observations.
• Figure 3: Errors as a function of time step for learning the gain with the online method. Within each time step, 100 gradient descent steps are performed. The Frobenius norm of the difference between the learned gain and the steady-state gain (red solid line). The KL divergence between the learned filtering distribution and the Kalman filter distribution (green dashed line). The mean-square prediction error in the filter mean compared to the true trajectory (blue solid line).
• Figure 4: Gains for the Lorenz '96 system. On the left is the extended Kalman filter (ExKF) gain, averaged over 50 iterations. On the right is the learned gain.
• Figure 5: Contours for the cost function as a function of localization radius and inflation with the Lorenz '96 system.

Theorems & Definitions (4)

• Remark 2.2
• Remark 2.3
• Remark 2.4
• Remark 2.5