【学习笔记】分子3D构象生成的扩散模型GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation

1 摘要

  从分子图预测分子构象是化学信息学和药物发现中的一个基本问题。最近，机器学习方法取得了重大进展，尤其是深度生成模型。受经典非平衡热力学中的扩散过程的启发，在这篇论文中，我们提出了一个新的生成模型，名为GEODIFF，用于分子构象预测。GEODIFF将每个原子视为一个粒子，并学习将扩散过程（即从噪声分布转换为稳定构象）直接反转为马尔可夫链。然而，对这样的生成过程进行建模是非常具有挑战性的，因为构象的可能性应该是旋转平移不变的。我们从理论上证明了用等变马尔可夫核演化的马尔可夫链可以通过设计诱导不变分布，并进一步提出了马尔可夫核的构建块，以保持期望的等变性质。通过优化（条件）似然的加权变分下界，可以以端到端的方式有效地训练整个框架。在多个基准上的实验表明，GEODIFF优于或可与现有的最先进的方法相媲美，尤其是在大分子上。
代码： https://github.com/MinkaiXu/GeoDiff

论文地址： https://arxiv.org/abs/2203.02923

2 简介

  图表示学习在从分子的各种任务中为分子建模取得了巨大成功，其中分子通常被表示为原子键图。然而，分子的一种更内在、更具信息性的表示是3D几何，也被称为构象，其中原子被表示为它们的笛卡尔坐标。如何预测稳定的分子构象仍然是一个具有挑战性的问题。基于分子动力学（MD）或马尔可夫链蒙特卡罗（MCMC）的传统方法在计算上非常昂贵，尤其是对于大分子。
  最近，机器学习方法取得了重大进展，特别是在深度生成模型方面。在本文中，我们提出了一种称为GEODIFF的模型，这是一种基于去噪扩散模型的原则性概率框架。该方法受到非平衡热力学中扩散过程的启发。我们将原子视为热力学系统中的粒子，这些粒子在与热浴接触时从原始状态逐渐扩散到有噪声的分布。在每个时间步长，随机噪声被添加到原子位置。我们的高级思想是学习反转扩散过程，从噪声分布中恢复目标几何分布。受去噪扩散模型在图像生成方面的最新进展的启发，作者将不同时间步长的噪声几何体视为潜在变量，并将正向扩散和反向去噪过程公式化为马尔可夫链。我们的目标是学习过渡核，以便反向过程可以从从噪声分布中采样的混沌位置恢复真实的构象。
  然而，将现有方法扩展到几何生成是困难的，在构象生成任务中直接应用扩散模型会导致生成质量差。如上所述，分子构象是旋转-平移不变量，即估计的（条件）可能性应不受平移和旋转变换的影响。为此，作者首先从理论上阐述从旋转-平移不变量先验分布开始并随着旋转-平移等变马尔可夫核演化的马尔可夫过程，将其诱导为旋转-平移不变密度函数。作者进一步提供了实用的参数化来定义旋转平移不变量先验分布和施加等变约束的马尔可夫核。此外，其推导了分子构象条件似然的加权变分下界，该下界也具有转平移不变性，可以有效地优化。
  GEODIFF的一个独特优势是，它直接作用于原子坐标，并完全绕过了训练和推理中使用的中间元素。这种方法使模型可以端到端地自然训练。同时，一阶段采样方式避免了累积任何中间误差，从而导致更准确的预测结构，而不是从键长或角度求解几何结构。此外，GEODIFF具有很高的模型容量来近似复杂的构象分布。因此，该模型可以更好地估计高度多模态分布，并生成具有高质量和多样性的结构。

3 模型框架

3.1 定义

符号定义

  在本文中，每个具有n个原子的分子表示为无向图$\mathcal{G}=⟨\mathcal{V},\mathcal{E}⟩$，其中 V = { v i } i = 1 n \mathcal{V}=\{v_i\}_{i=1}^n V={vi​}i=1n​表示原子的顶点集， E = { e i j ∣ ( i , j ) ⊆ ∣ V ∣ × ∣ V ∣ } \mathcal{E}=\{e_{i j}\mid(i,j)\subseteq|\mathcal{V}|\times|\mathcal{V}|\} E={eij​∣(i,j)⊆∣V∣×∣V∣}表示原子间键的边集。每个节点$v_i\in \mathcal{V}$包含原子属性,如元素种类。每个边缘$e_{ij}\in \mathcal{E}$描述$v_i$和$v_j$之间的对应连接，并标记其类型。此外，作者为未连接的边指定了一个虚拟类型。对于分子，V中的每个原子都被坐标向量$c\in {\mathbb{R}}^3$嵌入到三维空间中，即构象可以表示为矩阵$\mathcal{C}=[{\mathbf{c}}_1,{\mathbf{c}}_2,\cdots ,{\mathbf{c}}_n]\in {\mathbb{R}}^{n\times 3}$。

问题定义

  分子构象生成的任务是一个条件生成问题，为所提供的图G生成稳定的构象。给定多个图G，并且对于每个G，给定其构象C作为来自潜在玻尔兹曼分布的i.i.d样本，我们的目标是学习生成模型$p_{\theta }(\mathcal{C}|\mathcal{G})$。 该模型易于从中提取样本，以近似玻尔兹曼函数。

等方差

  等方差在原子系统的机器学习中普遍存在，例如，原子偶极子或力的矢量应该相对于构象坐标相应地旋转。将这种归纳偏差集成到建模3D几何的模型参数化中已显示出有效性，这对模型的泛化能力至关重要。
  形式上，函数$\mathcal{F}:\mathcal{X}\to \mathcal{Y}$对于$G$是等变的，即$$\mathcal{F}\circ T_g(x)=S_g\circ \mathcal{F}(x)(1)$$  其中$T_g$和$S_g$是元素$g\in G$的变换，分别作用于向量空间X和Y。在这项工作中，我们使用了$SE(3)$，即3D空间中的旋转矩阵。在后续需要估计不受平移和旋转变换影响的可能性。

3.2 模型定义

图1 GEODIFF的扩散和反向过程

  设$C^0$表示真实构象，${\mathcal{C}}^t(t=\mathrm{1...}T)$是具有相同维度的潜在变量序列，其中T是扩散过程的步数。然后，扩散概率模型可以被描述为具有两个过程的潜变量模型：正向扩散过程和反向生成过程。如图1所示，对于扩散过程，来自固定后验分布的噪声逐渐增加，直到构象被破坏。对于生成过程，从标准高斯分布中采样初始状态，并通过马尔可夫核逐步细化构象。

扩散过程

  作者将粒子C建模为一个不断发展的热力学系统。随着时间的推移，平衡构象$C^0$将逐渐扩散到下一个混沌状态$C^t$，并在T次迭代后最终收敛为白噪声分布。在扩散模型中，这种正向过程被定义为固定的（而不是可训练的）后验分布$q({\mathcal{C}}^{1:T}|{\mathcal{C}}^0)$。具体来说，我们将其定义为根据固定方差调度的马尔可夫链${\beta }_1,\dots ,{\beta }_T$：$$q({\mathcal{C}}^{1:T}|{\mathcal{C}}^0)=\prod _{t=1}^{T}q({\mathcal{C}}^t|{\mathcal{C}}^{t-1}),q({\mathcal{C}}^t|{\mathcal{C}}^{t-1})=\mathcal{N}({\mathcal{C}}^t;\sqrt{1-{\beta }_t}{\mathcal{C}}^{t-1},{\beta }_{t}I)(2)$$  需要注意的是，在这项工作中，我们没有对扩散过程施加特定的（不变性）要求，只要它可以有效地绘制噪声样本来训练生成过程。
  设${\alpha }_t=1-{\beta }_t$和${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$，正向过程的一个特殊性质是，任意时间步长t的$q({\mathcal{C}}^t|{\mathcal{C}}^0)$可以以闭合形式$q({\mathcal{C}}^t|{\mathcal{C}}^0)=\mathcal{N}({\mathcal{C}}^t;\sqrt{{\bar{\alpha }}_t}{\mathcal{C}}^0,(1-{\bar{\alpha }}_t)I)$计算。这表明，在T足够大的情况下，整个正向过程可以将$C^0$转换为白噪声，因此将$p(C^T)$设置为标准高斯分布。

反转过程

  我们的目标是在给定特定分子图G的情况下，学习从白噪声$C^T$中恢复构象$C^0$。这种生成过程是上述扩散过程的反向动力学，从噪声粒子$C^T\sim p(C^T)$开始。我们将这种反向动力学公式化为具有可学习转换的条件马尔可夫链：$$p_{\theta }({\mathcal{C}}^{0:T-1}|\mathcal{G},{\mathcal{C}}^T)=\prod _{t=1}^T{p}_{\theta }({\mathcal{C}}^{t-1}|\mathcal{G},{\mathcal{C}}^t),p_{\theta }({\mathcal{C}}^{t-1}|\mathcal{G},{\mathcal{C}}^t)=\mathcal{N}({\mathcal{C}}^{t-1};{\mu }_{\theta }(\mathcal{G},{\mathcal{C}}^t,t),{\sigma }_t^{2}I)(3)$$  其中${\mu }_{\theta }$是用于估计均值的参数化神经网络，${\sigma }_t$可以是定义的方差。初始分布$p(C^T)$为标准高斯分布。给定图$G$，其3D结构首先通过从$p(C^T)$采样噪声粒子$C^T$生成初始构象，然后通过逆马尔可夫核$p_{\theta }({\mathcal{C}}^{t-1}|\mathcal{G},{\mathcal{C}}^t)$迭代精化。
  在公式化了反向动力学之后，边际似然可以通过$p_{\theta }({\mathcal{C}}^0|\mathcal{G})=\int p({\mathcal{C}}^T)p_{\theta }({\mathcal{C}}^{0:T-1}|\mathcal{G},{\mathcal{C}}^T)\text{d}{\mathcal{C}}^{1:T}$计算。此处有一个重要的问题，平移和旋转的可能性应该是不变的，这是3D对象生成的关键归纳偏差。

3.3 旋转不变性

  在这个小节中，我们将详细说明如何参数化马尔可夫核$p_{\theta }({\mathcal{C}}^{t-1}|\mathcal{G},{\mathcal{C}}^t)$，以及如何通过考虑旋转平移不变性。
  我们考虑建立对旋转和平移变换不变的概率密度$p_{\theta }({\mathcal{C}}^0)$。这个概率密度需要不受平移和旋转影响的可能性。形式上，设$T_g$旋转平移操作，则可以有以下陈述：

• 若$p(x_T)$是旋转平移不变的密度函数，则$p(x_T)=p(T_g(x_T))$。相似的，若马尔可夫变换$p(x_{t-1}|x_t)$是旋转平移不变的，即$p(x_{t-1}|x_t)=p(T_g(x_{t-1})|T_g(x_t))$，因此我们得到的密度$P_{\theta }(x_0)=\int p(x_T)p_{\theta }(x_{0:T-1}|x_T)\mathrm{d}{\mathbf{x}}_{1:T}$也是旋转平移不变的。

  上述命题表明，从不变的标准密度沿着等变高斯马尔可夫核开始的动力学可以产生不变的密度。接下来，我们开始介绍GEODIFF的实际实现。

不变初始密度$p(C^T)$

  我们首先引入不变分布$\rho (C^T)$，它也将用于等变马尔可夫链。将具有零质心（CoM）的系统，称为无CoM系统。我们将$p(C^T)$定义为“无CoM标准密度${\rho }^{(C)}$”，即零质心系统。将粒子移动到零CoM总是可以确保平移不变性。

等变马尔可夫核$p({\mathcal{C}}^{t-1}|\mathcal{G},{\mathcal{C}}^t)$

  我们考虑将所有的中间状态$C^t$转化为零质心系统。具体来讲，给定均值${\mu }_{\theta }(\mathcal{G},{\mathcal{C}}^t,t)$和方差${\sigma }_t$，$C^{t-1}$的可能分布可以通过$\hat{\rho }(\frac{{\mathcal{C}}^{t-1}-{\mu }_{\theta }(\mathcal{G},{\mathcal{C}}^t,t)}{{\sigma }_t})$计算。无CoM保证了马尔可夫核中的平移不变性。
  接下来，我们关注旋转等变性。关键是确保均值${\mu }_{\theta }(\mathcal{G},{\mathcal{C}}^t,t)$对于$C^T$是旋转不变的。我们将${\mu }_{\theta }$参数化：$${\mu }_{\theta }({\mathcal{C}}^t,t)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathcal{C}}^t-\frac{{\beta }_t}{\sqrt{1-{\tilde{\alpha }}_t}}{ϵ}_{\theta }(\mathcal{G},{\mathcal{C}}^t,t)\right)(4)$$  其中${ϵ}_{\theta }$是具有可训练参数$\theta$的神经网络。直观地，模型${ϵ}_{\theta }$学习预测干扰构象所需的噪声。因此，现在问题转化为将${ϵ}_{\theta }$构造为旋转平移不变。作者设计了一个等变卷积层，称为图场网络（GFN）。在第l层中，GFN采用节点编码${\mathbf{h}}^l\in {\mathbb{R}}^{n\times b}$（b表示特征维度）和相应的坐标编码${\mathbf{x}}^l\in {\mathbb{R}}^{n\times 3}$作为输入，输出$h^{l+1}$和$x^{l+1}$如下所示：$${\mathbf{m}}_{ij}={\Phi }_m\left({\mathbf{h}}_i^l,{\mathbf{h}}_j^l,\Vert {\mathbf{x}}_i^l-{\mathbf{x}}_j^l{\Vert }^2,e_{ij};{\theta }_m\right)(5)$$$${\mathbf{h}}_i^{l+1}={\Phi }_h({\mathbf{h}}_i^l,\sum _{j\in \mathcal{N}(i)}{\mathbf{m}}_{ij};{\theta }_h)(6)$$$${\mathbf{x}}_i^{l+1}=\sum _{j\in \mathcal{N}\left(i\right)}\frac{1}{d_{ij}}\left({\mathbf{c}}_i-{\mathbf{c}}_j\right){\Phi }_x\left({\mathbf{m}}_{ij};{\theta }_x\right)(7)$$  其中$\Phi$是前馈网络，$d_{ij}$表示原子间距离。$N(i)$表示第i个节点的邻域，包括连接原子和半径阈值$\tau$内的其他原子，这使得模型能够捕捉长程相互作用，并支持具有断开分量的分子图。初始嵌入编码$h^0$是原子编码和时间步长编码的组合，$x^0$是原子坐标。
  作者所提出的GFN与其他GNN之间的主要区别在于方程(7)，其中x被更新为由${\Phi }_x({\mathbb{R}}^b\to \mathbb{R})$加权的径向方向的组合。这样的矢量场$x^L$具有旋转平移不变性质。形式上，我们有：

• 将${ϵ}_{\theta }(\mathcal{G},\mathcal{C},t)$参数化为L个GFN层的组合，并将L个更新后的$x^L$作为输出。那么噪声向量场${ϵ}_{\theta }$相对于3D系统$C$是旋转平移不变的。

  直观地说，给定$h^l$已经不变,$x^l$等变，则消息嵌入m也将是不变的，因为它只依赖于不变特征。由于x是用不变特征加权的相对差$c_i-c_j$更新的，因此它将是平移不变的和旋转等变的。总的来说，用L个GFN层组成的${ϵ}_{\theta }$可以实现与$C^t$的等价。

3.4 改进损失函数

  在制定了生成过程和模型参数化之后，现在我们考虑反向动力学的实际训练目标。由于直接优化精确对数似然是难以解决的，因此我们转而最大化通常的变分下界（ELBO）：$$\begin{array}{rl}\mathbb{E}\left[\log p_{\theta }({\mathcal{C}}^0|\mathcal{G})\right]& =\mathbb{E}\left[\log {\mathbb{E}}_{q({\mathcal{C}}^{1:T}|{\mathcal{C}}^0)}\frac{p_{\theta }({\mathcal{C}}^{0:T}|\mathcal{G})}{q({\mathcal{C}}^{1:T}|{\mathcal{C}}^0)}\right]\\ & \ge -{\mathbb{E}}_q[\sum _{t=1}^T{D}_{\text{KL}}(q({\mathcal{C}}^{t-1}|{\mathcal{C}}^t,{\mathcal{C}}^0)\Vert p_{\theta }({\mathcal{C}}^{t-1}|{\mathcal{C}}^t,\mathcal{G}))]:=-{\mathcal{L}}_{\text{ELBO}}(8)\end{array}$$  其中$q({\mathcal{C}}^{t-1}|{\mathcal{C}}^t,{\mathcal{C}}^0)$可以转化为$\mathcal{N}(\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathcal{C}}^0+\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathcal{C}}^t,\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t)$。
  最近有研究表明，在方程(4)中的参数化条件下，扩散模型的ELBO（Evidence Lower Bound）可以通过计算高斯分布之间的KL散度（Kullback-Leibler divergence）来进一步简化，具体实现时KL散度可以表示为均值${ϵ}_{\theta }$和$ϵ$之间加权${\mathcal{L}}_2$距离。因此可以有： L ELBO = ∑ t = 1 T γ t E { C 0 , G } ∼ q ( C 0 , G ) , ϵ ∼ N ( 0 , I ) [ ∥ ϵ − ϵ θ ( G , C t , t ) ∥ 2 2 ] (9) \mathcal{L}_{\text{ELBO}}=\sum_{t=1}^T\gamma_t\mathbb{E}_{\{\mathcal{C}^0,\mathcal{G}\}\sim_q(\mathcal{C}^0,\mathcal{G}),\epsilon\sim\mathcal{N}(0,I)}\bigg[\|\epsilon-\epsilon_\theta(\mathcal{G},\mathcal{C}^t,t)\|_2^2\bigg]\tag{9} LELBO​=t=1∑T​γt​E{C0,G}∼q​(C0,G),ϵ∼N(0,I)​[∥ϵ−ϵθ​(G,Ct,t)∥22​](9)  其中${\mathcal{C}}^t=\sqrt{{\bar{\alpha }}_t}{\mathcal{C}}^0+\sqrt{1-{\bar{\alpha }}_t}ϵ$，权重${\gamma }_t=\frac{{\beta }_t}{2{\alpha }_t(1-{\bar{\alpha }}_{t-1})}(t>1)$，${\gamma }_1=\frac{1}{2{\alpha }_1}$。我们的目标是从$q({\mathcal{C}}^{t-1}|{\mathcal{C}}^t,{\mathcal{C}}^0)$中独立地采样不同时间步的混乱构型（中间状态），并使用${ϵ}_{\theta }$来对噪声向量$ϵ$进行建模。为了获得更好的实验性能，建议将所有权重${\gamma }_t$设置为1，这与最近的噪声条件评分网络的目标一致。
  由于${ϵ}_{\theta }$被设计为等变，因此自然要求其监督信号$ϵ$与$C^t$等变。注意，一旦实现了等变，ELBO也将变为等变。然而，正向扩散过程中的$ϵ$没有施加这样的等变，违反了上述性质。在这里，我们提出了两种方法来获得修改后的噪声向量$\widehat{ϵ}$，在方程(9)中的${\mathcal{L}}_2$距离计算中进行替换后，它实现了所需的等变。
方法一：对齐方法
  由于$ϵ$可以通过$\frac{{\mathcal{C}}^t-\sqrt{{\overline{\alpha }}_t}{\mathcal{C}}^0}{\sqrt{1-{\overline{\alpha }}_t}}$计算，我们可以通过旋转和平移操作将$C^0$ 转换为${\hat{\mathcal{C}}}^0$，使其与 $C^t$ 对齐，然后通过$\frac{{\mathcal{C}}^t-\sqrt{{\bar{\alpha }}_t}{\hat{C}}^0}{\sqrt{1-{\bar{\alpha }}_t}}$计算$\widehat{ϵ}$。由于对齐的构象${\hat{\mathcal{C}}}^0$与$C^t$是等变的，因此处理后的$\widehat{ϵ}$也同样等变。具体来讲，通过首先将$C^0$转换为与$C^t$相同CoM，然后通过Kabsch对齐算法求解最优旋转矩阵来实现对齐。
方法二：链式规则方法
  另一个有意义的观察结果是，通过重参数化高斯分布$q({\mathcal{C}}^t|{\mathcal{C}}^0)$为${\mathcal{C}}^t=\sqrt{{\bar{\alpha }}_t}{\mathcal{C}}^0+\sqrt{1-{\bar{\alpha }}_t}ϵ$，$ϵ$可以将其视为加权得分函数$\sqrt{1-{\bar{\alpha }}_t}{\nabla }_{{\mathcal{C}}^t}q({\mathcal{C}}^t|{\mathcal{C}}^0)$。接着我们通过用链式规则将该得分函数${\nabla }_{{\mathcal{C}}^t}q({\mathcal{C}}^t|\cdot )$分解为${\partial }_{{\mathcal{C}}^t}{\mathbf{d}}^t{\nabla }_{{\mathbf{d}}^t}q({\mathcal{C}}^t|\cdot )$，将其设计为等变函数，其中$d^t$可以是结构$C^t$的任何不变特征，如原子间距离。作者认为，作为不变变量相对于等变变量的梯度，偏导数${\partial }_{{\mathcal{C}}^t}{\mathbf{d}}^t$始终与$C^t$等价。在这项工作中，在$d$也遵循高斯分布的共同假设下，我们的实际实现是首先将${\nabla }_{{\mathbf{d}}^t}q({\mathcal{C}}^t|{\mathcal{C}}^0)$计算为$\frac{{\mathbf{d}}^t-\sqrt{{\bar{\alpha }}_t}{\mathbf{d}}^0}{1-{\bar{\alpha }}_t}$，然后将修改后的噪声向量$\widehat{ϵ}$计算为$\sqrt{1-{\bar{\alpha }}_t}{\partial }_{{\mathcal{C}}^t}{\mathbf{d}}^t(\frac{{\mathbf{d}}^t-\sqrt{{\bar{\alpha }}_t}{\mathbf{d}}^0}{1-{\bar{\alpha }}_t})=\frac{{\partial }_{{\mathcal{C}}^t}{\mathbf{d}}^t\cdot ({\mathbf{d}}^t-\sqrt{{\bar{\alpha }}_t}{\mathbf{d}}^0)}{\sqrt{1-{\bar{\alpha }}_t}}$。

3.5 采样过程

  通过学习到的反向动力学${ϵ}_{\theta }(\mathcal{G},{\mathcal{C}}^t,t)$，${\mu }_{\theta }(\mathcal{G},{\mathcal{C}}^t,t)$可以由方程(4)计算。因此，给定一个图$\mathcal{G}$，其几何构象$C^0$是通过首先对混沌粒子$C^T\sim p(C^T)$进行采样，然后逐渐采样$C^{t-1}\sim p_{\theta }({\mathcal{C}}^{t-1}|\mathcal{G},{\mathcal{C}}^t)$来生成的。这个过程是马尔可夫过程，它将先前的有噪位置逐渐向平衡状态转变。下图展示了采样过程的伪代码。

图2 采样过程

  初始构象来源于高斯噪声。在每个采样循环中，我们首先将构象转化为零质心，然后通过参数化的模型计算均值，从而得到当前构象分布的均值和方差，进行采样，直到生成最终构象。

4 实验

5 代码

我们主要从计算损失的函数开始。

6 思考与改进