本文介绍由爱荷华州立大学的Zhaoning Yu 与Hongyang Gao发表于ICML的研究工作。作者提出了一种新的分子图表示学习方法，构建了一个包含基序节点和分子节点的异构图，并提出HM-GNN（Heterogeneous Motif-GNN）模型学习异构图中每个节点的特征表示。该模型支持多任务学习，适用于小规模数据集。此外，为解决潜在的效率问题，作者采用边采样方法减少计算资源。作者在多种分子数据集上对HM-GNN进行了评估，结果表明该模型优于现有的先进模型。该模型可以作为一种新的多任务学习方法来用于图学习研究。
1
简介
图神经网络已广泛应用于分子图的特征表示学习。然而，现有的大多数方法都是单独处理分子图，忽略了它们之间的联系，并且现有分子数据集规模普遍较小，图神经网络训练时很容易出现过拟合的现象。
针对上述问题作者提出了一种新的分子图构建方式，首先生成一个异构基序图将每个分子图联系起来，然后提出HM-GNN学习异构基序图中的节点特征，最后利用边采样器提高网络训练效率。作者的贡献如下：
- 提出了一种新的分子图构建方式，该方法可以把多个分子图联系起来，对于小规模数据也支持多任务学习
- 提出了HM-GNN来学习异构基序图中的节点特征。HM-GNN包含两个部分的图嵌入学习：①异构图中每个分子节点的基序级特征嵌入，②单个分子图的原子级特征嵌入，最后将两部分的特征嵌入连接起来，并输入多层感知器(MLP)进行预测。
- 分子较多的情况下生成的异构基序图过于复杂，影响训练效率。针对此问题采用边采样器在不降低模型的性能的同时加快训练速度。
2
方法
2.1分子基序词典
图1展示了由分子图构建基序词典的过程。首先确定分子图中所含的键和环，然后利用TF-IDF去除携带分子表征信息较少的结构，最终留下的结构作为后续建立异构基序图的参考词典。
图1 基序词典构建示意图
2.2异构基序图的构建
异构基序图包含分子节点和基序节点，图中有两种边，分别是“分子-基序”和“基序-基序”。首先在分子与由该分子中提取的基序之间添加“分子-基序”边；其次，如果两个基序在任意分子图中共享至少一个原子，那么在它们之间添加一条“基序-基序”边，最终得到一个包含所有分子以及基序词典中所有基序的异构基序图，如图2所示。
图2 异构基序图
在异构基序图中不同边的权重有所不同。对于分子-基序边，使用基序的TF-IDF值作为权重，如公式（1）所示，其中代表基序在分子中出现的次数，表示分子的个数，表示包含基序的分子个数。对于基序-基序边，使用PMI作为权重，如公式（2）所示。
2.3异构基序图神经网络（HM-GNN）
首先为每个基序节点和分子节点初始化特征，对于基序节点，使用独热编码生成其特征；对于分子节点，使用词袋模型生成其特征向量。在此基础上，利用H-GNN学习异构图中每个分子的基序级embedding。同时，利用原子作为节点，化学键作为边，将每个分子转换成一个图，利用GNN来学习该图包含原子级信息的embedding。最后将两个层级的embedding信息进行拼接得到最终的特征编码并输入到MLP中进行预测，如图3所示。
图3 HM-GNN示意图
2.4基于异构基序图的多任务学习
大多数分子数据集相对较小，使用图神经网络时很容易发生过拟合。多任务学习能有效降低过拟合并提高模型泛化能力。然而，由于不同数据集间缺乏明确的联系，很难直接使用多任务学习。异构基序图的构建可以很方便地连接一组分子数据，从而进行多任务学习以提升模型性能。给定个分子的数据集，首先从个分子数据集中构建基序词典，其中基序只需要在部分数据集中共享。然后构建异构基序图，使用HM-GNN学习每个分子的特征表示输入到单独的MLP中进行预测。在这个过程中，基序节点被认为是连接不同数据集的分子的连接器。
2.5通过边采样提升训练效率
异构基序图的边可能会随着分子数量的增多而成倍增加，从而影响训练速度，因此可以通过边采样器来减少异构基序图边的数量。在异构基序图中随机选择一个分子节点作为起始节点，进行广度优先搜索，每层均随机采样固定大小的边。其中，分子节点的第一跳均为基序节点，其过程如图4所示。图中红色节点为分子节点，蓝色节点为基序节点，经过3跳，每跳保留两边，最终得到简化后的异构基序图。
图4 边采样示意图
3
实验
3.1不同分子集上的模型性能
在五种生物信息学数据集上(PTC、MUTAG、NCI1、Mutagenicity分子数据集及 PROTEINS蛋白质数据集)通过5层GNN和2层MLP来验证HM-GNN对于分子图特征提取的性能，与其他主流方法的对比结果如表1所示。实验结果显示在五种数据集上HM-GNN的准确率均高于其他方法，这表明基序节点可以帮助GNN更好地学习分子图的基序级特征表示。在蛋白质数据集上HM-GNN模型也表现最佳，这表明蛋白质分子中的基序也包含有效的结构信息。与使用基序信息的GSN相比， HM-GNN通过基序-基序和基序-分子间的相互作用更好地学习了基序embedding。
表1各种图模型的分类准确率
3.2大规模数据集的性能研究
为了在大规模数据集上进行评估，作者使用了ogbg-molhiv和ogbg-molpcba数据集，前者为二元分类数据集，后者为多元分类数据集。在ogbg-molhiv数据集上的模型评价指标选取ROC-AUC，在ogbg-molpcba数据集上的模型评价指标选取AP。表2表明了HM-GNN在两种数据集上均优于其他方法，说明该方法在大规模数据集上具有良好的泛化能力。
表2大规模分子数据集上的分类准确率
3.3消融实验
异构基序图消融实验 为证明HM-GNN中基序级分子特征学习的有效性，将异构基序图和相应的GNN从HM-GNN中移除，从而简化为GINs。比较HM-GNN和GINs在PTC、MUTAG和PROTEINS数据集的准确率，结果如表3所示，证明了基序级特征对分子特征表示学习的重要性。
基序-基序相互作用的消融研究 在异构基序图中，将“基序-基序”边删去，得到异构基序图的变体，在三个数据集上评估变体与原本的异构基序图之间的性能，结果如表4所示，去除基序边后模型性能下降，证明“基序-基序”边的重要性。
表3 仅保留GIN的消融实验
表4 去除“基序-基序”边的消融实验
3.4 小规模分子图多任务学习
为证明多任务学习的有效性，使用PTC _MM和PTC_FR分别与PTC结合，创建新数据集PTC PTC_MM和PTC PTC_FR。表5汇总了新数据集词汇量的大小以及与原始PTC数据集的重叠率。选取不同比例的数据作为训练集，其余的数据集作为测试集，以此用小规模数据验证HM-GNN相较于其他方法的过拟合情况。实验结果表明PTC结合PTC_FR比PTC_MM性能提升更大，由于PTC与PTC_FR基序词重叠更多。因此，将数据集与相似基序词相结合有利于小规模分子数据集上的多任务学习。
表5多任务学习的性能指标
3.5 边采样实验
图5展示了边采样时起始节点的数量与内存使用率以及预测精度的关系。随着起始节点数量的增加，内存使用率线性提升，当起始节点数量为25000时，模型精度基本趋于稳定。因此设置合适的起始节点数值，可以在不减少模型精度地情况下，提升模型的计算效率，减少训练的资源使用率。
图5 边采样起始节点数对性能和资源的影响
3.6 基序词典的大小实验
在构建异构基序图的过程中，通过TF-IDF的值来去除一部分不重要的基序，不同的TF-IDF的阈值选取也会导致不同的异构基序图，从而影响模型最终性能，图6显示了不同阈值对模型的性能影响，在阈值为80%时模型性能最优。
图6 不同TF-IDF阈值对性能的影响
4
总结
在这项工作中，作者首先提出了一种新的异构基序图构建方法用于分子图特征表示学习，同时通过选择不同的TF-IDF阈值来筛选携带信息较少的基序，然后提出了HM-GNN学习异构基序图中的节点特征，并结合分子图得到分子的embedding表示，最后采用边采样器来提升模型训练效率。实验表明，该模型均优于目前其他的先进模型，并且通过实验证明了在小规模分子数据集上多任务学习的可行性。
参考资料
论文链接：https://arxiv.org/abs/2202.00529
源码
https://github.com/ZhaoningYu1996/HM-GNN
