ICLR2020 | GraphAF：基于FLOW的分子图自回归生成模型

作者 | 庞超 编辑 | 庞超

今天给大家介绍的是北京大学和上海交通大学的Chence Shi等人在2020年的ICLR上发表的会议论文GraphAF: A flow-based autoregressive model for molecular graph generation。分子的图生成作为药物发现的基本问题，正在引起越来越多的关注。这个问题非常具有挑战性，因为它不仅需要产生化学上有效的分子结构，而且还需要同时优化它们的化学性质。受深度生成模型最新进展的启发，本文提出了一种基于Flow的图生成自回归模型，称为GraphAF。GraphAF结合了自回归和基于Flow的方法的优点，可以高效并行计算训练，允许利用化学领域知识进行有效性检查。实验结果表明，即使没有化学知识规则，GraphAF也能产生68%的化学有效分子。在通过强化学习对目标导向的性质优化模型进行微调后，GraphAF实现了最先进的性能。

一、研究背景

设计具有期望性质的新分子结构是药物发现和材料科学等学科中的基本问题。这个问题非常具有挑战性，因为化学空间本质上是离散的，而且搜索空间巨大。得益于此领域的大量数据，机器学习在分子设计可以大展拳脚。近年来，越来越多的研究者开始开发可以自动生成化学有效的分子结构，同时优化其性能的机器学习算法。

目前出现的先进的机器学习算法包括变分自编码器，生成对抗网络与循环神经网络等，图1中展示了以前最先进的分子图生成模型。其中训练过程的比较只在自回归模型之间进行。

图1. 分子图生成模型

最近，除了上述三种生成模型外，Flow已经取得了显著进展，并成功应用于密度估计、变分推理和图像生成。基于Flow的方法定义了潜在分布(例如高斯分布)和真实世界的高维数据之间的可逆变换。这样的可逆映射允许计算精确的数据似然。同时，通过使用隐空间和观测空间之间的多层非线性变换，Flow对数据密度的建模能力很高。此外，根据不同应用场景的要求，可以设计不同的架构来促进快速训练或快速采样。

受上述模型的启发，作者提出了一种基于Flow的自回归模型，称为GraphAF。它具有较高的模型容量，因此能够模拟真实分子数据的密度。GraphAF的采样过程为一个自回归模型，它基于现有的子图结构动态生成节点和边，这样的顺序生成过程允许在每个生成步骤中利用化学领域知识和有效性检查，这保证了生成的分子结构的有效性。同时，与GCPN和MRNN这类训练期间的自回归模型不同，GraphAF定义了从分子图结构到基分布的前馈神经网络，因此能够并行计算精确的数据似然。因此，GraphAF的训练过程非常高效。

二、模型与方法

2.1 模型框架

作者将分子图生成问题形式化为一个顺序决策过程。设G=(A,X)表示一个分子图结构。从空图G1开始，在每个步骤中基于当前子图结构Gi，生成新节点Xi。然后，该新节点与现有节点之间的边根据当前图结构顺序生成。重复此过程，直到生成所有节点和边缘。图2中给出了一个示例。

图2. 生成过程的说明图示

GraphAF的目标是定义一个从基分布(多元高斯分布)到分子图结构G=(A,X)的可逆变换。注意作者在两个节点之间添加了一种额外类型的边，对应于两个节点之间没有边。由于节点类型Xi和边缘型Aij都是离散的，不适合基于Flow的模型，标准方法是采用反量化技术，通过添加实值噪声将离散数据转换为连续数据。作者采用这种方法，将离散图G=(A,X)预处理为连续数据z=(zA,zX)：

在形式上，作者将生成的条件分布定义为：

其中四个g为用于定义高斯分布均值和标准差的神经网络的参数。更具体地说，给定当前子图结构Gi，作者使用L层的Relational GCN学习节点嵌入HiL，和整个子图的嵌入hi，基于定义的高斯分布的均值和标准差分别生成节点和边缘：

其中sum表示sum-pooling操作，Hi,jL表示在嵌入HiL的第j个节点的嵌入。mµx，mαx是一种多层感知器(MLP)，它根据当前的子图嵌入来预测节点类型。而mµA，mαA是根据当前的子图嵌入和节点嵌入来预测边缘类型的MLP。

为了生成一个新的节点Xi和其连接到现有节点的边，只需从基高斯分布中采样随机变量ϵi和ϵij，并将其转换为离散特征。

2.2 高效的并行训练

在GraphAF中，由于f：ℇ→Z是自回归的，因此f-1：Z→ℇ过程的雅可比矩阵是三角矩阵，其行列式可以非常高效地计算。给定一小批训练数据G，通过下图的公式，可以有效地计算出给定顺序下每个分子的精确密度。最终的目标是最大化训练数据的似然。

在训练过程中，作者通过使用掩蔽来定义输入分子图G和输出隐变量之间的前馈神经网络来进行并行计算。为了进一步加快训练过程，训练图的节点和边缘根据广度优先搜索(BFS)顺序进行重新排序。由于BFS的性质，键只能存在于相同或连续的BFS深度内的节点之间。因此，节点之间的最大依赖距离受单个BFS深度中最大节点数量的限制。在作者的数据集中，任何单个BFS深度都不超过12个节点，这意味着只需要构建当前原子和最新产生的12个原子之间的边。

2.3 有效性约束

在化学中，存在着许多化学规则，它们可以帮助产生有效的分子。由于是顺序生成过程，GraphAF可以在每个生成步骤中利用这些规则。具体地说，可以在采样期间应用有效性约束，以检查当前化学键是否符合有效性规则，作者通过下式检验有效性：

其中X代表原子，Aij代表原子间的化学键。如果新添加的键打破了有效性约束，模型将拒绝键A，并在隐空间中采样一个新的ϵij并生成另一个新的键类型。如果图的尺寸达到最大尺寸或新生成的原子与先前的子图之间不产生键，生成过程将终止。最后，氢被添加到没有填满其化合价的原子中。

三、实验结果

根据现有的关于分子生成的工作，作者在三个标准任务上与最先进的方法进行比较实验。

1. 密度建模和生成，评估模型学习数据分布和生成真实与多样的分子的能力。

2. 性能优化，生成具有优化化学性质的新分子。对于这项任务，作者微调从密度建模任务预先训练的网络，以最大化所需的属性。

3. 约束性质优化问题。旨在修改给定的分子，以提高期望的性质，同时满足相似性约束。

实验结果如下：

图3. 生成分子有效性与新颖性(分子重建只在具有隐空间的模型中对比

图4. 惩罚logP与QED分数优化

图5. 约束性质优化

从实验结果可以看出，与基线模型相比，GraphAF模型达到了目前最先进的性能。

四、总结

作者提出了GraphAF，第一个基于Flow的自回归模型，以生成真实和多样的分子图。由于标准化Flow的灵活性，GraphAF能够模拟复杂的分子分布，并在实验中生成新的和100%有效的分子。此外，GraphAF的训练是非常高效的。为了优化生成分子的性质，作者通过强化学习对生成过程进行了微调。实验结果表明，GraphAF在标准任务上优于以往所有最先进的基线模型。未来作者计划在更大的数据集上训练GraphAF模型，并将其扩展到生成其他类型的图结构(例如，社交网络)。

代码

https://github.com/DeepGraphLearning/GraphAF

参考文献

Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian.Tang. 2020.

GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation.

ICLR 2020, Addis Ababa, Ethiopia, Apr.26-Apr. 30, 2020 (2020)