Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: Summary:
* This paper proposes a Rao-Blackwellized version of the straight-through gumbel-softmax gradient (STGS) estimator.
* The Gumbel-Rao estimator remains single-evaluation (but multiple sample), does not have higher variance than the original straight-through estimator.
* The estimator exhibits lower variance at lower temperatures in the experiments.

Contributions:
* Proposes a single-evaluation estimator that cannot have higher variance than the STGS gradient estimator.
* Demonstrates effectiveness of proposed estimator in terms of the variance of the gradient estimator and the ELBO on a toy task, a simple parsing task (ListOps), and a mixture model for MNIST.

Strengths:
* The method is simple and the computational overhead is very small compared to the original STGS estimator.
* The empirical results support lower variance claims and effectiveness at lower temperature.

Weaknesses:
* I am not convinced that the relative gains from training at lower temperatures are significant.
* The overall gains over ST-GS seem to be modest on MNIST as well as the L <= 50 setting in ListOps.
* In the ListOps experiments, lower temperatures barely achieved better accuracy.

Decision: Marginally below acceptance threshold
* Improving gradient estimators for discrete latent variable models is an important problem.
* The method is straightforward and the claims of performing better at lower temperatures are supported by empirical evidence.
* However, the overall performance on the ListOps dataset is lower than related work [1], and there does not appear to be a large gain from low temperatures.

Questions:
* The main argument of this paper hinges on the claim that lower temperatures result in lower bias of the gradient estimator. This claim seems reasonable, and is supported by figure 2b. Is there a proof or citation for it, and do we know more? It would be nice to know how variance and bias are traded off, as that would tell us how much (or how little) we could gain from training at lower temperatures.
* Is there an explanation for the difference in performance between the 99% accuracy obtained in Havrylov et. al. 2019 [1] and the performance obtained at low temperatures in this paper?
* How does this method perform versus the estimator proposed in Pervez et. al. [2], which is also single-evaluation?

Suggestions:
* The GR estimator is not guaranteed to have lower variance than ST-GS, just not higher.
* Is there an application where lower temperatures are necessary for training? That would strengthen the argument.

[1] Serhii Havrylov, German Kruszewski, and Armand Joulin. Cooperative Learning of Disjoint Syntax and Semantics. In Proceedings of NAACL 2019.

[2] Pervez, A., Cohen, T., & Gavves, E. 2020. Low Bias Low Variance Gradient Estimates for Hierarchical Boolean Stochastic Networks. ICML 2020.

Edited score after author comments.
review_title: Method allows training of gumbel straight-through at lower temperatures, but lower temperature gains seem small

final_score: rating: 7: Good paper, accept
review_content: 

This paper introduces the Rao-blackwellization technique to reduce the variance of the straight-through gumbel-softmax gradient (STGS) estimator wrt the parameters of discrete distributions. The proposed method introduces almost trivial computational costs (relative to function evaluations) and is empirically and theoretically shown to systematically improve STGS. 

I don’t have a lot of nitpicking to make for this paper, as it is quite well executed. The proposed method is very clean and the improvement over the STGS baseline is very consistent, and makes it even competitive with concrete-relaxation in the discrete latent variable model experiment. 

Details:

Why not show the curve of ELBO during training, but the arrival-time-at-certain-thresholds in Fig 2-c? 
Last paragraph of sec 5.4: The larger batch size here also reduces the variance of minibatch SGD, not just the variance of $\nabla_{GR}$ in (13). In fact each instance is a different approximate posterior, which has different base variance. This makes the discussion in 3.3 a bit misleading. 

Suggestions:
For figure 1, perhaps visualize the variance of both separately. An improvement by 2 is not that meaningful if the variances of both are >> 2. 

-- After rebuttal

Thank you for revising the paper. I've read the revised section, and stand by my original evaluation. 
review_title: cheap and effective variance reduction trick to improve a widely adopted gradient estimator (for discrete distributions)

final_score: rating: 8: Top 50% of accepted papers, clear accept
review_content: Summary: The paper presents a new way algorithm to compute the straight-through variant of the Gumbel Softmax gradient estimator. The method does not change the estimator's bias, but provably reduces its variance (with a small overhead, using Rao-blackwellization). The new estimator shows good performance on different tasks, and appears to lead to more efficient optimization for lower temperatures (lower bias).

Clarity: The paper is well written.

Originality: The use of Rao-blackwellization in the proposed way is, up to the best of my knowledge, novel.

Pros of the paper and significance: 
- Relaxation-based gradient estimators are widely used, and the proposed method may their variance quite significantly.
- The proposed algorithm has a clear justification from a theoretical perspective, and admits a simple implementation.
- The proposed algorithm does not require additional model evaluations, and thus may lead to large reductions in variance without incurring a high computational cost.
- The proposed method leads to more efficient optimization at lower temperatures (lower temperature translates to lower bias, but often higher variances).

Cons: I'd say one thing that could be included are additional baselines in the experimental section. There are other estimators that may be use. For instance, you could compare against VIMCO. While this is a different type of estimator (non single evaluation, not based on relaxations), it could be interesting to see how the results compare using this estimator too.

Recommendation: Accept (reasons in the "pros" list above).
review_title: Good work

Meta-review: The paper presents a variance reduction technique to the Straight-Through version of the Gumbel-Softmax estimator. The technique is relying on the truncated Gumbel of Maddison et al. I share the excitement of the reviewers about this work and I expect this technique to further influence the field. 
Decision: Accept (Oral)