Cold Brew: Distilling Graph Node Representations with Incomplete or Missing Neighborhoods

发布时间 2023-12-27 11:13:25作者: 馒头and花卷

Zheng W., Huang E. W., Rao N., Katariya S., Wang Z., Subbian K. Cold brew: Distilling graph node representations with incomplete or missing neighborhoods. ICLR, 2022.

图的结点的 degree 往往呈现服从幂律分布, 现在的训练方式往往会过于注重对 high-degree 的点的学习, 而对那些'冷启动'的点的效果则不怎么样. 本文在测试的给那些冷启动的点假设一些虚拟的邻居结点来缓解这个问题.

符号说明

  • \(\mathcal{V}\), node set, \(|\mathcal{V}| = N\);
  • \(\mathbf{A} \in \mathbb{R}^{N \times N}\), 邻接矩阵;
  • \(\mathbf{\tilde{A}}\), normalized 邻接矩阵;
  • \(\mathcal{N}_i\), \(v_i\) 的一阶邻居;
  • \(\mathbf{X} \in \mathbb{R}^{N \times d}\), node features;
  • \(\mathbf{Y} \in \{0, 1\}^{N \times C}\), node labels;

Cold Brew

  • 首先训练如下的 SE-GNN (SE: Structural Embedding):

    \[\mathbf{X}^{(l+1)} = \sigma \big( \mathbf{\tilde{A}}( \mathbf{X}^{(l)} \mathbf{W}^{(l)} + \mathbf{E}^{(l)} ) \big), \mathbf{X}^{(l)} \in \mathbb{R}^{N \times d_1}, \mathbf{W}^{(l)} \in \mathbb{R}^{d_1 \times d_2}, \mathbf{E}^{(l)} \in \mathbb{R}^{N \times d_2}. \]

  • 这里, 我们额外为每层设置一个可训练的 node embedding \(\mathbf{E}\), 并通过如下损失训练

    \[ CE(\mathbf{X}_{train}^{(L)}, \mathbf{Y}_{train}) + \eta \|\mathbf{E}\|_2^2. \]

  • 令最后的用于预测的 embedding 为 \(\mathbf{\bar{E}}\), 我们通过一个 MLP 来近似它:

    \[\mathbf{\hat{E}} := \mathbf{XW_1} \longrightarrow \mathbf{\bar{E}}. \]

    注: \(\mathbf{W_1}\) 用 MSE 训练得到的.

  • 接着, 为每个结点找到 \(K\) 个邻居, 并重加权得到新的结点表示:

    \[\tag{4} \mathbf{\tilde{e}} = softmax( \Theta_K( \mathbf{\hat{e}_i} \mathbf{\bar{E}}^T ) ) \mathbf{\bar{E}} \]

  • 然后再用一个变换去预测:

    \[ \hat{y} = \mathbf{W_2} \mathbf{\tilde{e}}. \]

    注: \(\mathbf{W}_2\) 是通过 CE 或者 MSE 训练得到的.

  • 故而, 重点应该是 (4), 它相当于重新学习了图的结构, 使得冷启动的结点也能有一些可用的邻居.

代码

[official]