Improving Molecular Graph Generation with Flow Matching and Optimal Transport
Published in Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS 2024), AIDrugX Workshop, 2024
Author: Tian Zhu*, Xiaoyang Hou*, Milong Ren, Dongbo Bu, Xin Gao, Chunming Zhang, Shiwei Sun
Abstract: Generating molecular graphs is crucial in drug design and discovery but remains challenging due to the complex interdependencies between nodes and edges. While diffusion models have demonstrated their potentiality in molecular graph design, they often suffer from unstable training and inefficient sampling. To enhance generation performance and training stability, we propose GGFlow, a discrete flow matching generative model incorporating optimal transport for molecular graphs and it incorporates an edge-augmented graph transformer to enable the direct communications among chemical bounds. Additionally, GGFlow introduces a novel goal-guided generation framework to control the generative trajectory of our model, aiming to design novel molecular structures with the desired properties. GGFlow demonstrates superior performance on both unconditional and conditional molecule generation tasks, outperforming existing baselines and underscoring its effectiveness and potential for wider application.