[论文解读] Geometric Latent Diffusion Models for 3D Molecule Generation

Minkai Xu, Alexander K. Powers+3 more|arXiv (Cornell University)|May 2, 2023

Model Reduction and Neural Networks被引用 32

一句话总结

GeoLDM 引入了一个潜在扩散框架，具有 SE(3)-等变潜在空间（不变标量和等变张量），用于3D分子生成，取得比以往方法更高的有效性与可控性。

ABSTRACT

Generative models, especially diffusion models (DMs), have achieved promising results for generating feature-rich geometries and advancing foundational science problems such as molecule design. Inspired by the recent huge success of Stable (latent) Diffusion models, we propose a novel and principled method for 3D molecule generation named Geometric Latent Diffusion Models (GeoLDM). GeoLDM is the first latent DM model for the molecular geometry domain, composed of autoencoders encoding structures into continuous latent codes and DMs operating in the latent space. Our key innovation is that for modeling the 3D molecular geometries, we capture its critical roto-translational equivariance constraints by building a point-structured latent space with both invariant scalars and equivariant tensors. Extensive experiments demonstrate that GeoLDM can consistently achieve better performance on multiple molecule generation benchmarks, with up to 7\% improvement for the valid percentage of large biomolecules. Results also demonstrate GeoLDM's higher capacity for controllable generation thanks to the latent modeling. Code is provided at \url{https://github.com/MinkaiXu/GeoLDM}.

研究动机与目标

• 通过在尊重回转平移等变性的几何感知潜在空间中建模，推进3D分子生成。
• 提出一个几何自编码器，生成既有不变量组件又有等变组件的潜在码。
• 在潜在空间中训练扩散模型以提高3D分子生成的有效性、多样性与可控性。
• 证明潜在 GeoLDM 在 QM9 和 DRUG 基准上提升生成指标并实现可控生成。
• 提供用于复现与进一步研究的公开代码。

提出的方法

• 构建一个几何自编码器，将3D分子几何映射到潜在空间，使用等变图神经网络（EGNN）实现每个节点的不变量（z_h）和等变（z_x）分量。
• 通过设计潜在转移为等变、潜在先验为不变量来确保 SE(3) 不变性，从而得到的边缘分布为 SE(3)-不变。
• 在结构化潜在空间上训练潜在扩散模型（LDM），去噪网络 epsilon_theta 以时间条件的EGNN实现。
• 采用两阶段训练：先训练几何 AE 以重建和正则化；然后固定编码器并训练潜在 DM；使用 p_theta(z_x,z_h) 进行生成，随后解码 p_xi(x,h|z_x,z_h)。
• 通过对属性向量 s 进行拼接而非交叉注意力，将对性质的条件化传递给去噪器和编码/解码器以实现可控生成。
• 可选地从训练分布 p(N) 采样分子尺寸 N 并据此生成。

Figure 1: Illustration of GeoLDM . The encoder ${\mathcal{E}}_{\phi}$ encodes molecular features ${\mathbf{x}},{\mathbf{h}}$ into equivariant latent variables ${\mathbf{z}}_{\textnormal{x}},{\mathbf{z}}_{\textnormal{h}}$ , and the latent diffusion transitions $q({\mathbf{z}}_{{\textnormal{x}},t},{\m

实验结果

研究问题

• RQ1 一个在几何感知潜在空间（同时包含不变量和等变分量）上工作的潜在扩散模型，是否可以提高3D分子生成质量？
• RQ2 在潜在空间中强制 SE(3) 等变性，是否比在原始原子特征空间中工作的潜在模型带来更高的化学有效性和结构多样性？
• RQ3 GeoLDM 能在目标化学性质上实现多大程度的可控生成，以及在标准3D分子基准上的表现如何？
• RQ4 在 QM9 和 DRUG 数据集上，GeoLDM 在无条件与有条件生成方面的表现如何？

主要发现

数据	原子 Sta (%)	分子 Sta (%)	有效性 (%)	有效性且唯一性 (%)	原子 Sta (%)	有效性 (%)
QM9	ENF	85.0	4.9	40.2	39.4	G-Schnet	95.7	68.1	85.5	80.3	GDM	97.0	63.2	-	-
QM9	GDM-aug	97.6	71.6	90.4	89.5	EDM	98.7	82.0	91.9	90.7
QM9	EDM-Bridge	98.8	84.6	92.0*	90.7	GraphLDM	97.2	70.5	83.6	82.7
QM9	GraphLDM-aug	97.9	78.7	90.5	89.5	GeoLDM	98.9	89.4	93.8	92.7
DRUG	GeoLDM	84.4	99.3

• GeoLDM 在基准测试中实现更高的生成质量，在 QM9/DRUG 的大型生物分子上比基线提高到最多 7% 的有效率。
• 同时具有不变量和等变潜在变量的潜在建模相对于仅使用标量的模型在有效性和多样性上表现更优。
• 有条件生成实验显示由于潜在空间设计和通过拼接进行的简单属性条件化而实现了更好的可控性。
• GeoLDM 在学习的边缘分布中展现出强烈的 SE(3)-不变性特征，有助于对转动和平移的泛化。
• 实证结果表明 GeoLDM 在多项指标（有效性、稳定性与 有效性×唯一性）上超过了若干前期3D分子生成方法。
• 作者在项目站点提供了用于复现与进一步探索的代码。

Figure 2: Molecules generated by GeoLDM trained on QM9 (left three) and DRUG (right four).