Nat. Mach. Intell2021 | MolCLR+：基于GNN的分子表征对比学习

原文标题：Molecular Contrastive Learning of Representations via Graph Neural Networks

代码：https://github.com/yuyangw/MolCLR

一、问题提出

首先，分子信息难以完全表示。例如，基于字符串的表示，如SMILES和SELFIES，不能直接对重要的拓扑信息进行编码。

其次，化学空间的大小是巨大的，例如，潜在的药理活性分子的大小估计在10^60量级。这给在潜在的化合物中推广任何分子表示都带来了很大的困难。

第三，分子学习任务的标记数据是昂贵的，远远不够，特别是当与潜在化学空间的大小相比。获得分子性质的标签通常需要复杂而耗时的实验室实验。

二、Methods

1、Graph Neural Networks

a molecule graph G is defined as G = (V,E), where V and E are nodes (atoms) and edges (chemical bonds),respectively

GNN第k层节点特征的聚合更新规则如式1所示:

h(k)v是节点v在第k层的特征， h(0)v由结点特征xv初始化

实际Readout为mean pooling

2、Contrastive Learning

其中zi和zj是从正数据对中提取的潜在向量，N是batchsize，sim(·)度量两个向量之间的相似性，τ是温度参数。余弦相似度：

整个pipeline(图1(a))由四个部分组成: data processing and augmentation、GNN-based feature extractor、non-linear projection head和normalized temperature-scaled cross-entropy (NT-Xent) contrastive loss。

过程：smiles Sn (batch_size=N)  → 分子图Gn → data processing and augmentation: 将Gn转换为两个不同但相关的分子图Gi、Gj，其中i = 2n−1和j = 2n。同一分子扩增的分子图记为正对，而由不同分子扩增的分子图记为负对 → GNN-based feature extractor：通过GNN提取特征，转换为隐藏state hi、hj → non-linear projection head：由一个隐层的MLP模型模拟，将表示hi和hj分别映射为潜在向量zi和zj → normalized temperature-scaled cross-entropy (NT-Xent) contrastive loss：将归一化温度交叉熵(NT-Xent)损失应用于2 N个潜在向量z，以最大化正对的一致性，同时最小化负对的一致性

对MolCLR预训练的GNN模型进行微调，用于分子性质预测，如图1(c)所示。与预训练模型类似，预测模型由一个GNN主干和一个MLP头组成，前者与预训练的特征提取器共享同一个模型，后者将特征映射到预测的分子性质。微调模型中的GNN骨干通过预训练模型中的参数共享初始化，而MLP头则是随机初始化。整个微调模型在目标分子性质数据库上以监督学习的方式进行训练。

3、Molecule Graph Augmentation

Atom Masking

分子图中的原子以给定的比例被随机遮蔽。当一个原子被屏蔽时，它的原子特征xv被masked token m所取代

Bond Deletion

键删除按一定比例随机删除原子之间的化学键，如图1(b)中黄色框所示。与用掩码标记替换原始特征的原子掩码不同，键删除是一种更严格的增强，因为它从分子图中完全删除了边缘。原子间化学键的形成和断开决定了化学反应中分子的属性。键的缺失模拟化学键的断裂，这促使模型了解一个分子在不同反应中的参与之间的相关性。

Subgraph Removal

子图去除可以被认为是原子掩蔽和键删除的结合。子图删除从随机选择的原子开始。去除过程是这样进行的:先遮蔽原子的邻原子，然后遮蔽邻原子的邻原子，直到被遮蔽原子的数量达到原子总数的给定比例。然后删除被屏蔽原子之间的键，这样被屏蔽原子和被删除的键形成原始分子图的诱导子图。该模型学会在剩余的子图中找到显著的基序，这些基序在很大程度上决定了分子的性质。

作者测试了四种增强策略。(1)结合原子掩蔽和键删除，两者比值p均设为25%。(2)去除随机比p在0% ~ 25%之间的子图。(3)去除固定比例为25%的子图。(4)三种增强方法的组成。

Training Details

实现一个以ReLU激活为GNN骨干的5图卷积层，使聚合与边缘特征兼容。在每个图上应用平均池作为读出操作，提取512维分子表示。具有一个隐藏层的MLP将表示映射到256维潜在空间。

权值衰减10−5、Adam优化器。在学习速率为5 × 10^(−4)的初始10个epoch之后，实现余弦学习衰减。batch_size为512，共50个epoch。

下游任务微调，在基本GNN特征提取器之上添加了一个随机初始化的MLP, 对分类和回归任务分别实现了交叉熵损失和L1损失。进行三次100 epoch的微调，以获得测试集上性能的平均值和标准差，只在训练集上训练模型，并在验证集上执行超参数搜索以获得最佳结果。

Pytorch Geometric：pytorch的一个图卷积神经网络的封装库。

三、实验

1、MoleculeNet classification

2、MoleculeNet Regression

与有监督的GCN和GIN相比，MolCLR预训练大大提高了所有基准测试的性能，这证明了MolCLR的有效性

在分类基准上，通过MolCLR, GCN的平均增益为12.4%，GIN为16.8%。同样，在回归分析中，GCN和GCN的平均改善率分别为27.6%和33.5%。一般来说，GIN通过MolCLR预训练表现出比GCN更大的改善。
这可能是因为GIN有更多的参数，能够学习更多具有代表性的分子特征。此外，MolCLR在大多数情况下比其他训练前/自我监督学习基线显示出更好的预测准确性。

3、消融实验

分子图增强可使平均ROC-AUC评分提高7.2%。将我们的分子图增强策略应用于有监督的分子性质预测任务，即使不进行预训练，也能大大提高性能。结果表明，分子图增强可以有效地帮助gnn学习鲁棒性和代表性特征。

4、Investigation of MolCLR Representation

可视化评估1：使用t-SNE嵌入检查预先训练的MolCLR学习的表示，将相似的分子表示映射到二维(2D)中的相邻点。通过t-SNE嵌入到2D的PubChem数据库验证集中的100K分子，根据分子量进行着色：

左下角显示的两个分子结构相似，其中一个halogen atom(氟或氯)与苯相连。这说明，即使没有标签，模型也能学习到分子之间的内在联系，因为具有相似性质的分子具有相近的特征。

可视化评估2：将MolCLR学习表征与传统分子指纹(FPs)进行了比较，例如ECFP和RDKFP。给定一个查询分子，通过MolCLR提取它的表示，并计算它与预训练数据库中所有参考分子的余弦距离：

然后根据表征距离（余弦距离）对所有参考分子进行排序，根据排序百分比将所有参考分子均匀分成20个集合bin。百分比阈值越低，预期查询中相似的分子就越多，因为MolCLR表示更接近。

4(a)展示每个bin内FP相似性的平均值和标准差。使用ECFP和RDKFP的相似度分布如图4(b)所示。ECFP往往比RDKFP获得更低的相似性，因为前者涵盖了与分子活性相关的更广泛的特征。但是，随着MolCLR表示距离的增加，ECFP和RDKFP的相似度都降低了。前5%的平均RDKFP相似度为~ 0.9，后5%降至~ 0.67。同样，平均ECFP相似性从前5%的~ 0.49下降到最低5%的~ 0.21

在MolCLR表示域中与查询分子最接近的9个分子如图4(c)所示，并标记了两种FPs相似性。这些分子具有较高的RDKFP相似性，从0.833到0.985，这进一步证明了MolCLR学习化学有意义的表示。