Large-scale variational Gaussian state-space models

We introduce an amortized variational inference algorithm and structured variational approximation for state-space models with nonlinear dynamics driven by Gaussian noise. Importantly, the proposed framework allows for efficient evaluation of the ELBO and low-variance stochastic gradient estimates without resorting to diagonal Gaussian approximations by exploiting (i) the low-rank structure of Monte-Carlo approximations to marginalize the latent state through the dynamics (ii) an inference network that approximates the update step with low-rank precision matrix updates (iii) encoding current and future observations into pseudo observations – transforming the approximate smoothing problem into an (easier) approximate filtering problem. Overall, the necessary statistics and ELBO can be computed in O(TL(Sr + S^2 + r^2)) time where T is the series length, L is the state-space dimensionality, S are the number of samples used to approximate the predict step statistics, and r is the rank of the approximate precision matrix update in the update step (which can be made of much lower dimension than L).