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

TL;DR

This paper introduces D3MES, a diffusion transformer with multihead equivariant self-attention, for accurate and versatile 3D molecule generation, addressing hydrogen attachment and multi-class generation challenges.

Contribution

The work presents a novel diffusion transformer model that improves 3D molecule generation by learning hydrogen attachment and enabling multi-class molecule synthesis.

Findings

Achieves state-of-the-art performance on key metrics

Demonstrates robustness and versatility in molecular design

Effective in large-scale early-stage molecule generation

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.