NMI 2022 | MoLFormer: 大规模化学语言表征能够捕捉分子结构与性质

今天介绍的这项工作来自 Nature Machine Intelligence。在标注分子数据稀缺、化学空间极其庞大的前提下，如何在缺少标签的情况下学到通用、可迁移的分子表征，是分子性质预测和药物发现的关键难题。该研究团队提出了基于分子语言建模的自监督框架 MoLFormer：把SMILES当作“化学语言”，在约 11 亿个无标注分子 （PubChem+ZINC） 上做掩码语言模型预训练。模型采用 Transformer 编码器，线性注意力和旋转位置编码，在有限算力下高效训练，预训练后通过冻结或微调编码器，迁移到多种下游分子性质预测任务。实验结果显示，MoLFormer在 MoleculeNet 上的多项分类与回归任务中整体表现优于基线模型。进一步的分子相似性度量与注意力权重分析表明，MoLFormer 学到的表征不仅能够刻画分子骨架，还对一定程度的空间关系具有敏感性。通过大规模自监督预训练的化学语言模型，可以获得兼具通用性、结构感知能力与任务可迁移性的分子表征，为后续的虚拟筛选、药物发现研究提供了一条具有代表性的技术路线。

获取详情及资源:

0 摘要

基于机器学习的模型可以实现快速而精确的分子性质预测，这对药物发现和材料设计具有重要意义。尽管已有多种监督学习模型取得了可观表现，但由于化学空间极其庞大、性质标注数据又十分有限，仅依赖监督学习依然面临巨大挑战。近几年，在大规模无标注语料上预训练的Transformer语言模型迅速发展，并在多种下游自然语言处理任务中取得了领先性能。受此启发，该研究提出了一种通过高效Transformer编码器模型MoLFormer学习分子表征的方法。MoLFormer采用旋转位置编码（rotary positional embeddings） 与线性注意力机制，并在PubChem与ZINC数据集中共约11亿个未标注分子的SMILES序列上进行高度并行的预训练。结果表明，相比现有的多种基线模型（包括监督和自监督的图神经网络以及分子语言模型），MoLFormer学到的分子表示在十个基准数据集中的多个下游任务上表现更优，在另外两个任务上也具备竞争力。进一步的分析，尤其是基于注意力权重的可视化，显示MoLFormer在化学SMILES上训练后，能够隐式地学习分子内部原子之间的空间关系。这些结果表明，大规模分子语言模型可以有效捕捉分子的化学与结构息，从而支持包括量子化学性质在内的多种分子性质预测。

1 引言

基于机器学习的模型，为高效预测分子性质提供了一条有前景且计算友好的途径，在药物发现和材料工程中具有重要应用。针对分子构建机器学习模型时，一种做法是直接基于预定义的化学描述符进行训练，例如无监督的分子指纹，或者由几何特征（如Coulomb矩阵）人工构造的派生特征。更近一步的工作则聚焦于自动学习特征：要么从天然的分子图（编码原子间连接关系），要么从分子结构的线性标记形式中学习，比如广泛使用的SMILES表征。SMILES通过对分子图进行深度优先的先序遍历，对每个原子、键、遍历决策以及断开的环路生成符号，因此最终得到的字符序列可以看作分子图某个生成树的“展开形式”。

在分子性质预测中，基于SMILES的学习已经得到广泛应用，因为相比图结构等其他表征方式，SMILES通常更加紧凑。此外，分子中的重要子结构，例如支链、环结构以及手性信息，在SMILES串中都有明确编码，而在图表示中则并非总是显式可见。然而，SMILES的语法复杂且高度约束——在给定字符表的所有可能序列中，绝大多数都不对应合法分子。也因此，出现了多种替代的字符串表征方式，例如SMARTS和SELFIES。相比SMILES，这些替代表征各自的优劣仍是一个活跃的研究方向。例如，有研究在学习到的表征空间上进行分子优化时发现，相比SELFIES，SMILES在优化能力和采样效率上并没有明显劣势，尤其是在语言模型足够强大的情况下。

尽管如此，字符串表征通常被认为不具备显式拓扑感知能力，而图结构则具备。基于这一点，一些工作认为深度化学语言模型可能主要在学习“分子字符串的语法”，而不是分子图隐含的拓扑结构。因此，虽然基于字符串的深度神经网络已被用于分子性质预测，但在大多数任务中，它们通常被图神经网络及其变体所超越。一般而言，图神经网络可以被视为一种“消息传递”框架：在图的连接结构约束下，在节点、边或整体图等不同粒度间进行局部邻域信息的聚合与更新。

对于用于分子性质预测的GNN或语言模型，监督训练面临的一个核心挑战是标注数据匮乏。分子标签的获取通常成本高昂，而需要标注的合理化学空间规模却是天文量级（约10⁶⁰–10¹⁰⁰）。在这种场景下，非常需要一种能够在无/自监督设置下学习到可泛化的分子表征，用于多种性质预测任务。

随着大规模Transformer基础模型的发展，一种通用范式逐渐成型：先在海量无标注语料上预训练“任务无关”的语言表征，再在具体下游任务上进行微调。这一思路也开始扩展到其他领域。针对分子性质预测，预训练语言模型和预训练GNN模型才刚刚出现。然而，在亿级分子上预训练得到的化学语言模型，究竟能在多大程度上捕捉到跨不同下游任务的“分子–性质”关系，仍缺乏系统探索。

沿着这一方向，该研究提出了基于分子SMILES的Transformer模型MoLFormer（molecular language transformer）。其中表现最优的变体被命名为MoLFormer-XL。MoLFormer-XL采用高效的线性注意力机制，在约11亿个分子构成的大规模语料上完成预训练（见Fig.1）。结果表明：在预测多种分子性质时（包括量子化学性质），基于SMILES的预训练Transformer编码器在多个基准数据集上，与现有的监督或无监督语言模型与GNN基线相比，整体表现具有竞争力，甚至在不少任务上达到更优水平。

该研究的主要贡献包括：

  1. 训练了一个大规模而高效的分子语言Transformer模型MoLFormer，在超过10亿个分子上完成预训练，同时仅依赖相对有限的硬件资源（最多16块V100 GPU）。这一可扩展性与训练效率得益于高效的线性时间注意力机制、自适应分桶(batch bucketing)策略，以及PyTorch Lightning与NCCL提供的并行能力。通过结合分桶和线性注意力，每块GPU可以处理约1600个分子；在16块GPU上预训练MoLFormer-XL四个epoch约需208小时。如果没有分桶和线性注意力，要在同样时间内完成训练，每块GPU最多只能处理不到50个分子，总体则需要上千块GPU。

  2. 系统比较了绝对位置编码与相对位置编码在分子SMILES中的表现，并提出了一种高效且精确的线性注意力近似版本，用于实现近期提出的相对位置RoFormer结构。

  3. 在十个基准数据集上进行了大规模实验与消融研究，覆盖量子化学、物理、 生物物理和生理等多种小分子性质预测任务。

  4. 结果表明，MoLFormer学习到的分子表征能够准确捕捉充足的化学与结构信息，从而支持对多种性质的预测；其性能在多数任务上优于或不逊于那些直接利用精确图拓扑结构（甚至包括键长等3D几何信息）的先进GNN模型。

  5. 进一步分析还显示，仅基于SMILES序列训练的MoLFormer，能够从中自动学到分子子结构信息以及原子间的空间距离关系。

总体来看，该研究系统评估了预训练化学语言模型在从量子化学到生理层面的广泛分子性质预测任务上的表征能力。尤其是，仅凭SMILES字符串来预测量子化学性质是高度非平凡的，因为这类性质在本质上极度依赖精确的三维分子构型，而这些3D信息在通常设定中是不可直接获得的“特权信息”。

图1 | 展示了MoLFormer框架概览。 该模型基于Transformer神经网络结构，在自监督范式下使用来自PubChem和ZINC两个公共化学数据库的大规模分子SMILES序列进行训练。MoLFormer采用高效的线性注意力机制与相对位置编码，旨在学习紧凑而有意义的化学分子表征。随后，这一基础模型通过在任务特定数据上的微调，被适配到不同的分子性质预测任务中。模型表征能力还通过两个方面进一步检验：一是利用MoLFormer编码恢复分子相似性，二是分析给定分子中原子间空间距离与注意力权重之间的对应关系。

2 结果

2.1 MoLFormer框架概述

MoLFormer的目标是从大规模化学SMILES数据中学习通用分子表征，并在多种下游分子性质预测任务上进行评估，如图1所示。为此，MoLFormer基于掩码语言模型（masked language model）框架：在训练过程中，随机遮蔽SMILES序列中一定比例的token，然后让模型去预测这些被遮蔽的token。通过这种掩码语言建模方式，模型可以在自监督设定下学习上下文信息。

为了获得更好的上下文建模能力并提升训练效率，MoLFormer采用了旋转位置编码（rotary positional embedding）替代绝对位置编码，同时结合线性注意力机制（linear attention）（更多模型结构与训练细节见 Methods 和补充材料）。训练过程中可以观察到，相比绝对位置编码，使用旋转位置编码在预训练阶段带来了更稳定的训练表现和更快的loss收敛（见图2）。

为了验证预训练好的 MoLFormer 是否可以作为一种通用、与具体任务无关的分子表征方式，研究团队在 MoleculeNet 中的多种具有挑战性的分类与回归任务上，对其在下游性质预测中的迁移表现进行了系统评估。各基准数据集的详细信息见补充材料中的相关章节。

图2 | 展示了绝对位置编码与旋转位置编码在训练与验证损失上的对比。 a,b，展示了在线性注意力MoLFormer中分别采用旋转（相对）位置编码和绝对位置编码时，在PubChem数据集上的训练损失（a）和验证损失（b）。可以看到，两种位置编码下的MoLFormer都呈现平滑的训练曲线，但采用旋转位置编码的线性注意力MoLFormer在训练损失和验证损失上都低于采用绝对位置编码的MoLFormer。

2.2 MoLFormer表征的获取

在MoLFormer中，每个分子的SMILES序列会先经过编码器，其最后一层隐状态中所有token对应的向量取平均，作为该分子的整体分子表征，该向量随后被统一用于所有下游性质预测任务。下游任务的使用方式可以分为两种：冻结（frozen）和微调（fine-tuned）。在冻结设置中，MoLFormer编码器产生的分子表征保持固定，仅在其上方单独训练一个全连接网络（MLP）来完成具体任务，而编码器参数不再更新；该策略下的模型结构与超参数通过网格搜索确定，具体搜索空间见Supplementary Table 1。在微调设置中，则将编码器与任务网络一起端到端训练：在编码器输出之上叠加一个两层全连接网络，隐藏层维度为768（与编码器输出维度一致），中间包含dropout（设为0.1）和GELU激活层；对于回归任务，最终输出层为单一神经元，用于预测目标分子性质。

2.3 MoLFormer表征在下游任务上的表现

该研究团队在MoleculeNet基准28提供的6个分类任务和5个回归任务上系统评估了MoLFormer分子表征的性能，并与已有基线方法进行了对比。研究中将基于约1.1亿分子（来自PubChem和ZINC全部分子）完成预训练的模型记为MoLFormer-XL。除非特别说明，MoLFormer-XL均采用带线性注意力机制和rotary位置编码的配置进行预训练，文中汇报的结果默认是在下游任务上对该模型进行微调后的性能。在下游性质预测任务中，该研究团队按照前一节所述的方式对MoLFormer进行微调；所有任务的训练集、验证集与测试集划分均严格遵循MoleculeNet基准的官方设置。

2.3.1 分类任务

该研究团队从MoleculeNet基准中选取了6个分类任务，并与9种基线方法进行了对比，其中包括4种监督学习方法和5种自监督/预训练方法。监督学习基线包括基于分子指纹的浅层机器学习模型（如表1中的随机森林RF和支持向量机SVM）以及若干图神经网络模型。在预训练/自监督基线中，一类方法基于图同构网络（GIN）在分子图上进行预训练，该网络在聚合过程中显式利用边特征，并通过多层感知机和加权节点特征求和实现消息传递；N-gram Graph是一种简单的无监督分子表示方法，先对图中节点进行嵌入，再通过在图上进行短游走，将顶点嵌入组合为紧凑的图表示；MolCLR是基于图同构网络的自监督对比学习框架，使用对比学习目标来区分正负样本对；GraphMVP-C则通过联合利用二维拓扑结构与三维几何视图之间的一致性与对应关系进行自监督预训练。此外，还考虑了三种显式利用几何信息的GNN基线模型，其中一种为监督式模型DimeNet，另外两种为自监督方法GeomGCL和GEM；ChemBERTa则是预训练于较小化学数据集上的分子语言模型。表1展示了在这六个分类基准任务上，在MoleculeNet设定下MoLFormer与上述各类基线方法的性能对比结果。

表1 | 展示了在多个分类基准上对比经微调的MoLFormer与现有监督学习及预训练/自监督基线方法的性能。

2.3.2 回归任务

该研究团队在 MoleculeNet 中更具挑战性的回归任务上评估了 MoLFormer-XL 的表现。在 QM9、QM8、ESOL、FreeSolv 和 Lipophilicity 五个回归基准上给出了结果（见表2）。其中，QM9 和 QM8 需要预测多种量子化学性质，在没有加入特殊三维几何信息的前提下，这类任务被认为非常具有挑战性。在所有任务中，该研究团队均采用 MoleculeNet 推荐的训练集、验证集和测试集划分方案。对比基线包括：分子图卷积网络 GC（一种图神经网络，先对节点及其邻居做均值池化再进行线性变换）、注意力指纹模型 A-FP，以及一种能够显式学习边特征（如原子对间距离）的 MPNN 变体。结果表明，在针对具体任务进行微调后，MoLFormer-XL 在这五个数据集上全部优于现有的监督式 GNN 基线模型 GC、A-FP 和 MPNN（在 QM8 和 QM9 上还加入了键/原子间距离作为额外几何信息）。补充表7进一步显示，在三个物理性质回归基准上，MoLFormer 也优于几何感知型GNN（如 DimeNet、GeomGCL 和 GEM）。结合分类任务上的结果，这些实验总体上证实了 MoLFormer-XL 在不同性质预测任务上的良好泛化能力。

表2 | 展示了在QM9、QM8、ESOL、FreeSolv和Lipophilicity回归基准上，微调后的MoLFormer与其他监督式GNN基线模型的性能对比结果。

2.3.3 QM9任务深入分析

补充表9进一步将MoLFormer-XL在QM9数据集中原子化能量和焓（即校正参考原子能量后的内能/焓，单位为eV）预测任务上的表现，和两个典型的三维监督式GNN模型SchNet和DimeNet进行了对比。仅基于SMILES训练的MoLFormer-XL在这四项任务上都不及这两种方法。不过，需要强调的是：这两种模型依托专门面向量子相互作用建模的架构，直接编码了三维结构信息，但相对于MoLFormer-XL，它们在预测精度上大约只优于8倍和10倍左右。结合表1和表2，这一结果一方面再次印证了：在大规模无标注分子库上，仅依赖SMILES等易获取信息学习通用分子表征具有相当强的能力；另一方面也明确了三维几何这一“特权信息”对于量子化学能量预测依然至关重要。此外，这一对比结果也为后续研究打开了新的方向：例如，如何系统评估MoLFormer在仅用SMILES时是否会“自发”产生几何感知能力，或者探讨在保留SMILES输入的前提下，引入部分或完整的三维几何信息，进一步增强MoLFormer的表达能力。

2.3.4 MoLFormer消融实验

研究团队围绕MoLFormer-XL进行多种消融分析，以帮助理解其优异性能背后的原因。整体上，这些消融设置大致可分为三类：(1) 预训练数据规模与性质、以及模型深度对性能的影响；(2) 下游任务中不进行微调（冻结特征）与进行微调（端到端调整）两种设置的对比；(3) 绝对位置编码与旋转位置编码（rotary positional embeddings）在模型中的作用差异。

2.3.5 MoLFormer数据与模型规模

该研究考察了预训练数据集规模对MoLFormer-XL在MoleculeNet若干下游任务上的影响。为此，选取了来自PubChem和ZINC的三种不同加权组合：一种只使用10% ZINC和10% PubChem的混合数据；一种使用100% PubChem加10% ZINC；以及一种只使用100% ZINC、0% PubChem的数据集。同时，该研究还通过在完整的ZINC+PubChem数据上预训练一个6层模型（记为MoLFormer-Base）来分析模型深度的影响。所有模型均采用旋转位置编码与线性注意力进行预训练，并与MoLFormer-XL进行对比；预训练与微调过程中使用相同的学习率、数据划分和优化策略等。扩展数据表1和2给出了这些实验结果。

总体上，MoLFormer-XL在平均性能上更优，但有两点现象值得注意：第一，只在第二大规模数据集（100% ZINC）上预训练的模型，性能持续低于其他所有预训练模型。可能的原因在于：相较于其他组合，ZINC数据集的词表明显更小，分子长度更短且长度方差也更小，导致模型可学习的结构多样性不足。第二，当MoLFormer-XL在某些任务上略逊时（例如在ESOL、QM8和FreeSolv上的表现，见表2），其差距通常非常微小。扩展数据表1和2还显示，在大多数任务中，MoLFormer-Base的表现弱于MoLFormer-XL，这表明更深的模型有助于学习更强的分子表征。

2.3.6 MoLFormer微调与冻结

扩展数据表3汇总了基于QM9基准的两组消融实验。简单来说，在所有预训练数据规模下，引入微调的实验都显著优于冻结编码器(frozen)的设定，因此该研究团队在其他基准任务中只重点考察了微调方案。这些结果从实证角度揭示了MoLFormer在模型规模与数据规模共同扩展时的表现规律。

表3 | 展示了基于QM9测试集中7,806个分子，在三类不同原子间距离区间下，对比MoLFormer模型的原子间空间距离矩阵与注意力矩阵之间余弦相似度的结果。

2.3.7 MoLFormer位置编码

位置编码的消融结果同样汇总于扩展数据表3。总体来看，当预训练数据集规模较小时，采用绝对位置编码的模型略优于使用旋转位置编码(rotary embeddings)且进行微调的MoLFormer；但一旦预训练分子数量超过10亿级，基于旋转位置编码并进行微调的MoLFormer开始反超绝对位置编码模型，表现更佳。

2.4 深入探究MoLFormer

2.4.1 MoLFormer分子相似性恢复

该研究团队进一步从“分子相似性恢复”的角度分析了MoLFormer的表现。具体而言，他们比较了两类相似性度量之间的相关性：一类是基于分子指纹、利用Tanimoto距离计算得到的传统化学相似性；另一类是基于MoLFormer-XL嵌入向量之间欧氏距离得到的相似性。同时，研究还考察了从PubChem中随机选取的一批分子，对每一对分子计算其最大公共子图所包含的原子数，并将这一结构相似性指标与对应嵌入空间中的欧氏距离进行相关性分析。扩展数据表4的结果表明，与ChemBERTa相比，MoLFormer-XL的嵌入与上述已知分子相似性度量之间具有更高的相关性，说明MoLFormer学到的分子表征能够有效反映化学结构层面的相似性，对分子相似性建模具有较强的表达能力。

2.4.2 MoLFormer注意力分析

该研究团队最后对MoLFormer-XL的注意力矩阵进行了分析，以探究其中所蕴含的化学信息。具体来说，他们利用QM9测试集中分子的空间几何结构，计算分子内原子两两之间的空间距离，并将其与模型注意力值之间的余弦相似度进行比较。空间距离来源于QM9基准数据集中提供的能量最小化几何构型。对比对象包括MoLFormer-XL以及一个在完整PubChem + ZINC数据集上、采用全注意力（full attention）和旋转位置编码（rotary positional embeddings）训练得到的MoLFormer变体。需要强调的是，这些MoLFormer模型在分析时并未在QM9数据集上进行微调。冻结参数的全注意力MoLFormer在QM9下游任务上的平均绝对误差（MAE）明显更高（MAE ≥ 12），尤其是在内部能量（U和U0）、焓（H）和自由能（G）等任务上表现较差。

在分析中，该研究团队将原子对之间的空间距离分为三类：短距离（≤2 Å，主要对应典型的共价键，譬如C–C单键约为1.5 Å）、中等距离（2–4 Å）以及长距离（≥4 Å），并在表3中分别汇总了对应的注意力与距离相似性。结果显示，无论是采用线性注意力还是全注意力、并结合旋转位置编码的MoLFormer，其注意力模式在短距离和中等距离范围内都与真实的原子间距离具有较高相似度，而在长距离范围内相似性较弱（约0.2）。这一现象表明，MoLFormer能够在一定程度上捕捉分子内部原子之间的空间关系，即便这些原子在SMILES序列中并非相邻。与全注意力版本相比，MoLFormer-XL（线性注意力）在中长程（medium和long-range）距离上的相似性略高，说明线性注意力版本在建模原子间空间关系方面反而更加有效。

图3进一步从可视化角度支撑了这一结论：该图展示了MoLFormer-XL在某一中间注意力层中的平均注意力系数，并将不同原子对之间的注意力与对应的共价键连通关系以及三维空间距离进行对比（同一分子在所有层上的完整注意力矩阵见补充图5和图6）。研究团队从QM9测试集中选取了两个分子，这两个分子的注意力模式与其中程空间距离具有较高的余弦相似度。可视化结果表明，经过不同注意力头聚合后的中间层旋转注意力，不仅能够很好地对应分子中的共价键连接模式，还在一定程度上反映了分子中“非键合”原子之间的空间关系特征。

综合这些注意力分析结果，可以看出MoLFormer-XL在相当程度上能够仅从SMILES序列中恢复分子的结构信息。这种能力很可能源于其在大规模化学SMILES语料上的预训练，使得模型可以学习到关于分子结构的基本规律，并据此捕捉包括量子化学性质和生理性质在内的多种下游分子性质。类似的现象也曾在蛋白序列建模的相关研究中被报道，这进一步印证了：在大规模数据上预训练的化学语言模型，其内部学到的表示中会自然涌现出与结构及多种性质相关的丰富信息。

图3 | 展示了在旋转位置编码（rotary embedding）下，分别采用全注意力或线性注意力的MoLFormer，对两个随机分子“CC1(C)C(C)(O)C1(C)O”(a)和“CC(C)C(C)(C)O”(b)所学习到的注意力图以及对应的分子结构（键连接关系和单位为Å的三维距离）的可视化。 注意力图（取值范围0到1，仅为清晰起见保留映射到组成原子的token），由某一中间注意力层的多头注意力平均池化得到，体现了对共价键连接模式以及原子间长程空间关系的感知。其中，线性注意力版本对于中程三维距离的刻画（红色圆圈标出）优于全注意力版本。

3 总结

该研究系统地探索了大规模无监督预训练分子语言模型在多种分子性质预测任务中的能力。与图结构表示不同，SMILES等分子语言并不显式编码分子拓扑结构。但在设计合理的自监督训练目标、结合大规模分子语料，以及采用具有强表达能力的架构（如带线性注意力机制的上下文Transformer语言模型）和高效并行训练策略的前提下，MoLFormer能够高效地学习到隐式的、富含信息的结构–性质关联。

具体而言，MoLFormer在多种分子回归与分类基准任务上，整体上优于现有的图神经网络等图结构基线模型。该研究表明，大规模自监督预训练的分子语言模型在从量子化学性质到生理性质的广泛分子性质预测任务中具有强大的适用性。此外，通过对模型注意力权重的分析可以看到，仅基于SMILES序列训练得到的MoLFormer，确实能够感知分子内部的原子间关系，甚至超越传统二维拓扑层面。

在大规模训练方面，MoLFormer还展示了对算力资源更高效、更环保的使用方式：在相同训练量级下，相比朴素方案，所需GPU数量从约1000块降至16块，减少了约60倍的资源消耗。MoLFormer在多靶点的大规模体外虚拟筛选(in silico screening)中具有直接应用潜力，这对于材料设计和药物发现具有积极的社会意义。同时，该研究也强调了潜在风险：如果缺乏严谨的实验与科学验证流程，这类技术的不当使用可能带来有害后果。已有研究表明，将高精度性质预测模型（如用于毒性预测）与生成模型结合，可能被利用来设计高毒性分子。这凸显了在使用此类强大技术时，亟需配套的负责任使用规范与伦理框架。

此外，该研究也为进一步探索MoLFormer的表征能力提出了方向，尤其是其能否仅从化学“语言”中更充分地学习分子结构信息，以及如何将方法扩展到本工作之外的更大、更复杂的分子体系。未来工作还可从多个方面改进MoLFormer，包括使用更大规模模型与更丰富的训练数据，引入更优或更具领域特色的自监督任务，以及探索例如SELFIES等其他字符串表示形式。