eXponential FAmily Dynamical Systems (XFADS): Large-scale nonlinear Gaussian state-space modeling

Thank you for taking the time to review our manuscript. Below we respond to your suggestions/questions starting with the weaknesses section.

The mentioned limitations of previous work's limitations are not convincing. Fast evaluation of the loss function is hardly the problem that previous DSSM approaches tried to tackle.

Thank you for your remark — one of the reasons fast evaluation of the loss and approximate posterior statistics is not a problem for previous DSSM approaches is that they restrict the approximate posterior covariance to be diagonal. When the diagonal restriction is dropped, then evaluating the KL terms requires linear algebraic operations (log-det, vector multiplication with matrix inverse, trace) that scale cubic with latent dimensionality, $L$. However, structure in the posterior covariance because of our particular sample approximation and low-rank updates make it possible to carry out those operations in linear time.

The related work section also seemed to stop past 2016

Thank you for pointing this out – we are aware of those works and there are many others as well we would have wished to include in the related works section (see [1a,2a,3a] below to name a few) but were ultimately limited by space.

We highlighted particular works we did because we aim at general purpose approach with minimum assumptions about the dynamics equipped a scalable inference algorithm; we're not focusing on arguing for any apriori latent structures or dynamics model that are tailored for particular problems. We demonstrated with simple yet general networks to showcase that with a good inference method you can still achieve sota results. Motivated by your comment, we will slightly condense the current related works section so that we can briefly touch upon recent works such as [1-5 and 1a-3a]

Eq. (6) is slightly incorrect: The expectation is taken only over z_{t-1}, however, the terms inside the expectation involve z_{t-1} and z_{t}.

We apologize that space was tight and we couldn’t add a few more lines in the main text, but the expectations for Eq(6) do contain the two slice marginals. Factorizing the variational posterior as $q(z_{1:T}) = q(z_1) \prod q(z_t \mid z_{t-1})$, we can write the standard ELBO as follows to arrive at Eq(6) involving the expected KL.

$\mathcal{L}(q) = \int q(z_{1:T}) \log \frac{p(y_{1:T}, z_{1:T})}{q(z_{1:T})} \, dz_{1:T}$ $= \sum \int q(z_t) \log p(y_t\mid z_t) \, dz_t - \sum \int q(z_{1:T}) \log \frac{q(z_t\mid z_{t-1})}{p(z_t\mid z_{t-1})} \, dz_{1:T}$ $= \sum \int q(z_t) \log p(y_t\mid z_t) \, dz_t - \sum \int q(z_{t-1}) q(z_t\mid z_{t-1}) \log \frac{q(z_t\mid z_{t-1})}{p(z_t\mid z_{t-1})} \, dz_{t-1, t}$ $= \sum \mathbb{E}\_{q_t} \left[ \log p(y_t \mid z_t) \right] - \sum \mathbb{E}\_{q_{t-1}} \left[ \mathbb{D}\_{\text{KL}}(q(z_t\mid z_{t-1}) || p(z_t \mid z_{t-1}) )\right]]$

For the same smoothing objective in an alternative paper see [4a] Eq.(6).

What exactly is the pseudo-observation \tilde{y}?

The pseudo observation $\tilde{y}\_t$ encodes current/future observations, $y_{t:T}$, into a Gaussian potential. Since it is a Gaussian potential, it is specified by parameters that interact linearly with the sufficient statistics – the form of those parameterizations are given explicitly by Eq (25) in the main text. In the text, we tried to make this distinction by writing “Importantly, pseudo-observations defined this way encode the current and future observations of the raw data – an essential component for transforming the statistical smoothing problem into an alternative filtering problem."

line 138: I can not follow this argument. Why is smoothing turned into a filtering problem? It becomes more confusing that the introduction of a different distribution pi(z_t) is supposed to help. Please, can you elaborate

Thank you for the question – we say smoothing is turned into a filtering problem because pseudo observations, $\tilde{y}\_t$, encode $y_{t:T}$ into a single Gaussian potential, meaning we can filter pseudo observations to obtain statistics of the smoothed posterior; in other words, we have that $\pi(z_t) \approx p(z_t \mid \tilde{y}\_{1:t}) \approx p(z_t \mid y_{1:T})$ . The reason that we introduce $\pi(z_t)$ is because we want tractable Gaussian marginal approximations, and $q(z_t \mid z_{t-1})$ is conditionally Gaussian.

eq (11): using the sum of (prior and likelihood) is not novel

Thank you for pointing this out, but we respectfully disagree that this particular parameterization is not novel. We were inspired by works like [5a] and also [4], to use this type of parameterization; but in those works the dynamics are conditionally linear and it is straightforward to apply Gaussian message passing and recover the posterior – here the dynamics we consider are arbitrary and nonlinear, which led us to develop the nonlinear filtering procedure given by Eqs.(27) & (28) to make this possible.

what is the difference between predictions from filtering and smoothing in Fig. 3?

You are correct the last filtering and smoothing state should be identical. Our purpose was to show what decoded hand position/latents would look like during filtering in the case of a streaming data setting.

[1a] Ansari et al, 2023. Neural Continuous-Discrete State Space Models for Irregularly-Sampled Time Series. [2a] Li et al, 2021. Scalable gradients for stochastic differential equations. [3a] Schimmel et al, 2022. iLQR-VAE : control-based learning of input-driven dynamics with applications to neural data. [4a] Krishnan et al 2016, Structured Inference Networks for Nonlinear State Space Models [5a] Johnson et al 2016, Composing graphical models with neural networks...

Thank you again for your thoughtful suggestions and useful comments. We hope if we have addressed them you might raise your score to reflect that. We are very motivated to position our paper as best possible, and appreciate all of your helpful input.

We thank the reviewer for their helpful suggestions. We apologize about the density of the paper at times and appreciate that you mentioned the possibility of moving some algorithmic discussion to the appendix in exchange for higher level intuitions and more in depth discussion pertaining to the experimental results. We wholly agree, and as we echoed in the global rebuttal statement we are taking those suggestions to heart and moving some details from the sampling

Based on what I wrote in weaknesses section, I would like some clarification on my comments and know how the authors plan to improve the clarity of the experimental section

In addition to what we wrote previously addressing this, motivated by your suggestions, we are restructuring the manuscript to enhance clarity of the experimental section by:

• Displaying figures in the same order which the experiments are presented in the text
• Explaining the data in more detail, such as the neural data we used to demonstrate model efficacy in monkey reaching and monkey timing (DMFC-RSG) tasks
• Interweaving more explanation of the figures into the main text so that important details are not only contained in the captions.
• Moving some algorithmic discussion to the appendix and, in turn, fleshing out the significance of the experimental results and giving more exposition.

Is the format of the citations as a superscript an acceptable one for neurips? It's not very common

Thank you for asking – we did consult the guidelines beforehand.

We have to again thank the reviewer for motivating us to move some algorithmic discussion to the appendix in exchange for space that can be spent on clarifying the practical implications demonstrated in the experimental results section. We hope if we’ve adequately addressed your concerns you might raise your score to reflect that. We are motivated to continue improving the manuscript, and appreciate all of the input you provided.

Thank you for taking the time to review the paper and your comments that will undoubtedly increase the quality of our manuscript.

The paper is dense and notation heavy at times.

We apologize that the paper could be dense at times. Motivated by this comment and others, in order to enhance the clarity of the manuscript's main take-aways, we will move non-essential algorithmic details to the appendix in favor of extending the discussions of the experimental results section.

It is unclear to me why the authors make a connection to causal amortized inference, and whether this is a natural connection for the scope of the paper.

Thank you for this question about causal amortized inference. One of the design choices we made, making it possible to perform causal inference (which we feel has important possible implications in real data-analysis where XFADS can be applied), was segmenting local/backward encoders – in that way filtering can be performed causally using the local encoder; because other sequential VAE models do not make that distinction, we find the ability to perform causal inference a feature of our model that distinguishes it from other approaches.

How easily extendable is the method to non-gaussian state-space models?

Thank you for the great question! The convenient part of working with general exponential family representations is that the inference algorithm we presented (sans the Gaussian specific parts) is agnostic to the choice of distribution so long as we have a way of evaluating $\lambda_{\theta}(z_{t-1})$ and $\mu_{\theta}(z_{t-1})$.

What do the learned low-rank covariance matrices look like?

Thank you for the question, we did not include these in the main paper due to space constraints, but we feel it would be beneficial to have some example learned covariances from an experiment. In the rebuttal PDF, we show an example trial from the bouncing ball experiment for the models that use nonlinear dynamics; its interesting to see how for example, contact with the wall results in spatially complex covariance structures in latent space that cannot be captured by the diagonal approximations.

How does amortization affect the quality of the approximate posterior? What if T is small?

In the case of smaller datasets, smaller T, or amortization networks that are not very expressive, the method could still suffer from overfitting/amortization gaps usually associated with VAEs.

In lines 89-92, the authors list the limitations of amortized inference for state space models. Are there limitations more specific to this particular method that the authors omit?

Thank you for pointing this out. We have added to the discussion section some comments on further limitations – “Furthermore, depending on generative model specifications, such as $L$, while the inference framework of Alg.1 is always applicable, modifications of the message passing procedure in Alg.2 might be necessary to maximize efficiency (e.g. if $L$ is small but $S$ is large)." In addition to that, in the related works section we will mention that there will be overhead incurred using low-rank approximations as compared to diagonal ones, but significant savings compared to methods like SVAE that use exact Gaussian message passing.

We appreciate the time and effort you put into your review. We hope that if you feel we have addressed the issues raised you might raise your score to reflect that.

Thank you for taking the time and effort and asking questions that we feel have helped make the manuscript stronger. Below we respond to them in order,

One suggestion is to keep only the core equations (such as the objective (21)(22), the low rank parameterization (24), etc.) and provide more high-level discussions on the intuition behind the model. The author could talk about the relationship between the pseudo-observations and the potentials in the SVAE paper, and give intuition on how to understand these quantities.

Thank you for this suggestion — we understand the paper is dense and building more intuition will benefit the reader. Motivated by your suggestion along with those of reviewer UuWx and others, we are moving some algorithmic details from section 3 to the Appendix in favor of higher level intuitions and clarity of the experimental results.

Some related work is missing. From my understanding, the authors use sampling-based linear Gaussian approximations for inference in the nonlinear dynamical system. This is closely related to the extended Kalman filter (EKF) which uses linear approximations to the dynamics based on the Jacobian, and the unscented Kalman filter (UKF) which uses anchor points instead of random samples to estimate the mean and covariance of the predicative distribution.

Thank you for bringing up other approximate filtering methods. This led to the discovery that the approximate predict step of Eq(23), when dynamics are nonlinear Gaussian, coincides with the statistically linearized Kalman filter [1] Thanks to your suggestion, we add an additional line after introducing the approximate filter to clarify this connection: “...otherwise, Monte Carlo integration can provide a differentiable approximation; for nonlinear and Gaussian dynamics this leads to a predict step equivalent to the statistically linearized Kalman filter [1].”

Somewhat limited experimental verification. While I do think that the existing results conveys the effectiveness of the proposed methods well, I do wish that more in-depth comparisons and discussions can be made. For example, how well do the low-rank update scale to higher-dimensional systems? Comparisons like this gives the reader a better sense of the tradeoffs of the method compared to others in the literature.

These are great suggestions and as per the earlier remark, we are moving non-essential algorithmic details into the appendix. We will be using the extra space to help build extra intuition – for example, elaborating on Fig.2(a,b) and emphasizing the favorable scaling illustrated in the figure. We have also run experiments using the latent SDE of [3] as an additional point of comparison(Fig.1 rebuttal PDF).

What are the practical computational efficiency of this method compared to the DKF or the SVAE in wall clock time? Are the computational overheads for the sampling-based inference significant in practice?

Thanks for this question -- it has motivated us to expand on this in the discussion section as well as add other methods wallclock time to Fig.1a for reference . To answer: compared to methods like DKF/DVBF that use diagonal approximation there is additional overhead since our method scales $\mathcal{O}(LSr)$ per step, however, there are significant savings compared to SVAE which scales $\mathcal{O}(L^3)$ per step due to its dense marginal covariances.

In figure 1, the DVBF and SVAE seems to be underperforming the DKF in the pendulum experiments, despite the pendulum motion closely resembling a linear dynamical system. Why is this the case?

It is true that the pendulum closely resembles an LDS, but in this case, the small angle approximation does not hold [2] which would make long-term forecasting difficult for a model using a linear dynamical system prior. This is one of the reasons that all methods perform very well in the smoothing metric but not prediction metric; especially SVAE where the linear dynamics approximation degrades as the horizon increases.

Can this method be efficiently parallelized?

Things can be easily parallelized across batches, but to parallelize across time would require a different inference scheme (since we sequentially propagate samples through the dynamics). However, one possibility is to use a parallelizable inference network architecture (such as S4 [4]), to produce marginals and train it as usual using Eq.(22)

How well do the low-rank update scale to higher-dimensional systems? I do wish that more in-depth comparisons and discussions can be made. For example, how well do the low-rank update scale to higher-dimensional systems? Comparisons like this gives the reader a better sense of the tradeoffs of the method compared to others in the literature

Thank you for these questions – we aimed to elucidate some aspects through conducting the experiments featured in the left half of figure 2. Motivated by your previous comment, we will now have more room to explain the significance of that experiment in showing: i) purposely exploiting covariance structures makes it possible to develop a filtering algorithm scaling linear in the latent dimension, L — compared to the O(L^3) cost per step incurred for typical Gaussian filters. and ii) how low-rank covariance parameterizations can also lead to a tighter ELBO.

We would like to thank the reviewer again for their time and useful suggestions that will serve to make our manuscript stronger. We hope that if we have addressed your concerns you might raise your score to reflect that, and we are eager to address any concerns or questions that might remain.

[1] Sarkka, Bayesian filtering and smoothing. [2] Zhao and Lidnerman, ICML 2023. Revisiting Structured Variational Autoencoders [3] Li et al, 2021. Scalable gradients for stochastic differential equations [4] Gu et al, 2021. Efficiently Modeling Long Sequences with Structured State Spaces