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
HSE University
HSE University, Skoltech! Moscow, Russia
Moscow, Russia
Moscow, Russia
k. struminsky&gmail.comm
artygadetakk@@yanee.u
rakitiindemis32@maail..
Danil Karpushkin
Dmitry Vetrov
AIRI,
Sber AI Lab, MIPT Moscow, Russia
HSE University, AIRI
Moscow, Russia vetrovd@yandex.oo
kardani @aaani@@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 probobbbbyyl 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 E, 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 PggPr(t (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
[Because Et(t)) = 0)
K, R, [m))m1)))
X x
(X/
k E K'}
log Pr (t>1 T = f A) + Fig poub (t²1, 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 [(z)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
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-RRINFOORE++ 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-Eddonds 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 ttrough 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-hheorttie 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 100::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.eefffnmmmm htt.://orrrmm//rrllgggg.... 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(4)):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.:2000000063, 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 followig 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 114, 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": "pJFg3wUhNsF", "reviewer": "Reviewer_9Dk3", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "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": "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.\n\nIt 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.\n\nI 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?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "None", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "VkIdZdBMyD", "reviewer": "Reviewer_Csyp", "summary": "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.", "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": "**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?\n\nThe 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.\n\n**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?\n\nThe method is technically sound and well explained. I particularly appreciated the min-k example in Algorithm 1.\n\nOne thing that isn't completely clear to me is the family of latent variable models to which this method is applicable.\n1) 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.\n2) 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.\n\nTogether a discussion of the generality of the method would be really helpful!\n\n**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.)\n\nThe paper is clearly written (except the point above) and well organized.\n\n**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?\n\nThe 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?\n\nI 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?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical issues.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8 hours"}, "extra_text": {}}, {"review_id": "LtYx_TiEZJb", "reviewer": "Reviewer_eF5a", "summary": "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.\n", "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": "**Originality:**\nTo 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.\n\n**Clarity:**\nWhile 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.\nHence, 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.\nMy 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.\nIf this is indeed the case, I'd strongly encourage the authors to make this clear and motivate this early on in the paper.\nTo 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.\n\nAnother 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.\n\nI also found the notation in Algorithm 1. confusing.\nIn 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.\n\n\n**Quality:**\nTo 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.\n\n**Significance:**\nMany systems and generative processes incorporate or can be described in terms of discrete latent variables, e.g. discrete control flow in probabilistic programming models.\nThe 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.\nMy concern is that the current manuscript is challenging to read and understand and as a result readers might not pick up on the idea.\n\n\n**Typos I noticed**\n  - Line 149: function of *the* trace\n  - Line 152: Lemma 2 -> Lemma 1?\n  - Line 307: For *the* encoder\n  - Line 321: *small*\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\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "adequate", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "-If9FjeqhYB", "reviewer": "Reviewer_ZnaB", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "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: To my knowledge, the idea is original.\n\nClarity 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:\n\n* 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.\n* Similarly, X is not explained well. \n* 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.\n* The need for partitioning the indices in algorithm 2 is not explained or really referenced elsewhere.\n* f_split, f_combine, and f_stop are not defined or really explained in depth.\n\nThese issues and more combine to make it very difficult to understand specifically what is being proposed.\n\nIn 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.\n\nFinally, 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.\n\nAdditionally, the current experimental results don't seem to strongly recommend the authors' proposed method.\n\nIf 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.\n\nEDIT: After rebuttal I have raised my score.", "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": "5 hours"}, "extra_text": {}}]

Author rebuttal:

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 (Poster)