Real-time variational method for learning neural trajectory and its dynamics
Authors: Matthew Dowling
Yuan Zhao
Il Memming Park
Venue: ICLR 2023 notable top 25% (2023)
Keywords: neural dynamics
neural trajectory
online variational inference

Abstract: Latent variable models have become instrumental in computational neuroscience for reasoning about neural computation.  This has fostered the development of powerful offline algorithms for extracting latent neural trajectories from neural recordings.  However, despite the potential of real-time alternatives to give immediate feedback to experimentalists, and enhance experimental design, they have received markedly less attention.  In this work, we introduce the exponential family variational Kalman filter (eVKF), an online recursive Bayesian method aimed at inferring latent trajectories while simultaneously learning the dynamical system generating them.  eVKF works for arbitrary likelihoods and utilizes the constant base measure exponential family to model the latent state stochasticity. We derive a closed-form variational analog to the predict step of the Kalman filter which leads to a provably tighter bound on the ELBO compared to another online variational method. We validate our method on synthetic and real-world data, and, notably, show that it achieves competitive performance.

Published as a conference paper at ICLR 2023
REAL-TIME VARIATIONAL METHOD FOR LEARNING NEURAL TRAJECTORY AND ITS DYNAMICS
Matthew Dowling
Yuan Zhao
Stony Brook University, New York, USA ma matthew.. dowl htwebroobooooo..ooo.oooo.o..oo.. edu
National Institute of Mental Health, USA yuan zhao@n gov
Il Memming Park
Champalimaud Research, Champalimaud Foundation, Portugal memming. par.e pattaaapalmaadd. org
ABSTRACT
Latent variable models have become instrumental in computational neuroscience for reasoning about neural computation. This has fostered the development of powerful offinn algorithms for extracting latent neural trajectories from neural recordings. However, despite the potential of real time alternatives to give imme- diate feedback to experimentalists, and enhance experimental design, they have received markedly less attention. In this work, we introduce the exponential family variational Kalman filter (KF)),, online recursive Bayesian method aimed inferring latent trajectories while simultaneously learning the dynamical system generating them. eVKF works for arbitrary likelihoods and utilizes the constant base measure exponential family to model the latent state stochasticity. We derive closed-fone variational analogue the predict step of the Kalman filter which leads to provably tighter bound on the ELBO compared to another online vari- ational method. We validate our method on synthetic and real-world data, and, notably, show that it achieves competitive performance.
INTRODUCTION
Population of neurons, especially in higher-order perceptual and motor cortices, show coordinated pattern of activity constrained to an approximately low dimensional 'neural manifold' (Sohn et al.) 2019; Churchland et al.) 2012 Saxena et al., 2022). The dynamical structure of latent trajectories evolving along the neural manifold is thought to be a valid substrate of neural computation. This idea has fostered extensive experimental studies and the development of computational methods to extract these trajectories directly from electrophysiological recordings. Great strides have been made in developing computational tools for the purpose of extracting latent neural trajectories in post hoc neural data analysis.
However, while recently developed tools have proven their efficacy in accurately inferring latent neural trajectories (Pandarinath et al. (2018) Pei et al. 2021 Yu et al. 2009 Zhao & Park (017).
learning their underlying dynamics has received markedly less attention. Furthermore, even less focus has been placed on real-time methods that allow for online learning of neural trajectories and their underlying dynamics. Real-time learning of neural dynamics would facilitate more efficient experimental design, and increase the capability of closed-loop systems where an accurate picture of the dynamical landscape leads to more precise predictions (Peixoto et al. 2021 Bolus et al., 202T).
In this work, we consider the problem of inferring latent trajectories while simultaneously learning the dynamical system generating them in an online fashion. We introduce the exponential family variational Kalman filter (eVKF), novel variational inference scheme that draws inspiration from the 'predict' and 'update" steps used in the classic Kalman filter (Anderson & Moore 1979). We theoretically justify variational inference scheme by proving it leads to tighter "filtering' evidence lower bound (ELBO) than a 'single step' approximation that utilizes the closed form solution of the proposed 'variational prediction' step. Finally, show how parameterization of the dynamics

---

Published as a conference paper at ICLR 2023
via a universal function approximator in tandem with exponential family properties facilitates an alternative optimization procedure for learning the generative model.
Our contributions are as follows: (i) We propose novel variational inference scheme for online learning analogous to the predict and update steps of the Kalman filter. (ii) We show the variational prediction step offers closed form solution when we restrict our variational approximations to constant base measure exponential families (Theorem [ (iii) We justify our two step procedure by showing that we achieve tighter bound on the ELBO, when compared to directly finding variational approximation to the filtering distribution (Theorem 2). (iv) We show that when using universal function approximators for modeling the dynamics, we can optimize our model of the dynamics without propagating gradients through the ELBO as is typically done in variational expectation maximization (vEM) or variational autoencoders (VARs) (Kingma & Welling 2014).
BACKGROUND
2.1 STATE-SPACE MODELS
In this paper. we consider observations (e.g. neural recordings), yı, arriving in a sequential fashion. It is assumed these observations depend directly latent Markov process (e.g. structured neural dynamics), 2t, allowing us to write the generative model in state-space form:
Ze Zt-1 1 po (Zz 2t-1)
(latent dynamics model)
Yr Zr ~ Pep (y z:)
(observation model)
where 2, E RR, E IR,, &b parameterize the observation model, and 0 parameterize the dynamics model. After observing Ve, any statistical quantities of interest related z, can be computed from the filtering distribution, p(z: y1:1). Since we are considering periodically sampled data streaming setting, it is important that we are able to compute p(zz Y1:t) in a recursive fashion, with constant time and space complexity.
In addition to inferring the filtering distribution over latent states, we will also be interested in learning the dynamics as the (prior) conditional probability distribution, pe 2z-1), which captures the underlying dynamical law that governs the latent state z and may implement neural computation. Learning the dynamics facilitates higher quality inference of the latent state, accurate forecasting, and generation of new data. In this paper we will be focused mainly on models where the dynamics are non-linear and parameterized by flexible function approximators. For example, may model the dynamics as 2t 2t-1 ~ (ff Q), with fg RL د R2 parameterized by a neural network.
2.2 KALMAN FILTER
Before diving into the general case, let's revisit the well-established Kalman filter (Särkkä 2013). Given linear Gaussian dynamics and observations, the state-space model description is given by
(e 2(-1) (z, Azz Q)
0= {A Q}
Pup (yz z2) = -y, Cz, + b, R.)
+b = (C, R]
The Kalman filter recursively computes the Bayes optimal estimate of the latent state z, Given the filtering posterior of previous time step, p(z;-1 V1a-1) = P(mm),,Pt-.). we first predict the latent state distribution (a. the filtering prior) time
1/ yia = Bp(a [Po (z, 2-11)]
(1)
=((z Am; AP AP.IA + QQ
(2)
Secondly, we update our belief of the current state with the observation y, by Bayes' rule
pize y1:)) x p(y: zz) p(z, yıı-1) N(z, m, P.)
(3)
In order to learn the underlying dynamics A, the linear readout C, state noise Q and observation noise R. the EM algorithm can be employed (Ghahramani & Hinton 1996). a calibrated measure of uncertainty over the model parameters is important, then a prior can be placed over those quantities, and approximate Bayesian methods can be used to find the posterior (Barber & Chiappal 2006).
When the dynamics nonlinear, then approximate Bayesian inference can be used compute the posterior over latent states (Kamthe et al. 2022 Hernandez et al., [0188] Pandarinath et al. 2018). Note that these methods are for learning the parameters in the offline setting.
2

---

Published as a conference paper at ICLR 2023
3 EXPONENTIAL FAMILY VARIATIONAL KALMAN FILTER (EVKF)
When the models are not linear and Gaussian, the filtering prior Eq. ① and filtering distribution Eq. B are often intractable. This is unfortunate since most models of practical interests deviate in one way or another from these linear Gaussian assumptions. Drawing inspiration from the predict and update procedure for recursive Bayesian estimation, we propose the exponential family variational Kalman filter (eVKF), a recursive variational inference procedure for exponential family models that jointly infers latent trajectories and learns their underlying dynamics.
3.1 EXPONENTIAL FAMILY DISTRIBUTIONS
We first take time to recall exponential family distributions, as their theoretical properties make them convenient to work with, especially when performing Bayesian inference. An exponential family
distribution can be written as
p(z)= h(z) exp (A t(z) - A())
(4)
where (z)) is the base measure, À is the natural parameter, (()) is the sufficient statisties, and A(A) is the log- pertt-rttiionn function (Wainwright& & Jordan (008). Many widely used distributions reside in the exponential family; Gaussian distribution, p(z) N ((m, P). for example, has f(z) = [z zz ] A= -- m -PP-'] and h(z) = (2 1/22, Note that the base measure h does not depend on for Gaussian distribution. We hereby call such an exponential family distribution constant base measure if its base measure, h, is constant w.r.t. z. This class encapsulates many well known distributions such the Gaussian, Bernoulli, Beta, and Gamma distributions.
An additional and important fact we use is that, for a minimal) T exponential family distribu- tion, there exists a one-to-one mapping between the natural parameters, X, and the mean pa- rameters, == Ep(x) [[(z]. This mapping is given by z = VxA(A), and its inverse by A V (PE:A())))) (log ((1))), though (p)(:)(())) [log((z; X(x)) is usually intractable (Seeger 2005).
If we have conditional exponential family distribution, (o; 24-1). then the natural parameters of 2z 2 are function of 2 1- In this case, we can write the conditional density function
Pe Zz. Zt- 1) h2,)) exp(Ao (2.-1) T ((z.)
(5)
where (-(+) maps - i to the space of valid natural parameters for z, This allows us to use expressive natural parameter mappings, while keeping the conditional distribution in the constant base measure exponential family.
Assume that time f, we have an approximation to the filtering distribution, q(z,-1), and that this approximation is a constant base measure exponential family distribution so that
(6) The primary goal of filtering is to efficiently compute good approximation q(z)) of p(z: y1:e), the filtering distribution at time l. As we will show, following the two-step variational prescription of, predict and then update, leads natural variational inference scheme and provably tighter ELBO than typical single-step variational approximation.
p Y1:-1) q(z,-1) h cxp(A t(z:1)) A()))
3.2 VARIATIONAL PREDICTION STEP
Now that we have relaxed the linear Gaussian assumption, the first problem we encounter is computing the predictive distribution (a.k.a. filtering prior)
fzz y1a-1) = Ep(xs ilyi. 1) [po | 11-
This is generally intractable, since the filtering distribution, p(z;-1 Y14-1), can only be found analytically for simple SSMs. Similar to other online variational methods (Marino et al.) [2088] Zhao
Park 2020 Campbell al. (02I). ). we substitute an approximation for the filtering distribution.
q( 22
p 1 Y1 1), and consider
B 1) [po (z: 1
(8)
minimality means that all sufficient statistics are linearly independent.
3

---

Published as a conference paper at ICLR 2023
Unfortunately, due to the nonlinearity in po(z;
z(-1). Eq. (8 is still intractable, making further approximation necessary. We begin by considering an approximation, q(z;), restricted to a minimal exponential family distribution with natural parameter X, i.c.
B 1) (Zo z.-1)) ) = ((z)) h exp t(z,) A()))
(9)
Taking a variational approach (Hoffman al., 2013), our goal is to find the natural parameter À that minimizes DRL. (ū(z+) ( [Po (z: 1t) Since this quantity cannot be minimized directly,
we can consider the following upper bound:
F = ==(((((z)) Eg(x))
pezz |z-1)] 2 DKL (q(Z)) [(( (p(ze ((-1))) (10)
Rather than minimizing JF with respect to À through numerical optimization, if we take q(z--)) q(z)), and (z z zz-1) to be in the same constant base measure exponential family, then we can show the following theorem which tells us how to compute the X that minimizes F.
Theorem 1 (Variational prediction distribution). If (e Zt-1), g(Zi-1), and q(Zz) are chosen to be in the same minimal and constant base measure exponential family distribution, & then (Z+) argminges. F()) has closed form
form solution given by q (Za) with natural parameters, X
Xg Eg 11 (70
(11)
Eq. ([I) demonstrates that the optimal natural parameters of are the expected natural parameters of the prior dynamics under the variational filtering posterior. While Xg cannot be found analytically, computing Monte-Carlo approximation is simple; we only have to draw samples from ((zz-1) 1) and then pass those samples through (-(.) This also reveals very nice symmetry that exists between closed form conjugate Bayesian updates and variationally inferring the prediction distribution. In the former case we calculate (2)2).)) [po (z: 2t11) while the latter we calculate E [>o 1)]
We summarize the eVKF two-step procedure in Algorithm T located in Appendix E.3
3.3 VARIATIONAL UPDATE STEP
Analogous to the Kalman filter, we update our belief of the latent state after observing yt. When the likelihood is conjugate to the filtering prior, can calculate Bayesian update in closed form by using (Z;) as our prior and computing p(z: y1:) Re (Z,, xx ply: z )z)()) where q(z)), with natural parameter A, belongs to the same family g(z,-1). In the absence of conjugacy, we variational inference to find ((x;) by maximizing the evidence lower bound (ELBO)
X argmax (, 0) = argmax [Eg(..) [logp(y: 24]] - (q(,, {)]l((Z)))]
(12)
If the likelihood happens be an exponential family family distribution, then one way maximize Eq. (12) is through conjugate computation variational inference (CVI) (Khan & Lin 2017). CVI is appealing in this case because it is equivalent to natural gradient descent, and thus converges faster, and conveniently it operates in the natural parameter space that we are already working in.
3.4 TIGHT LOWER BOUND BY THE PREDICT-UPDAS PROCEDURE
natural alternative to the variational predict then update procedure prescribed is to directly find a variational approximation to the filtering distribution. One way is to substitute 1))(,.-))PP(()) 4
for q(z,) into the ELBO earlier (Maiino et al., 2018 Zhao & Park 2020). Further details are provided in Appendix B. but after making this substitution and invoking Jensen's inequality we get the following lower bound on the -ag-inaa likelihood at time
Mr = Ey(,,) [logp(y: 2c)] Eg((.)) [log g(z)) - Egize-.) [log p(z x=-1]]]
(13)
However, prove in Appendix B this leads provably looser bound on the evidence compared to eVKF, as we state in the following theorem.
Theorem 2 (Tightness of L.). If we set
A(1) = Li(g) - Me(g)
(14)
then, we have that
A(g) = Eyeg 1) [A(N0(Z(-1)] A(do) 2 0.
(15)
so that
logp(y.) Lx(g)= Mi((g)
(16)
4

---

Published as a conference paper at ICLR 2023
In other words, the bound on the evidence when using the variational predict then update procedure is always tighter than the one step procedure. Thus, not only do the variational predict then update steps simplify computations, and make leveraging conjugacy possible, they also facilitate better approximation to the posterior filtering distribution.
3.5 LEARNING THE DYNAMICS
Our remaining desiderata is the ability to learn the parameters of the dynamics model pe (Zi Z Z(-1). One way of learning 0, is to use variational expectation maximization; with X fixed, we find the 0* that maximizes the ELBO
0- argmax L(X" .)
(17)
0
argmin Dx (g(ze: (o) Ae))
(18)
0
This objective may require expensive computation in practice, e.g. the log-determinant and Cholesky decomposition for Gaussian I and (o- However, since we chose to and to be in the same exponential family, then described in the following Proposition, tractable square loss function as an optimization objective.
consider the more computationally
Proposition (Optimal 0). If the mapping from Zt-1 to the natural parameters of zu given by Ao do (Z),,,,, a universal function approximator with trainable parameters, then setting
0- = argmin
2 =||**
Xee1?
(19)
is equivalent to finding 0- argmaxg O).
The proposition indicates that we find the optimal 0" that matches the natural parameters of predictive distribution to that the filtering distribution. The proof can found Appendix [C Empirically, we have found that even for small neural networks, following Eq. (19), works better in practice than directly minimizing the KL term.
3.6 CORRECTING FOR THE UNDERESTIMATION OF VARIANCE
It might be illuminating to take linear and Gaussian dynamical system, and compare the variational approximation of KF to the closed form solutions given by Kalman filtering. Given (z,2, z = [N Az; Q), the mapping from realization of Z;-1 to the natural parameters of ze is given by (z;- = 1 A
vcc(Q Q With this mapping, we can determine, in closed form, the prediction distribution given by eVKF, Assuming that q(z--1) = A((z, m.-1. P.-1). we can find the optimal variational prediction distribution by plugging Ao (o(zz11). into
Eq. 11 to find
(2:) = Nzz 1m Q)
(20)
However, we know that the prediction step of the Kalman filter returns
p(z,) = (2; Am,-1, Q AP:-IAT)
(21)
Though this issue has been examined when applying VI time series models, in Turner & Sahani (2011). it demonstrates that eVKF underestimates the true variance by an amount A For this example, we see that because the second natural parameter does not depend on least second order moments of 2,-1, the uncertainty provided by P,-1 will not be propagated forward. At least for the linear and Gaussian case, we can correct this with a post-hoc fix by adding AP--AA to the variance of the variational prediction. If we consider nonlinear Gaussian dynamics with pz(z: 1 = (Z, mo(2;-1), Q), then there does not exist an exact correction since the true prediction distribution will not be Gaussian. Empirically, we have found that adding an extended Kalman filter (Särkkä 2013) like correction of M,- M to the prediction distribution variance, where M,-1 Vmm (m,-1), helps to avoid overconfidence. In the Appendix E.4 we show a case where not including an additional variance term gives unsatisfactory results when dynamical transitions Gamma distributed.
5

---

Published as a conference paper at ICLR 2023
4 RELATED WORKS
Classic recursive Bayesian methods such as the particle filter (PF), extended Kalman filter (EKF), and unscented Kalman filter (UKF) are widely used for online state-estimation (Särkkä 2013 Typically, these methods assume a known generative model, but unknown parameters can also be learned by including them through expectation maximization (EM), or dual filtering (Haykin 2002 Wan Van Der Merwe [2000 Wan & Nelson 1997). While the PF can be used to learn the parameters of the dynamies in an online fashion, as in Kantas et al. (2015), it suffers from the well known issue of "weight degeneracy" limiting its applicability to low dimensional systems. While methods from the subspace identification literature are frequently employed to estimate the underlying dynamics in an offline setting, they are often limited to linear systems (Buesing et al. 2012).
Marino et al. 2018) and Zhao & Park (220) (VJF), in contrast to eVKF, perform a single step approximation each time instant, which leads to provably looser bound the ELBO as stated in Theorem 2 Zhao et al. (S02) use particle filtering to infer the filtering distribution and derives a surrogate ELBO for parameter learning, but because of weight degeneracy it is hard to scale this method to higher dimensional SSMs. Campbell et al. 2021(OVS) use a backward factorization of the joint posterior. Note that it updates the second most recent state with the most recent observation so that it is technically smoothing rather than filtering, furthermore, the computational complexity of this method can be prohibitive in the online setting as is evident from Table 2
EXPERIMENTS
5.1 SYNTHETIC DATA AND PERFORMANCE MEASURES
We first evaluate and compare eVKF to other online variational methods as well as classic filtering methods using synthetic data. Since the ground truth is available for synthetic examples, we can measure the goodness of inferred latent states and learned dynamical system in reference to the true ones.
To measure the filtering performance, we use the temporal average log density of inferred filtering distribution evaluated at the true state trajectory: T T- log q ). where X; are the optimal variational parameters of the approximation of the filtering distribution at time
To assess the learning of the dynamics model, we sample points around the attractor mani- fold, evolve them one step forward, and calculate the KI. divergence to the true dynamics: S E' DxL (po- Z,) [po Z))), where Z are the perturbed samples around the attractor manifold (e.g. stable limit cycle) of the true dynamics, Pa- is the learned distribution over the dynamics, and po is the true distribution over the dynamics. This helps us evaluate the learned dynamics in the vicinity of the attractor where most samples originate from.
The above divergence measures only the local structure of the learned dynamical system. To evaluate the global structure, we employ the Chamfer distance (Wu et al. 2021)
DCD((S S1 -SI'' XxeS,
min ||x y-lyl + SS2 -1 E min XES,
|||2
(22)
yES,
yES2
where S₁ and S2 are two distinct sets of points. Usually, this metric is used to evaluate the similarity of point clouds. Intuitively, low Chamfer distance would mean that trajectories from the learned dynamics would generate a manifold (point cloud) close to the true dynamics a signature that the attractor structure can be generated. Since the Chamfer distance not symmetric, we symmetrize as (CD(S 52) (Sc 52) + DCD (S2l(S))) and take the logarithm.
Chaotic recurrent neural network dynamics. We first evaluate the filtering performance of eVKF. We consider the chaotic recurrent neural network (CRNN) system used in Campbell et al. 202
Zhao et al. (0022)
pe Z141
zz) = (Zc+1 2, + AT-
(7 W tanh(z.) - zz), Q)
and vary the latent dimensionality. Since we restrict ourselves to filtering, we fix the model parameters at their true values. In addition to the online variational methods, we also include classical filtering algorithms: ensemble Kalman filter (enKF) and bootstrap particle filter (BPF) (Douc et al. 2014).

---

Published as a conference paper at ICLR 2023
METHOD
2
L
L
64
EVKF (OURS)
0.047
6.4e-4
0.150
5.8e-4
0.250 + 1.5e-3
0.450
5.8e-3
OVS
0.103 6.4e-4
0.178
5.8e-4
0.302
1 1.5e-3
0.323 1 1.5e-3
VJF
0.105 2.8e-2
0,288 + 4.0e-2 2
0.400 1 1.le-2
0.711 1 4.4e-
ENKF (1,000)
0.115 1 3.3e
0,437
+ 6.0e
2
0.619
1 8.2e-2
0.620
1 2.8e-2
BPF (10,000)
0.047 + 6.7e-4
0,422 = 9.3e
0.877 = 2.5e-2
1.660 = 4.2e-2
Table 1: RMSEs of state estimation for Chaotic RNN dynamics. We show the mean + one standard deviation (over 10 trials) of latent state RMSEs. The latent dimensionality varies from 2 up 64. Those in the parentheses are the size of ensemble and the number of particles.
A)
D)
True
eVKF
trajectory end
B)
E)
Average RMSE
velocity
slow
fast
C)
.4
eVKF
OVS
.2
SVMC
VJF
2000
time (bins)
Van der Pol oscillator with Poisson observations. A) The filtering distribution inferred by eVKF over time, shading indicates the 95% credible interval. B) Zoomed in view at the beginning observations. We plot the mean, and trajectories evolved from the filtered mean steps ahead using "snapshot" of the dynamics at that time, their ending positions are given by the x ss C) Same as before, but at the ending observations. eVKF has learned the dynamies, leading better filtering capabilities. D) True Van der Pol velocity field compared to the dynamics inferred by eVKF. E) Moving average RMSE of the filtering mean to the true dynamics, averaged over 10 trials, error bars indicate two standard errors.
Figure 1:
Table shows the RMSES (mean + standard deviation over 10 trials of length 250) under increasing latent dimensionality. Surprisingly, eVKF offers competitive performance to the BPF for the 2D case, a regime where the BPF is known to excel. The results show eVKF offers satisfactory results compared to the classic filtering algorithms that OVS performs better the case = 64, however, this is at the cost significantly higher computational complexity, shown in Table 2
well as similar online variational algorithms. We
Learning nonlinear dynamics. In this experiment we evaluate how well eVKF can learn the dynamics of nonlinear system that only have knowledge of through sequential stream of observations Y2, ... and so on. These observations follow Poisson likelihood with intensity given by linear readout of the latent state. For the model of the dynamics we consider noise corrupted Van der Pol oscillator so that the state-space model for this system is given by
Z++1.1 Zt,1 A oe
21412 21
To LA 2t,1) 22,2 2t,1) ae
(23)
Yr 2, Poisson(y; A exp CZz: b))
(24)
where exp(-) is applied element wise, A is the time bin size, and (N, In order to focus on learning the dynamical system, we fix vi- = (C, b} at the true values, and randomly initialize the parameters of the dynamics model so that we can evaluate how well eVKF performs filtering and learning the dynamics. We train each method for 3500 data points, freeze the dynamics model, then infer the filtering posterior for 500 subsequent time steps. In Table 2 we report all metrics in addition to the average time per step for both the Poisson and Gaussian likelihood cases. Figure TE see that eVKF quickly becomes the lowest RMSE filter and remains that way for all 4000 steps.
7

---

Published as a conference paper at ICLR 2023
GAUSSIAN LIKELIHOOD
POISSON LIKELIHOOD
METHON
((g 7
KT.$
log(CHAMPER)
TIME (MS)
log ((x))
KI.4
log(CHAMPER) 4
TIME (MS)
EVKF
1.15
6.817
5.66 = 0.93
104
0.57
7.131
2.30 +: 0.32
13
OVS
-0,92
13.48
7.76 = 0.16
6270
-0.21
9.132
3.76 0.32
4150
VIF
+3.58
134.3
6.61 0.26
-1.24
325.5
3.99 + 0.23
100
SVMC
84.83
5.85 0.39
314
410.2
4,06 1: 0.22
730
Table 2: Metries of inference for Van der Pol dynamics. We report the log-likelihood of the ground truth under the inferred filtering distributions, the KL of one-step transitions, the log symmetric Chamfer distance of trajectories drawn from the learned prior to trajectories realized from the true system, and computation time per time step. SVMC uses 5000 particles.
CONTINUOUS BERNOULLI
GAUSSIAN
METHOD
log ((2)) f
KL 4
log(CHAMFER) 1
log 4(z)) f
log(CHAMPER) +
EVKF
2.01
0.057
-0.19 + 0.25
6.15
2.55 + 0.36
OVS
3.61
45.74 + 16.9
VJF
1.94
2.78
3.22 + 0.50
-20.3
3,43 + 0.13
SVMC (5000)
2.37
3.24 1 0.45
3.66 1 0.22
Table Metries of inference for continuous Bernoulli dynamies. We use both CB and Gaussian approximations for the methods that are applicable. eVKF achieves the highest log-likelinn -eglikklihood of latent trajectories, lowest KL-divergence of the learned dynamics, and lowest Chamfer distance. The downside of using Gaussian approximations is most apparent when we look at the Chamfer distance, which is always worse within each method. Note, do not calculate the KL measure when Gaussian approximations are used.
To examine the computational cost, we report the actual time per step. Note that OVS took multi-fold amount of time per step.
The constant base measure exponential family opens up inter- esting possibilities of modeling dynamics beyond additive, independent, Gaussian state noise. Such dynamics could be bounded (i.e. Gamma dynamics) exist over compact space (i.e. Beta dynam- ies). In this example, we consider nonlinear dynamics that are conditionally continuous Bernoulli (CB) (Loaiza Ganem & Cunningham 2019) distributed, i.e.
Continuous Bernoulli dynamics.
(2 2;) II, (B( fo (ze)s)
p(yn,: (2.) N C 2t, r
(25)
where fo 1] + [0, 1 and
N. We choose factorized variational filtering distribution such that ((z)) = 11. CB(z,, At,)). where Ar,i is the i-th natural parameter at time t. In Fig. 2 we show that eVKF is able to learn accurate representation of the dynamics underlying the observed data. Fig. 2 also demonstrates that a CB prior over the dynamics is able to generate trajectories much
representative of the true data compared to Gaussian approximation. These results show CB dynamics could be a proper modeling choice if a priori the dynamics are known to be compact, and exhibit switching like behavior. In Table 3 we report the performance of e VVFF and the other methods synthetic data generated from the state-space model above when using both CB and Gaussian approximations. Notably, we see the Chamfer metric is lower within each method when using CB approximation, showing that even though the true filtering distribution might not exactly be a CB distribution, it is still a good choice.
5.2 ELECTROPHYSIOL.OGGI
RECORDING DURING REACHING TASK
To evaluate eVKF with real-world neural data, considered clectrophysiological recordings taken from monkey motor cortex during a reaching task (Churchland et al.) 2012)). This dataset has typically been used to evaluate latent variable modeling of neural population activity (Pei et al. 2021).
In each trial of the experiment, a target position is presented to the monkey, after which it must wait a randomized amount of time until "Go" cue, signifying that the monkey should reach toward the target. We first take 250 random trials from the experiment, and use latent states inferred by Gaussian process factor analysis (GPFA) Yu et al. 2009) to pretrain eVKF's model of the dynamics. Then,
8

---

Published as a conference paper at ICLR 2023
A)
B)
predicted
True
eVKF
E
fo ())
fo (z(1))
Ezz fz
(Z(1))
Gauss. CB
mean
true
dim1
dim2
Figure 2: Continuous Bernoulli dynamics. A) Velocity field for both E(z, (Z(2--1) and (z--)) from the synthetically created continuous Bernoulli dynamics, and those inferred by eVKF. We see that in mean there are limit cycle dynamics, but for the states actually saturate at the boundary there have to be strong attractor dynamics in parameter space. B) Inferred filtering distributions when using Gaussian approximations compared continuous Bernoulli approximations: Gaussian distributions are able to infer the latent state well but they cannot generate similar trajectories, as we see from trajectories propagated forward through the learned dynamics (shaded in gray)
A) True
B)
eVKF
C)
GPFA
D)
12
100 ms
30
firing rate (Hz)
20
4
neuron 72
neuron 99
12
10
R 0.433 eVKF
R² 0.449
GPFA
8
heuron 67
heuron 54
Figure 3: A) True hand movements from fixation point to target. B) The hand position given by the velocity that linearly decode using eVKF's inferred firing rates. C) Same as previous, but for GPFA. We see that the R value, and decoded hand positions using eVKF are competitive with GPFA. D) Single trial (thin lines), and condition average (bold lines) firing rates for select neurons and tasks, aligned to the movement onset (demarcated with green dots)
use eVKF perform filtering and update the dynamics model on a disjoint set of 250 trials. order to determine if cVKF learns a useful latent representation, we examine if the velocity of the monkey's 's movement can be linearly decoded using the inferred filtering distribution. In Fig. 33, we show the decoded hand position from the smoothed firing rates inferred by eVKF in parallel to the result of GPFA. eVKF able to achieve competitive performance even though GPFA is a smoothing method. In Fig. BC, we plot the single trial firing rates of some neurons over selected reaching conditions, showing that even for single trials, eVKF can recover firing rates decently.
CONCLUSION
We tackled the problem of inferring latent trajectories and learning the dynamical system generating them in real-time- for Poisson observation, processing took 10 ms per sample. We proposed a novel online recursive variational Bayesian joint filtering method, eVKF, which allows rich and flexible stochastic state transitions from any constant base measure exponential family for arbitrary observation distributions. Our two-step variational procedure is analogous to the Kalman filter, and achieves tighter bound on the ELBO than the previous methods. We demonstrated that eVKF performs par with competitive online variational methods of filtering and parameter learning. For future work, we will focus on extensions to the full exponential family of distributions, characterizing the variance lost in Future work will also incorporate learning the parameters of the likelihood & into eVKF, rather than focusing only the dynamics model parameters and filtering states.
generality, and improving performance as latent dimensionality is scaled up.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
MD and IP were supported by an NSF CAREER Award (IIS-1 185536 and NIH RF1DA056404. YZ was supported in part by the National Institute of Mental Health Intramural Research Program (ZIC-MH002968). We thank the anonymous reviewers for their helpful feedback and comments, and Josue Nassar for helpful suggestions for improving the manuscript.
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Brian D. O. Anderson and John B. Moore. Optimal Filtering. Prentice-Hall, Englewood Cliff', N.J., 1979. ISBN 978-0-13-638122-8
David Barber and Silvia Chiappa. Unified inference for variational bayesian linear gaussian state- space models. In B. Schölkopf, J. Platt, and T. Hoffman (eds.), Advances in Neural Information Processing Systems, volume 19. MIT Press, 2006.
Michael F. Bolus, Adam A. Willats, Christopher J. Rozell, and Garrett B. Stanley. State-space optimal feedback control of optogenetically driven neural activity. Journal of neural engineering. 18(3): 036006, 2021.
Lars Buesinng Jakob H. Macke, and Maneesh Sahani. Spectral learning of linear dynamics from generalised- -ineeer observations with application to neural population data. Advances in neural information processing systems, 25, 2012.
Andrew Campbell, Yuyang Shi, Thomas Rainforth, and Arnaud Doucet. Online variational filtering and parameter learning. Advances in Neural Information Processing Systems, 34:1 14:18633-18645,
2021.
Mark M. Churchland, John P. Cunningham, Matthew T. Kaufman, Justin D. Foster, Paul Nuyujukian, Stephen Ryu, and Krishna V. Shenoy. Neural population dynamics during reaching. Nature, 487 (7405):51-56, 2012.
Randal Douc, Eric Moulines, and David Stoffer. Nonlinear time series: theory, methods and applications with R examples. CRC press, 2014.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Higmoid-weightee linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
Zoubin Ghahramani and Geoffrey E. Hinton. Parameter estimation for linear dynamical systems. Technical report, 1996.
Simon S. Haykin. Adaptive filter theory. Pearson, 2002.
Daniel Hernandez, Antonio Khalil Moretti, Ziqiang Wei, Shreya Saxena, John Cunningham, and Liam Paninski. Nonlinear evolution via pratialllydependern linear dynamics for electrophysiology and calcium data. arXiv preprint
2018.
Matthew D. Hoffman, David M. Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 14(4):1303-134) 2013.
Sanket Kamthe, So Takao, Shakir Mohamed, and Mare Deisenroth. Iterative state estimation in non- linear dynamical systems using approximate expectation propagation. Transactions on Machine Learning Research, 2022.
Nikolas Kantas, Arnaud Doucet, Sumeetpal S. Singh, Jan Maciejowski, and Nicolas Chopin. On particle methods for parameter estimation in state-space models. Statistical science, 30(3)3328-351, 2015.
Mohammad Khan and Wu Lin. Conjugate- Conjugate-Computia Computation Variational Inference Converting Variational Inference in Non-Conjugate Models to Inferences in Conjugate Models. In Aarti Singh and Jerry Zhu (eds.), Proceedings of the 20th International Conference on Artificial Intelligence and Statistics,
volume 54 of Proceedings of Machine Learning Research, pp. 878 887. PMLR, 20-22 Apr 2017.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 1rXiv:1412.6980 2014.
Diederik P Kingma and Max Welling, Auto-Encoding variational bayes. In International Conference on Learning Representation, May 2014.
Gabriel Loaiza-Ganem and John P. Cunningham. The continuous bernoulli: fixing a pervasive error in variational autoencoders. Advances in Neural Information Processing Systems, 32, 2019.

---

Published as a conference paper at ICLR 2023
Joseph Marino, Milan Cvitkovic, and Yisong Yue. A general method for amortizing variational filtering. In Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Chethan Pandarinath, Daniel J. O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufman, Stephen Ryu, Leigh R. Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10)8805-81) 2018.
Felix Pei, Joel Ye, David M. Zoltowski, Anqi Wu, Raeed H. Chowdhury, Hansem Sohn, Joseph E. O'Doherty, Krishna V. Shenoy, Matthew T. Kaufman, Mark Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan Pillow, Il Memming Park, Eva L. Dyer, and Chethan Pandarinath. Neural latents benchmark '21: evaluating latent variable models of neural population activity. In Advances in Neural Information Processing Systems (NeurIPS), Track on Datasets and Benchmarks, 2021.
Diogo Peixoto, Jessica R. Verhein, Roozbeh Kiani, Jonathan Kao, Paul Nuyujukian, Chandramouli Chandrasekaran, Julian Brown, Sania Fong, Stephen I. Ryu, Krishna V. Shenoy, et al. Decoding and perturbing decision states in real time. Nature, 591(7851))604 609, 2021.
Simo Särkkä. Bayesian filtering and smoothing. Cambridge University Press, 2013. ISBN 978110761929.
Shreya Saxena, Abigail A. Russo, John Cunningham, and Mark M. Churchland. Motor cortex activity across movement speeds is predicted by network-level strategies for generating muscle activity. Elife, 11:667620, 2022.
Aaron Schein, Hanna Wallach, and Mingyuan Zhou. Poisson-gamma gamma dynamical systems. Advances in Neural Information Processing Systems, 29, 2016.
Matthias Seeger. Expectation propagation for exponential families. Technical report, 2005.
Hansem Sohn, Devika Narain, and Nicolas Meirhaegheand Mehrdad Jazayeri. Bayesian computation through cortical latent dynamics. Neuron, 103(5):934 947, sep 2019. ISSN 0896-6273. doi: https://doi.org// org/10. 101 101 /011///.eurro.2220 www. w0w.:m...0900666 0.....
R. Turner and M. Sahani. Two problems with variational expectation maximisation for time-series models. In D. Barber, T. Cemgil, and S. Chiappa (eds.), Bayesian Time series models, chapter 5, pp. 109-130. Cambridge University Press, 2011.
Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Leurning, 1(1-2):1-305, 2008. ISSN 1935-8237.
doi: 10.1561/2200000001.
Eric Wan and Alex T. Nelson. Neural dual extended kalman filtering: applications in speech enhancement and monaural blind signal separation. In Neural Networks for Signal Processing VII. Proceedings of the 1997 IEEE Signal Processing Society Workshop, pp. 466 IEEE, 1997.
Eric Wan and Rudolph Van Der Merwe. The unscented kalman filter for nonlinear estimation. In Proceedings of the IEEE 2000 Adaptive Systems for Signal Processing, Communications, and Control Symposium (Cat. No. 00EX373), pp. 153-158. Ieee, 2000.
Tong Wu, Liang Pan, Junzhe Zhang, Tai Wang, Ziwei Liu, and Dahua Lin. Density-aware chamfer distance comprehensive metric for point cloud completion. arXiv preprint arXiv:2111.12702. 2021.
Byron M. Yu, John P. Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V. Shenoy, and Maneesh Sahani. Gaussian- processs factor analysis for dow-dimensional single-trial analysis of neural population activity. Journal of Neurophysiology, 102(1)x614-33,, jul 2009. doi: 10.1152/jn. 90941. 20088. PMID: 19357332.
Yuan Zhao and Il Memming Park. Variational latent gaussian process for recovering single-tria trial dynamics from population spike trains. Neural Computation, 19(5):1293-13 1316, may 2017, doi: 10.1 1 162/neco, a_ 00953.
12

---

Published as a conference paper at ICLR 2023
Yuan Zhao and II Memming Park. Variational online learning of neural dynamics. Front. Comput. Neurosci., 14:71, 1, October 2020.
Yuan Zhao, Josue Nassar, Ian Jordan, Monica Bugallo, and II Memming Park. Streaming variational monte carlo. IEEE Trans. Pattern Anal, Mach. Intell., February 2022.
13

---

Published as a conference paper at ICLR 2023
APPENDIX
A Proof of Theorem
B Proof of Theorem
Proof of Proposition
D Variance Correction for Nonlinear Gaussian Dynamies
E Experimental Details
A PROOF OF THEOREM
Theorem (Variational prediction distribution). If po (t-1), q(z:-1). and g(ze) are chosen to be in the same minimal and constant base measure exponential family distribution, Eo. then q (Z+) = argmin F(ā) has closed solution given by q (zi) with natural parameters, Xg
Xg (9( (2+
(11)
Proof: The upper bound we want to minimize is given by
== -((((z))) )(g [log (z [z=-)]]
(26)
where H(g) is the entropy of q. For exponential family distributions, we recall that the negative entropy coincides with the conjugate dual of the log partition function, or -((gp) A (p) wright & Jordan 2008). Then, we have that,
(Wain
F = -Biçca.)) By(e.-.) (Z ((x)) - A(A(z:1)) - log (22)) + (p)
(27)
-Bgg Eg(n Do (Z)-1) T t(z,) A(X(Z,1)) log h| A
(28)
--- Byg(z--)) (Z + A²)) + constants
(29)
= =- T Xo + A + constants
(30)
where in the first line we use the fact that Ey(z,) [t(z,)] = p. In the second line, we use the fact that Po; 22-1) ) has a constant base measure. In the third line we separate out terms that are constant with respect to Å. In the fourth line
the definition Åg Ey(x.-.) [)o(2-1)]
Since our equation is in terms of the mean parameters of 7, thanks to the minimality of q. we can just consider the optimal variational parameters in their mean parameterization. Then, by considering maximization of -F rather than minimization of F, we can write that the optimal variational parameters satisfy
p argmax [ X (p) a
(31)
Taking derivatives of the right hand side and setting it equal to O, we have that
oo Va- A" A " = 0
(32)
do X = 0
(33)
X = -u where in the second line we use the fact that A = V A (Ht). As stated in Theorem
(34)
we have that
>g as claimed.
B PROOF OF THEOREM 2
Theorem (Tightness of L.). If we set
A(g) L.(9) Me(g)
(14)
then, we have that
A(g) Eg(ze-) [A()o(Z--1)) A(do) 2
(15)
that
logply.) 2 Lx(g) MM(g)
(16)
14

---

Published as a conference paper at ICLR 2023
Proof:
We write the two ELBOs
Le(g) = (z.) log ply: zi) - (
(35)
Me(g) = (z.)) log p(ye Zz) Eq(ze) [log g(t)) - E 1)) logp(z; ==-1)]
(36)
This means the difference, (44) = (g) M2(g), can be written as
(q) ((x),) [log 7(z)) - B)(x [log p(z: 21-1)]l
(a.),) fog ((ze) + Dg ((2z) A(Xo) - Ba((a-1)) [log h(zz) + 20( t(zi)
(Ag Eg(er-1) (Zo Bggaa)) [[(2))] + (Bo(z.-1) (A(A0(2z-11)))) -
Ao
-B((e-.) 1) [A Ao (2=- A(Xa)
Invoking Jensen's 's inequality, and the fact that the bog-partition function is convex in its arguments, we can write that
Bg((e-1) 1) (o(Ao(21-11)) A( (0(2=-1)))
(37)
which means that
A (g) Eug(e._)) (A(20(Z(-1)) A((o)) 20
(38)
as claimed.
C PROOF OF PROPOSITION 1
Proposition (Optimal 0). If the mapping from Zi-1 to the natural parameters of zu given by 10( a universal function approximator with trainable parameters, O, then setting
0 argmin I| 2
(19)
is equivalent finding 0" argmax L 0).
Proof: By the fact that _(X*, 0) = Ey(,.) logp(y, 2e) - 1xx (x(x))))i2))) we have that Vol(>* 0) = Dx 4o (2+)) so that
(oDk((((2;))((tzz) VoEy(en) [t(z,) (A' - 2o) + A(2o)]
(39)
--Voou + Votol A(Xo)
(40)
-ooop' + Ve To
(41)
VVoeol -x)
(42)
whereas, for the alternative objective we have that
Vo Xo 1
(Vo - X²)
(43)
Assume that (-(-) is a flexible enough function approximator that [Vo has full column rank.
Then the gradient the KL. term is O, either [eoèo] is 0, in which case Eq. 43) is O, or (Ao - X*) is O. which by continuity of the mapping from natural to mean parameters implies that (oo "1 is Showing equivalence stationary points of the two objectives.
D VARIANCE CORRECTION FOR NONLINEAR GAUSSIAN DYNAMICS
For the case of nonlinear Gaussian dynamics, we could consider directly linearizing the dynamics in order to forego solving variational problem (e.g. linearizing the dynamics of Eq. [ about the mean of 2 to evaluate the expectation directly). Concretely, consider nonlinear Gaussian dynamics specified via p(z, 2-11) N(z m= Q): assuming have variational approximation
15

---

Published as a conference paper at ICLR 2023
eVKF
OVS
VJF
SVMC
Figure 4: Learned phase portraits and inferred distribution of the latent states for the Van der Pol system with Poisson observations. For comparison, plot the ground truth latent state in black. We see that eVKF infers much smoother latent states than the other methods.
to the filtering distribution at time given by q(zt-1) N, Pt-1), then the prediction step could be approximated as
p(z1 V1:-1) = (x [pa(ze Z(-1)]
(44)
(2)(2-)) [N(z] mz( Q)]
(45)
1 [W(z]] mm + M,- 12 m,,),,
(46)
N(z
+Q)
(47)
= q(z))
(48)
where M,-1 (2
This prediction distribution coincides exactly with the one returned by the extended Kalman filter (Särkkäl 2013), as well as the one prescribed by eVKF. Similar procedures to facilitate tractable inference nonlinear Gaussian models is covered extensively
in Kamthe et al. (2022).
E
EXPERIMENTAL DETAILS
VAN DER POI.
with y 1.5, T1 12 0.1, 0.1. For the Poisson likelihood example, we can take advantage of CVI as mentioned in the main text. For this model, the expected log-likeli lideliod has analytical solution,
Data was generated according to Eq. 23
(x(,.) logp(y: 2.) = Σ Yne (C, m; mC
bb,) b,.) - Ac P.C.C.C..
(49)
To parameterize the dynamics, Po (z, 2t-1), ). we use a single layer MLP with 32 hidden units and SiLU (Elfwing et al. 2018) nonlinearity. During training we use Adam (Kingma & Ba 2014), and update the dynamics every 150 time steps. In total we use 3500 time points for training the dynamics model for all methods. For measuring the time step as in Table 2 the experiments were run computer with an Intel Xeon E5-2690 CPU at 2.60 GHz.
E.2 CONTINUOUS BERNOULLI
For the continuous Bernoulli example, we can take advantage of CVI as mentioned in the main text. For this, we require derivatives of the expected log- ligelio,,,ood, which for a Gaussian likelihood, p(yt, 2 CC, z, In), we have that
E (((x)) log p(y: 2:) = Σ R,
(C Var(z..)) + + [2,]' C, 2yn., C, (2:)
(50)
2r
where C, = C.C.C. We use Adam and update our dynamics model every 100 time steps. For this example, the synthetic data is length 500 sequence.
16

---

Published as a conference paper at ICLR 2023
eVKF
OVS
VJF
SVMC
Gaussiar
MYLAMMLA JUNNY
MALAM WWW.MAAA JUWN AATUMMA MAMAMMAMAM
CB
MIWAMIM WAWW AOTTTTTAAA
M WW
Figure 5: Learned phase portraits and inferred distribution of the latent states for the continuous Bernoulli example. The open source code for running OVS not immediately compatible with non-Gaussian approximations. Top: Phase portraits and filtered latent states when approximations are constrained to be Gaussian. Bottom: Same as top, but for approximations constrained be CB.
E.3 EVKF ALGORITHM
Below present the algorithm for using VVKF perform inference. Instead of updating O every data point, we could accumulate gradients for fixed number of steps, so that the variance of gradient steps reduced.
Algorithm eVKF
Input: yt E R.N. 9 (dynamics parameters)
for each y, or until done do
X; 11) [2)
predict update
X; argmaxx, E (g(x)(A)) [logp(y; 2:] (xl A)lla(ze X0)
l [A, Eggg.-1) (2:
0 + 0,-1 - Vole
end for
E.4 EXAMPLE OF GAMMA DYNAMICS
We consider synthetic data where observations conditionally Poisson and dynamic transitions Gamma distributed, so that the state-space model description is
p(ze+1 2,) Gamma (Ze+1
(nf((z)),,n fz.)) plye 2 Poisson(y: A exp( b))
(51)
(52)
Similar to standard Gaussian dynamical system, under this specification, we have that 1 [2142] = f(z)) like the gamma dynamical system presented in Schein et al. (2016). but unlike that work both and 3 are functions of that the variance is constant. We choose a variational approximation that factors as a product of Gamma distributions so that q(2.) = II Gamma(zi: ai; B.). Since use the canonical link function, the expected log-
17

---

Published as a conference paper at ICLR 2023
variance corrected
variance uncorrected
3
Wmw filtering mean
true
Figure 6: On the left: phase portrait and filtered latent states of system with gamma distributed
transitions using the prediction step variance correction. On the right: same as the left, but without using the prediction step correction.
likelihood can be calculated in closed form since
(x)))) logply: ze) = E,, (C + bn)ye,n eexp(bx)
III
1-1 (1- Cn.1
-aes
(53)
3r,
which allows to use CVI for inference. We take f(-) to be 64 hidden unit MLP with SiLU nonlinearity and softplus output. To see how correct for the underestimation of variance, notation will be less cluttered if our discussion is in terms of means/variances; so, Z1+1 ~ g(z++1) have mean 1 then the corrected variance, $ $0,411 should be equal to the prior transition variance (i.e.
1/bb) plus the correction term so that $4.41 = =1/bo + (st O Vomo,141) where Sy is the standard deviation of z, ~ q(z)). As shown in Figure without this correction, the quality of inference is noticeably worse.
18


Reviews:
[{"review_id": "x1rDgRo_3p", "reviewer": "Reviewer_8dja", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors consider state estimation problems in the setting where the underlying dynamics model is unknown. They introduce the exponential family variational Kalman filter (eVKF), which is a variational inference scheme for state-space systems that consists of a prediction step and an update step, much like the Kalman filter. The approximate posteriors are assumed to be in the exponential family to make computations tractable. In addition, the latent models can be learnt via an EM algorithm in a tractable fashion. They prove that the two-step predict-update inference procedure gives tighter bounds compared to directly finding an approximate filtering distribution using variational inference and show that it works well on several experiments.", "summary_of_the_review": "The paper is of good quality and the methods presented are backed up by rigorous mathematical arguments. Aside from minor issues and the fact that similar ideas exist in the literature, the ideas and experiments are executed well here. The paper definitely deserves to be accepted at ICLR.", "strength_and_weaknesses": "__Strengths:__\n\nOverall the paper is well-presented and the methods are clearly explained. The ability of the model to learn nonlinear dynamics well in an online setting presents exciting possibilities.\n\n__Weaknesses:__\n\nWhile the ideas presented in this paper are neat, the need to correct the filtering distribution in (23) is a slight sore point as it doesn't naturally follow from the derivations. I believe that this may be possible to overcome by minimising (12) directly instead of using the upper bound (11), where (9) can be computed approximately, e.g. by linearising the dynamics model $m_\\theta$ (as done for example in [1]).\n\n[1] Kamthe S, Takao S, Mohamed S, Deisenroth M. Iterative State Estimation in Non-linear Dynamical Systems Using Approximate Expectation Propagation. Transactions on Machine Learning Research. 2022.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "nkajwU-Q_i", "reviewer": "Reviewer_Hbu9", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper considers the problem of optimal filtering and simultaneous parameter estimation for discrete-time systems with exponential family latent dynamics. The authors propose a new variational inference algorithm, by splitting the prediction and update step in the filtering algorithm. The variational approximation for the prediction step is based on an upper-bound minimization and the update step is based on a lower-bound (ELBO) maximization. Parameter learning is performed using variational expectation maximization, which yields an easy objective under the exponential family assumptions.\nThe algorithm is empirically evaluated on three synthetic filtering problems and one real-data example.", "summary_of_the_review": "The paper discusses a new filtering and parameter learning algorithm using variational inference. The paper has convincing results, however, a heuristic is used to improve the results, where it is not very clear to me how and when it should be applied for setups different from those discussed in the paper. ", "strength_and_weaknesses": "Filtering and smoothing are probably one of the most applicable topics for time series data. Therefore, this paper which proposes a new filtering algorithm is of high interest to the community. The developed approach is interesting and novel to me. The algorithm development is easy to follow. The experimental evaluation is carried out nicely and a comparison to standard baselines and some newer state-of-the-art algorithms are given. \nOne thing that I find problematic is the use of the upper bound in the prediction step. This upper bound as described in Sec 3.6. leads to an underestimation of the variance, which is heuristically corrected. This correction step, however, does not follow a principled derivation and as such, it remains unclear if it is sensible to use in every setup. Also note, that in my version the discussion section Appendix D.4 is missing.\nOtherwise, I think this is a solid piece of work.\n\nA question to the authors is: When using the standard ELBO objective in the prediction step, e.g., with stochastic variational inference, is the problem of underestimating the variance still present?\n\nAs an aside: I am not so happy with the name eVKF as the filter does not assume the usual assumption that Kalman-type filters have (some form of linearity and Gaussianity). Maybe, it is a good idea to think about a slightly different name.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "YsTWa3MF9t", "reviewer": "Reviewer_b7sM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes an exponential family variational Kalman filter (eVKF) to model and predict the neural trajectory. The basic idea is to restrict the filtering posterior and latent dynamics to a special exponential family with constant base measure and minimal. The restricted distribution form could lead to efficient filtering and update. The proposed two step procedure is proofed to be able to bring tightness of ELBO. The empirical study shows the effectiveness of the proposed idea. ", "summary_of_the_review": "This paper is with clear motivation and the idea is reasonable with guarantee. The experiments show the effectiveness but without comparing with related works like existing variational Kalman filter. ", "strength_and_weaknesses": "Strength\n\n1. The specific exponential family is adopted to ease the filtering posterior distribution and latent dynamic learning. \n2. The guaranteed tightness ELBO comparing straightforward variational update. \n3. The consistent better performance. \n\nWeakness\n\n1. The specific exponential family, i.e., with constant base measure and minimal, is used without further discussions on the possible side effects. The constant base measure may not a strong restricted condition because many well known distributions are within this class, but what effects from the minimal?  especially for neural trajectory modeling and prediction? \n\n2. There are various works about variational kalman filtering in the literature like [1, 2]. Why these methods are not compared in experiments and even not discussed in related works? \n\n3. eVKF is claimed to be a 'novel variational inference scheme' but it is just a special case of VKF under a specific exponential family. It cannot be considered as a novel one. \n\n\n\n[1] Das, Niladri, Jed A. Duersch, and Thomas A. Catanach. \"Variational Kalman Filtering with Hinf-Based Correction for Robust Bayesian Learning in High Dimensions.\" arXiv preprint arXiv:2204.13089 (2022).\n\n[2] Huang, Yulong, Yonggang Zhang, Peng Shi, and Jonathon Chambers. \"Variational adaptive Kalman filter with Gaussian-inverse-Wishart mixture distribution.\" IEEE Transactions on Automatic Control 66, no. 4 (2020): 1786-1793.\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7hoYdwDskM4", "reviewer": "Reviewer_6sm4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents a variational approach for filtering and estimating the state parameters in an online manner.  In particular, the prediction distribution and filtering distribution are assumed to be within the constant base measure exponential family. With this assumption, many computations benefit from the conjugacy, such as the natural parameter of variational prediction distribution is in closed form, and the learning of parameters for the dynamic model can be simplified. The proposed approach also leads to a tighter lower bound than previous works. Experimental results include both synthetic data, using appropriate performance measures, and a compelling real-world use case. ", "summary_of_the_review": "The writing of the paper is high-quality and the analysis seems very interesting. ", "strength_and_weaknesses": "Pros: \n\nThis paper is very well-written and presented in an organized, self-contained way. \n\nThis paper offers an approximate solution that extends to non-Gaussian non-linear models even in latent spaces, while still staying relatively tractable. \n\nThe evaluation of eVKF with real-world data and comparisons to GPFA is helpful. \n\nCons: \n\nAlthough the online real-time learning aspect is well-motivated, it is unclear what factors contribute to achieving it, in making things possible. What are some of the bottlenecks of previous methods? Is backpropagation really bad for the real-time setting considered? The claim of real-time also needs to consider the data rate. What is the data rate per sample if the processing time is about 10ms per sample?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%