PaperlayerView exact revision on arXiv

Source-linked AI summary

Structured Inference Networks for Nonlinear State Space Models

Rahul G. Krishnan, Uri Shalit, David Sontag

arXiv:1609.09869v2stat.MLcs.AIcs.LG

TL;DR

The paper addresses efficient learning and inference for broad linear and nonlinear Gaussian state space models, including deep neural-network variants. It jointly learns the generative model and a recurrent structured inference network that mimics the posterior, and reports better held-out likelihood with structured approximations. The approach also supports scalable modeling of high-dimensional and missing-data settings, including electronic health records.

  • Problem

    Efficient learning of nonlinear Gaussian state space models from complex, high-dimensional time series remains challenging.

  • Method

    The method jointly learns a generative model and a recurrently parameterized structured inference network that compiles approximate posterior inference.

  • Results

    Structured posterior approximations outperform mean-field approximations, while the learned models support applications including high-dimensional electronic health records with missing data.

  • Takeaways & Limitations

    Compiling structured inference into a neural network provides a scalable approach for learning and inference in rich sequential latent-variable models.

  • Takeaways & Limitations

    Gaussian variational approximations limit inferential expressivity except in linear dynamical systems, where the posterior is Gaussian.

Abstract

from arXiv · show

Gaussian state space models have been used for decades as generative models of sequential data. They admit an intuitive probabilistic interpretation, have a simple functional form, and enjoy widespread adoption. We introduce a unified algorithm to efficiently learn a broad class of linear and non-linear state space models, including variants where the emission and transition distributions are modeled by deep neural networks. Our learning algorithm simultaneously learns a compiled inference network and the generative model, leveraging a structured variational approximation parameterized by recurrent neural networks to mimic the posterior distribution. We apply the learning algorithm to both synthetic and real-world datasets, demonstrating its scalability and versatility. We find that using the structured approximation to the posterior results in models with significantly higher held-out likelihood.

1 Introduction

The paper proposes scalable learning and inference for linear and nonlinear Gaussian state space models, including deep Markov models, by compiling structured posterior approximations into recurrent inference networks.

  • Efficient learning of nonlinear models from complex, high-dimensional time series remains challenging despite the broad use of sequential-data models.The paper targets this challenge with a unified algorithm for a broad class of Gaussian state space models.
  • The compiled inference algorithm handles high-dimensional observed and latent spaces without compromising inference or learning quality.
  • Deep Markov models replace linear emission and transition distributions with multilayer perceptrons while retaining the Markovian structure of hidden Markov models.This combines deep neural-network representational power with a state-space model structure.
  • The proposed recurrent structured inference networks are evaluated with fixed generative models, known-form parameter estimation, and learned deep Markov models.
  • The paper reports that future observations should inform latent-state inference and that structured variational approximations outperform mean-field approximations.
  • The learned deep Markov model supports counterfactual treatment queries on electronic health records and identifies effects of certain medications on patient health.The evaluation includes polyphonic music and high-dimensional health-record data with missing observations.

2 Background

The background defines Gaussian state space models through latent-state transitions and observation emissions, then introduces variational inference networks for approximating otherwise intractable posteriors.

  • Gaussian state space models represent sequential data with vector-valued latent variables z_t and observations x_t.
  • The transition distribution is Gaussian with mean and covariance determined by the previous latent state and elapsed time, while observations depend on the current latent state.
  • The formulation includes linear and nonlinear state space models, depending on the chosen functions and parameters for transitions and emissions.
  • Variational inference introduces a neural inference network that parameterizes an approximate posterior q_φ(z|x) when the true posterior is typically intractable.
  • Gaussian variational posteriors use neural networks to produce observation-dependent means and covariances, with analytic KL terms available when the prior is Gaussian.

3 A Factorized Variational Lower Bound

The paper derives a structured variational posterior that mirrors the exact posterior factorization, then optimizes a factorized variational lower bound with stochastic backpropagation.

  • Posterior factorization: The exact posterior factorizes so each latent state depends on its previous latent state and observations from the current time onward.Past observations are summarized by the previous latent state.
  • Posterior factorization: The variational approximation directly mimics this structure, using qφ(z1|x1, . . . , xT) and qφ(zt|zt−1, xt, . . . , xT).Neural networks parameterize the conditional mean and covariance.
  • Variational objective: The variational lower bound combines expected emission log-likelihoods with KL terms for the initial and transition distributions.The bound is optimized with respect to generative-model parameters θ and inference-network parameters φ.
  • Variational objective: Factorizing the KL divergence gives analytic individual KL terms and more stable gradients instead of requiring Monte Carlo estimation of the full KL.Sampling the unfactorized KL would produce high-variance estimates and gradients.
  • Optimization: Learning uses stochastic backpropagation through samples from the recognition network and updates θ and φ with gradient ascent over mini-batches.The implementation uses a single recognition-network sample to estimate expectations and aggregates gradients across mini-batches.

4 Structured Inference Networks

Structured inference networks use recurrent neural networks to parameterize Gaussian variational posteriors while incorporating information from both past and future observations.

  • Inference-network construction: The variational approximation models each posterior mean and diagonal covariance with recurrent neural networks, chosen for scalability to large datasets.The diagonal-covariance parameterization is generally approximate rather than equal to the true posterior.
  • Inference-network variants: The Deep Kalman Smoother implements the posterior-inspired form by combining the previous latent state with future observations xt, . . . , xT.It corresponds to the structured ST-R inference network.
  • Inference-network variants: The evaluated inference networks vary in mean-field versus structured factorization and in whether they use left, right, or both temporal directions.MF-LR and ST-LR are illustrated with a bidirectional recurrent neural network.
  • Combiner function: Forward and backward recurrent hidden states act as Gaussian messages summarizing past and future observations before the combiner produces posterior parameters.The combiner multiplies the messages and performs a variance-weighted average of their means.
  • Relation to prior approaches: Compared with approaches that condition posterior means only on current observations or use past information, the proposed networks incorporate future information relevant to each latent state.Related approaches either omit future observations or interleave local-parameter prediction with message-passing optimization.

5 Deep Markov Models

Deep Markov models replace linear emission and transition functions with neural networks while retaining the Markovian structure of state space models.

  • Model definition: When model functional forms are unknown, the emission, transition, and initial-state functions are parameterized with deep neural networks.This yields a class of nonlinear generative models for complex, high-dimensional data.
  • Emission model: For binary observations, a two-layer MLP with sigmoid output parameterizes the mean probabilities of independent Bernoulli emissions.The MLP uses element-wise nonlinearities such as ReLU, sigmoid, or tanh.
  • Transition model: The transition function uses a gated parameterization that mixes a linear transformation with a nonlinear proposed mean.The gating unit and proposed mean are computed by MLPs, allowing different latent dimensions to favor linear or nonlinear transitions.
  • Transition model: The transition mean and covariance share the nonlinear hidden representation, while the linear component is initialized as the identity with zero bias.The emission and transition parameters together form θ.

6 Evaluation

The evaluation tests compiled structured inference on synthetic state-space models, then examines deep Markov models on polyphonic music and electronic health records. Across these settings, structured inference supports accurate posterior approximation, parameter learning, richer generative models, and counterfactual prediction.

  • Synthetic Data: Synthetic experiments evaluate inference on linear and nonlinear Gaussian state-space models, including parameter estimation with N = 5000 sequences of length T = 25.The evaluation compares variational bounds and RMSE with respect to the true latent states.
  • Synthetic Data: DKS and ST-LR converge to the RMSE of exact Kalman smoothing, while their variational lower bounds become tight.Both structured methods outperform the past-only mean-field method MF-L in the fixed linear model.
  • Synthetic Data: Structured variational approximations match smoothed Unscented Kalman Filter inference on held-out nonlinear synthetic data and successfully estimate model parameters.These experiments test nonlinear generative models and learning a subset of parameters in a fixed model.
  • Polyphonic Music: ST-LR and DKS substantially outperform MF-LR and ST-L on polyphonic music, despite DKS using half as many parameters as ST-LR and MF-LR.The table reports test negative log-likelihood and variational bounds for inference networks trained with a fixed DMM structure.
  • Polyphonic Music: Increasing generative-model complexity improves results across all datasets, while DMM and DMM-Aug outperform several recurrent and state-space baselines across most polyphonic music datasets.DMM beats RNNs and HMSBN everywhere, while DMM-Aug improves over LV-RNN except on JSB.
  • EHR Patient Data: On electronic health records, the DMM handles noisy, high-dimensional, intermittently missing observations by marginalizing missing values and supports model-based counterfactual prediction.The medical model represents medications as external actions that modulate the patient’s latent state.

7 Discussion

The algorithm builds structured inference networks that mimic the posterior while remaining scalable in sequence length and training-set size. RNN-based inference also supports flexible conditioning and missing-data handling.

  • The algorithm’s space complexity is independent of sequence length T and training-set size N, unlike classical variational inference methods.This yields substantial computational savings for long sequences and large datasets.
  • The proposed variational networks mimic the true posterior’s structure, while alternative structured approximations could condition on future summaries.The paper specifically evaluates recurrent-network-based posterior approximations.
  • Using RNNs only for inference permits increased inference-network capacity and conditioning on different relevant modalities without concern about overfitting the generative model.
  • The DMM’s probabilistic semantics make marginalizing unobserved data straightforward, supporting learning with missing observations.

A Lower Bound on the Likelihood of data

The likelihood objective is optimized through a variational lower bound whose KL term exploits the factorization of the variational distribution and prior. Gaussian structure permits a closed-form KL evaluation, while test likelihood uses Monte Carlo estimation.

  • The likelihood lower bound is derived by combining the variational distribution with the generative model’s likelihood and prior terms.
  • The variational distribution and prior are factorized across the initial state and subsequent conditional transitions, simplifying the KL divergence.
  • Held-out log likelihood is estimated with Monte Carlo, and Eq. 10 can be computed stably using the log-sum-exp trick.
  • The KL divergence between Gaussian variational and prior distributions can be evaluated in closed form from their means and covariances.The variational parameters are produced by the variational model, while the prior parameters come from the generative model.
  • The prior uses an initial zero mean and unit covariance, followed by transition-dependent means and learned diagonal covariances scaled by a scalar parameter.

C Polyphonic Music Generation

The paper evaluates deep Markov models on polyphonic music and extends their emission distribution with NADE. Samples are shown for JSB Chorales, while DMM-Aug-NADE performs comparably to RNN-NADE on several datasets.

  • DMM-Aug-NADE replaces the DMM-Aug emission distribution with a Neural Autoregressive Density Estimator.
  • DMM-Aug-NADE performs comparably to the state-of-the-art RNN-NADE on JSB, Nottingham, and Piano.
  • Table 4 reports test negative log-likelihood on the Polyphonic Music Generation dataset, where lower values are better.
  • Figure 7 shows two samples from the DMM trained on JSB Chorales.

D Experimental Results on Synthetic Data

Synthetic experiments test compiled inference on held-out linear systems and nonlinear dynamical systems, including posterior visualizations against a smoothed UKF. Structured networks generalize and converge more effectively when they match the posterior.

  • The synthetic experiments use inference networks with RNN size 40.
  • On held-out linear SSM data with N=500 and T=25, inference networks—especially structured ones—generalize compiled inference while preserving training-set patterns.Performance is evaluated using RMSE and the variational upper bound.
  • In a nonlinear synthetic system with N=5000 and T=25, posterior-matching inference networks achieve faster convergence and better training and validation accuracy.
  • Figure 10 visualizes posterior means and covariances estimated by inference networks on nonlinear synthetic training data and compares them with a smoothed Unscented Kalman Filter.Posterior estimates are averaged over ten inference runs.

E Generative Models of Medical Data

The section describes implementation details and visualizes samples generated by a DMM trained on patient data. These samples capture temporal medical patterns, including correlations, biomarker variation, and rare comorbidities.

  • Marginalizing out Missing Data: Missing A1C and Glucose observations are tracked with indicators during batch learning to support marginalization.The implementation describes missing observations as being ignored during learning.
  • Sampling a Patient: The DMM generates patient samples from a model trained on medical data.The section presents these samples as part of its visualization of the generative model.
  • Sampling a Patient: Generated samples capture correlations within timesteps and variation in A1C and Glucose levels across timesteps.These patterns reflect both within-timestep relationships and temporal biomarker variation.
  • Sampling a Patient: The model also captures rare comorbidities observed among diabetic patients.
  • Sampling a Patient: Figure 11 plots observations over time, with color intensity representing values between zero and one.The x-axis denotes time and the y-axis denotes observations.
Loading 1609.09869v2…