arXiv:2607.18074cs.LG2026-07

通过流式计算优化图神经网络的张量卷积,显著提升速度并降低内存占用。

Sobek: Streaming Equivariant Tensor Product Convolutions

  • 重新设计执行顺序,将中间结果直接流入接收端,避免存储边级临时数据
  • 在75组对比中均更快,最高提速49.7倍,峰值内存减少99%
  • 适合需要大规模图神经网络推理与训练的科研和工业场景

等变图神经网络反复在图边上应用边条件张量积卷积。传统实现需存储每个边的权重、消息和伴随信息,导致张量积工作区和内存流量随图规模和算子宽度快速增长,限制了可处理问题的规模,并可能无法充分利用GPU。我们证明这些边级中间结果是执行调度造成的产物,而非等变算子的本质要求。通过重新组织径向投影、球谐耦合与图聚合的顺序,可将边局部乘积直接合并到接收端的有界状态中。由此提出的流式框架保留了全连接多重性混合能力,支持前向、反向及双重反向传播。我们在Sobek中实现了该框架,一个生成式CUDA后端,并在不同边扩展范围与特征结构下进行评估。在两个算子族和所有三种微分阶数下,Sobek在全部75组容量匹配对比中均更优,速度提升1.2倍至49.7倍,峰值分配内存减少高达99%。其还能执行超出OpenEquivariance两数量级的工作负载,同时保持接近峰值吞吐率。结果表明,边尺度张量积工作区是传统调度的属性,而非等变卷积本身所致。

原文摘要 · Abstract (English)

Equivariant graph neural networks repeatedly apply edge-conditioned tensor-product convolutions over graph edges. Conventional implementations materialize edge-specific weights, messages, and adjoints, causing tensor-product workspace and memory traffic to grow rapidly with graph size and operator width. This limits feasible workloads and can prevent larger problems from fully utilizing the GPU. We show that these edge-sized intermediates are artifacts of the execution schedule, not requirements of the equivariant operator. By reassociating radial projection, spherical-harmonic coupling, and graph aggregation, edge-local products can be consumed directly into bounded receiver-side state. The resulting streaming formulation preserves fully connected multiplicity mixing and extends through forward, backward, and double backward. We implement this formulation in Sobek, a generated-CUDA backend, and evaluate it across edge-scaling regimes and varied feature structures. Across two operator families and all three differentiation orders, Sobek is faster in all 75 capacity-matched comparisons, with speedups ranging from $1.2\times$ to $49.7\times$, and reduces peak allocated memory by up to 99\%. It also executes workloads up to two orders of magnitude beyond OpenEquivariance's capacity while retaining near-peak throughput. These results show that edge-scaled tensor-product workspace is a property of the conventional schedule, not of equivariant convolution itself.

图神经网络张量卷积流式计算内存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。