MolCLR:一个用于分子表征学习的自监督框架
作者：
Yuyang Wang1,2, Jianren Wang3 , Zhonglin Cao1 , and Amir Barati Farimani1,2,4,*
1Department of Mechanical Engineering, Carnegie Mellon University, Pittsburgh, PA 15213, USA
2Machine Learning Department, Carnegie Mellon University, Pittsburgh, PA 15213, USA
3Robotics Institute, Carnegie Mellon University, Pittsburgh, PA 15213, USA
4Department of Chemical Engineering, Carnegie Mellon University, Pittsburgh, PA 15213, USA
*corresponding author: Amir Barati Farimani (barati@cmu.edu)
摘要 ABSTRACT
分子机器学习(ML)在分子性质预测和药物发现方面具有重要的应用前景。然而，获取标记分子数据是昂贵和耗时的。由于标注数据有限，监督学习的ML模型要推广到巨大的化学空间是一个巨大的挑战。在这项工作中，我们提出了MolCLR:通过图神经网络(GNNs)对 表征 进行分子对比学习(Molecular contrast Learning of representation)，这是一个利用大量未标记数据(10M唯一分子)的自我监督学习框架自监督学习 | Self-supervised Learning。在MolCLR预训练中，我们建立分子图和开发GNN编码器来学习 可区分的表征。提出了三种分子图增广方法:原子掩蔽、键删除和子图删除。所提的分子图增强方法保证了增强时同一分子的一致性最大化以及不同分子一致性的最小化。实验表明，我们的对比学习框架显著提高了gnn在不同分子性质基准上的性能，包括分类和回归任务。得益于对大型未标记数据库的预训练，经过微调后，MolCLR甚至在几个具有挑战性的基准测试中达到了最先进的水平。此外，进一步的研究表明，MolCLR学会了将分子嵌入表征中，以区分化学上合理的分子相似性。
引言 Introduction
在功能化合物和新型化合物的设计中，分子表征是基本和必不可少的。由于可能的稳定化合物数量巨大，开发一个信息丰富的表征模型来概括整个化学空间是一个巨大挑战。传统的分子表征，如扩展连接性指纹(ECFP)已经成为计算化学的标准工具。最近﹐随着机器学习方法的发展﹐数据驱动的分子表征学习及其应用越来越被重视（如化学性质预测、化学建模、和分子设计）。
然而，由于以下三个主要的挑战﹐导致学习这样的表征很困难：
-
很难彻底表达分子信息。
例如基于字符串的表示法，如SMILES和SELFIES，不能直接编码重要的拓扑信息，为了保留丰富的结构信息，最近的许多工作利用图神经网络(GNNs)。
-
化学空间的规模是巨大的。
例如，潜在的药理活性分子的大小估计在数量级。
这给任何分子表征在潜在的化合物中进行概括带来了很大的困难。
-
分子学习任务的标签数据很昂贵，且远远不够﹐特别是与潜在的化学空间的大小相比，获得分子属性的标签通常需要昂贵的成本。
大多数分子学习基准中的标签数量是远远不够的，在这样有限的数据上训练的机器学习模型很容易过拟合。
随着机器学习特别是深度神经网络(DNNs)的发展和成功，分子表示学习在过去的十年中迅速发展。在传统的化学信息学中，分子以独特的指纹(FP)载体表示，如ECFP。考虑到FPs, DNNs是用来预测某些属性的。除了FP外，基于字符串的表示(例如SMILES)也被广泛应用于分子学习。建立在RNNs基础上的语言模型直接适合smiles32,33的学习表示。随着最近基于transform的架构的成功，这样的语言模型也被用于学习来自SMILES的表示。MPNN和D-MPNN实现了从分子图中聚合信息的消息传递体系结构。此外，SchNet模拟了GNN中分子之间的量子相互作用。DimeNet根据原子间的角度进行信息转换，实现方向信息的集成。
得益于可用分子数据的增长，自监督/预训练分子表征学习也被研究。像BERT这样的自我监督语言模型已经被用于学习以SMILES作为输入的分子表示。在分子图上，N-Gram Graph通过短步组装顶点嵌入来构建图的表示，无需训练。Hu等人提出了节点级和图级的GNN预训练任务。然而，图级的预训练是基于监督学习任务，这是受有限标签的约束。等人将对比学习扩展到非结构化图数据，但该框架并不是专门为分子图学习设计的，只是在有限的分子数据上进行训练。
本文提出的MolCLR通过图神经网络对表征进行分子对比学习，以解决上述挑战。MolCLR是一个自监督的学习框架，它是在具有1千万个独特分子的大型无标记分子数据集上训练出来的。通过对比损失，MolCLR通过对比积极的分子图对和消极的分子图对来学习表示。引入了三种分子图的增强策略：原子屏蔽、键删除和子图删除。广泛使用的GNN模型、图卷积网络（GCN）和图同构网络（GIN），被开发为MolCLR中的GNN编码器，以从分子图中提取信息表示。然后在MoleculeNet的下游分子特性预测基准上对预训练的模型进行微调。与通过监督学习训练的GCN和GIN相比，本文所提的MolCLR明显提高了分类和回归任务的性能。受益于对大型数据库的预训练，MolCLR在多个分子基准测试中超过了其他自监督学习和预训练策略。此外，在一些任务上，MolCLR可以与监督学习的基准线相媲美，甚至超过了监督学习，其中包括复杂的分子图卷积操作或特定领域的特征化。MolCLR表示法和传统的FPs之间的进一步比较表明，MolCLR可以通过对大量未标记数据的预训练来学习区分分子的相似性。
Figure 1.MolCLR的概述。
(a) MolCLR预训练:将小批量N个分子数据中的SMILES
S
n
S_n
Sn?转换为分子图
G
n
G_n
Gn?。对每个图应用两个随机分子图数据增广操作，得到两个相关的掩蔽图:
G
2
n
?
1
G_{2n-1}
G2n?1?和
G
2
n
G_{2n}
G2n?。一个基本特征编码器建立在图形卷积和读出操作的基础上，提取分子特征
h
2
n
?
1
h_{2n-1}
h2n?1?,
h
2
n
h_{2n}
h2n?。利用对比损失来最大化MLP project head的潜在向量
z
2
n
?
1
z_{2n-1}
z2n?1?和
z
2
n
z_{2n}
z2n?之间的一致性。
(b)分子图增强策略:原子掩蔽、键删除、子图删除。
?整个MolCLR框架:首先通过MolCLR对GNNs进行预训练，学习代表特征。对下游分子性质预测的微调共享GNN编码器的预训练参数，并随机初始化一个MLP头。然后遵循监督学习来训练模型。
综上所述，(1)我们提出了一个用于分子表示学习的自监督学习框架MolCLR。(2)引入3种分子图增强策略生成对比对，即原子掩蔽、键删除和子图删除。(3)得益于对大量未标记数据的预训练，通过MolCLR训练的简单GNN模型与监督学习相比，在所有分子基准上都有显著的改进。(4)与无法利用未标记数据的更复杂的GNN模型相比，MolCLR甚至在几个分子基准上通过微调将简单的GNN模型提升到最先进的(SOTA)水平
Results
MolCLR Framework
我们的MolCLR模型是建立在对比学习框架之上的。正增广分子图对的潜在表示与负对的表示进行了对比。整个pipline(图1(a))由四个部分组成:数据处理和增强、基于GNN的特征提取器、非线性project head和归一化temperature-scaled交叉熵(NT-Xent)对比损失。
给定大小为N的mini-batch SMILES数据 s n s_n sn?，构建相应的分子图 G n G_n Gn?，其中每个节点代表一个原子，每条边代表原子间的化学键。利用分子图增广策略，将 G n G_n Gn?转化为两个不同但相关的分子图: G i G_i Gi?和 G j G_j Gj?，其中 i = 2 n ? 1 , j = 2 n i = 2n-1, j = 2n i=2n?1,j=2n。由同一分子扩充而来的分子图表示为正对，而由不同分子扩充而来的分子图表示为负对。特征提取器 f ( ? ) f(·) f(?)由GNNs建模，并将分子图映射到表示 h i , h j ε R d h_i,h_j \varepsilon\mathbb{R}^d hi?,hj?εRd中。在我们的例子中，我们用平均池作为特征提取器来实现GCN和GIN。非线性project head g ( ? ) g(·) g(?)采用带隐层的MLP模型，将表示 h i h_i hi?和 h j h_j hj?分别映射为隐向量 z i z_i zi?和 z j z_j zj?。将归一化temperature 交叉熵(NT-Xent) loss应用于2N个潜向量z s，使正对的一致性最大化，负对的一致性最小化。框架是在来自PubChem40的10M未标记数据上进行预训练的。
经过MolCLR预训练的GNN模型被微调用于分子性质预测，如图1?所示。与训练前模型类似，预测模型由一个GNN骨干和一个MLP head组成，前者与训练前的特征提取器使用相同的模型，后者将特征映射到预测的分子性质中。通过预训练模型的参数共享对微调模型中的GNN骨干进行初始化，而对MLP头进行随机初始化。然后在目标分子属性数据库上以监督学习的方式训练整个微调模型。更多细节可以在方法一节中找到。
Molecule Graph Augmentation |分子图增强
We employ three molecule graph data augmentation strategies (Figure 1(b)) for input molecules in MolCLR: atom masking,bond deletion, and subgraph removal.
我们对输入分子在MolCLR中采用了三种分子图数据增强策略(图1(b)):原子掩蔽、键删除和子图删除。
原子屏蔽(Atom Masking)分子图中的原子以给定的比例被随机屏蔽。当一个原子被遮蔽时，它的原子特征被一个掩码标记 取代，该标记与图1 b中红框所示的分子图中的任何原子特征相区别。通过屏蔽，模型被迫学习内在的化学信息（比如由某些共价连接的原子的可能类型。
**键删除(Bond Deletion)**如图1(b)中的黄色方框所示, "键删除"以一定比例随机删除原子之间的化学键。与原子屏蔽方法中用掩码标记替代原始特征的原子掩码不同，键删除是一种更严格的增强，因为它从分子图中完全删除了边缘。原子之间化学键的形成和断裂决定了化学反应中分子的属性。键的删除模拟了化学键的断裂，促使模型学习一个分子在各种反应中的关联性。
**子图删除(Subgraph Removal)**子图删除可以被认为是原子屏蔽和键删除的结合。子图的去除从一个随机挑选的原点开始。去除过程通过屏蔽原原子的邻域，然后是邻域的邻域，直到被屏蔽的原子数量达到原子总数的一定比例。然后，被屏蔽的原子之间的化学键被删除，这样被屏蔽的原子和被删除的化学键就形成了原分子图的子图。如图1(b)中的蓝色方框所示，被移除的子图包括被遮蔽原子之间的所有化学键。通过匹配被移除的不同子结构的分子图，该模型学会了在剩余的子图中找到显著特征，这在很大程度上决定了分子的特性。
Molecular Property Predictions | 分子性质的预测
Table 1. Test performance of different models on seven classifification benchmarks. The fifirst seven models are supervised learning methods and the last four are self-supervised/pre-training methods. Mean and standard deviation of test ROC-AUC (%) on each benchmark are reported.
*Best performing supervised and self-supervised/pre-training methods for each benchmark are marked as bold.
表1。在七个分类基准上测试不同模型的性能。前七个模型是有监督的学习方法，后四个是自我监督/训练前的方法。报告每个基准测试ROC-AUC(%)的平均值和标准偏差。
**各基准表现最佳的监督及自我监督/训练前方法以粗体标记。
Table 2. Test performance of different models on six regression benchmarks. The fifirst seven models are supervised learning methods and the last four are self-supervised/pre-training methods. Mean and standard deviation of test RMSE (for FreeSolv,ESOL, Lipo) or MAE (for QM7, QM8, QM9) are reported.
*Best performing supervised and self-supervised/pre-training methods for each benchmark are marked as bold.
表2。在六个回归基准上测试不同模型的性能。前七个模型是有监督的学习方法，后四个是自我监督/训练前的方法。报告了测试RMSE(对于FreeSolv, ESOL, Lipo)或MAE(对于QM7, QM8, QM9)的平均值和标准差。
- *各基准表现最佳的监督及自我监督/训练前方法以粗体标记。
为了证明MolCLR的有效性，我们对来自MoleculeNet的多个具有挑战性的分类和回归任务的性能进行了基准测试。分子数据集的详细信息见补充表1和补充表2。表1显示了与监督和自我监督/训练前基线模型相比，我们的MolCLR模型在分类任务上的测试ROC-AUC(%)。报告了三个个体运行的平均和标准偏差。MolCLRGCN和MolCLRGIN分别表示以GCN和GIN作为特征提取器的MolCLR预训练。表1的观察结果如下:
- 与其他自监督学习或预训练策略相比，我们的MolCLR框架在7个基准测试中有5个达到了最佳性能，平均提高了4.0%。这样的改进说明了我们的MolCLR是一个强大的自我监督学习策略，它很容易实现，并且对特定领域的复杂性要求很少。
- 与表现最好的监督学习基线相比，MolCLR也表现出了竞争性能。在一些基准测试中(例如，ClinTox, BACE, MUV)，我们的训练前模型甚至超过了SOTA监督学习方法，其中包括复杂的聚合操作或特定领域的特征。例如，在ClinTox上，与有监督的D-MPNN相比，MolCLR使ROC-AUC提高了2.7%。
- 值得注意的是，MolCLR在有限数量的分子数据集上表现得非常好，如ClinTox、BACE和SIDER。该性能验证了MolCLR学习信息表示，可以在不同的数据集之间传输。
表2包含了MolCLR和基线模型在回归基准上的测试性能。FreeSolv、ESOL和Lipo采用均方根误差(RMSE)作为评价指标，而QM7、QM8和QM9采用均方根绝对误差(MAE)进行测量，均遵循MoleculeNet的建议。与分类相比，回归任务更具挑战性，因为后者只考虑手工定义的离散标签。表2的观察结果如下。
- MolCLR在6个基准中有5个超过了其他训练前基准，在其余的ESOL基准上取得了几乎相同的性能。与Hu等人45同样实现GIN作为编码器相比，MolCLRGIN在所有6个回归数据库上的性能都优于Hu等人45。以QM7和QM9为例，对Hu等的改进率分别为20.9%和45.8%。
- 与监督学习模型相比，MolCLR在大多数情况下达到竞争性能。例如，在Lipo数据库上，MolCLR得到了与性能最好的监督D-MPNN相似的结果。此外，GCN和GIN通过对所有回归基准进行MolCLR预训练获得了更好的预测性能。虽然，在QM9中，MolCLR不能与监督的SchNet19和MGCN52竞争。因为这两个模型是专门为量子交互作用设计的，并且利用了额外的3D位置信息。值得注意的是，尽管SchNet和MGCN在涉及量子力学特性的数据集(即QM7、QM8和QM9)上表现出了优越的性能，但在其他基准上，它们并没有显示出优于其他监督学习基线的优势。此外，在具有挑战性的QM9基准测试中，MolCLR预训练仍然被证明是有效的。与未经过训练的GCN和GIN相比，MolCLR仍能显著提高性能，分别提高38.7%和50.3%。同时，在QM9上，MolCLR的学习效果优于其他自监督学习基线，验证了MolCLR的有效性。由于QM9的性质具有不同的单位和量级，QM9的详细结果报告在补充表3中。
从表1和表2可以看出，与监督的GCN和GIN相比，MolCLR预训练在所有基准上的性能都有很大的提高，这说明了MolCLR的有效性。在分类基准上，通过MolCLR的平均增益，GCN为12.4%，GIN为16.8%。同样，在回归分析中，GCN和GCN的平均改善率分别为27.6%和33.5%。总的来说，通过MolCLR前训练，GIN比GCN表现出更多的进步。这可能是因为GIN有更多的参数，能够学习更有代表性的分子特征。此外，在大多数情况下，与其他训练前/自我监督学习基线相比，MolCLR显示出更好的预测精度。需要强调的是，MolCLR受益于对大型未标记数据库的预训练，而其他监督/自我监督学习基线则不然。利用未标记数据为MolCLR在化学空间和各种分子性质之间的泛化提供了比其他基线更大的优势。补充表4和补充图1进一步研究了训练前数据库对MolCLR的影响。这种泛化能力为预测药物发现和设计中潜在的分子性质提供了希望。
Optimal Molecule Graph Augmentations | 最优分子图增强
为了系统地分析分子图增强策略的效果，我们比较了原子掩蔽、键删除和子图删除的不同组成。图2(a)所示为每种数据增强策略在不同基准上的ROC-AUC(%)平均值和标准差。考虑了四种增强策略。(1)原子掩蔽和键缺失的整合，两者的比率p均设为25%。(2)随机比率p为0% ~ 25%的子图删除。(3)以固定的25%比例去除子图。(4)三种增广方法的组合。具体来说，首先应用随机比率为0%到25%的子图删除。然后，如果掩模原子的比例小于25%，我们继续随机掩模原子，直到达到25%的比例。同样，如果键删除率小于25%，则删除更多的键，以达到设定的比率。
如图2(a)所示，在四种组合中，25%的子图删除率平均达到最好的性能。与策略(1)相比，去除子图是原子掩蔽和键删除的内在结合，而且去除子图进一步解锁了局部子结构，因此，该算法的优异性能可归因于此。在BBBP数据集中，固定25%比例的子图删除效果较差，这是因为BBBP中的分子结构比较敏感，拓扑结构的微小变化会导致性质的巨大差异。此外，值得注意的是，在大多数基准测试中，与单子图删除增强相比，这三种增强(策略(4))的组合会损害ROC-AUC。一个可能的原因是，这三种增强策略的组合可以去除分子图中广泛的子结构，从而消除重要的拓扑信息。一般来说，子图删除在大多数基准测试中都能获得更好的性能。然而，它也表明，最优的分子图扩充是任务独立的。
Figure 2. Investigation of molecule graph augmentations on classifification benchmarks. (a) Test performance of MolCLR models with different compositions of molecule graph augmentation strategies. (b) Test performance of GIN models trained via supervised learning with and without molecular graph augmentations. The height of each bar represents the mean ROC-AUC(%) on the benchmark, and the length of each error bar represents the standard deviation.
图2.基于分类基准的分子图增强的研究。
(a)不同分子图增强策略组成的MolCLR模型的测试性能。
(b)测试通过有和没有分子图增强的监督学习训练的GIN模型的性能。每个误差条的高度代表基准上的ROC-AUC(%)的平均值，每个误差条的长度代表标准差。
Molecule Graph Augmentation on Supervised Learning | 有监督学习的分子图增强
我们工作中的分子图扩展策略，即原子掩蔽、键删除和子图删除，可以作为一个通用的数据增强插件来实现，适用于任何基于图的分子学习方法。为了验证分子图增强对监督分子任务的有效性，我们从随机初始化开始训练有增强和没有增强的GIN模型。具体来说，实现了固定比例为25%的子图掩蔽作为增强。图2(b)记录了测试ROC-AUC(%)相对于七个分子性质分类基准的平均值和标准差。在所有七个基准测试中，使用增强训练的GINs都超过了没有增强训练的模型。分子图增强可以提高平均ROC-AUC分数7.2%。将我们的分子图增强策略应用于有监督的分子性质预测任务，即使没有预先训练，也能大大提高性能。结果表明，分子图增广可以有效地帮助gnn学习鲁棒性和代表性特征。例如，子图删除匹配部分观察到的分子图。因此，模型学习在剩余的子图中寻找显着的基序，这对分子性质的学习非常有利。
Investigation of MolCLR Representation | MolCLR表征的研究
我们使用t-SNE embedding检查由预训练的MolCLR学习的表示。t-SNE算法将相似的分子表示映射到二维(2D)中的相邻点。图3所示是PubChem数据库验证集中的100K分子，通过t-SNE嵌入到2D中，根据分子量着色。我们还在图中加入了一些随机选择的分子，来说明通过MolCLR预训练学习到的相似/不相似分子有哪些。如图3所示，对于具有相似拓扑结构和官能团的分子，MolCLR学习了接近的表示。例如，上面显示的三个分子具有与芳基相连的羰基。左下角显示的两个分子结构相似，其中一个卤素原子(氟或氯)与苯相连。这说明，即使没有标签，该模型也能学习分子之间的内在联系，因为具有相似性质的分子具有相近的特征。更多的MolCLR表示的可视化可以在补充图2中找到。
- Embedding就是从原始数据提取出来的Feature，也就是那个通过神经网络映射之后的低维向量。
为了进一步评估MolCLR，我们将MolCLR学习表征与传统的分子指纹(FPs)进行了比较，如ECFP5和RDKFP。特别地，给定一个查询分子，我们通过MolCLR提取其表示，并计算其与我们的训练前数据库中所有参考分子的余弦距离。两个表示 ( u , v ) (u, v) (u,v)之间的余弦距离定义为 1 ? u ? v ∥ u ∥ ∥ v ∥ 1-\frac{u \cdot v}{\|u\|\|v\|} 1?∥u∥∥v∥u?v?。然后根据表示距离对所有参考分子进行排序，并根据排序百分比均匀地分为20个桶。百分比阈值越低，查询中期望的相似分子越多，因为MolCLR表示越接近。在每个bin中，随机选择5000个分子，并计算它们与查询的dice FP相似性。图4显示了一个查询分子的例子(PubChem ID 42953211)。如图4(a)所示为每个箱内FP相似度的均值和标准差。使用ECFP和RDKFP的相似度分布如图4(b)所示。ECFP比RDKFP具有更低的相似性，因为前者涵盖了更广泛的与分子活性相关的特征。结果表明，随着MolCLR表示距离的增加，ECFP和RDKFP的相似度都降低了。在前5%中，RDKFP的平均相似度为0.9，而在后的5%中，RDKFP的平均相似度为0.67。同样，平均ECFP相似性从前5%的0.49下降到后5%的0.21。虽然随着百分比阈值的增加会有波动，但总体趋势在MolCLR学习表征和化学指纹(FPs)之间是一致的。即MolCLR生成的表示 间的距离有效地反映了分子相似性。此外，图4?给出了在MolCLR表示域中与查询分子最接近的9个分子，两个帧数相似度都标注了。这些分子具有较高的RDKFP相似性，从0.833到0.985，进一步证明了MolCLR学习化学意义表征。观察到，这些选择的分子共享相同的官能团，包括烷基卤化物(氯)，叔胺，酮，芳香化合物。噻吩结构也可以在所有分子中找到。值得注意的是，图4?中第一行的第二个分子与查询分子完全相同，除了氯的位置，因此相似性最高。结果表明，通过对大量的未标记数据进行对比学习，MolCLR能够自动将分子嵌入到具有代表性的特征中，并以化学上合理的方式对化合物进行区分。在补充图3中可以找到更多查询分子的例子。
Figure 3. Visualization of molecular representations learned by MolCLR via t-SNE. Representations are extracted from the validation set of the pre-training dataset, which contains 100k unique molecules. Each point is colored by its corresponding molecular weight. Some molecules close in the representation domain are also shown.
图3。MolCLR通过t-SNE学习的分子表征的可视化。表示 从训练前数据集的验证集中提取，该数据集包含100k个唯一的分子。每个点都用相应的分子量着色。一些分子在表示域密切的也显示。
Figure 4. Comparison of MolCLR-learned representations and conventional FPs using the query molecule (PubChem ID42953211). (a) Change of ECFP and RDKFP similarities with respect to the distance between MolCLR representations. (b)Distribution of ECFP and EDKFP similarities with the query molecule. ? The query molecule and 9 closest molecules in MolCLR representation domain with RDKFP and ECFP similarities labeled.
图4。使用查询分子(PubChem ID 42953211)比较molclr学习表示和传统FPs。
(a) ECFP和RDKFP相似性相对于MolCLR表示间距离的变化。
(b)查询分子的ECFP和EDKFP相似性分布。
?在MolCLR表示域中标注RDKFP和ECFP相似性的查询分子和9个最接近的分子。
Conclusion
在这项工作中，我们研究了分子表示的自我监督学习。具体来说，我们提出了通过gnn和三种分子图增强策略(原子掩蔽、键删除和子图删除)来进行表示的分子对比学习。通过对比来自增广的正对和负对，MolCLR使用一般GNN骨干学习信息表示。实验表明，经过MolCLR预训练的GNN模型在各种分子基准上都取得了很大的改进，与以监督学习方式训练的模型相比，具有更好的泛化性能。
通过MolCLR学习的分子表示法证明了在有限的数据下可转移到分子任务，以及在大的化学空间上的泛化能力。在未来的工作中，有许多值得研究的方向。例如，改进GNN骨干(例如基于transform的GNN架构)可以帮助提取更好的分子表示。此外，对自我监督学习表征的形象化和解释也很有趣。这样的研究可以帮助研究人员更好地了解化合物，有利于药物的发现。
Methods
Graph Neural Networks | 图神经网络
在我们的工作中，一个分子图 G G G被定义为 G = ( V , E ) G = (V,E) G=(V,E)，其中 V V V和 E E E分别是节点(原子)和边(化学键)。现代图神经网络(GNNS)利用邻域聚合操作，迭代更新节点表示。GNN第k层节点特征的聚合更新规则如式
a v ( k ) = AGGREGATE ? ( k ) ( { h u ( k ? 1 ) : u ∈ N ( v ) } ) , h v ( k ) = COMBINE ? ( k ) ( h v ( k ? 1 ) , a v ( k ) ) \boldsymbol{a}_{v}^{(k)}=\operatorname{AGGREGATE}^{(k)}\left(\left\{\boldsymbol{h}_{u}^{(k-1)}: u \in \mathscr{N}(v)\right\}\right), \boldsymbol{h}_{v}^{(k)}=\operatorname{COMBINE}^{(k)}\left(\boldsymbol{h}_{v}^{(k-1)}, \boldsymbol{a}_{v}^{(k)}\right) av(k)?=AGGREGATE(k)({hu(k?1)?:u∈N(v)}),hv(k)?=COMBINE(k)(hv(k?1)?,av(k)?)
其中 h v ( k ) h_v^{(k)} hv(k)?为节点v在第k层的特征， h v ( 0 ) h_v^{(0)} hv(0)? 由节点特征 x v x_v xv?初始化。 N ( v ) N (v) N(v)为节点v的所有邻居的集合。为了进一步提取图级特征 h G h_G hG?，读出操作将图G中的所有节点特征进行整合，如式2所示
h G = READOUT ? ( { h u ( k ) : v ∈ G } ) \boldsymbol{h}_{G}=\operatorname{READOUT}\left(\left\{\boldsymbol{h}_{u}^{(k)}: v \in G\right\}\right) hG?=READOUT({hu(k)?:v∈G})
在我们的工作中，我们基于GCN和GIN构建了GNN编码器。在进行线性变换之前，GCN通过在节点本身及其邻接节点上引入一个均值池来集成聚合和组合操作。而GIN利用了MLP和节点特征加权求和的聚合方法。两者都是简单而通用的图卷积运算。此外，我们实现了广泛使用的平均值池作为读数。
Contrastive Learning | 对比学习
对比学习旨在通过对比正数据对和负数据对来学习表征。SimCLR48证明了图像的对比学习可以从数据增强和大批量的组合中获得极大的好处。SimCLR基于InfoNCE loss，提出NT-Xent loss如式3所示:
L i , j = ? log ? exp ? ( sim ? ( z i , z j ) / τ ) ∑ k = 1 2 N 1 { k ≠ i } exp ? ( sim ? ( z i , z k ) / τ ) \mathscr{L}_{i, j}=-\log \frac{\exp \left(\operatorname{sim}\left(\boldsymbol{z}_{i}, \boldsymbol{z}_{j}\right) / \tau\right)}{\sum_{k=1}^{2 N} \mathbb{1}\{k \neq i\} \exp \left(\operatorname{sim}\left(\boldsymbol{z}_{i}, \boldsymbol{z}_{k}\right) / \tau\right)} Li,j?=?log∑k=12N?1{k?=i}exp(sim(zi?,zk?)/τ)exp(sim(zi?,zj?)/τ)?
其中 z i z_i zi?和 z j z_j zj?为正数据对提取的潜向量，N为批量大小， s i m ( ? ) sim(·) sim(?)度量两个向量之间的相似性， τ τ τ为温度参数。在我们的MolCLR中，我们遵循NT-Xent损失对GNN编码器进行预训练，并将余弦相似度实现为 sim ? ( z i , z j ) = z i T z j ∥ z i ∥ 2 ∥ z j ∥ 2 \operatorname{sim}\left(z_{i}, z_{j}\right)=\frac{z_{i}^{T} z_{j}}{\left\|z_{i}\right\|_{2}\left\|z_{j}\right\|_{2}} sim(zi?,zj?)=∥zi?∥2?∥zj?∥2?ziT?zj??。补充表5对τ在MolCLR训练前的进一步研究。尽管对比学习框架已经被应用到各个领域，包括非结构化图形、句子嵌入和机器人规划。对于分子图的对比学习还没有进行全面而详尽的研究。
Datasets
预训练数据集。对于MolCLR的预训练，我们使用了由ChemBERTa从PubChem收集的1000万个独特的未标记分子SMILES。然后利用RDKit构建分子图并从SMILES串中提取化学特征。在分子图中，每个节点代表一个原子，每条边代表一个化学键。我们以95/5的比例随机将训练前数据集分割为训练集和验证集。
下游数据集。为了测试我们的MolCLR框架的性能，我们使用了来自MoleculeNet的13个数据集，总共包含44个二元分类任务和24个回归任务。这些任务涵盖了多个领域的分子特性。对于除QM9之外的所有数据集，我们使用scaffold拆分来创建一个80/10/10 train/valid/test拆分，如中建议的那样。与常见的随机划分不同，基于分子亚结构的支架分裂使得预测任务更具挑战性，但也更加现实。QM9遵循随机分割设置作为大多数相关工作的实现进行比较。
**训练细节。**分子图上的每个原子按原子序数和手性类型嵌入，每个键按类型和方向嵌入。我们将ReLU激活作为GNN骨干，实现了一个5层图卷积，并遵循Hu等人的修改，使聚合与边缘特征兼容。对每个图应用平均池化作为读出操作，提取512维的分子表示。一个具有隐藏层的MLP将表示映射到一个256维的潜在空间。带重量衰减 1 0 ? 5 10^{-5} 10?5的Adam优化器用于优化NT-Xent损失。在以学习速率初始的10个epoch之后， 5 ? 1 0 ? 4 5 *10^{-4} 5?10?4，执行一个余弦学习衰减。以512批处理大小训练模型，共训练50个epoch。
对于下游任务的微调，我们在基本GNN特征提取器上添加一个随机初始化的MLP。对分类任务和回归任务分别实现了Softmax交叉熵损失和L1损失。在每个任务上，我们对预训练模型进行了三次100 epoch的微调，以得到测试集性能的平均值和标准差。我们只在训练集上训练模型，并在验证集上执行超参数搜索以获得最佳结果。整个框架是基于Pytorch Geometric实现的。补充表6提供了更多的微调细节。
Baselines
有监督学习模型。我们综合评估了我们的MolCLR模型的性能，并将其与监督学习方法进行了比较。对于浅层机器学习模型，采用了以分子指纹为输入的随机森林(Random Forest, RF)和支持向量机(Support Vector machine, SVM)。还包括多个GNNS。考虑了具有边缘特征的GCN和GIN进行聚合。此外，我们还将几个在几个分子基准上达到SOTA的GNN模型(SchNet19、MGCN52和D-MPNN20)作为基准来实现。这些GNN是专门为分子设计的。例如，SchNet和MGCN明确地模拟了分子内部的量子相互作用。
自监督学习模型。为了更好地证明MolCLR框架的有效性，我们进一步将其他训练前或自我监督学习模型作为基线。Hu等人提出了分子图的节点级和图级预训练。需要指出的是，虽然节点级的预训练是基于自监督的，但图级的预训练是基于某些分子性质标签的监督。N-Gram图也被实现，它直接通过分子图计算一个紧凑的表示。