大规模图神经网络训练

许多图应用处理的是超大规模数据。社交网络、推荐系统和知识图谱的节点和边数量级都在数亿甚至数十亿个节点。例如，Facebook友情网络最近的一个快照包含8亿个节点和超过1000亿条边。

DGL中的采样方法

超大图是一个经典挑战，在图神经网络（GNN）中更是如此。在GNN模型中，计算一个节点的嵌入取决于其邻居的嵌入。这种依赖关系导致随着GNN层数的增加，涉及到的节点数量呈指数级增长。

一种典型的技术是采样，它有许多变体，其中一些适用于GNN，并且DGL支持其中的一部分。基本思想是修剪节点依赖，以减少计算量，同时仍能准确估计GNN的嵌入。这些技术的精确公式放在教程的最后（并鼓励用户阅读！），这里我们简要概述其要点：

• 邻居采样。这是最基础的：简单地随机选取少量邻居。
• 逐层采样。简单邻居采样的问题在于，随着层数加深，所需节点数量会继续呈指数级增长。这里的想法是将一层视为整体，并限制每层的采样数量。
• 随机游走采样。顾名思义，通过执行随机游走来选择样本。

采样估计器的方差可能很高。显然，包含更多邻居会降低方差，但这违背了采样的目的。一个有趣的解决方案是控制变量采样，这是蒙特卡罗方法中使用的标准方差降低技术。

基本思想很简单：给定一个随机变量$X$我们希望估计其期望值$\mathbb{E} [X] = \theta$，我们希望估计其期望$Y$与其高度相关$X$，控制变量方法寻找另一个随机变量$\mathbb{E} [Y]$可以很容易计算。控制变量需要维护额外状态，但除此之外，它在减少样本量方面是有效的。实际上，对于Reddit数据集，一个邻居就足够了！

对超大图的系统支持

采样提供了一种从数据并行视角处理超大图的不错可能性。DGL增加了两个组件，见下图：

• 采样器：采样器从给定（超大）图构造小型子图（NodeFlow）；NodeFlow的教程可以在这里找到。可以在本地、远程或多台机器上启动多个采样器。

• 图存储：这准备了I/O基础架构，以支持大量训练器进行横向扩展。图存储包含图嵌入及其结构。目前的实现基于共享内存，支持在多GPU和/或非统一内存访问（NUMA）机器上进行多进程训练。

共享内存图存储提供与DGLGraph类似的编程接口。未来的版本中，DGL还将支持分布式图存储，可在多台机器上存储图嵌入。

现在，让我们展示这些支持有多好。

NUMA机器上的多进程训练和优化

加速几乎是线性的，并且在Reddit数据集上仅需约20秒就能在20次迭代后收敛到96%的准确率。

• X1.32xlarge实例有4个处理器，每个处理器有16个物理CPU核心。
• 我们需要注意NUMA架构，否则增加CPU处理器几乎没有速度提升。更多详情请参阅我们的教程。

分布式采样器

可以移动采样器。对于计算密集度不高的GCN变体，在小型（且更便宜的）机器上运行采样器并在NUMA机器上训练更快且更具成本效益（额外20%-40%的速度提升）。

扩展到超大图

我们使用RMAT生成了三个大型幂律图。每个节点关联有100个特征，并计算具有64维的节点嵌入。下表显示了使用邻居采样的GCN训练速度和内存消耗。

节点数	边数	每epoch时间 (s)	内存 (GB)
5M	250M	4.7	8
50M	2.5B	46	75
500M	25B	505	740

我们可以看到，DGL可以在X1.32xlarge实例上扩展到包含5亿节点和250亿条边的图。

下一步计划？

我们期待您的反馈，并乐于了解您在超大图上的用例和模型！本博客中展示的所有这些新功能都将包含在我们即将发布的v0.3版本中。请试用并提供反馈。

我们正在继续改进在超大图上训练图神经网络的基础设施，包括：

• 在该基础设施中，我们将添加分布式图存储，以支持图神经网络的完全分布式训练。
• 我们将尝试各种策略来加速分布式训练。例如，我们将尝试快速且可扩展的图划分算法以减少网络通信。
• 我们还将添加更多其他采样策略的演示。例如，我们将使用随机游走采样器扩展PinSage。

GNN中的采样技术

我们以图卷积网络为例来展示这些采样技术。给定一个图$G=(V, E)$，表示为邻接矩阵$A$，表示为邻接矩阵$H^{(0)} = X \in \mathbb{R}^{|V| \times m}$，图卷积网络（GCN）计算隐藏特征$H^{(l+1)}$，以及节点特征$(l+1)$第-th层，通过计算每个$v$ $h_v^{(l+1)} = \sigma ( z_v^{(l)} W^{(l)}) \qquad z_v^{(l)} = \sum_{u \in \mathcal{N}(v)} \tilde{A}_{uv} h_u^{(l)}$，图卷积网络（GCNs）计算隐藏特征$\mathcal{N}(v)$表示...的邻域$v$, $\tilde{A}$的第$A$例如$D^{-1} A$, $\sigma(\cdot)$层通过计算每个节点$W^{(l)}$是...的可训练参数$l$：$L$层的GCN模型中，...的计算$H^{(L)}$其中$L$-跳邻居传播。这在mini-batch训练中过于昂贵，因为节点的感受野随着层数呈指数级增长。$L$.

邻居采样 与使用节点$L$节点的-跳邻居$v$的所有$\hat{\mathcal{N}}^{(l)}(v)$估计聚合$z_v^{(l)}$-跳邻居不同，邻居采样随机采样少量邻居$\mathcal{N}(v)$在$l$来估计其所有邻居$\hat{z}_v^{(l)}$

$$\hat{z}_v^{(l)} = \frac{\vert \mathcal{N}(v) \vert }{\vert \hat{\mathcal{N}}^{(l)}(v) \vert} \sum_{u \in \hat{\mathcal{N}}^{(l)}(v)} \tilde{A}_{uv} \hat{h}_u^{(l)} \\ \hat{h}_v^{(l+1)} = \sigma ( \hat{z}_v^{(l)} W^{(l)} )$$

设$D^{(l)}$的聚合$l$第-th层，则每个节点的感受野大小可以在...下控制$\prod_{i=0}^{L-1} D^{(l)}$层GCN中，使用无偏估计器

设$v$独立地采样少量邻居$\hat{\mathcal{N}}(v)$是在第$u$除了$v$层为每个节点采样的邻居数量，那么每个节点的感受野大小可以控制在$\vert \hat{\mathcal{N}}(v) \vert > 1$。在逐层采样中，采样过程在每一层只执行一次，其中每个节点$v$之下通过邻居采样。$\mathcal{S}^{(l)}$逐层采样 在节点级别的采样中，每个父节点$p^{(l)}(v)$, $\hat{z}_v^{(l)} = \sum_{u \in \mathcal{S}^{(l)}} \frac{1}{p^{(l)}(u)} \tilde{A}_{uv} \hat{h}_u^{(l)}$感受野大小可以直接通过...控制$\vert \mathcal{S}^{(0)} \vert$.

独立地采样少量邻居$s$以及一个衰减因子$\alpha$，这些邻居对除$s$是始于...的轨迹$s$之外的其他父节点$\alpha$。个性化PageRank（PPR）分数$\pi(s, t)$不可见。如果$s$终止于$t$.

，感受野大小仍然呈指数级增长。在逐层采样中，采样过程在每一层只执行一次，每个节点$v$选择前-$k$以概率$v$作为其邻居$\hat{\mathcal{N}}(v)$被采样到$u \in \hat{\mathcal{N}}(v)$由...加权$\pi(v, u)$,

$$\hat{z}^{(l)}_v = \frac{\sum_{u \in \hat{\mathcal{N}}(v)} \pi(v, u) \tilde{A}_{uv} \hat{h}_u^{(l)}}{\sum_{u \in \hat{\mathcal{N}}(v)} \pi(v, u)}$$

。感受野大小可以直接通过$L$-跳邻居，PinSage选择$k$控制。

控制变量 虽然是无偏的，但采样估计器，例如$\hat{z}_v^{(l)}$在邻居采样中存在高方差问题，因此仍然需要相对大量的邻居，例如$D^{(0)}=25$在邻居采样中的方差很高，因此仍然需要相对大量的邻居，例如$D^{(1)}=10$在GraphSage论文中。通过控制变量法，一种在蒙特卡洛方法中使用的标准方差缩减技术，实验中一个节点只需2个邻居就足够了。

给定一个随机变量$X$我们希望估计其期望值$\mathbb{E} [X] = \theta$，我们希望估计其期望$Y$与其高度相关$X$，控制变量方法寻找另一个随机变量$\mathbb{E} [Y]$可以很容易地计算。控制变量估计量是

$$Z = X - Y + \mathbb{E} [Y]$$

$\mathbb{VAR} [Z] = \mathbb{VAR} [X] + \mathbb{VAR} [Y] - 2 \cdot \mathbb{COV} [X, Y]$可以很容易计算。控制变量估计器是$\mathbb{VAR} [Y] - 2\mathbb{COV} [X, Y] < 0$，则$\mathbb{VAR} [Z] < \mathbb{VAR} [X]$.

如果$\bar{H}^{(l)}$未被采样的节点，修改后的估计量$\hat{z}_v^{(l)}$是

$$\hat{z}_v^{(l)} = \frac{\vert \mathcal{N}(v) \vert }{\vert \hat{\mathcal{N}}^{(l)}(v) \vert} \sum_{u \in \hat{\mathcal{N}}^{(l)}(v)} \tilde{A}_{uv} ( \hat{h}_u^{(l)} - \bar{h}_u^{(l)} ) + \sum_{u \in \mathcal{N}(v)} \tilde{A}_{uv} \bar{h}_u^{(l)}$$