Pith. sign in

REVIEW 1 major objections 8 minor 22 references

Edge-sized tensor-product workspace in equivariant GNNs is a scheduling artifact: reassociating radial, coupling, and aggregation sums yields the same operator with bounded receiver-side state through forward, backward, and double backward.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-01 16:10 UTC pith:AMERQ2LK

load-bearing objection The paper's central streaming result is correct and useful, but the scope claims in the abstract need tightening and the benchmark would benefit from an additional baseline. the 1 major comments →

arxiv 2607.18074 v1 pith:AMERQ2LK submitted 2026-07-20 cs.LG

Sobek: Streaming Equivariant Tensor Product Convolutions

Vladimir Choro\v{s}ajev , C\'edric B\'eny This is my paper
classification cs.LG
keywords equivariant graph neural networkstensor product convolutionstreaming executionmemory-efficient kernelsClebsch-Gordan couplingautomatic differentiationGPU memory optimizationirreducible representations
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Equivariant graph neural networks typically spend most of their memory on tensors that are used only once: per-edge tensor-product weights and per-edge messages. The paper's central claim is that these edge-sized intermediates are artifacts of the evaluation schedule, not requirements of the equivariant operator. Because edge weights are generated by a linear radial projection of a shared learned matrix, the sums can be reassociated so that each edge's contribution reduces directly into a bounded receiver-side state, and the learned channel mixing is applied once per receiver. The resulting streamed formulation is algebraically identical to the conventional one, extends to backward and double backward, and keeps tensor-product scratch independent of edge count. If correct, this removes the dominant memory term that currently forces a tradeoff among representation width, interaction range, batch size, and channel connectivity in equivariant models.

Core claim

The paper shows that the edge-specific weight tensor W and edge-message tensor M in edge-conditioned equivariant convolutions are scheduling artifacts. Substituting the radial factorization W_{ij,pi,a} = sum_q phi_{ij,q} A_{q,pi,a} into the convolution and reassociating the sums over edges, paths, and radial components (Eqs. 15-16) lets each edge contribution be consumed directly into a bounded receiver-side state S_g (Eq. 33); the learned channel mixing is then applied once per receiver as m = S A (Eq. 34). The two schedules are identical in exact arithmetic, so forward, backward, and double backward all preserve the operator's semantics while never materializing a complete edge-weight, edg

What carries the argument

The central mechanism is the reassociation identity that pulls the radial factor phi inside the receiver sum: instead of computing per-edge weights w_{ij,pi,a} = sum_q phi_{ij,q} A_{q,pi,a} and then forming edge messages, each edge contributes phi times a Clebsch-Gordan contraction directly into a bounded receiver-owned state S_g. The learned channel mixing A is applied once per receiver via the matrix product m = S A. This receiver-factorized dataflow, together with a fixed scratch-budget policy, is what converts edge-scaled workspace into O(1) tensor-product scratch.

Load-bearing premise

The entire bounded-workspace argument assumes every edge's tensor-product weights are a linear projection of one shared learned matrix onto the edge's radial embedding (W = phi A); if edge weights are generated nonlinearly or independently per edge, the receiver-factorized reassociation is not algebraically equivalent and the bounded-workspace guarantee does not follow.

What would settle it

Run the streamed implementation on a graph whose edge weights are generated by a nonlinear function of the radial embedding, violating the linear factorization premise, and compare against a materialized reference: if results diverge beyond the stated tolerance, the equivalence claim is bounded to the linear case. Alternatively, instrument peak memory on a graph with tens of millions of edges and check whether any tensor of shape |E| x N_w is ever allocated; if the fixed scratch budget is exceeded by such a tensor, the bounded-workspace guarantee is broken.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • Tensor-product workspace no longer grows with the number of edges; peak scratch stays at a fixed budget while graphs scale to millions of edges, removing the dominant O(|E| N_w) allocation.
  • The streamed operator preserves fully connected multiplicity mixing, so it computes exactly the same equivariant map as a conventional uvw convolution, not a cheaper approximation.
  • All three differentiation orders—forward, first backward, and double backward—admit the same streaming schedule, so the memory savings are not lost when training or computing Hessians.
  • Eliminating edge-sized write/read traffic also improves throughput: in all 75 capacity-matched comparisons the streamed implementation was faster (1.2x to 49.7x), and workloads beyond the reference's out-of-memory point still run at near-peak throughput.
  • The freed memory headroom can be spent on longer cutoffs, wider representations, or larger batches, decoupling model design from the peak memory of a transient edge tensor.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Editorial: the same 'pull the edge factor inside, mix once at the receiver' reassociation should apply to any bilinear edge-conditioned layer whose edge dependence is a linear projection of a shared matrix—for instance factorized equivariant attention—so the bounded-workspace principle may generalize beyond tensor-product convolutions.
  • Editorial: the paper's capacity-matched protocol ties both systems to the reference backend's largest batch, which likely understates the streamed implementation's advantage at larger, better-occupying workloads; comparing at equal memory budgets would probably widen the gap.
  • Editorial: 'bounded tensor-product workspace' means independent of edge count, not O(1) total memory: node features, edge embeddings, and spherical harmonics still grow with graph size, so the practical memory law is bounded scratch plus unavoidable graph and feature storage.
  • Editorial: a testable extension is to implement the same streaming schedule for SO(2)-aligned kernels to see whether lower angular complexity composes with the lifetime transformation; the paper flags this as promising but does not implement it.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

1 major / 8 minor

Summary. The paper presents Sobek, a generated-CUDA backend for edge-conditioned equivariant tensor-product convolutions. It observes that conventional schedules materialize edge-weight and edge-message tensors of size O(|E| N_w) and O(|E| C_out), and shows that by reassociating the sums over edges, radial basis components, and tensor-product paths, these intermediates can be consumed into a bounded receiver-side accumulator. For the radial-factorized weight model W_{ij,π,a} = Σ_q φ_{ij,q} A_{q,π,a} (Eq. 12), the receiver-side state S (Eq. 33) can be accumulated without an output-multiplicity axis, and a single channel GEMM applies the learned mixing per receiver (Eqs. 34–35). The authors derive streaming forward (Eqs. 15–16), first backward (Eqs. 18–22), and double backward (Eqs. 23–31) schedules, implement them in generated CUDA, and report benchmarks against OpenEquivariance. In 75 capacity-matched comparisons across uvw and uvu modes and all three differentiation orders, Sobek is faster (1.2–49.7×), reduces peak allocation by up to 99%, and extends to workloads beyond the baseline's OOM point. The paper's central claim is that edge-scaled tensor-product workspace is a schedule artifact, not an operator requirement.

Significance. If correct, this is a significant systems result. It identifies and removes a dominant memory bottleneck in a widely used equivariant GNN primitive without changing the operator's semantics, analogously to FlashAttention for attention. The derivation is clean and the paper is explicit that equivalence holds only up to floating-point summation order, with documented tolerances. Strengths include coverage of fully connected uvw and channelwise uvu, forward/backward/double backward in one framework, a deterministic kernel generator with a fixed scratch budget, and an unusually careful benchmark protocol (fresh processes, warm-up, median timing, capacity matching, confirmed OOM, cross-backend validation). The main caveat is that the receiver-factorized form (Eqs. 33–35) is specific to the linear radial factorization Eq. 12; the paper should clarify the scope of its "not of equivariant convolution itself" claim. With that clarification, the result is publishable.

major comments (1)
  1. [§3.1, §4.2, Abstract, Eq. (12)] The receiver-factorized state S and the one-GEMM-per-receiver channel mixing (Eqs. 33–35) rely on the linear factorization W=φA. For nonlinear edge-weight functions (e.g., an MLP over φ) or node-feature-dependent weights, the S/A reassociation does not hold and the channel GEMM once per receiver is not algebraically equivalent. The concluding claim that edge-scaled workspace is 'not of equivariant convolution itself' is therefore broader than the equations prove. Please add an explicit scope statement: direct per-edge evaluation and accumulation avoids materializing W and M for any edge-local weight function, but the specific bounded S/A scheme requires Eq. (12). This will align the abstract and conclusion with the proven result.
minor comments (8)
  1. [§5.1] The phrase 'strict double-peer validation' is misleading: deterministic repeats of the same implementation are not independent checks. The substantive correctness check is the cross-backend agreement within the stated tolerances. Please rephrase.
  2. [Tables 3 and 5] The header 'Graphs Boeq S|oeq' is confusing. Use e.g., 'B_oeq' and 'Peak GiB (Sobek | OpenEquivariance)' for clarity.
  3. [Figure 2 caption] The purple bands showing SOBEK's cross-profile range are also repeated on the OpenEquivariance axes; the visual distinction is hard to follow. Consider labeling the ranges directly in each panel.
  4. [§3.3] The double-backward derivation defines B^x and B^Y but does not give explicit formulas for their directional derivatives. Since double backward is a central contribution, please provide the explicit bilinear identities or a pointer to an appendix.
  5. [Appendix B.1] The sentence 'If a uvw contract contains an output group of dimension at least five, groups of dimension at least three are split into kernels covering at most four output components' is grammatically unclear. Please rewrite.
  6. [§4.1] Define 'ir_mul' and 'mul_ir' on first use for readers not familiar with e3nn layout conventions.
  7. [§5.1] The phrase 'double-peer' is unusual; define it explicitly on first use. Also, the base tolerance rtol=atol=5e-4 is not 'strict' in a numerical-analysis sense; clarify that 'strict' refers to the protocol, not the tolerance.
  8. [Abstract] The claim of workloads 'up to two orders of magnitude' beyond OpenEquivariance would benefit from specifying that this refers to particular profiles/passes, e.g., double backward on the 128/64/64/64 profile reaches ~170x the baseline capacity.

Circularity Check

0 steps flagged

No significant circularity: the streaming derivation is a self-contained algebraic reassociation; the only caveat is a scope limitation, not circularity.

full rationale

The central claim that complete edge-weight and edge-message tensors are scheduling artifacts is established by the paper's own algebra: Eq. (12) defines the radial-factorized weight model W_{ij,π,a} = Σ_q φ_{ij,q} A_{q,π,a}; substituting it into Eq. (11) and reassociating the sums gives Eqs. (15)–(16), and the receiver-factorized form in Eqs. (33)–(35) is an exact rearrangement of the same sum. No parameter is fitted from the benchmark data, no hidden constant enters the derivation, and no uniqueness theorem or prior result by the same authors is invoked to force the conclusion. The double-backward equations (23)–(31) are obtained by differentiating the same streamed local contractions, not by importing an external result. The numerical validation is a cross-backend check against OpenEquivariance with strict double-peer agreement; using the same library as the comparison baseline is an experimental-design choice, not a circular derivation. The one caveat worth noting is that the bounded-workspace guarantee is conditional on the linear radial-factorized form of Eq. (12), so the abstract's broad phrase 'not of equivariant convolution itself' slightly overstates scope for nonlinear edge-conditional weight models. That is a correctness/scope concern, not a circularity, because the paper states the factorization and the derivation uses it explicitly. No self-citations appear in the reference list, and the cited higher-order decomposition from OpenEquivariance is external to the authors. The derivation is self-contained within the stated operator family, so the circularity score is 0.

Axiom & Free-Parameter Ledger

2 free parameters · 6 axioms · 0 invented entities

The central derivation adds no fitted constants: Sobek's workspace boundedness is an exactly equivalent reassociation of the conventional schedule combined with a user-chosen scratch budget. Axioms are standard representation theory and the factorized radial-projection operator family that defines the paper's scope; the main scope assumption (Eq. 12) is also the paper's acknowledged boundary for the streaming argument. No new physical or mathematical entities are postulated.

free parameters (2)
  • receiver scratch budget B = 128 MiB
    Hand-chosen bound on floating receiver-side state (Eq. 40). The measured peak-memory figures (0.1-0.4 GiB at the last shared point) depend on this budget; the algebraic claim of boundedness does not, only the observed memory-law constants do.
  • cross-backend validation tolerances = rtol = atol = 5e-4; large/dense 2e-3 / 2e-2
    Chosen thresholds that determine when two schedules are deemed to 'match exactly' (Section 5.1). The numerical-equivalence claim is evaluated against these tolerances rather than bitwise equality.
axioms (6)
  • standard math Clebsch-Gordan coefficients define the coupling of O(3) irreps; the operator is a sum over CG paths with edge-dependent weights.
    Used throughout (Eq. 5, Eq. 32); standard representation theory of O(3), foundational for the whole equivariant-GNN literature the paper builds on.
  • domain assumption The radial projection factorizes as W = φA with per-edge embedding φ of width K and a shared learned matrix A.
    Eqs. 12-13; this is the operator family the paper analyzes (weight generation via a learned radial network), not a universal property of all equivariant convolutions.
  • domain assumption Edge embeddings and spherical harmonics can be either stored or recomputed; exact execution must retain them somewhere.
    Section 2.5; the memory accounting treats |E|-scaled φ and Y as irreducible, whereas the paper's claim concerns the Nw-magnified intermediates W and M.
  • standard math First and second derivatives are computed by the standard chain-rule/adjoint equations applied to the bilinear tensor product; no higher-order terms are dropped.
    Sections 3.2-3.3; reverse-mode autodiff of a bilinear map; standard.
  • standard math The receiver-factorized formulation is exactly equivalent to the conventional schedule under exact arithmetic because finite summation is associative.
    Section 3.1 (Eq. 15 = Eq. 16); relies only on reassociating a convergent finite sum; FP behavior differs and is handled by tolerance.
  • domain assumption OpenEquivariance is a representative optimized baseline for the matched deterministic forward/backward/double-backward contract.
    Section 5.1; the choice not to benchmark cuEquivariance or FlashTP assumes OpenEquivariance is a fair representative for this contract, which limits the scope of the empirical generalization.

pith-pipeline@v1.3.0-alltime-deepseek · 16174 in / 19129 out tokens · 188107 ms · 2026-08-01T16:10:55.303492+00:00 · methodology

0 comments
read the original 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.

Figures

Figures reproduced from arXiv: 2607.18074 by C\'edric B\'eny, Vladimir Choro\v{s}ajev.

Figure 1
Figure 1. Figure 1: Receiver-factorized SOBEK dataflow for one output group. Lanes distinguish parameter-, edge-, and node￾indexed values. Purple, blue, green, and magenta mark learned parameters, bounded receiver state, channel matrix products, and edge-local CG operations, respectively; rounded scatter/gather nodes change graph ownership. Dashed gray boxes are edge-local transient values, while the dashed purple box is the … view at source ↗
Figure 2
Figure 2. Figure 2: Peak allocated CUDA memory under graph-count (top) and degree (bottom) scaling, with differentiation [PITH_FULL_IMAGE:figures/full_fig_p011_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Absolute throughput by profile row, with differentiation order on columns. Solid and dashed blue curves [PITH_FULL_IMAGE:figures/full_fig_p012_3.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

22 extracted references · 3 linked inside Pith

  1. [1]

    Bronstein, Joan Bruna, Taco Cohen, and Petar Veliˇckovi´c

    Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veliˇckovi´c. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.arXiv preprint arXiv:2104.13478, 2021

    Pith/arXiv arXiv 2021

  2. [2]

    Clebsch–gordan nets: A fully fourier space spherical convolutional neural network

    Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch–gordan nets: A fully fourier space spherical convolutional neural network. InAdvances in Neural Information Processing Systems, volume 31, pages 10138–10147, 2018

    2018

  3. [3]

    Fuchs, Daniel E

    Fabian B. Fuchs, Daniel E. Worrall, V olker Fischer, and Max Welling. SE(3)-transformers: 3d roto-translation equivariant attention networks. InAdvances in Neural Information Processing Systems, volume 33, pages 1970–1981, 2020

    1970

  4. [4]

    Bekkers, and Max Welling

    Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J. Bekkers, and Max Welling. Geometric and physical quantities improve E(3) equivariant message passing. InInternational Conference on Learning Representations, 2022

    2022

  5. [5]

    Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds, 2018

    Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds, 2018

    2018

  6. [6]

    Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E

    Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials.Nature Communications, 13(1):2453, 2022

    2022

  7. [7]

    Kovács, Gregor N

    Ilyes Batatia, Dávid P. Kovács, Gregor N. C. Simm, Christoph Ortner, and Gábor Csányi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. InAdvances in Neural Information Processing Systems, volume 35, pages 11423–11436, 2022

    2022

  8. [8]

    Equiformer: Equivariant graph attention transformer for 3d atomistic graphs

    Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In International Conference on Learning Representations, 2023

    2023

  9. [9]

    Owen, Mordechai Kornbluth, and Boris Kozinsky

    Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics.Nature Communications, 14(1):579, 2023

    2023

  10. [10]

    The price of freedom: Exploring expressivity and runtime tradeoffs in equivariant tensor products

    Yuqing Xie, Ameya Daigavane, Mit Kotak, and Tess Smidt. The price of freedom: Exploring expressivity and runtime tradeoffs in equivariant tensor products. InProceedings of the 42nd International Conference on Machine Learning, volume 267 ofProceedings of Machine Learning Research, pages 68599–68625. PMLR, 2025

    2025

  11. [11]

    An efficient sparse kernel generator for O(3)-equivariant deep networks

    Vivek Bharadwaj, Austin Glover, Aydın Buluç, and James Demmel. An efficient sparse kernel generator for O(3)-equivariant deep networks. InProceedings of the 2025 SIAM Conference on Applied and Computational Discrete Algorithms (ACDA), pages 32–46. Society for Industrial and Applied Mathematics, 2025

    2025

  12. [12]

    cuEquivariance

    NVIDIA. cuEquivariance. Software, 2025. CUDA library for equivariant neural networks. 16 APREPRINT- JULY20, 2026

    2025

  13. [13]

    Wood, Abhishek Das, and Tess Smidt

    Yi-Lun Liao, Brandon M. Wood, Abhishek Das, and Tess Smidt. EquiformerV2: Improved equivariant transformer for scaling to higher-degree representations. InInternational Conference on Learning Representations, 2024

    2024

  14. [14]

    e3nn: Euclidean neural networks.arXiv preprint arXiv:2207.09453, 2022

    Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks.arXiv preprint arXiv:2207.09453, 2022

    Pith/arXiv arXiv 2022

  15. [15]

    Tensor product

    e3nn developers. Tensor product. e3nn documentation, 2026. Accessed 2026-07-11

    2026

  16. [16]

    Kernel operations on the GPU, with autodiff, without memory overflows.Journal of Machine Learning Research, 22(74):1–6, 2021

    Benjamin Charlier, Jean Feydy, Joan Alexis Glaunès, François-David Collin, and Ghislain Durif. Kernel operations on the GPU, with autodiff, without memory overflows.Journal of Machine Learning Research, 22(74):1–6, 2021

    2021

  17. [17]

    Fu, Stefano Ermon, Atri Rudra, and Christopher Ré

    Tri Dao, Daniel Y . Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. InAdvances in Neural Information Processing Systems, volume 35, pages 16344–16359, 2022

    2022

  18. [18]

    FlashAttention-2: Faster attention with better parallelism and work partitioning.arXiv preprint arXiv:2307.08691, 2023

    Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning.arXiv preprint arXiv:2307.08691, 2023

    Pith/arXiv arXiv 2023

  19. [19]

    Speeding up MACE: Low-precision tricks for equivarient force fields, 2025

    Alexandre Benoit. Speeding up MACE: Low-precision tricks for equivarient force fields, 2025

    2025

  20. [20]

    Seung Yul Lee, Hojoon Kim, Yutack Park, Dawoon Jeong, Seungwu Han, Yeonhong Park, and Jae W. Lee. FlashTP: Fused, sparsity-aware tensor product for machine learning interatomic potentials. InProceedings of the 42nd International Conference on Machine Learning, volume 267 ofProceedings of Machine Learning Research, pages 33143–33156. PMLR, 2025

    2025

  21. [21]

    Lawrence Zitnick

    Saro Passaro and C. Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant GNNs. In Proceedings of the 40th International Conference on Machine Learning, volume 202 ofProceedings of Machine Learning Research, pages 27420–27438. PMLR, 2023

    2023

  22. [22]

    Equivariant neural operator learning with graphon convolution

    Chaoran Cheng and Jian Peng. Equivariant neural operator learning with graphon convolution. InAdvances in Neural Information Processing Systems, volume 36, 2023. 17

    2023