Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
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.

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.
limitations_and_societal_impact: N/A
review_title: null

final_score: rating: 4: Ok but not good enough - rejection
review_content: summary: 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.

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

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_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
main_review: 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. 

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_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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).

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.
main_review: ### 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).
limitations_and_societal_impact: 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.
review_title: null

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)