Straight-Through Estimator as Projected Wasserstein Gradient Flow
This provides a theoretical foundation for a widely used method in machine learning for handling discrete variables, though it is incremental in nature.
The paper tackled the lack of theoretical justification for the Straight-Through estimator by interpreting it as a projected Wasserstein gradient flow, establishing convergence properties and proposing a variant that performs better on distributions like Poisson.
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.