【图表示学习】图预训练模型 GraphBERT

Graph-BERT:Only Attention is Needed for Learning Graph Representations

一、简介

依赖于图结构的图神经网络(GNN)面临的一些问题，例如：suspended animation problem和over-smoothing problem。而且由于图本质上是相互连接的，阻碍了图上的并行运算。论文提出了图预训练模型Graph-BERT，该模型只基于注意力机制，且不依赖图的连接。

二、符号

1. 小写字母表示标量$(\text{e.g.,}x)$，小写黑体表示列向量$(\text{e.g.,}\text{x})$，大写黑体表示矩阵$(\text{e.g.,}\text{X})$，大写手写字体表示集合或者高阶张量$(\text{e.g.,}\mathcal{X})$。
2. 给定矩阵$\text{X}$，其中$\text{X}(i,:)$和$\text{X}(:,j)$表示$i$行和$j$列，$\text{X}(i,j)$表示第$i$行第$j$列的元素。
3. ${\text{X}}^T$和${\text{x}}^T$表示矩阵和向量的转置。
4. $\Vert \text{x}{\Vert }_p=({\sum }_i|\text{x}(i){|}^p{)}^{\frac{1}{p}}$表示向量$\text{x}$的$L_p$范数。
5. $\Vert \text{X}{\Vert }_F=({\sum }_{i,j}|\text{X}(i,j{)}^2|{)}^{\frac{1}{2}}$表示矩阵的$\text{Frobenius}$范数。
6. $\text{x}\otimes \text{y}$表示向量对应元素相乘后的向量，即$\text{element-wise product}$。
7. $\text{x}\bigsqcup \text{y}$表示向量拼接。

三、无连接子图采样

1.图的定义

图(graph)可以表示为5元组
$$G=(\mathcal{V},\mathcal{E},w,x,y)$$
其中，$\mathcal{V}$和$\mathcal{E}$表示节点集合和边集合；$w:\mathcal{E}\to \mathbb{R}$表示边到权重的映射；$x:\mathcal{V}\to \mathcal{X}$表示节点到对应特征的映射；$y:\mathcal{V}\to \mathcal{Y}$表示节点到标签的映射。

2.基本符合和假定

a. 给定无权图$G$，对于所有的边$\forall (v_i,v_j)\in \mathcal{E}$，则$w(v_i,v_j)=1$；否则，$\forall (v_i,v_j)\in \mathcal{V}\times \mathcal{V}\setminus \mathcal{E}$，则$w(v_i,v_j)=0$。(即两节点间存在边，则权重为1；否则，权重为0)

b. $\mathcal{X}$和$\mathcal{Y}$表示特征空间和标签空间，论文中为了简化表示则令$\mathcal{X}={\mathbb{R}}^{d_x}$和$\mathcal{Y}={\mathbb{R}}^{d_y}$。

c. 给定一个节点$v_i$，其特征和标签可以表示为${\text{x}}_i=x(v_i)\in {\mathbb{R}}^{d_x}$和${\text{y}}_i=y(v_i)\in {\mathbb{R}}^{d_y}$。

3.无连接子图采样

$\text{Graph-BERT}$的输入不是完整的图，而是通过采样的无连接子图，这样就可以实现对图的并行化处理。

3.1 亲密度矩阵

论文提出了采样方式主要是基于亲密度矩阵$\text{S}\in {\mathbb{R}}^{|\mathcal{V}|\times |\mathcal{V}|}$，其每个元素$\text{S}(i,j)$表示节点$v_i$和$v_j$的亲密度分数。具体来说，论文使用$\text{PageRank}$来计算节点间的亲密度。

3.2 节点上下文

给定一个图$G$和亲密度矩阵$\text{S}$，则节点$v_i$的上下文为集合 Γ ( v i ) = { v j ∣ v j ∈ V ∖ { v i } ∧ S ( i , j ) ≥ θ i } \Gamma(v_i)=\{v_j|v_j\in\mathcal{V}\setminus\{v_i\}\land\textbf{S}(i,j)\geq\theta_i\} Γ(vi​)={vj​∣vj​∈V∖{vi​}∧S(i,j)≥θi​}，其中${\theta }_i$表示节点$v_i$的最小亲密度阈值。(直观来讲，$v_i$的上下文就是与其亲密度高于${\theta }_i$的所有节点)

论文中则令${\theta }_i$表示$\text{sorted}(\text{S}(i,:))$的第$k$个分量。

3.3 针对节点$v_i$的无连接子图采样

令$g_i$表示节点$v_i$对于的采样无连接子图，其可以表示为$g_i=({\mathcal{V}}_i,\varnothing )$，其中节点集合为 V i = { v i } ∪ Γ ( v i ) \mathcal{V}_i=\{v_i\}\cup\Gamma(v_i) Vi​={vi​}∪Γ(vi​)。(直观来说，$g_i$中的节点由$v_i$和其上下文组成，而且$g_i$中并没有边)

3.4 在图$G$上采样

给定图$G$，针对图中的每个节点进行无连接子图采样得到一个子图集合 G = { g 1 , g 2 , … , g ∣ V ∣ } \mathcal{G}=\{g_1,g_2,\dots,g_{|\mathcal{V}|}\} G={g1​,g2​,…,g∣V∣​}。这样，图$G$就转换为子图集合$\mathcal{G}$。如果$\mathcal{G}$尺度太大，还可以对$\mathcal{G}$进行采样，例如$\mathcal{B}\subseteq \mathcal{G}$。

四、输入节点的Embedding

虽然图中的节之间并没有顺序，但是为了方便模型的输入，仍然会将子图中的节点转换为有序的列表。具体来说，将无连接子图$g_i\in \mathcal{B}$中的所有节点${\mathcal{V}}_i$转换为节点链表$[v_i,v_{i,1},\dots ,v_{i,k}]$，对于任意两个节点$\forall v_{i,j},v_{i,m}\in {\mathcal{V}}_i$，若$\text{S}(i,j)>\text{S}(i,m)$，则节点$v_{i,j}$会排在$v_{i,m}$之前。

1. 节点原始特征Embedding

对于子图$g_i$中的每个节点$v_j\in {\mathcal{V}}_i$，将原始特征${\text{x}}_j$转换为新嵌入特征
$${\text{e}}_j^{(x)}=Embed({\text{x}}_j)\in {\mathbb{R}}^{d_h\times 1}$$

2. Weisfeiler-Lehman绝对角色Embedding

Weisfeiler-Lehman算法可以根据节点在图中的角色来标记节点。更正式来说，对于节点$v_j\in {\mathcal{V}}_i$的$\text{WL}$编码可以表示为$\text{WL}(v_i)\in \mathbb{N}$，即为每个节点分配一个代表角色的非负整数。使用$\text{Transformer}$中的Position Embedding方法对节点的角色$\text{WL}(v_i)$进行嵌入，即
$$\begin{gathered} {\text{e}}_j^{(r)}=\text{Position-Embed}(WL(v_j)) \\ =\left[sin\right(\frac{WL(v_j)}{10000^{\frac{2l}{d_h}}}),cos(\frac{WL(v_j)}{10000^{\frac{2l+1}{d_h}}}){]}_{l=0}^{\lfloor \frac{d_h}{2}\rfloor } \end{gathered}$$
其中，${\text{e}}_j^{(r)}\in {\mathbb{R}}^{d_h\times 1}$。

3. 基于亲密度的相对位置Embedding

基于Weisfeiler-Lehman算法得到的节点角色是全局性质的，本文使用前面的序列化方式来捕获局部信息。具体来说，将图$g_i\in \mathcal{B}$中所有节点转换为序列$[v_i,v_{i,1},\dots ,v_{i,k}]$，基于这个序列为每个节点分配一个位置。其中节点$v_i$的位置为$P(v_i)=0$，节点$v_{i,1}$的位置为$P(v_{i,1})=1$，以此类推。(注：相同的节点在不同子图中的相对位置不同)

基于上面的相对位置，仍然使用$\text{Transformer}$中的$\text{Position-Embed}$进行相对位置的嵌入
$${\text{e}}_j^{(p)}=\text{Position-Embed}(P(v_j))\in {\mathbb{R}}^{d_h\times 1}$$

4. 基于跳数的相对位置Embedding

前面两种位置的嵌入分别是全局和局部的，这个基于跳数的相对位置嵌入可以看做是前面两种位置嵌入的平衡。正式来说，对于子图$g_i$中的任意节点$v_j\in {\mathcal{V}}_i$，其与节点$v_i$的相对距离为两节点在原始图中的跳数，表示为$H(v_j;v_i)$。将这个相对位置进行嵌入
$${\text{e}}_j^{(d)}=\text{Position-Embed}(H(v_j;v_i))\in {\mathbb{R}}^{d_h\times 1}$$

五、Graph-Transformer和Graph-BERT

1.$\text{Graph-Transformer}$

对于子图$g_i$中的节点$v_j$，其向量表示由上面4中嵌入向量聚合而成，即
$${\text{h}}_j^{(0)}=\text{Aggregate}({\text{e}}_j^{(x)},{\text{e}}_j^{(r)},{\text{e}}_j^{(p)},{\text{e}}_j^{(d)})$$
本文中聚合函数$\text{Aggregate}$为向量求和。整个子图$g_i$中的节点会被组织为矩阵
$${\text{H}}^{(0)}=[{\text{h}}_i^{(0)},{\text{h}}_{i,1}^{(0)},\dots ,{\text{h}}_{i,k}^{(0)}{]}^T\in {\mathbb{R}}^{(k+1)\times d_h}$$
定义一种用于图的$\text{Transformer}$，称为$\text{G-Transformer}$。其原理如下：
$$\{\begin{array}{ll}{\text{H}}^{(l)}& =\text{G-Transformer}({\text{H}}^{(l-1)})\\ & =\text{softmax}(\frac{{\text{QK}}^T}{\sqrt{d_h}}\text{V}+\text{G-Res}({\text{H}}^{(l-1)},{\text{X}}_i)),\end{array}$$
$\text{Q、K、V}$是基于${\text{H}}^{(l-1)}$的：
$$\{\begin{array}{l}\text{Q}={\text{H}}^{(l-1)}{\text{W}}_Q^{(l)}\\ \text{K}={\text{H}}^{(l-1)}{\text{W}}_K^{(l)}\\ \text{V}={\text{H}}^{(l-1)}{\text{W}}_V^{(l)}\end{array}$$
上式中，${\text{W}}_Q^{(l)},{\text{W}}_K^{(l)},{\text{W}}_V^{(l)}\in {\mathbb{R}}^{d_h\times d_h}$。

其中，$\text{G-Res}({\text{H}}^{(l-1)},{\text{X}}_i)$是图残差项，其中${\text{X}}_i\in {\mathbb{R}}^{(k+1)\times d_x}$是子图$g_i$所有节点的原始特征。

2.$\text{Graph-BERT}$

$\text{Graph-BERT}$是多个$\text{G-Transformer}$堆叠而成，具体来说：
{ H ( 0 ) = [ h i ( 0 ) , h i , 1 ( 0 ) , … , h i , k ( 0 ) ] T H ( l ) = G-Transformer ( H ( l − 1 ) ) , ∀ l ∈ { 1 , 2 , … , D } z i = Fusion ( H ( D ) ) \begin{cases} \textbf{H}^{(0)}=[\textbf{h}_i^{(0)},\textbf{h}_{i,1}^{(0)},\dots,\textbf{h}_{i,k}^{(0)}]^T \\ \textbf{H}^{(l)}=\text{G-Transformer}(\textbf{H}^{(l-1)}),\forall l\in\{1,2,\dots,D\} \\ \textbf{z}_i=\text{Fusion}(\textbf{H}^{(D)}) \end{cases} ⎩⎪⎨⎪⎧​H(0)=[hi(0)​,hi,1(0)​,…,hi,k(0)​]TH(l)=G-Transformer(H(l−1)),∀l∈{1,2,…,D}zi​=Fusion(H(D))​
在NLP任务中，Transformer模型的目标是学习所有输入token的表示。本文的$\text{Graph-BERT}$目标是获得子图$g_i$中节点$v_i$的表示。函数$\text{Fusion}(\cdot )$的作用是将子图$g_i$中所有节点的表示进行平均，${\text{z}}_i\in {\mathbb{R}}^{d_h\times 1}$是目标节点$v_i$的最终状态。

六、Graph-BERT上的预训练任务

基于模型Graph-BERT，提出了两个预训练任务。

1.节点属性重构(Node Raw Attribute Reconstruction)

针对子图$g_i$中的节点$v_i$，$\text{GRAPH-BERT}$会产生相应的向量表示${\text{z}}_i$。节点属性重构任务就是，使用${\text{z}}_i$来恢复节点$v_i$的原始特征，即生成原始特征${\hat{x}}_i=\text{FC}({\text{z}}_i)$。定义损失函数来缩小生成原始特征${\hat{x}}_i$与真正原始特征$x_i$的差距
$$l_1=\frac{1}{|\mathcal{V}|}\sum _{v_i\in \mathcal{V}}\Vert x_i-{\hat{x}}_i{\Vert }_2$$

2.图结构恢复(Graph Structure Recovery)

为了使模型能够捕获到图的结构信息，本文还设计了图结构恢复任务。对于任意两个节点$v_i$和$v_j$，计算其向量表示的余弦相似度，即${\hat{s}}_{i,j}=\frac{z_i^T{z}_j}{\Vert z_i\Vert \Vert z_j\Vert }$。用这个相似度来表示两个节点的连接程度。

计算图$G$中所有节点间的余弦相似度，并组成相似度矩阵$\hat{\text{S}}\in {\mathbb{R}}^{|\mathcal{V}|\times |\mathcal{V}|}$，其中$\hat{\text{S}}(i,j)={\hat{s}}_{i,j}$。

定义损失函数缩小$\hat{\text{S}}$和亲密度矩阵$\text{S}$之间的查询
$$l_2=\frac{1}{|\mathcal{V}{|}^2}\Vert \text{S}-\hat{\text{S}}{\Vert }_F^2$$

七、模型迁移与微调Graph-BERT

将$\text{Graph-BERT}$应用在节点分类和图聚类任务中。

1. 节点分类

基于节点$v_i$的向量表示${\text{z}}_i$，计算其分类标签${\hat{y}}_i=\text{softmax}(\text{FC}(z_i))$。设输入的batch为$\mathcal{T}$，计算交叉熵损失函数
$$l_{nc}=\sum _{v_i\in \mathcal{T}}\sum _{m=1}^{d_y}-y_i(m)\log {\hat{y}}_i(m)$$
其中，$d_y$是节点的类别数量。

2.图聚类

假设存在$l$个簇 C = { C 1 , C 2 , … , C l } \mathcal{C}=\{\mathcal{C}_1,\mathcal{C}_2,\dots,\mathcal{C_l}\} C={C1​,C2​,…,Cl​}。对于簇${\mathcal{C}}_j\in \mathcal{C}$，以计算簇中心变量$u_j={\sum }_{v_i\in {\mathcal{C}}_i}{z}_i\in {\mathbb{R}}^{d_h}$。图聚类的目的就是将相似节点聚在相同的簇中，将不同的节点划分到不同簇中。图聚类的目标函数可以表示为：
$$\underset{u_1,\dots ,u_l}{\min }\underset{\mathcal{C}}{\min }\sum _{j=1}^l\sum _{v_i\in \mathcal{C}}\Vert z_i-u_j{\Vert }_2$$