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

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是整个预训练框架的说明。

3、 K-Bert

K-Bert是在BERT基础上提出的新的预训练策略，让模型能够直接从SMILES中提取分子特征。作者在原文中提出了三个预训练任务，包括：原子特征预测任务、分子特征预测任务和对比学习任务。原子特征预测任务允许模型学习基于图的方法中手动提取的信息（初始原子信息），分子特征预测任务允许模型学习基于描述符的方法中手动提取的信息（分子描述符/指纹），而对比学习任务允许模型使同一分子的不同 SMILES 字符串的embedding更相似，从而使 K-BERT 能够识别同一分子的不同SMILES字符串。具体，预训练任务2，即分子特征预测（图11B）为：对RDKit计算所得的分子特征进行预测。本研究采用MACCS指纹，同样也可将该任务视为多任务分类任务（可以更换为其他的指纹/描述符）。

图11 Knowledge-based BERT（K-BERT）预训练策略

输入表征：每个SMILES都是用Schwaller等人提出的标记化方式来标记成不同的token。然后将token（如’O’、‘Br’和‘[C@H]’）编码成K-BERT的输入。

Pre-training：每个重原子的原子特征和每个分子的分子特征都通过RDKit计算，并分别用于预训练任务2。

Fine-tuning：如图11D所示，K-BERT中有6个transformer encoder，从预训练模型中载入前5个transformer encoder的参数，第6层transformer encoder和预测层重新随机初始化。然后在下游任务数据上，重新训练模型。

Data augmentation: 每个分子的SMILES，通过rdkit随机扩充到5个不同的SMILES。在训练集中，每个SMILES都被当做是单独（不同的）分子。而在测试集和验证集中，同一分子的不同SMLES都被看成是该分子，5个不同SMILES的预测结果均值作为该分子的预测结果。

4、 MolCLR

分子机器学习为高效的分子特性预测和药物发现带来了希望。然而，标记后的分子数据的获取可能是昂贵和耗时的。由于标记数据有限，有监督学习的机器学习模型要归纳到巨大的化学空间面临着巨大的挑战。MolCLR（通过图神经网络进行分子对比学习表征）是一个自监督的学习框架，利用了大量的无标记数据（约1000万个独特的分子）。在MolCLR的预训练中，建立了分子图并开发了图神经网络编码器来学习可区分的表征。具体提出了三种分子图的增强：原子掩码、键删除和子图删除。引入对比损失，MolCLR通过对比积极的分子图对和消极的分子图对来学习表示。广泛使用的GNN模型、图卷积网络（GCN）和图同构网络（GIN），被开发为MolCLR中的GNN编码器，以从分子图中提取信息表示。然后在MoleculeNet的下游分子特性预测基准上对预训练的模型进行微调。与通过监督学习训练的GCN和GIN相比，本文所提的MolCLR明显提高了分类和回归任务的性能。受益于对大型数据库的预训练，MolCLR在多个分子基准测试中超过了其他自监督学习和预训练策略。此外，在一些任务上，MolCLR可以与监督学习的基准相媲美，甚至超过了监督学习，其中包括复杂的分子图卷积操作或特定领域的特征化。MolCLR表示法和传统的FPs之间的进一步比较表明，MolCLR可以通过对大量未标记数据的预训练来学习区分分子的相似性。

MolCLR预训练:将小批量N个分子数据中的SMILES Sn转换为分子图Gn。对每个图应用两个随机分子图数据增广操作，得到两个相关的掩码图:G_2n−1和G_2n。一个基本特征编码器建立在图形卷积和读出操作的基础上，提取分子特征h_2n−1 , h_2n。利用对比损失来最大化MLP project head的潜在向量z_2n−1 和z_2n之间的一致性。

图12 MolCLR结构

如图12中所示的分子图增强方法具体包括：

原子掩码(Atom Masking)。分子图中的原子以给定的比例被随机掩码。当一个原子被掩码时，它的原子特征被一个掩码标记取代，该标记与图12（b）中红框所示的分子图中的任何原子特征相区别。通过掩码处理，模型被迫学习内在的化学信息（比如由某些共价连接的原子的可能类型）。

键删除(Bond Deletion)。如图21(b)中的黄色方框所示，"键删除"以一定比例随机删除原子之间的化学键。与原子掩码方法中用掩码标记替代原始特征的原子掩码不同，键删除是一种更严格的增强，因为它从分子图中完全删除了边缘。原子之间化学键的形成和断裂决定了化学反应中分子的属性。键的删除模拟了化学键的断裂，促使模型学习一个分子在各种反应中的关联性。

子图删除(Subgraph Removal)。子图删除可以被认为是原子掩码和键删除的结合。子图的去除从一个随机挑选的原点开始。去除过程通过掩码原子的邻域，然后是邻域的邻域，直到被掩码的原子数量达到原子总数的一定比例。然后，被掩码的原子之间的化学键被删除，这样被掩码的原子和被删除的化学键就形成了原分子图的子图。如图12(b)中的蓝色方框所示，被移除的子图包括被掩码原子之间的所有化学键。通过匹配被移除的不同子结构的分子图，该模型学会了在剩余的子图中找到显著特征，这在很大程度上决定了分子的特性。