State-space graphical models and the variational autoencoder framework provide a principled apparatus for learning dynamical systems from data. State-of-the-art probabilistic approaches are often able to scale to large problems at the cost of flexibility of the variational posterior or expressivity of the dynamics model. However, those consolidations can be detrimental if the ultimate goal is to learn a generative model capable of explaining the spatiotemporal structure of the data and making accurate forecasts. We introduce a low-rank structured variational autoencoding framework for nonlinear Gaussian state-space graphical models capable of capturing dense covariance structures that are important for learning dynamical systems with predictive capabilities. Our inference algorithm exploits the covariance structures that arise naturally from sample based approximate Gaussian message passing and low-rank amortized posterior updates -- effectively performing approximate variational smoothing with time complexity scaling linearly in the state dimensionality. In comparisons with other deep state-space model architectures our approach consistently demonstrates the ability to learn a more predictive generative model. Furthermore, when applied to neural physiological recordings, our approach is able to learn a dynamical system capable of forecasting population spiking and behavioral correlates from a small portion of single trials.

翻译：状态空间图模型与变分自编码器框架为从数据中学习动态系统提供了原理性工具。当前先进的概率方法通常能够扩展到大规模问题，但代价是变分后验的灵活性或动态模型表达能力的降低。然而，如果最终目标是学习一个能够解释数据时空结构并做出准确预测的生成模型，这些折衷可能是有害的。我们提出了一种用于非线性高斯状态空间图模型的低秩结构化变分自编码框架，该框架能够捕获密集协方差结构，这对于学习具有预测能力的动态系统至关重要。我们的推断算法利用了基于样本的近似高斯消息传递和低秩摊销后验更新中自然产生的协方差结构——有效地执行近似变分平滑，其时间复杂度在状态维度上呈线性增长。与其他深度状态空间模型架构的比较表明，我们的方法始终展现出学习更具预测性生成模型的能力。此外，当应用于神经生理记录时，我们的方法能够学习一个动态系统，该系统能够从单次试验的小部分数据中预测群体发放活动及行为相关信号。