用于3D分子生成的等变扩散模型ICML2022
现代深度学习方法开始对分子科学产生重要影响。在Alphafold在蛋白质折叠预测方面取得成功的背后，越来越多的工作开发了深度学习模型来分析或合成分子。分子存在于物理3D空间中，因此受制于几何对称性，如平移、旋转。这些对称性被称为三维欧几里得群：E(3)。对于泛化能力，利用分子数据中的这些对称性是重要的。
来自：Equivariant Diffusion for Molecule Generation in 3D
目录
- 背景概述
- EDM: E(3)等变扩散模型
- 个人理解
背景概述
目前已经有研究将等变神经网络和自回归模型、流模型相结合用于 3D 分子生成，但存在采样和训练昂贵的问题。在这项工作中，设计了 E(3)等变的扩散模型（EDM）。EDM 学习对连续坐标和原子类别都起作用的逆扩散过程。该方法不需要特定的原子顺序（与自回归模型相比），并且可以比 Normalizing flows 更有效地进行训练。举个例子，当在 QM9 数据集上训练时，EDMs 产生的稳定分子比 E-NFs 多16倍，同时只需要一半的训练时间。这使得 EDM 能够在更大的药物数据集（如 GEOM-Drugs）上训练。
- 图1：EDM的概述。为了生成一个分子，将一组正态分布的点去噪成一个由3D原子坐标 x x x和原子类型 h h h组成的分子。由于模型是旋转等变的，所以当分子旋转 R R R时，概率密度不变。
EDM: E(3)等变扩散模型
- 图2：为了生成分子，坐标 x x x和特征 h h h是通过从标准正态噪声 z T z_{T} zT开始，逐步去噪到 z t z_{t} zt得到。这是通过从分布 p ( z t − 1 ∣ z t ) p(z_{t-1}|z_{t}) p(zt−1∣zt)中采样实现的。为了训练模型，在步骤 t t t中使用 q ( z t ∣ x , h ) q(z_{t}|x,h) q(zt∣x,h)将噪声添加到数据 x x x和 h h h，然后训练网络学习去噪。
- EDM将分子建模为点云，化学键被忽略，生成点云后，化学键由原子间距离和原子类型预测得到。
扩散过程服从正态分布： q ( z t ∣ x , h ) = N ( z t ∣ α t [ x , h ] , σ t 2 I ) q(z_{t}|x,h)=N(z_{t}|\alpha_{t}[x,h],\sigma_{t}^{2}I) q(zt∣x,h)=N(zt∣αt[x,h],σt2I)其中， z t = [ z t ( x ) , z t ( h ) ] z_{t}=[z_{t}^{(x)},z_{t}^{(h)}] zt=[zt(x),zt(h)]，分别表示第 t t t步扩散时的原子坐标和特征。
生成过程需要学习逆向的条件分布： p ( z s ∣ z t ) = N ( z s ∣ μ t → s ( [ x ^ , h ^ ] , z t ) , σ t → s 2 I ) p(z_{s}|z_{t})=N(z_{s}|\mu_{t\rightarrow s}([\widehat{x},\widehat{h}],z_{t}),\sigma_{t\rightarrow s}^{2}I) p(zs∣zt)=N(zs∣μt→s([x ,h ],zt),σt→s2I) [ x ^ , h ^ ] = z t / α t − ϵ ^ t ⋅ σ t / α t [\widehat{x},\widehat{h}]=z_{t}/\alpha_{t}-\widehat{\epsilon}_{t}\cdot\sigma_{t}/\alpha_{t} [x ,h ]=zt/αt−ϵ t⋅σt/αt通过神经网络学习扩散过程中向样本添加的噪声 ϵ t \epsilon_{t} ϵt，神经网络的预测为 ϵ ^ = [ ϵ ^ ( x ) , ϵ ^ ( h ) ] \widehat{\epsilon}=[\widehat{\epsilon}^{(x)},\widehat{\epsilon}^{(h)}] ϵ =[ϵ (x),ϵ (h)]。只要神经网络满足等变性，则整个生成过程具有等变性。因此神经网络采用EGNN实现。
EDM的优化和采样算法如下：
个人理解
个人认为EDM和GEODIFF相比，存在一些差距，GEODIFF本身包含了条件生成，而EDM更像是简单的扩散模型思想结合等变神经网络。而且EDM关于等变网络的设计似乎是不完善的，因为GEODIFF的作者在GEODIFF论文中描述过，仅通过等变网络学习噪声（噪声没有等变方差）会违背等变网络的性质。
另外，EDM还需要通过原子信息预测边的信息，这可能会是影响预测结果稳定性的一个因素。