深势科技发布三维分子预训练模型Uni-Mol，有望加速药物设计的发展
用微观的视角来看现实世界会是怎样的？那里有数以万计的生物分子，与细胞中的蛋白质、核酸等不断地自由组合与拆解中，再生成新的产物。
通常，药物分子性质由分子、骨架和整体分子所决定，其分子的结构性质则直接影响着药物性质和生物利用度。分子有千万种组合的方式，就药物设计领域而言，药物分子的多样结构也使候选药物的化学空间变得更大。
而从生物学的角度来看，分子的性质和药物的作用主要取决于分子的三维结构。因此，在药物设计中，分子表征学习模型和临床试验显得尤为重要。目前的分子表征学习（MRL）模型一般用一维序列或二维图结构作为输入，并且在大多数性能预测任务中表现良好。
然而，这也限制了它们在实际应用中结合分子 3D 信息的能力，尤其是 3D 相关任务的建模。因此，如何进一步提高现有 MRL 模型的性能和扩展其应用范围，是该领域一直以来亟待解决的问题。
近日，深势科技团队发布了目前药物设计领域第一个通用的大规模三维 MRL 框架 UNI-MOL，以扩大分子表征学习（MRL）的应用范围和表征能力。5 月 26 日，相关论文以《UNI-MOL：通用的三维分子表征学习框架》（UNI-MOL: A UNIVERSAL 3D MOLECULAR REPRESENTATION LEARNING FRAMEWORK）为题在CHEMRXIV上预发表。
该论文显示，UNI-MOL 包含两个有相同架构的模型，一个是由 2.09 亿个分子三维构象数据训练的分子预训练模型，一个是由 320 万个蛋白口袋数据训练的口袋预训练模型。UNI-MOL 在 14/15 分子性质预测任务中表现优于 SOTA。值得关注的是，在 3D 相关的任务，包括蛋白质-配体结合位点预测、分子构象生成等下游任务中表现尤其出色。
与现有大多数分子表征学习模型的不同之处在于，UNI-MOL 跳出大多模型所采用的一维序列或二维图结构，而是直接利用分子三维结构作为模型输入和输出，利用分子的三维信息训练模型。
这一三维 MRL 框架 UNI-MOL 由三个部分组成，即基于 TRANSFORMER（基于自注意力机制的一个深度学习模型）的骨干处理 3D 数据、两个分别学习分子表征和口袋表征的预训练模型、以及面向各类下游任务的微调策略。
表征学习（或预训练，自监督学习）指的是在大量无标注数据上学习数据高维表征的一种学习范式，例如自然语言处理（NLP）的 BERT 和 GPT、计算机视觉（CV）的 VIT 等。
表征学习存在一个共同点，即无标签的数据极为丰富，而有标签的相关数据则相对有限。在药物设计领域的应用亦是如此。一般而言，机器多在大规模的无标签数据上进行预训练，然后，在有标签的数据上进行微调，帮助模型从有限的监督信号中提取更多的信息。
为了利用大规模无标签数据进行预训练，该团队创建了两个大型数据集，并且也将开源 209M 的分子三维构象数据集和 3M 候选蛋白质口袋数据集，分别用于分子和蛋白口袋的两个预训练模型。
分子预训练数据集是基于多个可购买的分子数据集构造。数据集包含大约 1900 万个分子，共 2.1 亿个 3D 分子构象。蛋白质口袋预训练数据集来自于蛋白质数据库库（RCSB PDB），研究人员从数据库中生成一个由 320 万个候选口袋组成的数据集用于口袋预训练。
在实际应用中，蛋白质口袋直接参与了许多药物设计任务，对候选蛋白质口袋的预训练可以提高与蛋白质-配体结构和相互作用的相关任务性能。
图丨UNI-MOL 模型架构。左图：整体预训练架构；中间：模型输入，包括原子和和原子间空间位置编码；右图：原子对的表征和它的更新过程（来源：CHEMRXIV）
为了进一步验证 UNI-MOL 模型的有效性，研究人员对多个下游任务进行了广泛实验，包括分子性质预测、分子构象生成、口袋性质预测和蛋白-配体结合位点预测。
分子性质预测是新药发现的重要一环。MOLECULENET 是一个被广泛使用的分子性质基准预测,包括聚焦于分子不同层次属性的数据集。该研究实验的 15 个数据集均来自于 MOLECULENET，研究团队采用了骨架划分的方式。
在 3D 结构强相关的回归任务上，例如在水化自由能（ESOL, FREESOLV），亲脂性（LIPO），物化性质（QM7、QM8、QM9）上，与之前的 SOTA 相比有平均提升 21% 的效果。
在分子构象生成（MOLECULAR CONFORMATION GENERATION)实验中，UNI-MOL 评价 AI 模型生成构象多样性的指标 COVERAGE 和精准指标 MATCHING 基本上全面超越现有的基准线。
此外，研究团队进行了口袋性质预测实验。可药物性，指的是候选蛋白口袋与特定分子配体产生稳定结合的能力，是候选蛋白口袋的关键属性之一。UNI-MOL 在口袋药性数据集 NRDLD 上的准确率、召回率、精确度和F1得分是最高的。
在基准数据集中，预训练的 UNI-MOL 模型在四个分数上都优于非预训练的模型。这表明，对候选蛋白质口袋的预训练可以改善口袋属性预测任务。
图丨蛋白质-配体结合位点预测模型框架，编译器使用两路预训练 UNI-MOL 分表表征分子和口袋，解码器使用同样结构的随机初始化的 UNI-MOL（来源：CHEMRXIV）
蛋白质-配体结合的预测，是基于结构的药物设计中最重要的任务之一。UNI-MOL 结合了分子和口袋预训练模型来学习基于距离矩阵的打分函数，后对复杂的构象进行采样和优化，并在基准数据集上确保了结果的可泛化性。通过结合分子模型和口袋预训练模型，UNI-MOL 在蛋白质-配体结合任务中显著超过了普遍使用的对接工具。
打分函数即通过定量化评估药物与靶标蛋白的相互作用为药物研发中的药效评估提供理论依据，该方法目前广泛应用于基于结构的计算辅助药物设计。
通常来说，对药-靶的相互作用的评估步骤有两步，第一步是对接过程（DOCKING PROCESS），主要指构象搜索，找出潜在的结合位点（BINDING POSE）；第二步是打分过程（SCORING PROCESS），通常指打分，以预测药-靶结合力。
而 UNI-MOL 在 DOCKING POWER 和 BINDING POSE 两项评估结合最关键的指标上均表现出色，在打分函数测评上超越了一系列主流对接处理工具，以及基于 AI 的打分函数模型。在 BINDING POSE 预测能力上面，对于 CASF-2016 基准数据集预测的准确结合构象的比例超过目前主流的工具约 35%。
图丨蛋白质-配体结合位点预测模型在 CSAF-2016 上的测试结果（来源：CHEMRXIV）
从研究结果来看，3D 分子表征学习框架 UNI-MOL 十分理想，该框架的提出或有利于药物设计领域的进一步拓展和深耕。目前，药物设计领域一直缺乏高质量的公开数据集，许多公开数据集无法满足实际需求。而该项研究成果或有助于未来药物设计领域高质量基准的增多，或将加速新药发现与药物设计的发展。
