【AI4S】通过单一基础模型双向生成分子的结构和特性

2024年3月14日，韩国科学技术院Jinho Chang等人在Nature Communications上发表文章Bidirectional generation of structure and properties through a single molecular foundation model。

作者提出了一个多模态分子预训练模型，该模型结合了分子结构和生化特性的模态（structure-property multi-modal，SPMM），在一个共同的嵌入空间中对齐结构-性质特征，这使得模型能够考虑分子结构和性质之间的双向信息。这些信息产生了协同知识，从而能够通过单一分子基础模型处理多样化的分子预测下游任务。实验表明模型性能超越了现有方法。

摘要

近期，人工智能领域中基础模型的突破性进展推动了大规模化学预训练模型的兴起。尽管人们对能够为下游任务提供丰富信息表征的大型分子预训练模型日益感兴趣，但在分子领域开展多模态预训练方法的尝试却相对有限。为此，我们提出了一种融合分子结构与生化特性两大模态的多模态分子预训练模型，并从近年来多模态学习技术的最新进展中汲取灵感。我们的模型设计了一套高效的数据处理流程及训练目标，旨在将分子的结构与属性特征映射到统一的嵌入空间中，从而实现分子结构与属性之间的双向信息交互。这种创新性的知识协同机制，使我们得以通过单一模型同时应对多模态和单模态的下游任务。通过广泛的实验验证，我们证明了该模型具备解决多种重要化学挑战的能力，包括条件分子生成、性质预测、分子分类以及反应预测等。

引言

捕捉化学实体与其性质之间的复杂关系，正是众多化学挑战的核心所在。在过去十年中，人工智能已成为化学研究中一种极具潜力的工具，能够高效预测分子、聚合物和蛋白质之间诸多难以通过实验直接获得的生物化学特性与相互作用。在化学领域，多种基于深度学习的方法已成功运用深度神经网络，从原始分子数据中提取出诸如内在属性、生物活性及化学反应等关键特征。此外，利用循环网络、变分自编码器以及图网络等技术，人们还对从头设计新分子展开了深入研究。更近来，为克服传统监督学习方法中需针对每项任务单独学习特征的局限性，研究人员提出了无监督学习策略，旨在更好地表征化学输入信息。这些前沿方法与“基础模型”理念不谋而合——即通过大规模数据集进行训练的基础模型，正逐渐被视为深度学习领域的一种全新范式。

具体而言，一种通过自监督方式对神经网络进行预训练以获得更优特征表示的概念，已被成功应用于多个化学领域。N-Gram Graph和GROVER分别采用了图神经网络和图Transformer网络，从分子图中提取出预训练模型。而ChemBERTa-2则利用7,700万种分子训练了一个RoBERTa模型，构建了分子基础模型，其核心是让模型学会预测200种不同的化学属性值。

与此同时，在计算机视觉领域，诸如视觉-语言预训练（VLP）等多模态预训练方法，已在需要同时理解图像与文本的下游任务中展现出卓越性能。目前，大多数先进的VLP模型均基于Transformer架构及其交叉注意力机制，以捕捉不同模态之间的关联性。此外，一些研究还引入了对比学习方法，通过将具有相同上下文的特征紧密对齐，并拉远语义上无关的特征，从而实现图像特征与语言特征在共同特征空间中的有效对齐。正是得益于VLP技术，诸如视觉问答、图文检索、文本驱动的图像生成以及图像驱动的文本生成等多种任务才得以实现——这些任务若仅依赖单一模态的基础模型，则根本无法完成。

受多模态学习成功启发，近期一些研究尝试通过融合不同数据表示中的知识，以获取更优质的分子特征。例如，Winter 等人训练了一个从简化分子输入行输入系统（SMILES）到国际化学标识符（InChI）密钥的转换模型，从而得到一个包含两种分子表示共性、富含语义信息的特征向量。此外，Zhu 等人则采用了一种自监督的 BYOL 训练方法，将 SMILES 和分子图这两种不同的分子表示形式结合起来，构建了一个双视图模型。然而，这些研究仅将多模态技术用于增强单模态任务中的分子特征，并未深入探索不同模态之间的相互作用。更为重要的是，由于 SMILES、InChI 和分子图表示几乎包含了关于分子中原子间连接关系的相同信息，因此很难期待通过这些不同分子表示形式间的多模态学习产生全新的涌现特性。

在本工作中，我们关注分子结构与其相关性质之间的跨模态理解，这有助于解决诸多应用中的实际任务，例如性质预测、条件分子设计等。与以往利用已标注的性质作为标签来提取通用特征的多任务学习方法不同，我们的方法将一组性质视为一种独立的模态，用以表征输入分子，并提出：引入这种性质模态的分子多模态学习，能够生成更具信息量的特征。具体而言，我们提出了一种分子结构-性质多模态基础模型（SPMM），该模型可在计算机中开展各类化学实验，且经过大规模分子结构及相应性质向量的预训练。通过采用Transformer架构，我们得以分别借助自注意力机制和交叉注意力机制，实现模内特征提取以及模间特征融合。

实验结果表明，通过单一基础模型同时学习结构特征及其关联属性信息，能够获得更优的表示，进而针对多种下游任务进行微调。具体而言，由于模型对结构和属性采取了对称处理方式，它现在能够仅凭一个预训练模型即可实现双向生成与预测，而这一点在以往是无法做到的。

方法

图1a展示了SPMM的整体模型架构及训练目标。SPMM框架在双流VLP模型的基础上进行了扩展。双流VLP模型分别用单模态编码器对每种模态的输入进行编码，随后通过另一个编码模块实现跨注意力机制：以一种模态的特征作为查询，另一种模态的特征则分别充当键和值。当输入一个待训练的分子时，SPMM会将其SMILES字符串和属性向量（PV）作为多模态数据输入，如图1a所示。其中，SMILES和PV分别送入对应的单模态编码器，这些编码器执行自注意力操作，使嵌入后的输入同时成为键、查询和值。经过两个单模态特征的提取后，对比学习策略将SMILES和PV特征映射到同一嵌入空间中，确保包含相同上下文的信息特征得以对齐。这种方法已被证实能有效提升模型性能，因为它简化了跨模态编码过程，并引导单模态编码特征更准确地反映输入的语义信息。

随后，编码后的SMILES和PV特征会被送入融合编码器，该编码器在SMILES与PV特征之间执行跨注意力操作。由于对比学习已将SMILES编码器和PV编码器的输出对齐到同一特征空间，这一融合编码器能够通过交替使用查询与键/值输入，灵活地完成跨注意力任务。此外，融合编码器还经过预训练，分别采用SMILES的下一个词预测（NWP）、属性的下一个属性预测（NPP），以及SMILES与PV匹配损失（SPM）作为训练目标。其中，基于给定Transformer输入预测下一个组件，是目前广泛使用的自监督学习目标；而我们的NWP和NPP任务，则帮助模型深入理解SMILES标记与属性之间的上下文关系，同时借助另一种模态的语义特征增强学习效果。更为重要的是，SPM任务还能判断给定的SMILES和PV组合是否代表同一种分子。

经过训练后，SPMM可用于多种需要同时理解SMILES和属性信息的双向下游任务，例如属性预测（从SMILES到属性）以及基于属性的分子生成（从属性到SMILES，也称为逆QSAR），如图1b所示。此外，所采用的预训练目标使预训练后的SPMM同样适用于单模态任务，如分子分类和反应预测（见图1c）。在这些单模态任务中，预训练的SPMM表现与当前最先进模型相当，这表明该模型作为基础模型具有出色的泛化能力。

a SPMM 的模型架构与预训练目标。SPMM 采用分子的简化分子输入行输入系统（SMILES）及其属性向量（PV）。对比损失将两个单模态编码器的输出特征对齐到同一嵌入空间中。融合编码器则通过预测下一个词、预测下一个属性，以及 SMILES-PV 匹配损失进行训练，以学习两种模态之间的关系。[CLS]S 和 [CLS]P 分别代表用于 SMILES 和 PV 模态的特殊标记。b 需要多模态理解的潜在下游任务场景，包括从 PV 到 SMILES 的生成任务，以及从 SMILES 到 PV 的生成任务。c 单一模态输入的潜在下游任务场景，包括属性预测、正向反应与逆向反应预测。

如图1b所示，模型通过自回归采样生成PV或SMILES。以SMILES为例，具体来说，从想要生成的模态的[CLS]令牌开始，模型预测第一个组件，并重复使用之前的输出来预测下一个组件，直到完成或遇到停止标志。在使用SMILES编码器对输入分子进行编码后，将[CLS]S令牌的特征通过分类/回归头传递以获得输出。分类/回归头由带有一个隐藏层的前馈神经网络组成。用给定的训练集对模型进行微调，得到验证集上损失最小的检查点，并记录该检查点在测试集上的性能。

具体而言，首先，原始SMILES字符串由标记器进行标记，并由SMILES编码器使用[CLS]S标记和[SEP]标记嵌入。这里，[CLS]令牌是附加在每个输入序列开头的特殊令牌。尽管[CLS]令牌本身不包含任何含义，但该模型的双向注意力机制允许[CLS]令牌包含整个输入的上下文信息。一旦模型被预训练，给定序列的[CLS]令牌输出可以被视为输入表示向量，并用于分类/回归下游任务。

在SMILES标记化中，标记器将给定的SMILES标记为片段，这些片段包含在包含300个子词的准备好的标记字典中。该字典是通过字节配对编码（BPE）算法从预训练数据SMILES语料库中获得的，该算法从一组简单的字符开始，迭代地将最频繁的标记对追加为合并子词。BPE算法被广泛应用于各种语言模型，它提供了一个包含常见官能团和子结构(如苯环、羰基、双字母原子和氨基)的子词字典。与将每个字符视为单独标记的朴素字符标记化相比，合并的子词有助于模型对化学组的化学推断，并减少标记的总数。

在这项工作中，作者为每个包含53种分子性质的分子构建了一个PV，并将其视为一个长度为53的句子。RDKit中的这些属性涵盖了从简单的属性(如环数和摩尔质量)到复杂的属性(如溶解度、TPSA和可成药性)的广泛范围。Transformer架构将PV的每个元素视为执行注意机制的令牌。虽然与自然语言相比，SPMM词汇表的大小更有限，它们的顺序也更固定，但它提供了关于53个属性的更精确和紧凑的信息。将PV视为一种语言的一个好处是，不必收集构建有效PV的所有元素。

将SMILES和属性值（property values）视为一种语言进行处理

分子可以通过多种格式表示，例如fingerprints、SMILES、InChI等字符串，或分子图。由于这些不同形式的表示几乎包含了关于完整分子结构的相同信息，我们选择使用SMILES来描述分子结构。SMILES是一种用字符序列表示分子连接结构的方法。许多研究者将SMILES视作一种语言数据的变体，并借鉴语言模型的概念，将其应用于SMILES数据的化学任务中。

a 输入分子的SMILES表示通过词典进行分词，该词典是利用BPE算法获得的。分词结果中的每个标记会被替换为其对应的词嵌入，随后连同位置嵌入一并送入SMILES编码器。b 分子的化学属性列表经过准备后，会通过一个线性层。输出的向量会被随机替换为特殊的[UNK]标记，然后将处理后的结果送入PV编码器，并加入位置嵌入。

图a展示了针对输入SMILES的嵌入过程。原始的SMILES字符串经分词器分词后，再由SMILES编码器进行嵌入，并附加了[CLS] S标记和[SEP]标记。其中，[CLS]标记是每个输入序列开头特别添加的一个特殊标记。尽管[CLS]标记本身并不携带任何语义信息，但模型的双向注意力机制使其能够捕捉到整个输入的上下文信息。在模型完成预训练后，给定序列对应的[CLS]标记输出即可被视为一种输入表示向量，进而用于后续的分类或回归任务，正如许多基于BERT的图像处理变体以及VLP所采用的方式一样。

在SMILES分词中，分词器会将给定的SMILES分解为片段，这些片段均包含在预先构建的由300个子词组成的词汇字典中。该字典是通过BPE算法从预训练数据SMILES语料库中获得的——BPE算法从一组简单字符开始，通过迭代地将最频繁出现的词对合并为一个新子词，逐步扩展词汇表。作为一种被广泛应用于多种语言模型的技术，BPE算法生成的子词字典涵盖了常见的官能团和亚结构，例如苯环、羰基、双原子元素以及氨基等。与简单地按字符进行分词（即将每个字符视为独立的标记）相比，这种合并后的子词有助于模型更准确地推断化学基团，并有效减少了整体的标记数量。

图b展示了针对输入PV的嵌入流程。PV中的每个属性元素均为数值，并已通过该属性的均值和标准差进行归一化处理。这53个属性的排列顺序是预先确定的，正如补充表4所示，SPMM的性能并不会因PV中特定属性顺序的变化而受到影响。接下来，我们将PV中的每个数值利用线性层编码为一个特征向量，作为数值的编码表示。随后，我们随机将50%的属性特征替换为[UNK]标记——这一特殊标记专门用于模拟属性信息未知的情况。由于仅使用部分属性即可完整描述分子结构，这种做法完全可行。通过随机屏蔽部分属性特征，模型得以避免过度依赖于个别属性，同时实现了数据增强的效果，从而有效提升模型的泛化能力。尽管我们在本研究中所用的每种属性均可轻松且全面地由计算机生成，但在实际应用中，其他属性可能并不具备同样的条件。即便某些训练分子的部分属性信息未知，SPMM依然能够通过将这些未知属性替换为[UNK]标记来进行训练。此外，在随机屏蔽的数值编码基础上，我们还添加了一种可学习的位置编码，其原理与BERT中的位置编码类似。由于PV本身仅包含数值信息，这种位置嵌入能够明确指示每个数值具体对应的是哪种属性。同时，由于这些属性的排列顺序是预先设定好的，这种位置嵌入实际上等同于为每种属性分配一个唯一的索引，并在此基础上进一步嵌入该索引值。最后，我们将处理后的结果送入PV编码器，并在开头添加一个[CLS] P标记。

预训练目标

对比学习旨在通过将来自不同模态的特征对齐到同一特征空间中，以学习更优质的单模态表示。当SMILES S和PV P的[CLS]标记编码特征分别表示为$S_{cls}$和$P_{cls}$时，我们计算相似度函数sim(S, P)和sim(P, S)如下：

其中，$h_S$ 和 $h_P$ 分别是用于 SMILES 和属性向量的线性投影 + 归一化层。现在，对于给定的 $S$ 和 $P$ 对，按以下方式计算 SMILES-PV 以及 PV-SMILES 的跨模态相似度：

其中，$M$和$N$分别是用于损失计算的SMILES和PV的总数。这里，$\tau$是一个可学习的温度参数，它通过放大相似性差异，起到强化区分度的作用。同样地，模内相似性也可以用相同的方法进行计算。

总体对比损失通过交叉熵损失$H$与$one-hot$相似度$y$定义，其中$y$对源自同一分子的样本对取值为1，否则为0。

在最近将对比损失应用于VLP之后，构建了$SMILES$和$PV$队列，用于存储最新的$k$个SMILES和PV实例，并将其用于计算对比损失。将队列大小$k$设定为24,576。

下一词预测（NWP）训练模型，使其能够在给定前0至n-1个SMILES标记及其对应PV的情况下，准确预测第$n+1$个SMILES标记。预测下一个标记是训练语言模型的常见目标，这一方法尤其常用于GPT的预训练过程。具体实现时，可通过为每个SMILES序列单独设计一条流程，在SMILES编码器和融合编码器的自注意力机制中应用因果掩码来完成。设$S=s_0,s_1,\dots ,s_n$表示输入的SMILES序列，$P$为对应的PV，$p^{NWP}(s_{ₙ}|s_{₀:n-1},P)$则表示在给定$P$及前$0$至$n-1$个SMILES标记条件下，模型预测的第$n$个标记的概率分布。NWP的损失函数定义如下：

其中，$y_n^{NWP}$是第 n 个 SMILES 基础符号 $s_n$ 的独热编码标签。

我们为属性向量应用了与“Next Property Prediction”（NPP）相似的NWP概念。NPP通过使用对应的SMILES和前一阶段的属性，使模型能够预测下一个属性值。由于每个属性元素都是数值，用均方误差损失取代了NWP中的交叉熵损失。设$S$和$P=p_0,p_1,\dots ,p_n$表示输入的SMILES-PV对，而${\hat{p}}_n(p_{0:n-1},S)$则表示模型在属性值序列中应用因果掩码及融合编码器后所预测的下一个属性值，NPP的损失函数如下所示：

在NPP中，如果将该属性值替换为[UNK]标记，模型将无法预测其数值。
SMILES-PV Matching（SPM）用于学习给定的SMILES-PV对（S, P）是否匹配。我们从融合编码器的输出中分别提取[CLS]S和[CLS]P标记的特征，并将它们拼接后送入一个线性层构成的SPM头部。当$p^{SPM}（S,P）$为SPM head的输出时，SPM损失可定义为：

其中，$y^{SPM}$是一个用于binary标签 $SMILES-PV$ 匹配的独热编码向量；当 $S$ 和$P$ 来自同一分子时，该标签为 1，否则为 0。为了构建 SPM 的负样本，我们从另一种模态中随机选取每个 SMILES 和 PV 实例的“负”配对，并将其作为负样本。这种负样本的选择采用了硬负挖掘策略(hard-negative mining)，即对于那些与 Eq. (2) 公式计算出的相似度较高但并非正样本匹配的实例，它们更有可能被选作负样本。这一方法使训练过程更具挑战性，从而迫使模型学会区分相似实例之间的细微差异。

在对比学习中，使用独热标签可能过于严格，因为它会将所有来自其他配对的实例一视同仁地视为负样本。然而，某些分子指纹（PVs）可能与多个SMILES匹配，而不仅仅是一个与其配对的SMILES。此外，即使是同一个SMILES，也可能对应不同的分子指纹，因为分子指纹中存在50%的掩码机制——例如，“摩尔质量=[UNK]，logP=2.1，原子数=12”和“摩尔质量=78，logP=2.1，原子数=[UNK]”都能准确描述苯，即便它们分别源自不同分子。类似的问题同样存在于下一句预测任务（NWP）中：有时，下一个标记可能存在多个合理的选择，但若仍采用独热标签作为真实值，则可能会忽略这一情况。

为了解决这一问题，构建了动量教师模型（momentum teacher model），并利用其输出进行对比学习和NWP任务。该动量教师通过提供一种反映教师模型理解方式的伪标签，实现了知识蒸馏。具体而言，用于对比学习和NWP任务的标签，分别与动量模型的输出$s_{\ast ,momentum}(\ast \in s2p,p2s,s2s,p2p)$以及$p_{momentum}^{NWP}(s_n|s_{0:n-1},P)$混合，并通过一个可调超参数$\alpha$进行控制。关于如何将动量模型应用于对比学习和NWP任务的具体公式，详见式$(8)(9)$及式$(10)(11)$。

在每个批次中更新学生模型的参数$w_{model}$后，动量教师模型的参数$w_{momentum}$则通过指数移动平均（EMA）方法，基于$w_{m}odel$及一个EMA超参数$\lambda$，按照公式（12）进行更新。

整体预训练目标是对比、NWP、NPP 和 SPM 损失的综合损失。

注意到，当没有特定的注意力掩码时（例如GPT中使用的因果注意力掩码），Transformer模型的自注意力机制和交叉注意力机制的输出对每个特征向量的位置是不变的。这意味着，在预训练阶段，唯一与预先确定的顺序或PV的53种属性相关的目标任务是NPP任务；而在其他预训练目标中，当我们随机打乱属性的顺序及其对应的位置嵌入时，SPMM的表现完全相同。此外，如果我们将NPP任务替换为一种对顺序不敏感的目标任务（如BERT中的掩码语言建模），PV中元素的预先设定顺序将完全不会影响SPMM的输出。补充表3展示了当NPP任务被替换为掩码属性建模时，SPMM的表现情况，结果表明，这种替换并未提升模型性能。

结果

作者将SPMM与一些具有代表性的方法进行了比较。如表1所示，正向反应预测中，将底物SMILES作为输入，预测产物SMILES，而逆向反应预测则相反。SPMM通过SMILES-PV-SMILES的流程，实现SMILES-PV的双向学习，并用于SMILES预测。将预测的前K个分子的准确率作为评价指标，与一些基于SMILES的模型（string-based）或基于分子图的模型（graph-based）进行对比。在正向预测中，SPMM在最精确的（K=1）预测上超越了现有方法。在逆向预测中，SPMM在K=5和K=10预测中超越了现有方法。

作者设计了消融实验。在单个分子性质预测（通过SMILES预测PV）任务上（#task=1），将数据集大小（#data）不同的数据集上的结果进行对比，同时将采用预训练的SPMM和未采用预训练的SPMM（w/o pre-train）加入对比，将RMSE作为评价指标（越低越好）。结果表明预训练对SPMM的重要性。

作者还进行了案例分析。通过可视化预训练SPMM的注意力得分，进一步分析了学习到的SMILES和PV之间的跨模态理解。在图2中，当给出SMILES及其PV输入时，作者绘制了预训练SPMM的最后一个融合层的交叉注意力得分。可见，得分遵循了化学性质与分子片段之间的直观关系。与氢键相关的属性，即氢键配体数（#H bonding donor）和氢键受体数（#H bonding acceptor），对带有氧原子和氮原子的标记显示出很高的注意力分数。属性“RingCount”（环数）侧重于与环相关的令牌，而对侧基团的注意力较弱，属性“NumAromaticRings”（芳环数）仅对芳香环的组成部分给予较高的注意力分数。当不同的SMILES符号在分子中扮演相似的角色时，如分子7中不同位置的“c1ccccc1”，注意力模式也相似。这一结果表明，SPMM可以捕捉分子结构和化学性质之间的关系，而不需要明确给出它们之间的监督信息。

总结

本文提出了一种结合结构和性质信息的多模态化学基础模型SPMM。提出的模型允许双向生成/预测分子结构和性质，以及下游任务，如反应预测。在此过程中，作者引入了一种将属性集合视为一种语言的方法，以便模型可以独立学习SMILES令牌与每个属性之间的关系。实验证明了预训练的SPMM在SMILES和PV域之间的相互作用问题上表现出显著的性能。SPMM不仅适用于多模态挑战，甚至它的单模态特征也适用于SMILES，它还提供了一种有用的表示，可以对许多分子下游任务进行微调。

尽管SPMM具有显著的性能，但它还有可改进的地方。虽然SMILES可以包含分子二维结构的全部细节，但关于原子和键如何连接的信息只是隐式存在的。此外，分子结构的轻微改变可能会导致SMILES发生巨大变化。图格式是另一种广泛使用的分子表示方式，它包含邻接矩阵的显式信息，可以作为SMILES的替代方法。同时，图格式可以对分子的立体结构信息建模，而考虑立体化学信息在各种生化任务中起着至关重要的作用。

总之，SPMM可以提供一个能够同时包含每个输入域及其多模态域的预训练模型，这具有巨大的潜力。未来，通过使用更广泛和更丰富的分子模式，以及推广到可能的不同生化领域，如聚合物和蛋白质，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