Efficient Learning of Discrete-Continuous Computation Graphs

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.


Reviews and author responses:

 If I understand correctly, your point is that $\epsilon_j$ will be larger as $\beta$ grows? I agree with that, but $\epsilon_i$ will also have mean $\beta \gamma$, so it will also grow as $\beta$ grows, and at the same rate. Since the softmax function is invariant to constant shifts of all inputs (e.g. adding $\beta \gamma$ to $\theta_i$ and also all $\theta_j$ for $j \ne i$ does not affect the output), I wouldn't expect this to make any difference.

But, going back to my previous comment, maybe you can say something about the quantity $\epsilon_j - \epsilon_i$? It's zero in expectation, but should have higher variance as $\beta$ becomes larger. It seems to me that this would imply that the Gumbel-SoftMax output vectors don't *always* have more uniform entries for higher $\beta$, but that they *sometimes* have more uniform entries, and the chance that they have more uniform entries grows as the saturation of $\theta$ grows. I guess maybe that's compatible with what you're saying? If $\theta$ is saturated, then by default things are very discrete, so sometimes having near-uniform entries is better than never having near-uniform entries? Even so I'd suggest trying to be a bit more precise about what you mean by "this also implies more uniform entries of the relaxed Gumbel-SoftMax vectors", perhaps by specifying how you are measuring/defining "more uniform".

 Dear authors.

Thank you for your immediate response.

I also disagree with that the domain of neuro-symbolic learning is narrow, in fact, I believe that it takes one huge part of the research community (and I do agree that the authors made their boundary on the work in the introduction part). My concern is that (1) the paper tackles a problem from the domain of neuro-symbolic learning, but (2) the detailed issues are raised from general discrete-continuous computational graphs, not from the neuro-symbolic domain (if there's neuro-symbolic-domain-specific learning issue on the discrete-continuous computational graph, please let me know), and hence, (3) I'm naturally wondering that if the proposed method would work better on the other general problem, for example, multiple discrete-discrete layer setting such as Kool et al., 2019. 

Additionally, my question on the comparison between score-based methods and the proposed method is not addressed.

Please let me know if there's any misunderstanding on your work.

Best,
Xmig

 Dear reviewer Xmig,

Thank you for your answer. We are open for suggestions to change the title and to add a more thorough discussion of the shortcomings of prior work in the introduction.

We have two comments to make here:

(a) Neuro-symbolic learning is a broad and, in our opinion, overloaded term. Perhaps a better change of title would be to make sure we are addressing “deep”/“multi-layered”/“complex”  “stochastic softmax tricks-based” computation graphs? Most prior work has focused on models that have one consecutive discrete probability distribution (e.g. discrete VAEs). We don’t think that neuro-symbolic would narrow it down sufficiently. 

(b) we want to point out that we make the concrete problem we address absolutely  clear both in the abstract and introduction. There we write that we are addressing computation graphs with multiple discrete components along a computation path using softmax tricks and specifically mention the vanishing gradient problem. 

Thanks again for your response and please let us know what you think. 

 Dear Paper10964 Authors.

Thank you for your response and I have read your latest update (https://openreview.net/forum?id=TLIHuw3gcQB&noteId=gAyv899t3Pt).
I do believe the work is very novel, and I'm happy with your plan for the improvement (1) and (2).

However, at the same time, I'm still concerning with the performance of the work in the multiple discrete stochastic layer settings.
After re-reading this paper at a glance, I think this gap originates from the position of the paper.
The authors' motivation of the work starts from the neuro-symbolic learning systems, however, the title and the methodology section are rather stated in broader concept rather than the neuro-symbolic domain.
Hence, my suggestions on the paper to justify the coverage are (1) narrowing down the title, for example, adding "in Neuro-Symbolic Systems" at the end, and (2) adding preliminary section which discusses the limitations of the previous work in the neuro-symbolic domain (perhaps, the direct problem caused from the "vanishing gradient"?)

Thanks again for the work and your response.

Best,
Reviewer Xmig

 Thanks to the extremely helpful reviews and especially the review and discussion with reviewer RKc5, we will make the following improvements to the current manuscript:

(1) In addition to the existing ablation experiments, we also apply ablation experiments to other problems/datasets. Specifically, we are including new ablation results on the MNIST addition experiments. These experimental results are currently collected. 

(2) We plot a probability simplex for various values of $\tau$ and $\beta$ and add an experiment as suggested by reviewer RKc5 where we increase the value of $\beta$ for the setting where one $\theta_i$ has an increasingly higher positive value. This will support empirically (and better than the current experiment shown in Figure 3) the fact that higher values for $\beta$ make samples more uniform (both, the relaxed vectors through SoftMax and the distribution of the max probabilities). 

(3) We clarify what we mean by vanishing gradient (analogous to what is known as vanishing gradient for activation functions in standard deep neural networks), addressing the change in norm of the Jacobian and showing that the norm is higher for more uniform distributions.

(4) We clarify the presentation of the general form of the addressed problems, by improving Figure 1 and its caption and by being more specific about the model architectures and annealing schemes used in the experiments. For instance, we will move important information from the appendix to the main paper as discussed in the answers to the reviews. 

Thanks again to all of you. We are very happy about the feedback we obtained for the paper (irrespective of it being accepted/rejected :-))

 > But there are two different ways that outputs could be more uniformly distributed! One way would be for them to be closer to a vector with uniform entries, which would perhaps improve gradient flow based on your first point. But the other would be for them to be close to a discrete sample from a uniform distribution, which would have individual entries that are very different from each other! And I believe increasing the scale parameter of the Gumbel noise does this second thing; it makes them closer to discrete samples from a uniform distribution. By the same reasoning in your first point, this actually makes the problem worse, because a larger  would lead to outputs that are LESS like a vector with uniform entries.

Indeed, these are two different notions of "samples being uniform." The argument with the Gumbel-max trick in the previous post shows that the classes with max probability are more uniformly distributed when $\beta > 1$. However, what we are trying to argue is that this also implies more uniform entries of the relaxed Gumbel-SoftMax vectors *compared to noise samples from Gumbel(0, 1)*. When the $\theta_i$ for one particular class $i$ is larger than all others (later in training, the point close to saturation), then sampling from a Gumbel with larger scale parameter $\beta$ does *also* lead to (relaxed Gumbel-SoftMax) vectors with more uniform entries (*compared to sampling with* $\beta=1.0$). In expectation, a sampled $\epsilon_j$ ($j \neq i$) for $\beta > 1$ is larger than a sampled $\epsilon_j$ ($j \neq i$) for $\beta = 1$ because the mean of a RV with distribution Gumbel$(0, \beta)$ is $\beta\gamma$ with $\gamma$ the (positive) Euler-Mascheroni constant. 

We believe that the experiment you suggest above is a great idea. 

 Thanks for the reply. I appreciate the effort, but I think I have found a problem with your sketch.

Your first point (as I understand it) is that the Jacobian of softmax operation has its largest norm when the inputs are close to a uniform vector, e.g. all of the entries are close to one another. This seems true to me. 

Your second point, if I understand correctly, is that increasing the noise scale $\beta$ produces outputs that are "more uniformly distributed". But there are two different ways that outputs could be more uniformly distributed! One way would be for them to be closer to a vector with uniform entries, which would perhaps improve gradient flow based on your first point. But the other would be for them to be close to a discrete sample from a uniform distribution, which would have individual entries that are very different from each other! And I believe increasing the scale parameter of the Gumbel noise does this second thing; it makes them closer to discrete samples from a uniform distribution. By the same reasoning in your first point, this actually makes the problem worse, because a larger $\beta$ would lead to outputs that are LESS like a vector with uniform entries.

It seems like your argument is conflating the two types of uniform-ness? So it doesn't seem like this shows that TempMatch improves gradient flow at the end of training. If anything, it suggests that TempMatch will generally make gradients smaller, which is concerning.

Either way, this analysis isn't quite what I asked for, because it isn't actually considering the behavior of the full TempMatch operation. It seems that you've broken it down into two pieces, and then analyzed the softmax and the Gumbel-max separately. But of course the thing you actually care about is the Gumbel-softmax, which you haven't provided an analysis of; it seems that you're missing a key step of the proof and are still relying on intuition to connect the two pieces (and as I've shown above, that intuition doesn't really seem to make sense). A full analysis should consider the Jacobian norm of the full Gumbel-softmax operation, likely taking an expectation over the random noise.

I'd also be happy with a numerical approximation of such an analysis, if a proof ends up being too difficult given the time constraints. What happens if you fix $\tau = 1$, consider a few different values of $\beta$ (perhaps $0, 0.2, 0.4, 0.6, 0.8, 1$), then smoothly vary one of the logits of $\theta$ by some quantity $\delta$ as you describe above? You could then make a plot of what the norm of the expected Jacobian through the entire operation is by just drawing a bunch of samples and computing the Jacobian for each. Or you could plot the expected norm of the Jacobian instead, if that's what you are trying to measure. This would both concretely state what quantities you are interested in optimizing and what quantities you vary, and also demonstrate how the level of saturation changes the impact of $\beta$. Maybe it would also be helpful for coming up with a proof. I made a quick sketch of what a plot like this might look like, and I think it's consistent with what we discussed, in that larger $\beta$ causes slightly larger gradient norms for saturated $\theta$ but (much) smaller gradient norms for non-saturated $\theta$. So I'm hopeful that the technique is still sound, and just the analysis is incomplete.

(Some speculative thoughts that might be helpful, after thinking about this a bit more. Maybe a way of thinking about what TempMatch is doing is that, indeed, it is drawing samples that are much closer to discrete samples as $\beta$ increases. Each of the samples, then, is probably pretty far from a constant vector, so it probably has a lower softmax Jacobian norm by itself. But, if $\theta$ is also very saturated, then sometimes by chance they are saturated in opposite directions? So then $\beta$ occasionally randomly pulls you back to the "near-constant-vector" regime where there's a larger softmax Jacobian norm. This would explain why large $\beta$ is bad at first, because if $\theta$ is already small, no matter which direction you jump you probably make the softmax farther from its large-Jacobian point. But if $\theta$ is very big, jumping in a random direction sometimes (but still rarely) makes you lucky and lets you get back to class probabilities that are close to one another.)

 Dear reviewer RKc5,

Thank you for your continuous feedback and engagement with us. We really appreciate it. 

Let us try to outline the plans we have for a more thorough formalization and the ways in which we will integrate it into the paper. We only provide a sketch here due to the difficulties to write formal math in OpenReview. 

We will address the issue with three parts.

First, we reiterate the fact that backpropagation performs matrix-matrix multiplications involving the Jacobians of the individual functions under considerations. Hence, in order to analyze the vanishing gradient problem, we need to investigate the behavior of norms of the Jacobians when we move towards saturation of the SoftMax probabilities, that is, when for an input one of the class probabilities approaches 1.0 (something we observe empirically). 

Second, we show for several norms (e.g., the Frobenius norm) two properties: (a) the norms go to 0 when we approach saturation: if we increase one class probability by a constant $\delta>0$, the norm gets smaller; and (b) for the uniform distribution (of the SoftMax) the matrix norms of the Jacobian obtain their maximum (up to a constant). Hence, "keeping" the SoftMax distribution more uniform for longer during training keeps the norms of the Jacobian larger, therefore increasing the gradient signal for parameters of upstream functions. 

Third, we improve the discussion of the way in which TempMatch achieves this. This should also make the difference between lowering $\tau$ (of the standard Gumbel-softmax trick) on the one hand, and increasing $\beta$ (the scale parameter) while keeping $\tau$ fixed) on the other hand, more explicit and understandable.

Let us start with the Gumbel-max trick: here we take the parameters $\theta$ and perturb each $\theta_i$ by a sample $\epsilon_i$ from $\mathtt{Gumbel}(0, \beta=1)$. We then sample a $z$ by taking the one-hot representation of the argmax (the maximum configuration). We know that when we do this, we take faithful samples from the categorical distribution $p(z; \theta) = $ SoftMax$(z; \theta)$. 

Now, the Gumbel-SoftMax trick relaxes the argmax by dividing both vectors $\theta$ and $\epsilon$ by a parameter $\tau$ and by taking the SoftMax. For $\tau \rightarrow 0$ we have the Gumbel-max trick again. 

What we can show (as a simple extension of the Gumbel-max trick) is that if we perturb the parameters of each $\theta_i$ by a sample $\epsilon_i$ from $\mathtt{Gumbel}(0, \beta)$ and then take the argmax, we take faithful samples from the categorical distribution  $p(z; \theta/\beta) =$ SoftMax$(z; \theta/\beta)$, that is, from the SoftMax distribution for $\theta$ at temperature $\beta$. The larger $\beta$, the more uniformly distributed (over the classes) are the discrete samples from this distribution. 

With TempMatch we increase $\beta$ during training and (implicitly) relax the discrete Gumbel-max samples using the Gumbel-SoftMax trick (dividing both $\theta$ and $\epsilon$ by a constant $\tau$ and taking the SoftMax). By doing this we ultimately obtain relaxed (using Gumbel-SoftMax trick with $\tau$) samples from a categorical distribution whose class probabilities are more uniformly distributed. 

*Important:* Given the above explanation, one might be tempted to think that an alternative way to look at TempMatch is as a method that  instead simply divides the parameters $\theta$ by a constant $\beta$. Indeed, this would be true if we used the Gumbel-max trick. But since we use the Gumbel-SoftMax relaxation it is *not* equivalent. Indeed, TempMatch leads to more discrete Gumbel-softmax relaxations. 

Thank you again for your extremely helpful remark! And let us know if you have more questions. 

 Dear reviewer Xmig,

You gave the paper the lowest score of all reviewers. We are wondering if our response to your review has changed your mind about the score or if there are other concerns we could address / questions we could answer. 

Thank you.

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).

The 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).

These 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.

Overall, 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.  ### Vanishing gradient hypothesis and justification
The 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.

The "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?)

Empirically, the authors provide two results in support of their hypothesis.
- 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.)
- 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).

Overall 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.

### Proposed approach
The authors propose two strategies for training models with long chains of discrete operations.

**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

$
\text{Softmax}\big( (\theta_i + \beta \epsilon_i)/\tau\big)
$

which I believe is equivalent to

$
\text{Softmax}\big( ( (\theta_i/\beta) + \epsilon_i)/(\tau/\beta)\big)
$

Thus, increasing $\beta$ is exactly the same as just decreasing the temperature $\tau$ but simultaneously scaling the logits to be less confident.

The 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.

Assuming 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.)

I 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?

**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?

Overall, 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.)

### Originality and significance
Although 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.

I 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.)

### Experiment design and results

**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.)

I 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?

**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?

I'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?

**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?

### Clarity
Although 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.

Detailed comments:

Section 2:
- 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.")
- 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.
- 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)
- 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.
- TempMatch
  - 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?
  - Line 128: I don't understand why we want to continue training the model once it's already solved the task?
  - 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.
- DropRes
  - 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?
  - 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?)

Section 4:
- I found the introductory paragraph very helpful here in terms of contextualizing the goals of the experiments. Thanks!
- ListOps
  - 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."
  - 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).
  - Line 248 is missing a description of what you mean by training on the ground truth graph.
- Knowledge base
  - 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].
    - What is "FB15K-237"?
    - 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?
    - 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.
    - What do you mean by "train 1vsAll"?
    - Line 281: what is "the Gumbel hard setting"?

Small nitpicks:
- line 32: comma missing after "computation graphs"
- line 64: extra comma after "Both"
- line 306: "less" should be "fewer"

---

Edit: 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). This work is a methodological contribution that does not seem to have any direct negative societal impacts.

In 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.

 After further reflection, I have decided to go ahead and raise my score from 5 to 6. In addition to the promised formalization, I think the improvements the authors intend to make to the clarity of the paper, and the additional experiments they plan to add, will strengthen the submission. I still feel that the hand-waviness of the claims around vanishing gradients are the largest remaining weakness of this paper, and hope that the next revision includes more precise claims similar to those I suggested above.

 You mention in your previous response that you will include a formalization of your vanishing gradient claim and how TempMatch addresses it. What is that formalization, and how do you plan to incorporate it into the paper? This is my primary remaining concern, and I would be willing to raise my score if the authors can provide a precise and accurate statement of how the behavior of the gradients is affected by TempMatch (especially since TempMatch sometimes makes gradients worse). Something like: mathematically define the property you mean by "vanishing gradient", then show that under some specific conditions, the original Gumbel-Softmax has that property, whereas TempMatch does not.

 **Scope and limitations of DropRes**
> _Perhaps the idea is that you take the constant $2$, look it up in $Num$ to find an embedding, then map that embedding back to a distribution over tokens using $Pred$, then look that embedding up in $Num$ again? Or perhaps the figure is misleading and you just look up $Num[2]$ for literal token nodes in the computation graph, without having a $z$. Or, perhaps the $2$ that appears in the middle of figure 1 isn't actually the token $2$ but instead an approximate embedding of the token $2$?_

For the sake of readability, the figure is fairly simplified. For example, the following ideas are happening in the background; also see lines 486-505 in the appendix.

Following Paulus et al. [20], the generated latent graph consists of continuously relaxed edges using Gumbel-softmax. Thus, the latent graph is a complete directed graph during training. Even tokens like $2$ have a (small) incoming edge from the tokens $9$ and $MAX$.

To integrate this property, we update _all_ token representations at each iteration simultaneously. Overall, given the latent graph, we start with a lookup embedding of the tokens $2$, $9$ and $MAX$ and update these representations at each of the 5 (= max depth) iterations with regard to their incoming edges. Using our proposed discretization, we update the representations towards approximate embeddings of the numeral tokens.
The figure illustrates the model during test time, where we use the straight-through estimator to obtain an exact graph and exact embeddings for the intermediate representations.

**Precise wording and definition of TempMatch**
> _Perhaps it would be clearer to say "the parameter $\\beta$ should be optimized relative to $\\tau$," then, instead of "the two parameters $\\tau$ and $\\beta$ should be matched, that is, optimized relative to each other."_

We agree and will clarify the definition of TempMatch with a precise wording together with the proposed schedule.

**General comments regarding clarity**
> _I'm not sure what you mean by "the arborescence prior still holds", since your model could easily generate things that aren't (continuous relaxations of) arborescences. It seems like the prior you use is weaker than the arborescence prior._

We will rename it to a _directed pseudoforest_ prior. Since each arborescence is a directed pseudoforest, the optimal solution is an element of the subset induced by our prior. 

> _Thank you for the clarifications. It would be useful to make this more explicit in the paper, so that readers know where to look for notation definitions._

We will make it more explicit in the paper where to find these notations or define them ourselves. 

> _Interesting, thanks. I think it would strengthen the paper to describe these experiments, perhaps in an appendix._

As discussed with other reviewers, we will include additional ablation experiments for additional tasks in the updated version of the paper. We agree that further experiments focusing on the generalization properties would further strengthen the paper.

**Vanishing Gradients**
> _I appreciate the detailed reply, and definitely think adding a more formal explanation will be useful._
> _Going back to the sigmoid analogy, it seems that what you are doing is less like replacing sigmoids with ReLUs, and more like "squishing" the sigmoid so that it has nonzero gradients over a longer span, but smaller gradients at zero. So effectively you're taking this:_
```
            │
            │    xxxxxxxxx
            │ xxx
            │xx
            xx
           xx
         xxx│
 xxxxxxx    │
 ───────────┴───────────────
```
> _and transforming it to this:_
```
                       │
                       │           xxxxxxxxxx
                       │     xxxxxxx
                       │  xxxx
                       xxx
                     xxx
                  xxxx │
            xxxxxxx    │
  xxxxxxxxxxx          │
───────────────────────┴─────────────────────
```

> _Before the model saturates, this probably hurts, because it makes gradients smaller. (This was my intuition in my initial comment.) But after it saturates (when the input is far from zero), this would increase the gradient, so that's why you only increase $\\beta$ later in training. Is that right? (It seems like this transformation still might lead to fairly small gradients. Perhaps there's room for an analysis of when this is likely to help and when it will hurt.)_
> _(Somewhat of a side question: are you aware of any prior work that uses this kind of rescaling to counteract vanishing gradients in sigmoid networks? I wonder if a similar transformation would be useful there as well.)_
> _I still feel like the toy problem isn't showing the right things to be convincing. It seems like the exact value of $\\theta$ doesn't really matter (especially once it's big enough to saturate the softmax), so I'm not sure having it increase tells you much about the training dynamics of the sequential setting. And in terms of the gradient, it seems like you might be able to get similar behavior by just decreasing the learning rate on $\\theta$? I think the proposed plot of gradient norms for a sequential model could be much more convincing here._

First, we agree with the reviewer. We mention in lines 117-121 that for a higher $\\beta$ (scale) the gradients exhibit more variance (are on average smaller) before saturation. After saturation, the distribution is shifted in such a way that it increases the average gradient norm since it is more likely to also “sample” relaxed states that have lower probability. The analogy with sigmoid activation is in reference to the problem of vanishing gradients when multiple sequential layers are used. It also crossed our mind that by adding and increasing the magnitude of some noise (which is essentially the case when increasing the scale parameter) we could make sigmoid units saturate less. This is interesting future work.

We will add a plot of gradient norms for a sequential model. Thanks again for the suggestion.


 Thank you for the clarifications, and I am glad my feedback was useful. I have a few follow-up questions and comments.

**Scope and limitations of DropRes**
> what does it mean to have a computation graph that samples the number 2 from the number 2?

> We explain this in detail in Lines 224-231: Given the embedding layer $num$ that maps the 10 numeral tokens to their respective embeddings and a Gumbel softmax distribution $z$ over the 10 digits, we obtain the ‘sampled’ new vector $y = z^T num$.

I think my confusion here relates again to the notation in Figure 1 (bottom half). That figure seems to show a specific token 2, which is fed through a neural network to produce a distribution over $z$. But my understanding is that 2 already has an embedding (so you can just look up $Num[``2"]$), and I'm not sure what $z$ corresponds to. There's also no $y$ in that figure.

Perhaps the idea is that you take the constant "2", look it up in $Num$ to find an embedding, then map that embedding back to a distribution over tokens using $Pred$, then look that embedding up in $Num$ again? Or perhaps the figure is misleading and you just look up $Num[``2"]$ for literal token nodes in the computation graph, without having a $z$. Or, perhaps the "2" that appears in the middle of figure 1 isn't actually the token "2" but instead an approximate embedding of the token "2"?

**Precise wording and definition of TempMatch**
> In general, we leave $\tau$ fixed and increase $\beta$ from 0 to $\tau$.

Perhaps it would be clearer to say "the parameter $\beta$ should be optimized relative to $\tau$," then, instead of "the two parameters $\tau$ and $\beta$ should be matched, that is, optimized relative to each other."

**General comments regarding clarity**
> Since each arborescence is a directed pseudoforest, however, the arborescence prior still holds.

I'm not sure what you mean by "the arborescence prior still holds", since your model could easily generate things that aren't (continuous relaxations of) arborescences. It seems like the prior you use is weaker than the arborescence prior.

>What do you mean by "train 1vsAll"?

> We follow the terminology and standard procedure of prior work (e.g., Ruffinelli et al. [22]). ...

Thank you for the clarifications. It would be useful to make this more explicit in the paper, so that readers know where to look for notation definitions.

> 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. ...

> This is based on generalization experiments similar to the ones we conducted for the multi-hop experiments. These experiments are constrained to a subset of the natural numbers using the modulo function.

Interesting, thanks. I think it would strengthen the paper to describe these experiments, perhaps in an appendix.

**Vanishing Gradients**

I appreciate the detailed reply, and definitely think adding a more formal explanation will be useful.

Going back to the sigmoid analogy, it seems that what you are doing is less like replacing sigmoids with ReLUs, and more like "squishing" the sigmoid so that it has nonzero gradients over a longer span, but smaller gradients at zero. So effectively you're taking this:
```

            │
            │    xxxxxxxxx
            │ xxx
            │xx
            xx
           xx
         xxx│
 xxxxxxx    │
 ───────────┴───────────────
```
and transforming it to this:
```
                       │
                       │           xxxxxxxxxx
                       │     xxxxxxx
                       │  xxxx
                       xxx
                     xxx
                  xxxx │
            xxxxxxx    │
  xxxxxxxxxxx          │
───────────────────────┴─────────────────────
```
Before the model saturates, this probably hurts, because it makes gradients smaller. (This was my intuition in my initial comment.) But after it saturates (when the input is far from zero), this would increase the gradient, so that's why you only increase $\beta$ later in training. Is that right? (It seems like this transformation still might lead to fairly small gradients. Perhaps there's room for an analysis of when this is likely to help and when it will hurt.)

(Somewhat of a side question: are you aware of any prior work that uses this kind of rescaling to counteract vanishing gradients in sigmoid networks? I wonder if a similar transformation would be useful there as well.)

I still feel like the toy problem isn't showing the right things to be convincing. It seems like the exact value of $\theta$ doesn't really matter (especially once it's big enough to saturate the softmax), so I'm not sure having it increase tells you much about the training dynamics of the sequential setting. And in terms of the gradient, it seems like you might be able to get similar behavior by just decreasing the learning rate on $\theta$? I think the proposed plot of gradient norms for a sequential model could be much more convincing here.


 ## Vanishing Gradients

> _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 [Figure 3] 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?_

> _Line 128: I don't understand why we want to continue training the model once it's already solved the task?_

The point of the experiment is to look at the behavior of one discrete-continuous component while keeping in mind that the setting we care about is the one where multiple such components are connected sequentially. The gradient norms are small (you describe it here as “the task is solved”) from downstream components, and we want to be able to have such gradient signals still reach upstream parameters.  

Here we want to point the reviewer again to the intuition we provide in lines 90-98, where we draw a parallel to sigmoid activation functions in standard multi-layer neural networks. By replacing sigmoids with ReLUs, we have that the gradient is the identity (for positive values). Hence, having ReLUs in multiple sequential layers does mitigate the vanishing gradient problem. We observe the same behavior for the methods we propose. 

> $softmax((\\phi + \\beta \\epsilon) / \\tau) = softmax((\\phi / \\beta + \\epsilon) / (\\tau / \\beta))$. _Thus, increasing $\\beta$ is exactly the same as just decreasing the temperature $\\tau$ but simultaneously scaling the logits to be less confident._

This derivation is correct, at least for $\\beta > 0$. For $\\beta = 0$, we have the standard $softmax(\\phi / \\tau)$.

> _Assuming 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 $\\phi_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._

We use the parameters $\\tau$ and $\\beta$ in the way we do because we think it makes the behavior of increasing $\\beta$ (while keeping $\\tau$ fixed) easier to understand. Intuitively, we want to make the model more frequently draw (relaxed) discrete samples that are low probability but which then incur a higher downstream loss and, therefore, allow the model to continue to learn. This is also demonstrated by the example plotted in Figure 3: while the probability has saturated, the model with higher scale parameters $\\beta$ obtains more pronounced gradient signals precisely because more often than for a smaller value for $\\beta$, a (relaxed) configuration is sampled that incurs a downstream loss. While your alternative formulation is not wrong, we believe it makes it more difficult to explain the behavior of the model when $\\beta$ is increased. 

> _Overall 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._

> _Perhaps the derivation would be something like: assume that the $\\phi$ 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 $\\phi$? The norm of the expected gradient? Something else?_

The expected gradient norm with respect to the logits is vanishing, especially in the setting of multiple consecutive Gumbel-softmax nodes. We agree with the reviewer that a proper formalization of our claim would improve the overall clarity. Here, we will focus specifically on the setting of multiple consecutive Gumbel-softmax components and the upstream gradient norm. Again, placing a higher probability on one of the categories (which we observed happens most commonly) makes the other categories’ gradients go to zero. This behavior is counteracted by having a noise distribution with a higher variance (higher scale parameter $\\beta$) which leads to exploring (relaxed) discrete variable assignments that incur a higher downstream loss. Besides the formalization, we will also include a plot of gradient norms of a simple architecture with several sequential Gumbel-softmax components. 


## Originality and Significance

The reviewer calls attention to prior work that has discussed long chains of stochastic nodes.

> _I 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.)_

We thank the reviewer for these references and will adjust the related work session accordingly. Yes, our focus is on efficiently training such models and not on providing unbiased gradient estimates. 

---
[1] T. Trouillon, J. Welbl, S. Riedel, É. Gaussier, and G. Bouchard. Complex embeddings for simple link prediction. In Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, pages 2071–2080.


 First of all, we want to thank you for the extensive and balanced review. Your comments will help us to improve the paper significantly. While critical of several aspects of the paper (especially its clarity) the review is also encouraging _“With revision I think this could be a very strong publication, [...].”_ -- Thank you!

## Clarity
**Introductory example**

> _Figure 1, although referenced in many different places, does not seem to have enough information to understand what is going on._

We agree with the reviewer about the importance of easy-to-understand introductory examples and we will provide Figure 1 with a more detailed explanation and caption. The new explanation will include a short description of the tasks as well as the connection to the discretization that we make use of in the following sections.

**Scope and limitations of DropRes**

We will work to improve the discussion of the scope and limitations of our proposed approaches. For instance, while we define the generic functions we consider in our work (in Figure 2 and Lines 71-75), we agree that we should add _“what this is meant to represent [...] and what it can't represent”_.

> _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?_

Exactly. We will use this insight as an example together with the revised version of the introductory example (Figure 1) to clarify when using DropRes is sensible and when it isn’t.

> _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?_

We explain this in detail in Lines 224-231: Given the embedding layer $num$ that maps the 10 numeral tokens to their respective embeddings and a Gumbel softmax distribution $z$ over the 10 digits, we obtain the ‘sampled’ new vector $y = z^\\intercal num$.

**Precise wording and definition of TempMatch**

> _"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_
 
> _it would be helpful if you actually wrote out the proposed learning schedule you suggest using._

We provide such a proposed schedule in the appendix in line 484, $\\beta_t = \\tau(1-e^{-t\\gamma})$. We will move this equation to the main paper, together with a more precise definition.

> _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._

In general, we leave $\\tau$ fixed and increase $\\beta$ from 0 to $\\tau$.

> _What do you mean by "the inverse scale parameter" here, and how is it different from what you've already said?_

We mean to express that the scale parameter is increased over time. We will use a single precise definition to improve the readability. 

**General comments regarding clarity**

> _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._

You are correct, of course. We will rephrase this part accordingly and clarify that we refer here to smooth maps. 

> _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. […] I'd suggest renaming this to a "directed pseudoforest" prior (see Wikipedia) or something similar, to distinguish it from a proper arborescence constraint._

The reviewer is correct. Since each arborescence is a directed pseudoforest, however, the arborescence prior still holds. We agree that renaming the prior to distinguish it from the one used in Paulus et al. [20] is beneficial.

> _I 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?_

Correct. This experiment aims to emphasize the benefit of discretizing the list operations which has not been done before. We run the GNN on the true parse (given as part of the dataset) and only train the parameters of the list operations.

> _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._

We will add a more detailed introduction of the task and the architecture.

> _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)._

In lines 224-231 we describe these discrete-continuous components mathematically. We agree, however, that an intuitive explanation would aid clarity here. We believe that a revised and improved introductory example (Figure 1; see comment above) will also help.

> _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?_

> _What do you mean by "train 1vsAll"?_

We follow the terminology and standard procedure of prior work (e.g., Ruffinelli et al. [22]). Training “1vsAll” means that given a triple consisting of a source, relation type, and target, we predict the target using only the source and the relation type. The benchmarks assume there exists exactly one correct target. If other targets are correct, they are filtered out before the evaluation measures are applied. 

> _I'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?_

“Gumbel hard”  refers to the Gumbel-Softmax + Straight-Through estimator where a discrete configuration is sampled in the forward pass. We will clarify this.

> _What is "FB15K-237"?_

This is introduced in the appendix with lines 599-600. Since all experiments on FB15K-237 are in the appendix, we propose to remove this reference from the main paper.

> _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?_

Given a triple $(s, r, o)$ from the Knowledge Graph and the set of entities $E$, $(s, r, o)$ is the only true triple from the set of triples $\\{(s, r, o’): o’ \\in E\\}$.

> _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._

We use the ComplEx [1] scoring function as it pertains to the source and relation type to compute the logits. We will describe this in more detail in the updated version of the paper. 

> _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?_

This is based on generalization experiments similar to the ones we conducted for the multi-hop experiments. These experiments are constrained to a subset of the natural numbers using the modulo function.

> _Small nitpicks_ (lines 32, 64, 306)

We thank the reviewer again for the detailed engagement with our work and will correct these spelling mistakes.

 We want to thank the reviewer for the detailed and motivating review.

> _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._

Decreasing $\\tau$ leads to more discrete samples. Decreasing $\\tau$, however, does not change the probability of sampling low probability (relaxed) discrete variable assignments. Increasing $\\beta$ (while keeping $\\tau$ fixed), on the other hand, results in a more discrete but also more uniform distribution over relaxed discrete variable assignments. See also our discussion about this in lines 110-114. This leads to relaxed variable assignments with a higher downstream loss being more likely to be sampled from the distribution. 

> _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._

We agree with the reviewer’s comment as well as with the other reviewers that such an ablation study would provide further insights. We will include additional ablation experiments for additional tasks in the updated version of the paper.

> _In Table 4, why not include paths $\\leq$ 5?_

We use the same test set for the extrapolation tasks in Table 4 as for the path query task in Table 3. The difference in the extrapolation tasks is the training set. The training set varies between all training paths of size $\\leq$ 4 and all training paths of size $\\leq$ 1, i.e., triples only. Thus, training on all paths of size $\\leq$ 5 is exactly the path query task of Table 3. We will clarify this in the paper. 


 We would like to thank the reviewer for the detailed examination of the quality of our proposed methods on the one hand and the significance of the experiments on the other.

> _If 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. For example, how does the performance varies as the selection of $\\beta$ in TempMatch, or $\\alpha$ in DropRes. In this case, the quality of the paper matters._

> _What is the specific choice of DropRes rate $\\alpha$ for each experiment?_

If the reviewer is referring to exact schedules which we use to increase the dropout probability $\\alpha$ or to increase the Gumbel scaling parameter $\\beta$, we have chosen here the most straightforward schedules. We increase $\\alpha$ linearly with $\\alpha_0=0, \\alpha_t=\\min(1, ct)$, for some $c>0$. We increase $\\beta$ by using the schedule $\\beta_t = \\tau(1- e^{-t\\gamma})$ for some $\\gamma>0$. The exact parameters are detailed in the appendix (for example lines 484-485). We agree with the reviewer’s comments that the annealing schedules are of importance and belong in the main paper.


> _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?_

> _How is the probabilistic simplex varies as the epoch grows when using TM?_

Is our assumption correct that the reviewer is referring to the probabilistic simplex space as found in Maddison et al. [14], Figure 2? If so, we believe it would be an insightful addition to the paper since it could provide a better intuition of the difference between annealing $\\tau$ and increasing $\\beta$ for a fixed $\\tau$. We will add the respective experiments to the appendix. As discussed in lines 113-121, we expect the distribution to be more uniform for higher values of $\\beta$ yielding the desired training behavior.


> _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)_

> _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)_

A full ablation study (w/ TM only, w/ DR only, w/ both TM & DR, and w/o both TM & DR) can be found in the main paper regarding the ListOps task (see Table 2). As other reviewers also suggested, further ablation studies, e.g., focusing on the loss curves would be insightful additions to the paper. We will add these experiments to the updated version of the paper.



> _How do TM and DR affect the performance in discrete VAE settings?_

We actually have tried our methods for computation graphs with single discrete nodes particularly for the Discrete VAE setting. We did not see improvements in this case, what we believe is the expected outcome. Sigmoid activation functions also tend to cause vanishing gradient problems mainly in the setup of deep neural networks in which multiple sigmoid activations are used consecutively. Similarly, our methods have the most impact in stochastic computation graphs with multiple consecutive discrete nodes.



> _There is no confidence interval in the multi-hop reasoning over KG experiments._

We followed the experimental setup and evaluation protocol of prior work, which typically does not provide confidence intervals. For the sake of completeness and good practice, we will add all missing confidence intervals.

> _However, the numberings of references in appendix does not match the ones from the main paper._

Good catch! Thank you.


> _Considering the multiple discrete stochastic layers, my guess is that the work is well-fit with combinatorial optimization problems such as [1]._

This seems to be very interesting related work. Solving or approximating combinatorial optimization problems with models based on complex stochastic computation graphs is interesting future work. We will add this to the related work and discussion/future work sections.

---

[1] W. Kool, H. Van Hoof, and M. Welling. Attention, learn to solve routing problems! 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019.


 We would like to thank the reviewer for the helpful review.

> _Overall I find that the paper proposes two reasonably justified yet rather incremental improvements on the problem of learning in the presence of discrete nodes._

We respectfully disagree with the proposed approaches being incremental. The proposed methods might be relatively straightforward / simple. However, the proposed methods are novel and allow us to efficiently train more complex discrete-continuous computations graphs, a problem that has not seen much attention in the community so far. We also provide (which we believe to be extensive) novel empirical evidence demonstrating that the proposed methods work well in several use cases. In fact, we show that some discrete-continuous computations graphs resulting from meaningful use cases can only be efficiently trained with the proposed methods.


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.  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.

Both 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.

The 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.

The paper is accompanied by code which should make the reported experiments easier to reproduce.

Overall I find that the paper proposes two reasonably justified yet rather incremental improvements on the problem of learning in the presence of discrete nodes. N/A

The paper proposes an efficient training method on the complex discrete stochastic computational graph, other than the single discrete stochastic layer.
The 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.
The athuors provide three experiments that backups their suggested methods.
  Originality, Quality and Significance:
My opinion on the originality, the quality, and the significance of this paper is complicated, hence I rather write them down in one part.
If 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.
For example, how does the performance varies as the selection of \beta in TempMatch, or \alpha in DropRes.
In this case, the quality of the paper matters.
Or, 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.
In this case, the significance becomes the problem.
Regarding the above, I think this paper is located in the fuzzy area.
However, the idea and results seem promising, so the paper will be in a better form after revision.

Clarity:
While several definitions on the abbreviations are missing, the paper is well-written in general.
However, the numberings of references in appendix does not match the ones from the main paper.


Comments:
- The hypothesis on the gradient vanishing seems true, but it would be better if there are experimental backups.
- Is there any particular reason for 
- 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.)
- 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.
- Considering the multiple discrete stochastic layers, my guess is that the work is well-fit with combinatorial optimization problems such as [1].
- The code is provided together with the generated ListOps dataset for reproducibility.
- There is no confidence interval in the multi-hop reasoning over KG experiments.

Questions:
- What is the specific choice of DropRes rate \alpha for each experiment?
- 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?
- How is the probabilistic simplex varies as the epoch grows when using TM?
- 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)
- 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)
- How do TM and DR affect the performance in discrete VAE settings?

[1] Attention, Learn to Solve Routing Problems!, Kool et al., ICLR 2019.
 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.

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.  STRENGTHS

- The paper is well motivated and addresses a general and important problem (discrete reasoning).

- 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.

- 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.


WEAKNESSES 

- 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.

- 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. 


QUESTIONS/COMMENTS

- In Table 4, why not include paths <= 5? 

- Despite the issues above, I think the submission is a clear accept given its positive finding and the amount of substance. 
 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. 

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