大规模的化学语言 transformer 模型捕捉分子结构和性质

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

最近，在大型无标签语料库上预训练的无监督基于 transformer 的语言模型在许多下游自然语言处理任务中产生了最先进的结果。

受此启发，来自美国 IBM 研究院的研究人员提出了通过训练高效的 transformer 编码器模型 MOLFORMER 获得的分子嵌入，该模型使用旋转位置嵌入。该模型对来自 PubChem 和 ZINC 数据集的 11 亿个未标记分子的 SMILES 序列采用线性注意机制，并结合高度分布式训练。

研究表明，在来自十个基准数据集的几个下游任务上，学习到的分子表示优于现有基线，包括监督和自我监督的图神经网络和语言模型。进一步的分析，特别是通过注意力的镜头，表明经过化学 SMILES 训练的 MOLFORMER 确实学习了分子内原子之间的空间关系。这些结果提供了令人鼓舞的证据，表明大规模分子语言模型可以捕获足够的化学和结构信息来预测各种不同的分子特性，包括量子化学特性。

该研究以「Large-scale chemical language representations capture molecular structure and properties」为题，于 2022 年 12 月 21 日发布在《Nature Machine Intelligence》上。

机器学习 (ML) 已成为一种有吸引力的、计算效率高的预测分子特性的方法，对药物发现和材料工程具有重要意义。分子的 ML 模型可以直接在预定义的化学描述符上进行训练。然而，最近的 ML 模型侧重于从编码连接信息的自然图或分子结构的线注释中自动学习特征，例如流行的 SMILES 表示。

SMILES 学习已被广泛用于分子特性预测。然而，SMILES 语法复杂且有限制；适当字符集上的大多数序列不属于明确定义的分子。

用于分子特性预测的 GNN 和语言模型的监督训练面临的一个挑战是标记数据的稀缺性。分子的标签注释通常很昂贵，而且由需要注释的似是而非的化学物质组成的空间的大小是天文数字（10^60 到 10^100），这一事实使这个问题更加复杂。

这种情况产生了对分子表示学习的需求，这种学习可以推广到非/自我监督环境中的各种属性预测任务。基于大型 transformer 基础模型的成功，使用学习任务不可知语言表示的范例，通过对大型未标记语料库进行预训练并随后将其用于对感兴趣的下游任务进行微调，已扩展到其他领域。

用于预测分子特性的预训练语言模型和 GNN 最近才开始出现。然而，在数十亿个分子的大型语料库上训练的预训练语言模型在多大程度上能够捕获各种下游任务中的分子-属性关系仍未得到探索。

在此，研究人员提出了称为 MOLFORMER（分子语言 transformer）的分子 SMILES transformer 模型。将性能最佳的 MOLFORMER 变体命名为 MOLFORMER-XL。MOLFORMER-XL 是使用在 11 亿个分子的大型语料库上训练的有效线性注意机制获得的。

结果表明，分子 SMILES 的预训练 transformer 编码器在预测各种分子特性（包括量子力学特性）方面与现有的监督或无监督语言模型和 GNN 基线相比具有竞争力。

主要贡献如下：

• 研究人员在超过十亿个分子上训练了一个大规模高效的分子语言模型 transformer (MOLFORMER)，硬件资源相对有限（最多 16 个 V100 图形处理单元 (GPU)）。可扩展性和加速归功于高效的线性时间注意力、批处理的自适应分桶（ bucketing）以及 PyTorch Lightning 和 NCCL 中提供的开源并行化。通过结合分桶和线性注意力，能够实现每个 GPU 1,600 个分子的批量大小。使用 16 个 GPU，需要 208 小时才能完成 MOLFORMER-XL 的四个预训练阶段。要在没有分桶和线性注意力的情况下在相同的时间内完成训练，将限制在每个 GPU 少于 50 个分子，并且需要超过 1,000 个 GPU 来完成该任务。
• 探讨了表示分子 SMILES 时绝对位置嵌入和相对位置嵌入之间的差异。还为最近提出的相对位置 RoFormer 提供了一种新的、高效且准确的线性注意力近似。
• 对来自十个基准数据集的几个分类和回归任务进行了广泛的实验和消融研究，涵盖了来自 MoleculeNet 的小分子化学品的量子力学、物理、生物物理学和生理学特性预测。
• 结果提供了令人鼓舞的证据，表明 MOLFORMER 表示可以准确地捕获足够的化学和结构信息来预测各种化学性质。此外，MOLFORMER 的性能优于或相当于最先进的 GNN，这些 GNN 从精确的图形拓扑信息和其他信息（例如，键距离）中学习。
• 提供了进一步的分析，以证明 MOLFORMER 可以仅从 SMILES 注释中捕获子结构以及分子内的空间原子间距离。

本研究探讨了预训练化学语言模型在预测从量子化学到生理学的广泛下游分子特性方面的代表性能力。特别是，单独从 SMILES 字符串预测量子化学特性并非易事，因为这些特性在很大程度上取决于准确的三维 (3D) 分子几何信息，这些信息被认为是特权信息，通常不可用具体而言，MOLFORMER 在各种分子回归和分类基准上优于现有的基于图形的基线。这项工作验证了大规模自监督预训练分子语言模型在预测从量子化学到生理学的整个范围内的分子特性方面的能力。此外，通过分析学习到的注意力，表明在 SMILES 序列上训练的 MOLFORMER 确实知道分子内的原子间关系，甚至超出了二维拓扑。

最后，在大规模学习端，展示了 MOLFORMER 对计算资源的高效和环保使用，将执行训练所需的 GPU 数量减少了 60 倍（1,000 对 16）。

MOLFORMER 具有在不同靶标上更快地进行分子计算机筛选的直接潜力，这对材料设计和药物发现应用具有积极的社会影响。然而，应该注意的是，在湿实验室中未经适当的实验和科学验证而滥用此类技术可能会产生有害影响。

此外，目前的工作需要进一步探索 MOLFORMER 在其直接从化学语言中学习结构分子信息的能力的背景下的表征能力，并且可以扩展到本工作中研究的有机小分子之外。

未来的工作还将致力于通过采用更大的模型和更多的训练数据、使用改进的和/或特定领域的自监督任务以及使用其他基于字符串的表示（例如 SELFIES）来改进 MOLFORMER。