Source-linked AI summary
Deep Kalman Filters
Rahul G. Krishnan, Uri Shalit, David Sontag
TL;DR
The paper addresses causal questions about interventions from noisy, high-dimensional temporal data, including medical records where each patient provides only one intervention-outcome pair. It develops a variational deep generative framework for Kalman filters and finds it promising for modeling nonlinear dynamics and counterfactual analysis.
Problem
Medical counterfactual inference must estimate an intervention’s outcome when each patient contributes only one observed intervention-outcome pair.
Method
Variational inference and recognition networks train a unified, differentiable framework spanning linear and nonlinear Kalman filters with deep transition and emission functions.
Results
Nonlinear state-space models show promising effectiveness for counterfactual analysis, while the model captures noisy digit patterns and action effects in synthetic data.
Takeaways & Limitations
The learned latent state supports forward sampling under different actions and can represent patients for similarity queries.
Takeaways & Limitations
The quality of temporal and spatial invariance in filters learned from Healing MNIST remains an open question for future work.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper develops deep, nonlinear generative temporal models to address limitations of classical Kalman filters and enable counterfactual inference from noisy sequential data. It evaluates the approach on Healing MNIST and longitudinal patient records.
- Motivation: EHR sequences create opportunities for machine learning questions about treatment effects, patient similarity, and clinical outcomes.The paper represents diagnoses, surgeries, laboratory values, and prescriptions as time-ordered records.
- Motivation: Classical Kalman filters use linear dynamics, emissions, and action effects, limiting their ability to model complex real-world phenomena.Existing nonlinear modifications make learning more difficult.
- Contributions: Variational inference and deep neural networks provide one algorithm for tractably learning a broad class of Kalman filters with complex transition dynamics and emissions.Learning optimizes a bound on the data likelihood.
- Counterfactual inference: Counterfactual medical inference is difficult because each patient typically provides only one observed intervention-outcome pair.The model instead represents longitudinal observations and actions to estimate effects under alternative interventions.
- Evaluation: Healing MNIST tests noisy, perturbed digits, while EHR data cover 8,000 diabetic and pre-diabetic patients over 4.5 years.The experiments examine action effects on digits and anti-diabetic medication effects on patients.
2 Background
The background formulates Kalman filtering as latent-state sequence modeling with observations and actions, then replaces linear components with neural networks. Variational inference makes posterior learning tractable despite nonlinearities.
- Kalman filters: A Kalman filter models latent states, observations, and observed actions across a sequence.In the medical setting, these correspond to patient state, diagnoses and laboratory results, and treatments or procedures.
- Kalman filters: Classical Kalman filters use linear state transitions and action effects with Gaussian noise, alongside linear observation generation.The transition includes the previous action as an additive control input.
- Nonlinear models: The deep formulation replaces linear transformations with neural-network parameterizations, increasing modeling flexibility.The resulting models are harder to learn because their posterior becomes intractable.
- Variational inference: Variational autoencoders optimize a variational lower bound using a recognition network to approximate the intractable posterior.Stochastic backpropagation addresses the expectation’s dependence on variational parameters.
- Counterfactual estimation: Counterfactual estimation asks for the probability of an outcome under circumstances different from those empirically observed.A medical example is predicting blood sugar under a different medication.
3 Related Work
Related work spans probabilistic sequence models, nonlinear dynamical systems, variational sequence models, and temporal representations. The paper distinguishes its goal of counterfactual modeling from optimal control and notes domain-specific differences in action effects.
- Sequence modeling: Prior work models sequences with dynamic factor graphs, LSTMs, temporal variational autoencoders, and stochastic recurrent networks.These approaches use deterministic or stochastic latent representations for sequence prediction and reconstruction.
- Nonlinear dynamical systems: Nonlinear dynamical-system methods have used black-box classifiers or nonlinear dynamic factor analysis to construct latent representations.The latter scales quadratically with latent dimension.
- Variational generative models: Other work learns sequential likelihood bounds with discrete latent variables, whereas this paper develops continuous-state models.The distinction concerns the type of latent random variables used in the generative model.
- Control and counterfactuals: Temporal generative models have also been used for optimal control by learning locally linear latent spaces.This paper instead aims to model data for counterfactual inference.
- Medical setting: Medical actions may have weak and delayed effects compared with the immediate, dominant effects common in control settings.Endogenous and environmental factors can be relatively stronger in patient dynamics.
- Medical sequence models: Earlier patient modeling work found latent states corresponding to relevant laboratory-test levels, including A1c.That work used a partially observable Markov process for diabetic patients.
4 Model
The Deep Kalman Filter models noisy observations and action-dependent latent dynamics with neural networks, subsuming linear and nonlinear Kalman-filter variants. A single differentiable framework supports their learning by backpropagation.
- Model assumptions: The model assumes observations arise from a latent state evolving over time, with observed actions affecting that state nonlinearly.Observations are noisy and nonlinear functions of the latent state.
- Generative model: Latent states have Normal distributions whose mean and covariance depend nonlinearly on the previous state, action, and elapsed time.Observation parameters are generated from the corresponding latent state through a neural network.
- Generative model: Deep neural networks parameterize transition, covariance, and emission functions, with diagonal covariance ensuring a positive-definite covariance matrix.Binary observations can use Bernoulli emissions.
- Model family: Restricting the functional forms recovers classical Kalman filters, while arbitrary differentiable functions produce broader linear or nonlinear latent-state models.All such models can be trained with backpropagation.
- Inference: The recognition model supplies a parametric approximation to the posterior over latent sequences during learning.Figure 1 contrasts this sequence-level approximation with the recognition model used for static generative models.
5 Learning using Stochastic Backpropagation
The method learns temporal generative models by optimizing a variational lower bound with structured recognition networks and stochastic backpropagation. Markov structure simplifies posterior approximation and supports counterfactual inference after learning.
- 5.1 Maximizing a Lower Bound: The model maximizes conditional data likelihood through a variational lower bound, balancing reconstruction with KL divergence from the prior.The objective uses the expected log likelihood of observations minus the KL divergence between approximate posterior and prior.
- 5.2 On the choice of the Optimal Variational Model: The recognition model is conditioned on observations and actions, with alternatives ranging from independent MLPs to recurrent and bidirectional recurrent networks.The compared models are q-INDEP, q-RNN, and q-BRNN, differing in which observations and actions they use.
- 5.1 Maximizing a Lower Bound: Stochastic backpropagation updates generative and recognition-network parameters while estimating gradients of expectation terms under the approximate posterior.The learning algorithm infers latent states, reconstructs inputs, computes gradients, and updates parameters with ADAM.
- 5.2 On the choice of the Optimal Variational Model: Markov structure factorizes the posterior so each latent state depends on the preceding latent state, current and future observations, and later actions.This structure simplifies the posterior approximation and motivates using a recurrent network summarizing future observations alongside the previous latent state.
- 5.2 On the choice of the Optimal Variational Model: The variational lower bound becomes exact only when the approximate posterior equals the true posterior, a condition not generally expected for Gaussian variational distributions.Normalizing flows are discussed as a possible way to make the approximation equal or arbitrarily close to the true posterior.
- Counterfactual Inference: After learning, the temporal generative model estimates a latent state from pre-intervention history and forward-samples outcomes under contrasted actions for counterfactual inference.The procedure applies both the intended action and an alternative action after inferring the state at the intervention time.
6 Experimental Section
Experiments evaluate Healing MNIST under noise, rotations, and counterfactual actions, then compare recognition models and nonlinear Kalman-filter variants. Bidirectional recognition performs best overall, while the model reconstructs digit structure, styles, and action effects but struggles under extreme noise or limited sequences.
- Healing MNIST dataset: Healing MNIST combines rotated digits, superimposed three-square sequences, and up to 20% bit-flip noise to test temporal modeling under complex perturbations.Rotations are treated as actions and images as observations; the dataset emphasizes high-dimensional observations, nonlinear posteriors, and filtering through noise.
- Comparing recognition models: q-BRNN outperforms the other recognition variants by summarizing past and future observations to approximate each latent posterior.The comparison includes q-INDEP, q-LR, q-RNN, and q-BRNN recognition models.
- Comparing recognition models: q-INDEP performs poorly in sample quality and held-out log-likelihood, capturing rotation effects but not the three-square structure.Its posterior for z_t is based only on x_t, which poorly represents the sequential data.
- Comparing recognition models: q-LR achieves better test log-likelihood than q-RNN, but its samples fail to capture the three-consecutive-block structure because it carries no memory across time.The authors suggest this limitation matters when sequences contain multiple ordered patterns.
- Large Healing MNIST: The model reconstructs noisy digits and learns within-digit styles, rotation effects, and varying action intensities, including counterfactual samples for unseen digit styles.For some examples, noise is too high or too few sequences are available for accurate digit inference.
- Generative models of medical data: Medical experiments compare linear and nonlinear emission and transition functions on EHR data from 8,000 diabetic and pre-diabetic patients.The nonlinear model is used for glucose and A1c counterfactual analyses involving anti-diabetic medication.
7 Discussion
The paper presents nonlinear state-space models as promising tools for counterfactual analysis and patient representation. It identifies latent-space identifiability and disease-trajectory interpretation as future questions.
- The model approximates unseen patients’ latent states and forward-samples outcomes under different actions.
- The learned latent space can embed patients and support patient-similarity questions.
- The model’s temporal and spatial invariance and stochastic latent dynamics remain open areas for investigation.
A Lower Bound on Likelihood
The appendix expresses the variational objective through a factorized posterior and prior, then estimates test-set likelihood with Monte Carlo methods. Numerical stability can be improved with log-sum-exp.
- The variational distribution factorizes across the initial latent state and sequential transitions.
- The KL divergence between posterior and prior is decomposed into an initial-state term and expected sequential transition terms.
- Test-set marginal likelihood is evaluated using a Monte Carlo estimate based on the recognition network.
- The derivation follows variational-inference techniques described by Rezende et al. and Kingma and Welling.
- The test log likelihood is computed using the appendix’s likelihood estimator and can be evaluated numerically with log-sum-exp.
B KL divergence between Prior and Posterior
The appendix derives the sequence-level KL divergence by combining Gaussian KL terms with the model’s variational posterior and generative prior parameterizations. The prior mean depends on the previous latent state and action.
- The KL divergence between two multivariate Gaussians is written using their means, covariances, determinants, and trace term.
- The sequence-level KL is obtained by combining the factorized decomposition with the Gaussian KL formula.
- The variational posterior’s mean and covariance are outputs of the variational model.
- The initial prior has mean 0 and covariance 1, while later prior means are generated by G(z_t−1, u_t−1).
- The later prior covariance is a learned diagonal matrix with a scalar parameter Δ.
- Gradients are taken with respect to posterior means, posterior covariances, and the transition function G.
C Additional Experimental Results
Additional experiments train a Large Healing MNIST model on multiple styles of digits 0 and 2. Reconstructions and samples under constant rotations provide visualizations of learned sequence behavior.
- The variant uses 100 different kinds of digits 0 and 2 for Large Healing MNIST training.
- Figure 6 shows reconstructions and samples from a model trained on digits 0 and 2.
- Figure 6a pairs training sequences with mean probabilities of their reconstructions.
- Figure 6b shows mean probabilities sampled under different constant rotations.