[Paper Review] Deep Kalman Filters

Rahul G. Krishnan, Uri Shalit+1 more|arXiv (Cornell University)|Nov 16, 2015

Machine Learning in Healthcare27 references142 citations

TL;DR

This paper introduces Deep Kalman Filters, a variational inference-based framework that uses deep neural networks to learn nonlinear state-space models for sequential data. It enables effective counterfactual inference in complex, noisy environments—demonstrated on synthetic 'Healing MNIST' data and real electronic health records—showing that anti-diabetic drugs significantly reduce A1c and glucose levels in diabetic patients under simulated interventions.

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.

Motivation & Objective

• To develop a unified, scalable method for learning complex, nonlinear Kalman filters using deep neural networks and variational inference.
• To enable counterfactual inference in longitudinal data, particularly in healthcare settings where only one intervention outcome is observed per patient.
• To model the latent evolution of patient states under real-world interventions such as drug prescriptions, using noisy, high-dimensional EHR data.
• To evaluate the impact of different recognition models and model architectures on sequence modeling and counterfactual prediction performance.
• To demonstrate the utility of continuous state-space models for causal inference in longitudinal medical data.

Proposed method

• Proposes a variational inference framework to optimize a lower bound on the log-likelihood of sequential observations, enabling end-to-end training of deep nonlinear Kalman filters.
• Uses deep neural networks to parameterize the transition function Gα, emission function Fκ, and observation likelihood, allowing flexible, nonlinear dynamics.
• Employs a recognition model (e.g., q-BRNN) to infer posterior distributions over latent states, enabling efficient approximate inference.
• Applies Pearl’s do-operator by setting the lab indicator variable to 1 during counterfactual sampling, simulating intervention scenarios.
• Uses a fixed diagonal covariance matrix Sβ for the emission noise, simplifying inference while maintaining model flexibility.
• Trains the model on both synthetic data (Healing MNIST) and real EHR data from 8,000 diabetic patients over 4.5 years.

Experimental results

Research questions

• RQ1 Can deep neural networks be effectively combined with Kalman filters to model complex, nonlinear temporal dynamics in high-dimensional sequential data?
• RQ2 How well can the proposed model perform counterfactual inference in the absence of observed outcomes for alternative interventions?
• RQ3 Does the model capture long-range dependencies and structural invariances in perturbed sequential data, such as rotated or noisy MNIST digits?
• RQ4 Can the model accurately estimate the causal effect of anti-diabetic medications on patient outcomes like A1c and glucose levels using real EHR data?
• RQ5 How do different model architectures (linear vs. nonlinear transition/emission functions) affect performance in sequence modeling and counterfactual prediction?

Key findings

• The model with nonlinear transition and emission functions achieved the highest test log-likelihood, outperforming linear and non-linear emission-only variants.
• On the Healing MNIST dataset, the model successfully captured both short- and long-range effects of applied actions, such as rotation and noise, on digit sequences.
• In the EHR setting, patients not receiving anti-diabetic drugs (Metformin) were significantly more likely to have high A1c levels (above 8%) and high glucose levels post-intervention, as shown by counterfactual sampling.
• The proportion of patients with high glucose levels increased from ~20% to ~60% in the 'without' condition compared to the 'with' condition after Metformin prescription.
• The model demonstrated robustness in inferring latent states under noisy, high-dimensional EHR data, enabling reliable counterfactual comparisons.
• The use of a parametric posterior (q-BRNN) enabled accurate posterior approximation and effective forward sampling under different intervention conditions.