Training Discrete Deep Generative Models via Gapped Straight-Through Estimator

TL;DR

This paper introduces the Gapped Straight-Through (GST) estimator, a novel method for training discrete deep generative models that reduces gradient variance without resampling, improving performance on benchmark tasks.

Contribution

The paper proposes the GST estimator, a new variance reduction technique inspired by Straight-Through Gumbel-Softmax, eliminating resampling overhead in training discrete models.

Findings

GST outperforms strong baselines on MNIST-VAE.

GST achieves better results on ListOps.

The ablation study confirms the importance of key properties of the estimator.

Abstract

While deep generative models have succeeded in image processing, natural language processing, and reinforcement learning, training that involves discrete random variables remains challenging due to the high variance of its gradient estimation process. Monte Carlo is a common solution used in most variance reduction approaches. However, this involves time-consuming resampling and multiple function evaluations. We propose a Gapped Straight-Through (GST) estimator to reduce the variance without incurring resampling overhead. This estimator is inspired by the essential properties of Straight-Through Gumbel-Softmax. We determine these properties and show via an ablation study that they are essential. Experiments demonstrate that the proposed GST estimator enjoys better performance compared to strong baselines on two discrete deep generative modeling tasks, MNIST-VAE and ListOps.