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

推荐一篇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}} cupleft{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}}

来代替。

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

pytorch gpu image

0 人点赞