Coupled Gradient Estimators for Discrete Latent Variables

Abstract: Training models with discrete latent variables is challenging due to the high variance of unbiased gradient estimators. While low-variance reparameterization gradients of a continuous relaxation can provide an effective solution, a continuous relaxation is not always available or tractable. Dong et al. (2020) and Yin et al. (2020) introduced a performant estimator that does not rely on continuous relaxations; however, it is limited to binary random variables. We introduce a novel derivation of their estimator based on importance sampling and statistical couplings, which we extend to the categorical setting. Motivated by the construction of a stick-breaking coupling, we introduce gradient estimators based on reparameterizing categorical variables as sequences of binary variables and Rao-Blackwellization. In systematic experiments, we show that our proposed categorical gradient estimators provide state-of-the-art performance, whereas even with additional Rao-Blackwellization previous estimators (Yin et al., 2019) underperform a simpler REINFORCE with a leave-one-out-baseline estimator (Kool et al., 2019).


Reviews and author responses:

 Thank you! Yes, we intend to substantially revise the text for clarity and reorganize how the related works section is done.

Yes, you are correct. Thank you for catching that typo in the equation; we can confirm it is implemented properly in the code.


The paper concerns the estimation of gradients of expectations when sampling categorical random variables.
The work is closely related to recent works ARM and DisARM for binary variables, and ARS, ARSM for categorical random variables.
ARM was an estimator that samples a pair of binary random variables and couples them via antithetic sampling. This was achieved by sampling a continuous variable from a logistic distribution in an antithetic manner, and then transforming this continuous variables via a step function to obtain binary variables. The final form of the estimator in ARM still depended on the sampled continuous variable. DisARM noticed that this is unnatural, as in principle only the final binary variable should matter, and they proposed a Rao-Blackwellized estimator that integrates out the continuous random variable and obtains a strictly better gradient estimator compared to ARM. Both ARM and DisARM were meant for binary variables. 

ARS and ARSM instead worked for categorical variables. The current work first proposes to Rao-Blackwellize ARS and ARSM leading to ARS+ and ARSM+. Then they propose 3 ways to extend DisARM to the categorical case: 1) they perform a new derivation of DisARM based on importance sampling, then perform an antithetic coupling of the categorical variables by constructing them via a stick-breaking process from binary variables, and computing the importance weight; 2) instead of computing probabilities in the categorical space, the problem is reparameterized into a vector of binary variables motivated from the stick-breaking construction, and the DisARM (recall that the original one could be used with binary variables) estimator is applied in this binary space together with a stick-breaking coupling (they also use knowledge of the stick-breaking coupling to further improve the estimator); 3) construct the categorical variable via a binary tree, and apply the DisARM estimator on these binary variables together with similar improvements as in the stick-breaking construction case.

Moreover, they compared with REINFORCE with a leave-one-out baseline (RLOO). Previously, it was shown that ARM actually does not improve over the standard REINFORCE with a leave-one-out baseline, while DisARM did improve. In the current work, the authors showed that ARS and ARSM also do not improve over RLOO when the same number of function evaluations are used (5 evaluations for ARS and 11 for ARSM on a 5 category task), while the new categorical DisARM-IW, DisARM-SB and DisARM-tree methods did improve compared to a 2-sample RLOO estimator.

They performed experiments of training categorical variational autoencoders on binarized MNIST, FashionMNIST and Omniglot. The experiments showed a slight improvement in terms of learning speed and gradient accuracy (maybe around 25% -- 50% reduction in gradient variance compared to the 2-sample RLOO). They also tested the orderings of the binary variables in the stick-breaking estimator, and found that the dependence on ordering is inconsistent.  In general the quality of the work was good. The experimental analysis seemed fair, though I would have liked to see how the new estimators scale with the estimator sample size compared to RLOO. The paper presented a sensible technique for extending the coupling technique to categorical random variables. However, the improvement in performance seems small (25 - 50% reduction in gradient variance in practice is not a huge deal), so I am skeptical that that the method will make a major difference in practice. The clarity was mostly good, but could be improved by adding some more definitions, e.g. the antithetic binary variables are not properly defined, and some derivations in the appendix could be extended with the intermediate steps. Moreover, more related literature could have been added.

Currently, I gave it a 6 and may improve to 7 if some points in my review are addressed. 

**Experiments**

I think the comparisons performed in the paper are fair, and it was nice to see analysis from multiple perspectives, e.g. the gradient variance, ordering of the categories, etc. However, for the ARS and ARSM estimator comparisons, RLOO used more samples than were used for DisARM, which makes me question whether RLOO would also not outperform DisARM when the sample size is increased. For example, currently DisARM uses one pair of samples (i.e. 2 samples), when we double this, we expect the variance to decrease by 2 times; however, the same is not true for RLOO---when we double the sample size for RLOO, the LOO baseline will now be using 3 samples, instead of just 1, so I would expect the variance of the gradient estimator to decrease by more than 2 times, due to the better baseline. Therefore, if the sample size is increased, it is plausible that RLOO will outperform the new DisARM estimators, even when matching the sample sizes between RLOO and DisARM. It would have been nice to see an experimental analysis of this. Of course, in this case one could try to construct a coupling between more than 2 samples to try to improve over the standard RLOO, and this could be left for future work.

**Other literature**

The RLOO estimator was used earlier than Kool et al. (2019) by Mnih, A. and Rezende, D. (2016). "Variational inference for monte carlo objectives", and I find it inappropriate to only cite Kool. 

For Rao-Blackwellization, the paper by Titsias et al.  (2015). "Local expectation gradients for black box variational inference." was one of the first to use this method with binary random variables.

Regarding importance sampling there are also several works, e.g.
Ruiz, F. J., Titsias, M. K., & Blei, D. M. (2016). Overdispersed black-box variational inference.
Parmas, P., & Sugiyama, M. (2021, March). A unified view of likelihood ratio and reparameterization gradients.
Jie, T., & Abbeel, P. (2010). On a connection between importance sampling and the likelihood ratio policy gradient.

In addition to the mentioned approaches to deal with discrete variables, biased estimators are also an approach that could be mentioned.
Some methods:
Straight-through estimators:
Yoshua Bengio, Nicholas L ́eonard, and Aaron Courville. Estimating or propagating gradients through stochas-tic neurons for conditional computation.
Path-sample analytic estimators:
Shekhovtsov, A., Yanush, V., & Flach, B. (2020). Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks.

**Clarity**

I think it would be useful to add:
- A definition of antithetic binary random variables, and also explain how these are sampled and transformed from a logistic distribution.
- What is SB? I was able to figure out that it means stick-breaking, but you
should explicitly define the abbreviation when you first use
"stick-breaking".
-Explicitly specify that z_k is an integer in [1,...,C]. Otherwise
one might assume a one-hot encoding.
-In the equation between lines 78--79, please explain the notation in the
text.
-And other points...

----------------------------------------------------
Update:
The authors added experiments looking at the scaling with number of samples. The performance improvement reduced, but the new estimator still outperformed RLOO. I think the paper is a good addition to the literature on coupled gradient estimators for discrete random variables, so I increased my score to 7. Yes.

 OK, thanks. It seems that the advantage becomes smaller than in the 2-sample setting, but it still improves over RLOO.
I think the paper is a good addition to the literature on coupled gradient estimators with discrete random variables, so I'll increase my score.
I hope the authors will further improve the clarity and discussion of related works.

I think I spotted another typo: I guess on line 133, the sum should be from j=i to C not from j=i+1?

 For these multisample settings, DisARM-IW, DisARM-SB, DisARM-Tree all reduce variance relative to RLOO (on all 3 datasets) and perform similarly. Relative to the RLOO variance, the reduction ranges from up to 4% to 10% reduction in variance depending on the dataset and step in the training trajectory. The relative reduction in variance between 10 Samples/5 Pairs and 20 Samples/10 Pairs is similar.

The guidelines for the author response have changed substantially from previous years. From the FAQ (https://neurips.cc/Conferences/2021/PaperInformation/NeurIPS-FAQ)

> *May I include a link in the author response?*  External links are discouraged, but you may include a link if it is required to respond to a question from a reviewer. However, please remember that author responses may not contain any identifying information.  Violating the double-blind reviewing policy in the author response may be grounds for rejection.  Furthermore, if you are including links to any external material, it is your responsibility to guarantee anonymous browsing.  As with supplemental material, reviewers are not required to read linked material in full, but may do so at their discretion.

 We're not supposed to follow external links, so I would appreciate if you could explain the result here.

 You can view the plots of variance for those experiments here: https://mega.nz/folder/fM5iVI5C#E_SOTWNBsaNqIms9ygmSAA

Similarly to the main paper, the variance of all the gradient estimators was computed along the training trajectory of the RLOO estimator to make a fair comparison.

**Summary**
- The paper introduces a gradient estimator for discrete latent variable models. 
- The estimator is a natural extension of estimators that have been recently proposed for binary variables to categorical variables. 

**Edit after Author Response**
Thank you for taking the time to address all my concerns convincingly. I also read through the discussions with other reviewers. In light of your response, and the discussions I raise my score to 7. Even though, the empirical improvements are not very large, I think the insights, methods and experiments are a valuable contribution to the community and justify acceptance. This is a good paper.   **Strengths**
- Originality. The new derivation of DisARM/ U2G is insightful. 
- Quality. The submission is technically sound. Claims are well-supported, derivations are followable. 
- Clarity. The paper is well-written and has a clear, sensible structure. I enjoyed reading the paper. 
- Significance. Gradient estimators for categorical distributions are important. 

**Weaknesses**
- Quality. The experimental section could be more extensive. 
- Significance. Experimentally, the improvements of the proposed method appear marginal to me. 

**Decision**
- I like this submission and think it is above the acceptance threshold. 
- I wished the experimental improvements were larger and the experimental section was richer. 
- I also make some suggestions below which I think would be interesting to explore as part of the paper. 

**Suggestions and/ or questions**

- Does the estimators favorable properties over RLOO arise solely from avoiding “duplicate” samples in expectation or is there more behind it? In the former case, it would be interesting to highlight, how often duplicate samples with RLOO arise in the experiments to highlight that it is a real issue. 

- Are there some settings where the gains over RLOO are expected to be larger? This would strengthen the experimental case for this method, to me the improvements appear marginal. 

- The model was only evaluated on categorical VAEs. If possible, I would like to see the estimators being evaluated in a RL setting, for example discrete-action policy optimization as in Yin et al. (2019), or on the structured prediction prediction task considered in various papers, e.g. Maddison et al. (2016).

- The paper could also compare to other baselines, for example the unordered set estimator from Kool et al. (2020). In general, the case for multi-sample estimators could be made more strongly by comparing to single-sample relaxed estimators and keeping the number of function evaluations constant. However, I do think that RLOO is a strong baseline to compare against. 

- Would it be possible to do hypothesis testing (t-test) to highlight the statistical significance of the improvements of the proposed estimators over the baseline? Statistical significance of the improvements of the method is a concern I have, and this will hopefully highlight that the improvements are significant.

- The authors use Rao-Blackwellization as in Dong et al. (2020) to improve their gradient estimator. Rao-Blackwellization has also been used for improving other gradient estimators for discrete distributions, for example in Liu et al. (2018), Kool et al. (2020) or Paulus et al. (2021), and I think it's worth highlighting that in related work.

**References**
- Liu et al. (2018). Rao-Blackwellized stochastic gradients for disccrete distributions. 
- Kool et al. (2020) Estimating gradients for discrete random variables by sampling without replacement. 
- Paulus et al. (2021). Rao-Blackwellizng the Straight-Through Gumbel-Softmax estimator.  N/A.

 Do you have numbers for how much the gradient variance was?

 We ran experiments with 5 antithetic pairs (10 samples) and 10 antithetic pairs (20 samples) and found that the proposed estimators outperform RLOO with a comparable number of independent samples. Even with the increasing number of samples, the performance improvement of the proposed estimators w.r.t. RLOO still holds.

* **10 Samples/5 Pairs**
 
ELBO    |  DisARM-Tree    |  DisARM-SB      |   DisARM-IW     |        RLOO
--------|-----------------|-----------------|-----------------|----------------
Dynamic MNIST | -91.36&pm;0.11  | -91.67&pm;0.08  | -91.85&pm;0.04  | -92.35&pm;0.10
Fashion MNIST | -231.01&pm;0.14 | -231.38&pm;0.14 | -231.14&pm;0.17 | -231.78&pm;0.16
Omniglot | -109.00&pm;0.10 | -108.89&pm;0.10 | -108.79&pm;0.14 | -109.73&pm;0.07
  
* **20 Samples/10 Pairs**
 
ELBO    |  DisARM-Tree    |  DisARM-SB      |   DisARM-IW     |        RLOO
--------|-----------------|-----------------|-----------------|----------------
Dynamic MNIST | -90.80&pm;0.07  | -91.01&pm;0.06  | -90.95&pm;0.08  | -91.86&pm;0.15
Fashion MNIST | -230.95&pm;0.21 | -231.21&pm;0.18 | -231.25&pm;0.29 | -231.59&pm;0.21
Omniglot | -108.25&pm;0.03 | -108.15&pm;0.06 | -108.08&pm;0.06 | -108.95&pm;0.17

This paper presents the multiple methods of estimating the gradients of discrete random variables. Particularly, they expand DisARM by adding the variations of importance sampling, stick-breaking construction, and tree-based decision makings. The performance is compared to previous methods, such as RLOO, with a standard benchmark dataset in VAE settings.   1)
Given DisARM/U2G, the extension with the importance sampling in Eq 3 seems to be straightforward, and this variant is not congruent with the two other variants with SB and Tree. I think that this could be regarded extension of DisARM/U2G rather than considering this as a separate original version.

2)
Most development has been made to construct the coupling between $z_k$ and $\hat{z}_k$. Utilizing the stick-breaking prior in this regard has been considered in the previous literature, as well, though they are not in the suggested form and purpose [1]. Anyhow, this could be briefly mentioned.

3)
The two variations looks essentially same mechanism of creating the mapping function of $h$. Their differences are subtle, and it is worth enough to create a diagram explanation on their abstrate support sets. Eq 4 and Eq 5 are all same except the conditioned regions. These support regions could be visualized in both variations to express the difference between the two variations. Or, is it that different? I wonder.

4)
Experiments are not satisfactory. There have been many papers in this gradient estimations. The large categorization would be score-based and reparameterization-based approaches. There has been models like NVIL, REBAR, RELAX, STORB, etc. (These are easily searchable through acronym in google) Currently, RLOO is the only baseline in Table 1 and 2. Moreover, the original version of DisARM/U2G can also be compared in Table 1 and 2. 
Datasets are also limited to Omniglot and MNIST, which could be understandable given the prior papers with the same practice, but it seems that this line of work requires further expansion in the dataset, as well.

[1] Nalisnick, Eric, and Padhraic Smyth. "Stick-breaking variational autoencoders." arXiv preprint arXiv:1605.06197 (2016).
 Have the authors adequately addressed the limitations and potential negative societal impact of their work? If not, please include constructive suggestions for improvement.


There is no discussion at all. I understand that authors may not discuss the societal impact, but they can surely include the limitations that they expect from their models.




 Hm.I see. Eq 4 and Eq 5 had to be categorized to anticipate the categorical selection, and the estimator is consistent given the same selection. 

Then, my question becomes the contribution of DisARM-IW, and it seems that it is a consensus between reviewers that it shows the merits over the RLOO.

Thanks for the clarification, and I decided to raise my rating to be 6. Good work.

 We appreciate the reviewer’s time and hope to provide further clarification.

> Now, I understand the history. So, DisARM/U2G is already the first one to use a specific coupling in their own importance sampling, and this paper is only adjusting the coupling function.

DisARM/U2G is derived from a different perspective (continuous augmentation, then Rao-Blackwellization). They do not invoke importance sampling nor the relationship with RLOO and the extension to categorical variables is not clear. One of our contributions is to show that it can be derived in another way which naturally extends to categorical variables.

> Importance Sampling is a common method to increase the parameter inference quality with limited samples, so applying the importance sampling weight is a common extension in any probabilistic modeling, such as VAE to importance weighted VAE. Anyhow, now I see that importance sampling treatment is also not the contribution of this paper.

While we agree that importance sampling is well-known, we emphasize that we do not naively apply importance sampling (this leads to an impractical estimator). DisARM-IW does use per-dimension importance sampling, however, showing that this is unbiased is a non-trivial derivation that leverages the coupling proposal. This is one of our contributions.

> As the gradient method will move across the divided domains, aren't the gradients suddenly changed inbetween the change domain boundaries? Are the gradient smoothly changed? You may provide some theoretic analyses on its smoothness if you want to claim that the gradient method will not be affected by this hard regioning.

No, the gradient estimators are functions of the continuous model parameters $\theta$ and the sampled discrete latent variables $z, \tilde{z}$. The functions are smooth with respect to the model parameters, whereas the “divided domains” arise from the *discrete* latent variable samples. 

> I assume that you know the Lagrange method. The hard constraint can be added to the object function by applying the Lagrange multiplier. Also, the hard regioning can be soften by applying the assignment weights to its regions, and apply the weights to the summed objectives.

This is a misunderstanding. The “hard constraint” is not an optimization constraint. They are simply the cases when defining a function on the discrete latent variables.


 Thank the authors for their response. I raise my rating.

yes, JSA and the proposed method are different. However, for the experiments conducted in this paper, namely training variational auto-encoders with discrete latent variables, they solve the same problem (ELBO is a surrogate objective for log p(x)). From this view, I think, it should inform the readers about this comparison. Experiments beyond training latent variable models (e.g. RL) may be more supportive of the proposed method.

The equation under Line 78 contains a typo:

(alpha -> alpha

Training models with discrete latent variables is an important problem in machine learning. Among previous methods, REINFORCE with a leave-one-out-baseline (RLOO) outperforms Augment-REINFORCE-Merge (ARM) in the binary setting; DisARM/U2G (independently in Dong et al. (2020) and Yin et al. (2020)) outperforms RLOO, but is limited to binary random variables. This work develops a novel derivation of DisARM/U2G based on importance sampling and statistical couplings, which is simpler and more direct while providing a natural extension to the categorical case. Gradient estimators based on reparameterizing categorical variables as sequences of binary variables and Rao-Blackwellization are obtained. Experiments show the effectiveness of the proposed method.  Strengths and weaknesses:

This paper extends DisARM/U2G, which is limited to the binary setting, to the categorical setting. It appears to be a new work.

The paper is generally well-written. But there are a few places, which need further clarification, as detailed below.

A main concern is about performance evaluation and missing of related work. It has been shown in [a] that the joint stochastic approximation (JSA) method significantly outperforms a number of methods including VIMCO, Gumbel-Softmax and ARM/ARSM in both binary and categorical settings for learning discrete latent variable models. Evaluation experiments in [a] are conducted by training VAEs with discrete latent variables, which is also taken in the current paper. The model structures evaluated in the current paper are also similar to [a], which follows from the work of ARSM (Yin et al., 2019). It seems that the proposed method in the current paper is more complicated than [a], in both methodology and implementation. From a quick reading, the results reported in Table 1 are much worse than the Table 2 in [a]. Further, since the proposed method uses the Stick-breaking construction, the cardinality of the discrete latent variable needs to be the integer power of 2. Comments and comparisons are better to be added in this paper.

[a] Z. Ou, Y. Song. Joint Stochastic Approximation and Its Application to Learning Discrete Latent Variable Models. UAI, 2020.

Other problems:

Better to make clear the dimension/size of z, b, etc.

Line 51: "it is straightforward if more challenging to apply the ideas of Dong et al. (2020) to ARS and ARSM"
not clear

The equation under Line 78 contains typo?

For the equation under Line 95, \tilde{z} is used without definition. yes

 * Difference between DisARM-Tree and DisARM-SB

I do not think that both should be same. I am only claiming that they are at least similar, and in some conditions, they are going to be same. 

* Some clarification of DisARM/U2G

Now, I understand the history. So, DisARM/U2G is already the first one to use a specific coupling in their own importance sampling, and this paper is only adjusting the coupling function. 

Importance Sampling is a common method to increase the parameter inference quality with limited samples, so applying the importance sampling weight is a common extension in any probabilistic modeling, such as VAE to importance weighted VAE. Anyhow, now I see that importance sampling treatment is also not the contribution of this paper. 

* Hard regioning

As the gradient method will move across the divided domains, aren't the gradients suddenly changed inbetween the change domain boundaries? Are the gradient smoothly changed? You may provide some theoretic analyses on its smoothness if you want to claim that the gradient method will not be affected by this hard regioning.

* Softening the hard regioning

I assume that you know the Lagrange method. The hard constraint can be added to the object function by applying the Lagrange multiplier. Also, the hard regioning can be soften by applying the assignment weights to its regions, and apply the weights to the summed objectives. 


Anyhow, this paper seems to be determined by assessment of the contribution from the coupling function. I agree that it is a contribution, but it is not strong from my perspective at this moment (this view can be different from different eyes). 
It could have been stronger contribution if there were further theoretic analyses on its general characteristics, i.e. comparison between the regions, smoothness, unified objectives....

At the moment, I update my score to be five, but still this paper may require further investigation to be a whole contribution of what the authors claim for,


 > [Difference between DisARM-Tree and DisARM-SB]

Upon some reflection, the reviewer is correct that in some ways they are quite similar. DisARM-Tree uses a balanced binary tree, whereas DisARM-SB also can be viewed from the tree-structured process where the tree is instead a stick. In this case, Eq 4 is recovered from Eq 5. This is a useful perspective and we will rewrite that section in the next revision. The structure of the tree *does* affect the process (for example, paths in the balanced tree are log_2 C, whereas they can be as long as C in the stick) and a priori it is not clear why they should behave in the same way. If it is clear to the reviewer why they should be the same, then we would appreciate an explanation. However, we do observe that on the datasets and settings we evaluated, we do not find a large difference in performance.

> Additionally, only adding the importance sampling, DisARM-IW, is not bad and sometimes best. Then, wouldn't it be better to just use Eq 3, DisARM-IW, instead of introducing hard condition cases, such as Eq 4 and Eq 5?

Yes, we recommend DisARM-IW in the conclusion, however, we felt that it was useful to present the results because a priori it was not clear which of the estimators would be better.

>4. I don't think that adding DisARM/U2G to the experiment is a good idea, As you suggested, it is designed for the Bernoulli case. 
Actually, this raises some novelty issue of this paper. I think that extending DisARM/U2G to the multi-dimensional distribution could be contribution, but I have to admit that it is a weak one.

While we disagree with the reviewer on novelty and the strength of the contribution, we want to be clear on the content of the contribution. DisARM/U2G applies to multi-dimensional Bernoulli latent variables, and the primary contribution of this work is the derivation and evaluation of estimators for the multi-dimensional Categorical setting. Was this unclear?

> DisARM-IW is an importance sampling imposed on the DisARM/U2G, and I think it is straightforward in the probabilistic modeling world. [...] DisARM-SB and DisARM-Tree are the main result from Stick-Breaking Coupling, but it increases the complexity compared to DisARM-IW, and their performance advantages are minor.

Thank you for laying out your understanding so that we can clarify misunderstandings. 
* DisARM-IW is an estimator based on importance sampling applied to RLOO. 
* Naively applying importance sampling would result in an impractically high variance estimator. DisARM-IW uses a *coupling* as the proposal distribution and takes advantage of the coupling properties to form an unbiased estimator with low variance. DisARM-IW is *not* the naive application of importance sampling to RLOO.
* DisARM/U2G is DisARM-IW with a specific choice of coupling (antithetic Bernoulli). 
* In our categorical experiments, DisARM-IW uses the stick-breaking coupling as the choice of coupling. So both DisARM-IW and DisARM-SB use the stick-breaking coupling.

Can the reviewer clarify what they feel is “straightforward in the probabilistic modeling world”?

> I still wonder whether such hard regioning of Eq 4 and Eq 5 will damage the parameter inference, which fundamentally depends on the gradient method, which is one reason why I would not prefer Eq 4 and Eq 5.

Given the discrete nature of the problem, “hard regioning” is to be expected. Essentially any unbiased estimator (including *all* of the estimators evaluated) have this “hard regioning”.
Can you elaborate on why this would “damage the parameter inference”? Recall that all methods provide *unbiased* estimates of the same gradient.

> It would have been better if you can claim theoretic analyses, i.e. if you apply the Lagrange method or others to soften Eq 4 or Eq 5, can you claim that DisARM-IW could be one case of DisARM-SB? If multi-dimensional case bothers the understanding, such theoretic analyses can be done at the binary Bernoulli case.

We agree that it would be desirable to have a unified estimator that subsumes all of them, however, we do not yet know of a non-trivial way to do this (trivially any convex combination of unbiased estimators is unbiased and it is straightforward to adjust the weights of the convex combination in an optimal way to reduce variance c.f. REBAR/RELAX). Can the reviewer explain what they mean by using the Lagrange method to soften Eq 4 or Eq 5? We do not understand how that is applicable in this case.


 3.
Can you reply a further discussion on the support region of Eq 4 and Eq 5?
Are the support regions very different?

Also, the experiment results show little difference between DisARM-Tree and DisARM-SB.
It is understandable because of the similarity between Eq 4 and Eq 5, and I wonder whether the variation of Tree and SB is meaningful, or not.

Additionally, only adding the importance sampling, DisARM-IW, is not bad and sometimes best.
Then, wouldn't it be better to just use Eq 3, DisARM-IW, instead of introducing hard condition cases, such as Eq 4 and Eq 5?

4.
I don't think that adding DisARM/U2G to the experiment is a good idea, As you suggested, it is designed for the Bernoulli case.

Actually, this raises some novelty issue of this paper. I think that extending DisARM/U2G to the multi-dimensional distribution could be
contribution, but I have to admit that it is a weak one.

Given this thought, your treatment of Stick-breaking coupling and its binary representation requires further detailed explanations and
analyses on its characteristics. 


When I layout what I know,
- DisARM-IW is not bad given Table 1 and 2. Most of experiments favors DisARM-IW, or their confidence interval overlaps significantly.
- DisARM-IW is an importance sampling imposed on the DisARM/U2G, and I think it is straightforward in the probabilistic modeling world.
- DisARM is expanded to be the categorical case, and I think that this is necessary development, but could be a weak contribution.
- Stick-breaking coupling is something new, but it is difficult to follow the scheme with the current writing.The joint density model should be more explicitly illustrated, and some illustrations could be helpful.
- DisARM-SB and DisARM-Tree are the main result from Stick-Breaking Coupling, but it increases the complexity compared to DisARM-IW, and their performance advantages are minor.
- I still wonder whether such hard regioning of Eq 4 and Eq 5 will damage the parameter inference, which fundamentally depends on the gradient method, which is one reason why I would not prefer Eq 4 and Eq 5.


It would have been better if you can claim theoretic analyses, i.e. if you apply the Lagrange method or others to soften Eq 4 or Eq 5, can you claim that DisARM-IW could be one case of DisARM-SB? If multi-dimensional case bothers the understanding, such theoretic analyses can be done at the binary Bernoulli case.

Thanks

 As requested by one of the reviewers, we provide a discussion of the limitations of our methods here and will include it in the next revision of the manuscript.

Parameterizing Categorical variables in terms of Bernoulli variables requires imposing structure on the Categorical space (e.g. an ordering or a tree), which is not fully satisfactory because in most settings there is no such natural structure for Categorical spaces. Developing coupling-based estimators that do not rely on such a structure would be interesting future work which might lead to further improvements.

While we see that our coupling-based estimators generally outperform or perform at least as well as RLOO in our experiments, using coupled samples instead of independent samples is not guaranteed to lead to better performance in every case. Learning the couplings would provide a way of ensuring an improvement over RLOO.

Finally, the estimators we propose in this paper, like all multisample estimators, can be used for RL only if the environment is simulated or we have a model of it, as they require being able to perform multiple rollouts from the same state.


 Thank you for your review.

While the Joint Stochastic Approximation paper does provide an interesting approach to training models with discrete latent variables, it solves a different problem from the one we consider and thus is not a relevant baseline for us. JSA trains generative models by directly optimizing log p(x) for the model rather than a variational bound on it, like we do in the paper. Furthermore, we use variational training of latent variable models merely as a benchmark for unbiased gradient estimators for general expectations, and the JSA method does not provide such an estimator. We do not compare to Reweighted Wake Sleep for the same reasons as for JSA: these methods do not apply beyond training latent variable models (e.g. not to RL), do not provide a general gradient estimator for expectations, and do not optimize the same objective we do. We will emphasize this point in the next revision.

> From a quick reading, the results reported in Table 1 are much worse than the Table 2 in [a].

The datasets differ between the two tables. Table 2 in [a] uses the *statically* binarized MNIST dataset whereas Table 1 uses the *dynamically* binarized MNIST dataset. 

> Further, since the proposed method uses the Stick-breaking construction, the cardinality of the discrete latent variable needs to be the integer power of 2.

The stick-breaking construction does not place any restrictions on the cardinality of the latent variables. On the other hand, our implementation of the tree-structured construction does indeed require the variable cardinality to be a power of 2, but this is not a limitation of the method, just of our initial implementation.

> The equation under Line 78 contains a typo?

Can you clarify where the typo is?

We will address the other notational issues you pointed out.


 Thank you for your thorough review.

We will incorporate your suggestions for improving clarity and adding references in the next revision.

> [RLOO vs DisARM when scaling the number of samples]

This is a great question. The natural extension of DisARM to more samples follows the approach described in [1] (also [2]). Their approach focuses on multisample bounds, however, the same idea applies to multiple samples targeting the original objective. In this case, DisARM also uses multiple samples to construct a baseline too, so it may behave competitively with RLOO. We are running these experiments and will post the results shortly.

References
1. Dong et al., “DisARM: An Antithetic Gradient Estimator for Binary Latent Variables.”, NeurIPS 2020.
2. Dimitriev and Zhou., “ARMS: Antithetic-REINFORCE-Multi-Sample Gradient for Binary Variables.”, ICML 2021.


 Thank you for your review.

> 1) Given DisARM/U2G, the extension with the importance sampling seems to be straightforward […and .. ] could be regarded extension of DisARM/U2G rather than considering this as a separate original version.

Yes, it is an extension of DisARM/U2G. However, we would argue that the derivation is non-trivial. While the equation above Eq. 3 (using the global importance weight) is indeed straightforward, Eq. 3 crucially relies on the coupling to derive *per-dimension* importance weights, which are essential for a scalable method. Without the importance sampling perspective, it is not clear how to extend DisARM/U2G to the categorical setting.

> 2) [use of stick-breaking in VAEs]

We will cite this paper for prior use of stick-breaking in VAEs, though their setting of using stick-breaking to infer the latent space dimensionality is very different from ours.

> 3) [...] create a diagram explanation on their abstrate support sets. Eq 4 and Eq 5 [...]

Yes, the mapping function and the underlying Bernoulli probabilities are quite different. Thank you for the suggestion to include a diagram to illustrate the differences; we will include this in the next revision. 

> 4) [need more baselines methods, compare to original version of DisARM]

We used RLOO because it is a strong baseline that has been shown to be as least as good as the more complex methods such as RELAX/REBAR/STORB in the high-dimensional setting we are interested in (see [1] and [2]). For completeness, we are running additional baselines with RELAX and will add those results shortly.

We can include the results from the original version of DisARM/U2G. However, our concern was that the original version of DisARM/U2G was for Bernoulli latent variable models, so it is not applicable to the models considered in Tables 1 and 2. As a result, it may confuse readers. Do you feel that is not the case?

Please see our global comment above covering the limitations of our work.

References
1. Richter et al. "VarGrad: a low-variance gradient estimator for variational inference.", NeurIPS 2020
2. Kool at al. "Estimating Gradients for Discrete Random Variables by Sampling without Replacement", ICLR 2020


 Thank you for your detailed review.

> Does the estimators favorable properties over RLOO arise solely from avoiding “duplicate” samples in expectation or is there more behind it?

Thank you for this important question. There is more to it than avoiding "duplicate" samples because in high-dimensional spaces we will almost never see such duplicates. This is supported by the results from Kool et al. (2020) (Figures 2(b) and 4(b)) which show that in high dimension (similar to our setting) the performance and even the gradient variance of RLOO are *identical* to those of its extension that avoids duplicate samples. The coupling-based estimators generate sufficiently different, rather than merely non-identical, samples. We will clarify this point in the paper.

> Are there some settings where the gains over RLOO are expected to be larger? 

DisARM showed that with deeper Bernoulli hierarchical models the gap increased, so we might expect a similar trend to hold. Exploring this further could be an interesting future direction.

> The model was only evaluated on categorical VAEs.

Thank you for the suggestion. Time permitting, we will add experiments on structured prediction tasks in the revision. 

> The paper could also compare to other baselines […]

We used RLOO because it is a strong baseline that has been shown to be as least as good as the more complex methods such as RELAX/REBAR/STORB/UnOrd in the high-dimensional setting we evaluated (see [3] and [4]). We will add this reasoning to the revised manuscript.

> […] do hypothesis testing (t-test) to highlight the statistical significance of the improvements […]

With 500K training steps, as shown in the paper, we did the t-test for statistical significance of the improvement between various proposed estimators and the RLOO baseline.
 
t-statistics | DisARM-IW | DisARM-Tree |  DisARM-SB
-------------|-----------|-------------|------------
Dynamic-MNIST| 1.86014   |  2.74509    |   2.3831
FashionMNIST | 1.58487   |  1.86537    |   2.34371
Omniglot     | 0.508458  |  0.419834   |   0.863679

Training for longer (1 million steps), we find more significant results for Dynamic-MNIST and FashionMNIST

t-statistics| DisARM-IW | DisARM-Tree |DisARM-SB
---------|-----------|-------------|----------------
dynamic  |  3.59388  |   3.75532   |       3.63411
fashion  |  3.09949  |    2.8506   |       3.90916
omniglot | 0.386578  |  0.321789   |      0.545761
 
***
 
ELBO     | DisARM-IW | DisARM-Tree |DisARM-SB | RLOO
---------|-----------|-------------|-----------------|-----
Dynamic-MNIST  |-101.09&pm;0.12 |-100.67&pm;0.26 |-101.04&pm;0.14 |-101.84&pm;0.17
FashionMNIST  |-238.45&pm;0.17 |-238.38&pm;0.26 |-237.98&pm;0.25 |-239.54&pm;0.31
Omniglot |-118.34&pm;0.28 |-118.93&pm;0.26 |-118.23&pm;0.14 |-118.67&pm;0.79


> Rao-Blackwellization has also been used for improving other gradient estimators for discrete distributions…

Thank you for the references; we will include them in the revised manuscript.

References
1. Grathwohl et al. “Backpropagation through the void: Optimizing control variates for black-box gradient estimation.”, ICLR 2018.
2. Tucker et al. “The Mirage of Action-Dependent Baselines in Reinforcement Learning.”, ICML 2018.
3. Richter et al. "VarGrad: a low-variance gradient estimator for variational inference.", NeurIPS 2020.
4. Kool at al. "Estimating Gradients for Discrete Random Variables by Sampling without Replacement.", ICLR 2020.


Meta-review: Four knowledgeable reviewers recommend accepting this submission. I agree. This submission makes a valuable contribution by proposing new inference techniques for models with discrete latent variables.

Decision: accept