大规模图神经网络应用和最新范式的探索

一、解决大图内存/计算问题的三个范式

在两年前做的tutorial里面，我们有介绍过关于大规模神经网络，并且对20年以前的大规模图神经网络的进展有过一些介绍。在那个时候，考虑的是这样三个范式：layer wise，node wise layer wise和graph wise sampling。

现在来看，归根结底是要去减少图数据在内存和计算上的需求。最简单的方法是对图进行采样。回顾一下当年的一些总结，从14年的图神经网络开始走进人们的视野，到17年GCN的爆火，其实一直以来，对于大规模图神经网络的研究都是一个非常连续的过程。大家都是在朝着如何构造更好的采样和如何减少采样造成的偏差两个方向思考问题，也涌现出了非常多的优秀工作。

我们真的解决了大规模GNN的问题吗？我的答案是解决了，但没有真正解决。首先，确实解决了在实际工业中的应用，尤其是基于子图采样的方法，永远都可以采样出一个子图，Apply一个很复杂的模型，最后得到一个合适的预测。这个在腾讯的一些业务场景，比如推荐，已经有了很好的实践。

但是这个问题并没有真正的解决，因为这个方法其实是回避了核心问题，不能真正在大图做GNN更新。在真正做实践的时候会发现，由于各个地方的系统可能不一样，数据存储格式不一样，图采样的效率本身会依赖于系统实现。而图采样的时间消耗，很可能比训练的消耗更大。另外，这种采样会带来精度下降和信息缺失的风险。尤其是在制药和生物的一些场景里面，是不能随便的对比进行采样的。

那么近两年大规模图神经网络的进展到底怎样呢？可以总结为一句话，“我不想去做采样，但是要把大规模的GNN给做了”。

二、针对大规模GNN的优化

这里仿照之前WWW的GNN Tutorial做了这样的一个图，这个总结可能不是非常的全面，是我个人对这块领域的总结。接下来就挑一些重要的点来进行介绍，这两年大家到底做了什么事情？简单来说，首先我们用了图系统里面的一些分布式图系统计算的一些概念。把传统的GAS范式进行了扩展，成了SAGA。基于这个范式，就会有很多需要系统优化的点，那么具体优化的点可能就存在于：第一，图划分和图划分的优化；第二，对于节点特征传输的优化；第三，对于流水线和通讯的优化。接下来就每个单独的进行简要的介绍。

1、对传统图计算模型GAS扩展

首先，什么是SAGA，在说什么是SAGA的时候，我们就会讨论什么是GAS。GSA是当年图计算里面一个分布式图计算里面一个非常经典的范式，它把整个图计算划分成了三个步骤：Gather，Apply和Scatter 。

• Gather是什么意思呢？Gather就是说去通过边来收集邻居的信息。
• Apply是什么意思呢？Apply就相当于是把收集到的信息去计算出更新的这个节点信息。
• Scatter的意思就是把更新后的节点信息更新到这个边上面。

很多图算法都可以通过GAS的方式进行Formalize，比如PageRank。也就是说面对基于消息传播的图神经网络的时候，就可以基于GAS方式进行扩展，叫做SAGA。

其中可以分为四个步骤：Scatter、Apply Edge、Gather和Apply Vertex。

Scatter 和原来的GAS的Scatter 是一样的，就是把节点的数据先输送到边上面。这一个步骤是GPU Intensive。

然后再把这个节点的这个特征输到边上去，以后可能会对输到这个边上面的消息进行一些处理，这个处理有可能，比如说是GAT里面的计算权重，或者其他一些复杂操作，会Apply一个神经网络去处理，这个肯定是GPU intensive。

第三个， Gather传过来消息了，并且可能是做了处理的消息，那就要通过邻居的关系，把消息进行汇聚，得到新的消息，那么这肯定就是CPU intensive。

第四个步骤，得到汇聚后的消息后，可能还需要一个额外的Update，可能是通过一个神经网络进行，那么这个Update的结果其实就是最终的更新后的节点的表示。那么这个步骤叫做Apply Vertex，那么这其实就是GPU Intensive。

那么其实可以从这个范式里面可以看到：

第一，我们可以吸取以前GS的一些优良的系统，优化的一些策略，比如一些图划分、Pipeline、调度的策略。

第二，这里也挑战，比如前文提到这是一个GPU和CPU交替intensive的一个任务，那么怎么优化？比如说GPU和CPU之间的数据传输，甚至是不同集群之间的数据传输，其实是一个非常重大的问题，也是近年来研究的热点。那么接下来介绍一下它们之间的关系。

2、图划分&图划分优化

（1）单机上的图划分

首先，对于第一个优化点，图划分和图划分的优化，本质上就是其实因为没办法把图放到显存和单机上面，因此就需要把图划分到不同的Server上面去。对于Newer Graph来说，采用这种Locality Aware的图划分策略，相当让连在同一个节点上的边能尽可能在同一个窗口，这样再去做更新的时候，就不用访问其他的窗口，就可以优化内存的访问。

（2）基于模型图划分

第二个，ROC工作其实引入了Linear Regression Model，这个Linear Regression Model会去预测每一次Server里面的执行时间，然后通过这个执行时间去更新下一轮Partition的图Partition的策略。这里是基于模型的Cost、Running Time的图的划分。

（3）纵向划分节点特征

P3这篇干脆就回避了划分的一些缺点，因为如果是基于节点或边的分割，会导致额外的信息通讯和信息损失，所以P3的划分就不是基于图的，而是基于特征的。这个motivation的特征维度特别大，没办法全部放到同一个单机上面，但这个节点本质上是一个Adjacency List，是可以放到每一个单独的Server上面去的，因此可以对于每一个Server节点的特征进行动纵向划分。Per dimension把这些划分放到各个机器上面去，这样既保证了图的完整性，也保证了Somehow减少信息通讯的代价。

3、对于节点特征传输的优化

很多工作都利用了节点特征传输的优化，会发现有些时候节点的原始特征是比GNN的Dimension大很多，尤其是当节点是一个C的特征，通过之后的Embedding，这个节点特征会非常非常大。所以如果基于节点特征在节点的机器和机器之间做通信是非常不利的，不经济的。怎么尽量去减少这种Move节点的这样的通信，或者说计算，也就是是对于节点特征传输的一个优化？那么这个地方也是有三项工作。

PA graph其实就是一个静态缓存的机制。比如我们知道某个显存可能只有20个G做计算，那显存如果还有10个G，就会去缓存一些节点的原始特征。这个策略，本质上是基于静态的缓存机制，它会去把边的点的Out Degree进行排序，会把Out Degree多的点扔到缓存里面去，那显然Out Degree多的点很容易有更高的概率参与计算，所以说开始会把节点特征传输的通信量减少。

在DistGNN里面，设计了一个更复杂的开始Blocking的机制。这个开始Blocking的机制会把节点分为两类——要更新的目标节点和目标节点邻居节点。传统的方法是基于目标节点进行遍历，然后去拉它的邻居节点DistGNN。这里说的机制则是反向去遍历邻居节点，然后对目标节点特征进行alternately的更新。这样做好处是目标节点的特征会放到CPU的缓存里面，每次去找邻居节点，只需要读一次，不需要缓存，每次去更新目标节点的时候，都是在CPU缓存里面更新。这样子就是可以减少反复拉取邻居节点的开销。

对于P3来说，这个设计就会更复杂一些，因为P3的目标是不希望产生任何原始特征之间拉取的通信，所以设计了一个hybrid parameter的机制。对于输入层的GNN来说，模型会并行在每一个单机，也就是每一个单机有一部分特征。计算出sub partial activation，然后每个机器在本地计算完partial activation，会把这个partial activation汇聚成一个输入层的GNN的输出，得到这个输入层的GNN输出以后，再走正常的Data Parallelism，去做这样计算的动机是原始特征的维度特别大，那么第一层GNN的参数和通信量都会非常大。而除了第一层以外，其他GNN的黑洞Size其实都会比较小，这种时候如果在本地算好了第一层再去后面去算更深层的DNN，通信代价能大为减少。

4、流水线、通信优化

这是一个非常巧妙的设计，对于这种流水线通信的优化，就更不用说了。流水线通信优化，其实也是传统的图计算里面做的非常多的，简单来说，同步更新的一些信息需要算完一轮特征以后再更新下一轮的特征，但有些时候可以不这样去做，而是一个异步更新，并且在异步更新的过程中，使数据聚合的过程中，或者在更新的过程中使用一些老的特征，同时也会设计一个老化机制。

这样相当于是在这个分布式系统里面的半同步的机制，通过半同步机制，可以很好的去设计这个流水线，并且减少多机通信的代价。这里因为时间关系，就不具体展开讲了。

要注意的是，可以用的老的特征，但不能用太老的特征，如果这个特征太老，也会暂停更新，等其他节点的特征更新到最近的epoch以后，才继续更新。

其实像Dorylus这种工作，本身就是在多个机器上面，而且每个机器的这个什么角色还不一样，分为Graph Server，Lambda thread和Permit Server，这种情况下，对Pipeline流水线的设计是有很多细节的。总的来说，现有的方法，也是慢慢从单机多GPU到多机混合的趋势，从经验上的静态划分到这种动态的划分，以及引入更多系统层面的Pipeline的优化。

同时，像SANCUS的工作，其实也还进一步去证明了异步的更新是可以保证模型收敛和通讯优化的。就是说慢慢的大家从一些简单的实践到复杂实践，从一些没有理论保证的实验，到一些理论保证实践。这一块的发展还是非常的迅速的。

四、未来方向

接下来谈一下对未来方向的一个理解。首先我们发现，实际上SAGA的这个范式并不适用于所有图神经网络。如果不是基于Message Passage神经网络，其实它就不符合SAGA范式，同样的，就没办法利用到已有的GAS范式里面系统优化的一些trick，来对系统的整个代价进行优化，比如说对于Graph Transformer这样的模型，最近也非常的火，从20年开始到现在也有很多这样的模型出现。

那么，能不能在这样的模型里做Full Graph的优化、或者说Full Graph的训练，其实是一个非常有挑战性的问题。而且这并不是说没有应用，比如对于蛋白质建模来说，实际上如果蛋白质的这个序列足够长，把这个蛋白质作为一个Graph的话，它的训练代价会非常的大。显存可能会出现暴涨，那么怎么在这种非Message Passing的框架下面去对这种大规模GNN做系统优化，是一个非常重要的课题。

第二个就是很多时候，以前的这个图神经网络里面是没有包含几何信息的，什么叫几何信息，就是说节点可能是有空间位置信息的，比如坐标、速度，或者说一些其他的几何信息，这些东西实在很多，尤其是AI for Science领域是非常常见的。

比如我们对于粒子进行建模。对于这种催化剂系统的模拟，都会包含几何信息，而几何信息本身的查询和更新就是数据库领域的一个非常重要的课题，这个和Spatial Database有很强的联系，那么对于这种类型的数据，我们在已有的等变神经网络成果上面，能不能对它进行系统化或者规模化？因为实际上这个系统化，规模化也是非常急需的，就比如之前做的一个催化剂的比赛OCP，这个比赛里面都包含几百万个催化剂系统，数据量是T级别的，实际上对于模型的训练和推理都有非常大的挑战，那么对于这种几何信息的输出神经网络，是不是有很好的解决的方法，这也是未来研究的方向。

五、Q&A

Q：目前要把图神经网络从学术界的东西变成工业界真正能使用的东西，还有很大的差距。其中不仅是算法本身的优化，还有现在些创业企业都在做graph，做computing platform，然后做一个真正的end to end。大家想要从系统优化到上层算法，再到application，完整的做一套系统出来，然后来服务drug discovery、金融等领域。比如用这种更先进方法来从广告品中挖掘到更多信息。您对算法架构整个这块都有很多研究了，您对graph plus，也就是graph platform computing as a service，怎么看？

A：我觉得这是非常有前景的一个方向，简单来说，实际上graph computing本身的platform，在AI时代之前就有很多人研究了。因为在graph上面有很多传统的图算法，也是需要去做分布式的计算，这种分布式的计算就是service的，但是以前我们做的可能都是一些基础的计算。而现在像图神经网络等技术兴起后，我们发现这个系统里面会面临更多的挑战。

比如说以前的数据只需要在CPU上面算就可以了。现在类似于GNN这样的结构，我们一定会涉及到一种混合架构，第一做起来很难；第二，这个东西如果做出来了，门槛很高，所以说这一块是非常有前景的一个方向。

而且这一块市场的玩家目前不是很多，尤其是对于大企业来说，这一块要去集合做AI的和做系统的人，这两块人凑在一起，在大企业里面其实是比较困难的，因为毕竟大企业都是以商业为主轴的。所以我们正是需要一些创业的公司把这个东西做出来以后，真正去服务一些实际的应用。比如你说的drug discovery，或者说一些AI + Science，或者说可能的一些在城市道路，甚至社交网络上面的应用。其实这块东西属于门槛高，有前景，处于还需要人进一步去开拓和做这样的一个状态。