NIPS'20 | 史上最大GNN: Self-Supervised Graph Transformer

去年以来，随着新冠疫情的爆发，AI制药领域受到广泛关注。顶级会联网公司如字节/阿里/腾讯/华为纷纷成立自己的AI药物研发机构; 各大投资机构如红杉资本/鼎晖资本/高榕资本/DCM也都在疯狂下注。       

清华大学与腾讯AI Lab联合，为了解决AI制药中的分子表示问题，提出了大分子数据上的自监督图transformer，这是分子表示学习中用到的最大的数据集和最大的GNN模型。文章发表在顶级会议NIPS 2020 上。

在AI制药领域中，获得分子的表示是非常重要的。之前的研究人员是将分子抽象成图，并且利用GNN来进行分子的表示。但是在实际应用中，以下两种情形阻碍了GNN的应用：

（1）在监督学习中，已标记的分子不足。

（2）对新合成的分子泛化能力不足。

为了解决以上问题，我们提出了一种新的架构——GOVER（Graph Representation frOm self superVised mEssage passing tRansformer.）自监督消息传递transformer上的图表示。GROVER在节点、边、图的层面上的自监督任务，使得GROVER可以从大量的未标记分子数据中学习丰富的结构和语义信息。

为了将复杂的信息编码，GROVER将消息传递网络聚合成transformer架构，这样分子的编码更具有表现力。GROVER的灵活性使其可以自监督地在大分子数据集上训练，因为不再需要标记，所以之前的两个问题就得到了解决。

我们在一千万个未标记的分子上用一亿个参数来预训练GROVER，这是分子表示学习中用到的最大的数据集和最大的GNN模型。相对之前的最先进的模型，我们在11个基准上获得了极大的性能提升（平均高于6%）。

1.前言

近年来，研究人员在探索怎样利用深度学习来加快药物发现的进度，以及通过便利分子的识别来降低代价。

图神经网络（GNN）虽然在许多任务上取得了不错的结果。但是仍然被两个问题所困扰：

（1）在监督学习中，已标记的分子不足。

（2）对新合成的分子泛化能力不足。

造成这两个问题的原因在于，获取分子特性的标签需要进行耗时耗资源的湿实验室试验。而许多公公送的分子基准也缺少充足的标签。在这些基准上做深度学习，极容易出现过拟合现象，因而泛化能力不足。

为了解决这些问题，人们尝试了许多工作，但是这些工作仍然不能明确地进行结构信息的编码，因为这些信息不是拓扑感知的。

还有一些工作在分子图表示上建立提前训练好的模型，但是我们认为这些提前训练好的方式不是最优的。

在这篇文章中，我们提出了GROVER，建立了两种自监督的任务。

在节点-边层面的任务上，GROVER不是单独地预测节点/边类型，GROVER随机遮罩目标节点/边缘的局部子图，并根据节点表示预测上下文属性。这样GROVER不仅考虑了目标节点/边，还考虑了它的遮罩环境，这样就能减少歧义。

对于图层面的任务，通过合并通过合并领域知识，GROVER可以提取分子图中存在的语义主题，并在图表示中预测这些分子主题的出现。GROVER可以通过大量的未标记分子数据学习到丰富的结构、语义信息。为了进行更复杂信息的编码，GROVER将消息传递网络和transformer架构融合起来，来获得分子的更好的表示。GROVER的灵活性让它可以在大分子数据上被高效地训练，而不需要任何的监督。我们还用提前训练好的GROVER模型来进行下游分子的预测任务，并做了微调。

2.GOVER预训练模型

GROVER由两个模组组成，节点GNN transformer和边GNN transformer。这里仅介绍节点GNN transformer（简称GTransformer），如图1所示

图1

GTransformer的关键部分是图多头注意力组件，这个部分需要获得向量形式的输入，但是原始输入通常是未向量化的结构化数据。于是我们设计了dyMPN来从图的节点中提取向量，输入到注意力区。

GTransformer的高表现力可以归因于其双层信息提取框架。由GNN模型的输出可以获得局部子图结构，提取可得第一层信息；同时，在由V构成的完全连接图上，可以将Transformer编码器视为GAT的变体，这样就可以提取节点的全局联系信息，提取可得第二层信息。

预训练模型的成功关键取决于自监督任务的设计。我们在预训练中不使用监督标签，并在这两个级别上提出新的自我监督任务：上下文属性预测和图级主题预测，如图2所示。

图2

3.实验

我们根据MoleculeNet的10种流行基准和几种最新技术（STOA）对GROVER进行了综合评估。

表1

表1展示了数据集上模型得到的结果。其中灰色部分是之前的模型得到的最好的结果，蓝色标记GROVER得到的最好的结果。可以看到：

（1）GROVER模型在所有数据集上始终表现出最佳性能。所有数据集的总体相对改进为6.1％（分类任务为2.2％，回归任务为10.8％）。

（2）其中GROVERbase模型在8/11的数据集上表现比之前的模型要好 ，而GROVERlarge模型在全部的数据集上表现得比之前的模型要好。

（3）在只有642个标记分子的小数据集FreeSolv上，GROVER相对之前的模型取得了23.9%的提升。

然后进行了消融实验：

1.自监督的预训练

表2

从表2可以看到自监督的预训练相对无自监督预训练的模型平均获得了3.8%的AUC提升。可以认为自监督的预训练策略可以学到隐藏的信息，并且提升预测性能。

2.GTransformer

图4

从图4可以看出，有GTransformer的GROVER相对其他模型表现的更好，证明GTransformer是有效的。

4.总结和未来工作

我们在这项工作中探索了大规模预训练GNN模型的潜力。我们的模型GROVER可以学习大量未标记图信息。在11个具有挑战性的方面相对之前的先进模型而言，取得了巨大的进步（平均超过6%）。

但是在以下几个方面GROVER仍然有改进的空间：

（1）更多的自监督任务。精心设计的自我监督任务是GNN成功的关键预训练。除了本文介绍的任务外，其他有意义的任务也将有所助益训练前的表现。

（2）更多下游任务。希望探索更大范围的下游任务，例如不同种类图形上的节点预测和链接预测任务。

（3）更宽更深的模型。大规模的模型可以捕获更丰富的语义信息来进行更复杂的任务。因此可以尝试更大规模的数据集和模型。

您的“点赞/在看/分享”是我们坚持的最大动力！

坚持不易，卖萌打滚求鼓励 (ฅ>ω<*ฅ)