analiza aiqbench
Sobek introduces a streaming formulation for equivariant graph neural networks that eliminates unnecessary intermediate memory allocations inherent in conventional execution schedules. By reassociating radial projection and spherical-harmonic coupling operations directly into bounded receiver-side state, this approach achieves up to 49x speedups while reducing peak memory usage by nearly 100%. This advancement significantly expands the feasible scale of equivariant workloads on GPUs without sacrificing throughput or numerical accuracy across forward and backward passes.
Sobek: Streaming Equivariant Tensor Product Convolutions
arXiv:2607.18074v1 Announce Type: new 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.
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.
Machine Learning (cs.LG)
arXiv:2607.18074 [cs.LG]
(or arXiv:2607.18074v1 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2607.18074 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Text integral sub licența deschisă a publicației · 285 cuvinte
Subiecte conexe
Product Interaction: An Algebraic Formalism for Deep Learning Architectures
arXiv · 17 oct. 2026
LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs
arXiv · 14 oct. 2026
Inside NVIDIA’s cuDNN Graph API: Fusion, Autotuning, and Plan Reuse with cuDNN Frontend
MarkTechPost · 15 sept. 2026
Moment kernels: a simple and scalable approach for equivariance to rotations and reflections in deep convolutional networks
arXiv · 3 aug. 2026