Amortized backward variational inference in nonlinear state-space models

TL;DR

This paper introduces an amortized variational inference method for nonlinear state-space models, providing linear error bounds under mixing assumptions and demonstrating improved efficiency and accuracy over existing methods.

Contribution

It proposes a novel amortized inference framework using neural networks for variational kernels in nonlinear state-space models, with theoretical error guarantees and empirical improvements.

Findings

Error in variational expectations grows at most linearly with observations

Neural network-based amortized inference improves efficiency and accuracy

Analytical marginalization leads to faster smoothing algorithms

Abstract

We consider the problem of state estimation in general state-space models using variational inference. For a generic variational family defined using the same backward decomposition as the actual joint smoothing distribution, we establish for the first time that, under mixing assumptions, the variational approximation of expectations of additive state functionals induces an error which grows at most linearly in the number of observations. This guarantee is consistent with the known upper bounds for the approximation of smoothing distributions using standard Monte Carlo methods. Moreover, we propose an amortized inference framework where a neural network shared over all times steps outputs the parameters of the variational kernels. We also study empirically parametrizations which allow analytical marginalization of the variational distributions, and therefore lead to efficient smoothing algorithms. Significant improvements are made over state-of-the art variational solutions, especially when the generative model depends on a strongly nonlinear and noninjective mixing function.