基于3D几何的预训练分子图表示

编译| 沈祥振 审稿| 郭梦月

本文介绍了唐建课题组的Shengchao Liu等人基于3d几何研究的成果：鉴于现实世界场景中 3D 信息的缺乏极大地阻碍了对分子几何图表示的学习，本文提出了一种图预训练框架Graph Multi-View Pre-training (GraphMVP)。GraphMVP通过利用 2D 拓扑结构和 3D 几何视图之间的对应性和一致性来执行自我监督学习(SSL)。该方法有效地学习了一个 2D 分子图编码器，而该编码器之前已通过更丰富和更具辨别力的 3D 几何知识进行了预训练。作者进一步提供理论见解来证明 GraphMVP 的有效性。最后，综合实验表明 GraphMVP优于现有的图 SSL 方法。

1研究介绍

如何有区分度地表示具有矢量化嵌入的分子是当今药物发现中一个基本但开放的挑战。潜在问题之一是如何构建一个目标函数来监督训练，作者的想法是利用 3D 几何和 2D 拓扑之间的一致性来研究自我监督的分子表示学习。

然而，立体化学结构的获取通常非常昂贵，使得此类 3D 几何信息在下游任务中稀缺。为此，作者提出了 GraphMulti-View Pre-training (GraphMVP) 框架，其中 2D 分子编码器使用 3D 几何知识进行预训练，然后在没有 3D 信息的情况下对下游任务进行微调。学习范式是在预训练期间将 3D 分子几何学知识注入 2D 分子图编码器，这样即使没有可用的 3D 信息，下游任务也可以从隐式 3D 几何学中受益。

研究贡献包括：(1)首次将 3D 几何信息合并到图 SSL 中；(2)为预训练提出对比SSL和生成 SSL 的结合模型；(3)在所有 SSL baseline中实现了 SOTA 性能。

2模型与方法

构造基本思想

作者通过利用 3D 和 2D 分子图上的两个前置任务来实现上述目标：一个对比 SSL 和一个生成 SSL，它们分别在分子间和分子内水平作用。对比 SSL 在分子间水平上创建监督信号：如果一对 3D 和 2D 图来自同一分子，则它们为正，否则为负；然后它将对齐正对并同时对比负对。另一方面，生成式 SSL以分子内方式获得监督信号：它学习 2D/3D 表示，可以为每个分子本身重建其 3D/2D 对应视图。为了应对测量分子 3D 和 2D 空间重建质量的挑战，作者进一步提出了一种新的替代目标函数，称为变体表示重建 (VRR)，用于生成 SSL 任务，它可以有效地计算连续表示空间中的这种质量。

这两个 SSL 任务获得的知识是相辅相成的，因此GraphMVP 框架将它们整合在一起，形成更具判别力的 2D 分子图表示。作为通用的 SSL 预训练管道，GraphMVP 有两个阶段：预训练和微调。在预训练阶段，通过辅助任务对提供 3D 和 2D 分子结构的数据收集进行 SSL；在微调期间，预训练的 2D GNN 模型在随后的特定下游任务上进行微调。

GraphMVP预训练

自监督学习 (SSL) 基于视图设计，每个分子都有2d和3d两个自然视图。GraphMVP预训练阶段概述如下:

图1 GraphMVP 预训练阶段概述

对比SSL的主要思想是首先从数据间水平定义正负视图对，然后同时对齐正视图并对比负视图。对于每个分子，首先从,2D 和 3D 视图中提取表示，即 hx 和 hy。然后为对比学习创建正负对：同一分子的 2D-3D 对 (x,y) 被视为正，否则为负。最后，SSL对齐正对并对比负对。

生成SSL旨在通过重建每个数据点本身来学习有效的表示。具体到药物发现，我们有一个 2D 图和每个分子的一定数量的 3D 构象，我们的目标是学习一个强大的 2D/3D 表示，可以在最大程度上恢复其 3D/2D 对应物。通过这样做，生成 SSL 可以强制 2D/3D GNN 对最固有的几何/拓扑信息进行编码，这可以从逻辑上提高下游性能。如图1，它本质上是从采样的 2D 分子图表示 () 中重建 3D 构象异构体 (y)。

上述过程称为变分分子重建，在具体实现上，实际采用了一种类似于 VAE 的轻量级生成 SSL。除此之外，生成SSL还有一个步骤称为变分表示重建 (VRR)。在数据空间上进行图重构并不容易：由于分子（例如原子和键）是离散的，在分子空间上建模和测量会带来额外的障碍。为应对这项挑战就有了VRR，是一种替代损失，可以将重建从数据空间转移到表示空间，也就是说这样就无需将潜在代码解码到数据空间，我们可以直接将其投影到 3D 表示空间。

3实验测试结果

实验设置

数据集：在同一数据集上基于GEOM进行模型预训练，然后对广泛的下游任务进行微调。

2D GNN：使用GraphIsomorphism Network (GIN)作为主干模型。

3D GNN：选择 SchNet进行几何建模。

分子性质预测的主要结果

作者对 10 个 SSL baseline和随机初始化方法进行了全面的比较。

GraphMVP-G 和 GraphMVP-C是GraphMVP的两个变体，区别在于GraphMVP-G更注重生成SSL(Generative)，而GraphMVP-C更注重对比SSL(Contrastive)。

表1 分子特性预测任务的结果

我们观察到 GraphMVP 的性能明显优于随机初始化的方法，平均性能大大优于现有的 SSL 方法。此外，GraphMVP-G 和 GraphMVP-C性能在此基础上依次提高，支持以下说法：3D 几何是 2D 拓扑的补充。

掩模比和构象异构体数量的影响

这里作者进一步分析了 GraphMVP 中掩模比 M 和构象数 C 对结果产生的影响，具体探索了 M ∈ {0,0.15,0.3} 和 C ∈ {1,5,10,20} 的范围，并报告了平均性能

如表 2 所示，从效率和有效性的角度来看，作者鼓励在尝试更多的构象之前调整掩模比。

表2 M的消融研究(C设置为5)

表3 C的消融研究(M设置为0.15)

目标函数的影响

在这里作者部署消融研究来探索四个目标函数的影响，首先是它们的独立影响，然后是成对组合的影响。这些目标函数是：InfoNCE、EBM-NCE、VRR(variational representation reconstruction) 和RR(representation reconstruction)，具体结果如表4所示。

表4 目标函数的消融实验

更广泛的下游任务

在这里，作者测试了 4 个额外的回归属性预测任务和 2 个药物目标亲和力任务。

表5 四个分子特性预测任务和两个 DTA 任务的结果

如表 5 所示，一致的性能增益验证了GraphMVP 的有效性。

实际案例分析

当任务目标在 2D 拓扑方面具有挑战性但使用 3D 几何时很简单（如图 2 所示），作者研究了 GraphMVP 如何提供帮助。为此，设计了两个案例研究来证明 GraphMVP 如何将知识从 3D 几何转移到 2D 表示中。

图2 GraphMVP 预训练阶段概述

第一个案例研究是 3d 直径预测。使用 2d 图来预测 3d 直径。第二个案例研究是远程供体-受体检测。分子具有一种特殊的几何结构，称为供体-受体键，作者想用2d分子图来检测这种特殊的结构。最后作者验证了 GraphMVP 对这两个案例研究的持续改进，并在附录 G.6 中提供了更详细的讨论和解释。

4总结

从领域的角度来看，GraphMVP是结合 3D 信息来增强 2D 图表示学习的首次尝试，并且能够通过考虑建模中的随机性来利用 3D 构象。从技术创新的角度来看，GraphMVP引入了对比 SSL 和生成SSL的结合，并提出VRR 作为一种生成 SSL 方法，能够缓解分子生成中的潜在问。实证结果和理论分析都可以有力地支持上述领域和技术贡献。

参考资料

论文链接：https://arxiv.org/abs/2110.07728

代码链接：https://github.com/chao1224/graphmvp