Measure transport with kernel mean embeddings
L. Wang, N. Nüsken
TL;DR
The paper proposes kernel mean embedding (KME) dynamics to transport probability measures from a prior to a tempered posterior in continuous time, relaxing the strict equality constraint of classical mean-field dynamics. By embedding distributions into a reproducing kernel Hilbert space and enforcing equivalence of kernel mean embeddings, the authors derive a mean-field ODE for particle transport, alongside a kernelised continuity equation and an integral equation for the velocity field. A special case with the quadratic kernel reduces to the Kalman-Bucy filter, linking KME-dynamics to classical Kalman theory, while general kernels offer non-Gaussian flexibility. The framework also includes a score-estimation mechanism via kernel identities and a weighted (importance sampling) variant to compensate for numerical errors, with a variational interpretation that connects to kernelised diffusion maps and diffusion-type PDEs. Numerical experiments on toy problems and Lorenz models illustrate the method's robustness and its Kalman-adjusted variant's practical advantages for data assimilation, suggesting a scalable, nonparametric alternative to ensemble Kalman filtering in nonlinear settings.
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).
