eXponential FAmily Dynamical Systems (XFADS): Large-scale nonlinear Gaussian state-space modeling
State-space graphical models and the variational autoencoder framework provide a principled apparatus for learning dynamical systems from data. State-of-the-art probabilistic approaches are often able to scale to large problems at the cost of flexibility of the variational posterior or expressivity of the dynamics model. However, those consolidations can be detrimental if the ultimate goal is to learn a generative model capable of explaining the spatiotemporal structure of the data and making accurate forecasts. We introduce a low-rank structured variational autoencoding framework for nonlinear Gaussian state-space graphical models capable of capturing dense covariance structures that are important for learning dynamical systems with predictive capabilities. Our inference algorithm exploits the covariance structures that arise naturally from sample based approximate Gaussian message passing and low-rank amortized posterior updates -- effectively performing approximate variational smoothing with time complexity scaling linearly in the state dimensionality. In comparisons with other deep state-space model architectures our approach consistently demonstrates the ability to learn a more predictive generative model. Furthermore, when applied to neural physiological recordings, our approach is able to learn a dynamical system capable of forecasting population spiking and behavioral correlates from a small portion of single trials.
Paper
Similar papers
Peer review
Summary
This work develops a deep state-space model (DSSM) for state inference in and system identification of non-linear dynamical systems. The objective function for learning the DSSM parameters is based on a smoothing formulation and thus the evidence lower bound (ELBO). The prior and likelihood parts are separately parametrized and joined in the natural parameter space, where the posterior natural parameters are simply the sum of prior and likelihood natural parameters. Furthermore, a low-rank approximation of the approximate likelihood (via encoder) are used to improve computational efficiency. Experiments are performed on synthetic data (pendulum and bouncing ball) as well as neuroscientific data (motor cortex readings of monkeys). The authors compare mostly against pretty old baselines from before 2016 (DVBF, DKF, SVAE).
Strengths
- The paper is very well written and quite easy to follow for the most part (see weaknesses). - Hyperparameter, notation and algorithm details as well as derivations are given in the supplementary material, which is very helpful.
Weaknesses
- 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. The main challenges of the VAE-based approaches is that they often only learn to auto-encode and do not properly identify the correct dynamics [1]. Indeed, tackling this problem seems to be a goal of the present work as well, but it is not argued in the introduction. - The related work section also seemed to stop past 2016. There are multiple subsequent works that address and improve upon the problem of learning dynamics from high-dimensional data, see e.g. [2-5], all of which share several ideas with this work. Soundness: I have difficulties with the technical soundness of the paper, especially about the following: - 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}. A correct smoothing-based objective involves pairs two-slice marginals, see e.g. [6]. Unfortunately, the paper is based on this equation and it does not seem to be just a typo as subsequent equation (e.g. (19)) have the same error. - What exactly is the pseudo-observation \tilde{y}? What values does it take compared to an actual y? Previous work [2,4,5] use the term pseudo-observation to refer to an additional latent variable that encodes information about a single time-step obervation y_t. This makes it a bit confusing. As far as I understood, the \tilde{y} is not ever actually instantiated and instead used as a proxy for a likelihood approximation. - line 138: I can not follow this argument. Why is smoothing turned into a filtering problem? Filtering and smoothing solve different problems. Do you instead mean that you have a backward-forward smoothing approach? If so, this needs to be explained and defined better. I am having trouble with the equations (14), (15) and several subsequent equations involving \tilde{y}, again mostly because I do not know what exactly these pseudo-observations are. - It becomes more confusing that the introduction of a different distribution pi(z_t) is supposed to help. Please, can you elaborate on this part? How does the introduction of this distribution in the filtering formulation yield a smoothing/ELBO-based objective? Novelty: - It is not clear how this work compares to some of the more recent works that try to resolve issues with works prior to 2016. I referenced below already some closely related works that I am aware of, but there are several others missing in related work. It is not clear how this work is positioned in comparison to those. - eq (11): using the sum of the two terms (prior and likelihood) is not novel, see e.g. [4], section 3.3, where a Gaussian likelihood approximation is combined with the Gaussian prior dynamics, leading to a sum of corresponding natural parameters. Experiments: - The experimental evaluation is rather weak. There are no proper benchmarks against more recent works. The qualtative results visualized in Figure 3 do not look that great and it is difficult to estimate how hard it is to predict this kind of data. [1] Karl et al. 2016, Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data. [2] Fraccaro et al. 2017, A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning. [3] Becker-Ehmck et al. 2019, Switching Linear Dynamics for Variational Bayes Filtering. [4] Kurle et al. 2020, Deep Rao-Blackwellised Particle Filters for Time Series Forecasting. [5] Klushyn et al. 2021, Latent Matters: Learning Deep State-Space Models. [6] Simo Sarkka, Bayesian filtering and smoothing. ------- Minors and suggestions that did not influence the review score - The complexity statement in the introduction does not explain what the variables T, L, S, and R are. In my opinion, it would be better to describe it in words in the introduction with an emphasis on how it compares to other approaches. - line 45: filtering is an even more common approximation than smoothing. Filtering, smoothing and even fixed-lag smoothing are all valid applications, corresponding to whether inference needs to be online, offline or semi-online. - line 48: vEM needs a citation. For instance, VEM is mentioned in the dissertation of Beal, but there is likely a better reference therein. - line 51: to make the difference to the VAE approach more clear, you might consider emphasizing that the alternative is a *coordinate ascent* approach with *alternating* optimization steps. - line 60: which favorable properties are the ones relevant to this work? I suppose the additive property or anything else? It should be made clear what exactly you refer to. - line 68) typo: ii) -> iii) - Equation (3) typo: parameter - line 78: citation style changes - dKF, dVBF, vEM all written with lower case first letter, but other works (including the original ones) use capitalization.
Questions
- what is the difference between predictions from filtering and smoothing in Fig. 3? The last filter state should be identical to the last smoothing state.
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
Limitations have been addressed to some extent.
Summary
Update post author rebuttal. Thanks for the clarifications in the common and personal replies. I will raise the score to Accept, trusting that you will make the improvements you mention. ________ The paper presents a class of non linear state-space models whose dynamics are defined as exponential family distributions. The variational approximation leverages distributions in the same exponential family as the prior, and its parameterization is defined as the sum of prior parameters with a learnable data dependent term. Efficient filtering/smoothing is obtained with an approximate message passing algorithm that exploits the low rank structure of the variational distribution. The model is tested on a number of simulated and real world applications and shows good predictive performances.
Strengths
* The paper introduces a theoretically principled model that is quite flexible and can be used in a wide range of applications. As such I believe it can have impact in the neurips community (as long as the code is released) * The model combines ideas from previous work, but is overall novel to the best of my knowledge * The model can handle well missing data * The choice of inference network allows scalable inference * The theoretical section is dense but well explained (unlike the experimental one as noted below) * The predictive performances are better than similar SOTA models
Weaknesses
**Main comment** I found the experiment section to be too rushed and therefore hard to understand. The authors should improve its clarity (in case you need space to fit in the 9 pages, you can move some of the the theoretical results in section 4 to the appendix). There are several issues that taken together make the experimental quite hard to follow: 1. In figures 1 and 2 your method is presented as "ours (n)" and "ours (c)", while figure 3 uses "ours (a)" and "ours (c)" . These a, c and n versions are however not defined anywhere, so I'm not sure what I am looking at. 2. Overall the way the figures are combined is messy. Not all of them are mentioned and described in the main text, where they appear in the following quite random order: 2b, 2a, 1, 3c, 3b, 2d. Especially considering the fact that you leave a lot of the explanation to the caption of the figures, it's very hard for the reader to understand which parts of the figures are relevant to look at while reading the different parts of the experimental section. I suggest to rethink the structure of the figures, possibly making a single figure per experiment and making sure they are sufficiently discussed in the main text. 3. Figure 2c mentions the "DMFC RGS" dataset which is not defined in section 5, so not sure what experiment that refers to. Figure 2c is also never mentioned in the main text. 4. In the caption of figure 2a "(top) for fixed S and (bottom) for fixed r" seems wrong. **Other minor points** * Line 36 in the introduction: the complexity term used terms that are not yet defined * The method name (XFADS) only appears in the title and the discussion, but it is never introduced in the Method section, or used in the experimental section. Either you use it, or you can avoid introducing it * Line 145 - you reference the non existing equation (94) * Line 330, missing "the" in "if the number of samples" * Is the format of the citations as a superscript an acceptable one for neurips? It's not very common
Questions
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
Rating
7
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
Yes
Summary
This paper introduces a method for scalable nonlinear Gaussian state-space modeling that relies on variational autoencoders and a low-rank covariance matrix assumption for efficient inference by optimizing an approximate variational lower bound. The authors describe the computational benefits of their method and inference scheme and showcase its effectiveness on multiple real-world applications, drawing connections to neuroscience.
Strengths
The paper reads clearly and is well-organized. The authors produce convincing experiments, and a novel inference method that scales linearly with the dimensionality of the state space. The authors approximate the filtering distribution with a differentiable approximation for gradient-based updates. The appendix is concise, with relevant and clear background information. The authors clearly motivate their work using real-world applications and a thorough literature review.
Weaknesses
No major complaints. The paper is dense and notation heavy at times. 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.
Questions
How easily extendable is the method to non-gaussian state-space models? More discussion on general methodology? What do the learned low-rank covariance matrices look like? How does amortization affect the quality of the approximate posterior? What if T is small?
Rating
6
Confidence
3
Soundness
4
Presentation
4
Contribution
3
Limitations
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?
Summary
The paper presents a novel state space model (SSM) framework for learning dynamical systems with nonlinear latent dynamics. The proposed method borrows inspirations from structured variational autoencoders (SVAEs) and sample-based nonlinear Bayesian filtering, using low-rank corrections to the prior to capture information from the observations and compute approximations to the variational posterior. This framework addresses limitations in current methods which fall short on model expressivity or predictive power.
Strengths
This paper builds upon previous work very well while addressing their weaknesses adequately. The authors show a high level of technical understanding regarding Bayesian filtering and variational inference, and from what I have read the methods proposed are technically sound (I did not go over all of the math in the appendix). As far as I’m aware, the proposal of a structured inference framework for exponential family SSMs that allows for learned non-linear dynamics is a good contribution that fills a gap in the literature. The paper is also professionally written and contains few mistakes or typos. Overall a good paper with solid contributions.
Weaknesses
- Main technical ideas are conveyed in a convoluted way. The main paper is unnecessarily dense, making it hard to read and understand. 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. For example, 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. Technical details such as the sample-based approximate inference and the corresponding low-rank updates are pretty standard in my opinion and should be left in the appendix. - 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. There have also been attempts to combine neural networks with these nonlinear Bayes filtering methods such as (Liu et al., 2024). - 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. Overall, I consider these points to be relatively minor and do not detract from the paper's contributions. Liu, Wei, et al. "Neural extended Kalman filters for learning and predicting dynamics of structural systems." *Structural Health Monitoring* 23.2 (2024): 1037-1052.
Questions
- Can this method be efficiently parallelized? - 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? - 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? - How well do the low-rank update scale to higher-dimensional systems?
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
4
Limitations
The authors have adequately addressed the limitations of their work.
Please respond to author comments
Dear Reviewers, Please read the authors’ rebuttal and respond to their comments. The discussion period ends on Aug 13 and I would like to have time for back and forth, as necessary.
Score raised
Thanks for the clarifications in the common and personal replies. I will raise the score to Accept, trusting that you will make the improvements you mention.
Thank you very much for your thoughtful consideration and for raising the score to Accept. We are grateful for your positive feedback and are committed to making the improvements you mentioned. Your insights have been incredibly valuable in refining our work, and we will ensure that the final version addresses all points raised in the review.
I thank the authors for their detailed, thorough rebuttal, and address some of their comments below. > 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. Thank you for the clarification. > 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})$. Thank you for the clarification. If I understand correctly, fast inference relies on the reparameterization trick for approximating the reconstruction term in the evidence lower bound. For exponential family distributions that are not reparametrizable, what variance reduction techniques do you recommend? > 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. Very interesting. I request the authors include at least one example figure in the appendix (no need to compare to other methods here) so the readers may observe the low rank, non-diagonal covariance structures. > 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. Thank you for the additional comments regarding the limitations of the work.
We thank the reviewer for their prompt response and helpful suggestions. We address your comments below. > Thank you for the clarification. If I understand correctly, fast inference relies on the reparameterization trick for approximating the reconstruction term in the evidence lower bound. For exponential family distributions that are not reparametrizable, what variance reduction techniques do you recommend? Thank you for the question. You are right that fast inference relies on the reparameterization trick. In cases where the standard reparameterization trick cannot be applied, an alternative (equivalent to extended kalman filter in the Gaussian case) would be to linearize the dynamics function in mean parameter space, $\mu\_{\theta}(z_{t-1})$, about the sufficient statistics $\mathcal{T}(z_{t-1})$ so that $\mu\_{\theta}(z_{t-1}) \approx F \mathcal{T}(z_{t-1}) + f$, letting us evaluate the predict step equation $\mathbb{E}[\mu\_{\theta}(z_{t-1})] \approx \mathbb{E}[F \mathcal{T}(z_{t-1}) + f] = F \mu\_{t-1} + f$. Alternatively, it would also be viable to use implicit reparameterization gradients[1] when evaluating the ELBO. Motivated by this question, we will include additional discussion related to these alternatives in the appendix. [1] Figurnov et al. 2018. Implicit reparameterization gradients > Very interesting. I request the authors include at least one example figure in the appendix (no need to compare to other methods here) so the readers may observe the low rank, non-diagonal covariance structures. Thank you for the original suggestion — we will make sure to include in the appendix of the manuscript the bouncing ball covariances from the rebuttal PDF as well as visualizations of the latent state covariance for some of the other experiments. Thank you again for your thoughtful questions and constructive feedback. We appreciate your time and effort in reviewing our response and believe the revisions will strengthen the manuscript.
Response to rebuttal
Thank you for your reply. > 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. I agree with this and the favorable scaling is indeed a great contribution. However, I still think that the introduction does not do a great job in introducing the problems of this research area, how this work fits in this area and how it addresses short-comings of the previous works. What is missing is how previous works such as DVBF and DKF are insufficient (due to the diagonal approximation), that subsequent work such as [2,4,5] addresses this issue via conditional linearity and closed-form Kalman filtering/smoothing, how that subsequent work is insufficient, and how this work addresses this limitation. > [...] the expectations for Eq(6) do contain the two slice marginals. [...] You are right, the second expectation I thought was missing is of course contained in the KL. So I agree that Eq. (6) is indeed correct. > The pseudo observation encodes current/future observations into a Gaussian potential. Then I do not understand Eq. (12). Surely, you can approximate the likelihood of the actual data into a Gaussian potential. This is also what is used in the two-filter smoothing formula for linear dynamical systems: you collect the current and future data into a Gaussian potential in the form of the backwards-message, typically denoted by \beta = p(z_t | y_{t:T}). It is also possible to define the smoothing distribution via a backward-forward recursion, using that same \beta. The \beta message will be represented using the natural parameters of this Gaussian potential. Is this what you are doing here? What I don't understand is why you need any pseudo-targets for this and whether these pseudo-targets that take concrete values actually exist. Why not just write \beta or something like g(z_t; y_{t:T}). Or am I missing something? Are these pseudo-targets used for anything? To underscore my point, note also that the RHS of Eq. (12) does not contain any \tilde{y}. > we say smoothing is turned into a filtering problem because pseudo observations \tilde{y}_t, encode y_{1:t} into a single Gaussian potential, meaning we can filter pseudo observations to obtain statistics of the smoothed posterior; As mentioned above, this looks to me like you are doing a backward-forward algorithm for smoothing. But you are right then with the formulation that in this case smoothing is turned into a filtering problem, however I still find pseudo-targets confusing here and I think it is easier to work with the Gaussian potentials or backward messages directly. > Thank you for pointing this out, but we respectfully disagree that this particular parameterization is not novel. [...] 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. Agreed.
We thank the reviewer for their comments and respond below. > However, I still think that the introduction does not do a great job in introducing the problems of this research area, how this work fits in this area and how it addresses short-comings of the previous works. What is missing is how previous works such as DVBF and DKF are insufficient (due to the diagonal approximation), that subsequent work such as [2,4,5] addresses this issue via conditional linearity and closed-form Kalman filtering/smoothing, how that subsequent work is insufficient, and how this work addresses this limitation. We generally agree with the reviewer that some of the relevant advances need to be contrasted to form a better bigger picture for the reader. There are indeed many methods we have not included due to space. The switching LDS style systems such as [4] (or similarly rSLDS and its extensions) are good approximate representations of nonlinear state-space models, though we argue that it is not the naturally interpretable one when continuous state space is assumed (as in our neuroscience experiments). Since [4] is a particle filter like solution, it potentially has a tighter bound like [Zhao et al. 2022] which is another variational filter with constant time complexity. While [5] shares some goals of our work, similar to [4] it faces scalability issue in L. As the reviewer points out, our contribution on the scalability front is solid; for example in [5], the experiments were very small with L=3, 4, 5, and T = 15, 30, while we have L = 40 and T = 35, 130. Moreover, in contrast to previous works, we demonstrate the efficacy of our approach for causal filtering, which is immensely important for conducting causal investigations in neuroscientific settings. To the best of our knowledge, this is the only approach that can do so in large L scenarios. Unfortunately neither [4] nor [5] have publicly available code – we will contact the authors so that we can make proper comparisons; our code is already publicly available, and if accepted, the camera ready will have a link to our codebase to support future investigations and reproducibility. We would like to re-emphasize that our variational inference network structure and the associated novel ELBO (Eq.22) are non-trivial contributions that work together to enable the scalability, predictive performance, principled masking, and the causal real-time filtering. While we agree that SVAE and [2] both can have non-diagonal covariance, they are both based on a linear dynamical system (LDS) at the core – limiting their expressive power substantially; neuroscience data similar to those examined in the manuscript often exhibit dynamics with topological features LDS cannot capture (e.g. multiple fixed points). We will ensure that the related works section of the revised manuscript clearly articulates how these additional works fit into the bigger picture and also disambiguate their differences as you suggest. > Then I do not understand Eq. (12). Surely, you can approximate the likelihood of the actual data into a Gaussian potential. This is also what is used in the two-filter smoothing formula for linear dynamical systems: you collect the current and future data into a Gaussian potential in the form of the backwards-message, typically denoted by \beta = p(z_t | y_{t:T}). It is also possible to define the smoothing distribution via a backward-forward recursion, using that same \beta. The \beta message will be represented using the natural parameters of this Gaussian potential. Is this what you are doing here (yes)? What I don't understand is why you need any pseudo-targets for this and whether these pseudo-targets that take concrete values actually exist. Why not just write \beta or something like g(z_t; y_{t:T}). Or am I missing something? Are these pseudo-targets used for anything? To underscore my point, note also that the RHS of Eq. (12) does not contain any \tilde{y}. We believe we are on the same page and apologize for the confusion. The pseudo-observation $\tilde{y}\_t$ is indeed not instantiated, and is a representation for the natural parameters of the Gaussian potential given by $p(z_t | y_{t:T})$ which we call $\tilde{\lambda}\_\phi(y_{t:T})$. The $\beta\_{t+1}$ and $\alpha\_t$, parameterizing natural parameter updates as in Eq(24), when additively combined form this backward message as described in Eq (23). We realize this notational overhead might cause confusion and if accepted, we will make sure to reduce the notational overhead in the camera ready version. We hope this clears up any confusion. Thank you again for your suggestions and comments; we greatly appreciate the time you have taken to help improve our manuscript.
[Zhao et al. 2022]
We forgot to include the full citation of the relevant paper we mentioned in the previous response: - Zhao, Y., Nassar, J., Jordan, I., Bugallo, M., & Park, I. M. (2022). Streaming Variational Monte Carlo. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(1), 1150–1161. https://doi.org/10.1109/TPAMI.2022.3153225
Last day for discussion
Reviewer e1Je, today is the last day for discussion. I hope you will respond to the authors' rebuttal.
I would like to thank the authors for their thorough response. I am glad that the authors find my suggestions on the paper organization and related works helpful. Trusting that the authors will make the promised changes and provide more discussion on the benefit and scaling properties of the low-rank approximation, I have raised my score to a 7.
Decision
Accept (poster)