[论文解读] MUDiff: Unified Diffusion for Complete Molecule Generation

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

Machine Learning in Materials Science被引用 15

一句话总结

MUDiff 通过在二维图（边）和三维坐标（原子）上同时扩散，联合生成完整的分子表示，并使用新颖的 MUformer 以旋转平移等变方式进行去噪。

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.

研究动机与目标

• 推动联合生成二维（图连接性）和三维（几何）分子数据以获得完整的分子表示。
• 开发一个扩散框架，去噪连续分量（原子特征、坐标）和离散分量（边类型）两者。
• 引入 MUformer，一个等变图变换器，在整合二维和三维信息的同时保持三维旋转和平移对称性。
• 在生成与学习过程中实现对缺失的二维或三维数据的鲁棒性。
• 相比现有方法，展示生成分子在稳定性和多样性方面的提升。

提出的方法

• 提出 MUDiff：一个扩散模型，对原子特征和坐标施加连续噪声，对边类型施加离散噪声，联合去噪所有分量。
• 引入 MUformer，一个具有不变性通道和等变通道的统一变换器，在旋转平移约束下处理二维和三维分子数据。
• 定义训练目标，预测原子特征和坐标的噪声并对边类型进行分类。
• 采用将二维邻域、三维邻域和全局图特征融合的编码方案，以获得鲁棒表示。
• 应用 3D 特定的径向基函数和余弦截断函数来捕捉空间信息并确保等变性。
• 提供从完全嘈杂的潜在表示逐步去噪到完整分子的采样过程。

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

实验结果

研究问题

• RQ1 联合对二维图结构和三维坐标进行扩散是否比单一表示模型能产生更稳定且有效的分子？
• RQ2 如何设计一个对三维旋转平移等变同时有效整合二维和三维分子信息的变换器？
• RQ3 模型在生成或训练时对缺失的二维或三维数据是否具有鲁棒性？
• RQ4 通过联合建模二维和三维结构，在稳定性、唯一性等方面能实现何种性能提升？

主要发现

• MUDiff 生成的分子在稳定性方面比现有方法高出 7.9%（Sec 6.2）。
• MUDiff 在分子唯一性方面比现有方法提高了 2%（Sec 6.2）。
• MUformer 实现了原子特征、坐标和边类型的同时预测，且具备旋转平移等变性。
• 即使在仅有限的三维结构数据下训练，模型仍然有效，与在完整三维数据上训练的方法相比也具有竞争力（Sec 6.1）。
• 该方法在缺失二维或三维信息时也能独立工作，实现鲁棒的完整表示。