【论文阅读】FlexGraph: A Flexible and Efficient Distributed Framework for GNN Training

阅读思考问题：
Please briefly describe how hierarchical dependency graphs are built in FlexGraph, and point out the specific stage in the NAU abstraction where this process takes place.

请简要描述在FlexGraph中如何构建分层依赖图，并指出在NAU抽象中的具体阶段发生此过程。

框架概述

GNN：图神经网络，目标是通过迭代地聚合顶点邻居的特征，从输入的高维特征中学习图中每个顶点的低维特征。

FlexGraph：一个用于训练GNN模型的分布式框架，具有灵活的领域定义和分层聚合方案。相比之下现在的GNN框架通常是为具有固定定义和聚合方案的GNN设计，不能很好地同时支持不同类型GNN模型。

设计特点：底层的FlexGraph是一个简单的GNN编程抽象，称为NAU和一个紧凑的数据结构，用于建模各种聚合操作。为了实现更好的性能，FlexGraph配备了混合执行策略，以在聚合邻域特征时根据不同的上下文选择适当和有效的操作，应用程序驱动的工作负载平衡策略，以平衡GNN训练工作负载并减少同步开销，以及管道处理策略，以重叠计算和通信。

评估效果：使用现实生活中的数据集和GNN模型GCN，PinSage和MAGNN，我们验证了NAU使FlexGraph比以前的框架更具表现力（例如，DGL和Euler），它们采用类似GAS的编程抽象，例如，它可以处理DGL和Euler无法处理的MAGNN。评估进一步表明，FlexGraph在GCN和PinSage上的训练时间平均为DGL和Euler等最先进的GNN框架的8.5倍。

背景介绍

大量真实世界的数据集可以自然地用图表示。典型的例子是知识图，社交网络，交通网络和网络图。最近，随着深度学习技术的快速发展，图形神经网络（GNN），一个直接将神经网络应用于图形结构数据的新兴模型家族，已经证明在处理各种图形相关任务方面非常有效。GNNs结合了标准神经网络（NN）操作和迭代图传播。已经开发了许多GNN框架来解决关于训练GNN模型的表达性，效率，可扩展性和实现的挑战。

大多数GNN模型都适合“神经消息传递”框架，其中每个GNN层都与包括两个步骤的图传播过程绑定。对于图中的每个顶点v，它首先聚合v的“邻居“（第1层）的特征，以使用NN操作计算邻域表示（参见图1a）。然后，它将v自身的特征与邻域表示相结合，以通过更新操作调整v的特征（见图1b）。由于GNN涉及NN和图传播操作，即特征从“邻居”传播，算法开发者必须将图编码为稀疏张量，并在现有的面向张量的深度学习框架中使用张量操作手动模拟图传播，例如，TensorFlow 、PyTorch 和MXNet。这给GNN模型的实现带来了挑战。

受顶点中心编程模型的启发，例如，应用于图形处理系统的GAS（Gather-Apply-Scatter），已经提出了一些GNN的编程抽象来简洁地支持GNN模型。类似于GAS，这种抽象将GNN层中的计算分成多个阶段。对于每个阶段，用户指定需要在图中的一组顶点及其相邻边上进行的计算。所有当前最先进的GNN框架，例如，DGL ，PyG ，NeuGraph ，ROC 和Euler 采用类似GAS的抽象。他们能够表达一些流行的GNN模型，例如，GCN 、GIN 和G-GCN ，它们可以被归类为DNFA模型，因为这些模型使用直接邻居（即，1-图中的跳跃邻居）并实施平面聚合（即，单步操作）。

除了DNFA模型，我们发现GNN模型有其他各种定义的“邻居”和邻居聚合方案。一些定义间接邻居（即，图中除了1跳邻居之外的结构化数据），并应用称为INFA模型的平面聚合，例如，PinSage ；另一类INHA模型采用间接邻居和分层聚合（即，当聚集来自"邻居"的特征时的多步操作），例如，，P—GNN 和JK—Net 。INFA和INHA模型在促进许多与图相关的任务方面与DNFA模型一样重要。例如，最近的工作表明，像PinSage这样的INFA模型已经广泛用于工业中的推荐系统，例如，Pinterest与最先进的DNFA模型相比，INHA模型P—GNN在蛋白质—蛋白质相互作用网络的蛋白质功能分析中实现了更好的性能。不幸的是，目前的GNN框架主要是为DNFA模型设计的，缺乏对INFA和INHA模型的支持（见2.3节）。

本文对各种GNN模型进行了综合分析，并建立了相应的GNN分类。在此基础上，提出了一种基于阶段的GNN编程抽象NAU，它将GNN层的计算分为NeighborSelection、Aggregation和Update三个阶段。NAU能够以统一的方式训练DNFA，INFA和INHA模型。与直接利用输入图来捕获顶点之间的依赖关系的GAS类抽象不同，NAU采用分层依赖图（HDG）来编码不同类型的“邻居”，而不限于直接邻居。用户可以定义每个顶点如何选择它的“邻居”，NAU相应地构建HDG。此外，用户可以为HDG的每个级别指定聚合函数，并且以自下而上的方式自动聚合邻域信息。因此，NAU支持灵活的邻域定义和GNN的分层聚合方案。

为了提高层次聚集的性能，我们提出了一种混合执行策略。它区分了层次方案中的聚合，并根据它们的不同上下文，通过图处理，稀疏NN操作或密集NN操作来完成它们。这种策略的思想是利用高效的图形处理和NN操作。

此外，为了在分布式环境中有效地训练GNN模型，我们开发了一种应用程序驱动的方法来平衡工作负载并减少同步开销。它通过最小化估计的训练和通信成本将输入图的顶点划分为不相交的集合。我们进一步提供了一个管道处理策略，利用<\u>部分聚合提前聚合<\u>功能，并重叠部分聚合和通信，以提高分布式训练的吞吐量。

我们在NAU上实现了一个分布式GNN框架FlexGraph，该框架采用了上述所有优化策略。我们评估了FlexGraph在四个现实生活和合成图上训练DNFA模型GCN，INFA模型PinSage和INHA模型MAGNN，并将其与最先进的NN框架PyTorch和GNN框架DGL，DistDGL和Euler进行了比较。实验结果表明，对于MAGNN，只有FlexGraph可以有效地支持其在大型图上的训练，这是由于NAU的表达能力增加以及有效的存储和执行优化。对于PinSage，FlexGraph在训练时间上平均优于其他人25.30倍（高达119.33倍）。对于已经在现有框架中得到良好支持的GCN，FlexGraph仍然可以实现至少1.50倍的加速比。

本文贡献如下：

1. GNN的分类，考虑了邻域定义和聚合方案，以揭示GNN框架在训练GNN模型时的关键表现力和性能挑战（第2节）
2. GNN编程抽象NAU支持DNFA、INFA和INHA模型的训练（第3节）
3. 分层聚合的混合执行方案（第4节），以及应用程序驱动的工作负载平衡策略和高效分布式训练的管道处理策略（第5节）
4. 对GNN框架FlexGraph进行了广泛的评估，证明了其有效性（第7节）

GNN分类

在本节中，我们为各种GNN模型提供了一个新的二维分类，包括它们的邻域定义和聚合方案。我们首先回顾GN

2N（2.1节），然后给予分类（2.2节）。我们还确定了在此分类下支持不同GNN模型的挑战（第2.3节）。

GNN

图神经网络（GNNs）是一系列机器学习算法，它们将神经网络（NN）操作应用于图结构数据。最初，图中的每个顶点都与一个高维特征向量相关联。GNN模型的目标是学习每个顶点的低维特征表示，其可以进一步馈送到各种下游任务中，例如，顶点分类、链接预测和顶点聚类。

二维分类

有很多GNN模型被提出用于不同的应用（另见最近的调查）。然而，据我们所知，GNN模型并没有同时考虑GNN的静态和操作特性的分类。一般来说，对于如何定义邻域没有约束，即，使用静态表达式的顶点的“邻居”以及邻居聚合在GNN模型中的行为。因此，我们可以根据GNN模型的静态邻域定义和操作聚集方案将其分为三类。

• DNFA(direct neighbors with flat aggregation具有平坦聚合的直接邻居)：每个顶点的邻域由其所有1跳邻居组成，即，输入图中的直接邻居;并且聚合函数仅进行单步操作。
• INFA (indirect neighbors with flat aggregation平面聚集间接近邻)：一个顶点的"邻居"可以是来自附近子图的顶点(k跳)，而不是它实际的1跳邻居。
• INHA（indirect neighbors with hierarchical aggregation间接邻居与分层聚合）：邻居”可以是任何图形结构的数据。在邻域聚合的操作中，多步聚合操作，即，允许分层聚合。可看论文中MAGNN的解释。
• DNHA（direct neighbors with hierarchical aggregation直接邻居与分层聚合）：现有GNN模型都不满足，现有的GNN框架大多只关注DNFA，本文只讨论前三种。

挑战

GNN模型不仅需要NN操作，还需要图传播。因此，希望有一个框架，同时实现两种类型的操作，并具有良好的性能;并且更有效地支持不同类别的GNN模型（第2.2节）。

A Flexible GNN Framework

FlexGraph采用分层依赖图（HDG）来适应GNN模型中“邻居”和分层聚合方案的多样化定义（第3.1节）。FlexGraph是一个新的GNN编程抽象，称为NAU（3.2节）。在这个新的抽象中，FlexGraph首先为输入GNN模型构建HDG;在HDG的指导下，FlexGraph聚合来自“邻居”的特征，并以统一的方式迭代更新所有顶点特征。使用NAU，不同类别的GNN模型可以在FlexGraph中自然地表达（第3.3节）。

层次依赖图 HDG

给定GNN模型M和输入图G，M的层次依赖图（HDG）w.r.t.G中的顶点v是有向无环图（DAG），记为HDG（v）。它描述了在GNN模型M中v的特征是如何从它的"邻居"中聚合的。G中顶点v的所有HDG（v）的集合称为HDG。

DAG HDG（v）使用分层结构，其中顶层和底层都包含来自输入图G的顶点（参见图3）。每个HDG（v）由两部分组成：模式树和一组邻居实例。

1. HDG（v）的模式树T对由GNN模型M定义的邻居类型的分层（树）结构进行编码。具体地，顶点v被视为T的根，T的每个叶表示M中的邻居类型。当T只有一个邻居类型时，我们规定T = v。 2. HDG（v）的邻居实例，记为Nv，是由模型M定义的v的“邻居”的集合。Nv中的每个邻居实例引用HDG（v）中的一个顶点，并链接到T中编码其类型的叶子。如果相邻实例不遵循平面结构，即，如果它超出了输入图G的一个顶点，则HDG（v）将G中所有与其相关的顶点连接到该实例。

编程抽象NAU

FlexGraph采用了新的GNN编程抽象NAU。与类似GAS的抽象相反，NAU引入了NeighborSelection阶段来构建HDG，以捕获顶点之间的分层依赖关系。HDG就位后，在聚合阶段，NAU要求用户为HDG中的每个顶点指定聚合函数作为UDF（用户定义函数）。也就是说，用户为（i）模式树中的每个顶点和（ii）Nv中的每个邻居实例提供一个UDF，而不管它们是否具有平面结构。然后，NAU通过以自下而上的方式应用这些UDF，自动执行基于其HDG的GNN模型的分层聚合。通过这种方式，NAU支持灵活的邻域定义和各种GNN的分层邻域聚合方案。

NAU的工作流程。NAU由每个GNN层的三个阶段组成，即，NeighborSelection、Aggregation和Update（参见图4）

邻居选择这一阶段构建HDG以满足GNN模型中各种邻域定义的要求。取输入图g，模式树如schema，和顶点UDF nbr_udf作为输入，该阶段输出用于GNN模型的HDG。这里，UDF nbr_udf自定义每个顶点如何从输入图g中检索其“邻居”。

• 为了构建HDG，要求用户为与模式树中的叶子相关联的每个邻居类型实现一个UDF。这里，"邻居"可以是单个顶点或一组顶点，例如，输入图形的路径。图5分别显示了GCN、PinSage和MAGNN的UDF。可以看到，每个UDF都将GNN模型的schema_tree和顶点v作为输入。它返回v的"邻居"及其相应的邻居类型作为输出。例如，PinSage的UDF首先从v开始num_traces许多长度为n_hops的随机游走；然后选择具有top_k最高访问计数的顶点作为v的"邻居"。这些“邻居”有一个扁平的结构。这可能涉及复杂的图形计算

• 聚合给定来自前一层的特征feas（k−1）和由NeighborSelection构建的HDG，此阶段计算邻域表示nbr_feas（k）。使用HDG，允许用户以自底向上的方式指定单步聚合操作之外的多步聚合操作。在各种GNN模型中，在HDG的相同级别处的顶点通常（如果不总是）共享相同的操作语义，即，聚合函数因此，Aggregation函数的典型实现遵循逐层范例，如图6所示。从底层开始，每次它首先获得特定层i上HDG的子间隙g_s，其中包括层i和i-1上的顶点，以及连接它们的边。然后，它应用用户定义的聚合函数aggr_udf_i，该函数将子图g_s和第i层顶点的特征feas_i作为输入，并计算第i-1层顶点的特征。最后，它返回HDG的根的特征作为邻域表示。

• Update该阶段通过NN操作将旧特征feas（k−1）和邻域表示nbr_feas（k）组合。它只使用NN操作为下一层输出新的表示feas（k）。