[Paper Review] State Space Expectation Propagation: Efficient Inference Schemes for Temporal Gaussian Processes

William J. Wilkinson, Paul E. Chang+2 more|arXiv (Cornell University)|Jul 12, 2020

Gaussian Processes and Bayesian Inference4 citations

TL;DR

This paper introduces State Space Expectation Propagation (SSEP), a unifying framework that formulates approximate Bayesian inference in temporal and spatio-temporal Gaussian processes as iterative Kalman smoothing updates. By recasting expectation propagation (EP) within a state space model, the method achieves linear-time inference, combines EP’s accuracy with the efficiency of linearisation, and enables scalable, high-dimensional inference across diverse likelihoods and data structures, with empirical validation on audio, classification, and spatial point processes.

ABSTRACT

We formulate approximate Bayesian inference in non-conjugate temporal and spatio-temporal Gaussian process models as a simple parameter update rule applied during Kalman smoothing. This viewpoint encompasses most inference schemes, including expectation propagation (EP), the classical (Extended, Unscented, etc.) Kalman smoothers, and variational inference. We provide a unifying perspective on these algorithms, showing how replacing the power EP moment matching step with linearisation recovers the classical smoothers. EP provides some benefits over the traditional methods via introduction of the so-called cavity distribution, and we combine these benefits with the computational efficiency of linearisation, providing extensive empirical analysis demonstrating the efficacy of various algorithms under this unifying framework. We provide a fast implementation of all methods in JAX.

Motivation & Objective

• Address the computational intractability of non-conjugate likelihoods in Gaussian processes by unifying approximate inference methods under a state space framework.
• Overcome the cubic scaling of standard GP inference in large time series by leveraging stochastic differential equation (SDE) representations and Kalman filtering.
• Enable efficient, scalable inference for high-dimensional spatio-temporal problems by treating one dimension as sequential, even when not inherently temporal.
• Provide a flexible, unified implementation framework that supports multiple inference schemes (EP, EKF, UKF, VI) with minimal code changes.
• Demonstrate the superiority of iterated linearisation-based inference (e.g., EEP) over standard EP and VI in accuracy and stability, especially under non-Gaussian likelihoods.

Proposed method

• Reformulate Gaussian process models as continuous-time stochastic differential equations (SDEs), enabling exact linear-time inference via Kalman filtering and smoothing.
• Express expectation propagation (EP) as a Kalman smoother with cavity distributions, replacing moment matching with linearisation to unify with classical smoothers.
• Introduce an iterative inference scheme—Extended Expectation Propagation (EEP)—that refines linearisations over multiple passes, improving accuracy over standard EKF.
• Apply statistical linearisation via cubature points for likelihoods where Taylor expansion is inaccurate, balancing robustness and efficiency.
• Extend the framework to multidimensional inputs by treating one dimension as sequential (e.g., spatial grid as time), enabling linear-time inference over that axis.
• Implement all methods in JAX for high-performance, differentiable, and composable inference with automatic differentiation and GPU acceleration.

Experimental results

Research questions

• RQ1 Can expectation propagation be unified with classical Kalman smoothing and filtering methods under a single state space inference framework?
• RQ2 How does replacing EP’s moment matching with linearisation affect inference accuracy and computational efficiency in non-conjugate GP models?
• RQ3 To what extent can the state space framework scale to high-dimensional spatio-temporal problems with non-Gaussian likelihoods?
• RQ4 In what scenarios does iterated linearisation (e.g., EEP) outperform standard EP and variational inference in terms of accuracy and stability?
• RQ5 Can the framework be generalized beyond one-dimensional time series to multidimensional problems by treating one dimension as sequential?

Key findings

• The proposed State Space Expectation Propagation (SSEP) framework unifies EP, EKF, UKF, and variational inference under a single Kalman smoothing formulation, enabling consistent and efficient inference.
• Extended Expectation Propagation (EEP) with 3 iterations achieves superior performance on the banana classification task, closely matching the VGP baseline and outperforming EP and VI in both accuracy and stability.
• For the audio classification task with 6-dimensional likelihood sites, EEP is faster than cubature-based methods due to fewer likelihood evaluations per data point.
• On the 125,000-point rainforest spatial point process, EEP scales effectively with a 250-dimensional site parameter, while most methods become impractical due to high-dimensional updates.
• The method achieves linear-time inference over the sequential dimension in multidimensional problems, enabling scalable inference on large grids (e.g., 500×250 spatial grid with 500 time steps).
• Cubature-based methods are more robust than linearisation when likelihoods are highly nonlinear and cavity variances are high, but they scale poorly to high dimensions, highlighting a trade-off between robustness and scalability.