Collapsed amortized variational inference for switching nonlinear dynamical systems

1 INTRODUCTION . Consider watching from above an airplane flying across country or a car driving through a field . The vehicle ’ s motion is composed of straight , linear dynamics and curving , nonlinear dynamics . This is illustrated in fig . 1 ( a ) . In this paper , we propose a new inference algorithm for fitting switching nonlinear dynamical systems ( SNLDS ) , which can be used to segment time series data as sequences of images , or lower dimensional signals , such as ( x , y ) locations into meaningful discrete temporal “ modes ” or “ regimes ” . The transitions between these modes may correspond to the changes in internal goals of the agent ( e.g. , a mouse switching from running to resting , as in Johnson et al . ( 2016 ) ) or may be caused by external factors ( e.g. , changes in the road curvature ) . Discovering such discrete modes is useful for scientific applications ( c.f. , Wiltschko et al . ( 2015 ) ; Linderman et al . ( 2019 ) ) as well as for planning in the context of hierarchical reinforcement learning ( c.f. , Kipf et al . ( 2019 ) ) . There has been extensive previous work , some of which we review in Section 2 , on modeling temporal data using various forms of state space models ( SSM ) . We are interested in the class of SSM which has both discrete and continuous latent variables , which we denote by st and zt , where t is the discrete time index . The discrete state , st ∈ { 1 , 2 , . . . , K } , represents the mode of the system at time t , and the continuous state , zi ∈ RH , represents other factors of variation , such as location and velocity . The observed data is denoted by xt ∈ RD , and can either be a low dimensional projection of zt , such as the current location , or a high dimensional signal that is informative about zt , such as an image . We may optionally have observed input or control signals ut ∈ RU , which drive the system in addition to unobserved stochastic noise . We are interested in learning a generative model of the form pθ ( s1 : T , z1 : T , x1 : T |u1 : T ) from partial observations , namely ( x1 : T , u1 : T ) . This requires inferring the posterior over the latent states , pθ ( s1 : T , z1 : T |v1 : T ) , where vt = ( xt , ut ) contains all the visible variables at time t. For training purposes , we usually assume that we have multiple such trajectories , possibly of different lengths , but we omit the sequence indices from our notations for simplicity . This problem is very challenging , because the model contains both discrete and continuous latent variables ( a so-called “ hybrid system ” ) , and has nonlinear transition and observation models . The main contribution of our paper is a new way to perform efficient approximate inference in this class of SNLDS models . The key observation is that , conditioned on knowing z1 : T as well as v1 : T , we can marginalize out s1 : T in linear time using the forward-backward algorithm . In particular , we can efficiently compute the gradient of the log marginal likelihood , ∇ ∑ s1 : T log p ( s1 : T |z̃1 : T , v1 : T ) , where z̃1 : T is a posterior sample that we need for model fitting . To efficiently compute posterior samples z̃1 : T , we learn an amortized inference network qφ ( z1 : T |v1 : T ) for the “ collapsed ” NLDS model p ( z1 : T , v1 : T ) . The collapsing trick removes the discrete variables , and allows us to use the reparameterization trick for the continuous z . These tricks let us use stochastic gradient descent ( SGD ) to learn p and q jointly , as explained in Section 3 . We can then use q as a proposal distribution inside a Rao-Blackwellised particle filter ( Doucet et al. , 2000 ) , although in this paper , we just use a single posterior sample , as is common with Variational AutoEncoders ( VAEs , Kingma & Welling ( 2014 ) ; Rezende et al . ( 2014 ) ) . Although the above “ trick ” allows us efficiently perform inference and learning , we find that in challenging problems ( e.g. , when the dynamical model p ( zt|zt−1 , vt ) is very flexible ) , the model ignores the discrete latent variables , and does not perform mode switching . This is a form of “ posterior collapse ” , similar to VAEs , where powerful decoders can cause the latent variables to be ignored , as explained in Alemi et al . ( 2018 ) . Our second contribution is a new form of posterior regularization , which prevents the aforementioned problem and results in a significantly improved segmentation . We apply our method , as well as various existing methods , to two previously proposed lowdimensional time series segmentation problems , namely a 1d bouncing ball , and a 2d moving arm . In the 1d case , the dynamics are piecewise linear , and all methods perform perfectly . In the 2d case , the dynamics are piecewise nonlinear , and we show that our method infers much better segmentation than previous approaches for comparable computational cost . We also apply our method to a simple new video dataset ( see fig . 1 for an example ) , and find that it performs well , provided we use our proposed regularization method . In summary , our main contributions are • Learning switching nonlinear dynamical systems parameterized with neural networks by marginalizing out discrete variables . • Using entropy regularization and annealing to encourage discrete state transitions . • Demonstrating that the discrete states of nonlinear models are more interpretable . 2 RELATED WORK . In this section , we briefly summarize some related work . 2.1 STATE SPACE MODELS . We consider the following state space model : pθ ( x , z , s ) = p ( x1|z1 ) p ( z1|s1 ) [ T∏ t=2 p ( xt|zt ) p ( zt|zt−1 , st ) p ( st|st−1 , xt−1 ) ] , ( 1 ) where st ∈ { 1 , . . . , K } is the discrete hidden state , zt ∈ RL is the continuous hidden state , and xt ∈ RD is the observed output , as in fig . 2 ( a ) . For notational simplicity , we ignore any observed inputs or control signals ut , but these can be trivially added to our model . Note that the discrete state influences the latent dynamics zt , but we could trivially make it influence the observations xt as well . More interesting are which edges we choose to add as parents of the discrete state st. We consider the case where st depends on the previous discrete state , st−1 , as in a hidden Markov model ( HMM ) , but also depends on the previous observation , xt−1 . We can trivially depend on multiple previous observations ; we assume first-order Markov for simplicity . This means that state changes do not have to happen “ open loop ” , but instead may be triggered by signals from the environment . We can also condition zt on xt−1 , and st on zt−1 . It is straightforward to handle such additional dependencies ( shown by dashed lines in fig . 2 ( a ) ) in our inference method , which is not true for some of the other methods we discuss below . We still need to specify the functional forms of the conditional probability distributions . In this paper , we make the following fairly weak assumptions : p ( xt|zt ) = N ( xt|fx ( zt ) , R ) , ( 2 ) p ( zt|zt−1 , st = k ) = N ( zt|fz ( zt−1 , k ) , Q ) , ( 3 ) p ( st|st−1 = j , xt−1 ) = Cat ( st|S ( fs ( xt−1 , j ) ) , ( 4 ) where fx , z , s are nonlinear functions ( MLPs or RNNs ) , N ( · , · ) is a multivariate Gaussian distribution , Cat ( · ) is a categorical distribution , and S ( · ) is a softmax function . R ∈ RD×D and Q ∈ RH×H are learned covariance matrices for the Gaussian emission and transition noise . If fx and fz are both linear , and p ( st|st−1 ) is first-order Markov without dependence on zt−1 , the model is called a switching linear dynamical system ( SLDS ) . If we allow st to depend on zt−1 , the model is called a recurrent SLDS ( Linderman et al. , 2017 ; Linderman & Johnson , 2017 ) . We will compare to rSLDS in our experiments . If fz is linear , but fx is nonlinear , the model is sometimes called a “ structured variational autoencoder ” ( SVAE ) ( Johnson et al. , 2016 ) , although that term is ambiguous , since there are many forms of structure . We will compare to SVAEs in our experiments . If fz is a linear function , the model may need to use lots of discrete states in order to approximate the nonlinear dynamics , as illustrated in fig . 1 ( d ) . We therefore allow fz ( and fx ) to be nonlinear . The resulting model is called a switching nonlinear dynamical system ( SNLDS ) , or Nonlinear RegimeSwitching State-Space Model ( RSSSM ) ( Chow & Zhang , 2013 ) . Prior work typically assumes fz is a simple nonlinear model , such as polynomial regression . If we let fz be a very flexible neural network , there is a risk that the model will not need to use the discrete states at all . We discuss a solution to this in Section 3.3 . The discrete dynamics can be modeled as a semi-Markov process , where states have explicit durations ( see e.g. , Duong et al . ( 2005 ) ; Chiappa ( 2014 ) ) . One recurrent , variational version is the recurrent hidden semi-Markov model ( rHSMM , Dai et al . ( 2017 ) ) . Rather than having a stochastic continuous variable at every timestep , rHSMM instead stochastically switches between states with deterministic dynamics . The semi-Markovian structures in this work have an explicit maximum duration , which makes them less flexible . A revised method , ( Kipf et al. , 2019 ) , is able to better handle unknown durations , but produces a potentially infinite number of distinct states , each with deterministic dynamics . The deterministic dynamics of these works may limit their ability to handle noise . 2.2 VARIATIONAL INFERENCE AND LEARNING . A common approach to learning latent variable models is to maximize the evidence lower bound ( ELBO ) on the log marginal likelihood ( see e.g. , Blei et al . ( 2016 ) ) . This is given by log p ( x ) ≤ L ( x ; θ , φ ) = Eqφ ( z , s|x ) [ log pθ ( x , z , s ) − log qφ ( z , s|x ) ] , where qφ ( z , s|x ) is an approximate posterior.1 Rather than computing q using optimization for each x , we can train an inference network , fφ ( x ) , which emits the parameters of q . This is known as `` amortized inference '' ( see e.g. , Kingma & Welling ( 2014 ) ) . If the posterior distribution qφ ( z , s|x ) is reparameterizable , then we can make the noise independent of φ , and hence apply the standard SGD to optimize θ , φ . Unfortunately , the discrete distribution p ( s|x ) is not reparameterizable . In such cases , we can either resort to higher variance methods for estimating the gradient , such as REINFORCE , or we can use continuous relaxations of the discrete variables , such as Gumbel Softmax ( Jang et al. , 2017 ) , Concrete ( Maddison et al. , 2017b ) , or combining both , such as REBAR ( Tucker et al. , 2017 ) . We will compare against a Gumbel-Softmax version of SNLDS in our experiments . The continuous relaxation approach was applied to SLDS models in ( Becker-Ehmck et al. , 2019 ) and HSSM models in ( Liu et al. , 2018a ; Kipf et al. , 2019 ) . However , the relaxation can lose many of the benefits of having discrete variables ( Le et al. , 2019 ) . Relaxing the distribution to a soft mixture of dynamics results in the Kalman VAE ( KVAE ) model of Fraccaro et al . ( 2017 ) . A concern is that soft models may use a mixture of dynamics for distinct ground truth states rather than assigning a distinct mode of dynamics at each step as a discrete model must do . We will compare to KVAE in our experiments . In Section 3 , we propose a new method to avoid these issues , in which we collapse out s so that the entire model is differentiable . The SVAE model of Johnson et al . ( 2016 ) also uses the forward-backward algorithm to compute q ( s|v ) ; however , they assume the dynamics of z are linear Gaussian , so they can apply the Kalman smoother to compute q ( z|v ) . Assuming linear dynamics can result in over-segmentation , as we have discussed . A forward-backward algorithm is applied once to the discrete states and once to the continuous states to compute a structured mean field posterior q ( z ) q ( s ) . In contrast , we perform approximate inference for z using one forward-backward pass and then exact inference for s using a second pass , as we explain in Section 3 .

Summary: In this paper, the authors consider the problem of learning model parameters of a switching nonlinear dynamical system from a dataset. They propose a new variational inference algorithm for this model-learning problem that marginalizes all discrete random variables in the model using the forward-backward algorithm and, in so doing, converts the model to one with a differentiable density, so that the gradient of the variational objective can be estimated with the low-variance reparameterization estimator. The authors also point out an issue in choosing a variational objective; the standard ELBO objective is not suitable for their learning problem, because it leads to a model that does not use discrete random variables meaningfully. To overcome this issue, they suggest a new improved objective and a learning procedure, which encourage the learned model to use discrete variables for capturing different modes of dynamics. The proposed variational inference algorithm was applied to three datasets, and in all these cases, it showed promising results.