Motif-based Graph Self-Supervised Learning forMolecular Property Prediction

2021-10-15

字数统计: 3.6k | 阅读时长≈ 15 分钟

Motif-based Graph Self-Supervised Learning forMolecular Property Prediction，2021，NIPS

总结：作者提出了一种基于motif的多层级生成式自监督预训练方法——MGSSL，用于分子图属性预测。具体来说，作者先利用一种分子碎片化算法将原始图打碎，得到不同类型的motifs，然后把这些motifs看作节点，并连接到一起构成一棵motif树（这里和我之前看过的一篇超图方法比较类似，先构建子图，然后将子图看做节点，在子图间添加边，构成一张超图）。然后利用GNNs学习到的节点嵌入对这棵树进行重构。另外，为了捕获分子图中多尺度信息，作者设计的MGSSL框架是多层级的。即在进行motif-tree生成的同时，对节点类型和边类型进行预测，最小化两者的交叉熵损失。

1. 简介

1.1 摘要

Predicting molecular properties with data-driven methods has drawn much attention in recent years. Particularly, Graph Neural Networks (GNNs) have demonstrated remarkable success in various molecular generation and prediction tasks. In cases where labeled data is scarce, GNNs can be pre-trained on unlabeled molecular data to first learn the general semantic and structural information before being finetuned for specific tasks. However, most existing self-supervised pre-training frameworks for GNNs only focus on node level or graph-level tasks. These approaches cannot capture the rich information in subgraphs or graph motifs. For example, functional groups (frequently-occurred subgraphs in molecular graphs) often carry indicative information about the molecular properties. To bridge this gap, we propose Motif-based Graph Self-supervised Learning (MGSSL) by introducing a novel self-supervised motif generation framework for GNNs. First, for motif extraction from molecular graphs, we design a molecule fragmentation method that leverages a retrosynthesis-based algorithm BRICS and additional rules for controlling the size of motif vocabulary. Second, we design a general motif-based generative pre-training framework in which GNNs are asked to make topological and label predictions. This generative framework can be implemented in two different ways, i.e., breadth-first or depth-first. Finally, to take the multi-scale information in molecular graphs into consideration, we introduce a multi-level self-supervised pre-training. Extensive experiments on various downstream benchmark tasks show that our methods outperform all state-of-the-art baselines.

最近几年，利用数据驱动的方法预测分子属性备受关注。尤其是GNNs，在各种分子生成、预测任务中取得了很大成功。在有标签数据稀少的场景下，GNNs可以现在无标签分子数据上进行预训练，学习通用的语义和结构信息，然后在具体任务上进行微调。但是现有的大多数自监督预训练GNNs框架只关注节点层级或者图层级任务，这些方法不能捕获丰富的子图或者motif信息。比如，functional groups（分子图中经常出现的子图）经常携带有感分子性质的的指示信息。为了弥补这一缺陷，我们提出了Motif-based兔子监督学习方法MGSSL，一种新的自监督motif生成GNNs框架。首先，为了从分子图中提取motif信息，我们设计了一种分子碎片化方法，利用一种基于逆合成算法BRICS和额外的规则控制motif的大小。然后我们设计了一种通用的基于motif的生成式预训练框架，要求GNNs进行拓扑和标签预测。该框架可以以广度优先和深度优先两种方式实现。最后，为了考虑分子图的多尺度信息，我们引入了多层次自监督预训练。各种下游标准任务上的大量实验证明了我们提出的方法有SOTA方法。

1.2 本文工作

背景： GNNs最近几年被广泛用于各种分子预测任务，并且取得了很大成功，但是GNNs也存在一个弊端，即data-hungry。为了解决这个问题，受NLP和CV中相关方法的启发，研究人员提出了自监督GNNs。在标签稀少场景下，GNNs首先在大规模无标签数据集上进行预训练，然后再针对具体任务进行微调。

动机： 作者认为现有的GCL方法是次优的，因为他们没有捕获graph motifs中蕴含的丰富信息。虽然也有些方法用motif来搞对比学习，但是它们都没有考虑motif的拓扑信息。

本文工作：作者提出了Motif-based Graph Self-Supervised Learning方法MGSSL和多层级自监督预训练框架。

2. 方法

作者提出的MGSSL框架如下图所示：

框架包含三部分：分子碎片化，motif生成和多层级自监督预训练。

2.1 分子碎片化

所谓碎片化，分成两步：

把一张分子图拆分成若干个碎片，每个碎片即代表一个motif。
然后把每个motif当做一个节点，在motif之间加上边，构成一个motif tree。

用符号表示即： $G=(V,E)$ 表示一张分子图， $\mathcal{T(G)=(V,E,X)}$ 表示一棵标签树，其中 $\mathcal V=\{M_1,...,M_n\}$ 表示节点集合， $\mathcal E$ 和 $\mathcal X$ 分别表示边集合和motifs嵌入，每个motif $M_i=(V_i,E_i)$ 是 $G$ 的一个子图。

上述第1步拆分分子图方法有很多，但是需要实现下面4个目标：

所有motif $M_i$ 的并集要和图 $G$ 等价，即 $\cup_iV_i=V且\cup_iE_i\cup\mathcal E=E$ 。
不同motif之间不能有交集，即 $M_i\cap M_j=\empty$ 。
生成的motif必须有语义意义，比如化学领域中的functional group。
生成的motif必须在数据集中是频繁出现的，这样基于这些motif预训练的GNNs才能推广到下游具体任务。

作者提出的基于BRICS碎片化方法如下图2所示：

作者首先使用Breaking of Retrosynthetically Interesting Chemical Substructures算法将图打碎。BRICS定义了16条规则，将分子中和化学反应相匹配的化学键打破，然后在断点两端添加一个虚拟原子（上图2中蓝色小点），表示不同片段可以通过该点连接到一起。BRICS算法旨在保留具有结构和功能价值的分子组分，比如芳香环。

由于BRICS定义的规则中化学反应有限，所以有时候分子图拆分的不够细，只能得到很少片段；
虽然BRICS生成了很多motifs，但是这些motifs很多都是同一个underlying structure的变体。这导致最终得到的motif vocabulary虽然很大（超过100K个片段），但是这些片段在数据集中出现的次数不超过5次。

为了解决上述问题，作者对BRICS算法进行了改进。具体来说，作者对BRICS输出的片段进行二次分解，分解规则如下：

当一端原子在环上，而另一端不在环上时，打破化学键；
选择具有三个或者更多邻居的原子作为新的motif，并打破相邻的键。

第一条规则可以减少环变体的数量，第二条规则可以打破侧链。实验表明，这两条规则可以有效减小motif vocabulary的大小，并提高生成的motifs在数据集中出现的频次。

2.2 moitf（tree）生成

本文作者提出的MGSSL框架是一种生成式自监督学习方法。通过2.1小节介绍的分子碎片化算法，我们可以得到一棵真实的motif-tree。这一步的工作就是，利用GNNs学习到的节点嵌入生成一棵motif-tree，GNNs学习的目标就是让生成的tree和2.1节碎片化得到的tree尽可能相似。

用符号语言描述就是：

给定一张分子图 $G=(V,E)$ 和GNN模型 $f_\theta$ ，首先将分子图拆分成motif tree $\mathcal{T(G)=(V,E,X)}$ 。
然后用GNN模型对motif tree进行最大似然建模 $p(\mathcal T(G);\theta)$ ，表示motifs是如何被标记和连接的。
预训练GNNs的目标就是最大化motif trees的似然估计，即 $\theta^{*}=\operatorname{argmax}_{\theta} p(\mathcal{T}(G) ; \theta)$ 。

这里和前人graph generation工作类似，作者使用自回归方式来构建motif tree。

注：自回归模型是一种常用的深度生成模型，大概思想就是用前i步生成的数据指导第i步数据的生成。详细内容可以参考这篇博客。

具体来说，用 $\pi$ 表示motif的顺序， $i^\pi$ 表示id为i的motif在置换 $\pi$ 中的位置。这样 $p(\mathcal T(G);\theta)$ 就可以表示成所有可能的置换 $\pi$ 下的期望：

p(\mathcal{T}(G) ; \theta)=\mathbb{E}_{\pi}\left[p_{\theta}\left(\mathcal{V}^{\pi}, \mathcal{E}^{\pi}\right)\right]

其中 $\mathcal V^\pi$ 表示在 $\pi$ 这种排列下各个motif的标签， $\mathcal E^\pi$ 表示在 $\pi$ 这种排列下motifs之间相连的边。为了方便起见，作者假设所有可能的排列 $\pi$ 概率相同，在下文公式中会忽略掉上标 $\pi$ 。给定一种排列顺序，生成motif tree $\mathcal T(G)$ 的概率可以分解成：

\log p_{\theta}(\mathcal{V}, \mathcal{E})=\sum_{i=1}^{|\mathcal{V}|} \log p_{\theta}\left(\mathcal{V}_{i}, \mathcal{E}_{i} \mid \mathcal{V}_{<i}, \mathcal{E}_{<i}\right)

在第i步，使用前面1~i步已经生成的motif来指导第i个motif的生成。上面公式描述的其实就是一个自回归生成过程。

下面作者提出了BFS和DFS两种生成方式（和树的广度优先、深度优先一样），如下图3所示：

在DFS规则下，先进行topological预测，即该节点是否有孩子节点。如果需要为其生成孩子节点，则再预测新生成的节点的label。然后重复这一过程。
在BFS规则下，按层生成新的motif节点，对于每一层，MGSSL分别进行一次拓扑预测和结构预测。

在每一步，新加入的motif节点会接收来自其他已经生成的所有motif节点的信息进行预测。在motif tree的构建过程中，信息通过message向量 $h_{ij}$ 进行传播。

具体来说， $\hat{\mathcal E_t}$ 表示在t时刻motif tree中的边信息， $x_i$ 表示visited motif i在时刻t的embedding。message向量 $h_{i,j}$ 的计算方法如下：

\mathrm{h}_{i, j}=\operatorname{GRU}\left(x_{i},\left\{\mathrm{~h}_{k, i}\right\}_{(k, i) \in \hat{\mathcal{E}}_{t}, k \neq j}\right)

其中GRU表示Gated Recurrent Unit（一种比较古老的序列预测模型，我没有仔细研究）：

\begin{aligned} &s_{i, j}=\sum_{(k, i) \in \hat{\mathcal{E}}_{t}, k \neq j} \mathrm{~h}_{k, i} \\ &z_{i, j}=\sigma\left(\mathrm{W}^{z} x_{i}+\mathrm{U}^{z} s_{i, j}+b^{z}\right) \\ &r_{k, i}=\sigma\left(\mathrm{W}^{r} x_{i}+\mathrm{U}^{r} \mathrm{~h}_{k, i}+b^{r}\right) \\ &\tilde{\mathrm{h}}_{i, j}=\tanh \left(\mathrm{W} x_{i}+U \sum_{k=\mathcal{N}(i) \backslash j} r_{k, i} \odot \mathrm{h}_{k, i}\right) \\ &\mathrm{h}_{i, j}=\left(1-z_{i j}\right) \odot s_{i j}+z_{i j} \odot \tilde{\mathrm{h}}_{i, j} \end{aligned}

得到消息向量后，开始进行拓扑预测和标签预测：

拓扑预测

当MGSSL访问motif i时需要进行二分预测，即motif i是否有孩子节点。预测概率计算方法如下：

p_{t}=\sigma\left(U^{d} \cdot \tau\left(W_{1}^{d} x_{i}+W_{2}^{d} \sum_{(k, i) \in \hat{\mathcal{E}}_{t}} h_{k, i}\right)\right)

标签预测

下式计算的是父亲motif i生成的孩子motif j的标签

q_{j}=\operatorname{softmax}\left(U^{l} \tau\left(W^{l} h_{i j}\right)\right)

假设 $\hat p_t\in\{0,1\}$ 和 $\hat q_j$ 分别表示真实拓扑结构和标签信息，那么motif generation损失就定义上述两个预测的交叉熵损失：

\mathcal{L}_{\text {motif }}=\sum_{t} \mathcal{L}_{\text {topo }}\left(p_{t}, \hat{p}_{t}\right)+\sum_{j} \mathcal{L}_{\text {pred }}\left(q_{j}, \hat{q}_{j}\right)

2.3 多层级自监督预训练

为了捕获分子图中的多尺度信息，MGSSL被设计成如上图1所示的层次框架，包含Atom-level和Motif-level两个任务。所谓Atom-level，即利用GNNs学习到的节点/边嵌入，来预测原子类型和边类型，两者的交叉熵损失分别定义为 $\mathcal L_{atom}$ 和 $\mathcal L_{bond}$ 。