D3MES: Diffusion Transformer with multihead equivariant self-attention for 3D molecule generation

2025-01-13

AI Summary

Existing 3D molecular generation methods face significant bottlenecks in modeling complex molecules, jointly generating across chemical categories (e.g., drug-like and materials-related small molecules), and accurately adding hydrogen atoms—limiting their utility in drug and materials discovery. To address these challenges, we propose Diffusion Transformer: the first end-to-end architecture integrating class-conditional diffusion with SE(3)-equivariant multi-head self-attention, enabling dehydrogenated molecular graph representation learning and joint generation of diverse molecular classes. The model inherently satisfies SE(3) equivariance while incorporating category-aware conditioning, supporting conformational diversity modeling. On benchmarks including GEOM-Drugs and QM9, it consistently outperforms state-of-the-art methods, achieving substantial improvements in FCD, COV, and MAT metrics. It demonstrates strong robustness and generalization, making it suitable for large-scale early-stage molecular generation and property-aware screening.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Understanding and predicting the diverse conformational states of molecules is crucial for advancing fields such as chemistry, material science, and drug development. Despite significant progress in generative models, accurately generating complex and biologically or material-relevant molecular structures remains a major challenge. In this work, we introduce a diffusion model for three-dimensional (3D) molecule generation that combines a classifiable diffusion model, Diffusion Transformer, with multihead equivariant self-attention. This method addresses two key challenges: correctly attaching hydrogen atoms in generated molecules through learning representations of molecules after hydrogen atoms are removed; and overcoming the limitations of existing models that cannot generate molecules across multiple classes simultaneously. The experimental results demonstrate that our model not only achieves state-of-the-art performance across several key metrics but also exhibits robustness and versatility, making it highly suitable for early-stage large-scale generation processes in molecular design, followed by validation and further screening to obtain molecules with specific properties.

Problem

• Molecular Structure Generation
• 3D Modeling
• Chemical Research Limitations

Innovation

• D3MES
• Multi-Type Molecular Generation
• Hydrogen Atom Addition