Reproduce “Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Max B. Paulus, Chris J. Maddison, Andreas Krause
Researchers from ETH Zürich and the Vector Institute introduce the Gumbel-Rao (GR) estimator, which applies Rao-Blackwellization to the Straight-Through Gumbel-Softmax (ST-GS) gradient estimator to provably reduce its mean squared error without increasing computational cost. The method consistently improved test accuracy and convergence speed in models with discrete latent variables across various tasks, particularly at lower temperatures and smaller batch sizes.
Key insights
- Rao-Blackwellization provides a provable reduction in the mean squared error (MSE) of the ST-GS gradient estimator, ensuring an improved gradient estimate without introducing additional bias.
- The Gumbel-Rao approach effectively reduces variance at low temperatures and small batch sizes, regimes where ST-GS often performs poorly, allowing for more stable training closer to true discrete sampling.
- By modifying only the backward pass, the method retains the benefits of straight-through estimation, such as exploiting sparsity and avoiding error accumulation from relaxed forward passes.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.