MG-BERT：面向药物分子性质预测的大规模原子表征预训练模型

今天给大家介绍的是国防科技大学计算机学院张小琛博士、吴诚堃副研究员、中南大学曹东升教授及浙江大学侯廷军教授等人联合发表在Briefing in Bioinformatics上的一篇文章。作者将自然语言处理领域常用的BERT模型应用到分子图上进行无监督预训练，然后根据带标签的分子性质预测任务进行微调。实验表明，预训练过程可以很好的捕捉分子中原子“上下文”的不同，并产生对应的隐含表征，为下游的性质预测任务性能的提高奠定了坚实的基础。

1 研究背景

药物研发周期长，投入大且极具风险性。精确且高效的分子性质预测模型可以很大程度减少对于实验的依赖，减少成本，加快进度。传统的基于分子指纹与描述符的方法需要大量专业的知识进行优化设计，缺乏通用性与扩展性。目前，深度学习模型可以自动的从原始数据中提取相关特征并在大量任务中取得重大突破。然而由于分子性质预测数据的稀缺性，限制了深度学习模型在分子性质预测任务中的表现。

作者受到目前备受关注的预训练模型的启发，提出了基于分子图的MG-BERT模型，并通过随机覆盖分子图中的部分原子然后使用MG-BERT模型进行覆盖原子的恢复。在这个过程中，作者发现模型可以很好的捕捉分子中不同的原子环境信息，并对下游模型产生积极的泛化。

2 方法概述

2.1 模型概述

文章模型基本采用了自然语言处理中的BERT模型的基本结构，在BERT模型的注意力机制中加入分子邻接矩阵信息，只允许有原子键的原子之间交换信息。这样就赋予了模型在分子图上进行学习的能力，修改过的模型被称为MG-BERT (Molecular Graph BERT)模型。图2展示了MG-BERT模型的结构。模型除了需要输入每一个分子的原子列表外，还需要分子邻接矩阵以控制信息交换只发生在有键连接的原子之间。除此之外还加入了一个具有全局视野的超级节点，这个节点可以与其它所有的原子交换信息，并在微调阶段作为分子的整体表示。值得注意的是，本文并没有如传统的图神经网络使用多个原子特征进行原子表征嵌入，只使用了原子类型这一信息，这样可以大大降低模型的复杂度。

图2 MG-BERT模型结构

模型除了使用分子的构成信息之外，没有使用键的信息，这导致一些不同的分子可以转换成相同的图结构，如图所示。为了解决这一问题，文章提出把每一个原子的连接的氢原子显示地表达出来，这样模型可以推算出原子成键信息。文章中也做实验验证了这一点，如图3所示，在预训练阶段加入氢原子的模型可以得到更好的预训练精度，在下游任务上也有更好的泛化能力。

图3 氢原子对转化分子图的影响

图4 预训练结果

文章选择了基于Morgan指纹的XGBoost算法，以及常用的GCN，GAT，基于SMILES的CDDD指纹，SMILES-BERT作为基线模型进行对比。如图5所示，可以看出MG-BERT模型在多个不同任务下均取得了超越基线模型的良好表现。

图5 在具体分子性质预测任务上与基线模型对比结果

比较有意思的是，文章对为什么预训练会有用尝试进行了解答。文章随机选择了1000了分子，不进行覆盖的情况下送入预训练好的模型之中，并收集预训练模型输出的每一个原子表征，最后进行t-SNE可视化。结果如图6所示，可以看到，预训练模型可以根据每一个原子所处的环境不同，生成不同的原子表征。这跟传统的Morgan指纹有异曲同工之妙。可以预见，这些信息对于下游的分子性质预测可以产生一定的促进作用。

图6 预训练生成的原子表示可视化

除了具有良好的分子性质预测性能，模型的注意力机制也赋予了模型一定的可解释性。超级节点通过注意力机制从其他原子节点收集信息来产生分子表示，不同的注意力权重反映了每一个原子在最终的表示中的重要程度。通过可视化超级节点的权重，可以分析每一个原子对于模型最终结果的影响程度。文章分别选择了logD任务以及毒性相关的Ames任务进行微调，并展示了权重分布情况。如图7所示，模型对性质相关的原子或者基团赋予了更大的权重。

图7 注意力权重可视化

参考资料

Xiao-Chen Zhang, Cheng-Kun Wu, Zhi-Jiang Yang, Zhen-Xing Wu, Jia-Cai Yi, Chang-Yu Hsieh, Ting-Jun Hou, Dong-Sheng Cao, MG-BERT: leveraging unsupervised atomic representation learning for molecular property prediction, Briefings in Bioinformatics, 2021;, bbab152,

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

代码

https://github.com/zhang-xuan1314/Molecular-graph-BERT