Research | 使用无监督学习、多任务学习的分子性质预测新方法

2022年12月15日，中南大学湘雅药学院曹东升团队，国防科技大学吴诚堃团队，浙江大学侯廷军团队以及湖南大学曾湘祥教授团队合作在Research期刊上发表论文“Pushing the Boundaries of Molecular Property Prediction for Drug Discovery with Multitask Learning BERT Enhanced by SMILES Enumeration”。

准确预测小分子的药理特性在药物发现中扮演重要角色。传统的特征工程方法严重依赖手工描述符或指纹，需要大量的人类专家知识。随着人工智能技术的快速进步，数据驱动的深度学习方法显示出基于特征工程的方法无可比拟的优势。然而，现有的深度学习方法在应用于预测分子特性时通常存在标记数据稀缺和不同任务之间无法共享信息的问题，从而导致泛化能力较差。为此文章提出了一种新的多任务学习框架，名为MTL-BERT，利用大规模预训练、多任务学习和SMILES（简化分子输入行输入规范）枚举来缓解数据稀缺问题。MTL-BERT首先通过自监督预训练利用大量未标记数据来挖掘SMILES字符串中丰富的上下文信息，然后利用它们的共享信息同时为多个下游任务微调预训练模型。同时，在预训练、微调和测试阶段将SMILES枚举用作数据增强策略，以大幅增加数据多样性并帮助从复杂的SMILES字符串中学习关键相关模式。实验结果表明，经过少量额外微调的预训练MTL-BERT模型可以在60个实用分子数据集的大部分上实现比最先进方法更好的性能。此外，MTL-BERT模型利用注意力机制来关注SMILES字符特征，这些特征对于模型可解释性的目标属性至关重要。

引论

尽管目前生物技术取得了重大进展，但药物开发仍然是一个漫长、昂贵且复杂的过程，通常需要10到15年和数十亿美元。在药物发现过程中应用计算方法，特别是分子性质预测方法，可以减轻对耗时费力的实验的过度依赖。因此，在过去几十年中，预测化合物各种性质的计算方法的发展引起了广泛关注。

对于分子特性预测，主要挑战之一是学习分子结构的表达表示。传统的分子表示方法依赖于复杂的手工特征，如分子指纹或描述符，经常遭受冗长的设计程序和有限的适应性。近年来，蓬勃发展的深度学习方法提供了一种数据驱动的方式，可以在端到端训练中从原始数据中自动学习分子表征。目前已经有很多研究将成熟的深度学习模型应用于低级分子表示，包括分子图和SMILES表征。

通常，深度学习模型需要大量标记数据才能实现高效和良好的泛化能力。例如，在图像分类任务中，人们通常会收集数百万张图像来训练深度学习模型。不幸的是，对于大多数与生物活性相关的任务，由于昂贵且耗时的实验室实验，标记数据非常有限。训练数据的稀缺性大大增加了过度拟合的风险，并降低了数据密集型深度学习方法的泛化能力。

数据稀缺问题的一个可行解决方案是利用其他标记任务甚至未标记数据中包含的信息和知识。一些生化数据，如吸收、分布、代谢、排泄和毒性(ADMET)特性，是高度相关的。因此，多任务学习可用于利用不同分子特性任务之间的相关性并提高模型性能。除了从标记数据中学习表示之外，未标记数据包含丰富的知识和微妙的模式，这些知识和模式对表示学习非常重要，可以通过无监督学习加以利用。

缓解数据稀缺问题的另一种方法是数据扩充。对于给定的分子可以通过不同的起始原子和遍历顺序由不同的SMILES字符串表示。因此，每个训练样本都可以通过不同数量的SMILES表示进行扩展，以增加数据多样性并帮助学习隐藏在SMILES字符串的复杂语法中的关键相关模式。此外，SMILES枚举也可用于测试阶段，以纠正潜在的预测偏差，从而做出稳健而准确的预测。

在这项研究中，本文开发一种名为MTL-BERT的新型多任务学习框架，通过将大规模无监督学习、多任务学习和SMILES枚举相结合来缓解数据稀缺问题。所提出的MTL-BERT模型首先在大量未标记的分子数据上进行预训练，以挖掘SMILES字符串中丰富的上下文信息。在微调阶段，预训练模型在多个任务上联合训练，以挖掘和共享相似任务中的相关信息。来自多个任务的训练数据也可以通过相互施加约束来作为归纳偏差，从而提高预测准确性和学习速度。此外，SMILES枚举继续用作数据增强策略，以大幅增加预训练、微调训练和测试阶段的数据多样性。通过将这3种策略结合在一起，MTL-BERT可以为数据不足的分子特性预测任务提供解决方案。

材料与方法

数据集

本文首先从ChEMBL数据库收集了170万个未标记分子被用作预训练数据，以学习SMILES字符串中的上下文信息。本文随机保留了整个预训练数据的10%来评估预训练模型。另外，本文从ADMETlab和MoleculeNet两个分子性质预测基准收集了涵盖关键ADMET端点和各种常见分子特性的60个数据集（16个用于回归，44个用于分类），用于训练和评估MTL-BERT。

模型架构

如图1所示，MTL-BERT模型利用Transformer编码器部分作为特征提取器，利用多头自注意力机制同时捕获长期和短期依赖。

图1 MTL-BERT的预训练（A）和微调过程（B）示意图

基于masked SMILES recovery的预训练策略

本文的预训练策略遵循BERT的策略来屏蔽输入的SMILES。首先，SMILES中15%的token会被随机选择，对于只有少数token的分子，至少会选择一个token。对于每个选中的标记，它有80%的几率被特殊掩码标记替换，有10%的几率被字典中的其他标记随机替换，还有10%的几率保持不变。原来的SMILES作为groundtruth来训练模型，loss只在maskedtokens处计算。

微调模型用于多分子性质任务预测

在微调阶段，本文根据每个输入SMILES字符串之前的任务数量填充了几个任务标记。每个标记的模型输出与任务相关的可训练前馈神经网络配对，然后用于不同的分子特性预测任务。在微调阶段，本文只对SMILES字符添加位置编码以保持与预训练阶段一致。此外，本文还添加了一个注意力掩码，如图1所示，以防止任务直接交换信息。防止任务字符之间的直接信息交换允许任务直接从SMILES字符中学习，而不会造成不同任务之间的干扰，尤其是当任务数量很大时。此外，预训练和微调之间的不一致可以大大减少，因为SMILES字符在预训练阶段看不到任务字符。此外，它还可以为模型提供更好的可解释性。

结果

整体训练测试框架

MTL-BERT的概览流程如图2所示。MTL-BERT模型首先通过掩码标记预测任务对大量未标记的分子数据进行预训练，以挖掘SMILES字符串中的上下文信息。在预训练阶段，首先使用不同的起始原子和遍历顺序枚举SMILES字符串。然后，这些SMILES字符串被标记化并进一步随机屏蔽以进行预训练预测。SMILES增强策略可以显着增加数据多样性并有效地从SMILES字符串中学习语义信息。接下来，将用于多分子特性预测任务的每个数据集以8:1:1的比例随机分为训练、验证和测试数据集。之后，训练、验证和测试数据集被连接成多任务训练风格。然后，这些数据集通过随机SMILES枚举增加了20倍。在预测阶段，本文对来自同一分子的枚举SMILES进行所有预测的融合操作，以获得最终预测。

图2 模型训练测试流程图

模型结构研究

为了找出哪种MTL-BERT结构可以更好地完成分子特性预测任务，本文设计并比较了3种不同大小的结构。表中列出了3种类型的MTL-BERT结构的模型参数和性能。预训练恢复精度和平均微调性能被用作评估指标。本文选择了中等的MTL-BERT结构，因为它需要更少的训练成本并且可以获得更好的预测性能。

表1 模型超参数调优结果

策略有效性验证

本文将本文开发的MTL-BERT模型与单任务BERT(STL-BERT)模型和规范的基于SMILES的BERT(Cano-BERT)模型进行了比较。STL-BERT应用了与MTL-BERT相同的设置，但分别对每个任务的预训练模型进行了微调。Cano-BERT去除了MTL-BERT模型中的SMILES枚举步骤，仅使用canonicalSMILES对BERT模型进行预训练，并对每个任务分别微调预训练的BERT模型。结果如图3所示，可以看出与STL-BERT模型相比，Cano-BERT模型在所有任务中都表现出明显的性能损失，甚至在某些数据集中表现出超过10%的退化。此外，尽管都使用SMILES枚举，但MTL-BERT模型在大多数数据集上的表现优于单任务STL-BERT模型。简而言之，这些结果充分证明了多任务处理和数据增强策略的有效性。

图3 模型策略有效性实验结果

与其它机器学习的比较

本文选择了5个最先进的分子特性预测模型作为综合评估本文提出的MTL-BERT模型的基线。第一个是基于直径为4的扩展连接指纹(ECFP4-XGBoost)的XGBoost模型，这是分子特性预测任务的经典范例。此外，还包括3个具有代表性和广泛使用的GNN作为基线：图注意力网络、图卷积网络和AttentiveFP。最后一种方法基于连续和数据驱动的描述符（CDDD），它由一个固定的RNN（递归神经网络）编码器组成，该编码器已经在大量未标记的SMILES字符串上进行了预训练，并带有一个全连接神经网络。对于每项任务，本文使用网格搜索算法为基线模型找到最佳超参数设置。

结果如图4所示，显然，MTL-BERT的性能在绝大多数任务上都优于基线模型。除了更好的预测性能外，MTL-BERT模型不需要对每个任务进行复杂的超参数搜索，而且运行非常高效，这充分证明了其作为分子性质预测的良好选择的潜力。

图4 模型与基线模型对比结果

通过t-SNE分析来自预训练MTL-BERT模型的SMILES标记的表示

为了分析MTL-BERT模型在预训练阶段学到了什么，本文可视化了预训练MTL-BERT模型生成的SMILES字符串标记的表示，并试图找到隐藏在SMILES字符串空间中的一些有趣模式。具体来说，随机选择1,000个分子（包括大约35,000个标记）并无掩蔽地馈入预训练的MTL-BERT模型，并收集Transformer编码器层的输出以进行可视化分析。这样就为每个原子生成了一个256维的向量，总共得到了大约35000个向量。经典的降维方法t-SNE用于可视化这些高维向量。如图5所示，不同类型的标记清楚地聚集在一起并且可以很容易地区分，表明生成的表示包含原子类型的信息。进一步观察表明，同一类型的原子可以分为几个不同的组。生成的原子表示似乎比简单的原子类型包含更丰富的信息。为了进一步观察详细的模式，本文制作了一些字符“O”和字符“c”的局部放大图。在这些放大的图中，本文标记了聚集在一个小区域中的几个标记表示，并显示了相应的分子。第一张和第二张放大图上的“O”标记均代表与苯环相连的硝酸基团中的氧原子，但不同图上的氧原子位置不同。在第三张放大图中，“O”符号代表与苯环相连的羰基中的氧原子。这些结果清楚地表明，来自附近位置的标记的分子邻域环境是相似的，而来自较远位置的标记则明显不同。其他3张放大图显示了类似的结果，可以进一步证明这一发现。

图5 预训练模型字符向量聚类图

MTL-BERT的attention分析

MTL-BERT利用注意力机制从所有SMILES标记中聚合信息，以构建与任务相关的分子表示。因此，注意力权重在某种意义上代表了每个SMILES字符对最终分子表示的特定贡献，并且可以被视为对目标属性的相关度量。因此，MTL-BERT提供了一种自然的方法来发现分子子结构与分子性质之间的关系，这对于分析和优化分子至关重要。

为了验证MTL-BERT模型是否能够合理分配注意力权重，本文对LogS和AMES任务的测试集中的一些分子进行了分析。LogS任务与分子的水溶性有关。在图6A中，似乎更多的注意力集中在极性基团上，这是决定分子在水中溶解度的重要因素。AMES任务与分子的致突变性有关。根据图6B，注意力主要集中在叠氮化物、亚硝胺、酰氯和亚硝酸盐基团上，由于它们在诱变剂中的频繁出现，它们已被证明是诱变结构警报。结果表明，MTL-BERT能够为特定任务分配注意力权重，这在一定程度上对于药物化学家探索分子性质与亚结构之间的关系非常重要。

图6 模型性质预测时注意力分配图

结论

本文同时使用大规模预训练从未标记数据中学习，多任务的迁移学习从其他任务中学习，数据增强来增加标记数据的多样性。这3种策略的结合可以充分利用可用的数据信息，尽可能缓解数据稀缺问题。大规模预训练可以有效地挖掘出隐藏在未标记分子数据中的丰富知识和复杂模式。由于分子特性通常高度相关，因此多任务学习来探索不同任务之间的相似性，以通过迁移学习来提升整体性能。此外，多任务学习还可以通过一次训练多个任务来大大提高预测性能并降低计算成本。SMILES枚举作为一种有效的数据增强策略，可以显着增加训练数据的多样性，并有助于关注更重要的特征，这些特征可以用于预训练和微调以从 SMILES 字符串中提取关键相关特征。此外，SMILES 枚举也可用于测试阶段，以纠正模型潜在的预测偏差并执行更稳健的预测。

参考资料

Zhang, X.C., Wu, C.K., Yi, J.C., Zeng, X.X., Yang, C.Q., Lu, A.P., Hou, T.J. and Cao, D.S., 2022. Pushing the Boundaries of Molecular Property Prediction for Drug Discovery with Multitask Learning BERT Enhanced by SMILES Enumeration. Research, 2022, p.0004.