Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces
arXiv:2110.15072
Abstract
Structured latent variables allow incorporating meaningful prior knowledge into deep learning models. However, learning with such variables remains challenging because of their discrete nature. Nowadays, the standard learning approach is to define a latent variable as a perturbed algorithm output and to use a differentiable surrogate for training. In general, the surrogate puts additional constraints on the model and inevitably leads to biased gradients. To alleviate these shortcomings, we extend the Gumbel-Max trick to define distributions over structured domains. We avoid the differentiable surrogates by leveraging the score function estimators for optimization. In particular, we highlight a family of recursive algorithms with a common feature we call stochastic invariant. The feature allows us to construct reliable gradient estimates and control variates without additional constraints on the model. In our experiments, we consider various structured latent variable models and achieve results competitive with relaxation-based counterparts.
arXiv commentAccepted as a conference paper at NeurIPS 2021
References in corpus (14)
- Show, Attend and Tell: Neural Image Caption Generation with Visual Attention
- Adaptive Computation Time for Recurrent Neural Networks
- Attention, Learn to Solve Routing Problems!
- REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models
- Backpropagation through the Void: Optimizing control variates for black-box gradient estimation
- Non-Monotonic Sequential Text Generation
- Learning to Compose Words into Sentences with Reinforcement Learning
- Differentiable Perturb-and-Parse: Semi-Supervised Parsing with a Structured Variational Autoencoder
- Gradient Estimation with Stochastic Softmax Tricks
- Revisiting Reweighted Wake-Sleep for Models with Stochastic Control Flow
- Differentiable Top-k Operator with Optimal Transport
- Coupled Gradient Estimators for Discrete Latent Variables
- Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator
- Discovering Non-monotonic Autoregressive Orderings with Variational Inference