背景介绍
ScalableGCN是一种由阿里妈妈提出的在大规模图上加速Mini-Batch GCN训练速度方法。
它的主要思路是利用前向计算和反向计算的Cache,在mini-batch之间共享中间层表示的计算结果,同时维护每个顶点的异步更新的通路,达到在与GCN层数成线性关系的时间内训练GCN模型的目的。
ScalableGCN的官方介绍页如下:
ScalableGCN
模型思路
GCN采用聚合上一层顶点的邻居节点embedding(定义为AGGREGATE操作,简称AGG)并与自身节点上一层的embedding做卷积(定义为COMBINE操作,简称COMB)得到当前层的embedding。
在原始GCN训练中,聚合阶数越多、节点邻居数量越多,则AGG和COMB操作次数也越多,训练速度也越慢。
不同于在mini-batch中裁剪邻居节点(GraphSage)或者共享采样(FastGCN)的方法,考虑到mini-batch间,节点的低阶embedding存在被大量重复引用的情况,ScalableGCN设计了embedding缓存机制来缓存低阶embedding信息。以下图为例, 计算2阶a的embedding,对比GCN,ScalableGCN的计算过程只需要关心每一阶a的embedding计算,不需要对邻居节点(b,c)计算embedding。
ScalableGCN
计算得到a的第一阶embedding,将h1(a)更新至embedding cache——h1'中,即h1'(a)=h1(a);在第二阶计算中,h1'(b)和h1'(c)直接从embedding缓存中读取,即用h1'(b)和h1'(c)替代GCN计算中的h1(b)、h1(c)。
说完前向说反向,原始GCN的反向传播过程如下,即通过链式法则计算各个操作的梯度,前向计算loss后和梯度信息更新参数W。
GCN Backward
在ScalableGCN中,由于h1(b)和h1(c)都被缓存h1'(b)和h1'(c)替代,即使用原来的方案,部分g0(b),g0(c)以及g0(d),g0(e),g0(f)和g0(g)都不会和loss一起更新参数W,为解决这个问题,ScalableGCN开辟了与低阶embedding缓存表相对应的低阶embedding梯度缓存表,即在当前batch中,h1‘(b)和h1’(c)对应的g1‘(b)和g1’(c)会累加到梯度缓存表g1'中,而g1'(a)的内容会被“释放”和loss一起作用于g2(a)。
ScalableGCN Backward
Euler中的实现代码细节详见ScalableSageEncoder
(PS:如果本人理解有错,请批评指正,感谢)