Nat. Commun. | 通过单一分子基础模型实现结构和属性的双向生成

今天为大家介绍的是来自Jong Chul Ye团队的一篇论文。人工智能中基础模型的近期成功促使了大规模化学预训练模型的出现。尽管对于提供下游任务有用信息的大型分子预训练模型的兴趣日增，但在分子领域上对多模态预训练方法的尝试却很有限。为了解决这一问题，作者在此介绍了一个多模态分子预训练模型，该模型结合了结构和生化属性的模态，灵感来自于近期多模态学习技术的进展。作者提出的模型数据处理和训练目标在一个共同的嵌入空间中对齐了结构/属性特征，使得模型能够考虑分子结构和属性之间的双向信息。

捕获化学实体及其属性之间的复杂关系是许多化学挑战的本质。在过去的十年中，人工智能作为化学研究中估计许多生化属性和分子、聚合物和蛋白质之间的相互作用的有希望的工具而兴起，这些相互作用在实验上难以获得。化学领域的各种基于深度学习的方法采用深度神经网络来从原始分子数据中提取所需特征。此外，使用循环网络、变分自编码器、图网络等对全新分子设计进行了广泛研究。最近，为了克服以监督方式为每个任务单独学习特征的限制，科研人员已经提出了通过无监督学习方法更好地表示化学输入的方法。这些最新的方法与基础模型的概念走在同一条道路上，这些基础模型是用大型数据集训练的，并且通常被视为深度学习的一种新范式。

图 1

在这项工作中，作者对分子结构与其关联属性之间的跨模态理解感兴趣，这有助于解决许多应用中的有意义任务，如属性预测、条件分子设计等。相较于使用预备的属性作为标签来提取通用特征的多任务学习方法，作者的方法将一组属性视为代表输入分子的独立模态，并认为带有此属性模态的分子的多模态学习可以提供更多有信息量的特征。具体而言，作者提出了一种分子结构-属性多模态基础模型（SPMM），该模型允许进行各种计算机实验，它是用广泛的分子结构和其属性向量进行预训练的。通过采用Transformer架构，可以分别用自注意力和跨注意力机制完成模态内特征提取和模态间融合。实验结果显示，通过单一基础模型同时学习结构特征和来自关联属性的信息，为模型提供了一个更好的表示，该表示可以为各种下游任务进行微调。具体来说，通过对结构和属性进行对称处理，该模型能够用单一预训练模型进行双向生成和预测，这在以前是不可能的。

图1a展示了SPMM的整体模型架构和训练目标。SPMM的框架扩展了双流VLP模型的结构。双流VLP模型使用单模态编码器为每种模态编码输入，然后使用另一个编码器模块执行跨注意力操作，使用一种模态特征作为查询（query），另一种模态特征作为键（key）/值（value）。当给定一个训练分子时，SPMM将分子的SMILES字符串和其属性向量（PV）作为多模态数据输入，如图1a所示。SMILES和PV通过它们相应的单模态编码器传递，执行自注意力操作，其中嵌入的输入成为键、查询和值。在获得两个单模态特征后，对比学习通过同化包含相同上下文的特征，将SMILES和PV特征对齐到相同的嵌入空间。然后，编码的SMILES和PV特征通过融合编码器传递，该编码器在SMILES和PV特征之间执行跨注意力操作。这个单一的融合编码器可以使用其查询和键/值输入的交替来执行跨注意力操作，因为对比学习将SMILES编码器和PV编码器的输出对齐到相同的特征空间。融合编码器通过下一个单词预测（NWP）进行SMILES的预训练，下一个属性预测（NPP），和SMILES-PV匹配损失（SPM）。从给定的变换器输入预测下一个组件是一种常用的自监督学习目标，作者的NWP和NPP任务使模型学习SMILES令牌和属性之间的上下文关系，并借助另一种模态的语义特征。此外，SPM预测给定的SMILES和PV对是否代表同一分子。

一旦训练完成，SPMM可以用于需要理解SMILES和属性的各种双向下游任务，如属性预测（SMILES到属性）和基于属性的分子生成（属性到SMILES，也称为逆向QSAR）如图1b所示。此外，作者使用的预训练目标还允许预训练的SPMM应用于单模态任务，如分子分类和反应预测（见图1c）。预训练的SPMM在这些单模态任务中的表现与最先进的模型相当，这表明了该模型作为基础模型的泛化能力。

模型对SMILES和属性有双向理解

一旦SPMM被预训练，模型可以仅使用给定的PV输入生成SMILES，这对于许多化学任务，如全新分子设计，是一个关键挑战。作为药物发现的主要方法之一，已经提出了各种方法来生成具有所需属性的分子。在迄今为止提出的方法中，同时可控制的属性数量并不是很大。同时，输入属性向量的长度不能改变。每当目标属性发生变化时，都需要对模型进行重新训练以适应新的期望条件。与此相反，预训练的SPMM可以接受预训练中使用的53个属性作为输入条件，并生成满足所有这些属性的分子，无需针对每种属性组合进行单独的额外训练。此外，对于不想控制的属性，可以让模型忽略这些条件，通过使用预训练中使用的[UNK]标记来替换它们。这非常有用，因为在实践中控制所有53个输入属性并不是一个常见的场景，同时这也不容易，因为属性之间是相关和纠缠的。

图 2

为了展示SPMM的分子生成能力，作者准备了一系列的PV到SMILES生成场景，并让预训练的SPMM使用输入的属性自回归地生成SMILES。生成分子的有效性、唯一性和新颖性是SPMM分子生成能力的定量指标。此外，作为一种定性指标，以查看生成的SMILES与属性输入的匹配程度，作者测量了输入条件与生成分子属性之间的归一化均方根误差。对于第一个PV到SMILES生成场景，我们准备了来自PubChem的1000个SMILES的PV，这些PV不包含在预训练数据集中，并将它们输入到预训练的SPMM中以生成相应的SMILES。这里，采样过程以确定性方式进行：从SMILES的[CLS]标记（[CLS]S）开始，模型预测下一个标记的概率分布，并选择概率最高的选项。表1的第一行显示了其结果。在为1000个PV进行确定性PV到SMILES生成的输出中，99.5%的生成输出是有效的SMILES。53个归一化属性的平均RMSE为0.216，这意味着生成样本的属性与属性输入一致。像药物发现这样的应用领域经常需要为单一期望的目标属性条件生成多个分子。这可以通过从模型的概率分布中随机采样下一个标记来完成，而不是使用概率最高的标记。为了验证模型从单一PV输入生成多个分子的能力，作者在固定PV上用随机采样生成了1000个SMILES。图2显示了从单一PV输入生成的1000个分子的属性分布。每个属性分布的模式都落在输入属性值上（图2a）。在只给出一些属性的情况下，模型只考虑已知属性，而其他遮盖的属性不受限制（图2b, c）。即使完全没有属性信息，SPMM也能生成分子；当所有输入属性都被替换为[UNK]标记时（图2d），模型执行无条件的分子生成，输出遵循预训练数据集的分布。

图 3

上述结果表明，SPMM能够使用任意PV输入进行分子生成，这使得简单的分子设计和编辑成为可能。作为分子编辑的可能示例，图3包含了SPMM对五个PV输入进行随机分子生成的输出，这些输入都来源于分子1的PV，但其中四个进行了某些值的更改。生成的分子在保持未修改属性的同时，跟随输入的修改。SPMM甚至能够生成具有超出领域条件的分子，如“log P = 7”（注意预训练数据集中约有5%的log P > 7）。

图 4

随后作者将ZINC15数据集中的1000个SMILES（这些SMILES不包含在预训练数据集中）输入预训练的SPMM，并生成了它们对应的PV。图4是12个被选中的属性（我们用于预训练的53个属性中的一部分）的实际属性值与生成输出之间的散点图。很明显，SPMM预测的属性值与实际值非常接近，大部分数据点位于y = x线上。尽管由于50%随机属性遮盖的原因，模型在预训练期间实际上从未见过完整填充的PV，模型仍能自回归地预测全部53个属性。53个属性的平均r^2分数是0.924。

图 5

为了对迄今为止所呈现的预训练SPMM的性能进行解释，作者进一步分析了SMILES和属性向量之间学习到的跨模态理解，通过可视化预训练SPMM的注意力分数来实现。基于Transformer的模型具有直观的注意力可视化优势，显示了模型如何考虑输入查询和键之间的关系。在图5中绘制了当给定SMILES及其属性向量输入时，预训练的SPMM最后一个融合层的跨注意力分数。由于跨注意力有多个头部，作者取了它们注意力分数的平均值。有趣的是，跨注意力分数的方面遵循了化学属性和分子片段之间直观的关系。与氢键相关的属性（如‘NumHDonors’、‘NumHAcceptors’）显示出对含氧和氮原子的标记有高注意力分数。属性‘RingCount’关注与环有关的标记，而对侧基显示出弱的注意力，属性‘NumAromaticRings’只对芳香环的组成部分给出高注意力分数。当不同的SMILES标记在分子中扮演类似角色时，如分子7中的‘c1ccccc1)’和‘c1ccccc1’，它们的注意力模式也相似。这个结果表明，SPMM能够捕捉分子结构和化学属性之间的关系，而无需在它们之间明确给出监督。

编译 | 曾全晨

审稿 | 王建民

参考资料

Chang, J., Ye, J.C. Bidirectional generation of structure and properties through a single molecular foundation model. Nat Commun 15, 2323 (2024).

https://doi.org/10.1038/s41467-024-46440-3