GF-VAE: A Flow-based Variational Autoencoder for Molecule Generation

Generating novel molecules with desired properties is a fundamental problem in modern drug discovery. This is a challenging problem because it requires the optimization of the given objectives while obeying the rules of chemical valence. An effective approach is to incorporate the molecular graph with deep generative models. However, recent generative models with high-performance are still computationally expensive. In this paper, we propose GF-VAE, a flow-based variational autoencoder (VAE) model for molecular graph generation. Specifically, the model equips VAE a lightweight flow model as its decoder, in which, the encoder aims to accelerate the training process of the decoder, while the decoder in turns to optimize the performance of the encoder. Thanks to the invertibility of flow model, the generation process is easily accomplished by reversing the decoder. Additionally, the final generated molecules are processed by validity correction. Therefore, our GF-VAE inherits the advantages of both VAE and flow-based methods. We validate our model on molecule generation and reconstruction, smoothness of learned latent space, property optimization and constrained property optimization. The results show that our model achieves state-of-the-arts performance on these tasks. Moreover, the time performance of GF-VAE on two classical datasets can achieve 31.3% and 62.9% improvements separately than the state-of-the-art model.