ARM: AUGMENT-REINFORCE-MERGE GRADIENT FOR STOCHASTIC BINARY NETWORKS
To backpropagate the gradients through stochastic binary layers, we propose the augment-REINFORCE-merge (ARM) estimator that is unbiased, exhibits low variance, and has low computational complexity. Exploiting variable augmentation, REINFORCE, and reparameterization, the ARM estimator achieves adaptive variance reduction for Monte Carlo integration by merging two expectations via common random numbers. The variance-reduction mechanism of the ARM estimator can also be attributed to either antithetic sampling in an augmented space, or the use of an optimal anti-symmetric "self-control" baseline function together with the REINFORCE estimator in that augmented space. Experimental results show the ARM estimator provides state-of-the-art performance in auto-encoding variational inference and maximum likelihood estimation, for discrete latent variable models with one or multiple stochastic binary layers. Python code for reproducible research is publicly available.
Overview.
The authors present an algorithm for lowering the variance of the score-function gradient estimator in the special case of stochastic binary networks. The algorithm, called Augment-REINFORCE-merge proceeds by augmenting binary random variables. ARM combines Rao-Blackwellization and common random numbers (equivalent to antithetic sampling in this case, due to symmetry) to produce what the authors claim to be a lower variance gradient estimator. The approach is somewhat novel. I have not seen other authors attempt to apply REINFORCE in an augmented space and with antithetic samples / common random numbers, and Rao-Blackwellization. This combination of techniques may be a good idea in the case of Bernoulli random variables. However, due to a number of issues discussed below, this claim is not possible to evaluate from the paper.
Issues/Concerns
- I assess the paper in its current form as too far below the acceptable standard in writing and in clarity of presentation, setting aside other conceptual issues which I discuss below. The paper contains many typos and a few run-on sentences that span 5-7 lines. This hinders understanding substantially. A number key terms are not explained, irregularly. Although the paper assumes that readers do not know the mean and a variance of a Bernoulli random variable, or theof  definition of an indicator function, it does not explain what random variable augmentation means. The one sentence that comes close to explaining it seems to have a typo: "From (5) it becomes clear that the Bernoulli random variable z ∼ Bernoulli(σ(φ)) can be reparameterized by racing two augmented exponential random variables ...". It is not clear what is meant by "racing," here, and I do not find it clear from equation (5) what is going on. Unfortunately, in the abstract, the paper claims that variance reduction is achieved by "data augmentation," which has a very specific meaning in machine learning unrelated to augmented random variables, further obfuscating meaning. Similarly, the term "merge" is not explained, despite the subheading 2.3.
- Computational issues are not addressed in the paper. Whether or not this method is useful in practice depends on computational complexity
- No effort is made to diagnose the source of the variance reduction, other than in the special case of analytically comparing with the Augment-REINFORCE estimator, which does not appear in any of the experiments. 
- No effort is made to empirically characterize the variance of the gradient estimator, unlike Tucker et al (2017) and Grathwohl et al. (2018).
- The algorithm presented in the appendix appears to only address single-layer stochastic binary networks, which are uninteresting in practice.
- Figure 2 (d), (e), and (f) all show that ARM was stopped early. Given that RELAX and REBAR overfit, this is a little troubling. Overal, these results are not very convincing that ARM is better, particularly in the absence of variance analysis (empirically, or other than w.r.t. the same algorithm without the merge step). All algorithms should be run for the same number of steps, particularly in cases where they may be prone to overfitting.
- Figure 1 I believe contains an error for the REINFORCE figure. In my own research I have run these experiments myself, with a value of p close to the one used by the authors. REBAR and RELAX both reduce to a REINFORCE gradient estimator with a control variate that is differentiably reparametrizable, and so the erratic behaviour of the REINFORCE estimator in this case is likely wrong.
- There is a mysterious sentence on page 6 that refers to ARM adjusting the "frequencies, amplitudes, and signs of its gradient estimates with larger and more frequent spikes for larger true gradients"
-The value to the community of another gradient estimator for binary random variables is low, given the plethora of other methods available. Given the questions remaining about this methodology and its experiments, I recommend against publication on this basis also.
- Table 2 compares results that mix widely different architectures against each other, some taken directly from papers, others possibly retrained. This is not a valid comparison to make when evaluating a new gradient estimator, where the model must be fixed. 
* EDIT: I have re-evaluated the careful and comprehensive response to my concerns by the authors. I thank them for their effort in this. As many of the concerns were related to communication and have been addressed in the most recent draft, I think it is appropriate to move my review upwards. The revisions make this paper quite different from the original, and I am happy to re-evaluate on that basis--this is a peculiarity of the ICLR open review procedure, but I consider it a strength. 
I note that "data augmentation" in machine learning appears to have collided with a term in the Bayesian statistics literature, and the authors have provide a number of citations to support this. I strongly recommend "variable augmentation" going forward, as that is an accurate description (you are augmenting a random variable, rather than the input data domain). This appears to be one of the growing pains of the field of ML which has distinct and often orthogonal concerns to classical statistics around density approximation and computational issues.*