S ELF -S UPERVISED I NFERENCE IN S TATE -S PACE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
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.
An audio denoising experiment that uses real-world noise showing practical applicability of the methods.
Our contributions can be summarized as follows.
[Figure] 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.
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.
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.
A linearized smoothing approach is presented that does not require explicit additional parameterization and learning of the smoothing distribution.  R ELATED W ORK
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. 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 stateest. uncertainty noise dir. opt.
self-sup.
Ljung(1979) ✓ / × ✓ ✓ ✓ × ×
Hochreiteret.
al.
(1997) ✓ ✓ ✓ / × ✓ ✓ × Choetal.(2014) ✓ ✓ ✓ / × ✓ ✓ ×
Wahlstr¨ometal.(2015)
✓ ✓ ✓ / × × ✓ × Watteretal.(2015) ✓ × ✓ ✓ × ✓
Archeretal.(2015) ✓ / × × ✓ ✓ × ✓ Krishnanetal.(2017) ✓ × ✓ ✓ × ✓ Fraccaroetal.(2017) ✓ / × × ✓ ✓ × ✓ Karletal.(2017)
✓ × ✓ ✓ × ✓ Naessethetal.(2018) ✓ × ✓ ✓ × ✓ Yingzhenetal.
(2018) ✓ × ✓ × × ✓
Rangapurametal.(2018) ✓ / × ✓ (1D) ✓ × ✓ ✓ Doerretal.(2018) × ✓ ✓ ✓ ✓ ✓
Satorrasetal.(2019) ✓ ✓ × ✓ ✓ × Haarnojaetal.(2016) ✓ ✓ ✓ ✓ ✓ × Beckeretal.(2019) ✓ ✓ ✓ ✓ ✓ ×
Ours ✓ / × ✓ ✓ ✓ ✓ ✓ e k −e k e k + x k −x k x k + . . .
y k −y k y k
+ . . .
State-space model with deeper latent structure.
G ENERATIVE M ODEL A SSUMPTIONS
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 y 0:
K :=
( y 0 , . .
, y K ) , where each y k ∈ R M , k = 0 , . . .
, K .
We abbreviate: y l : k :=
( y l , . . .
, y k ) and y <
k := y 0: k −and y ≤ k :=
y 0:
k and y − k :=
( y 0: k − 1 , y k
+1: K ) .
We then assume that y 0: K is the result of some possibly non-linear probabilistic latent dynamics, i.e., of a distribution p ( x 0: K ) , whose variables are given by x 0: K :=
( x 0 , . .
, x K ) with x k ∈ R N .
Each y k is assumed to be drawn from some shared noisy emission probability p ( y k
| x k ) .
The joint probability is then assumed to factorize as:
[Equation]
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
[Equation]
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.
In contrast to other approaches (e.g., Krishnan et al. (2015)
; Johnson et al.
; Krishnan et al.
(2017)) where one tries to model the latent dynamics with transition probabilities p ( x k | x k − 1 ) and possibly non-linear emission probabilities p ( y k | x k ) , we go the other way around.
We assume that all the non-linear dynamics are captured inside the latent distribution p ( x 0: 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:
[Equation]
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 x k .
To see this, let z k ∈ R N − M denote non-emitted state variables.
We then put x k :=
[ y k z k ] ⊤ , where y k is computed by applying the nonlinear emission to z k .
We thus include the measurements in the modeled latent state x k .
Then we can put
H :
= (cid:2)
I M 0 M × ( N − M ) (cid:3) .
Second, techniques proposed by Laine et al.
(2019) allow for nonGaussian noise models, relaxing the need for assumption eq.
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).
P ARAMETERIZATION
In this section, we show how we parameterize the inference model.
A lot of the paper’s work relies on established Bayesian filtering machinery.
However, for completeness, we like to prove how all the update steps remain valid while using neural networks for function estimation.
Given our noisy measurements y 0: K = ( y 0 , . . .
, y K )
we want to find good estimates for the latent states x 0: K = ( x 0 , . . .
, x K ) , which generated y 0: K .
For this, we want to infer the marginal conditional distributions p ( x k | y ≤ k ) or p ( x k | y <k ) (for forecasting), for an online inference approach ( filtering ); and p ( x k | y 0: K ) or p ( x k | y − k ) , for a full inference approach ( smoothing ).
In the main body of the paper, we only consider filtering.
Smoothing can be performed similarly, which is detailed in the supplementary material (appendix A).
We start with the following advantageous parameterization:
[Equation]
) where ˆ F k | <k :
= ˆ F k | <k ( y <k ) , ˆ e k | <k :
= ˆ e k | <k ( y <k ) and ˆ Q k | <k :
= ˆ Q k | <k ( y <k ) are parameterized with neural networks.
Next, we have available
[Equation]
i.e., the previous time-step’s conditional marginal distribution of interest.
For k = 1 , this is some initialization.
Otherwise, it is the result of the procedure we are currently describing.
We use this distribution to evaluate the marginalization
[Equation]
with
[Equation]
Note that the distributions under the integral eq.
(6) are jointly Gaussian only because of the parameterization eq.
Hence, we can evaluate the integral analytically.
Finally, to obtain the conditional p ( x k | y ≤ k )
= p ( x k | y k , y <k ) we use Bayes’ rule:
[Equation]
Equation (3) and the result eq.
(6) allow us to also get an analytic expression for
[Equation]
with the following abbreviations:
[Equation]
We introduce the Kalman gain matrix similar to the classical formulas:
[Equation]
Note that taking the matrix inverse at this place in eq.
(11) is more efficient than in the standard
Gaussian formulas (for reference presented in appendix B) if M ≤ N , which holds for our experiments.
GRU f ∑ h → k + e k +1 F k + ∏ y k
̂ x k | ≤ k ̂ x k +1 | ≤ k
GRU f ∑ h → k + e k +2 F k + ∏ y k + ̂ x k +1 | ≤ k + ̂ x k +2 | ≤( k +1)
GRU f ∑ h → k e k F k ∏ y k − ̂ x k −1 | ≤ k − ̂ x k | ≤( k −1)
We note that we could also have parameterized
[Equation]
(12) with ˆ x k | <k ( y <k ) and ˆ P k | <k
( y <k ) directly estimated by a neural network.
This has the advantage that we do not rely on a local linear transition model.
However, it also means that we are estimating x k without any form of temporal regularization.
Additionally, it is harder to incorporate prior knowledge about the transitions maps into such a model.
Nonetheless, we detail this parameterization further in the supplementary material (appendix C) and include its performance in our experiments in section 8.
To conclude the section, we like to discuss some of the limitations of the approach.
The Gaussianity assumption of eq.
(4) ensures but also restricts eq.
(8) and eq.
(9) to these forms.
That is, we make a direct assumption about the form of the posterior p ( x k | y ≤ k ) .
Defending our case, we like to point out that methods such as variational inference (Krishnan et al., 2017) or posterior regularization (Ganchev et al., 2010) also make assumptions (e.g., mean-field Gaussian) about the posterior.
Since we did not explicitly specify a factorization of p ( x 0: K ) , we cannot ensure that the distributions we obtain from the above procedure form a posterior to the ground truth generative process.
This does not mean, however, that we cannot perform accurate inference.
Arguably, not making explicit assumptions about the generative process is preferred to making wrong assumptions and using those for modeling, which can be the case for variational auto-encoders.
The local linearity assumption is justifiable if the length between time-steps is sufficiently small.
However, note that the model is more flexible than directly parameterizing eq.
(12) (see appendix C) since it reduces to that case by putting F k := 0 for all k .
F ITTING
We have shown in the previous section how parameterization of a local linear transition model leads to recursive estimation of p ( x k | y <k ) and p ( x k | y ≤ k ) for all k using classical Bayesian filtering formulas.
The inference is only effective if the estimates ˆ F k | <k , ˆ e k | <
k and ˆ Q k | <k from eq.
(4) are accurate.
We can use the parameterization p ( x k | y <k ) of eq.
(6), the emission model p ( y k | x k ) =
N ( y k | H x k , R ) from eq.
(3), and the factorization from eq.
(1) to see that an analytic form of the log-likelihood of the data emerges:
[Equation] [Equation]
If we put ˆ y k | <k :
=
H
ˆ x k | <k ( y <k ) and ˆ M k | <k :
=
H ˆ P k | <k ( y <k ) H ⊤ + R , then the maximumlikelihood objective leads to the following loss function, which we can minimize using gradient descent methods w.r.t.
all model parameters:
[Equation]
Note that each term in the sum above represents a one-step-ahead self-supervised error term.
We thus minimize the prediction residuals ˆ y k | <k
( y <k ) − y k in a norm that is inversely scaled with the above covariance matrix, plus a regularizing determinant term, which prevents the covariance matrix from diverging.
The arisen loss function is similar to the noise2noise (Lehtinen et al., 2018;
Krull et al., 2019; Batson & Royer, 2019; Laine et al., 2019) objective from computer vision literature, combined with a locally linear transition model like Becker et al.
We show in appendix D that this objective will yield correct results (meaning estimating the ground-truth x k ) if the noise is independent with E [ y k |
Hx k ] =
Hx k .
A similar procedure in the causality literature is given by Sch ¨ olkopf et al.
An algorithmic presentation of performing inference and fitting is presented in appendix E.
Note that after fitting the parameters to the data, eq.
(6) can directly be used to do one-step ahead forecasting.
Forecasting an arbitrary number of time steps is possible by plugging the new value ˆ x K +1 | K via y K
+1 :
=
H
ˆ x K +1 | K back into the recurrent model, and so on.
This is not a generative model but merely a convenience that we deemed worth mentioning.
L INEARIZED S MOOTHING
So far, we have only discussed how to perform filtering.
Recall that for smoothing, we are instead interested in the quantity p ( x k | y 0: K ) .
A smoothing strategy highly similar to the methods described earlier can be obtained by explicitly parameterizing such a model, which we detail in the supplementary material.
Here, we introduce a linearized smoothing procedure.
The essential advantage is that no additional model has to be trained, which can be costly.
Several algorithms stemming from the Kalman filter literature can be applied, such as the RTS smoother (Rauch et al., 1965) and the two-filter smoother (Kitagawa, 1994).
To enable this, we need to assume that the conditional mutual information I ( x k − 1 ; y k : K | x k , y 0: k − 1 ) is small for all k = 1 , . .
, K .
In other words, we assume that we approximately have the following conditional independences:
[Equation]
To explain the motivation for this requirement, consider the model in fig.
If the states of y 0: k −and x k are known, then the additional information that y k : K has about the latent variable x k −would need to be passed along the unblocked deeper paths like y k +1 ← x k +1 ← e k +1 ← e k ← x k − 1 .
Then the assumption of small I ( x k − 1 ; y k : K | x k , y 0:
k − 1 ) can be interpreted as that the deeper paths transport less information than the lower direct paths.
If we consider all edges to the x k ’s as linear and the edges to the e k ’s as non-linear maps, the above could be interpreted as an information-theoretic version of expressing that the non-linear correction terms are small compared to the linear parts in the functional relations between the variables.
We will now show that under the earlier assumptions and eq.
(16) we get a Gaussian approximation: p ( x k | y 0: K ) ≈ N (cid:16) x k | ˆ z k , ˆ G k (cid:17) .
We will do backward induction with ˆ z K :
= ˆ x K |≤ K and ˆ G K :
= ˆ P K |≤ K .
To propagate this to previous time steps k − 1 we use the chain rule:
[Equation]
where the second term is known by backward induction and for the first term we make use of the approximate conditional independence from eq.
(16) to get
[Equation]
The latter was shown to be Gaussian in section 4:
[Equation]
By use of the usual formulas for Gaussians and the reverse Kalman gain matrix ˆ J k − 1 | k we arrive at the following update formulas, k = K, . .
, 1 , with ˆ z K :
= ˆ x K |≤ K and ˆ G K :
= ˆ P K |≤ K :
[Equation]
As such, we can perform inference for all k = 0 , . .
, K with p ( x k | y 0: K ) ≈ N (cid:16) x k | ˆ z k , ˆ G k (cid:17) .
Algorithmically, the above is presented in appendix E.
R ECURRENT N EURAL N ETWORK
Before going into the experiments section, we briefly explain how we specifically estimate the functions ˆ F k | < [Equation]
( y [Equation] ˆ e [Equation]
y [Equation] and ˆ Q k | <k ( y <k ) that parameterize the transition probability p ( x k | y <
k , x k − 1 ) (eq.
The choice of the model here implicitly makes further assumptions about the generative model.
If we consider neural networks, the temporal nature of the data suggests recurrent neural networks (Graves et al., 2013), convolutional neural networks (Kalchbrenner et al., 2014), or transformer architectures (Vaswani et al., 2017).
Additionally, if the data is image-based, one might further make use of convolutions.
For our experiments, we use a Gated Recurrent Unit (GRU) network (Cho et al., 2014), that recursively encodes hidden representations.
Therefore, we put
[Equation]
where ˆ L k is a Cholesky factor and f is a multi-layer perceptron decoder.
E XPERIMENTS
We perform three experiments, as motivated in section 1.
Technical details on the setup of the experiments can be found in the supplementary material (appendix F).
We refer to the model detailed in section 4 as the recursive filter , as it uses the Bayesian update recursion.
For smoothing experiments, we use recursive smoother .
The model obtained by parameterizing p ( x k | y <k ) directly (eq.
(12)) is referred to as the recurrent filter or recurrent smoother , as it only employs recurrent neural networks (and no Bayesian recursion) to estimate said density directly.
8. L INEAR D YNAMICS
In the linear Gaussian case, it is known that the Kalman filter will give the optimal solution.
Thus, we can get a lower bound on the test loss.
In this toy experiment, we simulate particle tracking under linear dynamics and noisy measurements of the location.
We use Newtonian physics equations as prior knowledge.
We generate trajectories T = { x 0: K , y 0: K } with x k ∈ R 6 and y k ∈ R 2 according to the differential equations:
[Equation]
We obtain sparse, noisy measurements y k
=
Hx k
+ r with r ∼ N ( 0 , R ) .
H is a selection matrix that returns a two-dimensional position vector.
We run this experiment in a filtering setting, i.e.,
[Figure] [Figure]
we only use past observations.
We compare against (1) the raw, noisy measurements which inherently deviate from the clean measurements, (2) the Kalman filter solution where we optimized the transition covariance matrix using clean data (hence supervised), (3) the optimal solution, which is a Kalman filter with ground truth parameters performing exact inference.
To estimate the transition maps for the Kalman filter, we use the standard Taylor series of e A · ∆ t up to the first order.
Additionally, we use this expert knowledge as an inductive bias for the recursive filter’s transition maps.
In fig. 4 we depict the test mean squared error (MSE, lower is better) as a function of the number of training samples.
Given enough data [Figure] the self-supervised models approximate the optimal solution arbitrarily well.
Our recursive model significantly outperforms both the Kalman filter and the inference model in the low-data regime by using incorporated expert knowledge.
Additionally, we report that the recursive model’s distance to the ground truth latent states is much closer to the optimal solution than both the inference model and Kalman filter.
Specifically, we report average mean squared errors of 0 .
685 for the inference model, 0 .
241 for the Kalman filter, 0 .
161 for the recursive model compared to 0 .
135 for the optimal Kalman filter.
Finally, it is worth noting that the recursive model has much less variance as a function of its initialization.
8. L ORENZ E QUATIONS
We simulate a Lorenz system according to
[Equation]
We have H
=
I , x ∈ R 3 and y ∈ R 3 .
The Lorenz equations model atmospheric convection and form a classic example of chaos.
Therefore, performing inference is much more complex than in the linear case.
This time, we perform smoothing (see appendix A) and compare against (1) the raw measurements, (2) a supervised Extended Kalman smoother (Ljung, 1979), (3) the variational inference approach of Krishnan et al.
(2017), (4) the supervised recursive model of Satorras et al.
Our models include a recurrent smoother, a recursive smoother, and the recursive filter with linearized smoothing (section 6).
Transition maps for the Extended Kalman smoother and the recursive models are obtained by taking a second-order Taylor series of e A · ∆ t .
For the supervised extended Kalman filter, we again optimize its covariance estimate using ground truth data.
In fig. 5 we plot the test mean squared error (MSE, lower is better) as a function of the number of examples available for training.
It is clear that our methods approach the ground truth states with
[Table]
more data.
This is in contrast to the Extended Kalman smoother, which barely outperforms the noisy measurements.
We also see that the recursive models significantly outperform the recurrent model in the low-data regime.
The recursive filter with linearized smoothing performs comparably to the other models and even better in low-data regimes.
We hypothesize that this is because the required assumption for the linearized smoother holds (section 6) and regularizes the model.
The variational method of Krishnan et al.
(2017) performs poorly in low-data regimes.
Most notably, the supervised method of Satorras et al.
(2019) outperforms our models only slightly.
8. A UDIO D ENOISING
Next, we test the model on non-fabricated data with less ideal noise characteristics.
Specifically, we use the SpeechCommands spoken audio dataset (Warden, 2018).
Performing inference on spoken audio is challenging, as it arguably requires understanding natural language.
To this end, recent progress on synthesizing raw audio has been made (Lakhotia et al., 2021).
However, this requires scaling to much larger and more sophisticated neural networks than presented here, which we deem out the current work’s scope.
Therefore, we take a subset of the entire dataset, using audio from the classes “tree”, “six”, “eight”, “yes”, and “cat”.
We overlay these clean audio sequences { x (1)0: K , . . .
, x ( N ) 0: K } with various noise classes that the dataset provides.
That is, for every noise class
C we obtain a set of noisy sequences { y (1)0: K , . . . , y ( N ) 0: K } C . We also consider a “combined” class in which we sample from the union of the noise sets.
The task is to denoise the audio without having access to clean data.
We evaluate the models on non-silent parts of the audio, as performance on those sections is the most interesting.
Notably, none of these noise classes is Gaussian distributed.
We show the mean squared error on the test set of all models per noise class in table 2.
Our models outperform the Kalman filter, Noise2Noise (Lehtinen et al., 2018), and SIN (Krishnan et al., 2017) unsupervised baselines.
We suspect that the relatively poor performance of SIN is due to its generative Markov assumption, regularizing the model too strongly.
The poor performance of Noise2Noise is due to the fact that it does not use the current measurement y k to infer x k .
Like before, note that the Kalman filter is “supervised” as we optimized its covariance matrix using clean data.
The supervised RKN (Becker et al., 2019) outperforms our models on most noise classes, but notably not on white noise.
Most of these noise classes have temporal structure, making them predictable from past data.
This is confirmed by observing these mean squared error values over the course of training.
Initially, the values were better than reported in table 2, but the model increasingly fits the noise over time.
Thus, although two of the main assumptions about the model (independent Gaussian noise) are violated, we still can denoise effectively.
Since the RKN’s targets are denoised (hence “supervised”), it does not have this problem.
In practice, obtaining clean data can be challenging.
C ONCLUSION
We presented an advantageous parameterization of an inference procedure for nonlinear state-space models with potentially higher-order latent Markov chains.
The inference distribution is split into linear and nonlinear parts, allowing for a recursion akin to the Kalman filter and smoother algorithms.
Optimization is performed directly using a maximum-likelihood objective that backpropagates through these recursions.
Smoothing can be performed similarly, but we additionally proposed linearized smoothing that can directly be applied to the filtering distributions.
Our model is simple and builds on established methods from signal processing.
Despite this, results showed that it can perform better or on par with fully supervised or variational inference methods.
E THICS S TATEMENT
The paper presents a simple method to perform inference using noisy sequential data.
Applications can be found throughout society, e.g., tracking particles, denoising images or audio, or estimating system states.
While many such examples are for good, there are applications with ethically debatable motivations.
Among these could be tracking humans or denoising purposefully corrupted data (e.g., to hide one’s identity).
R EPRODUCIBILITY S TATEMENT
We are in the process of releasing code for the current work.
For clarity and reproducibility, the presented methods are available as algorithms in the supplementary material.
Furthermore, we made explicit wherever we needed to make an approximation or an assumption
A P ARAMETERIZED S MOOTHING
In the main body of the paper we showed how to parameterize the model for recursive estimation of p ( x k | y ≤ k ) , k = 0 , . .
, K .
Additionally, we provided a linearized smoothing procedure that yields p ( x k | y 1: K ) .
The disadvantage clearly is the linearization.
Here, we show how we can recursively estimate p ( x k | y 1: K ) in a similar sense to the filtering setting.
First, we put
[Equation]
[Equation]
where ˆ F k |− k ( y − k ) and ˆ e k |− k
( y − k ) and ˆ Q k |− k ( y − k ) are two-sided recurrent neural network outputs similar to section 7.
We compute the distribution of interest as follows:
[Equation]
where we make the approximation to compute eq.
(27) efficiently and recursively.
It is justified if
I ( x k − 1 ; y k | y − k ) < ϵ for small ϵ .
That is, the additional information that y k conveys about x k −is marginal if we have all other data.
Let
[Equation]
be previous time-step’s posterior.
Then put
[Equation]
where we left out the arguments for the following quantities estimated by an RNN.
[Equation]
Applying the integral in eq.
(30) we get
[Equation]
where we put
[Equation]
A data likelihood can be computed as follows
[Equation]
where we made the same approximation eq.
(30) as before.
It evaluates to
[Equation]
with
[Equation]
For fitting to the data we now would use a maximum-pseudo-likelihood (Gourieroux et al., 1984) approach by maximizing:
[Equation]
leading to minimizing the following self-supervised loss function:
[Equation]
[Equation]
where ˆ y k |− k :
=
H ˆ x k |− k
( y − k ) and ˆ M k |− k :
= H ˆ P k |− k
( y − k ) H ⊤ + R .
B G AUSSIAN C ONDITIONING F ORMULAS
Since many of the calculations used in this work are based on the Gaussian conditioning formulas, we provide them here.
If
[Equation] [Equation]
(45) then p [Equation] [Equation] (47) with Σ =
[Equation]
H [Equation] −
[Equation]
C D IRECT P ARAMETERIZATION OF p ( x | y <k )
We show here how to directly parameterize p ( x | y <k ) and p ( x | y − k ) .
This parameterization is referred to as the recurrent model in our experiments.
The procedure is rather straightforward.
For filtering, we put p ( x k | y <k ) =
N (cid:16) x k | ˆ x k |− k
( y <k ) , ˆ P k | <k ( y <k ) (cid:17) .
We model:
[Equation]
where L k is a cholesky factor and f is a multi-layer perceptron.
The argument h → k is recursively given by
[Equation]
where we employ a Gated Recurrent Unit (GRU) network (Cho et al., 2014).
Note that this parameterization is equivalent to the model described in the main paper with F k := 0 .
For smoothing, p
[Equation]
x
[Equation]
[Equation]
where L k is a cholesky factor and f is a multi-layer perceptron.
[Equation]
Once p ( x | y <k ) (filtering) or p ( x | y − k ) (smoothing) is obtained, all the procedures for inference and optimization described in the main paper and appendix A remain the same.
D B IAS -V ARIANCE -N OISE D ECOMPOSITION OF THE S ELF -S UPERVISED G ENERALIZATION E RROR
Any estimate ˆ x k = ˆ x k ( y − k ) for x k that is not dependent on y k will give us a bias-variancenoise decomposition of the generalization error.
Note that this setting covers both the filtering and smoothing case.
Define “optimal model” ˆ y ∗ k :=
E [ y k | y − k ] , then under the specified generative model (section 3) we have [Equation] [Equation]
[Equation]
Thus,
[Equation] [Equation]
Note that ˆ y ∗ k = E
[ Hx k | y − k ]
= H E
[ x k | y − k ] .
Then, define our model ˆ y k :
=
H ˆ x k .
The reducible part of the error becomes
[Equation]
For this reason, the model output ˆ x k approaches the optimal model under minimization of the selfsupervised error (perturbed by H ).
Additionally, the model ˆ y k approaches Hx k (the uncorrupted measurement) under this criterion.
[Equation]
If we then take ˆ y k :
=
H ˆ x k then the reducible part of the error approximates the true x k (perturbed by H ).
[Equation]
E A LGORITHMS
p ( x k | y ≤ k ) = N (cid:16) x k | ˆ x k |≤ k , ˆ P k
≤ k (cid:17)
[Equation] [Equation]
ˆ
[Equation]
e
[Equation]
for k in 0 , . . .
K do
[Equation]
ˆ
[Equation]
L
[Equation]
[Equation]
[Equation]
det
[Equation]
k ˆ K
[Equation]
k ˆ P
[Equation]
P
[Equation]
H
[Equation]
P [Equation] k ˆ x k |≤ k :=
ˆ x k | <k + ˆ K k
( y k − ˆ y k
| <k )
h k +1 :=
GRU ϕ ( h k , y k )  
ˆ e k +1 | <
( k +1) ˆ
F k +1 | < ( k +1)
ˆ L k +1 | <
( k +1)   :
= f ϕ ( h k +1 )
ˆ Q k +1 | <
( k +1) :
= ˆ L k +1 | < ( k +1)
ˆ L ⊤ k +1 | < ( k +1)
ˆ P k +1 | <
( k +1) :
= ˆ F k +1 | < ( k +1)
ˆ P k |≤ k ˆ F ⊤ k +1 | < ( k +1)
+ ˆ Q k +1 | <
( k +1)
ˆ x k +1 | <
( k +1) :
= ˆ F k +1 | < ( k +1)
ˆ x k |≤ k + ˆ e k +1 | <
( k +1) end
For the training case we use backpropagation through the above loop to compute ∇ ϕ L K .
Algorithm 2: ) input : Data (time-series) y
: K = ( y 0 , . . .
, y K ) , , emission matrices H and R , initialized parameters ϕ 0 , number of training rounds I .
output: Model parameters ϕ ∗ for inference at test-time.
for i in 0 , . . .
, I do
Obtain L ( i ) K and ∇ ϕ L ( i ) K from algorithm 1.
Run preferred optimizer step to update parameters ϕ with ∇ ϕ L ( i ) K (and L ( i ) K ).
end
[Equation] [Equation]
h
[Equation]
ˆ
[Equation]
ˆ x 0 |− 0 :
=
ˆ e
[Equation]
for k in 0 , . . .
K do
[Equation]
h
[Equation]

[Equation]
F [Equation]  [Equation]
[Equation]
k [Equation] k +1 ) ˆ Q k +1 |−
( k +1) :
= ˆ L k +1 |−
( k +1)
ˆ L ⊤ k +1 |−
( k +1)
ˆ P k +1 |−
( k +1) :
=
ˆ F k +1 |−
( k +1)
ˆ P k | 0: K ˆ F ⊤ k +1 |−
( k +1)
+
ˆ Q k +1 |−
( k +1)
ˆ x k +1 |− ( k +1) :
=
ˆ F k +1 |−
( k +1)
ˆ x k | 0: K
+ ˆ e k +1 |−
( k +1) end
For training case we also use backpropagation through the above loop to compute ∇ ϕ L K
Algorithm input : , emission matrices initialized parameters training rounds I output for i in 1 , . .
, I do
Obtain ∇ ϕ L ( i ) K from algorithm 3.
Run preferred optimizer step.
end
ˆ [Equation] [Equation]
end input : Training data (time-series)
{+y 
 : K = ( y 0 , . . .+}
, y K ) , emission function N ( Hx k , R ) , initialized parameters ϕ 0 , number of training iterations n .
output
Optimized parameters ϕ ∗ do h 0 :
=
h K +1 :=
L ( i ) :
=for k in 0 to K do h
→ k :=
GRU ϕ ( h → k − 1 , y k − 1 ) h ← k :=
GRU ϕ ( h ← k − 1 , y k +1 )
(cid:20)
ˆ x k |− k
ˆ
L k |− k (cid:21) :
= f ϕ ( h ← k , h → k ) ˆ P k |− k :
=
ˆ L k |− k
ˆ L ⊤
k |− k
ˆ y
[Equation]
end
Compute ∇ ϕ L ( i ) K and apply SGD step with respect to all model parameters, which amounts to backpropagation through the above calculations.
end
For the filter variant, the steps that involve the backward direction ( ← ) are left out.
input : Test data
: K = ( trained parameters ϕ emission function N
Hx , R ) .
output posteriors
= h 0 : = h 1 :=for k in 0 to K do
[Equation]
end
For the filter variant, the steps that involve the backward direction ( ← ) are left out.
F E XPERIMENTS : D ETAILS F. L INEAR D YNAMICS
As specified in the main paper, the dynamics are according to
[Equation]
Since these are linear transitions, we can calculate any transition directly using x ( t + ∆ t ) =
e A ∆ t x ( t ) .
[Equation]
(cid:21) (67)
We used c = 0 .
06 , τ = 0 .
17 , ∆ t := 1 and covariance
[Equation]
The matrix exponential is computed using Bader et al.
The parameters for the emission distribution:
[Equation]
We simulate a K := 131 , 072 trajectory for training, K := 16 , 384 trajectory for validation and
K := 32 , 768 for testing.
The ˜ F that is used in the (non-optimal) Kalman filter and recursive model is computed as follows:
[Equation]
[Equation]
L ORENZ E QUATIONS
We simulate a Lorenz system according to
[Equation]
We integrate the system using dt = 0 .
00001 and sample it uniformly at ∆ t = 0 .
We use ρ = 28 , σ = 10 , β = 8 / 3 .
The transition in ∆ t arbitrary time-steps is linearly approximated by a Taylor expansion and used in the Kalman smoother and recursive models.
[Equation]
We simulate K := 131 , 072 steps for training, K := 32 , 768 for testing and K := 16 , 384 for validation.
We have H :
=
I and thus x ∈ R 3 and y ∈ R 3 .
We use R := 0 .
5 2 I .
G P ARAMETERIZED S MOOTHING : A LTERNATIVE P OSTERIOR E VALUATION
We would like to point out that the distribution p ( x k | y 0: K ) can be obtained without making assumption eq.
(30), which we stretch out here.
Initial experiments showed that using these calculations the model did not converge as smoothly as when using the ones stated before.
However, it could be of interest to further investigate.
Returning to the posterior of interest
[Equation]
[Equation]
We have
[Equation]
as the previous time-step’s posterior.
[Equation]
where ˆ F k |− k ( y − k )
[Equation]
ˆ e k |− k
( y − k ) and ˆ Q k |− k ( y − k ) are estimated by a neural network.
Combining this with noise model eq.
(3) we get:
[Equation]
where we directly applied the Woodbury matrix identity to obtain Kalman gain matrix
[Equation]
Then, applying the integral we get:
[Equation]
with:
[Equation]
and
[Equation]