eXponential FAmily Dynamical Systems (XFADS): Large-scale nonlinear Gaussian state-space modeling
Authors: Matthew Dowling
Yuan Zhao
Il Memming Park
Venue: NeurIPS 2024 poster (2024)
Keywords: variational inference
nonlinear state-space model
dynamical system
Primary area: probabilistic_methods
TL;DR: We introduce a structured variational approximation and inference algorithm for efficient Bayesian inference in nonlinear state-space models.

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
Matthew Dowling
Champalimaud Research, Champalimaud Foundation, Portugal beethew..dwwing@reeeeee fchampal imaud. org
Yuan Zhao
National Institute of Mental Health, USA
yuan www.wnih.ahhh.cmm .oooo gov
II Memming Park
Champalimaud Research, Champalimaud Foundation, Portugal memming. park@research. fchampal imaud. org
Abstract
State-space graphical models and the variational autoencoder framework provide 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 flex- ibility 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 low-rank structured variational autoencoding framework for nonlinear Gaussian state-space graphical models capable of captur- ing dense covariance structures that are important for learning dynamical systems with predictive capabilities. Our inference algorithm exploits the covariance struc- tures 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 dimension- ality. In comparisons with other deep state-space model architectures our approach consistently demonstrates the ability leam Furthermore, when applied to neural physiological recordings, our approach is able to learn dynamical system capable of forecasting population spiking and behavioral correlates from small portion of single trials.
predictive generative model.
Introduction
State-space models (SSM) are invaluable for understanding the temporal structure of complex natural phenomena through their underlying dynamics While engineering physies problems often assume the dynamical laws of the system of interest are known to high degree of accuracy, in unsupervised data-driven investigation, they have 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 inference network to amortize posterior computation of the unknown latent 0 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 ELBO) while preserving the temporal structure of the posterior.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 an approximate filtering problem for pseudo-observations that encode a representation of current and future data, and, (iii) An inference algorithm that scales L(T + S + r²)) - 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 space models are probabilistic graphical models where observations y, in discrete time conditionally independent given a continuous latent state, z,, evolving according to Markovian dynamics, so that the complete data likelihood for T consecutive observations factorizes as,
plyI.r, ZIr)) (2( (2))
II
1-2 Pe (yel 2:) po (Ze 2z-1)
where z, RL are real-valued latent states, 0 parameterizes the dynamics and initial condition, and s parameterizes the observation model. When the generative model, (0, vb), is known, the statistical inference problem is to compute the smoothing posterior, p p(Z [yiT)) Otherwise, (0, +/) have to be learned from the data known as system identification
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 lower bound to the log-marginal likelihood the BLBO with respect to the parameters of an approximate posterior, ((21.)) Rs p(21.T] yI:)); then, with the approximate posterior fixed, the ELBO maximized with respect to parameters of the generative mode 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
Model specifications. Although our approach is applicable to any exponential family sate-space process, given their ubiquity, we focus on dynamical systems driven by Gaussian noise so that,
pe (z, 24-1) (2 (2 Qa)
where ma RRL + R. might be a nonlinear neural network function with learnable parameters 0, and Qo E RL2L is learnable state-noise covariance matrix. Given the favorable properties of exponential family distributions 15, especially in the context of variational inference, write the prior dynamics in their exponential family representation (natural parameter form),
po Z, 2t 1) = h(zz) exp (T(z:) A0(2 - A(A
(2)
where h is the base measure, T(z:) the sufficient statisties, A(-) the log-partition function, and Aa(-) a map R² + 2 that transforms 2:-1 to natural parameters for Zt. For a Gaussian distribution,
the sufficient statisties can be defined T T(z)) zz Z so that (- for is given by,
Ao (z(Z=-)) = [Q01 "mmalze (Zi-1)
(dynamics model in natural paramter form)
Qo
As will simplify subsequent analysis, the mean parameter mapping corresponding to this natural parameter mapping (guaranteed to exist as long as the exponential family is minima 12 is given by
mo (Zi-1) ( + Qo)
mean
H4 (z(-1) = Epo (x:/mee))))) [T(z)) = -
parameter form
(4)
Furthermore, we make the following assumptions i) the state-noise, Qo. is diagonal or structured for efficient matrix-vecto multiplications. ii) m-(-), ), is a nonlinear smooth function. ii) the likelihood,
Pv (yel lzz), may be non-conjugate. iv) L may be large enough that L is comparable to T.

---

Amortized inference for state-space models. A useful property of SSMs is that, conditioned on 2 and YL.T, is independent of Y1æ- 1. i.e, p(z, I.:1) = p(zz 1:-1. (516 suffices to construct an approximate posterior that factorizes forward in time,
It thus
((zT.) - ((Z)) II
gZe (Z:-)
(5)
and introduce learnable function approximators to amortize inference by mapping - and y:.T to the parameters of g(2,) kt-1). This makes it simple to sample 21.7 from the approximate posterior (using the reparameterization trick) and evaluate the ELBO (a.k.a. negative variational free energy),
((q)= = Σ E. [logp(y. z,)] - Eq.,, [KL (9(z, |z,- 1)l po (z 21-1))1 & logp(yirr)
(6)
where 11xx (-I1-) is the Kullback-Leibler (KL) divergence and E, = ative model and inference network parameters can be learned through stochastic backpropagation. Many works for Gaussian q(zz| =;- such as Krishnan et al. [. Alaa and van der Schaar Girin
that the gener-
et al. [T]] Hafner et al. [9 construct inference networks that parameterize the variational posterior as
[z 1 = (Ima(Z)-1..i P, P (1(2)-1.11r)))..
(7)
There are limitless ways to construct m+ and (-) so &b can be learned through gradient ascent on the ELBO, but straightforward and illustrative approach 617 is transform future data, YeT, using a recurrent neural network (RNN), or any efficient autoregressive sequence to sequence model, and then mapping the preceding latent state, 2T-1, using a feed-forward neural network, so that a complete inference network description could be.
(m (Z Vc.r), P. Po(2
NN((2-1: Uz]);
u, = -22S((
where (- is parametric saneence-t--squenee function that maintains hidden state u, and takes as input yt, and NN(.) is a parametric function designed to output approximate posterior parameters. hakkwadd-forwatt algorithm, meaning that data y1:T are mapped to u1:T in reverse time, and then samples are drawn from gZz Z,-1) forward in time.
This leads
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 fighvvarancce 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 (SVAE) 20 makes possible to efficiently evaluate the ELBO while preserving the temporal structure of the posterior by restricting the prior to a linear dynamicul system (L.SS) and then constructing the approximation as (Z1:T) (Z) I[ exp(t(z.) (vy))) 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 2 expanded on the SVAE framework by exploiting the LDS structure and associative scan operations to improve its scalability.
The deep Kalman filter (DKF) 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 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 Difficulties of learning the generative model that arise as a result of more standard VAE 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 SVAE 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).
3

---

4 Method
An alternative to constructing variational approximations through specification of conditional dis- tributions, as in Eq. 7 involves the use of data-dependent Gaussian potentials, that we refer to as pseudo-observation::
p ze) exp()o(yi:) T((z)) = exp(X, T(zi) exp (k. z: | 113)
(9)
These Gaussian potentials can then be combined with the prior through Bayes' rule, yielding the approximation
((ZI.)) P(ye (P)
(10)
p(yi.))
benefit of this formulation, that inherently imposes the latent dependency structure of the generative model onto the amortized posterior. This parameterization, was introduced in Johnson al. 20, where an important point highlighted, that the Gaussian potentials can encode any arbitrary subset of observations; for example, p( Zz) could be made to depend on Yt alone, or even the entire dataset, y1:7- Regardless of that particular design choice, the corresponding ELBO for the variational approximation of Eq. (10) is
( = Σ Egg p(y: 2:)] - Eg. [log p( 22)] + log p(V1:T)
(11)
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 the right-hand side, while the filtered statisties 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 -oservvtoosrrvno relies first obtaining the filtered marginals, logical starting point is to develop a method for approximating these filtered marginals. To this end, we propose differentiable approximate message passing algorithm specifically designed to compute filtered posterior statistics in models characterized by nonlinear latent dynamics and observations represented Gaussian potentials. Building this foundation, 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 two step proce- 2ure2, In the predict step, our belief of the latent state integrated through the dynamics, yielding ((2 [1a-1) pe (4-1)9(2;-1 F1a-1) dz,-1 (a. the predictive distribution). Then, applying Bayes' rule, q(z [1::) x p(y: Zi (g(z 11:--1). we update our belief. Evidently then, developing approximate filtering algorithm that exploits the conjugacy of the pseudo observa- tions can be recast as the problem: given an approximation π (Zz-1) Rs g(zz-1 V144-), find an approximation to the predictive distribution, (2z)) ((z (1--11)- The recursion would continue forward by updating our belief analytically, setting क (zz) x p 2e (z,), then finding a Gaussian approximation of (Z(+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),
Variational filtering
क (2t) = argmin DKL (E.,,, (z: 12-11 (e-1)]|| [(2+))
p = Ex,_, [z0(2)1)] A X A
(12)
(ii) T2,)) argmin (RL(((z)) p(5. Z (Z)))
(13)
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, p., 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,
log logp(ýn:7) Elog log /
p(yr [ Zeyelze)oe (z( 151a-1) dz ~ Elog p(y: (z) dz,
(14)
where the last integral can be evaluated analytically as a result of the Gaussian forms of the ap- proximations and pseudo observations. While formulating step (ii) as 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 de min- imizing LR (B [ gt), by using the forward KL divergence (as in step (i)), leads to an intractable problem because the backward Markov transitions, qt|t+1. are not conditionally Gaussian. Conversely, fixed point solution of the reverse KL objective (as in step(ii)), DxL ( (gell 19
necessitates an iterative procedure, which can be computationally expensive.
Smoothing as filtering. In light of these difficulties, we offer simple solution that exploits the flexibility in choosing the pseudo observation data dependence: define the parameters, k, and K,, of each pseudo observation, Yr, to be function of current and future data, Yr.r, so that,
p (2) X (xp(A<(Verr) T(z:))
(15)
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,) ~ g(z, V1:)) = p(z; Vı.r). This solution circumvents the challenges associated with backward message computation and only requires single pass through the pseudo observations to obtain approximate smoothed posterior statistics. Substittinng Te, as an approximation to qt. in Eq. (TI.. leads to the following approximation of the ELBO.
Variational smoothing ELBO
((T)
Σ Es, [log ply 22)] - Es, [log p(y: [2)] log Ex, [p(Fr/z))
(16)
= Σ Ex, [log 2c] - DKL (2:))
(17)
By expressing the approximate ELBO compactly as Eq. 177. we highlight that it promotes learning models where the posterior at time i aligns closely with the one-step posterior predictive at that time (which depends 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- ligeiiolioood term can be approximated using the reparamcterization trick However, a point of practical importance should be raised now: every filtering step and evaluation of the KI. term has a time complexity of ((L)), which may become a bottleneck for large L. In the following discussion, the Gaussian potential inference network that produces Ay.r. to reduce the computational burden that filtering and evaluating [(T) pose in the large L regime.
will explore effective strategies to parameterize
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) local encoder, (-), for current observation, and ii) a backward encoder, (+(+), ), for future observations, i.c.,
Sp(yer) = (y.) + (y++1) (or for the sake of brevity) A, = ac + B.41
(18)
Furthermore, by building the dependence of (-) on Yr+1-7 through their representation so that (y = 3 (04+1::)), a missing observation at time is handled by setting a; = 0. While data dependent natural parameter update of faithfully represents a missing observation in the absence of data, the prior should not be updated- - alternatively setting y: 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²) memory requirements, it is appealing to consider low-rank parameterization for the local and backward encoders we consider
1:711:7
a(y:) a(x))
b
b(aEr)
a: = A,A.T = (A()))(()))))
G
B.B.
(((ar))((arr))
(19)
5

---

a)
pendulum
b)
bouncing ball
R angle
position
true
0 44
8
0
8
ours (n)
ours (c)
DKF
DVBF
SVAE
prediction
25
25
starts
smoothed
predicted
Figure 1: Smoothing and predictive performance on bouncing ball and pendulum. To the left of the red line are samples from the posterior during the data window projected to image space, to the right of the red line are samples unrolled from pæ(z; Zz-1). a) while all methods are adept at smoothing the context window, our methods predictive performance better by noticeable margin measured by the R². h) similar results hold for the bouncing ball dataset.
where A, E L2r with B, E 4xra parameterize low-rank local/backward precision updates, and a, R& with b, E RR parameterize local/backward precision- precision-scaled mean updates. Using these descriptions and the additive decomposition 18 the parameters of single pseudo observation are,
K: K.K.
(RtK,)=
k(YT)) K(YET ( ((ye))())))))))
+ b, [A, B. A B.
A,
(20)
where K E I R.r if Ta + T3 and k E R The low-rank structure of the natural parameter updates will be key component 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 (B), a differentiable architecture producing our and B1.r could be,
a₁ NN(y.)
B. = S2S([ a.)),
(21)
which overall defines the map YI.T + (I.T)) In addition, the separation of local and backward encoders can reduce the complexity of the backward encoder for L < N. Those familiar with sequen- tial Monte-Carlo (SMC) methods can view the backward encoder similar twisting functions used to combine future information with filtered state beliefs to produce smoothing approximations 2324
Exploiting structure for efficient filtering. A benefit of using the forward KL to design 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 2 ~ π (Z4-1). gives the approximation,
ma(zi-1) p = 1/s Σ --1 | (mo (2i-)) (m (z +Q0]
(22)
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,
m, = 1/ss Σ. ma (zi-1)
P. M.M.T + Qo
(23)
where MΦ is the L S matrix of samples passed through the dynamics function, then centered by the mean, defined for convenience as,
M 1/Vs [mo
m., (m -m,2"),,((((_)))) m.] E RRSS
(24)
Writing the covariance estimate as it is in Eq. 233 reveals that can alternatively represented by the pair (M,. Qo). In the regime where L S, significant computational savings 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, (and its inverse, after application of the Woodbury identity). Consequently, the natural parameters of R,, after updating ,,, are
P. P mm;- = P m, k,
P = P K.K.
(25)

---

a)
fixed
10
b)
[(T))
C) data
samples
d)
DMFC RSG
100
2.0
LIE/2
S/H
L = 50
L = 100
10 Hz
BPS
wall-clock time (s)
50
1.8
diagonal
.2
ours (n)
.1
1.6
S/H/3
LE/
ours (c)
fixed 5/5
40
40
5 Hz
DKF
2.5
40
15
DVBF
.12
25
2.0
15
25
25
50
S
S/H14
SVAE
10 Hz
.10
1.5
smoothed
10
100
10
0
10
prediction
L
Kalman
epoch
time
1.3s
starts
predicted
Figure 2: a) Empirical time complexity scaling. Since complexity is function of L, S. and we vary L (top) for fixed r = 10 and (bottom) for fixed S = 5; we examine several values of the variable not fixed. Examining wall-clock time shows empirically our implementation scales linearly in L; on the (bottom) we plot wall-clock time for a Kalman filter implementation, showing the standard cubic dependence on L. b) (top) Negative ELBO as a function of training epoch when N = L (bottom)
when N = L/5; the left column shows the case L = 50 and the right when L = 100. Different colors indicate different settings of the local/backward encoder rank; zooming in for L = 100, shows low-rank updates can match diagonal ones. c) Peristimulus time histogram (PSTH) for the DMFC RSG dataset for different trial condition averages; we consider a context window of 1.3s and a prediction window of 1.3s. d) BPS for each method for context/prediction windows.
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)) cost. Instead, recog- nizing the potential computational advantages of exploiting these structures, and never instantiating the predictive/updated covariance and precision matrices, makes it possible to develop an approxi- mate filtering algorithm, where in the case L is significantly larger than S r, has complexity of OL((Sr r + +))) per step. More details regarding time complexity are in App. B.5
Efficient sampling and ELBO evaluation. When E,,, [log 22)] can not be evaluated in closed form, Monte-Carlo integration can be used as differentiable approximation. To sample from xzz)) without explicitly constructing Pt, we can take 2 ~ No P.) and w ~ NO,, and set,
z
z K,Y.Y (K ž w().
(26)
While more details are provided in App. B.4 this can be done efficiently since samples can be drawn cheaply from म (2;) using Eq. (23). Whereas Monte-Carlo approximations of the expected log-like log-likel 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,
TeK (T [ ((:)) 3 [m, m.) TP.' (m, m.) + +r((P.'P.) log([ log((P.//PP.)) L.]] (27)
scales O(L3). However, since matrix vector multiplies with P,-' can be performed efficiently and the trace/log-determintt terms can be rewritten using the square-root factors acquired during the forward pass, as we describe in App. C.1 is possible evaluate the KL in OLSr + LSS + Lr²) time. After complete forward pass through the encoded data, we acquire the samples z 2 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. |
Causal amortized inference for streaming data. In constructing a fully differentiable variational approximation, the parameters of the approximate marginals were effectively amortized according to recursion in the natural parameter space by iterating Eqs. (12) and (13). This recursion can be recognized more easily by introducing the function, (-()). and writing
x; = 5o ()- + ar + Bit1 with (A. = VA- (fa(z. d :)sool (00 (2+-1) dzi-1) (4-1))
(28)
7

---

Algorithm 1 End-to-end learning
Input: Yı.T
while not converged do
for t = T to do
a: Ny()):)
local encoder
B. S2S acl) # backward encoder
k; a, + b,
K, [A, B.]
end for
1:T mI:7, mI.T, T Alg.
1:S
kI.T, KI.T) DKL (Teelell Tz)]
((T) E Σ loogri p(y: / 22i)
(Φ. 0, +) (Φ, 0. v)) V.(T)
end while
Output: 1 mI.7, mi.T, Y:r
VA : (m, -P mm T)) + (P-1m, P-1), and (-) is the convex conjugate of the log- partition function So that, (-) can be thought of as mapping A forward in time by first taking the expectation of (4) with respect to 11 and then applying the mean-to-natural coordinate transformation.
One limitation of amortizing inference through the recursion 28 is its inability to produce approxima- tions for the filtering distributions, p(ze yı:r), which can valuable 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 by-product of obtaining smoothed marginal statistics. For example, an alternative pruuencc-to-sequenee map for X could be defined by,
A = -o (>: - 3₂) + Oz + 3
(29)
smoothed future filtered
so that > Bi+1 obey the recursion X, To (}(-1)+ + ar and are natural parameters of an approximate filtering distribution, #(z,) = p(z, yıa). Consequently the approximations to posterior and predictive distributions will have more complicated relationship than they previously did; while efficient sampling and ELBO evaluation are more intricate as result - linear time scaling in the dimee-dimension can still achieved with additional algebraic manipulations, as we show in
App. C.2
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 be to achieve satisfactory results, and iii) how does convergence using low-rank approximations compare to diagonal approximations, considering that they require larger number of parameters. We expect that full-rank approximations would perform best (given sufficient amount of data), since the true posterior will have dense second-order statistics due 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.
should the rank of precision updates
We simulated data from 50D and 100D linear dynamical systems and compared the convergence between dense and diagonal approximations (Fig. p); we examined the ELBO for different rank parameterizations two regimes i) observations and states are of the same dimensionality, N L (Fig. p top), and ii) observations are lower dimensional than states, N L/5 (Fig. D bottom). While not surprising that dense variational Gaussian approximations achieve superior performance, message passing in latent Gaussian models with dense covariance scales like 0 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

---

a)
true
b)
smoothed
filtered
predicted
monkey reaching
R²
bps
ours (n)
ours (c)
R
.89
.88
.74
DKF
d)
predictions start
DVBF
.8
R
smoothed
SVAE
predicted
LFADS
.2
movement
smoothed
predicted
200
600
time (ms)
240 ms
Figure 3: Predict behavior from a causally inferred initial condition. a) Actual reaches. b) (top) Reaches linearly decoded from smoothed (R² = 0.89), causally filtered (R² = 0.88), & predicted (R2 0.74) latent trajectories starting from an initial condition causally inferred during the prepara- tory period. (bottom) Top 3 principal latent dimensions per regime for three example trials. e) bps R of predicted hand velocity using rates inferred from the 700ms context window and the 500ms prediction window. d) Velocity decoding R² using predicted trajecto- ries as a function how far into the trial the latent state was filtered until was only sampled from the autonomous dynamics; by the the movement onset, behavioral predictions using latent trajectory predictions are nearly on par with behavior decoded from the smoothed posterior.
compared to L. To empirically examine time complexity scaling function of L, S, and in Fig. 2h we plot wall-clock times for fixed while varying S and L (bottom), and for fixed while varying 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) pendulum system and ii) a bouncing ball 8888 Each dataset consists of sequences of observations that are 16 x 16 pixel images that are governed by a low-dimensional dynamical system. 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 (DVBF) fi) deep Kalman filter(DKF) iii) structured VAE(SVAE) We denote 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 Rt of velocity / position predictions during the context (smoothing) and forecast (prediction) windows. Fig- [ shows that 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 dynamies generating plausible hypothesis about underlying neural computation. In addition to DKF, DVBF, and SVAE, we include the LFADS = First, we considered recordings from motor cortex of monkey performing a reaching task 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 ral and R22 for decoding hand velocity. Similar to the previous experiment, we evaluate the performance in two regimes: i) 700ms context window and ii) a 500ms prediction window following an initial context window of 200ms. Fig. 3 shows that, while all the methods excel at smoothing in the context window, method makes more informative

---

predictions in terms of R and BPS. Next, we examined how well the monkey's 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. 33 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. Fig. 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 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 to reproduce that interval (termed the 'st'''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 dynamies to make predictions onward: in Fig. 21 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/preiiction 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. [2c. 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/smoohhnng algorithm, variational learning objec- tive, 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 varia- tional filtering algorithm, while especially useful for the nonlinear dynamics we considered, could also be applied to high-dimensionall linear systems where exact computations might be infeasible for large L. Although our variational objective loses the property of bowur-bounding the original data log-marginal likelihood, experiments showed that 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 meaningful nonlinear dynamical model; using 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. [ 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.
Acknowledgements
We would like to thank the anonymous reviewers and Scott Linderman for their insightful comments and constructive feedback, which greatly improved the quality of this paper. We thank Mahmoud Elmakki for help with organizing the code base and developing helpful demos. Yuan Zhao was supported in part by the National Institute of Mental Health Intramural Research Program (ZIC- MH002968). This work was supported by NIH RF1-DA056404 and the Portuguese Recovery and Resilience Plan (PPR), through project number 62, Center for Responsible AI, and the Portuguese national funds, through FCT Fundação para a Ciência Tecnologia in the context of the project UIDB/04443/200.
10

---

References
[1] Kailath, T., Sayed, A. H. & Hassibi, B. Linear estimation. Prentice Hall, 2000.
[2] Särkkä, S. Bayesian filtering and smoothing. Cambridge University Press, 2013, ISBN 9781107619289.
[3] Anderson, B. D. O. & Moore, J. B. Optimal Filtering. Prentice-Hall, Englewood Cliff', N.J., 1979. ISBN 978-0-13-63812244
[4] Kingma, D. P.& & Welling, M. Encoodin variational bayes. In International Conference on Learning Representation, May 2014.
[5] 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.
[6] Krishnan, R. G., Shalit, U. & Sontag, D. A. Structured inference networks for nonlinear state space models. In AAAI Conference on Artificial Intelligence, 2016.
[7] Pandarinath, C., O' Shea, D. J., Collins, 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)800-8815 2018.
[8] Van Overschee, P. & De Moor, B. Subspace identification for linear systems. Springer US, Boston, MA, 1996. ISBN 978 97814613806 1911777466116644
[9] Turner,
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.
[10] Kingma, D. P. & Ba, Adam: method for stochastic optimization. arXiv preprint arXiv: 69800 2014.
[11] Rezende, D. J., Mohamed, S. & Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. In International conference machine learning, pages 1278-1286. PMLR, 2014.
[12] Wainwright, M. J. & Jordan, M. I. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, (1-22)11330) 2008
[13] Diaconis, P. & Ylvisaker, D. Conjugate priors for exponential families. The Annals of statistics, pages 269-281, 1979.
[14] Seeger, M. Expectation propagation for exponential families. Technical report, University of California at Berkeley, 2005.
[15] 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.
[16] Beal, M. J. Variational algorithms for approximate Bayesian inference. University of London, University College London (United Kingdom), 2003.
[17] Alaa, A. M. & Schaar, M.van der Attentive state-space modeling of disease progression. Advances neural information processing systems, 32, 2019.
[18] Girin, L., Leglaive, S., Bie, X., Diard. J., Hueber, T. & Nameda-Pineda, X. Dynamical variational autoencoders: A comprehensive review. arXiv preprint arXiv:2008, 12595, 2020.
[19] 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.
[20] 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.0663 1603.06277, 2:2016, 2016.
[21] Zhao, & Linderman, S. Revisiting structured variational autoencoders. In International
Conference Machine Learning, pages 42046-42057. PMLR, 2023.

---

[22] Douc, R., Moulines, E. & Stoffer, D. Nonlinear time series: theory, methods and applications with R examples. CRC press, 2014.
[23]
Whiteley, N. & Lee, A. Twisted partiele filters. 2014.
[24] Lawson, D., Raventós, A., Warrington, A. & Linderman, S. Sixo: Smoothing inference with twisted objectives. Advances in Neural Information Processing Systems, 35:38844-38833 2022.
[25] Cseke, B. & Heskes, T. Approximate marginals in latent gaussian models. The Journal of Machine Learning Research, 12:417-454, 2011.
[26] Botev, A., Jaegle, A., Wirnsberger, P., Hennes, D. & Higgins, I. Which priors matter? Bench- marking models for learning latent dynamics. arXiv preprint arXiv:2111. 11.0558, 2021.
[27] Jiang, X., Missel, R., Li, Z. & Wang, L. Sequential latent variable models for few-shot high- dimensional time-series forecasting. In The Eleventh International Conference on Learning Representations.
[28] Missel, R. Torch-Neural-SSMM htttps://github.. com/ cww.qu-gg/////ooo..eo/eo......ae qu-gg/ torchssm 2022.
[29] Krishnan, R. G., Shalit, U. & Sontag, D. Deep Kalman filters. arXiv preprint arXiv: 1511,0112, 2015.
[30] 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-55,, 2012.
[31] Pei, F., Ye, J., Zoltowski, D. M., Wu, A., Chowdhury, R. H., Sohn, H., O 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 Neural Information Processing Systems (NeurIPS),
Track on Datasets and Benchmarks, 2021.
[32] Sohn, H., Narain, D. & Jazayeri, N. M. M. Bayesian computation through cortical latent dynamics. Neuron, 103(5):934-947, sep 2019.
[33] Khan, M. F. & Rue, H. The bayesian learning rule. arXiv preprint arXiv:2107.04502, 2021.
[34] Cong, Y., Chen, B. & Zhou, M. Fast simulation of hyperplane-truncated multivariate normal distributions. 2017.
[35] Elfwing, S., Uchibe, & Doya, K. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
[36] Jiang, X., Missel, R., Li, Z. & Wang, L. Sequential latent variable models for few-shot high- dimensional time-series forecasting. In The Eleventh International Conference Learning Representations, 2023.
[37] Li, X., Wong, T.-K. L.,, Chen, & Duvenaud, D. Scalable gradients for stochastic differential equations. In International Conference Artificial Intelligence and Statistics, pages 3870-3888 PMLR, 2020.
12

---

Contents
Nomenclature
14
Variational filtering
14
B.1
Variational predict step
14
Variational update step
15
Variational smoothing step
15
Efficiently sampling structured marginals
16
B.5
Efficient filtering
16
Evaluating the KL using low-rank structure
16
C.1
Smoothing inference network
16
C.2
Causal/streaming inference network
17
Comparison method details
18
DJI SVAE
18
D.2 DVBF
19
D.3 DKF
19
Experimental details
19
High-dimensional linear dynamical system
19
Time complexity
19
Pendulum
19
Bouncing ball
20
MC Maze
E.6
DMFC RSG
22
Collated results
23
Useful expressions
23
Forward KL fixed point
25
Linear Gaussian (information form) smoothing
25
13

---

A Nomenclature
Symbol
Description
SSM
state-space model
LGSSM
linear and Gaussian state-space model variational approximation, π (2()) = p(zz (y:r)
T (2+)
X
natural parameters of IT (z())
mean parameters of i (Ze)
m, / Pt.
mean covariance of (2:)
Oa
local natural parameter update, (yr)
3++1
backward natural parameter update, (y
t (zz)
variational approximation with mean parameters fa = Bm.., (2)))
X
natural parameters iT (zz)
mean parameters of क (2z)
im / Pt
mean / covariance of (z,)
parameters of the dynamies and initial condition
Pe [ze] 11)
prior over state transitions
Po (Z))
prior over initial condition
parameters of the observation model/likelihood observation model/likelihood
Pep (yt (2z)
B Variational filtering
Principles of Bayesian inference make it straightforward to write down algorithm recursively computing the filtering posterior, p(z; y1:1)² Given, p(z;1 F1a-1), updating our belief to p(z y1a)) after observing y; can be broken down into two steps: first, marginalize p(z 1y.e-1) through the dynamics to obtain the predictive distribution,
pize V1+-1) 1pplx-llyu--. (p 1z-1))
(predict step)
(30)
Then, we update our belief by incorporating through Bayes" rule,
plze lyaa) x Pvp 2. ()) y1:-1) (update step)
(31)
With all of the filtered/predictivee beliefs, the smoothing step is given by,
E Eptrasl (2 2++1 12))
p(z ly:r) = Ex(re)) [p(z] 2141, p(zz v1:t)
(32)
Ep(nivi))
(o (Z(+1 2:)]
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 carried out exactly, and for nonlinear non non-Gaussian observations neither can the update step.
Alternatively, by considering variational analogues of the predict / update steps, we can develop recursive and fully differentiable procedure for finding approximations π (2,) = p(z; y1:)). In developing the variational analogues, it is assumed the approximations belong to exponential family of distributions, (i.e. æ E Q where Q is an exponential family distribution) not necessarily Gaussian.
B.1 Variational predict step
Similar to developing recursive algorithm in the exact case, given (Zz-1) p(z--1 Iia- we approximately marginalize π (z.-1) through the dynamics, by solving the following variational (forward KL/ moment-matching) problem,
म (2,) argmii DKL. (Ex(x,))) [Zo 2:11]] क (z:))
(33)
#
So that (zel 2--1) E Q, the optimization problem is minimized when the mean parameters of (z,), denoted fa, set the expected mean parameter transformation under π (2t-1),
P = Ex(z.-)) 0w (z,-1)]
(variational predict step)
(34)
14

---

For LGSSM with (o/z/zze-1 = (F Q), using the fact that,
Fzz-1 12, F + +o))
(e
(35)
means that if (Zt-1) (m, Pr-1), setting the mean and variance of न (Z;) to.
m, Fm,-1
(36)
P, FP, F F + Qo
(37)
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, (z+) prior for the latest observation, Yt, and solve the following variational (reverse KL) problem,
π (z:)
argmin (KL(T(Z)) (y: (z)))
(38)
#EQ
If we denote the natural parameters of (ze) by A, then the optimal A, satisfy the implicit equation M
A par RT,, [l 2:) + X
(variational update step)
(39)
This usually requires iterative optimization procedure, except when the likelihood is conjugate to л (2+) in which case, the likelihood must take the following form with respect to 24.
Ps (yell 21) x exp(T(z.)
(40)
so that, as expected, the natural parameters of the solution are given as Bayes' rule would suggest by adding the data dependent update the natural parameters of the prior so that,
A d + A
(41)
For LGSSM with po(y) = NCzz, R), this results in the following updates,
h; R
(42)
J. J, R-
(43)
which reassuringly recover the information form of the Kalman filter update step 16
B.3 Variational smoothing step
Letting ((zz) s p(zz Via) and q(zx+1) p(z+41 \y:r), we calculate the statistics of q(z,) approximating the smoothed marginal posterior, by minimizing the following objective with respect to the marginal distribution, q(z.).
(ă Dx (?(z.) Ex(m+)) (g(z, 2(+1)l)
(44)
po (zez+ (2x)
IDx (ū(,,) I| ((2.)) Eg(2.)
log E Bzj(z=))
(45)
E (izz) [po
(Ze41 (2
pe (2 Zc++1 (2z)
Dik q(z:)) - By(e.)
log Egitee))
Es(()
(46)
7(2 ((2++1)
taking natural gradients, we find that fixed point, the smoothed marginal posterior parameters satisfy the implicit relationship.
A Xr + V Egize) [A (Ne(ze) + de+1 det1) A (No
(47)
15

---

B.4 Efficiently sampling structured marginals
Algorithm 2 Nonlinear variational filtering
Input: kıT,, K
for f = 1 to T do
m; =S-1 Emo (Z
M = 55-1/2
mm 22
m. ... (z m
T, = Cholesky(Is + Mi Q MM;)
h = P; 'n # using Q. M, Y and Eq. 50
Y Cholesky(I, + K, P,K,
using Q. M.] and Eq. 23
h h; k;
mi = P,h # using M;, K, Y and Eqs. and
w (OO,, I1+5)
2 = p///w, w. # using [Q, Mq] and Eq. 23
w (O, I,)
z m, z
Y (K) z + w;)
end for
1: Output: 1.7 mi., mi.1, 1:T
While P, has a structured representation, drawing samples from (2,) straightforward because we do not have a structured representation for a square-root of P.. However, using the factorization of P, in Eq. (23). it is possible to sample from (O,,PP Py) efficiently since p!/² M, Q1/2). Now, combining the fact that the posterior marginal can be written as,
т (2t) (m, P.)
(m, (P-' KKKK)--
(48)
with the result from Cong et stating that sampling zi (2:) is equivalent to sampling
34
w/ ~ (O, and z ~ NO, P.). and then setting
z m, पे K,Y,Y, (K2 23 w(),
(49)
makes it possible to efficiently draw samples from the posterior marginal.
B.5 Efficient filtering
Evaluating h, P.- m, and MVMs with P can be carried out in O(LS S²) time, after an initial cost of O 5t) to factorize Y, Y.T. (Is M 1 by applying the Woodbury identity to 23
- P--=
P.
Q Oo M, M Qg
(50)
Since this specifies all quantities that characterize मं (2)), following by adding the information from the pseudo observation them,
13 next is to update our belief
h h + k,
P P T K,
(51)
As result, the multiplication in m, P,h, requires OLLS Lr) time, and, analogous (50). the square root factor Y, in
P. P - -.K.TPKKT,TT
(52)
requires 0(r + L.S + 52r) time where Y, r. = (I, + K.P.K,)
C Evaluating the KL using low-rank structure
C.1 Smoothing inference network
Efffccintt training of the generative model and inference networks require efficient numerical evaluation of the ELBO. We take ad advantage of the structured precision matrices arising from low-rank updates
16

---

and sample approximations. The expected log likelihood can be evaluated using a Monte Carlo approximation from samples during the filtering pass. The KL term,
Pr
(Tx((1 (z:)) = 1/2 (m me)
P. (m - m.) + tr(P, P.) + log
L
(53)
P
can be evaluated in closed form and we can expand each term as, Log-determinant. Writing
log Pr
log P. + K,
(54)
log
P:| log K, P.K.
(55)
gives
P Pr|
log
log [ K. P,K,
(56)
2
(57)
Trace Writing,
tr (P,, tr( - K, +
(58)
= - (r, /2 K. (I+ P.K.) KK, p!/3)
(59)
L - tr Y K, ,,K
(60)
which by taking Y2 to be an X r square root such that,
r.r. (I+ K, (. P.K,)
(61)
further simplifies to tr tr(P,'P,) L- L K,Y, M;)
122)
(62)
L - LrM, K,Y, x K. M2) UN( K, Q Q T
(63)
note that the size of the triple product, M; K, Y, is STT.
Causal/streaming inference network C.2
When the real-time parameterization of the inference network is used the expressions become slightly more complicated due the more intricate relationship between the posterior at time / and the posterior predictive at time L.
Log-determinant. We need to first find log P.rrl log [Pr| so begin with using the matrix- determinant lemma to write,
P arrl = [M M + Q
(64)
Is M Q
Mirrl |Q]
(65)
then expand the smoothed covariance to write.
Prl = (P, + B,BBB)
(66)
IP. + B.B.B. 1-
(67)
= (Lrr +B, P.B.l x D-
(68)
[1, + + P. -P.
(69)
and another time to write,
1- + A,
(70)
= I A P.A. x PP-1
(71)
17

---

and another time,
P.l M + +Q
(72)
IIs M Q M x [Q
(73)
When combined we are finally able to write
log P aa log Prl = log Is + M Q 1 log Lrg + B, P.B.
(74)
+ log I,, + A, P. log IIs + M Q M
(75)
For the initial condition we have, log
Pi| log PP = log I, + + P Bi + IL. PAA1
(76)
Trace. For the trace,
tr(P,7P.) = tr(Q ur(Q M (
Mi Q (M
M 1/7
Q-11PPP
(77)
tt(QQ 'P,) tr( urQQ Mi 1 u Y d/2
M T Q P
(78)
tr(P,Q-1) tr(P T.rl T M Mi7 Q-
(79)
where we expand the first rhs term for a numerically efficient implementation by writing
tr(P,Q-)) tr(P tr(P,A,, + A, P.A.) AA,
(80)
tr(P,B,, + P,B, 'B, P,Q
(81)
To reduce notational clutter we define,
T.Y, = (Is + M; Q
(82)
Y 1/ r T 1T (Is + Q MM( )
(83)
Y r = =I... + +A|P,A,) P
(84)
r,Y, = (I,, + B, P.B,)
(85)
and so the trace
tr(P tr(P.- P.) L MM Q Q 'M tr(Q 1/2 P 1/2)
(86)
(I(OQ-1/2,,,,T B. 1/2)
(87)
-tr[P.Q MM Y Tgrl T M Q-11
(88)
which is in form that is easy handle using fast MVMs with Pt, Pt, P..
For the initial condition term we use the fact that P1 is diagonal,
tr tr(P,PTT) = L- tr(Pp//2 A
TTTP! A
tr(P, - B
(89)
D Comparison method details
D.1 SVAE
For the 20 the latent state prior is a linear dynamical system parameterized as,
Po (ze 2;-1) = -N(z, Fz,1, Q)
(90)
Using conjugate potentials, with likelihood p(ý: (2)) = exp(t(z,) (y:)). the approximate posterior given by (ZIrT) III ze (ZLTT) so that its statistics can found by applying Kalman filtering/smoothingg to the sscudo-observation. In this case, the ELBO can be evaluated as
C(g) Σ E. [logp(y: 2t)] Eg. [lg 2:)] + log Egg [p(5, 2x)]
(91)
where qe := (2z,) = pzz 11.--1 is the filtering predictive distribution. These expressions can be evaluated closed form and written concisely terms of natural/mean parameters and fog-partition functions
((g) Σ Eg. [log ply: 1x)] - μ. a + A(², + ax)) A(x))
(92)
where f4, AA + ar + B.+1). Using the identities given App. these expressions can be written in more familiar mcan/covariance parameters.
18

---

D.2 DVBF
For the DVB parameterize the latent state prior using nonlinear dynamical system of the same form as Eq. (T. Then, using an inference network that encodes data in reverse-time to produce the parameters of a diagonal Gaussian distribution, W, g(w,) (m,, diag(s;)), sample
the latent trajectory forward using the recursion, Z; = (z. + Q Q'/²w,. Parameters of the
generative model/inference network learned jointly by minimizing the ELBO,
L(g) Σ Bys(s)) [logP (v)/)))] Ixx (g( (g(we))l
(93)
where, p(w.) -.w.)=N((O.))))
D.3 DKF
For the DKF the latent state is also parameterized using nonlinear dynamical system of the same form Eq. We follow the parameterization outlined in the text with S2S(.) implemented using recurrent neural network. We sample trajectories using the inference network and jointly train all parameters on the ELBO,
(gg) = Σ
Eug. [logply: [2)] Eg
14-1)]lpozz) 2t-1))] S log p(y:() (94)
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 (MIP), with SiLU nonlinearity (5 that gets used repeatedly:
MLPP (rin( Fliddden: nouu)
(Linear(ma) (iin Thidden)) SiLU(), Linear (Thhiden Nour))
((g)
Σ Eyg [ly, 2x]
Egg
(g(z zi-1) pe [z-1))]] s logp(y1:r)
(95)
E.1 High-dimennionall 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 random seeds for two scenarios i) N = L and ii) N = L/5. For each scenario, we also vary the rank the local/backward encoder precision updates.
E.2 Time complexity
We generate random LDS systems of appropriate dimension and measure the average time to complete one forward pass and take gradient step. The system used for benchmarking wall-clock time RTX 4090 with 128GB of RAM with an AMD 5975WX processor.
E.3
Pendulum
We consider the pendulum system from ing/validationn/estingg All methods trained for 5000 epochs for 3 different random seeds. We consider a context window of 50 images and a forecast window of 50 images. decoder was fit from the latent state on the training during the context window; then, for held data, examine performance of the decoder during the context and forecast windows.
26 We generate 500/150/150 trials of length 100 for train-
The generative model is parameterized as:
L =
256
T 50
likelihood
19

---

Py (z)) (z, C+ (Z,) + b, R)
Cy M4, 128, 256)
For each method, inference amortized using the following neural network architectures:
our inference network
Ya =
rg = 4
Oxgi MLP(256, 128, 20)
Ae: (128, Linear(128, 20)]
DKF inference network
-u GRU((188)
(m,, log Pœ): MLP(132, 128,8)
DVBF inference network
ug: [12)([188]]
(Mg. log Φ)): MLP(132, 128,8)
SVAE inference network
ag: MLPP(25,, 128, 20)
Optimization and training details:
optimizer: Adam(Ir 0.001)
batch size: 128
ours(n)
ours(c) time
Figure 4: Learned covariance nonlinear SSMs. (top) single trial of pendulum. Below are the posterior covariances output by the causal and non-causal variants of XFADS.
E.4 Bouncing ball
We consider a bouncing ball dataset commonly used as a baseline to benchmark the performance of inference and learning deep state-space 26 For this dataset we take 500 150/1500111 trials of length 75 for trainingyvalidaaio//eesnnnnn All methods are trained for 5000 epochs for different random seeds. The generative model parameterized as:
L 8
256
T 50
likelihood
-Po(yelz)) (2, (z) + b, R)
C2 MLP(8, 128, 256)
For each method, inference amortized using the following neural network architectures:
our inference network
Ta = 8
rg = 4
Ogi MLP(226, 128, 72)
20

---

3 [GRU(128), Linear 40)]
DKF inference network
Up: GRU((28)
(mø, log P PΦ): 128, 16)
DVBF inference network
-us [GRU(128)]
-(Lo, log Φx): MLP((13,, 128, 16)
* SVAE inference network
agi MLP(256, 128, 72)
Optimization and training details:
optimizer: Alam(Ir 0.001)
batch size: 128
ours(n) I
ours(c)
DVBF
time
Figure 5: Learned covariance of nonlinear SSMs. (top) single trial bouncing ball. Below are the posterior covariances output by the nonlinear SSMs considered qualitatively, we observe more complex covariance structures arise when the ball hits the wall that diagonal approximations cannot capture.
E.5
MC Maze
The monkey reaching dataset 30 was first real dataset examined in the main text. For this dataset, partitioned 1800/200/200 vraining/valddatiooneessin trials sampled at 20ms per bin. All methods are trained for 1000 epochs for different random seeds. The generative model is parameterized as:
7. 40
N 182
T 35
likelihood
(yo /z,) Poisson(z, 2, (2,))) + b)
CC Linear(40, (40, 182)
For each method, inference amortized using the following neural network architectures:
our inference network
Ta 15
ra 5
Oggg: MLP((12, 128, 640)
Bz: (18), Linear(128, 240)]
DKF inference network
upi GRU(128)
lmg, P&) P (128 128, 80)
21

---

DVBF inference network
Us: [GRU(128)]
(m², log P a)): MLPP 128, 80)
SVAE inference network
as: (18P 256, 1640)
Optimization and training details:
optimizer: Alam(Ir 0.001)
batch size: 128
R
0.14
0.44
0.62
0.71
0.77
time to movement onset
-240 ms
-200 ms
-160 ms
-120 ms
-80 ms
true reach
predictions
start
Speed (a.u)
true speed
0 time (ms)
0
model speed
movement
200 ms
onset
Figure 6: XFADS predictive capabilities on real data. On the left are example reaches the monkey made during several trials. Using the learned model from the monkey reaching experiment, filtered neural activity starting from -240ms until -00s before movement onset (each vertical panel represents the result from filtering more and more data) and unroll the final filtered states through the dynamics (starting from the dashed red line) to make predictions.
E.6 DMFC RSG
The second real dataset examined was the timing interval reproduction task of samples at 10ms bins. For this dataset, we partitioned 700/1 700/150/150 traiiing/vliddtiioneestng trials. All methods are trained for 1000 epochs for 3 different random seeds. The generative model is parameterized as:
L 40
. 54
T 130
likelihood
(y |2,) Poisson(z, z, (z,) + b)
Cv Linear(40. 54)
For each method, inference is amortized using the following neural network architectures:
our inference network
ra 15
Tg = 5
Ccgi: MLP(54, 128, 640)
GB: (18), Linear(128, 240)
22

---

DKF inference network
u2= GRU((28)
(mz, log Pœ): M12( 128,80)
DVBF inference network
-u [GRU((288]
(Hg. log aa): MLP(128, 128, 80)
SVAE inference network
ag: M25, 128, 1640)
Optimization and training details:
optimizer: Adam(Ir 0.001)
batch size: 128
E.7
Collated results
pendulum
bouncing ball
monkey reaching
smoothing
prediction
smoothing
prediction
smoothing
prediction
ours (n)
99.7 (.01)
72.7 (3.4)
80.7 (.70)
34.6 (1.5)
88.5 (.1)
52.7 (1.7)
ours (c)
99.6 (.01)
69.7 (3.0)
80.4 (.60)
29.1 (4.7)
87.8 (.1)
53.3 (1.1)
DKF
99.2 (.70)
19.1 (9.5)
65.6 (4.9)
0.32 (10.)
81.2 (.3)
1.7 (9.1)
DVBF
90.7 (2.8)
-19.6 (9.8)
75.3 (.70)
1.3 (.20)
86.3 (.1)
33.4 (.5)
SVAE
98.4 (.40)
-39.7 (10.)
76.5 (2.1)
-23.3 (.70)
87.5 (.1)
-2.4 (13.)
L-SDE
92.1 (2.3)
13.8 (12.)
81.3 (.70)
23.1 (1.3)
76.6 (.4)
18.7 (1.0)
LFADS
91.1 (.1)
23.2 (1.1)
Figure 7: Collated results including the latent SDE (L-SDE) 37
We ran an additional baseline on pendulum/bouncing ball/monkey reaching datasets. L-SDE achieves results with similar level of performance as the other models. Listed are the R² values for decoding angular velocity/x-y position/reach velocity from the latent representations Iearned for each dataset in both smoothing and prediction regimes.
F Useful expressions
mean/natural parameter inner product One common expression that frequently arises the inner product between mean and natural parameter, i.e.
n a
(96)
where in the Gaussian case if the mean/natural parameter coordinates are,
m 1 (P, m.m,)
a.
at
A,A,T
(97)
then,
He Q; m, a
A m.
tr -
(98)
23

---

difference of log partition functions by.
Another common expression that frequently arises is given
A(), + - A((i))
(99)
so that if,
P--mm m
A
(100)
P.r
then,
A(A: ost:)(At)
(101)
- (P, + A.A, A.AT)mi ) log P + A,A
(102)
m Pi p.'mm + 2 log |
2 || 11 | p + ||A, m 1| log I + A A P:A.))
(103)
conditional linear Gaussian log partition in terms of the natural parameters,
For Gaussian distribution, the log-partition function (h, vec(J)), given by
A(A) Th - log [J]
(104)
so that, when
Q- A 0) 0
0 (011)
F
f
(105)
and u = (u, vec(ug) is arbitrary constant, can write ((F((z) + + + u) a [(z(](() + for some a and show this, we can just expand the log partition function
A(Ft(z)+f-u) + 2 A Q -Q- u2) - Az + z
A Q (Q-(0-----__ + u2) -1 u (106)
+ J (Q - + U2) '' log u1 Q
a f(z) +
(107)
then a and b can be identified as
A A Q + +U2) -2)02Luuu u A A 1 'Q-1+ -2) Q--Q--A A) b= JuT =2) -'u log -Q- U2
a =
(108)
Then, for LGSSM, we can evaluate the following expression which describes the difference between predictive and smoothed marginals,
A Ao 00, + V E B((
[A (>o(zz) + Ae+1 )r-1) A (Ao(2.))]
(109)
-Q- + P - - - Betl))
ai =a
(110)
A S
a + B
(111)
where S:11 1
1Q- + +P-11 - -P

---

G Forward KL fixed point
Using the moment matching property of the fixed point solution to the forward KI. objective, we can write
/ (2(+1) Eeg (po(z [2,]] dz
pi+1
(112)
t(ze11)
(113)
/ h(ze)h(2zza) exp (t(z)+1) )o(z)) + t((z)) - A(2,) - A(Xo(zi))) dzp
dz
/ h(zz) exp (t(zi) X A(Ai))
(114)
[f t(zc+1))(z(t)) exp (t(z++1) (ooaz+1)>oozi)) A(No(zi))) dz dz
B,, [zi)
(115)
This result holds for any setting where stochastic transitions, Pt+1/t. and the approximate marginal distributions, gt, are restricted to the same exponential family distribution.
H Linear Gaussian (information form) smoothing
Using the derived relation,
P-1=P-- C CRR- F S
(116)
means that,
P. P
1
(117)
so that using the definition of S, (rewritten for convenience) then plugging in, we get,
S. = Q------'Q---P P Q
(118)
= Q - Q- Q--1 R F S+11
(119)
which gives a recurrence backward in time for S1.7 starting from ST+1 = 0. Now since,
h, - h4 = F F
Q
( +P ( - -h+11)
(120)
and,
Q QS.Q
(121)
have that,
F
- =
R F s(+1)
(122)
which simplifies into a recursion for $1.7 backward time starting from ST+1 0, given by.
= Q -=-0(
S. FC R + F FT s(41)
(123)
= (I- -I- y44 + FT s())
(124)
25

---

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: try to clearly state 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 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 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? [Yes]
Answer:
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 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, c.g., if the approach was only tested few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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, 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 impor- tant 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 of assumptions and a complete (and correct) proof?
Answer: [Ys]]
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 cross- referenced.
26

---

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 they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' 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 ex- perimental results the paper to the extent that 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, 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 dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to reproduce the model (e.g., with an open- 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 welcome 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., 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 instruc- tions 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 NeurlPS code and data submission guidelines (https://nips.ce)
for more details.
27

---

reproducible, they While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
htt.://bbbbcgg.iiiiio/o/e/esebaa Coo///ooeeSSeeaaaeee wwwesssooolliooooooo for more details.
The authors should provide instructions on data access and preparation, including how 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 subset of experiments 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 supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetalls
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary 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 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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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, confi- dence 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).
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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful 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 the text how they were calculated and reference the corresponding figures tables in the text.
28

---

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 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 well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics attere///earriis.cc/pllll///eeie
Answer: [Yes]
Justification: we adhere to the code of ethics. Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
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 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 direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On 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 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 unintentional) misuse of the technology.
If there 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 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 models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA]
29

---

Justification: the paper poses no such risks Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 the asset used and, 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 released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's '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 the dataset/code/moodl 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 anonymize your assets (if applicable). You 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:
30

---

The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included 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 equivalent approval/review based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: this paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending 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, not include any information that would break anonymity (if applicable), such the institution conducting the review.
31


Reviews:
[{"review_id": "RMCrQMK86x", "reviewer": "Reviewer_e1Je", "summary": "The paper presents a novel state space model (SSM) framework for learning dynamical systems with nonlinear latent dynamics. The proposed method borrows inspirations from structured variational autoencoders (SVAEs) and sample-based nonlinear Bayesian filtering, using low-rank corrections to the prior to capture information from the observations and compute approximations to the variational posterior. This framework addresses limitations in current methods which fall short on model expressivity or predictive power.", "questions": "- 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?", "limitations": "The authors have adequately addressed the limitations of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This paper builds upon previous work very well while addressing their weaknesses adequately. The authors show a high level of technical understanding regarding Bayesian filtering and variational inference, and from what I have read the methods proposed are technically sound (I did not go over all of the math in the appendix). As far as I’m aware, the proposal of a structured inference framework for exponential family SSMs that allows for learned non-linear dynamics is a good contribution that fills a gap in the literature. The paper is also professionally written and contains few mistakes or typos. Overall a good paper with solid contributions.", "weaknesses": "- Main technical ideas are conveyed in a convoluted way. The main paper is unnecessarily dense, making it hard to read and understand. One suggestion is to keep only the core equations (such as the objective (21)(22), the low rank parameterization (24), etc.) and provide more high-level discussions on the intuition behind the model. For example, the author could talk about the relationship between the pseudo-observations and the potentials in the SVAE paper, and give intuition on how to understand these quantities. Technical details such as the sample-based approximate inference and the corresponding low-rank updates are pretty standard in my opinion and should be left in the appendix.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time and effort and asking questions that we feel have helped make the manuscript stronger.  Below we respond to them in order,\n\n>   One suggestion is to keep only the core equations (such as the objective (21)(22), the low rank parameterization (24), etc.) and provide more high-level discussions on the intuition behind the model. The author could talk about the relationship between the pseudo-observations and the potentials in the SVAE paper, and give intuition on how to understand these quantities. \n\nThank you for this suggestion — we understand the paper is dense and building more intuition will benefit the reader. Motivated by your suggestion along with those of reviewer UuWx and others, we are moving some algorithmic details from section 3 to the Appendix in favor of higher level intuitions and clarity of the experimental results. \n\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.\n\nThank you for bringing up other approximate filtering methods.  This led to the discovery that the approximate predict step of Eq(23), when dynamics are nonlinear Gaussian, coincides with the statistically linearized Kalman filter [1]  Thanks to your suggestion, we add an additional line after introducing the approximate filter to clarify this connection: “...otherwise, Monte Carlo integration can provide a differentiable approximation; for nonlinear and Gaussian dynamics this leads to a predict step equivalent to the statistically linearized Kalman filter [1].”\n\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\nThese are great suggestions and as per the earlier remark, we are moving non-essential algorithmic details into the appendix.  We will be using the extra space to help build extra intuition – for example, elaborating on Fig.2(a,b) and emphasizing the favorable scaling illustrated in the figure.  We have also run experiments using the latent SDE of [3] as an additional point of comparison(Fig.1 rebuttal PDF).\n\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\nThanks for this question -- it has motivated us to expand on this in the discussion section as well as add other methods wallclock time to Fig.1a for reference .  To answer: compared to methods like DKF/DVBF that use diagonal approximation there is additional overhead since our method scales $\\mathcal{O}(LSr)$ per step, however, there are significant savings compared to SVAE which scales $\\mathcal{O}(L^3)$ per step due to its dense marginal covariances.\n\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\nIt is true that the pendulum closely resembles an LDS, but in this case, the small angle approximation does not hold [2] which would make long-term forecasting difficult for a model using a linear dynamical system prior.  This is one of the reasons that all methods perform very well in the smoothing metric but not prediction metric; especially SVAE where the linear dynamics approximation degrades as the horizon increases.\n\n> Can this method be efficiently parallelized?\n\nThings can be easily parallelized across batches, but to parallelize across time would require a different inference scheme (since we sequentially propagate samples through the dynamics).  However, one possibility is to use a parallelizable inference network architecture (such as S4 [4]), to produce marginals and train it as usual using Eq.(22)\n\n> How well do the low-rank update scale to higher-dimensional systems? I do wish that more in-depth comparisons and discussions can be made. For example, how well do the low-rank update scale to higher-dimensional systems? Comparisons like this gives the reader a better sense of the tradeoffs of the method compared to others in the literature\n\nThank you for these questions – we aimed to elucidate some aspects through conducting the experiments featured in the left half of figure 2.   Motivated by your previous comment, we will now have more room to explain the significance of that experiment in showing: i) purposely exploiting covariance structures makes it possible to develop a filtering algorithm scaling linear in the latent dimension, L — compared to the O(L^3) cost per step incurred for typical Gaussian filters.   and ii) how low-rank covariance parameterizations can also lead to a tighter ELBO.\n\nWe would like to thank the reviewer again for their time and useful suggestions that will serve to make our manuscript stronger.  We hope that if we have addressed your concerns you might raise your score to reflect that, and we are eager to address any concerns or questions that might remain.\n\n[1] Sarkka, Bayesian filtering and smoothing.\n[2] Zhao and Lidnerman, ICML 2023. Revisiting Structured Variational Autoencoders\n[3] Li et al, 2021. Scalable gradients for stochastic differential equations\n[4] Gu et al, 2021. Efficiently Modeling Long Sequences with Structured State Spaces", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HSRDIAi108", "reviewer": "Reviewer_S7ai", "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.", "questions": "How easily extendable is the method to non-gaussian state-space models? More discussion on general methodology? \nWhat do the learned low-rank covariance matrices look like? \nHow 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?", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review the paper and your comments that will undoubtedly increase the quality of our manuscript.\n\n> The paper is dense and notation heavy at times. \n\nWe apologize that the paper could be dense at times.  Motivated by this comment and others, in order to enhance the clarity of the manuscript's main take-aways, we will move non-essential algorithmic details to the appendix in favor of extending the discussions of the experimental results section.\n\n> 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.\n\n\nThank you for this question about causal amortized inference. One of the design choices we made, making it possible to perform causal inference (which we feel has important possible implications in real data-analysis where XFADS can be applied), was segmenting local/backward encoders – in that way filtering can be performed causally using the local encoder; because other sequential VAE models do not make that distinction, we find the ability to perform causal inference a feature of our model that distinguishes it from other approaches.  \n\n> How easily extendable is the method to non-gaussian state-space models?\n\nThank you for the great question! The convenient part of working with general exponential family representations is that the inference algorithm we presented (sans the Gaussian specific parts) is agnostic to the choice of distribution so long as we have a way of evaluating $\\lambda_{\\theta}(z_{t-1})$ and $\\mu_{\\theta}(z_{t-1})$.  \n\n\n> What do the learned low-rank covariance matrices look like? \n\nThank you for the question, we did not include these in the main paper due to space constraints, but we feel it would be beneficial to have some example learned covariances from an experiment.  In the rebuttal PDF, we show an example trial from the bouncing ball experiment for the models that use nonlinear dynamics; its interesting to see how for example, contact with the wall results in spatially complex covariance structures in latent space that cannot be captured by the diagonal approximations.\n\n\n> How does amortization affect the quality of the approximate posterior? What if T is small?\n\nIn the case of smaller datasets, smaller T,  or amortization networks that are not very expressive, the method could still suffer from overfitting/amortization gaps usually associated with VAEs. \n\n> 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?\n\nThank you for pointing this out.  We have added to the discussion section some comments on further limitations – “Furthermore, depending on generative model specifications, such as $L$, while the inference framework of Alg.1 is always applicable, modifications of the message passing procedure in Alg.2 might be necessary to maximize efficiency (e.g. if $L$ is small but $S$ is large).\"  In addition to that, in the related works section we will mention that there will be overhead incurred using low-rank approximations as compared to diagonal ones, but significant savings compared to methods like SVAE that use exact Gaussian message passing.\n\nWe appreciate the time and effort you put into your review.  We hope that if you feel we have addressed the issues raised you might raise your score to reflect that.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ej5Cqcd35y", "reviewer": "Reviewer_UuWx", "summary": "Update 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.", "questions": "Based on what I wrote in weaknesses section, I would like some clarification on my comments and know how the authors plan to improve the clarity of the experimental section", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper introduces a theoretically principled model that is quite flexible and can be used in a wide range of applications. As such I believe it can have impact in the neurips community (as long as the code is released)\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", "weaknesses": "**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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful suggestions.  We apologize about the density of the paper at times and appreciate that you mentioned the possibility of moving some algorithmic discussion to the appendix in exchange for higher level intuitions and more in depth discussion pertaining to the experimental results.  We wholly agree, and as we echoed in the global rebuttal statement we are taking those suggestions to heart and moving some details from the sampling\n\n> Based on what I wrote in weaknesses section, I would like some clarification on my comments and know how the authors plan to improve the clarity of the experimental section\n\nIn addition to what we wrote previously addressing this, motivated by your suggestions, we are restructuring the manuscript to enhance clarity of the experimental section by:\n\n- Displaying figures in the same order which the experiments are presented in the text\n- Explaining the data in more detail, such as the neural data we used to demonstrate model efficacy in monkey reaching and monkey timing (DMFC-RSG) tasks\n- Interweaving more explanation of the figures into the main text so that important details are not only contained in the captions.\n- Moving some algorithmic discussion to the appendix and, in turn, fleshing out the significance of the experimental results and giving more exposition.\n\n\n> Is the format of the citations as a superscript an acceptable one for neurips? It's not very common\n\nThank you for asking – we did consult the guidelines beforehand.\n\n\nWe have to again thank the reviewer for motivating us to move some algorithmic discussion to the appendix in exchange for space that can be spent on clarifying the practical implications demonstrated in the experimental results section.  We hope if we’ve adequately addressed your concerns you might raise your score to reflect that.  We are motivated to continue improving the manuscript, and appreciate all of the input you provided.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "95ios5AeJ0", "reviewer": "Reviewer_ehDs", "summary": "This 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).", "questions": "- what is the difference between predictions from filtering and smoothing in Fig. 3? The last filter state should be identical to the last smoothing state.", "limitations": "Limitations have been addressed to some extent.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- The mentioned limitations of previous work's limitations are not convincing. Fast evaluation of the loss function is hardly the problem that previous DSSM approaches tried to tackle. The main challenges of the VAE-based approaches is that they often only learn to auto-encode and do not properly identify the correct dynamics [1]. Indeed, tackling this problem seems to be a goal of the present work as well, but it is not argued in the introduction.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript.  Below we respond to your suggestions/questions starting with the weaknesses section.\n\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. \n\nThank you for your remark — one of the reasons fast evaluation of the loss and approximate posterior statistics is not a problem for previous DSSM approaches is that they restrict the approximate posterior covariance to be diagonal.  When the diagonal restriction is dropped, then evaluating the KL terms requires linear algebraic operations (log-det, vector multiplication with matrix inverse, trace) that scale cubic with latent dimensionality, $L$.  However, structure in the posterior covariance because of our particular sample approximation and low-rank updates make it possible to carry out those operations in linear time.\n\n> The related work section also seemed to stop past 2016\n\nThank you for pointing this out – we are aware of those works and there are many others as well we would have wished to include in the related works section (see [1a,2a,3a] below to name a few) but were ultimately limited by space.\n\nWe highlighted particular works we did because we aim at general purpose approach with minimum assumptions about the dynamics equipped a scalable inference algorithm; we're not focusing on arguing for any apriori latent structures or dynamics model that are tailored for particular problems. We demonstrated with simple yet general networks to showcase that with a good inference method you can still achieve sota results.  Motivated by your comment, we will slightly condense the current related works section so that we can briefly touch upon recent works such as [1-5 and 1a-3a]\n\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}.\n\nWe apologize that space was tight and we couldn’t add a few more lines in the main text, but the expectations for Eq(6) do contain the two slice marginals.  Factorizing the variational posterior as $q(z_{1:T}) = q(z_1) \\prod q(z_t \\mid z_{t-1})$, we can write the standard ELBO as follows to arrive at Eq(6) involving the expected KL.\n\n$$  \\mathcal{L}(q) = \\int q(z_{1:T}) \\log \\frac{p(y_{1:T}, z_{1:T})}{q(z_{1:T})}  \\, dz_{1:T}$$\n$$= \\sum \\int q(z_t) \\log p(y_t\\mid z_t) \\, dz_t - \\sum \\int q(z_{1:T}) \\log \\frac{q(z_t\\mid z_{t-1})}{p(z_t\\mid z_{t-1})}  \\, dz_{1:T}$$\n$$= \\sum \\int q(z_t) \\log p(y_t\\mid z_t) \\, dz_t - \\sum \\int q(z_{t-1}) q(z_t\\mid z_{t-1}) \\log \\frac{q(z_t\\mid z_{t-1})}{p(z_t\\mid z_{t-1})}  \\, dz_{t-1, t}$$\n$$= \\sum \\mathbb{E}\\_{q_t}   \\left[ \\log p(y_t \\mid z_t) \\right] -  \\sum \\mathbb{E}\\_{q_{t-1}}  \\left[ \\mathbb{D}\\_{\\text{KL}}(q(z_t\\mid z_{t-1}) || p(z_t \\mid z_{t-1}) )\\right]] $$\n\nFor the same smoothing objective in an alternative paper see [4a] Eq.(6).\n\n> What exactly is the pseudo-observation \\tilde{y}? \n\nThe pseudo observation $\\tilde{y}\\_t$ encodes current/future observations, $y_{t:T}$, into a Gaussian potential.  Since it is a Gaussian potential, it is specified by parameters that interact linearly with the sufficient statistics – the form of those parameterizations are given explicitly by Eq (25) in the main text.  In the text, we tried to make this distinction by writing “Importantly, pseudo-observations deﬁned this way encode the current and future observations of the raw data – an essential component for transforming the statistical smoothing problem into an alternative ﬁltering problem.\"\n\n> line 138: I can not follow this argument. Why is smoothing turned into a filtering problem?  It becomes more confusing that the introduction of a different distribution pi(z_t) is supposed to help. Please, can you elaborate\n\nThank you for the question – we say smoothing is turned into a filtering problem because pseudo observations, $\\tilde{y}\\_t$, encode $y_{t:T}$ into a single Gaussian potential, meaning we can filter pseudo observations to obtain statistics of the smoothed posterior; in other words, we have that $\\pi(z_t) \\approx  p(z_t \\mid \\tilde{y}\\_{1:t}) \\approx p(z_t \\mid y_{1:T})$ .  The reason that we introduce $\\pi(z_t)$ is because we want tractable Gaussian marginal approximations, and $q(z_t \\mid z_{t-1})$ is conditionally Gaussian.\n\n> eq (11): using the sum of (prior and likelihood) is not novel\n\nThank you for pointing this out, but we respectfully disagree that this particular parameterization is not novel.  We were inspired by works like [5a] and also [4], to use this type of parameterization; but in those works the dynamics are conditionally linear and it is straightforward to apply Gaussian message passing and recover the posterior – here the dynamics we consider are arbitrary and nonlinear, which led us to develop the nonlinear filtering procedure given by Eqs.(27) & (28) to make this possible.\n\n> what is the difference between predictions from filtering and smoothing in Fig. 3? \n\nYou are correct the last filtering and smoothing state should be identical.  Our purpose was to show what decoded hand position/latents would look like during filtering in the case of a streaming data setting.\n\n[1a] Ansari et al, 2023. Neural Continuous-Discrete State Space Models for Irregularly-Sampled Time Series. \n[2a] Li et al, 2021. Scalable gradients for stochastic differential equations.\n[3a] Schimmel et al, 2022. iLQR-VAE : control-based learning of input-driven dynamics with applications to neural data.\n[4a] Krishnan et al 2016, Structured Inference Networks for Nonlinear State Space Models\n[5a] Johnson et al 2016, Composing graphical models with neural networks...\n\nThank you again for your thoughtful suggestions and useful comments.  We hope if we have addressed them you might raise your score to reflect that.  We are very motivated to position our paper as best possible, and appreciate all of your helpful input.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
First, we would like to thank all of the reviewers for their time, effort, and helpful comments regarding the submitted manuscript – we feel many of your suggestions have led us to changes and additions that better position the paper.  

Many reviewers were positive about the clarity of writing, but felt that the paper was a bit dense at times and that room taken up by some algorithmic details could be used instead to deliver more insightful discussion in the experimental results section.  We agree with this sentiment and are very happy it was brought to our attention – after reading through all of the reviews, we will make the following changes that will help to elucidate the significance of the experimental results section.

- Some equations from the sample approximation structure will be delegated to the appendix which will help to i) create space for higher level insights and discussion and ii) keep focus of the reader away from these lower level details that might detract from the main message.
- We will use the extra space for extra discussion regarding the results;  for example, Fig.2a/b demonstrate the favorable scaling of our inference algorithm, how low-rank approximations scale, and convergence of the ELBO for different parameter settings – now we are able to drive these points home.
- As an additional point of comparison, we have run pendulum/bouncing-ball/monkey reaching experiments using the latent SDE method of [1] – for a total of 5 latent variable models that we compare against.


In the rebuttal figure, the additional results include:
- Collated results featuring the new results from applying the latent-SDE method to experiments considered in the paper
- Additional analysis of the model learned from the monkey reaching data set, showing the potential impact of having a highly predictive model when applied to motor control or brain computer interfaces. 
- Examples of the learned posterior covariances of the nonlinear state-space methods in comparison with our own on the bouncing ball dataset.

We would also like to take some room to point out what we believe are important technical contributions, and the expected impact this work could have.

## technical contributions

To the best of our knowledge, the state-space variational autoencoding framework developed in this manuscript is the only one to allow the recovery of non-trivial covariance over space while scaling well to higher dimensions when general form (possibly nonlinear) dynamics are used.  Gaussian message passing naively scales O(L^3) per time step, but exploiting the sample approximation and low-rank update structures let us develop an O(LSr) complexity algorithm – as a further point, infinite horizon Kalman filters (which assume steady-state covariance) have been classically used when L becomes large and still scale O(L^2).  By being able to filter in O(LSr) time per step, we open the door for high-dimensional filtering with tunable fidelity knobs given by S (# of samples in predict step) and r (rank of update step precision update) without restrictions to trivial covariance structures.

We also view introducing the distinction of specialized local/backward natural parameter encoders in the context of state-space models as a novel contribution for several reasons.  i) This makes it possible to use the causal variant of the model after it has been trained, and deploy it in streaming data settings to perform online filtering with no modifications needed.  ii) As also discussed in the manuscript, this allows missing data to be seamlessly handled as a result of their natural parameter representation. iii) This distinction also makes it possible to avoid parameterizing a neural network for the local encoder when the observation model is conjugate to the dynamics, since the optimal local encoder could be found in closed form as $\alpha_t = \nabla_{\mu_t}\mathbb{E}\_{q_t}[\log p(y_t \mid z_t)]$.  We hope also this distinction can help lead to better architectural design choices for inference networks in the context of amortized inference.

## expected impact

Our algorithm and parameterization of the variational posterior allows for causal inference that is suitable for real-time applications and supports recovery of temporally causal states.  For example, in the manuscript we show how on numerous examples, the inference framework we propose learns models that are much more predictive than other deep state-space models. Applied to the monkey reaching data in Fig.3 we showed how when applied to real neural recordings our learned model was able to successfully *predict* reaching behavior (which it was not trained on).  We support these results further in the rebuttal PDF, showing accurate movement speed predicted from the model long before time of the movement onset.

Thank you again to all of the reviewers for their comments and carefully thought out suggestions.  We hope that we can have a productive dialogue and continue to address any remaining concerns that the reviewers might have and position this work as best as possible.


Best,
The authors

Meta-review: This manuscript presents a method for efficient, structure-exploiting, amortized variational inference in sequential variational autoencoders. The approach is based on a smoothing formulation of the variational posterior, along with some clever approximations that capitalize on the exponential family structure of the variational potentials. Moreover, the results on a variety of real and synthetic datasets look very impressive. 

The math and notation are a bit heavy in places, and I encourage the authors to spend time editing for clarity. There are also a few places where large jumps are taken from one equation to the next, and I think the paper would draw in a broader audience if the main text was more approachable. Overall, though, I think this is a very nice contribution to the field.
Decision: Accept (poster)