Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast, Expressive SE$(n)$ Equivariant Networks through Weight-Sharing in Position-Orientation Space

Erik J Bekkers, Sharvaree Vadgama, Rob D Hesselink, Putri A van der Linden, David W Romero

TL;DR

Based on the theory of homogeneous spaces, geometrically optimal edge attributes to be used within the flexible message-passing framework are derived and formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally.

Abstract

Based on the theory of homogeneous spaces we derive geometrically optimal edge attributes to be used within the flexible message-passing framework. We formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally. We define equivalence classes of point-pairs that are identical up to a transformation in the group and derive attributes that uniquely identify these classes. Weight sharing is then obtained by conditioning message functions on these attributes. As an application of the theory, we develop an efficient equivariant group convolutional network for processing 3D point clouds. The theory of homogeneous spaces tells us how to do group convolutions with feature maps over the homogeneous space of positions $\mathbb{R}^3$, position and orientations $\mathbb{R}^3 {\times} S^2$, and the group $SE(3)$ itself. Among these, $\mathbb{R}^3 {\times} S^2$ is an optimal choice due to the ability to represent directional information, which $\mathbb{R}^3$ methods cannot, and it significantly enhances computational efficiency compared to indexing features on the full $SE(3)$ group. We support this claim with state-of-the-art results -- in accuracy and speed -- on five different benchmarks in 2D and 3D, including interatomic potential energy prediction, trajectory forecasting in N-body systems, and generating molecules via equivariant diffusion models.

Fast, Expressive SE$(n)$ Equivariant Networks through Weight-Sharing in Position-Orientation Space

TL;DR

Based on the theory of homogeneous spaces, geometrically optimal edge attributes to be used within the flexible message-passing framework are derived and formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally.

Abstract

Based on the theory of homogeneous spaces we derive geometrically optimal edge attributes to be used within the flexible message-passing framework. We formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally. We define equivalence classes of point-pairs that are identical up to a transformation in the group and derive attributes that uniquely identify these classes. Weight sharing is then obtained by conditioning message functions on these attributes. As an application of the theory, we develop an efficient equivariant group convolutional network for processing 3D point clouds. The theory of homogeneous spaces tells us how to do group convolutions with feature maps over the homogeneous space of positions , position and orientations , and the group itself. Among these, is an optimal choice due to the ability to represent directional information, which methods cannot, and it significantly enhances computational efficiency compared to indexing features on the full group. We support this claim with state-of-the-art results -- in accuracy and speed -- on five different benchmarks in 2D and 3D, including interatomic potential energy prediction, trajectory forecasting in N-body systems, and generating molecules via equivariant diffusion models.
Paper Structure (46 sections, 6 theorems, 46 equations, 2 figures, 5 tables)

This paper contains 46 sections, 6 theorems, 46 equations, 2 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Mathematical Prerequisites
  4. Motivation 1: Group convolution is all you need
  5. Motivation 2: Efficiency and expressivity --The Homogeneous space $\mathbb{R}^3 \times S^2$
  6. Motivation 3: Conditional message passing as generalized convolution
  7. Fast, Expressive $\mathrm{SE}(n)$ Equivariant Neural Networks
  8. Weight-sharing and optimal invariant attributes
  9. Separable group convolution in position-orientation space
  10. The P$\Theta$NITA Architecture
  11. Related Work
  12. Experiments
  13. Conclusion
  14. The Fiber Bundle vs Point Cloud viewpoint
  15. Regular vs Steerable Group Convolutions
  16. ...and 31 more sections

Key Result

Lemma 1

For any chosen representatives $g_i \in G$ of $x_i \in X \equiv G / H$ such that $x_i {=} g_i \, x_0$, and any $x_j \in X$, the following mapping from the space of equivalence classes of point pairs $X {\times} X / \sim$ to the space $H \backslash X$ of orbits of $H$ in $X$ is a bijection:

Figures (2)

  • Figure 1: Separable group convolutions on position orientation space $\mathbb{R}^3 \times S^2$. Efficiency is obtained due to parallelizing the most expensive step, step 1 (message passing), over orientations and channels. Steps 2 and 3 are efficient as well as spherical convolutions are batched over positions and channels, and channel mixing is batched over positions and orientations.
  • Figure 2: Molecule generation via denoising diffusion models trained on QM9. Negative Log Likelihood (NLL), atom, and molecule stability.

Theorems & Definitions (17)

  • Definition 3.1: Equivalent point pairs
  • Definition 3.2: Equivalence class of point pairs
  • Definition 3.3: Weight-sharing in message passing
  • Lemma 1: Equivalence class correspond to $H$-orbits in $X$
  • proof
  • Theorem 1: Bijective attributes for homogeneous spaces of $\ \mathrm{SE}(n)$
  • proof
  • Corollary 1.1
  • proof
  • proof
  • ...and 7 more