Molecular Representation Learning via Heterogeneous Motif Graph Neural Networks

基本信息

发表于ICML 2022

博客贡献人

徐宁

作者

Zhaoning Yu ， Hongyang Gao

摘要

研究分子图的特征表示学习问题。图神经网络在分子图的特征表示学习中有着广泛的应用。然而，现有的大多数方法都是单独处理分子图，而忽略了它们之间的连接，如基序级关系。为了解决这一问题，我们提出了一种新的分子图表示学习方法。特别地，我们构建了一个包含基序节点和分子节点的异构基序图。每个基序节点对应于从分子中提取的基序。然后，我们提出了一种异构基序图神经网络(HM-GNN)来学习异构基序图中每个节点的特征表示。我们的异构基序图也能够有效地进行多任务学习，特别是对于小分子数据集。为了解决潜在的效率问题，我们提出使用边缘采样器，它可以显著减少计算资源的使用。实验结果表明，该模型的性能优于以往的先进模型。在多任务设置下，我们的方法在组合数据集上的良好性能为小分子数据集提供了新的学习模式。最后，我们证明了我们的模型在显著减少计算资源的情况下获得了相似的性能。

问题定义

在如今下游应用中，例如分子性质预测中，现有的大部分GNN都只考虑了最基本的分子图拓扑，通过邻域特征聚合和池化的方式获取分子结构特征信息。并没有考虑到分子图间的联系，特别是针对于分子图中有关motif的共享。

方法

方法描述

本文的搭建模型所进行的工作主要分为提取motif词汇、构造分子-motif异构图、搭建motif异构图神经网络（HM-GNN）、基于异构motif图进行多任务学习和为了提升训练效率采用的边缘采样策略。

分子图的motif字典

在分子图中，motif是指重复出现并且具有统计学显著性存在特殊性质的子图，利于分子图中的化学键，环状结构可以被当作motif。

图1.motif提取示意图

为了便于后续构建基于motif的异构图，需要先从数据集中提取出motif并构建成字典。这项工作中，首先确定提取motif的方法，然后遍历数据集中的所有分子图。例如上图一中，对对苯甲磺酸分子motif的提取，首先从原子图中提取出六个键和环结构，然后去掉重复的部分后加入到motif字典中。同时为了去除motif字典中的噪声motif（只含有较少的分子特征信息），采用TF-IDF（词频-逆文档频率）算法对motif进行排序，最后只保留最基本的motif作为最终motif字典。

TF-IDF(词频-逆文档频率)算法：词频表示为分子中出现motif的频率，逆文档频率指包含motif的分子数量。
T F − I D F i j = C ( i ) j ( log ⁡ 1 + M 1 + N ( i ) + 1 ) {TF-IDF}_{ij}=C(i)_j(\log \frac {1+M}{1+N(i)} +1) TF−IDFij=C(i)j(log1+N(i)1+M+1)
其中 C ( i ) j C(i)_j C(i)j表示是motif i出现在分子j中的次数，M指的是这个数据集包含的分子图个数，N(i)指的是有多少分子图包含motif i。最终以包含对应motif的分子的TF-IDF平均值作为motif的TF-IDF值进行排序。

其次模型中对于字典中motif的选择也较为重要，因为环状结构可以表示较多的motif，因此本文中选择选取所有的环状结构和不在环状结构上的边作为motif构成字典，既保证了字典的motif足够基础也控制了字典的大小，保证了提取motif的时间复杂度不会太高。

构造分子-motif异构图

分子-motif异构图中包含有两种不同的节点，分别是分子节点和motif节点，其中分子节点表示数据集中的分子图，motif节点表示上述建立的motif字典中的一个motif。
同时异构图中也包含着两种边，一种是motif节点到motif节点的边，另一种是motif节点到分子节点的边。并且因为不同类型的边和不同的motif对分子的贡献也不同，因为也需要对不同的边设置不同的权重。

分子-motif边：如果分子中包含该motif，则对应分子节点和motif节点连接。使用TF-IDF算法计算边的权重。
motif-motif边：如果两个motif在分子中共享了至少一个原子，将两个motif节点连接。使用PMI算法计算边的权重。对应权重计算公式如下：
A i j = { P M I i j i , j 节点都是 m o t i f s T F − I D F i j i , j 节点其中一个是 m o t i f 0 A_{ij}=\left\{ \begin{array}{rcl} PMI_{ij} & &{ i,j 节点都是motifs} \\ TF-IDF_{ij} & &{i,j节点其中一个是motif}\\ 0 \end{array} \right. Aij=⎩ ⎨ ⎧PMIijTF−IDFij0i,j节点都是motifsi,j节点其中一个是motif

TF-IDF算法计算方式已在构建motif字典时给出，PMI算法计算方式如下：

P M I i j = log ⁡ p ( i , j ) p ( i ) p ( j ) PMI_{ij}=\log \frac{p(i,j)}{p(i)p(j)} PMIij=logp(i)p(j)p(i,j)
其中 p ( i , j ) p(i,j) p(i,j)是分子中motif i和motif j同时出现的概率， p ( i ) , p ( j ) p(i),p(j) p(i),p(j)分别表示motif i和moitf j出现的概率，若存在motif的PMI值为负，则权重计为零。

图2.分子-motif异构图

异构神经网络

本文主要搭建了一个同时学习原子级和motif级的图特征表示的神经网络。
如下图所示，模型上层即采用传统的GNN模型对分子图提取出原子级别的特征。然后在下层对上述构造分子-motif异构图提取motif级别的特征，因为该异构图包含了整个数据集中的分子和motif信息，所以若正常按GNN提取特征，维度将远大于上层的传统分子图，所以采取对异构图中的motif节点使用one-hot向量表示节点特征，对于其中的分子节点采用词袋模型（bag-of-words）填充特征向量。将每个motif看作单词，每个分子看作一个文档，通过词袋模型即可得到分子节点的特征向量，采用这种方式即可得到每个分子的motif级特征嵌入。最后将上下层的特征嵌入拼接在一起作为最终的分子图特征嵌入作为后面下游应用的输入。

图3.HM-GNN模型示意图

基于异构motif图进行多任务学习

因为在分子领域存在很多小数据集，每个数据集可能只有几百或者几千个分子图，使用GNN在其上面训练很容易出现过拟合的现象。而使用本文的异构图，可以做到同时训练不同的数据集，通过不同的数据集间共享motifs的方法组合成一个大的分子数据集进行训练，达到多任务训练的训练模式。当然，这一训练的前提就是不同的数据集需要共用很多的motif，这样才能构建成一个比较大的异构图。
在这种多任务训练中，构造motif字典时，其中motif只需要在部分的数据集中共享。然后就能在所有的数据集上构建异构motif图，其中motif便是作为连接不同数据集的分子的连接器。

边采样策略

可知用上述方法构建的异构图的大小是关联于数据集的大小，若涉及的数据集非常大，则对应的异构图也会非常大，对计算资源的需求也会加大，模型的训练效率也会降低，也会存在很难把异构图放在模型中并同时训练图上的点。
因此需要采用采样的方法在原始的异构图上采样出一个子图作为输入。同时根据异构图的特点，即异构图上既存在motif 节点，也有要预测的分子节点。所以构建了一个边采样的方法。如下图所示，采样策略为：

1.随机选择一些需要预测的分子节点作为起始节点（红色节点），然后采用广度优先算法进行选择，每层都采样固定数量的边。
2.在第一跳中，都是分子节点到motif节点的边，为保证信息不产生丢失，所以会选取所有分子节点到motif的边。
3.从第二跳开始，只对图中motif到motif的边进行采样选取，以此来尽可能的保留更多的motif信息。

图4.边采样示意图

实验

实验设置

本文的实验，主要为模型性能在不同数据集上的研究，消融实验，计算效率的研究，motif字典大小的性能差距等方面进行实验。

数据集

TUDataset的五个流行的生物信息学基准数据集：
PTC：一个包含344种化学化合物的数据集，报告了19个不同标签的雄性和雌性大鼠的致癌性。
MUTAG：一个包含188个致突变芳香族和异芳香族硝基化合物的数据集，其任务是预测它们对伤寒沙门氏菌的致突变性。
NCI1：由国家癌症研究所(NCI)公开发布。它是一个筛选的复合平衡数据集的子集，设计用于抑制或抑制一组具有37个离散标签的人类肿瘤细胞株的生长。
Mutagenicity：是药物用化合物的数据集，可分为两类:诱变剂和非诱变剂。
PROTEINS：其中节点是二级结构元素(SSE)。如果两个节点是氨基酸序列或三维空间中的相邻节点，则它们之间存在边。它有3个离散的标签，分别表示螺旋、切片或旋转。

小数据集：
PTC MM：小鼠数据集
PTC FR：大鼠数据集

对比方法

PatchySAN
GCN
GraphSAGE
DGCNN
GIN
PPGN
CapsGNN
WEGL
GraphNorm
GSN

实验结果及分析

分子图数据集性能研究

根据TUDataset中的五个流行生物信息学图基准数据集进行图分类任务来评估模型，包含四个分子数据集和一个蛋白质数据集。

表1.分子图数据集上性能比较

由上表一可以看出本文模型在5个数据集上的表现都优于baseline（baseline的数据来自原论文），可说明由motif字典构成的motif节点可以帮助GNN更好的学习分子图的motif特征表示，在蛋白质数据集上的优异表现，也可说明蛋白质中的motif也包含有有效的结构信息。
同时也引用了motif信息的GSN模型，虽然没有考虑到motif和motif之间的信息交流和分子与分子之间的信息交流，但实验结果也好于其他的baseline。

HM-GNN的消融实验

表2.HM-GNN消融实验结果表

通过在HM-GNN中去除异构图和相应的GNN，即只保留上层提取原子级特征信息的GIN，作为消融实验对比项，在三个数据集上进行对比其图分类精度，由上表可看出，包含motif级特征学习的HM-GNN性能明显优于GIN，说了HM-GNN中motif级特征学习的有效性。

motif-motif相互作用的消融研究

表3.motif-motif相互作用的消融研究结果表

在这个实验中，通过在异构图中去除motif到motif的边形成一个变体Variant与HM-GNN在三个数据集上进行性能的比较，可以看出Variant的实验效果相比于HM-GNN下降了大约1%，说明了异构图中motif与motif之间的信息交流是有效的。

小型分子数据集的多任务学习

表3.多任务学习实验结果表

这个实验主要验证了HM-GNN模型在小数据集上做多任务学习的有效性的实验结果。实验中所使用的三个数据集分别为PTC，PTC+PTC_MM,PTC+PTC_FR,这三个数据集都较小只含有几百个图。表中Vocab Size表示数据集中含有的motif字典的词汇量，Overlap Ratio表示的两个组合的数据集中moitf的重叠程度，后面的比例即表示训练集所使用的数据集的比例。
由此可以看出当两个数据集共享的motif越多，其最后训练出来的模型效果越好，即使只使用其中10%的数据集作为训练集其效果也超过大部分的baseline，说明了通过motif作为不同小数据集之间的联系的方式做多任务学习对于解决小数据集训练过拟合问题的有效性。

motif字典大小研究

在这一部分的实验中，通过实验来研究在motif字典中所保留的motif比例的不同对模型性能的影响。
如下图所示，刚开始模型的性能随着字典中motif保留率的增加而提高，保留率越高，词汇量中的motif就越多，图中的分子可以有更多的连接（motif）与其他分子通信，使得motif级特征传播的效果更好。当保留率大于80%时，模型性能开始下降，表明最后20%基序是有噪声的，会影响模型的泛化和鲁棒性。也说明了本文使用TF-IDF算法过滤字典中噪声motif的有效性。

图5.motif字典大小实验结果图

总结

亮点

提出了包含motif节点和分子节点的异构motif图利用了motif级的图特征。
通过motif连接不同分子，实现了分子图之间的信息传递。
说明了基于小分子数据集的多任务学习的有效性。
通过特定的采样方式即缩小的异构图的大小，也尽可能的保留的大部分的motif信息。

[不足]

创新点较为松散，其中的小数据集上的多任务学习与其他点的联系在逻辑上不太紧密。

BibTex


@InProceedings{pmlr-v162-yu22a,title =     {Molecular Representation Learning via Heterogeneous Motif Graph Neural Networks},author =       {Yu, Zhaoning and Gao, Hongyang},booktitle =     {Proceedings of the 39th International Conference on Machine Learning},pages =     {25581--25594},year =      {2022},editor =    {Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan},volume =    {162},series =     {Proceedings of Machine Learning Research},month =     {17--23 Jul},publisher =    {PMLR},pdf =      {https://proceedings.mlr.press/v162/yu22a/yu22a.pdf},url =     {https://proceedings.mlr.press/v162/yu22a.html},
}