UBC Math Department Colloquium: Eviatar Bach

Filtering is a core task in many areas of science and engineering, including weather and climate prediction, where it is known as data assimilation. It involves estimating the conditional probability distribution of a dynamical system's states given partial and noisy observations – a problem that is generally intractable for high-dimensional, nonlinear systems. The ensemble Kalman filter (EnKF) approximates the filtering distribution with an ensemble of interacting particles, employing a Gaussian ansatz for the joint distribution of the state and observation at each observation time. The EnKF is robust, but the Gaussian ansatz limits accuracy.

We address this shortcoming by using machine learning to map the forecast distribution and observation to the filtering distribution. We propose two types of cost functions for learning this map: one based on variational inference and one on strictly proper scoring rules. Both cost functions are minimized uniquely at the true filtering distribution. By time-averaging these cost functions over long trajectories in ergodic dynamical systems, the map can be learned and subsequently used for future filtering; this is a form of amortized Bayesian inference.

We demonstrate the approach using a set transformer neural architecture, which is invariant to ensemble permutations and allows for different ensemble sizes. The learned filtering algorithms outperform state-of-the-art methods in filtering chaotic systems, using both probabilistic and deterministic performance metrics. They also perform well in challenging highly non-Gaussian and multimodal problems where the EnKF fails. Once learned at a given ensemble size, the learned map can be applied to other ensemble sizes with minimal fine-tuning. As further validation, this approach accurately recovers known analytic solutions in the case of linear Gaussian problems.