Revisiting Reweighted Wake-Sleep for Models with Stochastic Control Flow

Abstract

Stochastic control-flow models (SCFMs) are a class of generative models that involve branching on choices from discrete random variables. Amortized gradient-based learning of SCFMs is challenging as most approaches targeting discrete variables rely on their continuous relaxations---which can be intractable in SCFMs, as branching on relaxations requires evaluating all (exponentially many) branching paths. Tractable alternatives mainly combine REINFORCE with complex control-variate schemes to improve the variance of naive estimators. Here, we revisit the reweighted wake-sleep (RWS) (Bornschein and Bengio, 2015) algorithm, and through extensive evaluations, show that it outperforms current state-of-the-art methods in learning SCFMs. Further, in contrast to the importance weighted autoencoder, we observe that RWS learns better models and inference networks with increasing numbers of particles. Our results suggest that RWS is a competitive, often preferable, alternative for learning SCFMs.

arXiv commentTuan Anh Le and Adam R. Kosiorek contributed equally; accepted to Uncertainty in Artificial Intelligence 2019

References in corpus (10)

• Auto-Encoding Variational Bayes
• Categorical Reparameterization with Gumbel-Softmax
• Stochastic Backpropagation and Approximate Inference in Deep Generative Models
• Attend, Infer, Repeat: Fast Scene Understanding with Generative Models
• Backpropagation through the Void: Optimizing control variates for black-box gradient estimation
• Deep Probabilistic Programming
• Auto-Encoding Sequential Monte Carlo
• Learning the Structure of Deep Sparse Graphical Models
• Inference Networks for Sequential Monte Carlo in Graphical Models
• DVAE#: Discrete Variational Autoencoders with Relaxed Boltzmann Priors

Cited by in corpus (17)

• Probabilistic Neural Architecture Search
• A Tutorial on Deep Latent Variable Models of Natural Language
• Decision-Making with Auto-Encoding Variational Bayes
• Reweighted Expectation Maximization
• Variational Inference for Deblending Crowded Starfields
• Learning Proposals for Probabilistic Programs with Inference Combinators
• Collapsed Amortized Variational Inference for Switching Nonlinear Dynamical Systems
• On importance-weighted autoencoders
• Learning to learn generative programs with Memoised Wake-Sleep
• Teaching deep neural networks to localize single molecules for super-resolution microscopy
• Amortized Population Gibbs Samplers with Neural Sufficient Statistics
• Gaussian Process Bandit Optimization of the Thermodynamic Variational Objective
• Capsule Networks -- A Probabilistic Perspective
• Amortised Learning by Wake-Sleep
• Imitation Learning of Factored Multi-agent Reactive Models
• Hybrid Memoised Wake-Sleep: Approximate Inference at the Discrete-Continuous Interface
• Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces