Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This problem addresses the problem of differentiating expectations of discrete random variables. The authors proceed from a derivation of DisARM that applies importance sampling to the RLOO gradient, extending this estimator beyond 2 categories. A core component of this new multi-sample gradient estimator is a coupling which leads to (negatively) correlated samples of the categorical. The main proposed coupling uses stick-breaking. The authors propose two estimators that do not use importance sampling, but extend DisARM to the categorical case by reparametrising a categorical using binary random variables. This approach is amenable to standard DisARM. The authors consider two such binarizations: first using the stick-breaking construction again, second using a decision tree. Finally, the authors propose ARS+ and ARSM+ plus, which are Rao-Blackwellized versions of ARS and ARSM. Experimental evaluation focuses on categorical latent variable VAEs.
main_review: Novelty/related work
-----------------------------
The work constitutes a natural extension of existing work, focusing on moving from binary to categorical random variables. Regarding related work, a piece that seems to be missing is ARMS [1]: a recent (ICML 2021) paper that generalizes RLOO and DisARM, but remains focused on the binary case. To me, the key mathematical contributions seem to be
1. Extending antithetic coupling based methods to categoricals (DisARM-IW)
2. Proposing an explicit coupling using stick breaking (DisARM-IW)
3. Alternatively, two binarizations of categoricals that can then use vanilla DisARM (-SB and -Tree)
4. Rao-Blackwellization of ARS and ARSM to give ARS+ and ARSM+

With all these different variations floating around, I do think the authors should do a little bit more work to clarify the relations between these different methods. For instance, it wasn’t 100% clear to me whether there is any relationship between ARS+ and ARSM+ and the DisARM-* methods. My intuition was no, because the DisARM-* methods don’t use any continuous augmentation that could be integrated out. I think clarifying this would be very helpful, perhaps with some kind of table.

Quality
---------
The paper is of good quality. Although I haven’t read into the appendix in great detail, I have good confidence in the underlying mathematics of the approach. One thing that could have further grounded this work is stronger theoretical analysis. When considering the properties of a good coupling, the authors say (112-114) “moving mass away from this configuration will reduce variance. Furthermore, for the estimator to be valid, the coupling must put non-zero mass on all $z_k \ne \tilde{z}_k$ configurations.” But this is not followed up by more careful analysis. For instance, the variance of the new estimator is not calculated / bounded. Furthermore in 279-280 “DisARM-IW is simpler to implement, more natural to understand, and easier to generalize, we recommend it in practice”. This kind of comparison between different proposed methods would again have been stronger with some kind of analysis or calculation of the actual variance, to guide a selection of the right estimator. 

On the empirical side, I would say the authors have been fairly detailed here. They have looked at three sensible datasets, and compared against suitable baselines. It was a tiny bit confusing to separate out the comparison as DisARM-* versus RLOO, and secondarily RLOO versus ARS(+) and ARSM(+). But I think the logic makes sense here: the estimators in the first group only require 2 function evaluations. One direction that the experiments could have been pushed is to consider other models beyond the VAE. Could these estimators be applied to belief nets, or would be number of discrete latents be problematic?

On strengths and weaknesses, one thing that could be discussed in more detail is the relative computational cost of all of the estimators being discussed in this paper (as there are many). For the DisARM-Tree estimator, is there any cost to calculating the tree probabilities (if C was extremely large, perhaps this becomes problematic)?

Clarity
---------
I found the paper to be well-written. However, there are some places where I think the authors could be clearer. As mentioned previously, a more precise and self-contained comparison between the different estimators proposed in this paper (plus some baselines) would help. I understand that this information is present in the paper, but it can be hard to locate. For instance, a table with concrete information such as complexity, continuous augmentation (y/n) and so on could be a real asset. One thing that should be stressed is the distinction between DisARM-IW with a stick breaking coupling and DisARM-SB, there is definitely a possibility for confusion here. 

Significance
----------------
This paper deals with an important area of ML research. It makes some sensible and logical next steps in the development of low variance estimators for discrete random variables. The empirical gains are not colossal, but they do provide new state of the art performance. I would say this paper is definitely a valuable contribution, without being a huge leap forward.

[1] Dimitriev, Alek, and Mingyuan Zhou. "ARMS: Antithetic-REINFORCE-Multi-Sample Gradient for Binary Variables." arXiv preprint arXiv:2105.14141 (2021).
limitations_and_societal_impact: I think the authors have addressed wider societal impact adequately. On limitations, see my main review.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: The authors consider optimization problems that can be formulated as maximizing parameterized functions that can be expressed as an expectation of a function f with respect to some factorial catergorical distribution.  These kinds of optimization problems arise naturally in standard mean-field ELBO optimization problems.  The proposed approaches aim to go beyond existing work by introducing a stick-breaking based coupling technique to reduce the variance of the sampling based estimator.  This is done by first presenting an importance-sampling characterization of DisARM, from which the newly proposed techniques are derived.
main_review: Originality:  While the work does build on top of existing work, the perspective is new.

Quality:  1) The claim that, "In the case of nonlinear models (Figure 1), all three proposed estimators perform similarly, with lower
240 gradient variance and better performance than the baseline estimator (RLOO) across all datasets," does not seem to match the figures, e.g., not on Omniglot.

2) The data sets selected for comparison are all quite similar.  It is unclear to me that similar behavioral trends would hold for a wider variety of data sets.  Can the authors explain if they looked at other data sets and why these data sets were chosen?

3) The Rao-blackwellized portion of the experiment section feels like a bit of an afterthought and most of it is relegated to the appendix.  I think I'd prefer a more detailed study of disarm variants.

Clarity:  1) While this paper is a bit out of area for me, it is incredibly difficult to read as many assertions are made without concreate explanations or citations of why something is true, e.g., "Monte Carlo estimates of the first term can have large variance."  This should have a citation.  There are many more examples of this, too many to reasonably include here which greatly limited my ability to evaluate and certify the results.

2) The whole discussion of related work seems like a bit of a mess to me as the paper bounces back and forth from discussing related work, a background section on related work, a section that builds on related work, and then a separate related work section. I think these sections can be combined and streamlined so as to minimize repetition of information and allow for additional explanations where needed.

Significance:  The performance improvements of the new approach range from mild to moderate, but the perspective may be of interest more broadly.
limitations_and_societal_impact: In my opinion, no detailed discussion of limitations is provided.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper reformulates the DisARM/U2G (Dong et al., 2020; Yin et al., 2020) gradient estimators using importance sampling. While these estimators are originally applicable to binary latent variables, the new formulation allows for the extension to the categorical latent variables case by making use of couplings (common random numbers). The paper also shows that another family of estimators ARS & ARSM (Yin et al., 2019) is not performant compared to RLOO (leave-one-out REINFORCE) (Kool, et. al, 2009) even when its variance is reduced further with Rao-Blackwellization.

main_review: **Originality:** While the paper builds on previous work, the reformulation of DisARM/U2G (Dong et al., 2020; Yin et al., 2020) with importance sampling and the use of couplings are novel. The experimental evaluation is standard for this type of problem and uses the usual datasets and models. The paper does a decent job at positioning its contributions within the literature and differentiates itself from its direct comparisons; however, it still misses some relevant references, e.g. Lee et al. 2018, Peters and Welling 2018, Shayer et al. 2018, Cong et al. 2019, Richter, et. al. 2020.

**Quality:** I think this work is technically solid. The mathematical derivations seem to be correct to me (although I did not verify all of them in detail) and the motivation was clear. Although, I think that the paper is lacking some more explanation and intuition on why this type of estimator reduces the variance (I had to go to Owen 2009 to get some intuition). 

One thing I was not able to verify is how computationally expensive this estimator is? In the experiments section, the authors state that: “These estimators require only 2 expensive function evaluations regardless of the number of categories (C) ”. How expensive is this compared to the 2-sample RLOO estimator?  Does the SB and Tree coupling constructions have any additional overhead? I think these questions should be addressed in the experimental section. Currently, the experiments only show the ELBO with respect to the training steps. I think a similar comparison needs to be performed but with respect to wall-clock time to show the effectiveness of this estimator in practice. 

Furthermore, it is also good to include a comparison to other popular gradient estimators for categorical variables such as REBAR (Tucker, et. al. 2017) and RELAX (Grathwohl, et. al. 2018). While these estimators are expected to perform well because of their adaptivity, they are quite expensive to compute. I think that benchmarking w.r.t. wall-clock time might reveal that DisARM-IW is competitive in comparison.

Finally, I found the benchmarking of ARS+ & ARSM+ separately to be strange, since the main thesis of this paper is that these do not close the gap to RLOO motivating the need for DisARM-IW. I think the comparison with DisARM needs to be direct in this instance. Please see further comments below on this point.

**Clarity:** I think the paper is well-written and the technical details are mostly clear. However, I found the discussion of ARS & ARSM a little bit arbitrary. These estimators are only considered in the experiments and in section 3.2 on their Rao-Blackwellization. 

Further, these are not compared to the contributed estimator (DisARM-IW/SB/Tree). I understand that their inclusion is to demonstrate they are not as effective in reducing the variance as the proposed estimator or RLOO, but couldn’t this comparison be made directly? At present, this paper reads as two distinct contributions (A major one: DisARM-IW and a minor one: Rao-Blackwellization of ARS & ARSM) in a single manuscript. I think the manuscript will be much clearer if the discussion of ARS & ARSM is integrated better in the exposition (including the experiments). Alternatively, this discussion could be demoted to the appendix or submitted as a workshop paper in the future.

Since this paper is relatively heavy on notation, I would suggest the following three points to improve its clarity:
* Needed: An algorithmic description (pseudocode) of DisARM-IW and the computation of the different couplings (some of this can be added to the supplement).
* Nice to have: following the standard notation of representing vectors in boldface.
* Nice to have: definitions of undefined symbols, e.g. $1_C$.

**Significance:** I think this paper provides a significant contribution to a few ML communities (mainly probabilistic modelling and RL). The introduction of the idea of coupling for variance reduction is an important contribution to these communities and can create many avenues for more research into variance reduction for Variational Inference and Policy Gradients. However, the impact of this paper is hampered by the limited experimental evaluation. That being said, I am willing to reconsider my evaluation based on the authors’ rebuttal/clarifications and the discussion with the other reviewers.

limitations_and_societal_impact: There was not much discussion of the technical limitations of this work in the manuscript. It would be great if the authors can comment on those, especially if there is any extra computational overhead.

The potential societal impact of this work was sufficiently considered.
review_title: null

final_score: rating: 4: Ok but not good enough - rejection
review_content: summary: The paper proposes three versions of DisARM/U2G estimators.
The authors extend the binary version to the categorical random variables, and the extentions are in the importance sampling sense (IW) and stick-breaking coupling perspecitve (SB and Tree).
They also verify their method through VAE experiments.
main_review: Originality:
Extending the idea of binary random variable gradient estimator to the categorical case is natural.
Also, constructing categorical variable through stick-breaking or tree construction is also a natural idea.
However, combining those two ideas in the research on the gradient estimotors for the categorical variable seems novel.

Quality:
The paper seems technically sound.
The appendix substantiate the unbiasness of the proposed gradient estimators.

Clarity:
I can't say that the paper is well-written (it seems that it can be written in some better way), the paper is written in relatively clear manner.

Significance:
My concern on this paper is in two parts.
The first one is that the experiment is too limited.
The baselines are limited and the experimental settings, where the reparameterization estimators are excluded and a lot of score function estimators are missing. (including UnOrd [1])
Also, the experimental task is limited to VAE case.
As the research on the gradients over categorical random variables grows, more extreme experiments are gradually required to differentiate the effect of the proposed methods.
I suggest authors demonstrate the idea in the broader experimental settings.
The second concern is that the method (or the implementation) is more complex than the sole baseline, RLOO, but the performance gap is too marginal.
I guess the authors should provide an example that the proposed complex method gives a clear benefit compared to the baseline.
Therefore, I lean my score to reject.

Questions:
- How does the running time of each model as the number of latent categories changes? (compared to RLOO) Especially, the stick-breaking construction might take a long time.
- From the learning curves in Fig 1 and 2, it seems that the models are not reached the optimal point, and it seems they still have a learning chance. Am I missing something?
- From the right-most figures of Fig 1, the variance of gradients reduces up to a certain point and keeps increases later. What causes such curve shapes?
- How does the performance change (including RLOO) if one sets the extreme prior in VAE?

[1] Estimating gradients for discrete random variables by sampling without replacement., Kool et al., ICLR 2020.

limitations_and_societal_impact: The authors argue that they pointed out the limitation of their work in the conclusion (but I can not find it), and there is no negative societal impact.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper considers the problem of designing a gradient estimator for categorical latent variables that is unbiased and does not depend on continuous relaxation. It proposes three different types of such estimators (called DisARM-IW, DisARM-{SB,Tree}, and {ARS,ARSM}+). DisARM-IW is based on a new derivation, which uses a coupling, of a previous estimator called DisARM; DisARM-{SB,Tree} are obtained by reparameterizing with binary variables and then applying DisARM; and {ARS,ARSM}+ are obtained by Rao-Blackwellizing previous estimators called {ARS,ARSM}. DisARM-{IW,SB,Tree} are shown to outperform a state-of-the-art estimator called RLOO on several benchmarks.

main_review: ## Summary of Review

The paper considers an important problem and proposes new methods, some of which outperform the state-of-the-arts. One of the proposed methods (DisARM-IW) looks novel to me, but the other methods look less novel. The paper is technically sound. The paper needs more improvement in writing.

## Strengths

(1) This paper considers an important problem and proposes several solutions. Among the proposed solutions, DisARM-IW and DisARM-{SB,Tree} are shown to outperform a state-of-the-art estimator.

(2) DisARM-IW is novel to me in that its derivation looks new and it can be seen as a generalization of a previous estimator called DisARM. 

## Weaknesses

(1) The estimators other than DisARM-IW look incremental and much less novel to me. DisARM-{SB,Tree} are obtained by applying reparameterization based on the stick-breaking or tree construction and then simply passing the reparameterized form to the well-known estimator DisARM. {ARS,ARSM}+ are obtained by simply following an approach similar to DisARM. I think little novelty in DisARM-{SB,Tree} is less problematic as they outperform RLOO.

(2) The writing needs more improvement. First, even though an antithetic Bernoulli coupling is used importantly in several parts of this paper, it is not explained at all in detail. Second, it was difficult to grasp the overall picture of, and connections between, proposed estimators. I was misled that DisARM-IW and DisARM-{SB,Tree} are similar types of estimators, which is not the case: the former does not use DisARM at all, while the latter uses DisARM importantly. In addition I felt that {ARS,ARSM}+ and Sec 3.2 are completely separate from the other estimators DisARM-* and Sec 3-3.1. This makes the paper less coherent and makes it difficult to understand the main point of the paper.

## Questions and comments

- Line 127-128. p(b_k,i) > 0.5 should be p(b_k,i=1) >= 0.5 (note >= instead of >)? If this is the case, isn’t it a problem if q(z_k=i) / \sum_{j=i}^C q(z_k=j) = 0.5 (Line 131), as this implies p(z_k>i, \tilde{z_k}>i) = 0 and violates the support condition?
- Line 106. Include the proof or calculation (to text or appendix) showing that Eq (3) recovers DisARM for a particular coupling, as this is one of the major contributions of this paper.

- Line 78-79. \sigma is undefined.
- Line 86-87. j is undefined.
- Line 133. \sum_{j=i+1}^C … should be \sum{j=i}^C?

-----
**Updates after the author response.** Thank the authors for answering my questions. The response gives a more clarified view of the paper. I am still positive on the paper and will keep the same score. 
limitations_and_societal_impact: Limitations and societal impact are discussed in the paper.
review_title: null

Meta-review: The submission is concerned with low-variance gradient estimators for discrete random variables, extending recent work on binary to categorical. The reviewers agreed the work here is novel and interesting, but shared concerns around the lack of baselines and clarity.  The authors promised to improve clarity, and add more experimental results [but did not post these in the rebuttal]. An additional expert review was sought and the AC agrees with the expert reviewer that the submission’s overall contribution is enough to justify acceptance. The authors are strongly encouraged to improve the writing and add more experiments, as promised to the reviewers.
Decision: Accept (Poster)