DBsurf: A Discrepancy Based Method for Discrete Stochastic Gradient Estimation

TL;DR

DBsurf is a novel gradient estimator for discrete distributions that reduces sampling discrepancy, leading to lower variance and improved performance in tasks like VAE training and neural architecture search.

Contribution

Introduces DBsurf, a reinforce-based estimator with a new sampling method that minimizes distribution discrepancy, enhancing gradient accuracy in discrete stochastic models.

Findings

DBsurf achieves lowest variance in benchmark least squares tasks.

DBsurf outperforms existing estimators in training VAEs across datasets.

DBsurf enables a simple, efficient NAS with state-of-the-art results.

Abstract

Computing gradients of an expectation with respect to the distributional parameters of a discrete distribution is a problem arising in many fields of science and engineering. Typically, this problem is tackled using Reinforce, which frames the problem of gradient estimation as a Monte Carlo simulation. Unfortunately, the Reinforce estimator is especially sensitive to discrepancies between the true probability distribution and the drawn samples, a common issue in low sampling regimes that results in inaccurate gradient estimates. In this paper, we introduce DBsurf, a reinforce-based estimator for discrete distributions that uses a novel sampling procedure to reduce the discrepancy between the samples and the actual distribution. To assess the performance of our estimator, we subject it to a diverse set of tasks. Among existing estimators, DBsurf attains the lowest variance in a least squares problem commonly used in the literature for benchmarking. Furthermore, DBsurf achieves the best results for training variational auto-encoders (VAE) across different datasets and sampling setups. Finally, we apply DBsurf to build a simple and efficient Neural Architecture Search (NAS) algorithm with state-of-the-art performance.