[论文解读] Sketch-GNN: Scalable Graph Neural Networks with Sublinear Training Complexity

Mucong Ding, Tahseen Rabbani+3 more|arXiv (Cornell University)|Jun 21, 2024

Advanced Graph Neural Networks被引用 7

一句话总结

Sketch-GNN 在紧凑的邻接与节点特征草图上训练 GNN，随着图规模的增长，训练时间和内存呈亚线性增长，同时保持有竞争力的准确性。它使用多项式张量草图和可学习的 LSH 来近似非线性激活和消息传递。

ABSTRACT

Graph Neural Networks (GNNs) are widely applied to graph learning problems such as node classification. When scaling up the underlying graphs of GNNs to a larger size, we are forced to either train on the complete graph and keep the full graph adjacency and node embeddings in memory (which is often infeasible) or mini-batch sample the graph (which results in exponentially growing computational complexities with respect to the number of GNN layers). Various sampling-based and historical-embedding-based methods are proposed to avoid this exponential growth of complexities. However, none of these solutions eliminates the linear dependence on graph size. This paper proposes a sketch-based algorithm whose training time and memory grow sublinearly with respect to graph size by training GNNs atop a few compact sketches of graph adjacency and node embeddings. Based on polynomial tensor-sketch (PTS) theory, our framework provides a novel protocol for sketching non-linear activations and graph convolution matrices in GNNs, as opposed to existing methods that sketch linear weights or gradients in neural networks. In addition, we develop a locality-sensitive hashing (LSH) technique that can be trained to improve the quality of sketches. Experiments on large-graph benchmarks demonstrate the scalability and competitive performance of our Sketch-GNNs versus their full-size GNN counterparts.

研究动机与目标

• 在不需要全图内存或对图大小呈线性依赖的情况下，实现对非常大图的可扩展 GNN 训练。
• 在训练之前压缩邻接矩阵和特征矩阵的草图化框架。
• 将多项式张量草图扩展以近似 GNNs 中的非线性激活。
• 开发可学习的局部敏感哈希以在训练过程中自适应草图并提升准确性。

提出的方法

• 将消息传递表示为在 CXW 上的非线性运算，并通过多项式张量草图来近似 CXW。
• 使用计数草图和张量草图将 C 和 XW 草绘成较小的矩阵，从而实现亚线性内存和时间。
• 将非线性展开为幂级数并学习系数 ck 来近似激活。
• 引入基于 SimHash 的可学习 LSH，以在训练过程中自适应草图以应对不断演化的隐藏表示。
• 提供一个双边草图框架，使跨层高效计算草绘表示成为可能。
• 加入跳跃连接以缓解跨层在草图化流程中的误差累积。

实验结果

研究问题

• RQ1 是否可以通过对邻接和特征矩阵进行草图化，在图大小方面实现亚线性时间和内存的 GNN 训练？
• RQ2 在不还原到完整空间的情况下，如何在草图化框架中近似 GNN 的非线性激活？
• RQ3 可学习的 LSH 是否可以在训练过程中提升草图质量并保持竞争力的准确性？
• RQ4 在常见的 GNN 架构中，Sketch-GNN 在内存、时间与准确性之间的理论与经验权衡是什么？

主要发现

• Sketch-GNN 在图规模方面实现亚线性训练复杂性，同时在大型基准数据集上保持竞争力的准确性。
• 该框架将 n×n 的邻接矩阵和 n×d 的特征矩阵草绘成更小的草绘矩阵，从而在每层实现 O(c) 的内存和时间（c 与 n 无关）。
• 多项式张量草图（PTS）被用来在结合 CX 和 XW 的张量/计数草图的情况下近似非线性激活。
• 可学习的 LSH 方法在训练期间调整哈希表以提升草图质量并降低性能损失。
• 跳跃连接有助于缓解草图化流水线中跨层的误差累积。
• 实验结果显示在与完整图 GNN 的对比中具有竞争力的性能，并且在大图上具有良好的可扩展性。