UNI-MOL: A UNIVERSAL 3D MOLECULAR REPRESENTATION LEARNING FRAMEWORK

UNI-MOL: A UNIVERSAL 3D MOLECULAR REPRESENTATION LEARNING FRAMEWORK

Neurips23

推荐指数：#paper/⭐⭐⭐#​（工作量不小)

动机

在大多数分子表征学习方法中，分子被视为 1D 顺序标记或2D 拓扑图，这限制了它们为下游任务整合 3D 信息的能力，这使得 3D几何预测/生成几乎不可能。为此，作者提出了一种3D分子表征学习方法

模型框架

​

贡献：

1. backborn:基于 Transformer 的模型，可以有效地捕获输入的 3D 信息，并直接预测 3D 位置。
2. 预训练。两个大规模数据集：一个 209M 分子构象数据集和一个 3M 候选蛋白质口袋数据集，分别用于分子和蛋白质口袋的预训练 2 个模型。以及两个预训练任务：3D 位置恢复和掩蔽原子预测，用于有效学习 3D 空间表示。
3. 微调。适用于各种下游任务的多种微调策略。例如，如何在分子性质预测任务中使用预训练的分子模型;如何在蛋白质-配体结合姿势预测中结合两个预训练模型。

backborn的设计：

常用有两种的设计：GNN与transformer

由于GNN在捕获领域信息有优势，而局部连接的图缺乏捕捉原子之间长距离相互作用的能力。作者认为长距离信息更有用，因此选择Transformer 作为 Uni-Mol 中的主干模型。因为它完全连接了节点/原子，因此可以学习可能的长程相互作用。

​

整体结构概述

如图 2 所示，Uni-Mol 主干是基于 Transformer 的模型。它有两个输入，原子类型和原子坐标。模型中维护了两种表示形式（atom 和 pair）。原子表示由 Embedding 层从 atom 类型初始化;对表示由原子坐标计算的不变空间位置编码初始化。特别是，基于原子之间的成对欧几里得距离，对表示对全局旋转和平移是不变的。这两种表示在 self-attention 模块中相互通信。

编码3D位置：

编码 3D 位置 由于其排列不变性，Transformer 在没有位置编码的情况下无法区分输入的位置。与 NLP/CV 中使用的离散位置不同 ，3D 空间中的位置(坐标)是连续值。此外，位置编码过程在全局旋转和平移下需要保持不变。已经提出了几种 3D 空间位置编码来解决这个问题，我们没有兴趣重新发明一种新的编码。因此，我们对现有的编码进行了基准测试，并使用了一个简单而有效的编码：原子对的欧几里得距离，然后是一个对类型感知的高斯核。

此外，由于不变的 3D 空间位置编码是在对级编码的，作者还在 Transformer 中维护了对级表示，以增强 3D 空间表示。具体来说，对表示被初始化为上述空间位置编码。

然后，为了更新对表示，作者通过 self-attention 中的多头Query-Key 产品的结果使用atom到pair 通信。形式上，ij 对表示的更新表示为

q i j l + 1 = q i j l + { Q i l , h ( K j l , h ) T d ∣ h ∈ [ 1 , H ] } , \boldsymbol{q}_{ij}^{l+1}=\boldsymbol{q}_{ij}^l+\{\frac{\boldsymbol{Q}_i^{l,h}(\boldsymbol{K}_j^{l,h})^T}{\sqrt{d}}|h\in[1,H]\}, qijl+1​=qijl​+{ d ​