Large-scale variational Gaussian state-space models
Matthew Dowling Yuan Zhao Il Memming Park
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$).