Source-linked AI summary
Black box variational inference for state space models
Evan Archer, Il Memming Park, Lars Buesing, John Cunningham, Liam Paninski
TL;DR
The paper addresses approximate inference for latent time-series models whose posteriors are generally intractable and whose temporal correlations are poorly captured by standard variational approximations. It introduces a structured Gaussian black-box posterior with stochastic variational optimization, achieving accurate recovery in closed-form settings and strong performance against specialized inference in PLDS. The approach remains limited by its Gaussian posterior family because many latent time-series models have non-Gaussian posteriors.
Problem
General latent time-series models often have intractable posteriors, while standard variational approximations may inadequately capture their temporal correlation structure.
Method
The paper combines a structured Gaussian variational posterior that expresses temporal dependencies with SGVB and neural networks mapping observations into the smoothed posterior.
Results
The method reaches solutions comparable to specialized PLDS VBEM before one EM iteration completes and accurately recovers nonlinear posterior-mean relationships.
Takeaways & Limitations
The same black-box inference approach can approximate posteriors across a wide class of latent-variable generative models while scaling linearly in time and space with sequence length.
Takeaways & Limitations
The approach is limited by its Gaussian approximate posterior because most latent-variable time-series models have non-Gaussian posteriors.
Abstract
from arXiv · showhide
Latent variable time-series models are among the most heavily used tools from machine learning and applied statistics. These models have the advantage of learning latent structure both from noisy observations and from the temporal ordering in the data, where it is assumed that meaningful correlation structure exists across time. A few highly-structured models, such as the linear dynamical system with linear-Gaussian observations, have closed-form inference procedures (e.g. the Kalman Filter), but this case is an exception to the general rule that exact posterior inference in more complex generative models is intractable. Consequently, much work in time-series modeling focuses on approximate inference procedures for one particular class of models. Here, we extend recent developments in stochastic variational inference to develop a `black-box' approximate inference technique for latent variable models with latent dynamical structure. We propose a structured Gaussian variational approximate posterior that carries the same intuition as the standard Kalman filter-smoother but, importantly, permits us to use the same inference approach to approximate the posterior of much more general, nonlinear latent variable generative models. We show that our approach recovers accurate estimates in the case of basic models with closed-form posteriors, and more interestingly performs well in comparison to variational approaches that were designed in a bespoke fashion for specific non-conjugate models.
1 INTRODUCTION
The paper targets scalable approximate inference for latent time-series models whose posteriors are difficult to compute, while preserving temporal dependence. It introduces a structured black-box variational approach and demonstrates strong performance on PLDS and nonlinear dynamics.
- Motivation: Exact posterior inference is generally unavailable for non-Gaussian or nonlinear time-series models, while sampling methods often scale poorly.Closed-form filtering and smoothing are available only for special structures; general models require numerical approximation or sampling.
- Approach: The method uses a structured Gaussian posterior with temporal dependencies and a fast inference algorithm that does not depend explicitly on the generative model form.A neural network maps observations into the smoothed posterior, while the block tri-diagonal inverse covariance supports linear scaling in time-series length.
- Motivation: The approach focuses on computationally efficient inference for structured, interpretable parameterizations of high-dimensional time-series.The motivating examples include neural spike-train recordings and latent trajectories with low-dimensional dynamical structure.
- Results: In PLDS inference, the black-box method reaches solutions comparable to specialized VBEM before VBEM completes one EM iteration.It also accurately recovers nonlinear relationships in the posterior mean for a one-dimensional nonlinear dynamical system.
2 STOCHASTIC GRADIENT VARIATIONAL BAYES
The paper formulates variational inference by optimizing an ELBO over model and variational parameters, then uses reparameterized stochastic gradients for efficient black-box optimization. The estimator supports flexible neural parameterizations and has lower empirical variance than earlier sampling-based estimators.
- Variational objective: Variational inference replaces the intractable posterior pθ(z|x) with a tractable qφ(z|x) and jointly optimizes θ and φ through the ELBO.The ELBO combines the approximating posterior's entropy with the expected joint log probability.
- Variational objective: The parameterization qφ(z|x) may depend explicitly on observations x, unlike approaches that omit this dependence.This allows the variational posterior to be conditioned on the observed time-series.
- SGVB: The reparameterization trick expresses samples as a deterministic function of x and random noise, enabling stochastic gradient estimation.All randomness in q comes from the auxiliary random variable ϵ.
- SGVB: The SGVB estimator uses iid noise samples and has much lower empirical variance than previous sampling-based gradient estimators.The estimator is formed by averaging gradients evaluated at reparameterized samples.
- SGVB: The estimator is black-box because it requires evaluating the objective at samples rather than using the particular functional form of the generative model.Different modeling choices can still affect computation time and convergence rate in practice.
- SGVB: Deep neural networks can learn a function that maps observations directly into the posterior, avoiding additional gradient steps after training.This provides a flexible amortized parameterization of the variational transformation.
3 VARIATIONAL APPROACH TO STATE-SPACE MODELING
The paper uses a structured Gaussian variational posterior to capture temporal correlations in latent time series while keeping inference scalable in sequence length. Its block tri-diagonal precision mirrors Kalman smoothing structure and supports differentiable, efficient computation.
- Most previous variational posteriors inadequately capture the temporal correlation structure of latent state-space models.
- The proposed multivariate Gaussian posterior uses a block tri-diagonal inverse covariance to express time-local dependencies with parameters scaling linearly in T.This structure borrows from the Gaussian posterior of linear dynamical systems.
- The structured parameterization retains a dense covariance while representing only its diagonal and first block off-diagonal precision matrices.The covariance itself need not be explicitly represented during learning.
- Block structure makes Cholesky factorization, matrix inversion, and sampling linear in the time-series length.The lower-triangular Cholesky factor is computed in linear time and space in T.
- Unlike standard forward-backward use in Kalman smoothing, the method differentiates variational parameters through the factorization Σ^-1 = RRT.
4 PARAMETERIZATION OF THE SMOOTHING POSTERIOR
The smoothing posterior is implemented through neural-network parameterizations of its mean and structured precision, with two alternatives that differ in how Gaussian factors interact. In a Kalman experiment, both alternatives closely recover the closed-form posterior means.
- Parameterization of the smoothing posterior: The general smoothing posterior is parameterized with neural networks whose choices can materially affect practical performance.The networks define functions of the observations for posterior parameters such as the mean and covariance.
- Diagonal and block off-diagonal parameterization: VILDSblk parameterizes blocks separately and adds αI to the precision to enforce positive-definiteness.Here α > 0 is fixed.
- Product-of-Gaussians approximate posterior: The product-of-Gaussians form q(z|x) ∝ r1(z|x)r0(z) treats r0 as a prior-like factor in the final posterior.
- Product-of-Gaussians approximate posterior: VILDSmult requires D^-1 and C^-1 to be block tri-diagonal and allows D to initialize means with a specified degree of smoothness.This differs from the formulation in Section 4.1.
- Empirical comparison: In a T = 5000, two-dimensional Kalman experiment, VILDSmult and VILDSblk posterior means agree closely with the closed-form Kalman filter means.The comparison shows 300 time-points along each state-space dimension.
5 EXPERIMENTS
The experiments evaluate VILDS on linear-Gaussian, Poisson, and nonlinear dynamical systems. VILDS recovers accurate posterior structure and reaches stronger evidence lower bounds than the compared variational methods.
- 5.1 KALMAN FILTER MODEL: VILDS recovers the analytic Kalman-filter posterior in a 2-dimensional latent system with 100-dimensional linear observations.The experiment used 5000 simulated time-points and neural-network parameterizations for the approximate posterior.
- 5.2 POISSON LDS (PLDS): The PLDS experiment compares VILDS with a specialized VBEM method because Poisson observations yield an intractable posterior.
- 5.2 POISSON LDS (PLDS): VILDS reaches a higher ELBO than MF and VBDual before VBDual completes a single expectation-maximization iteration.VILDS posterior means are also smoother than those learned with the mean-field posterior.
- 5.3 NONLINEAR DYNAMICS SIMULATION: VILDS recovers the nonlinear relationship in a toy one-dimensional nonlinear dynamical system from 5000 simulated samples.The generative-model parameters were fixed while only the variational parameters were learned.
6 CONCLUSION
The paper concludes that a block tri-diagonal Gaussian posterior enables smoothed trajectory inference with linear scaling. It demonstrates broad applicability while identifying Gaussianity as the main limitation.
- The proposed Gaussian posterior expresses smoothed time-series trajectories through block tri-diagonal covariance structure.
- Inference scales linearly in both time and space complexity with the time-series length T.
- The SGVB framework supports approximate inference for a wide class of latent-variable generative models.
- The approach is limited by its Gaussian approximate posterior because most latent-variable time-series models have non-Gaussian posteriors.The paper points to non-Gaussian variational posteriors as a possible direction for future work.
- Detailed comparisons with a closely related method were left for future work.
A SMOOTHING APPROXIMATE POSTERIOR WITH EXPLICIT FORWARD/BACKWARD
This section derives a structured Gaussian smoothing posterior by combining Gaussian factors. Its block structure enables efficient forward/backward computation and imposes stability assumptions on the dynamics.
- The posterior mean and covariance can be computed with the standard Kalman forward-backward algorithm.
- The matrix Q represents innovation-noise covariance, while A is the linear dynamics matrix.
- The posterior is written as q(z|x) ∝ r0(z)r1(z), yielding a multivariate normal with parameters determined by eq. 19.
- Efficient computation of the posterior parameters exploits the block bi-diagonal structure of R.
- Stable dynamics assume that the eigenvalues of A have magnitude less than one, although this constraint was not enforced in the reported examples.