[Paper Review] Sublinear-Time Approximate MCMC Transitions for Probabilistic Programs

Yutian Chen, Vikash K. Mansinghka+1 more|arXiv (Cornell University)|Nov 6, 2014

Markov Chains and Monte Carlo Methods10 references3 citations

TL;DR

This paper proposes a sublinear-time approximate Metropolis-Hastings algorithm for latent variable updates in probabilistic programs by subsampling edges in a dynamically constructed probabilistic execution trace (PET) graph, enabling efficient inference in highly coupled models like regressions and state-space models. The method achieves up to 10x speedup over exact MCMC with negligible bias and maintains accuracy across diverse applications including Bayesian logistic regression, nonparametric mixtures, and stochastic volatility models.

ABSTRACT

Probabilistic programming languages can simplify the development of machine learning techniques, but only if inference is sufficiently scalable. Unfortunately, Bayesian parameter estimation for highly coupled models such as regressions and state-space models still scales poorly; each MCMC transition takes linear time in the number of observations. This paper describes a sublinear-time algorithm for making Metropolis-Hastings (MH) updates to latent variables in probabilistic programs. The approach generalizes recently introduced approximate MH techniques: instead of subsampling data items assumed to be independent, it subsamples edges in a dynamically constructed graphical model. It thus applies to a broader class of problems and interoperates with other general-purpose inference techniques. Empirical results, including confirmation of sublinear per-transition scaling, are presented for Bayesian logistic regression, nonlinear classification via joint Dirichlet process mixtures, and parameter estimation for stochastic volatility models (with state estimation via particle MCMC). All three applications use the same implementation, and each requires under 20 lines of probabilistic code.

Motivation & Objective

• To address the scalability bottleneck in Bayesian inference for highly coupled models such as regressions and state-space models, where standard MCMC transitions scale linearly with data size.
• To generalize approximate MCMC techniques beyond i.i.d. data subsampling by instead subsampling dependencies (edges) in a dynamically constructed graphical model (PET).
• To integrate sublinear approximate MH transitions into a general-purpose probabilistic programming framework, enabling interoperability with existing inference algorithms.
• To provide a rigorous proof of ergodicity for the approximate Markov chain under milder conditions than prior work, showing bias vanishes as the approximation parameter approaches zero.
• To empirically validate the method’s efficiency and accuracy across diverse, real-world probabilistic programming applications with minimal code overhead.

Proposed method

• The method constructs a probabilistic execution trace (PET) graph that tracks dependencies between random variables during program execution, enabling factorization of the joint density.
• Instead of subsampling data items, the algorithm subsamples edges in the PET graph, stochastically ignoring dependencies to reduce computational cost.
• Approximate MH transitions are applied selectively to latent variables with high dependency degrees (e.g., parameters in regression or state-space models), while exact transitions are used for others to maintain balance.
• The algorithm uses a sequential test to determine whether the approximate transition is sufficiently accurate, based on a control parameter ε that governs the approximation level.
• The approach is integrated into a higher-order probabilistic programming system, allowing seamless use with existing inference strategies like particle Gibbs and HMC.
• Computation time is balanced across variables using a step_z parameter to ensure that faster approximate updates to parameters do not dominate the overall inference schedule.

Experimental results

Research questions

• RQ1 Can sublinear-time approximate MCMC transitions be designed for highly coupled latent variables in probabilistic programs, beyond the scope of i.i.d. data subsampling?
• RQ2 Does edge subsampling in a dynamically constructed PET graph preserve the ergodicity and convergence properties of the original MCMC chain under mild conditions?
• RQ3 Can a single implementation of sublinear approximate MH be effectively applied across diverse probabilistic models, including parametric and nonparametric models, with minimal code changes?
• RQ4 To what extent does the method improve inference efficiency without introducing significant bias, especially in models with strong latent state dependencies?
• RQ5 How does the performance of sublinear approximate MH compare to exact MH in terms of effective sample size per second and convergence speed across different model classes?

Key findings

• The sublinear approximate MH algorithm achieves up to 10x speedup over exact MH in Bayesian logistic regression and joint Dirichlet process mixtures, with no significant loss in accuracy.
• In the stochastic volatility model, the method achieves approximately 2x higher efficiency than exact MH, despite slower mixing in latent states, with samples closely matching the ground truth distribution.
• The empirical results confirm sublinear per-transition scaling, demonstrating that computational cost grows sublinearly with the number of observations.
• The histogram of sampled parameters (φ and σ) from the sublinear method closely matches those from exact MH, indicating negligible bias when ε = 1e-3.
• Effective sample size (ESS) per second is significantly higher for the sublinear method, with gains of 2x–10x depending on the model complexity and data size.
• The method maintains accuracy across all three applications—Bayesian logistic regression, nonparametric mixtures, and stochastic volatility—using under 20 lines of probabilistic code in each case.