Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

技术摘要：利用 GPU 将 Weisfeiler–Leman 表达能力分析扩展至大规模图

问题陈述
Weisfeiler–Leman (1-WL) 测试及其相关的稳定着色（stable coloring）是理解图神经网络（GNN）表达能力的基础。1-WL 稳定着色为消息传递 GNN 区分图结构的能力提供了上界；接收到相同稳定颜色的节点在这些网络中是不可区分的。然而，为现实世界的大规模图计算这种稳定着色面临两个关键瓶颈：

1. 串行执行： 基于划分细化（partition refinement）的经典算法（例如基于 splitter 的 $O(m \log n)$ 方法）本质上是串行的，无法利用现代 GPU 架构的大规模并行性。
2. 全局内存需求： 这些算法要求整个图同时驻留在内存中。这种“全局性”特征阻碍了对超出单个设备甚至大型 CPU 集群容量的 Web 级图（例如拥有数百亿条边的图）进行分析。

现有的并行化或线性化尝试通常依赖于浮点运算，这在大规模应用时会引入数值不稳定性，或者未能解决大于设备内存的图的内存约束问题。

方法论
作者提出了一种两管齐下的方法，通过利用 1-WL 稳定着色的线性代数解释，将计算的可扩展性与内存的可扩展性解耦。

1. 随机线性代数细化（计算可扩展性）
不同于串行划分细化，作者利用等价划分（equitable partitions）的不变性属性重新表述了该问题：当且仅当任何块常数向量（block-constant vector）在邻接矩阵作用下的像仍为块常数时，该划分才是稳定的。

• 算法： 该方法迭代地采样在当前划分上对称的随机向量，将其与邻接矩阵相乘（$z = Aw$），并根据所得向量的对称性来细化划分。
• 算术域： 为了避免浮点幂迭代（power iteration）在处理大规模图时可能导致的数值不稳定性（这会导致错误的合并），该算法完全在模环 $\mathbb{Z}/2^{64}\mathbb{Z}$ 内运行，使用无符号 64 位整数算术。
• 理论保证： 通过利用邻接矩阵行和受限于 $2^{64}/2$ 的结构约束，作者推导出了一个概率正确性保证。一个随机探测检测到任何缺失细化的失败概率至多为 $10^{-9}$，从而绕过了需要域（field）的经典 Schwartz–Zippel 引理。
• GPU 映射： 核心操作简化为单次稀疏矩阵-向量乘法（SpMV）及随后的字典序排序，这两者在 GPU 硬件上都具有极高的效率。

2. 保正性的批处理（内存可扩展性）
为了处理无法放入设备内存的图，作者引入了一种批处理方案，将边集分解为不相交的子集，同时保持全局稳定着色的正确性。

• 内部节点与边界节点： 对于给定的边批次，节点被分类为“内部”节点（所有出边都位于该批次内）或“边界”节点（出边跨越批次边界）。
• 局部细化： 算法独立处理每个批次。至关重要的是，每个批次的初始划分将所有内部节点置于单个块中，而将所有边界节点置于单元素块中。这确保了细化过程仅合并那些在当前批次内拥有完整邻域信息的节点。
• 迭代收缩： 局部细化结果被合并为全局划分。随后，算法构建一个商图（quotient graph，其中节点代表当前划分的块）并重复此过程。
• 终止条件： 过程持续迭代，直到缩减后的图能够放入单个批次中（从而得到证明的最粗稳定着色），或者达到不动点。如果在达到内存容量限制前达到了不动点，则结果是一个有效的稳定着色，它是最粗稳定着色的一个细化。

核心贡献

1. 随机细化算法： 一种用于 1-WL 稳定着色的蒙特卡洛框架，在整数算术下运行并具有严格的概率保证（误差 $<10^{-9}$），避免了浮点方法的精度损失。
2. 批处理方案： 一种分解策略，允许在处理独立边子集并利用商图合并结果的情况下，对超过设备内存的图进行 1-WL 计算，并证明能保持稳定着色的属性。
3. GPU 实现： 一个基于 CUDA 的实现，将这些算法直接映射到 GPU 原语（SpMV、排序、规约），实现了相对于 CPU 基准测试的大幅加速。

实验结果
作者在 11 个中等规模图和 3 个 Web 级图（33–480 亿条边）上评估了其方法。

• 正确性： 在全图可放入内存的中等规模图上（$N_b=1$），随机算法恢复了精确的最粗稳定着色，与确定性基准一致。
• 数值稳定性： 所提出的整数算术方法成功区分了在处理大规模图时被先前浮点幂迭代方法错误合并的节点（例如，it-2004 数据集显示了超过 800,000 个块的差异）。
• 批处理质量： 当使用多个批次（$N_b=2, 4$）来模拟内存约束时，生成的划分在经验上与最粗稳定着色的差距在 5% 以内。对于若干数据集，该差距为零。
• 性能：
  • 速度： 单 GPU 实现相比于传统的单线程 CPU 划分细化实现了高达 138× 的加速，相比于多核 CPU 基准实现了高达 449× 的加速。
  • 可扩展性： 在拥有超过 300 亿条边的 Web 级图上，所有 CPU 基准（包括多核符号归约）要么超时，要么内存溢出。所提方法成功计算了这些实例的稳定着色，实现了约 39–46% 的块数量减少。

意义与主张
论文声称，这项工作是首个成功在超过 300 亿条边的 Web 级图上计算 1-WL 稳定着色的研究，此前由于内存和计算限制，这一任务是无法实现的。其意义在于：

• 赋能表达能力分析： 允许在真实的、大规模数据集而非仅限于中小规模图上进行 GNN 表达能力的实证分析。
• 架构契合： 证明了 1-WL 问题——传统上被视为串行且全局的问题——可以通过针对 GPU 架构定制的线性代数公式和批处理策略进行有效的并行化和分布式处理。
• 鲁棒性： 提供了一种既具有数值稳定性（通过整数算术）又具备内存可扩展性（通过批处理）的方法，克服了先前方法的特定失效模式。

作者总结道，他们的方法使得在以前无法达到的规模下进行 1-WL 表达能力分析成为可能，同时指出，稳定着色与等价划分之间的对应关系，为将这些技术应用于下游图机器学习任务和非线性动力系统的模型简化提供了潜在的未来路径。