Efficient 3D Molecular Design with an E(3) Invariant Transformer VAE
Orion Dollar, Nisarg Joshi, Jim Pfaendtner, David A. C. Beck
Abstract
This work introduces a three-dimensional (3D) invariant graph-to-string transformer variational autoencoders (VAE) (Vagrant) for generating molecules with accurate density functional theory (DFT)-level properties. Vagrant learns to model the joint probability distribution of a 3D molecular structure and its properties by encoding molecular structures into a 3D-aware latent space. Directed navigation through this latent space implicitly optimizes the 3D structure of a molecule, and the latent embedding can be used to condition a generative transformer to predict the candidate structure as a one-dimensional (1D) sequence. Additionally, we introduce two novel sampling methods that exploit the latent characteristics of a VAE to improve performance. We show that our method outperforms comparable 3D autoregressive and diffusion methods for predicting quantum chemical property values of novel molecules in terms of both sample quality and computational efficiency.