【知识图谱应用】Molecular Contrastive Learning with Chemical Element Knowledge Graph （论文笔记）

文章概述

• 关键词：化学分子表示学习、知识图谱构建、对比学习
• 内容：分子的表示学习对于化学领域下游任务例如药物设计、分子属性预测等非常关键，为了更好的表示分子，图对比学习由于不用人工标注和利用自监督信号而成为一个富有前景的范式；然而以前的工作忽略了将基本的领域知识包含进图的语义中，并且忽略了那些具有共同属性但在分子中并没有键链接的原子之间的关联性。因此，文章构建了一个化学元素知识图谱(Chemical Element Knowledge Graph)用于总结元素之间微小的关联，并提出一个新的知识增强对比学习框架（KCL）用于分子的表示学习。

研究背景

-因：预测分子属性作为化学和药物学中最基础的任务，在近年深度学习的发展下，许多监督学习模型已经通过预测分子属性来学习分子的表示。但是监督学习模型需要大量标记数据，由于实验数据以及标注工作的成本昂贵，以及分子的数量之多以及多样性，监督模型几乎不能泛化到一些未可见的案例，这也就抑制了模型在化学以及药物学领域的发展。
-发展：为了缓解上述问题，有一些工作尝试通过masked attribute prediction (Hu 2020)、graph-level motif prediction (Rong 2020) 、graph context prediction（Liu 2019）；另一部分工作根据来自CV领域的对比学习框架，致力于利用图增强（node dropping、edge perturbation、subgraph extraction）构建相似和不相似的view pairs。
-果：本文选择具有更少的参数，以及更简单的预定义任务的对比学习框架来研究分子的表示学习。
-困难：图对比学习与图像对比学习不同，图的语义和结构信息根据领域的不同差距很大，因此构建一个通用的增强scheme非常困难，尤其是化学分子，删除或者增加一个键会巨大改变他们的性质和属性（You 2020）；更重要的是图对比学习模型主要注重图的结构，并没有考虑图本身语义当中一些基础的领域知识；而且在建模原子时只考虑了有边链接的结构，而忽略了哪些没有键连接但是具有相同属性的原子之间的关联。
-工作的引出以及大概的研究方法：为了克服上述的困难，文章将领域知识包含进分子的图形表示当中。由于化学的领域知识非常关键，文章假设元素的属性能够影响分子的性质。1. 为了获取领域知识并建立原子之间的关联，文章首先基于Periodic Table of Elements构建了化学分子知识图谱；2. 然后根据知识图谱增强分子图的表示，帮助建立具有相同属性但是并不直接相连的原子之间的关联。这种方式下，增强后的分子图不仅仅包含结构拓扑信息也包含了基本元素领域知识；3. 综上所述，我们提出了一个新的知识增强的对比学习框架使用三个模块提升分子的表示。1)-knowledge-guided graph augmentation 利用知识图谱引导图增强保存图的拓扑结构并建立原子之间的关联；2)knowledge-aware graph representation使用通用的图编码器编码原始分子图，同时设计一个Knowledge-aware Message Passing Neural Network (KMPNN) 编码器用于增强之后的分子图提供异质注意的信息传播；3)-contrastive objective训练编码器最大化正例和非负例之间差异的一致性。

研究方法

Knowledge-guided Graph Augmentation

• 化学分子知识图谱构建：从the Periodic Table of Elements中爬取所有的化学元素以及他们的属性，每个元素最多包含15种属性（metallicity, periodicity, state, weight, electronegativity, electron affinity, melting point, boiling point, ionization, radius, hardness, modulus, density, conductivity, heat, and abundance）；其中在将元素和属性转换为三元组的过程中，对于一个元素的连续的属性可将其转换为两种离散的属性而便于三元组的产生。
• 图增强：大多数已存在的增强方法干扰了分子中的化学语义并且忽略了图语义中基本领域知识的影响，文章为了解决上述问题提出了一个基于化学元素知识图谱的知识引导图谱增强模块，元素与属性之间的连接方向从属性到元素，原子之间的连接是双向的；我们获取到一个增强的分子图，原始的分子结构被保留；原子相关属性的邻居拓朴结构被引进。保留拓扑结构的同时增强后的分子图也考虑了基本的存在于元素中的领域知识以及有共同属性但并无直接关联的原子之间的微小关联。这种被增强之后的分子图因此包含了更丰富且更复杂的信息，在之后的对比学习中被当作positive sample。
• Knowledge-aware 图表示：

  • Knowledge Feature Initialization：不同于原子和键的随机初始化，为了获取增强分子图属性与关系的初始特征，文章采用了常用的知识图谱嵌入方法RotateE, 训练训练元素知识图谱。通过这种方式，最初的特征便能够抓住三元组的结构信息。

  • KMPNN Encoder：增强后的分子图是复杂的不规则结构数据，包含两种信息类型（隐藏于分子键中的结构知识以及从化学元素知识图谱中提取的领域知识），因此文章设计了一个KMPNN编码器，利用$f^{\prime }(.)$学习图级别的表示。这个编码器背后的关键思想是为不同类型的邻居提供两种不同的消息传递方式，并根据他们的重要性分配不同的注意力。
    

  • GNN-based Encoder：对f(-)没有网络结构的约束。架构对$f(.)$的约束。我们选择了简单的方法，采用常用的GCN来获得$h_G=f(G)$，这是经过加权求和后的输出及最大池化读出。

• 对照性目标：
  • Projection Head：非线性的转换$g(.)$，映射头将原始的和增强后的表示映射到另一个用于计算对比损失的空间中。在KCL框架中，一个两层感知机用于获得$z=g(h_G)$, $z^{\prime }=g(h_{G^{\prime }})$。值得注意的是，预训练完成后，我们将投影头扔掉，只将编码器用于下游任务。
  • Negative Mining：文章采用另一种hard negative mining scheme, 将于anchor instance相似的分子当做负例。具体的，文章使用Morgan Fingerprints表示每个分子，这种方式注意到分子中每个原子周围圆形子结构的存在，并将其编码为固定长度的binary vector，然后通过Tanimoto coefficient计算分子相似性。为了确保每个分子都有对应的反例，这里利用相似性挑选了样本并且选择一系列最相似的原子作为负例。
  • Contrastive Loss：这里使用图引导的增强方式增强了$N$个相似的分子图，最终输出2N个图。给定一个正向对，我们将同一minibatch中剩余的另外$2(N-1)$个图作为hard negative samples。使用NT-Xent作为目标方法。

实验

实验步骤

• 预训练数据收集：文章从ZINC2015数据集中获取250K未标记的分子样例来预训练KCL。
• Fine-tuning 任务与数据集：这里使用来自MoleculeNet的8个数据集作为benchmark。其包含了各种分子任务（量子力学、物理化学、生物物理学和生理学）。在每个数据集上以三个随机播种的随机拆分或scaffold拆分上独立运行，训练/验证/测试的比例为8：1：1。
• Baselines:
  • Supervised learning methods：GCN和Weave是两种类型的图卷积方法。MPNN和它的变体DMPNN,CMPNN,CoMPT考虑了边特征以及通过信息传递增强键与原子之间的信息交互。
  • Pre-trained methods：N-GRAM预测节点属性来实现节点的嵌入。GROBER是包含节点级别和图级别pretext任务的预训练模型。
  • Graph contrastive learning baselines：InfoGraph最大化分子与图之间的信息. MICRO-Graph是基于motif的对比方法。GraphCL通过手选ad-hoc增强构建图数据对比实例。JOAO自动化增强选择。MoCL利用利用领域知识在两个层面上辅助表示学习。
• Evaluation Protocol：先预训练模型然后再下游任务中评估这个学到的模型
  • Fine-tune protocol: 为了实现模型的所有潜力，在利用KCL编码器给定图嵌入，我们使用一个额外的MLP来预测分子的属性，并fine-tune 编码器和MLP中的参数。
  • Linear Protocol：为了对比我们的模型和对比学习的baselines，我们固定来自预训练模型的图嵌入，并且训练了一个线性分类器。
• 实现细节：使用Adam优化器，初始学习率为0.0001，batch size 256。对于预训练模型，running batch 20。$T=0.1$。对于下游任务，在验证集上采用early stopping。 基于验证集采用随机搜索获取最好的超参数。