Brief Bioinform｜MESPool：分子边缘收缩池化用于分子表示学习和性质预测

2023年12月2日，西安交通大学龙建纲教授团队，联合西安交通大学孟德宇教授，在Briefings in Bioinformatics上发表文章MESPool: Molecular Edge Shrinkage Pooling for hierarchical molecular representation learning and property prediction。

作者提出了一种基于边缘(或化学键)的分子边缘收缩池化(Molecular Edge Shrinkage Pooling，MESPool)方法。MESPool保留了关键的边，并将官能团内的其他边进行收缩，从而能够在保留原始图的连接的情况下搜索分子中的关键结构。实验结果表明MESPool的性能优于以前的方法。

背景

分子性质预测是药物发现的一项基本任务，在计算机辅助药物发现的工作流中起着至关重要的作用，因为许多方法依赖于预测的分子性质来评估、选择和生成分子。

近年来，人工智能驱动的分子性质预测方法已成为先导化合物发现与优化领域的热点。同时，图神经网络(GNN)在图表示学习方面已经显示出强大的能力，并被进一步应用于分子图数据。识别与任务相关的结构对于分子性质预测非常重要。在GNN中，图池化可以对节点进行分组并分层表示分子图。

然而，之前的池化方法或者会丢失节点信息，或者无法准确保留原始图的连接。因此，现有的方法很难识别连续的子结构，而且在分子图上缺乏可解释性。

方法

边缘及其特征在许多现实世界的图数据中起着至关重要的作用。对于分子图，边特征描述化学键类型、共轭、环和立体信息。此外，边可以表示具有两个连接节点的图中的最小子结构。MESPool的主要思想是保留关键边，并将其他结构收缩为超级节点(见图1A, B)。与之前的方法不同，MESPool最初选择边而不是节点作为单元进行池化。在这个框架中，边缘选择比节点选择更能区分连接的子结构(见图1C)，这有助于分层表示。

图1 MESPool结构图

以往的图池化方法根据节点选择方法大致可以分为两类，即稀疏池化和密集池化。稀疏池化的目的是保留每层中与任务相关的节点，剔除不相关的节点，过程中会丢失部分结构信息，但保留原有连接。同时设置阈值(最小分数)，自适应调整节点池数。密集池化的目的是对节点进行聚类，并在每一层对图进行分层表示。密集池化不会丢弃节点信息，能够更好地保留原始图结构，但池化图的重建连接与原始图的关联不强，且密集池化比稀疏池化需要更多的计算资源。

MESPool可以看作是稀疏池化和密集池化的混合，它像稀疏池化一样通过评分来选择和分割单元，但同时，在保留低评分的节点(单元)的情况下，像密集池化一样将它们的簇减少为超级节点。如图1C所示，MESPool同时进行边缘选择和节点选择，在能够准确识别结构信息的同时，其参数数量不会随着图的大小而递增，这使得MESPool能够兼具优秀的预测性能和计算效率。

如图1A所示，在MESPool的选择过程中，MESPool对无向边进行评分，以表示相应单元的权重，相邻的低评分单元将被视为一个子集(池化子结构)。因此，在对边进行评分时，有必要考虑邻接信息，因此，在选择过程中采用边缘消息传播，使得边缘特征可以通过分数进一步更新。在此过程中，MESPool不仅保持了原始图的连通性，而且不丢失任何节点信息。此外，由于骨架被收缩为节点，因此在池化图上侧链之间的距离缩短，因为池化后的官能团的结合和相互作用可以在更深的网络中表示。

池化集拆分后，采用强连接分量查找算法，用于将连接的池节点分组为子集。作者提出了一种边缘特征图同构网络（EGIN）应用于图的卷积和读出。在EGIN层中，首先通过连接其起始节点特征来更新边缘特征，并将更新后的边缘特征作为图同构网络（GIN）中从邻居节点到中心节点的加权消息⁠，并使用连接函数来组合不同维度的中心节点特征。该算法以层次化的形式进行（图1B），使得子集可以进一步缩小为具有更新节点特征的超级节点。在每个块的末尾应用求和读出函数，聚合节点特征以获得分层图表示。最后，将输出的分层图表示连接并传递给线性层进行分类。

在子结构被压缩成超级节点后，在更深的网络中，它们对整个图的影响将被减小，可以认为收缩的部分是任务无关的。然而，在一些特殊情况下，会有一些节点对与多个边相连或与多个自环相连的一个节点。因此，通过将连接相同节点的重复边相加为一条边来优化连接(图1D)。

初始的原子特征为120维，以独热编码的形式来编码原子类型（100维），原子是或不是芳环的一部分（1维），成键氢原子的数目（5维），形式电荷数（5维），手性原子的手性类型（4维），杂化原子的杂化形式（5维）。初始的边缘特征为13维，以独热编码的形式来编码化学键类型（5维），键是共轭的还是非共轭的（1维），化学键是或不是环的一部分（1维），立体键的立体类型（6维）。在实验中，采用随机骨架的数据集分割方法。根据分子的骨架进行数据集划分，使得预测比随机划分更具挑战性。每个方法的超参数(包括隐藏层维度、池化率、dropout比率和学习率)通过网格搜索独立调整。

结果

作者将MESPool与一些具有代表性的方法进行了比较，如表1所示，加粗的数字表明RMSE的均值最低（最优），下划线的数字表明RMSE的标准差最低（最鲁棒）。在FreeSolv和Lipophilicity这两个数据集上，MESPool的RMSE均低于现有方法，且具有较低的标准差，鲁棒性较强。在ESOL数据集上，本文提出的EGIN的RMSE低于现有方法，且具有最低的标准差。

表1 与其他方法对比

作者设计了消融实验，以EGIN为主干模型，研究不同的池化机制（TopK, SAG, MESPool）在EGIN上的影响。BACE1是产生神经毒性的β-淀粉样肽所需的β-分泌酶。抑制BACE1是阿尔茨海默病的重要治疗手段之一。BACE数据集的二元分类标签是定性的结果。抑制剂标记为1，非抑制剂标记为0。GRL-8234是一种良好的BACE1抑制剂，图2显示了GRL-8234的池化结果。TopK和SAG均错误预测GRL-8234的评分低于0.5，可视化结果随着网络层数的增加，结构信息的丢失越来越明显，留下的是缺乏化学合理性的稀疏结构。相比之下，MESPool预测的准确性源于结构选择的合理性。MESPool识别出了GRL-8234中具有关键作用的异邻苯二甲酰亚胺骨架(黄色部分)和3-甲氧基苯基(粉色部分)，而对次要的结构（蓝色部分）做了池化。消融实验表明MESPool的边缘选择与收缩机制对准确识别药物分子及其关键子结构十分重要。

图2 消融实验

作者还进行了案例分析，选择两种抗新冠药物，瑞德西韦和VV116进行讨论(图3)，MESPool对这两种药物的预测评分分别为0.99和0.94，这两种药物都是RNA依赖性-RNA聚合酶(RdRp)抑制剂，它们是在核苷类似物核心GS-441524上设计的。GS-441524是一种药物前体，能够扩散到细胞内，通过磷酸化缓慢转化为单磷酸核苷，并与磷酸激酶一起加工成活性的三磷酸核苷衍生物，抑制RdRp。瑞德西韦是GS-441524的单磷酸盐，其附加的官能团加速了磷酸化过程。另一方面，三异丁酸酯VV116通过7-氘化GS-441524的酯化反应获得了良好的口服生物利用度。从图3C可以看出，池化层识别了瑞德西韦和VV116上的核苷类似物核心（瑞德西韦分子的浅粉色部分，VV116分子的粉色部分）。此外，瑞德西韦保留了单磷酸结构，添加的苄基和极性/非极性混合官能团分别被收缩。此外，对VV116上的3个异丁酰基也进行了识别和收缩。

图3 案例分析

总结

本文提出了一种新的用于分子表示学习的图池化方法MESPool，该方法通过选择边缘来减少结构，并能够自适应调整池化比例。与以前的方法最大的不同是，MESPool可以直接在原始图上搜索任务相关的结构，这使得MESPool具有合理性。MESPool表现出比现有方法更好的性能。同时，池化结果具有良好的化学直觉，符合药物设计逻辑。

在许多与分子相关的任务中，学习找到关键的子结构是一个重要的课题。除了性质预测外，还有药物-药物相互作用、药物-蛋白质相互作用和分子生成等。MESPool有望应用到更多的研究领域。此外，通过预训练等方法引入化学先验信息，加强MESPool的可解释性，进一步识别新的官能团，也是未来关键的研究方向。

参考文献

[1] Xu et al. MESPool: Molecular Edge Shrinkage Pooling for hierarchical molecular representation learning and property prediction. Brief Bioinform. 2023