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 discrete-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
Accunty
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
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:/ 3)/ Σ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 two 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-Throue 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) reparameterizatioo 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 ciscrete-continuou 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 10.8100 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 -hhooretc 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. 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 226, 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": "fhBccrWyKv4", "reviewer": "Reviewer_b78f", "summary": "The paper under review deals with the problem of parameter learning in computation graphs that contain both continuous and discrete nodes. Training in this setting typically employs the Gumbel-Softmax relaxation of the discrete distribution. The paper proposes two tricks (1) scaling the Gumbel noise sample by an amount matching the Gumbel-Softmax temperature parameter and (2) introducing residual dropout connections which allow the computation graph to exclude the discrete connections and better propagate gradients.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper falls into the by-now large literature of introducing discrete nodes in computational graphs via the Gumbel-Softmax relaxation and the reparameterization trick. The paper focuses specifically on the problem of learning in the presence of multiple discrete nodes in the graph. Gradient learning in this setting is challenging because of the accumulation of vanishing gradients in serial connection to each other.\n\nBoth techniques proposed in the paper make intuitive sense to me. They both introduce some leaky behavior to the discrete nodes and counteract the vanishing gradient problem. The paper does not provide some strong theoretical support/foundation behind either of the two proposed techniques but gives intuitive arguments motivating them.\n\nThe paper evaluates the proposed technique in three tasks which employ discrete-continuous graphs in a rather small-scale / confined setting. The results look good, yet I am not very familiar with the corresponding application domains to be able to judge the significance of the presented results.\n\nThe paper is accompanied by code which should make the reported experiments easier to reproduce.\n\nOverall I find that the paper proposes two reasonably justified yet rather incremental improvements on the problem of learning in the presence of discrete nodes.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "PfSgBjxoaCY", "reviewer": "Reviewer_Xmig", "summary": "The paper proposes an efficient training method on the complex discrete stochastic computational graph, other than the single discrete stochastic layer.\nThe authors suggest two methods: 1) TempMatch, which manipulates the temperature and the Gumbel scale parameter of the Gumbel-Softmax; and 2) DropRes, a stochastic residual connection.\nThe athuors provide three experiments that backups their suggested methods.\n", "questions": "", "limitations": "", "rating": 4, "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, Quality and Significance:\nMy opinion on the originality, the quality, and the significance of this paper is complicated, hence I rather write them down in one part.\nIf the goal of the paper is to provide some techniques for better learning with the Gumbel-Softmax (which I do think that the idea is novel), the paper should provide an ablation study on the parameter selections in the proposed methods.\nFor example, how does the performance varies as the selection of \\beta in TempMatch, or \\alpha in DropRes.\nIn this case, the quality of the paper matters.\nOr, if the goal of the paper is to provide better results in certain tasks, the detailed explanations on the tasks or architectures of models (or the modification on the original models) should be included in detail in the main paper.\nIn this case, the significance becomes the problem.\nRegarding the above, I think this paper is located in the fuzzy area.\nHowever, the idea and results seem promising, so the paper will be in a better form after revision.\n\nClarity:\nWhile several definitions on the abbreviations are missing, the paper is well-written in general.\nHowever, the numberings of references in appendix does not match the ones from the main paper.\n\n\nComments:\n- The hypothesis on the gradient vanishing seems true, but it would be better if there are experimental backups.\n- Is there any particular reason for \n- The protocol of the DropRes rate \\alpha choice seems not grounded where the protocol is crucial to the quality of this paper. Is this should be considered as a part of hyper-parameter tuning, for example, the temperature annealing of the Gumbel-Softmax? How's the result differ as DropRes connection samples vary? Are the results stable, and how's the stability of the result? (regarding the variance of the result, for example.)\n- Even though the authors are dealing with stochastic computational graphs with multiple discrete latent layers, my opinion is that the experiments should be compared against the score function-based gradient estimators and the naive Gumbel-Softmax with temperature annealing.\n- Considering the multiple discrete stochastic layers, my guess is that the work is well-fit with combinatorial optimization problems such as [1].\n- The code is provided together with the generated ListOps dataset for reproducibility.\n- There is no confidence interval in the multi-hop reasoning over KG experiments.\n\nQuestions:\n- What is the specific choice of DropRes rate \\alpha for each experiment?\n- How is the probabilistic simplex space different from the original Gumbel-Softmax as utilizing TM in 2-simplex, rather than 1-simplex case proposed in the paper?\n- How is the probabilistic simplex varies as the epoch grows when using TM?\n- How is the behavior of the loss curves different? (w/ TM only, w/ DR only, w/ both TM & DR, and w/o both TM & DR)\n- How is the behavior of the variance of gradients different? (w/ TM only, w/ DR only, w/ both TM & DR, and w/o both TM & DR)\n- How do TM and DR affect the performance in discrete VAE settings?\n\n[1] Attention, Learn to Solve Routing Problems!, Kool et al., ICLR 2019.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors argue that they pointed out the limitation of their work in one line in the conclusion (where I do not think that the line is the limitation of the work), and there is no negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "E7ohqJvqOMZ", "reviewer": "Reviewer_nNqs", "summary": "The paper presents two novel techniques to train models with deep Gumbel-Softmax (GS) layers. Both aim to make the GS gradient larger. The first is gradually increasing the Gumbel distribution scale \"beta\" to be relatively larger than the GS temperature \"tau\" so that samples are more uniform and the gradient doesn't saturate. The second (more important) is a stochastic residual connection that kicks in with probability 1-alpha where alpha gradually increases to 1. These are shown to be more effective than the more conventional annealing of tau and non-stochastic residual connection, respectively. The paper achieves impressive performance on 3 different tasks that involve discrete reasoning.", "questions": "", "limitations": "", "rating": 7, "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\n- The paper is well motivated and addresses a general and important problem (discrete reasoning).\n\n- The proposed techniques are simple but effective. By focusing on making the gradient in Eq. (5) larger, the paper teaches the reader a lot about what troubles/remedies one might find when using deep GS layers. Changing beta in Gumbel and making the residual connection stochastic are neat novel ideas.\n\n- The experiments are extensive (thanks for providing the code which is a great reference when the paper is not clear). In particular, the results on multi-hop KB reasoning are impressive because the model beats a much bigger baseline CoKE. To me this seems to validate and justify the utility of introducing explicit discrete layers which makes me feel relieved.\n\n\nWEAKNESSES \n\n- Increasing beta is intuitive but it's not formally justified. There's no reason that this should be more effective than decreasing tau in making the gradient (5) larger, other than the handwavy argument that the latter saturates. The paper also doesn't really address the concern about the mismatch between the Gumbel trick and the model distribution. In short, decreasing tau seems to be still the \"right\" thing to do in theory, but the new scheme is more effective empirically. But it's still significantly less important than the residual connection so I suppose it's a minor issue.\n\n- I feel there should be explicit comparisons with existing techniques (tau annealing and non-stochastic residual connection) in each task. Figure 5 and 6 only show this for 1 task respectively, and that makes me wonder if this story still holds in other tasks and how much it matters. \n\n\nQUESTIONS/COMMENTS\n\n- In Table 4, why not include paths <= 5? \n\n- Despite the issues above, I think the submission is a clear accept given its positive finding and the amount of substance. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "For limitations, see the issue with formal justification above. But it's not a theory paper so it's not a big deal, I feel. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "2uag5ZKFBfS", "reviewer": "Reviewer_RKc5", "summary": "This paper focuses on training models that incorporate long chains of discrete sampling operations interleaved with learnable neural components. In particular, the authors focus on neuro-symbolic systems and other similar approaches, where the goal is to learn to assemble and compose interpretable discrete operations on discrete objects (for instance, parsing an expression like \"max [ sum [ 1 3 ] 2 ]\" into a syntax tree, then treating each node of the tree as an interpretable mathematical function on the integers, and using that to evaluate the expression).\n\nThe authors start from the Gumbel-Softmax trick, and observe that it may introduce vanishing gradients (which they provide evidence for using a simple toy problem), which they hypothesize will lead to difficulty training with long chains of discrete sampling operations. To combat this, they introduce two strategies: first, they suggest modifying the scale of the Gumbel noise during training (which I believe is equivalent to rescaling the model's logits), and second, they propose stochastically adding \"skip\" connections that allow communication to bypass discrete operations (in particular, to allow gradient flow through overly-confident samples).\n\nThese two ideas seem quite straightforward to implement, and lead to a range of improvements across a few different experimental domains. The authors modify a variety of existing models by adding additional discrete sampling steps and using their training strategies, and show that these models can reach similar or slightly better accuracy while also being easier to interpret and generalizing better to longer inputs.\n\nOverall, this approach seems very promising and could lead to exciting future work. However, I think it has some significant shortcomings regarding the clarity of the proposed approach and missing justifications for the claims. I found some of the sections hard to follow, and many of the ideas could benefit from more thorough explanation and introductory examples. I also think some of the claims the authors make are too strong (e.g. they say that they analyze behavior and show that training difficulties are due to vanishing gradients, but it seems like this is justified only using intuition and a few small experiments, and I'm not convinced that they have shown what they claim). With revision I think this could be a very strong publication, but I don't think it is quite there yet.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Vanishing gradient hypothesis and justification\nThe authors state that they \"show both analytically and empirically that it is challenging to optimize the parameters of these models, mainly due to vanishing gradients\". While this seems like a reasonable hypothesis, I'm not entirely convinced that the authors have actually shown it to be true.\n\nThe \"analytically\" part of the claim seems to be based on just writing down the expression for the gradient of Gumbel-Softmax, and then asserting that this has a vanishing gradient problem. This doesn't seem like a sufficient justification to me. It's not immediately obvious what part of the shown expression is problematic, or why their proposed modifications would necessarily fix the problem. I can make some guesses as to what the authors mean, but I don't think it's reasonable to expect the reader to believe the claim from just what is shown. (Perhaps the derivation would be something like: assume that the $\\theta$ parameters are saturated to some degree, show that it's still saturated with high probability after Gumbel noise, then show that for every $i, j$ pair at least one term in (5) is thus \"small\" with high probability? I'm not completely sure what the authors mean by vanishing gradient here, anyway, is it supposed to be the expected gradient norm with respect to $\\theta$? The norm of the expected gradient? Something else?)\n\nEmpirically, the authors provide two results in support of their hypothesis.\n- First, they show plots of gradient values for a toy problem where the Gumbel-Softmax outputs are directly supervised with a cross entropy loss. This seems like a strange experiment to me, since they are essentially treating a (relaxed) discrete sample as a distribution itself; as temperature goes to zero this cross entropy loss would almost always become infinite. I also can't really tell what they are trying to show. If they change the size of the Gumbel noise (which I comment on more below), that's just like changing the distribution to be less confident, and so it seems like all this plot is saying is something like \"if you prevent the model from being confident, the model takes longer to train and so its gradients take longer to reach zero\"? I don't really know why this would be expected to be representative of a real world experiment, it seems like \"taking longer to train\" is a bad thing? (They also have a plot saying that the larger values of $\\beta$ lead to $\\text{softmax}(\\theta_1)$ being closer to 1, but that seems misleading to me, since as the authors note later, changing $\\beta$ changes the distribution being sampled from and thus I think it would be more fair to plot $\\text{softmax}(\\theta_1 \\beta)$ here.)\n- Second, they show ablation results with the gradient norm on a real-world task with five layers, for which their residual connections seem to increase gradients through previous layers. This does seem more convincing to me, at least in support of the residual connections. But I'm not sure if this plot also includes the modified $\\beta$ terms? It would be interesting to see the version of this plot with all four versions (GS alone, GS+TM, GS+RB, GS+TM+RB).\n\nOverall I'd guess that vanishing gradients probably do cause difficulty, but I think that the claims the authors make are nevertheless too strong here. A proper analytical demonstration should include at least formalizing the claim (and probably also proving it to be true), and I think the experimental results could be more clear here as well.\n\n### Proposed approach\nThe authors propose two strategies for training models with long chains of discrete operations.\n\n**TempMatch**: The authors suggest changing the scale parameter $\\beta$ of the Gumbel-(soft)max distribution during training. As the authors note, this actually changes the distribution that the system samples from, with higher values of $\\beta$ leading to more uniform samples that ignore the logits. Mathematically, I think this is actually equivalent to just changing the temperature and rescaling the logits: letting $\\epsilon \\sim \\text{Gumbel}(0,1)$, they propose using\n\n$\n\\text{Softmax}\\big( (\\theta_i + \\beta \\epsilon_i)/\\tau\\big)\n$\n\nwhich I believe is equivalent to\n\n$\n\\text{Softmax}\\big( ( (\\theta_i/\\beta) + \\epsilon_i)/(\\tau/\\beta)\\big)\n$\n\nThus, increasing $\\beta$ is exactly the same as just decreasing the temperature $\\tau$ but simultaneously scaling the logits to be less confident.\n\nThe authors then assert that vanishing gradients can be addressed by \"increasing the scale parameter ... $\\beta$ relative to the temperature parameter $\\tau$\" and later that \"the two parameters $\\tau$ and $\\beta$ should be matched, that is, optimized relative to each other\". I'm not sure what they actually mean by this, or what the evidence for this claim is. Is the idea to increase/decrease both $\\beta$ and $\\tau$ proportionally? Or increase $\\beta$ while decreasing $\\tau$? Increase the ratio between $\\tau$ and $\\beta$? It seems like in the experiments they actually just increase $\\beta$ while leaving $\\tau$ fixed, so I'm not sure what they mean by \"relative\" here.\n\nAssuming they mean to just fix $\\tau$ and increase $\\beta$, this would be equivalent to decreasing the effective temperature $\\tau/\\beta$ while simultaneously decreasing the model's confidence $\\theta_i/\\beta$. I don't follow why this should prevent vanishing gradients. Intuitively, I would expect decreasing the temperature and simultaneously decreasing the influence of the model on the output would lead to smaller gradients, not larger ones. Perhaps that intuition is wrong, but if so the authors need to actually show why. (Again, having an actual formalization of the vanishing gradient claim would help here.)\n\nI also don't really understand how large $\\beta$ should get. If $\\beta \\to \\infty$, the model is essentially ignored, so presumably we don't want to increase $\\beta$ that much?\n\n**DropRes**: The other strategy the authors propose is to stochastically insert skip connections \"around\" discrete operations with some probability $\\alpha$, which is annealed to zero during training. This strategy was explained quite clearly, but it wasn't immediately obvious to me how to interpret these skip connections. It (at least initially) seemed pretty limiting to only be able to do this to discrete operations that maintain the shape of the input. After reading through the experiments section, I think I have a better sense of this: if I understand correctly, the idea is that the discrete operations in question are only discrete as a sort of inductive bias, e.g. instead of allowing the model to pass an arbitrary real-valued vector, it is forced to sample the embedding of one of the 10 numerical digits. Here, the skip connections essentially allow it to occasionally \"cheat\" and pass additional information on the side, allowing it to learn even if the discrete operations have become incorrectly very confident. I think this means the DropRes idea can only apply in some limited situations. For instance, is it true that you can't apply DropRes to the prediction of the edges in the graph for the ListOps experiment, since those are structural changes and not just a discretization inductive bias?\n\nOverall, I'd appreciate it if the scope and limitations of the DropRes idea could be stated up front. For instance, at the beginning of 2.2, the authors could provide an example where DropRes makes sense, and an example where it doesn't. An introductory example like that would be very helpful for understanding what it is meant to accomplish. (Later in the section, the authors point to Figure 1, but as I discuss later in my review, that wasn't enough for me to understand what was going on.)\n\n### Originality and significance\nAlthough the two strategies the authors propose seem fairly simple, I'm not aware of those strategies being considered before, at least not from the perspective of enabling deeper training of stochastic computation graphs. The empirical results are also quite impressive. Given this, I do think the paper could have a fair bit of impact on training mixed discrete/continuous models. I'd also be very interested to see more work in this space, regarding the training dynamics of deep models with discrete sampling steps.\n\nI do think it is important to give more credit to some of the existing gradient estimators that DO discuss long chains of stochastic nodes, contrary to the claim here that \"Prior work has focused on computation graphs with a single discrete component on each of the graph’s execution paths\". For instance,  Tucker et al. [27] describe how to apply their approach to multilayer models (in section 3.3 and appendix F), and Grathwohl et al. [7] discuss gradient estimates for a sequential Markov decision problem with a sequence of stochastic discrete actions (in section 3.4). While the authors do cite both of these prior works as single-node estimators, I think it would be appropriate to also mention that they did consider sequential computation graphs as well. (The key difference seems to be that this current work is focused on *efficiently* training these graphs, not just on constructing an unbiased gradient estimator for them.)\n\n### Experiment design and results\n\n**ListOps:** Although the authors claim to have \"simplified\" the arborescence prior in Paulus et al. [20], I think they may have accidentally broken the arborescence property. In particular, although it is true that every node in an arborescence has exactly one parent, it is NOT true that every directed graph where every node has exactly one outgoing edge is an arborescence! So simply having a Gumbel-Softmax edge choice over all possible parent nodes does not necessarily lead to an arborescence. I'd suggest renaming this to a \"directed pseudoforest\" prior (see [Wikipedia](https://en.wikipedia.org/wiki/Pseudoforest#Directed_pseudoforests)) or something similar, to distinguish it from a proper arborescence constraint. (Or, you could go ahead and actually implement the arborescence prior from Paulus et al. [20], but I don't think that's necessary as long as the names are precise.)\n\nI don't understand what you mean by training on the ground truth graph. Is the idea that you don't even predict a graph at all, you just have a GNN on the true parse? But you still do discrete operations on the nodes?\n\n**Multi-Hop Reasoning over Knowledge Graphs:** Do you assume a many-to-one relationship here? In other words, given a source node and a relation, do you assume there's at most one target?\n\nI'm not sure what you mean by \"we use the Gumbel hard setting also during training\". Is this referring to the Gumbel-Softmax + Straight-Through estimator, where the forward pass uses Gumbel-Max? Or something else?\n\n**MNIST:** I'm interested to know what the evidence for \"the learned operation is modular and likely able to extrapolate to multiple sequential addition operations\" is. Was this based on manual inspection of the results, or did you actually run an extrapolation experiment?\n\n### Clarity\nAlthough the introduction was quite clear, I found section 2 and section 4 to be quite hard to follow. Overall, I think the paper could benefit from giving more clear intuition and specific examples of where each of the proposed techniques could be applied.\n\nDetailed comments:\n\nSection 2:\n- Figure 1, although referenced in many different places, does not seem to have enough information to understand what is going on. In particular, it was very difficult to tell what the different boxes, operations, and colors were supposed to represent, especially since the actual tasks were not described until the experiments section. Perhaps the figure caption could be expanded to briefly describe in words what the task is and how it happens (something like \"Given a textual sequence [MAX 2 9], the model samples a discrete set of edges in a graph representing the operations. Next, each node computes a discrete output as a function of (embeddings of) its inputs.\")\n- Line 57: \"Any deterministic neural network can be represented by a composition of basic differentiable functions.\": While most neural networks can be represented this way, there's been a lot of interest in more exotic networks that include complex control flow, differential equation solvers, etc., so I'd suggest wording this differently.\n- Figure 2 and surrounding paragraphs: I wasn't sure at this point why you were limiting attention to this type of function. It might be worth giving an example here of what this is meant to represent (e.g. sampling a discrete integer answer for a mathematical operation) and what it can't represent (e.g. sampling the set of edges of a graph)\n- Equation (5): As I mentioned above, it's very unclear what you mean by vanishing gradients here, or how this equation indicates they are present.\n- TempMatch\n  - Line 119: As mentioned above, I don't know what you are proposing by \"relative to\". It would be helpful if you actually wrote out the proposed learning schedule you suggest using. I can tell you are increasing $\\beta$, but it seems like you mean something more specific than that?\n  - Line 128: I don't understand why we want to continue training the model once it's already solved the task?\n  - Line 130: What do you mean by \"the inverse scale parameter\" here, and how is it different from what you've already said? Are you saying something new here? In general this section feels like a bunch of assertions, and I wasn't able to follow the justification for them.\n- DropRes\n  - As I mentioned above, it wasn't obvious to me what it means to drop out a discrete sampling operation. I think this only makes sense in a few limited (although perhaps common) cases?\n  - Relatedly, when you mention applying this to Figure 1, it's not clear what pieces you are referring to. I take it that you apply DropRes to the Max node, and maybe the constant nodes, but not the \"generate graph\" node? (For that matter, what does it mean to have a computation graph that samples the number 2 from the number 2? Is this just translating embedding spaces somehow? Is \"2\" an embedding, or a one hot vector?)\n\nSection 4:\n- I found the introductory paragraph very helpful here in terms of contextualizing the goals of the experiments. Thanks!\n- ListOps\n  - It was very difficult to follow this section without having Paulus et al. [20] open in another window. For readers who may be unfamiliar with that work, could you include a brief introduction to what the task and architecture actually is? Something like \"Following Paulus et al. [20], we train a model that first maps from the text description to a directed acyclic graph, then runs a GNN on that graph to compute the answer.\"\n  - Similarly, it would be helpful to give some intuition for what it means to model the nodes as discrete-continuous components (in particular, that each node represents the output of a particular operation).\n  - Line 248 is missing a description of what you mean by training on the ground truth graph.\n- Knowledge base\n  - Again, I found this very hard to read without referring to the prior work, and it wasn't clear what parts of your notation you took from [8] and what you took from [22].\n    - What is \"FB15K-237\"?\n    - What does it mean for a triple to be \"true\"? Is it that this is the only triple in the knowledge base with those two first elements?\n    - Intuitively, how are you computing logits? From the appendix it looks like you are using something from ComplEx, but you don't explain what that means, and don't cite it until after describing computing logits on line 272.\n    - What do you mean by \"train 1vsAll\"?\n    - Line 281: what is \"the Gumbel hard setting\"?\n\nSmall nitpicks:\n- line 32: comma missing after \"computation graphs\"\n- line 64: extra comma after \"Both\"\n- line 306: \"less\" should be \"fewer\"\n\n---\n\nEdit: I have increased my score from 5 to 6 in light of the clarifications made by the authors and the planned improvements in the next revision (see discusson below).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This work is a methodological contribution that does not seem to have any direct negative societal impacts.\n\nIn terms of limitations, the authors briefly mention that \"learning models with discrete nodes is still considerably more challenging than models without\". I think it would also be good to discuss how their proposed strategies make particular assumptions about the structure of the computation graph that might not hold for all discrete-continuous computation graphs.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Arguing that the task of learning parameters in stochastic computation graphs containing sequences of discrete random variables is difficult due to vanishing gradients resulting from the use of the Gumbel-softmax relaxation, the authors propose two techniques to counteract this difficulty: (i) dropout residual connections (DropRes) that bypass the discrete sampling / discretization step with some probability and (ii) annealing the Gumbel scale parameter while keeping the relaxation temperature fixed. The authors demonstrate the effectiveness of the techniques at training discrete/continuous models to perform unsupervised parsing of lists of operations and multi-hop reasoning in knowledge graphs.

The reviewers agreed that the techniques are novel, and the paper is interesting and potentially influential. The experimental results are also impressive, with the proposed methods being clearly helpful at achieving good performance on fairly sophisticated tasks. The experiments could be more systematic and informative however by including more ablations. Including at least one modern REINFORCE+variance reduction technique as a baseline would be nice, as such methods do not suffer from vanishing gradients (suffering from higher gradient variance instead).

The paper could be also improved by making the motivating argument about vanishing gradients more formal and improving the clarity of writing throughout, taking into account the reviewer suggestions.
Decision: Accept (Poster)