Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces

Abstract: Structured latent variables allow incorporating meaningful prior knowledge into deep learning models. However, learning with such variables remains challenging because of their discrete nature. Nowadays, the standard learning approach is to define a latent variable as a perturbed algorithm output and to use a differentiable surrogate for training. In general, the surrogate puts additional constraints on the model and inevitably leads to biased gradients. To alleviate these shortcomings, we extend the Gumbel-Max trick to define distributions over structured domains. We avoid the differentiable surrogates by leveraging the score function estimators for optimization. In particular, we highlight a family of recursive algorithms with a common feature we call stochastic invariant. The feature allows us to construct reliable gradient estimates and control variates without additional constraints on the model. In our experiments, we consider various structured latent variable models and achieve results competitive with relaxation-based counterparts.


Reviews and author responses:

 We thank the reviewers for the discussion. We appreciate the hard work giving thoughtful feedback, reading our detailed responses, and engaging in further debate. In the future revision, our highest endeavor will be to refine the unclear aspects pointed out by the reviewers, supplement the experiments' section with the mentioned tables and polish the overall narration.

This paper proposes a novel way to construct low-variance gradient estimators for differentiating through the output of recursive algorithms. By using random variables present in intermediate steps of the algorithm (called "the trace"), the authors propose a low-variance score function estimator formed using the gradient of the density of the trace. They explore this empirically in several settings and demonstrate improved performance.  Originality: To my knowledge, the idea is original.

Clarity and Quality: The main issue with this paper is its lack of clarity and polish. There are numerous things that were poorly explained or not explained at all. Here is a list:

* When introducing the algorithm, the definition of T is not given explicitly or implicitly. After reading the paper several times I have come to think that it is the 'key' of the item that is removed from the set at each recursive step, but that is not stated anywhere explicitly.
* Similarly, X is not explained well. 
* I think a better introduction to section 3 (lines ~102) would go a long way if you introduced each of the objects and the overall framework instead of diving in immediately to discuss the exponential variables.
* The need for partitioning the indices in algorithm 2 is not explained or really referenced elsewhere.
* f_split, f_combine, and f_stop are not defined or really explained in depth.

These issues and more combine to make it very difficult to understand specifically what is being proposed.

In addition to the clarity issues, there is an overall lack of polish. There are several spelling and grammatical errors, and the title is also a bit awkward. I would recommend "Latent structure learning through stochastic invariance" or "latent structure learning through a stochastic invariant", but the current phrasing doesn't sound right to my ear.

Finally, the experimental evaluation of these methods leaves a bit to be desired. Because the authors are proposing a gradient estimator, test-set performance is only part of the picture. To clearly make the case for the estimator I would like to see comparisons of the variance of these different gradient estimators over training, as well as train-set performance over time. The reason for seeing the training set performance is that better gradient estimators should result in better optimization, and test set results can be clouded with issues like overfitting.

Additionally, the current experimental results don't seem to strongly recommend the authors' proposed method.

If the authors substantially revise their draft to make it more clear and add the additional metrics I requested, I will update my score to reflect the new paper.

EDIT: After rebuttal I have raised my score. Yes.

 Thank you for your rebuttal. Your points make sense and I appreciate the additional results. I have updated my score, please do make the changes you describe.

 We would like to thank you for your helpful suggestions that will be beneficial to improve the presentation in the future revision of the paper. Indeed, the analysis of the introduced bias of the relaxation-based gradient estimators will give additional insights and further promote the proposed unbiased gradient estimators. We will try to do our utmost to include all the discussed points in the future revision of our manuscript.

 Thank you for the given suggestions. We will make sure to include the discussion of the limitations behind computing score function w.r.t. $X$ and clarify the restrictions in the future paper revision. In addition, we will do our best to incorporate the feedback to improve the overall presentation.

 Thank you for the feedback. It indicates the unclear parts of the draft, which we will improve upon in the revision.

*"As the authors mention, in [1] "In each section, they show that the additional prior knowledge improves the performance of the models." which I agree with, and it doesn't seem to happen here."*

Indeed, we do not introduce any additional prior knowledge compared to SST [1] inasmuch we use the same structured latent variables (spanning trees, rooted arborescence, etc.). Since SST already showed improvements compared to unstructured models, we considered only the corresponding models from [1] experiments and did not include unstructured models in the experiment.

*”… It would be more beneficial to better tune the experiments to showcase the benefit of the model.”*

We would like to emphasize that for each model we did hyperparameter search to tune the experiments (details are in Appendix E). The experimental section in the paper and the numbers given in response to reviewer ZnaB show that our gradient estimator results in better objective values thus is more effective as an optimization method. Constructing specific objective functions or regularization techniques consistent with metrics of interest is another broad area of research and was out of the scope of our work. 

[1] Paulus M. B. et al. Gradient Estimation with Stochastic Softmax Tricks //NeurIPS 2020. – 2020.

 I'll keep my score but highly suggest the authors incorporate these explanations (especially of the first two points) in the paper to improve clarity.

The authors propose score-function gradient estimators for structured latent variable models in which the structured variables are described as the output of a stochastic invariant perturbed algorithm. They derive different score-function estimators based on the algorithm's input $E$ and execution trace $T$ which can be computed even if the standard score function $\nabla_\lambda \log p_X(X=x; \lambda)$ is not available.
  **Originality:**
To the best of my understanding the presented work is novel and methodologically significantly differs from existing work. Most notably seems to be the work by Paulus et. al. [1], which combines the perturbation model framework and relaxation for gradient estimation. While the authors take inspiration from [1] they take a different approach in that they do not make use of relaxations but develop score-function gradient estimators for a class of 'stochastic invariant' perturbation models.

**Clarity:**
While the technical aspects of the paper are described in a clear manner I found it challenging to navigate and understand. After reading the first half of the paper I was under the impression that the authors develop lower variance score function estimators based on the Gumbel-Max trick as an alternative to continuous relaxations, but also as an alternative to standard score function estimators, which often have high variance.
Hence, after the authors explain in section 4 that the variance of the standard score function estimator is actually superior to the proposed variants I was a bit confused at first.
My current understanding is that the paper offers a way to compute score function estimators for 'stochastic invariant' discrete variable models even when the standard score-function estimator can not be computed, i.e. when $\nabla_\lambda \log p_X(X=x; \lambda)$ is not available/is intractible.
If this is indeed the case, I'd strongly encourage the authors to make this clear and motivate this early on in the paper.
To this end, it might also be a good idea to move the related work section right after the introduction. This way the presented work can be put into context and it's clear what the actual contributions over existing work are right from the start.

Another source of confusion was that the authors started using the acronym 'SST' in the experiment section and 'SST (Our Impl.)' in Table 2 without introducing it. Based on the label in the table I was wondering if there is a variant of the proposed estimator with the acronym SST that I somehow missed. After checking some of the related work, I discovered that [1] also uses the acronym 'SST', which stands for 'stochastic softmax tricks', and that it's indeed a baseline.

I also found the notation in Algorithm 1. confusing.
In the algorithm block $T = \arg \min E_j$ seems to denote an index, while in the writing and caption $T$ denotes an execution trace. I do not understand how these two notions would map onto each other. If they don't, using different letters might be more clear.


**Quality:**
To the best of my knowledge the paper is technically sound and the empirical results support the authors claim that the proposed method achieves results competitive with their relaxation-based counterparts.

**Significance:**
Many systems and generative processes incorporate or can be described in terms of discrete latent variables, e.g. discrete control flow in probabilistic programming models.
The presented work offers a technically interesting approach which allows to compute score-function estimators and corresponding baselines for 'stochastic invariant' perturbation models. Moreover, the presented results show that these estimators can perform on par with reparameterized relaxation-based estimators. As such, I do believe that this work is of interest to the broader probabilistic machine learning community.
My concern is that the current manuscript is challenging to read and understand and as a result readers might not pick up on the idea.


**Typos I noticed**
  - Line 149: function of *the* trace
  - Line 152: Lemma 2 -> Lemma 1?
  - Line 307: For *the* encoder
  - Line 321: *small*

*[1] Max B Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris J Maddison. Gradient estimation with stochastic softmax tricks. arXiv preprint arXiv:2006.08063, 2020*


 adequate

 Thank you for your detailed answers and for addressing my concerns regarding the clarity of the manuscript.

After reading the other reviews and responses, it seems to me that the rebuttal addresses most concerns. For my part, I'm happy with the rebuttal and still think the idea to use the exponential-min trick to derive score-function estimators is neat and novel. I will increase my score, recommending to ‘accept’ the paper, accordingly.

Here are some additional thoughts:

- As *Reviewer Csyp* noted, most readers, myself included, will be familiar with defining a generative model by defining some ancestral sampling process which allows for easy evaluation of the score function regardless of the size of it's domain. I think some additional discussion and a running example, such as the one you provided in your reply (generate the adjacency matrix of a tree), would be helpful for the reader to understand the difference in representation. It would also be great to have some motivation for when such a redundant representation might be needed or is of advantage. 

- While the variance of the score-function based estimators might not be directly comparable to the variance of the estimators which are based on the relaxed objective, there still seems to be a strong connection as the scale of the problem should not change significantly, no? Additionally quantifying the amount of bias might be interesting, i.e. comparing the gradient estimates to score-function gradient estimates computed with a very large amount of samples (treated as ground truth). The introduced bias is often used as an additional argument against using continuous relaxations so it would be interesting to see if the relaxed gradient estimates are indeed significantly biased.


The authors propose an approach to learn structured latent variables through score function estimators. This approach would add the possibility of incorporating prior knowledge on deep learning models. The results with respect to RELAX are better on a number of different applications. However, the baseline SST is shown to still outperform in particular in structure recovery.  The paper is very dense from a theoretical point of view and not always clear to non-familiar audience. While I'm not familiar with the literature and did not check every mathematical detail, I appreciate the effort of the authors in being complete as much as possible.

It is very hard to read the experiments section. I appreciate the effort of the authors to report each application they tested their method on, but the results for each are not properly explained and I had to resort to the supplementary material to understand them. Also I understand the baseline is the RELAX algorithm, while the other rows in table 2-3-4 are the variants of the proposed methods, is that correct? Further, the authors mention that "Overall, score function methods performed better than their relaxation-based counterpart, achieving higher values of ELBO on the test set, but slightly worse performance in terms of structure recovery metrics." Is there any reason why the structure recovery of the score function is worse? I thought that one benefit of the score function was to incorporate prior knowledge on the problem, which would more better benefit the structure recovery more than better likelihood.

I understand the approach allows to incorporate additional prior knowledge in the model. However, I'm not sure how this translates when applying the method, especially for each application. The algorithm requires an additional variable R, which is an auxiliary variable accumulating the information from the recursive calls. How this relates with the prior knowledge that the method allows to include in the inference algorithm?

 None

 I thank the authors for their time in replying and clarifying a couple of points.

I am still not totally convinced from the experiment design and the results. Based on the experiments I don't see it evident that adding the additional prior knowledge benefit the results (in particular the structure recovery). As the authors mention, in [1] "In each section, they show that the additional prior knowledge improves the performance of the models." which I agree with, and it doesn't seem to happen here. 

"We took the experiment design and the objective from [1]. In practice, a more suitable objective or an additional training signal in a semi-supervised setup may resolve the issue." I agree with this. It would be more beneficial to better tune the experiments to showcase the benefit of the model.

However I can appreciate the fact that it is a relevant and detailed work, so I'm willing to reassess my score based on the authors' reply and other reviews.

 We would like to thank the reviewers for their time spent providing us with a detailed evaluation of our work and fruitful suggestions. The general concerns regard the lack of clarity and polish of the manuscript. We will make our best effort to take into consideration all of the suggestions and improve the narration in our paper. Below we personally address reviewers’ individual concerns.

 *"If the authors substantially revise their draft to make it more clear and add the additional metrics I requested, I will update my score to reflect the new paper.*"

We provide the metric in our response. Unfortunately, the reviewing format does not allow us to change the draft. We thank you for the suggestions, which we will certainly incorporate in the paper revision. Instead of updating the draft, we will try to clarify some points in the response.

 *"When introducing the algorithm, the definition of T is not given explicitly or implicitly. After reading the paper several times I have come to think that it is the 'key' of the item that is removed from the set at each recursive step, but that is not stated anywhere explicitly. "*

The main idea was to use Algorithm 2 as a reference template. We will define the undefined terms as various parts of the template. We will clarify this in the revision.

For example, given an Algorithm 2 instance, the execution trace is the union of all variables $T_1, ..., T_m$ obtained at each recursion level.

In the case of top-k, the trace coincides with a list of items removed at each recursion level. In general, it may not be the case. For example, the Edmonds algorithm may not remove indices $T_1, ..., T_m$. It finds a loop among the edges $T_1, ..., T_m$, removes all edges between the loop vertices, and contracts the loop into a single vertex. 

*" Similarly, X is not explained well."*

Again, we will define $X$ as an output of an instance of Algorithm 2.

*"The need for partitioning the indices in algorithm 2 is not explained or really referenced elsewhere."*

Multiple partition sets allow representing Edmonds algorithm as an instance of Algorithm 2. The binary tree variable also uses multiple partition sets. We will discuss this in more detail.

*"f_split, f_combine, and f_stop are not defined or really explained in depth."*

Again, these functions are the elements of Algorithm 2. We try to explain the intuition in the paragraph starting on line 135 and provide multiple examples in section F of the supplementary materials. We will summarise the role of these functions:

* Having a stop condition computed by **f_stop** is necessary for any recursive function;
* We restrict Algorithm 2 to a single recursive call. Function **f_combine** represents all the computations taking place after the recursive call;
* For the sake of generality, we allow multiple simultaneous Exponential-Min tricks. Function **f_split** defines sets across which we find the minimum elements. The function plays a role in the Edmonds algorithm. In particular, the algorithm finds minimum across each column of the adjacency matrix simultaneously.

The signature of these functions (the variables they take in and return) restricts the algorithm and allows us to make claims about the variables $E, T$, and $X$. Each instance of Algorithm 2 may define these functions in its way.

*"Finally, the experimental evaluation of these methods leaves a bit to be desired. Because the authors are proposing a gradient estimator, test-set performance is only part of the picture. To clearly make the case for the estimator I would like to see comparisons of the variance of these different gradient estimators over training, as well as train-set performance over time. The reason for seeing the training set performance is that better gradient estimators should result in better optimization, and test set results can be clouded with issues like overfitting."*

Below we provide the objective values on the training set. After comparing against the respective metrics on the test set from the paper, we conclude that there is no overfitting on Graph Layout and ListOps datasets. On the L2X dataset, there is slight overfitting and, indeed, the unbiased T-REINFORCE+ estimator performs better in terms of objective (structure recovery metrics are not available because there are no labels for train data).

Train results on L2X Aroma aspect (mean MSEs with STDs)

| | | k=5 | k=10 | k=15 |
|---|:---:|:---:|:---:|:---:|
| Simple | SST | 3.22 +- 0.17 | 2.96 +- 0.14 | 2.54 +- 0.13 |
| | E-REINFORCE+ | 3.45 +- 0.19 | 3.38 +- 0.23 | 2.85 +- 0.21 |
| | T-REINFORCE+ | 3.38 +- 0.13 | 2.93 +- 0.19 | 2.43 +- 0.12 |
| | RELAX | 3.36 +- 0.12 | 3.13 +- 0.18 | 2.78 +- 0.17 |
| Complex | SST | 2.69 +- 0.14 | 2.24 +- 0.18 | 2.18 +- 0.11 |
| | E-REINFORCE+ | 2.85 +- 0.22 | 2.56 +- 0.25 | 2.44 +- 0.19 |
| | T-REINFORCE+ | 2.5 +- 0.16 | 2.19 +- 0.15 | 2.09 +- 0.1 |
| | RELAX | 2.53 +- 0.15 | 2.21 +- 0.17 | 2.19 +- 0.14 |

Train results on Graph Layout

| Estimator | ELBO | | Precision | | RELAX | |
|:---:|:---:|---|:---:|---|:---:|---|
| | MEAN +- STD | MAX | MEAN +- STD | MAX | MEAN +- STD | MAX |
| SST | -1846.93 +- 1124.23 | -1357.03 | 88 +- 22 | 96 | 94 +- 4 | 96 |
| T-REINFORCE+ | -1584.11 +- 572.16 | -1193.11 | 70 +- 31 | 91 | 86 +- 8 | 91 |
| RELAX | -2086.93 +- 573.72 | -1207.83 | 43 +- 31 | 90 | 81 +- 6 | 90 |

Train results on ListOps

| Estimator | Accuracy | | Precision | | Recall | |
|:---:|:---:|---|:---:|---|:---:|---|
| | MEAN +- STD | MAX | MEAN +- STD | MAX | MEAN +- STD | MAX |
| SST | 79.31 +- 8.17 | 94.73 | 57.15 +- 19.92 | 82.51 | 30.58 +- 19.03 | 73.28 |
| E-REINFORCE+ | 60.64 +- 2.51 | 65.21 | 41.12 +- 6.68 | 45.95 | 40.99 +- 6.75 | 45.68 |
| T-REINFORCE+ | 88.69 +- 3.02 | 93.06 | 78.6 +- 7.37 | 80.68 | 61.78 +- 14.52 | 80.68 |
| RELAX | 79.92 +- 9.35 | 88.64 | 54.84 +- 17.51 | 75.27 | 53.73 +- 17.18 | 75.27 |

We cannot provide a graph comparing variances. Instead, we report the mean gradient std at three phases of the training loop (beginning, middle, ending). We report the average std across the train dataset and parameter dimension. For each training item, we take 10 samples to compute std (increasing the number does not change the results). As expected, SST has a lower variance. RELAX has a higher variance at the beginning due to the underfitting of the control variate. It is also worth mentioning that SST and estimators based on the score function are not directly comparable because SST computes the gradient of the relaxed objective.

| Mean Grad Std | beginning | 25k | 50k |
|:---:|:---:|:---:|:---:|
| SST | 0.1674 | 0.0343 | 0.0379 |
| T-REINFORCE+ | 1.6320 | 1.4409 | 0.9924 |
| RELAX | 3.0592 | 1.0292 | 0.8874 |

*"Additionally, the current experimental results don't seem to strongly recommend the authors' proposed method."*

While we agree that the proposed gradient estimators do not significantly outperform continuous relaxation baseline (SST), this was not the main goal that we pursued. We emphasize that while the proposed gradient estimators perform similarly on problems that admit continuous relaxations, they have wider applicability as was shown in the non-monotonic generation experiments. There were many attempts to tackle the problem of gradient estimation for structured discrete random variables as was discussed in the related work. In our work, we develop a general learning approach for structured discrete random variables including but not limited to subset selection, binary trees, spanning trees, rooted arborescence, and permutations.


 *"My current understanding is that the paper offers a way to compute score function estimators for 'stochastic invariant' discrete variable models even when the standard score-function estimator can not be computed, i.e. when $\nabla \log p_X(X = x ; \lambda)$ is not available/is intractable. ... If this is indeed the case, I'd strongly encourage the authors to make this clear and motivate this early on in the paper. To this end, it might also be a good idea to move the related work section right after the introduction. This way the presented work can be put into context and it's clear what the actual contributions over existing work are right from the start."*

Indeed, this is the case. We consider a specific distribution family, for which the standard score function estimator is intractable. Our goal was to define a tractable and efficient score-based gradient estimator. We achieve the goal by changing the action space from the variables $X$ to execution traces $T$. After defining a score function for $T$, we consider the existing methods for variance reduction: LOO-baseline and sample-dependent baselines (RELAX). The score function based on $T$ is a more practical alternative to the score function in the continuous action space $E$. The latter is always tractable, but it has higher variance and significantly degrades quality in practice (as Tables 2, 3, 4 indicate).

We will reorganize the main section to highlight the main idea and to outline the scope of applicability of the methods.

*"Another source of confusion was that the authors started using the acronym 'SST' in the experiment section and 'SST (Our Impl.)' in Table 2 without introducing it."*

Thank you, we will define the abbreviation in the main text before the experimental section.

*"I also found the notation in Algorithm 1. confusing. In the algorithm block $T = \arg\min E_j$ seems to denote an index, while in the writing and caption $T$ denotes an execution trace. I do not understand how these two notions would map onto each other. If they don't, using different letters might be more clear."*

In all algorithm instances, we define the execution trace as the sequence of argminima obtained at all the recursion steps (lines 119-120). Algorithm 1 finds the global argminimum at each recursion step. The execution trace for Algorithm 1 is then the sequence of $k$ indices $(T^1, ..., T^k)$, corresponding to the $k$ minimal elements of the input array, sorted in ascending order. We will reorganize the notation in this algorithm to fix the overlap and make the correspondence more clear.

 *”is [the computation of the score function] difficult in most cases?”*

While ancestral sampling is a viable approach, it has shortcomings.

For example, we can generate an adjacency matrix of a tree by generating binary components of the adjacency matrix step by step (with a total of $n^2$ steps). Adding the next edge will require to ensure the global constraints (i.e., the current graph is a forest and can contain $n - 1$ edges in the end). In general, verifying whether a partial sample has a feasible extension may also be computationally demanding.

An alternative and arguably more natural way to construct a tree is to add $n - 1$ one by one (only need to verify whether the graph is a forest, the number of edges is always correct). This is an ancestral sampling algorithm for a different variable representation. However, unlike the previous algorithm and similarly to our approach, the score function includes redundant information, namely the order in which we add the edges.

*”The family of generative processes this is applicable to seems restricted on the surface: I need to use exponentials, and then I need to make sure I do an argmin, and then to subtract the minimum exponential from the others, and then have a recursive call.”*

The generative process is indeed restricted. In the above listing, there is only one essential restriction. The others restrictions do not lead to loss of generality.

Continuing the previous answer, we argue that our approach is an ancestral sampling instance for a redundant representation of a random variable (the execution trace). Some aspects do not bring limitations. The recursion defines the factorization. In particular, with each recursion level, we proceed to the next set of conditional probabilities. At each recursion level, we sample some categorical variables. We parametrize the variables using the exponential-min trick and end up with exponential distributions and arg-minima.

The main limitation is the relation between the input $E$ and the recursive call input $E’$. The choice $E_j’ = E_j - E_{T_i}$ (Fig. 1, Alg. 2) leads to a very specific conditional distribution on the next level. To represent an arbitrary conditional distribution it would suffice to define $E’$ as a new vector of exponential variables with some new parameters conditioned on $T_1, …, T_m$.

*”As a side note, I found it cool that a simple leave-one-out estimator is better than a complex action-dependent RELAX-based estimator, is there any intuition behind this?”*

The RELAX estimator optimizes the control variate simultaneously with the main objective. However, the optimal control variate is different for each set of objective parameters. Our intuition is that the control variate optimization does not keep up with the objective optimization. At the same time, the sample mean of the objective which is used as the control variate in the leave-one-out estimator is always an unbiased expectation estimate for the current set of parameters. You also can refer to recent work [1], which provides further insights behind competetive performance of the LOO estimator.

*”I notice that there is only one baseline algorithm (the SST). Are they the only one applicable? Or are they already the best of all the ones applicable? Would wake-sleep methods like reweighted wake-sleep be applicable?”*

The majority of current relaxation-based estimators concentrate on permutations and subset selection. SST has better performance, so we did not include the other baselines. We are not aware of relaxation-based baselines for the spanning trees and the arborescences. Other score-based methods are applicable.


[1] Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco J. R. Ruiz, Ömer Deniz Akyildiz. VarGrad: A Low-Variance Gradient Estimator for Variational Inference. NeurIPS 2020.

 *“Also I understand the baseline is the RELAX algorithm, while the other rows in table 2-3-4 are the variants of the proposed methods, is that correct?“*

The baseline algorithm is SST [1]. It is a relaxation-based gradient estimation method for structured latent variables. In our work, we consider a closely related probabilistic distribution family. Both works consider a deterministic algorithm mapping continuous random variables $E$ to a discrete random variable $X$. Unlike SST, we approach learning with score function-based gradient estimators. We start with a trivial score function estimator (E-REINFORCE+). Then we introduce an intermediate variable T to improve the trivial estimator (T-REINFORCE+) and derive the conditional distribution $E | T$ used to define the RELAX [2] control variate in our scenario.

*””Overall, score function methods performed better than their relaxation-based counterpart, achieving higher values of ELBO on the test set, but slightly worse performance in terms of structure recovery metrics." Is there any reason why the structure recovery of the score function is worse?“*

We speculate that objective and auxiliary metrics may be inconsistent. A near-optimal objective may not necessarily lead to near-optimal structure recovery. In particular, the quoted caption notes that better test objective values (along with better train objective values, as the answer to reviewer ZnaB shows) do not lead to better structure metrics. We took the experiment design and the objective from [1]. In practice, a more suitable objective or an additional training signal in a semi-supervised setup may resolve the issue.

*“I thought that one benefit of the score function was to incorporate prior knowledge on the problem, which would more better benefit the structure recovery more than better likelihood.”*
 
The ability to incorporate prior knowledge through the generation process is a prominent feature of latent variable models. Score function estimators constitute one of the approaches to fitting latent variable models.

Our work mainly focuses on a specific latent variable family and comparison with a relaxation-based fitting alternative. Still, the results in Figure 2 indicate that prior knowledge allows improving non-monotonic generation. The semi-supervised “ss” model generates more correct sequences than the “l2r” model using left-to-right generation order. In our experience, weak supervision allows steering the training towards the desired optimum. A similar trick is used in [3]. Unlike the “smart” model, the semi-supervised model does not require specifying generation order for each training sample. 

See the experimental section in [1] for more illustrations on incorporating prior knowledge. In each section, they show that the additional prior knowledge improves the performance of the models.

*”The algorithm requires an additional variable R, which is an auxiliary variable accumulating the information from the recursive calls. How this relates with the prior knowledge that the method allows to include in the inference algorithm?”*

Again, we incorporate the prior through the design of a latent variable model. We use the letter $X$ to denote the latent variable (we do not introduce a distinguished letter for observed variables to avoid overcrowding). By specifying the details of Algorithm 2, we define the variable $X$.

Besides using the word variable for the elements of a latent variable model, we use the word "variable" to refer to the algorithm inputs (arguments). We introduce the auxiliary argument $R$ in Algorithm 2 for technical purposes. $R$ is short for “reduction” whereas the argument $K$ is short for “keys.” As an illustration, in the top-k algorithm, $R$ plays the role of $k$: to compute top-k with $k = R$ we find the top element, exclude it, and then reduce the problem to top-k with $R’=k-1$. In different algorithms instances, $R$ plays different technical roles, on which we elaborate in the supplementary materials.

[1] Paulus M. B. et al. Gradient Estimation with Stochastic Softmax Tricks //NeurIPS 2020. – 2020.

[2] Grathwohl W. et al. Backpropagation through the Void: Optimizing control variates for black-box gradient estimation //International Conference on Learning Representations. – 2018.
    
[3] Corro C., Titov I. Differentiable Perturb-and-Parse: Semi-Supervised Parsing with a Structured Variational Autoencoder //International Conference on Learning Representations. – 2018.


This paper cleverly uses the Exponential-min trick in Lemma 1 to derive a score-function based gradient estimator for learning latent variable models. It is applicable to a special class of latent variable generation (in settings where we learn a generative model-recognition model pair, the proposed method puts constraints on the recognition model) described in Algorithm 2, which successively transforms a sequence of Exponential random variables to a set of "trace" variables which are used to finally obtain the desired latent variables. This function is recursive and the "stochastic invariant" refers to the fact that each recursive call gets as arguments stochastic variables that always have an Exponential distribution. The proposed method is supposed to be useful especially when we cannot directly evaluate the score function of the latent variables (which would have lower variance) -- this occurs for example, when it is produced by an implicit generative process. There are 4 experiments comparing the different versions of the proposed gradient estimator with the recently proposed "stochastic softmax tricks" gradient estimator.  **Originality**: Are the tasks or methods new? Is the work a novel combination of well-known techniques? (This can be valuable!) Is it clear how this work differs from previous contributions? Is related work adequately cited?

The method is new. The work is a novel combination of the Exponential-min trick to a score function estimator. I'm not an expert in perturbation methods, but as far as I understand, this work differs from previous perturbation methods applied to latent variable model learning in that it's used to improve a score function based estimator. And yes, the related work provides enough pointers to the most important related work.

**Quality**: Is the submission technically sound? Are claims well supported (e.g., by theoretical analysis or experimental results)? Are the methods used appropriate? Is this a complete piece of work or work in progress? Are the authors careful and honest about evaluating both the strengths and weaknesses of their work?

The method is technically sound and well explained. I particularly appreciated the min-k example in Algorithm 1.

One thing that isn't completely clear to me is the family of latent variable models to which this method is applicable.
1) In lines 38-40 and in the paragraph in lines 56-62, it is explained why the computation of the score function is difficult -- is it really that difficult in most cases? If I were to write a complex generative model with discrete latent variables, I would define how those discrete latent variables are ancestrally sampled and so even if the space of all possibilities is combinatorial, I'd end up knowing how to evaluate the probability of the sampled variables alongside their score function.
2) The family of generative processes this is applicable to seems restricted on the surface:  I need to use exponentials, and then I need to make sure I do an argmin, and then to subtract the minimum exponential from the others, and then have a recursive call.

Together a discussion of the generality of the method would be really helpful!

**Clarity**: Is the submission clearly written? Is it well organized? (If not, please make constructive suggestions for improving its clarity.) Does it adequately inform the reader? (Note that a superbly written paper provides enough information for an expert reader to reproduce its results.)

The paper is clearly written (except the point above) and well organized.

**Significance**: Are the results important? Are others (researchers or practitioners) likely to use the ideas or build on them? Does the submission address a difficult task in a better way than previous work? Does it advance the state of the art in a demonstrable way? Does it provide unique data, unique conclusions about existing data, or a unique theoretical or experimental approach?

The results are important as it pushes the limits of learning discrete latent variable models. These are potentially important for more interpretable, explainable and generalizable systems. Yes, other researchers can definitely build on these ideas. It performs better than the recent "stochastic softmax tricks" method in most cases (inclusion of "failure" cases is appreciated as it paints a more realistic expectation of how this algorithm performs). As a side note, I found it cool that a simple leave-one-out estimator is better than a complex action-dependent RELAX-based estimator, is there any intuition behind this?

I notice that there is only one baseline algorithm (the SST). Are they the only one applicable? Or are they already the best of all the ones applicable? Would wake-sleep methods like reweighted wake-sleep be applicable? Yes, the authors adequately addressed the limitations and potential negative societal impact of their work.

Meta-review: This is an interesting paper that introduces a method for taking gradients in models containing discrete structured latent variables.  The approach is based on using the exponential-min trick and a score function estimator that is carefully applied using intermediary "trace" variables in the generation process, allowing the intractability of the score function estimator for the output variables to be overcome, while providing lower variance than applying the score trick to the "raw" random variables.

Overall the reviewers were generally positive about the paper, praising its novelty, quality, and potential significance, and all backing its acceptance.  The main question marks were over the clarity of the paper, which indeed could certainly do with some non-negligable work, but I feel that these issues are not so fundamental that they cannot be addressed in time for the camera-ready version of the paper, a viewpoint which the reviewers seem to agree with.  As such, I recommend its acceptance to the conference, but strongly encourage the authors to work on refining the paper as per the reviewers' feedback.

Some comments of my own from reading the paper that may assist further with making improvements:
- Though not discussed in the response, Reviewer ZnaB is correct that “stochastic invariant” is grammatically incorrect and I agree that it should be “stochastic invariance”.  The term does not seem to be used entirely consistently as well, as it is mostly discussed as a property, but then Eq 8 seems to suggest another meaning (though maybe this is a typo?).
- It would be good to use a formal definition for “stochastic invariant” on lines 131-132 as it is easily missed in passing that this is where it actually being defined.  This definition possibly also needs to come earlier in the paper as well as the term is used a few times before its defined.
- Though the low-level writing of the paper is generally fine, there is definitely quite low of room for improvement on more high-level aspects.  In particular, the writing could be improved by making clear why things are being introduced before they are (in the context of the overall goals of the paper): at the moment many things are just explained without introduction and the reason we care about them (or indeed exactly what they are) only comes later.  This makes it difficult to follow as the reader is being fed technical information before they appreciate the context.  An example of this is Algorithm 2, which is discussed at length without a clear statement of what it actually is, while how both Algorithms 2 and 3 fit into the high-level aims of the paper is never stated explicitly.  More generally, the writing regularly falls into the trap of explaining things as if the reader already understands what is going on and knows what is coming later in the paper; it misses a lot of the “glue” that keeps the reader on track about where the paper is going, how things fit together, and how the low-level ideas link into the high-level aims.  It also often leaves too much implicit (e.g. what the baselines are and what the proposed models are in the experiments — note I would consider E-REINFORCE to be a baseline here, which should be made clear), rather than directly and explicitly stating what things are and what they are used for.  Some repetition of key points would also help, as there are a few key things that are only briefly stated once, and if you miss them the paper quickly becomes difficult to understand.
- Related to the above, and as noted by the reviewers, the paper could also do a better job of making clear what its contributions are in the introduction and why these are useful.  The motivation at the moment is a bit back and forth, and the contributions need stating explicitly and directly (as per the discussions with Reviewer EF5a), rather than in the implicit manner they are currently conveyed.
- One potential application I envisage of the work that is not mentioned in the paper is probabilistic programming, as here one often has complex branching behaviour that prohibits relaxation-based approaches. It might be good to add some mention of this if space can be found, while the concept of a “trace” the paper used is commonplace in the probabilistic programming literature and may be worthy of mention in the related work.
- I think it would strengthen the paper to include more baselines.  Even if SST is considered the “best”, this would still provide more evidence and make the competitiveness of the approach clear.  More discussion of the tested approaches is also needed in general and some repetition of exactly what everything is (e.g. what RELAX is can be difficult to immediately see, and the fact a previous approach has this name can be misleading), without which it can be surprisingly difficult to figure out what everything is.
- A possible missing reference you may wish to add: “Reparameterization Gradient for Non-differentiable Models” Lee et al. NeurIPS 2018.

Decision: accept