Knowledge-based BERT: 像计算化学家一样提取分子特征的方法

今天介绍一篇浙江大学智能创新药物研究院侯廷军教授团队、中南大学曹东升教授团队和腾讯量子计算实验室联合在Briefings in Bioinformatics发表的一篇论文“Knowledge-based BERT: a method to extract molecular features like computational chemists”。本文提出了一种新的预训练策略，通过学习由计算化学家预定义的分子特征和原子特征，使得模型能够像计算化学家一样从SMILES中提取分子特征。K-BERT在多个成药性数据集上表现了优异的预测能力。此外，由K-BERT 生成的通用指纹 K-BERT-FP 在 15个药物数据集上表现出与 MACCS 相当的预测能力。并且通过进一步预训练，K-BERT-FP还可以学习到传统二进制指纹（如MACCS和ECFP4）无法表征的分子大小和手性信息。

研究背景

基于机器学习算法的分子性质预测模型已成为在药物发现早期阶段对没有前途的先导分子进行分类的重要工具。与主流的基于描述符和图的分子特性预测方法相比，基于 SMILES 的方法无需人类专家知识即可直接从 SMILES 中提取分子特征，但它们需要更强大的特征提取算法和更大量的数据进行训练，这使得基于 SMILES 的方法并不如前两者流行。

Knowledge-based BERT预训练策略

作者在BERT基础上提出了新的预训练策略，让模型能够直接从SMILES中提取分子特征。作者提出了三个预训练任务：原子特征预测任务、分子特征预测任务和对比学习任务。原子特征预测任务允许模型学习基于图的方法中手动提取的信息（初始原子信息），分子特征预测任务允许模型学习基于描述符的方法中手动提取的信息（分子描述符/指纹），而对比学习任务允许模型使同一分子的不同 SMILES 字符串的embedding更相似，从而使 K-BERT 能够识别同一分子的不同SMILES字符串。

预训练任务1-原子特征预测（图1A）：对RDKit计算所得的分子中每个重原子的原子特征进行预测。原子特征将包括度、芳香性、氢、手性和手性类型等，因此可看作是一个多任务分类任务；

预训练任务2-分子特征预测（图1B）：对RDKit计算所得的分子特征进行预测。本研究采用MACCS指纹，同样也可将该任务视为多任务分类任务（可以更换为其他的指纹/描述符）；

预训练任务3-对比学习（图1C）：对于canonical SMILES输入，通过SMILES随机化得到多种不同的SMILES形式。该预训练任务的目标是最大化同一分子不同SMILES字符串嵌入的余弦相似度，最小化不同分子间嵌入的相似度，使得模型能够更好地“理解”SMILES。

图1. Knowledge-based BERT（K-BERT）预训练策略

模型训练和评估

输入表征：每个SMILES都是用Schwaller等人提出的标记化方式来标记成不同的token。然后将token（如’O’、‘Br’和‘[C@H]’）编码成K-BERT的输入。

Pre-training：每个重原子的原子特征和每个分子的分子特征都通过RDKit计算，并分别用于预训练任务1和2。使用RDKit 计算 CHEMBL 中每个分子的一个canonical SMILES和4个随机生成的 SMILES，用于预训练任务3。CHEMBL中大约 180 万个分子被用于预训练K-BERT，目标为最小化3个预训练任务的损失函数。

Fine-tuning：如图1D所示，K-BERT中有6个transformer encoder，我们从预训练模型中载入前5个transformer encoder的参数，第6层transformer encoder和预测层重新随机初始化。然后在下游任务数据上，重新训练模型。

Data augmentation: 每个分子的SMILES，通过rdkit随机扩充到5个不同的SMILES。在训练集中，每个SMILES都被当做是单独（不同的）分子。而在测试集和验证集中，同一分子的不同SMLES都被看成是该分子，5个不同SMILES的预测结果均值作为该分子的预测结果。

ROC-AUC被用于评估分类模型的表现。R²、MAE、RMSE被用于评估回归模型的表现。

数据集与实验任务

成药性小数据集：作者测试了K-BERT在15个成药性小数据集上的表现，数据集分子数都在2000以下。具体数据集如下：Pgp-substrate (Pgp-sub)、human intestinal absorption (HIA)、 human oral bioavailability 20% (F20%)、human oral bioavailability 30% (F30%)、CYPsubstrate (CYP1A2-sub、CYP2C19-sub、CYP2C9-sub、CYP2D6-sub和CYP3A4-sub)、half-life (T1/2)、 drug-induced liver injury (DILI)、FDA maximum recommended daily dose (FDAMDD)、skin sensitization (SkinSen)、carcinogenicity (Carcinogenicity)和respiratory toxicity (Respiratory)。

Malaria 数据集：Malaria 数据集是搜集自Malaria Treatment Response Portal的一个子集，里面的分子都包含手性信息。这个数据集被用于评估K-BERT是否能够学习到手性信息。

CHIRAL1 数据集：CHIRAL1数据集是 Lyu 等人报道的多巴胺受体D4对接筛选数据的一个子集。CHIRAL1中的每个分子只有一个四面体中心，根据中心的手性分为R和S。在本研究中，共有204778个分子用于进一步的预训练，使得K-BERT能学习到手性信息。

实验结果

K-BERT在成药性小数据集上的表现

作者首先评估了K-BERT在15个成药性数据集上的表现。K-BERT取得了优异的表现（表1），在8个数据集中取得了最好的结果。

表1. K-BERT和其他方法在15个成药性数据集上的表现。

Pre-training能够提升模型提取分子特征的能力

数据增强在基于SMILES的模型中表现出了极大的优势。本文发现，数据增强和预训练都是增强模型从SMILES中提取分子特征的能力。如表2所示，作者采用了不同的策略训练模型。由于这里会比较数据增强，而对比学习在预训练过程中有类似数据增强的操作，为了公平，表2中均未采用对比学习预训练任务。K-BERT-WCL显著优于K-BERT-WP，说明预训练能够提升模型提取分子特征的能力。同时，K-BERT-WP-AUG表现优于K-BERT-WP也能够说明，数据增强也能够帮助模型更好的理解SMILES进而提升模型性能。K-BERT-WCL和K-BERT-WCL-AUG表现差不多，这说明数据增强对已经过预训练的模型提升帮助非常有限。这也符合预期，通过预训练，模型已经能够较好的理解SMILES规则。这时候用同一分子的不同SMILES进行数据增强，相当于在对同一分子进行多次训练，自然难以提升模型性能。

表2. K-BERT在不同Pre-training和Fine-tuning策略下的表现。

对比学习任务能够让模型更好的“理解”SMILES

作者比较了同一分子不同SMILES通过模型生成的embedding的平均Tanimoto相似度。结果如图2所示，经过对比学习预训练任务，embedding的相似度得到了显著提升。这说明，对比学习能够帮助模型识别同一分子的不同SMILES字符串。此外，作者以分子‘C=CCC(O)CC(C)(C)C’（不在预训练数据集中）为例，通过RDkit随机生成十个SMILES字符串，并对分子中的不同原子embedding进行了t-SNE可视化，结果如图2所示。结果显示，经过对比学习预训练任务，模型能够识别不同SMILES中同一化学环境下的原子。

图2. 50个分子的embedding平均Tanimoto相似度比较

图3.不同原子embedding的t-SNE可视化

K-BERT能够生成通用的分子指纹K-BERT-FP

K-BERT生成的分子embedding能够作为一种通用的分子指纹K-BERT-FP（不局限于某一任务）。作者比较了K-BERT-FP和MACCS在成药性数据集上的表现（图4），结果显示K-BERT-FP和MACCS取得了可比较的预测能力。

图4. K-BERT-FP和MACCS的比较

K-BERT能够捕捉MACCS不能捕捉的分子大小信息

作者为了证明K-BERT-FP不是MACCS的简单复制，而是能够捕捉到一些MACCS不能捕捉的信息（如分子大小信息）。作者首先比较了K-BERT-FP和MACCS对DrugBank数据集中分子的TMAPs可视化结果。结果如图5所示，K-BERT-FP和MACCS都能够较好的对DrugBank数据集中的分子进行可视化，K-BERT-FP并没有比MACCS更好的组织能力。这可能是由于大分子中的一些分子碎片信息，隐性的包含了分子大小信息，使得MACCS也能够反应分子大小信息。为了进一步比较，作者构建了一个数据集Sim-Sub-Dataset，这个数据集都是基于相似碎片重复生成的（图6）。由于MACCS只对表征是否包含某一分子碎片，而未表征分子碎片的数量，因此MACCS不能够反应此类分子的分子大小情况。作者比较了K-BERT-FP和MACCS在预测该数据集分子权重的能力。结果如表3所示，K-BERT-FP显著优于MACCS，这说明K-BERT-FP能够捕捉到MACCS不能捕捉的信息。

图5. K-BERT-FP和MACCS对DrugBank的TMAPs可视化

图6. Sim-Sub-Dataset的生成方式

表3. K-BERT-FP和MACCS在Sim-Sub-Dataset上的表现

经过进一步预训练，K-BERT能够捕捉MACCS不能捕捉的手性信息

作者首先比较了K-BERT和MACCS在表征CHIRAL1数据集上2500个手性异构体的能力。结果如图7A和图7B所示，K-BERT和MACCS并不能较好的区分CHIRAL1上的手性异构体。为了使得K-BERT-FP中包含手性信息，作者将K-BERT在手性数据集CHIRAL1上进行了进一步预训练，K-BERT-FP-CHIRAL1的分子预测预训练任务仍然是MACCS分子指纹，而K-BERT-FP-CHIRAL1-R-S的分子预测预训练任务改为预测分子的手性R/S。结果如图7C和图7D所示，K-BERT-FP-CHIRAL1-R-S能够很好的区分手性异构体。此外，同一组分子的异构体都能从另一组中找到，这说明K-BERT-FP-CHIRAL1-R-S在表征分子手性信息的同时，仍然蕴含着分子的结构信息。同时，作者比较了不同指纹在手性数据集Malaria上的预测表现（基于XGBoost建模）。计算结果显示，K-BERT-FP-CHIRAL1-R-S优于其他指纹，这说明通过定制化的预训练任务，K-BERT能够关注手性信息，进而提高模型对手性特征的提取能力。

图7. CHIRAL1数据集上2500个手性异构体的TMAPs可视化. (A). MACCS TMAP color encoded by R/S chirality; (B). K-BERT-FP TMAP color encoded by R/S chirality. (C). K-BERT-FP-CHIRAL1 TMAP color encoded by R/S chirality. (D). K-BERT-FP-CHIRAL1-R-S TMAP color encoded by R/S chirality.

总结

作者提出了一种能够像药物化学家一样提取分子特征的预训练策略K-BERT，K-BERT能够较好地从SMILES字符串中提取分子特征，且在成药性预测数据集上表现了较强的预测能力。此外，作者还发现K-BERT能够生成一种通用的分子指纹K-BERT-FP，且K-BERT-FP能够捕捉MACCS不能捕捉的分子大小信息。经过进一步预训练，K-BERT-FP还能够捕捉到手性的信息。这表明，通过对特定任务的了解，可以设置不同的预训练任务来使K-BERT-FP捕捉特定的分子特征信息。

参考资料

Zhenxing Wu, Dejun Jiang, Jike Wang, Xujun Zhang, Hongyan Du, Lurong Pan, Chang-Yu Hsieh, Dongsheng Cao, Tingjun Hou, Knowledge-based BERT: a method to extract molecular features like computational chemists, Briefings in Bioinformatics, 2022;, bbac131,

https://doi.org/10.1093/bib/bbac131