Measure transport with kernel mean embeddings

TL;DR

This paper introduces KME-dynamics, a novel particle system for Bayesian inference that extends Kalman filters using kernel mean embeddings, improving accuracy in nonlinear, non-Gaussian scenarios.

Contribution

We develop KME-dynamics, a new framework that incorporates kernel mean embeddings into continuous-time particle systems, connecting to optimal transport and score-based modeling.

Findings

KME-dynamics outperforms ensemble Kalman filter in experiments.

Hybrid Kalman-adjusted KME-dynamics shows promising results.

Connections established to diffusion maps and variational formulations.

Abstract

Kalman filters constitute a scalable and robust methodology for approximate Bayesian inference, matching first and second order moments of the target posterior. To improve the accuracy in nonlinear and non-Gaussian settings, we extend this principle to include more or different characteristics, based on kernel mean embeddings (KMEs) of probability measures into reproducing kernel Hilbert spaces. Focusing on the continuous-time setting, we develop a family of interacting particle systems (termed $\textit{KME-dynamics}$) that bridge between prior and posterior, and that include the Kalman-Bucy filter as a special case. KME-dynamics does not require the score of the target, but rather estimates the score implicitly and intrinsically, and we develop links to score-based generative modeling and importance reweighting. A variant of KME-dynamics has recently been derived from an optimal transport and Fisher-Rao gradient flow perspective by Maurais and Marzouk, and we expose further connections to (kernelised) diffusion maps, leading to a variational formulation of regression type. Finally, we conduct numerical experiments on toy examples and the Lorenz 63 and 96 models, comparing our results against the ensemble Kalman filter and the mapping particle filter (Pulido and van Leeuwen, 2019, J. Comput. Phys.). Our experiments show particular promise for a hybrid modification (called Kalman-adjusted KME-dynamics).