Online Variational Filtering and Parameter Learning

1 Introduction . 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 can not 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 ] . 2 Background . 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 . 3 Online Variational Filtering and Parameter Learning . 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 ) . 3.1 Backward Decomposition of the Variational Smoothing Distribution . 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 ) .

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.