GEODIFF：用于分子构象生成的几何扩散模型

编译 | 刘名权 审稿 | 陶雯

本文介绍一篇ICLR 2022在审论文《GEODIFF：A GEOMETRIC DIFFUSION MODEL FOR MOLECULAR CONFORMATION GENERATION》。根据分子图来预测分子构象是化学信息学和药物发现中的一项基础工作。随着深度生成模型的兴起，这一工作取得重大进展。在这篇论文中，作者提出了新的生成模型GEODIFF。受热力学粒子扩散模型的启发，GEODIFF将每个原子当作一个粒子，GEODIFF所需要学习的部分——分子构象生成即扩散模型的逆过程。对比最先进的一些生成模型，GEODIFF在多个基准上展示出其竞争力，对于大分子尤其明显。

1介绍

在分子建模领域中，图表示法在性质预测和分子生成等多个任务中取得了巨大成功。作者使用的三维几何（也叫构象）是另一种表示方法，在图的基础上它使用笛卡尔坐标来表示原子。三维结构决定了分子的生物和物理性质，因此它在药物计算和材料设计中发挥着重要作用。

由于分子构象具有旋转-平移不变性（roto-translational invariant），一些机器学习方法通过学习具有旋转-平移不变性的中间几何变量（如原子距离等）取得引人注目的结果，从而规避了对原子坐标进行建模。但是，使用中间几何变量有着固有的限制。最理想的解决方法仍然是对原子坐标进行建模的同时考虑分子的旋转平移不变性。基于此，作者提出了GEODIFF，其本质上是一种基于去噪扩散模型的概率框架。GEODIFF的独特之处在于训练和预测过程中都直接作用于原子坐标而不是中间元素。因此，它具有以下几个优点。首先，模型可以直接进行端到端的训练而不需要其他复杂的技巧（如bilevel programming等）；其次，模型不是从键长或角度出发求解几何，而是使用one-stage采样的方式避免了累积任何的中间误差，从而能更精确的预测分子结构；最后，GEODIFF具有较好的模型能力来近似复杂的构象分布。因此，该模型可以更好的估计多模态分布并且生成高质量的、多样的分子结构。

2模型

目标

分子构象生成任务本质上是条件生成问题，作者的目标是对于给定的分子图G生成稳定的分子构象C。

理论

扩散过程：作者将这一过程定义为具有固定后验概率分布的马尔可夫隐变量模型。公式如下所示。

其中β是固定的数值。此外作者还证明出时间步t条件概率的近似形式：

逆扩散过程：对于给定的图G，从白噪声

开始，逐步去噪，最后恢复构象

。

等变逆向生成过程：作者证明了一个重要命题，如下所示。假设给定满足SE(3)不变性的密度函数和满足SE(3)等变性的马尔可夫核，那么边缘似然密度函数具有SE(3)不变性。

作者借鉴了CoM系统的思想得到了具有旋转-平移不变性的初始密度函数。同样，作者也将该思想用到了马尔可夫核的构建过程中，保证了马尔可夫核的平移不变性，因此只需要考虑它的旋转等变性。由于马尔可夫核的旋转-平移等变性体现在其均值，因此作者对均值进行了如下参数化的神经网络，它将尝试学习出decorrupt构象的噪声。问题进一步转变为构建一个具有旋转-平移等变性的神经网络，在这里作者采用的是GFN网络。

改进训练目标

由于难以直接优化对数似然，作者采用了最大化变分下界（ELBO）的方法。为了使得ELBO过程具有所需要的等变性，作者提供两种方式修正观测噪音信号。第一种是alignment方式，还有一种chain-rule的方式。

采样

采样过程本质上是马尔可夫链：

3实验结果

作者使用了GEQM-QM9和GEOM-Drugs两个数据集，并选择CVGAE，GRAPHDG和CGCF等六种最新的模型作为基线。

构象生成任务

为了测试不同模型所生成构象的质量和多样性。作者评估了建立在两原子坐标矩阵的Frobenius范数基础上的四项指标。

表1和表2：COV百分比越高或者MAT得分越低则生成的构象更真实，字母R和P分别代表Recall和Precision两种评价方式，GEODIFF-A和GEODIFF-C分别是使用assignment和chain-rule两种方式修正的ELBO进行训练的模型。

结果显示，GEODIFF在所有的数据集和所有的指标上都比其他的模型表现的要好，尤其是在GEO-Drugs这一更具挑战的数据集上，远胜其他模型。此外，GEODIFF-C要比GEODIFF-A表现的稍微好一点。因此，作者使用GEODIFF-C进行接下来的实验，如图2。

最后，作者使用force field（FF）对结构进行优化，并与同样使用FF进行优化的RDKIT模型进行比较，结果如表3所示。

结果表明GEODIFF+FF有着优越的多样性（Recall metrics）同时显著地提高了准确率（Precision metrics）。

性质预测任务

这项任务使用分子集成性质作为指标进行测试。作者使用PSI4来计算每个构象的能量E和HOMO-LUMO gap 并且比较了平均能量、最小能量、平均gap、最小gap、最大gap。

如表4所示，数值代表着所计算的性质和基本事实之间的平均绝对误差。由于性质对于几何结构非常敏感，GEODIFF的卓越表现证明了它能够更准确的预测不同分子的构象。

4总结

作者提出了一种用于分子构象生成的新型概率模型GEODIFF。GEODIFF在去噪扩散模型中使用了几何表示，将逆向生成这一动态过程参数化为一条马尔可夫链，并使用等变马尔可夫核将旋转平移不变性加入到密度函数中。在优化似然函数时，从变分下界中推导出易处理的具有不变性的优化目标。多个任务的综合实验表明，GEODIFF与现有的最先进的模型相比具有竞争力。在未来的工作中，可以使用扩散模型的其他成果来进一步改进或者加速模型，以及将模型运用到其他更具挑战性的结构（例如蛋白质）。

参考资料

https://openreview.net/forum?id=PzcvxEMzvQC