[Paper Review] MUDiff: Unified Diffusion for Complete Molecule Generation

Chenqing Hua, Sitao Luan+5 more|arXiv (Cornell University)|Apr 28, 2023

Machine Learning in Materials Science15 citations

TL;DR

MUDiff jointly generates complete molecular representations by diffusion on both 2D graphs (edges) and 3D coordinates (atoms), using a novel MUformer to denoise in a roto-translation equivariant fashion.

ABSTRACT

Molecule generation is a very important practical problem, with uses in drug discovery and material design, and AI methods promise to provide useful solutions. However, existing methods for molecule generation focus either on 2D graph structure or on 3D geometric structure, which is not sufficient to represent a complete molecule as 2D graph captures mainly topology while 3D geometry captures mainly spatial atom arrangements. Combining these representations is essential to better represent a molecule. In this paper, we present a new model for generating a comprehensive representation of molecules, including atom features, 2D discrete molecule structures, and 3D continuous molecule coordinates, by combining discrete and continuous diffusion processes. The use of diffusion processes allows for capturing the probabilistic nature of molecular processes and exploring the effect of different factors on molecular structures. Additionally, we propose a novel graph transformer architecture to denoise the diffusion process. The transformer adheres to 3D roto-translation equivariance constraints, allowing it to learn invariant atom and edge representations while preserving the equivariance of atom coordinates. This transformer can be used to learn molecular representations robust to geometric transformations. We evaluate the performance of our model through experiments and comparisons with existing methods, showing its ability to generate more stable and valid molecules. Our model is a promising approach for designing stable and diverse molecules and can be applied to a wide range of tasks in molecular modeling.

Motivation & Objective

• Motivate joint generation of 2D (graph connectivity) and 3D (geometric) molecular data for a complete molecular representation.
• Develop a diffusion-based framework that denoises both continuous (atom features, coordinates) and discrete (edge types) components.
• Introduce MUformer, an equivariant graph transformer that preserves 3D roto-translation symmetry while integrating 2D and 3D information.
• Enable robustness to missing 2D or 3D data during generation and learning.
• Demonstrate improved stability and diversity of generated molecules compared to existing methods.

Proposed method

• Propose MUDiff: a diffusion model that adds continuous noise to atom features and coordinates and discrete noise to edge types, jointly denoising all components.
• Introduce MUformer, a unified transformer with invariant and equivariant channels to process 2D and 3D molecular data under roto-translation constraints.
• Define training objectives that predict noise for atom features and coordinates and classify edge types.
• Employ an encoding scheme that fuses 2D neighborhood, 3D neighborhood, and global graph features for robust representations.
• Apply 3D-specific radial basis and cosine cutoff functions to capture spatial information and ensure equivariance.
• Provide sampling procedures that progressively denoise from a fully noisy latent to a complete molecule.

Figure 1: The figure showcases our MUformer for processing 2D and 3D molecular data. Within the Transformer backbone, two channels exist: purple for 2D data and brown for 3D data. The blue part encodes 2D molecular structures, while the green part handles atom-level information and the red part proc

Experimental results

Research questions

• RQ1 Can joint diffusion over 2D graph structures and 3D coordinates produce more stable and valid molecules than single-representation models?
• RQ2 How can a transformer be designed to be equivariant to 3D roto-translation while effectively integrating 2D and 3D molecular information?
• RQ3 Is the model robust to cases where either 2D or 3D data are missing during generation or training?
• RQ4 What performance gains (stability, uniqueness) are achieved by jointly modeling 2D and 3D structures?

Key findings

• MUDiff generates molecules that are 7.9% more stable than existing methods (Sec 6.2).
• MUDiff increases molecular uniqueness by 2% compared to existing methods (Sec 6.2).
• MUformer enables simultaneous prediction of atom features, coordinates, and edge types with roto-translation equivariance.
• The model remains effective even when trained with limited 3D structures, showing competitive performance against methods trained on complete 3D data (Sec 6.1).
• The approach supports independent functioning when either 2D or 3D information is missing, enabling robust complete representations.