Introduction

在GNN的neighborhood aggregation中，对于拥有很少邻居的节点，在聚合过程中是否充分从邻居中获得了信息是一个问题。为解决该问题， 本文提出为每个节点做局部增强，即以中心节点为条件，学习邻居节点表示的分布。为了在局部邻域中生成一些样本来提升中心节点的neighborhood aggregation，本文提出一种数据增强框架：LA-GNNs， 以局部结构和中心节点特征为条件，生成neighborhood features。具体来说，在pre-training 阶段，通过一个生成模型，以中心节点的特征为条件来学习邻居特征的条件概率分布。然后利用这个邻居特征分布来生成中心节点的增强邻居特征。另外，通过pre-training来学习邻居增强特征生成器的过程是与下游任务无关的，所以该生成器生成的增强特征可以应用于其他GNN模型。

Local Augmentation for Graph Neural Networks (LAGNN)

Motivation

GNN在message passing的过程利用局部信息聚合来得到node representations。 但是对于邻居数量较少的节点，从邻居中得到的信息可能会不足。为了为节点$v$的邻域中$\mathcal{N}_v$生成更多样本，就需要知道邻居表示的分布。 由于一个节点邻居分布是与中心节点相关，所以我们要以中心节点$v$的representation为条件，学习它的邻居表示分布。

Approach

本文利用Conditional Variational Auto-Encoder (CVAE) 来学习给定中心节点$v$，邻居$u \in \mathcal{N}_v$的节点特征的条件分布。给定中心节点特征$\boldsymbol{X}_v$，关于中心节点的邻居分布为$p_\theta(\boldsymbol{X}_u | \boldsymbol{X}_v)$。定义隐变量$\mathbf{z}$，则先验可以定义为$p_\theta(\mathbf{z}|\boldsymbol{X}_v)$。结合隐变量$\mathbf{z}$，邻居特征$\boldsymbol{X}_u$的分布可以改写为如下形式： $$ \begin{aligned} \log p_\theta(\boldsymbol{X}_u | \boldsymbol{X}_v) &= \log \frac{p_\theta(\boldsymbol{X}_u , \boldsymbol{X}_v)}{p_\theta( \boldsymbol{X}_v)}= \frac{p_\theta(\boldsymbol{X}_u , \boldsymbol{X}_v)p_\theta(\mathbf{z}|\boldsymbol{X}_u , \boldsymbol{X}_v)}{p_\theta( \boldsymbol{X}_v)p_\theta(\mathbf{z}|\boldsymbol{X}_u , \boldsymbol{X}_v)} \\ &=\log \frac{p_\theta(\boldsymbol{X}_u , \boldsymbol{X}_v, \mathbf{z})}{p_\theta( \boldsymbol{X}_v)p_\theta(\mathbf{z}|\boldsymbol{X}_u , \boldsymbol{X}_v)} \\ &= \log \frac{p_\theta(\boldsymbol{X}_u , \mathbf{z}| \boldsymbol{X}_v)}{p_\theta(\mathbf{z}|\boldsymbol{X}_u , \boldsymbol{X}_v)}\\ \end{aligned} $$ 假设隐变量$\mathbf{z}$的分布为$q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)$， 左右两边对分布$q_\phi$计算期望，左边： $$ \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log p_\theta(\boldsymbol{X}_u | \boldsymbol{X}_v) dz = \log p_\theta(\boldsymbol{X}_u | \boldsymbol{X}_v) $$ 右边： $$ \begin{aligned} &\int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta(\boldsymbol{X}_u , \mathbf{z}| \boldsymbol{X}_v)}{p_\theta(\mathbf{z}|\boldsymbol{X}_u , \boldsymbol{X}_v)} dz \\ =& \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \left(\frac{p_\theta(\boldsymbol{X}_u , \mathbf{z}| \boldsymbol{X}_v)}{ q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)} \cdot \frac{ q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)}{p_\theta(\mathbf{z}|\boldsymbol{X}_u , \boldsymbol{X}_v)}\right) dz \\ =& \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta(\boldsymbol{X}_u , \mathbf{z}| \boldsymbol{X}_v)}{ q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)} dz + \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \frac{ q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)}{p_\theta(\mathbf{z}|\boldsymbol{X}_u , \boldsymbol{X}_v)}dz \\ =& \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta(\boldsymbol{X}_u , \mathbf{z}| \boldsymbol{X}_v)}{ q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)} dz + K L\left(q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) || p_\theta\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)\right) \\ \geq& \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta\left(\boldsymbol{X}_u, \mathbf{z} \mid \boldsymbol{X}_v\right)}{q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)} \mathrm{d} \mathbf{z} = ELBO \end{aligned} $$ Evidence Lower Bound (ELBO) 可以写为 $$ \begin{aligned} L_{ELBO} &= \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta\left(\boldsymbol{X}_u, \mathbf{z} \mid \boldsymbol{X}_v\right)}{q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)} \mathrm{d} \mathbf{z} \\ &= \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta\left(\boldsymbol{X}_u, \boldsymbol{X}_v, \mathbf{z}\right)}{q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) p_\theta\left(\boldsymbol{X}_v\right)} \mathrm{d} \mathbf{z} \\ &= \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta\left(\boldsymbol{X}_u \mid \boldsymbol{X}_v, \mathbf{z}\right) p_\theta\left(\boldsymbol{X}_v, \mathbf{z}\right)}{q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) p_\theta\left(\boldsymbol{X}_v\right)} \mathrm{d} \mathbf{z} \\ &= \int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log \frac{p_\theta\left(\boldsymbol{X}_u \mid \boldsymbol{X}_v, \mathbf{z}\right) p_\theta\left(\mathbf{z} \mid \boldsymbol{X}_v\right)}{q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)} \mathrm{d} \mathbf{z} \\ &= -K L\left(q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) || p_\theta\left(\mathbf{z} \mid \boldsymbol{X}_v\right)\right)+\int_z q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) \log p_\theta\left(\boldsymbol{X}_u \mid \boldsymbol{X}_v, \mathbf{z}\right) \mathrm{d} \mathbf{z} \\ &= -K L\left(\underbrace{q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right)}_{Encoder} || \underbrace{ p_\theta\left(\mathbf{z} \mid \boldsymbol{X}_v\right)}_{\text{Normal Distribution}}\right) + \mathbb{E}_{\mathbf{z} \sim q_\phi\left(\mathbf{z} \mid \boldsymbol{X}_u, \boldsymbol{X}_v\right) }\log p_\theta\left(\boldsymbol{X}_u \mid \boldsymbol{X}_v, \mathbf{z}\right) \end{aligned} $$ 在CVAE pre-training的过程中，第一项KL中CVAE Encoder 的一对邻接节点对，对于该节点对，输出一组分布参数均值$\mu$和方差$\sigma$，作为隐变量$z$的分布参数，第一项的优化目标使得编码器输出的分布接近Normal Distribution。然后利用reparameterization trick可微的从生成的$\mathbf{z}$分布中采样一个encoding:

def reparameterize(means, log_var):
    std = torch.exp(0.5 * log_var)
    eps = torch.randn_like(std)
    return means + eps * std   // z

若当前输入节点对为$(\boldsymbol{X}_v, \boldsymbol{X}_u)$，从输出的分布中采样一个encoding $\mathbf{z}$然后将$\mathbf{z}$输入decoding中，使得用$\mathbf{z}$和中心节点$\boldsymbol{X}_v$可以重构邻接节点$\boldsymbol{X}_u$。用所有邻接节点对训练encoder参数$\phi$和generator参数$\phi$。 这样在下游预测任务前，直接从Normal Distribution 随机采样$\mathbf{z}$，拼接中心节点$\boldsymbol{X}_v$输入generator中就可以为节点$v$生成增强邻居。过程如下图所示：

Introduction

对于监督对比学习（Supervised Contrastive Learning, SupCon）, SupCon loss旨在表示空间中拉近属于同一个class的数据点，分离不同类的数据点。 但是SupCon难以处理高类内方差，类间相似度较大的数据集。为了解决该问题，本文提出了Cluster-aware supervised contrastive learning loss (ClusterSCL)。什么是高类内方差，高跨类相似度问题？如图1(a)所示，节点$u_1$和$u_3$ 是同类节点，$u_2$和$u_4$是同类节点。他们是同类节点但在不同的社区中，所以类内方差较大，即同一个类内的节点跨越了多个community。 另外$u_1$和$u_2$， $u_3$和$u_4$，是不同类的节点对， 但他们处在同一个社区中，导致在MPNN过程中，这些处在同一个community中的不同类节点被拉近，导致跨类相似度较高的问题。

如果对节点$u_2$计算SupCon时，如图1(b)所示，SupCon会使得同类节点被拉近，如$u_2$和$u_4$会被拉近。但是$u_3$和$u_4$处在同一个社区中（structurally similar）那么MPNN会使得$u_3$和$u_4$被拉近，所以SupCon在拉近$u_2$和$u_4$的同时，会间接拉近不同类节点$u_2$和$u_3$。同时，对于构成negative pairs的不同类节点，例如$u_1$和$u_2$，SupCon会推远$u_1$和$u_2$，但是$u_1$和$u_5$ structurally similar, 因此会推远$u_1$和$u_2$会间接导致$u_2$和$u_5$这两个同类节点被推远。因此对于一个cluster内节点不同类，且不同cluster中存在同类节点的情况，会导致复杂的决策边界，即在拉近同类但不同社区的节点时，也会间接拉近不同类不同社区的节点。在推远不同类同社区的节点时，也可能间接推远同类同社区的节点。

为了解决上述问题，最直接的方法是对于每个cluster，如图1(a)的Community 1，不考虑其他cluster，只对当前cluster内节点做SupCon。但是这么做忽略了跨cluster的同类节点交互，如$u_1$和$u_3$，$u_2$和$u_4$，这些跨cluster的positive pairs可能包含有益的信息。为了解决这个问题，本文提出cluster-aware data augmentation (CDA) 聚类感知的数据增强，来为每个节点生成augmented positives and negatives，如图1(b)中ClusterSCL所示。对于每个节点$u$，为它生成positive 和negative samples, 生成的samples 位于或接近$u$所在的cluster。Recall SupCon存在的问题：

• SupCon会使得$u_2$和$u_4$被拉的太近，从而间接导致$u_2$和$u_3$被拉近，所以对于high intra-class variances，要求不同cluster的同类节点如$u_2$和$u_4$不要被拉太近；
• SupCon会使得$u_1$和$u_2$被推远，从而间接导致$u_2$和$u_5$被推远，所以对于high inter-class similarity，要求同一个cluster内的不同类节点如$u_1$和$u_2$不要被拉的太远。

Method

Two stage training with Supervised Contrastive Loss

SupCon encourages samples of the same class to have similar representations, while pushes apart samples of different classes in the embedding space.

First Stage: 计算node embeddings $H = g_\theta(G)$， 然后用SupCon Loss来训练 $g_\theta$ 。即已经知道训练集中的节点label，基于这些节点label，SupCon在embedding space中把同label的节点拉近，不同label的节点分开。$g_\theta$用于得到node embeddings.

Second Stage：基于学习好的$g_\theta$， 用$\hat{Y}=f_\phi\left(g_\theta(G)\right)$来得到logits/prediction。即用cross-entropy loss来训练$f_\theta$。

SupCon

从同一个类中采样的节点构成positive pairs。batch中随机采样的节点对为negative pairs。给定$N$个随机采样的节点，对于每个节点，从其对应的class中随机采样一个不为它的节点作为positive pairs。所以一个batch有$N$对positive pairs，共$2N$个节点。

用$I \equiv\{1,2, \ldots, 2 N\}$ 表示一个batch中的node indices。$s_i \in I$表示这$2N$个节点中与节点$v_i$属于同一类的节点的indices。如下式所示， 在一个batch中，令$S_i \subset I$表示$2N$个节点中可以与节点$v_i$构成positive pairs的节点集合。相比其他节点，SupCon的objective是拉近positive pairs。 $$ \max \sum_{i \in I} \frac{1}{\left|S_i\right|} \sum_{s_i \in S_i} \log \frac{\exp \left(\mathbf{h}_i^{\top} \mathbf{h}_{s_i} / \tau\right)}{\sum_{j \in I \backslash\{i\}} \exp \left(\mathbf{h}_i^{\top} \mathbf{h}_j / \tau\right)} $$ 但是如果batch中与$v_i$同类的节点$v_j$和它不属于同一个cluster，$v_j$所属的cluster不同类的节点较多，那么拉近他们的距离也会间接拉近$v_i$与不同类节点间的距离。

ClusterSCL

为了解决上述问题，本文提出ClusterSCL。

Cluster-aware Data Augmentation (CDA)

定义隐变量 $c_i$，该隐变量取值范围为$c_i \in \{1,2, \ldots, M\}$ 表示节点$v_i$属于哪一个cluster。给定两个anchor node $v_i$,$v_j$，CDA使用线性插值法为$v_j$生成augmentation： $$ \tilde{\mathbf{h}}_j=\alpha \mathbf{h}_j+(1-\alpha) \mathbf{w}_{c_i} \tag{4} $$ 其中$c_i$指示了节点$v_i$所在的cluster。$\mathbf{w}=\left\{\mathbf{w}_m\right\}_{m=1}^M$ 表示每个cluster的cluster prototypes，即每个cluster的中心，serve to characterize the cluster。$\mathbf{w}_{c_i}$表示节点$v_i$所在cluster的中心表示。$\tilde{\mathbf{h}}_j$包含了节点$v_j$的信息。并且，由于$\mathbf{w}_{c_i}$是$v_i$所在cluster的prtotype，所以通过调整$\alpha$，可以使$\tilde{\mathbf{h}}_j$位于$v_i$所在cluster的附近或内部。

（1）如果$(v_i, v_j)$是一个batch中的positive pair，$v_i$是anchor节点，如果$v_j$位于$v_i$所在的cluster $c_i$内，那么就需要学到的$\mathbf{h}_j$与$\mathbf{h}_i$尽可能靠近。在SupCon中，需要设置较大的$\alpha$使得$\tilde{\mathbf{h}}_j$保留更多$\mathbf{h}_j$ 。对于anchor节点$\mathbf{h}_i$，将它与$\tilde{\mathbf{h}}_j$拉近的时候，由于$\tilde{\mathbf{h}}_j$保留了更多$v_j$特征，所以$\mathbf{h}_j$也会被和$\mathbf{h}_i$拉近。如图2(a)所示。

（2）如果$(v_i, v_j)$是positive pair，如果$v_j$位于$v_i$所在的cluster $c_i$外时，如果$v_j$周围有negative samples （不一定在该batch中），那么直接拉近$(v_i, v_j)$会间接导致潜在的negative samples也会被拉近，因此对于位于$v_i$所在cluster外的节点$v_j$，要求它最终的表示$\mathbf{h}_j$不能被拉的太近，此时就需要小一些的$\alpha$，使得$\tilde{\mathbf{h}}_j$保留少一些$v_j$的信息，那么在SupCon拉近$\mathbf{h}_i$和$\tilde{\mathbf{h}}_j$的过程不会导致$\mathbf{h}_j$被拉近太多。因为$ \mathbf{w}_{c_i}$占据了$\tilde{\mathbf{h}}_j$的大部分，且它与$\mathbf{h}_i$已经很接近。如图2(b)所示。

对于negative pair $(v_i, v_j)$。如果$v_j$位于$v_i$所在的cluster $c_i$内，如果直接推远$\mathbf{h}_i$和$\mathbf{h}_j$，会导致如果$v_j$的邻居有$v_i$的positive sample，那么这个positive sample也会被间接推远。所以$\tilde{\mathbf{h}}_j$应该保留较少的$\mathbf{h}_j$，即$\alpha$应该小。但是本文不考虑negative pairs的这种情况了，直接套用posiive的CDA原则。

综合上面的（1）（2）即对于close positive pairs，要让他们尽可能接近，即$\alpha$要大，对于distant positive pairs，要让他们不要太接近，$\alpha$要小一些。所以$\alpha$应与positive pair之间的相似度相关。所以$\alpha$定义如下： $$ \alpha=\frac{\exp \left(\mathbf{h}_i^{\top} \mathbf{h}_j\right)}{\exp \left(\mathbf{h}_i^{\top} \mathbf{h}_j\right)+\exp \left(\mathbf{h}_i^{\top} \mathbf{w}_{c_i}\right)} $$ 上式分子越大$\alpha$越大，说明对于anchor node $v_i$， 如果的positive sample $v_j$位于它的cluster内（$v_i$与$v_j$相似）,$\mathbf{h}_j$的augmentation $\tilde{\mathbf{h}}_j$要保留越多自身信息。

Integraging Clustering and CDA into SupCon Learning

目标是给定节点$v_i$以及它所在的cluster 隐变量$c_i$，$v_i$的positive samples $s_i$的cluster-aware SupCon定义为条件概率： $$ \begin{aligned} p\left(s_i \mid v_i, c_i\right) &=\frac{\exp \left(\mathbf{h}_i^{\top} \tilde{\mathbf{h}}_{s_i} / \tau\right)}{\sum_{j \in V \backslash\{i\}} \exp \left(\mathbf{h}_i^{\top} \tilde{\mathbf{h}}_j / \tau\right)} \\ &=\frac{\exp \left(\mathbf{h}_i^{\top}\left(\alpha \mathbf{h}_{s_i}+(1-\alpha) \mathbf{w}_{c_i}\right) / \tau\right)}{\sum_{j \in V \backslash\{i\}} \exp \left(\mathbf{h}_i^{\top}\left(\alpha \mathbf{h}_j+(1-\alpha) \mathbf{w}_{c_i}\right) / \tau\right)} \end{aligned} \tag{6} $$ 即对于positive pair $(v_i, s_i)$，最大化$\mathbf{h}_i$和$s_i$的augmentation $\tilde{\mathbf{h}}_{s_i}$之间的一致性，$\alpha$可以依据$s_i$和$c_i$的关系来调整$\mathbf{h}_{s_i}$对于SupCon的贡献，使得$\mathbf{h}_{i}$与$\mathbf{h}_{s_i}$在位于不同cluster的情况下不会被拉的太近。其中$c_i$是隐变量。

首先定义关于节点$v_i$的cluster分布，即$v_i$属于每个$c_i$的概率。给定anchor node $v_i$，它属于cluster $c_i \in \{1,2, \ldots, M\}$的概率定义为： $$ p\left(c_i \mid v_i\right)=\frac{\exp \left(\mathbf{h}_i^{\top} \mathbf{w}_{c_i} / \kappa\right)}{\sum_{m=1}^M \exp \left(\mathbf{h}_i^{\top} \mathbf{w}_m / \kappa\right)} \tag{7} $$ $\mathbf{w}_{c_i}$是cluster $c_i$的prototype表示，$v_i$属于在embedding空间中与它相似的cluster的概率更高。ClusterSCL旨在最大化给定锚节点$v_i$，锚节点与其positive sample $s_i$的likelihood： $$ p\left(s_i \mid v_i\right)=\int_{c_i} p\left(c_i \mid v_i\right) p\left(s_i \mid v_i, c_i\right) d c_i= \sum^M_{m=1} p(m|v_i)p(s_i|v_i,m) \tag{8} $$ 即给定anchor $v_i$，$s_i$是$v_i$的postive sample的概率为 当$v_i$属于cluster $m$的情况下，$s_i$是其positive sample的概率， over all clusters $m \in M$。

在likelihood Eq.(8)中，anchor node $v_i$ 为待优化参数，它的positive sample $s_i$为观测数据，$c_i$为隐变量。

Maximize likelihood Eq.(8) via EM

Objective: $\mathrm{maximize} \log p(s_i| v_i)$，即最大化positive pairs的条件概率 given anchor node $v_i$。

E-step：$\mathbb{E}_{p(c_i|s_i, v_i)} \log p(s_i, c_i|v_i)$

M-step: $\widehat{v}_i = \arg \max_{v_i} \mathbb{E}_{p(c_i|s_i, v_i)} \log p(s_i, c_i|v_i)$

可见，如果要通过EM算法来优化得到anchor node $v_i$的表示，需要计算后验 $p(c_i|s_i, v_i)$： $$ \begin{aligned} p(c_i|s_i, v_i) & = \frac{p(c_i,v_i, s_i)}{p(s_i,v_i)} \\ & = \frac{p(v_i)p(s_i,c_i | v_i)}{p(s_i |v_i) p(v_i)}\\ & = \frac{p(s_i,c_i|v_i)}{p(s_i |v_i)} \\ & = \frac{p(s_i,c_i|v_i)}{\sum^M_{m=1}p(m|v_i)p(s_i | v_i, m)}\\ & = \frac{\frac{p(s_i,c_i,v_i)}{p(v_i)} = \frac{p(c_i,v_i)}{p(v_i)} \frac{p(s_i,c_i,v_i)}{p(c_i,v_i)} = p(c_i | v_i)p(s_i|c_i,v_i)}{\sum^M_{m=1}p(m|v_i)p(s_i | v_i, m)} \\ & = \frac{p(c_i | v_i)p(s_i|c_i,v_i)}{\sum^M_{m=1}p(m|v_i)p(s_i | v_i, m)}
\end{aligned}\tag{9} $$ 后验中，$p(c_i | v_i)$，$p(m|v_i)$是$v_i$的cluster 分布，在Eq.(7)中给出定义。但是，对于$p(s_i|c_i,v_i)$和$p(s_i | v_i, m)$，Eq.(6)中给出了它的定义，$p\left(s_i \mid v_i, c_i\right) =\frac{\exp \left(\mathbf{h}_i^{\top} \tilde{\mathbf{h}}_{s_i} / \tau\right)}{\sum_{j \in V \backslash\{i\}} \exp \left(\mathbf{h}_i^{\top} \tilde{\mathbf{h}}_j / \tau\right)}$， 可以看出分母部分需要计算$\mathbf{h}_i$与所有节点，并且还要over all $M$ cluster，因此后验难以计算。为了解决这个问题，我们可以maximize evidence lower bound (ELBO) of $\log p(s_i | v_i)$： $$ \begin{aligned} \log p(s_i | v_i) &= \log \frac{p(s_i,v_i)}{p(v_i)} = \log \frac{p(s_i,v_i)p(c_i | v_i,s_i)}{p(v_i)p(c_i | v_i,s_i)} \\ &= \log \frac{p(v_i,s_i,c_i)}{p(v_i)p(c_i | v_i,s_i)} \\ &= \log \frac{p(s_i|v_i,c_i)p(v_i,c_i)}{p(v_i)p(c_i | v_i,s_i)} = \log \frac{p(s_i|v_i,c_i)p(c_i|v_i)p(v_i)}{p(v_i)p(c_i | v_i,s_i)} \\ & = \log \frac{p(s_i|v_i,c_i)p(c_i|v_i)}{p(c_i | v_i,s_i)} \\ & = \log p(s_i|v_i,c_i)-\log p(c_i | v_i,s_i) + \log p(c_i|v_i) \\ & \text{引进一个关于隐变量$c_i$的分布$q(c_i)$，可以是任意形式，这里定义为$q(c_i|v_i,s_i)$}\\ & = \log p(s_i|v_i,c_i) - \log \frac{p(c_i | v_i,s_i)}{q(c_i|v_i,s_i)} + \log \frac{p(c_i|v_i)}{q(c_i|v_i,s_i)} \\ &\text{左右两边分别对$q(c_i|v_i,s_i)$求期望} \\ \text{左边} &= \int q(c_i|v_i,s_i) \log p(s_i | v_i) d c_i = \int q(c_i|v_i,s_i) d c_i \cdot \log p(s_i | v_i) = \log p(s_i | v_i) \\ \text{右边} &= \int q(c_i|v_i,s_i) \log p(s_i|v_i,c_i) dc_i - \underbrace{\int q(c_i|v_i,s_i) \log \frac{p(c_i | v_i,s_i)}{q(c_i|v_i,s_i)} dc_i}_{-\mathrm{KL}(q(c_i|v_i,s_i)||p(c_i | v_i,s_i))} + \underbrace{\int q(c_i|v_i,s_i) \log \frac{p(c_i|v_i)}{q(c_i|v_i,s_i)} dc_i}_{-\mathrm{KL}(q(c_i|v_i,s_i)||p(c_i|v_i))} \\ &= \int q(c_i|v_i,s_i) \log p(s_i|v_i,c_i) dc_i + \mathrm{KL}(q(c_i|v_i,s_i)||p(c_i | v_i,s_i)) - \mathrm{KL}(q(c_i|v_i,s_i)||p(c_i|v_i)) \\ & \geq \underbrace{\mathbb{E}_{q(c_i|v_i,s_i)} \log p(s_i|v_i,c_i) - \mathrm{KL}(q(c_i|v_i,s_i)||p(c_i|v_i))}_{ELBO} \end{aligned} \tag{10} $$ 因此， we have： $$ \log p(s_i | v_i) \geq ELBO = \mathbb{E}_{q(c_i|v_i,s_i)} \log p(s_i|v_i,c_i) - \mathrm{KL}(q(c_i|v_i,s_i)||p(c_i|v_i)) \tag{10} $$ 所以目标为找到node embeddings $\mathbf{h}_i, \forall i$， 最大化ELBO。接下来可以定义任意关于隐变量$c_i$的vartational distribution $q(c_i)$。 这里将关于$c_i$的variational distribution定义为用mini-batch近似后验的形式： $$ q\left(c_i \mid v_i, s_i\right)=\frac{p\left(c_i \mid v_i\right) \tilde{p}\left(s_i \mid v_i, c_i\right)}{\sum_{m=1}^M p\left(m \mid v_i\right) \tilde{p}\left(s_i \mid v_i, m\right)} \tag{11} $$ 其中$\tilde{p}\left(s_i \mid v_i, c_i\right)=\exp \left(\mathbf{h}_i^{\top} \tilde{\mathbf{h}}_{s_i} / \tau\right) / \sum_{j \in I \backslash\{i\}} \exp \left(\mathbf{h}_i^{\top} \tilde{\mathbf{h}}_j / \tau\right)$，与Eq.(6)不同的是$\tilde{p}\left(s_i \mid v_i, c_i\right)$的分母只计算mini-batch $I$内的negative samples。并且用$\tilde{p}\left(s_i \mid v_i, c_i\right)$来替换了Eq.10 中的$ p(s_i|v_i,c_i)$，文中证明了这种替代的合理性。

在E-step，定义了variational distribution $q\left(c_i \mid v_i, s_i\right)$和likelihood objective的ELBO，在M-step最大化ELBO。给定一个mini-batch $I$，目标函数为最大化$I$中每对positive pairs 的ELBO： $$ \max \quad \mathcal{L}_{\mathrm{ELBO}}(\theta, \mathbf{w} ; I) \approx \frac{1}{|I|} \sum_{i \in I} \frac{1}{\left|S_i\right|} \sum_{s_i \in S_i} \mathcal{L}_{\mathrm{ELBO}}\left(\theta, \mathbf{w} ; v_i, s_i\right) $$

数据 $X = \{x_1, \cdots x_N\}$, 模型参数为$\theta$，Likelihood 定义为 $P(X | \theta)$：当参数为$\theta$时，观测到给定数据$X$的概率。 $$ P(X|\theta) = L(\theta | X) = P_\theta(X) \tag{1} $$ 最大似然估计 （Maximum Likelihood Estimation, MLE）: $$ \theta_{\mathrm{MLE}} = \arg \max_\theta P(X|\theta) \tag{2} $$

最大似然估计：给定一组样本$X$，模型的参数$\theta$是研究对象。若能找到参数$\theta_{\mathrm{MLE}}$，使得样本发生的可能性最大，则此估计值$\theta_{\mathrm{MLE}}$为参数$\theta$的最大似然估计。

举例来说，如果模型是单个Gaussian Distribution下，参数为Gaussian Distribution的参数（均值$\mu$, 标准差$\Sigma$， $\theta = {\mu, \Sigma}$）. 给定一组数据$X$, 要计算$X$来自什么样的Gaussian，即：$P(\cdot | \theta) = f_\theta(\cdot) = \mathcal{N}(\cdot | \mu,\Sigma)$是一个Gaussian Distribution函数，目标为： $$ \theta_{\mathrm{MLE}} = \mu^\star, \Sigma^\star = \arg \max_{\mu,\Sigma} \sum^N_{i = 1} \log \mathcal{N}(x_i|\mu,\Sigma) \tag{3} $$ 即MLE的目标是找到最佳的高斯分布，是的从该分布中采样出数据$X$的概率最高。

如果只需要用一个Gaussian来拟合$X$的分布的话，这个Gaussian可以很容易用求导的方式获得$\theta_{\mathrm{MLE}}$的解析解： 对$\mu$求导：$\frac{\partial P(X|\mu,\Sigma)}{\partial \mu}$；对$\Sigma$求导：$\frac{\partial P(X|\mu,\Sigma)}{\partial \Sigma}$，令导数为0，即可求得最佳的$\mu$，$\Sigma$，使得对应的高斯分布符合数据$X = {x_1, \cdots x_N}$的分布。

但是，要用更复杂的模型（更多参数）来更准确的拟合$X$的分布，例如Gaussian Mixture Model，即多个Gaussian的组合，其模型参数为： $$ \theta = \{\underbrace{\mu_1, \cdots,\mu_K}_{\text{每个Gaussian的 mean参数}}, \underbrace{\Sigma_1,\cdots, \Sigma_K}_{\text{每个Gaussian的 std参数}}, \underbrace{\alpha_1, \cdots, \alpha_{K-1}}_{\text{每个Gaussian的权重}} \} \tag{4} $$ 假设是一个$K$个Gaussian的Gaussian Mixture Model，那么$\sum^K_{k = 1}\alpha_k = 1$。

给定数据$X = \{x_1, \cdots x_N\}$，若要用$K$维Gaussian Mixture Model来拟合该数据，就要优化所有$K$个Gaussian的均值参数，标准差参数，和权重参数，使得混合高斯分布采样出$X$的概率最大，即： $$ \begin{aligned} \theta_{\mathrm{MLE}} &= \mu_1^\star,\cdots \mu_K^\star,\Sigma_1^\star, \cdots, \Sigma_K^\star, \alpha^\star_1,\cdots,\alpha^\star_{K-1} \\ &=\underset{\theta}{\arg\max} \sum^N_{i = 1} \log \sum^K_{k=1} \alpha_k \mathcal{N}(x_i|\mu_k,\Sigma_k) \end{aligned} \tag{5} $$ 如果要得到上式的解析解，要对$\mu_1,\cdots,\mu_K, \Sigma_1, \cdots,\Sigma_K, \alpha_1, \cdots, \alpha_{K-1}$求导，再令导数为0来求解，这非常困难，由此引出EM算法。

期望最大算法

求解MLE问题时，在最大化log-likelihood: $$ \theta_{\mathrm{MLE}} = \arg \max_\theta \log P(X|\theta) \tag{6} $$ 难以直接对$\theta$求导来得到解析解时（如高斯混合模型情况），可以使用EM算法来迭代求解： $$ \theta^{(t+1)}=\underset{\theta}{\arg \max} \int_z \log P(X,z | \theta) \cdot P(z|X,\theta^{(t)}) dz \tag{7} $$ $X$为观测数据， $z$为latent variables（隐变量），隐变量必须不会影响$X$的边缘分布,即 $P(X) = \int_z P(X|z) P(z) dz$。

而公式(7)中 $$ \begin{aligned} &\int_z \underbrace{\log P(X,z | \theta)}_{\text{每个z对应的值}} \cdot \underbrace{P(z|X,\theta^{(t)})}_{\text{z的分布}} dz\\ =&\mathbb{E}_{P(z|X,\theta^{(t)})} \left[\log P(X,z | \theta)\right] \end{aligned} \tag{8} $$ 所以，期望最大化算法求参数$\theta$的迭代公式可改写为: $$ \theta^{(t+1)}=\arg \max_\theta \mathbb{E}_{P(z|X,\theta^{(t)})} \left[\log P(X,z | \theta)\right] \quad \text{期望最大化} \tag{9} $$

其中$P(z|X,\theta^{(t)})$为后验分布posterior。

EM算法收敛性证明

因为EM算法通过迭代的方式优化模型参数$\theta$，使得对数似然$\log P(X|\theta)$最大。通过公式(7)，可以保证在$\theta^{(t+1)}$参数下的模型比$\theta^{(t)}$参数下的模型更拟合数据分布。通过公式(7)迭代更新参数$\theta^{(t)} \to \theta^{(t+1)}$，可以使得$\log P(X|\theta)$变大。收敛性即证明： $$ \log P(X|\theta^{(t)}) \leq \log P(X|\theta^{(t+1)}) \tag{10} $$ 证明. $$ \begin{aligned} &\because P(X,z) = P(z|X) P(X)\quad \text{always true}, \text{then}\quad P(X) = \frac{P(X,z)}{P(z|X)} \\ &\therefore P(X|\theta) = \frac{P(X,z|\theta)}{P(z|X,\theta)} \\ &\therefore \log P(X|\theta) = \log P(X,z|\theta) - \log P(z|X,\theta) \end{aligned} \tag{11} $$ 上式左右两边对分布$P(z|X,\theta^{(t)})$求期望： $$ \mathbb{E}_{z \sim P(z|X,\theta^{(t)})} \underbrace{\left[\log P(X|\theta)\right]}_{\text{与z无关}} = \mathbb{E}_{z \sim P(z|X,\theta^{(t)})} \left[\log P(X,z|\theta) - \log P(z|X,\theta)\right] \tag{12} $$ 上式左边$=\log P(X|\theta)$，右边： $$ \begin{aligned} &\mathbb{E}_{P(z|X,\theta^{(t)})} \left[\log P(X,z|\theta) - \log P(z|X,\theta)\right] \\ =& \underbrace{\int_{z} P(z|X,\theta^{(t)}) \log P(X,z|\theta) dz}_{Q(\theta,\theta^{(t)})} - \underbrace{\int_z P(z|X,\theta^{(t)}) \log P(z|X,\theta) dz}_{H(\theta,\theta^{(t)})} \end{aligned} \tag{13} $$ 注意到$Q(\theta,\theta^{(t)}) = \int_{z} P(z|X,\theta^{(t)}) \log P(X,z|\theta) dz$ 就是EM算法的迭代更新函数，即$\theta^{(t+1)} = \arg \max_\theta Q(\theta,\theta^{(t)})$。结合公式(12)和公式(13)： $$ \log P(X|\theta) = Q(\theta,\theta^{(t)}) - H(\theta,\theta^{(t)}) \tag{14} $$

因此log-likelihood under $\theta^{(t)}$ and $\theta^{(t+1)}$： $$ \begin{aligned} \log P(X|\theta^{(t+1)}) &= Q(\theta^{(t+1)},\theta^{(t)}) - H(\theta^{(t+1)},\theta^{(t)}) \\ \log P(X|\theta^{(t)}) &= Q(\theta^{(t)},\theta^{(t)}) - H(\theta^{(t)},\theta^{(t)}) \end{aligned} \tag{15} $$

首先，根据EM的迭代求解公式，$\theta^{(t+1)}$由 $$ \theta^{(t+1)} = \arg \max_\theta Q(\theta,\theta^{(t)}) \tag{16} $$ 得到，所以$Q(\theta^{(t+1)},\theta^{(t)}) \geq Q(\theta,\theta^{(t)})$一定成立。所以下式成立： $$ Q(\theta^{(t+1)},\theta^{(t)}) \geq Q(\theta^{(t)},\theta^{(t)}) \tag{17} $$ 对于$H(\theta,\theta^{(t+1)})$，首先介绍Jensen Inequality:

Jensen Inequality：

If $g(x)$ is a convex function on $R_X$, and $\mathbb{E}[g(x)]$ and $g(\mathbb{E}[X])$ are finite, then $\mathbb{E}[g(x)] \geq g(\mathbb{E}[X])$。

显然$\log$是concave，所以$\mathbb{E}[\log(\cdot)]\leq \log(\mathbb{E}[\cdot])$。同理$-\log$是convex，所以$\mathbb{E}[-\log(\cdot)]\geq -\log(\mathbb{E}[\cdot])$。

下面，计算$H(\theta^{(t)},\theta^{(t)})-H(\theta,\theta^{(t)})$： $$ \begin{aligned} &H(\theta,\theta^{(t)}) = \int_z P(z|X,\theta^{(t)}) \log P(z|X,\theta) dz \\ &H(\theta^{(t)},\theta^{(t)})-H(\theta,\theta^{(t)})\\ =&\int_z P(z|X,\theta^{(t)}) \log P(z|X,\theta^{(t)}) dz - \int_z P(z|X,\theta^{(t)}) \log P(z|X,\theta) dz \\ =& \underbrace{\int_z P(z|X,\theta^{(t)}) \log \frac{P(z|X,\theta^{(t)})}{P(z|X,\theta)} dz}_{\mathrm{KL}(P(z|X,\theta^{(t)})|P(z|X,\theta))}\\ =& -\int_z P(z|X,\theta^{(t)}) \log \frac{P(z|X,\theta)}{P(z|X,\theta^{(t)})} dz \\ =& \mathbb{E}_{P(z|X,\theta^{(t)})} \left[-\log \frac{P(z|X,\theta)}{P(z|X,\theta^{(t)})}\right]\\ \geq & -\log \mathbb{E}_{P(z|X,\theta^{(t)})} \left[\frac{P(z|X,\theta)}{P(z|X,\theta^{(t)})}\right] \\ =& -\log \int_z \frac{P(z|X,\theta)}{P(z|X,\theta^{(t)})} \cdot P(z|X,\theta^{(t)}) dz \\ =& -\log \int_z P(z|X,\theta) dz \\ =& - \log 1 \\ =& 0 \end{aligned}\tag{18} $$ 因此，下式成立： $$ \begin{aligned} &\therefore H(\theta^{(t)},\theta^{(t)})\geq H(\theta,\theta^{(t)}) \\ &\therefore H(\theta^{(t+1)},\theta^{(t)}) \leq H(\theta^{(t)},\theta^{(t)}) \\ &\because Q(\theta^{(t+1)},\theta^{(t)}) \geq Q(\theta^{(t)},\theta^{(t)})\\ &\therefore Q(\theta^{(t+1)},\theta^{(t)}) - H(\theta^{(t+1)},\theta^{(t)}) \geq Q(\theta^{(t)},\theta^{(t)}) - H(\theta^{(t)},\theta^{(t)}) \\ &\therefore \log P(X|\theta^{(t+1)}) \geq \log P(X|\theta^{(t)}) \end{aligned} \tag{19} $$ 所以通过EM算法的迭代得到新的$\theta^{(t+1)}$增大likelihood，使得模型更加拟合数据。

EM算法公式推导

EM算法Maximize Likelihood Estimation迭代公式： $$ \begin{aligned} \theta^{(t+1)}&=\underset{\theta}{\arg \max} \int_z \log P(X,z | \theta) \cdot P(z|X,\theta^{(t)}) dz\\ &=\arg \max_\theta \mathbb{E}_{P(z|X,\theta^{(t)})} \left[\log P(X,z | \theta)\right]
\end{aligned} \tag{20} $$

• E-Step： $\mathbb{E}_{P(z|X,\theta^{(t)})} \left[\log P(X,z | \theta)\right]$
• M-Step：$\arg \max_\theta \mathbb{E}_{P(z|X,\theta^{(t)})} \left[\log P(X,z | \theta)\right]$

其中，$X$是观测数据，$z$是隐变量，$(X,z)$为完整数据，$\theta$为待优化模型参数，$P(\cdot|X)$为后验。

上一节通过收敛性证明，验证了上式每次迭代都朝着最大化log-likelihood的方向。本节推导EM的迭代公式。

公式(11)中得到： $$ \log P(X|\theta) = \log P(X,z|\theta) - \log P(z|X,\theta) \tag{21} $$ 引入一个关于隐变量$z$的分布$q(z)$，可以定义为任意关于$z$的非0分布。上式可以改写为： $$ \log P(X|\theta) = \log \frac{P(X,z|\theta)}{q(z)} - \log\frac{ P(z|X,\theta)}{q(z)} \tag{22} $$ 左右两边对$q(z)$求期望： $$ \text{左边} = \mathbb{E}_{q(z)} \log P(X|\theta) = \int_z q(z) \log P(X|\theta) dz = \log P(X|\theta) \underbrace{\int_z q(z) dz}_{=1} = \log P(X|\theta) \tag{23} $$

$$ \begin{aligned} \text{右边}&=\mathbb{E}_{q(z)} \left[\log \frac{P(X,z|\theta)}{q(z)} - \log\frac{ P(z|X,\theta)}{q(z)}\right] \\ &= \underbrace{\int_z q(z) \log \frac{P(X,z|\theta)}{q(z)} dz}_{ELBO=\text{Evidence Lower Bound}} \underbrace{- \int_z q(z) \log\frac{ P(z|X,\theta)}{q(z)} dz}_{\mathrm{KL}(q(z)||P(z|X,\theta))} \end{aligned} \tag{24} $$

所以 $$ \log P(X|\theta) = ELBO + \mathrm{KL}(q(z)||P(z|X,\theta)) \tag{25} $$ 其中$P(z|X,\theta)$为后验（posterior）。而$\mathrm{KL}(q(z)||P(z|X,\theta)) \geq 0$, 当分布$q(z) = P(z|X,\theta)$时，等号成立。所以 $$ \log P(X|\theta) \geq ELBO = \int_z q(z) \log \frac{P(X,z|\theta)}{q(z)} dz \tag{26} $$ 因此，最大化log-likelihood $\log P(X|\theta)$问题可以转化为最大化$\log P(X|\theta)$的下界ELBO，即： $$ \hat{\theta} = \arg \max_\theta \log P(X|\theta) \Longleftrightarrow \hat{\theta} = \arg \max_\theta ELBO \tag{27} $$

$$ \begin{aligned} \hat{\theta} &= \arg \max_\theta ELBO \\ & = \arg \max_\theta \int_z q(z) \log \frac{P(X,z|\theta)}{q(z)} dz \quad \text{令关于}z\text{的分布}q(z) = P(z|X,\theta^{(t)})\\ &= \arg \max_\theta \int_z P(z|X,\theta^{(t)}) \left[\log P(X,z|\theta) - \underbrace{P(z|X,\theta^{(t)})}_{\text{与}\theta \text{无关，去掉不影响结果}}\right] dz \\ &= \arg \max_\theta \int_z P(z|X,\theta^{(t)}) \log P(X,z|\theta) dz\\ &= \text{公式(7)} \end{aligned} \tag{28} $$

我把本文整理成了PDF

参考

https://youtube.com/playlist?list=PLOxMGJ_8X74bhcPbpiX642NIfPlkpD1BC

https://zhuanlan.zhihu.com/p/78311644