Deep Kalman Filters

Rahul G. Krishnan, Uri Shalit, David Sontag

Seminal November 2015

TL;DR Deep Kalman Filters replace a Kalman filter's linear transition and emission with neural networks and learn them by maximizing a variational lower bound, using an RNN recognition network for inference. On claims data from 8,000 diabetic patients, sampling forward without anti-diabetic drugs gives far more patients with high glucose and A1c than sampling with them.

What they do

A classical Kalman filter is linear with Gaussian noise. Deep Kalman Filters replace its transition, action and emission functions with neural networks. Every variant is learned the same way: stochastic backpropagation on a variational lower bound of log p(x | u). The goal is counterfactual inference on patient records: infer a latent state, then sample forward under the drugs actually given and under an alternative.

How it works

Model. z_t is Gaussian, with mean and diagonal covariance computed by networks from z_{t-1}, u_{t-1} and the time gap. A network of z_t sets the parameters of x_t’s distribution (e.g. Bernoulli). Linear maps give back the classical filter.
Bound. The Markov structure splits the KL into one term per time step. The authors say this gives more stable gradients than sampling the full KL.
Recognition network. They prove that the posterior factorizes as p(z_t z_{t-1}, x_t..x_T, u). They compare four encoders: a per-frame MLP, a 3-frame MLP, a forward RNN and a bidirectional RNN.
Counterfactuals. Infer z up to time t, then sample forward under each action. Sicker patients are tested more often. To handle this, the model includes a lab-indicator variable, which the do-operator sets to 1.

Results

Healing MNIST. Rotated 1s and 5s with up to 20% bit-flip noise and a three-square pattern that lasts several frames (140,000 length-five sequences in the large set). The bidirectional RNN has the best test log-likelihood. The 3-frame MLP beats the forward RNN on likelihood, but its samples miss the squares.
Patient model. On 8,000 patients (4.5 years in 3-month steps, nine drugs as actions), models with non-linear transitions reach the highest test log-likelihood.
Counterfactuals. For 800 patients first prescribed Metformin, removing the medication leaves many more patients with high glucose and high A1c (above 8%).

Limitations

The causal reading of the model is assumed rather than learned, and the paper reports no check of the counterfactuals against observed outcomes. A Gaussian q cannot in general match the true posterior, so the bound is generally not tight. When the noise is heavy and the model sees too few frames, it cannot recover the digit. Every comparison is between the paper’s own variants.

Figures

From Figure 1: Learning in a Deep Kalman Filter. Latent states z evolve under actions u and emit observations x; a re... — From Figure 1: Learning in a Deep Kalman Filter. Latent states z evolve under actions u and emit observations x; a recognition network q(z | x, u) approximates the posterior and is used for inference during learning.

From Figure 2: Large Healing MNIST. Pairs of noisy training sequences (TS) and the mean probabilities of their recons... — From Figure 2: Large Healing MNIST. Pairs of noisy training sequences (TS) and the mean probabilities of their reconstructions (R).

From Figure 4: Test log-likelihood on 8,000 diabetic and pre-diabetic patients for models with linear (L) or non-line... — From Figure 4: Test log-likelihood on 8,000 diabetic and pre-diabetic patients for models with linear (L) or non-linear (NL) emission (Em) and transition (Tr) functions.

From Figure 4: Proportion of patients inferred to have high (top quantile) glucose with and without anti-diabetic dru... — From Figure 4: Proportion of patients inferred to have high (top quantile) glucose with and without anti-diabetic drugs, starting from the first Metformin prescription.

Abstract

Kalman Filters are one of the most influential models of time-varying phenomena. They admit an intuitive probabilistic interpretation, have a simple functional form, and enjoy widespread adoption in a variety of disciplines. Motivated by recent variational methods for learning deep generative models, we introduce a unified algorithm to efficiently learn a broad spectrum of Kalman filters. Of particular interest is the use of temporal generative models for counterfactual inference. We investigate the efficacy of such models for counterfactual inference, and to that end we introduce the "Healing MNIST" dataset where long-term structure, noise and actions are applied to sequences of digits. We show the efficacy of our method for modeling this dataset. We further show how our model can be used for counterfactual inference for patients, based on electronic health record data of 8,000 patients over 4.5 years.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong