The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products

• URL: http://arxiv.org/abs/2506.13523v2
• Date: Tue, 15 Jul 2025 07:36:30 GMT
• Title: The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products
• Authors: YuQing Xie, Ameya Daigavane, Mit Kotak, Tess Smidt,
• Abstract summary: $E(3)$-equivariant neural networks have demonstrated success across a wide range of 3D modelling tasks.<n>A fundamental operation in these networks is the tensor product, which interacts two geometric features in an equivariant manner to create new features.
• Score: 0.6005053595985627
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: $E(3)$-equivariant neural networks have demonstrated success across a wide range of 3D modelling tasks. A fundamental operation in these networks is the tensor product, which interacts two geometric features in an equivariant manner to create new features. Due to the high computational complexity of the tensor product, significant effort has been invested to optimize the runtime of this operation. For example, Luo et al. (2024) recently proposed the Gaunt tensor product (GTP) which promises a significant speedup. In this work, we provide a careful, systematic analysis of a number of tensor product operations. In particular, we emphasize that different tensor products are not performing the same operation. The reported speedups typically come at the cost of expressivity. We introduce measures of expressivity and interactability to characterize these differences. In addition, we realized the original implementation of GTP can be greatly simplified by directly using a spherical grid at no cost in asymptotic runtime. This spherical grid approach is faster on our benchmarks and in actual training of the MACE interatomic potential by 30%. Finally, we provide the first systematic microbenchmarks of the various tensor product operations. We find that the theoretical runtime guarantees can differ wildly from empirical performance, demonstrating the need for careful application-specific benchmarking. Code is available at https://github.com/atomicarchitects/PriceofFreedom.

Related papers

• Tensor Decomposition Networks for Fast Machine Learning Interatomic Potential Computations [63.945006006152035]
  tensor decomposition networks (TDNs) achieve competitive performance with dramatic speedup in computations.<n>We evaluate TDNs on PubChemQCR, a newly curated molecular relaxation dataset containing 105 million DFT-calculated snapshots.
  arXiv  Detail & Related papers  (2025-07-01T18:46:27Z)
• An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks [0.5737287537823071]
  Rotation equivariant graph neural networks yield state of the art performance on spatial deep learning tasks.<n>Key to these models is the Clebsch-Gordon (CG) tensor product, a kernel that contracts two dense feature vectors with a highly-structured sparse tensor to produce a dense output vector.<n>We introduce a GPU sparse kernel generator for the CG tensor product that provides significant speedups over the best existing open and closed-source implementations.
  arXiv  Detail & Related papers  (2025-01-23T08:20:47Z)
• Multivector Neurons: Better and Faster O(n)-Equivariant Clifford Graph Neural Networks [17.716680490388306]
  In this work, we test a few novel message passing graph neural networks (GNNs) based on Clifford multivectors. We push the state-of-the-art error on the N-body dataset to 0.0035; an 8% improvement over recent methods.
  arXiv  Detail & Related papers  (2024-06-06T13:17:44Z)
• Equivariant Graph Neural Operator for Modeling 3D Dynamics [148.98826858078556]
  We propose Equivariant Graph Neural Operator (EGNO) to directly models dynamics as trajectories instead of just next-step prediction. EGNO explicitly learns the temporal evolution of 3D dynamics where we formulate the dynamics as a function over time and learn neural operators to approximate it. Comprehensive experiments in multiple domains, including particle simulations, human motion capture, and molecular dynamics, demonstrate the significantly superior performance of EGNO against existing methods.
  arXiv  Detail & Related papers  (2024-01-19T21:50:32Z)
• Low-Rank Tensor Function Representation for Multi-Dimensional Data Recovery [52.21846313876592]
  Low-rank tensor function representation (LRTFR) can continuously represent data beyond meshgrid with infinite resolution. We develop two fundamental concepts for tensor functions, i.e., the tensor function rank and low-rank tensor function factorization. Our method substantiates the superiority and versatility of our method as compared with state-of-the-art methods.
  arXiv  Detail & Related papers  (2022-12-01T04:00:38Z)
• Softmax-free Linear Transformers [90.83157268265654]
  Vision transformers (ViTs) have pushed the state-of-the-art for visual perception tasks. Existing methods are either theoretically flawed or empirically ineffective for visual recognition. We propose a family of Softmax-Free Transformers (SOFT)
  arXiv  Detail & Related papers  (2022-07-05T03:08:27Z)
• Equivariant vector field network for many-body system modeling [65.22203086172019]
  Equivariant Vector Field Network (EVFN) is built on a novel equivariant basis and the associated scalarization and vectorization layers. We evaluate our method on predicting trajectories of simulated Newton mechanics systems with both full and partially observed data.
  arXiv  Detail & Related papers  (2021-10-26T14:26:25Z)
• Frame Averaging for Invariant and Equivariant Network Design [50.87023773850824]
  We introduce Frame Averaging (FA), a framework for adapting known (backbone) architectures to become invariant or equivariant to new symmetry types. We show that FA-based models have maximal expressive power in a broad setting. We propose a new class of universal Graph Neural Networks (GNNs), universal Euclidean motion invariant point cloud networks, and Euclidean motion invariant Message Passing (MP) GNNs.
  arXiv  Detail & Related papers  (2021-10-07T11:05:23Z)
• MTC: Multiresolution Tensor Completion from Partial and Coarse Observations [49.931849672492305]
  Existing completion formulation mostly relies on partial observations from a single tensor. We propose an efficient Multi-resolution Completion model (MTC) to solve the problem.
  arXiv  Detail & Related papers  (2021-06-14T02:20:03Z)
• UNiTE: Unitary N-body Tensor Equivariant Network with Applications to Quantum Chemistry [33.067344811580604]
  We propose unitary $N$-body tensor equivariant neural network (UNiTE) for general class of symmetric tensors. UNiTE is equivariant with respect to the actions of a unitary group, such as the group of 3D rotations. When applied to quantum chemistry, UNiTE outperforms all state-of-the-art machine learning methods.
  arXiv  Detail & Related papers  (2021-05-31T00:48:18Z)
• Deep Polynomial Neural Networks [77.70761658507507]
  $Pi$Nets are a new class of function approximators based on expansions. $Pi$Nets produce state-the-art results in three challenging tasks, i.e. image generation, face verification and 3D mesh representation learning.
  arXiv  Detail & Related papers  (2020-06-20T16:23:32Z)
• OpEvo: An Evolutionary Method for Tensor Operator Optimization [6.273446055072434]
  We propose a novel evolutionary method, OpEvo, which efficiently explores the search spaces of tensor operators. Our comprehensive experiment results show that OpEvo can find the best configuration with the lowest variance and least efforts in the number of trials and wall-clock time.
  arXiv  Detail & Related papers  (2020-06-10T05:33:33Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.