Fast, Expressive $\mathrm{SE}(n)$ Equivariant Networks through Weight-Sharing in Position-Orientation Space
Authors: Erik J Bekkers
Sharvaree Vadgama
Rob Hesselink
Putri A Van der Linden
David W. Romero
Venue: ICLR 2024 poster (2024)
Keywords: Equivariance
Point Clouds
Message Passing Neural Network
Molecules
Diffusion Model
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We propose efficient SE(d) equivariant networks using group convolutions over position-orientation space and achieve state-of-the-art performance on 2D and 3D data.

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.

*Code available at [https://github.com/ebekkers/ponita](https://github.com/ebekkers/ponita)*

Published as a conference paper at ICLR 2024
FAST, EXPRESSIVE SE(n) EQUIVARIANT NETWORKS THROUGH WEIGHT-SHARING IN DOSITION-ORIENTATION SPACE
Erik J. Bekkers" Sharvaree "Sharvare Vadgama !* Rob D. Hesselink', Putri A. van der Linden', David W. Romero
Davidd
University of Amsterdam Vrije Universiteit Amsterdam 3 NVIDIA Research
Equal contribution.
Code available at https://gttub..cmmbbbeeeessppppp
ABSTRACT
Based on the theory of homogeneous spaces we derive geometrically optimal edge attributes to be used within the flexible pressage-passing framework. We formal- ize the notion of weight sharing in convolutional networks as the sharing mes- sage functions over point-pairs that should be treated equally. We define equiva- Ience 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, develop 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 R position and orientations I3 x5" and the group SE(3) itself. Among these, x is an optimal choice due the ability to represent directional information, which IR3 methods cannot, and it significantly enhances computational efficiency compared indexing features on the full SE(3) group. We We support this claim with state-of-the-ane 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.
INTRODUCTION
Inspired by the foundational role of convolution operators in deep learning and the 'convolution is all you need' theorem (Cohen et al. (2019, Thm 3.1); Bekkers (2019, Thm 1)) -which asserts that any layer that linear and equivariant must be group convolution, we propose an efficient and expressive group convolutional approach for constructing neural networks equivariant to SE(n): the group of n--imennsioal translations and rotations. While this theorem is a theoretical result, several studies provide empirical truth to the statement well. For example, ConvNeXt (Liu et al., 2022b) challenges the need for Transformers (Vaswani et al., 2017) in vision tasks, and Romero et al. (2021); Poli et al. (2023) show that convolutions sufficient model long context in sequences, e.g., language, without the need for transformers or recurrent networks.
In our work, we revisit the influential inductive bias of weight sharing in convolutions (LeCun al., 1998), classically defined as the sharing of a convolution kemel (linear transformation) over all the neighborhoods in an image. In the discrete image setting, the kernel is given as set of weights and is appropriate to refer to the convolution kernel as the weights. However, in the continuous case, and more general processing frameworks such as in message passing networks (MPNs) (Gilmer al., 2017a), this terminology no longer literally applies. Therefore, we formalize the notion of weight sharing to gain insights into how use this inductive bias in a more general setting-
Traditionally, weight sharing refers to the idea of using the same linear transformation matrix over neighborhoods identical up to translation. We generalize this notion by the construction of equiva- lence classes of neighboring point-pairs, in which we say neighbor pairs are equivalent if they are identical up to a transformation in a group G. Weight sharing then becomes the notion of sharing message functions in MPNs over the equivalence classes, which is achieved by conditioning the functions on attributes that act as identifiers for the equivalence classes. Sec. 3.1 we formalize this construction and present all the attributes one needs for SFin) equivariant weight sharing.

---

Published as a conference paper at ICLR 2024
3: channel interactions
2: orientation interactions
spherical signal f(p n)
Pj
1: spatial interactions
Figure I: Separable group convolutions on position orientation space IR3 x S Efficiency is ob- tained 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 batched over positions and orientations.
As an application of the theory, we construct expressive and efficient SE(n) equivariant group con- volutional networks for the processing of 3D point clouds. Our simple, fully convolutional archi- tecture achieves wate-of-the-ar results on three different benchmarks: interatomic potential energy prediction, trajectory forecasting in N-body systems, and generating molecules via equivariant dif- fusion models. The contributions of this paper are summarized follows:
We formalize the notion of weight sharing in Sec. 3.1.
We derive optimal attributes (Thm. 1) for weight sharing over equivalence classes of point pairs and show that these are all you need to build equivariant universal approximators (Cor. 1.1).
We present fast, expressive equivariant architecture based purely on convolutions (Fig. 1). Our method reaches state-of-the-ar performance on three equivariant benchmarks.
2 BACKGROUND
2.1 MATHEMATICAL PREREQUISITES
Groups and homogencous spaces. A group is an algebraic construction defined by a set G and binary operator G G - G, known the group product. This structure must satisfy the following axioms: closure, where Vh. hg E G; the existence of both an identity and an inverse g- element such that g- g=e; and associativity, where V gg,h,ie (g h)) i=g (h i). We denote the group product between two elements g.g' G by juxtaposition, i.e., as g g.6
is useful to think of the elements of G as transformations. The group product then tells how transformation g' followed by another transformation g can be represented by a single transfor- mation g We focus on the Special Euclidean motion group SE(n) consisting of distance and orientation preseevvnn transformations. Elements (x R) E SE(n) are parameterized by trans- lation vectors x E In and rotation matrices R E (n).). Here, S0(n) denotes the set of nxn matrices with determinant 1, which forms group in itself with matrix multiplication as a group product. The SE(n) group product between two roto- aio-ttranslations g=(x, R) and g'= R') is given by (x, R. (x', R') RRx= + x, RRR,, and its identity element is given by I0,
A group can act on spaces other than itself via group action T G X - X, where X is the space on which G acts. For simplicity, we denote the action of g E G on x E X as g gx. Such a transformation is called group action if it is homomorphic to G and its group product. That is, it follows the group structure: (g x=g (g' z) Vg. g' G,I E X, and e ex==. For example, consider the space of 3D positions X = R e.g., atomic coordinates, acted upon by the group G-SE(3). A position E 3 is roto- translated by the action an element (x, R) E SE(3) as (x, R. R
A group action is termed transitive if every element x E X can be reached from arbitrary origin Io E X through the action of some g E G, i.c., x= A space X equipped with transitive action of G is called homogeneous space of G. Finally, the orbit G x {gx 9 E G) of an element
2

---

Published as a conference paper at ICLR 2024
under the action of a group G represents the set of all possible transformations of by G. For homogeneous spaces, X-G 2o for any arbitrary origin zo E X.
Quotient spaces. The aforementioned space of 3D positions X=R serves homogeneous space of = SE) as every element p can be reached by noto-transsation from 0, i... for every p there exists (x, R) such that p=(x* Xx R. 0 x=x. Note that there are several elements in SE(3) that transport the origin O to P. as any action with a translation vector x-p suffices regardless of the rotation R. This is because any rotation R' E SO(3) leaves the origin unaltered.
We denote the set of all elements in G that leave an origin xn E X unaltered the stabilizer subgroup Staber In subsequent analyses, we use the symbol H denote the stabilizer subgroup of chosen origin H in G as g H {gh h h E H}. In the example of positions p E X= we concluded that we can associate a point p with many group elements
in a homogeneous space, i.c., ( We further denote the left coser of
SE(3) that satisfy p=g0, 0. In general, letting Ir. be any group element s.t. r= To, then any group element in the left set 92 H is also identified with the point p. Hence, any I E X can be identified with a left coset H and vice versa.
Left cosets g H then establish an equivalence relation among transformations G. We say that two elements 9: of E G are equivalent, i.e., 9 ~ y if and only if g zo=g' zp. That is, if they belong to the same coset g H. The space of left cosets is commonly referred to as the quotient space G/H. H.
We consider feature maps f X - RC as multi-channel signals over homogeneous spaces X. Such maps are of interest as they often form the hidden representations in various deep learning tasks. In this work, we treat point clouds as sparse feature maps, e.g., sampled only at atomic positions. In the general continuous setting, we denote the space of feature maps over X with X. Such feature maps undergo group transformations through regular group representations (>(gg) X + X' parameterized by g. and which transform functions! J E X via [0x (g)[[(()=](()
2.2 MOTIVATION 1: GROUP CONVOLUTION IS ALL YOU NEED
In deep learning, we often employ learnable operators + - y, such as self-attention or convo- lution layers, to iteratively transform feature maps. Such an operator is termed G-equivariant if it commutes with group representations on the input and output feature maps: x² (²)( o Φ o xx (g). Group equivariance ensures that operators preserve the geometric structure of the data, meaning that derived features follow the same transformation laws as the input.
An important result in the field of equivariant deep learning and signal processing is that if want to be linear and group equivariant, then it must be a group convolution:
[+ = S ()
(1)
Essentially, the group convolution performs template matching of a kernel is against patterns in by taking -in-inner products of the shifted kernel k(gz".)) and f. Recall that g, denotes any group element such that y= yo. and any other group element in the set gy H is also valid. This implies that Eq. 1 is only valid if the kernel is invariant to left actions of H, i.e. if Vhen : =(h--z)=k((r) These findings are summarized in various seminal equivariant deep learning works (Cohen et al., 2019, Thm. 3.1 convolution is all you need), (Bekkers, 2019, Thm 1), (Kondor & Trivedi, 2018).
In light of the "convolution is all you need" claim and the pivotal role of convolutions in areas like computer vision and signal processing, one might question the need for anything more complex than convolutions. We therefore explore the potential of a straightforward, yet theoretically grounded, fully convolutional approach to equivariant deep learning.
2.3 MOTIVATION 2: EFFICIENCY AND EXPRESSIVITY -THE HOMOGENEOUS SPACE TR33 x S2
It is worth noting that the domain of the output signal Y in Eq. 1 does not have to match the domain of the input signal X. The aforementioned theorems show that SE(3) equivariant convolutions on require isotropic, rotation invariant kernels, as seen
feature maps with domain R "===E(3)/SOG) in SchNet (Schütt et al., 2023). However, maximal expressivity is gained when the kernel has no constraints. This is achieved when the domain Y is the group itself, i.e. SE(3) as H=(e} is then the trivial group. Generating nigher-dimensions referred to as lifting. It is done by performing template matching (Eq. 1) over both translations and rotations, as opposed to just translations. Despite enhanced expressivity, however, subsequent layers must compute integrals over the whole SF(() space, which can be computationally restrictive.
feature maps, from X-R to Y=SE(3). is typically
3

---

Published as a conference paper at ICLR 2024
In this work, we opt for middle ground and define feature maps the homogeneous space of posi- tions and orientations X x SEE)3 /SO(2)) We denote elements in R X as tuples (p. o) of positions p IR and orientations E S We set the origin as z0= eg). where e, (0= T)T. Every orientation can be obtained from a rotation R. E SO(3) that maps the unit vector e, to o. Hence, 5 is homogeneous space of S3(3) with stabilizer subgroup H =t=abbbo)3) (e=)=S0(2): the group of rotations around the z-axis. Therefore, each position-oriccntatoon tuple corresponds to an equivalence class of transformations in SE(3), formalized as (p, o) (p, Ro) H in SE(3)/SO(2).
While methods based on I positions are computationally efficient, intermediate feature maps may fail to capture directional information due to kernel constraints. On the other hand, full SE(3) group convolutions are computationally demanding but excel at capturing directional features. Compared to the latter, our method offers significant improvement computational efficiency without com- promising the capability to represent directional features. This claim is proven in (Gasteiger et al., 2021, Thm. 3) by relying on the fact that equivariant predictions can be obtained by combining an equivariant basis with coefficients obtained from an invariant network (Villar et al., 2021). In DimeNet and GemNet (Gasteiger et al., 2019; 2021), the basis is determined by the direction of edges in an atomic point cloud, and coefficients are derived through invariant message passing. Rather than assigning a sparse basis of reference directions in S to each point position, our method employs dense basiy and assign a spherical signal to each point in R$, see Fig. 1.
2.4 MOTIVATION 3: CONDITIONAL MESSAGE PASSING AS GENERALIZED CONVOLUTION
Message passing. Our primary focus lies in the processing of point clouds on homogeneous spaces. Point cloud methods are conventionally discussed within the broader context of message passing (Gilmer et al., 2017a). We intentionally adopt this framework as it serves familiar and versatile paradigm suitable for the description of various deep learning layers, including convolutions.
We model point-clouds graphs G=(V, with nodes i V and edges (i, E. Each node i has an associated coordinate I; E X in a homogeneous space X. We consider features over such graphs as sparse discretizations of dense feature maps f : X EC with node features (zi, fi) fxx)). Considering optional pair-wise edge attributes aij. message passing layers update feature nodes as:
1. compute the message maj from node to i:
mij = @m (fx, aiy)
(2)
2. Perform a permutation invariant message aggregation:
m, = > Seeioo) mij,
(3)
3. Update the node features:
fr = 6s (fs,
(4)
Here, (i I (1,j) E E} denotes the set of neighbors of the node i.
Convolutional message passing. Convolutions can be described in the message passing framework as approximations over a sparse set of points on which continuous signals are sampled:
[{
ESENra) K9
(5)
Here, we can recognize the message passing form, in which @m @m, 9 xz) = k(g x;) is a lin- ear transformation performed by matrix-vector multiplication with the kernel ((o,,ta)) E Coux that depends on 97,'24,. z4. Following Brandstetter al. (2021), we interpret this as using a message function that is conditioned on the attribute a.j. denoted as 6m (fr, fyi axj) to emphasize this de- pendency. We could directly consider the use of the invariant attribute a. = x, in conditional message passing as an extension of convolution. However, it's crucial to recognize that, despite its invariance under global group actions, the value of this attribute depends on the selected representa- tive Gz.- Specifically, any 9=; h. where belongs to the stabilizing group H, can result in a different a-==h" 9z, 2j. In essence, any asz E H Dry within the orbit of H should be treated equivalently.
3 FAST, EXPRESSIVE SE(n) EQUIVARIANT NEURAL NETWORKS
3.1 WEIGHT-SHAAING AND OPTIMAL INVARIANT ATTRIBUTES
Problem statement. Our objective is to find an invariant attribute a:j that can be associated with two points my) in homogeneous space X of group G. such that it satisfies these criteria:
4

---

Published as a conference paper at ICLR 2024
(i) Invariance the global action of G. Any pair in the equivalence class [2,, {g 24 g2g
9 E G} must be mapped to same attribute Oij-
iii Uniqueness. Each attribute dij should be unique for the given equivalence class.
This problem boils down to finding a bijective map [y] + Gij. Note that invariance alone is insufficient to build an expressive network. For example, mapping 23] + 0 is invariant but trivial, and thus not useful. Bijectivity ensures that each attribute is unique and it fully characterizes the space of all possible equivalence classes of point pairs. In essence, bijective attribute mapping over equivalent point pairs and obtain full expressiveness.
is all you need to enable The notion of weight sharing is formalized by the following three definitions:
Definition 3.1 (Equivalent point pairs). Two point pairs (Iz, Iz). (2), x/) Xxx in a homoge- neous space X of a group G are called equivalent iff they can be transformed into each other through left multiplication by an element 9 E G. We define this equivalence relation as:
(g)) ~ (1)
3gEG (x), = (g Rg 9 z5)
Definition 3.2 (Equivalence class of point pairs), The equivalence relation ~ of definition 3.1 de- fines an equivalence class of point pairs, denoted as:
[zg 2z] x() E X x (z),,2)) (z,,
Here, [2t, represents set of point pairs that should be treated as equivalent, with (2;, .,, serving as its representative. The space of equivalence classes of point pairs is denoted as X XxX/ ~.
Definition 3.3 (Weight-sharing in message passing). A message passing layer (Eq. 2-4) is said to share weights over equivalent point pairs it processes equivalent point pairs in the same way. That is, if its message function ém,(f. (f,, fy; [2t, xzl) is conditioned on the equivalence class [z),
See (Koishekenov & Bekkers, 2023) for options for conditioning message functions. Our goal now to parameterize the space of equivalence classes x X ~ with concrete attributes that condition message functions. To that end, reduce the problem of identifying equivalence classes of point pairs to the task of identifying orbits of single points X, akin to LieConv (Finzi et al., 2020a). We show that equivalence classes [x,,, 2y] correspond to orbits in H using the following lemma:
Lemma (Equivalence class correspond to H -oriits in X). For any chosen representatives g, G of o4 E X G/H such that x;=9; To. and any 2j E X, the following mapping from the space of equivalence classes of point pairs X X to the space X of orbits of H in X is bijection:
[x, + Hg, Tj. Proof. See the appendix, section B.1.
[2,, E X XXXX ~9 9, 9 xj E H
(6)
The subsequent theorem characterizes bijective mappings from equivalence classes æyl crete attributes a. that we can compute for homogeneous spaces of SE(n). Since these mappings bijective, the attributes serve as unique identifiers of the equivalence classes. Therefore, they are sufficient for the construction of expressive yet efficient equivariant message passing networks.
Theorem 1 (Bijective attributes for homogeneous spaces of SE(n)). Consider three pairs of homogeneous spaces X=SE(n) H and stabilizer subgroups H: (i) position space: X=R" H=SO(n); (ii) position-orientation space: X=Rd x S5=1, H=SO(n-1); and (iii) the whole group: X=SE(n), [=] For n=2,3, the following mappings from equivalence classes [zt, x,] XxX/ ~  as defined in definitions 3.1 and 3.2, to explicit attributes are bijective:
2 and 3
P,l + aij ||p, -Pill.
(7)
R² x 5t and SE(2): (p: (p).o.))pp.)))) (p)
H ai (R,, (p: arccos arccos oTo,),
(8)
(p P.)
R x S²:
(p:. (,). + aij [I(P; -P:) - o ( - -P)o::1
(9)
arccos 0,
SE(3):
[(,, R.), (p,, R,)] H ag (R (p, -P),, R. RR)
(10)
Proof. A proof is provided in the appendix, section B.2.
The three components of dzy in Eqs. 8, 9 share the same interpretation: the first two decompose displacement p P, into shifts parallel and orthogonal to o, the third represents the angle between o and oj. As for any coordinate system, these mappings are not unique: rewrite Eq.
5

---

Published as a conference paper at ICLR 2024
in polar coordinates aij (Ilp, Pill, arccos o (p: p.), arccos of the invariants used in DimeNet (Gasteiger et al., 2019). Eqs. 8-10 implement 9; 9j. The following is proven in Appx. B.4 Corollary 1.1. Message passing networks in geometric graphs over I S and SE(n), with message function conditioned on the attributes of Thm. 1, are equivariant universal approximators.
3.2 SEPARABLE GROUP CONVOLUTION IN POSITION-ORIENTATIOO SPACE
Based on the bijective equivalence class embedding from Sec. 3.1, we define a group convolution as Jgs Jss k([pp. o). (p', f')l) o' ''p'''''
foul (p.o) 0) -
(11)
Here, f k([(p, o), (p', o'))) E Couut returns a linear transformation matrix for each equivalence class of point pairs, resulting an output signal fout x + RCoue with Cout channels.
S22 RCi. represents the input feature map with Can channels. the kernel
We use the fact that the third component of the invariant attribute in Eq. 9 only depends on orienta- tions separate the R$ S2 convolution into three parts: spatial convolution, spherical convolution, and channel mixing, as illustrated in figure 1. This approach combines the efficiency strategies of Knigge et al. (2022); Chollet (2017), to separate group convolutions over different group parts, and channel interactions from spatial interactions -generally homogencous space interactions-, respec- tively. Following this strategy, we factorize the convolution kernel into three parts:
k((P,,), (p',
(3) (oTo') (oT('' - p), llo I (p' - l)
(12)
with k:(1) RxR3 - C the channel-wise spatial mixing kernel, k(22) [0, m] RC the channel-win wiie orientation mixing kernel, and K(3) RC XCin the channel mixing kernel. Substi- tuting the factorized kernel in Eq. 11 allows us to decompose it into three distinct modules:
SpatialCCons
SphericalCCony
Linear
fout,
(13)
with Linear
the usual linear layer -implementing (39) and
Spat So 1 GConnv
(1) 1io
Σ
k
(oT - pP)) 10. (P, -P:)l) @ 83
(14)
SEN(()
N-1
SphericalGConv
Jio
(1 Σ d=
@ (1)
(15)
Here @ denotes element-wise multiplication and we turned the integrals of Eq. 11 into discrete sums via Eq. 5. In the discrete implementations, the position-orientation space feature maps are of dimension [P, NN CJ. with /' the number of points in the point cloud, N the number of orientations on a precomputed spherical grid that is shared over all positions, and with C the number of channels. We use fj.o to denote the feature vector at position pj and orientation o.. Unlike traditional point- cloud methods, we introduce an additional axis to index features over different orientations.
Each of the operations in Eq. 13 is highly efficient. The most expensive module is SpatialGConv, as it requires aggregating over spatial neighborhoods. We minimize the computation cost of this operation by avoiding channel mixing in this step, essentially performing a depth-wise separable convolution (Chollet, 2017). SphericalGConv is efficient as it can be batched over positions and channels. Moreover, since the spherical grids are shared over all positions, the kernel (2) can be precomputed, resulting in a kernel of shape [N, N', C). Eq. 15 is then implemented as an einsum over the N' axis of the kernel and input with shape [P, (P,'CC).
While the feature maps in Eq. 14, 15 are discretized, the convolution kernels are still continuous. We parameterize these using Multilayer Perceptrons (MIPs) that take attributes as input. Follow- ing standard practices, we first embed the attributes, which we choose to do using Polynomial Embedding (PE) (cf. Appx. C). That is, we sample the kernel via a
PE
MLP
k.
3.3 THE PONITA ARCHITECTURE
We use the separable R x group convolution in ConvNeXt (Liu al., 2022b) layer structure:
SpalialGCunv
Spher MnnerivalCCrrv
Layertiozm
Linear
Aclival
Linear
fout

---

Published as a conference paper at ICLR 2024
and use it to construct simple fully convolutional neural network given by:
f + NodeFmbed
+
ConvNeXt x L Readout + four
(16)
We call this position-orientationn space network based on invariant attributes PONITA. Since PONITA uses position-oriennation space feature maps, we require a method to embed inputs onto spherical signals, and readout outputs from it. We do so via the following node-wis wise modules for embedding and predicting scalar fields 8 - R and vector fields v - R2:
N-1
SphereToScalar
8
h
ScalarToSphere
fi = 86
N
SphereToVec
Vi Si Ooz
VecToSphere
fi = =Voov
The VeeToSphere module is essentially a spherical harmonic embedding of frequency one, fol- lowed by an inverse spherical Fourier transform: [x(=)))) fers a generalized approach for embedding higher-order tensors onto the sphere by using them as coefficients in an inverse Fourier transform, c.g., using the e e3n library (Geiger & Smidt, 2022).
This perspective of-
4 RELATED WORK
The theoretical part of our work is inspired by other works on the construction of equivariant neural networks on homogeneous spaces, such as (Cohen et al., 2019; Kondor & Trivedi, 2018; Bekkers, 2019). Each of these papers contains a version of the "convolution is all you need" theorem, which we take as a motivating starting point in this paper (cf. section 2.2). For a comprehensive unifying coverage of equivariant methods, we recommend (Weiler et al., 2023).
Group convolutions are typically defined in regular or steerable / tensor field form. The regular group convolution viewpoint -which adopt in this paper- naturally extends convolution tem- plate matching of a kernel over an underlying signal, but now matching all transformation in group G instead of translation only (Cohen & Welling, 2016; Bekkers, 2019). On gridded data, such as im- ages, regular group convolutions either achieve exact equivariance by considering symmetry groups of the grid (Cohen & Welling, 2016; Worrall & Brostow, 2018), or numerically approximate equiv- ariance by relying on interpolation (Bekkers et al., 2018; Kuipers & Bekkers, 2023), basis functions (Weiler et al., 2018 8; Sosnovik al., 2019), or continuous MLP-based kernels (Finzi al., 2020a; Knigge et al., 2022). The steerable approach bypasses the need for discretizations by working fields of vectors that transform under (irreducible) representations of the SO(n) group Weiler al. (2021). I). Under the 'convolution is all you need' paradigm, regular and steerable methods are equivalent and they relate to each other via SO(n) Fourier transforms (Kondor et al., 2018; Cesa et al., 2021). In Appx. A we make precise how to interpret PONITA in either the steerable or regular form.
In the field of deep learning- bbasdd molecular modeling, in which physical modeling based on Spher- ical Harmonics, Clebsch-Gordan Gordan tensor products, and representation theory of SO(n) is more com- mon, there is a tendency to rely on the tensor field paradigm. Influential contributions in this field Tensor Field Network (Thomas et al., 2018) and Cormorant (Anderson al., 2019). Limitations of such methods are the high computational overhead due to the specialized Clebsch- sor products and the need for solid understanding of representation theory. An advantage of the tensor field approach is that one obtains exact equivariance, whereas the regular group convolution approach requires discretizations of the group, which leads equivariance only up the discretiza- tion resolution. In this work, however, we show that these numerical errors are not detrimental to performance and opt for simplicity and familiarity with the standard convolutional paradigm.
ten-
Another class of equivariant methods is based on message passing with invariant geometric at- tributes, such as DimeNet (Gasteiger et al., 2019), GemNet (Gasteiger et al., 2021), and SphereNet (Liu et al., 2022a). These are closest to the regular group convolution paradigm and could be framed as non-linear group convolutions (Brandstetter et al., 2021) over the homogeneous space of positions and orientations. Our method is different in that it is based on the standard convolution paradigm instead of intricate wiring of interaction modules, and we use dense grids over the sphere instead of sparse equivariant grids defined by edge directions. A notable result from (Gasteiger et al., 2021) is that networks over position-orin oriintaaiinnnnnnnnnnnnoooo space are universal equivariant approximators and thus as expressive as equivariant methods that use feature representations the full group SF(n).
7

---

Published as a conference paper at ICLR 2024
Our method connects to a line of papers on theory and algorithms for the processing of functions over orientation space via orientation scores. See e.g.. (Duits & Franken, 2010a;b) and (Janssen et al., 2018) for respectively 2D and 3D applications (Duits et al., 2021). Noteworthy, Portegies et al. (2015) derive invariant attributes for position- position--rieenatioon pairs derived via the logarithmic map on (E. In particular, these works address the processing of Diffusion- Diffusioo--wightted MRI data -which are signals over H. x 5², for which our work provides simple equivariant learning recipe.
5 EXPERIMENTS
In this section, we evaluate our approach. Comprehensive implementation details, including archi- tecture specifications and optimization techniques, can be found in Appx. C and D.
Benchmark 1: Predicting interatomic potential energy and forces on rMD17. rMD17 (Chris- tensen & Von Lilienfeld, 2020) is a dataset comprising molecular dynamics trajectories of ten small molecules. Each molecule is represented as an atomic point cloud consisting of 3D atom posi- tions and atomic numbers. The objective is to predict the energy conserinn force field for each molecule. The regression targets are the total energy of each molecule and the force on each atom. Following common practice and the laws of physics, we predict the energy with an invariant model and compute the forces each atom as the gradients of the predicted energy with respect to the positions of atoms. We utilize the PONITA model as given in Eq 16, with a linear embedding layer that takes one-hot encodings of the atom numbers as input and maps them onto the sphere via the Scalar ToSphere module. Following MACE (Batatia et al., 2022b), DimeNet and GemNet (Gasteiger et al., 2019; 2021), we predict a node-wise energy E., after every layer using linear readout layer and obtain the total energy as the sum of all predicted energies EiE..,, Elo
Our results are reported in Tab. 1. As baseline, construct the same architecture as PONITA, except that internal feature maps live solely in position space R We label it PNITA. In this case, the pairwise attribute is simply the distance (Eq. 7), which makes PNITA similar SchNet (Schütt et al., 2023). The inability of such models to learn direction-sensitiie representations explains the large performance gap relative to PONITA and other methods that capture directional information.
Tab. also shows that PONITA either matches or outperforms the state-of-the-art on several molecules while being remarkably faster. PONITA is about 3.5 times faster than the established NequlP method (Batzner et al., 2022): a type of steerable tensor field network (Thomas et al., 2018). We benchmarked their optimized codebase within our training framework and recorded a training time per epoch of 20.7, compared to 5.7 seconds for PNITA and 6.1 for PONITA. Additionally, PONITA stands out for its simplicity. It is a general purpose convolutional network, whose design does not require in-depth knowledge of representation theory and molecular modeling like UNiTE (Qiao ct al., 2022), MACE (Batatia et al., 2022b), NequIP (Batzner ct al., 2022), Allegro (Musaclian al., 2023) and BOTNet (Batatia al., 2022a), nor does it involve intricate wiring of interaction layers as in GemNet (Gasteiger et al., 2021) and ViSNet (Wang et al., 2022).
Benchmark 2: Generating molecules with equivariant diffusion models trained QM9. Dif- fusion models have proven to be very effective for the generation of data such as images (Song et al., 2021) and molecules (Hoogeboom et al., 2022). Prior work Gebauer al. (2019); Simonovsky & Komodakis (2018); Simm et al. (2021) has demonstrated the importance of leveraging molecular symmetries for generalization. Through weight-sharing, PONITA is designed to represent such sym- metries, making it a promising candidate for this application. Recently advancements like EDMs (Hoogeboom et al., 2022), and MiDi (Vignac et al., 2023) have employed denoising diffusion for molecular data generation. These models utilize EGNN (Satorras et al., 2021) as an equivariant denoising diffusion model, which operates on both atomic positions and types.
We extend EDMs by building a similar architecture for molecular generation, incorporating a joint diffusion process for both continuous coordinates and discrete features. However, we substitute the EGNN backbone with our proposed method PONITA. For the readout layer in Eq. 16, use the Sphere2Vec module to predict denoising displacement vectors. We train on the QM9 dataset (Ra- makrishnan R., 2014), standard dataset containing molecular properties, one-hot representations of atom types (H, C, N, O, F) and 3D coordinates for 130K molecules with up to 9 heavy atoms.
Tab. 2 reveals that PONITA outperforms several baselines in terms of atom stability -determined by atoms with the right valency-, and molecule stability -determined by the proportion of molecules for which all atoms are stable-. We attribute this significant improvement to PONITA 's ability to
8

---

Published as a conference paper at ICLR 2024
Table 1: Mean absolute errors (MAE) of energy (E) (kcal/mol) and force (F) (kcal/mol/A) for 10 small organic molecules on rMD17 compared with the state-of-the-art.
Molocule
UNITE
GemNet
NequiP
BOTNet
Allegro
MACE
ViSNet
PNITA
PONITA
(317 worrpec))
157 enokgoch)
seceppohhy
Aspirin
E
2,4
9.5
8.2
2.3
2.3
2.3
6.6
2.2
6.6
1.9
16.3
4.7
5.8.0.18
1.7
E
1.1
0.7
0.7
1.2
1.2
0.7
3.2
0.7_.0
Azobenoene
4.2
2.9
3.3
2,6
3.0
2.5
12.2
2.3
0.07
0.04
0.03
0.3
0.4
0.03
0.2
0.17
Benzene
0.73
0.5
0.3
0.3
0.2
0.3
0.2
0.4
0.3g0.00
Ethanol
0.62
0.4
0.4
0.3
0.7
0.4_0.0
F
3.7
3.6
2.8
3.2
2.1
2.1
2.3
4.1
2.500000
Malonaldehyde
0.8
0.8
0.9
0.6g0.03
F
6.6
6.6
5.1
5.8
3.6
4.1
3.9
5.1
4.0g..30
Naphthalone
0.46
0.2
0.2
0.2
0.4
0.2
1.1
0.3000.00
F
2.6
1.9
1.3
1.8
0.9
1.6
1.3
5.6
1.3g0.00
1.9
1.4
1.3
1.5
1.3
1.1
2.8
1.1 _..00
Paracetamol
F
7.1
5.9
5.8
4.9
4.8
4.5
11.4
4.3:0022
Salicylic acid
E
0.73
0.7
0.8
0.9
0.9
0.7
1.7
0.7.0000
F
3.8
5.3
4.0
4.3
2.9
3.1
3.4
8.6
3.3g0.07
E
0.45
0.3
0.3
0.4
0.5
0.3
0.6
0.3
Toluene
F
2.5
2.2
1.6
1.9
1.8
1.5
1.1
3.4
1.3g0.07
Uracil
E
0.58
0.4
0.4
0.6
0.5
0.3
0.9
0.4 0.4_a.04 00.00
3.1
2.1
2.1
2.4g0.00
Table 2:
Molecule generation via denoising diffusion models trained on QM9. Negative Log Likelihood (NLL), atom, and molecule stability.
Table 3: Mean squared error on N- body trajectory prediction.
Method MSE
seclepoch
Models
NLL
Atom stability
Mol stability
SE(3)-Tr.
.0411
E-NF#
59.7
85.0
4.9
TFN
.0555
G-
N.A.
95.7
68.1
NMP
.0107
GDM
-94.7
97.0
63.2
Radial Field
.0101
GDM-aug"
-92.5
97.6
71.6
EGNN
.0070 400022
EDM
-100.7
98.720.1
82.0.o.
SEGNN G+P
.0013 00015
1.59
CGENN
.0039
PONITA
137.4
98.9
87.8
PONITA
.0013
0.66
learn orientation-senntae representations: a feature absent in other equivariant diffusion models such as G-Schnet (Gebauer et al., 2019), E-NF (Garcia Satorras et al., 2021) and therein introduced non-equivariant baseline Graph Diffusion models (GDM) trained with and without augmentation.
Benchmark 3: Predicting trajectories in N-body systems. Finally, we benchmark PONITA on the charged N-body particle system experiment proposed in Kipf et al. (2018). It involves five parti- cles that carry either positive or negative charge, each having initial position and velocity in a 3D space. The objective is to predict all particle positions after a fixed number of time steps. We adapt the experimental setup from (Satorras et al., 2021; Brandstetter et al., 2021) and use PONITA for tra- jectory prediction. Following the baseline methods, PONITA takes two scalars: velocity magnitude and charge, and two vectors: initial velocity and direction to the center of the system, input, and outputs a displacement vector. We measure the performance in Mean Squared Error (MSE) and com- pare against (B33))Trrasformeer (Fuchs al., 2020), Tensor Field Networks (Thomas al., 2018), Neural Massage Passing (Gilmer et al., 2017b), Radial Field (Köhler et al., 2020), EGNN (Satorras al., 2021), SEGNN (Brandstetter et al., 2021) and CGENN (Ruhe al., 2023). Tab. 3 shows the results and the reproduced train- hineetiieeeppch for SEGNN, showing that PONITA is #2.5 times faster. Once again, our versatile equivariant architecture, PONITA, delivers outstanding performance.
Benchmark 4: Superpixel MNIST and Benchmark 5: QM9 regression. In App. E. and E.2 we further show the potency of PONITA with state-of-the-art results on benchmarks 2D and 3D and further compare an edge-index-induced R x 52 point cloud vs a grid-based (fiber bundle) approach.
CONCLUSION
In summary, our method simplifies the construction of equivariant neural networks by following the paradigm of weight-sharing, achieving ttateoof-thear results in various benchmarks. In Thm. 1, we derived geometric attributes that fully characterize the space of equivalence classes of point-pairs in homogeneous space of SE(n). As such, these attributes are all you need to facilitate SE(n) equivariant weight-sharing.
an application of the theory, we propose PONITA: an efficient convolutional architecture for equivariant processing of 3D point clouds. We show that our general- purpose equivariant approach serves as a compelling alternative to prevailing tensor field methods.

---

Published as a conference paper at ICLR 2024
REFERENCES
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Jimmy Aronsson. Mathematical Foundations of Equivariant Neural Networks. Chalmers University of Technology, PhD Thesis, 2023.
Pedro HC.Avelar Anderson R Tavares, Thiago LT da Silveira, Clíudio R Jung, and Luís Lamb. Su- perpixel image classification with graph attention networks. In 2020 33rd SIBGRAPI Conference Graphics, Patterns and Images (SIBGRAPI) pp. 203-209, IEEE, 2020.
Ilyes Batatia, Simon Batzner, Dávid Péter Kovács, Albert Musaelian, Gregor NC Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi. The design space of (3)equivariant atom- centered interatomic potentials. arXiv preprint arXiv:2205.0662 06643,, 2022a.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:11423- 1436, 2022b.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Ko- rnbluth, Nicola Molinari, Tess Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1): 2453, 2022.
Erik Bekkers. B-spline enns on lie groups. Representations, 2019.
In International Conference Learning
Erik J Bekkers, Maxime W Lafarge, Mitko Veta, Koen AJ Eppenhof, Josien PW Pluim, and Remco Duits. Roto-translation covariant convolutional networks for medical image analy- sis. In Medical Image Computing and Computer Assisted Intervention-MICCAI 2018: 21st International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part I, pp. 440- 448. Springer, 2018
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https:////// wandb. com/ Software available from wandb.com.
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling. Geomet- ric and physical quantities improve (3) equivariant message passing. International Conference Learning Representations, 2021.
Gabriele Cesa, Leon Lang, and Maurice Weiler. program to build e (n)-equivariant steerable cnns.
In International Conference Leamning Representations, 2021.
François Chollet. Xception: Deep learning with depthwise separable convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1251-1258, 2017.
Anders S Christensen and O Anatole Von Lilienfeld. On the role of gradients for machine learning of molecular energies and forees. Machine Learning: Science and Technology, 1(4):045018, 2020.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pp. 2990229999999 2999.9999..00 PMLR, 2016.
Taco S Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant cnns homo- geneous spaces. Advances in neural information processing systems, 32, 2019.
Pim De Haan, Maurice Weiler, Taco Cohen, and Max Welling, Gauge equivariant mesh cnns:
Anisotropic convolutions on geometric graphs. In International Conference on Learning Representations, 2020.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794 2021.
10

---

Published as a conference paper at ICLR 2024
R Duits, BMN Smets, AJ Wemmenhove, JW Portegies, and EJ Bekkers. Recent geometric flows in multi mullti-orientation image processing via a cartan connection. Handbook of Mathematical Models and Algorithms in Computer Vision and Imaging: Mathematical Imaging and Vision, pp. 1-60, 2021.
Remco Duits and Erik Franken. Left-invariant parabolic evolutions on se(2) and contour enhance- ment via invertible orientation scores part i: Linear left-invariant diffusion equations se(2). Quarterly of Applied Mathematics, 68(2):255-29) 2010a.
Remco Duits and Erik Franken. Left-invariant parabolic evolutions on se(2) and contour enhance- ment via invertible orientation scores part ii: Nonlinear left- inttrinvari diffusions invertible ori- entation scores. Quarterly applied mathematics, 68(2):293-33) 2010.
Nadav Dym and Haggai Maron. On the universality of rotation equivariant point cloud networks. In International Conference on Learning Representations, 2020.
Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Matthias Fey, Jan Eric Lenssen, Frank Weichert, and Heinrich Müller. Splinecnn: Fast geomet- ric deep learning with continuous b-spline kernels. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 869-877, 2018.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolu- tional neural networks for equivariance to lie groups arbitrary continuous data. In International Conference on Machine Learning, pp. 3165-3176. PMLR, 2020a.
Mare Anton Finzi, Roberto Bondesan, and Max Welling. Probabilistic numeric convolutional neural networks. In International Conference on Learning Representations, 2020b.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se trnntffsfsssmrss: 3d roto- translation equivariant attention networks. Advances in neural information processing systems, 33:1970-1981, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling, E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181- 4192, 2021.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molec- ular graphs. In International Conference on Learning Representations, 2019.
Johannes Gasteiger, Shankari Giri, Johannes T Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115. 2020.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wort- man Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: / //openeeview..... /orr//oooi wttt://frmmmmomomm -s H__________________________
Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry- generation 3d point sets for the targeted discovery of molecules. Advances in neural information processing systems, 32, 2019.
Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks, 2022. URL htttttsss///xii org/abs 2207.222222 09453.
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp- 1263-1272. PMLR, 2017a.
11

---

Published as a conference paper at ICLR 2024
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1263-1272. PMLR, 06-11 Aug 2017b. URL https:// proceedings mmlr www.ppresssveeee g1 m 7a .htll.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems. 33:6840-6851. 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR. 2022.
Michiel HJ Janssen, Augustus JEM Janssen, Erik J Bekkers, J Oliván Bescós, and Remco Duits. Design and processing of invertible orientation scores of 3d images. Journal of mathematical imaging and vision, 10:11427-458, 2018.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-21777, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint utXiv:1412.69888 2014.
Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International conference on machine learning, pp. 2688- 2697. PMLR, 2018.
David M Knigge, David W Romero, and Erik Bekkers. Exploiting redundancy: Separable group convolutional networks on lie groups In International Conference on Machine Learning, pp. 11359 1138. PMLR, 2022.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pp. 5361-536777 PMLR, 2020.
Yeskendir Koishekenov and Erik Bekkers. An exploration of conditioning methods in graph neural networks. In ICLR 2023-Machine Learning for Drug Discovery workshop. 2023.
Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International Conference on Machine Learning, pp. 2747-2755. PMLR. 2018
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. ftebsch-gordhe nets: fully fourier space spherical convolutional neural network. Advances in Neural Information Processing Systems, 31, 2018.
Thijs P Kuipers and Erik J Bekkers. Regular se (3) group convolutions for volumetric medical image analysis. In Medical Image Computing and Computer Assisted Intervention-MICCAI 2023. arXiv preprint arXiv:2306.1 2023.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):278-23244 1998.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Representations, 2022a. URL https::////.......t. / /wwpreevieeeeeeeeeeeeeeeeee www..net/oooummiffieem wttp://rrnniiiiissdiiioo 9t..
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976-11986. 2022b.
Federico Monti, Davide Boscaini, Jonathan Masci, Emanuele Rodola, Jan Svoboda, and Michael M Bronstein. Geometric deep learning on graphs and manifolds using mixture model enns. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5115-5124, 2017.
12

---

Published as a conference paper at ICLR 2024
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Ko- mbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury. Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances neural information processing systems, 32, 2019.
Michael Poli, Stefano Massaroli, Eric Nguyen. Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302. 10866, 2023.
Jorg Portegies, Gonzalo Sanguinetti, Stephan Meesters, and Remco Duits. New approximation of scale space kernel on se (3) and applications in neuroimaging. In International Conference Scale Space and Variational Methods in Computer Vision, pp. 40-52. Springer, 2015.
Zhuoran Qiao, Anders S Christensen, Matthew Welborn, Frederick R Manby, Anima Anandku- mar, and Thomas F Miller III. Informing geometric deep learning with electronic interactions to accelerate quantum chemistry. Proceedings of the National Academy of Sciences, 119(31): e2205221119, 2022.
Rupp M.and Von Lilienfeld O.A. Ramakrishnan R., Dral P.O. Quantum chemistry structures and properties of 134 kilo molecules, 2014.
David W Romero, Anna Kuzina, Erik J Bekkers, Jakub Mikolaj Tomczak, and Mark Hoogen- doorn. Ckconv: Continuous kernel convolution for sequential data. In International Conference on Learning Representations, 2021.
David Ruhe, Johannes Brandstetter, and Patrick Forré. Clifford group equivariant neural networks. arXiv preprint arXiv:2305. 11141, 2023.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference machine learning, pp. 9323-93333 PMLR, 2021.
Kristof Schütt,, Stefaan SP Hessmann, Niklas WA Gebauer, Jonas Lederer, and Michael Gastegger. Schnetpack 2.0: A neural network toolbox for atomistic machine learning. The Journal of Chemical Physics, 158(14), 2023.
Gregor N. C. Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hernández-Lobato. Symmetry-awarr actor-critic for 3d molecular design. In International Conference on Learning Representations, 2021. URL https: wwwpreviee.eeece////pp/e. net ftt::/ooommmmdmdddde
Martin Simonovsky and Nikos Komodakis. Graphvae: Towards generation small graphs using variational autoencoders. In Artificial Neural Networks and Machine Leaminng ICANN 2018: 27th International Conference on Artificial Neural Networks, Rhodes, Greece, October 4-7, 2018, Proceedings, Part 27, pp. 412 422. Springer, 2018
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference machine learning, pp. 2256-2265. PMLR, 2015.
Yang Song. Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. Advances in Neural Information Processing Systems, 34:1415-1428, 2021.
Ivan Sosnovik, Michat Szmaja, and Amold Smeulders. Scale-equivariant steerable networks. In International Conference on Learning Representations, 2019.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich- Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions low dimensional domains. Advances in Neural Information Processing Systems, 33:7537-7547, 2020.
13

---

Published as a conference paper at ICLR 2024
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and eequivvainnnniiiiiaann neural networks for 3d point 2018.
clouds. arXiv preprint
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Clement Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for molecule generation. arXiv preprint arXiv:2302.09048. 2023.
Soledad Villar, David W Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars
universal: Equivariant machine learning. structured like classical physics. Advances in Neural Information Processing Systems, 34:28848 28863, 2021.
Ian Walker and Ben Glocker. Graph convolutional gaussian processes. In International Conference on Machine Learning, pp. 6495-6504, PMLR, 2019.
Yusong Wang, Shaoning Li, Xinheng He, Mingyu Li, Zun Wang, Nanning Zheng, Bin Shao, Tong Wang, and Tie Yan Liu. Visnet: scalable and accurate geometric deep learning potential for molecular dynamics simulation. arXiv preprint arXiv:2210.16513 2022.
Maurice Weiler and Gabriele Cesa. General (2)-equivariant steerable cnns. Advances in neural information processing systems, 32, 2019.
Maurice Weiler, Fred A Hamprecht, and Martin Storath. Learning steerable filters for rotation equiv- ariant cnns. Proceedings the IEEE Conference on Computer Vision and Pattern Recognition, pp. 849-858 2018.
Maurice Weiler, Patrick Forré, Erik Verlinde, and Max Welling, Coordinate independent convolu- tional networks-isometry and gauge equivariant convolutions on riemannian manifolds. arXiv preprint arXiv:2106.06020, 2021.
Maurice Weiler,
Erik Verlinde, Equivariant and Coordinate Independent Convolutional Networks.
Patrick
Foréé
and Max
Welling.
2023.
URL
https: ce wweiler.gitlle wt/cio///ooooooeeeeee pdf.
Daniel Worrall and Gabriel Brostow. Cubenet: Equivariance to 3d rotation and translation. In Proceedings the European Conference on Computer Vision (ECCV), pp. 567-584, 2018.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: geo- metric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.
14

---

Published as a conference paper at ICLR 2024
A THE FIBER BUNDLE VS POINT CLOUD VIEWPOINT
A.1 REGULAR VS STEERABLE GROUP CONVOLUTIONS
In the field of group equivariant deep learning, one often adopts a steerable or regular convolutional viewpoint equivariant neural networks (Weiler & Cesa, 2019). Let us consider the case of (E()) equivariance, and note that the group SE(n) is a semi-direct product between the translation group and the rotation group SO(n). We then consider functions (feature fields) over domains X that are homogeneous spaces of SEn),, namely the case X = R". X IR" x S and X (n). The difference between steerable and regular group convolutions then lies in the type of feature fields that are considered.
With regular group convolutions one considers multi-channel scalar fields f X R of the homogeneous space X. And convolutions are of the form
pout = kgg ff (x)dx,
(A1)
with By representative of the point y such that By zo with zo some arbitrary origin in X. Regular feature fields, and convolution kernels, transform via the regular representation of SEn)) via
(p'()))((x) ==
Given this, regular group convolutions can be simply considered template matching between a roto-translated convolution kernel (by transforming it with the regular representation) and the underlying signal, where similarity is measured by the L2 inner product. When X SF((n) with a non-trivial group, such as = SO(3) when X = RD or H = SO(2) when = Rd x S², then the kernel should satisfy kernel constraint WReH k(x) khaz), see (Bekkers, 2019). As discussed in the main body of this paper, this constraint is automatically satisfied if we parametrize the kernel by the invariant attributes d:y of Theorem 1. The advantage of regular group convolutions is that the codomain of the feature fields are just scalars and one can use any point-wise wise activation function without breaking equivariance. A drawback of such methods though is that one needs to discretize the domain X, including the fibers 5 or SO(n), which leads to numerical inexactness of the equivariance property!.
With steerable group convolutions one considers feature fields f R + Vp over a base space IR with a codomain which is a vector space V, on which a representation p of SO(n) can act. Such feature fields transform via the so-called induced representation of SE(n), defined as
(Innd SEE() (l()))((x) = p(R) f(g))((g"))"
Le., the induced representation [Tndso(n) p) does not only transform the feature map's domain but also its codomain via the representation p. For example, if one roto-translates a vector field v R3 + R3, the vectors vxx) in the field at location should be moved to the new locations 'x, but also their values should be rotated via Rv(g The steerable convolutions are then simply given by
foul (y) = / k(x yf- RR
(A2)
Now, however, the kernel needs to satisfy the kernel constraint
{(R x) = Pour (R)) (R))
with Pin and Pout the representations that define the vectors spaces of the codomains of the input and output feature fields respectively, see e.g. (Weiler & Cesa, 2019; Cesa al., 2021). An ad advantage of the steerable approach is that one can obtain exact equivariance by choosing to let the vector spaces direct sum of irreducible representa-
V, transform via representations p @iPt that are given tions. Such irreducible representations do not require a grid and rotations can be computed exactly. A drawback of such an approach is, however, that one can no longer apply arbitrary point- wise activation functions but one needs specialized operations on Vp in order not to break equivariance. This limits expressivity (Weiler & Cesa, 2019). Another drawback, arguably, is that this approach requires considerable understanding of representation theory for proper use in practice.
Our experiments, however, show that this approximate equivariance is detrimental to performance
15

---

Published as a conference paper at ICLR 2024
A.2 REGULAR GROUP CONVOLUTIONS IMPLEMENTED AS STEERABLE GROUP CONVOLUTIONS
It is important to note that in the regular group convolution case for SE(n) equivariance, one rec- ognizes a semi-direct product structure of the group SE(n) = (R"; x (n()), with (R", +) the translation group, which as space can be identified with position space R When one then consid- ers regular group convolutions over X = I" x S or X = R x SO(n), one could think of the feature fields f RR" Y + R as assigning to every position R a function fx := f(x,) -) over the space Y, with here Y = 5 or Y = (O))). These functions fx are to be considered vectors the infinite-dimensionaa vector space of square-integrable functions (Y). For this vector space we know the left-regular representations pfa which are simply given
( = fx fx(" o), and (PL)()0()) (R) (R(RRR).
Thus, one can think of regular group convolutions as working with feature fields f R V Vigen
Further, note that any (O(n) representation can be decomposed into direct sum of irreducible representations via change of basis. Le., there is linear change of basis such that Vie 222))) = V for some P = DiPi. with Pi an irreducible representation. This change of basis is given by the Fourier transform on Y. Namely, a function fx E Vigeor can be transformed to a vector of Fourier
coefficients via fx = [Sx] Vp, and back via the inverse Fourier transform fx = Jy [f.)
Regular group convolutions can thus either be implemented steerable convolutions with fields of signals over Y, i.e., f T > VoL "gg(r>" or via fields of irreducible representations f RR" - V
that be thought of as fields of Fourier coefficients. In this latter case, the kernel constraint is satisfied by parametrizing using the Clebsch-Gordee tensor product (Cesa et al., 2021; Brand- stetter et al., 2021). This approach is in fact how group equivariant neural networks are of- ten implemented (Weiler & Cesa, 2019; Thomas et al., 2018), however, it has the drawback that one either has use specialized activation functions use activation functions of the form severseecuries + Act i CentionFunceee
+ FourierTrans form, which leads to computational overhead and leads to discretization artifacts that break exact equivariance. In this paper, we present a fully regular group convolution viewpoint that does not require Fourier trans- forms or specialized activation functions.
A.3 FIBER BUNDLE VIEWPOINT FOR PONITA
From fiber -thlorettc viewpoint, the feature fields f R - Vp can be regarded as sections of vector bundles associated with the principal fiber bundle RR" x SO(n) + R" (Weiler et al., 2023; Aronsson, 2023). As such, without going into further technical details behind this theory, we merely use this connection to refer to the vector spaces Vp as fibers. The PONITA model processes feature with vector spaces of spherical signals Vg "1115-11 as fibers. Instead of
fields R Vp (1
sample describing these signals in basis of irreducible representations (the Fourier approach), them on a grid S C Sa-1 of orientations. See Figure 1 in which the grid is depicted as a set of yellow dots, and the actual spherical signal is depicted efficient use of the fiber bundle structure by separating convolutional operations over the field into spatial interaction part (convolving over IR") followed by a within-fibe fiber orientation interaction part (convolving over S"
gray-scale density. Our model makes
A.4 A POINT CLOUD VIEWPOINT FOR PONITA
The efficient implementation of PONITA is made possible by the fact that each node shares the same grid (fiber) of orientations S. Let us denote with Sx C - the grid of orientations available at position x. If we share the same grid over all positions, i.e., Sx = S, we can decide to spatially convolve per orientation in the grid as follows
/
(x)x) (x,
poul
k([(x, o), (x', o)))"" (x)''''
which is the continuous formulation of the discretized spatial convolution of Eq. (14). This purely spatial convolution is made possible by the fact that each position x' we can sample the same
16

---

Published as a conference paper at ICLR 2024
orientation E S. However, if we do not share the same grid S over all positions can longer separate the convolution. In this case, we adopt a point-cloud viewpoint.
There are natural instances in which space. Examples include triangulated shape meshes, which can be represented by the centroids the triangles p and their normal vectors o. See e.g. (De Haan al,,,2200 2020). In contrast (De Haan al., 2020) our method would take the curvature of the surface mesh into account (via the invariant of o,), as well as off plane distance, however, the method by De Haan al. (2020) takes anisotropic features within the plane into account whereas the second invariant in Eq. (9) restricts the kernel to be isotropic within the tangent plane.
has to compute with point clouds in position-orientatiin
Another instance of a the edges an atomic point cloud, as DimeNee/GemNNet (Gasteiger al., 2019; 2021). In those works, grid is assigned to each atom location x; based on the direction towards neighboring atom locations X, that are connected by covalent bonds. I.e., each atom obtains grid Sx. = { X; -x
space point cloud would be in a message passing over
-x,-
j E (i)), with (ii) the set of neighbor indices of node i.
Convolutions in these case can no longer be separable and include spatial and orientation interactions simultaneously:
four (x,, o,)
Σ
k[[(x,, o.), (x,, o))) (x,, o).
(EN(() DESx;
This is the discretized version of the full continuous convolution of Eq. (11).
A.5 IMPLEMENTATION
Both versions of position-orientation space convolutions -the fiber bundle approach as well as the point cloud approach- are made available in the public repository http://gtttuu....eeeeeeeeeeeeeee
B PROOFS
B. PROOF OF LEMMA
Lemma shows that equivalence classes of point pairs [r;, zyl (192,,922) I E G)} correspond orbits H 9; j. for any chosen şi such that x; gi Ho, with nj, 2D X. In the following proof, one might recognize that the space of equivalence classes X x X/ ~ can also be identified with the space (H, doubbe cosets in G. That is X X/ HAGH H The proof of Lemma 1 is as follows:
Proof. To prove the bijection given by equation 6 we rewrite the equivalence class into the form of double coset as follows. For any representatives (g., gj) we have
[ 25]
((g gzz) g E G} ((g h, mo.g 99) hy mo) 1 G, H1}
9 h 97'
{(h 9, -o hy mo. mo) g E G, hj E H}
gme
{(h 9. 'g hy ma, Gma) h hi H)
my= mo. by 2e =m
{(, 9, ---za, z,, mm) hy hy H}
G Cma=M
[h4 g, x, h4 HH}
H 9; mz.
where in the second to last step note that the second item in the tuple, G is constant (it represents the entire homogeneous space X = Gmg), and thus we can represent the tuple simply with the first element H 9; 2.. Thus any equivalence class [x,, 2y) can be represented by the orbit Hg, Ty with 9, any such that x; 9: Zo-
17

---

Published as a conference paper at ICLR 2024
B.2 PROOF OF THEOREM 1
Theorem lists four cases of equivalence classes and provides bijective map for each of them:
R² and R²:
[pi.Psl H dij =||p> -Pall,
(A3)
R² x S' and SE((): (i),, (p), 0,)]
H azj =RR, (p, -Pi), arccos
(A4)
of (P: p:)
R.a x 5
(p: 0z), (Pj: 0j]]
+ Gij =
(1(py - p.) - 0 (p) - -p:jo.ll
(A5)
arccos 0,
SE(3):
((P.. (P,, Ry)]
+
aij (R, (p) p.), RR)).
(A6)
We provide constructive proof for each of these cases by providing inverse mapping dij [pa, py]. This boils down to defining representative zij of the orbit of H g, az for given aij. The proofs are as follows.
Proof. For the I" case the orbit representative can be given given by
Fij = 0ij 0z
(A7)
with e₂ -0 -(8). The equivalence class is represented by the orbit H H
= Hl Pej. with Pij = Pj - Pt. Thus Pij is valid representative. Since H does not alter the norm of Pij and it acts transitively on spheres of a given radius, we can say that zij lies in the orbit since |æyl|| = ||,,,==|| ||P;yll- This proofs that the mapping in (A3) is bijective, with (A7) as inverse mapping.
To prove the ? S SE((), and SE(3) cases we first remark that IR2 S is equivalent to the SE(2). We can uniquely identify any (p.o) E R x 5" with a roto-translation (p. R.) E SO(2) via (p. = (p. R.) cz) with er = (b), noting that the stabilizer group Stabso/2) (0, cz)
I) } is trivial. Thus the spaces are equivalent. Since in these three cases the stabilizer group H = {e} is trivial, there is only one unique group element associated with each point x; = 9i and Iy = -j. with gt. By such that I; gi 2g and zj = gj Ty. And thus the orbit that represents the equivalence class [2i, = H 9. representative
-
x4
{gr 9;} consist of a single element. We can write its
Tig 9, 93
(A8)
which is precisely the relative group action for SF(2) in (A4) and for SE(() in (A6). This proofs that (A4) and A6) are bijective with (A8) as inverse mapping.
The representative for the R x S case is given by
Tzj
sin cas
))
(A9)
with a, b, e the components of the attribute where aij- = (a, b. c)²
Note that the orbits are given by H .R. R- the e axis act on three independent subspaces of T x St. The spatial part of which is given by
Further note that rotations in H around
( and e, (g) We project the chosen
V, (e and V, span(ea,, ey). with e
00
representative onto this basis, such that R- Pij ex + y By ez. First compute = as
== (R.,' ).y) I'e e, = P R! = Pfjo, = =a
T
The vector ae, is invariant to the action of Н. The component of the unit vector in S is also invariant, and thus uniquely determined as
(R, e n. n; = cos
and thus c = arccos n. ny. Since both quantities a and are obtained independent of a chosen rotation in H. we can pick any as long as the vector (8) lies in the orbit H R Ptj. Since H
is norm preserving, we only have to check that for the chosen b the norm [| 0) = coincides with the norm of the Vú, vector component of R Pij. The component of R- P:j in V,, is given by
18

---

Published as a conference paper at ICLR 2024
(Ps)) (pl, e, and its norm by (I (p = [[ R,, (p Ro, e= = (1(p))) - -p,1o,,).1 =b,
where in the second step used that the norm is invariant under any rotation in SO(3), and thus we are free to multiply the vector in the norm with R In conclusion, the components a and c were uniquely determined the relative group action (R Pij. Ro'0,) and had to be chosen to correspond to the radius of the orbit in the zy plane. This proves that (A5) is bijectíve, with (??) as inverse mapping.
For all cases, we have now provided valid representatives, and thus for each provided a mapping from attribute to orbit. And thus the provided equivalence class attribute mappings [yy + ag are bijective.
B.3 PROOF OF COROLLARY 1.1
Corollary implies that the PONITA architectures are equivariant universal approximators. order to prove this we first show that each of the steps in the architecture indeed equivariant, and then make use of the results in Dym & Maron (2020); Villar et al. (2021); Gasteiger et al. (2021) to prove universality.
B.3.1 EQUIVARIANCE, MESSAGE PASSING LAYERS WITH INVARIANT ATTRIBUTES
First, we make precise the notion of equivariance in the point cloud setting (see Sec. A for the bundle point cloud viewpoint). Let consider a point cloud in space X which denote with G = EV, nodes and F (f E V i V) denotes the set node features that live in vector space Vp, with the representation of (0(n) that acts on Vp. Let the action of g = (x, R) E SE(n) on the graph be denoted with
in which X C X = {z, E V} denotes the set of points associated with the
9 E
with
gX (yx; i V.J.
9F (p(R) fi V}.
Then graph message passing layer M MN is said be equivariant to the group SE(n) if
MPN 9) = MPN(G).
Proposition 1.1. Message passing layers described by Eqs 2-4 as conditioned on the altributes in Theorem are SE()) equivariant.
Proof. The layers perform regular group convolution (cf Sec. A) in which case the feature spaces are considered to be vectors of scalars in R on which rotations act trivially, i.e., (R) fi- Since the positions of the point clouds are un-altered in the message passing steps the output point cloud X' equals the input point cloud X, and since the attributes [g Ii, ætş] H aij are invariant for all SEE(), all steps in the message passing scheme are invariant to the group actions. Let us denote the input and output graphs as Gin (V, E, specifically have that
Fin) Fin) and Gout = (V, E, Fwmt). Then we
G g Gout
(V, E, X.F"
So, the node features are invariant and the graph itself (point cloud) is not updated by the message passing layer and thus is trivially equivariant.
B.3.2 EQUIVARIANCE OF LIFTING THE GRAPH FROM Ra TO Rd x 5n
Now consider a lifting transform that assigns an (approximately) uniform grid of elements in S to each based point in a point cloud in X = R". Let such a grid be denoted with C S For now let consider the availability of an infinite dimensional grid S S in the fiber bundle viewpoint Sec. A, and that we are thus able to assign to every point p E and every o E 5 feature vector through the point-wise spherical signals fx -
19

---

Published as a conference paper at ICLR 2024
with X R" to graph G = ( via the ScalarToSphere and VecToSphere modules. The lifted graph G can either be seen point cloud in X C R" S"-1 with scalar features fi RR over the positions in X' C R and with spherical signals Jt E pl (28--11 as node features. We
We then consider lifting of =
(steerable) feature field
denote the latter with G = EV,
We adopt the bundle viewpoint and consider lifting to the graph Go. Scalar features s; at node will be lifted via
oo) ScalarToSphere(sj()) S;
i.c., it generates constant signals over 5 with value 34. Vector features v, at node i are lifted to spherical signal via
fi(o) = pecTooppherevvllo vvoo.
Proposition 1.2. Lifting graph G graph G via the VecToSphere and ScalarToSphere operations, is equivariant via
-V.
(. (V.
That is, the nodes in the lifted graph are equally transformed by the action of g E SE(n) and the spherical signals Ja E L2(5 are permuted via the left-regular representation p via g FP {p(g) ii V/.
Proof. For the ScalarTosphere operation we have that the spherical signals fi(o) 8; are invariant since the scalars in the original graph are unaffected by the group action. For the VecToSphere operation we have that
Rv; + VecToSphere [Rv;](o) = v R o (=1 (R) VecTossphere(v)((o)
and thus rotation of the input vectors leads a permutation by (s2 (R) of the spherical signals.
B.3.3 EQUIVARIANCE PREDICTING SCALARS AND VECTORS
From the lifted graph G we can predict at each position x; X scalar or vector using the SphereToScale ar and SphereToVec modules. In their continuous forms they are given by
/ gn-1 f((o),
8 = SphereToScalax(ff]
(A10)
SSa f(())dd.
V, [==hereeeooecc(]]
(A11)
g
Proposition 1.3. The process of predicting scalars via Eq. (A10) is invariant and predicting vectors via Eq. (Al1) is (m)) equivariant via
SphereTovec(p(R)) SphereToeecl]].
Proof. Since the scalars are invariant to rotations, the SphereToScalar module is trivially in- variant. For the SphereToVes ToVec module we have that if the input spherical signals were to be transformed via p(R), than have
p(R) + SphereToVec(p(R1[]
(o)
Sn-1
[i(Ro))ooo
Sn-1
Ro
Sn-1
R
fi(o)odo
9n-1
R Sphereoveec(f]]
where in the step = used the change of variables
R o.
20

---

Published as a conference paper at ICLR 2024
B.4 PROOF OF COROLLARY 1.1
Corollary 1.1 states the following: Message passing networks in geometric graphs over I" x S and (n))), with message functions conditioned the attributes of Thm. are equivariant univer- sal approximators. The proof is as follows.
Proof. In the above, we have shown that the steps in the PONITA architecture are equivariant, and it allows to predict scalar vector fields. As discussed in Sec. A, the PONITA architecture and its layers fit the steerable tensor-field network class of equivariant graph neural networks. Dym & Maron (2020) have proven such types of equivariant graph neural networks to be equivariant universal approximators. As such, (Dym Maron, 2020, Theorem 2) proves that the bundle im- plementation of PONITA, as described in the main body of this paper, is indeed equivariant universal approximator. the case of the point cloud approach (Sec. A.4), universality is given by (Gasteiger et al., 2021, Theorem 3), which is based the results in (Villar et al., 2021).
L
C IMPLEMENTATION
C.I USED LIBRARY
Our implementations are done in Pytorch (Paszke et al., 2019), using Pytorch-Geometricss message passing and graph operations modules (Fey & Lenssen, 2019), and made use of WandB (Biewald, 2020) for logging. Our code is available at https:/ https:://githtt com/ ebekkers/ponita..
C.2 SPHERICAL GRIDS
Uniform grids on S², As described in the main body, we assign spherical grid to each node in the graph. The grid consists of N points which cover the sphere as uniformly as possible. Note that an exact uniform grid can only be achieved five cases, corresponding to the five platonic solids with N = 4, 6, 8, 12, 20. To achieve an (approximate) uniform grid for any possible N, and thus have full flexibility in specifying angular resolution, we generate grids via a repulsion method. This method randomly initializes points the sphere and then via gradient descent iteratively pushes them away from each other until convergence. We used the repulsion model code of Kuipers & Bekkers (2023), which was used in their work to generate uniform grids SO(3).
Random grids. To minimize a potential bias to particular directions in the grid, we randomly perturb the grids with random rotation matrix in each forward pass of the method. Herein follow the approach of Knigge et al. (2022) for achieving equivariance to continuous groups, sampled on discrete grids. Every graph in the batch obtains its own randomly rotated spherical grid. This is akin to rotational data augmentation.
Grid resolution We found that grid resolution N 12 was generally sufficient for all tasks. However, we got slightly better results with N = 20, and thus ran all results with this setting. Increasing beyond N did not seem to improve results.
C.2.1 ATTRIBUTE EMBEDDING
The continuous kernels used in PONITA be considered as Neural Fields. For Neural Fields, it common to first embed/vectorize the input coordinates. A common approach for this is the Ran- dom Fourier Feature embedding, which samples random Fourier basis functions on the provided input coordinates (Tancik et al., 2020). We experimented with this approach as well but found that simple Polynomial Embedding worked slightly better and was less sensitive chosen hyperpa- rameters (polynomial degree). The polynomial embedding modules take as input vector of values (a, y, 2....) and generate polynomial combinations of the inputs up a maximum degree. Le. Pol 1 ynomi a Rmmeed(r, y, 2) (xr, y, 2, ny, Z.2 yz, Tay, a..2...)).
In all experiments, the kernel used shared basis which was computed at the start of the forward pass
aij
PE
L.i near
ActFn
Linear
ActFF
bij
21

---

Published as a conference paper at ICLR 2024
and each group convolution layer then obtains its layer-specifie kernel from this shared basis via
bi Linear
- kij
In all layers, we used the Gaussian Error Linear Unit (GeLU) as an activation function. all experiments, we used 256-dimensional basis.
D EXPERIMENTAL DETAILS
In this section, we provide the hyperparameters used for training the models and provide additional information regarding the benchmarks and the results.
D.I MD17 EXPERIMENTS
Training settings. The rMD17 results were obtained with PONITA and PNITA with L 5 layers, 128 hidden features. The polynomial degree was set to The models were trained for 5000 epochs, with batch size of 5. We used the Adam optimizerKingma & Ba (2014), with learning rate of 5e - 4, and with a CosineAnealing learning rate schedule with a warmup period of 50 epochs. We used N 20 grid points the sphere. The networks optimized the following loss
OF
L= >=]|E E112 AF
3M Σ
(A12)
Spr,a
with E 12, The loss is weighted sum of energy and force loss terms, Here M is the number of atoms in the system, É is the predicted potential energy and λø and Ap are the energy- and feigh-wes,,,,ngs, respectively. We set 7 500. The embedding layer is linear embedding the one-hot encoded atom numbers, followed by the Sca the main body, the readout layer was applied after every ConvNeXt block and consisted of a single Linear layer.
module. As explained in
D.2 DENOISING DIFFUSION MODELS
D.2.1 SUMMARY OF THE DENOISING DIFFUSION PROBLEM
In diffusion processes, distributions are learned through reverse diffusion process, i.e. denoising process. Consider point clouds æ (ar....., .MM E R Mx3 with corresponding features h (hI,..., ,,,s) E I af xnf the diffusion process that adds noise for time / = 0,... . is given by multivariate normal distribution:
ạ (2. Φ) N' (z+ 0i: ai1).
where a₂ E R controls signal retention and a, controls noise added. Sohl-Dickstein et al. (2015); Ho et al. (2020) shows a special case of variance-preserving noising process where = / o.. Kingma et al. (2021): Hoogeboom et al. (2022) define signal to noise ratio as SN R(L) = " into account that the diffusion process is Markov, the entire noising process is then written as:
o7 Taking
T
q (Zo, .1,..., 2T x) (zo x) II
q (zz Zt-1).
The generative process diffusion model is defined with respect to the true denoising process, and the variable is approximated using neural network. Following the findings of Ho et al. (2020); Kingma et al. (2021), the variational lower bound on the log-likelihood term of x given by
I 11
L. E N(
SNR(I /)/ SNR(()) &
where we predict a gaussian noise and use the parameterization 2; = a,æ + G.E, then the neural network á outputs (2,,t).
22

---

Published as a conference paper at ICLR 2024
Köhler et al. (2020) showed that invariant distributions composed with an equivariant invertible function result in an invariant distribution. Additionally, Xu et al. (2022) proved that for x ~ p(æ) that is invariant to group G, and the transition probabilities of a Markov chain defined p(y x) are equivariant, then the marginal distribution of y at any time step is invariant to group transformations. In the case of the denoising diffusion model, need an invariant distribution and the neural network parameterizing the denoising diffusing process to be equivariant. This results in the marginal distribution of the denoising model to be an invariant distribution. For detailed proof, we refer to Xu et al. (2022).
In order to define the diffusion process, we take 0, VI- of as per Sohl-Dickstein et al. (2015) and let at (1 (1-2s) ff)-f(t)+ + where f(t) = (1- (t/T))), such that values decrease monotonically, starting 2 1 and Oer R 0, similar to EDMs for fair comparison. To avoid numerical instabilities, during sampling, we follow clipping procedure similar to Dhariwal & Nichol (2021) and compute Orje-1 = 01/0x-1, where a = 1 and values of dfj--1 are clipped from below by ,001 similar to Hoogeboom et al. (2022).
D.2.2 EXPERIMENTAL DETAILS
Training settings. All hyperparameters were set the same as in the original EDM implementation
Hoogeboom al. (2022). That is, we used = 9 layers, C' 256 hidden features. The models were trained for 1000 epochs (deviationg from the 3000 the EDM baseline does), with batch size
of 64. The polynomial degree of the basis was set to 3. We used the Adam optimizerKingni & Ba (2014), with learning rate of 5e 4, and without any learning rate scheduler. We used N 20 grid points on the sphere.
The embedding layer is linear embedding of the one-hot encoded atom numbers, followed by the Sca Paalareto ToSppere module. The readout was a single linear layer, applied to the output of the last ConvNext block, and produced single channel velocity signal the sphere. This signal was converted to a vector via the SphereToVec module.
D.3 N-BODY EXPERIMENT
Training settings. We trained PONITA with == layers, C= 128 hidden features. The poly- nomial degree was set to 3. The models were trained for 10, 000 epochs, with batch size of 100. We used the Adam optimizerKingma & Ba (2014), with a learning rate of 5e - 4. and with a Coealing learning rate schedule with warmup period of 100 epochs. We used N = 20 grid points on the sphere.
The embedding layer took two vector features inputs initial velocity and direction vector point- ing from the particle position to the average position of all particles. These two vectors were em- bedded as two scalar fields over the sphere via the VecToSphere module. The embedding also took two scalar features as input: charge (-1 or +1) and the norm of the initial velocity. Both were embedded on the sphere via the ScalarToSphere module. The 4 features were mapped vector of size C' (the hidden dimension size). The readout layer was applied after every ConvNext block, and conisted of Linear layer followed by the SphereToVe module. The final predicted velocity was the average of the velocity predictions after each layer.
In addition to the invariant geometric attributes, the basis functions also took the product the charges of the sending and receiving node as input.
E ADDITIONAL EXPERIMENTS
E SUPER-PIXEE MNIST
To verify the impact of lifting point clouds to position orientation space in 2D, we benchmark PONITA on superpixel mnist. Superpixel mnist is dataset of 2D point clouds of MNIST digits, in which each of the original images is segmented into 75 superpixels (clusters of similar pixels). We use the data loader provided by the library torch-geomest Fey & Lenssen (2019). We utilize the exact same architecture as in all the other experiments presented so far, except that now we utilize the three 2D invariants given in (8).
23

---

Published as a conference paper at ICLR 2024
We used the following hyper parameters. We trained for 50 epochs with batch size of 96, cosine learning rate decay, and initial learning rate 5e - 4. For PONITA we used = 10 grid points to sample the circle.
We compare our results with PNITA and PONITA in Table against the recent state-of-the- art method of Finzi et al. (2020b) called Probabilistic Numeric Convolutional Neural Networks (PNCNN), as well as classic graph NN approaches MONET (Monti et al., 2017), SplineCNN (Fey al., 2018), GCCP (Walker & Glocker, 2019), and GAT (Avelar al., 2020). Our results are aver- aged over runs with different seeds and are reported with the standard deviation in test error over these runs.
Apart from showing that with PONITA state-of-the-an results are obtained with the simple group convolutional architecture PONITA, the results in Table 4 also show the significant benifit of uti- lizing group convolutions (PONTTA) over planar convolutions with isotropic kernels (PNITA) ). The results further show that a simple fully convolutional architecture sufficient obtain outstanding results.
Table 4:
Classification Error on the 75-SuperPixel MNIST problem.
Method
Error rate
MONET
8.89
SplineCNN
4.78
GCGP
4.2
GAT
3.81
PNCNN
1.24
PNITA
3.04
PONITA
1.17+.11
Table 5: QM9 regression results.
Target
Unit
DimeNet++
PNITA
PONITA (point cloud)
PONITA (fiber bundle)
D ag
0.0286
0.0207
0.0115
0.0121
0.0469
0.0602
0.0446
0.0375
EHOMO
meV
27.8
26.1
18.6
16.0
ELUM
meV
19.7
21.9
15.3
14.5
Ac
meV
34.8
43.4
33.5
30.4
(R²)
oB
0.331
0.149
0.227
0.235
ZPVE
meV
1.29
1.53
1.29
1.29
Uo
meV
8.02
10.71
9.20
8.31
U
meV
7.89
10.63
9.00
8.67
H
meV
8.11
11.00
8.54
8.04
G
meV
0.0249
0.0112
0.0095
0.00863
Co
molk
cal
0.0249
0.0307
0.0250
0.0242
E.2
QM9 REGRESSION
To further demonstrate the potency of PONITA on 3D point cloud tasks we validate it on the QM9 benchmark (Ramakrishnan R., 2014) on the task of predicting molecular properties from input atomic point clouds and their covalent bonds.
Since QM9 provides a connectivity derived from the covalent bonds, we have an option to treat the molecules as point clouds in position-orientaii
space. Namely, we can treat every edge, which connects neighboring point X; to central node x;, as a point in opsiiiinn-orinnattin space R3 x S by treating taking starting point of the edge as position and the normalized direction vector as orientation. That is, for every edge e = (j, E È we have identify a point x = (pr, P Pap e
R3 22, The covalent bonds thus naturally generate point-cloud in position-orientationn space.
24

---

Published as a conference paper at ICLR 2024
Since this point cloud no-longer has fiber bundle structure (see App. A) we longer separate the convolution over positions and orientations separately and so we do the Nomain-miixngg in one single convolution, followed by channel mixing linear layer. So now the convolution is separated over 2 steps, instead of 3 as for the other experiments, as described Section 3.2. Otherwise, the PONITA architecture is held exactly the same. In this experiment either the atomic point cloud as input (PNITA) ) versus using the covalent-bond cloud (edges) as input (PONITA), The latter is close to the approach of DimeNet and GemNet (Gasteiger et al., 2019; 2021) which also perform message passing operations between edges. We therefore compare our approach against DimeNet++ (Gasteiger al., 2020).
compare convolutions using
We did hyperparmeter search for C' and L for PNITA and PONITA separately and found that for PNITA best results were obtained with C 128 and T. 5 and for PONITA with 256 and = 9. We further trained for 1000 epochs with a batch size of 96, with cosine learning rate decay scheduler and initial learning rate of 5e For all atoms we constructed fully connected graph. For reference, we also trained fiber bundle PONITA with C' 256, L 9, and N 16.
The results are presented Table
25


Reviews:
[{"review_id": "5aFyXNG8PY", "reviewer": "Reviewer_3W6S", "summary": "When considering a homogeneous space $X\\simeq G/H$ and features that are functions $f:X\\to \\mathbb{R}^C$, equivariant linear layers are convolutions (Equation 1). The authors propose to extend such layers to the case where $X$ is replaced by a discrete geometry (a graph of points of $X$), which allows for equivariant graph convolutions. They remark that the direct generalization (Equation 5) is not intrinsic as it depends on a represent of $x\\simeq [g_x]$; they characterize filters that are intrinsic and give applications to quotients of SE(3) and SE(2).", "questions": "None", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is very well written, the problem well posed and meaningful, and the solutions natural. I find the paper very interesting.", "weaknesses": "I don't see any", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oPI5z69k11", "reviewer": "Reviewer_MJgQ", "summary": "This paper studies the message passing scheme for signals define on space $\\mathbb{R}^3\\times S^2$ with SE(3) symmetry, which then can be used to construct neural architectures for pointcloud data, by lifting the pointcloud from $\\mathbb{R}^3$ to $\\mathbb{R}^3\\times S^2$.   In particular, they give the explicit form of an invariant and bijective edge embedding $a_{i,j}$ for two points $x_i,x_j\\in \\mathbb{R}^3\\times S^2$. This edge embedding serves to weight the pairwise message passing.  \n\nThey argue that the proposed method is more efficient than expressive $SE(3)$ group convolution (which requires integral on $SE(3)$) and meanwhile more expressive than convolution on $\\mathbb{R}^3$ (e.g., SchNet). Experiments on potential prediction, molecule generation as well as trajectory prediction show the better performance and efficiency of the method.", "questions": "- There are two P$\\Theta$NITA columns in Table 1. Should be a typo?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- the analysis on invariant feature $a_{i,j}$ seems technically sound and solid to me\n- the experimental part is extensive", "weaknesses": "- As there are many approaches for lifting a pointcloud, it may not be clear why the space chosen may find the best trade-off between efficiency and expressiveness. \n- So I think it could make the work stand out if we can find some applications that require to directly deal with signals on $\\mathbb{R}^3\\times S^2$. \n- Also just a minor personal idea: I think the word \"weight sharing\" makes me think of constructing linear transformation invariant/equivariant to certain symmetry. While here we are actually seeking for the full invariants describing the equivalence class.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6DFoLjoO3I", "reviewer": "Reviewer_7BQC", "summary": "This work derives point-pair attributes uniquely identifying the equivalence class of the point-pair and uses these attributes for equivariant group convolutional network. Its application on 3D point clouds achieves state-of-the-art performance and scalability.", "questions": "1. Does this work have provable expressivity? For example, can it universally approximate continuous equivariant function?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Conceptually simple, powerful, and efficient model verified with extensive experiments.\n2. A detailed introduction of previous work.", "weaknesses": "1. There are so many contents from previous works that I can hardly tell the novelty. From my point of view, only Theorem 1 is new, but it is still quite easy to prove.\n\n2. Though this work is titled with $SE(n)$, it only discusses $n=3$ case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)