NeurIPS-21 | MGSSL: 基于官能团的分子属性预测图网络自监督预训练

本文介绍一篇来自中国科学技术大学刘淇教授课题组和腾讯量子实验室联合发表的一篇文章。该文章提出了基于官能团的分子属性预测图网络自监督预训练方法MGSSL。MGSSL结合化学领域知识，在大量无标签分子数据中划分官能团和构建官能团树，并进一步通过官能团生成式自监督任务，使得预训练的图网络可以学习到官能团的结构和语义信息，提高下游分子属性预测任务的效果。

1.研究背景

分子属性预测任务对于药物合成和筛选具有重要意义，例如新冠病毒药物筛选。传统通过实验和理论计算的方法的得到分子属性耗时且昂贵。近年来，基于深度学习的分子属性预测方法大大减少了预测时间和成本。为了保留分子的结构信息，分子通常可以表示成图，并利用图网络进行分子属性预测。但是训练图网络用于分子属性预测需要大量有标签分子数据，否则图网络模型易于过拟合并且泛化性能差。

为了克服有标签分子数据匮乏的问题，可以考虑引入类似于CV和NLP中自监督预训练的方法：先在大量无标签分子数据上预训练图网络，然后再在下游有标签分子数据上微调。如何对于图数据特别是分子数据，设计合适的自监督预训练方法是当下的研究热点和难点。目前的图自监督预训练方法大致可以分为两类：基于对比学习的方法和基于预测的方法。（1）基于对比学习的方法，让来自同一个图的视图表示接近，不同图的视图表示相异；（2）基于预测的方法从数据本身特点构建预测任务，例如通过周围的结点预测被遮掩住结点的类型。

然而当前分子图网络模型的自监督预训练无法有效利用分子官能团的信息，使得预训练效果受到限制。分子官能图指的是分子中经常出现并且对于分子属性具有决定作用的片段。例如，具有（-OH）的小分子通常水溶性比较好。因此，设计基于官能团的图网络自监督预训练方法。

2.MGSSL 方法

图1 MGSSL方法示意图

MGSSL 创新的提出了一种自监督官能团生成预测任务，使得预训练的图网络可以有效捕捉分子图中官能团的结构与语义信息。如图1所示，MGSSL 同时考虑了原子层级和官能团层级的自监督任务。在原子层级，MGSSL使用了遮掩并预测原子/化学键种类的方法学习到原子层级的信息。在官能团层级，MGSSL自监督框架由三部分构成：由化学领域知识启发的分子官能团划分，官能团生成预测和多层级自监督预训练。

图2 分子官能团划分

如图2所示，为了划分分子官能团，我们首先使用了基于化学反应模板的 BRICS（Breaking of Retrosynthetically Interesting Chemical Substructures）方法。然而通过BRICS划分得到的官能团词典较为冗余，我们还定义了两个规则去除冗余的官能团，最终构建 motif tree：（1）当一个化学键连接的两个原子，一个在环上，一个不在环上时，断开化学键；（2）当不在环上的原子有三个或者三个以上邻居时，断开与邻居间的化学键，并且选取所述不在环上的原子作为新的官能团。在划分官能团后，我们可以构建官能团树并且得到预训练数据集的官能团词典。

图3 官能团树生成（深度优先顺序和广度优先顺序）

如图3所示，利用构建的官能团树，可以建立官能团生成的自监督预训练任务，官能团生成的每一步都进行拓扑预测和官能团种类预测。其中官能团的生成顺序，主要包括深度优先顺序和广度优先顺序。最终MGSSL自监督预训练的损失函数包括原子层级和官能团层级的损失函数。

3.实验结果

表1分子属性预测实验结果

在论文中，作者主要考虑分子属性预测任务。自监督预训练在ZINC15 中250k 个分子数据上进行。下游分子属性预测数据集使用的是MoleculeNet中8个双分类数据集。实验中数据集划分使用的是scaffold-split。MGSSL的对比方法主要有Infomax, Attribute masking, GCC, GPT-GNN 和 Grover。

下游分子属性预测实验结果如表1所示，MGSSL 在MoleculeNet 8个数据集中的7个得到了最好的表现，验证了基于分子官能团的自监督预训练的有效性。

图4 不同预训练方法的训练和测试曲线

在图4中，我们展示了不同预训练方法的训练和测试ROC-AUC随着微调epoch的变化。可以看出，对比其他自监督预训练和没有预训练直接微调，MGSSL预训练过的图网络具有更快的收敛速度。一旦使用MGSSL预训练好之后，图网络可以以很小的预训练开销快速迁移到各种下游任务中去。

表2 对比MGSSL在不同图网络模型的效果提升

在表2中，我们进一步对比了使用MGSSL在不同图网络模型的效果提升。可以看到MGSSL是通用的图网络自监督预训练方法，可以提升GCN,GIN, RGCN, DAGNN, 和GraphSAGE在分子属性预测任务中的表现。

4.总结

这篇论文提出了一种基于官能团生成的图网络自监督预训练方法MGSSL。大量实验表明MGSSL可以使得预训练图网络学习到官能团结构和语义信息，提升下游分子属性预测的表现。未来的研究工作包括：（1）设计更多基于官能团的图网络自监督学习任务 （2）探索基于官能团的自监督学习在其他领域的应用。