药物设计领域的BERT？三维分子表征学习框架Uni-Mol，一个模型刷爆所有下游任务

2022 年 5 月 27 日 PaperWeekly

预训练模型正在席卷 AI 领域。从大规模无标注数据中提取表征信息，再在小范围标注的下游任务上进行监督学习，正在成为很多领域的事实解决方案。NLP 中有 BERT、GPT-3，CV 中有 ViT，而这样的模式如何助力药物设计，也一直都是人们密切关注的问题。药物分子与图片、语言文字的不同之处在于， “什么是最好的分子表征”依旧是一个人们未能形成共识的问题。 主流分子预训练模型均从一维序列或二维图结构出发，但分子结构本身是在三维空间中表示的。能否直接从三维信息出发构建预训练模型、获得更好的分子表征，是一个重要而有意义的问题。

近日，曾连续发布 Uni-Fold、Uni-FEP、Uni-EM 等“Uni-”系列产品的深势科技团队，发布了首个三维分子预训练模型 Uni-Mol。Uni-Mol 直接将分子三维结构作为模型输入，而非采用一维序列或二维图结构。从三维信息出发的表征学习让 Uni-Mol 在几乎所有与药物分子和蛋白口袋相关的下游任务上都超越了 SOTA（state of the art），也让 Uni-Mol 得以能够直接完成分子构象生成、蛋白-配体结合构象预测等三维构象生成相关的任务，并超越现有解决方案。相关成果以《Uni-Mol: A Universal 3D Molecular Representation Learning Framework》为题，出现在 ChemRxiv 上（点击文末“阅读原文”查看）。

Uni-Mol框架

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

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

1.1 处理3D空间信息的Transformer

1.1.1 旋转平移不变的空间位置编码

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

其中为通道数，为原子对，为原子对之间的欧氏距离，为边类型，为高斯核函数，其参数为。注意这里边类型不是化学键类型，它只与原子对的原子类型有关。是一个仿射变换，参数为和，它在原子对的欧氏距离和其对应的边类型之间建立联系。

1.1.2 原子对表征

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

其中是原子对表征，是 attention heads 的数目，是隐藏层维度，是投影矩阵。

此外，为了利用原子表征中的三维信息，作者还引入了原子对到原子的通信，将原子对表征作为自注意力机制的 bias 项，可表示成如下公式：

1.1.3 具有SE(3)等变性的coordinate head

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

其中为分子中的原子数目，是模型层数，是输入坐标，是投影矩阵。

1.2 预训练任务设计

1.2.1 预训练数据集

为了利用大规模无标签数据进行预训练，深势团队自己构造了两个大规模的有机小分子和蛋白口袋 3D 结构数据集。通过统一的预训练模型框架，结合有效的预训练任务策略，在大规模分布式集群上进行了预训练模型的训练。

分子预训练数据集是基于多个可购买分子数据集构造。经过归一化和去重，数据集包含大约 1900 万个分子，共 2.1 亿的 3D 分子构象。通过使用 rdkit 结合分子力场优化，高效生成分子构象。每个分子随机生成 10 个构象。由于某些分子 rdkit 生成 3D 构象失败，因此还额外对每一个分子生成了一个 2D 构象（基于分子图）帮助预训练。

蛋白质口袋预训练数据集来自于蛋白质数据库（RCSB PDB [http://www.rcsb.org]）。库中有 180K 的结晶真实蛋白 3D 结构。为了构造合理的用于预训练的候选口袋数据，深势团队首先通过补全蛋白侧链和极性氢来进行蛋白准备，然后使用口袋检测工具 Fpocket 检测蛋白质上的潜在药效性口袋，同时保留了其中的水分子，通过上述方法，深势团队构造一个由 320 万个候选蛋白口袋组成的 3D 构象数据集。

1.2.2 自监督策略

与 BERT 类似，Uni-Mol 中也使用了对原子掩码的预测任务。对于每个分子/口袋，通过添加了一个特殊的原子 [CLS]，其坐标是所有原子的中心，用 [CLS] 的表征代表整个分子/口袋的表征。然而，由于 3D 空间位置编码是有化学键信息泄露的，模型很容易依据相互间的距离推测出被掩盖的原子类型。因此单独对于原子掩码进行预测并不能帮助模型学习有用的信息。

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

1）还原被掩盖的原子间欧氏距离基于原子对表征，预测被掩盖的原子对的欧氏距离。

2）直接预测被掩盖的原子坐标通过设计合理的结构，从而保证模型更新对于平移、旋转具有等变性（SE(3) 等变性），去直接预测被掩盖的原子的正确坐标。

两个预训练模型都使用了上述的自监督任务帮助训练，由于蛋白口袋与许多药物设计任务直接相关，作者认为在候选蛋白质口袋数据上的预训练可以提高与蛋白质-配体结构及相互作用有关的任务的表现。图2是整个预训练框架的说明。

实验

2.1 分子性质预测

首先在备受 AI 从业者关注的分子性质预测任务上进行了实验。实验使用的 15 个数据集均来自于 MoleculeNet，划分方式上对齐了之前的工作，采用了骨架划分。从结果上来看，Uni-Mol 在 14/15 个数据集上取得 SOTA，尤其是在 3D 结构强相关的回归任务上，例如水化自由能（ESOL，FreeSolv），亲脂性（ Lipo），物化性质（QM 系列）上面相对于之前的 SOTA 平均有 21% 的效果提升！

2.2 分子构象生成

区别于以往的分子构象生成 baseline，Uni-Mol 是对 RDKit 生成的构象进行优化，在评价 AI 模型生成构象多样性的指标 Coverage 和精度指标 Matching 上，Uni-Mol 基本上全面超越现有的 baseline。

同时深势科技团队也提出对于该领域，目前使用的公开数据集主要关注低能的真空、水相模拟构象，而缺少真实的和蛋白结合的药效构象，因此分子构象生成的数据标准也是未来一个很重要的研究方向。

2.3 口袋性质预测

可药性，即候选蛋白质口袋与特定分子配体产生稳定结合的能力，是候选蛋白质口袋最关键的性质之一。由于有标签的数据很有限，这项任务非常具有挑战性。例如常用的 NRDLD 数据集，只包含 113 条数据。因此，除了 NRDLD 之外，作者还构建了一个回归数据集，用于模型性能测试。在表 4 中可以看到，Uni-Mol 表现卓越。

2.4 蛋白质-配体结合位点预测

▲ 图3. 蛋白质-配体结合位点预测模型框架，编码器使用两路预训练Uni-Mol分表表征分子和口袋，解码器使用同样结构的随机初始化的Uni-Mol

蛋白质-配体结合的预测是药物设计中最重要的任务之一。Uni-Mol 结合了分子和口袋预训练模型来学习基于距离矩阵的评分函数，之后对复杂的构象进行采样和优化。在基准数据集上，作者使用 CASF-2016 作为测试集，使用 PDBbind General set 作为训练集，并且和测试集进行了去重，确保结果的可泛化性。

在 docking power 和 binding pose 两项评估结合最关键的指标上，Uni-Mol 均表现非常出色，在打分函数 docking power 测评上超越了一系列主流的 docking 工具以及 AI based 打分函数模型，更令行业兴奋的是，在最为直接的 binding pose 预测能力上面，对于 CASF-2016 基准数据集（RMSD<2.0 一般认为是可接受的 pose 预测结果）预测的准确结合构象的比例超过目前主流的 docking 工具约 35%，这无疑是巨大的飞跃。