用于3D分子生成的等变扩散模型ICML2022

现代深度学习方法开始对分子科学产生重要影响。在Alphafold在蛋白质折叠预测方面取得成功的背后，越来越多的工作开发了深度学习模型来分析或合成分子。分子存在于物理3D空间中，因此受制于几何对称性，如平移、旋转。这些对称性被称为三维欧几里得群：E(3)。对于泛化能力，利用分子数据中的这些对称性是重要的。

来自：Equivariant Diffusion for Molecule Generation in 3D

背景概述

目前已经有研究将等变神经网络和自回归模型、流模型相结合用于 3D 分子生成，但存在采样和训练昂贵的问题。在这项工作中，设计了 E(3)等变的扩散模型（EDM）。EDM 学习对连续坐标和原子类别都起作用的逆扩散过程。该方法不需要特定的原子顺序（与自回归模型相比），并且可以比 Normalizing flows 更有效地进行训练。举个例子，当在 QM9 数据集上训练时，EDMs 产生的稳定分子比 E-NFs 多16倍，同时只需要一半的训练时间。这使得 EDM 能够在更大的药物数据集（如 GEOM-Drugs）上训练。

• 图1：EDM的概述。为了生成一个分子，将一组正态分布的点去噪成一个由3D原子坐标$x$和原子类型$h$组成的分子。由于模型是旋转等变的，所以当分子旋转$R$时，概率密度不变。

EDM: E(3)等变扩散模型

• 图2：为了生成分子，坐标$x$和特征$h$是通过从标准正态噪声$z_T$开始，逐步去噪到$z_t$得到。这是通过从分布$p(z_{t-1}|z_t)$中采样实现的。为了训练模型，在步骤$t$中使用$q(z_t|x,h)$将噪声添加到数据$x$和$h$，然后训练网络学习去噪。
• EDM将分子建模为点云，化学键被忽略，生成点云后，化学键由原子间距离和原子类型预测得到。

扩散过程服从正态分布：$$q(z_t|x,h)=N(z_t|{\alpha }_t[x,h],{\sigma }_t^{2}I)$$其中，$z_t=[z_t^{(x)},z_t^{(h)}]$，分别表示第$t$步扩散时的原子坐标和特征。

生成过程需要学习逆向的条件分布：$$p(z_s|z_t)=N(z_s|{\mu }_{t\to s}([\hat{x},\hat{h}],z_t),{\sigma }_{t\to s}^{2}I)$$$$[\hat{x},\hat{h}]=z_t/{\alpha }_t-{\widehat{ϵ}}_t\cdot {\sigma }_t/{\alpha }_t$$通过神经网络学习扩散过程中向样本添加的噪声${ϵ}_t$，神经网络的预测为$\widehat{ϵ}=[{\widehat{ϵ}}^{(x)},{\widehat{ϵ}}^{(h)}]$。只要神经网络满足等变性，则整个生成过程具有等变性。因此神经网络采用EGNN实现。

EDM的优化和采样算法如下：

个人理解

个人认为EDM和GEODIFF相比，存在一些差距，GEODIFF本身包含了条件生成，而EDM更像是简单的扩散模型思想结合等变神经网络。而且EDM关于等变网络的设计似乎是不完善的，因为GEODIFF的作者在GEODIFF论文中描述过，仅通过等变网络学习噪声（噪声没有等变方差）会违背等变网络的性质。

另外，EDM还需要通过原子信息预测边的信息，这可能会是影响预测结果稳定性的一个因素。