Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This paper proposes a framework for learning neural network-based state-space models from raw data. In comparison to the standard RNN-based approach [1], there are several modifications: (i) a modified ELBO formulated in the context of a constrained optimization (CO) framework [2]. CO prioritizes reconstruction term at the beginning of the training; (ii) hierarchical prior [3] instead of a simple Gaussian; (iii) a modified Kalman-VAE framework [4] instead of RNN-based transition model. Experimental evidence on pendulum data and the reacher environment indicates that the resulting framework outperforms baselines in predictive modeling and ELBO values.


[1] A Recurrent Latent Variable Model for Sequential Data

[2] Taming VAEs

[3] Learning Hierarchical Priors in VAEs

[4]  A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning
main_review: *Originality*: The resulting framework is a combination of three existing ideas, as indicated in the summary above. Additional modifications were introduced to adapt the CO framework to the sequential setting and to replace LSTMs in the Kalman-VAE model in the spirit of [5]. Overall, the originality is somewhat limited. This is generally not a big issue in my mind, as long as combining the ideas from the previous work is well-justified.

*Motivation:* My main issue with this paper is motivation. The authors claim that each introduced modification is beneficial for improving the predictive model of the described dynamical system. However, I found it difficult to understand why would introduced changes make a substantial difference. For example, the authors claim that 'RNN-based transition models as in [6, 7] can lead to a non-Markovian state-space, where the latent variables do not capture the entire information about the system’s state.'. However, I do not find this well-justified and it is completely unclear to me why this would be the case. The authors also claim that 'As a consequence of the RNN-based transition models, the angular velocity is not encoded in zt but in the RNN." which is also unclear. It is also unclear to me why CO is necessary -- why it would be beneficial to optimize the reconstruction term (the distortion) first and then the KL term (the rate). Using a more complex prior can certainly improve ELBO, but this is a very minor modification that leads to non-surprising results.
I would be willing to improve my score if the authors convinced me why the proposed modifications are necessary and beneficial.
At the present state, the overall justification of the proposed framework boils down to experimental results.

*Clarity:* The paper is fairly well-written overall and relatively easy to follow. It is also well-structured.

*Quality:* The methodology seems technically sound.  Combined with the supplementary material, there are several interesting experimental studies. Including more baselines would make the work more convincing.

*Significance:* The general topic is relevant. The presented results are performed on the synthetic data so it is difficult to judge the practical significance. One big drawback is that the source code is missing, especially since the overall framework is fairly complex and hence difficult to implement from scratch.

[5] Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data

**Update:** Raised the score to 7.
limitations_and_societal_impact: The authors have addressed the limitations. There are no obvious concerns in terms of societal impact.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper proposes various improvements to deep sequential latent variable models (or deep state space models). Namely, 1) the constrained optimization method from the Taming VAEs paper [Rezende & Viola, 2018], 2) a method for improving empirical priors over the first time step’s latent variable (referred to as VHP), and 3) a new model that combines the locally linear mappings of Karl et al., 2017 and the KVAE variable separation from Fraccaro et al., 2017. These improvements are demonstrated on pendulum and reacher datasets (both with images and joint angles), showing improvements in system identification, prediction accuracy, and the ability to use the state representations from MBRL. Overall, the paper provides several helpful tips for constructing and training deep state space models.
main_review: **Originality**: The paper centers on applying several ideas developed for static latent variable models to the dynamic setting. Likewise, the proposed model (extended Kalman VAE) draws heavily from Fraccaro et al.’s KVAE and locally linear mappings that have been used previously for sequential latent variable models, e.g., Karl et al. 2017. The main novelty of this paper is in demonstrating that the aforementioned techniques are also useful in learning sequential models. While I believe that these findings are useful, they are not particularly surprising, as the previous techniques are not specific to static models. One point of originality that I appreciated was the application of this technique to the reacher environment, both with pixel observations and for the purposes of model-based control. This is helpful for connecting the sequential generative modeling community with the RL community.


**Quality**: The paper is of fairly high quality, combining multiple complex techniques and demonstrating improvements over several baseline models. Implementing and evaluating these techniques is non-trivial and at the forefront of current research. The experiments in the paper also help to illuminate some overlooked aspects of sequential latent variable models, particularly tied to representation learning, e.g., the importance of constrained optimization and the Markovian latent spaces. The experiments analyze multiple metrics across two dynamics environments.


**Clarity**: The writing is overall quite clear. My main concern here is regarding the complexity of the resulting model and inference/learning procedure. The paper combines multiple complex ideas: constrained optimization, variational hyperprior (VHP), locally linear dynamics, and the already fairly complex Kalman VAE. Given the number of elements in this paper, the main message of the paper is rather dispersed. Including diagrams of the various techniques and model/inference dependencies would be immensely helpful for parsing the paper, even if these were included in the appendix. Additionally, some of the main points of the paper, e.g., the benefit of constrained optimization for representation learning, are only apparent from results tables. The figures in the main paper currently only demonstrate the ability to learn disentangled latent representations. Figure 7 from the appendix, which compares the latent representations between constrained optimization and annealing, strikes me as a stronger demonstration; I would consider moving this to the main paper.


**Significance**: The significance of the paper is one of my larger concerns. Again, given the number of components to the final model, it is somewhat difficult to assess where the benefits are coming from and the take-aways for the generative modeling community. While some of the qualitative and quantitative results are clear, e.g., the benefits of constrained optimization for system identification, the significance of other quantitative results, e.g., around VHP, are less obvious. For instance, how meaningful is a difference of 1 (assuming nat) in test ELBO? The authors also present results using their model for model-based RL. Yet, it is unclear how this differs in capabilities from other models, i.e., any generative model could, in principle, be used for model-based planning. The authors rely on particular assumptions around encoding reward functions in latent space and conditioning their parametric policy on the latent variable. Thus, while I appreciate the model-based RL experiments, this doesn’t strike me as a significant contribution.

**Other Points**:

- In equation 1, the Markov assumption is not necessarily standard for modern models. I would consider a more general formulation in which this assumption is dropped.
- Line 59: the recognition model shouldn’t necessarily need to learn the dynamics. E.g., Marino et al. present an amortized filtering method that does not require learning the dynamics within the inference model.
- Eq. 2: The variety of script F’s and L’s is somewhat confusing.
- Eq. 5: Previous methods (as well as this formulation) don’t actually use EM. Instead, they use amortized variational inference, amortizing the E-step optimization in an inference model.
- Line 91: It would be useful to define what “over-regularization” means here.
- Line 102: The acronym VHP is not defined (variational hyper-prior).
- Line 103: Why just the initial time step?
- Line 106: When will this bound be tight?
- Line 116: Is this two-stage procedure well-justified? It seems somewhat ad-hoc.
- Section 4.3 was difficult to follow.
- Section 5: There are some previous works in sequential VAEs that are not cited, e.g., Chung et al., 2015.
- Line 221: these are not necessarily the papers I would cite for stochastic gradient variational Bayes, as they make the additional assumption of amortizing inference. Other papers, such as BBVI (Ranganath et al.), may be more appropriate.
- Figure 1: Wasn’t lambda supposed to be in [0,1]?
- Table 1c: Why aren’t there other models here?
- Figure 3: I like this example. but isn’t it a bit unfair to compare while conditioning on 1 time step? In practice, such models would be used after a “burn in” of multiple time steps.
- Section 6.4: the benefits of this particular model for model-based RL could be explained better. It’s unclear from the writing. Also, the goal encoding process is unclear.

Marino et al., 2018. A general method for amortizing variational filtering.;
Chung et al., 2015. A recurrent latent variable model for sequential data.;
Ranganath et al., 2013. Black box variational inference.

---
**Update**: I have read the other reviews and the authors' response. Each of the other reviewers raised separate concerns regarding the paper. Reviewer oVeY raised concerns regarding the motivation of the ideas in the paper. I agree with the authors, though, regarding the justification of the Markov assumption and constrained optimization. Reviewer SDjn felt that the paper was somewhat disjoint in its presentation. The (lack of) cohesion of the paper, in my view, is not such a concern as to prevent readers from understanding the concepts and take-home points. Finally, reviewer Qj5e raised various concerns regarding the model/evaluation. The authors provided a detailed response to these questions, none of which struck me as serious issues with the paper. Thus, I do not see the points raised by the other reviewers as meaningfully detracting from the paper, and I am keeping my score at 7. This paper explores several useful ideas for dynamics modeling with deep latent variable models, which are demonstrated with a thorough empirical evaluation. The research community would benefit from the ideas in this paper, thus, in my view, warranting publication.
limitations_and_societal_impact: The limitations of the paper are not addressed in-depth. The authors claim that they have shown that “prediction accuracy strongly depends on the model itself.” I would like to see some discussion around how other models could be used for model-based RL and whether this technique could be scaled to more complex settings, such as natural video or audio.

There do not appear to be specific negative societal impacts associated with this paper.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The authors make three contributions. 1) They apply the constrained optimization framework for VAE type models to the deep state space model case. 2) They add in the hierarchical prior approach to model the initial distribution for the first state. 3) They propose a model with linear transition matrices given by a neural network weighting learned basis matrices. Auxiliary observation variables are also introduced to allow for a linear observation model along with an extra encoder/decoder. This means filtering and smoothing distributions can be computed analytically.
main_review: Overall, I think this paper has some interesting ideas, however, when reading the paper it is hard to distil the paper to a contribution that can be easily built upon and used by other researchers. The paper feels like it should be split into two papers meaning each contribution in isolation can be more thoroughly explained and evaluated.

Beginning with the constrained optimization framework, judging by the experiments included, this looks to be very helpful for learning consistent and interpretable transition functions. In experiments it looks to consistently and significantly improve performance, not only in the author’s proposed model, but in all the baselines too. I think this is a very interesting observation and feel like it should be explored further. For example, are there any heuristics for setting the desired reconstruction quality D_0  parameter and what effects does this hyperparameter have on the quality of learned models. Also, it may be interesting to look at how the rate-distortion framework from http://proceedings.mlr.press/v80/alemi18a/alemi18a.pdf can be interpreted from this sequential learning viewpoint.

Along with this, the authors also propose a new model class - the EKVAE - which I feel could be a paper in its own right. The CO framework can be integrated with a wide variety of models so I am unsure about why these two contributions are in the same paper, it makes the overall message hard to discern. It would be good to have some more intuition on the design choices for the EKVAE for example why it was chosen that each matrix should be a weighting of a set of base matrices, did this give better performance than other models in preliminary experiments?

Finally the hierarchical prior does seem to help models learn a better initialization for their generative model. Though in some cases, e.g. Fig. 8 it seems the model can still learn a useful transition model even with the naïve prior. This then brings into question whether the added complexity of the hierarchical prior is useful. I think this should be taken into account in the experiments as it is important to separate out the benefit of the prior and the transition model. For short sequences e.g. pendulum then the marginal better results with the prior could be just due to the first two frames as opposed to a higher quality transition model. For the longer reacher dataset if it was just the first two frames then the marginal difference between having the prior and without would be even less but these results are not included in the table. Overall, I think it would be good to discuss whether the prior just helps the first 1 or 2 states or actually results in a higher quality transition model.

Regarding clarity, I think the paper is well written and organized. Each of the claims are demonstrated in their own experiment though I feel each claim could do with more expansive investigation (see above comments).

**Edit after rebuttal**

I have read the author responses and other reviews. I believe the paper has good ideas and encourage the authors to continue working on this problem but I believe the paper is not quite ready for publication yet in its current format. This is due to the ideas being presented in a way that makes it more difficult for other researchers to be build on them.  I think the paper would greatly benefit from a clearer focus on single elements of the proposed design or by achieving improvements on a significant problem that could not be achieved if any one of the proposed elements was not included.

**2nd update**
After further consideration of author responses and other reviews, I have increased my score to 6. I believe the ideas presented are novel and with a bit extra motivation of how the proposed methods in the paper are related and both necessary in a camera ready, I think it warrants publishing.

limitations_and_societal_impact: The limitations of the specific contributions could be explained more thoroughly e.g. dependence on hyperparameters, linearity assumptions etc.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper has two main contributions. 
1. The authors adapt REWO for helping with disentanglement of learned representations and improving predictive accuracy
2. The authors replace the RNN encoding of the nonlinearity in [6] by the locally linear model of [11-12] while keeping closed form inference

The authors then illustrate how combining these together leads to better accuracy in learned latent dynamics for a set of models, in particular they show that their approach can recover the unobserved velocity as a latent variable in motion models.
main_review: Overall I believe this paper shows promising empirical results, however it should be carefully revisited in order to be publishable.

Section 3

The statement "Thus, one can assume that initial-state samples uniformly cover the manifold of the learned latent representation" l97-98 is not clear. How is this assumption/statement a consequence of batching/splitting time series?
I would argue that Algo 1 should be reproduced in this section, as it stands, you list it as one of your main contributions, but the details are absent from the core of the paper.


Section 4

The authors mention that the main reason they resort to learning a surrogate Jacobian matrix is essentially computational complexity as using AD to compute it would require $D_z$ backward passes. I agree that this is indeed computationally expensive. However, as it stands their proposed parametrisation is not related to extended linearisation but rather to sigma-points one, which, similarly to the authors suggestion considers a linearisation $F$ as a weighted sum of "base" matrices. The main difference between the two approaches being that while the authors consider the weights to depend on $z,u$ while sigma-point methods implicitely consider that the "base" matrices $F^{(m)}$ do instead. 

This consideration makes the claim that $F_{\psi} \approx J_f$ feel wrong to me and makes the name EKVAE a bit misleading. If anything the authors are trying to jointly learn the best "basis" of transition matrices $F^{(m)}$ at the same time as they try to learn how to project $f$ on it. This is indeed a valid idea, but would in my opinion deserve to be compared with standard sigma-points techniques applied to an explicit transition model. I expect these would be highly competitive for at least three reasons: sigma-points methods are highly parallelisable on GPU and don't require expensive Jacobian calculations, the models considered are fairly low-dimensional both in $z$ and $a$ (as per Appendix 8 if I understand it correctly), and the model posterior distribution should be unimodal by design.

In Eq 15, the authors mention that they use MC to integrate $z_{t-1}$ out in order to compute the observation log-likelihood increment. However this contradicts the fact that they are assuming that $z_t, z_{t-1}$ are jointly Gaussian, so that they should be able to marginalise it in closed form too?

The authors mention that they chose to use a fixed matrix $H = \delta_{ij}$ instead of learning it as was done in e.g. [6]. I agree that this is useful for disentanglement as it reduces unidentifiability within the model. It is however unclear if the resulting improvement of EKVAE over KVAE shown in section 6 comes from this only or from better encoding of the state transitions (except if the authors chose to run KVAE experiments with H (C in [6]) fixed too?).

Section 6

The most concerning aspect of this section consists in the choice of the hyperparameter $\mathcal{D}_0$, for which the authors give no principled rule or heuristics. Was it optimised for or curated in some way? If so this would make the comparison with fixed annealing schedules if these were not curated. Additionally the impact of changing $\mathcal{D}_0$ should be studied at least on the simplest model so as to provide the reader with a clear understanding of performance improvement/deterioration.

The phrasing of what the experiments exactly are is a bit unclear (e.g. l276-278), this would be explained better by wrtting it more formally: what exactly are the authors sampling from? Additionally I don't understand why the authors are conditioning on the latent $z_1$ to measure predictive accuracy when conditioning on the observations up to time $t$ in order to predict the future would make more sense. Moreover, the way I understand this is that the conditioned $z_1$ was fixed once and for all and not sampled from, which makes the values in the tables a, b, c random. If I am correctly understanding it makes it hard to know if the improvement is statistically significant.


Comments below did not affect my evaluation but would need to be corrected in the event of acceptance.

Typos:
l18: "sequence data" should probably be "sequential data"
l90:  It is common practise
l116: In the initial phase (upper case)
l119: the VHP $(\psi_0, \phi_0)$  and the transition model $\phi_0)$  ---> these are the *parameters* of the models you want to optimise
l189: range of application; result in a poorly trained 

There are a certain number of issues with the bibliography, a non-exhaustive list being:
- Kalman and Tung are missing first names
- [5] is broken
- several middle names are missing the full stop, e.g. [15], [18]

## Update:
After having read other reviews and authors responses (in particular the back and forth on my comments), I have improved my rating from 4 to 6. I still believe there are a number of issues with the presentation of the paper which would warrant a more thorough correction than the rebuttal format allows. However, I now believe that this work could be published given the minor additional corrections suggested, although would be more beneficial to the research community if it was partially reorganised (I have comments similar to that of reviewer SDjn). 
limitations_and_societal_impact: The authors do not study the impact of the hyperparameters on the learned models, the most critical one being in my opinion $\mathcal{D}_0$, as per the authors description, is critical in learning predictive model. 

I believe the authors methodology only applies to model with unimodal latent posteriors (due to their linearisation format)

The authors only consider somewhat short time series (30 time steps at most). It is not clear how this gradient based learning would work with longer datasets.
review_title: null

Meta-review: This paper develops a methodology for learning deep state-space models to obtain accurate temporal predictions of observed dynamical systems. This is achieved by combining various techniques: a constrained optimization (CO) method (along the lines of the GECO method of Rezende & Viola, 2018), a method to improve the prior for the first latent (VHP) and a modified extended Kalman-VAE (EKVAE) framework (akin to Karl et al., 2017). The authors show that the combination of these techniques provide improvements over baselines in predictive modeling and ELBO values.

After the rebuttal, two reviewers have increased their score (Reviewers oVeY and Qj5e). The final scores are 7, 7, 6, 6. 
The main concerns raised by the reviewers in the discussion, in particular Qj5e and SDjn (as well as by Reviewer xnvk in their review) is that the paper lacks focus. While CO applied to state-space models is interesting, the addition of th EKVAE and VHP seems more marginal and makes the resulting model and inference/learning procedure fairly complex. The message of the paper is not quite clear and it is difficult to parse as a result. Because of the different elements, it was felt that crucial details were only mentioned in passing.  The paper would benefit from having an improved presentation (see suggestions by Reviewer xnvk) and, if the authors want to keep  CO, VHP and EKVAE in the paper, a significant application where all the proposed elements are necessary (CO, VHP and EKVAE).

This paper was discussed at length between the AC and SAC.


Decision: Accept (Poster)