eXponential FAmily Dynamical Systems (XFADS): Large-scale nonlinear Gaussian state-space modeling

TL;DR

This work introduces 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.

Abstract

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.

Figures (7)

• Figure 1: Smoothing and predictive performance on bouncing ball and pendulum. To the left of the red line are samples from the posterior during the data window projected to image space, to the right of the red line are samples unrolled from $p_{\boldsymbol{\theta}}(\vz_t\!\mid\!\vz_{t-1})$. a) while all methods are adept at smoothing in the context window, our methods predictive performance is better by a noticeable margin as measured by the $R^2$. b) similar results hold for the bouncing ball dataset.
• Figure 2: a) Empirical time complexity scaling. Since complexity is a function of $L$, $S$, and $r$, we vary $L$ (top) for fixed $r=10$ and (bottom) for fixed $S=5$; we examine several values of the variable not fixed. Examining wall-clock time shows empirically our implementation scales linearly in $L$; on the (bottom) we plot wall-clock time for a Kalman filter implementation, showing the standard cubic dependence on $L$. b) (top) Negative ELBO as a function of training epoch when $N=L$ (bottom) when $N=L/5$; the left column shows the case $L=50$ and the right when $L=100$. Different colors indicate different settings of the local/backward encoder rank; zooming in for $L=100$, shows low-rank updates can match diagonal ones. c) Peristimulus time histogram (PSTH) for the DMFC RSG dataset for different trial condition averages; we consider a context window of $1.3$s and a prediction window of $1.3$s. d) BPS for each method for context/prediction windows.
• Figure 3: Predict behavior from a causally inferred initial condition.a) Actual reaches. b) (top) Reaches linearly decoded from smoothed ($R^2=0.89$), causally filtered ($R^2=0.88$), & predicted ($R^2=0.74$) latent trajectories starting from an initial condition causally inferred during the preparatory period. (bottom) Top 3 principal latent dimensions per regime (smoothing/filtering/prediction) for three example trials. c) bps / $R^2$ of predicted hand velocity using rates inferred from the $700$ms context window and the $500$ms prediction window. d) Velocity decoding $R^2$ using predicted trajectories as a function of how far into the trial the latent state was filtered until it was only sampled from the autonomous dynamics; by the the movement onset, behavioral predictions using latent trajectory predictions are nearly on par with behavior decoded from the smoothed posterior.
• Figure 4: Learned covariance of nonlinear SSMs. (top) single trial of a pendulum. Below are the posterior covariances output by the causal and non-causal variants of XFADS.
• Figure 5: Learned covariance of nonlinear SSMs. (top) single trial of a bouncing ball. Below are the posterior covariances output by the nonlinear SSMs considered -- qualitatively, we observe more complex covariance structures arise when the ball hits the wall that diagonal approximations cannot capture.