复旦大学王满宁教授团队提出MoleSG模型，通过非重叠掩模的互补多模态自监督学习进行分子性质预测

自监督学习在分子表征学习中起着重要的作用，因为标记的分子数据通常在许多任务中受到限制，例如化学性质预测和虚拟筛选。然而，现有的分子预训练方法大多集中在分子数据的一种模态上，并且没有充分挖掘两种重要模态(SMILES和图)的互补信息。

近年来，越来越多的研究表明，多模态预训练包含更多的信息，往往能达到更好的效果，因此具有更大的潜力。然而，大多数现有的多模态融合方法往往依赖于对比学习，但它们忽略了不同模式之间的细粒度交互。尽管UniMAP等多模态融合框架采用基于掩码重建的生成式预训练，但其只进行简单的掩码重建，没有具体设计掩码策略，因此仍然不能充分利用互补信息交互。

2024年5月27日，复旦大学王满宁教授团队在Briefings in Bioinformatics上发表文章Complementary multi-modality molecular self-supervised learning via non-overlapping masking for property prediction。

作者提出了一个有效的多模态自监督学习框架，结合SMILES和图（graph）的信息，因此称为MoleSG。具体而言，MoleSG首先对SMILES数据和图数据进行标记，以便由基于Transformer的统一骨干网络进行处理，通过掩膜重建策略进行训练。此外，作者还引入了一种专门的非重叠掩膜策略，以引导这两种模式之间的细粒度交互。实验结果表明，MoleSG在一系列分子性质预测任务中超越了现有方法。

如图1所示，MoleSG通过统一的特征提取骨干网络对SMILES和图这两种模态进行掩模重构，共同学习特征。也就是说，在预训练过程中，随机掩膜一些图上的节点特征，然后掩膜一部分SMILES标记以对应于图中剩余的未掩模原子，使模型可以执行非重叠掩模，以方便两个模态之间的信息交互。

图1 MoleSG结构图

在预训练期间，作者采用对称联合编码器-解码器框架进行进一步的特征提取。该框架由用于两种模式的两个独立分支和用于特征融合的共享主干组成。独立的编码器分支将两种不同模式的数据编码成统一的形式，这适合于基于Transformer的序列理解，可以学习模态内部和模态之间的依赖关系，并为随后的独立解码器学习输出特征。最后，基于骨干网的输出，对原始的SMILES序列和图形进行重构。在微调过程中，作者利用预训练的图编码器作为分子表示网络，并添加相应的输出头来预测一系列分子性质。

为了促进跨不同模式的细粒度特征的交互，作者使用两个独立的编码器将两个完全不同模式的数据转换为相同维度的嵌入，以便由Transformer模型进一步处理。对于SMILES序列，作者按照ChemBERTa首先转换被掩膜的SMILES令牌，通过对数据集中的所有标记进行全面分析来扩展其词汇表。然后，采用RoBERTa中使用的一系列多头注意块的Transformer模型计算它们对应的嵌入。对于图，作者采用了CoMPT作为图编码器，它通过相关核加强了节点和边之间的消息交互。经过图编码器处理后，模型得到了标记嵌入。

作者设计了一个基于Transformer的统一骨干网络，以促进两种模式之间的特征交互。通过Transformer的注意力机制，模型能够学习跨两种模态的不同输入标记嵌入之间的相关性。在处理了两个模态特定的编码器之后，模型添加了可训练的参数，然后把两种模态的嵌入连接起来，形成统一的骨干识别模态，连接嵌入并送入骨干网络。在这里，作者使用RoBERTa的Transformer模型作为骨干网络，其具有一系列多头自注意力块，可以促进同一模态内和不同模态的令牌嵌入之间的信息交互。在主干中提取特征后，对输出特征可以拆分为SMILES特征和图特征，用于对应模态的掩码重建任务。SMILES解码器是RoBERTa中的LMhead（语言模型头），以预测掩码令牌，而图解码器是一个图同构网络GIN，重建被掩码节点特征。整个任务的总体损失如下是以上两种模态的重建损失之和。

SMILES标记可分为三类：(1)原子，包括C和N等单字符原子，以及Ca和Au等多字符原子，以及[Cl-]和[Fe+3]等离子；(2)化学键，例如双键用“=”表示；(3)其他符号，如数字“1”和“2”表示环中原子的位置，括号“(”和“)”表示含侧链。考虑到在SMILES中经常省略单键，实现化学键的两种模式之间的一对一对应是不现实的。因此，在本文中，重点是对准原子索引。因此，作者收集表示原子的令牌，并为它们分配索引，以建立图中的原子与过滤后的SMILES令牌中的原子之间的一致对应关系。接着随机掩膜图上的原子特征和SMILES序列上的原子标记。为了引导两种模态之间更好的相互作用，将两种模态中掩蔽原子之间的重叠比设置为0。具体来说，基于原子索引的一对一对应关系，将图中被遮挡原子的位置定位到SMILES序列上，从而避免了在SMILES序列上掩盖已经在图上被掩盖的原子。作者使用AdamW优化器训练MoleSG进行迭代，将图的掩码比率设置为25%，SMILES的掩码比率设置为15%。

在预训练阶段，作者从ZINC15中采样了25万个未标记的分子，ZINC15是一个用于药物发现和计算化学研究的化合物的综合数据集。在微调阶段，作者利用了来自MoleculeNet的多个基准数据集，涵盖了各个领域的分子数据。作者以8:1:1的比例将每个基准数据集划分为训练集、验证集和测试集，在结果中报告每个基准的三个随机种子结果的均值和标准差。

作者将MoleSG与一些具有代表性的方法进行了比较，如表1所示。作者采用了生理学领域的五个数据集，以及生物物理领域的三个数据集，将ROC曲线下的面积（AUC）作为评测指标。MoleSG在以上8个数据集上的性能均超越了现有方法。

表1 与其他方法对比

作者设计了消融实验来验证模型设计的有效性。为了进一步揭示多模态预训练方法的优势，作者将多模态预训练与单模态预训练进行对比，结果如表2所示。在这个实验中，每个编码器都增加了一个输出头。“SMILES scratch”和“Graph scratch”代表了两个从头开始训练的网络。“SMILES预训练”和“Graph预训练”中编码器的初始权值是使用相同的MoleSG框架进行单模态预训练，同时不使用另一模态得到的。“Ours SMILES”和“Ours graph”中编码器的初始权值来自多模态预训练MoleSG的相应编码器。从这些结果可以看出，作者提出的方法在所有下游任务上都达到了最佳性能。此外，值得注意的是，单一模式的预训练可能会导致性能下降。然而，通过充分利用不同模态之间的互补信息，MoleSG可以提高所有下游任务的性能，显示出更大的实际应用潜力。

表2 消融实验

作者还进行了案例分析。MoleSG特征提取能力的可视化结果如图2所示，可以看出MoleSG在BBBP和BACE分类任务中具有很强的特征判别能力。作者将MoleSG与从头开始训练的模型(没有预训练)、单模态预训练的模型(即图预训练)和对比预训练的模型⁠进行对比。在微调期间，这些对比方法都使用图编码器。从图2中可以看出，与单模态预训练和对比预训练相比，MoleSG具有更好的特征辨别能力。

图2 案例分析

注意到，由于预训练与微调数据集的概率分布不一致，对编码器和骨干网络进行微调可能会导致性能损失。因此，作者进行了实验来验证这一假设。在这里只考虑两种组合。一种组合是对单个图编码器进行微调，而另一种组合是对图编码器和主干网络都进行微调。作者在所有基准测试上执行实验，结果如表3所示，仅使用图编码器在所有任务中都获得了更高的性能。

表3 不同微调策略对比

在本文中，作者解决了从两种互补的模式中学习细粒度信息的挑战：即SMILES和图。为了更好地从这两种模式之间的相互作用中捕获丰富的分子特征，作者设计了一个高效的多模态预训练框架MoleSG，该框架利用统一的特征处理网络融合这两种模式。此外，作者提出了一种非重叠掩模策略，以促进两种模式之间的信息交换。在下游任务上的大量实验表明，作者的方法达到了更好的性能。未来的工作有两个潜在的方向。首先，类似的多模态预训练方法可以用于蛋白质表示学习，因为蛋白质也同时具有序列表示和图表示。其次，非重叠掩膜策略可以扩展到其他多模态数据的联合预训练研究中。

参考文献

Shen et al. Complementary multi-modality molecular self-supervised learning via non-overlapping masking for property prediction. Brief Bioinform. 2024