KDD 2022 | KPGT：基于图Transformer的知识引导的预训练模型用于分子性质预测

本文介绍一篇清华大学交叉信息研究院发表在KDD 2022上的文章，题为KPGT: Knowledge-guided Pretraining of Graph Transformer。该文章提出了一个基于图Transformer的知识引导的预训练模型，KPGT利用了一个高通量的线性图Transformer（LiGhT），用于突出分子中化学键的重要性，在此基础上提出了一种知识引导的预训练策略，利用额外的分子知识来指导模型从大量未标记的分子图中捕获丰富的结构和语义信息。

1. 前言

分子性质预测在药物和材料领域对于发现具有所需特性的新型分子具有重要意义。将分子以分子图的形式表示，其中节点对应原子，边对应化学键。随着深度学习的迅速发展，图神经网络（GNN）已被广泛用于分子图表示学习。然而，有标签的数据量有限，但化学空间巨大，用监督学习策略训练的深度学习模型往往表现不佳，尤其是对分布外数据样本的预测。在自监督学习方法在自然语言处理（NLP）和计算机视觉（CV）领域取得巨大成功后，最近许多工作提出采用自监督学习策略，通过利用大规模未标记分子对GNN进行预训练，与监督学习方法相比，在下游分子性质预测任务中取得了优异的预测性能。尽管取得了不错的成果，但预测性能仍未达到理想情况。本文认为目前分子图的自监督学习方法主要遇到了以下两个问题：

1.1 预训练任务定义不明确

自监督学习方法的性能很大程度上取决于预训练任务的设计。到目前为止，分子图上的自监督学习方法可以大致分为两类，即生成式方法和对比式方法。生成式方法沿用了NLP领域的掩码语言模型，通过掩码分子图的边、节点或子图等，然后学习检索原始图。分子图的对比式方法通过节点的替换、删除和边扰动等策略进行增强，然后将增强后的图与嵌入空间中相应的原始分子图进行匹配学习。

然而，自然语言的词语掩码和图增强不会影响原始输入的基本语义，而分子图的微小改变会很大程度影响相应分子的性质（如图1）。因此，目前关于分子图的自监督学习方法只能捕捉到图的结构信息和分子的简单构造规则（如价键规则），但不能从化学结构中提取出与分子特性相关的丰富语义，而这些语义对下游的学习任务可能更为重要。

图1 分子中微小改变对分子性质会产生较大影响

1.2 模型容量有限

由于巨大的化学空间和广泛的分子性质，需要一个高容量的模型从未标记的分子中获取足够的信息。目前，在分子性质预测领域，基于Transformer的模型取得了很好的结果。以GIN为例，作为主干网络，只能提供有限的模型容量，可能无法捕获所需的各种分子特征。

为了解决上述两个问题，本文提出了一个新的自监督学习框架，即KPGT。

2. 方法

2.1 分子线性图

首先，将原始的分子图转化为分子线性图（如图2）。对于一个给定的分子图，可以表示为G = (V, E)，其中V表示节点（即原子），E表示边（即化学键）。原始分子图中如果有化学键将两个原子相连（如原子1、2），就在分子线性图中构建一个新的节点（如②），其中包含原子1、2和化学键。如果两个化学键有共同的节点（如原子2），就连接这两个化学键所对应的新节点（如①②之间的连线），以此类推即可获得分子线性图。

图2 将原始分子图转化为分子线性图

2.2 LiGhT

基于分子线性图，作者提出了线性图Transformer (LiGhT)，这是一种基于Transformer的新架构，用于编码分子线性图和学习分子性质预测的全局特征表示（如图3）。

图3 LiGhT 架构图

LiGhT就是在Transformer的基础上，引入了距离编码和路径编码，用于捕获分子的结构信息。其中关于Transformer架构部分不再赘述，此处介绍路径编码和距离编码。

（1）路径编码：对分子线性图中的每一对节点推导出最短路径，然后将这个最短路径特征编码为一个注意力标量（如公式1），其中表示两个节点之间的最短路径，SP为最短路径函数，和都是可训练的映射矩阵，Np表示最短路径的长度。

                                公式1

（2）距离编码：对分子线性图中的每一对节点推导出最短路径距离，也将其编码为一个注意力标量（如公式2），其中di,j 表示两个节点之间的最短路径距离，SPD为最短路径距离函数，和 都是可训练的映射矩阵。

                       公式2

之后利用注意力矩阵将编码的结构信息融入到模型中（如公式3）。

                    公式3

2.3 知识引导的预训练策略

之后，采用知识引导的预训练策略（如图4），在预训练框架中引入了分子的额外知识，作为掩码图中丢失的语义来指导掩码图的节点预测。之前的利用分子指纹或分子描述符的模型可以较好地对分子性质进行预测，因此说明分子指纹和分子描述符可以提供有价值的分子局部或全局信息。

为了将这些额外的知识纳入预训练框架，为每个分子线性图定义一个特殊节点，称为知识节点(K节点)，并在K节点和分子线性图之间建立连接。利用分子描述符和指纹描述符初始化K节点的原始特征。因此，在预训练过程中，额外知识可以通过Transformer的自注意模块被其它节点所关注，从而指导掩码节点的预测。

图4 知识引导的预训练策略

3. 结果与讨论

3.1 下游分子性质预测

作者在特征提取设置和迁移学习设置两种评价方案下进行了下游分子性质预测的测试（表1和表2）。在特征提取设置中，首先固定预训练模型，并将其作为特征提取器来获取数据样本的图表示，然后训练预测器进行预测；在迁移学习设置中，对模型的所有参数进行微调。

表1 特征提取和迁移学习设置下，不同方法在分类数据集上的AUROC性能。

表2 在特征提取和迁移学习设置下，不同方法在回归数据集上的RMSE性能。

从测试结果来看，在特征提取和迁移学习设置下，KPGT在所有分类和回归任务上都取得了优于其它baseline模型的结果。在特征提取设置中，所有任务的总体相对优化为5%，分类任务为3.5%，回归任务为8.9%；对于迁移学习设置，在所有任务上的总体相对改善为3.9%，在分类任务上为2.2%，在回归任务上为8.3%。这一显著的改进表明KPGT可以作一个用于准确预测分子性质的工具。

3.2 消融实验

评估LiGhT的表达能力：将KPGT中的LiGhT替换为其他主干网络，包括GIN、vanilla Transformer（没有路径编码和距离编码）和Graphormer。为了进行公平的比较，将所有模型限制为具有几乎相同数量的参数，其结果见表3。证明了在分子图的自监督学习上，LiGhT是一个比较可靠的框架。

表3 在迁移学习设置下，不同架构在分类和回归数据集上的比较结果，分别用AUROC和RMSE进行衡量

评估自监督学习策略：比较了经过预训练的LiGhT(KPGT)和未经过预训练的LiGhT的预测性能，两者都有相同的超参数设置，其结果见表4。经过预训练的模型取得了更优的预测性能，证明自监督学习策略的有效性。

表4 预训练的LiGhT (KPGT）和未预训练的LiGhT在分类和回归任务上的AUROC和RMSE的比较结果。

评估额外知识：为了验证简单地使用KPGT中相同的额外知识是否可以取得类似结果，引入了两种baseline模型，其中也加入了相同的额外知识。首先，直接连接分子描述符和分子指纹作为输入，然后通过一个2层的MLP进行预测(MD+FP)。GROVER框架中已经将分子描述符连接到图表征上，进一步将分子指纹连接到输出中，并通过一个2层的MLP进行最终的预测(grover+FP)（如表5）。实验结果表明本文提出的自监督学习策略可以潜在地学习分子语义信息之外的知识。

表5 在特征提取设置下KPGT和baseline模型之间的比较结果，后者也包含了相同的附加知识，在分类和回归数据集上，分别以AUROC和RMSE进行衡量。

评估不同掩码率：本文的自监督学习策略中，在分子线性图中掩盖了一部分节点和额外知识。之后评估了不同的掩码率的影响，分别对KPGT进行15%、30%、50%和60%掩码率的预训练，并报告了下游任务的预测结果（见表6）。实验表明，在掩码率为50%时，模型预测性能最好。

表6 在迁移学习设置下，不同掩码率的KPGT在分类和回归数据集上的消融结果，分别用AUROC和RMSE进行测量。

4. 总结

本文首先指出了当前分子图自监督学习方法的两个主要问题，即预训练任务定义不明确和模型容量有限。之后提出了一个新的自监督学习框架KPGT，它由两个主要部分组成，即LiGhT和知识引导的预训练策略。实验结果表明，KPGT在分子性质预测任务上比目前SOTA方法具有更高的准确性。

原文链接：

Li H, Zhao D, Zeng J. KPGT: Knowledge-Guided Pre-training of Graph Transformer for Molecular Property Prediction[J]. arXiv preprint arXiv:2206.03364, 2022.

(https://doi.org/10.48550/arXiv.2206.03364)

作者：孔祥泰

责任编辑：李叙潼，郑明月

--------- End ---------

感兴趣的读者，可以添加小邦微信（zhiyaobang2020）加入读者实名讨论微信群。添加时请主动注明姓名-企业-职位/岗位 或

姓名-学校-职务/研究方向。

- 历史文章推荐 -

【AI药物设计】