Composing graphical models with neural networks for structured representations and fast inference

Matthew J. Johnson; David Duvenaud; Alexander B. Wiltschko; Sandeep R. Datta; Ryan P. Adams

Source-linked AI summary

Composing graphical models with neural networks for structured representations and fast inference

Matthew J. Johnson, David Duvenaud, Alexander B. Wiltschko, Sandeep R. Datta, Ryan P. Adams

arXiv:1603.06277v5Mar 20, 2016stat.ML

TL;DR

The paper addresses how to combine interpretable graphical-model structure with flexible neural observation models while retaining efficient inference. It introduces structured variational autoencoders, which use recognition networks to produce conjugate potentials for graphical-model approximations. The framework jointly learns model components with a scalable objective and supports message passing, natural gradients, and reparameterization-based optimization.

Problem
Flexible neural observation models and structured latent variables are difficult to combine because efficient graphical-model inference depends on conjugacy.
Method
Structured variational autoencoders use graphical-model variational families and recognition networks that output conjugate potentials instead of complete variational distributions.
Results
The framework provides a single objective that jointly learns all model components with stochastic variational inference, message passing, natural gradients, and the reparameterization trick.
Takeaways & Limitations
SVAEs provide a general framework for rich latent representations, flexible nonlinear observations, and fast variational inference across graphical-model structures.
Takeaways & Limitations
The paper develops three specific examples from a broader model family, and some local latent structures are kept simple in the presentation.

Abstract

from arXiv · show

We propose a general modeling and inference framework that composes probabilistic graphical models with deep learning methods and combines their respective strengths. Our model family augments graphical structure in latent variables with neural network observation models. For inference, we extend variational autoencoders to use graphical model approximating distributions with recognition networks that output conjugate potentials. All components of these models are learned simultaneously with a single objective, giving a scalable algorithm that leverages stochastic variational inference, natural gradients, graphical model message passing, and the reparameterization trick. We illustrate this framework with several example models and an application to mouse behavioral phenotyping.

1 Introduction

The paper targets the tension between flexible learned representations and interpretable, tractable probabilistic structure. It proposes a framework combining graphical-model latent structure with neural-network observation models and recognition networks.

Graphical models offer structured representations but can impose rigid assumptions and require feature engineering.
Deep learning learns flexible representations automatically but does not directly encode interpretable or tractable probabilistic structure.
Mouse behavioral video motivates representations that capture low-dimensional nonlinear structure and behavioral dynamics for phenotyping.
The framework couples nonlinear likelihoods with structured latent-variable representations for modeling complex data.
The proposed approach combines flexible deep-learning feature models with structured Bayesian or nonparametric priors.

2 Latent graphical models with neural net observations

The examples compose structured latent variables with neural observation models to represent clusters, temporal dynamics, and hybrid discrete-continuous behavior. This combination preserves interpretable latent structure while supporting nonlinear data distributions.

2.1 Warped mixtures for arbitrary cluster shapes: Discrete mixture models discover clusters or units, but these structures are difficult to represent directly in neural networks.
2.1 Warped mixtures for arbitrary cluster shapes: A Gaussian mixture model can fail to represent natural spiral clustering because its inflexible Gaussian observation model poorly fits the data semantics.
2.1 Warped mixtures for arbitrary cluster shapes: Neural-network density models fit data density flexibly but do not explicitly represent discrete mixture components.
2.1 Warped mixtures for arbitrary cluster shapes: Composing a latent GMM with nonlinear observations learns discrete clusters alongside non-Gaussian cluster shapes.
2.2 Latent linear dynamical systems for modeling video: For video, a Gaussian linear dynamical system gives low-dimensional latent states and dynamics with a rich nonlinear image generator.
2.3 Switching linear dynamical systems for behavioral phenotyping: A latent switching linear dynamical system combines discrete dynamics-selected units, continuous latent states, and neural-network video observations.

3 Structured mean field inference and recognition networks

Inference becomes difficult when flexible neural likelihoods break the conjugacy needed by efficient graphical-model algorithms. The paper addresses this by combining structured variational families, recognition networks, and conjugate inference machinery.

Rich latent structure and flexible likelihoods create an inference challenge because message passing and structured mean field rely on conjugate exponential-family likelihoods.
Conjugate observation models preserve exponential-family conditionals, enabling efficient graphical-model inference algorithms.
For linear Gaussian dynamical systems, the optimal latent variational factor remains a Gaussian dynamical system and can be computed with message passing.
Variational autoencoders use recognition networks to output Gaussian approximations from observations when nonlinear posteriors are difficult to compute directly.

4 Structured variational autoencoders

Structured variational autoencoders combine neural-network observation models with graphical-model latent structure by using recognition networks that output conjugate potentials. The resulting objective and algorithm support simultaneous learning with tractable graphical-model inference and stochastic gradient optimization.

Recognition networks output conjugate graphical-model potentials, which are then used by inference algorithms instead of non-conjugate observation likelihoods.
The SVAE algorithm samples minibatches, computes evidence potentials, performs graphical-model inference, and estimates gradients of a mean-field variational objective.
The resulting algorithm is described as simple, efficient, and applicable to varied learning problems and graphical-model structures.
The model class combines exponential-family latent variables and priors with general, potentially neural-network, observation likelihoods while retaining conjugate structure for inference.
Recognition-model potentials approximately optimize local variational factors, making inference efficient while preserving exponential-family and graphical-model structure.
The SVAE objective lower-bounds the mean-field objective and can become tight when the recognition model represents the expected log likelihood.
Gradient-based optimization maximizes a lower bound on model log evidence, while reparameterization and automatic differentiation compute gradients for observation and recognition parameters.
Local optimization and message passing also provide efficient expectations and natural-gradient estimates for latent-model parameters.

5 Related work

Related work combines variational autoencoders with sequential or continuous-state models and develops general methods for nonconjugate variational inference. This work distinguishes itself by using recognition networks and amortized inference to preserve conjugate structure, while noting complementary alternatives.

Krishnan et al. combine variational autoencoders with continuous state-space models but do not extend to general graphical models or discrete latent variables.
Archer et al. use structured Gaussian variational families for continuous state-space models without considering parameter learning or discrete latent variables.
Other sequential VAE extensions focus on recurrent neural networks rather than probabilistic graphical models.
Alternative variational-inference methods exploit conjugate exponential-family structure while handling arbitrary nonconjugate factors through proximal-gradient or related approaches.
This work instead handles intractable terms with recognition networks and amortized inference, combining inference networks with conjugate stochastic variational inference.

6 Experiments

The experiments show that LDS and SLDS SVAEs learn flexible representations of image and video data, support prediction, and reveal structured behavioral states.

The SVAE was applied to synthetic and real data to learn feature representations and latent structure.
6.1 LDS SVAE for modeling synthetic data: The LDS SVAE accurately represented bouncing-dot images and made long-term predictions with uncertainty.
Figures 4, 5, 6, and 7 visualize synthetic-data predictions, experimental results, depth-video predictions, and inferred behavior states.
6.1 LDS SVAE for modeling synthetic data: Natural-gradient updates learned faster and were less dependent on parameterization than standard-gradient updates.At an untuned stepsize of 0.1, standard-gradient updates at 0.1 and 0.05 produced indefinite matrix parameters.
6.2 LDS SVAE for modeling mouse video: On mouse depth video, the learned image manifold captured smooth variation in body pose, and the model produced predictions paired with real data.The experiment used a 10-dimensional latent space and MLP observation and recognition models.
6.3 SLDS SVAE for parsing behavior: An SLDS SVAE clustered short-timescale patterns into behavioral units represented by discrete latent states.The model was fit with 30 discrete states.

7 Conclusion

Structured variational autoencoders combine graphical-model structure with neural-network flexibility for representation learning and inference. The appendix then introduces gradient notation used later.

Structured variational autoencoders combine rich graphical-model latent representations with fast variational inference.
Neural networks provide flexible nonlinear observation models and recognition networks that map observations to conjugate graphical-model potentials.
A.1 Gradient notation: The appendix defines gradient, Jacobian, and Hessian notation for continuously differentiable functions.

A.2 Local and partial optimizers

This section defines unconstrained and constrained local partial optimizers and states the corresponding first-order optimality conditions.

A partial optimizer maximizes an objective over one variable block while holding the other block fixed.
A local partial optimizer only needs to maximize the objective within some neighborhood of the selected variable value.
For an unconstrained local partial optimizer, the gradient with respect to the optimized variable is zero.
For constrained optimization, a full-rank constraint Jacobian yields a Lagrange multiplier condition and orthogonality to feasible variations.
The constraint regularity condition is unnecessary when the constraints are linear.
Gradient-based algorithms guarantee stationary limit points, while block coordinate ascent additionally makes each coordinate block a partial optimizer.

A.3 Partial optimization and the Implicit Function Theorem

The section develops gradients for objectives involving partial optimization and introduces exponential-family machinery for the SVI objective. It also states the local nature and assumptions of these constructions.

A.3 Partial optimization and the Implicit Function Theorem: For a composite objective g(x)=f(x,y*(x)), the chain rule includes both direct x-dependence and the sensitivity of y*(x).
A.3 Partial optimization and the Implicit Function Theorem: When y*(x) is a local partial optimizer, the sensitivity term vanishes under the stated unconstrained or regularly constrained conditions.
A.3 Partial optimization and the Implicit Function Theorem: The Implicit Function Theorem provides a local differentiable representation of y*(x) when h(x,y)=0 has a nonsingular y-gradient.
A.3 Partial optimization and the Implicit Function Theorem: The construction assumes local uniqueness and differentiability, and discusses sensitivity almost everywhere when multiple optimizers occur.
Exponential-family notation: Exponential families are parameterized by sufficient statistics, natural parameters, and a log partition function over an appropriate base measure.
Exponential-family notation: Derivatives of the log partition function generate expected sufficient statistics and higher cumulants.
Exponential-family notation: Conjugate priors are constructed by defining prior statistics from the likelihood’s natural parameters and log partition function.
C.1 SVI objective: The SVI objective locally partially optimizes q(x), using first-order stationary conditions rather than requiring a global optimizer.

C.3 Stochastic natural gradients for large datasets

This section develops stochastic natural-gradient inference for structured latent-variable models with conjugacy, using local updates that support scalable optimization on large datasets.

Stochastic natural gradients for large datasets: Natural-gradient SVI estimates gradients by sampling data-point terms from the full-data objective.The sampled index is drawn from {1, 2, . . . , N}, with probability-weighted contributions yielding an unbiased Monte Carlo estimate.
Conditionally conjugate local optimization: Conditionally conjugate local variables z and x can be optimized by alternating updates of their variational parameters.The procedure computes expected sufficient statistics for each variational factor.
Model class: The framework covers models including latent Dirichlet allocation, switching linear dynamical systems, mixture models, and hidden Markov models with exponential-family emissions.Its conditionally conjugate structure is powerful but restrictive.
Conditionally conjugate local optimization: The block coordinate ascent updates converge to a stationary point that is partially optimal in each factor’s parameters.Each update uses quantities already computed in gradient-based optimization and requires no step-size choice.
SVAE extension: SVAE generalizes the stochastic-gradient setup beyond complete-data conjugate models to more general likelihood models.The section defines unbiased stochastic estimates of gradients, including natural gradients for variational parameters with conjugacy structure.

D.1 SVAE objective

This section defines the SVAE objective by combining conjugate variational factors with a recognition-model surrogate for potentially non-exponential-family likelihoods.

Recognition model and surrogate: The surrogate replaces the potentially general likelihood term with a tractable function related to the exponential-family latent model.This construction preserves exploitable conjugacy structure for local optimization.
Variational family: The SVAE uses exponential-family variational factors for global parameters, likelihood parameters, and local latent variables.The factors are parameterized by natural parameters ηθ, ηγ, and ηx.
Recognition model and surrogate: A recognition model r(y; φ) maps observations into parameters that determine a local variational factor through a surrogate objective.The local parameter η∗x(ηθ, φ) is chosen as a local partial optimizer of that surrogate.
SVAE objective: The SVAE objective lower-bounds the partially optimized mean-field variational objective for any recognition-model function class.The bound can be tight when some φ∗ makes the recognition model match the optimal local factor.
SVAE objective: Maximizing the SVAE objective optimizes a lower bound on model evidence while minimizing KL divergence from the variational family to the target posterior.The framework therefore motivates using a rich recognition-model function class.

D.2 Estimating the natural gradient e∇ηθLSVAE

This section derives efficient stochastic gradient estimators for the SVAE, including a natural gradient for global conjugate parameters and reparameterized gradients for neural components.

Natural gradient with respect to ηθ: The SVAE natural gradient with respect to ηθ combines the SVI natural gradient with an additional recognition-model-dependent term.Unlike SVI, the additional term is generally nonzero because the local factor is optimized through the SVAE surrogate.
Single local factor: For a single unconstrained local factor q(x), stationary conditions simplify the natural-gradient calculation through cancellation of Fisher-information factors.The resulting expression is derived from the local surrogate optimum.
Limitation: The natural-gradient simplification does not necessarily hold when ηx is constrained by additional factorization structure.It always applies for Gaussian q(x) with fixed covariance because third and higher-order cumulants vanish.
Structured local factors: With additional local factorization, as in GMM and SLDS models, the natural gradient uses the chain rule and the reparameterization trick.Applying (∇2 log Zθ(ηθ))^-1 adds computational cost, which is small when ηθ has low dimension.
Recognition and likelihood gradients: Unbiased estimates of ∇φLSVAE and ∇ηγLSVAE are obtained by sampling γ and x and differentiating a Monte Carlo estimate of the objective.Only the expected likelihood term requires the reparameterization-trick approximation.
Structured local optimization: The same efficient local partial optimization applies when local variables are separated into structured groups z and x under conjugacy assumptions.Alternating updates compute the local optimizers used by the SVAE objective.

E Experiment details and expanded figures

The experiment trains an LDS SVAE on synthetic one-dimensional dot videos and visualizes inferred behavioral states from depth video.

Synthetic 1D dot video: The synthetic LDS SVAE uses 80 image sequences of length 50, with one sequence per update, to predict future frames from a prefix.The image and recognition models are one-hidden-layer MLPs with 50 units, and the latent state dimension is 8.
Expanded figures: Figure 8 shows four example depth-video frame sequences for each inferred behavior state, selected near that state’s median duration.Frames are temporally subsampled by four, padded on both sides, and marked with a lower-right square when the state is active.