Online Variational Filtering and Parameter Learning
Authors: Andrew Campbell
Yuyang Shi
Tom Rainforth
Arnaud Doucet
Venue: NeurIPS 2021 Oral (2021)
Keywords: Approximate Inference
Variational inference
State-space models
Filtering
Time series
Online learning

Abstract: We present a variational method for online state estimation and parameter learning in state-space models (SSMs), a ubiquitous class of latent variable models for sequential data. As per standard batch variational techniques, we use stochastic gradients to simultaneously optimize a lower bound on the log evidence with respect to both model parameters and a variational approximation of the states' posterior distribution. However, unlike existing approaches, our method is able to operate in an entirely online manner, such that historic observations do not require revisitation after being incorporated and the cost of updates at each time step remains constant, despite the growing dimensionality of the joint posterior distribution of the states. This is achieved by utilizing backward decompositions of this joint posterior distribution and of its variational approximation, combined with Bellman-type recursions for the evidence lower bound and its gradients. We demonstrate the performance of this methodology across several examples, including high-dimensional SSMs and sequential Variational Auto-Encoders.

Online Variational Filtering and Parameter Learning
Andrew Campbell
Yuyang Shi
Tom Rainforth
Arnaud Doucet
Department of Statisties, University of Oxford, UK {campbell, yshi, rainforth, doucet) o ac.uk .u
Abstract
We present variational method for online state estimation and parameter learning in state-space models (SSMs), a ubiquitous class of latent variable models for sequential data. As per standard batch variational techniques, gradients to simultaneously optimize a lower bound on the log evidence with respect both model parameters and variational approximation of the states' posterior distribution. However, unlike existing approaches, our method is able operate in entirely online manner, such that historic observations do not require revisitation after being incorporated and the cost of updates at each time step remains constant, despite the growing dimensionality of the joint posterior distribution of the states. This is achieved by utilizing backward decompositions of this joint posterior distribution and of its variational approximation, combined with Bellman-type recursions for the evidence lower bound and its gradients. We demonstrate the performance of this methodology across several examples, including high-dimensional SSMs and sequential Variational Auto-Encoders.
use stochastic
Introduction
Many tasks in machine learning with time series Suca- as video prediction [16, 27], speech enhancement [35] or robot localization [9, 18, 28]-often need be performed online. Online techniques are also necessary in contexts as diverse as target tracking [4], weather prediction [13] and financial forecasting [42]. popular class of models for these sequential data are SSMs which, when combined with neural network ideas, can also be used to define powerful sequential Variational Auto-Encoders (VAEs); see e.g. [8, 15, 16, 30]. However, performing inference in SSMs challenging problem and approximate inference techniques for such models remain active research area.
Formally, an SSM is described by a latent Markov process and an observation process. Even if the model parameters are assumed known, online inference of the states of the latent process complex problem known as filtering. Standard approximations such as the extended Kalman Filter (KF), ensemble KF, and unscented KF can be used, but only provide an ad hoc Gaussian approximation to the filter [13, 17, 37]. More generally, assumed density filtering techniques [3, 7, 32] can provide other simple parametric family approximations. These approximate filtering methods can be used, in turn, develop online parameter learning procedures by either augmenting the state with the static parameters or using gradient- based approaches. However, such approaches are notoriously unreliable [19]. Particle Filtering (PF) methods, on the other hand, provide a more principled approach for online state and parameter estimation with thcoretical guarantecs [11, 12, 19, 41], but the variance of PF estimates typically scales exponentially with the state dimension [6].
Although they typically do not return consistent estimates, variational techniques provide attractive alternative for simultaneous state estimation and parameter learning which scales better to high- dimensional latent states than PF, and are not restricted simple parametric approximations, Many
"Equal contribution
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

such methods have been proposed for SSMs over recent years, e.g. [2, 10, 20, 22, 23, 29, 33, 36]. However, they have generally been developed for batch inference where one maximizes the Evidence Lower Bound (EIBO) for a fixed dataset. As such, they are ill-suited for online learning as, whenever a new observation is collected, one would need to update the entire joint variational states distribution whose dimension increases over time. Though a small number of online variational approaches have been developed [30, 38, 45], these rely on significant restrictions of the variational family, leading to approximations that cannot faithfully approximate the posterior distribution of the latent states.
The main contribution of this paper is novel variational approach to perform online filtering and parameter learning for SSMs which bypasses those restrictions. As per standard batch variational inference, we simultaneously maximize an ELBO with respect to both model parameters and variational approximation of the joint state posterior. However, our method operates in an entirely online manner and the cost of updates at each time step remains constant. Key to our approach is backward decomposition of the variational approximation of the states posterior, combined with a representation of the ELBO and gradients as expectations of value functions satisfying Bellman-typp recursions akin to those appearing in Reinforcement Learning (RL) [39].
2 Background
State-Space Models. SSMs are defined by latent X'-valued Markov process (24))21 and Y-valued observations (21) which are conditionally independent given the latent process. They correspond to the generative model
yrlær ~ go (ze): where E Iede is parameter of interest and we consider here Rdd. For y y1st = (311....3))) we thus have
I ~ Ho(x)),
2e41 I. ~ fo(# 12.))
(o 5²) no (In (yn (m)) II
k= (o( (xxk)
Assume for the time being that 0 is known. Given observations (Be)es) and parameter values 0, one can perform online state inference by computing the posterior of x, given y which satisfies
po(xx)l")) (Vv (Te (zx.)p
pa (01)=/// fa ( 1
1)dre-1s
(1)
po (ye by -11
with po (011111 == Mo (z()). The log evidence ( == log Po (v') is then given by
(0)) Ei- - log (yn 1) where pe (yix 1 I go(nn (= (rx))))(((l)l
(2)
Here (x ( is known as the filtering distribution. While the recursion (1) and the sequential
decomposition (2) are at the core of most existing online state and parameter inference techniques [19, 37, 45], we will focus here on the joint state posterior distribution, also known as the smoothing distribution, of the states Fia given y² and the corresponding representation of the log evidence Po = po (2 y²) p
(e()) = log po (v²) = log (J (o y v')dz1:).
Variational Inference. For general SSMs, the filtering and smoothing distributions well as the log evidence are not available analytically and need to be approximated. For data y standard variational approaches stochastic gradient techniques to maximize the following ELBO
(0 Φ) (m))))))))) [log (po( //) 9 (z (zial)"))] I (0).
Maximizing this ELBO w.r.t. the parameter á of the variational distribution qº (xi:elyy") corresponds to doing variational smoothing while maximizing it w.r.t. 0 corresponds to doing parameter learning.
As the true smoothing distribution satisfics ps (2 lyy)) po( II Pe (xxly", 1). one typi- cally uses 9 (æ1:: ly') = o (21 (e')) IT>= (zxlv' -1) 1) for the variational smoothing distribution; see e.g. [2. 20, 22, 43]. However this approach is not suitable for online variational filtering and parameter learning. Firstly. the resulting marginal (xx/p/ lB/") of (Z1=11)") approximating Pe Pe( (sl"") typically not available analytically. Secondly, when the new observation y is collected, this approach would require recomputing
entirely new variational smoothing distribution with di- mension that increases with time. One can attempt to bypass this problem by restricting ourselves to q (I = q (y²) Hi 4 (x*ly*, , as per [30]. However, the switch from conditioning on y² to ² prohibits learning an accurate approximation of po ( (y²) as this formulation does not condition on all relevant data.

---

3 Online Variational Filtering and Parameter Learning
Our online variational filtering and parameter estimation approach exploits backward factorization of the variational smoothing distribution qº (x1=el))) The ELBO and its gradients w.r.t. 0 and ø are computed in an online manner as increases by using a combination of ideas from dynamic programming and RL, with computational time that remains constant at each time step. To simplify notation, henceforth we will write of (I1:)) = (x1ally)"
3.1 Backward Decomposition of the Variational Smoothing Distribution
The key property that we will be exploiting is that po(xialy") satisfies
So ( (ox
Po Po(2naly') (o( Po(xll xx+1)) po xk41)
(4)
Po Po(Ix+10
Equation (4) shows that, conditional upon y², (=+) is reverse-time Markov chain of initial distribution pe (=)ll')) and backward Markov transition kernels po ( ,f++1); Crucially the backward transition kernel at time depends only on the observations until time
see e.g. [12, 19].
To exploit this, we consider variational smoothing distribution of the form
(m)a) = (f(
(5)
where o (Ire) and oh (2 [Tx-1) 1) are variational approximations of the filtering distribution pa (xxlly") and the backward kernel pe (zx xx+1) respectively. Using (5), one now has
(,(0, Φ) (0) - (1, (Ix) (Po((ll)))))))
- 1=1
E 9 (=)-))
( (xoooll""
mm+1))]
where KL is the Kullback-Leible divergence and 9: (Ix+1) is the marginal distribution of rx+1 under the variational distribution q. (IIIe) defined in (5). Considering this variational distribution thus makes it possible to learn an arbitrarily accurate variational approximation of the true smoothing distribution whilst still only needing to condition on y at time k and not on future observations.
Additionally, it follows directly from (5) that easily update g 1 (#1:0-1 from 9 (x1:c) using
Letz ( ( (1
(21411) oe (41.)m44 (#441/1)) for (2441 =
(6)
of (xe)
Here mit1 (#141)01) can be viewed as an approximation of the Markov transition density dett (2)41/=1) x gftt (2al ( it typically not proper Markov transition density; # 1 as f se41 (xe/ (#+41)drt++ + q. (xx).
i.e. [ (1
Let us assume that o (Ix) = # (æx) and gx (1a)) = dh (Ix- |xx), then (5) and (6) suggest that we only need to estimate oe at time as does not impact the backward Markov kernels prior time However, we also have to be able compute estimates of (0, Φ) and (0 Φ) to optimize parameters in constant computational time at each time step, without having to consider the entire history of observations y². This detailed the next subsections where we show that the sequence of ELBOs (6, Φ) h21 and its gradients (VeLi(0, σ) h22 and (V,Li h) can all be computed online when using the variational distributions {g. (21 defined in (5).
3.2 Forward recursion for the ELBO
We start by presenting a forward-only recursion for the computation of (C, h>2.. This recursion illustrates the parallels between variational inference and RL and is introduced to build intuition. Proposition 1. The ELBO L, σ) satisfies for 2
6,(0,)) 6) -E,)( (x,) (x=))( for ((:0(()) = Egg (m))-))=))) [log (po(aLe w/)/ (#1=))
3

---

with the convention (x1) :- r..0 (xo, x1) == (21. (x1)) Additionally, we have
v [+1 (1(+1) = =E Wii1 (2012ee))) (x r (2), x(-1)>, where
(7)
A. 0(41) == log (Jo 1 go) (ve |1+ m (2+4) (x))
(8)
All proofs are given in Appendix C. Proposition shows that we can compute C, o). for 2 1, online by recursively computing the functions V: using (7) and then taking the expectation of V w.r.t. of (az) obtain the ELBO time t. Thus, given V.l.0, we can compute V and (1+1(0,0) using only Ut+1, with cost that remains constant in t.
This type of recursion is somewhat similar to those appearing in RL. We can indeed think of the which
ELBO (0, +) as the expectation of sum of "rewards" Tk we compute recursively using the "value" function V, However, while in RL the value function is given by the expectation of the sum of future rewards starting from Tt, the value function defined here is the expectation of the sum of past rewards conditional upon arriving in Vy. This yields the required forward recursion instead of backward recursion. We expand on this link in Section 4 and Appendix A.3.
9.0 given in (8) from k 1
3.3 Forward recursion for ELBO gradient w.r.t. 0
similar recursion can be obtained to compute (V,L, 122 This recursion will be at the core of our online parameter learning algorithm. Henceforth, we will assume that regularity conditions allowing both differentiation and the interchange integration and differentiation are satisfied. Proposition The ELBO gradient
(oL,(0, satisfies &
VoL,(0, (e,)
0.0 (x+)); where $ VoV. (x.).
Additionally, if we define 8441 (fe, = Vor 0.0 (Iz, x(+1) = V log fo ( ( then
$44 (xx-1) = B (m
(xi) 1 (z), (1e,
Proposition 2 shows that we can compute (VoCe 0) online by propagating { h22 using (9) and taking the expectation of the vector S w.r.t. 9, (xx) to obtain the gradient at time t. Similar ideas have been previously exploited in the statistics literature to obtain forward recursion for the score vector (0) so as to perform recursive maximum likelihood parameter estimation; see e.g. [19, Section 4]. In this case, one has (0) (xx(=xl)) (ar)] where St satisfies recursion similar to (9) with qe (zy (z+11) replaced by Po (z, lB", x1+1).
3.4 Forward recursion for ELBO gradient w.r.t. +
We finally establish forward-only recursions for the gradient of the ELBO w.r.t. & which will allow us to perform online variational filtering. We consider here the case where, for all k. 9 (xx) = % (zz)) and of (Ix-1|ZZ) ofth (#x-1(zz)) so L,(0, () Lr Φ1a) and (20 v Bi (zx). At time step t, we will optimize w.r.t de and hold all previous Φ1a-1 constant. Our overall variational posterior (5) is denoted (x11.)) The alternative approach where á is shared through time (amortization) is investigated Appendix D.
Since the expectation is taken w.r.t. 2h (I11)) in L. optimizing é is slightly more difficult than for 0. However, we can still derive a forward recursion for the o gradients and we will leverage the reparameterization trick reduce the variance of the gradient estimates; i.e. we assume that ze ()) d. (Ix) and 1 E E,-1, 2,) q: (I1-1 when 4 X(e), & X(e).
Proposition The ELBO gradient Vo. (1,, Φ1:)) satisfies for 2
V C,,, V E
V 0 (xx)) ExCe)) V v0 (x(())))))))))
(m.)

---

Additionally, one has
V Verti 1.,01e41 (04110441:4444))
/+1
T0 (xal011e Fe+ ( 24410011100
= =Esees)
d
V Voun 0,00+ t - - (0
T
where T. (xt) = a
vr (zx) satisfies the forward recursion
8es
1-1 Theeet (I++1) = -x(ee)
Ie.00: (xc(Out1: T++1)) dz, (
a04+1
9.00a11 (2: (x+ (0r+1 Ce It-1), s4t1) +(+1)) (10)
V a1ti r+1
Oz 10 Jacobians of appropriate dimensions. 024t1
Here,
dor.1
3.5
Estimating the ELBO and its Gradients
As we consider the case te (xx) = 9 (Ix), Ox (xx-1 = 9 (xx)) we have Sf.0 (x:) = Sh (zz). At time t optimize de and hold Φ1--11 constant. Practically, we are not able to compute in closed-form the functions y. V. (xz), St squbie (xt) and T. (xt) appearing in the forward recursions of (0, o1:)), (O,, )1:e) and V (0,, @1a) respectively. However, we can exploit the above recursions approximate these functions online using regression is commonly done in RL. We then show how to use these gradients for online filtering and parameter learning.
We approximate 50.0ut with S Equation (9) shows that S441 can be learned using S through regression of the simulated dataset? ( (x;t) + 4441 (x), r(+1)) with (r/,,//1)
d1+1 (22, for = 1,..,, N (see Appendix A. 1 for derivation). We can use neural networks to model S, or Kernel Ridge Regression (KRR). Note that the use of KRR estimate gradients for variational learning has recently been demonstrated by [26].
We similarly approximate T 7.0de:++ with Te+1- As before, we can model T441 using neural networks or KRR. We use recursion (10) and T. to create the following dataset for regression"
0.d fiits Tr(xe(01+1 2(,)))1) az (0 + V ai T(41 (Ie(0:+1i ci, x11)) x =11)
02
where 2/41 ~ 01+1 (I4+1) and of A(c) for = 1,..., N. The choice of the distribution over the inputs, ri+1, for each simulated dataset is arbitrary but it will determine where the approximations are most accurate. We choose 4441 (#++1) to best match where we expect the approximations to be evaluated, more details are given in Appendix A.
Note that if one is interested in computing online an approximation of the ELBO, we can again similarly approximate -:0 (xx) using regression obtain (x(x)) by leveraging (7). We will call the resulting approximate ELBO the Recursive ELBO (RELBO). We could also then differentiate (x(e)) w.r.t. 2, to obtain an alternative method for estimating T. -bia and optimizing dt. However, as we are ultimately interested accurate gradients, this approach does not exploit the readily available gradient information during the regression stage.
By approximating T with Ti+1 and S,41 df with S introduce some bias into gradient estimates. We can trade bias for variance by using modified recursions; e.g.
S &. (It+1) = -B, (a
8..0nn- (xx- + Ei She (II, TR+1) 1-11
As L increases, we will reduce bias but increase variance. Such ideas are also commonly used in RL but we will limit ourselves here to using L
We define So 0, a == 0. sf (3ro, .1) = Va log μs (x1))g (ya (=/==)) and of (xo, 21) == 9f (x1).
Similarly, define To : ab = @ and r
8.00.1 as in Proposition 1.
5

---

Algorithm 1: Online Variational Filtering and Parameter Learning.
for f ,....,? do Initialize Oe e.g. Φi + Φ1-1 /+ Update db using M stochastic gradient steps
for m = ,,...,,. M do
Sample -i- t5 ~ of" (11-1, using reparameterizzttoo trick for = 1,... N
de + Ur + 7oo E (x)-1)
dri-1 + V Var doc
end
Update T(( and ((( as in Section 3.5
(2x)) t ( 0m-11 (
der
(141 m ( - (5+-1,
Update 0 using a stochastic gradient step
Sample 1i-11, of (x), $ ~ qi - 1 for == 1,..,N N
0: + 01-1 + ni 1 + de -e- (xt-,,)) (
end
3.6 Online Parameter Estimation
Assume, for the sake of argument, that one has access to the log evidence 4,(0) and that the observations arise from the SSM with parameter 0* Under regularity conditions, the average log- likelihood & (0)/t /t converges as - oo towards function {(O) which is maximized 0*; see e-g. [12, 41J. We can maximize this criterion online using Recursive Maximum Likelihood Estimation (RMLE) [19, 24, 40, 41] which consists of updating the parameter estimate using
0.+1- 02 + 7
0t (Epeo.. (00 1 Ivf (z, + $141 (Zt,
(11)
The difference of two expectations the r.h.s. of (11) is approximation of the gradient of the log prgeiiici log pe (V(+11ly') evaluated at 6,. The approximation given by V log Pe.., (y*+1) v log Poo..-1 (y') with the notation V log Po.. (1*+1) corresponding to the expectation of the sum of terms s Rk+1 (2 Tk+1) w.r.t the joint posterior states distribution defined by using the SSM with parameter 03 at time k +++1 1.
We proceed similarly in the variational context and update the parameter using
O2+1 =0 0+++
(E
9.11 (#1))44)) Sx (xx)) + ( #e.11 - of (x)) (s.(zr)))
Here (z,, approximates Se(z,) satisfying (1+4( == oct (zel2-.)) (Se(rr) + (2 xe41)].
deti
We compute S, in Section 3.5 with simulated dataset using S.11 and 0.-1.
Putting everything together, we summarize our method Algorithm using a simplified notation to help build intuition, full description is given in Appendix A.2. It is initialized using initial parameters Φ1, 0g. We re-iterate the algorithm's 's computational cost does not grow with t. We need only store fixed size T and Š models as well as the most recent &i and 0g parameters. When performing backpropagation, T and Š summarize all previous gradients, meaning we do not have to 1. Therefore, we only ineur standard backpropagation computational
roll all the way back cost w.r.t. di and 0. To scale to large d.,, we can use mean field q. (ar). gt linear in dg.
di (x1-1/2;) keeping costs
4 Related Work
As mentioned briefly in Section 3.2, our recursive method has close links with RL, which we make explicit here. In 'standard' RL, we have value function and corresponding Bellman recursion given by
Σ r(sk a2)
T
V/RL(S) == E
VRL (RL( Esurrar [r(se, az) + VRL (S141)]
(12)

---

where (st, a.) is the state-action pair at time t. Whereas the RL value function summarizes future rewards and so recurses backward in time, our "value' function summarizes previous rewards and recurses forward in time. Writing this using the state-action notation, one obtains
Vi((s)) := E Σ r(sk: an)
Ve41($(41) = Butyiis [r(s++1, a4+1) + Va())))
(13)
Note we have defined an anti-causal graphical model, with on depending on 8141 and a++1- we further let 8t Tt, O2 = Σt-1 ~ ( ((De)1/zre) P(*t|8e4) 0x-1) = 8(s 04+1) and r(s,, az) = (t-1) as in equation (8), then the recursion in equation (13) for V2 corresponds to the recursion in equation (7) and Lr = E 9g (zv) [Vr(zT)]
}]- We then differentiate this recursion with respect to 0
and 24+1 to get our gradient recursions (9) and (10) respectively. Full details are given in Appendix A.3. A similar correspondence between state actions and successive hidden states was also noted in [43] which explores the of RL ideas in the context of variational inference. However, [43] exploits this correspondence to propose a backward in time recursion for the ELBO of the form (12) initialized at the time of the last observation of the SSM When a new observation is collected at the next time step, this backward recursion has to be re-run which would lead to a computational time increasing linearly at each time step. In [20], the links between RL and variational inference are also exploited. The likelihood of future points in an SSM is approximated using temporal difference learning [39], but the proposed algorithm is not online.
To perform online variational inference, [45] proposes to use the decomposition (2) of the log evidence (0) and lower bound each term log po (y 1)) appearing in the sum using
fo ØR go (vk (+ )o (2 -11-111 11
log
log po (ox
(14)
9 (EE-1
gg (k))
Unfortunately, the term on the Lh.ss of (14) cannot be evaluated unbiasedly as it relies on the intractable filter -o(==-ily*~ so [45] approximates it by po(#x-1ly**1) so o (1x-1) to obtain the following Approximate ELBO (AELBO) by summing over k:
(000 Σ
k== 1E (an (mx mmm [log 9.00 (xx_1,
(15)
Additionally, [45] makes the assumption q (xx- .x.) o (xx- )a (sx)) and we will refer to (15) this case as AELBO-1. While [16] does not consider online learning, their objective is actually a generalization of [45], with a (xx=1, = of (Ix 19 ( l and we will refer to (15) in this case AELBO-2. It can be easily shown that AELBO-2 is only equal to the true ELBO given in Proposition in the (unrealistic) scenario where of (Ixx) = po (xx for all k. Moreover, in both cases the term pe (xx-1ly~~11 is replaced by ( causing a term involving 0 to be ignored in gradient computation. The approach developed here be thought of as way to correct the approximate ELBOs computed in [16, 45] in a principled manner, which takes into account the discrepancy between the filtering and approximate filtering distributions, and maintains the correct gradient dependencies in the computation graph. Finally [44] relies on the PF to do online variational inference. However the variational approximation of the filtering distribution is only implicit as its expression includes an intractable expectation and, like any other PF technique, its performance is expected to degrade significantly with the state dimension [6].
5 Experiments
5.1 Linear Gaussian State-Space Models
We first consider linear Gaussian SSM for which the filtering distributions can be computed using the KF and the RMLE recursion (11) can be implemented exactly. Here the model is defined as
fo( = (x, FT.-1, U), go (3a (x.) = -Nis; Gxt, V),
where F E &Rdexd= G Rd,xd U Rdxxds V E Rd. xd, 0 = (F,G). We let
d (z,) N -N diag(o))))) d. (0)1)))) - (z--1 Wi2, + b, diag(o?))
7

---

10a
0.20
0.20
10a
Ours
0.15
AELBO-1
0.15
10
Mean Absolute
AELBO-2
10
0.10
RMLE
0.10
10-1
10
0.05
10-3
0.00
0.00
6 Time step
5
25k
50k
25k Time step
50k
Time step
(a)
Figure I: (a) (L(g) (z 2t (1 x vs time step of the SSM. Between each time step, we plot the progress of the KL over 5000 iterations of inner loop oi optimization. (b) Mean Absolute Error for model parameters F (left) and G (right) vs time step (AELBO-) off the scale).
with Φi = {Mr, log os, We, be, log When fo(x)(==-1) is linear Gaussian and we use a Gaussian for (x we can select of (x)-1 [e(=a|=r-1)9 é:-1 (#t-1) as noted in [34]. In this experiment however, we aim to show our full method can converge to known ground truth values hence still fully parameterize a," (It-1 (22) as well as setting the matrices F, U,G to be diagonal, so that po (xx)lyy') and pe p0 1ire--ly"111 xx) are in the variational family.
For d, dy 10, we first demonstrate accurate state inference by learning d, at each time step whilst holding 0 fixed at the true value. We represent T; (xx) non-parametrically using KRR. Full details for all experiments are given in Appendix B. Figure Ia illustrates how, given extra computation, our variational approximation comes closer and closer to the ground truth, the accuracy being limited by the convergence of each inner stochastic gradient ascent procedure. We then consider online learning of the parameters F and G using Algorithm 1, comparing our result to RMLE and variation of Algorithm using AELBO-1 and 2 (see Section 4). Our methodology converges much closer to the analytic baseline (RMLF) _F) than AELBO-2 [16] and exhibits less variance, even though the variational family sufficiently expressive for AELBO-2 to learn the correct backward transition. In addition, we find that AELBO-1 [45] did not produce reliable parameter estimates in this example, as it relies on a variational approximation that ignores the dependence between #g-1 and æg. As expected, our method performs slightly worse than the analytic RMLE, inevitably small errors will be introduced during stochastic optimization and regression.
5.2 Chaotic Recurrent Neural Network
We next evaluate the performance of our algorithm for state estimation in non-linear, high-dimensional SSMs. We reproduce the Chaotic Recurrent Neural Network (CRNN) example in [44], but with state dimension d, 5, 20, and 100. This non-linear model is an Euler approximation of the continuous-time recurrent neural network dynamics
f N (I1 - (-2--1 + W tanh(z;-1)),U))
and the observation model is linear with additive noise from a Student's t-distribution. We compare our algorithm against ensemble KF (EnKF), bootstrap PF (BPF), as well as variational methods using AELBO-1 and AELBO-2. We let q. (1)) = (1 MM (z:), diag(??)) and qt'h (z,) 1-layer Multi-Layer Perceptron (MLP) with 100 neurons for each 9t (xt-1/0z)) We generate a dataset of length 100 using the same settings as [44], and each algorithm run 10 times to report the mean and standard deviation. We also match approximately the computational complexity for all methods. From Table 1, we observe that the EnKF performs poorly on this non-linear model, while the PF performance degrades significantly with d,, as expected. Among variational methods, AELBO-1 does not give accurate state estimation, while AELBO-2 and our method achieve the lowest error in terms of RMSE. However, our method achieves the highest ELBO; i.e. lowest KL between the variational approximation and the true posterior since
N (14; 14, diag(o?)) where we
'Code available https: pettps:/giiie http//ggtthttttbt
fil_

---

0 is fixed effect not represented using just the RMSE. We confirm this is the case in Appendix B.2 by comparing our variational filter means 14, against the 'ground truth' posterior mean for d = 5 computed using PF with 10 million particles. Furthermore, our method is also able to accurately estimate the true ELBO online. Figure 2a shows that our online estimate the ELBOO, RELBO (Section 3.5), is very close to the true ELBO, whereas AELBO-2 is biased and consistently overestimates it. Further, AELBO-1 is extremely loose meaning its approximation of the joint state posterior is very poor.
Using this CRNN problem, we also investigate the penalty we pay for our online method. At time t, we train t (are), q(" (x1-1 and hold th (
fixed. This because all information to learn the true factor po( 2x) is available at time as does not depend on future not in the variational family (as this CRNN observations. However, when po p2 1- ,x4) example), 06 (Iz-1|0Z)) will aim to be most accurate in the regions of I'k in which it is typically evaluated. The evaluation distribution over zy does depend on future variational factors and could shift over time. This may result in learning different variational joint q. "a (I11.) between when using our online BLBO and when just using the full offline ELBO (3). We quantify this difference by training the same joint variational distribution using either our online ELBO the offline ELBO (3). Figure 2b plots the final mean and standard deviations of the marginals of the trained q." (#1:1) ) in both cases. We see that these quantities are very close, suggesting this effect is not an issue on this example. This may be due to the evaluation distribution over dy not changing a significant enough amount to cause appreciable changes in the learned variational factors. We show this result holds over dimensions and seeds in Appendix B.2.
1.2
0.8
-10
ELBO
-20
0.4
-30
True hidden state
RELBO (ours)
Ours True ELBO
Online
0.0
40
Offline
AELBO-2
AELBO-2 True ELBO
0
20
40
60
80
100
20
40
60
80
100
Time step
Time step
(a)
(b)
Figure 2: (a) Estimates and true values of the ELBO the Chaotic RNN task. RELBO uses KRR for V. whilst for the other methods we eq. (15). (b) Comparison between joint variational distributions trained online and offline on the CRNN task. The colored lines show the mean of de (x)1) whilst the shaded region represent 11 std. The true hidden state is also shown in black.
5.3 Sequential Variational Auto-Encoder
We demonstrate the scalability of our method on a sequential VAE application. In this problem, an agent observes long sequence of frames that could, for example, come from robot traversing new environment. The frames are encoded into a latent representation using a pre-trained decoder. The agent must then learn online the transition dynamics within this latent space using the single stream of input images. The SSM is defined by
fo fo( (xx NN/ (Z+-)).U).
( = N N16 NN (z+), V),
where d, = 32, NN/ is a residual MIP and NN convolutional neural network. NN/ is learned online whilst NN is fixed and is pre-trained on unordered images from similar environments using the standard VAE objective
[21]. We perform this experiment on a video sequence from a DeepMind Lab environment [5] (GNU GPL license). We use the same q," parameterization as for the CRNN but with a 2 hidden layer MLP with 64 neurons. KRR is used to learn I; whereas we use an MLP for learning S. We found that MLPs scale better than KRR as dg is high. Our online algorithm is run on sequence of 4000 images after which probe the quality of the learned NN The results

---

Table Root Mean Squared Error between filtering mean and true state and the average true ELBO for the methods in varying dimensions on the Chaotic RNN task.
EnKF
BPF
AELBO-1
AELBO-2
Ours
Filter RMSE
0.1450z0.0026
0.1026#0.0001
0.1284=0.0035
0.1035=0.0012
0. 0.1032+0.0009
ELBO (nats)
-20.52246.766
-30.94442.2223
-15.845+1.7385
Time per step
1.0998
0.9268
1.5067
2.2270
2.6899
Filter RMSE
0.1541=0.0016
0.109240.0014
0.1355#0.0012
0.1086#0.0004
0.1082:0.0003
20
ELBO (nats)
928.80=10.464
-393.68+3.9052
-340.36+3.9730
Time per step
5.1879
3.8932
2.3587
2.7000
3.5935
Filter RMSE
0.157120.0017
0.2493+0.0122
O. 0.123920.0006
0. 0.1070#0.0001
0.1068+0.0001
100
ELBO (nats)
-4247 2247.9420.905
2209.7741.814
-1794.725.4173
Time per step
6.4546
4.6184
3.2697
4.5539
5.9263
(a) Before training
(b) After training
Frames predicted by rolling out NN! from two different starting points, before and after training. Between each frame, 3 transition steps taken. Before training, meaningful change is predicted but after training NN1 predicts plausible movements.
Figure 3:
shown in Figure 3. Before training, NN: predicts no meaningful change but after training it predicts movements the agent could realistically take. We quantify this further in Appendix B.3 by showing that the approximate average log likelihood & (0) computed using Monte Carlo increases through training. thereby confirming our method can successfully learn high-dimensional model parameters of the agent movement in fully online fashion.
Discussion
We have presented novel online approach for variational state estimation and parameter learning. In our experiments, this methodology outperforms standard filtering approaches for high dimensional SSMs and is also able to learn online high dimensional model parameters sequential VAE. However, it is not without its limitations. As with any stochastic variational inference technique, we can only obtain an accurate posterior approximation the variational family is expressive enough and our stochastic gradient method finds reasonable minimum. We also need to use flexible function approximators to keep the bias in our gradient estimates small. Finally, although our method is online, it can be quite computationally expensive in absolute terms requires solving an optimization problem and solving a regression task for each time step.
neural network in
To reduce this computational cost, one can amortize the optimization cost for & by learning network taking a representation of observations up to now, y', and outputting oe as illustrated in Appendix D. Further work will investigate ways to also amortize the cost of function regression across time through meta-learning strategy. From a theoretical point of view, it would be useful to establish conditions under which the proposed variational filter is exponentially stable and to study asymptotic properties of the parameter learning procedure.
10

---

Acknowledgments and Disclosure of Funding
The authors are grateful Adrien Corenflos, Desi Ivanova and James Thortoo for their com- ments. Andrew Campbell acknowledges support from the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/L) Arnaud Doucet is partly supported by the EPSRC grant EP/R034710// He also acknowledges support of the UK Defence Science and Technology Laboratory (DSTL) and EPSRC under grant EP/R013616/ This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative.
References
[1] Andrychowicz, M.. Denil, M., Gómez, S., Hoffman, M. W., Pfau, D., Schaul, T., Shillingford, B., and de Freitas, N. (2016). Learning learn by gradient descent by gradient descent. In Advances in Neural Information Processing Systems.
[2] Archer, E., Park, I. M., Buesing, L.,, Cunningham, J., and Paninski, L. (2016). Black box varia- tional inference for state space models. International Conference on Learning Representations Workshop Track.
[3] Azimi-Sadjadi, B. and Krishnaprasad, P. (2005). Approximate nonlinear filtering and its applica- tion in navigation. Automatica, 41 (1669555-9
[4] Beard, M., Vo, B. T., and Vo, B.-N. (2020). A solution for large-scale multi-object tracking. IEEE Transactions on Signal Processing, 68:2754-2769.
[5] Beattic, C., Leibo, J. Z., Teplyashin, D., Ward, T., Wainwright, M., Küttler, H., Lefrancq. A., Green, S., Valdés, V. Sadik, A., et al. (2016). DeepMind lab. arXiv preprint 0rXiv:1612.03801.
[6] Bengtsson, T., Bickel, P., Li, B., al. (2008). dimesooffii revisited: Collapse of the particle filter in very large scale systems. In Probability and Statistics: Essays in Honor of David A. Freedman, pages 316-334. Institute of Mathematical Statistics.
[7] Brigo, D., Hanzon, B., and LeGland, (1998). A differential geometric approach to nonlinear filtering: the projection filter. IEEE Trunsactions Automatic Control, 43(2):247-252.
[8] Chung, J., Kastner, K., Dinh, L.,, Goel, K., Courville, A., and Bengio, Y. (2015). A recurrent latent variable model for sequential data. In Advances in Neural Information Processing Systems.
[9] Corentlos, A., Thornton, J., Deligiannidis, G., and Doucet, A. (2021). Differentiable particle filtering via entropy-regulariized optimal transport. In International Conference on Machine Learning.
[10] Courts, J., Hendriks, J., Wills, A., Schön, T. B., and Ninness, B. (2021). Variational state and parameter estimation. IFAC Symposium on System Identification.
[11] Del Moral, P. (2004). Feynman- -ac Formulae: Genealogical and Interacting Particle Systems with Applications. Springer.
[12] Douc, R., Moulines, E.,, and Stoffer, D. (2014). Nonlinear Time Series: Theory, Methods and Applications with R Examples. CRC Press.
[13] Evensen, G. (2009), Data Assimilation: The Ensemble Kalman Filter. Springer Science & Business Media.
[14] Fellows, M., Mahajan, A., Rudner, T. G., and Whiteson, S. (2019). VIREL: A variational inference framework for reinforcement learning. In Advances in Neural Information Processing Systems.
[15] Fraccaro, M., Sønderby, S. K., Paquet, U., and Winther, O. (2016). Sequential neural models with stochastic layers. In Advances in Neural Information Processing Systems.
[16] Gregor, K., Papamakarios, G., Besse, F., Buesing, I.., and Weber, T. (2019). Temporal difference variational auto-encoder. In International Conference Learning Representations.

---

[17] Ito, K. and Xiong, K. (2000). Gaussian filters for nonlinear filtering problems. IEEE Transac- tions on Automatic Control, 45
[18] Jonschkowski, R., Rastogi, D., and Brock, O. (2018). Differentiable particle filters: End-to-end learning with algorithmic priors. In Robotics: Science and Systems.
[19] Kantas, N., Doucet, A., Singh, S. S., Maciejowski, J., and Chopin. N. (2015). On particle methods for parameter estimation in state-space models. Statistical Science, 30(3):328-3511
[20] Kim, G.-H., Jang, Y., Yang, H., and Kim, K.-E. (2020). Variational inference for sequential data with future likelihood estimates. In International Conference on Machine Learning.
[21] Kingma, D. and Welling, M. (2014). Auto-encoding variational Bayes. In International Conference on Learning Representations.
[22] Krishnan, R. G., Shalit, U., and Sontag, D. (2017). Structured inference networks for nonlinear state space models. In AAAI Conference on Artificial Intelligence.
[23] Le, T. A., Igl. M., Rainforth, T., Jin, T., and Wood, F. (2018). Auto-encoding sequential Monte Carlo. In International Conference on Learning Representations.
[24] LeGland, F. and Mevel, L. (1997). Recursive estimation in hidden Markov models. In Proceedings of the 36th IEEE Conference on Decision and Control, volume 4, pages 1468-3473.
[25] Levine, S. (2018). Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.0090
[26] Li, K. W., Moskovitz, T., Kanagawa, H., and Sahani, M. (2020). Amortised learning by wake-sleep. In International Conference on Machine Learning.
[27] Li, Y. and Mandt, S. (2018). Disentangled sequential autoencoder. In International Conference on Machine Learning.
(28) Ma, X., Karkus, P., Hsu, D., Lee, W. S., and Ye, N. (2020). Diseriminative particle filter reinforcement learning for complex partial observations. In International Conference on Learning Representations.
[29] Maddison, J., Lawson, J., Tucker, G., Heess, N., Norouzi, M., Mnih, A., Doucet, A., and Teh, Y. W. (2017). Filtering variational objectives. Advances in Neural Information Processing Systems.
[30] Marino, J.. Cvitkovic, M., and Yue, Y. (2018). A general method for amortizing variational filtering. In Advances in Neural Information Processing Systems.
[31] Marschall, O., Cho, K., and Savin, C. (2020). A unified framework of online learning algorithms for training recurrent neural networks. Journal of Machine Learning Research, 21:1-34.
[32] Minka, T. P. (2001), Expectation propagation for approximate Bayesian inference. In Uncer- tainty Artificial Intelligence.
[33] Naesseth, C. A., Linderman, S. W., Ranganath, R., and Blei. D. M. (2018). Variational sequential Monte Carlo. In International Conference on Artificial Intelligence and Statistics.
[34] Pfrommer, D. and Matni, N. (2022), Linear variational state space filtering. arXiv preprint wxXi0:2001.01333.
[35] Richter, J., Carbajal, G., and Gerkmann, T. (2020). Speech enhancement with stochastic temporal convolutional networks. In Interspeech.
[36] Ryder, T., Prangle, D., Golightly, A., and Matthews, I. (2021). The neural moving average model for scalable variational inference of state space models. In Uncertainty Artificial Intelligence.
[37] Särkkä, S. (2013). Bayesian Filtering and Smoothing. Cambridge University Press.
[38] Smídl, V. and Quinn, A. (2008). Variational Bayesian filtering. IEEE Transactions on Signal Processing, 50200035000
12

---

[39] Sutton, R. S. and Barto, A. G. (2018), Reinforcement Learning: An Introduction. MIT press.
[40] Tadić, V. B. (2010), Analyticity, convergence, and convergence rate of recursive maximum- likelihood estimation hidden Markov models. IEEE Transactions on Information Theory, 56(12):640664432 6432.
[41] Tadié, V. B. and Doucet, A. (2021). Asymptotic properties of recursive particle maximum likelihood estimation. JEEE Transactions on Information Theory, 67(3):1825-1848
[42] Tsay, R. S. (2005). Analysis of Financial Time Series, volume 543. John Wiley & Sons.
[43] Weber, T., Heess, N., Eslami, A., Schulman, J., Wingate, D., and Silver, D. (2015), Reinforced variational inference. In Advances Neural Information Processing Systems Workshop.
[44] Zhao, Y., Nassar, J., Jordan, I., Bugallo, M., and Park, I. M. (2019). Streaming variational Monte Carlo. arXiv preprint arXiv: 1xXxv:1066001544
[45] Zhao, Y. and Park, I. M. (2020). Variational online learning of neural dynamics. Frontiers in Computational Neuroscience, 14.
13


Reviews:
[{"review_id": "v4Z8WV2PGNL", "reviewer": "Reviewer_YSkX", "summary": "This paper develops an algorithm for online variational inference and parameter learning in state-space models, where both model parameters and variational parameters are learned online with each new data point in the time series. The authors accomplish by approximating a backwards decomposition of the full posterior $p(x_{1:t}|y_{1:t}) = p(x_t|y_t) \\prod_{k=1}^t p(x_k | y_{1:k}, x_{k+1})$ with a factored variational distribution $q_t(x_{1:t}) = q_t(x_t) \\prod_{k=1}^{t-1} q_t(x_k | x_{k+1})$. This factored approximation is used to derive a (forward) recursive equation for computing the ELBO at time $t$, along with gradients of the ELBO with respect to the model and variational parameters. These recursive equations allow both the ELBO and its gradients to be estimated (using recursive learning of regressors) in an online manner, hence enabling online learning and inference. The authors evaluate this approach on two synthetic examples, and one example of video data, demonstrating that online inference successfully minimizes KL divergence to the true posterior as each new data point is observed, and that online learning successfully converges to plausible model parameters / transition dynamics, including in the video modeling task.", "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": "This was a nice paper to read. I thought it was clearly described and motivated, made use of the backwards decomposition of the posterior distribution in an insightful, original, and principled manner, and clearly demonstrated the utility of their approach relative to existing online baselines in a series of experiments. I just have a few clarificatory questions:\n\n1. How fast is the regression step for $\\hat T_t$ and $\\hat S_t$? It was not immediately clear to me that this would take constant time, or exactly what kind of fitting was going on (it might be helpful to give a concrete example in the main text). I looked at the appendix, and it seems like it requires sampling some 500 and data points, and multiple iterations of fitting. This seems like it might be slightly too slow for some applications (e.g. real time robot tracking), and may require variable time if convergence of the regressor is desired.\n2. Why is regression necessary at all, as opposed to Monte Carlo estimation of $T_t^{\\theta,\\phi}$ and $S_t^{\\theta,\\phi}$? Is it because they need to learned as functions of $x_t$ in order to be queried at different (sampled) values of $x_t$? Would it be possible to use (recursive) Monte Carlo approximations instead, avoiding the need for regression?\n3. Are there potential issues with issues with the convergence of online learning, given that the recursive equations for $T_t^{\\theta,\\phi}$ and $S_t^{\\theta,\\phi}$ are not strictly followed once the parameters $\\theta$ and $\\phi$ are updated? e.g. $\\hat S_{t-1}$ is learned using $\\theta_{t-2}$, but $\\hat S_t$ is learned using $\\theta_{t-2}$. Perhaps there is RL theory addressing this?\n4. How does this approach compare in speed / accuracy / computational cost versus amortized variational inference (e.g. learning an RNN as a recognition network)?\n\nI also have the following minor comments:\n\n- In the Abstract, I would personally avoid calling the decomposition / factorization \"non-standard\", because this gave me the impression that it might be unprincipled. \"Backwards decomposition\" would be more informative, and avoid that impression. Having read through the math fairly carefully, the proposed method is definitely principled in a way that I cannot say for many other papers that integrate variational inference with state space models.\n- In Algorithm 1, there is a minor inconsistency between the update rules for $\\phi_t$ and $\\theta_t$ in terms of how the time-index subscripts are used.\n\nOverall, I thought this was a solid paper. It could be made even stronger with a little more clarity on the regression steps, the inclusion of the AELBO baselines for the SVAE task, and additional comparisons of speed / accuracy to amortized inference methods, but the approach is adequately evaluated, well-motivated, and principled, and hence deserving of acceptance in my evaluation.\n\n== Post-Rebuttal ==\n\nThank you to the authors for their detailed responses and clarifications. I believe that they have adequately responded to the questions raised by myself and other reviewers, and I will maintain my score of 8.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed limitations and society impacts. With the extra space in the camera ready version, I would also suggest discussing their method in comparison with amortized variational inference in SSMs (e.g. using RNNs), which are also able to perform constant-time inference of new latent states after batch-wise training on a dataset of sequences. This means additional training cost, and potential overfitting, but potentially lower cost at inference time. It would be good to discuss the costs and benefits of each approach.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "rsTtwhSckwT", "reviewer": "Reviewer_iuPB", "summary": "This paper proposes a variational filtering scheme for state space models that enables recursion to successively estimate the ELBO to achieve efficient inference. Experiments are performed on linear Gaussian filtering, RNN estimation, and sequential VAEs.\n", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "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": "tldr: I tend to weakly accept this paper, although I would strongly encourage the authors to make the connection to reinforcement learning way more explicit in the paper. \n\nedit: after the author(s)' rebuttal, I tend to accept the paper. Thank you for the clarifications.\n\n## Originality\n\n*Are the tasks or methods new?* Yes, I haven't seen any sort of recursive filtering mechanism for state space models like this before.\n\n*Is the work a novel combination of well-known techniques?* Yes, this is probably a novel combination of variational filtering with the variational interpretation of reinforcement learning models in the sense of [Levine, '18](https://arxiv.org/abs/1805.00909) and [Fellows et al, '18](https://arxiv.org/abs/1811.01132). Note that this is a positive comment.\n\n*Is it clear how this work differs from previous contributions?* Yes, it's pretty clear how it differs from other approaches in the state-space models literature.\n\n*Is related work adequately cited?* Yes, overall, I think so. See discussion + the papers listed above as well.\n\n## Quality\n\n*Is the submission technically sound?* Yes, it seems to be. The proofs seem to be mostly algebraic book-keeping (which is fine). \n- I would have liked to have more details about the function approximators that use $\\hat S_{t+1}$ and $\\hat T_{t+1}$ in lines 161-180. \n\n*Are claims well supported (e.g., by theoretical analysis or experimental results)?* \nYes, see comment above. The experimental results suggest that the method works pretty well.\n\n*Are the methods used appropriate?* \nYes, overall. Not knowing the literature well I'd like to know how KRR is used in line 165.\n\n*Is this a complete piece of work or work in progress?* This looks like a complete work.\n\n*Are the authors careful and honest about evaluating both the strengths and weaknesses of their work?* Mostly, see discussion.\n\n## Significance\n\n*Are the results important?* Yes, better variational inference techniques for state space models are likely to encourage their usage in practice, while the technical connection to the Bellman recursion (I think) is probably important.\n\n*Are others (researchers or practitioners) likely to use the ideas or build on them?* Yes, probably.\n\n*Does the submission address a difficult task in a better way than previous work?* Yes.\n\n*Does it advance the state of the art in a demonstrable way?* Yes.\n\n*Does it provide unique data, unique conclusions about existing data, or a unique theoretical or experimental approach?* A unique theoretical / practical approach.\n\n## Discussion\n\n- While I agree that the connection to reinforcement learning is there (I think this is closely related to the Bellman recursion and dynamics models), the authors really need to do some more work to formalize this connection as otherwise, it's really lost on the reader and is a bit of a weakness. \n    - RL can be constructed as probabilistic (and really variational) inference as in [Levine, '18](https://arxiv.org/abs/1805.00909) and succeeding works, so this work could help bridge the gap between these types of interpretations of RL and state space models generically. \n    - However, none of the connections are made explicit so it's a bit tricky to tell right now what the impact is. Improving these connections would significantly improve the paper.\n    - specifically, my concerns are about the paragraphs in lines 122 (I think you're referring to the bellman recursion but not sure), \n\n### Questions\n\n- While I understand that the time complexity is constant per iteration, it still seems like there is $t$ separate operations that much be computed. Does this mean that as the number of time steps increases, the method slows down over time?\n    - Or is there only a single pass being computed through the sequence?\n\n- In Fig. 1b why does the quality of the variational approximation get worse as the number of steps?\n    - Is this overfitting or the quality of the variational approximation breaking down?\n\n- How are the number of steps of SGD / Adam chosen in the inner loop?\n\n- How are the summary models (either krr or mlp) initialized and with what sort of data? These are the models using $\\hat S$ and $\\hat T$?\n    - Why doesn't computation for the KRR slow down over time?\n\n- Is fig 1a just the first 10 time steps? Is there a tradeoff between \"time per iteration\" and \"accuracy\" here? If so, how is the tradeoff chosen?\n\n### Writing Comments\n\n20: network[s]\n\n24: [estimating?] \"online inference\" \n\n29: ref for \"notoriously unreliable\" as a non-expert in this field, i'd like to see a reference to verify.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"ethical_concerns": "n/a", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "kb7M1W18qXH", "reviewer": "Reviewer_dXyo", "summary": "The paper borrows the ideas from reinforcement learning to come up with an online inference and learning for state-space models. It focuses on the smoothing distribution instead of the usual filtering distribution for online learning. When performing variational EM with the new incoming data, the ELBO is recursively updated, and the gradients are estimated with a regressor. Experiments show good performance on various datasets.", "questions": "", "limitations": "", "rating": 7, "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": "The paper proposes a novel online inference scheme for state-space models. The paper is generally clear and well-written. Here are some unaddressed points in my mind:\n- It seems the (auxiliary) neural network for gradient estimation is learned in advance. If this is the case, in a non-stationary environment, will this estimation become outdated as the network is stationary? \nTo train the auxiliary neural networks for gradient estimation, an additional dataset is also needed. This seems not a desideratum of an online setup where an additional dataset is usually not available.\n- Does it also work for model-based reinforcement learning since it borrows ideas from RL applications? If so, how does it perform?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "It would be great to state which requirement does not coincide with the usual online learning setup.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "cdnYu5jfXJ0", "reviewer": "Reviewer_FJzo", "summary": "This paper presents a variational inference algorithm for state-space models, that is suitable for online learning: at each time step, the variational posterior approximation can be updated to account for the new data in constant time. Parameters of the generative model can also be learned online. The algorithm is based on variational families $q_t^\\phi$ that factor as $q^{\\phi_t}(x_t \\mid y_{1:t}) \\prod_{i=t}^2 q^{\\phi_{i}}(x_{i-1} \\mid x_i, y_{1:i-1})$. The key point is that when a new datapoint $y_{t+1}$ is observed, the optimal backward kernels $p(x_{i-1} \\mid x_{i})$ (which the $q^{\\phi_i}(x_{i-1} \\mid x_i)$ approximate) do not change for $i = 1,\\dots,t$. The authors use this to justify freezing the parameters $\\phi_i$ after time step $i$, so that at each time step, only the parameters $\\phi_{t+1}$ must be learned. Besides this factorization, the other trick enabling online learning is a scheme for estimating gradients of the ELBO with respect to $\\phi_{t+1}$ in constant time. This requires an additional approximation step, leading to biased gradient estimates. However, in experiments on several models and datasets, the authors show that their proposed technique can produce reasonable results in practice.\n", "questions": "", "limitations": "", "rating": 8, "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": "This is a nicely written paper that presents (to my knowledge) a novel technique for online inference in state-space models, and validates it empirically in three compelling experiments. The exposition is clear and provides good intuition for the key insights behind Algorithm 1. I think the paper is of definite value and interest to the NeurIPS community and recommend acceptance.\n\t\nHowever, I have a less clear picture of what exactly is happening theoretically. One question is what penalty  we should expect to incur (e.g., compared to batch variational inference) by using this online variant. Intuitively, we might expect that in the \"online\" Algorithm 1, the learned parameters $\\phi$ would, assuming a sufficiently high number $M$ of gradient updates per time step, recover the optimal \"batch\" solution using the same variational family—i.e., the parameter settings you'd find if you optimized $L_t$ with respect to all of $\\phi$ at once. But I believe this is not the case in general, because although $p_\\theta(x_k \\mid y^k, x_{k+1})$ does not depend on observations past time $k$, the optimal $q^{\\phi}(x_k \\mid y^k, x_{k+1})$ might. This is because, if $q^\\phi(x_k \\mid y^k, x_{k+1})$ cannot perfectly represent the true backward kernel, it will prioritize accurate backward transitions for the most probable values of $x_{k+1}$. Since the posterior on $x_{k+1}$ does change as new data comes in, the optimal $q$'s priorities also change. The experiments suggest that this may not be a problem in practice, perhaps because $q(x_k \\mid x_{k+1})$ is sufficiently expressive to learn an accurate backward kernel for all $x_{k+1}$, or because in these particular datasets, the filtering and smoothing distributions over $x_{k+1}$ were sufficiently similar. But it might be nice to see some discussion of this issue, and/or an empirical validation that the online and batch variants learn similar variational approximations in practice.\n\t\nA similar issue arises, I think, with the biased gradient estimates — L296-297 point out the importance of using sufficiently expressive function approximators to reduce the bias, but it also seems important that the training data used to optimize the approximators be from a good distribution. If I'm following correctly, the \"training data\" distribution for the approximators $\\hat{T}$ and $\\hat{S}$ are not equal to their \"test time\" distributions, which incorporate an additional data point. Could this introduce appreciable error in practice?\n\t\nOverall, I suspect both these concerns are nitpicks -- but perhaps they could be signposted as assumptions the algorithm makes. To reiterate my overall positive impression of the work, I think the paper presents a clever idea with clear exposition and convincing experiments. \n\t\nOne very minor comment: in discussing [30], L222-223 state that \"any\" particle filtering technique will scale poorly when the state is high-dimensional, but this seems a harder claim to defend when the technique uses learned / well-adapted proposal distributions. (In any case, it is not a reason to dismiss PF in favor of VI; if you have a variational approximation to the posterior, you could use it within a particle filter as the proposal, and the accuracy should only improve.)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors include a limitations section, which I think is adequate, but could be expanded to discuss the issues raised in the main review. I did not see discussion of the potential negative societal impact, but the work is theoretical enough that I am not sure it is necessary to include.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {"ethical_concerns": "I don't see any ethical concerns -- although inference in state-space models can be applied to questionable ends (e.g. object/person tracking in surveillance videos), this is a methods paper, and the example applications presented here do not strike me as problematic."}}]

Author rebuttal:

Meta-review: The authors propose a general scheme for approximation in state space systems where the variation distribution is conditioned _backwards_ in time.  This leads to an ELBO (and gradients) that can be computed online, without having to adjust previous terms. 

The reviewers unanimously praised the clarity of the writing, and I was impressed that each reviewer was left with a very clear picture of the paper's contribution and relevance - a sign that the idea has been presented with great clarity. 

One clarification that arouse during the discussion was around a regression that occurs at time step $t+1$, and whether the input/output pairs here are drawn in such a way that could introduce error. I'm happy with the authors' clarification, which I anticipate in the camera ready manuscript. 

One in-depth discussion between a reviewer and the authors covered connections to reinforcement learning: the authors connect Bellman recursion to the ELBO updates. This discussion could make for a neat appendix to the paper. 
 
Decision: Accept (Oral)