ICLR 2022 under review | 从零开始生成三维分子几何结构的自回归流模型

作者 | 庞超 编辑 | 李仲深

今天给大家介绍的是ICLR2022上underreview的文章《An autoregressive flow model for 3d molecular geometry generation from scratch》。虽然目前已经开发了多种方法来生成分子图，但从零开始生成分子的三维几何结构问题并没有得到充分的探索。在这项工作中，作者提出了G-SphreNet，一种生成三维分子几何的自回归流模型。G-SphereNet采用了一种一步步将原子放置在三维空间上灵活的顺序生成方案，它并不直接生成三维坐标，而是通过生成距离、角度和扭转角来确定原子的三维位置，从而确保不变性和等变性。此外，作者建议使用球形信息传递和注意力机制进行条件信息提取。实验结果表明，G-SphreNet在随机分子几何结构生成和目标分子发现任务方面优于以往的方法。

一、介绍

无法从完整的二维分子图中获得全部的分子信息，因为分子的三维几何结构决定了分子的许多性质。三维分子几何结构表达了原子的三维坐标，而这对于量子性质的准确预测非常重要。因此，二维分子图难以用于生成具有某些理想量子性质的新分子。开发一个生成模型，从二维分子图中产生三维分子几何结构是一个很有前途的解决方案。

在这项工作中，作者提出了G-SphereNet，一个受SphereNet启发的三维分子几何结构生成模型。在G-SphereNet中，通过在三维空间中依次放置原子来生成三维分子几何结构。原子的空间位置通过距离、角度和扭转角来隐式确定，以确保不变性和等变性。此外，G-SphereNet还采用了SphereNet和注意力机制来提取条件信息。实验结果表明，G-SphereNet在三维分子几何结构生成任务上的表现优于以往的方法。

二、问题描述与背景知识

这项工作考虑的是从头开始生成三维分子几何结构的问题。设

是一个三维分子几何结构的集合，函数S(G)∈R计算G的一个特定的量子性质。作者考虑的两个生成任务定义如下。

• 从G中学习一个随机生成模型pθ(·)，使该模型可以从高维概率分布pθ(G)中采样一个有效三维分子几何结构G。

• 学习一个目标分子发现模型pθ(·)，以最大化（或最小化）预期的量子性质分数

。

一个流模型定义了一个带参数的可逆映射fθ：z∈Rd→x∈Rd，其中数据点x和潜在变量z都是随机变量。假设z是从一个已知的先验分布pZ中采样的，并且fθ是可逆的，可以如下计算x的对数似然：

其中

是雅可比矩阵矩阵。为了在给定的数据集

上训练流模型，从上式中计算数据点的对数似然并通过梯度下降最大化。因此，易于处理和廉价的计算detJ对于高效的训练是必须的。在大多数流模型中，fθ的一个常见选择是仿射耦合映射，在这种情况下，计算detJ非常容易，因为J是一个上三角矩阵。

三、模型与方法

设k是原子类型的数量。作者使用G=(A, R)来表示一个具有n个原子的分子的三维几何结构，其中A∈{0,1}n×k是原子类型矩阵，R∈Rn×3是原子坐标矩阵。矩阵A中的每一行都是一个独热向量，a[j, u]=1表示第j个原子为u类型。在矩阵R的第j行处的行向量表示第j个原子的三维笛卡尔坐标。

作者认为三维分子几何结构的生成是一个连续的决策过程。从一个初始具有一个碳原子的分子几何G1开始，通过在每一步中添加一个新的原子生成完整的几何结构。具体来说，在第i步，设由前面的i-1步生成的中间三维分子几何结构是Gi=(Ai, Ri)，它有i个原子。新原子的原子类型ai∈{0, 1}k由基于隐变量zia的生成模型ga生成。生成模型gr决定了新原子的三维笛卡尔坐标ri∈R3。ga和gr都是自回归函数。整个顺序生成过程可以用以下方程来描述：

自回归流模型

作者采用自回归流模型，在每一步生成新原子的原子类型ai。由于原子类型是离散数，作者采用去量化方法，通过添加实值噪声将它们转换为连续的数字。为了生成ai，首先从标准高斯分布N(0, I)中采样隐变量zia∈Rk，然后将zia通过仿射变换映射到

尺度因子s和位移因子t都依赖于从中间结构Gi中提取的条件信息。

应该对Ri上的任何刚性变换保持不变，即如果我们在三维空间中旋转或平移Ri，它不应该改变。因此，作者使用一个对称不变模型来计算s和t。

然而，不能像原子类型ai的生成那样生成三维坐标ri。如果用自回归流模型直接计算ri，坐标的等变性和似然的不变性都不满足。所以作者建议用对称不变元素来确定新原子的三维相对位置。与G-SchNet类似，首先选择一个焦点作为新原子位置的参考点。然后，模型依次生成距离、角度和扭转角。作者证明了这样一个在三维空间中放置新原子的过程可以严格满足等变性与不变性。图1给出了G-SphereNet的整个生成过程。

图1. G-SphereNet中顺序生成过程的示意图

生成新原子的原子类型和三维坐标需要在每一步从中间几何结构中捕获条件信息。可描述的条件信息应该包含几何结构的全面三维结构特征，并对任何刚性变换不变。为了实现这一目标，作者使用SphereNet捕获条件信息。

SphereNet将输入的分子几何结构处理成一个截止图，其中节点和边缘嵌入用球形基函数初始化，并通过球形消息传递进行更新。SphereNet可以在所有粒度层次上对对称不变的三维结构信息进行编码，并在多个量子性质预测任务中具有强大的特征提取能力。基于这些优势，作者使用SphereNet模型作为主干特征提取器，从中间分子几何结构中捕获条件信息。

然而，实验表明，只使用SphereNet经常导致新原子在三维空间中的错误放置。作者认为这是因为节点嵌入只包含局部三维信息，这不足以准确地生成新原子的三维位置。为了解决上述问题，作者用多头注意网络提取的全局特征来增强节点嵌入。多头注意力机制的使用有助于提取更全面的三维条件信息，在作者的消融研究中被证明可以大大提高生成性能。

四、实验结果

作者选择了E-NFs与G-SchNet两个模型作为对比基线，在生成分子的有效性，键长分布的最大平均差异（MMD），HOMO-LUMO能极差和各向同性极化度四个指标上进行了比较。实验结果证明G-SphereNet在这几项任务中有较好的性能表现。

表1. 不同方法在随机分子三维结构生成任务上的比较。通过有效性百分比和键长分布的MMD距离来评价其性能。这里↑表示此值越高越好，而↓则相反。

表2. 在发现目标分子时，不同方法的性能比较。这里的↓意味着目标是最小化性质分数，而↑是最大化性质分数。

五、总结

G-SphereNet与以往的三维分子结构生成方法相比有许多优点。首先，G-SphereNet更容易生成有效的几何结构。其次，G-SphereNet中三维位置的生成比G-SchNet中更灵活。在G-SchNet中，新原子必须被放置在一个候选网格点上，而在G-SphereNet中新原子可以放置焦点原子的任何相对位置。第三，与E-NFs相比，G-SphereNet更高效，也更有效。G-SphereNet通过生成距离、角度和扭转角来获得三维位置。因此，G-SphereNet可以避免E-NFs的复杂操作。此外，E-NFs一次生成几何结构中所有原子的坐标，而G-SphereNet一次生成一个原子的坐标。虽然G-SphereNet可能较慢，但作者认为顺序生成方式有助于模型更有效地捕捉原子和几何密度之间的依赖性。实验结果也表明，G-SphereNet可以生成更多有效的分子几何结构。

欢迎关注公众号，获取更多论文解读

转载或合作请联系邮箱 intbioinfo@163.com