Title: COLLAPSED AMORTIZED VARIATIONAL INFERENCE FOR SWITCHING NONLINEAR DYNAMICAL SYSTEMS

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.

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 s t and z t , where t is the discrete time index. The discrete state, s t ∈ {1, 2, . . . , K}, represents the mode of the system at time t, and the continuous state, z i ∈ R H , represents other factors of variation, such as location and velocity. The observed data is denoted by x t ∈ R D , and can either be a low dimensional projection of z t , such as the current location, or a high dimensional signal that is informative about z t , such as an image. We may optionally have observed input or control signals u t ∈ R U , which drive the system in addition to unobserved stochastic noise. We are interested in learning a generative model of the form p θ (s 1:T , z 1:T , x 1:T |u 1:T ) from partial observations, namely (x 1:T , u 1:T ). This requires inferring the posterior over the latent states, p θ (s 1:T , z 1:T |v 1:T ), where v t = (x t , u t ) 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 z 1:T as well as v 1:T , we can marginalize out s 1:T in linear time using the forward-backward algorithm. In particular, we can efficiently compute the gradient of the log marginal likelihood, ∇ s 1:T log p(s 1:T |z 1:T , v 1:T ), where z1:T is a posterior sample that we need for model fitting. To efficiently compute posterior samples z1:T , we learn an amortized inference network q φ (z 1:T |v 1:T ) for the "collapsed" NLDS model p(z 1:T , v 1: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(z t |z t-1 , v t ) 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:

EQUATION

where s t ∈ {1, . . . , K} is the discrete hidden state, z t ∈ R L is the continuous hidden state, and x t ∈ R D is the observed output, as in fig. 2 (a). For notational simplicity, we ignore any observed inputs or control signals u t , but these can be trivially added to our model.

Note that the discrete state influences the latent dynamics z t , but we could trivially make it influence the observations x t as well. More interesting are which edges we choose to add as parents of the discrete state s t . We consider the case where s t depends on the previous discrete state, s t-1 , as in a The diamonds represent deterministic nodes computed with RNNs: h x t is a bidirectional RNN applied to x 1:T , and h z t is a unidrectional RNN applied to h x t-1 and z t-1 .

hidden Markov model (HMM), but also depends on the previous observation, x t-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 z t on x t-1 , and s t on z t-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:

EQUATION

where f x,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 ∈ R D×D and Q ∈ R H×H are learned covariance matrices for the Gaussian emission and transition noise.

If f x and f z are both linear, and p(s t |s t-1 ) is first-order Markov without dependence on z t-1 , the model is called a switching linear dynamical system (SLDS). If we allow s t to depend on z t-1 , the model is called a recurrent SLDS (Linderman et al., 2017; Linderman & Johnson, 2017) . We will compare to rSLDS in our experiments.

If f z is linear, but f x 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 f z 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 f z (and f x ) to be nonlinear. The resulting model is called a switching nonlinear dynamical system (SNLDS), or Nonlinear Regime-Switching State-Space Model (RSSSM) (Chow & Zhang, 2013) . Prior work typically assumes f z is a simple nonlinear model, such as polynomial regression. If we let f z 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; θ, φ) = E q φ (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.

2.3. MONTE CARLO INFERENCE:

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(s t |s t-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(z t |z t-1 , v), so instead we use variational inference to learn to approximate this.

3. METHOD:



3.1. INFERENCE:

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, and q φ (z|x) is defined below. To compute q φ (z|x), we first process x 1:T through a bidirectional RNN, whose state at time t is denoted by h x t . (As noted in Krishnan et al. (2017) , due to the Markov assumptions of our model, we only need a backward RNN to summarize x t:T , rather a bidirectional RNN to summarize x 1:T , but we use the latter for simplicity.) We then use a forward (causal) RNN, whose state denoted by h z t , to compute the parameters of q(z t |z 1:t-1 , x 1:T ), where the hidden state is computed based on h z t-1 and h x t . This gives the following approximate posterior:

q φ (z 1:T |x 1:T ) = t q(z t |z 1:t-1 , x 1:T ) = t q(z t |h z t ). See fig. 2(b)

for an illustration. We can draw a sample z 1: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 A

t (j, k) = p(s t = j|s t-1 = k, x t-1 ), B t (k) = p(x t |z t )p(z t |z t-1 , s t = k) for t > 1, B 1 (k) = p(x 1 |z 1 )p(z 1 |s 1 = k), and π(k) = p(s 1 = k).

We can then compute the following posterior marginals:

γ 2 t (j, k) = p(s t = k, s t-1 = j|x 1:T , z 1:T ) and γ 1 t (k) = p(s t = k|x 1:T , z 1:T )

, which can be used to compute the gradients of the log likelihood, as we discuss below.

3.2. LEARNING:

The evidence lower bound (ELBO) for a single sequence x is given by

L ELBO (θ, φ) = E q φ (z|x)p θ (s|x,z) [log p θ (x, z)p θ (s|x, z) -log q φ (z|x)p θ (s|x, z)]

(5)

EQUATION

Because q φ (z) is reparameterizable, we can approximate the gradient as follows:

EQUATION

where z is a sample from the variational proposal z ∼ q φ (z

1 |x 1:T ) T t=2 q φ (z t |z t-1 , x 1: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 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 p prior (s t = k) = 1/K (Burke et al., 2019) . We call this a cross-entropy regularizer:

EQUATION

EQUATION

Our overall objective now becomes

EQUATION

)

Note that 0 ≤ L CE ≤ ∞, 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(s t = k|s t-1 = j, x t-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.

4. EXPERIMENTS:

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, ŝt = arg max q(s t |x 1: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 F 1 score across frames. Here the F 1 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 F 1 by only considering the frames where the ground truth state changes. This measure compliments the frame-wise F 1 , 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.

4.1. 1D BOUNCING BALL:

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 f x and f z are both MLPs. We found that regularization was not necessary in this experiment. We also consider the case where f x and f z 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 Gumbel-Softmax 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.

4.2. 2D REACHER TASK:

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 F 1 scores. We also report the F 1 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. 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 u t . We replace this with three latent discrete control states: go straight, turn left, and turn right. These correspond to fixed, but unobserved, input signals u t (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 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 F 1 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.

Table 2 : Quantitative comparisons (in %) for S(N)LDS on Dubins path. For SLDS, F 1 scores with both greedy 1-to-1 matching (Greedy) and optimal merging (Merging) are provided. The switching point F 1 scores are estimated with both precise matching (Tol 0) or allowing at most 5-step displacement (Tol 5).

METRIC SLDS (Greedy) SLDS (Merging) SNLDS:

F 1 (Switching point, Tol 0) 3.5 ± 1.0 4.4 ± 3.1 11.3 ± 5.7 F 1 (Switching point, Tol 5) 33.7 ± 2.5 67.0 ± 3.4 82.5 ± 1.9 F 1 (Frame-wise)

29.4 ± 3.6 61.5 ± 8.0 84.3 ± 7.2

4.4. ANALYSIS OF THE ANNEALING SCHEDULE:

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 τ . In the second run (red), the cross entropy regularization coefficients starts to decay at step 20, 000, while temperature annealing starts at step 40, 000. In the third run (blue), the coefficient decay starts at step 50, 000, while temperature annealing starts at step 100, 000.

Figure 4 demonstrates the effect of 3 different annealing schedules on the relative log likelihood (defined as L t -L min , where L min = min t L t;1,2,3 across all three runs, and L t is the negative log-likelihood.), and the F 1 score. The green curve starts annealing right away; we see the F 1 score is flat, since only one discrete state is used. The red curve starts annealing later, which improves F 1 . 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 F 1 score as a metric to determine the best annealing schedule. However, it seems that the schedules that improve F 1 the most also improve likelihood the most.

5. CONCLUSION:

We have demonstrated that our proposed method can effectively learn to segment high dimensional sequences into meaningful discrete regimes. 

A APPENDIX A.1 DERIVATION OF THE GRADIENT OF THE ELBO:

The evidence lower bound objective (ELBO) of the model is defined as:

EQUATION

EQUATION

EQUATION

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:

EQUATION

where zt ∼ q(z t ) 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

EQUATION

For the gradient with respect to φ, we can use the reparameterization trick to write

EQUATION

Therefore, the gradient is expressed as:

EQUATION

EQUATION

To compute the derivative of the log-joint likelihood ∇ θ,φ log p θ (v), where we define v = (x 1:T , z 1:T ) as the visible variables for brevity. Therefore

EQUATION

where we used the fact that log p(v) = log p(v, s) -log p(s|v) and

EQUATION

)

For ∇ log p(v, s), we use the Markov property to rewrite it as:

∇ log p(v, s) = T t=2

∇ log p(x t |z t )p(z t |z t-1 , s t )p(s t |s t-1 , x t-1 ) + ∇ log p(x 1 |z 1 )p(z 1 |s 1 )p(s 1 ),

with the expectation being:

EQUATION

)

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 h x t from x 1:T ; Recursively sample z t ∼ q(z t |z t-1 , x 1:T ) using forward RNN applied to z t-1 and h x t ; Compute parameters A, B and π given x, z; Use forwards-backwards to compute γ 1 1: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:

EQUATION

where s t ∈ {1, . . . , K} is the discrete hidden state, z t ∈ R L is the continuous hidden state, and x t ∈ R D 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)

where

EQUATION

EQUATION

where h t is the hidden state of a deterministic recurrent neural network, f RN N (•), which works from left (t = 0) to right (t = T ), summarizing past stochastic z 1:t-1 . We also feed in h b t , which is a bidirectional RNN, which summarizes x 1:T . The Gumbel-Softmax distribution q Gumbel-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 L ELBO (θ, φ) = E q φ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 h x t from x 1:T ; Recursively sample z t ∼ q(z t |z t-1 , x 1:T ) using forward RNN applied to z t-1 and h x t ; Recursively sample s t ∼ q Gumbel-Softmax (s t |g(h b t , s t-1 ), τ ) using feedforward network applied to s t-1 and h x t ; 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 Figure 6 : Illustration of the observations in reacher experiment. This is 2-D rendering of the observational vector, but the inputs to the model are sequences of vectors, as in Kipf et al. (2019) , not images.

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(z t |x 1: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(s t |s t-1 ) at each timestep. For SNLDS, we use a single-layer MLP as the continuous hidden state transition functions p(z t |z t-1 , s t ), 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 model 3 is a simplified flight, or vehicle, trajectory that is the shortest path to reach a target position, given the initial position (x 0 , y 0 ), the direction of motion θ 0 , the speed constant V , 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.

:

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.

https://en.wikipedia.org/wiki/Dubins_path

https://en.wikipedia.org/wiki/Dubins_path

Figures and Tables:

Figure 1: (a): Trajectory of a particle moving along a counter-clockwise route. The direction of motion is indicated by the arrow and the brightness, from lower to brighter intensity. (b) Ground truth segmentation into "regimes". Blue is moving straight, yellow is turning counter-clockwise, red is turning clockwise. (c) Segmentation learned by our SNLDS model. (d) Segmentation learned by baseline SLDS model. Note that to model the nonlinear dynamics, the SLDS model needs to use more segments.

Figure 2: Left: Illustration of the generative model. Dashed arrows indicate optional connections.Right: Illustration of the inference network. Solid black arrows share parameters θ with the generative model, solid blue arrows have parameters φ that are unique to q. The diamonds represent deterministic nodes computed with RNNs: h x t is a bidirectional RNN applied to x 1:T , and h z t is a unidrectional RNN applied to h x t-1 and z t-1 .

Figure 3: Segmentation on bouncing ball (left) and reacher task (right). From top to bottom: Row 1. ground truth of latent discrete states; Rows 2, 3, 4, 5, 6, 7. the posterior marginals, p(s t = k|x 1:T , z 1:T ), of SNLDS, SLDS, rSLDS, SVAE, KVAE, and Gumbel-Softmax SNLDS respectively, where lighter color represents higher probability. CompILE is not included because it represents a different model family that directly predicts the segment boundary without calculating posterior marginals at each time step.

Figure 4: Comparing the relative negative log-likelihood (left) and the frame-wise F 1 scores (right) on Dubins paths with 3 different annealing schedules.In the first run (green), the regularization coefficient and temperature start to decay at the very beginning of training. In the second run (red), the cross entropy regularization coefficients starts to decay at step 20, 000, while temperature annealing starts at step 40, 000. In the third run (blue), the coefficient decay starts at step 50, 000, while temperature annealing starts at step 100, 000.

Figure 5: Left Column: SNLDS Segmentation on bouncing ball task with an RNN continuous transition function. Top left: illustration of input sequence and reconstruction. Center Left (green): ground truth of the latent discrete states, corresponding to two directions of motion. Lower left (blue): the posterior marginal of p(s t = k|x 1:T , z 1:T ) of SNLDS at 100, 1000, 2000 and 10000 training steps, where lighter color represents higher likelihood. Right Column: Training progress of relative negative log-likelihood (Orange) and frame-wise F1 score (Blue) for SNLDS. Relative negative log-likelihood is calculated as ln(nllk -min(nllk) + 1.), where nllk is negative log-likelihood. The scale emphasizes that the loss still improves even late during training.

Figure 7: Image sequence reconstruction for Dubins path. The sequence is averaged with early timepoints scaled to low intensity, late timepoints unchanged to indicate direction.



Table content: <table><tr><td>DATASET</td><td colspan="2">Bouncing Ball</td><td colspan="2">Reacher Task</td></tr><tr><td>METRIC</td><td>F 1 (S.P.)</td><td>F 1 (F.W.)</td><td>F 1 (S.P.)</td><td>F 1 (F.W.)</td></tr><tr><td>SLDS (Ours)</td><td>100.</td><td>100.</td><td colspan="2">59.6 ± 3.2 81.0 ± 3.4</td></tr><tr><td>SNLDS (Ours)</td><td>100.</td><td>100.</td><td colspan="2">78.1 ± 4.2 89.0 ± 2.0</td></tr><tr><td>rSLDS</td><td>100.</td><td>100.</td><td colspan="2">47.2 ± 3.2 69.8 ± 3.5</td></tr><tr><td>SVAE</td><td>100.</td><td>100.</td><td colspan="2">35.3 ± 2.6 62.3 ± 4.9</td></tr><tr><td>KVAE</td><td>100.</td><td>100.</td><td colspan="2">21.5 ± 8.0 33.7 ± 7.5</td></tr><tr><td colspan="5">Gumbel-Softmax SNLDS 97.6 ± 1.8 93.8 ± 4.0 5.0 ± 8.7 14.2 ± 9.3</td></tr><tr><td>CompILE</td><td>-</td><td>-</td><td>74.3 ± 3.3</td><td>-</td></tr><tr><td colspan="2">4.3 IMAGE DATASET (DUBINS PATH)</td><td/><td/><td/></tr></table>

Future work includes applying this to harder image sequences, and to hierarchical reinforcement learning. Danilo J. Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Intl. Conf. on Machine Learning (ICML), 2014. URL https://arxiv.org/abs/1401.4082. Kenneth Rose. Deterministic annealing for clustering, compression, classification, regression, and related optimization problems. Proc. IEEE, 80:2210-2239, November 1998. URL http: //scl.ece.ucsb.edu/html/papers%5FB.htm.

Table content: <table><tr><td>George Tucker, Andriy Mnih, Chris J. Maddison, John Lawson, and Jascha Sohl-Dickstein. REBAR:</td></tr><tr><td>low-variance, unbiased gradient estimates for discrete latent variable models. In Advances in</td></tr><tr><td>Neural Info. Proc. Systems (NIPS), 2017. URL https://openreview.net/forum?id=</td></tr><tr><td>ryBDyehOl.</td></tr><tr><td>Alexander B Wiltschko, Matthew J Johnson, Giuliano Iurilli, Ralph E Peterson, Jesse M Katon,</td></tr><tr><td>Stan L Pashkovski, Victoria E Abraira, Ryan P Adams, and Sandeep Robert Datta. Mapping</td></tr><tr><td>sub-second structure in mouse behavior. Neuron, 88(6):1121-1135, 2015. URL https://www.</td></tr><tr><td>cell.com/neuron/fulltext/S0896-6273(15)01037-5.</td></tr></table>

Peer review (iclr2020):
Thank you for an interesting read. As far as I understand, the paper claims two contributions: 1. A combination of collapsed variational inference and amortised inference for SNLDS, to make the training pipeline fully differentiable; 2. An improved loss function upon the variational lowerbound (ELBO) to force the model to use the discrete states. ======= novelty ======= The 1st idea is combinatorial: 1. The forward-backward algorithm is a standard inference method for HMM-like sequence models; collapsed variational inference has been investigated extensively in 2000s when hierarchical Bayes models were actively developed; amortised inference is widely used in variational auto-encoders. 2. The combination of the above two inference methods on S(N)LDS is new to the best of my knowledge. However, this combination has been proposed on a similar model called Kalman VAE ( Fraccaro et al. 2017), where the sequence model can be viewed as a *soft* version of SLDS. The 2nd idea is interesting but not very well explained to some extent: 1. The goal of the modified objective function is to encourage the model to use the discrete states (instead of pushing all useful information to the continuous states). It is interesting as it regularises the *exact posterior* of the discrete states conditioned on the *approximately inferred* continuous states. 2. I believe the entropy regulariser is non-differentiable as it is based on a *histogram* estimate of the temporally averaged discrete state distribution. How exactly is this regulariser implemented? 3. I agree adding the KL regulariser can avoid the iterating assignment pathology, however, is random assignment of the regime preferred in any case? From the introductory example, I think contiguous segments are preferred. ======= significance ======= Experiments consider 3 synthetic examples for sequence segmentation (so that ground truth is available). The proposed approach performs significantly better which is a good sign. The paper also provides useful analysis on the effects of balancing parameter tempering which is always welcome. However, two baselines are missing: 1. To claim the significance of the collapsed variational inference approach, a non-collapsed inference version of the proposed SNLDS model needs to be compared. The authors did discuss this and mentioned possible workarounds (e.g. using the Gumbel-softmax trick for discrete state inference), but the comparison is not reported. If compared, this will serve well as an ablation study for the inference method. 2. The paper also provides comparisons across models, but I do think the Kalman VAE model needs to be compared to SNLDS. Both models are more flexible than the original SLDS, but the complexity is added in different ways. Since I think the inference mechanisms are similar (both using forward-backward inference for top-level latents and amortised inference for bottom-level latents), this comparison would provide a better ablation study on the modelling side. ======= clarity ======= 1. The paper presentation is overall clear to me, although I found the many sentences in parenthesis a bit distracted, so I would suggest maybe using footnotes for them instead. 2. For readers who are less familiar with HMMs/forward-backward algorithms, the papers can be difficult to understand, as it skips all the detailed computation of the gamma terms. I would suggest adding the details in the appendix, and/or visualise the intuition using e.g. message passing on factor graphs. 3. I found the related work well presented with most relevant papers, although I do think the Kalman VAE approach is highly relevant which needs to be cited and discussed. ======== references ======== Fraccaro et al. (2017). A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning. NeurIPS 2017