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 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.