State Space Expectation Propagation: Efficient Inference Schemes for Temporal Gaussian Processes
This work addresses the challenge of scalable inference for researchers and practitioners in machine learning and statistics, offering incremental improvements by unifying and optimizing existing methods for temporal Gaussian processes.
The paper tackles the problem of efficient Bayesian inference in non-conjugate temporal and spatio-temporal Gaussian process models by formulating it as a parameter update rule during Kalman smoothing, unifying methods like expectation propagation and classical Kalman smoothers. It combines the benefits of EP with the computational efficiency of linearisation, demonstrating efficacy through extensive empirical analysis and providing a fast JAX implementation.
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.