Online Variational Filtering and Parameter Learning
Year: 2022

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.

heading: 1 Introduction
text: Many tasks in machine learning with time series data—such as video prediction [16, 27], speech enhancement [35] or robot localization [9, 18, 28]—often need to be performed online. Online techniques are also necessary in contexts as diverse as target tracking [4], weather prediction [13] and financial forecasting [42]. A 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 is a challenging problem and approximate inference techniques for such models remain an 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 is a 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, to 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 theoretical guarantees [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 an attractive alternative for simultaneous state estimation and parameter learning which scales better to highdimensional latent states than PF, and are not restricted to simple parametric approximations. Many
∗Equal contribution
35th Conference on Neural Information Processing Systems (NeurIPS 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 (ELBO) 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 a 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 a 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 a backward decomposition of the variational approximation of the states posterior, combined with a representation of the ELBO and its gradients as expectations of value functions satisfying Bellman-type recursions akin to those appearing in Reinforcement Learning (RL) [39].

heading: 2 Background
text: State-Space Models. SSMs are defined by a latentX -valued Markov process (xt)t≥1 and Y-valued observations (yt)t≥1, which are conditionally independent given the latent process. They correspond to the generative model
x1 ∼ µθ(x1), xt+1|xt ∼ fθ(xt+1|xt), yt|xt ∼ gθ(yt|xt), where θ ∈ Rdθ is a parameter of interest and we consider hereX = Rdx . For yt := y1:t = (y1, ..., yt), we thus have
pθ(x1:t, y t) = µθ(x1)gθ(y1|x1) ∏t k=2 fθ(xk|xk−1)gθ(yk|xk).
Assume for the time being that θ is known. Given observations (yt)t≥1 and parameter values θ, one can perform online state inference by computing the posterior of xt given yt which satisfies
pθ(xt|yt) = gθ(yt|xt)pθ(xt|yt−1)
pθ(yt|yt−1) , pθ(xt|yt−1) =
∫ fθ(xt|xt−1)pθ(xt−1|yt−1)dxt−1, (1)
with pθ(x1|y0) := µθ(x1). The log evidence `t(θ) := log pθ(yt) is then given by `t(θ) = ∑t k=1 log pθ(yk|yk−1), where pθ(yk|yk−1) = ∫ gθ(yk|xk)pθ(xk|yk−1)dxk. (2) Here pθ(xk|yk) 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 x1:t given yt and the corresponding representation of the log evidence
pθ(x1:t|yt) = pθ(x1:t, yt)/pθ(yt), `t(θ) = log pθ(yt) = log (∫ pθ(x1:t, y t)dx1:t ) .
Variational Inference. For general SSMs, the filtering and smoothing distributions as well as the log evidence are not available analytically and need to be approximated. For data yt, standard variational approaches use stochastic gradient techniques to maximize the following ELBO
Lt(θ, φ) := Eqφ(x1:t|yt) [ log ( pθ(x1:t, y t) / qφ(x1:t|yt) )] ≤ `t(θ). (3)
Maximizing this ELBO w.r.t. the parameter φ of the variational distribution qφ(x1:t|yt) corresponds to doing variational smoothing while maximizing it w.r.t. θ corresponds to doing parameter learning.
As the true smoothing distribution satisfies pθ(x1:t|yt) = pθ(x1|yt) ∏t k=2 pθ(xk|yt, xk−1), one typi-
cally uses qφ(x1:t|yt) = qφ(x1|yt) ∏t k=2 q
φ(xk|yt, xk−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 qφ(xt|yt) of qφ(x1:t|yt), approximating pθ(xt|yt), is typically not available analytically. Secondly, when the new observation yt+1 is collected, this approach would require recomputing an entirely new variational smoothing distribution with a dimension that increases with time. One can attempt to bypass this problem by restricting ourselves to qφ(x1:t|yt) = qφ(x1|y1) ∏t k=2 q
φ(xk|yk, xk−1) as per [30]. However, the switch from conditioning on yt to yk prohibits learning an accurate approximation of pθ(x1:t|yt) as this formulation does not condition on all relevant data.

heading: 3 Online Variational Filtering and Parameter Learning
text: Our online variational filtering and parameter estimation approach exploits a backward factorization of the variational smoothing distribution qφ(x1:t|yt). The ELBO and its gradients w.r.t. θ and φ are computed in an online manner as t increases by using a combination of ideas from dynamic programming and RL, with a computational time that remains constant at each time step. To simplify notation, henceforth we will write qφt (x1:t) = q φ(x1:t|yt).

heading: 3.1 Backward Decomposition of the Variational Smoothing Distribution
text: The key property that we will be exploiting is that pθ(x1:t|yt) satisfies
pθ(x1:t|yt) = pθ(xt|yt) t−1∏ k=1 pθ(xk|yk, xk+1), pθ(xk|yk, xk+1) = fθ(xk+1|xk)pθ(xk|yk) pθ(xk+1|yk) . (4)
Equation (4) shows that, conditional upon yt, (xk)tk=1 is a reverse-time Markov chain of initial distribution pθ(xt|yt) and backward Markov transition kernels pθ(xk|yk, xk+1); see e.g. [12, 19]. Crucially the backward transition kernel at time k depends only on the observations until time k.
To exploit this, we consider a variational smoothing distribution of the form
qφt (x1:t) = q φ t (xt) ∏t−1 k=1 qφk+1(xk|xk+1), (5)
where qφt (xt) and q φ k+1(xk|xk+1) are variational approximations of the filtering distribution pθ(xt|yt) and the backward kernel pθ(xk|yk, xk+1) respectively. Using (5), one now has
Lt(θ, φ) = `t(θ)− KL(qφt (xt)||pθ(xt|yt)) − ∑t−1
k=1 Eqφt (xk+1)
[ KL(qφk+1(xk|xk+1)||pθ(xk|y k, xk+1)) ] ,
where KL is the Kullback–Leibler divergence and qφt (xk+1) is the marginal distribution of xk+1 under the variational distribution qφt (x1:t) 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 yk at time k and not on future observations. Additionally, it follows directly from (5) that we can easily update qφt+1(x1:t+1) from q φ t (x1:t) using
qφt+1(x1:t+1) = q φ t (x1:t)m φ t+1(xt+1|xt), for m φ t+1(xt+1|xt) :=
qφt+1(xt|xt+1)q φ t+1(xt+1)
qφt (xt) . (6)
Here mφt+1(xt+1|xt) can be viewed as an approximation of the Markov transition density qφt+1(xt+1|xt) ∝ q φ t+1(xt|xt+1)q φ t+1(xt+1) but it is typically not a proper Markov transition density;
i.e. ∫ mφt+1(xt+1|xt)dxt+1 6= 1 as ∫ qφt+1(xt|xt+1)q φ t+1(xt+1)dxt+1 6= q φ t (xt).
Let us assume that qφk (xk) = q φk k (xk) and q φ k (xk−1|xk) = q φk k (xk−1|xk), then (5) and (6) suggest that we only need to estimate φt at time t as yt does not impact the backward Markov kernels prior to time t. However, we also have to be able to compute estimates of ∇φLt(θ, φ) and ∇θLt(θ, φ) to optimize parameters in a constant computational time at each time step, without having to consider the entire history of observations yt. This is detailed in the next subsections where we show that the sequence of ELBOs {Lt(θ, φ)}t≥1 and its gradients {∇θLt(θ, φ)}t≥1 and {∇φLt(θ, φ)}t≥1 can all be computed online when using the variational distributions {qφt (x1:t)}t≥1 defined in (5).

heading: 3.2 Forward recursion for the ELBO
text: We start by presenting a forward-only recursion for the computation of {Lt(θ, φ)}t≥1. This recursion illustrates the parallels between variational inference and RL and is introduced to build intuition. Proposition 1. The ELBO Lt(θ, φ) satisfies for t ≥ 1
Lt(θ, φ) = Eqφt (xt)[V θ,φ t (xt)] for V θ,φ t (xt) := Eqφt (x1:t−1|xt)
[ log ( pθ(x1:t, y t) / qφt (x1:t) )] ,
with the convention V θ,φ1 (x1) := r θ,φ 1 (x0, x1) := log(pθ(x1, y1)/q φ 1 (x1)). Additionally, we have
V θ,φt+1(xt+1) =Eqφt+1(xt|xt+1)[V θ,φ t (xt) + r θ,φ t+1(xt, xt+1)], where (7) rθ,φt+1(xt, xt+1) := log ( fθ(xt+1|xt)gθ(yt+1|xt+1) / mφt+1(xt+1|xt) ) . (8)
All proofs are given in Appendix C. Proposition 1 shows that we can compute Lt(θ, φ), for t ≥ 1, online by recursively computing the functions V θ,φt using (7) and then taking the expectation of V θ,φ t w.r.t. qφt (xt) to obtain the ELBO at time t. Thus, given V θ,φ t , we can compute V θ,φ t+1 and Lt+1(θ, φ) using only yt+1, with a cost that remains constant in t.
This type of recursion is somewhat similar to those appearing in RL. We can indeed think of the ELBO Lt(θ, φ) as the expectation of a sum of “rewards” rθ,φk given in (8) from k = 1 to k = t which we compute recursively using the “value” function V θ,φt . However, while in RL the value function is given by the expectation of the sum of future rewards starting from xt, the value function defined here is the expectation of the sum of past rewards conditional upon arriving in xt. This yields the required forward recursion instead of a backward recursion. We expand on this link in Section 4 and Appendix A.3.

heading: 3.3 Forward recursion for ELBO gradient w.r.t. θ
text: A similar recursion can be obtained to compute {∇θLt(θ, φ)}t≥1. 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 of integration and differentiation are satisfied.
Proposition 2. The ELBO gradient∇θLt(θ, φ) satisfies for t ≥ 1
∇θLt(θ, φ) = Eqφt (xt)[S θ,φ t (xt)], where S θ,φ t (xt) := ∇θV θ,φ t (xt).
Additionally, if we define sθt+1(xt, xt+1) := ∇θr θ,φ t+1(xt, xt+1) = ∇θ log fθ(xt+1|xt)gθ(yt+1|xt+1) then
Sθ,φt+1(xt+1) = Eqφt+1(xt|xt+1) [ Sθ,φt (xt) + s θ t+1(xt, xt+1) ] . (9)
Proposition 2 shows that we can compute {∇θLt(θ, φ)}t≥1 online by propagating {Sθ,φt }t≥1 using (9) and taking the expectation of the vector Sθ,φt w.r.t. q φ t (xt) to obtain the gradient at time t. Similar ideas have been previously exploited in the statistics literature to obtain a forward recursion for the score vector ∇θ`t(θ) so as to perform recursive maximum likelihood parameter estimation; see e.g. [19, Section 4]. In this case, one has ∇θ`t(θ) = Epθ(xt|yt)[Sθt (xt)] where Sθt satisfies a recursion similar to (9) with qφt+1(xt|xt+1) replaced by pθ(xt|yt, xt+1).

heading: 3.4 Forward recursion for ELBO gradient w.r.t. φ
text: 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, qφk (xk) = qφkk (xk) and q φ k (xk−1|xk) = q φk k (xk−1|xk) so Lt(θ, φ) = Lt(θ, φ1:t) and V θ,φ t (xt) = V θ,φ1:t t (xt). At time step t, we will optimize w.r.t. φt and hold all previous φ1:t−1 constant. Our overall variational posterior (5) is denoted qφ1:tt (x1:t). The alternative approach where φ is shared through time (amortization) is investigated in Appendix D.
Since the expectation is taken w.r.t. qφ1:tt (x1:t) in Lt, optimizing φt is slightly more difficult than for θ. However, we can still derive a forward recursion for the φ gradients and we will leverage the reparameterization trick to reduce the variance of the gradient estimates; i.e. we assume that xt(φt; t) ∼ qφtt (xt) and xt−1(φt; t−1, xt) ∼ q φt t (xt−1|xt) when t−1 ∼ λ( ), t ∼ λ( ).
Proposition 3. The ELBO gradient∇φtLt(θ, φ1:t) satisfies for t ≥ 1
∇φtLt(θ, φ1:t) = ∇φtEqφtt (xt)[V θ,φ1:t t (xt)] = Eλ( t)[∇φtV θ,φ1:t t (xt(φt; t))].
Additionally, one has
∇φt+1V θ,φ1:t+1 t+1 (xt+1(φt+1; t+1)) = Eλ( t) [ T θ,φ1:tt (xt(φt+1; t, xt+1(φt+1; t+1)))
dxt(φt+1; t, xt+1(φt+1; t+1)) dφt+1
+ ∇φt+1r θ,φt:t+1 t+1 (xt(φt+1; t, xt+1(φt+1; t+1)), xt+1(φt+1; t+1))
] ,
where T θ,φ1:tt (xt) := ∂ ∂xt V θ,φ1:tt (xt) satisfies the forward recursion
T θ,φ1:t+1 t+1 (xt+1) = Eλ( t) [ T θ,φ1:tt (xt(φt+1; t, xt+1)) ∂xt(φt+1; t, xt+1)
∂xt+1
+ ∇xt+1r θ,φt:t+1 t+1 (xt(φt+1; t, xt+1), xt+1)
] . (10)
Here, dxt(φt+1; t,xt+1(φt+1; t+1))dφt+1 , ∂xt(φt+1; t,xt+1) ∂xt+1 are Jacobians of appropriate dimensions.

heading: 3.5 Estimating the ELBO and its Gradients
text: As we consider the case qφk (xk) = q φk k (xk), q φ k (xk−1|xk) = q φk k (xk−1|xk), we have S θ,φ t (xt) = Sθ,φ1:tt (xt). At time t we optimize φt and hold φ1:t−1 constant. Practically, we are not able to compute in closed-form the functions V θ,φ1:tt (xt), S θ,φ1:t t (xt) and T θ,φ1:t t (xt) appearing in the forward recursions of Lt(θ, φ1:t),∇θLt(θ, φ1:t) and∇φLt(θ, φ1:t) respectively. However, we can exploit the above recursions to approximate these functions online using regression as is commonly done in RL. We then show how to use these gradients for online filtering and parameter learning.
We approximate Sθ,φ1:t+1t+1 with Ŝt+1. Equation (9) shows that Ŝt+1 can be learned using Ŝt through regression of the simulated dataset2 {xit+1, Ŝt(xit) + sθt+1(xit, xit+1)} with (xit, xit+1)
i.i.d.∼ q φt+1 t+1 (xt, xt+1) for i = 1, ..., N (see Appendix A.1 for derivation). We can use neural networks to model Ŝt or Kernel Ridge Regression (KRR). Note that the use of KRR to estimate gradients for variational learning has recently been demonstrated by [26].
We similarly approximate T θ,φ1:t+1t+1 with T̂t+1. As before, we can model T̂t+1 using neural networks or KRR. We use recursion (10) and T̂t to create the following dataset for regression3{
xit+1, T̂t(xt(φt+1; i t, x i t+1))
∂xt(φt+1; i t,x i t+1)
∂xit+1 +∇xit+1r θ,φt:t+1 t+1 (xt(φt+1; i t, x i t+1), x i t+1)
} ,
where xit+1 ∼ q φt+1 t+1 (xt+1) and i t ∼ λ( ) for i = 1, ..., N . The choice of the distribution over the inputs, xit+1, for each simulated dataset is arbitrary but it will determine where the approximations are most accurate. We choose qφt+1t+1 (xt+1) to best match where we expect the approximations to be evaluated, more details are given in Appendix A.1.
Note that if one is interested in computing online an approximation of the ELBO, we can again similarly approximate V θ,φ1:tt (xt) using regression to obtain V̂t(xt) by leveraging (7). We will call the resulting approximate ELBO the Recursive ELBO (RELBO). We could also then differentiate V̂t(xt) w.r.t. xt to obtain an alternative method for estimating T θ,φ1:t t and optimizing φt. However, as we are ultimately interested in accurate gradients, this approach does not exploit the readily available gradient information during the regression stage.
By approximating T θ,φ1:t+1t+1 with T̂t+1 and S θ,φ1:t+1 t+1 with Ŝt+1, we introduce some bias into our gradient estimates. We can trade bias for variance by using modified recursions; e.g.
S θ,φ1:t+1 t+1 (xt+1) = Eqφt−L+2:t+1t (xt−L+1:t|xt+1)
[ S θ,φ1:t−L+1 t−L+1 (xt−L+1) + ∑t k=t−L+1 s θ k+1(xk, xk+1) ] .
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 = 1.
2We define Ŝ0 := 0, xi0 := ∅, sθ1(x0, x1) = ∇θ logµθ(x1)gθ(y1|x1) and qφ11 (x0, x1) := q φ1 1 (x1). 3Similarly, we define T̂0 := 0, xi0 := ∅ and rθ,φ0:11 as in Proposition 1.
Algorithm 1: Online Variational Filtering and Parameter Learning. for t = 1, . . . , T do
Initialize φt e.g. φt ← φt−1 /* Update φt using M stochastic gradient steps */ for m = 1, . . . ,M do
Sample xit−1, xit ∼ qφtt (xt−1, xt) using reparameterization trick for i = 1, ..., N φt ← φt + γm 1N ∑N i=1{T̂t−1(x i t−1) dxit−1 dφt
+∇φtrt(xit−1, xit)} end /* Update T̂t(xt) and Ŝt(xt) as in Section 3.5 */ T̂t(xt) regression← T̂t−1(xt−1) ∂xt−1∂xt +∇xtrt(xt−1, xt). Ŝt(xt) regression← Ŝt−1(xt−1) + s θt−1 t (xt−1, xt). /* Update θ using a stochastic gradient step */ Sample xit−1, xit ∼ qφtt (xt−1, xt), x̃it−1 ∼ q φt−1 t−1 (xt−1) for i = 1, ..., N
θt ← θt−1 + ηt 1N ∑N i=1{Ŝt−1(x i t−1) + s θt−1 t (x i t−1, x i t)− Ŝt−1(x̃it−1)}
end

heading: 3.6 Online Parameter Estimation
text: Assume, for the sake of argument, that one has access to the log evidence `t(θ) and that the observations arise from the SSM with parameter θ?. Under regularity conditions, the average loglikelihood `t(θ)/t converges as t→∞ towards a function `(θ) which is maximized at θ?; see e.g. [12, 41]. We can maximize this criterion online using Recursive Maximum Likelihood Estimation (RMLE) [19, 24, 40, 41] which consists of updating the parameter estimate θ using
θt+1 = θt + ηt+1 ( Epθ0:t (xt,xt+1|yt+1)[St(xt) + s θt t+1(xt, xt+1)]− Epθ0:t−1 (xt|yt)[St(xt)] ) . (11)
The difference of two expectations on the r.h.s. of (11) is an approximation of the gradient of the log-predictive log pθ(yt+1|yt) evaluated at θt. The approximation is given by ∇ log pθ0:t(yt+1) − ∇ log pθ0:t−1(yt) with the notation ∇ log pθ0:t(yt+1) corresponding to the expectation of the sum of terms sθkk+1(xk, xk+1) w.r.t. the joint posterior states distribution defined by using the SSM with parameter θk at time k + 1.
We proceed similarly in the variational context and update the parameter using
θt+1 =θt+ηt+1 ( E q φt+1 t+1 (xt,xt+1) [ Ŝt(xt)+s θt t+1(xt, xt+1) ] −E q φt t (xt) [ Ŝt(xt) ]) .
Here Ŝt(xt) approximates St(xt) satisfying St+1(xt+1) := Eqφt+1t+1 (xt|xt+1) [St(xt)+s θt t+1(xt, xt+1)]. We compute Ŝt as in Section 3.5 with a simulated dataset using Ŝt−1 and θt−1.
Putting everything together, we summarize our method in Algorithm 1 using a simplified notation to help build intuition, a full description is given in Appendix A.2. It is initialized using initial parameters φ1, θ0. We re-iterate the algorithm’s computational cost does not grow with t. We need only store fixed size T̂ and Ŝ models as well as the most recent φt and θt parameters. When performing backpropagation, T̂ and Ŝ summarize all previous gradients, meaning we do not have to roll all the way back to t = 1. Therefore, we only incur standard backpropagation computational cost w.r.t. φt and θ. To scale to large dx, we can use mean field q φt t (xt), q φt t (xt−1|xt) keeping costs linear in dx.

heading: 4 Related Work
text: 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 a value function and corresponding Bellman recursion given by
V RLt (st) := E [ T∑ k=t r(sk, ak) ] , V RLt (st) = Est+1,at [ r(st, at) + V RL t+1(st+1) ] , (12)
where (st, at) 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
Vt(st) := E
[ t∑
k=1
r(sk, ak) ] , Vt+1(st+1) = Eat+1,st [r(st+1, at+1) + Vt(st)] . (13)
Note we have defined an anti-causal graphical model, with st depending on st+1 and at+1. If we further let st = xt, at = xt−1 ∼ qφt (xt−1|xt), P (st|st+1, at+1) = δ(st = at+1) and r(st, at) = r(xt, xt−1) as in equation (8), then the recursion in equation (13) for Vt corresponds to the recursion in equation (7) and LT = EqφT (xT ) [VT (xT )]. We then differentiate this recursion with respect to θ and xt+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 use 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 `t(θ) and lower bound each term log pθ(yk|yk−1) appearing in the sum using
Eqφk (xk−1,xk)
[ log
fθ(xk|xk−1)gθ(yk|xk)pθ(xk−1|yk−1) qφk (xk−1, xk)
] ≤ log pθ(yk|yk−1). (14)
Unfortunately, the term on the l.h.s. of (14) cannot be evaluated unbiasedly as it relies on the intractable filter pθ(xk−1|yk−1) so [45] approximates it by pθ(xk−1|yk−1) ≈ qφk−1(xk−1) to obtain the following Approximate ELBO (AELBO) by summing over k:
L̃t(θ, φ1:t) = ∑t
k=1 Eqφk (xk−1,xk)
[ log rθ,φk (xk−1, xk) ] . (15)
Additionally, [45] makes the assumption qφk (xk−1, xk) = q φ k−1(xk−1)q φ k (xk) and we will refer to (15) in this case as AELBO-1. While [16] does not consider online learning, their objective is actually a generalization of [45], with qφk (xk−1, xk) = q φ k (xk)q φ k (xk−1|xk), and we will refer to (15) in this case as AELBO-2. It can be easily shown that AELBO-2 is only equal to the true ELBO given in Proposition 1 in the (unrealistic) scenario where qφk (xk) = pθ(xk|yk) for all k. Moreover, in both cases the term pθ(xk−1|yk−1) is replaced by qφk−1(xk−1), causing a term involving θ to be ignored in gradient computation. The approach developed here can be thought of as a 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].

heading: 5 Experiments

heading: 5.1 Linear Gaussian State-Space Models
text: We first consider a 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
fθ(xt|xt−1) = N (xt;Fxt−1, U), gθ(yt|xt) = N (yt;Gxt, V ),
where F ∈ Rdx×dx , G ∈ Rdy×dx , U ∈ Rdx×dx , V ∈ Rdy×dy , θ = {F,G}. We let qφtt (xt) = N ( xt;µt, diag(σ2t ) ) , qφtt (xt−1|xt) = N ( xt−1;Wtxt + bt, diag(σ̃2t ) ) ,
with φt = {µt, log σt,Wt, bt, log σ̃t}. When fθ(xt|xt−1) is linear Gaussian and we use a Gaussian for qφt−1t−1 (xt−1), we can select q φt t (xt−1|xt) ∝ fθ(xt|xt−1)q φt−1 t−1 (xt−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 qφtt (xt−1|xt) as well as setting the matrices F,G,U, V to be diagonal, so that pθ(xt|yt) and pθ(xt−1|yt−1, xt) are in the variational family. For dx = dy = 10, we first demonstrate accurate state inference by learning φt at each time step whilst holding θ fixed at the true value. We represent T̂t(xt) non-parametrically using KRR. Full details for all experiments are given in Appendix B.4 Figure 1a 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 a variation of Algorithm 1 using AELBO-1 and 2 (see Section 4). Our methodology converges much closer to the analytic baseline (RMLE) than AELBO-2 [16] and exhibits less variance, even though the variational family is 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 xk−1 and xk. As expected, our method performs slightly worse than the analytic RMLE, as inevitably small errors will be introduced during stochastic optimization and regression.

heading: 5.2 Chaotic Recurrent Neural Network
text: 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 dx = 5, 20, and 100. This non-linear model is an Euler approximation of the continuous-time recurrent neural network dynamics
f(xt|xt−1) = N ( xt;xt−1 + ∆τ −1 (−xt−1 + γW tanh(xt−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φtt (xt−1|xt) = N (xt−1; MLP φt t (xt), diag(σ̃ 2 t )) and q φt t (xt) =
N ( xt;µt, diag(σ2t ) ) where we use a 1-layer Multi-Layer Perceptron (MLP) with 100 neurons for each qφtt (xt−1|xt). We generate a dataset of length 100 using the same settings as [44], and each algorithm is 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 dx, as expected. Among variational methods, AELBO-1 does not give as 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
4Code available at https://github.com/andrew-cr/online_var_fil
θ is fixed - an effect not represented using just the RMSE. We confirm this is the case in Appendix B.2 by comparing our variational filter means µt against the ‘ground truth’ posterior mean for dx = 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 of the ELBO, 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 qφtt (xt), q φt t (xt−1|xt) and hold q φk k (xk−1|xk), k < t fixed. This is because all information to learn the true factor pθ(xk−1|yk−1, xk) is available at time k as it does not depend on future observations. However, when pθ(xk−1|yk−1, xk) is not in the variational family (as in this CRNN example), qφkk (xk−1|xk) will aim to be most accurate in the regions of xk in which it is typically evaluated. The evaluation distribution over xk does depend on future variational factors and so could shift over time. This may result in learning a different variational joint qφ1:tt (x1:t) between when using our online ELBO 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 or the offline ELBO (3). Figure 2b plots the final mean and standard deviations of the marginals of the trained qφ1:tt (x1:t) 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 xk 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.

heading: 5.3 Sequential Variational Auto-Encoder
text: We demonstrate the scalability of our method on a sequential VAE application. In this problem, an agent observes a long sequence of frames that could, for example, come from a robot traversing a 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
fθ(xt|xt−1) = N (xt; NNfθ (xt−1), U), g(yt|xt) = N (yt; NN g(xt), V ),
where dx = 32, NN f θ is a residual MLP and NN g a convolutional neural network. NNfθ is learned online whilst NNg 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φtt parameterization as for the CRNN but with a 2 hidden layer MLP with 64 neurons. KRR is used to learn T̂t whereas we use an MLP for learning Ŝt. We found that MLPs scale better than KRR as dθ is high. Our online algorithm is run on a sequence of 4000 images after which we probe the quality of the learned NNfθ . The results are
shown in Figure 3. Before training, NNfθ 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 `t(θ)/t computed using Monte Carlo increases through training, thereby confirming our method can successfully learn high-dimensional model parameters of the agent movement in a fully online fashion.

heading: 6 Discussion
text: We have presented a 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 of a neural network in a sequential VAE. However, it is not without its limitations. As with any stochastic variational inference technique, we can only obtain an accurate posterior approximation if the variational family is expressive enough and our stochastic gradient method finds a 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 as it requires solving an optimization problem and solving a regression task for each time step.
To reduce this computational cost, one can amortize the optimization cost for φ by learning a network taking a representation of observations up to now, yt, and outputting qφt as illustrated in Appendix D. Further work will investigate ways to also amortize the cost of function regression across time through a 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 the asymptotic properties of the parameter learning procedure.
Acknowledgments and Disclosure of Funding
The authors are grateful to Adrien Corenflos, Desi Ivanova and James Thornton for their comments. Andrew Campbell acknowledges support from the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1). Arnaud Doucet is partly supported by the EPSRC grant EP/R034710/1. He also acknowledges support of the UK Defence Science and Technology Laboratory (DSTL) and EPSRC under grant EP/R013616/1. This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative.

Review summaries:
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.

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.

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.