Smoothing for exponential family dynamical systems

Dear reviewer,

We would like to thank the reviewer for their extensive feedback and careful reading of our manuscript. Below, we respond in order of your questions, but before doing so, we would like to point to the additional real-data experiment motivated by your feedback about the experimental validation. Additionally, we have also removed the vBF material as per your recommendation in order to make better use of the available space to unpack certain concepts e.g. the predictive plug-in through an illustrating example.

Why does L have different parameters in (10) vs (11)?

This is an inconsistency on our side, thank you. We have fixed this in the revised manuscript. We tried to denote ELBOs where a parametric form for the approximation was not assumed yet using $\mathcal{L}(q)$.

Just below (14), you point out that as a result of the approximations it is no longer a valid ELBO. This seems bad (although not terminal), and I was surprised this issue is not referenced or addressed elsewhere. If this is being optimized end-to-end, and the objective is not a true bound, then you cannot make guarantees on what you are optimizing. Can the authors comment on this.

Thank you for asking -- by not valid, we meant that it would not be an exact ELBO of the marginal likelihood. While this approximation does break the bound, we are doing so for the sake of tractability; one nice intuitive feature of the plug-in approximations used is that in the linear and Gaussian case they are exact. We also want to point out many existing works that are based around approximating difficult to manage terms in an ELBO/log-marginal likelihood with ones that are easier, e.g. site approximations of expectation propagation, tractable approximations ELBO to circumvent stochastic gradients [Keeley et al 2020, Efficient non-conjugate Gaussian process factor models…], or even using Delta approximations [Wang and Blei. 2012. Nonconjugate variational inference].

Additionally, we have evaluated the tightness empirically in Fig.1. using the Kalman smoother and particle filter as baselines for the example linear and nonlinear dynamical systems respectively, we can see even in the nonlinear dynamics case the inferred posterior is close to the particle filter posterior. We have revised the figure legend and title to make this more apparent.

At the bottom of page 4, you say the following: The backward factorization will make gaining insight into the structure of the variational posterior analytically easier; whereas $L_F(q)$ is easier to use in an amortized inference setting where trajectories from the approximate posterior are drawn and used to form a stochastic approximation of the bound (Krishnan et al., 2016; Karl et al., 2016). Again, this is not obvious to me. Nowhere else in the paper is the interpretability of the ELBO considered. You also assert that LF is easier to evaluate compared to LB , but then later on, you evaluate LB , so I don't understand what the difference is.

Thank you for your question. We could have made the point clearer; we were trying to convey that the forward factorization and its ELBO makes it natural to design an inference network that amortizes $q(z_t | z_{t-1})$; however, amortizing $q(z_t | z_{t+1})$ and then evaluating L_B would be comparatively more difficult since we would have to use a Monte Carlo approximation of $E[p(z_t|z_{t-1})]$, and additionally the KL between those distributions would also require a further Monte Carlo approximation. We have tried to explain this further in the revised manuscript. However, what we do evaluate is $\hat{\mathcal{L}}_B$, which after substitution of the plug-in predictive, does not require making a MC approximation of the KL term. We have tried to expand this point more in the main text.

This method be applied to any model with exponential family potentials? Can one mix and match potential types? I’m thinking something like a Poisson-HMM would be a nice demonstrative example of this, and one where there are many readily available baselines you could compare to.

That is correct. Thank you for your suggestion -- in the revised manuscript, we the real data example shows this by using a Poisson likelihood to model neural spike-trains.

Dear reviewer,

We would like to kindly ask if the review uploaded isn't for a different paper than this one. If that is the case, we hope to receive your review so that we could address any concerns you may have.

Best regards, the authors

I sincerely want to apologize to the authors (AC, PCs, and the rest of the reviewers) for my mistake on the initial review, which I accidentally swapped with another submission with a "similar" ID number. I also want to apologize to the authors for receiving a confusing review and if this has caused a bad impression of the reviewing process of ICLR, which I remark is still rigorous and aware of the mistake I accidentally committed. Last but not least, I want to remark my positive consideration of the work which is thorough and of high technical value, even if I remarked some points in my review.

The correct review of this paper is now visible again and in its correct form.

Dear reviewer,

Thank you for taking the time to provide feedback on the manuscript. Motivated by concerns of the limited results, we have included an additional real-data experiment in order to demonstrate the efficacy of the ideas presented applied to learning a dynamical system from neural data. Motivated by other reviewer suggestions, we have also removed the vBF section in favor of using the additional space to add more details and have a more cohesive story. Additionally, we respond to your questions below

1. In this case, the ELBO would not be tight in the general case; but we find it reassuring that the plugin predictive approximation is exact for the linear Gaussian state-space model, and we have tried to include that in the text to aid in building intuition. We have evaluated the tightness empirically in Fig.1. And considering the Kalman smoother and particle filter marginal likelihoods for the example linear and nonlinear dynamical systems, we can see our ELBOs (purple dashed lines) are close to the marginals.

2. This is a good point. The objective ELBO defined by Eq.14 is still a lower bound of some marginal likelihood by its form. However, it corresponds to a prior different from (approximating) the original one defined by the generative model. Many other approaches in fact have to use approximation too, and thus implicitly introduced a new prior/marginal as well.

3. Theoretically it leads to an ELBO of a different marginal as we mentioned in 2) and in turn might increase or decrease the variational gap. Empirically, we have not found this to be a problem in the experiments performed.

Thank you again for your feedback. We feel that the additional experiment and approximation details motivated by your questions have helped increase the quality of our manuscript greatly. We hope that you could consider adjusting your score if you find it warranted given the extra improvements.

Thanks for your time,

The authors

Dear reviewer,

Thank you for taking the time to go through our manuscript, we appreciate your valuable feedback. We will try to address your questions in order. In particular, we would like to point to the additional real-data experiment motivated in part by the feedback you provided about the experimental validation; additionally, we have also added further intuition/motivation behind the plug-in predictive and a walk through of how it is exact in the linear gaussian case. Motivated by other reviewer suggestions, we have also removed the vBF section in favor of using the additional space to add more details and have a more cohesive story. Below, we have tried to answer your questions.

Is it possible to elaborate on what it means to be a “plug-in” predictive distribution? Is this different than a posterior predictive distribution?

Yes, exactly. It was intended as an acting posterior predictive distribution for a single-step into the future. We name it as we do “plug-in” $\bar{q}$ as an approximation to $E_{q_t}[p(z_t | z_{t-1})]$ (the best exponential family approximation as measured by the forward KL).

is it possible to state the number of total parameters that each of the models in the experiments had?

Certainly, we have added these in the experiments section.

One thing that was unclear to me was the claim that "The backward factorization will make gaining insight into the structure of the variational posterior analytically easier". Could this be elaborated on?

Thank you for the question, our wording here could be improved a lot, and we hope revisions to the manuscript clear it up. Since the backward factorization leads to the KLs of expected distributions, our insight was that the KL between exponential family distributions can usually be calculated in closed form, so that combined with the subsequent approximation we would be better equipped to gain analytical insights from the approximate ELBO.

The explicit backward conditional distribution of backward factorization leads to analytical backward information propagation as an additive contribution to the posterior natural parameter, unlike other methods that use deep NN amortization e.g. bidirectional RNN or transformer that blackboxed and mixed the backward information propagation with forward message and likelihood contribution.

3. how aBFS and vBF would scale to more complex problems. I think there is a lot of room for improvement here, and thus think the paper has a lot of unrealized potential in its current form.

Thank you for this question, it has helped motivate us to apply aBFS to real data in order to validate its efficacy. We hope the additional experiment applying aBFS to analyzing a more complex real dataset will help position the paper better.

Again we would like to thank you for your valuable feedback, which we feel has helped to greatly improve the manuscript through additional discussion and applying aBFS to a real dataset to demonstrate its efficacy. We hope that if we have addressed your concerns, you might increase your score to reflect that.

Thank you again for your time,

The authors