ICLR2023|一个泛用的3D分子表示学习框架

AITIME论道

于 2023-07-04 18:30:35 发布

阅读量385

点赞数

文章标签：学习

原文链接：https://mp.weixin.qq.com/s?__biz=MzIzOTAxOTk3Nw==&mid=2247515371&idx=1&sn=a4ded9085ef8ac924155c972605fcb2b&chksm=e9328e4dde45075b610306c16ac3315f297a22e9de2bba3740bd5f02e24184cc67ad1a11b63d&scene=126&sessionid=0

版权

欢迎向本公众号投稿文献解读类原创文章，投稿邮箱：1298332329@qq.com，请用markdown格式写稿件，以附件形式发送。海内外招生、访学、招聘等稿件，请联系微信：17067431029。

作者 | 黄　锋

审核 | 付海涛

今天给大家介绍一篇深势科技发表于ICLR2023的论文“Uni-Mol: A Universal 3D Molecular Representation Learning Framework”。这篇论文发表了一个泛用的预训练框架用于分子表示学习，并在大量下游任务中达到了SOTA。

摘要

分子表示学习(MRL)由于其在药物设计等应用中发挥重要作用而获得了极大的关注。在大多数MRL方法中，分子被视为1D序列符号或2D拓扑图，限制了它们为下游任务合并3D信息的能力，特别是使得3D几何预测/生成几乎不可能。本文提出了一种泛用的3D MRL框架Uni-Mol，极大地扩展了MRL方案的表示能力和应用范围。Uni-Mol包含两个具有相同SE(3)-Transformer架构的预训练模型:一个由209M个分子构象预训练的分子模型;一个由3M候选蛋白口袋数据预训练的口袋模型。此外，Uni-Mol包含几种微调策略，将预训练的模型应用于各种下游任务。通过适当地结合3D信息，Uni-Mol在14/15分子性质预测任务中优于SOTA。此外，UniMol在蛋白质配体结合姿势预测、分子构象生成等3D空间任务中表现优异。贡献：1)Uni-Mol是一个能够直接用于3D任务的分子预训练框架；2)基于大量的基准，作者提出一个简单有效的SE(3)-Transformer作为骨干网络，并为模型设计了一个3D预训练的策略；3)Uni-Mol在多种下游任务中超越了SOTA。

方法

上图左边展示了这篇论文提出的预训练框架，中间是模型输入，包含原子表示和成对表示，右边是模型主干网络的示意图。

骨干网络

Transformer和GNN是现如今两大最流行用于学习分子表示的编码器，作者考虑使用Transformer作为主干网络，因其能够捕捉到分子内部原子间的长程互作信息。首先是原子的位置编码，作者考虑了多种原子空间位置编码方法并选择了一个简单有效的方式，即原子对的欧氏距离(相关结果参看原文附录E.1)，这种位置编码被称为成对表示(pair representation)，其在输入到模型中以下种方式被更新：

式中的是原子对在层的成对表示，是注意力头的数目，是隐层的维数，而是第原子在第层头的QKV注意力机制下的Query(Key)。为了将3D信息充分融入原子表示中，作者引入一个对到原子(pair-to-atom)的交互机制来计算注意力：

这里的就是QKV注意力机制下的Value。为了使得模型能够直接输出3D位置，作者引入一个SE(3)等变的预测头：式中的是原子数目，是上述编码器的层数，是第原子的输入坐标，是对应的输出坐标，和是将成对表示转化为标量实数的投影矩阵。

预训练

为了进行预训练，作者生成了两个大规模的数据集，一个由有机分子的3D结构组成，另一个由候选蛋白质口袋的3D结构组成。然后，分别使用这两个数据集对两个模型进行预训练。由于口袋直接参与了许多药物设计任务，直观地，对候选蛋白质口袋的预训练可以提高与蛋白质配体结构和相互作用相关的任务的性能。对于分子预训练数据，其来源于多个公开的分子数据集总计19M分子，并使用RDKit中的ETKGD方法搭配Merck分子力场优化来为每个分子生成11个构象，因此总共209M分子构象。而蛋白口袋预训练集来自于Protein Data Bank(RCSB PDB)数据库，包含180K的蛋白3D结构。为了获取候选口袋，作者首先通过添加缺失侧链和氢键来清洗数据，然后使用Fpocket来检测出可能的蛋白结合口袋，最后通过残基数目来过滤到口袋，最终数据集中包含了3.2M的候选口袋。

为了让模型充分学习到3D结构信息，作者设计了一个3D位置恢复的自监督任务。主要的思想就是输入一个破坏的位置，恢复出正确的3D位置。采用这种方式学习，输入的破坏过的原子位置和正确的位置之间不能相差太多，否则学习难以进行；也不能相差太少使得学习过于简单。为此，作者设计一个贪婪策略来生成破坏的位置。算法如下图所示，给真实位置随机一个噪音范围，添加噪声过大时还可以继续采用重排的方式使得原子错位。作者通过各种基准测试实验发现r=1Å而不需要重排时为最佳方案(附录E.6)。最后考虑三种任务，其一是恢复原子对之间的欧氏距离；其二是直接的坐标预测；其三是掩码掉破坏的原子的原子类型，与BERT类似，使用[CLS]代替，其坐标使用所有原子的中心，最后预测正确的原子类型(具体设置参看附录C，关于自监督任务的基准测试参看附录E.5和E.6)。

微调

按下游任务的种类分为三类阐述微调部分。首先是非3D预测任务，诸如分子性质预测、分子相似性、口袋成药性预测、蛋白配体结合力预测等等。简单地，只需要将预训练好的模型，利用读出操作从原子表示得到分子表示，加一个线性的预测头对应各种下游任务进行微调即可，而对于是像口袋-分子对的预测任务，则可以将两者表示拼接后接一个线性预测头即可。再者是分子或口袋的3D预测任务，诸如分子构象预测等等。对于这类任务，在Uni-Mol中则简化为输入一个构象生成另一个新的构象的构象优化任务。最后是蛋白-配体对的3D预测任务，这是一类在基于结构的药物设计中最重要的任务之一。这个任务预测蛋白结合位点和分子配体的复杂结构。除了口袋和分子本身的构象变化外，还需要考虑分子在口袋中的位置，即刚性运动额外的6度自由度(3个旋转和3个平移)。在Uni-Mol中，分子表示和口袋表示被拼接起来输入到一个4层的解码器来学习重原子的成对距离。然后将预测的成对距离矩阵为评分函数，首先随机放置配体，然后通过直接反向传播当前成对距离矩阵和预测的距离矩阵之间的损失来优化配体的原子坐标。这个过程比传统的对接工具快上100倍，更多细节参看原文附录C.6。