文献总结｜通过大规模化学语言表示捕获分子结构和性质

doi.org/10.1038/s42256-022-00580-7

本文介绍于 IBM 研究中心于 2022 年发表在 Nature Machine Intelligence 上的一篇文章，文章原标题为 Large-scale chemical language representations capture molecular structure and properties，文章使用大量来自 PubChem 和 ZINC 的分子训练了基于 Transformer 的化学语言模型，该模型将原始的 SMILES 序列转化为向量，文章结果表明，使用这种方式实现的分子词嵌入在许多任务上都表现优异。

在深度学习领域，SMILES 是最为常用的分子表示形式，但越来越多研究结果表明，SMILES 无法反映分子的拓扑结果，在许多任务上都有局限性。这使得许多研究人员转向使用基于图的分子表示方式，然而由于缺少用于训练的标记数据，研究这一类模型也具有很大挑战。

预训练语言模型是近年 NLP 领域中的热点，先使用大量语料训练语言模型，再微调该预训练模型并用于解决具体任务。借鉴这一概念，文章的目标就是使用庞大的分子数据集（11 亿个分子）训练化学语言模型，让该模型能够捕获分子的深层特征，并用它来完成各种分子预测的下游任务。

方法

模型

文章所使用的模型称为 MoLFormer，该模型通过对 Transformer 改造得到，减小了计算开销，适用于处理大量数据。在 Transformer 中，输入序列中的每个位置上的 token 都会通过位置嵌入将位置信息加入到词向量中。由于 self-attention 机制，Transformer 在每个位置 $m$ 上需要计算：

$$\mathrm{Attention}_m(Q,K,V)=\frac{\sum_{n=1}^N\exp(\langle q_m,k_n\rangle)v_n}{\sum_{n=1}^N\exp(\langle q_m,k_n\rangle)}$$

指数计算的计算开销很大，可以用核函数近似的方法简化计算（Vanilla Transforme），得到线性注意力：

$$\mathrm{Attention}_m(Q,K,V)=\frac{\sum_{n=1}^N\langle\varphi(q_m),\varphi(k_n)\rangle v_n}{\sum_{n=1}^N\langle\varphi(q_m),\varphi(k_n)\rangle}$$

此外，文章还引入了旋转位置嵌入（Rotary Position Embeddings, RoPE）替换了原来的位置嵌入方式。RoPE 是一种新的位置编码技术，旨在理解文本中单词的顺序表示，在化学语言模型中，它可能能够更好地捕捉 token 之间的内在关系，表示出分子的拓扑结构。

在训练上，文章使用了掩码语言模型方法，选出 15% 的 token 作为噪声，其中 80% 的 token 将随机地被 [MASK] 代替，10% 随机替换为其他 token，其他 10% 不变。通过这种方法可以强迫模型不依赖于 token，而是依赖于上下文做出判断，增强模型的稳健性。

数据

来自 PubChem 的 1.1 亿个分子和来自 ZINC 的 10 亿个分子构成了训练数据，所有分子经过 RDKit 规范化后用于分词，总共得到 2357 个不重复的 token，设置 202 token 为模型输出的最长序列。

结果

MoLFormer 作为一种通用的预训练模型，在经过大数据集的训练后，再使用不同的标准分子数据集微调模型，完成 MoleculeNet 中的分类与回归任务，并将其结果与多种基线模型的表现对比，验证 MoLFormer 的编码效果。在 AUC-ROC 数据中，MoLFormer 在 6 组数据集上获得了 3 组最优表现，显著优于其他基线模型。

除此之外，文章还对比了 MoLFormer 和与之类似的预训练模型 ChemBerta 编码分子的能力。文章从 PubChem 数据集中随机挑选 10000 对分子，计算每对分子指纹相似性与预训练模型得到的词嵌入相似性之间的相关系数，计算每对分子最大公共子图（Maximum Common Subgraph, MCS）中原子数量与两词嵌入欧氏距离之间的相关系数。

实验结果中，在两个对比的指标上，MoLFormer 都具有更大的相关性，说明通过它得到的词嵌入能够更好地反映分子结构特征，即两分子的分子指纹相似性越高，表示两分子的向量也越相似；两分子的共有结构越大，表示两分子向量之间的距离也更小。

最后，文章分析了 MoLFormer 学习到的注意力矩阵。随机选取分子，绘制各原子间的化学键连接矩阵和 3D 距离矩阵，与相应的 MoLFormer 全注意力矩阵和线性注意力矩阵对比，可以看出，线性注意力矩阵与分子的化学键连接矩阵和 3D 距离矩阵具有部分相似的权重，这可能解释了 MoLFormer 编码后为什么能保留分子的结构特征。

结论

文章使用大规模数据训练了基于 Transformer 的化学语言模型，经过实验对比，这种通用的预训练模型可以很好地通过微调迁移到其他数据集上完成包括分类、回归在内的任务，并且取得了优异的表现。此外，文章还确定了通过该模型得到的词嵌入能够反映分子的结构，同时比以往类似模型具有更好效果。

文章主要对 Transformer 中的注意力机制做了改造，使用线性注意力减少计算开销，大大加快了训练速度，因此才能在 11 亿个分子的大规模数据上完成学习过程。文章还使用 RoPE 提高模型识别上下文的效果，这可能也是该模型具有更优表现的原因。目前预训练的通用语言模型如 GPT 和 Bart 等在 NLP 领域大放异彩，而对于这一类化学语言模型的研究还很少，它具有助力完成下游分子预测等任务的潜力。