Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces
Authors: Kirill Struminsky
Artyom Gadetsky
Denis Rakitin
Danil Karpushkin
Dmitry P. Vetrov
Venue: NeurIPS 2021 Poster (2021)
Keywords: structured discrete latent variable
gumbel-max trick
perturb-and-map
gradient estimation
score function

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.

Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces
Kirill Struminsky'
Artyom Gadetsky"1
Denis Rakitin'
HSE University Moscow, Russia
HSE University
HSE University, Skoltech! Moscow, Russia
Moscow, Russia
k. sttuminskyygmail.comm
artygadetaky@@yanex.u
rakitiindemis32@maail..
Danil Karpushkin
Dmitry Vetrov
AIRI,
Sber AI Lab, MIPT Moscow, Russia
HSE University, AIRI
Moscow, Russia vetrovd@yandex.oo
kardani @atani@@mall.. ru
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 define latent variable as a perturbed algorithm output and to use a differentiable surrogate for training. In general, the surrogate puts additional constraints the model and inevitably leads biased gradients. To alleviate these shortcomings, 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 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 the model. In our experiments, we consider various structured latent variable models and achieve results competitive with relaxation-based counterparts.
Introduction
To this day, the majority of deep learning architectures consists of differentiable computation blocks and relies on gradient estimates for learning. At the same time, architectures with discrete intermediate components are a good fit for incorporating inductive biases [ or dynamic control flow 26 [2] One of the approaches to train such architectures is to replace the discrete component with stochastic latent variable and optimize the expected objective.
In practice, the expectation has high computational cost, thus one typically resorts to stochastic estimates for the expectation and its gradient. Particularly, the two prevalent approaches estimate the gradient of the objective are the score function estimator [45] and the reparameterization trick
[6 40) for relaxed discrete variables 29 (5). The former puts mild assumptions on the distribution and the objective, requiring the gradient of log probabbbbyyl with respect to the distribution parameters be differentiable, and provides unbiased estimates for the objective gradient. However, the naive
Equual contribution
Corresponding author
Skolkovo Institute of Science and Technology
8 Artificial Intelligence Research Institute
Moscow Institute of Physics and Technology
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

estimate suffers from high variance and is less intuitive in implementation. In comparison, the reparameterized gradient estimates seamlessly integrate within the backpropagation algorithm and exhibit low variance out of the box. At the same time, the relaxation requires an architecture to be defined on the extended domain of the relaxed variable and introduces bias to the gradient estimate.
In the recent years, the attention of the community shifted towards models with structured latent variables. Informally, a structured variable models distribution over structured objects such as graphs 34 sequences Y or matchings [3. Such latent variable may alter the computation graph or represent a generative process of data. Often, a structured variable is represented as a sequence of categorical random variables with joint distribution incorporating the structure constraints (e.g., the fixed number of edges in an adjacency matrix of a tree). Recent works on structured latent variables address model training largely through the reparameterization trick using relaxed variables. In fact, the Gumbel-Softmal trick naturally translates to structured variables when arg max operator is applied over structured domain rather than component-we component-wise [4]. In contrast, score function estimators are now less common in structured domain, with few exceptions such as 50 [4]. The primary difficulty is the sample score function: neither Gibbs distributions, nor distribution defined through a generative process have general shortcut to compute it.
In our work, we develop a framework to define structured variables along with a low-variance score function estimator. Our goal is to allow training models that do not admit relaxed variables and to improve optimization by alleviating the bias of the relaxed estimators. To achieve the goal we define the structured variable as an output of an algorithm with a perturbed input. Then, we outline family of algorithms with a common property we call stochastic invariant. The property was inspired by the observation 34 Appendix, Sec. B], where the authors showed that the Kruskal's ? are recursively applying the Gumbel-Max trick. We construct new algorithms with the same property and show how to use the property to learn structured latent variables. In the experimental section, we report performance par with relaxation-based methods and apply the framework in a setup that does not allow relaxations.
algorithm (24] and the CLE algorithm
2
The Recursive Gumbel-Max Trick in Algorithms With Stochastic Invariants
The section below shows how to define a distribution over structured domain. Conceptually, we define structured random variable as output of an algorithm with random input (e.g., to generate a random tree we return the minimum spanning tree of graph with random weights). A common solution to incorporate such variable in a latent variable model is to replace the original algorithm with a differentiable approximation to allow gradient- the difficutly of computing the sample probability. In contrast, we outline family of algorithms for which we can get the probability of each intermediate computation step. To get the probabilities restrict our attention to algorithms with specific recursive structure and random inputs with exponential distribution. In the next section, we leverage the probabilities for approximate inference in latent variable models without the differentiable approximations of the algorithm.
learning 4 3T Such solution bypasses
2.1 The Gumbel-Max Trick in arg top k
We illustrate our framework with recursive algorithm generating a subset of a fixed size. The lemma below is well-known result used to generate categorical random variables using a sequence of exponential random variables.
Lemma 1. (the Exponential-Mis trick) If E, Exp (As), i (1....,d) d) are independent, then for X = argmin, E,
1. the outcome probability is Px = I; A) x Azi
2. random variables E E, - Ex, E fl......d d} are mutually independent given X with E X Exp(X.) when # X and E otherwise
The lemma is equivalent to The Gumbel-Max trick, defined for the variables G; := - log F4,, E {I, [L...., df and the maximum position argmax, G.. In the above case, a variable G. has Gumbel
6 "s convention, assume that 0 d Exp (oo).

---

distribution with the location parameter 0; log A₁ hence the name. Though the Gumbel-Mar trick formulation is more common in the literature. we formulate the framework in terms of the exponential distribution and the equivalent Exponential-Min trick. Although the two statements are equivalent and we use their names interchangeably, some of the examples have a natural formulation in terms of the exponential distribution.
Importantly, the second claim in Lemma Tallows applying the Exponential-Min trick succesively. We illustrate this idea with an algorithm for finding top-k elements. We present the recursive form of arg top k in Algorithm [ For each recursion level, the algorithm finds the minimum element, decrements k and calls itself to find the subset excluding the minimum variable. For reasons explained below, the algorithm subtracts the minimum from the sequence Ej E 5 Er before the recursion. This step does not change the output and may seem redundant.
Assuming the input of the algorithm is vector F. of independent exponential variables with rate parameters A, the first argument in the recursive call E' is again a vector of independent exponential variables (given 7) due to Lemma In other words, the input distribution class is invariant throughout the recursion. Subtraction of the minimum is not necessary, but it allows to apply Lemma directly and simplifies the analysis of the algorithms. Besides that, for each recursion level variable T has categorical distribution (conditioned on T found in the above calls) with output probabilities proportional Ag. k E K.
We use upper indices to denote the recursion depth and, with a slight abuse of notation, denote the concatenation of variables T for each recursion depth as T = (T) (T',.... ".). The output X is set and does not take into account the order in T. Intuitively, T acts as the execution trace of the algorithm, whereas X contains only partial information about T. The marginal probability of x: is the sum xx = I; X) = Exex (()) PT = f; X) over all possible orderings of x = {=1,.....0 xx3 denoted as X (x)). The direct computation of such sum is prohibitive even for moderate k.
The arg top illustration is a well-known extension of the Exponential-Min trick. In particular, the distribution of T' is known as the Plackett-Luce distribution 3ri and coincides with & categor- ical samples without replacement. Following the recursion, the observation probability factorizes according to the chain rule with i-th factor governed by equation Pr, (T; = t; (4-...,4-11i
X Σ5 - Ary"
We discuss the multiple applications of the trick in Section Next, we extend
Algorithm beyond subset selection.
2.2 General Algorithm With the Stochastic Invariant
In this section, we generalize Algorithm D The idea is to preserve the property of Algorithm Ihhat allows applying the Exponential-Min trick and abstract away the details to allow various instantiations of the algorithm. Algorithm is the generalization we put next to Algorithm for comparison. It has a similar recursive structure and abstracts away the details using the auxiliary subrouties: faup is the stop condition, fmup and fcombine handle the recursion and fspli is optional subroutine for the Exponential-Mo trick. Additionally, we replace & with an auxiliary argument R used to accumulate information from the above recursion calls. Below, we motivate the subroutines and discuss the properties of a arbitrary instance of Algorithm
After checking the stop condition with Ssup, Algorithm 2 applies the Exponential trick simulta- neously over m disjoint sets rather than the whole index set K. For example, such operation occurs when we find columnwise minimum in CLE algorithm 7. To allow the operation construct a partition of indices PI...., Pm and find the arg min across the partition sets. To generate the partition, we introduce new subroutine faplit taking the index set K and the auxiliary argument R as inputs. The partition size m may also be variable.
After the m simultaneous Exponential-Mir tricks, the generalized algorithm calls fmap to select a subset of indices K' Ç K and to accumulate the necessary information for the next call R'. Intuitively, the argument R' represents reduction to a smaller problem solved with recursive call. In the arg top example, K' K \ (T) and R' is the decrement not allow to capture such information with the other inputs E' and K' exclusively.
1. Note that Algorithm 2 does
3

---

Algorithm 1 (E, K, k) finds k smallest elements in a sequence E, where K the set of indices (keys) of E
Algorithm 2 Fistuct [, K, R) returns structured
variable X based on utilities E indexed by К and an auxiliary variable R
Input: E. K,K.k
Input: E, K,K,R R
Output: X if k = 0 then
Output: X
if Sstop K, R) then
return
return
end if
end if
(Find the smallest element)
Pi...., Pm - [spti R)
[ = K)
T < arg minzek Fg
for i to m do
for jorje K do
T. = arg minjeP. Ej for P. do
E E, ET
end for (Exclude arg min index T, decrement k} K' = K (T)k-1 E' { K*K
E. + E2 ET.
end for end for K" R'
fmap E' < E k E K"}
K. R. {Ti)
IK' Ç K)
(Solve for = k-1
X" Pupx E', K',k) return UT)
X +
Pamal struet E' K' R')
(Recursive call)
return fcombine (X', K, R, (T:)
Figure 1: The recursive algorithm for arg top lt and the general algorithm with the stochastic invariant put side-by-side. Both algorithm perform the Exponential-Min trick and proceed with recursion using a subset of variables. The output X combines the current trace ''' and the recursion output X'.
Finally, the algorithm calls Scombine to construct the structured variable using the recursive call output X' and the other variables. In the top-k example, foombine appends the minimum variable index I' to the set X'.
Now argue that Algorithm 2 preserves the invariant observed in Algorithm I Again, call the sequence of variables T = Ti..... Tm)) the trace of the algorithm. By design, if the input E is sequence of independent exponential random variables, then the recursion input E' conditioned on T is again a sequence of independent exponential distributions. For short, we call this property the stochastic invariant. The key to the stochastic invariant is the signature of the subroutines Algorithm| uses. The algorithm only accesses E values though the Exponential-Min trick. As result, the intermediate variables K' and R' as well as the output X depend E only through T. In other words, the execution trace is function of perturbation T T(F)) and the structured variable X X (2") is a function the trace. Additionally, due to Lemma [ the trace components Ti...., T'm have categorical distributions, whereas E%, k E K are exponential random variables. We prove these properties by induction w.r.t. the recursion depth in Appendix A.
2
Given the above, we derive two modifications of Algorithm 2 generalizing Lemma D and the Plackett- Luce distribution from the illustration. Algorithm 3 computes the -og-probability log Pr (t; X) of trace realization t. In Section g we use the algorithm output to construct gradient estimators. Again, the pseudo-code introduces index j to denote the recursion depth and assumes the input / = (t. ) the concatenation of trace variables for all recursion depths j 1,..., k. Similarly, in Appendix B we present an algorithm returning sample from E T given trace realization t.
2.3 Further Examples
This subsection contains an overview of algorithms with stochastic invariants along with the corre- sponding structured variables. We present the details and the pscudo-code in Appendix E.
Analogous the arg top & and the subset variable, the insertion sorting algorithm is an algorithm with the stochastic invariant. In the case of sorting, we do not omit the order of the trace variable and return the permutation X T. The resulting variable X has the Plackett-Lucs distribution. We use the variable as latent variable for iieertion-based non-monotonic generation [13] As an alternative to the Plackett- distribution, we consider a square parameter matrix and find matching between rows and columns. We perturb the matrix and iteratively find the minimum element

---

Algorithm 3 Flog.pnid prob (1, X, K, R)- returns log (1; X) for trace I, rates X, K and R as in Alg.
Input: 1, 2, K, R
Output: log (r( if' Ssopl (K, R) then
return end if
Pi Pm = Espli (, R)
for i 1 to m do
log Pgg (t); = log do - loo (EkEp. Ax)
[Index j in TV denotes the recursion level)
for P. \ {t] do
X/2 + A
end for X2 = to0 end for K', + Smup K, R, [m))m1)))
[Because Et(t)) = 0)
X x
(X/
k E K'}
log Pr (t T = f²; A) + Fig poub (t21, R')
[Compute log-prob of of == {t)))a))
return Ex= log (t;; + log IP (( (1>1
1 = A}
in the matrix. We exclude the row and the column containing the element and proceed to the next recursion step. Notably, in contrast to this algorithm, the Hungarian algorithm 25 for the minimum weight matching does not have the stochastic invariant.
As observe, Kruskal's algorithm 24 and Chu-LiuuFFdmonds 7 algorithm recursively apply the Exponential-Min trick, thus have the stochastic invariant. The former constructs the minimum spanning (MST) tree edge-by-edge. The corresponding trace variable T is a sequence of edges, whereas X is an unordered set of edges. Interestingly, we could not represent Prim's 's algorithm 38 for the MST as an algorithm with the stochastic invariant. The analog of Kruskal's s algorithm for directed graphs. It returns the minimum tree with a fixed root.
algorithm is an
Additionally, we construct latent binary tree variable specifically for the non non-monotonic generation model 44 In this model, each token in the sentence corresponds node in the tree. We assign weights to nodes and perturb the weights to obtain a sample. During the recursion, we set the minimum weight node to be the parent. We put the nodes the left-hand side to the left subtree and the nodes on the right-hand side to the right subtree.
3 Gradient Estimation for the Recursive Gumbel-Max Trick
In this section, we develop the gradient estimates for the structured distributions defined Subsec- tion 2.2 We start with a brief introduction of the two main approaches to gradient estimation for discrete categorical variables: the score function estimator 45 and the Gumbel-SSftmax estima- tor 12 15). Then, we propose a low-variance modification of the score function estimator for the structured variables based on the intermediate representation of the variable. Finally, we conclude with discussion of control variates use together with the proposed estimator.
3.1 Gradient Estimation for Categorical Variables
We consider gradient estimates an expected objective VxEx where discrete random variable X has parametric distribution Px A) with finite support. The basic score fimetton es- timator 45 also known as REINFORCE, defines an unbiased estimate for the gradient using sample as [(x)V, A log P X æ; A). The estimator does not make any assumptions about L(-). but requires an efficient sampling algorithm for X and the score function VA log PX = x; A). For a categorical random variable X with outcome probabilities Px (X = k; A) x d computation of Vx log x = x; X) is linear the number of outcomes Therefore, the gradient estimation fast when d is small. However, for structured variables, such as graphs or sequences, the number of outcomes d grows rapidly with the structure size. In this case, the estimator requires custom algorithms for sampling and estimating the score function.
5

---

The Gumbel-softmax estimator, introduced in 15 29 is an alternative estimator that defines continuous relaxation based on Lemma I On the forward pass, replaces the categorical variable X with differentiable surrogate ğ = soft. max (2). where the input y (- log c.,.... log ea) is component- wise transformation of exponential samples. Due to Lemma D the surrogate converges
to the one-hot encoding of a categorical sample as soft max converges to arg max when -> 0.
On the backward pass, the estimator uses the chain rule to construct the gradient Vx using x estimator naturally extends to structured variables [31 34]. Specifically, the component- cemponentt-wiseeeee optimization in Lemma cam be replaced with a linear program over a structured set to generate structured variables and a relaxation can be used to define gradients. In the experimental section, we consider Stochastic Softmax Tricks (SST), introduced in [34], the relaxation-based baseline for comparison.
the reparameterization trick [66 to define the partial derivative of a sample as ax.
deg
The
As opposed to the score function estimator, the Gumbel-Softma estimator requires differentiable loss C(.). Such requirement imposes an additional restriction on model architecture. The architec- ture must be defined for the relaxed samples as well as the hard samples, non-trivial requirement for the models where discrete variables define branching 26 or the inputs the model not allowed to see [19 13 In practice, the vanilla score function estimator has notoriously higher variance compared the Gumbel-Softnn estimator and requires control variate to improve the gradient descent convergence.
3.2 The Score Function Estimator for the Recursive Gumbel-Max Trick
22 we have introduced probabilistic model involving exponential variable E and the structured variable defined as an output of an algorithm with input E. Additionally, we have defined an intermediate trace variable T = T(E) such that X is function X ( In this subsection, we apply T to estimate gradients of Ex (X)..
In Subsection
In our setup, the score function lolog function estimator
e; X) is available out of the box. However, the score
GE = ((X(()(()))))))))))) -E
=
(I)
which we refer to as E-REINFORCE, is rarely used in practice. In fact, the variance of the score function estimator using E as an action space exceeds the variance the estimator using X. On the other hand, the score function for the structured variable V log Px = x;; X) involves marginalization w.r.t. E and may require significant computation resources to estimate.
To mitigate the variance of gg. we follow the observation of 43 Appendix, Sec. B] and define another estimator as
gr :T C(X()) log PT(T
which is essentially the score function estimator that uses 7 rather than F. Below we refer to it as T-REINFORCE. Such estimate can he seen as the score function estimator gE marginalized over E given 7 (Appendix A contains the detailed derivation)
V log Pr (T; A) EEjr [C(XVV, log (r;AA(E;))) T].
(3)
As result, the proposed gradient estimate @r unbiased
EgC(X)V, log PE (E;; X) = Erggg
V - log [T,
(4)
whereas the variance of the estimate does not exceed the variance of gF
varg [ge] = vary EE TgE] + Ey var [gæ] = vary [gr] + Ez var [gel 2 vary [gr]. (5)
In fact, in our experiments, we observed significant difference in optimization due to the reduced variance of the estimator.
As we have argued in Subsection 2.2 we can compute the score function for trace variable and apply the estimator gr practice. Similarly, marginalization with respect to T X leads to the score function estimator gx := L(x)VA log Px(X = x;; X) and reduces the variance even further gr vart [gT] 2 varx [9x]. Therefore, the standard score function estimator is preferable when V log Px (X z; X) is available. In other cases, gr is practical alternative.

---

make strong as- 3.3 Further Variance Reduction for the Score Function Estimator
In addition to the marginalization described above, we mitigate the variance of the score function estimator with control variates. We two strategies construct the control variates. The first strategy uses the algorithm for conditional reparameterization of F (Appendix B, Algorithm 6) and defines family of sample-dependent control variates for the score function estimator 44 M. The estimator generates a sample e, runs the corresponding algorithm to obtain and x = (t). adds control variate c(e) and uses independent sample from the conditional distribution E T to eliminate the introduced bias
(C(X(t)) c(e) V (T:(T=t;;)) = - Vxc(e).
(6)
In general, the above estimate extends to any pair of random variables (B, 2) such that B = ( and the conditional distribution Z B admits the reparameterization trick. [43], the control variate used the relaxed loss C(-), whereas [II] proposed to learn the control variate improve the training dynamic. In our experiments, we use the estimator of [ and refer to as RELAX.
The second family of control variates we consider uses K >> 1 samples ty...., tx to reduce the variance. Besides averaging the independent estimates, uses the objective sample mean L. >-8 c(x(t,)) reduce the variance even further:
K Σ (c(x(t)) - VA -ig
Despite being quite simple, the above leave-one-out estimator proved to be competitive with multiple recent works [6] 41] In our experiments, we refer to such estimator as T-RRINFOOCE++ To facilitate fair comparison, batch training setup we reduce the batch size proportionally to K.
4 Related Work
Early models with structured latent variables include HMMs 39 PCFGs 36
sumptions about the model structure, and typically use EM-algorithm variations for training. This paper continues the line of work on perturbation models B3 for distributions over combinatorial sets. Initially, perturbation models approximated Gibbs distributions with an efficient sampling procedure using the MAP oracle for the Gibbs distribution. Later, [S 29] proposed to relax the component-we wise optimization used in the Gumbel-Max trick to facilitate gradient-based -assed learning for Gibbs distributions. The combination of the two approaches, namely a perturbed model together with bespoke relaxation of the MAP oracle, allows designing learning algorithms for latent subsets Recently, [44] developed a systematic approach to design relaxations for perturbed models with linear MAP oracle. Unlike the previous works, 45 for learning.
4б). permutations [11, trees [4] and sequences
mainly focus on the score function estimators
We illustrate the framework with the well-known Gumbel-Top-k trick 49 Among the various applications. 47 used the trick to define differentiable relaxation for the subset model considered used the trick to define score function estimators for latent
in paper; meanwhile, [0 42
permutations. Besides that, 21 22 leveraged the trick for sampling without replacement for certain family of graphical models and design gradient estimator using the sampler (0). Importantly, [34] Appendix, Sec. B] showed that the Kruskal's 's algorithm and the Chu-Liu- Shu-Liu-Edmonds algorithm extend the Gumbel-Top-k trick. They used the observation to argue in favor of exponential perturbations. In turn, we generalize the observation and propose learning algorithm based on the generalization.
The conditional reparameterization scheme, proposed in our work, allows action-dependent control variates 43 ||] for learning the structured variables. However, similarly to 41 6 we observed that
often simpler leave-one-out baseline 20 has better performance. Besides the control variates, [5 recently adopted conditional reparameterization to construct an improved Gumbel Straight-Throup estimator [15] for the categorical latent variables.
We denote the analogue, which uses the exponential score instead of the score of the trace, by E- REINFORCE+
7

---

Table 1: Results of k-subset selection on Aroma aspect data. MSE xx and subset precision (%) is shown for best models selected on validation averaged across different random seeds.
k 10
k 15
MODEL
ESTIMATOR
PREC.
PREC.
PREC.
SST (Our Impl.)
3.6 0.1
3.21 0.12
29.5
2.77 0.09
28.1
SIMPLE
E-REINFORCE+
3.80
25 1.4
3.77 0.23
26.7 3.4
3.16 0.16
25.3
T-REINFORCE+
0.19
30.5 +:
3.14
2.9
2.69 0.11
27.6
RELAX
0.16
24 + 1.9
3.5 + 0.13
28.9 1.9
2.95 0.15
26.1 1.9
SST (Our Impl.)
0.09
56 2.1
2.55 0.08
49.4
2.51 0.05
0.9
COMPLEX
REIIFOORC++
3.08
4.3
2.92
4.
T-REINFORC++
2.75 0.08
55.8 3.4
2.48 0.05
48.6
2.4 0.03
44.2 1.3
RELAX
2.8
2.1
2.58 : 0.09
47.6
2.46 #: 0.08
42.1 1.5
Applications
In the section below, we study various instances of algorithms with stochastic invariants and compare them against relaxation- based Stochastic Softmax Tricks (SST) introduced in [44] SST offers a generalization of the well-known Gumbel-Softmax trick to the case different structured random variables. The experimental setup is largely inherited from [e apart from Subsection 5.4 where the specifics of the problem do not allow relaxation- applicability of the score function-bass basedd gradient estimators. The main goal of [4] was to show that introducing structure will lead to superior performance compared to unstructured baselines. turn, we focus studying the benefits that one could get from using score-function gradient estimators based on rather than as well as showing competitive performance compared to relaxation-based SSTs.
gradient estimators, thus showing broader
Concerning the efficiency of the proposed srrrefffuctioo-baseee gradient estimators, the recursive form of Algorithm 2 does not facilitate batch parallelization in general. Specifically, the recursion depth and the decrease the size of may differ within a batch. Therefore, Algorithms 2 may require further optimization. We discuss our implementations in Appendix, Section C, and provide speed benchmarks in Table shows that, in practice, the performance is not much inferior the relaxation-based competitor. The implementation and how-to-use examples are publicly available
5.1 Learning to Explain by Finding Fixed Size Subset
We evaluated our method on the experimental setup from I.2X on the BeerAdvocate (30 dataset. The dataset consists of textual beer reviews and numerical ratings for four beer aspects (Aroma, Taste, Appearance, Palate). The model utilizes encoder-decoder architecture. The encoder outputs parameters of top-k distribution over subsets of review given the entire review. The decoder predicts a rating given the review subset of size k. Following setup of 34 we use & = {5, 10, 15) and two one-layer CNN and Complex
CNN architectures: Simple models for each aspect using MSE loss. We train models using several score function based estimators and compare them with implementation of SST. We report means and standard deviations for loss and precision averaged across different random model initializations. Table shows the obtained results for Aroma aspect. The best metrics with respect means highlighted in bold. Detailed experimental setup and description models can be found Appendix D.I.
three-layer CNN and train our
5.2 Learning Latent Spanning Trees with Kruskal's Algorithm
Given system of interacting particles, the dependencies of their states be formally described as an undirected graph. We use Neural Relational Inference [17], initially representing relaxation- based approach, and examine its performance interconnections, when applying score function techniques instead. We build an experiment in line with [4], generating data by translating ground truth latent spanning tree (corresponding the connections in the system) into a sequence of positions of points on a real plane, representing
generative model and ability reconstruct system

---

Table 2: Graph Layout experiment results for T=10 iterations. Metrics are obtained by choosing models with best validation ELBO and averaging results across different random seeds on the test set.
T 10
ESTIMATOR
ELBO
EDGE PREC.
EDGE REC.
MEAN + STD
MAX
MEAN + STD
MAX
MEAN = STD
MAX
SST (Our Impl.)
-1860.79 = 1116.83
-1374.81
87 + 21
95
93 + 3
95
T-REINFORCE+
-1582.72 72 + 571.18
-1192.04
70 + 31
91
86 + 8
91 90
RELAX
18079.18 + 569.49
-1205.87
43 + 31
90
81 +
Table 3: Unsupervised Parsing ListOps. We report the average rest-performance of the model with the best validation accuracy across different random initializations.
ESTIMATOR
ACCURACY
PRECISION
RECALL
MEAN STD
MAX
MEAN
MAX
MEAN + STD
MAX
SST (Our Impl.)
78.42 8.14
93.78
56.84 20.08
82.40
30.18 + 19.10
73.11
E-REINFORCE+
60.25 2.29
64.47
40.87 6.90
45.74
40.74 + 6.93
45.46
T-REINFORCE+
87.34 3.00
91.97
77.93 7.36
79.65
61.10 14.11
79.65
RELAX
79.60 1 9.36
88.64
54.73 17.48
75.27
53.61 17.14
75.27
dynamics of the system over time. These points are obtained executing force-directed algorithm for T or
g
iterations and fed into the model.
The architecture of the model consists of a graph neural network (GNN) encoder, producing distri- bution over spanning trees, and GNN decoder, producing distribution over time series of points positions. Model is trained in a manner of variational autoencoders (VAEs), optimizing ELBO, lower bound on the joint migeliiklihood of the observed data points all timesteps.
We measure precision and recall with respect to the encoder samples and the ground truth dependency spanning tree. Table shows the results for T=10 iterations. Overall, score function methods counterpart, achieving higher values of ELBO on the
performed better than their relaxation-bal test set, but slightly worse performance in terms of structure recovery metrics. The results for T=20 and the detailed experimental setup are described in the Appendix D.2.
5.3 Unsupervised Parsing with Rooted Trees with CLE Algorithm
We study the ability of the proposed score function estimators to recover the latent structure of the data setting, where can be quite accurately described with an arborescence. Following details about data and models outlined by [d] we use simplified version of the ListOps (3) dataset. It consists of mathematical expressions (e.g. min med [3
2]), written in prefix form along with results of their evaluation, which are integers in [0,9]. Given a prefix expression, one can algorithmically recover its structure as a parse tree. We bound maximal length expressions and maximal depth of their parses along with removing the examples with summod operator. These limitations sufficiently decrease the amount of memory a model should have to calculate the result and facilitates the usage of GNNs which now become capable of evaluating expressions by bounded number of message passing steps.
Our model consists of two parts: an encoder and a classifier. The encoder is pair of LSTMs that generate parameters of the distribution over rooted arborescence on token nodes. The classifier is GNN, which passes a fixed number of messages over the sampled arborescence and feeds the resulting embedding of the first token into the final MLP. Models are trained simultaneously to minimize cross-ents cross-entropy. We examine the performance of the models by measuring classification accuracy along with precision and recall with respect to the edges ground truth parse trees. Table shows score function based estimators, particularly T-REINFORCE+, show more stable performance comparison to relaxation-basse estimator. Detailed description the experiment can be found in the Appendix D.3.

---

5,4
Non-monotonic Generation of Balanced Parentheses with Binary Trees
Percentage of New and Correct
Percentage of Old
0.014
0.95
0.012
0.90
0.010
127
0.85
0.008
10
20
30
40
50
10
20
30
40
50
Number of Parentheses Types
Number of Parentheses Types
Figure 2: Generative statisties of hon-monotonic language model with different orders
We apply our methods in a setting where the application of continuous relaxations is seemingly not possible. We construct simple dataset by sampling fixed number of balanced parentheses of various types (from 10 to 50 types) and model them with non- -onmonotooic architecture, defined in 44 Here, language model generates sequences by consecutively adding tokens in between previously generated ones. This process can be seen as modeling joint distribution over binary trees (treated orders of generation) and sequences (which generated along these orders).
We refer to this prohabilistic model as the decoder, fix architecture to be single-layer LSTM and teacher-forcing for training. More precisely, compare two different modes of training decoder: fixed order and semi-supervised trainable order. In the first case, the order is obtained by applying a deterministic function to the input sequence. We study two fixed orders: left-to-right, corresponding to a degenerate tree with only right child nodes, and a more natural smart order, described in Appendix D.4. In the semi-supervised (ss) case, 10% of sequences overall model is trained as a VAE. We train the decoder by directly backpropagating the output signal and obtain the gradients for the encoder using RELAX estimator.
paired with the smart order, and the
We choose models with the best validation perplexity and evaluate them by generating 200000 unique sequences and measuring portion of those which are present in the train dataset (Old) and those which are balanced and not seen during training (New and Correct), Results of this toy experiment show that it is possible to improve generative metrics by considering non-trivial orders of generation. Experiment is described details in the Appendix D.4.
Discussion
Below, we speculate about the pros and cons of the relaxation-based 34 and the score function-based approaches training latent structured variable models. While both build upon the Gunbel-Max trick, the generalizations develop different ideas. Our work summarizes the recent observations and formulates general algorithm for which the Gumbel-Max applies recursively. We utilize the properties of these algorithms to construct unbiased score cnnction-based gradient estimators. In contrast, Rd generalizes the relaxation of the Gumbel-Max trick different combinatorial structures and produces biased reparametrizatiio-bann based gradient estimators. While the relaxation- gradient estimators are biased and limited to differentiable objectives, they have lower variance and faster batch processing time (Tables 8 and 9). At the same time, score function- based estimators are unbiased and apply when the loss functions do admit relaxations (Section (.4). Occasionally, they lead to better optimization and lower objective values (Tables 10, 11 and 12). The choice of control variates also introduces a trade-off in our framework. As the experiments show, the objective is highly parallelable, the multi-sample T T-REINFORCE+ estimator is preferable. However, a single-sample RELAX estimator a more suitable choice when we cannot obtain multiple objective samples (i.e., in reinforcement learning). Finally, as a direction for future improvements, we suggest applying the conditional reparameterization scheme used in the control variate improve the Gumbel ttriugh estimators.
10

---

Acknowledgements
The authors thank the reviewers for the valuable feedback. The research was supported by the Russian Science Foundation grant no. 19-71-30020 and through the computational resources of HPC facilities at HSE University [23]
References
[1] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv preprint 1rX.h:/509.000 2015.
[2] Jianbo Chen, Le Song, Martin Wainwright, and Michael Jordan. Learning to explain: An informattonn-hheortiie perspective on model interpretation. In International Conference on Machine Learning, pages 883-892. PMLR, 2018.
[3] -aatise Cordonnier, Aravindh Mahendran, Alexey Dosovitskiy, Dirk Weissenborn, Jakob Uszkoreit, and Thomas Unterthiner. Differentiable patch selection for image recognition. arXiv preprint wrXiv:2204.03000 2021.
[4] Caio Corro and Ivan Titov. Differentiable perturb-and-parse: Semi-supervised parsing with structured variational auto-encoder. In International Conference on Learning Representations, 2019.
[5] Sanjoy Dasgupta, Christos H Papadimitriou, and Umesh Virkumar Vazirani. Algorithms. McGraw-Hill Higher Education New York, 2008
[6] Zhe Dong, Andriy Mnih, and George Tucker. Coupled gradient estimators for discrete latent variables. arXiv preprint 110::2006.08056, 2021.
[7] Jack Edmonds. Optimum branchings. Journal of Research of the national Bureau of Standards
B, 71(4):233-240, 1967.
[8] Thomas MJ Fruchterman and Edward M Reingold. Graph drawing by force-directed placement. Software: Practice and experience, 21(11):1129-111 1991.
[9] Yao Fu, Chuanqi Tan, Bin Bi, Mosha Chen, Yansong Feng, and Alexander M Rush. Latent template induction with gumbel-crfs. arXiv preprint arXiv:2011. 1..244 2020.
[10] Artyom Gadetsky, Kirill Struminsky, Christopher Robinson, Novi Quadrianto, and Dmitry Vetrov. Low-variance black-box gradient estimates for the plackett-luce distribution. Proceed- ings of the AAAI Conference on Artificial Intelligence, 34(06)::0126-10133 2020.
[11] Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Duvenaud. Backpropa- gation through the void: Optimizing control variates for black-box gradient estimation. arXiv preprint arXiv: /7 1.00/23, 2017.
[12] Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint 08983, 2016.
[13] Jiatao Gu, Qi Liu, and Kyunghyun Cho. Insertion-based decoding with automatically inferred generation order. Transactions of the Association for Computational Linguistics, 7:661-676, 2019.
[14] Serhii Havrylov, Germán Kruszewski, and Armand Joulin. Cooperative learning of disjoint syntax and semantics. arXiv preprint arXiv:I902.0910 2019.
[15] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprini arXiv: /6/ 1.01 2016.
[16] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
[17] Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel Neural relational inference for interacting systems. In International Conference on Machine Learning, pages PMLR, 2018.

---

[18] Jon Kleinberg and Eva Tardos. Algorithm design. Pearson Education India, 2006.
[19] Wouter Kool, Herke van Hoof, and Max Welling, Attention, learn solve routing problems! In International Conference on Learning Representations, 2018.
[20] Wouter Kool, Herke van Hoof, and Max Welling. Buy 4 reinforce samples, get baseline for free! http://openteviee.eefffommmm htt.://orrrmm//rrrlgggg.... 2019.
[21] Wouter Kool, Herke Van Hoof, and Max Welling. Stochastic beams and where to find them: The gumbel-top- trick for sampling sequences without replacement. International Conference on Machine Learning, pages 3499 3508. PMLR, 2019.
[22] Wouter Kool, Herke van Hoof, and Max Welling, Ancestral gumbel-top-k sampling for sampling without replacement. Journal of Machine Learning Research, 21(47):1-36, 2020.
[23] PS Kostenetskiy, RA Chulkevich, and VI Kozyrev. Hpc resources of the higher school of economics. Journal of Physics: Conference Series, 1740(1):012050, 2021.
[24] Joseph Kruskal. On the shortest spanning subtree of graph and the traveling salesman problem. Proceedings of the American Mathematical society, 7(1):48 40, 1956.
[25] Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):83-2) 07.883397719 1955.
[26] Tuan Anh Le, Adam Kosiorek, N Siddharth, Yee Whye Teh, and Frank Wood. Revisiting reweighted wake-sleep for models with stochastic control flow. In Uncertainty in Artificial Intelligence, pages 1039-1049, PMLR, 2020.
[27] Tao Lei, Regina Barzilay, and Tommi Jaakkola. Rationalizing neural predictions. arXiv preprint 160iv: 04155, 2016.
[28] Xuanlin Li, Brandon Trabucco, Dong Huk Park, Michael Luo, Sheng Shen, Trevor Darrell, and Yang Gao. Discovering non-monotonic autoregressive orderings with variational inference. In International Conference on Learning Representations, 2021.
[29] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation discrete random variables. arXiv preprint arXiv:1611.0077 2016.
[30] Julian McAuley, Jure Leskovec, and Dan Jurafsky. Learning attitudes and attributes from multi-aspect reviews. 2012 IEEE 12th International Conference on Data Mining, pages 1020-1025. IEBE, 2012.
[31] Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permuta- tions with gumbel-sinkhon networks. arXiv preprint arXiv:/802.08665, 2018.
[32] Nikita Nangia and Samuel R Bowman. Listops: diagnostic dataset for latent tree learning. arXiv preprint arXiv: arXiv:/804.000 /804.06028 2018.
[33] George Papandreou and Alan L Yuille. Perturb-and -nd--adddaappp random fields: Using discrete optimiza- tion to learn and sample from energy models. In 2011 International Conference on Computer Vision, pages 193-200. IEEE, 2011.
[34] Max B Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris Maddison. Gradient estimation with stochastic softmax tricks. arXiv preprint arXiv:2006. 000.:2000000033, 2020.
[35] Max B Paulus, Chris J Maddison, and Andreas Krause. Rao-blackwellizzing the straight-through gumbel-softmax gradient estimator. arXiv preprint arXiv:2010.048338 2020.
[36] Slav Petrov and Dan Klein. Discriminative log-linear grammars with latent variables. In Advances in neural information processing systems, pages 1153-1160. Citeseer, 2008.
[37] Robin I. Plackett. The analysis of permutations. Journal of the Royal Statistical Society: Series
C (Applied Statistics), 193-)2193-22,, 1975.
12

---

[38] Robert Clay Prim. Shortest connection networks and some generalizations. The Bell System Technical Journal, 36(6):1389-1411, 1957.
[39] Lawrence R Rabiner. A tutorial on hidden markov models and selected applications speech recognition. Proceedings of the IEEE, 77(2):257-266, 1989.
[40] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pages 1278-1286. PMLR, 2014.
[41] Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco JR Ruiz, and Ömer Deniz Aky- ildiz. Vargrad: A low-variance gradient estimator for variational inference. arXiv preprint arXiv:2010.10433 10436 2020.
[42] Valentino Santucci, Josu Ceberio, and Mareo Baioletti. Gradient search in the space of permuta- tions: an application for the linear ordering problem. In Proceedings of the 2020 Genetic and Evolutionary Computation Conference Companion, pages 7704-1711, 2020.
[43] George Tucker, Andriy Mnih, Chris J Maddison, John Lawson, and Jascha Sohl- Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. In Advances in Neural Information Processing Systems, pages 2627-2636, 2017.
[44] Sean Welleck, Kianté Brantley, Hal Daumé lii, and Kyunghyun Cho. Non-monotonic sequential text generation. In International Conference on Machine Learning, pages 6716-6726. PMLR, 2019.
[45] Ronald J Williams. Simple statistical foolowig-folloiing algorithms for connectionist reinforce- ment learning. Machine learning, 8(3-4):229-2555 1992.
[46] Sang Michael Xie and Stefano Ermon. Reparameterizable subset sampling via continuous relaxations. In IJCAI, 2019.
[47] Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable top-k operator with optimal transport. arXiv preprint arXiv:2002.06504 2020.
[48] Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In International conference on machine learning, pages 2048-2057. PMLR, 2015.
[49] John I Yellott Jr. The relationship between luce's choice axiom, thurstone's 's theory of compara- tive judgment, and the double exponential distribution. Journal of Mathematical Psychology, 15(2):109-144 144, 1977.
[50] Dani Yogatama, Phil Blunsom, Chris Dyer, Edward Grefenstette, and Wang Ling. Learning to compose words into sentences with reinforcement learning. arXiv preprint 1rXiv:1611.09100 2016.
13


Reviews:
[{"review_id": "w1LgzuSZL-P", "reviewer": "Reviewer_SiHy", "summary": "Endowing a latent representation with particular structural properties is a common way to incorporate inductive biases in deep learning. Structured latent variable models often require discrete components which can make optimisation of these models challenging. Typically, learning approaches use (i) differentiable surrogates which rely on relaxations (to continuous components) and thereby introduce additional constraints and biases in gradient estimation or (ii) score function estimators which lead to high variance in the gradient estimates. This submission proposes a framework providing low-variance score function estimators on the basis of a property termed *stochastic invariance* and an intermediate *trace* variable $T$. The invariance describes that the distribution of the perturbed (exponential) inputs $E$ does not change in the proposed recursive algorithms, i.e. stays invariant throughout the recursion. Empirical evaluation is performed on different common tasks with structured latent variables: (i) identifying fixed-size subsets of predictors which are most informative for target prediction (*BeerAdvocate* dataset), (ii) graph-represented particle interactions (*Graph Layout*), (iii) recovering of parse trees / arborescence (simplified *ListOps* dataset), and (iv) a generated dataset of parentheses (akin to language models) where relaxation-based approaches might not be feasible. The considered models are based on CNNs, GNNs, and LSTMs arranged in encoder-decoder or encoder-classifier architectures. In particular, models employing stochastic softmax tricks (SST) and a relaxed version of the used score function estimator are compared to the proposed method variants. The results indicate that the proposed approach can generally compete well or even outperform these baselines.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## 1. Strengths\n\n### Originality / Novelty:\n- This submission derives conditions and algorithms for which score function estimators (REINFORCE) can yield low-variance gradient estimates and compete with SOTA relaxation-based approaches to learning of structured latent variables. In this work, the perturbed (exponential) input $E$ is connected via an intermediate trace variable $T\\equiv T(E)$ to structured (target) variable $X\\equiv X(T)$. The proposed recursive algorithms make use of stochastic invariance which takes advantage of the Gumbel-Max trick. The key aspect of this work is that computations can be made w.r.t. trace $T$ and compute score functions on the basis of $T$ ($T$-REINFORCE+) with reduced variance. \n- The underlying idea of stochastic invariance and trace $T$ is not novel and the submission shares similarity to other work which was a main discussion point in a previous submission to ICML 2021. However, in my opinion this submission convincingly highlights the differences to related work and puts emphasis on practical usage of $T$ for score function estimation and conditional reparameterisation. As such I believe the work extends the current literature and yields some novel aspects. \n \n### Clarity:\n- I consider this submission well-structured, clearly presented and mostly self-contained. \n\n### Quality:\n- The empirical as well as formal and more technical aspects of this submission appear sound as far as I can judge.\n- In my opinion, the submission succeeds in demonstrating that the proposed approach indeed provides empirical advantages. It is shown that the $T$-REINFORCE+ approach generally outperforms $E$-REINFORCE+ and competes mostly well with the relaxation-based models. \n\n### Significance / Relevance:\n- The submission extends the literature concerning score function estimators in the above described ways. This work might contribute to making score function estimators more applicable and relevant for a larger audience. \n\n## 2. Weaknesses\n\n### Clarity:\n- Details on auxiliary variable $R$: I am unsure I fully understand the meaning and necessity of auxiliary variable $R$. It is stated that it *”[…] accumulate[s] the necessary information for the next call […]”* (lines 141-142), but I do not quite follow here. This is probably a more technical detail, but could the authors give some details on that? \n- *”Notably, the multi-sample estimator with a smaller batch is more stable than the single-sample estimator with a sample-dependent control variate.”* (lines 321-322) I am not sure I can follow in this conclusion here. Where was this shown or what are the authors referring to here?\n\n### Quality:\n- One goal of the research here is to *”[…] allow learning models that do not admit relaxed variables […]”* (l. 45) which is attempted to be addressed in the last experiment (section 6.4). As I have more limited experience in this direction, could the authors maybe comment on how relevant this task is. Is there for instance a class of problems where it is clear that relaxation-based approaches are not viable? In the end, it might be a less pressing matter as apparently the community also moved more towards relaxation-based approaches and identifying a surrogate formulation which is differentiable. I believe the submission would greatly benefit from such an exposition.\n\n## 3. Additional Feedback\n\nSome minor aspects the authors might want to take into consideration when revising the submission:\n\n- Figure 1, Algo.s 1 & 2: I might be mistaken, but should it not read $T\\Leftarrow$ argmin $E$ and $T_i\\Leftarrow$ argmin $E_i$ in these algorithms? It could have been meant to highlight what is done there more explicitly, but I rather stumbled upon that.\n- Lines 111-113: I might have missed the explanation to which is referred here, but I could not follow why $E’_j=E_j-E_T$ might be required / beneficial. Could the authors point out where this is explained or comment on that?\n- L. 128, missing dot at the end of the sentence.\n- L. 152: “[…] Lemma 2 […]” should probably be “[…] Lemma 1 […]”\n- L. 202: *“As we’ve argued […]”* -> “As we have argued […]”\n- SST abbreviation: It might be advisable for clarity to introduce the abbreviation of stochastic softmax tricks (SST) in the main text (currently in the appendix, L. 666).\n\n\n## 4. Recommendation\n\nI believe this is a well-written and self-contained submission which contributes new insights to score function estimators and learning structured latent variables. In my opinion the work highlights differences to other related papers and extends these. Therefore, I consider this submission a worthy contribution to NeurIPS and my initial recommendation is to accept this paper.\n\n### Post-Rebuttal Update:\nI would like to thank the authors for their thorough rebuttal and additional clarifications. I stand with my initial assessment and endorse accepting this submission for publication at NeurIPS.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors have adequately addressed the limitations and societal impact, in my opinion (see page 9 lines 326-329).", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "gVINElqcKL5", "reviewer": "Reviewer_NYB7", "summary": "This paper shows that, contrary to popular belief, score function gradient estimators for structured models are competitive to estimators based on continuous relaxations. The paper is based on the Gumbel-Max / competing exponential clock trick for structured models, and identifies that we can estimate gradients wrt the “trace” T of the program involved in the implicit mapping from the exponential clocks E to the observed X. This is most helpful when we can easily compute p(T), but not p(X).", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Originality*\n* This work integrates tricks for gradient estimation using REINFORCE in discrete models and extends them to more complex settings. The use of REINFORCE on traces seems novel, and its success relative to continuous relaxations is surprising.\n\n*Quality*\n* The submission appears to be technically sound, and the experiments do support the claim that REINFORCE can achieve success comparable to continuous relaxations in structured settings. However, the numbers reported for the baseline SST do seem to sometimes be quite a bit lower than those reported in Paulus et. al. [1], which is concerning. For example, in Table 2 the complex model for SST has a mean MSE of 2.93, while all the complex models in Paulus et. al. [1] have mean MSE's <= 2.7.\n* The experiment with non-monotonic generation of balanced parentheses  demonstrates that the proposed estimator is indeed more general than continuous relaxation-based methods.\n* An analysis of the different gradient estimator’s MSEs compared to a continuous relaxation on a tractable structured model (so you can compute the exact gradient) would also be nice to see.\n\n*Clarity*\n* The submission could be improved by better motivating the problem and rationally reconstructing the proposed solution. One possible way of doing this could be to start by showing how gradient estimation in structured distributions is hard. Explain why, i.e. computing p(X) can be expensive, and give an example structured distribution that shows this. Use this to motivate the trace as a transformation T(E) = (X, Y), where although we have unused information Y, it is easier to compute p(T).\n* The details on Gumbel-softmax following Lemma 1 do not seem relevant to the contribution of the paper (getting REINFORCE working) and can be removed.\n* The details on variance reduction should come after the introduction of the trace-based estimator.\n* The motivation for f_split in algorithm 2 should be in the main text, and required for the example.\n\n*Significance*\n* The results are important, and provide evidence that REINFORCE-based gradient estimators can be competitive with continuous relaxations.\n\n*Questions*\n* Why was stochastic invariant chosen as the keyword for the approach? Why not something like trace-based gradient estimators?\n* Does the T-REINFORCE+ estimator only use the LOO estimator, and no continuous relaxation? If so, does this mean that the control variates that use conditional reparameterization never outperformed the LOO estimator?\n* Instead of using a batch size of N/K and K samples, can the baseline SST be run with K noise samples in order to keep the batch size the same?\n* What are the speed numbers (time per iteration)? Since batching was difficult, it is clear these will not be great, but it would still be good to show them.\n* Can you detail exactly where the algorithm fails to apply to Prim's?\n\n*Overall*\n* Nice contribution, but experiments and writing need to be improved\n\n[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.\n\nEdit: Author response clarified experimental numbers. The remaining concern is writing, but the paper is understandable.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "* The speed limitations due to the difficulty of batching algorithm 2 was made clear in the text.\n* However, the limitations of the conditions of the algorithm itself are not as clear. While it was mentioned that Prim's algorithm cannot be handled, a better description of the limitations is needed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "I_1Uv_rVNDl", "reviewer": "Reviewer_npTm", "summary": "Mathematically, this paper proposes a framework for designing a structured latent variable X and also for training the related objective E_{X} [L(X)]. \n\nThe structured latent variable X could model ranking, subsets, sequences, spanning trees, rooted trees, etc. \n\nThe framework implicitly models X with Exponential random variables followed by a recursive calculation; the corresponding P(X; \\lambda) could be intractable. Here, after one step of the calculation, the resulting output will also be Exponentially distributed; the authors name such phenomenon (unchanged Exponential distribution) as \"Stochastic Invariant.\"\n\nTo train the objective E_{X} [L(X)], where P(X; \\lambda) is intractable, the authors propose to do REINFORCE with internal variable T from the above recursive calculation.\n\nExperiments on a variety of situations are conducted to evaluate the proposed framework.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality. The presented framework is believed novel. However, the relationship between the proposed framework and related work (such as those mentioned in Section 3.3) is not clear. The detailed discussions/comparisons are necessary to highlight the novelties here. For example, the proposed framework seems to include Kruskal's algorithm [21] and Chu-Liu-Edmonds [6] as special cases, right? \n\nQuality. The presented framework is technically sound. The theoretical analysis and experimental results are convincing.\n\nClarity. The writing is clearly not friendly to general machine learning audiences. I have to read the paper several times to understand the main idea. General audiences would appreciate simple and straight statements of the problem settings and contributions.\nDefinitions should be presented early. For example, what are structured latent variables, Stochastic Invariant, etc? For example, I don't understand the term of Stochastic Invariant until Line 116. Besides, these definitions are much easier to understand if their mathematical expressions are provided.\n\nSignificance. The presented framework is likely to be important and to be reused by other researchers.\n\nOther comments.\n\nWhen I read the submission for the first time, I am completely lost after seeing Table 1. Honestly, this table, placed here, is not easy to understand at all. \n\nNotations could be improved. For example, use a bold symbol to denote a vector; replace X^{-1} in Line 124 with another notation.\n\nIn Lines 135-145, the sub-routines, e.g., f_{map}, are not defined.\n\nLine 152. Lemma \"2\" is a typo. \n\nThe authors seem to have considered many extensions of their framework, for example, in Section 3.3 and the appendix. However, that is not clearly stated in the current submission. I would suggest leveraging a table to summarize the extensions and using mathematical definitions for clear demonstrations.  \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "BSLKADJPkIz", "reviewer": "Reviewer_ahTk", "summary": "The paper proposed a new gradient estimator, which is composed of Gumbel max top-k trick and the score function, for training discrete distribution. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The gradient estimator is new.\n\nClarity: The paper is hard to follow. \n1. I didn't find what's the meaning of \"SST\", which is treated as the baseline in all the experiments.  \n2. I couldn't understand what is \"stochastic invariant\", is the definition of \"stochastic invariant\" first appears in line 131: \"the first argument E of the algorithm always has\n132 exponential distribution (conditioned on the trace T from the above recursion levels)\"?\n3.  I think it would be better to move the introduction of  \"Execution Trace T\" to the beginning of section 3, otherwise, we cannot understand table 1.\n\n\nQuality:\n1. The proposed method needs a lot of computation, however, it has larger variance than the standard X space score function estimator. (line 205)  It is mentioned that the standard score function estimator is preferable when we can access the score function, otherwise, the proposed estimator is available. Could the author give a few practical examples where the score function is not available?\n\n\nSignificance:\nI don't know the significance since I don't even know what the baseline method SSL represents for.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the limitations are addressed in section 3.3 and 6.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper exploits a recent observation that exponential random variables, when filtered through an algorithm that makes local, greedy choices, can sometimes keep their distributions invariant. The paper shows that in such cases, this invariance can be exploited to provide lower variance REINFORCE gradient estimators. After the response period and discussion, most reviewers agreed that this is an interesting and valuable contribution. Some reviewers felt that the lack of clarity in the writing significantly diminished the contribution, but no reviewers felt that this disqualified the paper from publication.

I am convinced that this paper merits publication for two reasons. 
* I believe the stochastic invariant structure is of independent interest in cases beyond those considered here.
* The insight that stochastic invariants allow for lower variance gradient estimation is a valuable one.

For the authors to achieve the maximum impact possible with this paper, I suggest that they take seriously the feedback on clarity. In particular, that they consider *minor* revisions to address the motivation, definition of stochastic invariance, and details of the algorithm.
Decision: Accept (Poster)