Efficient Marginalization of Discrete and Structured Latent Variables via Sparsity

Abstract

Training neural network models with discrete (categorical or structured) latent variables can be computationally challenging, due to the need for marginalization over large or combinatorial sets. To circumvent this issue, one typically resorts to sampling-based approximations of the true marginal, requiring noisy gradient estimators (e.g., score function estimator) or continuous relaxations with lower-variance reparameterized gradients (e.g., Gumbel-Softmax). In this paper, we propose a new training strategy which replaces these estimators by an exact yet efficient marginalization. To achieve this, we parameterize discrete distributions over latent assignments using differentiable sparse mappings: sparsemax and its structured counterparts. In effect, the support of these distributions is greatly reduced, which enables efficient marginalization. We report successful results in three tasks covering a range of latent variable modeling applications: a semisupervised deep generative model, a latent communication game, and a generative model with a bit-vector latent representation. In all cases, we obtain good performance while still achieving the practicality of sampling-based approximations.

arXiv commentAccepted for spotlight presentation at NeurIPS 2020

References in corpus (23)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Auto-Encoding Variational Bayes
• Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms
• Categorical Reparameterization with Gumbel-Softmax
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• Semi-Supervised Learning with Deep Generative Models
• Learning to Communicate with Deep Multi-Agent Reinforcement Learning
• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification
• Neural Variational Inference and Learning in Belief Networks
• REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models
• Variational inference for Monte Carlo objectives
• Multi-Agent Cooperation and the Emergence of (Natural) Language
• Variational Bayesian Inference with Stochastic Search
• Monte Carlo Gradient Estimation in Machine Learning
• A Regularized Framework for Sparse and Structured Neural Attention
• SparseMAP: Differentiable Sparse Structured Inference
• Sparse projections onto the simplex
• Learning to Play Guess Who? and Inventing a Grounded Language as a Consequence
• A Tutorial on Deep Latent Variable Models of Natural Language
• ARM: Augment-REINFORCE-Merge Gradient for Stochastic Binary Networks
• Rao-Blackwellized Stochastic Gradients for Discrete Distributions
• Estimating Gradients for Discrete Random Variables by Sampling without Replacement

Cited by in corpus (4)

• Storchastic: A Framework for General Stochastic Automatic Differentiation
• NeuralPDR: Neural Differential Equations as surrogate models for Photodissociation Regions
• Evidential Sparsification of Multimodal Latent Spaces in Conditional Variational Autoencoders
• Sparse Communication via Mixed Distributions