Straight-Through Estimator as Projected Wasserstein Gradient Flow
Abstract
The Straight-Through (ST) estimator is a widely used technique for back-propagating gradients through discrete random variables. However, this effective method lacks theoretical justification. In this paper, we show that ST can be interpreted as the simulation of the projected Wasserstein gradient flow (pWGF). Based on this understanding, a theoretical foundation is established to justify the convergence properties of ST. Further, another pWGF estimator variant is proposed, which exhibits superior performance on distributions with infinite support,e.g., Poisson distributions. Empirically, we show that ST and our proposed estimator, while applied to different types of discrete structures (including both Bernoulli and Poisson latent variables), exhibit comparable or even better performances relative to other state-of-the-art methods. Our results uncover the origin of the widespread adoption of the ST estimator and represent a helpful step towards exploring alternative gradient estimators for discrete variables.
Cite
@article{arxiv.1910.02176,
title = {Straight-Through Estimator as Projected Wasserstein Gradient Flow},
author = {Pengyu Cheng and Chang Liu and Chunyuan Li and Dinghan Shen and Ricardo Henao and Lawrence Carin},
journal= {arXiv preprint arXiv:1910.02176},
year = {2019}
}
Comments
Accepted as NeurIPS 2018 Bayesian Deep Learning Workshop