MGM、MolGPT、PAR、Uni-Mol、K-Bert、MolCLR…你都掌握了吗？一文总结生物制药必备经典模型（三）

1990年代后期，计算生物学开始成为生物学中非常重要的一部分。在大热的AlphaFold掀起浪潮之前，就有科学家断言：所有生物学都是计算生物学。AI或者深度学习的出现，给计算生物学带来了新的巨大的发展空间。

对于生物学本身，传统的实验和分析手段已难以充分开发海量生物数据，确实需要计算生物学这种跨学科同时兼顾多个细分领域的综合性工具来解决问题。在具体实验方法上，当前绝大多数采用的都是基于已有数据库和资源、利用成熟工具来解决特定问题或自行设计统计分析、数值计算的方法，而计算生物学的出现让干湿实验结合的新方法开始走向主流（在生物上讲，干实验就是通过计算机模拟以及生物信息学方法来进行研究。湿实验就是通过在实验室里采用分子、细胞、生理学试验方法进行研究）。引入AI，实现了对在传统的湿实验环境中的假设的验证，干湿实验共同迭代加速，AI和传统科研结合带来的巨大潜能，有望带来一场全新的科学革命。

在计算生物学中，AI的应用主要有三类：一是，计算推演生物性质及原理，包括：蛋白质结构预测、致病机理研究、蛋白质相互作用预测（PPI）、抗体和抗原的表位预测、基于基因组学寻找疾病成因或寻找新型的生物标志物等。（生物标志物是指可以标记系统、器官、组织、细胞及亚细胞结构或功能的改变或可能发生的改变的生化指标，可用于疾病诊断、判断疾病分期或者用来评价新药或新疗法在目标人群中的安全性及有效性。）这些研究的成果后续可用于得到新的药物靶点等，为疾病治疗提供基本思路。二是搭建预测及判断模型，包括：AI制药中基于靶点的化合物性质预测（主要涉及小分子药物开发），疾病诊断/监控/治疗建模，涵盖细胞/器官/人体的生物模拟器等。其中，生物模拟器的本质功能是用于验证特定疗法有效性的生物模拟器，可以简单理解为生物医药领域的数字孪生。三是对生物体进行控制改造，包括：新疗法/药物开发、精准医疗和生物制造（以合成生物学为代表）。其中新疗法/药物开发是目前落地最成熟的场景。再往细来说，对癌症的个性化治疗和基因组学也将成为精准医疗中最先落地的场景。AI应用于新药开发，可以实现药物靶点发现、药物筛选和结构优化、合成路线等。

本文聚焦于生物制药中必备的TOP模型，具体来说就是第三类AI应用中的主要模型，可以运用到整个药物从研发、中试到生产的所有关键技术环节。上面提及的第二类AI应用主要是生物医药领域的数字孪生，不包含在本文的讨论范围内。

本文回顾的必备TOP模型主要包括蛋白质结构预测和蛋白设计、分子生成、分子表征和性质预测这三类应用，而化学合成/逆合成及其它大数据分析应用等，暂不包含在本文讨论的模型范围中。AI的各种模型和算法应用在生物制药领域，需要与对应的生物学、医学知识高度结合，因此，本报告中对必备TOP模型的介绍主要是从AI建模的角度对模型总体架构和整体设计思路进行介绍，各个模型设计的技术细节、模型调参等思路和技术创新点，需结合原文和所应用的医学场景深入理解。

二、分子生成

1、 MGM

分子设计是一个具有挑战性的问题，可应用于药物发现和材料设计。作者引入了一个掩码图（Masked Graph）模型，该模型通过捕捉未观察到的节点（原子）和边（键）的条件分布来学习图上的分布。

图1模型结构

图1给出了Masked Graph模型图，包括特征化的细节。首先将顶点和边嵌入图G/η中，分别得到连续表示𝐡_𝑣𝑖∈ℝ^𝑑0和𝐡_𝑒𝑖,𝑗∈ℝ^𝑑0，其中，d0是连续表示空间的维数。然后，将这些表征传递给一个信息传递神经网络（message passing neural network，MPNN）。使用MPNN作为模型的基本组成部分，因为它对图的同构性是不变的。一个MPNN层包括一个聚合步骤，将每个节点的相邻节点的信息聚合起来，然后是一个更新步骤，使用聚合的信息来更新每个节点的表述。将L层叠加起来建立一个MPNN；所有L层的参数都是相同的。对于除最后一层以外的所有层，从第l层输出的更新的节点和边缘表示被送入第l+1层。与原始版本的MPNN不同，还在每一层维护和更新每个边的表征。可以使用有效地模拟节点和边缘特征之间关系的图神经网络的任何变体，如MPNN。

图2 MPNN更新步骤

\MPNN的节点和边的更新步骤图如图20。在MPNN的每一层，首先通过计算每个节点vi的累积信息𝐮(𝑙)𝑣𝑖，使用聚合函数Jv和相邻节点之间的空间残差连接R来更新其隐藏状态。

同样，每条边𝐡𝑒𝑖,𝑗的隐藏状态也使用以下规则对所有j∈N(i)进行更新：

边缘附带节点的两个隐藏表征之和通过Je，一个两层全连接网络，在两层之间有ReLU激活，产生一个新的潜在边缘表征。然后，节点投影层𝐴𝑣:ℝ𝑑0→Λ𝑇和边缘投影层𝐴𝑒:ℝ𝑑0→Λ𝑅处理来自最后一层的节点和边缘表示，其中ΛT和ΛR分别为节点和边缘类型的概率简化。结果是所有𝑣和所有𝑒的分布p（v∣G\η）和p（e∣G\η）。

使用六个节点属性来表示每个节点，索引为{𝜅∈ℤ:1≤𝜅≤6}，每个节点都有自己的独热嵌入。图2中的每个节点对应于分子中的一个重原子。在前向传递过程中，这些嵌入中的每一个都乘以一个单独的权重矩阵𝑊𝜅∈ℝ𝑇𝜅×𝑑0，其中Tκ是属性κ的类别数。由此产生的连续嵌入被加在一起，形成一个节点的整体嵌入。每个属性的独热嵌入为：

• 原子类型：原子的化学符号（如C、N、O）。

• 氢数：与该原子结合的氢原子的数量。

• 电荷：原子上的净电荷，其中第一个指数代表数据集中原子上的最小电荷，最后一个指数代表最大电荷。

• Chirality type：未指定，顺时针四面体，逆时针四面体，其他。

• 环内：原子是或不是环状结构的一部分。

• Is-aromatic: 原子是或不是芳香环的一部分。

在用MPNN处理图形后，将每个节点的表示通过六个独立的全连接的两层网络，层与层之间有ReLU激活。对于每个节点，每个网络的输出是对其中一个属性的初始独热向量的类别分布。在训练过程中，计算预测的分布和被corruption过程掩盖的所有属性的ground-truth之间的交叉熵损失。对某一特定属性进行masked的节点的选择与对所有其他属性的选择无关。这样做的目的是为了让模型能够更容易地学习不同属性类型之间的关系。

对于边属性嵌入来说，使用与节点属性嵌入描述相同的框架。只使用一个具有权重矩阵W∈ℝ𝑅×𝑑0的边缘属性，其独热嵌入定义如下：

• 键的类型：无，单，双，三或芳烃键。

2、 MolGPT

应用深度学习技术从头生成分子（de novo generation of molecules），被称为逆向分子设计，在药物设计中获得了巨大的吸引力。用SMILES符号将分子表示为一串字符，这使得自然语言处理中最先进的模型，如Transformer，可以用于一般的分子设计。受已被证明能成功生成有意义文本的生成性预训练（generative pre-training，GPT）模型的启发，在本研究中利用 masked self-attention对下一个标记预测任务进行transformer-decoder 训练，以生成类药物分子。

训练和生成MolGPT模型的示意图如图3所示。对于无条件训练，分子SMILES首先使用SMILES标记器进行标记，然后在下一个标记预测任务中训练模型：

图3使用MolGPT模型进行训练和生成的流程

对于性质和骨架条件的训练，作者使用RDkit提取分子性质和骨架，并将它们作为条件与分子SMILES一起传递。生成过程则是首先给模型提供一个开始的标记，模型按顺序预测下一个标记，从而生成一个分子，然后，向模型提供一组分子性质和骨架条件以及开始标记来对分子进行采样。

模型的架构如图3所示，该模型基本上是生成预训练Transformer(GPT)模型的迷你版。与GPT1有大约110M的参数不同，MolGPT只有大约6M的参数。MolGPT由堆叠的解码器块组成，每个解码器块都是由masked的自注意力层和全连接神经网络组成。全连接神经网络的最后一层返回一个大小为256的向量，然后作为下一个解码器块的输入。MolGPT由八个这样的解码器块组成。为跟踪输入序列的顺序，将位置值嵌入分配给每个标记。在条件训练期间，提供分段标记以区分条件标记和SMILS标记。分段标记的嵌入表示特定输入是条件还是分子SMILES令牌，以便于通过模型区分两者。使用嵌入层将所有分子SMILES标记映射到256维向量。类似地，使用单独的可训练嵌入层来将位置标记和分段标记映射到256维向量。然后将这些SMILES的标记嵌入、位置嵌入和分段标记嵌入相加，为SMILES的每个标记产生大小为256的向量，然后将其作为输入传递给模型。

图4 MolGPT模型结构

三、分子表征和性质预测

1、Iterative Refinement LSTM

分子性质预测在药物发现中有着重要的作用，可以识别具有目标特性的候选分子。同时，分子性质预测也是一个小样本问题，因为标记数据很少。通常，在分子性质预测领域使用机器学习方法建立分子结构和特定特性之间的联系的方法由两部分组成：一个分子编码器，将分子结构编码为固定长度的分子表示；一个预测器，它根据分子表示估计某个属性的活性。

这项工作展示了如何利用一次性学习来显著降低药物发现应用中预测所需的数据量，具体引入了一个新的迭代细化长短期记忆（the iterative refinement long short-term memory）架构，当它与图卷积神经网络结合时，可以显著提高对小分子有意义的距离度量的学习。

注意力机制a取决于两个嵌入函数f : X → R^p和 g : X → R^p，它们将输入的样本（在小分子空间）嵌入到一个连续的表示空间。在过去的工作中，f和g都是标准卷积网络，而在这项工作中，f和g是图卷积网络。嵌入向量f(x)和{g(x_i)}之间的相似性是用相似性度量k(-, -)计算的。考虑到相似度的选择，一个简单的注意力机制a可以通过将相似度值归一化来定义

当f(x)与其他g(x_j )相比在度量k下更接近g(xi)时，a(x, x_i)较大。这种one-shot学习的表述被称为Siamese one-shot learning 方法。

上述简单的one-shot学习公式使用了独立的嵌入f(x)和g(x_i)，只有距离度量k将两个特征图的信息联系在一起。因此，特征图f(x)是在没有任何关于支持S中可用数据的背景下形成的。匹配网络架构（matching-network）则通过开发完整的上下文嵌入来解决这个问题，其中，嵌入f(x)=f(x | S)和g(x_i)=g(x_i | S)是使用x和S计算的。完整的上下文嵌入允许x和x_i的嵌入相互影响。从经验上看，这种增加允许更强的one-shot学习结果。为了构建f(x | S)和g(x_i | S)，匹配网络使用标准卷积神经网络计算嵌入f'(x)和g'(x)。嵌入g(x)的计算方法是将所有g(x_i)放入一个向量，并使用双向LSTM（BiLSTM）将所有元素连接起来，使每个g(x_i)受到所有g'(x_j )的影响。

为了计算嵌入f(x | S)，匹配网络利用了一个基于上下文的注意力LSTM模型（attLSTM）。这个模型计算其输入的独立顺序组合.

the iterative refinement LSTM架构保留了匹配网络的上下文感知设计，但解决了上面提到的支持-嵌入g的顺序依赖性以及查询和支持的非对称处理。核心思想是使用attLSTM生成查询嵌入f和支持嵌入g。匹配网络嵌入需要支持嵌入g（- | S）来定义f（- | S）。为了解决这个问题，图23的架构同时迭代演化了这两个嵌入。

在第一次迭代中，定义f(x) = f'(x)，g(S) = g'(S)。然后，使用注意力机制迭代更新嵌入f和g。这种结构使嵌入可以迭代地相互通知。下面的公式说明了在每个迭代中进行的全部更新

图5 用于药物发现的一次性学习的网络结构示意图

如图6，生成的双重嵌入被用来进行一次性学习。

图6 嵌入的迭代细化的图形描述。为了说明问题，输入/输出是二维的，q1和q2构成坐标轴。红色和蓝色的点代表正/负样本（仅用于说明）。原始嵌入g′(S)显示为方块。预期的特征r显示为空圆圈

在图卷积系统中，每个节点都有一个描述符的向量。然而，在预测时，我们将需要整个图的一个固定大小的向量描述符。引入一个新的图收集卷积层，它简单地将图中所有节点的所有特征向量相加，得到一个图特征向量（见图25）。这些图形卷积操作对于将小分子转化为连续的矢量表示很有用。这样的表征正在成为在深度网络中操纵小分子的一种强大方式。早期的one-shot学习使用卷积神经网络将图像编码为连续矢量，然后用来进行单次预测。通过类比，使用图卷积神经网络将小分子编码为适合one-shot预测的形式。

图7 本文描述的主要图形操作的图形表示。对于每个操作，被操作的节点都显示为蓝色，未被改变的节点显示为浅蓝色。对于图卷积和图池化，操作显示为单个节点v，这些操作是同时对图中的所有节点v进行的

2、 PAR

这项工作提出了一个属性感知关系网络（Property-Aware Relation，PAR）来处理分子性质预测问题，具体利用了相关的子结构和分子之间的关系在不同的分子特性中都会发生变化这一事实。此外，设计了一个自适应关系图学习模块，以联合估计分子关系图和完善分子嵌入的目标属性，从而使有限的标签能够有效地在类似分子之间传播。采用元学习策略，在任务中选择性地更新参数，以便分别建立通用知识和属性意识知识模型。

图7 PAR结构，我们绘制了一个来自Tox21的2-way 2-shot任务。PAR对一组任务进行了优化。在每个任务Tτ中，带虚线的模块在支持集Sτ上进行微调，带实线的模块是固定的。一个查询分子x_τ,i首先将使用基于图的分子编码器表示为g_τ,i，然后通过属性感知嵌入函数转换为p_τ,i。这个p_τ,i进一步与S_τ中的分子在关系图上的嵌入共同拟合为h_τ,i，作为最终的分子嵌入并用于类别预测

我们的目标是，从一组为数不多的分子特性预测任务{Tτ }中学习一个预测器，并对给定的几个标记分子进行泛化以预测新的特性。第τ个任务T_τ预测索引为i的分子x_τ,i在目标属性上是活跃的（y_τ,i = 1）还是不活跃的（y_τ,i = 0），前提是每类有少量的K标记样本。然后，这个Tτ被表述为一个双向的K-shot分类任务，支持集Sτ = {(xτ,i, yτ,i)}，包含2K个标记的样本，查询集Qτ = {(x_τ,j , y_τ,j )}，包含(Nτ)^q个未标记的待分类样本。

首先，使用一个基于图的分子编码器，为每个x_τ,i提取一个长度为d^g的分子嵌入，这个基于图的分子编码器的参数表示为Wg。然而，现有的基于图形的分子编码器不能捕捉到属性意识的子结构。特别是在跨任务学习时，一个分子可以被评估为多个属性。这导致了分子和属性之间的一对多的关系，这使得少数的分子属性预测变得特别困难。我们的目标是隐含地捕捉嵌入空间中的子结构，即Tτ的目标属性。令c表示类c∈{0, 1}的类原型，计算为

把这些类的原型作为Tτ的上下文信息，并把它们编码到x_τ,i的分子嵌入中，如下所示

其中，b_τ,i是使用按比例的点积自注意力来计算的，这样，每个g_τ,i都可以以维度的方式与类原型进行比较。然后得到属性感知的分子嵌入p_τ,i，即

MLP_Wp表示以Wp为参数的多层感知器（MLP），引入MLP_Wp以寻找一个低维空间，该空间编码与Tτ的目标属性更相关的子结构。这种情景化的p_τ,i是具有属性意识的，它可以对目标属性有更多的预测性。

图8 在Tox21的两个任务中共存的相同分子的关系图的示例。红色（蓝色）边缘意味着连接的分子在目标属性上都是活跃（不活跃）的

除了相关的子结构外，分子之间的关系也会在不同的属性中发生变化。如图27所示，具有共同属性的两个分子在另一个属性上可能是不同的。因此，进一步提出了一个自适应的关系图学习模块，以捕捉和利用分子之间的这种属性意识的关系图，从而使有限的标记能够在相似的分子之间有效地传播。在这个关系图学习模块中，我们交替地估计分子间关系图的邻接矩阵，并在学习到的关系图上完善分子嵌入，时间为T次。

在第t次迭代中，令(Gτ)^(t)表示关系图，其中，Vτ将Sτ中的2K个分子和Qτ中的一个查询分子作为节点。(Gτ)^(t)的相应邻接矩阵表示如下：

理想情况下，x_τ,i、x_τ,j 的属性意识分子嵌入p_τ,i、p_τ,j之间的相似性揭示了它们在当前属性预测任务下的关系。因此，初始设定(h_τ,i)^(0) = p_τ,i。首先使用当前的分子嵌入来估计A_τ^(t)。A_τ^(t)]ij的第(i, j)个元素记录了x_τ,i,和x_τ,j之间的相似性，其计算方法是：

然而，在一个双向K-shot任务中，一个查询分子在Gτ^(t)中只有K个真正的邻居。对于二元分类，在相反的类别中选择一个错误的邻居将严重恶化分子嵌入的质量，特别是当每个类别只提供一个标记的分子。为了避免错误邻居的干扰，进一步将Gτ^(t)简化为K-近邻（KNN）图，其中，将K设定为与S中每一类的标记分子数完全相同。对应前K个最大的[A_τ^(t)]_ij的指数的x_τ,i记录为N (t)(x_τ,i)。然后，我们看到

将每个节点嵌入h(t)与这个更新的关系图上的其他节点嵌入共同适应，编码为Aˆ(t)τ。令Hτ^(t)表征所有节点嵌入，其中，第i行对应于h_τ,i^(t)。Hτ^(t)的更新方式是

进一步设计一个 neighbor alignment regularizer 来惩罚关系图中错误邻居的选择，具体表述为

将yˆτ,i作为x_τ,i的类别预测（相对于主动/非主动），其计算公式为

最后，PAR的meta-training过程如下：

3、 Uni-Mol

深势科技团队发布了首个三维分子预训练模型 Uni-Mol。Uni-Mol 直接将分子三维结构作为模型输入，而非采用一维序列或二维图结构，在利用 2 亿个分子三维构象和 三百万个蛋白候选口袋数据进行预训练后，Uni-Mol 在几乎所有与药物分子和蛋白口袋相关的下游任务上都超越了 SOTA（state of the art），也让 Uni-Mol 能够直接完成分子构象生成、蛋白 - 配体结合构象预测等三维构象生成相关的任务，并超越现有解决方案。

图9 Uni-Mol框架示意图。Uni-Mol由两个模型组成：一个是由2.09亿分子三维构象训练的分子预训练模型；一个是由三百万候选蛋白口袋数据训练的口袋预训练模型。这两个模型独立用于不同的任务，在蛋白质-配体结合任务上两个都用

与现有大多数分子表征学习模型的不同之处在于，Uni-Mol 跳出大多模型所采用的一维序列或二维图结构，而是直接利用分子三维结构作为模型输入和输出，利用分子的三维信息训练模型。这一三维 MRL 框架 Uni-Mol 由三个部分组成：基于 Transformer的骨干处理 3D 数据、两个分别学习分子表征和口袋表征的预训练模型、面向各类下游任务的微调策略。

具体介绍Uni-Mol的3个部分。1）主干。在Transformer的基础上，增加了不变的空间位置编码和对级表示，以更好地捕捉三维信息。此外，引入一个 equivariant head 来直接预测三维位置。2）预训练。创建了两个大规模的数据集，一个是209M的 molecular conformation 数据集，一个是3M的候选蛋白口袋数据集，分别用于对分子和蛋白口袋的2个模型进行预训练。对于预训练任务，除了 masked atom prediction ，还使用了三维位置去噪任务来学习三维空间表征。3）微调。根据具体的下游任务，使用的预训练模型是不同的。例如，在分子特性预测任务中，只使用分子预训练模型；在蛋白质-配体结合姿势预测中，两种预训练模型都要使用。

图10 Uni-Mol模型架构。左图：整体预训练架构，包括输入和预训练设计的任务；中间：模型结构，包括模块连接和3D结构编码；右图：模型基本单元，包括表征以及更新过程

由于 Transformer 有置换不变性，它在没有位置编码的情况下无法区分输入的具体位置，而且位置编码需要在全局旋转和平移的情况下保持不变。不同于基于离散值的位置编码，分子的 3D 信息，即三维空间中的坐标是连续值，要保证它对旋转和平移的不变性。深势科技团队使用所有原子对的欧氏距离，融合分子图中边的类型，之后经过高斯核函数得到位置编码，形式上可以表达成如下公式：

通常 Transformer 只维护 Token（原子）级别的表征，在微调时的下游任务中也会调用它。然而由于分子的空间位置信息是在原子对级别上编码的，Uni-Mol 模型中也维护原子对表征，以便更好地学习分子的 3D 表征。具体实现中，原子对表征的初始化是上面提到的空间位置编码。之后为了更新原子对表征，深势科技团队通过自注意力机制中多头的 Query-Key 的乘积，进行原子到原子对的通信。形式上，原子对的更新可以表达成如下公式：

通过 3D 空间位置编码和原子对表征，Uni-Mol 可以学习到一个分子 3D 表征。但它仍然缺乏直接输出坐标的能力，而这在 3D 空间任务中是必不可少的。为此，模型中加入了一个简单的  SE(3)-equivariance head 。SE(3)-equivariance head 的设计可以表示为：

为了利用大规模无标记数据进行预训练，深势团队自己构造了两个大规模的有机小分子和蛋白口袋 3D 结构数据集。通过统一的预训练模型框架，结合有效的预训练任务策略，在大规模分布式集群上进行了预训练模型的训练。与 BERT 类似，Uni-Mol 中也使用了对原子掩码的预测任务。对于每个分子/口袋，通过添加了一个特殊的原子 [CLS]，其坐标是所有原子的中心，用 [CLS] 的表征代表整个分子/口袋的表征。然而，由于 3D 空间位置编码是有化学键信息泄露的，模型很容易依据相互间的距离推测出被掩码的原子类型。因此单独对于原子掩码进行预测并不能帮助模型学习有用的信息。

为了解决这个问题同时又能从 3D 信息中学习，Uni-Mol 设计了一个基于 3D 坐标的去噪任务。具体实现中，对于被掩码的 15% 的原子，给坐标同时加入 [-1Å , 1Å] 的均匀分布噪声，之后模型根据被污染的坐标计算出来空间位置编码。这样一来，对于原子掩码的预测任务就不再可有可无。此外，还引入两个额外的任务单元来直接对于原子坐标进行预测：1）还原被掩码的原子间欧氏距离。基于原子对表征，预测被掩码的原子对的欧氏距离。2）直接预测被掩码的原子坐标。图29是整个预训练框架的说明。