AAAI2022 | KCL: 化学元素知识图谱指导下的分子图对比学习

论文题目：Molecular Contrastive Learning with Chemical Element Knowledge Graph

本文作者：方尹、张强、杨海宏、庄祥、邓淑敏、张文、秦铭、陈卓、范骁辉、陈华钧（浙江大学）

接收会议：AAAI 2022

论文链接：https://arxiv.org/abs/2112.00544

数据集及代码：https://github.com/ZJU-Fangyin/KCL

欢迎转载，转载请注明出处

一、总述

近年来，生物医学领域开始将分子图表示学习应用于生物医药相关的多种下游任务，例如分子属性预测和药物设计。 图对比学习方法无需设计复杂的预训练任务，通过从大规模无监督数据中挖掘自身监督信息来学习图的表示。

与图像不同，分子图上的对比学习方法面临独特的挑战。首先，图的结构和语义信息在不同领域中差异显著，因此难以设计通用的图增强方案。尤其是对于分子图，化学键或官能团的添加或删除将在很大程度上改变分子的身份和特性。目前图对比学习的大多数方法主要关注图的拓扑结构，很少考虑图中蕴含的领域知识。

另一个易被忽略的问题在于，分子图中的原子通常被建模为只有彼此之间存在化学键时才会产生联系的个体，未能考虑原子之间隐含的关联性（例如，具有相同性质的原子之间的共性）。

化学元素知识图谱建立了未通过化学键相连但化学性质有所关联的原子之间的联系

为解决以上问题，我们通过领域知识来指导分子图对比学习。

首先，为了建立元素之间的微观联系及各元素的基本领域知识，我们基于化学元素周期表，构建了化学元素知识图谱 (Chemical Element Knowledge Graph)。 如上图所示，化学元素知识图谱描述了元素（图中绿色部分）之间的关系及其基本的化学性质（例如，周期性、金属性，图中红色部分）。

然后，我们通过化学元素知识图谱指导原始分子图的图增强过程，有助于在不相邻但具有相同性质的原子之间建立联系。通过这种方法，增强的分子图既包含拓扑结构知识，又具备元素基本的化学领域知识。

基于该化学元素知识图谱，本文创新性地提出一种知识增强的分子图对比学习框架 K nowledge-enhanced Contrastive Learning (KCL) 。 KCL 利用化学元素知识图谱指导原始分子图的增强过程，并针对分子增强图设计了知识感知的消息传递网络 KMPNN ，通过最大化正样本对之间的一致性和难负样本对之间的差异性构建对比损失以优化模型。实验结果表明， KCL 在涵盖不同分子属性的 8 个数据集上获得了 SOTA 性能。

二、方法

KCL框架图

具体来说，KCL框架分为三个模块。

（1）Knowledge-guided Graph Augmentation 

知识指导的图增强模块利用化学元素知识图谱指导原始分子图的增强过程，使分子增强图不仅包含拓扑结构知识，还包含元素的基本领域知识。

Chemical Element KG Construction: 我们从化学元素周期表中获取所有化学元素及其基本化学性质。每个元素具有15个以上的性质，包括金属性、周期性、状态、重量、电负性、电子亲和力、熔点、沸点、电离、半径、硬度、模量、密度、导电、热量和丰度。提取出的三元组以 (Gas, isStateOf, Cl)的形式存在于KG中，代表元素和性质之间存在指定的关系。

化学元素知识图谱的统计信息

Graph Augmentation: 对于原始分子图中的每一个原子，找出在化学元素知识图谱中将该原子作为尾实体的三元组。将这些三元组中的头实体作为新节点，关系作为头实体（性质）和尾实体（元素/原子）之间的边，获得分子增强图。分子增强图作为原始分子图的正样本，包含更丰富复杂的信息，能够捕捉原子之间的微观联系。

（2）Knowledge-aware Graph Representation 

知识感知的图表示模块针对分子增强图设计了知识感知的消息传递网络KMPNN，以更好的传递和融合增强图中两种不同类型的知识。

Knowledge Feature Initialization: 我们采用常用的KGE方法，RotateE，对分子增强图中性质和关系节点进行初始化。

KMPNN Encoder: KMPNN针对不同类型的邻居，提供了两种不同类型的消息传递，并根据邻居的重要性程度为其分配不同的注意力。通过KMPNN，可获得分子增强图的表示。算法1描述了KMPNN的编码过程：

知识感知的消息传递网络KMPNN的编码过程

GNN based Encoder: 对于原始分子图，采用GNN模型学习其表示。

（3）Contrastive Objective 

对比目标模块通过最大化正样本对之间的一致性和难负样本对之间的差异性构建对比损失，以优化表示模型。

Projection Head: 将原始分子图和分子增强图的表示映射到同一潜在特征空间，以便计算对比损失。

Negative Mining: 利用难负样本挖掘技术，选择在分子指纹空间中距离相近的分子图及其分子增强图作为负样本。

Contrastive Loss: 一个训练批次中某分子及其分子增强图所构成的正样本对的损失函数可表示为：

三、实验

（1）数据集 

为了验证KCL的效果，我们在MoleculeNet的8个基准数据集上进行评测，数据集具体信息如下所示：

数据集信息

（2）实验结果 

我们分别在fine-tune protocol和linear protocol两种设定下进行了实验。实验表明，KCL在两种设定下均比以往分子图表示学习方法效果更佳。

Fine-tune protocol效果

Linear protocol效果

四、总结与展望

本文旨在将基本化学领域知识融入分子图表示学习中。我们构建了化学元素知识图谱以建立元素之间的微观联系，并提出一种知识指导的分子图对比学习框架—KCL。实验证明了KCL在fine-tune protocol和linear protocol两种设定下的有效性，并表明KCL相对以前的方法具有更好的可解释性和表达能力。

我们将在以下几个方面扩展此工作。我们将引入不同粒度的领域知识以丰富化学元素知识图谱；我们将使用更深层次的知识表达手段，如OWL2，为化学元素知识图谱加入描述逻辑；我们将发布多语言版本的开放数据集，持续更新化学元素知识图谱。