Joint Stochastic Approximation and Its Application to Learning Discrete Latent Variable Models

Abstract

Although with progress in introducing auxiliary amortized inference models, learning discrete latent variable models is still challenging. In this paper, we show that the annoying difficulty of obtaining reliable stochastic gradients for the inference model and the drawback of indirectly optimizing the target log-likelihood can be gracefully addressed in a new method based on stochastic approximation (SA) theory of the Robbins-Monro type. Specifically, we propose to directly maximize the target log-likelihood and simultaneously minimize the inclusive divergence between the posterior and the inference model. The resulting learning algorithm is called joint SA (JSA). To the best of our knowledge, JSA represents the first method that couples an SA version of the EM (expectation-maximization) algorithm (SAEM) with an adaptive MCMC procedure. Experiments on several benchmark generative modeling and structured prediction tasks show that JSA consistently outperforms recent competitive algorithms, with faster convergence, better final likelihoods, and lower variance of gradient estimates.

arXiv commentAccepted by UAI2020

References in corpus (14)

• Auto-Encoding Variational Bayes
• Markov Chain Monte Carlo and Variational Inference: Bridging the Gap
• On the ergodicity properties of some adaptive MCMC algorithms
• Auxiliary Deep Generative Models
• Improving Variational Inference with Inverse Autoregressive Flow
• REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models
• Latent Normalizing Flows for Discrete Sequences
• MuProp: Unbiased Backpropagation for Stochastic Neural Networks
• Reweighted Wake-Sleep
• Estimating Gradients for Discrete Random Variables by Sampling without Replacement
• ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables
• Markovian Score Climbing: Variational Inference with KL(p||q)
• Direct Optimization through $\arg \max$ for Discrete Variational Auto-Encoder
• Joint Stochastic Approximation learning of Helmholtz Machines

Cited by in corpus (1)

• Energy-Based Models with Applications to Speech and Language Processing