Re-balancing Variational Autoencoder Loss for Molecule Sequence Generation
This addresses a specific bottleneck in molecule generation for computational chemistry, offering an incremental improvement over existing methods.
The paper tackles the posterior collapse problem in RNN-based VAEs for molecule sequence generation by identifying underestimated reconstruction loss as the cause and proposing a solution, achieving state-of-the-art reconstruction accuracy and reducing generation time from 1450s to 9s for 10,000 unique valid SMILES on the ZINC 250K dataset.
Molecule generation is to design new molecules with specific chemical properties and further to optimize the desired chemical properties. Following previous work, we encode molecules into continuous vectors in the latent space and then decode the vectors into molecules under the variational autoencoder (VAE) framework. We investigate the posterior collapse problem of current RNN-based VAEs for molecule sequence generation. For the first time, we find that underestimated reconstruction loss leads to posterior collapse, and provide both theoretical and experimental evidence. We propose an effective and efficient solution to fix the problem and avoid posterior collapse. Without bells and whistles, our method achieves SOTA reconstruction accuracy and competitive validity on the ZINC 250K dataset. When generating 10,000 unique valid SMILES from random prior sampling, it costs JT-VAE1450s while our method only needs 9s. Our implementation is at https://github.com/chaoyan1037/Re-balanced-VAE.