Recurrent Kalman Networks: Factorized Inference in High-Dimensional Deep Feature Spaces

1 INTRODUCTION . State-estimation in unstructured environments is a very challenging task as observations or measurements of the environment are often high-dimensional and only provide partial information about the state . Images are a good example : Even for low resolution , the number of pixels can quickly exceed tens or hundreds of thousands and it is impossible to obtain any information about the dynamics , such as velocities , from a single image . Additionally , the observations may be noisy or may not contain useful information for the task at hand . Such noise can for example be introduced by poor illumination or motion blur and occlusions can prevent us from observing some or all relevant aspects of the scene . In addition to state estimation , it is also often desirable to predict future states or observations , for example , in order to assess the consequences of future actions . To this end , an initial estimate of the current state is necessary which again has to be inferred from observations . In such environments we typically also have to deal with a high uncertainties in the state estimates . Being able to model this uncertainty is crucial in many decision making scenarios , e.g. , if we need to decide to perform an action now or wait until more information about the scene is available . Deep learning models have been very successful for time-series modelling in unstructured environments . Classical models such as LSTMs ( Hochreiter and Schmidhuber , 1997 ) or GRUs ( Cho et al. , 2014 ) perform well but fail to capture the uncertainty of the state estimate . Recent probabilistic deep learning approaches have used the Kalman filter ( KF ) as a tool to integrate uncertainty estimates into deep time-series modelling ( Haarnoja et al. , 2016 ; Watter et al. , 2015 ; Archer et al. , 2015 ; Fraccaro et al. , 2017 ; Krishnan et al. , 2017 ) . These approaches use the KF to perform inference in a low-dimensional ( latent ) state space that is typically defined by a deep encoder . However , using KF in such a state space comes with two main limitations . In order to be usable for non-linear dynamics , we have to introduce approximations such as the extended KF ( Haarnoja et al. , 2016 ) and variational inference methods ( Krishnan et al. , 2017 ; Fraccaro et al. , 2017 ) . Moreover , the KF equations require computationally expensive matrix inversions that are hard to scale to high dimensional latent spaces for more complex systems and computationally demanding to fully backpropagate in an end-to-end manner . Most of these methods are implemented as ( variational ) auto-encoders and are therefore also limited to predicting future observations or imputing missing observations and can not be directly be applied to state estimation . We introduce the Recurrent Kalman Network , an end-to-end learning approach for Kalman filtering and prediction . While Kalman filtering in the original state space requires approximations due to the non-linear models as well as matrix inversions that are hard to back-propagate and computationally expensive , the RKN uses a learned high-dimensional latent state representation that allows for efficient inference using locally linear transition models and a factorized belief state representation which avoids expensive and numerically problematic matrix inversions . Conceptually , this idea is related to kernel methods which use high-dimensional feature spaces to approximate nonlinear functions with linear models ( Gebhardt et al. , 2017 ) . However , in difference to kernel feature spaces , our feature space is given by a deep encoder and learned in an end-to-end manner . The RKN can be used for any time-series data set for which LSTMs and GRUs are currently the state of the art . In contrast to those , the RKN uses an explicit representation of uncertainty which governs the gating between keeping the current information in memory or updating it with the current observation . While the RKN shows a slightly improved performance in terms of state estimation errors , both LSTMs and GRUs struggle with estimating the uncertainty of the prediction while the RKN can provide accurate uncertainty estimates . In relation to existing KF-based approaches , our approach can be used for state estimation as well as for generative tasks such as image imputation . We also show that we outperform state of the art methods on a complex image imputation task . 1.1 RELATED WORK . Using encoders for time-series modelling of high-dimensional data such as images is a common approach . Such encoders can also be easily integrated in well known deep time-series models such as LSTMs ( Hochreiter and Schmidhuber , 1997 ) or GRUs ( Cho et al. , 2014 ) . These models are very effective but do not provide good uncertainty estimates as shown in our experiments . Pixel to Torques ( P2T ) ( Wahlström et al. , 2015 ) employs an autoencoder to obtain low dimensional latent representations from images together with a transition model . They subsequently use the models to perform control in the latent space . Embed to Control ( E2C ) ( Watter et al. , 2015 ) can be seen as an extension of the previous approach with the difference that a variational autoencoder ( Kingma and Welling , 2013 ) is used . However , both of these approaches are not recurrent and rely on observations which allow inferring the whole state from a single observation . They can therefore not deal with noisy or missing data . Another family of approaches interprets encoder-decoder models as latent variable models that can be optimized efficiently by variational inference . They derive a corresponding lower bound and optimize it using the stochastic gradient variational Bayes approach ( Kingma and Welling , 2013 ) . Black Box Variational Inference ( BB-VI ) ( Archer et al. , 2015 ) proposes a structured Gaussian variational approximation of the posterior , which simplifies the inference step at the cost of maintaining a tri-diagonal covariance matrix of the full state . To circumvent this issue , Structured Inference Networks ( SIN ) ( Krishnan et al. , 2017 ) employ a flexible recurrent neural network to approximate the dynamic state update . Deep Variational Bayes Filters ( DVBF ) ( Karl et al. , 2016 ) integrate general Bayes filters into deep feature spaces while the Kalman Variational Autoencoder ( KVAE ) ( Fraccaro et al. , 2017 ) employs the classical Kalman Filter and allows not only filtering but also smoothing . Variational Sequential Monte Carlo ( VSMC ) ( Naesseth et al. , 2017 ) uses particle filters instead , however they are only learning the proposal function and are not working in learned latent spaces . Yet , these models can not be directly used for state estimation as they are formulated as generative models of the observations without the notion of the real state of the system . Moreover , the use of variational inference introduces an additional approximation that might affect the performance of the algorithms . The approaches given in ( Archer et al. , 2015 ; Fraccaro et al. , 2017 ; Haarnoja et al. , 2016 ) directly use the Kalman update equations in the latent state , which limits these approaches to rather low dimensional latent states due to the expensive matrix inversions . The BackpropKF ( Haarnoja et al. , 2016 ) applies a CNN to estimate the observable parts of the true state given the observation . Similar to our approach , this CNN additionally outputs a covariance matrix indicating the models certainty about the estimate and allows the subsequent use of an ( extended ) Kalman filter with known transition model . In contrast , we let our model chose the feature space that is used for the inference such that locally linear models can be learned and the KF computations can be simplified due to our factorization assumptions . A summary of all the approaches , listing also their basic properties , can be seen in Table 1 . We compare the approaches whether they are scaleable to high dimensional latent states , whether they can be used for state estimation , whether they can provide uncertainty estimates , whether they can handle noise or missing data and whether the objective can be optimized directly or via a lower bound . All probabilistic generative models rely on variational inference which optimizes a lower bound , which potentially affects the performance of the algorithms . The P2T and E2C approaches rely on the Markov assumption and therefore can not deal with noise ( or need very large window sizes ) . Traditional recurrent models ( LSTMs , GRUs ) can be trained directly by backpropagation through time and therefore typically yield very good performance but are lacking uncertainty estimates ( which , however , can be artificially added as in our experiments ) . Our RKN approach combines the advantages of all methods above as it can be learned by direct optimization without the use of a lower bound and it provides a principled way of representing uncertainty inside the neural network . 2 FACTORIZED INFERENCE IN DEEP FEATURE SPACES . Lifting the original input space to a high-dimensional feature space where linear operations are feasible is a common approach in machine learning , e.g. , in kernel regression and SVMs . The Recurrent Kalman Network ( RKN ) transfers this idea to state estimation and filtering , i.e. , we learn a high dimensional deep feature space that allows for efficient inference using the Kalman update equations even for complex systems with high dimensional observations . To achieve this , we assume that the belief state representation can be factorized into independent Gaussian distributions as described in the following sections . 2.1 LATENT OBSERVATION AND STATE SPACES . The RKN encoder learns a mapping to a high-dimensional latent observation space W = Rm . The encoder also outputs a vector of uncertainty estimates σobst , one for each entry of the latent observation wt . Hence , the encoder can be represented by ( wt , σobst ) = enc ( ot ) . The latent state space Z = Rn of the RKN is related to the observation space W by the linear latent observation model H = [ Im 0m× ( n−m ) ] , i.e. , w = Hz with w ∈ W and z ∈ Z . Im denotes the m ×m identity matrix and 0m× ( n−m ) denotes a m× ( n−m ) matrix filled with zeros . The idea behind this choice is to split the latent state vector zt into two parts , a vector pt for holding information that can directly be extracted from the observations and a vector mt to store information inferred over time , e.g. , velocities . We refer to the former as the observation or upper part and the latter as the memory or lower part of the latent state . For an ordinary dynamical system and images as observations the former may correspond to positions while the latter corresponds to velocities . Clearly , this choice only makes sense for m ≤ n and in this work we assume n = 2m , i.e. , for each observation unit pi , we also represent a memory unit mi that stores its velocity information . For each sequence , we initialize z+0 with an all zeros vector and Σ + 0 with 10 · I . In practice , it is beneficial to normalize wt since the statistics of noisy and noise free images differ considerably . 2.2 THE TRANSITION MODEL . To obtain a locally linear transition model we learnK constant transition matrices A ( k ) and combine them using state dependent coefficients α ( k ) ( zt ) , i.e. , At = ∑K k=0 α ( k ) ( zt ) A ( k ) . A small neural network with softmax output is used to learn α ( k ) . Similar approaches are used in ( Fraccaro et al. , 2017 ; Karl et al. , 2016 ) . Using a dense transition matrix in high-dimensional latent spaces is not feasible as it contains too many parameters and causes numerical instabilities and overfitting , as preliminary experiments showed . Therefore , we design each A ( k ) to consist of four band matrices A ( k ) = [ B ( k ) 11 , B ( k ) 12 ; B ( k ) 21 , B ( k ) 22 ] with bandwidth b . This choice reduces the number of parameters while not affecting performance since the network is free to choose the state representation . We assume the covariance of the transition noise to be diagonal and denote the vector containing the diagonal values by σtrans . The noise is learned and independent of the state . Moreover , it is crucial to correctly initialize the transition matrix . Initially , the transition model should focus on copying the encoder output so that the encoder can learn how to extract good features if observations are available and useful . On the other hand it is crucial that A does not yield an instable system . We choose B ( k ) 11 = B ( k ) 22 = I and B ( k ) 12 = −B ( k ) 21 = 0.2 · I .

Summary: This paper proposes a method for inferring the latent state and making predictions based on a sequence of observations. The idea is to map the observation to a latent space where the relation to the latent state is linear, and the dynamics of the latent state are locally linear. Therefore, in this latent space a Kalman filter can be applied to infer the current state and predict the next state, including uncertainty estimates. Finally, the predicted latent state is mapped to a prediction for the observation or some other variable of interest.