Accelerating 3D Molecule Generative Models with Trajectory Diagnosis

Hanlin Wu (Institute of AI Industry Research (AIR), Tsinghua University) · Yuxuan Song (Tsinghua University; Bytedance) · Zhilong Zhang (Tsinghua University) · Hao Zhou (Bytedance AI Lab) · Wei-Ying Ma (Tsinghua University) · Dongzhan Zhou (Shanghai Artificial Intelligence Laboratory) · Jingjing Gong (Shanghai Innovation Institute ) · Yichun Wang (Industrial Engineering, Tsinghua University)
adjustment phaseatomic feature valuescompetitive performancecomputational efficiencyconsistency parameter objectivediffusion-based modelsgeneration inefficiencygenerative trajectorygeometric molecule generative modelsgeometric-informed priormolecular graph generationpermutation orderpermutation phasesampling stepsscalable molecular generationstate-of-the-art models

Geometric molecule generative models have found expanding applications across various scientific domains, but their generation inefficiency has become a critical bottleneck. Through a systematic investigation of the generative trajectory, we discover a unique challenge for molecule geometric graph generation: generative models require determining the permutation order of atoms in the molecule before refining its atomic feature values. Based on this insight, we decompose the generation process into permutation phase and adjustment phase, and propose a geometric-informed prior and consistency parameter objective to accelerate each phase. Extensive experiments demonstrate that our approach achieves competitive performance with approximately 10 sampling steps, 7.5 × faster than previous state-of-the-art models and approximately 100 × faster than diffusion-based models, offering a significant step towards scalable molecular generation.