Brief Bioinform｜KGE-UNIT：药物发现中基于知识图谱的分子相互作用预测与多任务统一学习框架

2024年2月12日，哈尔滨工业大学臧天仪教授、赵天意教授团队，在Briefings in Bioinformatics上发表文章KGE-UNIT: toward the unification of molecular interactions prediction based on knowledge graph and multi-task learning on drug discovery。

作者提出了一个结合知识图谱嵌入(KGE)和统一的多任务学习（unified multi-task learning，UNIT）框架，称为KGE-UNIT，用于同时预测药物-靶标相互作用(DTI)和药物-药物相互作用(DDI)，并提高每个任务的性能。KGE-UNIT通过KGE从药物知识图谱中提取异构特征，增强药物和靶标蛋白质节点的结构特征，从而提高特征的质量。此外，利用多任务学习，KGE-UNIT引入了一种创新的预测器，该预测器由基于任务感知的卷积神经网络(CNN)编码器和基于任务感知的注意力解码器组成，可以更好地融合多模态特征，捕捉分子任务的上下文交互并增强任务感知，从而提高性能。实验表明，所提出的模型优于现有的方法。

背景

分子相互作用的预测对药物发现至关重要。图神经网络（GNN）在分子相互作用预测中得到了广泛应用，然而，现有的基于GNN的方法往往侧重于单个预测任务，而忽略了它们之间的关系。此外，某些任务也可能由于数据量不足，从而导致性能受限。

方法

KGE-UNIT的工作流程(图1)包括三个主要部分：通过KGE和CNN进行多模态特征提取，使用基于CNN的编码器对特征进行集成和编码，以及使用任务感知的注意力解码器对任务感知特征进行解码。

图1 KGE-UNIT结构图

分子相互作用预测涉及分子的结构、功能和相互作用等多个维度的信息。构建知识图谱（KG）允许集成多源数据，从而更好地捕获分子之间的复杂关联。在这项研究中，作者整合了来自基因组学、蛋白质组学和代谢组学的异构数据，构建了一个KG，其中生物医学概念被表示为节点，相互作用/关联(如DTI、DDI和药物-疾病相互作用)被表示为边。例如，KG中的三元表示<DB15035, DTI, P04626>说明了药物DB15035与蛋白P04626之间的相互作用。因此，可以得到包含大量信息的KG三元组，包括拓扑结构和语义关系。

在构建KG之后，使用KGE模型来学习所有实体和关系的拓扑结构和语义关系。本研究选择了ConvE作为KGE模型，其中输入实体和关系之间的交互是由卷积和全连接层建模的。与其他KGE模型相比，ConvE在涉及分子相互作用预测的任务中表现出优越的适用性。这种高度的适用性主要归功于它使用卷积来增强特征学习能力。

结构特征提取器旨在从药物SMILES结构和蛋白质序列中提取药物和蛋白质实体的结构特征。为此，使用RDKit将药物SMILES结构转化为分子图，RDKit将原子表示为节点，化学键表示为边。此外，通过蛋白质的氨基酸残基的CTD（组成、转移和分布，Composition, Transition and Distribution）特征方法将蛋白质序列转化为高维特征向量。考虑到这些特征向量的高维性和潜在的噪声，应用主成分分析来降低维数，同时保留相关实体特征中包含的基本信息。然后，使用基于任务感知的CNN编码器分别从药物和蛋白质结构中提取局部化学背景和分子结构，使其能够融合局部和全局特征，从而提高分子关系预测中特征的质量。

为了整合多个分子相互作用任务，作者在KGE-UNIT中开发了一种新的编码器-解码器预测器。基于任务感知的CNN编码器融合了异构特征和结构特征。任务感知注意力解码器由任务交互注意块、任务感知注意块和任务特定预测头组成。任务交互注意块和任务感知注意块都由多头自注意力(MHSA)和多层感知器(MLP)组成。MHSA具有整合局部和全局信息的能力，从而促进生成内容更全面、更丰富的特征表示。每个注意力头的独立学习能力使它们能够关注输入序列中的不同位置，有效地捕获与靶标和药物分子中存在的氨基酸残基和原子相关的特征信息。通过并行进行缩放的点积注意力计算，MHSA对三个关键实体进行操作，即查询、键和值，它们都以矢量形式表示。

具体而言，任务交互注意块的目标是挖掘不同分子相互作用任务的相互作用信息，通过基于CNN的编码器对每个任务进行异构和结构特征的整合。任务感知注意块用于学习DDI和DTI任务的表示，同时考虑它们的交互上下文，理解任务之间的相互作用，并动态地为每个特定任务分配不同的注意力。每个任务由其相应的任务感知注意力块处理。这两个块具有相似的结构，但在查询、键和值的实体上有所不同。在任务交互注意块中，查询、键和值都通过MLP从原始特征中编码得到。在任务感知注意块中，将任务交互注意块的输出作为键和值，而基于CNN的编码器的输出作为查询。

在提取任务感知特征后，使用特定任务的预测头对分子相互作用进行预测。最后，采用DTI预测损失和DDI预测损失的加权和作为损失函数，对模型进行训练。

结果

作者将KGE-UNIT与一些具有代表性的方法进行了比较。如表1所示，采用AUROC和AUPR对比，KGE-UNIT在DTI和DDI预测的AUROC和AUPR中均超越了现有方法。

表1 与其他方法对比

作者设计了消融实验。为了探索多源特征（即基于KGE的特征）、结构特征和多任务预测如何提高KGE-UNIT的性能，对以下变体进行了研究:

无结构特征的KGE-UNIT联合多任务预测器(w/o SF+MP)仅使用ConvE分别预测DTI和DDI。应该注意的是，结果与表1和表2略有不同，因为表1上的ConvE结果(对应于表1)。(表2)是通过训练所有的DTI(表2)得到的。(DDI)样本。

无结构特征的KGE-UNIT (w/o SF)仅将从KG学习到的药物-药物对和药物-靶标对的多源特征作为预测器的输入。

无多源特征的KGE-UNIT (w/o HF)只使用从药物和蛋白质结构中学习到的结构特征作为预测器的输入。

无DTI任务的KGE-UNIT (w/o DTI)是专门针对DDI任务的单任务方法。

无DDI任务的KGE-UNIT (w/o DDI)是专门针对DTI任务的单任务方法。

消融实验结果如图2所示，图2A和图2B分别表示DTI和DDI的消融实验结果。显然，集成了所有模块的KGE-UNIT达到了最佳性能。多源特征和结构特征都有助于预测分子相互作用，但结构特征的影响更为显著。与DDI任务相比，在DTI任务中，使用多任务学习比使用单任务学习有更显著的性能提升。

图2 消融实验

作者还进行了案例分析。从训练好的KGE-UNIT模型中提取每个模块的学习表征，并使用t-SNE将其投影到二维空间中，如图3所示。可以观察到，尽管经过训练，KGE方法在区分分子对方面表现出有限的能力。相比之下，KGE-UNIT通过基于CNN的编码器和任务感知的注意解码器集成多源特征，有效区分交互和非交互对，进一步增强了KGE-UNIT的判别能力。在DTI任务中，（D）完整的KGE-UNIT优于（A）仅基于多源特征（B）仅基于结构特征（C）仅基于CNN编码器的模型。在DDI任务中，（H）完整的KGE-UNIT优于（E）仅基于多源特征（F）仅基于结构特征（G）仅基于CNN编码器的模型。

图3 案例分析

总结

在这项研究中，作者提出了一种结合知识图谱嵌入（KGE）和多任务学习优点的多类型分子相互作用预测新方法KGE-UNIT。通过利用多模态特征和使用有效挖掘任务交互信息并强调任务特定特征的预测器，该方法提高了药物分子相互作用的单个任务的性能。此外，KGE-UNIT的框架易于扩展，可同时解决更多的分子相互作用预测问题。实验不仅证明了研究结果的可靠性，而且强调了KGE-UNIT在识别现实世界药物相互作用方面的有效性。总之，与现有方法相比，KGE-UNIT提供了更好的性能和可扩展性。本研究为多任务学习的进一步发展和统一框架在分子相互作用预测领域的应用开辟了可能性。

然而，KGE-UNIT也存在一些局限性，需要在未来继续改进。首先，KGE-UNIT框架尚未考虑药物和蛋白质的3D结构的结合。理想情况下，基于三维结构的模型应该在包含测量的配体-受体亲和力和配体和蛋白质的共晶结构的高度可靠的数据集上进行训练。然而，由于这些数据集收集的成本高昂，并且需要对结构信息进行实验验证，这些数据集相对较小。因此，在处理有限的数据量时，如何充分发挥三维结构特征的优势是需要考虑的一个方向。其次，消融实验体现了KGE-UNIT框架内的分子结构特征的重要性。然而，值得注意的是，并非所有与药物相关的任务都能提供结构特征，尤其是对新发现的分子而言。结构特征的缺失会影响KGE-UNIT的性能，这种限制是该方法固有的。因此，如何整合更多的多模态特征和更多种类的关联来扩展异构数据以提高特征质量是需要考虑的一个方向。最后，在引入新的实体或关系时，需要重新训练模型是KGE的一个常见挑战，这消耗了大量的时间和计算资源。这一挑战也适用于KGE-UNIT模型。因此，如何通过预训练模型学习特征表示并减少训练过程是需要考虑的一个方向。

参考文献

[1] Zhang et al. KGE-UNIT: toward the unification of molecular interactions prediction based on knowledge graph and multi-task learning on drug discovery. Brief Bioinform. 2024