Manipulating 3D Molecules in a Fixed-Dimensional E(3)-Equivariant Latent Space
Zitao Chen, Yinjun Jia, Zitong Tian, Wei-Ying Ma, Yanyan Lan
TL;DR
MolFLAE presents a fixed-length, E(3)-equivariant VAE for 3D molecular data with a Bayesian Flow Network decoder, enabling unconditional generation and broad zero-shot editing in a latent space that is independent of atom counts. The framework supports analog design, structure reconstruction across shapes, and smooth latent interpolations, demonstrated on QM9, GEOM-Drugs, and a drug-optimization task targeting the human glucocorticoid receptor. Key contributions include the fixed-dimensional latent construction with learnable virtual nodes, partial latent-space disentanglement between geometry and substructure, and a practical hGR drug-design case showing balanced potency and hydrophilicity. The results underscore MolFLAE’s robustness, flexibility, and potential for real-world molecular editing and optimization, while outlining future directions for inpainting, superposition, and improved latent disentanglement.
Abstract
Medicinal chemists often optimize drugs considering their 3D structures and designing structurally distinct molecules that retain key features, such as shapes, pharmacophores, or chemical properties. Previous deep learning approaches address this through supervised tasks like molecule inpainting or property-guided optimization. In this work, we propose a flexible zero-shot molecule manipulation method by navigating in a shared latent space of 3D molecules. We introduce a Variational AutoEncoder (VAE) for 3D molecules, named MolFLAE, which learns a fixed-dimensional, E(3)-equivariant latent space independent of atom counts. MolFLAE encodes 3D molecules using an E(3)-equivariant neural network into fixed number of latent nodes, distinguished by learned embeddings. The latent space is regularized, and molecular structures are reconstructed via a Bayesian Flow Network (BFN) conditioned on the encoder's latent output. MolFLAE achieves competitive performance on standard unconditional 3D molecule generation benchmarks. Moreover, the latent space of MolFLAE enables zero-shot molecule manipulation, including atom number editing, structure reconstruction, and coordinated latent interpolation for both structure and properties. We further demonstrate our approach on a drug optimization task for the human glucocorticoid receptor, generating molecules with improved hydrophilicity while preserving key interactions, under computational evaluations. These results highlight the flexibility, robustness, and real-world utility of our method, opening new avenues for molecule editing and optimization.