基于化学元素知识图的分子对比学习

编译|厉小明 审稿|贺小龙

本文介绍一篇来自浙江大学计算机科学系、杭州创新中心、杭州西湖生命科学与生物医学实验室等联合发表的文章。该文章构建了一个化学元素知识图（KG）来总结元素之间的微观联系，并提出了一个用于分子表征学习的知识增强对比学习（KCL）框架。KCL由三个模块组成。第一个模块是知识引导图增强，对原有的基于化学元素KG的分子图进行扩充。第二个模块是知识感知图表示，对原始分子图使用通用图编码器来提取分子的表示，并使用知识感知消息传递神经网络（Knowledge-aware Message Passing Neural Network, KMPNN）对增强分子图中的复杂信息进行编码。最后一个模块是一个对比目标，以最大化分子图的这两种视图之间的一致性。

1.引言

准确预测分子性质是化学和制药界中基本任务的核心。通过监督模型的传统方法面临着有限标记数据的挑战，因为对数据进行注释既昂贵又耗时，且难以推广到一般情况，实用性很差。解决这些问题的方法之一是设计前置任务来学习无标签的节点或图的表示。另一条线是遵循计算机视觉领域的对比学习框架，旨在通过图增强构造相似和不同的视图对，包括节点丢弃、边缘扰动、子图提取和属性掩蔽。由于参数少，预定义任务简单，本文采取对比学习。

然而，图形对比学习有其独特的挑战。首先，图的结构信息和语义在不同领域之间存在显著差异，这使得设计通用的图扩充方案变得困难。第二，现有的图对比学习框架主要关注图的结构，没有考虑基本领域知识。第三，它们将分子图中的原子建模为只有在存在边（即化学键）时才能相互作用的个体，而没有考虑原子之间的相关性。

为了克服这些挑战，可以通过结合领域知识，假设元素的属性可以影响分子的属性。首先基于元素周期表构建一个化学元素知识图（KG），化学元素 KG 描述了元素（图1中用绿色表示）与其基本化学属性（图1中用红色表示的周期性和金属度）之间的关系。然后对原始分子图进行扩充，如图1所示，这有助于建立具有共同属性但没有直接通过键连接的原子之间的联系。这样，扩充的分子图不仅包含了结构拓扑，还包含了元素的基本领域知识。

图1 化学元素KG在原子之间建立关联

在此基础上，本文提出了一个新的知识增强对比学习(KCL)框架，包括三个模块：知识引导图增强、知识感知图表示和对比目标，以改进分子表示。

2.研究方法

图2展示了KCL的工作过程。（1）知识引导图增强在化学元素KG的引导下，将任意给定的分子图G转化为增广分子图G’。（2）知识感知图表示旨在分别从G和G’中提取分子的表示。（3）对比目标旨在将分子的表示投射到使用对比损失的空间，并让训练编码器最大限度地提高了正样本对之间的一致性和困难负样本之间的差异。

图2 KCL的一个示例

我们从元素周期表中提取所有化学元素及其属性。每个元素包含超过15种属性，包括金属度、周期性、状态、重量、电负性、电子亲和性、熔点、沸点、电离性、半径、硬度、模量、密度、电导率、热量和丰度。

然后在KG中构造提取的三元组（Gas, isStateOf, Cl），表示元素与属性之间存在特定的关系。然而，由于每个元素都有一些不同的连续属性，KG很难建模它们的连接。为了克服这个困难，我们对连续属性进行组织化并将它们转换为离散标签。化学元素KG的统计如表1所示。

表1 化学元素KG的统计数据

3.实验结果

本文主要探讨以下5个问题：

Q1：与最先进的分子性质预测方法相比，KCL的表现如何？

Q2：模块1中的知识引导图增强是否比一般增强学习到更好的表示？

Q3：模块2中的知识特征初始化和图形编码器如何影响KCL？

Q4：模块3中的自我监督对比学习和困难样本策略有多有用？

Q5：我们如何从特定领域的角度解释KCL（KMPNN）？

KCL预训练任务在ZINC15中250k个分子数据上进行。实验采用MoleculeNet的8个基准数据集。

评估过程遵循两个步骤，首先对模型进行预训练，然后在两种协议下对下游任务的学习模型进行评估。

• 微调协议： 为了实现模型的全部潜力，给定KCL编码器输出的图形嵌入，使用额外的MLP来预测分子的性质。 微调编码器和MLP中的参数。
• 线性协议： 为了比较KCL和对比学习基线，修正预训练模型中的图形嵌入，并训练一个线性分类器。

3.1 性能比较（Q1和Q2）

在微调协议下的性能。测试所提出的KCL方法是否比SOTA（state-of-the-art）方法性能更好。表2显示了监督学习基线和预训练方法的完整结果，其中带下划线的单元格表示以前的SOTA，带粗体的单元格表示KCL取得的最佳结果。Tox21、SIDER和ClinTox都是多任务学习任务，共包括42项分类任务。KCL模型有两个版本，使用GCN编码器的原始分子图和使用KMPNN作为编码器的增强分子图。

表2 KCL在微调协议下的性能

表2提供了以下观察结果：（1）KCL在所有数据集上始终取得了最好的性能。所有数据集的总体相对改善率为7.1%（分类任务2.6%，回归任务20.4%）。这表明了KCL在分子特性预测任务中的有效性。（2）在只有642个标记分子的小型数据集FreeSolv中，KCL比SOTA基线提高了16.8%。这证实了KCL的强度，因为它可以显著帮助标签信息非常有限的任务。

在线性协议下的性能。研究模块1中的知识引导图增强是否有助于学习更好的分子表示。表3显示了不同增强和对比学习方法的比较结果。为了与先前的工作保持一致并使比较公平，使用线性协议来评估分类数据集的性能。

这两种版本的KCL都比其他的图增强方法（表3中的第一组）产生了更好的结果。这验证了知识引导的图增强不会违反分子中的生物语义，因此比其他增强方法工作得更好。此外，KCL比之前的最佳对比学习方法（第二组）提高了7.0%，这证实了通过结合基本的化学领域知识和捕获原子之间的微观联系，可以获得更好的分子图表示。

表3 在线性协议下的性能

3.2 消融研究（Q3和Q4）

如图3所示，在所有架构中，具有知识特征初始化和困难样本挖掘方案（黄色条）的KCL表现出最好的性能。具有随机初始化和随机负采样的模型（以“w/o ALL”表示）几乎总是表现最差。这说明知识特征初始化和困难样本挖掘策略对KCL都是必要的，因为前者捕获知识特征三元组的结构信息，而后者引导编码器生成更有区别的表示。

图3 KCL在微调协议不同设置下的性能

我们用其他体系结构代替了GCN，KMPNN来探索图形编码器的影响。表4的结果表明，在原始分子图上应用不同的编码器（例如GIN，GAT）对性能没有显著影响。另外，与以前的异构图神经网络（RGCN）和通用消息传递框架（MPNN）相比，KMPNN通过在增强分子图上为不同类型的知识提供异构的注意信息传递，具有更好的表达能力。

为了研究自我监督策略的作用，我们比较了在微调协议下，KCL在有对比学习和无对比学习的情况下的表现。表4报告了比较结果，与未使用对比学习的模型相比，自监督对比学习可提高绩效，分类平均提高8.5%，回归平均提高56.9%。这证实了对比学习可以通过缩小潜在空间中结构视图和知识视图之间的距离来学习更好的表征，并提高下游任务的预测性能。

表4 KCL在有对比学习和无对比学习的情况下的表现

3.3 化学解释性分析（Q5）

最后，我们通过可视化分子中每条边的注意来探索模型的可解释性。具体地说，我们从KCL（KMPNN）的最后一层中提取并归一化原子对相邻原子的注意权重。

图4 BBBP数据集中不同类型邻居（属性和原子）的注意力可视化示例

图4展示了BBBP数据集中的一个示例。BBBP记录了一种化合物是否具有穿透血脑屏障的渗透性。如图左边所示，原子倾向于把更多的注意力放在它们的电子亲合力、电负性、金属度和电离性上。这些特性与原子失去电子的能力密切相关。原子获得或失去电子的能力将在很大程度上影响分子的极性，因此影响其渗透性。此外，更活跃的原子邻居更容易被注意到，如图4右侧所示。元素Cl具有较高的电负性，因此具有较强的获得电子的能力。此外，羟基促进亲水性，因此被赋予更高的重视。另一个有趣的观察结果是细粒度属性（例如重量、半径）受到的关注少于粗粒度属性（例如电子亲和性、电负性、金属丰度和电离性）。这是因为粗粒度属性比细粒度属性更抽象、信息量更大，因此包含更丰富的领域知识。这与层次机器学习相一致，在层次较高的粗粒度特征可以被视为目标预测方面的细粒度特征的总结。

4.总结和未来工作

本文旨在将基本领域知识融入到分子图表示学习中。构建元素KG来建立元素之间的微观联系，并建议利用KCL框架中的知识来增强分子图对比学习。实验证明了KCL在微调和线性协议下的有效性，并且KCL优于以前的方法，具有更好的解释和表示能力。未来的工作为：首先，注入更多层次的领域知识以丰富化学元素知识。此外，探索元素KG是否也能促进化学和生物领域的其他类别的任务。第三，开源化学元素KG，继续提高其质量，扩大其规模。

参考资料

https://arxiv.org/abs/2112.00544