作者 | 熊展坤 审核 | 黄紫阳
今天给大家介绍一篇来自浙江大学侯廷军教授课题组和中南大学曹东升教授课题组合作发表在2021年5月份发表在Briefings in Bioinformatics上的一篇文章《MG-BERT: leveraging unsupervised atomic representation learning for molecular property prediction》。本文使用了无监督原子表达学习来进行分子的性质预测。
1.摘要
动机:准确和有效的分子性质预测是药物设计的基本问题之一。传统的基于特征工程的方法在特征设计和特征选择过程中需要广泛的专业知识。随着人工智能技术的发展，数据驱动的方法在各个领域都比基于特征工程的方法具有无可比拟的优势。然而，在应用于分子性质预测时，人工智能模型往往缺乏标记数据，泛化能力较差。
结果:在本研究中，作者提出了分子图BERT (MG-BERT)，它将图神经网络(GNNs)的局部消息传递机制集成到强大的BERT模型中，以方便从分子图中进行学习。此外，本文提出了一种有效的自监督学习策略，即掩蔽原子预测，对MG-BERT模型进行预处理，以挖掘分子中的上下文信息。作者发现，MG-BERT模型可以在预处理后生成上下文敏感的原子表示，并将学习到的知识用于各种分子性质的预测。实验结果表明，在11个ADMET数据集上，添加少量微调的MG-BERT预训练模型的性能始终优于现有方法。此外，MG-BERT模型还使用了注意力机制来关注对目标属性至关重要的原子特性。并且，训练模型具有良好的可解释性。MG-BERT模型不需要任何手工制作的特征作为输入，并且由于其出色的可解释性和可靠性，为开发最先进的模型提供了一个新的框架，使其能够被广泛的应用于药物发现任务。
2.Introduction
药物发现是一个有风险、耗时长、耗费资源的过程，通常需要10-15年时间和数十亿美元的资金。为了提高药物发现的效率，研究者们在计算工具和生物信息学方法的发展方面投入了大量的努力。在这些方法中，计算性方法模型能够更加准确的预测分子性质，这对药物发现过程的影响更为显著和直接，因为它可以缓解过度依赖耗时和劳动密集型实验的情况，大大降低了经济成本和时间成本。高精度的分子特性预测模型已经成为药物发现过程中许多阶段不可或缺的工具。
优秀的分子表示对于分子预测来说是必不可少的，传统方法严重依赖于特征工程，将分子的相关结构信息或者物理化学性质编码为固定长度的向量。分子指纹和分子描述符是分子特征的两大类但是他们的可扩展性和通用性都很低。近年来，深度学习(deep learning, DL)方法在许多领域取得了重大突破，如计算机视觉、自然语言处理(natural language processing, NLP)等。DL背后的基本原理是:设计一个合适的深度神经网络(DNN)，并在大量原始数据上对其进行训练，从而自动学习表示，而不是依赖于人工制作的特征。
DL在各个领域的成功应用启发了它在分子性质预测方面的应用。许多分子属性预测的研究试图将DNN直接应用于低阶分子表示，如SMILES 字符串或分子图。SMILES字符串通过一行ASCII字符串描述分子的组成和化学结构。作为一种文本，一些合适的文本处理算法，如CNN、LSTM和Transformer，可以直接应用于构建预测模型。然而，这些算法需要学会从SMILES复杂的语法中解析出分子的有用特征，这大大增加了学习和泛化的难度。值得注意的是，基于自动编码器等模型的无监督方法已被应用于SMILES，从大量的无标记数据中学习有用的表示.这些模型可以通过训练大量的未标记数据来将离散分子嵌入到一个连续的向量空间中。对下游预测任务可以采用潜在表示。然而，基于SMILES的分子表示对于一般预测任务可能不是最优的，不能进一步优化。新兴的GNN可以直接从图数据中学习，这在分子性质预测方面具有很大的优势。然而，由于受过拟合和过平滑问题的限制，目前的GNN通常太浅(一般为2-3层)，削弱了它们提取深层信息的能力。
而且，DL模型在分子性质预测中面临的共同挑战是缺乏标记数据。众所周知， DL模型通常需要大量的标记数据来达到高有效性和泛化性。例如，在图像分类任务中，人们通常会收集数百万张图像来训练他们的DL模型。不幸的是，获取如此多的分子性质数据是不现实的，特别是ADMET端点数据，往往需要大量耗时、费力、昂贵的实验。这种困境使得DL模型往往容易过拟合，极大地损害了其泛化能力。
标签数据的缺乏促使了自监督或半监督学习方法在其他领域的发展。在自然语言处理领域，最近提出的BERT模型可以利用大量的未标记文本进行预先训练，显著提高各种下游任务的性能。BERT模型的成功可以归因于掩码预测，该模型根据同一句子中的其他可见词学习预测掩码或污染词。在这个过程中，该模型被驱动来挖掘句子中的上下文信息。这种上下文信息可以适合被下游任务使用，并大大提高其预测性能。受BERT模型的启发，有研究者提出了SMILES -BERT模型，他们将BERT模型直接应用于SMILES字符串。但是SMILES字符串中存在辅助字符，SMILES- BERT模型缺乏可解释性。此外，SMILES字符串复杂的语法也增加了模型学习的难度。
为了解决上述这些问题，作者提出了一种新的分子图BERT (MG-BERT)模型，将GNN的局部消息传递机制集成到强大的BERT模型中。所提出的MG-BERT模型能够克服一般GNN所面临的过平滑问题，并为分子表达的生成提供了足够的深度特征提取能力。作者进一步提出了掩蔽原子预测预训练作为自动挖掘分子上下文信息的有效策略。实验结果表明，MG-BERT模型能够在11个实际任务上生成上下文敏感的原子表达，并显著提高了分子属性预测任务的性能，且在这11个实际任务上，MG-BERT模型的性能始终优于最新的模型。此外，MG-BERT还可以通过注意力机制来关注原子和子结构，为分析和优化分子提供有价值的线索。
3.模型框架
与非结构化自然语言处理的原始BERT模型不同，MG-BERT根据分子图的特点进行了一些修改。在嵌入层中，单词标记被原子类型标记取代。由于分子中的原子没有顺序关系，因此不需要分配位置信息。在自然语言的句子中，一个词可能与其他任何词有关，因此需要global attention。然而，在分子中，原子主要是与它由键连接的的原子相连。为了有效地实现这种归纳偏置，作者将BERT中的global attention修改为基于化学键的local attention，只允许原子通过化学键交换信息。这种本地消息传递机制使MG-BERT成为GNN的一个新的变种。值得注意的是，MG-BERT可以克服BERT中的res-connection机制导致的过平滑问题，并有足够的能力提取分子图中的深层特征信息。如图1所示，作者使用邻接矩阵来控制分子中的信息交换。
为了获得graph-level的分子表示，便于后续微调阶段的预测任务，作者为每个分子添加了一个连接所有原子的超级节点。一方面，该超级节点可以与其他节点进行信息交换，在一定程度上可以很好地解决远程依赖问题;另一方面，这种超级节点输出可以看作是最终的分子表示，用于解决下游的分类或回归任务。
BERT利用两个学习任务对模型进行预训练，分别是MLM和NSP。MLM是一个填空任务，其中一个模型使用上下文词包围一个掩码标记来预测掩码词应该是什么。NSP任务是确定两个句子是否连续。由于分子缺乏像句子这样的持续关系，作者只使用掩蔽原子预测任务来预训练模型。作者提出的预训练策略与BERT非常相似。首先，随机选择一个分子中15%的原子，对于只有少量原子的分子至少选择一个原子。对于每个选定的原子，有80%的概率被[MASK]标识替换，有10%的概率被其他原子随机替换，有10%的概率保持不变。用原始分子作为训练模型的groud truth，只计算被掩盖原子的loss。
模型的框架如上图。模型主要分成两个阶段，预训练阶段和精调阶段（fine-tuning）。
（1）预训练阶段
在预训练阶段，作者使用RDKit根据每个分子的组成原子及其连接关系，将每个分子转换成二维无向图。然后在每个分子图上添加一个连接所有节点的超节点。在此之后，随机选择某些原子根据预训练策略来进行掩蔽。最后，将分子图发送给MG-BERT模型，以预测原子的类型。对于一些只有几个原子的分子，作者保证至少有一个原子被选择用于掩蔽。该模型采用标准批次梯度下降算法和Adam优化器进行训练。学习速率设置为1e−4，批处理大小设置为256。该模型预训练10个epoch。为了评价预训练的表现，采用预训练掩蔽策略对测试集中的分子进行掩蔽，然后计算回收率作为评价指标。
（2）精调阶段（fine-tuning）
预训练结束后，取下预训练头。在超级节点对应的Transformer编码器层输出中添加一个两层的任务相关的全连接神经网络。采用dropout策略来最小化过拟合。需要注意的是，dropout rate对最终预测有很大的影响，需要根据具体任务进行优化。根据作者的经验结果，dropout rate的推荐范围[0.0,0.5]。使用Adam优化器作为微调优化器，对每个任务进行有限的超参数扫描，批次大小从{16,32,64}中选择，学习速率从{1e-5,5e-5,1e-4}中选择。采用方差分析评价回归模型，采用ROC-AUC (ROC-AUC)评价分类模型。作者使用early stopping策略来避免过拟合，并将最大epoch设置为100。为了减少随机误差，作者对每个数据集进行10次随机数据分割训练，计算出的平均值和标准差作为最终性能。
4.实验结果
1.数据
作者从ChEMBL数据库中随机抽取170万compounds作为预训练数据。为了验证预训练模型，作者随机保留10%进行预训练评估。训练的个数最终达到153万。在微调阶段，对预训练模型进行进一步训练，以预测特定分子的性质。从ADMETlab和MoleculeNet中收集了16个数据集(8个用于回归和8个用于分类)，包括关键的ADMET端点和各种常见的分子特性，用于训练和评价MG-BERT。这16个数据集的详细信息如表所示。这些数据集中的所有分子都以SMILES字符串格式存储。数据集按8:1:1的比例分为训练数据集、有效数据集和测试数据集。值得注意的是，SMILES字符串的长度范围很广，从几个字符到超过100个字符。因此，采用SMILES长度分层抽样，使数据分割更加均匀。
2.MG-BERT模型结构选择
作者设计了三种MG-BERT模型，分别具有大中小三种规模，结果如下表所示，小型MG-BERT由于层数太少，不如其他两种模型。与中等MG-BERT模型相比，大MG-BERT模型在预训练任务上表现更好，而在分子属性预测任务上表现稍差。这种现象可能是由于大型MG-BERT模型由于模型参数过多而存在过拟合风险所致。由于MG-BERT模型在分子性质预测方面具有最好的性能，最终采用了MG-BERT模型。小MG-BERT模型由于层数太少，不如其他两种模型。与中等MG-BERT模型相比，大MG-BERT模型在训练前恢复任务上表现更好，而在分子属性预测任务上表现稍差。这种现象可能是由于大型MG-BERT模型由于模型参数过多而存在过拟合风险所致。由于中型MG-BERT模型在分子性质预测方面具有最好的性能，最终采用了中型MG-BERT模型。
3.预训练有效性验证
为了验证预训练的有效性，在相同的参数设置下，作者比较了预训练和未预训练的MG-BERT模型在分子性质预测方面的性能。从下表的对比结果可以看出，经过预处理的MG-BERT模型在所有数据集上都比未经过预处理的MG-BERT模型的泛化性能提高了2%以上，这清楚地说明了预处理策略的有效性以及预处理后模型良好的泛化能力。对于一些小数据集如Caco2和FDAMDD，预测性能提高了7%以上，表明预训练策略可以更有效地提高对小数据集的预测性能。这些结果表明，MG-BERT模型确实可以学习有用的知识，并将学习到的知识传递给下游任务。
4.氢原子对预训练准确度和预测任务的影响
在大多数分子性质预测模型中，氢原子通常被忽略。在本研究中，作者进行了一个控制实验来探索MG-BERT模型是否需要氢原子。MG-BERT模型在相同的超参数设置下，建立了基于不含所有氢原子的分子图的无氢模型。如下图所示，添加氢的MG-BERT模型的预训练精度可以达到98.31%，而无氢模型的预训练精度只能达到92.25%。
从下表所示的微调结果可以看出，MG-BERT含氢模型的性能远远好于无氢模型。特别是在一些回归任务中，含氢模型的性能优于无氢模型的4%以上。这背后的逻辑是，MG-BERT只利用分子的组成和连接信息。在这种情况下，氢原子可以帮助确定其他类型原子的化学键数。在屏蔽原子恢复任务中，键的数目对于确定屏蔽原子的类型是至关重要的。因此，无氢MG-BERT模型显示掩蔽原子回收率显著降低。此外，氢原子的缺失也会影响到预训练阶段的上下文信息挖掘过程，削弱预训练模型的泛化能力。
此外，如果去除氢原子，一些分子就会变得难以分辨。如图所示，如果去除氢原子，苯和环己烷可以转化为同一个图。然而，如果保留氢原子，它们将被转换成两个不同的图。因此，氢原子的缺失对微调模型的性能有很大的影响。
5.与其他机器学习方法的对比
基于不同的分子表征，作者选择了一些最先进的模型作为baseline，综合评价所提出的MG-BERT模型。第一个是基于ECFP4指纹的XGBoost模型(ECFP4-XGBoost)。第二个和第三个GNN两个最具代表性的图注意网络(GAT)和图卷积网络(GCN)。第四个是基于连续数据驱动的描述符(CDDD)，它由一个固定的基于RNN的编码器和一个完全连接的神经网络组成，该编码器已经对大量未标记的SMILES字符串进行了预训练。baseline方法中还包括了SMILES-BERT模型，它直接使用了SMILES字符串的原始BERT模型。预测结果见下表和下图。
ECFP4-XGBoost模型的性能在不同的数据集上表现出很大的差异。ECFP4很可能是一个固定长度的分子表示，这导致它所表示的信息可能适合或不适合特定的任务。当标记数据充足时，GNN模型(包括GAT和GCN)表现良好。然而，当标记数据稀缺时，它们的性能会变得更差，甚至比基于分子指纹的模型还要差。CDDD模型具有一定的竞争力。然而，CDDD模型的分子表示是通过SMILES编码和解码任务获得的，不能针对特定任务进一步优化。相比之下，SMILES BERT模型和MG-BERT模型能够在预训练阶段学习丰富的上下文敏感信息，并可以进一步优化具体任务。SMILES-BERT模型稍微落后于MG-BERT模型。这可能是因为从SMILES字符串学习要比从分子图学习复杂得多，这意味着SMILES- bert模型必须解析隐藏在SMILES字符串复杂语法中的分子信息。虽然MG-BERT模型可以直接从分子图中学习，而分子图是分子的自然表示。所提出的MG-BERT模型可以一致地优于其他方法。总体改善为28.1%(分类任务为7.02%，回归任务为21.28%)。值得注意的是，对于PPB数据集，MG BERT模型的改进超过了6%。根据配对t检验(P≤0.001)，MG-BERT模型相对于baseline方法的改进具有统计学意义(95%置信区间，CI)。这些结果有效的说明了MG-BERT在分子性质预测任务和在药物设计方向的潜力。
5.总结
本文提出的MG-BERT将GNN的本地消息传递机制集成到强大的BERT中。MG-BERT作为一种新型的GNN，能够克服过平滑问题，有足够的能力提取分子图中的深度特征。并且MG-BERT可以利用大量的未标记分子，通过掩蔽原子来挖掘分子图中的上下文信息，并将学到的知识转移到分子属性预测中。MG-BERT还可以在没有任何手工设计的特征的情况下，在分子性质预测上优于最先进的模型精心设计的特征，并通过注意力机制提供可解释性，根据与目标性质的相关性，合理地分配注意力系数权重到原子或子结构。在我们今后需要对药物或者化合物分子提取特征时，可以考虑使用这一方法来进行分子特征的提取。
参考文献
文章地址
https://doi.org/10.1093/bib/bbab152
代码地址
https://github.com/zhang-xuan1314/Molecular-graph-BERT.
