Collapsed amortized variational inference for switching nonlinear dynamical systems
Authors: Zhe Dong
Bryan A. Seybold
Kevin P. Murphy
Hung H. Bui
Venue:  (2020)

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 SGD. We show that this method can successfully segment time series data (including videos) into meaningful "regimes", due to the use of piece-wise nonlinear dynamics.

Under review as a conference paper at ICLR 2020
COLLAPSED AMORTIZED VARIATIONAL INFERENCE FOR SWITCHING NONLINEAR DYNAMICAL SYSTEMS
Anonymous authors Paper under double-bling review
ABSTRACT
We propose an efficient inference method for switching nonlinear dynamical sys- tems. The key idea learn an inference network which can be used proposal distribution for the continuous latent variables, while performing exact marginal- ization of the discrete latent variables. This allows 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.
the reparameterization
INTRODUCTION
Consider watching from above an airplane flying across country or a car driving through field. The vehicle's motion is composed of straight, linear dynamics and curving, nonlinear dynamics. This illustrated in fig. [a). 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 the agent (c.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): [inderman et al. (219 as
well as for planning in the context of hierarchical reinforcement learning (c.f., Kipf al.
2019
(a) Inputs
(b) Ground Truth
(c) SNLDS
(d) SLDS
Figure (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 turning clockwise. (c) Segmentation Iearned 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.
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 and 2,, where is the discrete time index. The discrete state, 3t E ..,, ,,KK, K}. represents the mode of the system at time i, and the continuous state, E IR represents other factors of variation, such location and velocity. The observed data is denoted by I; E R,, and can either be a low dimensional projection of 21, such as the current location, high dimensional signal that is informative about 2,, such image. We may optionally have observed input or control signals u, E R which drive the system in addition to unobserved stochastic noise. We are interested in learning generative model of the form

---

Under review as a conference paper at ICLR 2020
Po($1:1: 21:7,, (1:r[u1:1) from partial observations, namely (æ1:T, u1:r). This requires inferring the posterior over the latent states, Po($1:T; æ1:7 (err) where 22 = (Ze, Uz) 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 new way to perform efficient approximate inference in this class of SNLDS models. The key observation is that, conditioned on knowing 1::7 as well as V1:7, we can marginalize out $1:7 in linear time using the forward bard-bacwward algorithm. In particular, we can efficiently compute the gradient of the log marginal likelihood. V E log p($1:2 (i:T, V1:r). where &1.T is a posterior sample that we need for model fitting. To efficiently compute posterior samples 31. we learn an amortized inference network I Iø (z (1:))) ( for the "coolapsedd" NLDS model p(21.T P1 The collapsing trick removes the discrete variables, and allows us to use the reparamcterization trick for the continuous These tricks let us use stochastic gradient descent (SGD) to learn p and jointly, as explained in Section 3 We can then use q as a proposal distribution inside Rao- Blackwellised particle filter (Doucet al. 2000), although in this paper, we just use single posterior sample, as is common with Variational AutoEncoders (VAEs, Kingma & Welling (2014); Rezende 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(zz|zz-13 1%)) 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 be ignored, as explained in Alemi et al. (2018). Our second contribution is new form of posterior regularization, which prevents the aforementioned problem and results in significantly improved segmentation.
We apply our method, as well as various existing methods, to two previously proposed low- dimensional time series segmentation problems, namely a 1d bouncing ball, and a 2d moving arm. In the 1 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. [ for an example), and find that 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.
RELATED WORK
In this section, we briefly summarize some related work.
2.1 STATE SPACE MODELS
We consider the following state space model:
T
pe( z, 8) = ((01lzz)ppt p((z\z))p(2]ss)) (s))
II p( ( (2-1, (oo(ss/Ss-10 $1-1, x1-1)
(1)
where 8, E [I....,K] K} is the discrete hidden state, 2; E R is the continuous hidden state, and D, E RR' is the observed output, in fig. 2 a). For notational simplicity. we ignore any observed inputs or control signals u, but these can be trivially added to our model.
Note that the discrete state influences the latent dynamics 2,, but could trivially make it influence the observations H, as well. More interesting are which edges we choose to add as parents of the discrete state 8t. We consider the case where 8, depends the previous discrete state, 8 1, in
2

---

Under review as a conference paper at ICLR 2020
S+1
R-1
Generative
Figure 2: Left: Illustration of the generative model. Dashed arrows indicate optional connections. Right: Illustration of the inference network. Solid black arrows share parameters O with the generative model, solid blue arrows have parameters & that are unique to The diamonds represent deterministic nodes computed with RNNs: hi is bidirectional RNN applied to 11T, and hū is unidrectional RNN applied to h; and 21-1-
hidden Markov model (HMM), but also depends on the previous observation, 3t-1- 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 2, 2t-1, and s øę-1 It is straightforward to handle such additional dependencies (shown by dashed lines in fig. ((a) in our inference method, which is not true for some of the other methods discuss below.
We still need to specify the functional forms the conditional probability distributions. this paper, make the following fairly weak assumptions:
pl 21) ((aa:)/a(z:)
p( = 1)) = NVzilff (2) Q), p(se|s1-1 = )1 Cat(s))(((((())1.))))
(3)
where fa,..s are nonlinear functions (MLPs or RNNs), (V is multivariate Gaussian distribution, Cat(-) is categorical distribution, and 8(.) is softmax function. R E RXD and Q R learned covariance matrices for the Gaussian emission and transition noise.
If fz and are both linear, and p(st|(t-)) is first-order Markov without dependence 24-1, the model called a switching linear dynamical system (SLDS). we allow 8; to depend on ž;-1, the model called a recurrent SLDS (Linderman al. 2017 Linderman & Johnson 2017). We will compare to rSLDS in our experiments.
If f2 is linear, but fa 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 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. dd). We therefore allow fz (and fz) 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. is simple nonlinear model, such as polynomial regression. If we let fz be a very flexible neural network, there risk that the model will not need to use the discrete states at all. We discuss 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-Markon model (rHSMM, Dai et al. (017)). Rather than having stochastic continuous variable every timestep, rHSMM instead stochastically switches between states with deterministic dynamics. The semi-Mere Mari-Maaoviian structures this work have an explicit maximum duration, which makes them less flexible. A revised method, (Kipf al. 2019), is able to better handle unknown
3

---

Under review as a conference paper at ICLR 2020
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 given by log p(z) S [m; 0 Φ) = (x.a)) [log Po (2r, 2, 8) - log qΦ s [ where ds (z(2, s (x:) is an approximate posterion Rather than computing q using optimization for each æ. we can train an inference network. (a))). which emits the parameters of q. This is known as "amortized inference" (see e.g., Kingma &
Welling (2014))
If the posterior distribution (=. s xx) is reparameterizable, then we can make the noise independent of Φ. and hence apply the standard SGD optimize 0, Φ. Unfortunately, the discrete distribution p(s[2:) 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. (017b), or combining both, such as REBAR (Tucker et al. 201 2017). We will compare against a Gumbel-Softme version of SNLDS in our experiments. The continuous relaxation approach was applied to SLDS models in (Becker-Ehrie However, the relaxation can lose many of the benefits of having discrete variables (Le et al. 2019).
et al., 2019) and HSSM models in (Liu et al. 8a Kipf ct al.
2019).
Relaxing the distribution to a soft mixture of dynamics results in the Kalman VAE (K VAE) model Fraccaro al. (2017). A concern is that soft models may use a mixture of dynamics for distinet ground truth states rather than assigning a distinct mode of dynamics at each step as a discrete model must do. We will compare KVAE in our experiments. In Section 3 we propose new method to avoid these issues, in which we collapse out s so that the entire model is differentiable.
The SVAE model Johnson et al. (2016) also uses the forward-backwer backward algorithm to compute ((s|v); however, they assume the dynamics of z are linear Gaussian, so they can apply the Kalman smoother to compute (z/(v). Assuming linear dynamics can result in over-segmentation, as we have discussed. A continuous states to compute structured mean field posterior (s)) In contrast, we perform approximate inference for z using one forward-backwarrd pass and then exact inference for using second pass, as we explain in Section
algorithm is applied once to the discrete states and once to the
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.. Doscet & Johansen (201I)) When the model is nonlinear (as in our case), we may need a lot of particles to get 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., D.g.ett al. (2001)) has applied RBPF to SLDS models, leveraging the fact that 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(,, (8-1,8 v) in this case. However, this relies on the model being conditionally linear Gaussian. By contrast, we marginalize out p(s)2, 2, v). handle nonlinear models. In this case, it is hard to compute the optimal proposal distribution for sampling from p( p(zz Z v), instead we variational inference to learn approximate this.
can
METHOD
3.1 INFERENCE
We use the following variational posterior: (0.0(2, s|x) g0 (2xxxx))) (e) æ), where po (s)s2, 22z x) is the exact posterior (under the generative model) computed using the
algorithm,
In the case of sequential models, we can create tighter lower bounds using methods such as FTVO Maddison 70) although this is orthogonal to our work.
4

---

Under review as a conference paper at ICLR 2020
and (=(2))) is defined below. To compute (z 2)), we first process æ1.T through bidirectional RNN, whose state at time is denoted by h (As noted Krishnan et al. (017)) due to the Markov assumptions of our model, only need a backward RNN to summarize T rather bidirectional RNN to summarize 21:7, but we use the latter for simplicity.) We then use a forward (causal) RNN, whose state denoted by hf, compute the parameters of g(z; 21:-- æ1.r), where the hidden state is computed based on h and h." This gives the following approximate posterior: 90 (211) II, (1 (1.7) 11. ((z (ti). See fig. 2 for an illustration.
(z) sequentially, and then treat this as "soft evidence" for the We can draw a sample 21:7- ~ T (
backward algorithm are HMM model. The (sample dependent) parameters used in the 1 -1 = k) for t >> 1, given by A,(j, k) = p(s: S = k. 11-1,, B, (k) = p( (a x ())p((
p(8) k). We can then compute the following posterior marginals: 7/ (2, k) p(s; = kt St-1 -j T1:T, E1:T) and (k) p(s, k/ #1:T),
(k) = p( 21 (()))((2)) 81 k), and π (k)
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 is given by
CEL Φ) (x æ, [log po (a:, z po az, z) log (0gn (=)oa) (log 2) log (z(x))
(5)
(6)
Because (z)z) is reparameterizable, we can approximate the gradient as follows:
Vo.e((0, R Vo. Po 2) - V log (=(x)))
where á is sample from the variational proposal & q (2) x (.T) II 4 2,, u.T). The second term can be computed by applying backpropagation through time to the inference RNN. In the appendix, show that the first term given by
T
kV [log (. + Σ V llog B (b))(+))
(8)
3.k
3.3 ENTROPY REGULARIZATION AND TEMPERATURE ANNEALING
When using expressive nonlinear functions (e.g. an RNN or MLP) to model p(z)/ st), we found that the model only used a single discrete state, analogous to posterior collpase in VAES (see e.g., Alemi et al. 2011 8) ) 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 cach time step and uniform prior Pprior (st k) 1/K (Burke et al. 2019). We call this cross-entropy regularizer:
LeB Σ
DEL (pprior (x)
Our overall objective now becomes
C(O, ()) LFLBO ().0)))--000 (
(10)
Note that 0 < LcB co, need to choose the scale of appropriately. To further smooth the optimization problem, we apply lemperature annealing to the discrete state transitions, as follows:
p(s: k| = x1 S( p(* p(spp(s-))a) -R|s 1- 11), 1)), where 7 is the temperature.
At the beginning stage of training, B, T 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)) 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 approach works in continuous end-to-end fashion.
5

---

Under review as a conference paper at ICLR 2020
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 each time step, & arg max g(s,) (.: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 Fi score across frames. Here the Fi score is the harmonic mean of precision and recall, 2 X precision on X recall/ (prili/ppeiiionnno + recall), where precision is the percentage of the predictions that match the ground truth states, and recall the percentage the ground truth states that match the predictions. We also compute the switching-poing F1 by only considering the frames where the ground truth state changes. This measure compliments the frame-wise Fi, 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, it occurs within some small temporal interval around the ground truth as noted in Section 4.3
4.1 ID BOUNCING BALL
In this section, we use simple dataset from Johnson et al. (2016). The data encodes the location of ball bouncing between two walls 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 fz and fg are both MLPs. We found that regularization was not necessary in this experiment. We also consider the case where fx and fs are linear (which corresponds to an SLDS model), the rSLDS model of Linderman et al. (2017), , the SVAE model Johnson et al. (016)), the Kalman VAE (KVAE) model of FFaccaro et al. (217)) and Gumbel- Softmax version of SNLDS as described in Appendix A.2 We use the implementations of rSLDS, SVAE, and KVAEE provided by the authors.
As the data is generated by a simple wise linear dynamics and all models we tested learn a perfect segmentation, shown in Figure TigureBja) and Table This serves as a "sanity check" that we are able to use and implement the ISLDS, 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 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 Oceam's '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 dataset proposed in the ComplLE paper (Kipl et al. 2019). The observations are sequences of 36 dimensional vectors, derived from the 2d locations various static objects, and the 2d joint locations of a moving arm (see Appendix A.4 details and visualization). The ground truth discrete state for this task the identity of the target that the arm currently reaching for (i.e., its "ggal""..
We fit the same 6 models as above to this dataset. Since it is much harder problem, we found that needed to add regularization to model to encourage it to switch states. Figure 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 ISLDS, 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, scores. We also report the F, number from CompIL.E The Compll.F paper uses an iterative segmentation scheme that detect state changes,

---

Under review as a conference paper at ICLR 2020
but it does infer what the current latent state is, cannot include it in Figure b). As in Table | we find that our SNLDS method is significantly better than the other approaches.
(a) Bouncing Ball
Reacher
Ground Truth
SNLDS (ours)
SLDS (ours)
rSLDS
Gumbel Softmax
Segmentation bouncing ball (left) and reacher task (right). From top to bottom: Row 1. ground truth of latent discrete states; Rows 2, 4, 6, 7. the posterior marginals, p(8; k|a1.r 21:7), of SNLDS, SLDS, iSLDS, SVAF, KVAF, and Gunbel-Softmax SNLDDS respectively, where lighter color represents higher probability. CompILE is not included because it represents different model family that directly predicts the segment boundary without calculating posterior marginals at each time step.
Figure 3:
Table 1: Quantitative comparisons (in % =o) for segmentation bouncing ball and reacher task. We
report the Fi scores in percentage with mean and standard deviation over 5 runs. (S.P. for switching point, F.W. for frame-wise, the best mean is in bold.) The F1 score for CompILE is adapted from Kipf ct al. 2019), where only switching point F1 score is provided. The Fi score for KVAE is computed based on taking "argmax' on the 'dynamics parameter network' described Fraccaro al. (2017)
DATASET
Bouncing Ball
Reacher Task
METRIC
Fi (S.P.)
Fi (F.W.)
Fi (S.P.)
Fi (F.W.)
SLDS (Ours)
100.
100.
59.6 + 3.2
81.0 + 3.4
SNLDS (Ours)
100.
100.
78.1 + 4.2
89.0 + 2.0
rSLDS
100.
100.
47.2 3.2
69.8 = 3.5
SVAE
100.
100.
35.3
62.3 - 4.9
KVAE
100.
100.
21.5
33.7 7
Gumbel-Softmax SNLDS
97.6 + 1.8
93.8
03.8
5.0 - 8.7
14.2 9.3
ComplLE
74.3
4.3 IMAGE DATASET (DUBINS PATH)
In this section, apply our method to new dataset that is created by "rendering" point moving in the 2d plane. The motion follows the Dubins mode 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 external control signal uy We replace this with three latent discrete control states: go straight, turn left, and turn right. These correspond fixed, but unobserved, input signals u, (see Appendix A..] for details). After generating the motion, we create series of images, where we render the location of the moving object as small circle on a white background. Our goal in generating this dataset was assess how well we can recover latent dynamics from image data in 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 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 already showed in Section 4.2 that method is much better than these
hhtpps://en... wi
7

---

Under review as a conference paper at ICLR 2020
approaches, well as Kalman VAE and Gumbel-Softmax version of SNLDS, the simpler reacher task, we expect the same conclusion to hold on the harder task of segmenting videos.
Instead focus on comparing SNLDS with SLDS to be represented by a nonlinear model. The results of segmenting one sequence with these models using 5 states are shown in Figure We that the SLDS model has approximate the left and right turns with multiple discrete states, whereas the non-linear model learns a more interpretable representation.
the advantage of allowing each regime
We again compare the models using F1 scores in Table 2 states, we calculated the best single state to match the ground truth. The second used an oracle to optimally merge states match the ground truth. The SNLDS model significantly outperforms the SLDS in both scenarios.
Because the SLDS model used too many versions of the metric. The first greedy metric that optimally assigned
Table 2: Quantitative comparisons (in %) for S(N)LDS on Dubins path. For SLDS, F1 scores with both greedy 1-to-1 matching (Greedy) and optimal merging (Merging) are provided. The switching point F1 scores estimated with both precise matching (Tol 0) or allowing at most 5-step displacement (Tol5) 5).
METRIC
SLDS (Greedy)
SLDS (Merging)
SNLDS
F1 (Switching point, Tol 0)
1.0
4.4 4.4 3.1
11.3 5.7
Fi (Switching point, Tol 5)
33.7 + 2.5
67.0 + 3.4
82.5 - 1.9
Fi (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 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 AR-HMM and SLDS model, and then merge them; similarly, to the SVAE model, Johnson et al. (2016) first train with single latent state and then increase K.
We found a similar strategy was necessary for the Reacher and Dubins tasks, but we do this a smooth way using annealed regularizations. Early in training, train with large temperature and entropy coefficient B. 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 over time. We found best to first decay the entropy coefficient and then decay the temperature T.
Log Likellhood 2.45
Entropy
Loso/F!
Decays
Arnoaing
1.0
Se+5
0.6
Relative
le4s
2
0.6
0.4
5.644
0.2
0.0
100K 150K
200K
300K
50K
100K
150K
200K
250K
300K
Training Iterations
Training Iterations
Figure 4: Comparing the relative negative log- likeliiel (left) and the frame-wise F, scores (right) on Dubins paths with 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 4 demonstrates the effect of 3 different annealing schedules the relative log likelihood (defined as L, Lmins where Lmin min, L,,,,,3 across all three runs, and L, is the negative log-likelihood.), and the Fi score. The green curve starts annealing right away; we the F2 score
8

---

Under review as a conference paper at ICLR 2020
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 ground truth, cannot use the Fi score metric to determine the best annealing schedule. However, seems that the schedules that improve Fi the most also improve likelihood the most.
CONCLUSION
We have demonstrated that our proposed method can effectively learn segment high dimensional sequences into meaningful discrete regimes. Future work includes applying this to harder image sequences, and to hierarchical reinforcement learning.
REFERENCES
Alexander A Alemi, Ben Poole, lan Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing a broken ELBO. In Intl. Conf. Machine Learning (ICML), 2018. URL http://
arxiv.org/abs /1740/000
Philip Becker-Ehmek, Jan Peters, and Patrick Van Der Smagt. Switching linear dynamics for variational Bayes filtering. Intl. Conf. on Machine Learning (ICML), 2019. URL http
//proceedings.nne press
press wtt7/pccker-ehnckk1 html
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. J. of the Am. Stat. Assoc. (JASA),
Michael Burke, Yordan Hristov, and Subramanian Ramamoorthy. Hybrid system identiiicatiiiiinnnnnnnnnii using switching density networks. CoRR,
2019. URL/http:://arivvoor/////
1907. 04360]
Silvia Chiappa. Explicit-duration Markov switching models. Foundations and Trends in Machine Learning, 7(6), December 2014. http://s://
05800
Sy-Minn Chow and Guangjian Zhang. Nonlinear regime-switching state-space (RSSS) models.
Psychometrika, 78(4):740-766 768, 2013. URL
com/article/
15.1007/s113333 013 93
Hanjun Dai, Bo Dai, Yan-Ming Zhang, Shuang Li, and Le Song. Recurrent hidden semi-Markov model. In Inil. Conf. on Learning Representations (ICLR), 2017. URL /tL/stttp:///oeeeeeveee
net
Amaud Doucet and Adam M Johansen. A tutorial particle filtering and smoothing: Fifteen
years later. In Crisan and B Rozovsk (eds.), The Oxford Handbook of Nonlinear Filtering. Ox- ford University Press, 2011. (RRLhtt:: /ww.. stwtssttt......... ox
fchansen_tutor a PPP0111 1.dff
Arnaud Doucet, Nando de Freitas, Kevin P. Murphy, and Stuart Russell. Rao-Blackwellised particle filtering for dynamic Bayesian networks. In Proc. of the Conf. on Uncertainty in Al (UAI), 2000.
URL http:://d/ acm. wttgg//iiaiion... www.a:m.mmmmmmmmmmmmmmmmm
720075
Amaud Doucet, Neil Gordon, and Vikram Krishnamurthy. Particle filters for state estimation of jump Markov linear systems. IEEE Trans. Signal Processing, 49()):613622) 2001.
Thi Duong, Hung H. Bui, Dinh Q Phung, and Svetha Venkatesh. Activity recognition and abnormality detection with the switching hidden semi-Markov model. In Proc. IEEE Conf. Computer Vision and Puttern Recognition (CVPR), 2005. URL http::://eeeeexplrreeeeeee ieee
org/ org//ocuue 3oocmentt/146733
Marco Fraccaro, Simon Kamronn. Ulrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. In Advances in Neural Info. Proc. Systems (NIPS), 2017. htttttss:: htttppss//axxxv.or///bb..tt
1710. 05741

---

Under review as a conference paper at ICLR 2020
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In Intl. Conf. on Learning Representations (ICLR), 2017. URL /
Eorum?id=rkE3y83
Matthew Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta. Structured vaes: Composing graphical models with neural networks for structured representations and fast inference. In Advances in Neural Info. Proc. Systems (NIPS), 2016. URL https
//arxiv. wttp://ggg/looooooooooooo //aaxiv.org/abb//le......
06277
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Intl. Conf. 6980
Learning Representations (ICLR), 2015. URL ULL.ttp://azx....rr///bbb
Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In Indl Conf. on Learning 6114
Representations (ICLR), 2014. UtL.ttppss///xxx...rr//////////
Thomas Kipf, Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Sanchez-Gonzaler Edward Grefen- stette, Pushmeet Kohli, and Peter Battaglia. ComplLE: Compositional imitation learning and execution. In Intl. Conf. on Machine Learning (ICML), 2019. URL
abs/1812. 01483
Rahul G. Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlinear state space models. In Proc. Conf. Am. Assoc. Al (AAAI), 2017. URL httpp::/awww....//////.////////t
1609. 09869
Tuan Anh Le, Adam R Kosiorek, N Siddharth, Yee Whye Tch, and Frank Wood. Revisiting reweighted Wake-Sleep for models with stochastic control flow. In Proc. of the Conf. on Uncertainty in Al (UAI), 2019. http:://.x.i....rrrr./r abs/ 805. 64
Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Panin- ski. Bayesian learning and inference in recurrent switching linear dynamical systems. Conf. Al and Statistics (AISTA7S), 2017. URL http:///prcceeeengg ml www.m.............. press/v54/
inderman 1mm
mtml
Scott Linderman, Annika Nichols, David Blei, Manuel Zimmer, and Liam Paninski. Hierarchical recurrent state space models reveal discrete and continuous dynamics of neural activity in c. elegans. In biorxiv, 2019. URLbtte //tp//dxxxo doi org.10.0.0000.0.00 wtt.://......0001...55.55111111.
Scott W. Linderman and Matthew J. Johnson. Structure-xplooitig variational inference for recurrent switching linear dynamical systems. In IEEE 7th International Workshop on Computational Advances in Multi- Multi-Sensor Adaptive Processing, (CAMSAP). 2017.
Hao Liu, Lirong He, Haoli Bai, Bo Dai, Kun Bai, and Zenglin Xu. Structured inference for recurrent hidden semi-Markov model. In Intl. Joint Conf. Al (IJCAI), pp. 2447-2453, 2018a. URL
https: www. ijcai. org otg/ppoceeeinggs/ /2tp//211/3999
Rosanne Liu, Joel Lehman, Piero Molino, Felipe Such, Eric Frank, Alex Sergeev, and Jason Yosinski. An intriguing failing of convolutional neural networks and the coordconv solution. In Advances Neural Info. Proc. Systems (NeurIPS), volume 0bs//100.002247 2018b. URL
http:://rxii: org/orgggooo......o.o abs 1807.03241 03247
Chris Maddison, John Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Teh. Filtering variational objectives. In Advances in Neural Info. Proc. Systems (NIPS), 2017a.
www.riii.rrg./ieee abs al0///5.. 09279
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: continuous relaxation of discrete random variables. In Intl. Conf. on Learning Representations (ICLR), 2017b.
URL / net ftte/fforumm2 wtt.//ormmm//// forummmdd=11111.....
lain Murray and Zoubin Ghahramani. A note the evidence and Bayesian occam's razor. Technical report, Gatsby Computational Neuroscience Unit, University College London, 2005, http
/mll 1 g eng. cam ac. http:a/www....ooooaaaa//aa/a//aa
10

---

Under review as a conference paper at ICLR 2020
Danilo J. Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approxi- mate inference in deep generative models. In Int.. Conf. on Machine Learning (ICML), 2014. URL
https://arxix: 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: 7/scl. //scl.ece.cc.. ucsb.
htm.
George Tucker, Andriy Mnih, Chris J. Maddison, John Lawson, and Jascha Sohl-Dickstein. REBAR: low-variance, unbiased gradient estimates for discrete latent variable models. In Advances in Neural Info. Proc. Systems (NIPS), 2017. http:t/:: /ttp///pppr.....cooo/.o net forum? ?id- EYBDyeh01
Alexander B Wiltschko, Matthew J Johnson, Giuliano lurilli, Ralph E. Peterson, Jesse Katon, Stan L Pashkovski, Victoria E Abraira, Ryan P Adams, and Sandeep Robert Datta. Mapping sub-sccond structure in mouse behavior, Neuron, 1216)11 2015. URL httthttp::/// ://w:///w
A APPENDIX
A.1 DERIVATION OF THE GRADIENT OF THE ELBO
The evidence lower bound objective (ELBO) of the model is defined as:
C(0,0) Egg (2,,)))))))) [log pe (x, =, s) log qa, (s [log pe (xx =) (s|a,, z) log qs (z
(11)
(12)
= (x)a)) po 2)] + H (z)((z)))
(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 (z(x))) as:
T
H(ga( = H((o(z)))
Σ H(as(zz|21a-))
(14)
where 2, ~ g(zz) is sample from the variational posterior. In other words, compute the marginal entropy for the output of the RNN inference network at each time step, and then sample a single latent vector 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 0 is estimated
V (x)aa))))) [logPo (a:, z)] gx m) Ve log po (ax =)]
(15)
For the gradient with respect to Φb. we can use the reparameterization trick to write
VoBgs(tla) po z)] = E VV logpolo., (())
(16)
Therefore, the gradient is expressed as:
VoC (=)m) [Volog z)]
(17)
V V C = E [ log po (e(e, V VoH((
(18)
To compute the derivative of the log-joint likelihood V (o)) where we define (#1.7, Z1.T) as the visible variables for brevity. Therefore
V ((v) = Ep(elo) [V logp(v)]
(19)
= Ep(alv) Vlogp(u, s)] - Ep(ajo /V [Vlog((s)))
(20)
= Ep(alv)
V [logp(u,s)] 0
(21)
where we used the fact that logp(v) logp(u, logp(s/v) and
Ep(sto) (V [ogp(ssuu)] = /
p(s/v) Vp(s'u) v /((ll)) V p(s|v)
V1 = 0.
(22)
11

---

Under review as a conference paper at ICLR 2020
For V Vlogp(u s), use the Markov property to rewrite as:
V Vlogp(u,8) Σ "pgpp(akkebneeells--soossdosss
V [ogplers==
(23)
with the expectation being:
v log p(v) = Ep(elv) [Vog 8)]
Σ p(si kl//V/ooogpeeeaaaaaaaalsss k)p(s1 k)
T
Σ Ma
k)p(s: = k/8--1 20
j.k
T
ΣV
E
+ Σ log (k))
(24)
1.k
Therefore we reach the Eq.
In summary, one step of stochastic gradient ascent for the ELBO can be implemented Algorithm
Algorithm SVI for Training SNLDS
Use Bi-RNN to compute he from æ1:7;
Recursively sample 2; g( 1t-1, (1,T) using forward RNN applied zt-1 and hT;
Compute parameters A, B and π given æ, =; Use forwards barkddds to compute 7K.T. Vir-1 from (A, B, T);
Use y compute Vo. log p(æ, 2); Take gradient step.
A.2 GUMBEL-SOFTMAN SNLDS
Instead of marginalizing out the discrete states with the continuous relaxation via reparameterization, e.g. the Gumbel-Softmax trick (Jang et al. (017)). to infer the most likely discrete states. We call this Gumbel-Softmex SNLDS.
algorithm, one could use
We consider the same state space model as SNLDS:
p(ce st)(s std
(0 z, 8) p( 1()211)))
(25)
where 9t {I....,K] K} the discrete hidden state, 2N E R- the continuous hidden state, and 2e E R. is the observed output, as in Figure 2(a). The inference network for the variational posterior now predicts both s and z and defined as
Cd..dr (z, s/x) = do. (z|x))o, (s|x)
(26)
where
lo. (0. (2IT((11T))
Il
II a ((2)) 21
(27)
1o. (S1r//E1))
Il q((s 11 E1:T) II [ gh st-1),
(28)
where h is the hidden state of deterministic recurrent neural network. faNN(-), which works from left (t = 0) to right (t = T), summarizing past stochastic æ We also feed in h", which is bidirectional RNN, which summarizes #1.7. The Gumbel-SSfttax distribution 4Gumbel- -oottmax
12

---

Under review as a conference paper at ICLR 2020
takes the output of feed-forward network g(-) and a softmax temperature т, which is annealed according to a fixed schedule.
The evidence lower bound (ELBO) could be written
LELBO Φ) Bug (=) me) (s|=) [log Po ( s) log qΦ. (p, (s/z²)
(29)
One step of stochastic gradient ascent for the ELBO can be implemented as Algorithm
Algorithm SVI for Training Gumbel-Softmax SNLDS
Use Bi-RNN to compute h from & 1:77
Recursively sample 2 Recursively sample $; ~ (Gumbel-Sootmaa: (5 St-1), using feedforward network applied to 8t-1 and h;::
F1.T) using forward RNN applied to z and h;:
Compute the likelihood for eq. Take gradient step.
A.3 DETAILS ON THE BOUNCING BALL EXPERIMENT
The input data for bouncing ball experiment is set of 100000 sample trajectories, each which is of 100 timesteps with its initial position randomly placed between two walls separated by distance of 10.. The velocity of the ball for each sample trajectory sampled from (.50 0.5)). The exact position of ball is obscured with Gaussian noise (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 and 4. For SLDS, we use linear transition for continuous states. For SNLDS, 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.
GRU with 4 hidden units followed by linear
Reconstructed Data
Inputs
1.0
15
5
Log(Relative Likelihood)
20
40
60
80
100
Ground Truth
0.8
10
F1 Score
100
Step Number
0.4
10K
2000
4000
6000
8000
10000
Traing Steps
Figure 5: Left Column: SNLDS Segmentation bouncing ball task with RNN continuous transition function. Top left: illustration input sequence and reconstruction. Center Left (green): ground truth of the latent discrete states, corresponding the posterior marginal of p(s; = 1 21.r) of SNLDS at 100, 1000, 2000 and 10000 training steps, where lighter color represents higher likelihood. Right Column: Training progress of relative negative log-likelihone (Orange) and frame-wise Fl score (Blue) for SNLDS. Relative negative log-likeling likeliild is calculated as In(nlik- min(nllk) 1.), where nllk negative liglliiel.ooo. The scale emphasizes that the loss still improves even late during training.
directions of motion. Lower left (blue):
An example training SNLDS model on the Bouncing Ball task is provided as Figure 55 Early in training, the discrete states do not align well to the ground truth transitions. The three states
13

---

Under review as a conference paper at ICLR 2020
transition rapidly near one of the walls and the frame-wise FI score is near chance values. However, by ten thousand iterations, the model has learned to igmore one state and switches between the two states corresponding the ball bouncing from the wall. Notably the negative log- ligeelklioodd 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
. 7
.
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 Kipt ct al. (2019), not images.
The observations in the reacher experiment are sequences of 36 dimensional vectors, as described in Kipf al.] (201)). First 30 elements are the target indicator, a, and location, æ, y, for 10 randomly generated objects. 3 out of 10 objects start as targets, a 1. The (Ir, y) location for 5 of the non- target 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 warm- up. 10-5 to 10 over 5000 steps, from followed by a cosine decay, with decay rate of 2000 and minimum of 10 15- 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 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 = is 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 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 (==x(2).7)) The RNN inference networks consist of 32-unit bidirection L.TM and 64-unit forward LSTM. The emission network is 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 x 5 matrix for transition probability ((-1) )) at each timestep. For SNLDS, we use a single-layer MLP as the continuous hidden state transition functions (t- 8r), 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 mode is a simplified flight, or vehicle, trajectory that is the shortest path to reach target position, given the initial position (No; yo), the direction of motion Og, the speed constant V,
https://een wi wikipedia a.w.dia..rrgaaa http:/wwwii/ubiiiiuottt
14

---

Under review as a conference paper at ICLR 2020
and the maximum curvature constraint 0 u. The possible motion along the path is defined by i = V cos cos( üt = V sin(0.), 0, - 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 uniform distribution V [( 0.5]) (pixel/step), angular frequency from uniform distribution u/2z ~ 1100,, 0.15)) (/step), and initial direction Op ~ [0 2r)). 'hee generated trajectories always start from the center of image (0,0), The duration each regime is sampled from a Poisson distribution with mean 25 steps, with full sequence length 100 steps. The floating-poin point positional information rendered onto 28 x 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 similar to the reacher task except for the encoder and decoder networks. Each observation is encoded with CoordConv (Liu et al., 2018b) network before passing into RNN inference networks, the archicture is defined in Table 3 The emission network p(æz|Zz) also uses CoordConv network described in Table 4 The continuous hidden state in this experiment is 4 dimensional. The number of discrete hidden states is set to be 5, which is larger than ground truth The inference networks are a 32-unit bidirection LSTM and 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 MI.P as continuous hidden state transition function p(z =t- 8t), with [32, 32] hidden units and ReLU activation. For SLDS, we use linear transition for continuous states.
Table CoordConv encoder Architecture. Before passing into the following network, the image is padded from [28, 28, 1] to [28, 28, 3] with the pixel coordinates.
Layer
Filters
Shape
Activation
Stride
Padding
2
[5, 5]
relu
same
2
4
[5, 5]
relu
same
3
4
[5, 5]
relu
same
4
8
[5, 5]
relu
same
5
8
[7, 7]
relu
valid
6
8
2 (Kernel Size)
None
causal
Table 4: CoordConv decoder Architecture. Before passing into the following network, the input z; is tiled from (8) to (28, 28, 8), where is the hidden dimension, and is then padded [28, 28, 10) with the pixel coordinates.
Layer
Filters
Shape
Activation
Stride
Padding
14
[1, 1]
relu
valid
14
11,
relu
valid
3
28
[1, 1]
relu
valid
28
[1. 1]
relu
valid
[1.
relu
same
See Figure for 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 described in Section
A.6 REGULARIZATION AND MULTI-STER TRAINING
Training our SNLDS model with a powerful transition network but without regularization will the dynamics p(2/)2-11 st) with single state. With randomly initialized networks, state fits the
15

---

Under review as a conference paper at ICLR 2020
B
B
(a) Input
(b) SNLDS
(c) SLDS
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.
dynamics better at the beginning and the flow through that state than others. The best state is the only one that gets better.
algorithm will cause more gradients to
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
Multi-steps training has been used by previous models, and it serves the same purpose regularization. SVAE first trains 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 single AR-HMM for initialization, then fits 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 that it allows the model be trained with smooth transition between early and late training.
16


Reviews:
[{"review_id": "rylQ61g-cB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "SUMMARY:\nThis 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.\n\nSTRUCTURE:\nThe paper is well-written and easy to understand.\n\nNOVELTY:\nI found the technique of estimating gradients using forward-backward to be interesting and potentially useful in other domains when parts of generative models can be marginalized out using belief propagation.\nWhile the problem of unsupervised time-series segmentation is an important one, I'm not sure the proposed technique addresses it completely.\nThe main thing that seems to be doing the work is not the marginalization using forward-backward, but rather the annealing scheme which itself seems ad-hoc and it is not clear whether this is generalizable to other domains or it just happened to work on the problems in the paper.\nIt is not clear why maximizing the entropy of the variational transition should encourage meaningful clustering.\nWould this work even if the emission distribution is made much more powerful?\n\nEXPERIMENTS:\nThere are experiments on three synthetic datasets. While the proposed method beats the competing methods, it is unclear that collapsing is helpful. Also, no annealing was used in the baseline methods (like increasing K in SVAE or multi-step training).\n\nCONCLUSION:\nWhile the problem this paper is tackling is significant, it isn't clear that the proposed method tackles it. I would consider bumping up my score if \n- this method is demonstrated to work on a real dataset and/or\n- there is a better understanding of the principles behind why this annealing scheme helps.\nAlso, proper tweaking of the competing algorithms (similar to annealing) is needed to compare the proposed method fairly.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HyedYpoTKr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nI found the main idea and technique of the paper simple and nice. I am reasonably positive about the paper. The main text of the paper is written well, but the experimental result section seems to be rushed and needs to be polished slightly. I gave weak accept, but if the authors give a convincing answer for my question below, I may raise my score.\n\nI presume that the objective L(theta,phi) in (11) is optimized by a version of gradient ascent. Here is my question related to this:\n\n[Q] Why is H(O) in p5 differentiable with respect to theta and phi? \n\nI am asking this question because the distribution O is defined in terms of arg max, which is not a differentiable operator. Furthermore, the definition of O uses p(s_t|z,x), which uses the model parameters theta. Oh, by the way, I think that the definitions of H and L_CE should include the expectation with respect to q_theta(z|x). \n\nSome minor comments are added below.\n\n* formula (1), p2: p(x1|s1) should be replaced by p(x1|z1)p(z1|s1)\n\n* p3: There are no sub-figures labeled with (a) and (b) in Figure 2. I suggest to put (a) and (b) in front of the captions of the two diagrams in Figure 2. A similar comment applies to Figure 3, because the main text refers to something called Figure 3(a) and Figure 3(b). Also, the paper uses fig. 2(b) sometimes, and Figure 2(b) in other times. Using one convention consistently might help some readers.\n\n* p3: Cat(s_t | S(f_s(...)) ===> Cat(s_t | S(f_s(...)))\n\n* p3: SDLS ===> SLDS\n\n* p3: log p(x) <= L(...) ===> log p(x) >= L(...) \n\n* p4: I found the phrase \"so they need to perform multiple forward-backward (FB) passes\" vague. The algorithm in the paper uses FB twice, and \"multiple\" in the quoted phrase might mean 2, 3 or more. This makes it less clear whether the algorithm has any benefit over the existing approaches.\n\n* p6: This measures compliments ===> This measure complements\n\n* p6: within some small temporal around ... where noted ===> within some range around ... as noted\n\n* p6: are is constant ===> are constant\n\n* p6: The ground truth discrete states ===> The ground truth discrete state\n\n* p7: The resulting of ===> The result of ", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "BklPSoBsFB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Thank you for an interesting read.\n\nAs far as I understand, the paper claims two contributions:\n1. A combination of collapsed variational inference and amortised inference for SNLDS, to make the training pipeline fully differentiable;\n2. An improved loss function upon the variational lowerbound (ELBO) to force the model to use the discrete states.\n\n======= novelty =======\nThe 1st idea is combinatorial: \n1. 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. \n2. 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. \n\nThe 2nd idea is interesting but not very well explained to some extent:\n1. 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. \n2. 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? \n3. 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.\n\n======= significance =======\nExperiments 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.\n\nHowever, two baselines are missing:\n1. 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.\n2. 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.\n\n======= clarity =======\n1. 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. \n2. 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.\n3. 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.\n\n======== references ========\nFraccaro et al. (2017). A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning. NeurIPS 2017", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is an interesting paper on an important topic.  The reviewers identified a variety of issues both before and after the feedback period; I urge the authors to consider their comments as they continue to refine and extend their work.
Decision: Reject