NAS-X: Neural Adaptive Smoothing via Twisting

Dieterich Lawson, Michael Y. Li, Scott Linderman

Keywords: sequence models, probabilistic inference, reweighted wake-sleep, sequential monte carlo, smoothing, mechanistic models

TL;DR: We introduce a method for fitting sequential latent variable models that combines the benefits of reweighted wake-sleep and smoothing sequential Monte Carlo.

Abstract: Sequential latent variable models (SLVMs) are essential tools in statistics and machine learning, with applications ranging from healthcare to neuroscience. As their flexibility increases, analytic inference and model learning can become challenging, necessitating approximate methods. Here we introduce neural adaptive smoothing via twisting (NAS-X), a method that extends reweighted wake-sleep (RWS) to the sequential setting by using smoothing sequential Monte Carlo (SMC) to estimate intractable posterior expectations. Combining RWS and smoothing SMC allows NAS-X to provide low-bias and low-variance gradient estimates, and fit both discrete and continuous latent variable models. We illustrate the theoretical advantages of NAS-X over previous methods and explore these advantages empirically in a variety of tasks, including a challenging application to mechanistic models of neuronal dynamics. These experiments show that NAS-X substantially outperforms previous VI- and RWS-based methods in inference and model learning, achieving lower parameter error and tighter likelihood bounds.

Submission Number: 15111