Constrained Graph Variational Autoencoders for Molecule Design
Graphs are ubiquitous data structures for representing interactions between entities. With an emphasis on applications in chemistry, we explore the task of learning to generate graphs that conform to a distribution observed in training data. We propose a variational autoencoder model in which both encoder and decoder are graphstructured. Our decoder assumes a sequential ordering of graph extension steps and we discuss and analyze design choices that mitigate the potential downsides of this linearization. Experiments compare our approach with a wide range of baselines on the molecule generation task and show that our method is successful at matching the statistics of the original dataset on semantically important metrics. Furthermore, we show that by using appropriate shaping of the latent space, our model allows us to design molecules that are (locally) optimal in desired properties.
Summary: This paper describes a model for generating graph-structured data, with molecule generation being the example task. This model is based around a variational autoencoder whose encoder/decoder are designed to handle graph-structured data. The decoder builds a graph sequentially by starting from an arbitrary node and sampling edges to other nodes, which are placed in a queue; upon sampling an edge to a “stopping node,” the next node is taken from the queue and the process continues until there are no more nodes to expand. The distributions from which these samples are taken are a function of the graph state (notably, not the specific steps taken to arrive at the current state), where the state vectors are encoded using a gated graph neural network (GGNN). Additionally, masking functions can be specified that serve as hard constraints on the sorts of edges that may be sampled (in case these would lead to graphs that are disallowed, e.g. that would lead to impossible molecules). The encoder is also comprised of a GGNN. Training is accomplished using a standard VAE objective; however, for computational efficiency reasons the reconstruction loss is only estimated. Additionally, a term can be added to the loss which is a function of the latent representation that attempts to optimize some numerical quantity related to the graphs being generated; at test time, gradient ascent can be used to find a latent representation with a “better” value of this quantity. Experiments were run on three molecule datasets, with comparisons being made against a variety of graph generation baselines regarding generated molecule uniqueness, validity, and novelty, as well as how the distributions of generated molecule statistics compare against the training dataset.  Comments: (using the same citation numbering as in the paper) This work is primarily a synthesis of ideas and techniques explored in other works: using a GGNN to model graphs as in [21]; training using a VAE objective as in [27]; enforcing hard constraints on generated graphs as in [27]; and learning a model that is invariant to graph construction order as in [27]. The real strength of this paper is that it combines these ideas and makes tweaks that allow the model to function well in practice, which seems to be primarily a combination of the modeling strength of GGNN (demonstrated by comparing the GGNN-based models against the non GGNN-based models during experiments) with the masking of invalid edges during generation and construction-order invariance. The provided experiments provide sufficient evidence that the proposed model has absorbed the benefits of these prior models (though it would have been nice to see more comparisons against [21] to ensure that the better results of this model are not an artifact of the one dataset used).  There are a few things I felt were missing that would have strengthened this paper a bit. The optimization of task-specific metrics (as described in section 4.3) was somewhat underexplored - section 6.2 presents a small experiment demonstrating that it is possible to learn a network that can optimize for a given metric, but only a few samples are provided. A full quantitative evaluation of this portion of the model would have been interesting to see (e.g. an ablation comparing the average metric for generated molecules with and without this additional modeling step and whether doing so hurts performance on other metrics at all). Additionally, details of training were completely omitted, including training hyperparameters and how they were chosen - this hurts reproducibility. Some further ablations related to training details, e.g. whether or not the number of path samples taken of the training data, would also have been interesting to see. Overall, though, these omissions do not detract significantly, and I feel that this paper is very worthy of submission.  It should also be noted that this paper is very well written, with the specific details presented in an exceptionally clear manner. None of the details presented were confusing.  POST AUTHOR RESPONSE EDIT: Thanks for the response. My opinions remain more or less the same, as your response did not address my primary concerns - that a comparison against Li et al on more datasets would further justify your claims and that the experiments in section 6.2 are somewhat preliminary in nature. However, I do appreciate the additional ablation experiments you did run, and your discussion of the evaluation metrics was important and should probably be added to the discussion presented during the experiments.