Timezone: »
The Structured Variational Autoencoder (SVAE) was introduced five years ago. It presented a modeling idea---to use probabilsitic graphical models (PGMs) as priors on latent variables and deep neural networks (DNNs) to map them to observed data---as well as an inference idea---to have the recognition network output conjugate potentials to the PGM prior rather than a full posterior. While mathematically appealing, the SVAE proved impractical to use or extend, as learning required implicit differentiation of a PGM inference algorithm, and the original authors' implementation was in pure Python with no GPU or TPU support. Now, armed with the power of JAX, a software library for automatic differentiation and compilation to CPU, GPU, or TPU targets, we revisit the SVAE. We develop a modular implementation that is orders of magnitude faster than the original code and show examples in a variety of different settings, including a scientific application to animal behavior modeling. Furthermore, we extend the original model by incorporating interior potentials, which allows for more expressive PGM priors, such as the Recurrent Switching Linear Dynamical System (rSLDS). Our JAX implementation of the SVAE and its extensions open up avenues for many practical applications, extensions, and theoretical investigations.
Author Information
Yixiu Zhao (Stanford University)
Scott Linderman (Columbia University)
More from the Same Authors
-
2021 : Bayesian Inference in Augmented Bow Tie Networks »
Jimmy Smith · Dieterich Lawson · Scott Linderman -
2021 Poster: Generalized Shape Metrics on Neural Representations »
Alex H Williams · Erin Kunz · Simon Kornblith · Scott Linderman -
2021 Poster: Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems »
Jimmy Smith · Scott Linderman · David Sussillo -
2018 Poster: Point process latent variable models of larval zebrafish behavior »
Anuj Sharma · Robert Johnson · Florian Engert · Scott Linderman -
2018 Spotlight: Point process latent variable models of larval zebrafish behavior »
Anuj Sharma · Robert Johnson · Florian Engert · Scott Linderman -
2017 : Poster Session 1 »
Magdalena Fuchs · David Lung · Mathias Lechner · Kezhi Li · Andrew Gordus · Vivek Venkatachalam · Shivesh Chaudhary · Jan Hůla · David Rolnick · Scott Linderman · Gonzalo Mena · Liam Paninski · Netta Cohen