Efficient Learning of Discrete-Continuous Computation Graphs
Authors: David Friede
Mathias Niepert
Venue: NeurIPS 2021 Poster (2021)
Keywords: Discrete-Continuous Learning
Stochastic Computation Graphs
Gumbel-softmax Trick

Abstract: Numerous models for supervised and reinforcement learning benefit from combinations of discrete and continuous model components. End-to-end learnable discrete-continuous models are compositional, tend to generalize better, and are more interpretable. A popular approach to building discrete-continuous computation graphs is that of integrating discrete probability distributions into neural networks using stochastic softmax tricks. Prior work has mainly focused on computation graphs with a single discrete component on each of the graph's execution paths. We analyze the behavior of more complex stochastic computations graphs with multiple sequential discrete components. We show that it is challenging to optimize the parameters of these models, mainly due to small  gradients and local minima. We then propose two new strategies to overcome these challenges. First, we show that increasing the scale parameter of the Gumbel noise perturbations during training improves the learning behavior. Second, we propose dropout residual connections specifically tailored to stochastic, discrete-continuous computation graphs. With an extensive set of experiments, we show that we can train complex discrete-continuous models which one cannot train with standard stochastic softmax tricks. We also show that complex discrete-stochastic models generalize better than their continuous counterparts on several benchmark datasets.

are instances Efficient Learning of Discrete-Continuous Computation Graphs
David Friede
Mathias Niepert' mathias nieperteneclab.... eu
david@inf ormatik. uni-niiinn wwrlmmanneimm.ommm de
University of Mannheim Mannheim, Germany
NEC Laboratories Europe Heidelberg, Germany
Abstract
Numerous models for supervised and reinforcement learning benefit from com- binations of discrete and continuous model components. End-to-end learnable discrete-continuous models are compositional, tend to generalize better, and are more interpretable. A popular approach to building discrete-continuous compu- tation graphs is that of integrating discrete probability distributions into neural networks using stochastic softmax tricks. Prior work has mainly focused com- putation graphs with a single discrete component on each the graph's 's exccution paths. We analyze the behavior of more complex stochastic computations graphs with multiple sequential discrete components. We show that it is challenging to optimize the parameters of these models, mainly due small gradients and local minima. We then propose two new strategies to overcome these challenges. First, we show that increasing the scale parameter of the Gumbel noise perturbations dur- ing training improves the learning behavior. Second, we propose dropout residual connections specifically tailored to stochastic, ciscrrte-coninuuuus computation graphs. With an extensive set of experiments, we show that we can train complex discrete-continuous models which one cannot train with standard stochastic soft- max tricks. We also show that complex discrete-stochastic models generalize better than their continuous counterparts on several benchmark datasets.
caacrete-coninuuuus
Introduction
Neuro-symbolic learning systems aim combine discrete and continuous operations. The ma- jority of recent research has focused on integrating neural network components into probabilistic logics 28 20), that is, making logic-based reasoning approaches more amenable to noisy and high- dimensional input data. On the other end of the spectrum are the data-driven approaches, which aim learn modular and discrete program structure in end-to-end neural network based system. The broader vision followed by proponents of these approaches are systems capable of assembling the required modular operations to solve variety of tasks with heterogeneous input data. Reinforcement learning [33], neuro-symbolic program synthesis [25], and neural module networks [ of such data-integrated discrete-continuuus learning systems.
We focus on learning systems comprised of both symbolic and continuous operations where the symbolic operations modeled discrete probability distributions. The resulting systems can be described by their stochastic computation graphs, a formalism recently introduced to unify several related approaches [30] Figure iillustrates
instances of such stochastic computation graphs. For a given high-dimensional input such as a set of images or a list of symbols, computation graph, consisting of stochastic (discrete) and continuous (neural) nodes, is created to solve partieular task. Both, the mechanism to assemble the graphs (if not already provided) and the various operations are
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

3
plz, B)
ADD
pl: 8)
ADD
5
p(:: 6)
Mcc 8)
L
2
plz B)
[MAX 2 9]
GENERATE GRAPH
2
p(z: 0) -
MAX
9
pa: 0) MAX
P
L
M ()
Two instances of discrete-continuous computation graphs. Top: The MNIST addition task aims to learn neuro-symbolic program end-to-end that sums two or more digits only based on raw image inputs. The input images are mapped to the parameters of a categorical probability distribution (orange neural network). Samples from said distribution are mapped to learned vector representations of digits (green neural network). The addition module ADD takes the vector representations of two digits and learns to compute the sum of the corresponding digits. Bottom: The ListOps task aims to learn the solution to arithmetic expression. Given the syntactic input and the numerical solution, we learn the latent parse tree, which consists of discrete and continuous components modeling the operations MAX, MIN, and MED. More details are provided in the experimental section.
Figure I:
learned end-to-end. Discrete nodes in the computation graph are Gibbs distributions modeled at particular temperature which can also be used to model the argmax operation at zero temperature.
The majority of prior work has focused graphs with single discrete probability distribution along each execution path. Examples are the discrete variational autoencoder [14], learning to explain [4], and other applications of stochastic softmax tricks (26). We aim to analyze and improve the training behavior of complex computation graphs, that is, graphs with more than one discrete probability distribution in its execution paths. More concretely, we focus stochastic nodes are categorical random variables modeled with the Gumbel-softmax trick [9]] In
computation graphs where the
Section [] we show both analytically and empirically that it is challenging to optimize the parameters of these models, mainly due to insufficient gradient signals often caused by local minima and saturation. We propose two new methods for mitigating the causes of poor optimization behavior. First, we show in Section 2.1 that increasing the scale parameter /3 of the Gumbel noise perturbations improves the models' learning behavior. Increasing 3 increases the probability of escaping local minima during training. Second, we propose dropout residual connections for discrete-contiin contituuus computation graphs in Section 2.2 By randomly skipping some discrete distributions, we provide more informative gradients throughout the full computation graph. We show empirically for several complex dissrete-continun show that the resulting state the art approaches on several benchmark datasets.
models that the proposed methods are required for training. We also models generalize better and significantly outperform
2 Efficient Learning of Discrete-Continuoor Computation Graphs
Standard neural networks compose basic differentiable functions. These networks, therefore, can be fully described by directed acyclic graph (the computation graph) that determines the operations executed during the forward and backward passes. Schulman et al. [30 proposed stochastic computa- tion graphs as an extension of neural networks that combine deterministic and stochastic operations node in the computation graph can be either differentiable function or probability distribution. stochastic node is typically a random variable with a parameterized probability distribution Po(z) with parameters O. Suppose that f is smooth function (such the loss function of learning problem), then the gradient of J at the stochastic node is Vo [xx) Kingma and Welling [[S] proposed the reparameterization trick to overcome the problem of computing gradients with respect to the parameters of a distribution. The idea is to find function g and distribution such that one

---

Scale
0.0
1.0
0.5
1:0
gradie
-0.2
Scale
softmaxle
0.9
Scale
2.0
-0.4
0
0=5.
8 4.0.0
-0.6
2.0
2.0
4.0
4,0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000 3000
4000
5000
Step
Figure 2: The values of the parameters, their gradients, and the softmax probabilities for various values 3 (the scale parameter) for constant Gumbel-Softmax parameter т. Early during training, lower scale parameter (relative 7) works well and has less variance. Once the probabilities saturate, however, we can only continue to obtain sufficient gradient signal for larger values of B.
can replace x ~ po( by =(2, with ~ p(z). If this is possible, then one can write
VolExrpa(s) Es-p(=) Vott
0)) R Ves(o(z,,)) 0)) with 2, p(z).
(1)
With this paper address the problem of training stochastic computation graphs where the stochas- tic nodes are based on categorical (discrete) distributions approximated using stochastic softmax tricks [14 26]. More specifically, we consider
components modeling categorical variable X with k possible values 21,..., Zh. Each of the is the one-hot encoding of the category i.
We consider discrete-comtinuous function Ra IR with v f(u) defined follows
0 gw
(2a)
"
p(z; 0 ~ 7
p(z;;0) exp(O;)
(2b)
Σj=1 exp((,)
2 p(= 0)
(2c)
Figure 3:
Illustration of generic discrete- continuous component (function f).
v w= haw (z)
(2d)
We assume that the functions and (parameterized by w and w') are expressed using differentiable neural network components. In the majority of cases, h defined as =' for learnable matrix w'. mapping object z; to its learned vector representation. Figure illustrates generic discrete- continuous component.
Let Gumbel(O, B) be the Gumbel distribution with location and scale 3. Using the Gumbel-max trick, we can sample p(2; 0) follows:
with
arg max [0 + cz] where Ej ~ Gumbel 1).
jE(I....k)
The Gunbeelsoftmax trick is a relaxation of the Gumbel-max trick (relaxing the argmax into softmax with scaling parameter that allows
perform standard backpropagation:
2, exp ((0, + e.)r))
where 4 ~ Gumbel 1).
(4)
exp (0,, <5)/T)
Hence, instead of sampling discrete z, the Gumbel-Softmax trick computes a relaxed its place.
We are concerned with the analysis of the behavior the Gumbel-softma trick more complex stochastic computation graphs. In these computation graphs, multiple sequential softmax components occur on a single execution path. Throughout the remainder of the paper, we assume that during training, we the Softmax-tric using the Gumbel-max trick. Depending on the use case, we also sometimes compute the argmax at test time instead of samples from the distribution.
outlined above, while at test time, sample discretely
Let us first take a look at dv/ /θu, that is, the partial derivative of the output of the discrete-contimuous function with respect to its input. Using the chain rule (and with a slight abuse of notation to
For the sake of simplicity assume the same input and output dimensions.

---

1/4
0150
0101550
Acounty
Acquracy
2.000
60022000
0.75
0003073
TauAnn
TM
Base
TauAnnn
240.9
Step
Figure 4: A comparison of the learning behavior of different annealing schemes for the MNIST addition task, averaged over 8 trained models. Left: The base model (Base), with constant parameters T = and = 1, cannot achieve an accuracy of more than 0.7. Increasing 3 during training (TM) achieves better results than annealing (TauAnn). Middle: Two training runs for TM and TauAnn, respectively, where an accuracy of 0.9 was reached the latest. The accuracy curves of TM jump to the next plateau earlier and more consistently. The dotted grey line depicts the moment when run of TM achieves an accuracy of ~ 1.0. Right: For the same run, 2j and the absolute value of aL/00, are plotted. The downstream gradient signal for 0j is largely neutralized by small value of æj until the corresponding category is used and its average probability abruptly reaches 0.1.
simplify the presentation), we write 0v/Ou = (0v/02) (0z/00) (a0) By assumption dv/ 00 and 00 /u exist and are efficiently computable. Using the derivative of the softmax function, we have
0z. 00,
2,(1 2x)
i
(5)
2,2
i
We have empirically observed that during training, there is tendency of the models to not utilize all existing categories of the distribution and to wrongfully map different input representations to the same category. For instance, for the MNIST addition problem, all encoded images of two digits are mapped to the same category the distribution. unused categories upper bound on the gradient parameter 0, of the categorical distribution with respect to a loss function L. First, have that
these cases, the gradients of the parameters the vanishingly small. In order to analyze this behavior more closely, we derive an
((
8z 00,
Σ
Oz: ae,
2
1
2
Σ (
2
2
(6)
where I| is the Frobenius norm. The full derivation of this inequality can be found in the Appendix. Since the Frobenius norm is sub- cub-mutiplicative matrix norm, we can write
0L aL Oz 00,
(
dz 30,
v2
(0))
S
2j
Oz 00,
where the last incquality follows from eq.
16 Hence, small value of of (and consequently 2j) leads to a small gradient for 0;, As a consequence, such models are more prone to fall into poor minima during training. Figure iillustrates an example of such situation, which we encountered in practice. The small value of a specific 2, leads to a small gradient at f, (right) as well as to suboptimal plateauing of the accuracy curve (middle). In other words, the downstream gradients information for 0g is neutralized by a small zj. Note that this issue only occurs when L is more complex and not the categorical-cross entropy loss.
Another related problem is saturated probabilities in sequentially connected probabilistic components. Similar to the problem of sigmoid activation units in deep neural networks, which have been replaced by ReLUs and similar modern activation functions, saturated probabilities lead to vanishing gradients. Indeed, we observe in several experiments that the Gumbel-sofune distributions saturate and, as result, that sufficient gradient information cannot reach the parameters of upstream neural network components.
We propose two strategies to mitigate the vanishing gradient behavior in complex discrete-continuous computation graphs. First, we analyze the interplay between the temperature parameter and the scale parameter of the Gumbel distribution. We make the subtle difference between these parameters explicit, which may also be of interest for improving stochastic softmax tricks (26). By increasing

---

T-4.0 8- 1.0
T- 2.0 8-1.0
T-1.5 8-1.0
1.0 B- 1.0
T- 0.5 B-
T- 0.1 8-1.0
r=2.0 0=0.1
T= 2.0 8=5.5
T= T=2.0 2.0 #=100
2.0 B= 1.5
r= 2.0 B 2=
T= 2.0 8=
A $
Figure 5:
Two annealing schemes for the Gumbel-softnax trick for 0 (2, 0.5, 1). Top: The standard
procedure of annealing the softmax temperature as introduced in [9 (see eq. (4) Bottom: Starting with 3 = 0 and increasing it during training, the samples are increasingly more discrete but, compared the standard annealing approach + 0), more uniformly distributed over the corners of the probability simplex.
3 relative to r while keeping fixed, we increase the probability a gradient flow in the case of saturation. We show that the larger 3. the more uniformly distributed (over the classes) are the discrete samples from the distribution, increasing the chance to escape poor minima. Second, we introduce DROPRES connections allowing us lower-bound the gradients expectation and leading to more direct gradient flow to parameters of upstream model components.
2.1 TEMPMATCH: Temperature Matching
We explore the behavior two interdependent parameters: the Gumbel-softmax temperature T and the Gumbel scale parameter 3. First, have the temperature parameter from the Gumbel-softman trick (see eq. (4). The purpose of this parameter is to make the output of the softmax function more or less discrete, that is, we can adjust the scale parameter 33 the Gumbel distribution. For scale parameter 3, we sample the noise perturbation iid e; Gumbel(0, 3). If we use the Gumbel-max trick of eq. (3) implicitly generate samples from the distribution p(z;; 0) = exp(0:/ /)/ Σ5= (xp(0,/3), Increasing the scale parameter 3, therefore, makes the Gumbel-max samples more uniform and less concentrated. When using the Gumbel-softmax trick instead of the Gumbel-max trick, we obtain samples that are more uniformly distributed over the categories and more discrete. Indeed, in the limit of 3 - oc, obtain discrete samples uniformly distributed over the categories, independent of the logits. For + O, we obtain the standard softmax function. F iilustratts the impact the two parameters on the Gumbel-softmax distribution. Now, the problem of insufficient gradients caused by local minima can be mitigated by increasing the scale parameter 3 relative to the temperature parameter т. The annealing schedule follow is the inverse exponential annealing = r(1 e-1r) for some > 0. Increasing & increases the probability of generating samples whose maximum probabilities more uniformly distributed over the categories. This has desired effects. First, samples drawn during training have a higher probability of counteracting poor minima caused at an earlier stage of the training. Figure 4 (left) shows that higher values for allow the model to find its way of poor minima. A higher value of 3 makes the model more likely to utilize categories with small Os, allowing the model to obtain improved minima (Figure 4 (middle)). Second, gradients propagate to parameters of the upstream components even after downstream components are saturating. To illustrate the second effect, we conducted the toy experiment depicted in Figure The model here computes Softmax I + c) with parameters 0T = (01, 02) and с Gumbel(0, 3). The learning problem is defined through a cross-entropy loss between the output probabilities and constant target vector (1.0, T. We observe that early during training, lower values for 3 work well and exhibit less variance than higher values. However, once the probabilities and, therefore, the gradients start to saturate, higher value for enables the continued training of the model. While the example artificial, it supposed to show that larger values of the scale parameter 33, sustain gradients for upstream components even if downstream components have saturated.
or less concentrated on of the categories. Second, and this is new insight,
5

---

Epoch
Epoch 20
Epoch 80
Mean gradient
0.004
0.003
Residuals
None
Dropout
0.001
0.000
GNN round
GNN round
GNN round
Figure 6: The influence of dropout residuals on the mean absolute gradients for ListOps. The values on the x-axis are the GNN message-passian iterations: the higher the number, the closer to the loss function is the corresponding mitigates the vanishing gradient problem for components with greater distance to the loss.
component. Adding dropout residual connections
summary, increasing the Gumbel scale parameter 3 has positive effects on the training dynamics. We propose exponential increasing scheme relative to We show empirically that for more realistic learning problems choosing a higher value for /3 is beneficial, especially later during training.
2.2 DROPRES: Residual Dropout Connections
Residual (or shortcut) connections are crucial when training deep neural networks A standard residual connection for the
components we consider here would be achieved by replacing Equation (2d) with = u + h (z). By creating direct connection to the continuous input of the discrete distributions, the optimizing problem simplifies. Unfortunately, what sets our problem apart from other end-to-end learnable neural networks is that we want to completely remove the residual connections at test time to obtain pure discrete representations and, therefore, the desired generalization properties. To mitigate the problem of overfitting to the existence of residual connections while at the same time reducing vanishing gradients expectation, we propose DROPRES connections. These are residual connections sampled from Bernoulli distribution with parameter a. We replace Equation (2d) by
{ u + (2) v v={ how (2)
with probability 1
with probability
that is, drop out the residual connection with probability a. With probability (1 (1- ai) obtain a shortcul connection between the neural components, effectively bypassing the Gumbel-softmax, With probability the training path goes exclusively through the Gumbell-oftmax distribution. The expectation of the gradients, taken over a Bernoulli random variable with parameter o, is now
do au
do Oz 00
= (1 a) 1 Oz 00 Ou
We propose simple linear scheme to modify from 0 1 during training. Hence, the model obtains a stronger gradient signal for upstream models expectation in the early phase of training.
To illustrate the impact of dropout residual connections, we analyzed the gradients of message passing steps for the ListOps problem (see experiments) in a continuuss computation graph of the type depicted in Figure ( As we can observe in Figure [ if we do not use dropout residual connections, the greater the distance a discrete-continuous operation to the loss function (the lower the value on the x-axis), the smaller is the mean absolute value of the gradients reaching it. This illustrates the vanishing gradient problem. When using dropout residual connections, on the other hand, the mean absolute values of the gradients do not vanish proportional to their distance to the loss and are more evenly distributed.
To summarize, optimizing the parameters of models comprised of discrete distributions and continu- ous neural network components is challenging, mainly due to local minima and vanishing gradients. By setting the Gumbel scale parameter 3 = as well as the dropout probability of the residual connections a = at the beginning of training, we obtain a continuous and deterministic relaxation of the target model. Increasing the DROPRES parameter a over time makes the model increasingly use the discrete stochastic nodes. Increasing the Gumbel scale parameter 3 allows the model to escape local minima and small gradients.

---

3 Related Work
Various gradient estimators have been proposed for the problem of backpropagating through stochastic nodes. The REINFORCE. algorithm [40] utilizes the log derivative trick. The Straight-Throun estimator 22 back-propagates through hard samples by replacing the threshold function by the identity in the backward pass. More recent approaches are based reparameterization tricks [1] that enable the gradient computation by removing the dependence of the density on the input parameters. Maddison et al. [19] and Jang et al. [4 propose the Gumbel-Softmax trick, a continuous (but biased) reparameterizztioo trick for categorical distributions. Tucker et al. [88 and Grathwohl et al. [9] introduce parameterized control variates to lower the variance for these gradient estimators. They show that these estimators can, in theory, also be generalized to chains of stochastic nodes but do not discuss any learning behavior. Shayer et al. [32 modify the local reparameterization trick
16 to improve the training of discrete gradient estimators. Recently, Paulus et al. [26) proposed a framework that generalizes the Gumbel-Softmax trick for various discrete probability distributions. Our aim is not to improve gradient estimators for a single stochastic node but to analyze the hehavior of discrete-continuous computation graphs with multiple sequential discrete components. Pervez et al. 277 utilize harmonic analysis for Boolean functions to control the bias and variance for gradient estimates for Boolean latent variable models. We focus on efficiently training models with gradient estimates for categorical distributions, Huang et al. [13] propose to drop out full layers of ResNet to obtain models of stochastic depth during training. In contrast, we use dropout on the residual connections to approximate the associated discrete model. SparseMAP 24 is an approach to structured prediction and latent variables, replacing the exponential distribution (specifically, the softmax) with sparser distribution. LP-SSpareMMAP (23] is a continuation of SparseMAP using relaxation of the optimization problem rather than a MAP oracle. In addition, the idea to enforce more sparsity can also be exploited for efficient marginal inference in latent variable models [5] which can then be used to compute stochastic gradients.
The respective advantages of neural and symbolic methods have led to various attempts to combine these two. Early work like Towell and Shavlik (35) and Hölldobler et al. [2] have explored the possibilities of translating propositional logic programs into neural networks. Evans and Grefen- stette (8 and Dong et al. 7 build upon these ideas by utilizing the framework of inductive logic programming to learn rules from examples. Kool et al. [18] combine an attention model and the REINFORCE algorithm with greedy baseline to learn heuristics for combinatorial optimization problems. A different direction of work focuses on knowledge base reasoning. Rocktäschel and Riedel 28 combine Prolog's backward chaining algorithm with a differentiable unification whereas Yang et al. 41 learn to compose the inference task into differentiable operations. Serafini and Garcez 31 propose differentiable first-order logic language to differentiate deductive reasoning, and Manhaeve et al. (20) extend the probabilistic logical language ProbLog to interact with neural networks. Also related is prior work on automatic program induction. Andreas et al. [ combine neural networks with the compositional linguistic structure of questions, Mao et al. (21) additionally learn visual concepts and Tsamoura and Michael (37] introduce the concept of abduction to learn the inputs of symbolic program. All prior work integrates neural modules into classical or fuzzy logic or separates the neural and the symbolic parts in other ways. In contrast, we propose method that allows training complex discrete-continuon computation graphs with multiple sequential discrete distributions.
Experiments
The aim of the experiments is threefold. First, we want to analyze the behavior of complex stochastic computations graphs arising in typical application domains such as multi-hop reasoning in knowledge graphs and unsupervised parsing of lists of operations. Second, we want to evaluate the behavior of the stochastic computation graphs when incorporating the proposed methods (dropout residual connections and temperature matching) to improve the vanishing gradient problem. Third, we want to compare the resulting siscrete-continuous models with state of the art models which do not have stochastic components. Here, we are especially interested analyzing the generalization (extrapolation) behavior of the models under consideration. The implementations are in Py and can be found at GeForce RTX 2080 Ti GPU.
All experiments were run on a
7

---

Table 1: The results for unsupervised parsing on ListOps. Results taken from Paulus et al. [26 are marked with an asterisk (*) and are based on different (unpublished) dataset.
Task acc. (extrapolation)
Model
Task acc.
Edge pree.
Inter. acc.
d d=8
10
Und,*
91.2 + 1.8*
33.1 + 2.9* 75.0 + 7.0"
n.a.
Arb.* *
95.0 + 3.0"
n.a.
n.a.
LSTM
91.5 1 0.3
83.7 2.0
76.9 3.7
Arb., T = 2
96.8 + 0.3
77.4 + 1.8
84.3
1.4
75.4 + 1.7
Ours, T =
96.1 10.1+00 0.4
82.3 1.1
70.9 + 1.2
92.6 + 1.1
86.9 + 4.9
Ours, T 2
96.3 + 0.5
2.2
62.9 + 2.4
92.7 1.3
88.7 3.3
Arb. (GT) Ours (GT)
98.7 + 0.1
100.0
86.4
+ 1.1
71.0 = 2.1
99.8 1088100 0.1
100.0
99.9 + 0.0
99.8 + 0.1
99.9 = 0.1
Table 2:
The ablation study for the ListOps task. Temperature matching improves the training behavior of the model. The addition of dropout residual connections is erucial for efficient learning.
Task acc. (extrapolation)
Model
Task acc.
Edge prec.
Inter. acc.
d = 8
d= 10
Ours, 7 =1
96.1 + 0.4
82.3 + 1.1
70.9 + 1.2
92.6 + 1.1
86.9 + 4.9
()) TM
95.1 + 1.7
76.9 + 12.1
66.3 + 9.2
90.5 2.4
83.7 + 4.0
DR
74.8 11.5
43.9 + 30.3
29.4 = 18.6
9.3
66.7 : 8.8
()) DR, TM
56.2 + 5.0
15.1 = 7.7
12.7 2 3.2
53.1 1.1:00 6.6
49.2 + 7.1
Unsupervised Parsing on ListOps The Listops dataset contains sequences in prefix arithmetic syntax such as max[ 2 9 min[ 4 4 and its unique numerical solutions (here: 9) [22). Prior work adapted and used this dataset to evaluate the performance of stochastic softmax tricks [26] Following this prior work, we first encode the sequence into a directed acyclic graph and then run a graph neural network (GNN) on that graph to compute the solution. Since the resulting dataset was not published, we generated a dataset following their description. In addition to examples of depth 1 S d < 5, we further generated test examples of depth d = 8, 10 for the extrapolation experiments. More details on the dataset creation can be found in the supplementary material.
The arithmetic syntax of each training example induces directed rooted in-tree, from arborescence, which is the tree along which the message-passen passing steps of a GNN operate. We use the same bi-LSTM encoder as Paulus et al. (26) compute the logits of all possible edges and the same directed GNN architecture. The original model uses a directed version of Kirchoff's matrix-tree theorem as introduced by Koo et al. [17 induce the arborescence prior. We simplify this idea over all possible parent nodes. Here we make use
called
by taking, for each node, the Gumbel- of the fact that in an arborescece, each non-root node has exactly one parent. Note that this prior is slightly less strict than the arborescence prior. As in Paulus et al. (26). we exclude edges from the final closing bracket] since its edge assignment cannot be learned from the task objective.
In contrast to Paulus et al. (26). where only the edges of the latent graph are modeled with categorical variables, we also modelled the nodes of the latent graphs with discrete-continuous components (see Figure (bottom)). Let Num E IR 10x0dim be the embedding layer that maps the 10 numeral tokens their respective embeddings and let Pred dRdim + R be the clas- sification layer that maps the embedding æ E R of the final output to the logits of the 10 classes: Pred(x)= == Lin 10xdim (ReLU(Ling) dim x dim (x))) After each message- message-passing passing operation of the GNN, we obtain an embedding u E Rdim for each node. By choosing Iw = Pred, we obtain the 10 class logits 0 9w
Pred(u) 1210 for the numerals 0, ..., O.... Using the Gumbel- softmax trick with logits 0, scale B, and temperature т, we obtain E 10 We then compute h (2) = T Num. We apply this on all intermediate node embeddings simultancously and repeat this after each of the first 4 out of 5 message-passee ppssiin rounds of the GNN. We use dropout residual connection with increasing dropout probability a. At test time, the model performs discrete and interpretable operations. Figure deppict an example stochastic computation graph for the task. We used the same LSTM as in Paulus al. (26) and a re-implemented version of their model with our customized prior. For ablation purposes, run all our models with and without dropout residuals
8

---

Table 3:
Results for the path query bench- Table 4:
Results for the extrapolation bench- for the path query dataset [T].
mark [0].
mark 39
WordNet
Freebase
WordNet
Freebase
Model
MQ
H@ 10
MQ
H@ 10
Paths
Model
MQ
H@10
MQ
@@ 10
Bilinear-C
89.4
51.3
83.5
42.1
CoKE
93.6
65.5
93.1
71.2
S 4
DistMult-C
90.4
31.1
84.8
38.6
Ours
94.3
64.7
88.8
69.4
TransE-C
93.3
43.5
88.0
50.5
S 3
CoKE
92.6
65.0
90.6
64.6
Ours
93.4 4
65.0
88.3
68.5
Path-RNN
98.9
CoKE
90.8
49.5
89.4
59.5
ROP
90.7
56.7
2
Ours
90.2
62.0
87.6
68.6
CoKE
94.2
67.9
95.0
77.7
CoKE
73.5
16.7
72.7
37.3
Ours, r = 4 94.4
64.3
89.6
68.7
Ours
81.1
52.2
82.1
63.9
and temperature matching, respectively. All models are run for 100 epochs with a learning rate of 0.005 and we select 21, We keep all other hyperparameters We evaluate the task accuracy, the edge precision (using the correct graphs), the task accuracy on intermediate operations, and two extrapolation tasks. For the evaluation of the extrapolation task with depth d, we run the GNN d rounds instead of 5.
in Paulus al. (26).
Table I compares the proposed models with state of the art approaches. The best performing models are those with discrete-continuond components trained with dropout residuals and temperature matching. While our models do not improve the task accuracy itself, they exceed state of the art methods on all other metrics. Most noticeable is the improved generalization behavior: our models extrapolate much better to expressions with depth not during training. This becomes more visible when trained on the ground truth graph directly (instead of on the sequence) where the generalization accuracy is close 100%. Our model also achieves the best performance on the graph structure recovery with a precision of 82.3%. The mismatch between the task accuracy and the edge precision mainly due examples for which the correct latent graph is not required to solve the problem, e.g., for the task max 2 max 4 7 it does not matter whether the 7 points to the first or the second max[. Note that our method performs discrete operations test time which are therefore verifiable and more interpretable. Our model generates intermediate representations that can be evaluated in exactly the same way as the final outputs, see Column Inter. acc. in Table [
The ablation study in Table highhightss the challenge of learning discrete-continuous computation graphs with multiple sequential Gumbel-softmax components. It is entirely impossible train the model without dropout residuals and temperature matching. The analysis in Section Figure F reveal the reason. Sequential discrete distributions in the computation graph cause vanishing gradients. Temperature matching stabilizes learning somewhat but to avoid vanishing gradients entirely, the use of dropout residuals is necessary.
Multi-Hop Reasoning over Knowledge Graphs Here consider the problem of answering multi-hop (path) queries in knowledge graphs (KGs) [10). The objective is to find the correct object given the subject and a sequence of relations (the path) without knowing the intermediate entities.
We evaluate various approaches on the standard benchmarks for path queries [0]
We model each application of relation type an (intermediate) entity a discrete-continuous component that discretely maps to one of the KG entities. The discrete distribution z; 0), therefore, has as many categories as there are entities the KG. Let score r, o) R3xdim -> R be scoring function that maps the embeddings of triple to the logit (unnormalized probability) of the triple being true. Given the current (intermediate) entity embedding a and relation type embedding r we compute the logits 0 for all possible entities. Hence, the function & computes the logits for all possible entities. Using the Gumbel-softmax trick with parameter 0, scale B3, and temperature т, we obtain the sample E R The function h now computes x T E where E be the matrix whose rows are the entity embeddings. We use dropout residual connection between the input and output vectors of the
component during training. Note that we do not increase the number of
Revised results in Table and Table because of a flawed evlauation protocal in the original results.

---

Table A comparison of our proposed method with DeepProbLog and a CNN baseline.
Model
Task Acc.
Interpretable
Discrete (test time)
Learned Structure
Baseline
89.14 = 1.22
No
No
No
DeepProbLog
98.17 0.20
Yes
Yes
No
Ours, 8 98.10 0.15
Yes
Yes
Yes
parameters compared to the base scoring models. We use ComplEx (36 as the scoring function. We choose a dimension of 256, a batch size of 512 and learning rate of 0.001. We further train IvsAll with the cross-entropy loss for 200 epochs and with temperature of r = 4. We compare our best performing model on the original dataset from Guu et al. [I0] with their proposed evaluation protocol
and baselines against RNN models, Path-RNN (6. ROP [42] and the state-of-the-art transformer model CoKE [39]
for the standard task as well as in an extrapolation setting.
Table 3 lists the results for the proposed model in comparison to several baselines. Our proposed model performs significantly better than the baselines KGC models which do not have stochastic components (labeled with -C). On WordNet, our model can even keep pace with the state of the art transformer model CoKE.
Similar to the results on ListOps, the discrete-continuout models have the strongest generalization performance (see Table 4 Here, we the same test set before but reduce the length of the paths seen during training. Even when trained only on paths of length S 3. our model performs better on the path query task than most baselines trained on paths of all lengths.
End-to-End Learning of MNIST Addition The MNIST addition task addresses the learning problem of simultaneously (i) recognizing digits from images and (ii) performing the addition operation on the digit's numerical values [20]. The dataset was introduced with DeepProbLog. a system that combines neural elements with a probabilistic logic programming language. In DeepProbLog the program that adds two numbers is provided. We adopt their approach of using convolutional neural network (CNN) to encode the MNIST images. Contrary to their approach, we learn the operation without any prior knowledge about adding numbers in data-driven manner.
dim x2 dim (ReLU(Ling) 2im x2 dim ([æ; y)))) The addition operation is modeled as Add(2, y) = Ling B
and we use a single linear layer to compute the logits. Given u E Rddim the encoding of the digit obained from the CNN, the function g computes 0 = Lin(u), that is, the 19 logits for all possible numbers. Using the Gumbel-softmax trick with parameter 0, scale 3, and temperature т, we obtain I P e Ig where P is the weight matrix of Lin. We
E 19. The function now computes discretize the two outputs of the CNN, execute the addition layer and use the single linear layer P for the classification. This results into model that has in total fewer parameters than the baseline in 220 We use learning rate of 0.0001 and temperature of r == 8.
Table lists lists the results on this task. Our proposed model, even though trained without a given logic program and in a data-driven manner, achieves the same accuracy as DeepProbLog, Since we use hard sampling during test time, we can discretely read out the outputs of the internal CNN layer.
Conclusion
We introduced two new strategies overcome the challenges of learning complex discrete-continuous computation graphs. Such graphs can be effective links hetween neural and symbolic methods but suffer from vanishing gradients and poor local minima. We propose dropout residual connections and a new temperature matching schema for mitigating these problems. Note that even with our proposed methods, learning models with discrete nodes is still considerably more challenging than standard neural networks. Furthermore, it is not clear how to use dropout residuals for structural stochastic nodes or in such setups in which only hard decisions are feasible. In experiments across diverse set discrete-contnnoon
of domains, we demonstrate that our methods enable the training of complex models, which could not learned before. These models beat prior state of the art models without discrete components on several benchmarks and show remarkable gain in generalization behavior. Promising future work includes applying the proposed methods to reinforcement learning problems, effectively sidestepping the problem of using the high variance score function estimator.
10

---

6 Acknowledgements
Many thanks to the NeurIPS reviewers for their highly constructive feedback and suggestions. Their input has helped to improve the paper tremendously. Special thanks to reviewer RKc5 who provided an impressive review and engaged in long and fruitful discussion with us.
References
[1] J. Andreas, M. Rohrbach, T. Darrell, and D. Klein. Neural module networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
[2] Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1308.3432, 2013.
[3] A. Bordes, N. Usumier, A. Garcia-Duran, J. Weston, and O. Yakhnenko. Translating embeddings
for modeling multi-relational data. Advances in Neural Information Processing Systems 26: 27th Annual Conference on Neural Information Processing Systems 2013. Proceedings of meeting Lake Tahoe, Nevada, United Stutes, held December 5-8, pages 1-9, 2013.
[4] Chen, L. Song, M. Wainwright, and M. Jordan. Learning to explain: An information thhooretc perspective on model interpretation. In International Conference on Machine Learning, pages 883-892. PMLR, 2018.
[5] G. M. Correia, V. Niculae, W. Aziz, and F. Martins. Efficient marginalization of discrete and structured latent variables via sparsity. Advances in Neural Information Processing Systems, 2020.
[6] R. Das, A. Neclakantan, D. Belanger, and McCallum. Chains of reasoming over entities, relations, and text using recurrent neural networks. In Proceedings of the 15th Conference of the the pages 132-141, Valencia, Spain, Apr. 2017.
European Chapter the Association for Computational Linguistics: Volume Long Papers,
[7] H. Dong, J. Mao, T. Lin, C. Wang, L. Li, and D. Zhou. Neural logic machines. 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019.
[8] R. Evans and Grefenstette. Learning explanatory rules from noisy data. Journal of Artificial Intelligence Research, 61:1-64, 2018.
[9] W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - 30-May 3, 2018.
[10] K. Guu, J. Miller, and P. Liang. Traversing knowledge graphs in vector space. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, 2015.
[11] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778, 2016.
[12] S. Hölldobler, Y. Kalinke, and H.-P. Störr. Approximating the semanties of logic programs by recurrent neural networks. Applied Intelligence, 11(1):45-58, 1999.
[13] G. Huang, Y. Sun, Z. Liu, D. Sedra, and K. Q. Weinberger. Deep networks with stochastic depth. European conference computer vision, pages 646-661, 2016.
[14] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017.
[15] D. P. Kingma and M. Welling, Auto-encoding variational bayes. 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014.
[16] D. P. Kingma, T. Salimans, and M. Welling. Variational dropout and the local reparameterization trick. Advances in neural information processing systems, 28:2575-2583, 2015.

---

[17] T. Koo, A. Globerson, X. Carreras, and M. Collins. Structured prediction models via the matrix- tree theorem. In Proceedings of the 2007 Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning SEMNLP-CoNNLL pages 141-150, 2007.
[18) W. Kool, H. Van Hoof, and M. Welling. Attention, learn to solve routing problems! 7th International Conference on Learning Representutions, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019.
[19] C. J. Maddison, A. Mnih, and Y. W. Teh. The concrete distribution: continuous relaxation of discrete random variables. 5th International Conference Learning Representations, ICLR 2017, Toulon. France, April 24-26, 2017.
[20] R. Manhaeve, S. Dumancic, A. Kimmig, T. Demeester, and L. De Raedt. Deepproblog: Neural probabilistic logic programming. Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, Montréal, Canada, December 3-8, 2018, 31:3749-3755 2018.
[21] Mao, C. Gan, P. Kohli, J. B. Tenenbaum, and J. Wu. The neuro-symbolic concept learner: Interpreting scenes, words, and sentences from natural supervision. 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019.
[22] N. Nangia and Bowman. Listops: A diagnostic dataset for latent tree learning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Student Research Workshop, pages 92-99, 2018.
[23] V. Niculae and A. F. T. Martins. Lp-sparsemap: Diffcrentiable relaxed optimization for sparse structured prediction. In ICML, 2020.
[24] V. Niculae, A. F. T. Martins, M. Blondel, and C. Cardie. Sparsemap: Differentiable sparse structured inference. In ICML, 2018.
[25] E. Parisotto, A.-r Mohamed, R. Singh, L. Li, D. Zhou, and P. Kohli. Neuro-symbolic program (tl::1611.01555 2016.
synthesis. arXiv preprint
[26] M. Paulus, D. Choi, D. Tarlow, A. Krause, and C. Maddison. Gradient estimation with stochastic softmax tricks. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, December 6-12, 2020.
[27] A. Pervez, Cohen, and E. Gavves. Low bias low variance gradient estimates for boolean stochastic networks. Proceedings of the 37th International Conference on Machine Learning. 119:7632-7640, 13-18 Jul 2020.
[28] Rocktäschel and S. Riedel. End-to-end differentiable proving. Advances in Neural Informa- tion Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, Long Beach, CA, USA, December 4-9, 2017.
[29] D. Ruffinelli, S. Broscheit, and R. Gemulla. You [can] teach an old dog new tricks! on training knowledge graph embeddings. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020.
[30] J. Schulman, N. Heess, T. Weber, and P. Abbeel. Gradient estimation using stochastic com- putation graphs Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, Montreal, Quebec, Canada, December 7-12, 2015.
[31] L. Serafini and A. d. Garcez. Logic tensor networks: Deep learning and logical reasoning from data and knowledge. arXiv preprint 160.0466. 2016.
[32] O. Shayer, D. Levi, and E. Fetaya. Learning discrete weights using the local reparameterization trick. 6th International Conference on Learning Representations, ICLR 2018 Vancouver, BC, Canada, April 30- May 3, 2018
[33] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
12

---

[34] K. Toutanova and D. Chen. Observed versus latent features for knowledge base and text inference. In Proceedings of the 3rd Workshop on Continuous Vector Space Models and their Compositionality. pages 57-66, Beijing, China, July 2015. Association for Computational Linguistics.
[35] G. G. Towell and W. Shavlik. Knowledge-baid based artificial neural networks. Artificial intelli- gence,
19-166, 1994.
[36] Trouillon, J. Welbl, S. Riedel, É. Gaussier, and G. Bouchard. Complex embeddings simple link prediction. In Proceedings of the 33nd International Conference Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, pages 2071-2088 PMLR, 2016.
[37] E. Tsamoura and L. Michael. Neural-symbolic integration: compositional perspective. arXiv preprint arXiv:2010.I 1926, 2020.
[38] G. Tucker, A. Mnih, C. J. Maddison, D. Lawson, and Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems 30: Anmual Conference on Neural Information Processing Systems 2017, Long Beach, CA, USA, December 4-9, 2017.
[39] Q. Wang, P. Huang, H. Wang, S. Dai, W. Jiang, J. Liu, Y. Lyu, Y. Zhu, and H. Wu. Coke: Contextualized knowledge graph embedding. arXiv preprint arXiv: TrXXv:9911.02168 2019.
[40] R. J. Williams. Simple statistical gradient- followwig algorithms for connectionist reinforcement learning. Machine learning, 8(3-4)2229-25 256, 1992.
[41] Yang. Z. Yang, and W. W. Cohen. Differentiable learning of logical rules for knowledge base reasoning. Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, Long Beach, CA, USA, December 4-9, 2017.
[42] W. Yin, Yaghoobzadeh, and Schütze. Recurrent one-hop predictions for reasoning over knowledge graphs. arXiv preprint arXiv:1806.04522 2018.
13


Reviews:
[{"review_id": "yG1eknIVU_", "reviewer": "Reviewer_Uy2p", "summary": "Summary.\n\n- This paper introduces two techniques to help with learning neural network architectures that embed two or more layers of categorical probability distributions (parameters of those distributions are learnt!). \n- The first method is Temp-Match: The suggestion is to separately tune the variance of the Gumbel noise and the temperature parameter. \n- The second method are dropout-residual connections. The authors add residual connections to connect the input to the output stochastically for a given discrete-continuous component (Figure 2). \n\n", "questions": "", "limitations": "", "rating": 4, "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": "**Strengths**\n- Clarity. The paper is well-written overall and I like the structure.\n- Significance. Studying gradient estimation for discrete-continuous computation graphs is interesting.\n- Originality. I am not an expert on dropout and residual connections, but I have not seen dropout-residual connections as presented in this paper before. \n\n**Weaknesses**\n- Originality. I feel the relationship between the temperature parameter and the Gumbel noise variance is known widely, even though I don’t have a concrete reference on hand right now. \n- Quality. I do not find the analysis of the vanishing gradient problem when using Gumbel-Softmax convincing.\n- Quality, Significance. From the argument and evidence presented, I wouldn’t use Temp-Match.\n- Quality. I believe the experiments need to be revised to make a stronger case for the methods presented (e.g. more complex models with more stochastic layers, see below).\n\n**Decision**\n- Ok, but not good enough. \nAs is, unfortunately, I do not think the paper is ready for acceptance. \nI would like to make some suggestions for improvement below. \n\n**Suggestions**\n\n*As is, I do not find the analysis of the vanishing gradient problem when using Gumbel-Softmax convincing.*\n- The authors say that gradients may vanish whenever the probability saturates. \nBut when the probability saturates, the categorical distribution becomes degenerate, the stochasticity vanishes, the output is deterministic. \nSo my question: Isn’t this a setting in which I may not want to model a latent categorical distribution in the first place, and could use a deterministic function instead? \nMy issue with the toy example presented in Figure 3 is similar. \nThis toy example is presenting an edge case, where the ideal distribution is degenerate, since target is $(1, 0)$. \nWhy not present the same case as the REBAR paper, which seems very similar in set-up but didn’t pick an edge case. To what extent is the problem present there?\n\n*From the argument and evidence presented, I probably wouldn’t use Temp-Match.*\n- I am missing a discussion on the theoretical side, that highlights that for every $(\\tau, \\beta)$ pair, there is a pair $(\\alpha, \\tau’)$ with $\\beta = 1$ that gives rise to the same Gumbel-Softmax distribution as $(\\tau, \\beta)$. Here, $\\alpha$ is a parameter that scales the logits. The reason for this is basically, that $\\tau$ inversely scales both logits and Gumbels. To decouple scaling both logits and Gumbels together, one can introduce either an additional parameter to rescale the Gumbels (this is $\\beta$) or an additional parameter to rescale the logits (I used $\\alpha$ for that above). In practice, since the logits are outputs from a neural network with learnable parameters, the $\\alpha$ is implicit and is learnt. My questions: Why is Gumbel scaling better than logit scaling? Why should $\\alpha$ (or $\\beta$) not be learnt? \n- Further, Figure 3 seems to present some evidence for the benefits of tuning $\\beta$. But $\\tau$ is held fixed. Isn’t the argument that is required compare “tuning $\\tau$” vs “tuning $\\tau$ and $\\beta$” (both with a schedule ideally). There should also be a loss curve, to make the case that “\\tuning $\\tau$ and $\\beta$” vs “tuning $\\tau$” yields significant improvements that justify an extra hyperparameter. \n-Finally, Temp-Match is evaluated in an ablation study. I like that. But the standard deviations are so high, that I do not see any evidence for benefits from TempMatch. Or am I misinterpreting Table 2? \n\n*The focus of this paper is on complex networks with multiple discrete components, but the experimental section is underdelivering.* \n- To make a stronger case for the setting and methods of the paper, I would like to see settings and models that are more ambitious and practical. As a concrete suggestion, I would recommend looking into the model Gumbel-Tree-LSTM from Choi et al. (2017). This model is much better tailored to the ListOps dataset than the GNN model considered by the authors. \nIt also requires much more stochastic decisions, namely #number of tokens - 1, rather than a fixed set, so I feel is more relevant experimentally to the methods presented. Other papers that cite ListOps also include this model or variants that I believe are generally viewed as more suitable for unsupervised parsing on ListOps. I do not understand why the authors considered a GNN architecture instead and modified the ListOps dataset to make the GNN practical. \n\n**Minors.**\n- Gumbel hard setting -> I think generally this is referred to as Gumbel straight-through. See for example, Jang et al. (2016) or Paulus et al. (2021).\n\n**References.**\n- Jang et al. (2016). Categorical reparameterizatoin with Gumbel-Softmax.\n- Choi et al. (2017). Learning to compose task-specific tree structures. \n- Paulus et al. (2021). Rao-Blackwellizing the Straight-Through Gumbel-Softmax estimator. \n\n\nEDIT. \nI read through the response of the authors, but I am keeping my score unchanged. \nUnfortunately, I am of the opinion that more work is required to improve the paper before it is ready for submission along the lines I suggested above.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "I don't stop a clock, but probably 4 hours."}, "extra_text": {}}, {"review_id": "J3s0-ZRbALI", "reviewer": "Reviewer_DtA3", "summary": "Motivated by the vanishing gradient problem of multiple discrete operations that leverage the Gumbel Softmax trick for training, this paper proposes two new tricks.  First, the scale parameter of the Gumbel distribution is annealed across training rather than annealing the temperature parameter as is common in most prior works.  Second, residual connections are stochastically dropped via another annealing strategy so that gradients are non-vanishing in the beginning but eventually the residual connections are removed. Multiple empirical experiments demonstrate that the method is better at extrapolating than the baselines.\n", "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": "**Summary of review**\nThis paper introduces two novel tricks for complex discrete-continuous networks and shows strong results for extrapolation. While some writing could be improved and more explanation added, the paper seems to make useful contributions.\n\n**Strengths:**\n- The paper connects the interplay between the scale parameter of the Gumbel distribution and the temperature parameter of the softmax operation. In particular, it seems that annealing the Gumbel $\\beta$ parameter may be better than annealing the temperature parameter $\\tau$.\n\n- The paper introduces a simple but interesting modification to residual connections that combines the idea of residual networks and dropout.\n\n- The empirical results demonstrate superior performance for extrapolation on multiple experiments.\n\n**Weaknesses:**\n- It seems that the annealing schedules for both proposed tricks is quite important. Could more discussion be given to the selection of annealing schedules?\n\n- The problem with overfitting that motivates DropRes is not clearly explained.  It seems to merely be an empirical observation.  This should be explained more carefully and any intiution or explanation should be provided.\n\n- Related to above, how does the Gumbel scale parameter directly affect things?  It doesn't seem to be connected theoretically in say equations 3-5.  Is this just an empirical observation and basic intuition that they are related?  Is there a more fundamental or mathematical relationship between these parameters? \n\n**Other comments or questions**\n- How do these two tricks work for a single discrete operator?  The paper seems to emphasize multiple times that a sequence of discrete operations is the key problem being tackled in the paper.  Thus, I'm wondering how much it makes a difference for the simpler cases.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "- The paper seems adequate.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.25"}, "extra_text": {}}, {"review_id": "FCg25xNvA7S", "reviewer": "Reviewer_UQHi", "summary": "This paper tackles the learning of a computation graph with possibly multiple discrete and continuous nodes. The authors argue the challenge for effective learning is the gradient vanishing in Gumbel-softmax in discrete nodes, and propose two techniques, annealing parameters in Gumbel-Softmax and introducing residual dropout connections, to deal with the challenge. Empirical experiments show the effectiveness of the proposed method. ", "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": "This work addresses the task of learning in a computation graph with possibly multiple discrete and continuous nodes. Particularly, the authors consider discrete nodes instantiated as Gumbel-softmax, and argue that the vanishing gradient on Gumbel-softmax is the challenge posed to effective learning that should be addressed. This overall motivation is an important one, yet is it not un-explored. For example, the balance between approximation quality and saturation/vanishing gradients has already been noted [1] and compared [2], and improvements have been proposed [2]. These existing works are not present in the current draft, and the authors are encouraged to include them to give a better overview of the topic and the motivation behind it. The authors also argue for interpretability as an advantage for the proposed tasks throughout the draft (e.g. L4, L38, L253, Table 5), without elaboration. It may be better to demonstrate interpretability though either examples and/or visualization in one of the experiments to give readers a better understanding.\n\nTo tackle the challenge, the authors proposed two techniques, namely TM (Temperature Matching and Inverse Scale Annealing) and DR (Residual Dropout Connections), detailed in Section 2.1 and 2.2 respectively. From my understanding, TM is a newly proposed method --- actually the original Gumbel-Softmax paper [1] already adapts annealing in the training. Also, DR, or stochastic residual, has been proposed for vision tasks [3], while it's application to Gumbel-Softmax would be novel. Another question regarding DR would be, since only Gumbel-Softmax poses the issue of vanishing gradients, why not just skip-connecting the $\\theta \\to z$ part? Neitherness, the authors are encouraged to discuss these related works and compare the proposed method with them, as they are closely related to the claimed contributions.\n\nThe experiments are concluded thoroughly and the code is available (great job!) Probably due to the design of the methods, there are some hyperparameters to choose such as annealing rate (L484 \"The temperatures are updated 10 times per epoch\"), and schedule of decaying $\\alpha$ that seem to be hard-coded in the code. The experiment section mentioned that hyperparameters are chosen following previous methods, but it may be more useful to conduct study on how they work to give readers more rationale behind these choices.\n\n\nOverall, The work is clearly written and well organized andI think this is an interesting work Future works may draw inspirations (empirical methods that mitigate the gradient vanishing problem in Gumbel-Softmax) proposed in this manuscript. However, it still needs some improvements as detailed above, as a few discussions with existing works are missing.\n\nReferences:\n\n* [1] Jang et al. Categorical Reparameterization with Gumbel-Softmax. ICLR ;17 https://openreview.net/forum?id=rkE3y85ee\n* [2] Shayer et al. Learning Discrete Weights Using the Local Reparameterization Trick. ICLR '18 https://openreview.net/forum?id=BySRH6CpW \n* [3] Huang et al. Deep Networks with Stochastic Depth. ECCV '16. https://arxiv.org/pdf/1603.09382v2.pdf \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}, {"review_id": "6kdekTqTGHO", "reviewer": "Reviewer_Thd9", "summary": "The authors consider the problem of learning in models that combine discrete and continuous operations through the use of stochastic sampling of categorical variables. In the case where the learning makes use of the Gumbel-Softmax reparametrization of the categorical distributions, the authors propose two modifications to the architecture of the network in order to improve training: 1) tempering the categorical distributions towards more uniformity in the later stages of training, and 2) including a stochastic pass-through (residual connection). The authors demonstrate the benefits of the proposed modifications on a couple of common tasks in the domain.", "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": "\nThe authors propose techniques which are somewhat unsurprising, but they demonstrate good results on proposed applications. Unfortunately, the authors do not provide direct comparisons between baselines and their suggested techniques, making it difficult to identify the source of the improvement. This paper could be greatly improved by more directly focusing on the problem of gradient propagation through stochastic neurons, and in particular it would be helpful for the authors to directly compare with commonly used methods (e.g. see [1]). Given the lack of direct evidence for the effectiveness of the proposed methods, I do not believe that this paper is quite ready. Please see more detailed comments below.\n\nOn the empirical results. Given the main claims in the abstract (“we show that it is challenging to optimize [...], mainly due to vanishing gradients. We then propose two new strategies to overcome these challenges [...]”), I believe that the main contribution of this paper is to tackle the fairly general problem of gradient propagation through stochastic neurons. In this light, the empirical section of the paper should be reworked to better support these claims, with more emphasis on ablation and comparative studies (e.g. as presented in table 2, but for a large range of models and problems), rather than introducing the author’s own variations for standard problems, where the source of the improvement is not obviously identified. Additionally, such comparisons should attempt to include other commonly used strategies to propagating gradients through stochastic neurons: in particular, the author’s proposal for DropRes has some parallels to the straight-through estimators. Similarly, figure 4 should be extended for different combinations of models, problems and gradient estimators.\n\nOn the bias of the proposed methods. Compared to estimators of the gradient based on the score function, the reparameterized gradients for discrete stochastic neurons are slightly biased. The proposed methods in this paper further amplify this bias, by 1) increasing $\\beta$ (effectively scaling down $\\theta$ and $\\tau$), and 2) considering an uncorrected residual connection. Given the presence of this bias, it would be interesting to compare these methods to standard methods (e.g. based on the score-function estimator a.k.a. REINFORCE and their variance-reduced counterparts) to better understand the impact of these methods, and e.g. whether the improvements are due to better gradient propagation and training behavior, or some implicit regularization due to the introduced bias.\n\n\n[1] Estimating or propagating gradients through stochastic neurons for conditional computation. Bengio, Leonard and Courville. ArXiv 2013\n\nEdit after author response\n\nI thank the authors for their response. I believe that ablation studies should be the primary empirical justification for claims advanced in this paper, and I strongly encourage the authors to consider expanding the ablation studies and reducing the number of comparison with custom-tuned models, which although interesting, do not necessarily provide evidence for the current claims as they do not clearly demonstrate the source of the improvement.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Please see main review for limitations.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposed two ways to optimize the computation graph that involves with multiple sequential discrete components, namely the dropout residual connection and the tuning of Gumbel noise. While the reviewers generally agree that the problem studied is an important one, there are several common concerns.
- The technical part of the paper, especially its novelty, the motivation and the rigorous analysis behind it.
- More convincing experiments, with better positioning among existing works and better design of the experiments.

After the authors' response, all the reviewers participated in the committee discussion, and the recommended decision was made based on the consensus of all the reviewers and AC. Despite that the paper cannot be accepted in the current form, the reviewers are optimistic about the paper and encourage the authors to further improve the paper based on reviewers' comments, especially on the motivation/analysis of the proposed method and the empirical evaluations.
Decision: Accept (Poster)