Paper:
TREE-STRUCTURED VARIATIONAL AUTOENCODER
Many kinds of variable-sized data we would like to model contain an internal hierarchical structure in the form of a tree, including source code, formal logical statements, and natural language sentences with parse trees. For such data it is natural to consider a model with matching computational structure. In this work, we introduce a variational autoencoder-based generative model for tree-structured data. We evaluate our model on a synthetic dataset, and a dataset with applications to automated theorem proving. By learning a latent representation over trees, our model can achieve similar test log likelihood to a standard autoregressive decoder, but with the number of sequentially dependent computations proportional to the depth of the tree instead of the number of nodes in the tree.
1 INTRODUCTION
A significant amount of recent and ongoing work has explored the use of neural networks for modeling and generating various kinds of data. Newer techniques like the variational autoencoder (Rezende et al., 2014; Kingma & Welling, 2013) and generative-adversarial networks (Goodfellow et al., 2014) enable training of graphical models where the likelihood function is a complicated neural network which normally makes it infeasible to specify and optimize the marginal distribution analytically. Another family of techniques involves choosing an ordering of the dimensions of the data (which is particularly natural for sequences such as sentences) and training a neural network to estimate the distribution over the value of the next dimension given all the dimensions we have observed so far.
These techniques have led to significant advances in modeling images, text, sounds, and other kinds of complicated data. Language modeling with sequential neural models have halved perplexity (roughly, the error at predicting each word) compared to n-gram methods (Jozefowicz et al., 2016). Neural machine translation using sequence-to-sequence methods have closed half of the gap in quality between prior machine translation efforts and human translation (Wu et al., 2016). Generative image models have similarly progressed such that they can generate samples largely indistinguishable from the original data, at least for relatively small and simple images (Gregor et al., 2015; 2016; Kingma et al., 2016; Salimans et al., 2016; van den Oord et al., 2016), although the quality of the model here is harder to measure in an automated way (Theis et al., 2015).
However, many kinds of data we might wish to model are naturally structured as a tree. Computer program code follows a rigorous grammar, and the usual first step in processing it involves parsing it into an abstract syntax tree, which simultaneously discards aspects of the code irrelevant to the semantics such as whitespace and extraneous parentheses, and makes it more convenient to further interpret, analyze, or manipulate. Statements made in formal logic similarly have a hierarchical structure, which determines arguments to predicates and functions, the scoping of variables and quantifiers, and the application of logical connectives. Natural language sentences also contain a latent syntactic structure which is necessary for determining the meaning of the sentence, although
∗Majority of work done while at Google.
most sentences admit many possible parses due to the multiple meanings of each word and ambiguous groupings.
In this paper, we explore how we can adapt the variational autoencoder to modeling tree-structured data. In general, it is possible to treat the tree as a sequence and then use a sequential model. However, a model which follows the structure of the tree may better capture long-range dependencies: recurrent models sometimes have trouble learning to remember and use information from the distant past when it is relevant to the current context, but these distant parts of the input may be close to each other within the tree structure.
In our proposed approach, we decide the identity of each node in the tree by using a top-down recursive neural network, causing the distributed representation which decides the identity of each node in the tree to be computed as a function of the identity and relative location of its parent nodes.
By using the architecture of the variational autoencoder (Rezende et al., 2014; Kingma & Welling, 2013), our model can learn to capture various features of the trees within continuous latent variables, which are added as further inputs into the top-down recursive neural network and conditions the overall generation process. These latent variables allow us to generate different parts of the tree in parallel; specifically, given a parent node n and its children c1 and c2, the generation of (the distribution placed over different values of) c1 and its descendants is independent of the generation of c2 and its descendants (and vice versa), once we condition upon the latent variables. By structuring the model this way, while our model generates one dimension (the identity of each node within the tree) of the data a time, it is not autoregressive as the probability distribution for one dimension is not a function of the previously generated nodes.
We evaluate our model on a variety of datasets, some synthetic and some real. Our experimental results show that it achieves comparable test set log likelihood to autoregressive sequential models which do not use any latent variables, while offering the following properties:
• For balanced trees, generation requires O(log n) rather than O(n) timesteps required for a sequential model because the children of each node can be generated in parallel. • It is straightforward to resample a subtree while keeping the other parts of the tree intact. • The generated trees are syntactically valid by construction. • The model produces a latent representation for each tree, which may prove useful in other
applications.
2 BACKGROUND AND RELATED WORK
2.1 TREE-STRUCTURED MODELS
Recursive neural nets, which processes a tree in a bottom-up way, have been popular in natural language processing for a variety of tasks, such as sentiment analysis (Socher et al., 2013), question answering (Iyyer et al., 2014), and semantic relation extraction (Socher et al., 2012). Starting from the leaves of the tree, the model computes a representation for each node by combining the representations of its child nodes. In the case of natural language processing, each tree typically represents one sentence, with the leaf nodes corresponding to the words in the sentence and the structure of the internal nodes determined by the constituency parse tree for the sentence.
If we restrict ourselves to binary trees (given that it is possible to binarize arbitrary trees in a lossless way), the we compute the k-dimensional representation rn ∈ Rk for a node n by combining the representations of nodes nleft and nright:
rn = f(Wrnleft + V rnright)
where W and V are square matrices (in Rd×d) and f is a nonlinear activation function, applied elementwise. Leaf nodes are represented by embedding the content of the leaf node into a ddimensional vector, by specifying a lookup table from words to embedding vectors for instance.
Variations and extensions of this approach specify more complicated relationships between rn and the childen’s representations rnleft and rnright , or allow internal nodes to have variable numbers of children. For example, Tai et al. (2015) extend LSTMs to tree-structured models by dividing the vector
representation of each node into a memory cell and a hidden state and updating them accordingly with gating.
Neural network models which generate or explore a tree top-down have received less attention, but have been applied to generation and parsing tasks. Zhang et al. (2015) generate natural language sentences along with their dependency parses simultaneously. In their specification of a dependency parse tree, each word is connected to one parent word and has a variable number of left and right children (a depth-first in-order traversal on this tree recovers the original sentence). Their model generates these trees by conditioning the probability of generating each word on its ancestor words and the previously-generated sibling words using various LSTMs. Dyer et al. (2016) generate sentences jointly with a corresponding constituency parse tree. They use a shift-reduce parser architecture where shift is replaced with word-generation action, and so the sentence and the tree can be generated by performing a sequence of actions corresponding to a depth-first pre-order traversal of the tree. Each action in the sequence is predicted based upon the tree constructed so far, the words (the tree’s terminal nodes) generated so far, and the previous actions performed. Dong & Lapata (2016) generate tree-structured logical forms using LSTMs, where the LSTM state branches along with the tree’s structure; they focus on generating these logical forms when conditioned upon a natural-language description of it.
2.2 VARIATIONAL AUTOENCODERS
The variational autoencoder (Kingma & Welling, 2013; Rezende et al., 2014), or VAE for short, provides a way to train a generative model with a fixed prior p(z) and a neural network used to specify pθ(x | z). Typically, the prior p(z) is taken to be a standard multivariate normal distribution (mean at 0) with diagonal unit covariance. Naively, in order to optimize log p(x), we need to compute the following integral:
log pθ(x) = log ∫ z pθ(x | z)p(z)dz
which can be tractable when pθ(x | z) is simple but not when we want to use a neural network to represent it. Inference of the posterior p(z | x) also becomes intractable. Instead, we learn a second neural network qφ(z | x) to approximate the true posterior, and use the following variational bound:
log p(x) ≥ −DKL(qφ(z | x) ‖ p(z)) + Eqφ(z|x)[log pθ(x | z)]
where DKL represents the Kullback-Leibler divergence between the two distributions. Given that we represent qφ(z | x) with a neural network which outputs the mean and diagonal covariance for a normal distribution, we can analytically compute the KL divergence term and then use the reparameterization trick:
Eqφ(z|x)[log pθ(x | z)] = Ep( )[log pθ(x | z = µ+ σ · )]
where p( ) is a standard multivariate normal distribution, and µ and σ are outputs of the neural network implementing qφ(z | x). These two techniques combined allow us to compute stochastic gradients (by sampling , treating it as constant, and backpropagating through the model) and use standard neural network training techniques (such as SGD, Adagrad, and Adam) to train the model.
Another interpretation of the variational autoencoder follows from a modification of the regular autoencoder, where we would like to learn a mapping x → z from the data to a more compact representation z, and an inverse mapping z → x. In the VAE, we replace the deterministic x → z with a probabilistic q(z | x), and as a form of regularization, we ensure that this distribution is close to a prior p(z).
3 TREE-STRUCTURED VARIATIONAL AUTOENCODER
In this section, we describe how we combine the variational autoencoder and recursive neural networks in order to build our model.
3.1 TREES
We consider arbitrarily-branching typed trees where each node contains a type, and either child nodes or a terminal value. Each type may be a terminal type or a non-terminal type; nodes of terminal type contain a value, and nodes of non-terminal type have (zero or more) child nodes.
A non-terminal type T comes with a specification for how many children a node NT of type T should have, and the types permissible for each child location. We distinguish three types of child nodes:
• NT may have some number of singular child nodes. For the ith singular child, we specify SINGULARCHILD(T, i) = {T1, · · · , Tn} as the set of types that child node can have. If the singular child node is optional, we denote this by including φ in this set. SINGULARCHILDCOUNT(T ) gives the number of singular child nodes in T .
• NT may have an arbitrary number of repeated child nodes. Each repeated child node must have type belonging within REPEATEDCHILDREN(T ) = {T1, · · · }. If this set is empty, no repeated child nodes are allowed. These children may be of different types.
For each terminal type, we have a list of values that a node of this type can have. We also have a list of types that the root node can have.
The above specification serves as an extension of context-free grammars, which are commonly used to specify formal languages. The main difference is in optional and repeated children, which makes it easier to specify an equivalent grammar with fewer non-terminal types.
As an example, consider the for loop in the C programming language. A node representing this contains three singular children: an initializer expression, the condition expression (evaluated to check whether the loop should continue running), and the iteration statement, which runs at the end of each loop iteration. It also has repeated children, one child per statement in the loop body.
3.2 BUILDING A TREE
Now that we have specified the kinds of trees we consider, let us look at how we might build one. First, we describe the basic building block that we use to create one node, then we will look at how to compose these together to build an entire tree.
Assume that we know the node that we are about to construct should have type T , and that we have a hidden state h ∈ Rk which contains further information about the node.
• If T is a terminal type, we use WHICHTERMINALVALUET (h), producing a probability distribution over the set of possible values, and sample from this distribution to choose the value.
• If T is a non-terminal type, we use the following procedure GENERATENODE(T,h): 1. Compute m = SINGULARCHILDCOUNT(T ) + 1{REPEATEDCHILDREN(T ) 6= ∅}.
In other words, count the number of singular children, and add 1 if the type allows repeated children.
2. Compute g1, · · · ,gm = SPLITT (h). The function SPLITT (h) : Rk → Rk×· · ·×Rk maps the k-dimensional vector h into n separate k-dimensional vectors g1 to gm.
3. For each singular child i: (a) Sample Ti ∼ WHICHCHILDTYPET,i(gi) from a distribution over the types in
REQUIREDCHILD(T, i). (b) If Ti 6= ∅, use GENERATENODE(Ti,gi) to build the child node recursively.
4. If T specifies repeated children: (a) Compute gcur,gnext = SPLITREPEATEDT (gm). (b) Sample s ∼ STOPREPEATT (gcur) from a Bernoulli distribution. If s = 1, stop
generating repeated children. (c) Sample Tchild ∼ WHICHCHILDTYPET,repeated(gcur), a probability distribution
over the types in REPEATEDCHILDREN(T ). (d) Use GENERATENODE(Tchild,gcur) to build this child recursively. (e) Set gm := gnext and repeat this loop.
For building the entire tree starting from the root, we assume that we have an embedding z which encodes information about the entire tree (we describe how we obtain this in the next section). We sample Troot ∼ WHICHROOTTYPE(z), the type for the root node, and then run GENERATENODE(Troot, z).
3.3 ENCODING A TREE
Recall that the variational autoencoder framework involves training two models simultaneously: p(x | z), the generative (or decoding) model, and q(z | x), the inference (or encoding) model. The previous section described how we specify p(x | z), so we now turn our attention to q(z | x). Overall, we build the inference model by inverting the flow of data in the generative model. Specifically, we use ENCODE(n) to encode a node n with type T :
• If T is a terminal type, return EMBEDDING(v) ∈ Rk by performing a lookup of the contained value v within a table.
• If T is a non-terminal type: 1. Compute gi = ENCODE(ni) for each singular child ni of n. If ni is missing, then gi = 0.
2. If T specifies repeated children, set grepeated := 0 and nchild to the last repeated child of n, and then run: (a) Compute gchild = ENCODE(nchild). (b) Set grepeated := MERGEREPEATEDT (grepeated, gchild) ∈ Rk. (c) Move nchild to the previous repeated child, and repeat (until we run out of repeated children). 3. Return MERGET (g1, . . . , gm, goptional, grepeated) ∈ Rk.
Thus hroot = ENCODE(nroot) gives a summary of the entire tree as a k-dimensional embedding. We then construct q(z | x) = N (µ, σ) where µ = Wµhroot and σ = softplus(Wσhroot). Applying softplus(x) = log(1 + ex) as a nonlinearity gives us a way to ensure that σ is positive as required.
3.4 IMPLEMENTING SPLIT, MERGE, AND WHICH FUNCTIONS
In the previous two sections, we described how the model traverses the tree bottom-up to produce an encoding of it (q(z | x)), and how we can generate a tree top-down from the encoding (p(x | z)). In this section, we explicate the details of how we implemented the SPLIT, MERGE, and WHICH functions that we used previously.
COMBINE. We can consider SPLIT : Rk → Rk × · · · × Rk and MERGE : Rk × · · · × Rk → Rk functions to be specializations of a more general function COMBINE : Rk×· · ·×Rk → Rk×· · ·×Rk which takes m inputs and produces n outputs.
A straightforward implementation of COMBINE is the following: y1, . . . ,yn := COMBINE(x1, . . . ,xm)
[y1 · · · yn] = f(W [x1 · · · xm] + b) where we have taken xi and yi to be column vectors Rk, [x1 · · ·xm] stacks the vectors xi vertically, W ∈ Rn·k×m·k and b ∈ Rn·k are the learned weight matrix and bias vector respectively, and f is a nonlinearity applied elementwise.
For WHICH : Rk → Rd, which computes a probability distribution over d choices, we use a specialization of COMBINE with one input and one (d-sized rather than k-sized) output, and use softmax as the nonlinearity f .
While this basic implementation sufficed initially, we discovered that two modifications led to better performance, which we describe subsequently.
Gating. We added a form of multiplicative gating, similar to those used in Gated Recurrent Units (Chung et al., 2014), Highway Networks (Srivastava et al., 2015), and Gated PixelCNN (van den Oord et al., 2016). The multiplicative gate enables the COMBINE function to more easily pass through information in the inputs to the outputs if that is preferable to transforming the input. Furthermore, the multiplicative interactions used in computing the gates may help the neural network learn more complicated functions.
First, we compute candidate values for yi using a linear layer as before: [ŷ1 · · · ŷn] = f(W [x1 · · · xm] + b)
Then we compute multiplicative gates for each ŷi and each (xi,yj) combination, or (m+ 1)n gate variables (recall that m is the number of inputs and n is the number of outputs).
[gy1 · · · gyn ] = σ(Wgy [x1 · · · xm] + bgy )[ g(x1,y1) · · · g(x1,yn) ] = σ(Wg1 [x1 · · · xm] + bg1)
...[ g(xm,y1) · · · g(xm,yn) ] = σ(Wgm [x1 · · · xm] + bgm)
Then we compute the final outputs yi: yi = gyi ŷi + g(x1,yi) x1 + · · ·+ g(xm,yi) xm. σ is the sigmoid function σ(x) = 1/(1 + e−x) and is the elementwise product. We initialized bgi = 1 and bgy = −1 so that gyi would start out as a small value and g(xi,yj) would be large, encouraging copying of the inputs xi to the outputs yi.
Layer normalization. We found that using layer normalization (Ba et al., 2016) also helps stabilize the learning process. For our model, it is difficult to use batch normalization because the connections of each layer (the functions MERGE, SPLIT, WHICH) occur at variable points according to the particular tree we are considering.
Following the procedure in the appendix of Ba et al. (2016), we replace each instance of f(W [x1 · · · xm] + b) with f(LN(W1x1;α1) + · · · + LN(Wmxm;αm) + b) where Wi ∈ Rnk×k are horizontal slices of W and αi ∈ R are learned multiplicative constants. We use LN(z;α) = α · (z − µ)/σ where µ ∈ R is the mean of z ∈ Rk and σ ∈ R is the standard deviation of z.
3.5 WEIGHT SHARING
In the above model, each function with a different name has different weights. For example, if we have two types PLUS and MINUS each with two required children, then SPLITPLUS and SPLITMINUS will have different weights even though they both have the same signature Rk → Rk × Rk. However, this may be troublesome when we have a very large number of types, because in this scheme the amount of weights increases linearly with the number of types. For such cases, we can apply some of the following modifications:
• Replace all instances of SPLITT : Rk → Rk × · · · × Rk and SPLITREPEATED with a single SPLITREC : Rk → Rk×Rk. We can apply SPLITREC recurrently to get the desired number of child embeddings.
• Similarly, replace instances of MERGE with MERGEREC. • Share weights across the WHICH functions: a WHICH function which produces a distribu-
tion over T1, . . . , Tn contains weights and a bias for each Ti. We can share these weights and biases across all WHICH functions where Ti appears.
3.6 VARIABLE-SIZED LATENT STATE
In order to achieve low reconstruction error Eqφ(z|x)[log pθ(x | z)], the encoder and decoder networks must learn how to encode all information about a tree in z and then be able to reproduce the tree from this representation. If the tree is large, it becomes a difficult optimization problem to learn how to do this effectively, and may require higher-capacity networks in order to succeed at all which would require more time to train.
Instead, we can encode the tree with a variable number of latent state vectors. For each node ni in the tree, we specify q(zni | x) = N (µni , σni) where[
µni σni
] = ( id
softplus )[ Wµ Wσ ] ENCODE(ni)
Then when computing GENERATENODE(T,h), we first sample zni ∼ q(zni | x) at training time or zni ∼ p(z) at generation time, and then use ĥ = MERGELATENT(h, zni) in lieu of h. We fixed the prior of each latent vector zi to be the standard multivariate normal distribution with diagonal unit covariance, and did not investigate computing the prior as a function of other samples of zi as in Chung et al. (2015) or Fraccaro et al. (2016) which also used a variable number of latent state vectors.
4 EXPERIMENTS
4.1 TYPE-AWARE SEQUENTIAL MODEL
For purposes of comparison, we implemented a standard LSTM model for generating each node of the tree sequentially with a depth-first traversal, similar to Vinyals et al. (2015). The model receives each non-terminal type, and terminal value, as a separate token. We begin the sequence with a special 〈BOS〉 token. Whenever an optional child is not present, or at the end of a sequence of repeated children, we insert 〈END〉. This allows us to unambiguously specify a tree following a given grammar as a sequence of tokens.
At generation time, we keep track of the partially-generated tree in order to only consider those tokens which would be syntactically allowed to appear at that point. We also tried using this constraint at training time: when computing the output probability distribution, only consider the syntactically allowed tokens and leave the unnormalized log probabilities of the others unconstrained. However, we found that for our datasets, this did not help much with performance and led to overfitting.
4.2 SYNTHETIC ARITHMETIC DATA
To evaluate the performance of our models in a controlled way, we created a synthetic dataset consisting of arithmetic expressions of a given depth which evaluate to a particular value.
Grammar. We have two non-terminal types, PLUS and MINUS, and one terminal type NUMBER. PLUS and MINUS have two required children, left and right, each of which can be any of PLUS, MINUS, or NUMBER. For NUMBER, we allowed terminal values 0 to 9.
Generating the data. We generate trees with a particular depth, defined as the maximal distance from the root to any terminal node. As such, we consider 1 + (2 + 3) and (1 + 2) − (3 + 4) to both have depth 3. To get trees of depth d which evaluate to v, we first sampled 1,000,000 trees
uniformly at random from all binary tree structures up to depth d− 1, and randomly assigning each non-terminal node to PLUS or MINUS and setting each terminal node to a random integer between 0 and 9. Then we randomly pick two such trees, which when combined with PLUS or MINUS, evaluate to v to build a tree of depth d.
As training data, we generated 100,000 trees of depth 5, 7, 9, and 11. Within each set of trees, each quarter evaluates to −10, −5, 5, and 10 respectively. We use a test set of 1,024 trees, which we generated by first sampling a new set of 1,000,000 subtrees independently.
Results. Table 1 shows statistics on the datasets and the experimental results we obtained from training various models. The tree variational autoencoder model achieves better performance on deeper trees. In particular, the sequential model fails to learn well on depth 11 trees. However, it appears that a tree-structured model but with a fixed z performs similarly, although consistently worse than with the VAE.
4.3 FIRST-ORDER LOGIC PROOF CLAUSES
We next consider a dataset derived from Alemi et al. (2016): fragments of automatically-generated proofs for mathematical theorems stated in first-order logic. An automated theorem prover tries to prove a hypothesis given some premises, producing a series of steps until we conclude the hypothesis follows from the premises. Many theorem provers work by resolution; it negates the hypothesis and shows that a contradiction follows. However, figuring out the intermediate steps in the proof is a highly nontrivial search procedure. If we can use machine learning to generate clauses which are likely to appear as a proof step, we may be able to speed up automated theorem proving significantly.
Grammar. We generate first-order logic statements which are clauses, or a disjunction of literals. Each literal either contains one predicate invocation or asserts that two expressions are equal. Each predicate invocation contains a name, which also determines a fixed number of arguments; each argument is an expression. An expression is either a function invocation, a number, or a variable. A function invocation is structurally identical to a predicate invocation.
We consider the set of functions and predicates to be closed. Furthermore, given that each function and predicate has a fixed number of arguments, we made each of these its own type in the grammar. To avoid having a very large number of weights as a consequence, we applied the modifications described in Section 3.5.
Results. Table 2 describes our results on this dataset. We trained on 955,529 trees and again tested on 1,024 trees. The sequential model demonstrates slightly better log likelihood compared to the tree variational autoencoder model. However, on this dataset we observe a significant improvement
in log likelihood by adding the variational autoencoder to the tree model, unlike on the arithmetic datasets.
5 DISCUSSION AND FUTURE WORK
Conditioning on an outside context. In many applications for modeling tree-structured data, we have an outside context that informs which trees are more likely than others. For example, when generating clauses which may appear in a proof, the hypothesis in question greatly influences the content of the clauses. We leave this question to future work.
Scaling to larger trees. Currently, training the model requires processing an entire tree at once, first to encode it into z and then to decode z to reproduce the original tree. This can blow up the memory requirements, requiring an undesirably small batch size. For autoregressive sequence models, truncated backpropagation through time provides a workaround as a partial form of the objective function can be computed on arbitrary subsequences. In our case, adapting methods from Gruslys et al. (2016) and others may prove necessary.
Improving log likelihood. In terms of log likelihood, our model performed significantly better than an autoregressive sequential model only on one of the datasets we tested, and about the same or slightly worse on the others. Adding depth to the tree structure (Irsoy & Cardie, 2014), and a more sophisticated posterior (Rezende & Mohamed, 2015; Kingma et al., 2016; Sønderby et al., 2016), are some modifications which might help with learning a more powerful model. Introducing more dependencies between the dimensions of x during generation is another possibility but one which may reduce the usefulness of of the latent representation (Bowman et al., 2015).
A MODEL HYPERPARAMETERS
We used the Adam optimizer with a learning rate of 0.01, multiplied by 0.98 every 10000 steps. We clipped gradients to have L2 norm 3. For the synthetic arithmetic data, we used a batch size of 64; for the first-order logic proof clauses, we used a batch size of 256. For the tree-structured variational autoencoder, we set k = 256 and used the ELU nonlinearity wherever another one was not explicitly specified. For the sequential models, we used two stacked LSTMs each with hidden state size 256, no dropout. We always unrolled the network to the full length of the sequence during training, and did not perform any bucketing of sequences by length.
B KL DIVERGENCE DYNAMICS DURING TRAINING
Optimizing the variational autoencoder objective turned out to be a significant optimization challenge, as pointed out by prior work (Bowman et al., 2015; Sønderby et al., 2016; Kingma et al., 2016). Specifically, it is easy for the KL divergence term DKL(qφ(z | x) ‖ p(z)) to collapse to zero, which means that qφ(z | x) is equal to the prior and does not convey any information about x. This leads to uninteresting latent representations and reduces the generative model to one that does not use a latent representation at all.
As explained by Kingma et al. (2016), this phenomenon occurs as at the beginning of training it is much easier for the optimization process to move qφ(z | x) closer to the prior p(z) than to improve p(x | z), especially when qφ(z | x) has not yet learned how to convey any useful information. To combat this, we use a combination of two techniques described in the previous work:
• Anneal the weight on the KL cost term slowly from 0 to 1. Similar to Bowman et al. (2015), our schedule was a shifted and horizontally-scaled sigmoid function.1
• Set a floor on the KL cost, i.e. use −max(DKL(qφ(z | x) ‖ p(z)), λ) instead of DKL(qφ(z | x) ‖ p(z)) in the objective (Kingma et al., 2016). This change means that the model receives no penalty for producing a KL divergence below λ, and as the other part of the objective (the reconstruction term) benefits from a higher KL divergence, it naturally learns a more informative qφ(z | x) at least λ in KL divergence.
We found that at least one of these techniques were required to avoid collapse of the KL divergence to 0. However, as shown in Figure 2, we found that different combinations of these techniques could led to different overall results, suggesting that finding the desired equilibrium necessitates a hyperparameter search.
1To anneal from a to b, we used σ ( step − a+b
2
) /10 to weight the KL cost as a function of the number of
optimization steps taken.

Review:
This paper introduces a novel extension of the variational autoencoder to arbitrary tree-structured outputs. Experiments are conducted on a synthetic arithmetic expression dataset and a first-order logic proof clause dataset in order to evaluate its density modeling performance.
Pros:
+ The paper is clear and well-written.
+ The tree-structure definition is sufficiently complete to capture a wide variety of tree types found in real-world situations.
+ The tree generation and encoding procedure is elegant and well-articulated.
+ The experiments, though limited in scope, are relatively thorough. The use of IWAE to obtain a better estimate of log likelihoods is a particularly nice touch.
Cons:
- The performance gain over a baseline sequential model is marginal.
- The experiments are limited in scope, both in the datasets considered and in the evaluation metrics used to compare the model with other approaches. Specifically: (a) there is only one set of results on a real-world dataset and in that case the proposed model performs worse than the baseline, and (b) there is no evaluation of the learned latent representation with respect to other tasks such as classification.
- The ability of the model to generate trees in time proportional to the depth of the tree is proposed as a benefit of the approach, though this is not empirically validated in the experiments.
The procedures to generate and encode trees are clever in their repeated use of common operations. The weight sharing and gating operations seem important for this model to perform well but it is difficult to assess their utility without an ablation (in Table 1 and 2 these modifications are not evaluated side-by-side). Experiments in another domain (such as modeling source code, or parse trees conditioned on a sentence) would help in demonstrating the utility of this model. Overall the model seems promising and applicable to a variety of data but the lack of breadth in the experiments is a concern.
* Section 3.1: "We distinguish three types" => two
* Section 3.6: The exposition of the variable-sized latent state is slightly confusing because the issue of how many z's to generate is not discussed.
* Section 4.2-4.3: When generating the datasets, did you verify that the test set is disjoint from the training set?
* Table 1: Is there a particular reason why the variable latent results are missing for the depth 11 trees?