Nat Mach Intel｜MolFormer：大规模化学语言模型表征分子结构与性质

2022年12月21日，IBM的Jerret Ross等人在Nature Machine Intelligence上发表文章Large-scale chemical language representations capture molecular structure and properties。作者提出了MolFormer：一个基于Transformer的大规模化学语言模型表征分子结构与性质。

该模型使用旋转位置编码、线性注意力机制，结合高度分布式训练，对PubChem和ZINC数据集中11亿个未标记分子的SMILES序列进行训练。实验表明，学习的分子表征优于现有模型。进一步的分析证明了SMILES训练的MolFormer确实学习了分子内原子之间的空间关系。这些结果说明大规模分子语言模型可以捕获足够的化学和结构信息，以预测各种不同的分子性质。

文章链接：

https://www.nature.com/articles/s42256-022-00580-7

文章arXiv版本链接：

http://arxiv.org/abs/2106.09553

代码GitHub链接：

https://github.com/IBM/molformer

背景

基于机器学习的模型可以实现准确和快速的分子性质预测，这对药物发现和材料设计很有意义。各种有监督的机器学习模型已经证明了很有前途的性能，但广阔的化学空间和有限的属性标签使有监督的学习具有挑战性。

最近，在大型未标记语料库上预处理的基于无监督Transformer的语言模型在许多下游自然语言处理任务中产生了最先进的结果。受这一发展的启发，作者提出了通过训练有效的Transformer编码器模型MolFormer获得有效的分子表征。

方法

MolFormer的目标是从大规模化学SMILES数据中学习和评估各种下游分子性质预测任务的表征，如图1所示。

图1. MolFormer-XL模型

因此，该模型是使用掩模（masked）语言模型框架开发的，该框架随机屏蔽了SMILES序列中的词元（token）百分比，然后预测这些词元。掩模语言建模利用自监督并实现上下文学习（contextual learning）。

MolFormer是一个基于掩模自编码器（masked autoencoder）的Transformer模型。相比于基本的Transformer模型，Transformer-XL（Transformer for eXtra Long sequences，用于超长序列的Transformer）通过引入相对位置编码（relative position embedding）更好地进行上下文学习。在Transformer-XL的基础上，作者以旋转位置编码（rotary position embedding）实现相对位置编码，通过线性注意力机制（linear attention mechanism）替代基本的注意力机制，构建MolFormer-XL。

该模型在PubChem（包含约1亿个分子）和ZINC（包含约10亿个分子）这两个公共化学数据库中的大量化学分子集合相对应的序列SMILES上，以自监督的方式预训练（pre-training）基于Transformer神经网络的模型，学习有效和低维的化学分子表征。然后通过对任务特定数据的微调（fine-tuning），将模型以迁移学习（transfer learning）的方式，应用于不同的下游分子性质预测任务，计算分子相似性，以及分析给定分子的原子间空间距离和注意力值之间的对应关系。

对于预训练，由于使用掩模语言模型方法，最初，15%的词元被选择为可能去噪的部分。从该选择中，80%的词元将被随机选择并替换为[MASK]词元，10%的词元将被随机选择以替换为随机词元，而其余10%的词元将保持不变。通过整个PubChem+ZINC数据集进行训练，通过无限带宽（InfiniBand）结构连接的2台服务器上的16个GPU上，设置每GPU上1600个分子的批量大小（batch size）。通过分布式的并行训练策略，预训练将能够扩展到大型数据集（10亿个以上的样本）。使用线性注意力替代基本的注意力机制，也能有效降低训练成本。综合如上的策略，MolFormer能够将所需的GPU数量从大约1000个减少到16个，相当于原来的六十分之一。

结果

作者将所提出的PointVS（包括分别使用CoreR和Core80）与现有的多种方法进行了比较，对于每一种方法，分别对比了引入或者不引入晶体姿态作为特征的情况。结果如表1所示，表明PointVS的性能超过了现有模型，亦表明引入晶体姿态作为特征能够有效提升模型的性能。作者将片段细化（fragment elaboration）的结果与现有工具HotspotsAPI做对比， PointVS显著优于HotspotsAPI的结果。

表1. 分类问题上不同方法对比

表2. 回归问题上不同方法对比

图2在QM9数据集中的两个分子上做了可视化分析，显示了带旋转位置编码的MolFormer-XL中间注意力层中的平均学习注意力系数。其中，bond connectivity matrix表征两个原子间是否成键（黑色为成键，白色为不成键），3D distance matrix表征两个原子间的距离（单位Angstrom），full attention with rotary和linear attention with rotary分别表征使用基本的注意力或者线性注意力，结合旋转位置编码的注意力值。

比较不同原子标记对之间的注意力对应的共价键连接性和原子对之间的3D距离，注意力值显示出与中等范围空间距离的高余弦相似性，中间旋转注意力层上的头部聚集很好地对应于共价键，同时还捕获分子内非键合原子之间的空间关系的特征。

这些结果表明，MolFormer-XL能够从相应的SMILES中恢复分子结构信息。这种能力可能来自于对大量SMILES的预训练，并支撑MolFormer-XL学习化学物质的基本特性，包括结构信息和各种性质。在由大规模数据上预先训练的化学语言模型学习的表征中，这是对模型学习分子的结构化的多样性质的首次确认。

图2. 注意力可视化

总结

在这项工作中，作者将无监督的大规模预先训练的分子语言模型运用于各种分子性质预测任务。与图不同，SMILES等分子语言不明确编码分子拓扑结构。然而，通过在大规模语料库上进行精心设计的自我监督训练，例如基于上下文的Transformer的语言模型，具有线性注意力机制，以及并行化训练，作者所提出的MolFormer可以有效地学习隐式丰富的结构属性关系信息。这是第一个验证大规模自我监督能力的工作和预训练的分子语言模型，用于预测从量子化学到生物化学领域的分子性质。

此外，通过分析学习到的注意力，作者发现MolFormer确实在SMILES序列上进行了训练，即使在二维拓扑结构之外，也知道分子内的原子间关系。最后，在大规模学习方面， MolFormer展示了计算资源的高效和环保使用，减少了GPU需要执行60倍的训练（1000对16）。

MolFormer具有在不同靶点上快速进行分子筛选的潜力，这一点对于材料设计和药物发现很重要。此外，有必要进一步探索MolFormer直接从化学语言中学习的结构分子信息扩展到本工作中研究的小有机分子之外。

未来的工作还将通过以下方式改进MolFormer，包括使用更大的模型和更大的训练数据，使用改进的和/或特定领域的自我监督任务，以及使用其他基于字符串的表征等。

参考资料

[1]Ross et al. Large-scale chemical language representations capture molecular structure and properties. Nat. Mach. Intel. 2022

[2]Dai et al. Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. ACL. 2019

[3]Su et al. Roformer: Enhanced transformer with rotary position embedding. arXiv. 2021

[4]Katharopoulos et al. Transformers are RNNs: Fast autoregressive transformers with linear attention. ICML. 2020