ICML'21 GAS：一键将任意GNN扩展到任意大图上

Houye

发布于 2021-07-09 15:20:07

4360

发布于 2021-07-09 15:20:07

文章被收录于专栏：图与推荐

推荐一篇ICML 2021的文章，与之前推荐的ICML'21 1000层GNN！单块GPU即可训练！类似，本文也不是关注于GNN的架构设计，而是希望能够改善大规模GNN的训练问题。

作者包括Matthias Fey(著名的Pytorch Geometric的开发者)和斯坦福Jure等一众大佬。

image-20210705195205701

本文的核心idea是：通过保存历史计算过的节点表示，使得任意GNN可以扩展到任意大图上运行，且GPU消耗为常量。在具体实现时，通过调用本文提供的PyGAS框架，只要修改短短几行代码，即可一键起飞~

下图中的彩色部分为需要修改的所有代码。

image-20210705204338246

介绍

大图上的GNN基本很难以full-batch的形式运行(拿不到完整的邻接矩阵)。目前的解法主要是进行采样，比如逐点采样其局部邻居结构并聚合。

但是在采样聚合过程中，会有很多冗余的计算。比如聚合节点

v_1

和节点

v_2

的邻居并学习表示时，节点

v_3

和

v_4

会重复计算。

image-20210705204224465

那么一个简单的想法就是：如果

v_3

和

v_4

的表示会重复计算的话，那算完的表示直接存下来，下次直接调用就完事了呗。

下图为GAS的计算图，

\overline{\boldsymbol{H}}^{}

代表上次迭代更新的历史节点表示。可以看出这里

v_3,v_4

的表示没有迭代更新了，而是直接从

\overline{\boldsymbol{H}}^{}

中取出来。

image-20210705205643777

本文的核心公式如下：

\begin{array}{l} \boldsymbol{h}_{v}^{(\ell+1)}=\boldsymbol{f}_{\boldsymbol{\theta}}^{(\ell+1)}\left(\boldsymbol{h}_{v}^{(\ell)},\left\{\boldsymbol{h}_{w}^{(\ell)}\right\}_{w \in \mathcal{N}(v)}\right) \\ =\boldsymbol{f}_{\boldsymbol{\theta}}^{(\ell+1)}\left(\boldsymbol{h}_{v}^{(\ell)},\left\{ \boldsymbol{h}_{w}^{(\ell)}\right\}_{w \in \mathcal{N}(v) \cap \mathcal{B}} \cup\left\{\boldsymbol{h}_{w}^{(\ell)}\right\}_{w \in \mathcal{N}(v) \backslash \mathcal{B}}\right) \\ \approx \boldsymbol{f}_{\boldsymbol{\theta}}^{(\ell+1)}(\boldsymbol{h}_{v}^{(\ell)},\left\{\boldsymbol{h}_{w}^{(\ell)}\right\}_{w \in \mathcal{N}(v) \cap \mathcal{B}} \cup \underbrace{\left\{\overline{\boldsymbol{h}}_{w}^{(\ell)}\right\}_{w \in \mathcal{N}(v) \backslash \mathcal{B}}}_{\text {Historical embeddings }}) \end{array}

其实就做了一件事，把会重复计算的节点表示

\left\{\ {\boldsymbol{h}}_{w}^{(\ell)}\right\}_{w \in \mathcal{N}(v) \backslash \mathcal{B}}

用其历史计算过的表示

\left\{\overline{\boldsymbol{h}}_{w}^{(\ell)}\right\}_{w \in \mathcal{N}(v) \backslash \mathcal{B}}

来代替。

更多具体细节欢迎去围观原文。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2021-07-06，如有侵权请联系 cloudcommunity@tencent.com 删除

pytorch