Year: 2019

Abstract:
We propose an efficient inference method for switching nonlinear dynamical systems. The key idea is to learn an inference network which can be used as a proposal distribution for the continuous latent variables, while performing exact marginalization of the discrete latent variables. This allows us to use the reparameterization trick, and apply end-to-end training with stochastic gradient descent. We show that the proposed method can successfully segment time series data (including videos) into meaningful “regimes”, by using the piece-wise nonlinear dynamics.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: In this section, we briefly summarize some related work.

heading: 2.1 STATE SPACE MODELS
text: 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.

heading: 2.2 VARIATIONAL INFERENCE AND LEARNING
text: 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.

heading: 2.3 MONTE CARLO INFERENCE
text: There is a large literature on using sequential Monte Carlo methods for inference in state space models (see e.g., Doucet & Johansen (2011)). When the model is nonlinear (as in our case), we may need a lot of particles to get a good approximation, which can be expensive. We can often get better (lower variance) approximations by analytically marginalizing out some of the latent variables; the resulting method is called a “Rao Blackwellised particle filter” (RBPF).
Prior work (e.g., Doucet et al. (2001)) has applied RBPF to SLDS models, leveraging the fact that it is possible to marginalize out p(z|s,v) using the Kalman filter. It is also possible to compute the optimal proposal distribution for sampling from p(st|st−1,v) in this case. However, this relies on the model being conditionally linear Gaussian. By contrast, we marginalize out p(s|z,v), so we can handle nonlinear models. In this case, it is hard to compute the optimal proposal distribution for sampling from p(zt|zt−1,v), so instead we use variational inference to learn to approximate this.

heading: 3 METHOD

heading: 3.1 INFERENCE
text: We use the following variational posterior: qφ,θ(z, s|x) = qφ(z|x)pθ(s|z,x), where pθ(s|z,x) is the exact posterior (under the generative model) computed using the forward-backward algorithm,
1 In the case of sequential models, we can create tighter lower bounds using methods such as FIVO (Maddison et al., 2017a), although this is orthogonal to our work.
and qφ(z|x) is defined below. To compute qφ(z|x), we first process x1:T through a bidirectional RNN, whose state at time t is denoted by hxt . (As noted in Krishnan et al. (2017), due to the Markov assumptions of our model, we only need a backward RNN to summarize xt:T , rather a bidirectional RNN to summarize x1:T , but we use the latter for simplicity.) We then use a forward (causal) RNN, whose state denoted by hzt , to compute the parameters of q(zt|z1:t−1,x1:T ), where the hidden state is computed based on hzt−1 and h x t . This gives the following approximate posterior:
qφ(z1:T |x1:T ) = ∏ t q(zt|z1:t−1,x1:T ) = ∏ t q(zt|hzt ). See fig. 2(b) for an illustration.
We can draw a sample z1:T ∼ qφ(z|x) sequentially, and then treat this as “soft evidence” for the HMM model. The (sample dependent) parameters used in the forward-backward algorithm are given by At(j, k) = p(st = j|st−1 = k,xt−1), Bt(k) = p(xt|zt)p(zt|zt−1, st = k) for t > 1, B1(k) = p(x1|z1)p(z1|s1 = k), and π(k) = p(s1 = k). We can then compute the following posterior marginals: γ2t (j, k) = p(st = k, st−1 = j|x1:T , z1:T ) and γ1t (k) = p(st = k|x1:T , z1:T ), which can be used to compute the gradients of the log likelihood, as we discuss below.

heading: 3.2 LEARNING
text: The evidence lower bound (ELBO) for a single sequence x is given by
LELBO(θ,φ) = Eqφ(z|x)pθ(s|x,z) [log pθ(x, z)pθ(s|x, z)− log qφ(z|x)pθ(s|x, z)] (5) = Eqφ(z|x) [log pθ(x, z)− log qφ(z|x)] (6)
Because qφ(z) is reparameterizable, we can approximate the gradient as follows:
∇θ,φL(θ, φ) ≈ ∇θ,φ log pθ(x, z̃)−∇φ log qφ(z̃|x) (7)
where z̃ is a sample from the variational proposal z̃ ∼ qφ(z̃1|x1:T ) ∏T t=2 qφ(z̃t|z̃t−1,x1:T ). The second term can be computed by applying backpropagation through time to the inference RNN. In the appendix, we show that the first term is given by
T∑ t=2 ∑ j,k γ2t (j, k)∇ [logBt(k)At(j, k)] + ∑ k γ11(k)∇ [logB1(k)π(k)] (8)

heading: 3.3 ENTROPY REGULARIZATION AND TEMPERATURE ANNEALING
text: When using expressive nonlinear functions (e.g. an RNN or MLP) to model p(zt|zt−1, st), we found that the model only used a single discrete state, analogous to posterior collpase in VAEs (see e.g., Alemi et al. (2018)). To encourage the model to utilize multiple states, we add an additional regularizing term to the ELBO that penalizes the KL divergence between the state posterior at each time step and a uniform prior pprior(st = k) = 1/K (Burke et al., 2019). We call this a cross-entropy regularizer:
LCE = T∑ t=1 DKL(pprior(st)||p(st|z1:T ,x1:T )). (9)
Our overall objective now becomes
L(θ,φ) = LELBO(θ,φ)− βLCE(θ,φ). (10)
Note that 0 ≤ LCE ≤ ∞, so we need to choose the scale of β > 0 appropriately. To further smooth the optimization problem, we apply temperature annealing to the discrete state transitions, as follows: p(st = k|st−1 = j,xt−1) = S(p(st=k|st−1=j,xt−1)τ ), where τ is the temperature. At the beginning stage of training, β, τ are set to large values. Doing so ensures that all states are visited, and can explain the data equally well. Over time, we reduce the regularizers to 0 and temperature to 1, according to a fixed annealing schedule; this allows clusters to start to separate (c.f., Rose (1998)), as each regime learns its own local dynamical model. The overall approach is similar to multi-step pretraining, as used in prior papers such as the rSLDS paper (Linderman et al., 2017), but our approach works in a continuous end-to-end fashion.

heading: 4 EXPERIMENTS
text: In this section, we compare our method to various other methods that have been recently proposed for time series segmentation using latent variable models. Since it is hard to evaluate unsupervised learning methods, such as segmentation, we use three synthetic datasets, where we know the ground truth.
In each case, we fit the model to the data, and then estimate the most likely hidden discrete state at each time step, ŝt = argmax q(st|x1:T ). Since the model is unidentifiable, the state labels have no meaning, so we post-process them by applying the best permutation over labels so as to maximize the F1 score across frames. Here the F1 score is the harmonic mean of precision and recall, 2× precision× recall/(precision+ recall), where precision is the percentage of the predictions that match the ground truth states, and recall is the percentage of the ground truth states that match the predictions. We also compute the switching-point F1 by only considering the frames where the ground truth state changes. This measure compliments the frame-wise F1, because it measures the accuracy in time. Since matching the exact time of the switch point is very hard in the unsupervised setting with noisy observations, we also consider a detected change point as correct, if it occurs within some small temporal interval around the ground truth as noted in Section 4.3.

heading: 4.1 1D BOUNCING BALL
text: In this section, we use a simple dataset from Johnson et al. (2016). The data encodes the location of a ball bouncing between two walls in a one dimensional space. The initial position and velocity are random, but the wall locations are constant.
We apply our SNLDS model to this data, where fx and fz are both MLPs. We found that regularization was not necessary in this experiment. We also consider the case where fx and fz are linear (which corresponds to an SLDS model), the rSLDS model of Linderman et al. (2017), the SVAE model of Johnson et al. (2016), the Kalman VAE (KVAE) model of Fraccaro et al. (2017) and a GumbelSoftmax version of SNLDS as described in Appendix A.2. We use the implementations of rSLDS, SVAE, and KVAE provided by the authors.
As the data is generated by a simple piece-wise linear dynamics and all models we tested learn a perfect segmentation, as shown in Figure 3(a) and Table 1. This serves as a “sanity check” that we are able to use and implement the rSLDS, SVAE, KVAE and Gumbel-Softmax SNLDS code correctly. (See also Appendix A.3 for further analysis.)
Note that the “true” number of discrete states is just 2, encoding whether the ball is moving up or down. We find that our method can learn to ignore irrelevant discrete states if they are not needed. This is presumably because we are maximizing the marginal likelihood since we sum over all hidden states, and this is known to encourage model simplicity due to the "Bayesian Occam’s razor" effect (Murray & Ghahramani, 2005). By contrast, with the other methods, we had to be more careful in setting K.

heading: 4.2 2D REACHER TASK
text: In this section, we consider a dataset proposed in the CompILE paper (Kipf et al., 2019). The observations are sequences of 36 dimensional vectors, derived from the 2d locations of various static objects, and the 2d joint locations of a moving arm (see Appendix A.4 for details and a visualization). The ground truth discrete state for this task is the identity of the target that the arm is currently reaching for (i.e., its "goal").
We fit the same 6 models as above to this dataset. Since it is a much harder problem, we found that we needed to add regularization to our model to encourage it to switch states. Figure 3(b) visualizes the resulting segmentation (after label permutation) for a single example. We see that our SNLDS model matches the ground truth more closely than our SLDS model, as well as the rSLDS, SVAE, KVAE, and Gumbel-Softmax baselines.
To compare performance quantitatively, we evaluate the models from 5 different training runs on the same held-out dataset of size 32, and compute the F1 scores. We also report the F1 number from CompILE. The CompILE paper uses an iterative segmentation scheme that can detect state changes,
but it does infer what the current latent state is, so we cannot include it in Figure 3(b). As in Table 1, we find that our SNLDS method is significantly better than the other approaches.

heading: 4.3 IMAGE DATASET (DUBINS PATH)
text: In this section, we apply our method to a new dataset that is created by "rendering" a point moving in the 2d plane. The motion follows the Dubins model2, which is a simple model for piece-wise nonlinear (but smooth) motion that is commonly used in the fields of robotics and control theory because it corresponds to the shortest path between two points that can be traversed by wheeled robots, airplanes, etc. In the Dubins model, the change in direction is determined by an external control signal ut. We replace this with three latent discrete control states: go straight, turn left, and turn right. These correspond to fixed, but unobserved, input signals ut (see Appendix A.5 for details). After generating the motion, we create a series of images, where we render the location of the moving object as a small circle on a white background. Our goal in generating this dataset was to assess how well we can recover latent dynamics from image data in a very simple, yet somewhat realistic, setting.
The publicly released code for rSLDS and SVAE does not support high dimensional inputs like images (even though the SVAE has been applied to an image dataset in Johnson et al. (2016)), and there is no public code for CompILE. Therefore we could not compare to these methods for this experiment. Also, since we already showed in Section 4.2 that our method is much better than these
2 https://en.wikipedia.org/wiki/Dubins_path
approaches, as well as Kalman VAE and Gumbel-Softmax version of SNLDS, on the simpler reacher task, we expect the same conclusion to hold on the harder task of segmenting videos.
Instead we focus on comparing SNLDS with SLDS to see the advantage of allowing each regime to be represented by a nonlinear model. The results of segmenting one sequence with these models using 5 states are shown in Figure 1. We see that the SLDS model has to approximate the left and right turns with multiple discrete states, whereas the non-linear model learns a more interpretable representation.
We again compare the models using F1 scores in Table 2. Because the SLDS model used too many states, we calculated two versions of the metric. The first was a greedy metric that optimally assigned the best single state to match the ground truth. The second used an oracle to optimally merge states to match the ground truth. The SNLDS model significantly outperforms the SLDS in both scenarios.

heading: 4.4 ANALYSIS OF THE ANNEALING SCHEDULE
text: Many latent variable models are trained in multiple stages, in order to avoid getting stuck in bad local optima. For example, to fit the rSLDS model, Linderman et al. (2017) firstly pretrain an AR-HMM and SLDS model, and then merge them; similarly, to fit the SVAE model, Johnson et al. (2016) first train with a single latent state and then increase K.
We found a similar strategy was necessary for the Reacher and Dubins tasks, but we do this in a smooth way using annealed regularizations. Early in training, we train with large temperature τ and entropy coefficient β. This encourages the model to use all states equally, so that the dynamics, inference, and emission sub-networks stabilized before beginning to learn specialized behavior. We then anneal the entropy coefficient to 0, and the temperature to 1 over time. We found it best to first decay the entropy coefficient β and then decay the temperature τ .
Figure 4 demonstrates the effect of 3 different annealing schedules on the relative log likelihood (defined as Lt − Lmin, where Lmin = mint Lt;1,2,3 across all three runs, and Lt is the negative log-likelihood.), and the F1 score. The green curve starts annealing right away; we see the F1 score
is flat, since only one discrete state is used. The red curve starts annealing later, which improves F1. However, the best results are shown in the blue curve, which starts the decay even later.
On real problems, where we have no ground truth, we cannot use the F1 score as a metric to determine the best annealing schedule. However, it seems that the schedules that improve F1 the most also improve likelihood the most.

heading: 5 CONCLUSION
text: We have demonstrated that our proposed method can effectively learn to segment high dimensional sequences into meaningful discrete regimes. Future work includes applying this to harder image sequences, and to hierarchical reinforcement learning.

heading: A APPENDIX
text: A.1 DERIVATION OF THE GRADIENT OF THE ELBO
The evidence lower bound objective (ELBO) of the model is defined as:
L(θ, φ) = Eqθ,φ(z,s|x) [log pθ(x, z, s)− log qθ,φ(z, s|x)] (11) = Eqφ(z|x)pθ(s|x,z) [log pθ(x, z)pθ(s|x, z)− log qφ(z|x)pθ(s|x, z)] (12) = Eqφ(z|x) [log pθ(x, z)] +H(qφ(z|x)) (13)
where the first term is the model likelihood, and the second is the conditional entropy for variational posterior of continuous hidden states. We can approximate the entropy of qφ(z|x) as:
H(qφ(z|x)) = H(qφ(z1)) + T∑ t=2 H(qφ(zt|z̃1:t−1)) (14)
where z̃t ∼ q(zt) is a sample from the variational posterior. In other words, we compute the marginal entropy for the output of the RNN inference network at each time step, and then sample a single latent vector to update the RNN state for the next step.
In order to apply stochastic gradient descent for end-to-end training, the minibatch gradient for the first term in the ELBO (Eq. 13) with respect to θ is estimated as
∇θEqφ(z|x) [log pθ(x, z)] = Eqφ(z|x) [∇θ log pθ(x, z)] (15) For the gradient with respect to φ, we can use the reparameterization trick to write
∇φEqφ(z|x) [log pθ(x, z)] = E ∼N [∇φ log pθ(x, zφ( ,x))] (16) Therefore, the gradient is expressed as:
∇θL = Eqφ(z|x) [∇θ log pθ(x, z)] , (17) ∇φL = E ∼N [∇φ log pθ(x, zφ( ,x))] +∇φH(qφ(z|x)). (18)
To compute the derivative of the log-joint likelihood ∇θ,φ log pθ(v), where we define v = (x1:T , z1:T ) as the visible variables for brevity. Therefore
∇ log p(v) = Ep(s|v) [∇ log p(v)] (19) = Ep(s|v) [∇ log p(v, s)]− Ep(s|v) [∇ log p(s|v)] (20) = Ep(s|v) [∇ log p(v, s)]− 0 (21)
where we used the fact that log p(v) = log p(v, s)− log p(s|v) and Ep(s|v) [∇ log p(s|v)] = ∫ p(s|v)∇p(s|v)
p(s|v) = ∇
∫ p(s|v) = ∇1 = 0. (22)
For∇ log p(v, s), we use the Markov property to rewrite it as:
∇ log p(v, s) = T∑ t=2 ∇ log p(xt|zt)p(zt|zt−1, st)p(st|st−1,xt−1)
+∇ log p(x1|z1)p(z1|s1)p(s1), (23)
with the expectation being:
∇ log p(v) = Ep(s|v) [∇ log p(v, s)] = ∑ k p(s1 = k|v)∇ log p(x1|z1)p(z1|s1 = k)p(s1 = k)
+ T∑ t=2 ∑ j,k p(st−1 = j, st = k|v)∇ log p(xt|zt)p(zt|zt−1, st = k)p(st = k|st−1 = j,xt−1)
= T∑ t=2 ∑ j,k γ2t (j, k)∇ logBt(k)At(j, k) + ∑ k γ11(k)∇ logB1(k)π(k). (24)
Therefore we reach the Eq. 8.
In summary, one step of stochastic gradient ascent for the ELBO can be implemented as Algorithm 1.
Algorithm 1 SVI for Training SNLDS Use Bi-RNN to compute hxt from x1:T ; Recursively sample zt ∼ q(zt|zt−1,x1:T ) using forward RNN applied to zt−1 and hxt ; Compute parameters A, B and π given x, z; Use forwards-backwards to compute γ11:T , γ 2 1:T−1 from (A,B,π);
Use γ to compute ∇θ,φ log p(x, z); Take gradient step.
A.2 GUMBEL-SOFTMAX SNLDS
Instead of marginalizing out the discrete states with the forward-backward algorithm, one could use a continuous relaxation via reparameterization, e.g. the Gumbel-Softmax trick (Jang et al., 2017), to infer the most likely discrete states. We call this Gumbel-Softmax SNLDS.
We consider the same state space model as SNLDS:
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) ] , (25)
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 Figure 2(a). The inference network for the variational posterior now predicts both s and z and is defined as
qφz,φs(z, s|x) = qφz (z|x)qφs(s|x) (26)
where qφz (z1:T |x1:T ) = ∏ t q(zt|z1:t−1,x1:T ) = ∏ t q(zt|ht)δ(ht|fRNN (ht−1, zt−1,hbt)) (27)
qφs(s1:T |x1:T ) = ∏ t q(st|st−1,x1:T ) = ∏ t qGumbel−Softmax(st|g(hbt , st−1), τ) (28)
where ht is the hidden state of a deterministic recurrent neural network, fRNN (·), which works from left (t = 0) to right (t = T ), summarizing past stochastic z1:t−1. We also feed in hbt , which is a bidirectional RNN, which summarizes x1:T . The Gumbel-Softmax distribution qGumbel−Softmax
takes the output of a feed-forward network g(·) and a softmax temperature τ , which is annealed according to a fixed schedule.
The evidence lower bound (ELBO) could be written as
LELBO(θ,φ) = Eqφz (z|x)qφs (s|x) [log pθ(x, z, s)− log qφz (z|x)qφs(s|x)] (29)
One step of stochastic gradient ascent for the ELBO can be implemented as Algorithm 2.
Algorithm 2 SVI for Training Gumbel-Softmax SNLDS Use Bi-RNN to compute hxt from x1:T ; Recursively sample zt ∼ q(zt|zt−1,x1:T ) using forward RNN applied to zt−1 and hxt ; Recursively sample st ∼ qGumbel−Softmax(st|g(hbt , st−1), τ) using feedforward network applied to st−1 and hxt ; Compute the likelihood for eq. (29); Take gradient step.
A.3 DETAILS ON THE BOUNCING BALL EXPERIMENT
The input data for bouncing ball experiment is a set of 100000 sample trajectories, each of which is of 100 timesteps with its initial position randomly placed between two walls separated by a distance of 10.. The velocity of the ball for each sample trajectory is sampled from U([−0.5, 0.5]). The exact position of ball is obscured with Gaussian noise N (0, 0.1). The training is performed with batch size 32. The evaluation is carried on a fixed, held-out subset of the data with 200 samples. For the inference network, the bi-directional and forward RNNs are both 16 dimensional GRU. The dimensions of discrete and continuous hidden state are set to be 3 and 4. For SLDS, we use linear transition for continuous states. For SNLDS, we use GRU with 4 hidden units followed by linear transformation for continuous state transition. The model is trained with fixed learning rate of 10−3, with the Adam optimizer (Kingma & Ba, 2015), and gradient clipping by norm of 5. for 10000 steps.
An example of training a SNLDS model on the Bouncing Ball task is provided as Figure 5. Early in training, the discrete states do not align well to the ground truth transitions. The three states
transition rapidly near one of the walls and the frame-wise F1 score is near chance values. However, by ten thousand iterations, the model has learned to ignore one state and switches between the two states corresponding to the ball bouncing from the wall. Notably the negative log-likelihood changes by over 10 orders of magnitude before the model learns accurate segmentation of even this simple problem. We hypothesize that the likelihood is dominated by errors in continuous dynamics rather than in the discrete segmentation until very late in training.
A.4 DETAILS ON THE REACHER EXPERIMENT
The observations in the reacher experiment are sequences of 36 dimensional vectors, as described in Kipf et al. (2019). First 30 elements are the target indicator, α, and location, x, y, for 10 randomly generated objects. 3 out of 10 objects start as targets, α = 1. The (x, y) location for 5 of the nontarget objects are set to (0, 0). A deterministic controller moves the arm to the indicated target objects. Once a target is reached, the indicator is set to α = 0. (Depicted as the yellow dot disappearing in Figure 6.) The remaining 6 elements of the observations are the two angles of reacher arm and the positions of two arm segment tips. The training dataset consists of 10000 observation samples, each 50 timesteps in length.
This more complex task requires more careful training. The learning rate schedule is a linear warmup, 10−5 to 10−3 over 5000 steps, from followed by a cosine decay, with decay rate of 2000 and minimum of 10−5. Both entropy regularization coefficient starts to exponentially decay after 50000 steps, from initial value 1000 with a decay rate 0.975 and decay steps 500. The temperature annealing follows the same exponential but only starts to decay after 100000 steps. The training is performed in minibatches of size 32 for 300000 iterations using the Adam optimizer (Kingma & Ba, 2015).
The model architecture is relatively generic. The continuous hidden state z is 8 dimensional. The number of discrete hidden states is set to 5 for training, which is larger than the ground truth 4 (including states targeting 3 objects and a finished state). The observations pass through an encoding network with two 256-unit ReLU activated fully-connected nets, before feeding into RNN inference networks to estimate the posterior distributions q(zt|x1:T ). The RNN inference networks consist of a 32-unit bidirection LSTM and a 64-unit forward LSTM. The emission network is a three-layer MLP with [256, 256, 36] hidden units and ReLU activation for first two layers and a linear output layer. Discrete hidden state transition network takes two inputs: the previous discrete state and the processed observations. The observations are processed by the encoding network and a 1-D convolution with 2 kernels of size 3. The transition network outputs a 5 × 5 matrix for transition probability p(st|st−1) at each timestep. For SNLDS, we use a single-layer MLP as the continuous hidden state transition functions p(zt|zt−1, st), with 64 hidden units and ReLU activation. For SLDS, we use linear transitions for the continuous state.
A.5 DETAILS ON THE DUBINS PATH EXPERIMENT
The Dubins path model3 is a simplified flight, or vehicle, trajectory that is the shortest path to reach a target position, given the initial position (x0, y0), the direction of motion θ0, the speed constant V ,
3 https://en.wikipedia.org/wiki/Dubins_path
and the maximum curvature constraint θ̇ ≤ u. The possible motion along the path is defined by
ẋt = V cos(θt), ẏt = V sin(θt), θ̇t = u.
The path type can be described by three different modes/regimes: ‘right turn (R)’ , ‘left turn (L)’ or ‘straight (S).’
To generate a sample trajectory used in training or testing, we randomly sample the velocity from a uniform distribution V ∼ U([0.1, 0.5]) (pixel/step), angular frequency from a uniform distribution u/2π ∼ U([0.1, 0.15]) (/step), and initial direction θ0 ∼ U([0, 2π)). The generated trajectories always start from the center of image (0, 0). The duration of each regime is sampled from a Poisson distribution with mean 25 steps, with full sequence length 100 steps. The floating-point positional information is rendered onto a 28× 28 image with Gaussian blurring with 0.3 standard deviation to minimize aliasing.
The same schedules as in the reacher experiment are used for the learning rate, temperature annealing and regularization coefficient decay.
The network architecture is similar to the reacher task except for the encoder and decoder networks. Each observation is encoded with a CoordConv (Liu et al., 2018b) network before passing into RNN inference networks, the archicture is defined in Table 3. The emission network p(xt|zt) also uses a CoordConv network as described in Table 4. The continuous hidden state z in this experiment is 4 dimensional. The number of discrete hidden states s is set to be 5, which is larger than ground truth 3. The inference networks are a 32-unit bidirection LSTM and a 64-unit forward LSTM. The discrete hidden state transition network takes the output of observation encoding network in the same manner as the reacher task. For SNLDS, we use a two-layer MLP as continuous hidden state transition function p(zt|zt−1, st), with [32, 32] hidden units and ReLU activation. For SLDS, we use linear transition for continuous states.
See Figure 7 for an illustration of the reconstruction abilities (of the observed images) for the SLDS and SNLDS models. They are visually very similar; however, the SNLDS has a more interpretable latent state as described in Section 4.3.
A.6 REGULARIZATION AND MULTI-STEPS TRAINING
Training our SNLDS model with a powerful transition network but without regularization will fit the dynamics p(zt|zt−1, st) with a single state. With randomly initialized networks, one state fits the
dynamics better at the beginning and the forward-backward algorithm will cause more gradients to flow through that state than others. The best state is the only one that gets better.
To prevent this, we use regularization to cause the model to select each mode equally likely until the inference and emission network are well trained. Thus all discrete modes are able to learn the dynamics equally well initially. When the regularization decays, the transition dynamics of each mode can then specialize. The regularization helps the model to better utilize its capacity, and the model can achieve better likelihood, as demonstrated in Section 4.4 and Figure 4.
Multi-steps training has been used by previous models, and it serves the same purpose as our regularization. SVAE first trains a single transition model, then uses that one set of parameters to initialize all the transition dynamics for multiple states in next stage of training. rSLDS training begins by fitting a single AR-HMM for initialization, then fits a standard SLDS, before finally fitting the rSLDS model. We follow these implementations of both SVAE and rSLDS in our paper. Both multi-step training and our regularization ensure the hidden dynamics are well learned before learning the segmentation. What makes regularization interesting is that it allows the model to be trained with a smooth transition between early and late training.

Review summaries:
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.

This paper proposes a method to segment time series into discrete intervals in an unsupervised way. The data is modeled using a state space model where each state consists of a discrete and a continuous part. The discrete state denotes the segment the system is currently in and the continuous state which is conditioned on the discrete one denotes an uninterpretable feature vector. The transition distributions are non-linear. The observation at each time step is high-dimensional and produced by an emission distribution whose parameters are given by a neural network which takes in the continuous state. Learning and inference is done by maximizing the evidence lower bound (ELBO). Problems with the discreteness of latent variables is circumvented by marginalizing (collapsing) them out using the forward-backward algorithm. Problems with making discrete states meaningful when there are non-linear transitions/emissions is addressed by annealing. This annealing scheme forces the conditional distributions on the discrete state to have high entropy (be close to a uniform distribution) at the start by adding a term to the ELBO objective and the multiplier of this term is decreased as the training progresses. There are actually two terms to do this since one alone didn't work.