Graph Condensation for Graph Neural Networks

发布时间 2023-12-24 15:49:15作者: 馒头and花卷

Jin W., Zhao L., Zhang S., Liu Y., Tang J. and Shah N. Graph condensation for graph neural networks. ICLR, 2022.

图上做压缩的工作.

符号说明

  • \(\mathbf{A} \in \mathbb{R}^{N \times N}\), 邻接矩阵;
  • \(\mathbf{X} \in \mathbb{R}^{N \times d}\), 结点特征;
  • \(\mathbf{Y} \in \{0, \ldots, C - 1\}^N\), labels;
  • \(\mathcal{T} = \{\mathbf{A}, \mathbf{X}, \mathbf{Y}\}\), a graph dataset;

Motivation

  • 我们知道, 在不通过 graph sampling 等方法的时候, 在一个规模很大的图上 (\(N\) 很大) 进行训练是很吃资源的, 所以, 本文希望能够对图数据进行压缩, 即构造一个

    \[\mathcal{S} = \{\mathbf{A}', \mathbf{X}', \mathbf{Y}'\}, \quad \mathbf{A}' \in \mathbb{R}^{N' \times N'}, \mathbf{X}' \in \mathbb{R}^{N' \times D}, \mathbf{Y}' \in \{0, 1, \ldots, C-1\}^{N'} \]

    的小数据集 (\(N' \ll N\)), 希望通过 \(\mathcal{S}\) 训练得到的模型和通过 \(\mathcal{T}\) 训练得到的模型性能上是接近的.

  • 严格来说, 我们希望求解如下的一个 bi-level 的问题:

    \[\min_{\mathcal{S}} \mathcal{L}(\text{GNN}_{\theta_{\mathcal{S}}} (\mathbf{A}, \mathbf{X}), \mathbf{Y}), \quad \text{s.t.} \: \theta_{\mathcal{S}} = \text{argmin}_{\theta} \mathcal{L}(\text{GNN}_{\theta}(\mathbf{A}', \mathbf{X}'), \mathbf{Y}'). \]

  • 为了防止对某个初始化参数 over-fitting, 可以进而求解如下的问题:

    \[\min_{\mathcal{S}} \mathbb{E}_{\theta_0 \sim P_{\theta_0}} [\mathcal{L}(\text{GNN}_{\theta_{\mathcal{S}}} (\mathbf{A}, \mathbf{X}), \mathbf{Y})], \quad \text{s.t.} \: \theta_{\mathcal{S}} = \text{argmin}_{\theta} \mathcal{L}(\text{GNN}_{\theta}(\mathbf{A}', \mathbf{X}'), \mathbf{Y}'). \]

GCOND

  • 当然了, 直接取求解 bi-level 问题是困难的, 本文采取 Gradient Matching 来解决这一问题.

  • 假设在第 \(t\) 步, 我们有:

    \[\theta_{t+1}^{\mathcal{S}} \leftarrow \theta_t^{\mathcal{S}} - \eta \nabla_{\theta} \mathcal{L}( \text{GNN}_{\theta_t^{\mathcal{S}}}(\mathbf{A}', \mathbf{X}'), \mathbf{Y}'), \\ \theta_{t+1}^{\mathcal{T}} \leftarrow \theta_t^{\mathcal{T}} - \eta \nabla_{\theta} \mathcal{L}( \text{GNN}_{\theta_t^{\mathcal{T}}}(\mathbf{A}', \mathbf{X}'), \mathbf{Y}'). \]

  • Gradient matching 希望:

    \[ \min_{\mathcal{S}} \mathbb{E}_{\theta_0 \sim P_{\theta_0}}[ \sum_{t=0}^{T-1} D(\theta_t^{\mathcal{S}}, \theta_{t}^{\mathcal{T}})]. \]

  • 这么做的原因是, 倘若在 \(\mathcal{S}\) 上训练的参数的更新过程和在 \(\mathcal{T}\) 上的差不多, 那么很自然最后的模型也会差不多.

  • 我们可以进一步将上面的的目标转换为:

    \[\min_{\mathcal{S}} \mathbb{E}_{\theta_0 \sim P_{\theta_0}}\Bigg[ \sum_{t=0}^{T-1} D( \nabla_{\theta} \mathcal{L}( \text{GNN}_{\theta_t}( \mathbf{A}', \mathbf{X}' ), \mathbf{Y}' ), \nabla_{\theta} \mathcal{L}( \text{GNN}_{\theta_t}( \mathbf{A}, \mathbf{X} ), \mathbf{Y} ) ) \Bigg]. \]

    其中 \(D\) 为每一层的距离, 对于每一列有

    \[ dis(\mathbf{G}^{\mathcal{S}}, \mathbf{G}^{\mathcal{T}}) = \sum_{i=1}^{d_2}\bigg( 1 - \frac{ \mathbf{G}_i^{\mathcal{S}} \cdot \mathbf{G}_i^{\mathcal{T}} }{ \|\mathbf{G}_i^{\mathcal{S}}\| \| \mathbf{G}_i^{\mathcal{T}} \| } \bigg). \]

  • 剩下的问题是, \(\mathcal{S}\) 如何量化, 一种直接的策略是把 \(\mathbf{A}', \mathbf{X}', \mathbf{Y}'\) 设定为参数去学习. 但是这么做 1. 计算量大; 2. 训练难度也大 (个人感觉是灵活度太高了).

  • 所以, 作者首先固定 \(\mathbf{Y}'\), 让其和 \(\mathbf{Y}\) 保持同样的分布.

  • 接着, 作者令 \(\mathbf{X}'\) 自由学习, 并令

    \[ \mathbf{A}_{ij}' = \text{Sigmoid}\bigg( \frac{ \text{MLP}_{\Phi}([\mathbf{x}_i'; \mathbf{x}_j']) +\text{MLP}_{\Phi}([\mathbf{x}_j'; \mathbf{x}_i']) }{2} \bigg), \]

    容易发现 \(\mathbf{A}_{ij}'\) 是对称的.

代码

[official]