Self-Supervised Inference in State-Space Models

We perform approximate inference in state-space models with nonlinear state transitions . Without parameterizing a generative model , we apply Bayesian update formulas using a local linearity approximation parameterized by neural networks . It comes accompanied by a maximum likelihood objective that requires no supervision via uncorrupt observations or ground truth latent states . The optimization backpropagates through a recursion similar to the classical Kalman filter and smoother . Additionally , using an approximate conditional independence , we can perform smoothing without having to parameterize a separate model . In scientific applications , domain knowledge can give a linear approximation of the latent transition maps , which we can easily incorporate into our model . Usage of such domain knowledge is reflected in excellent results ( despite our model ’ s simplicity ) on the chaotic Lorenz system compared to fully supervised and variational inference methods . Finally , we show competitive results on an audio denoising experiment . 1 INTRODUCTION . Many sequential processes in industry and research involve noisy measurements that describe latent dynamics . A state-space model is a type of graphical model that effectively represents such noise-afflicted data ( Bishop , 2006 ) . The joint distribution is assumed to factorize according to a directed graph that encodes the dependency between variables using conditional probabilities . One is usually interested in performing inference , meaning to obtain reasonable estimates of the posterior distribution of the latent states or uncorrupt measurements . Approaches involving sampling ( Neal et al. , 2011 ) , variational inference ( Kingma & Welling , 2013 ) , or belief propagation ( Koller & Friedman , 2009 ) have been proposed before . Assuming a hidden Markov process ( Koller & Friedman , 2009 ) , the celebrated Kalman filter and smoother ( Kalman , 1960 ; Rauch et al. , 1965 ) are classical approaches to solving the posterior inference problem . However , the Markov assumption , together with linear Gaussian transition and emission probabilities , limit their flexibility . We present filtering and smoothing methods that are related to the classical Kalman filter updates but are augmented with flexible function estimators without using a constrained graphical model . By noting that the filtering and smoothing recursions can be back-propagated through , these estimators can be trained with a principled maximum-likelihood objective reminiscent of the noise2noise objective ( Lehtinen et al. , 2018 ; Laine et al. , 2019 ) . By using a locally linear transition distribution , the posterior distribution remains tractable despite the use of non-linear function estimators . Further , we show how a linearized smoothing procedure can be applied directly to the filtering distributions , discarding the need to train a separate model for smoothing . To verify what is claimed , we perform three experiments . ( 1 ) A linear dynamics filtering experiment , where we show how our models approximate the optimal solution with sufficient data . We also report that including expert knowledge can yield better estimates of latent states . ( 2 ) A more challenging chaotic Lorenz smoothing experiment that shows how our models perform on par with recently proposed supervised models . ( 3 ) An audio denoising experiment that uses real-world noise showing practical applicability of the methods . Our contributions can be summarized as follows . 1 . We show that the posterior inference distribution of a state-space model is tractable while parameter estimation is performed by neural networks . This means that we can apply the classical recursive Bayesian updates , akin to the Kalman filter and smoother , with mild assumptions on the generative process . 2 . Our proposed method is optimized using maximum likelihood in a self-supervised manner . That is , ground truth values of states and measurements are not assumed to be available for training . Still , despite our model ’ s simplicity , our experiments show that it performs better or on par with several baselines . 3 . We show that the model can be combined with prior knowledge about the transition and emission probabilities , allowing for better applicability in low data regimes and incentivizing the model to provide more interpretable estimates of the latent states . 4 . A linearized smoothing approach is presented that does not require explicit additional parameterization and learning of the smoothing distribution . 2 RELATED WORK . Becker et al . ( 2019 ) provide a detailed discussion of recent related work , which we build on here and in table 1 . An early method that extends the earlier introduced Kalman filter by allowing nonlinear transitions and emissions is the Extended Kalman filter ( Ljung , 1979 ) . It is limited due to the naive approach to locally linearize the transition and emission distributions . Furthermore , the transition and emission mechanisms are usually assumed to be known , or estimated with Expectation Maximization ( Moon , 1996 ) . More flexible methods that combine deep learning with variational inference include Black Box Variational Inference ( Archer et al. , 2015 ) , Structured Inference Networks ( Krishnan et al. , 2017 ) , Kalman Variational Autoencoder ( Fraccaro et al. , 2017 ) , Deep Variational Bayes Filters ( Karl et al. , 2017 ) , Variational Sequential Monte Carlo ( Naesseth et al. , 2018 ) and Disentangled Sequential Autoencoder ( Yingzhen & Mandt , 2018 ) . However , the lower-bound objective makes the approach less scalable and accurate ( see also Becker et al . ( 2019 ) . Furthermore , all of the above methods explicitly assume a graphical model , imposing a strong but potentially harmful inductive bias . The BackpropKF ( Haarnoja et al. , 2016 ) and Recurrent Kalman Network ( Becker et al. , 2019 ) move away from variational inference and borrow Bayesian filtering techniques from the Kalman filter . We follow this direction but do not require supervision through ground truth latent states or uncorrupt emissions . Satorras et al . ( 2019 ) combine Kalman filters through message passing with graph neural networks to perform hybrid inference . We perform some of their experiments by also incorporating expert knowledge . However , contrary to their approach , we do not need supervision . Finally , concurrently to this work , Revach et al . ( 2021 ) develop KalmanNet . It proposes similar techniques but evaluates them in a supervised manner . The authors , however , do suggest that an unsupervised approach can also be feasible . Additionally , we more explicitly state what generative assumptions are required , then target the posterior distribution of interest , and develop the model and objective function from there . Moreover , the current paper includes linearized smoothing ( section 6 ) , parameterized smoothing ( appendix A ) , and the recurrent model ( appendix C ) . We also denote theoretical guarantees under the noise2noise objective . scalable state est . uncertainty noise dir . opt . self-sup . Ljung ( 1979 ) ✓/× ✓ ✓ ✓ × × Hochreiter et . al . ( 1997 ) ✓ ✓ ✓/× ✓ ✓ × Cho et al . ( 2014 ) ✓ ✓ ✓/× ✓ ✓ × Wahlström et al . ( 2015 ) ✓ ✓ ✓/× × ✓ × Watter et al . ( 2015 ) ✓ × ✓ ✓ × ✓ Archer et al . ( 2015 ) ✓/× × ✓ ✓ × ✓ Krishnan et al . ( 2017 ) ✓ × ✓ ✓ × ✓ Fraccaro et al . ( 2017 ) ✓/× × ✓ ✓ × ✓ Karl et al . ( 2017 ) ✓ × ✓ ✓ × ✓ Naesseth et al . ( 2018 ) ✓ × ✓ ✓ × ✓ Yingzhen et al . ( 2018 ) ✓ × ✓ × × ✓ Rangapuram et al . ( 2018 ) ✓/× ✓ ( 1D ) ✓ × ✓ ✓ Doerr et al . ( 2018 ) × ✓ ✓ ✓ ✓ ✓ Satorras et al . ( 2019 ) ✓ ✓ × ✓ ✓ × Haarnoja et al . ( 2016 ) ✓ ✓ ✓ ✓ ✓ × Becker et al . ( 2019 ) ✓ ✓ ✓ ✓ ✓ × Ours ✓/× ✓ ✓ ✓ ✓ ✓ Table 1 : We compare whether algorithms are scalable , state estimation can be performed , models provide uncertainty estimates , noisy or missing data can be handled , optimization is performed directly and if supervision is required . “ ✓/× ” means that it depends on the parameterization . 3 GENERATIVE MODEL ASSUMPTIONS . In this section , we explicitly state the model ’ s generative process assumptions . First , we assume that we can measure ( at least ) one run of ( noise-afflicted ) sequential data y0 : K : = ( y0 , . . . , yK ) , where each yk ∈ RM , k = 0 , . . . , K. We abbreviate : yl : k : = ( yl , . . . , yk ) and y < k : = y0 : k−1 and y≤k : = y0 : k and y−k : = ( y0 : k−1 , yk+1 : K ) . We then assume that y0 : K is the result of some possibly non-linear probabilistic latent dynamics , i.e. , of a distribution p ( x0 : K ) , whose variables are given by x0 : K : = ( x0 , . . . , xK ) with xk ∈ RN . Each yk is assumed to be drawn from some shared noisy emission probability p ( yk | xk ) . The joint probability is then assumed to factorize as : p ( y0 : K , x0 : K ) = p ( x0 : K ) K∏ k=0 p ( yk | xk ) . ( 1 ) Further implicit assumptions about the generative model are imposed via inference model choices ( see section 7 ) . Note that this factorization encodes several conditional independences like yk ⊥⊥ ( y−k , x−k ) | xk . ( 2 ) Typical models that follow these assumptions are linear dynamical systems , hidden Markov models , but also nonlinear state-space models with higher-order Markov chains in latent space like presented in fig . 2 . In contrast to other approaches ( e.g. , Krishnan et al . ( 2015 ) ; Johnson et al . ( 2016 ) ; Krishnan et al . ( 2017 ) ) where one tries to model the latent dynamics with transition probabilities p ( xk | xk−1 ) and possibly non-linear emission probabilities p ( yk | xk ) , we go the other way around . We assume that all the non-linear dynamics are captured inside the latent distribution p ( x0 : K ) , where at this point we make no further assumption about its factorization , and the emission probabilities are ( wellapproximated with ) a linear Gaussian noise model : p ( yk | xk ) = N ( yk |Hxk , R ) , ( 3 ) where the matrix H represents the measurement device and R is the covariance matrix of the independent additive noise . We make a brief argument why this assumption is not too restrictive . First , if one is interested in denoising corrupted measurements , any nonlinear emission can be captured directly inside the latent states xk . To see this , let zk ∈ RN−M denote non-emitted state variables . We then put xk : = [ yk zk ] ⊤ , where yk is computed by applying the nonlinear emission to zk . We thus include the measurements in the modeled latent state xk . Then we can put H : = [ IM 0M× ( N−M ) ] . Second , techniques proposed by Laine et al . ( 2019 ) allow for nonGaussian noise models , relaxing the need for assumption eq . ( 3 ) . Third , we can locally linearize the emission ( Ljung , 1979 ) . Finally , industrial or academic applications include cases where emissions are ( sparse ) Gaussian measurements and the challenging nonlinear dynamics occur in latent space . Examples can be found in MRI imaging ( Lustig et al. , 2007 ) and radio astronomy ( Thompson et al. , 2017 ) .

Summary: A Kalman filter typically requires specifying a state update equation, along with noise covariance matrices for state transition and observation. Determining these can be a challenge in certain scenarios. The manuscript proposes to replace these unknowns with trainable neural networks. The input to the neural networks are the observations, and the output are the unknown parameters (or their factorizations, to ensure positive definiteness). The manuscript validates the approach through numerical experiments.