Graph Condensation for Graph Neural Networks
概
图上做压缩的工作.
符号说明
- \(\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]
