Title: eXponential FAmily Dynamical Systems (XFADS): Large-scale nonlinear Gaussian state-space modeling
Venue: NeurIPS 2024
Accepted: yes
Keywords: variational inference

nonlinear state-space model

dynamical system
Abstract: 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.

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

Abstract

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.

1 Introduction

State-space models (SSM) are invaluable for understanding the temporal structure of complex natural phenomena through their underlying dynamics 1-3 . While engineering or physics problems often assume the dynamical laws of the system of interest are known to a high degree of accuracy, in an unsupervised data-driven investigation, they have to be learned from the observed data. The variational autoencoder (VAE) framework makes it possible to jointly learn the parameters of the state-space description and an inference network to amortize posterior computation of the unknown latent state 4-7 . However, it can be challenging to structure the variational approximation and design an inference network that permits fast evaluation of the loss function (evidence lower bound or ELBO) while preserving the temporal structure of the posterior.

In this work, we develop a structured variational approximation, approximate ELBO, and inference network architecture for generative models specified by nonlinear dynamical systems with Gaussian state noise. Our main contributions are as follows,

(i) A structured amortized variational approximation that combines the prior dynamics with low-rank data updates to parameterize Gaussian distributions with dense covariance matrices,

(ii) Conceptualizing the approximate smoothing problem as an approximate filtering problem for pseudo-observations that encode a representation of current and future data, and,

(iii) An inference algorithm that scales O ( TL ( Sr + S 2 + r 2 )) - made possible by exploiting the low-rank structure of the amortization network as well as Monte Carlo integration of the latent state through the dynamics.

2 Background

State-space models are probabilistic graphical models where observations y t in discrete time are conditionally independent given a continuous latent state, z t , evolving according to Markovian dynamics, so that the complete data likelihood for T consecutive observations factorizes as,

where z t ∈ R L are real-valued latent states, θ parameterizes the dynamics and initial condition, and ψ parameterizes the observation model. When the generative model, ( θ , ψ ) , is known, the statistical inference problem is to compute the smoothing posterior, p ( z 1: T | y 1: T ) 2 . Otherwise, ( θ , ψ ) have to be learned from the data - known as system identification 8 .

Variational inference makes it possible to accomplish these goals in a harmonious way. The variational expectation maximization (vEM) algorithm iterates two steps: first, we maximize a lower bound to the log-marginal likelihood - the ELBO - with respect to the parameters of an approximate posterior, q ( z 1: T ) ≈ p ( z 1: T | y 1: T ) ; then, with the approximate posterior fixed, the ELBO is maximized with respect to parameters of the generative model 9 . For large scale problems, vEM can be slow due to the need to fully optimize the variational parameters before taking gradient steps on parameters of the generative model. Therefore, the variational autoencoder (VAE) is better suited for large scale problems for its ability to simultaneous learn the generative model and inference network - an expressive parametric function that maps data to the parameters of approximate posterior 10,11 .

Model specifications. Although our approach is applicable to any exponential family state-space process, given their ubiquity, we focus on dynamical systems driven by Gaussian noise so that,

where m θ : R L → R L might be a nonlinear neural network function with learnable parameters θ , and Q θ ∈ R L × L is a learnable state-noise covariance matrix. Given the favorable properties of exponential family distributions 12-15 , especially in the context of variational inference, we write the prior dynamics in their exponential family representation (natural parameter form),

where h is the base measure, T ( z t ) the sufficient statistics, A ( · ) the log-partition function, and λ θ ( · ) is a map R L ↦→ R L 2 + L that transforms z t -1 to natural parameters for z t . For a Gaussian distribution, the sufficient statistics can be defined as T ( z t ) glyph[latticetop] = [ z glyph[latticetop] t -1 2 z t z glyph[latticetop] t ] , so that λ θ ( · ) for (1) is given by,

As it will simplify subsequent analysis, the mean parameter mapping corresponding to this natural parameter mapping (guaranteed to exist as long as the exponential family is minimal 12 ) is given by

Furthermore, we make the following assumptions i) the state-noise, Q θ , is diagonal or structured for efficient matrix-vector multiplications. ii) m θ ( · ) , is a nonlinear smooth function. ii) the likelihood, p ψ ( y t | z t ) , may be non-conjugate. iv) L may be large enough so that L 3 is comparable to T .

Amortized inference for state-space models. A useful property of SSMs is that, z t conditioned on z t -1 and y t : T , is independent of y 1: t -1 , i.e., p ( z t | z t -1 , y 1: T ) = p ( z t | z t -1 , y t : T ) 5,16 . It thus suffices to construct an approximate posterior that factorizes forward in time,

and introduce learnable function approximators to amortize inference by mapping z t -1 and y t : T to the parameters of q ( z t | z t -1 ) . This makes it simple to sample z 1: T from the approximate posterior (using the reparameterization trick) and evaluate the ELBO (a.k.a. negative variational free energy),

where D KL ( ·|| · ) is the Kullback-Leibler (KL) divergence and E q t ≡ E z t ∼ q ( z t ; y 1: T ) , so that the generative model and inference network parameters can be learned through stochastic backpropagation. Many works for Gaussian q ( z t | z t -1 ) , such as Krishnan et al. 6 , Alaa and van der Schaar 17 , Girin et al. 18 , Hafner et al. 19 , construct inference networks that parameterize the variational posterior as

There are limitless ways to construct m φ ( · ) and P φ ( · ) so φ can be learned through gradient ascent on the ELBO, but a straightforward and illustrative approach 6,17 is to transform future data, y t : T , using a recurrent neural network (RNN), or any efficient autoregressive sequence to sequence model, and then mapping the preceding latent state, z t -1 , using a feed-forward neural network, so that a complete inference network description could be,

where S2S ( · ) is a parametric sequence-to-sequence function that maintains a hidden state u t and takes as input y t , and NN ( · ) is a parametric function designed to output approximate posterior parameters. This leads to a backward-forward algorithm, meaning that data y 1: T are mapped to u 1: T in reverse time, and then samples are drawn from q ( z t | z t -1 ) forward in time.

Possible drawbacks of this inference framework are i) missing observations obstruct inference (the example networks cannot naturally accommodate missing data); ii) sampling entire trajectories to approximate the expected KL term can potentially lead to high-variance gradient estimators, and iii) statistics of the marginals (e.g. second moments) can only be approximated through sample averages.

3 Related works

Many existing works also explore inference and data-driven learning for state-space graphical models within the VAE framework. We highlight the most closely related studies and note specific limitations that our work seeks to address. The structured variational autoencoder (SV AE) 20 makes it possible to efficiently evaluate the ELBO while preserving the temporal structure of the posterior by restricting the prior to a linear dynamical system (LDS) and then constructing the approximation as q ( z 1: T ) ∝ p θ ( z 1: T ) ∏ exp( t ( z t ) glyph[latticetop] ψ ( y t )) so that its statistics can be obtained using efficient message passing algorithms. However, the SVAE is not directly applicable when the dynamics are nonlinear since the joint prior will no longer be Gaussian (thereby not allowing for efficient conjugate updates). Recently, Zhao and Linderman 21 expanded on the SVAE framework by exploiting the LDS structure and associative scan operations to improve its scalability.

The deep Kalman filter (DKF) 6 uses black-box inference networks to make drawing joint samples from the full posterior simple. However, pure black-box amortization networks such as those can make learning the parameters of the generative model dynamics difficult because their gradients will not propagate through the expected log-likelihood term 5 . In contrast, we consider inference networks inspired by the fundamental importance of the prior for evaluating Bayesian conjugate updates. The deep variational Bayes filter (DVBF) also considers inference and learning in state-space graphical models 5 . Difficulties of learning the generative model that arise as a result of more standard V AE implementations defining inference networks independent of the prior are handled by forcing samples from the approximate posterior to traverse through the dynamics. Our work extends this concept, by directly specifying the parameters of the variational approximation in terms of the prior.

Our approach constructs an inference network infused with the prior similar to the SV AE and DVBF but i) avoids restrictions to LDS and ii) affords easy access to approximations of the marginal statistics (such as the dense latent state covariances) without having to average over sampled trajectories (or store them directly which would be prohibitive as the latent dimensionality becomes large).

4 Method

An alternative to constructing variational approximations through specification of conditional distributions, as in Eq. (7), involves the use of data-dependent Gaussian potentials, that we refer to as pseudo-observations:

These Gaussian potentials can then be combined with the prior through Bayes' rule, yielding the approximation

A benefit of this formulation, is that it inherently imposes the latent dependency structure of the generative model onto the amortized posterior. This parameterization, was introduced in Johnson et al. 20 , where an important point highlighted, is that the Gaussian potentials can encode any arbitrary subset of observations; for example, p (˜ y t | z t ) could be made to depend on y t alone, or even the entire dataset, y 1: T . Regardless of that particular design choice, the corresponding ELBO for the variational approximation of Eq. (10) is

For linear Gaussian latent dynamics, conjugate potentials could be efficiently integrated with the prior using exact message passing, yielding filtered and smoothed marginal statistics. The smoothed statistics can be used to evaluate the first two terms on the right-hand side, while the filtered statistics can be used to evaluate the final term, the log-marginal likelihood of the pseudo-observations.

However, this approach does not directly apply to nonlinear dynamical systems, where the variational posterior is no longer Gaussian. Since evaluating the smoothed marginals and the log-marginal likelihood of pseudo-observations relies on first obtaining the filtered marginals, a logical starting point is to develop a method for approximating these filtered marginals. To this end, we propose a differentiable approximate message passing algorithm specifically designed to compute filtered posterior statistics in models characterized by nonlinear latent dynamics and observations represented as Gaussian potentials. Building on this foundation, we then return to the subsequent challenges of efficiently computing smoothed posterior statistics and evaluating the ELBO.

Differentiable nonlinear filtering. Bayesian filtering is often conceptualized as a two step procedure 2 . In the predict step, our belief of the latent state is integrated through the dynamics, yielding q ( z t | ˜ y 1: t -1 ) = ∫ p θ ( z t | z t -1 ) q ( z t -1 | ˜ y 1: t -1 ) d z t -1 (a.k.a. the predictive distribution). Then, applying Bayes' rule, q ( z t | ˜ y 1: t ) ∝ p (˜ y t | z t ) q ( z t | ˜ y 1: t -1 ) , we update our belief. Evidently then, developing an approximate filtering algorithm that exploits the conjugacy of the pseudo observations can be recast as the problem: given an approximation π ( z t -1 ) ≈ q ( z t -1 | ˜ y 1: t -1 ) , find an approximation to the predictive distribution, ¯ π ( z t ) ≈ q ( z t | ˜ y 1: t -1 ) . The recursion would continue forward by updating our belief analytically, setting π ( z t ) ∝ p (˜ y t | z t )¯ π ( z t ) , then finding a Gaussian approximation of ¯ π ( z t +1 ) and so forth.

With the problem restated this way, we propose an approximate filtering solution designed to exploit two key factors at play i) the approximate beliefs are constrained to the same exponential family as the latent state transitions ii) the pseudo observations are encoded as conjugate potentials. Our approach involves recursively solving intermediary variational problems (their fixed point solutions on the right),

Steps (i) and (ii) can be thought of as variational analogues of the predict/update steps of Bayesian filtering, and importantly, finding their fixed point solutions does not require an iterative procedure because of our problem specifications. Reassuringly, iterating (i) and (ii) in the case of an LDS generative model would exactly recover the information form Kalman filtering equations. In the case of nonlinear dynamical systems, directly taking the expectation in (i) is intractable. We can overcome

this by employing the reparameterization trick to obtain a differentiable sample approximation of the parameters, ¯ µ t , of the fixed point solution. Naturally now, the statistics of the approximate filtered beliefs can be used to approximate the log-marginal likelihood of the pseudo observations as,

where the last integral can be evaluated analytically as a result of the Gaussian forms of the approximations and pseudo observations. While formulating step (ii) as a variational problem may appear superfluous given the conjugate structure, it hints at the possibility of approximating smoothed posterior marginal statistics within a variational framework. However, pursuing this idea further reveals significant challenges. Trying to develop a backward recursion for the distribution q t minimizing D KL ( E q t +1 [ q t | t +1 ] ∣ ∣ ∣ ∣ q t ) , by using the forward KL divergence (as in step (i)), leads to an intractable problem because the backward Markov transitions, q t | t +1 , are not conditionally Gaussian. Conversely, a fixed point solution of the reverse KL objective (as in step(ii)), D KL ( q t ∣ ∣ ∣ ∣ E q t +1 [ q t | t +1 ] ) , necessitates an iterative procedure, which can be computationally expensive.

Smoothing as filtering. In light of these difficulties, we offer a simple solution that exploits the flexibility in choosing the pseudo observation data dependence: define the parameters, k t and K t , of each pseudo observation, ˜ y t , to be a function of current and future data, y t : T , so that,

With this choice, filtered statistics of the latent state-relative to the pseudo-observations-can be used to approximate posterior smoothed marginals, i.e. π ( z t ) ≈ q ( z t | ˜ y 1: t ) ≈ p ( z t | y 1: T ) . This solution circumvents the challenges associated with backward message computation and only requires a single pass through the pseudo observations to obtain approximate smoothed posterior statistics. Substituting, π t , as an approximation to q t , in Eq. (11), leads to the following approximation of the ELBO.

Variational smoothing ELBO

By expressing the approximate ELBO compactly as Eq.(17), we highlight that it promotes learning models where the posterior at time t aligns closely with the one-step posterior predictive at that time (which depends on the generative model and the posterior at time t -1 ). The KL term in Eq.(17) can be evaluated in closed form, while the expected log-likelihood term can be approximated using the reparameterization trick 4 . However, a point of practical importance should be raised now: every filtering step and evaluation of the KL term has a time complexity of O ( L 3 ) , which may become a bottleneck for large L . In the following discussion, we will explore effective strategies to parameterize the Gaussian potential inference network that produces ˜ λ 1: T , to reduce the computational burden that filtering and evaluating ˆ L ( π ) pose in the large L regime.

Local and backward encoders. For state-space models, inferences about the latent state should be possible even with missing observations. To enable the amortized inference network to process missing observations in a principled way, we decompose the natural parameter update into two additive components: i) a local encoder, α φ ( · ) , for current observation, and ii) a backward encoder, β φ ( · ) , for future observations, i.e.,

Furthermore, by building the dependence of β φ ( · ) on y t +1: T through their representation as α t +1: T , so that β φ ( y t +1: T ) = β φ ( α t +1: T ) , a missing observation at time t is handled by setting α t = 0 . While a data dependent natural parameter update of 0 faithfully represents a missing observation - in the absence of data, the prior should not be updated - alternatively setting y t = 0 would introduce a harmful inductive bias into the inference network, since an observation of 0 can be arbitrarily informative. Given the impracticality of O ( TL 2 ) memory requirements, it is appealing to consider a low-rank parameterization for the local and backward encoders - we consider

true

ours (n)

ours (c)

DKF

DVBF

SVAE

pendulum

a)

b)

0

.4

.8

prediction

starts

smoothed

predicted

time

0

.4

.8

angle

position

0

25

0

25

where A t ∈ R L × r α with B t ∈ R L × r β parameterize low-rank local/backward precision updates, and a t ∈ R L with b t ∈ R L parameterize local/backward precision-scaled mean updates. Using these descriptions and the additive decomposition (18), the parameters of a single pseudo observation are,

where K ∈ R L × r if r = r α + r β and k ∈ R L . The low-rank structure of the natural parameter updates will be a key component to develop an efficient approximate message passing algorithm for obtaining sufficient statistics of the approximate posterior and evaluating the ELBO. Analogous to the inference network description (8), a differentiable architecture producing α 1: T and β 1: T could be,

which overall defines the map y 1: T ↦→ ( α 1: T , β 1: T ) . In addition, the separation of local and backward encoders can reduce the complexity of the backward encoder for L < N . Those familiar with sequential Monte-Carlo (SMC) methods 22 can view the backward encoder similar to twisting functions used to combine future information with filtered state beliefs to produce smoothing approximations 23,24 .

Exploiting structure for efficient filtering. A benefit of using the forward KL to design a variational analogue to the exact Bayesian predict step is immediate access to the fixed point solution. While nonlinear specification of the latent dynamical system make the expectation of Eq. (12) intractable, using the reparameterization trick with z s t -1 ∼ π ( z t -1 ) , gives the approximation,

where S is the total number of samples. Converting this finite sample estimate from mean parameter coordinates to a mean/covariance representation, we get that,

where ¯ M c t is the L × S matrix of samples passed through the dynamics function, then centered by the mean, defined for convenience as,

Writing the covariance estimate as it is in Eq. (23), reveals that it can alternatively be represented by the pair ( ¯ M c t , Q θ ) . In the regime where L > S , significant computational savings can be afforded by capitalizing on the low-rank structure of the covariance as estimated via the reparameterization trick. This structure can be exploited for efficient linear algebraic operations involving ¯ P t (and its inverse, after application of the Woodbury identity). Consequently, the natural parameters of π t , after updating ¯ π t , are

2.0

2.5

1.6

2.0

2

5

15

25

5

15

25

50

0

40

40

0

10

0

10

0

b)

a)

c)

fixed

fixed

10

75

25

=100

50

10

40

25

5

10

100

wall-clock time (s)

.1

.2

.10

.12

d)

DMFC RSG

BPS

0

.2

.4

ours (n)

ours (c)

DKF

DVBF

SVAE

smoothed

predicted

S/H/3

L/E/1

L/E/2

S/H/4

S/E/5

S/H/4

5 Hz

time

1.3s

prediction

starts

1.5

1.8

10 Hz

10 Hz

diagonal

samples

data

=50

Kalman

and also admit structured representations. Without both the sample approximation structure (during the variational predict step) and low-rank parameterization (during the variational update step), the cost of approximate filtering each time-step would be dominated by an O ( L 3 ) cost. Instead, recognizing the potential computational advantages of exploiting these structures, and never instantiating the predictive/updated covariance and precision matrices, makes it possible to develop an approximate filtering algorithm, where in the case L is significantly larger than S or r , has complexity of O ( L ( Sr + r 2 + S 2 )) per step. More details regarding time complexity are in App. B.5.

Efficient sampling and ELBO evaluation. When E π t [log p ( y t | z t )] can not be evaluated in closed form, Monte-Carlo integration can be used as a differentiable approximation. To sample from π ( z t ) without explicitly constructing P t , we can take ¯ z s t ∼ N ( 0 , ¯ P t ) and w s t ∼ N ( 0 , I L + S ) and set,

While more details are provided in App. B.4, this can be done efficiently since samples can be drawn cheaply from ¯ π ( z t ) using Eq. (23). Whereas Monte-Carlo approximations of the expected log-likelihood term might be unavoidable, the closed form solution for the KL between two Gaussian distributions should be used to avoid further stochastic approximations. The only difficulty, is that the time complexity of naively evaluating the KL term,

scales O ( L 3 ) . However, since matrix vector multiplies with ¯ P -1 t can be performed efficiently and the trace/log-determinant terms can be rewritten using the square-root factors acquired during the forward pass, as we describe in App. C.1, it is possible to evaluate the KL in O ( LSr + LS 2 + Lr 2 ) time. After a complete forward pass through the encoded data, we acquire the samples z 1: S 1: T and all necessary quantities for efficient ELBO evaluation. We detail the variational filtering algorithm in Alg. 2 in App. C.2 and the complete end-to-end learning procedure in Alg. 1.

Causal amortized inference for streaming data. In constructing a fully differentiable variational approximation, the parameters of the approximate marginals were effectively amortized according to a recursion in the natural parameter space by iterating Eqs. (12) and (13). This recursion can be recognized more easily by introducing the function, F θ ( · ) , and writing

Algorithm 1 End-to-end learning

∇ A ∗ : ( m , -1 2 ( P + mm glyph[latticetop] )) ↦→ ( P -1 m , P -1 ) , and A ∗ ( · ) is the convex conjugate of the logpartition function 12 . So that, F θ ( · ) can be thought of as mapping λ t -1 forward in time by first taking the expectation of (4) with respect to π ( z t -1 ; λ t -1 ) , and then applying the mean-to-natural coordinate transformation.

One limitation of amortizing inference through the recursion (28) is its inability to produce approximations for the filtering distributions, p ( z t | y 1: t ) , which can be valuable in streaming or online settings, as well as for testing hypotheses of causality. However, since (17) only depends on the posterior and posterior predictive marginal statistics, we have the freedom to alter our inference network in a way such that filtered marginal statistics are a by-product of obtaining smoothed marginal statistics. For example, an alternative sequence-to-sequence map for λ t could be defined by,

so that ˘ λ t ≡ λ t -β t +1 obey the recursion ˘ λ t = F θ ( ˘ λ t -1 ) + α t and are natural parameters of an approximate filtering distribution, ˘ π ( z t ) ≈ p ( z t | y 1: t ) . Consequently the approximations to posterior and predictive distributions will have a more complicated relationship than they previously did; while efficient sampling and ELBO evaluation are more intricate as a result - linear time scaling in the state-dimension can still be achieved with additional algebraic manipulations, as we show in App. C.2.

5 Experiments

Time complexity & low-rank precision updates. We first investigated the properties of low-rank variational Gaussian approximations in the large L regime. To guide us, we had several questions in mind such as: i) how does the performance of low-rank approximations compare to full-rank and diagonal covariance approximations, ii) how large compared to L should the rank of precision updates be to achieve satisfactory results, and iii) how does convergence using low-rank approximations compare to diagonal approximations, considering that they require a larger number of parameters. We expect that full-rank approximations would perform best (given a sufficient amount of data), since the true posterior will have dense second-order statistics due to the interactions of latent states in both the dynamics and observation models. However, it remains unclear how many dimensions are necessary for a low-rank approximation to achieve similar performance and whether this number will be practical.

We simulated data from 50 D and 100 D linear dynamical systems and compared the convergence between dense and diagonal approximations (Fig. 2 b ); we examined the ELBO for different rank parameterizations in two regimes i) observations and states are of the same dimensionality, N = L (Fig. 2 b - top), and ii) observations are lower dimensional than states, N = L/ 5 (Fig. 2 b - bottom). While not surprising that dense variational Gaussian approximations achieve superior performance, message passing in latent Gaussian models with dense covariance scales like O ( L 3 ) 25 and becomes prohibitive for large L ; thus it is reassuring that in both regimes, low-rank approximate posterior parameterizations achieve comparable results for precision matrix updates of relatively low rank

smoothed

predicted

bps

0

.4

.8

ours (n)

ours (c)

DKF

DVBF

SVAE

LFADS

0

.2

.4

.2

.8

200

600

time (ms)

movement

onset

smoothed

predicted

.89

.88

.74

smoothed

filtered

predicted

true

a)

b)

d)

240 ms

c)

monkey reaching

predictions

start

compared to L . To empirically examine time complexity scaling as a function of L , S , and r , in Fig. 2 a , we plot wall-clock times for fixed r while varying S and L (bottom), and for fixed S while varying r and L (top); reassuringly, inference time complexity scales O ( L ) .

Baseline comparisons - pendulum & bouncing ball. Next, we wondered how our approach fared against other modern deep state-space models when it came to learning complex dynamical systems from data. To explore this, we considered two popular datasets: i) a pendulum system 26 and ii) a bouncing ball 27,28 . Each dataset consists of sequences of observations that are 16 × 16 pixel images that are governed by a low-dimensional dynamical system. An interesting aspect of these datasets is that images can be reconstructed with impartial knowledge of the latent state, but for accurate long-term predictions, the dynamics will need to propagate features of the latent state that are irrelevant to the likelihood (e.g. pendulum angular velocity). For benchmarks, three other deep SSM approaches were included: i) deep variational Bayes filter(DVBF) 5 ii) deep Kalman filter(DKF) 29 iii) structured VAE(SVAE) 20 . We denote our causal amortization network with (c) and the noncausal version with (n).

We trained all models in context windows of 50 consecutive images and then sampled future 50 / 25 time-step latent states from the learned dynamical system for pendulum / bouncing ball. To measure quality of learned latent representation and dynamics, we fit angular velocity / position decoders from training set latent states inferred from pendulum / bouncing ball observations. Then, on held-out test data, we measured the R 2 of velocity / position predictions during the context (smoothing) and forecast (prediction) windows. Fig. 1 shows that all methods are able to reconstruct well in context windows, however, when prediction is concerned, where the underlying dynamics would need to be learned well for accurate forecasts, our method consistently performs better.

Neural population dynamics. We consider two neuroscientific datasets where previous studies have shown the importance of population dynamics in generating plausible hypothesis about underlying neural computation. In addition to DKF, DVBF, and SVAE, we include the LFADS method 7 . First, we considered recordings from motor cortex of a monkey performing a reaching task 30 and evaluate each methods' ability to forecast neural spiking and behavioral correlates. We measure the performance by bits-per-spike (BPS) using inferred spike-train rates 31 and R 2 for decoding hand velocity. Similar to the previous experiment, we evaluate the performance in two regimes: i) a 700 ms context window and ii) a 500 ms prediction window following an initial context window of 200 ms. Fig. 3 c shows that, while all the methods excel at smoothing in the context window, our method makes more informative

predictions in terms of R 2 and BPS. Next, we examined how well the monkey's behavior could be predicted given only causal estimates of the latent state; we trained a model using the causal amortized inference network, given by Eq. (29), then use learned inference network to infer latent states to predict behavior in three regimes: smoothing, filtering, and prediction. Fig. 3 b shows that hand velocity can be decoded nearly as well in the filtering regime (without access to future data) as in the smoothing regime. In Fig. 3 d , we plot how the quality of predictions change as filtered latent states are unrolled through the learned dynamics at different points in the trial; showing that forecasts starting prior to movement onset exhibit strong predictive capability.

Next, we investigated our method's performance with data exhibiting a more intricate trial structure. Specifically, we analyzed physiological recordings from the DMFC region of a monkey engaged in a timing interval reproduction task 32 . During this task, the monkey observes a random interval of time (termed the 'ready'-'set' period) demarcated by two cues, and the goal of the monkey is to reproduce that interval (termed the 'set'-'go' period). We perform a similar procedure as before, but for this experiment we use the period before 'set' as the context window, and use the learned dynamics to make predictions onward; in Fig. 2 d we show the BPS measured on test data during the context and forecast windows. To further investigate the predictive capabilities, we examined condition averaged PSTH produced by samples from the latent state posterior during the joint context/prediction windows. Using the trained model, we sample spike valued observations for the context/prediction windows and then computed condition averaged PSTHs; the results shown in Fig. 2 c , show that PSTHs sampled from the model remain true to the data, even during the lengthy prediction window.

6 Discussion

We presented a new approximate variational filtering/smoothing algorithm, variational learning objective, and Gaussian inference network parameterization for nonlinear state-space models. Focusing on approximations parameterized by dense covariances forced us to consider strategies that ensured computational feasibility for inference and learning with large L . The introduced approximate variational filtering algorithm, while especially useful for the nonlinear dynamics we considered, could also be applied to high-dimensional linear systems where exact computations might be infeasible for large L . Although our variational objective loses the property of lower-bounding the original data log-marginal likelihood, experiments showed that our method consistently outperforms approaches using potentially tight lower bounds. Quantifying this gap or considering potential corrections present interesting directions for future work. Furthermore, while the same variational approach is applicable to any exponential family dynamical system, specific distributions will have their own associated challenges, offering avenues for further research.

Given that neural computation is inherently nonlinear, system identification methods capable of modeling nonlinear dynamical systems are essential for advancing neuroscience. General SSMs can perform well on inferring smoothed latent state trajectories without learning a good model of the nonlinear dynamics. Our proposed method, XFADS, can not only perform efficient system identification and smoothing but also forecast future state evolution for population recordings-a hallmark of a meaningful nonlinear dynamical model; using a causal inference network, XFADS can be used for real-time monitoring, feedback control, and online optimal experimental design, opening the door for new kinds of basic and clinical neuroscience experiments. Future work will focus on developing network architectures for precision matrix updates that are more parameter efficient when the rank those updates become moderate. Moreover, while Alg. 1 remains applicable in the confines of the generative model constraints considered, in certain scenarios, such as when S > L , modifications will need to be made to Alg. 2 to minimize time complexity.

References

Kailath, T., Sayed, A. H. & Hassibi, B. Linear estimation . Prentice Hall, 2000.

Särkkä, S. Bayesian filtering and smoothing . Cambridge University Press, 2013. ISBN 9781107619289.

Anderson, B. D. O. & Moore, J. B. Optimal Filtering . Prentice-Hall, Englewood Cliffs, N.J., 1979. ISBN 978-0-13-638122-8.

Kingma, D. P. & Welling, M. Auto-Encoding variational bayes. In International Conference on Learning Representation , May 2014.

Karl, M., Soelch, M., Bayer, J. & Smagt, P.van der . Deep variational Bayes filters: Unsupervised learning of state space models from raw data. In International Conference on Learning Representations , 2017.

Krishnan, R. G., Shalit, U. & Sontag, D. A. Structured inference networks for nonlinear state space models. In AAAI Conference on Artificial Intelligence , 2016.

Pandarinath, C., O'Shea, D. J., Collins, J., Jozefowicz, R., Stavisky, S. D., Kao, J. C., Trautmann, E. M., Kaufman, M. T., Ryu, S. I., Hochberg, L. R. & others, . Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods , 15(10):805-815, 2018.

Van Overschee, P. & De Moor, B. Subspace identification for linear systems . Springer US, Boston, MA, 1996. ISBN 9781461380610,9781461304654.

Turner, R. E. & Sahani, M. Two problems with variational expectation maximisation for time-series models. In Barber, D., Cemgil, T. & Chiappa, S., editors, Bayesian Time series models , chapter 5, pages 109-130. Cambridge University Press, 2011.

Kingma, D. P. & Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Rezende, D. J., Mohamed, S. & Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning , pages 1278-1286. PMLR, 2014.

Wainwright, M. J. & Jordan, M. I. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning , 1(1-2):1-305, 2008.

Diaconis, P. & Ylvisaker, D. Conjugate priors for exponential families. The Annals of statistics , pages 269-281, 1979.

Seeger, M. Expectation propagation for exponential families. Technical report, University of California at Berkeley, 2005.

Khan, M. & Lin, W. Conjugate-Computation Variational Inference : Converting Variational Inference in Non-Conjugate Models to Inferences in Conjugate Models. In Singh, A. & Zhu, J., editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statistics , volume 54 of Proceedings of Machine Learning Research , pages 878-887. PMLR, 20-22 Apr 2017.

Beal, M. J. Variational algorithms for approximate Bayesian inference . University of London, University College London (United Kingdom), 2003.

Alaa, A. M. & Schaar, M.van der . Attentive state-space modeling of disease progression. Advances in neural information processing systems , 32, 2019.

Girin, L., Leglaive, S., Bie, X., Diard, J., Hueber, T. & Alameda-Pineda, X. Dynamical variational autoencoders: A comprehensive review. arXiv preprint arXiv:2008.12595 , 2020.

Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H. & Davidson, J. Learning latent dynamics for planning from pixels. In International conference on machine learning , pages 2555-2565. PMLR, 2019.

Johnson, M. J., Duvenaud, D., Wiltschko, A. B., Datta, S. R. & Adams, R. P. Structured VAEs: Composing probabilistic graphical models and variational autoencoders. arXiv preprint arXiv:1603.06277 , 2:2016, 2016.

Zhao, Y. & Linderman, S. Revisiting structured variational autoencoders. In International Conference on Machine Learning , pages 42046-42057. PMLR, 2023.

Douc, R., Moulines, E. & Stoffer, D. Nonlinear time series: theory, methods and applications with R examples . CRC press, 2014.

Whiteley, N. & Lee, A. Twisted particle filters. 2014.

Lawson, D., Raventós, A., Warrington, A. & Linderman, S. Sixo: Smoothing inference with twisted objectives. Advances in Neural Information Processing Systems , 35:38844-38858, 2022.

Cseke, B. & Heskes, T. Approximate marginals in latent gaussian models. The Journal of Machine Learning Research , 12:417-454, 2011.

Botev, A., Jaegle, A., Wirnsberger, P., Hennes, D. & Higgins, I. Which priors matter? Benchmarking models for learning latent dynamics. arXiv preprint arXiv:2111.05458 , 2021.

Jiang, X., Missel, R., Li, Z. & Wang, L. Sequential latent variable models for few-shot highdimensional time-series forecasting. In The Eleventh International Conference on Learning Representations .

Missel, R. Torch-Neural-SSM. https://github.com/qu-gg/torchssm , 2022.

Krishnan, R. G., Shalit, U. & Sontag, D. Deep Kalman filters. arXiv preprint arXiv:1511.05121 , 2015.

Churchland, M. M., Cunningham, J. P., Kaufman, M. T., Foster, J. D., Nuyujukian, P., Ryu, S. I. &Shenoy, K. V. Neural population dynamics during reaching. Nature , 487(7405):51-56, 2012.

Pei, F., Ye, J., Zoltowski, D. M., Wu, A., Chowdhury, R. H., Sohn, H., O'Doherty, J. E., Shenoy, K. V., Kaufman, M. T., Churchland, M., Jazayeri, M., Miller, L. E., Pillow, J., Park, I. M., Dyer, E. L. & Pandarinath, C. Neural latents benchmark '21: evaluating latent variable models of neural population activity. In Advances in Neural Information Processing Systems (NeurIPS), Track on Datasets and Benchmarks , 2021.

Sohn, H., Narain, D. & Jazayeri, N. M. M. Bayesian computation through cortical latent dynamics. Neuron , 103(5):934-947, sep 2019.

Khan, M. E. & Rue, H. The bayesian learning rule. arXiv preprint arXiv:2107.04562 , 2021.

Cong, Y., Chen, B. & Zhou, M. Fast simulation of hyperplane-truncated multivariate normal distributions. 2017.

Elfwing, S., Uchibe, E. & Doya, K. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks , 107:3-11, 2018.

Jiang, X., Missel, R., Li, Z. & Wang, L. Sequential latent variable models for few-shot highdimensional time-series forecasting. In The Eleventh International Conference on Learning Representations , 2023.

Li, X., Wong, T.-K. L., Chen, R. T. & Duvenaud, D. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics , pages 3870-3882. PMLR, 2020.

Contents

A Nomenclature

B Variational filtering

Principles of Bayesian inference make it straightforward to write down an algorithm recursively computing the filtering posterior, p ( z t | y 1: t ) 2 . Given, p ( z t -1 | y 1: t -1 ) , updating our belief to p ( z t | y 1: t ) after observing y t can be broken down into two steps: first, we marginalize p ( z t -1 | y 1: t -1 ) through the dynamics to obtain the predictive distribution,

Then, we update our belief by incorporating y t through Bayes' rule,

With all of the filtered/predictive beliefs, the smoothing step is given by,

However, these steps can usually not be evaluated in closed form when we depart from assumptions of Gaussianity and linearity. For nonlinear Gaussian dynamics the predict step can not be carried out exactly, and for nonlinear or non-Gaussian observations neither can the update step.

Alternatively, by considering variational analogues of the predict / update steps, we can develop a recursive and fully differentiable procedure for finding approximations π ( z t ) ≈ p ( z t | y 1: t ) . In developing the variational analogues, it is assumed the approximations belong to an exponential family of distributions, (i.e. π ∈ Q where Q is an exponential family distribution) - not necessarily Gaussian.

B.1 Variational predict step

Similar to developing a recursive algorithm as in the exact case, given π ( z t -1 ) ≈ p ( z t -1 | y 1: t -1 ) , we approximately marginalize π ( z t -1 ) through the dynamics, by solving the following variational (forward KL / moment-matching) problem,

So that if p θ ( z t | z t -1 ) ∈ Q , the optimization problem is minimized when the mean parameters of ¯ π ( z t ) , denoted ¯ µ t , are set to the expected mean parameter transformation under π ( z t -1 ) ,

For a LGSSM with p θ ( z t | z t -1 ) = N ( Fz t -1 , Q ) , using the fact that,

means that if π ( z t -1 ) = N ( m t -1 , P t -1 ) , setting the mean and variance of ¯ π ( z t ) to,

minimizes the forward KL objective, and reassuringly, recovers the familiar Kalman filter predict step.

B.2 Variational update step

For the variational analogue of the filtering update step, we use ¯ π ( z t ) as a prior for the latest observation, y t , and solve the following variational (reverse KL) problem,

If we denote the natural parameters of π ( z t ) by λ t , then the optimal λ t satisfy the implicit equation 33 ,

This usually requires an iterative optimization procedure, except when the likelihood is conjugate to π ( z t ) in which case, the likelihood must take the following form with respect to z t ,

so that, as expected, the natural parameters of the solution are given as Bayes' rule would suggest by adding the data dependent update to the natural parameters of the prior so that,

For a LGSSM with p ψ ( y t | z t ) = N ( Cz t , R ) , this results in the following updates,

which reassuringly recover the information form of the Kalman filter update step 16 .

B.3 Variational smoothing step

Letting ˘ q ( z t ) ≈ p ( z t | y 1: t ) and q ( z t +1 ) ≈ p ( z t +1 | y 1: T ) , we can calculate the statistics of q ( z t ) approximating the smoothed marginal posterior, by minimizing the following objective with respect to the marginal distribution, q ( z t ) ,

taking natural gradients, we find that at a fixed point, the smoothed marginal posterior parameters satisfy the implicit relationship,

B.4 Efficiently sampling structured marginals

Algorithm 2 Nonlinear variational filtering

While P t has a structured representation, drawing samples from π ( z t ) is not straightforward because we do not have a structured representation for a square-root of P t . However, using the factorization of ¯ P t in Eq. (23), it is possible to sample from N ( 0 , ¯ P t ) efficiently since ¯ P 1 / 2 t = [ ¯ M c t Q 1 / 2 ] . Now, combining the fact that the posterior marginal can be written as,

with the result from Cong et al. 34 , stating that sampling z s t ∼ π ( z t ) is equivalent to sampling w s t ∼ N ( 0 , I L + S ) and ¯ z s t ∼ N ( 0 , ¯ P t ) , and then setting

makes it possible to efficiently draw samples from the posterior marginal.

B.5 Efficient filtering

Evaluating ¯ h t = ¯ P -1 t ¯ m t and MVMs with ¯ P -1 t , can be carried out in O ( LS + S 2 ) time, after an initial cost of O ( LS 2 + S 3 ) to factorize ¯ Υ t ¯ Υ glyph[latticetop] t = ( I S + ¯ M c glyph[latticetop] t Q -1 θ ¯ M c t ) -1 , by applying the Woodbury identity to (23),

Since this specifies all quantities that characterize ¯ π ( z t ) , following (13), next is to update our belief by adding the information from the pseudo observation to them,

As a result, the multiplication in m t = P t h t requires O ( LS + Lr ) time, and, analogous to (50), the square root factor Υ t in

requires O ( r 3 + LSr + S 2 r ) time where Υ t Υ glyph[latticetop] = ( I r + K glyph[latticetop] ¯ P t K t ) -1

C Evaluating the KL using low-rank structure

C.1 Smoothing inference network

Efficient training of the generative model and inference networks require efficient numerical evaluation of the ELBO. We take advantage of the structured precision matrices arising from low-rank updates

and sample approximations. The expected log likelihood can be evaluated using a Monte Carlo approximation from samples during the filtering pass. The KL term,

can be evaluated in closed form and we can expand each term as,

Log-determinant. Writing gives

Trace. Writing, further simplifies to

note that the size of the triple product, ¯ M c glyph[latticetop] t K t Υ t , is S × r .

C.2 Causal/streaming inference network

When the real-time parameterization of the inference network is used the expressions become slightly more complicated due to the more intricate relationship between the posterior at time t and the posterior predictive at time t .

Log-determinant. We need to first find log | ¯ P t | T | -log | P t | so begin with using the matrixdeterminant lemma to write,

then expand the smoothed covariance to write,

and another time to write,

which by taking Υ t to be an r × r square root such that,

and another time,

When combined we are finally able to write

For the initial condition we have,

Trace. For the trace,

where we expand the first rhs term for a numerically efficient implementation by writing

To reduce notational clutter we define,

and so the trace is

which is now in a form that is easy to handle using fast MVMs with ¯ P t , ˘ P t , P t .

For the initial condition term we use the fact that ¯ P 1 is diagonal,

D Comparison method details

D.1 SVAE

For the SVAE 20 , the latent state prior is a linear dynamical system parameterized as,

Using conjugate potentials, with likelihood p (˜ y t | z t ) = exp( t ( z t ) glyph[latticetop] α ( y t )) , the approximate posterior is given by q ( z 1: T ) = ∏ p (˜ y t | z t ) p θ ( z 1: T ) so that its statistics can be found by applying Kalman filtering/smoothing to the pseudo-observations. In this case, the ELBO can be evaluated as

where ¯ q t := ¯ q ( z t ) = p ( z t | ˜ y 1: t -1 ) is the filtering predictive distribution. These expressions can be evaluated in closed form and written concisely in terms of natural/mean parameters and log-partition functions as

where µ t = ∇ A ( ¯ λ t + α t + β t +1 ) . Using the identities given in App. F, these expressions can be written in more familiar mean/covariance parameters.

D.2 DVBF

For the DVBF 5 , we parameterize the latent state prior using a nonlinear dynamical system of the same form as Eq. (1). Then, using an inference network that encodes data in reverse-time to produce the parameters of a diagonal Gaussian distribution, w t ∼ q ( w t ) = N ( m t , diag( s t )) , we sample the latent trajectory forward using the recursion, z t = m θ ( z t -1 ) + Q 1 / 2 w t . Parameters of the generative model/inference network are learned jointly by minimizing the ELBO,

D.3 DKF

For the DKF 6 , the latent state is also parameterized using a nonlinear dynamical system of the same form as Eq. (1). We follow the parameterization outlined in the text with S2S( · ) implemented using a recurrent neural network. We sample trajectories using the inference network and jointly train all parameters on the ELBO,

E Experimental details

In describing the neural network architectures used to parameterize the inference model, it will be useful to define the following multilayer perceptron (MLP), with SiLU nonlinearity 35 , that gets used repeatedly:

E.1 High-dimensional linear dynamical system

We simulated data from an LDS generative model (with dynamics restricted to the set of matrices with singular values less than 1 ) for latent dimensions L ∈ [20 , 50 , 100] over 3 random seeds for two scenarios i) N = L and ii) N = L/ 5 . For each scenario, we also vary the rank of the local/backward encoder precision updates.

E.2 Time complexity

Wegenerate random LDS systems of appropriate dimension and measure the average time to complete one forward pass and take a gradient step. The system used for benchmarking wall-clock time was an RTX 4090 with 128 GB of RAM with an AMD 5975WX processor.

E.3 Pendulum

We consider the pendulum system from 26 . We generate 500 / 150 / 150 trials of length 100 for training/validation/testing. All methods are trained for 5000 epochs for 3 different random seeds. We consider a context window of 50 images and a forecast window of 50 images. A decoder was fit from the latent state on the training set during the context window; then, for held out data, we examine performance of the decoder during the context and forecast windows.

The generative model is parameterized as:

L = 4

N = 256

T = 50

likelihood

-p ψ ( y t | z t ) = N ( z t | C ψ ( z t ) + b , R ) -C ψ : MLP(4 , 128 , 256)

For each method, inference is amortized using the following neural network architectures:

our inference network

-r α = 4

-r β = 4

-α φ : MLP(256 , 128 , 20)

-β φ : [GRU(128) , Linear(128 , 20)]

DKF inference network

-u φ : GRU(128)

-( m φ , log P φ ) : MLP(132 , 128 , 8)

DVBF inference network

-( µ φ , log σ φ ) : MLP(132 , 128 , 8)

SVAE inference network

-α φ : MLP(256 , 128 , 20)

Optimization and training details:

-optimizer: Adam(lr = 0 . 001) -batch size: 128

ours(n)

ours(c)

time

E.4 Bouncing ball

We consider a bouncing ball dataset commonly used as a baseline to benchmark the performance of inference and learning in deep state-space models 26,36,36 . For this dataset we take 500 / 150 / 150 trials of length 75 for training/validation/testing. All methods are trained for 5000 epochs for 3 different random seeds. The generative model is parameterized as:

L = 8 · N = 256 · T = 50 · likelihood -p ψ ( y t | z t ) = N ( z t | C ψ ( z t ) + b , R ) -C ψ : MLP(8 , 128 , 256)

For each method, inference is amortized using the following neural network architectures:

our inference network

-r α = 8

-r β = 4

-α φ : MLP(256 , 128 , 72)

-β φ : [GRU(128) , Linear(128 , 40)]

DKF inference network

-u φ : GRU(128)

-( m φ , log P φ ) : MLP(132 , 128 , 16)

DVBF inference network

-u φ : [GRU(128)]

-( µ φ , log σ φ ) : MLP(132 , 128 , 16)

SVAE inference network

-α φ : MLP(256 , 128 , 72)

Optimization and training details:

-optimizer: Adam(lr = 0 . 001) -batch size: 128

ours(n)

ours(c)

DKF

DVBF

LGL

X

time

E.5 MC_Maze

The monkey reaching dataset of 30 was the first real dataset examined in the main text. For this dataset, we partitioned 1800 / 200 / 200 training/validation/testing trials sampled at 20 ms per bin. All methods are trained for 1000 epochs for 3 different random seeds. The generative model is parameterized as:

L = 40

N = 182

T = 35

likelihood

-p ψ ( y t | z t ) = Poisson ( z t | C ψ ( z t ) + b ) -C ψ : Linear(40 , 182)

For each method, inference is amortized using the following neural network architectures:

our inference network -r α = 15 -r β = 5 -α φ : MLP(182 , 128 , 640) -β φ : [GRU(128) , Linear(128 , 240)] · DKF inference network -u φ : GRU(128) -( m φ , log P φ ) : MLP(128 , 128 , 80)

DVBF inference network

-u φ : [GRU(128)]

-( m φ , log P φ ) : MLP(128 , 128 , 80)

SVAE inference network

-α φ : MLP(182 , 256 , 1640)

Optimization and training details:

optimizer: Adam(lr = 0 . 001)

batch size: 128

0

0

0

0

0

0

1

predictions

start

200 ms

time (ms)

movement

true speed

model speed

speed (a.u)

true reach

time to

movement

onset

0.14

-240 ms

0.44

-200 ms

0.62

-160 ms

0.71

-120 ms

0.77

-80 ms

onset

E.6 DMFC_RSG

The second real dataset examined was the timing interval reproduction task of 32 samples at 10 ms bins. For this dataset, we partitioned 700 / 150 / 150 training/validation/testing trials. All methods are trained for 1000 epochs for 3 different random seeds. The generative model is parameterized as:

L = 40

N = 54

T = 130

likelihood

-p ψ ( y t | z t ) = Poisson ( z t | C ψ ( z t ) + b )

-C ψ : Linear(40 , 54)

For each method, inference is amortized using the following neural network architectures:

our inference network

r α = 15 - r β = 5 - α φ : MLP(54 , 128 , 640) - β φ : [GRU(128) , Linear(128 , 240)]

DKF inference network

-u φ : GRU(128)

-( m φ , log P φ ) : MLP(128 , 128 , 80)

DVBF inference network

-u φ : [GRU(128)]

-( µ φ , log σ φ ) : MLP(128 , 128 , 80)

SVAE inference network

-α φ : MLP(256 , 128 , 1640)

Optimization and training details:

optimizer: Adam(lr = 0 . 001)

batch size: 128

E.7 Collated results

F Useful expressions

mean/natural parameter inner product One common expression that frequently arises is the inner product between a mean and natural parameter, i.e.

where in the Gaussian case if the mean/natural parameter coordinates are,

then,

difference of log partition functions Another common expression that frequently arises is given by,

so that if, then,

conditional linear Gaussian log partition For a Gaussian distribution, the log-partition function in terms of the natural parameters, λ = ( h , vec( J )) , is given by

and u glyph[latticetop] = ( u glyph[latticetop] 1 vec( u 2 ) glyph[latticetop] ) is an arbitrary constant, we can write A ( F t ( z ) + f + u ) = a glyph[latticetop] t ( z ) + b for some a and b . To show this, we can just expand the log partition function

then a and b can be identified as

Then, for a LGSSM, we can evaluate the following expression which describes the difference between predictive and smoothed marginals,

where S t +1 = Q -1 -Q -1 ( Q -1 + P -1 t +1 -¯ P -1 t +1 ) -1 Q -1 .

so that, when

G Forward KL fixed point

Using the moment matching property of the fixed point solution to the forward KL objective, we can write

This result holds for any setting where stochastic transitions, p t +1 | t , and the approximate marginal distributions, q t , are restricted to the same exponential family distribution.

H Linear Gaussian (information form) smoothing

Using the derived relation, means that,

so that using the definition of S t (rewritten for convenience) then plugging in, we get,

which gives a recurrence backward in time for S 1: T starting from S T +1 = 0 . Now since,

and, we have that,

which simplifies into a recursion for s 1: T backward in time starting from s T +1 = 0 , given by,

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: we try to clearly state our contributions.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]

Justification: We discuss potential limitations of the work.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We list assumptions about the types of generative models we consider.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide experimental details in the appendix that include the architectures and training regimes.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide sufficient instructions and algorithmic details.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide details about splits, optimization, and parameterizations in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We discuss the neuroscientific implications of the results our model could be used to discover from neural data.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: we provide our computing setup.

Answer: [Yes]

Justification: We provide the specifications of our computing setup.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: we adhere to the code of ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the broader impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: the paper poses no such risks Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: we produced all code to run the experiments, we cite real data from where they were procured.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: the paper does not release new assets

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing or human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: this paper does not involve human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work develops a deep state-space model (DSSM) for state inference in and system identification of non-linear dynamical systems. \nThe 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).\n\nSTRENGTHS:\n- The paper is very well written and quite easy to follow for the most part (see weaknesses). \n- Hyperparameter, notation and algorithm details as well as derivations are given in the supplementary material, which is very helpful.\n\nWEAKNESSES:\n- 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.\n- 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. \n\nSoundness:\nI have difficulties with the technical soundness of the paper, especially about the following:\n- 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. \n- 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. \n- 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. \n- 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?\n\nNovelty:\n- 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. \n- 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. \n\nExperiments:\n- 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.\n\n\n[1] Karl et al. 2016, Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data.\n[2] Fraccaro et al. 2017, A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning.\n[3] Becker-Ehmck et al. 2019, Switching Linear Dynamics for Variational Bayes Filtering.\n[4] Kurle et al. 2020, Deep Rao-Blackwellised Particle Filters for Time Series Forecasting.\n[5] Klushyn et al. 2021, Latent Matters: Learning Deep State-Space Models.\n[6] Simo Sarkka, Bayesian filtering and smoothing.\n\n-------\nMinors and suggestions that did not influence the review score\n- 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. \n- 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. \n- line 48: vEM needs a citation. For instance, VEM is mentioned in the dissertation of Beal, but there is likely a better reference therein.\n- 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. \n- 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.\n- line 68) typo: ii) -> iii)\n- Equation (3) typo: parameter\n- line 78: citation style changes\n- dKF, dVBF, vEM all written with lower case first letter, but other works (including the original ones) use capitalization.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nLimitations have been addressed to some extent.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUpdate post author rebuttal.\nThanks 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.\n________\n\nThe paper presents a class of non linear state-space models whose dynamics are defined as exponential family distributions.\nThe 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.\nEfficient filtering/smoothing is obtained with an approximate message passing algorithm that exploits the low rank structure of the variational distribution.\nThe model is tested on a number of simulated and real world applications and shows good predictive performances.\n\nSTRENGTHS:\n* 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)\n\n* The model combines ideas from previous work, but is overall novel to the best of my knowledge\n\n* The model can handle well missing data\n\n* The choice of inference network allows scalable inference\n\n* The theoretical section is dense but well explained (unlike the experimental one as noted below)\n\n* The predictive performances are better than similar SOTA models\n\nWEAKNESSES:\n**Main comment**\n\nI 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).\nThere are several issues that taken together make the experimental quite hard to follow:\n\n1. 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.\n\n2. 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.\n\n3. 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.\n\n4. In the caption of figure 2a \"(top) for fixed S and (bottom) for fixed r\" seems wrong.\n\n\n**Other minor points**\n\n* Line 36 in the introduction: the complexity term used terms that are not yet defined\n\n* 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\n\n* Line 145 - you reference the non existing equation (94)\n\n* Line 330, missing \"the\" in \"if the number of samples\"\n\n* Is the format of the citations as a superscript an acceptable one for neurips? It's not very common\n\nQUESTIONS:\nBased 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\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

comments: 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?

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?
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\n- 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.\n- 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).\n- 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.\n\nOverall, I consider these points to be relatively minor and do not detract from the paper's contributions.\n\nLiu, Wei, et al. \"Neural extended Kalman filters for learning and predicting dynamics of structural systems.\" *Structural Health Monitoring* 23.2 (2024): 1037-1052.\n\nQUESTIONS:\n- Can this method be efficiently parallelized?\n- 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?\n- 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?\n- How well do the low-rank update scale to higher-dimensional systems?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-17"}