Sobek: Streaming Equivariant Tensor Product Convolutions

2026-07-20

AI Summary

This work addresses the prohibitive memory and communication overhead of conventional equivariant graph neural networks, which explicitly instantiate edge-dependent weights and intermediate tensors, limiting GPU efficiency. The authors propose a streaming equivariant convolution scheduling strategy that reorders radial projection, spherical harmonic coupling, and graph aggregation to directly fuse edge-local multiply-accumulate operations into bounded node states, thereby eliminating large intermediate tensors. The approach preserves full connectivity and multiplicity mixing across forward, first-order, and second-order backward passes, and is implemented via a generative CUDA backend supporting diverse operator families and feature structures. Experiments demonstrate consistent superiority across 75 capacity-matched comparisons: achieving 1.2–49.7× speedups, up to 99% peak memory reduction, scalability to workloads two orders of magnitude beyond OpenEquivariance, and near-peak throughput retention.

Abstract

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.

Problem

• equivariant graph neural networks
• tensor-product convolutions
• memory traffic
• edge-sized intermediates
• GPU utilization

Innovation

• streaming
• equivariant convolution
• tensor product
• graph neural networks
• memory efficiency