Learning a chemistry-aware latent space for molecular encoding and generation with a large-scale Transformer Variational Autoencoder

Abstract

Searching for molecules optimizing certain properties remains a key challenge due to the vastness of the chemical space, its discrete nature, and the limited availability of bioactivity data. One way to address these issues is to build a mapping from the chemical space to a continuous latent embedding space where efficient exploration and smooth interpolations become possible. Existing methods suffer from several limitations: failure to accurately reconstruct molecules, poorly structured latent space where whole regions cannot be decoded to valid molecular graphs or where distances fail to reflect chemical similarities, not to mention the unavailability of ready-to-use code or models trained on sufficiently large datasets, limiting their practical application. In this work, we provide a large-scale pre-trained Variational Autoencoder based on the Transformer architecture to convert small organic molecules to continuous fixed-size embeddings in a chemistry-aware structured latent space and back. With a 97% reconstruction rate and a 100% validity rate, the model provides a reliable mapping and is able to generate a high diversity of novel molecules. We introduce a training objective with a novel loss term explicitly enforcing embedding distances to reflect Tanimoto similarities between molecular fingerprints, leading to a smooth and well-structured latent space enabling seamless exploration and interpolation. We show that our model achieves solid performances on molecular optimization tasks compared to other generative models. We release an open-source, user-friendly implementation as a pip package and at github.com/3BioCompBio/chembed, along with pre-trained models, that can be readily used or fine-tuned for downstream drug design applications. Scientific contributionWe introduce a large-scale SELFIES Transformer VAE with a novel loss aligning latent space distances with chemical similarity. This yields a highly structured, chemistry-aware latent space, alongside high reconstruction accuracy. Our open-source release with pre-trained weights makes the model directly usable and adaptable for molecular design applications.