Gaussian Ensemble Belief Propagation for Efficient Inference in High-Dimensional Systems

TL;DR

Gaussian Ensemble Belief Propagation (GEnBP) unifies the Ensemble Kalman Filter with Gaussian Belief Propagation to enable efficient inference in high-dimensional graphical models. It employs Diagonal Matrix with Low-rank perturbation (DLR) to maintain low-rank Gaussian parameterisations and uses ensemble-based moments to propagate beliefs, allowing belief updates without dense covariance matrices. The approach scales to dimensions where traditional GaBP becomes intractable and demonstrates improved accuracy and speed in 1D transport and 2D Navier–Stokes CFD problems, outperforming GaBP and, in some regimes, the Laplace approximation. GEnBP offers a practical framework for high-dimensional data assimilation, spatiotemporal modeling, and system identification, with potential extensions to localisation, inflation, and domain adaptation.

Abstract

Efficient inference in high-dimensional models is a central challenge in machine learning. We introduce the Gaussian Ensemble Belief Propagation (GEnBP) algorithm, which combines the strengths of the Ensemble Kalman Filter (EnKF) and Gaussian Belief Propagation (GaBP) to address this challenge. GEnBP updates ensembles of prior samples into posterior samples by passing low-rank local messages over the edges of a graphical model, enabling efficient handling of high-dimensional states, parameters, and complex, noisy, black-box generation processes. By utilizing local message passing within a graphical model structure, GEnBP effectively manages complex dependency structures and remains computationally efficient even when the ensemble size is much smaller than the inference dimension -- a common scenario in spatiotemporal modeling, image processing, and physical model inversion. We demonstrate that GEnBP can be applied to various problem structures, including data assimilation, system identification, and hierarchical models, and show through experiments that it outperforms existing belief propagation methods in terms of accuracy and computational efficiency. Supporting code is available at https://github.com/danmackinlay/GEnBP

Figures (12)

• Figure 1: Prior and posterior samples for latent $\boldsymbol{\mathsf{q}}$ in the 1D system identification problem (Section \ref{['sec:ex-onedadvect']}). The GEnBP ensemble has size $N{=}64$, and $N{=}64$ samples are drawn from the GaBP posterior. The GEnBP prior comprises samples; the GaBP prior is sampled from a Gaussian density with the same moments.
• Figure 2: Influence of dimension $D_{\mathscr{Q}}$ in the 2D fluid dynamics model. Error bars indicate empirical 50% intervals from $n=10$ runs. In grey-shaded regions, GaBP ran out of memory.
• Figure 3: Influence of viscosity $\nu$ on a $32\times 32$ 2D fluid dynamics model. Error bars indicate empirical 50% intervals from $n=40$ simulations.
• Figure 4: Generative model for the system identification problem. The latent parameter $\boldsymbol{\mathsf{q}}$ influences all states, and observed states are shaded.
• Figure 5: Influence of dimension $D_{\mathscr{Q}}$ in the transport example. Error bars are empirical 90% intervals from $n=40$ runs. Missing log likelihood values indicate undefined values from non-positive definite covariance estimates.

Theorems & Definitions (12)

• Proposition 1: Belief Propagation on Factor Graphs
• Definition 1
• Definition 2: Gaussian Density Forms
• Proposition 2
• Proposition 3
• Proposition 4
• Proposition 4: Belief Propagation on Factor Graphs
• Proposition 5
• Proposition 5
• Proposition 6