F RAME A VERAGING FOR I NVARIANT AND E QUIVARIANT N ETWORK D ESIGN
A BSTRACT
Many machine learning tasks involve learning functions that are known to be invariant or equivariant to certain symmetries of the input data.
However, it is often challenging to design neural network architectures that respect these symmetries while being expressive and computationally efﬁcient.
For example, Euclidean motion invariant/equivariant graph or point cloud neural networks.
We introduce Frame Averaging (FA), a general purpose and systematic framework for adapting known (backbone) architectures to become invariant or equivariant to new symmetry types.
Our framework builds on the well known group averaging operator that guarantees invariance or equivariance but is intractable.
In contrast, we observe that for many important classes of symmetries, this operator can be replaced with an averaging operator over a small subset of the group elements, called a frame.
We show that averaging over a frame guarantees exact invariance or equivariance while often being much simpler to compute than averaging over the entire group.
Furthermore, we prove that FA-based models have maximal expressive power in a broad setting and in general preserve the expressive power of their backbone architectures.
Using frame averaging, we propose a new class of universal Graph Neural Networks (GNNs), universal Euclidean motion invariant point cloud networks, and Euclidean motion invariant Message Passing (MP)
GNNs.
We demonstrate the practical effectiveness of FA on several applications including point cloud normal estimation, beyond 2 -WL graph separation, and n -body dynamics prediction, achieving state-of-the-art results in all of these benchmarks.
I NTRODUCTION
Many tasks in machine learning (ML) require learning functions that are invariant or equivariant with respect to symmetric transformations of the data.
For example, graph classiﬁcation is invariant to a permutation of its nodes, while node prediction tasks are equivariant to node permutations.
Consequently, it is important to design expressive neural network architectures that are by construction invariant or equivariant for scalable and efﬁcient learning.
This recipe has proven to be successful for many ML tasks including image classiﬁcation and segmentation (LeCun et al., 1998; Long et al., 2015), set and point-cloud learning (Zaheer et al., 2017;
Qi et al., 2017a), and graph learning (Kipf &
Welling, 2016;
Gilmer et al., 2017; Battaglia et al., 2018).
Nevertheless, for some important instances of symmetries, the design of invariant and/or equivariant networks is either illusive (Thomas et al., 2018; Dym & Maron, 2020), computationally expensive or lacking in expressivity (Xu et al., 2018a; Morris et al., 2019; Maron et al., 2019; Murphy et al., 2019).
In this paper, we propose a new general-purpose framework, called Frame Averaging (FA), that can systematically facilitate expressive invariant and equivariant networks with respect to a broad class of groups.
At the heart of our framework, we build on a basic fact that arbitrary functions φ ∶ V → R , Φ ∶ V → W , where V,W are some vector spaces, can be made invariant or equivariant by symmetrization , that is averaging over the group (Yarotsky, 2021; Murphy et al., 2018), i.e.,
[Equation]
where G = { g } denotes the group, ψ ∶ V → R is invariant and Ψ ∶ V → W is equivariant with respect to G .
Furthermore, since invariant and equivariant functions are ﬁxed under group averaging, i.e., ψ = φ for invariant φ and Ψ =
Φ for equivariant Φ , the above scheme often leads to universal (i.e., maximally expressive) models (Yarotsky, 2021).
However, the challenge with equation 1 is that when the cardinality of G is large (e.g., combinatorial groups such as permutations) or inﬁnite (e.g., continuous groups such as rotations), then exact averaging is intractable.
In such cases, we are forced to approximate the sum via heuristics or Monte Carlo (MC), thereby sacriﬁcing the exact invariance/equivariance property for computational efﬁciency, e.g., Murphy et al.
(2018; 2019) deﬁne heuristic averaging strategies for approximate permutation invariance in GNNs; similarly, Hu et al.
(2021) and Shuaibi et al.
(2021) use MC averaging for approximate rotation equivariance in GNNs.
A concurrent approach is to ﬁnd cases where computing the symmetrization operator can be done more efﬁciently (Sannai et al., 2021).
The key observation of the current paper is that the group average in equation 1 can be replaced with an average over a carefully selected subset F( X ) ⊂
G while retaining both exact invariance/equivariance and expressive power.
Therefore, if F can be chosen so that the cardinality ∣F( X )∣ is mostly small, averaging over F( X ) results in both expressive and efﬁcient invariant/equivariant model.
We call the set-valued function F ∶ V → 2 G , a frame , and show that it can successfully replace full group averaging if it satisﬁes a set equivariance property .
We name this framework Frame Averaging (FA) and it serves as the basis for the design of invariant/equivariant networks in this paper.
We instantiate the FA framework by considering different choices of symmetry groups G , their actions on data spaces V,W (manifested by choices of group representations), and the backbone architectures (or part thereof) φ, Φ we want to make invariant/equivariant to G .
We consider: (i) MultiLayer Perceptrons
(MLP), and Graph Neural Networks (GNNs) with node identiﬁcation (Murphy et al., 2019; Loukas, 2020) adapted to permutation invariant Graph Neural Networks (GNNs); (ii)
Message-Passing GNN (Gilmer et al., 2017) adapted to be invariant/equivariant to Euclidean motions,
E ( d ) ; (iii) Set network, DeepSets and PointNet (Zaheer et al., 2017; Qi et al., 2017a) adapted to be equivariant or locally equivariant to E ( d ) ; (iv) Point cloud network, DGCNN (Wang et al., 2018), adapted to be equivariant to E ( d ) .
Theoretically, we prove that the FA framework maintains the expressive power of its original backbone architecture which leads to some interesting corollaries: First, (i) results in invariant universal graph learning models; (ii) is an E ( d ) invariant/equivariant GNN that maintain the power of message passing (Xu et al., 2018a; Morris et al., 2019); and (iii), (iv) furnish a universal permutation and
E ( d ) invariant/equivariant models.
We note that both the construction and the proofs are arguably considerably simpler than the existing alternative constructions and proofs for this type of symmetry (Thomas et al., 2018; Fuchs et al., 2020; Dym & Maron, 2020).
We experimented with FA on different tasks involving symmetries including: point-cloud normal estimation, beyond 2 -Weisfeiler-Lehman graph separation, and n -body dynamics predictions, reaching state of the art performance in all.
F RAME A VERAGING
In this section we introduce the FA approach using a generic formulation; in the next section we instantiate FA to different problems of interest.
2. F RAME AVERAGING FOR FUNCTION SYMMETRIZATION
Let φ ∶ V → R and Φ ∶ V → W be some arbitrary functions, where V,W are normed linear spaces with norms ∥⋅∥ V , ∥⋅∥ W , respectively.
For example, φ, Φ can be thought of as neural networks.
We consider a group G = { g } that describes some symmetry we want to incorporate into φ, Φ .
The way the symmetries g ∈ G are applied to vectors in V,W is described by the group’s representations ρ 1 ∶ G → GL ( V ) , and ρ 2 ∶ G → GL ( W ) , where GL ( V ) is the space of invertible linear maps
V → V (automorphisms).
A representation ρ i preserves the group structure by satisfying ρ i ( gh ) =
ρ i ( g ) ρ i ( h ) for all g,h ∈ G (see e.g., Fulton & Harris (2013)).
As customary, we will sometimes refer to the linear spaces V,W as representations.
Our goal is to make φ into an invariant function, namely satisfy φ ( ρ 1 ( g ) X )
=
φ ( X ) , for all g ∈ G and X ∈ V ; and Φ into an equivariant function, namely Φ ( ρ 1 ( g ) X )
=
ρ 2 ( g ) Φ ( X ) , for all g ∈ G and
X ∈ V .
We will do that by averaging over group elements, but instead of averaging over the entire group every time (as in equation 1) we will average on a subset of the group elements called a frame .
Deﬁnition 1.
A frame is deﬁned as a set valued function F ∶ V → 2 G ∖ ∅ .
A frame is G - equivariant if F( ρ 1 ( g ) X ) =
g F( X ) , ∀ X ∈ V, g ∈ G , where as usual, g F( X ) = { gh ∣ h ∈ F( X )} , and the equality should be understood as equality of sets. A frame
is bounded over a domain K ⊂ V if there exists a constant c > {+
+} so that ∥ ρ 2 ( g )∥ op ≤ c , for all g ∈ F( X ) and all X ∈ K , where ∥⋅∥ op denotes the induced operator norm over W .
Figure 1 provides an illustration.
How are equivariant frames useful?
Consider a scenario where an equivariant frame is easy to compute, and furthermore its cardinality, ∣F( X )∣ , is not too large.
Then averaging over the frame, denoted ⟨⋅⟩ F and deﬁned by
[Equation]
provides the required function symmetrization.
In Appendix A.1 we prove:
Theorem 1 (Frame Averaging) .
Let F be a G equivariant frame, and φ ∶ V → R , Φ ∶ V → W some functions.
Then, ⟨ φ ⟩ F is G invariant, while ⟨ Φ ⟩ F is G equivariant.
Several comments are in order: First, the invariant case (equation 2) is a particular case of the equivariant case (equation 3) under the choice of W = R and the trivial representation
ρ 2 ( g ) ≡ 1 .
Second, in this paper we only consider X and frame choices F for which F( X ) are ﬁnite sets.
Nevertheless, treating the inﬁnite case is an important future research direction.
Third, a trivial choice of an equivariant frame is F( X ) ≡ G , that is, taking the frame to be the entire group for all X ∈ V (for inﬁnite but compact G the sum in the FA in this case can be replaced with Harr integral).
This choice can be readily checked to be equivariant, and turns the FA equations 2, 3 into standard group averaging operators, equation 1.
The problem with this choice, however, is that it often results in an intractable or challenging computation, e.g., when the group is large or inﬁnite.
In contrast, as we show below, in some useful cases one can compute a manageable size frame and can use it to build invariant or equivariant operators in a principled way.
Let us provide a simple example for Frame
Averaging: consider V = R n , W = R , and G = R with addition as the group action.
We choose the group actions 1 in this case to be ρ 1 ( a ) x
= x + a 1 , and ρ 2 ( a ) b = b + a , where a,b ∈ R , x ∈ R n , and 1 ∈ R n is the vector of all ones.
We can deﬁne the frame in this case using the averaging operator
F( x ) = { 1 n 1 T x } ⊂ G = R .
Note that in this case the frame contains only one element from the group, in other cases ﬁnding such a small frame is hard or even impossible.
One can check that this frame is equivariant per Deﬁnition 1.
The FA of φ ∶ R n → R would be ⟨ φ ⟩ F ( x ) =
φ ( x − 1 n ( 1 T x ) 1 ) in the invariant case, and ⟨ φ ⟩ F ( x ) =
φ ( x − 1 n ( 1 T x ) 1 ) + 1 n 1 T x in the equivariant case.
Incorporating G as a second symmetry.
An important use case of frame averaging is with the backbones φ, Φ already invariant/equivariant w.r.t.
some symmetry group H and our goal is to make it invariant/equivariant to H × G .
For example, say we want to add G = E ( 3 ) equivariance to permutation invariant set or graph functions, i.e., H = S n .
We will provide sufﬁcient conditions for the
FA to provide this desired invariance/equivariance.
First, let us assume H is acting on V and W by the representations τ 1 ∶ H → GL ( V ) and τ 2 ∶ H → GL ( W ) , respectively.
Assume φ is H invariant and Φ is H equivariant.
We say that representations ρ 1 and τ 1 commute if ρ 1 ( g ) τ 1 ( h ) X
= τ 1 ( h ) ρ 1 ( g ) X for all g ∈ G , h ∈ H , and X ∈ V .
If ρ 1 and τ 1 commute then the map γ 1 ∶ H × G → GL ( V ) deﬁned by γ 1 ( h,g ) =
τ 1 ( h ) ρ 1 ( g ) is a representation of the group H × G .
Second, we would need that the frame F(
X ) is invariant to H , that is F( τ 1 ( h ) X )
=
F( X ) .
We show a generalization of Theorem 1:
Theorem 2 (Frame Average second symmetry) .
Assume F is H -invariant and G -equivariant.
Then,
If φ ∶ V → R is H invariant and ρ 1 ,τ 1 commute then ⟨ φ ⟩ F is G × H invariant.
If Φ ∶ V → W is H equivariant and ρ i ,τ i , i = 1 , 2 , commmute then ⟨ Φ ⟩ F is G × H equivariant.
Right actions.
Above we are of equivariance, e.g., if one of the actions is right.
For example, if g multiplies F( X ) from the right, then equivariance will take the form:
[Equation]
and accordingly
[Equation]
are G invariant and equivariant, respectively.
Efﬁcient calculation of invariant frame averaging.
There could be instances of the FA framework (indeed we discuss such a case later) where ∣F( X )∣ is still too large to evaluate equations 2,3.
In the invariant case, there is a more efﬁcient form of FA, that can potentially be applied.
To show it, let us start by deﬁning the subgroup of symmetries of X , i.e., its stabilizer.
The stabilizer of an element X ∈ V is a subgroup of G deﬁned by G X = { g ∈ G
∣
ρ 1 ( g ) X = X } .
G X naturally induces an equivalence relation ∼ on F( X ) , with g ∼ h ⇐⇒ hg − 1 ∈ G X .
The equivalence classes (orbits) are [ g ] = { h ∈ F( X )∣ g ∼ h } =
G X g ⊂ F( X ) , for g ∈ F( X ) , and the quotient set is denoted F( X )/ G X .
Theorem 3.
Equivariant frame F( X ) is a disjoint union of equal size orbits, [ g ] ∈ F( X )/ G X .
The proof is in A.3.
The ﬁrst immediate consequence of Theorem 3 is that the cardinality of F( X ) is at-least that of the stabilizer (intuitively, the inner-symmetries) of X , namely ∣F( X )∣ ≥ ∣ G X ∣ .
Therefore, there could be cases, such as when X describes a symmetric graph, where ∣F( X )∣ could be too large to average over.
A remedy comes from the following observation: for every h ∈ [ g ] , we have that
h
= rg , r ∈ G X , and φ ( ρ 1 ( h ) − 1 X )
=
φ ( ρ 1 ( g ) − 1 ρ 1 ( r ) − 1 X ) =
φ ( ρ 1 ( g ) − 1 X ) , since also r − 1 ∈ G X .
Therefore the summands in equations 2 are constant over orbits, and we get
[Equation]
where m F = ∣F( X )/ G X ∣ = ∣F( X )∣/∣ G X ∣ .
This representation of invariant FA requires only m F = ∣F( X )∣/∣ G X ∣ evaluations, compared to ∣F( X )∣ in the original FA in equation 2.
Approximation of invariant frame averaging.
Unfortunately, enumerating F( X )/
G X could be challenging in some cases.
Nevertheless, equation 6 is still very useful: it turns out we can easily draw a random element from F( X )/ G X with uniform probability.
This is an immediate application of the equal orbit size in Theorem 3:
Corollary 1.
Let F( X ) be an equivariant frame, and g ∈ F( X ) be a uniform random sample.
Then [ g ] ∈ F( X )/
G X is also uniform.
Therefore, an efﬁcient approximation strategy is averaging over uniform samples, g i ∈ F( X ) , i ∈
[ k ] ,
[Equation]
[Equation]
This approximation is especially useful, compared to the full-blown FA, when m F = ∣F( X )∣/∣ G X ∣ is small, i.e., when ∣ G X ∣ is large, or X has many symmetries.
Intuitively, the smaller m F the better the approximation in equation 7.
A partial explanation to this phenomenon is given in Appendix A.4, while an empirical validation is provided in Section 5.2.
2. E XPRESSIVE POWER
Another beneﬁt in frame averaging as presented in equations 2 and 3 is that it preserves the expressive power of the base models, φ , Φ , as exaplined next.
Consider some hypothesis function space H = { Φ } ⊂ C( V,W ) , where C( V,W ) is the set of all continuous functions V → W .
As mentioned above, the case of scalar functions φ is a special case where W = R , and ρ 2 ( g ) ≡ 1 .
H can be seen as the collection of all functions represented by a certain class of neural networks, e.g., Multilayer
Perceptron (MLP), or DeepSets (Zaheer et al., 2017), or Message Passing Graph Neural Networks (Gilmer et al., 2017).
We denote by ⟨H⟩ the collection of functions Φ ∈ H after applying the frame averaging in equation 3, ⟨H⟩ = {⟨ Φ ⟩ F ∣ Φ ∈ H} .
We set some domain K ⊂ V over which we would like to test the approximation power of ⟨H⟩ .
To make sure that FA is well deﬁned over K we will assume it is frame-ﬁnite , i.e., for every X ∈ K ,
F( X ) is a ﬁnite set.
Next, we denote K F = { ρ 1 ( g ) − 1 X ∣ X ∈ K,g ∈ F( X )} ; intuitively, K F ⊂ V contains all the points sampled by the FA operator.
Lastly, to quantify approximation error over a set A ⊂ V let us use the maximum norm ∥ Φ ∥ A,W = max X ∈
A ∥ Φ ( X )∥ W .
We prove that an arbitrary equivariant function Ψ ∈ C( V,W ) approximable by a function from H over K F is approximable by an equivariant function from ⟨H⟩ (Proof details are found on Appendix A.5).
Theorem 4 (Expressive power of FA) .
If F is a bounded G -equivariant frame, deﬁned over a frame-ﬁnite domain K , then for an arbitrary equivariant function
Ψ ∈ C( V,W ) we have
[Equation]
Where c is the constant from Deﬁnition 1.
This theorem can be used to prove universality results if the backbone model is universal, even for non-compact groups (e.g., the Euclidean motion group).
Below we will use it to prove universality of frame averaged architectures for graphs, and point clouds with Euclidean motion symmetry.
M ODEL INSTANCES
We instantiate the FA framework by specifying: i)
The symmetry group G , representations ρ 1 , ρand the underlying frame F ; ii)
The backbone architecture for φ (invariant) or Φ (equivariant).
3. P OINT CLOUDS , E UCLIDEAN MOTIONS .
Symmetry.
We would like to incorporate Euclidean symmetry to existing permutation invariant/equivaraint point cloud networks.
The symmetry of interest is G = E ( d ) =
O ( d ) ⋉ T ( d ) , namely the group of Euclidean motions in R d deﬁned by rotations and reﬂections O ( d ) , and translations
T ( d ) .
We also discuss G = SE ( d ) =
SO ( d ) ⋉ T ( d ) , where SO ( d ) is the group of rotations in R d .
We deﬁne V =
R n × d , and the group representation 2 is ρ 1 ( g ) X =
XR T + 1 t T , where R ∈ O
( d ) or R ∈ SO ( d ) , and t ∈ R d denotes the translation.
W,ρ 2 are deﬁned similarly, unless translation invariance is desired in which case we use the representation ρ 2 ( g ) X = XR T .
Frame.
We deﬁne the frame F( X ) in this case based on Principle Component Analysis (PCA), as follows.
Let t = 1 n X T 1 ∈ R
d be the centroid of X , and C =
( X − 1 t T ) T ( X − 1 t T )
∈ R
d × d the covariance matrix computed after removing the centroid from X .
In the generic case the eigenvalues of C satisfy λ 1 < λ 2 < ⋯ <
λ d .
Let v 1 , v 2 ,..., v d be the unit length corresponding eigenvectors.
Then we deﬁne F( X )
=
([ α 1 v 1 ,...,α d v d ] ,t ) ∣ α i ∈ {− 1 , 1 }} ⊂ E ( d ) .
The size of this frame (when it is deﬁned) is 2 d which for typical dimensions d = 2 , 3 amounts to frames of size 4 , 8 , respectively.
For G = SE ( d ) , we restrict F( X ) to orthogonal, positive orientation matrices; generically there are 2 d − 1 such elements, which amounts to 2 , 4 elements for d = 2 , 3 , respectively.
Proposition 1.
F( X ) based on the covariance and centroid are E ( d ) equivariant and bounded.
This choice of frame is deﬁned for every X ∈ V for which the covariance matrix C has simple spectrum (i.e., non-repeating eigenvalues).
It is known that within symmetric matrices, those with repeating eigenvalues are of co-dimension 2 (see e.g., Breiding et al.
Therefore, F( X ) is deﬁned for almost all X except rare singular points.
Where it is deﬁned, F is continuous as a direct consequence of perturbation theory of eigenvalues and eigenvectors of normal matrices (see e.g.,
Theorem 8.1.12 in Golub & Van Loan (1996)).
However, when very close to repeating eigenvalues, small perturbation can lead to large frame change.
In Appendix B we present an empirical study of frame stability and likelihood of encountering repeating eigenvalues in practice.
Since we would like to incorporate E ( d ) symmetries to an already S n invariant/equivariant architectures, per Theorem 2, we will also need to show that the ρ 1 (and similarly ρ 2 ) commute with τ ∶ S n → GL ( V ) deﬁned by τ ( h ) X
=
P X , where P = P h , h ∈ S n , is the permutation representation.
That is P h ∈ R n × n is the permutation matrix representing h ∈ S n , that is, P ij = 1 if i = h ( j ) and 0 otherwise.
Indeed τ ( h ) ρ 1 ( g ) X
= P ( XR T + 1 t T ) = ρ 1 ( g ) τ ( h ) X .
Furthermore, note that F( τ ( h ) X ) =
F( X ) , therefore F is also S n invariant.
Backbone architectures.
We incorporate the FA framework with two existing popular point cloud network layers: i)
PointNet (Qi et al., 2017a); and ii) DGCNN (Wang et al., 2018).
We denote both architectures by Φ d,d ′ ∶ R n × d → R n × d ′ for the S n equivariant version of these models.
To simplify the discussion, we omit particular choices of layers and feature dimensions, Appendix C.1 provides the full details.
We experimented with two design choices: First, consider frame averaged Φ 3 , 3 , i.e., ⟨ Φ 3 , 3 ⟩ F , yielding a universal, E ( 3 ) equivariant versions for PointNet and DGCNN, dubbed FAPointNet and FA-DGCNN.
A more complex design choice, taking inspiration from deep architectures with multiple modules, is to compose blocks of FA networks.
For example, to build a local E ( 3 ) equivariant version of PointNet, denote also Υ d,d ′ ∶ R d → R d ′ an MLP.
We decompose the input point cloud to k -nn patches X i ∈
R k × 3 , i ∈ [ n ] , where each patch is sorted by distance.
Next, feed each patch into an equivariant FA Υ 3 k, 3 d , each with its own frame F i , resulting in E ( 3 ) equivariant features in R 3 d for every point, Y ∈ R n × 3 d ; then applying equivariant FA Φ 3 d, 3 providing output in R n × 3 .
That is, Ψ ( X ) = ⟨ Φ 3 d, 3 ⟩ F ([⟨ Υ 3 k, 3 d ⟩ F 1 ( X 1 ) ,..., ⟨ Υ 3 k, 3 d ⟩ F n ( X n )]) , where brackets denote concat in the ﬁrst dimension.
We name this construction FA-Local-PointNet.
Universality.
We use Theorem 4 to prove that using a universal set invariant/equivariant backbone φ, Φ , such as DeepSets or PointNet (see e.g., Zaheer et al.
Qi et al.
(2017a)
; Segol &
Lipman (2019)) leads to a universal model.
Let H be any such universal set-equivariant function collection.
That is, for arbitrary continuous set function Ψ we have (in the notation of Section 2.2) inf Φ ∈H ∥ Ψ − Φ ∥ Ω ,W = 0 for arbitrary compact sets Ω ⊂ V .
If K ⊂ V is some bounded domain, then the choice of frame F described above implies that K F is also bounded and therefore contained in some compact set Ω ⊂ V .
Consequently, Proposition 1 and Theorem 4 imply Corollary 2.
A similar results holds for SE ( d ) , which provides a similar expressiveness guarantee as the one from (Dym &
Maron, 2020) analyzing Tensor Field Networks (Thomas et al., 2018; Fuchs et al., 2020).
Corollary 2 (FA-DeepSets/PointNet are universal) .
Frame Averaging DeepSets/PointNet using the frame F deﬁned above results in a universal E ( d ) × S n invariant/equivariant model over bounded frame-ﬁnite sets K ⊂ V .
3. G RAPHS , PERMUTATIONS
Symmetry and frame.
Let G = S n , and V = R n × d × R n × n , where X =
( Y , A ) ∈ V represents a set of node features Y ∈ R n × d , and an adjacency matrix (or some edge attributes)
A ∈ R n × n ; we assume undirected graphs, meaning A
=
A T .
Let X ∈ V , the representation ρ 1 is deﬁned by ρ 1 ( g ) X
= ( P Y , P AP T ) , where
P =
P g is the permutation matrix representing g ∈ S n .
We deﬁne F(
X ) to contain all g ∈ S
n that sort the rows of the matrix S ( X ) in column lexicographic manner, i.e., P S ( X ) is lexicographically sorted; the matrix S is deﬁned as follows.
Let L = diag
( A 1 ) − A be the graph’s Laplacian.
For every eigenspace (traversed in increasing eigenvalue order), spanned by the orthogonal basis u 1 ,..., u k , we add the (equivariant) column diag (∑ ki = 1 u i u Ti ) to S (Fürer, 2010).
Hence, the number of columns of S equals to the number of unique eigenvalues of L .
Proposition 2.
F( X ) deﬁned by sorting of S ( X ) is S n -equivariant and bounded.
Backbone architectures.
In this case we perform only invariant tasks, for which we chose two universal backbone architectures for φ : (i) MLP applied to vec ( Y , A ) and (ii) GNN+ID
(Murphy et al., 2019; Loukas, 2020), denoting a GNN backbone equipped with node identiﬁers as node features.
We perform FA with φ according to the frame constructed in Proposition 2.
Note that Theoremimplies, ∣F( X )∣ ≥ ∣ Aut ( G )∣ , since the stabilizer G X is the automorphism group of the graph, i.e., g ∈ G X
=
Aut ( X ) iff ρ 1 ( g ) X
= X .
This means that for symmetric graphs equation 2 can prove too costly.
In this case, we use the approximate invariant FA, equation 7.
Universality.
Let us use Theorem 4 to prove FA-MLP and FA-GNN+ID are universal.
Again, it is enough to consider the equivariant case.
Let H ⊂ C
( V,W ) denote the collection of functions that can be represented by MLPs or GNN+ID.
Universality results in (Pinkus, 1999; Loukas, 2020; Puny et al., 2020) imply that for any continuous graph (i.e., S n equivariant) function Ψ , inf Φ ∈H ∥ Ψ − Φ ∥ Ω ,W =for any compact Ω ⊂ V .
Let K ⊂ V be some bounded domain, then since G is ﬁnite then K F is also bounded and is contained is some compact set.
Proposition 2 and Theorem 4 now imply:
Corollary 3 (FA-MLP and FA-GNN+ID is graph universal) .
Frame Averaging MLP/GNN+ID using the frame F above results in a universal S n equivariant graph model over bounded domains K ⊂ V .
3. G RAPHS , E UCLIDEAN MOTIONS .
Symmetry and frame.
We consider the group G = E ( d ) acting on graphs, i.e., V = R n × d × R n × n , where X =
( Y , A ) ∈ V represents a set of node and edge attributes, as described above.
The group representation is ρ 1 ( g ) X
= ρ 1 ( g )( Y , A ) =
( Y R T + 1 t T , A ) .
We deﬁne the frame F( X ) =
F(
Y ) using the node features as in the point cloud case, Section 3.1.
Therefore Propositionimplies F is equivariant and bounded.
Next, also in this case we would like to incorporate E ( d ) symmetries to an already S n invariant/equivariant graph neural network architectures; again per
Theorem 2, we will also need to show that the ρ 1 (and similarly ρ 2 ) commutes with τ ∶ S n → GL ( V ) deﬁned by τ ( h ) X
= ( P Y , P AP T ) , where
P =
P h is the permutation matrix of h ∈ S n .
Indeed τ ( h ) ρ 1 ( g ) X
= ( P ( Y R T + 1 t T ) , P AP T ) =
ρ 1 ( g ) τ ( h ) X .
Furthermore, note that as in the point cloud case F( τ ( h ) X )
=
F( X ) , therefore F is also S n invariant.
Backbone architecture.
The backbone architecture we chose for this instantiation is the Message
Passing GNN in Gilmer et al.
(2017), an S n equivariant model denoted Φ d,d ′ ∶ R n × d × R n × n →
R n × d ′ × R n × n .
In this case we constructed a model, as suggested above, by composing l equivariant layers Ψ ( X ) = ⟨ Φ ( l ) 3 d ′ , 3 ⟩ F ○ ... ⟨ Φ ( i ) 3 d ′ , 3 d ′ ⟩ F ... ○ ⟨ Φ ( 1 ) 6 , 3 d ′ ⟩ F
( X ) .
The input feature size is 6 since we use velocities in addition to initial position as input ( n -body problem).
We name this model FA-GNN.
P REVIOUS WORKS
Rotation invariant and equivariant point networks.
State of the art S n invariant networks, e.g., (Qi et al., 2017a;b; Atzmon et al., 2018; Li et al., 2018; Xu et al., 2018b; Wang et al., 2018) are not invariant/equivariant to rotations/reﬂections by construction (Chen et al., 2019).
Invariance to global or local rotations can be achieved by modifying the 3D convolution operator or modifying the input representation.
Relative angles and distances across points (Deng et al., 2018; Zhang et al., 2019) or angles and distances w.r.t.
normals (Gojcic et al., 2019) can be used for rotation invariance.
Other works use some local or global frames to achieve invariance to rotations and translations.
Xiao et al.
Yu et al.
; Deng et al.
(2018) also use PCA to deﬁne rotation invariance, and can be seen as instances of the FA framework.
We augment this line of work by introducing a more general framework that includes equivariance to rotation/reﬂection and translation, more general architectures and symmetries as well as theoretical analysis of the expressive power of such models.
Equivariance is a desirable property for 3D recognition, registration (Ao et al., 2021), and other domains such as complex physical systems (Kondor, 2018).
A popular line of work utilizes the theory of spherical harmonics to achieve equivariance (Worrall et al., 2017; Esteves et al., 2018;
Liu et al., 2018; Weiler et al., 2018; Cohen et al., 2018).
Notably, Tensor Field Networks (TFN),
SE ( 3 ) transformers, and Group-Equivariant Attention (Thomas et al., 2018;
Fuchs et al., 2020;
Romero & Cordonnier, 2021) achieve equivariance to both translation and rotation, i.e., SE ( 3 ) , and are maximally expressive (i.e., universal) as shown in Dym & Maron (2020).
These methods, however, are speciﬁcally adapted to SE ( 3 ) and require high order SO ( 3 ) representations as features.
Recently Deng et al.
(2021) propose a rotation equivariant network by introducing tensor features, linear layers that act equivariantly on them and equivariant non-linearities etc.
However, their architecture is not proved to be universal.
Discrete convolutions (Cohen et al., 2019; Li et al., 2019;
Worrall & Brostow, 2018) have also been used for achieving equivariance.
In particular, Chen et al.
(2021) propose point networks that are SE ( 3 ) equivariant and use separable discrete convolutions.
Lastly, (Finzi et al., 2020) construct equivariant layers using local group convolution, and extends beyond rotations to any Lie group.
Graph neural networks.
Message passing (MP) GNNs (Gilmer et al., 2017) are designed to be S n equivariant.
Kondor et al.
(2018) introduces a broader set of equivariant operators in MP-GNNs, while
Maron et al.
(2018) provides a full characterization of linear permutation invariant/equivariant GNN layers.
In a parallel approach, trying to avoid harming expressively due to restricted architectures (Xu et al., 2018a; Morris et al., 2019), other works suggested symmetrization of non invariant/equivariant backbones.
Ranging from eliminating all symmetries by a canonical ordering (Niepert et al., 2016) to averaging over the entire symmetry group (Murphy et al., 2019; 2018), which amount to the trivial frame F ≡ ρ ( G ) , with the symmetry group G = S n .
As we have also shown, this approach comes at a cost of high variance approximations hindering the learning process.
Our FA framework reduces the variance both by choosing a canonical ordering and addressing the fact that it may not be unique.
Recently, a body of work studies GNNs with invariance/equivariance to E ( 3 ) (or a similar group) to deal with symmetries in molecular data or dynamical systems .
Many SE ( 3 ) equivariant con- structions (Anderson et al., 2019; Fuchs et al., 2020; Batzner et al., 2021; Klicpera et al., 2021) extend TFN (Thomas et al., 2018) and inherit its expensive higher order feature representations in the intermediate layers.
Finally, a recent work by Satorras et al.
(2021) provides an efﬁcient message passing construction which is E ( d ) equivariance but is not shown to be universal, thus far.
E XPERIMENTS
We evaluate our FA framework on a few invariant/equivaraint point cloud and graph learning tasks: point cloud normal estimation ( O ( 3 ) equivariant and translation invariant); graph separation tasks ( S n invariant); and particles position estimation, i.e., the n -body problem ( E ( 3 ) equivariant).
5. P OINT C LOUDS : N ORMAL E STIMATION
Normal estimation is a core geometry processing task, where the goal is to estimate normal data from an input point cloud, X ∈ R n × 3 .
This task is O ( 3 ) equivariant and translation invariant.
To test the effect of rotated data on the different models we experimented with three different settings: i) I / Itraining and testing on the original data; ii) I / SO ( 3 ) - training on the original data and testing on randomly rotated data; and iii) SO ( 3 )/ SO ( 3 ) - training and testing with randomly rotated data.
We used the ABC dataset (Koch et al., 2019) that contains 3 collections (10k, 50k, and 100k models each) of Computer-Aided Design (CAD) models.
We follow the protocol of the benchmark suggested in Koch et al.
(2019), and quantitatively measure normal estimation quality via 1 − ( n T ˆ n ) 2 , with n being the ground truth normal and ˆ n the normal prediction.
We used the same random train/test splits from Koch et al.
For baselines, we chose the PointNet (Qi et al., 2017a) and DGCNN (Wang et al., 2018) architectures, which are popular permutation equivariant 3D point cloud architectures.
In addition, we also compare to VN-PointNet and VN-DGCNN (Deng et al., 2021), a recent state of the art approach for SO ( 3 ) equivariant network design.
We also tested our FA models, FA-PointNet and FA-DGCNN as described in Section 3.1.
In addition, we tested our local O ( 3 ) equivariant model, FA-Local-PointNet.
See Appendix C.1 for the further implementation details.
The results in
Table 1 showcase the advantages of our FA framework:
Incorporating FA to existing architectures is beneﬁcial in scenarios (ii-iii), outperforming augmentation by a large margin.
In contrast to VN models, FA models maintain (and in some cases even improve) the baseline estimation quality (i), attributed to the expressive power of the FA models. [Table] [Figure] [Table]
[Figure]
[Table]
5. G RAPHS : E XPRESSIVE P OWER
[Table]
Producing GNNs that are both expressive and computationally tractable is a long standing goal of the graph learning community.
In this experiment we test graph separation ( S n invariant task): the ability of models to separate and classify graphs, a basic trait for graph learning.
We use two datasets: GRAPH8c
(Balcilar et al., 2021) that consists of all non-isomorphic, connected 8 node graphs; and EXP (Abboud et al., 2021) that consists of 3-WL distinguishable graphs that are not 2-WL distinguishable.
There are two tasks: (i) count pairs of graphs not separated by a randomly initialized model in GRAPH8c and EXP; and (ii) learn to classify EXP to two classes.
We follow Balcilar et al.
(2021) experimental setup.
As baselines we use GCN Kipf & Welling (2016), GAT Veliˇckovi´c et al.
GIN Xu et al.
(2018a),
CHEBNET Tang et al.
PPGN Maron et al.
(2019), and GNNML3 Balcilar et al.
(2021), all of which are equivariant by construction.
We compare to our FA-MLP, and FA-GIN+ID, as described in Section 3.2 that provide tractable option for universal GNNs.
We also compare to the trivial (i.e., entire group) frame averaging F ≡ G , denoted GA-MLP and GA-GIN+ID as advocated in (Murphy et al., 2019).
Appendix C.2 provides full implementation details.
[Figure]
The two ﬁrst columns in Table 2 show the results of the separation task (i).
As expected both FA and GA provide perfect separation, however, are not perfectly invariant unless one computes the full averages on F and G , respectively.
Since for some graphs (with many symmetries) full averages are not feasible, we compare the invariance error (see appendix for the exact deﬁnition) of GA with that of the approximate FA, equation 7, for 50 randomly permuted inputs.
Figure 2-left shows the mean, std, and 90% percentile of invariance error for increasing sample size k .
Note that approximate FA is more invariant even with as little as k = 1 samples, which explains why in the classiﬁcation task (ii), in Table 2, right column, FA is able to learn while GA fails, when both are trained with sample size k = 1 .
Figure 2-right compares m F and m G (the maximal number of unique elements in the FA, equation 6) for
GRAPH8c dataset.
The color and size of the points in the plot represents how many graphs in the dataset have the corresponding ( m G ,m F ) values.
This demonstrates the beneﬁt (in one example) of
FA over GA in view of Theorem 5 and approximation in equation 7.
Of course, other, more powerful frames can be chosen, e.g., using higher order WL (Morris et al., 2019) or substructure counting (Bouritsas et al., 2020) that will further improve the approximation of equation 7.
5. G RAPHS : n -B ODY P ROBLEM
[Table]
In this task we learn to solve the n -body problem ( E ( 3 ) equivariant).
The dataset, created in (Satorras et al., 2021; Fuchs et al., 2020), consists of a collection of n = 5 particles systems, where each particle is equipped with its initial position in Rand velocity in R 3 , and each pair of particles (edge) with a value indicating their charge difference.
The task is to predict the particles’ locations after a ﬁxed time.
We follow the protocol of (Satorras et al., 2021).
As baselines we use EGNN (Satorras et al., 2021), GNN (Gilmer et al., 2017), TFN (Thomas et al., 2018), SE( 3 )-Transformer
(Fuchs et al., 2020) and Radial Field (Köhler et al., 2019).
We test our FA-GNN, as described in
Section 3.3.
Table 3 logs the results, where the metric reported is the Mean Squared Error between predicted locations and ground truth.
As can be seen, FA-GNN improves over the SOTA by more than 20% .
Note that the number of parameters used in FA-GNN and EGNN is roughly the same.
More details are provided in Appendix C.
C ONCLUSIONS
We present Frame Averaging, a generic and principled methodology for adapting existing (backbone) neural architectures to be invariant/equivariant to desired symmetries that appear in the data.
We prove the method preserves the expressive power of the backbone model, and is efﬁcient to compute in several cases of interest.
We use FA to build universal GNNs, universal point cloud networks that are invariant/equivariant to Euclidean motions E ( 3 ) , and GNNs invariant/equivariant to Euclidean motions E ( 3 ) .
We empirically validate the effectiveness of these models on several invariant/equivariant learning tasks.
We believe the instantiations presented in the paper are only the ﬁrst step in exploring the full potential of the FA framework, and there are many other symmetries and scenarios that can beneﬁt from FA.
For example, extending the invariance (or equivariance) of a model from a subgroup H of G , to G .
Further Interesting open questions, not answered by this paper are: What would be a way to systematically ﬁnd efﬁcient/small frames?
How the frame choice effects the learning process?
How to explore useful FA architectures and modules?
A CKNOWLEDGMENTS
OP, MA and HB were supported by the European Research Council (ERC Consolidator Grant, "LiftMatch" 771136) and also by a research grant from the Carolito Stiftung (WAIC)
A P ROOFS A. P ROOF OF T HEOREM
Proof.
First note that frame equivariance is deﬁned to be F( ρ 1 ( g ) X ) =
g F( X ) which in particular means ∣F( ρ 1 ( g ) X )∣ = ∣F( X )∣ .
For invariance, let g ′ ∈
G
[Equation]
and for equivariance
[Equation]
A. P ROOF OF T HEOREM
Proof.
The proof above A.1, of Theorem [Equation] is a special case of the following where the second symmetry is chosen to be trivial.
In principle the proofs are quite similar.
First note that frame equivariance and invariance, together with ρ 1 ,τ 1 commuting mean
[Equation]
which in particular implies that ∣F( γ 1 ( h,g ) X )∣ = ∣F( X )∣ .
Let ( h ′ ,
g ′ ) ∈
H × G be arbitrary.
Then,
[Equation]
meaning that ⟨ φ ⟩ F is H × G invariant.
Next,
[Equation]
showing that ⟨ Φ ⟩ F is H × G equivariant.
A. P ROOF OF T HEOREM 3 AND C OROLLARY
Proof.
(Theorem 3)
First we show G X acts on F( X ) .
For an arbitrary h ∈ G X and equivariant frame F , we have F( X ) =
F( ρ 1 ( h ) X )
=
h F( X ) , where in the ﬁrst equality we used the fact that h ∈ G X and in the second the equivariance of F .
This means that if g ∈ F( X ) and h ∈ G X then also hg ∈ F( X ) , or in other words F( X ) is closed to actions from G X .
Furthermore, since a group acts on itself via bijections we have that the cardinality of all orbits is the same, ∣[ g ]∣ = ∣ G X g ∣ =
∣ G X ∣ , for all [ g ] ∈ F( X )/ G X .
Lastly, the equivalence relation g ∼ h ⇐⇒ hg − 1 ∈ G X shows that F( X ) is a union of disjoint orbits of equal cardinality.
Proof.
(Corollary 1) Theorem 3 asserts that all orbits [ g ] ∈ F( X )/ G X have the same number of elements.
Therefore, a random choice g ∈ F( X ) will have equal probability to land in each orbit.
A. T HE ROLE OF m F IN APPROXIMATION QUALITY OF EQUATION
To better understand the role of m F in approximating equation 6 we ﬁrst note that equation 7 can be written as ⟪ φ ⟫ F ( X )
=
∑ [ g ]∈F( X )/ G X ˆ µ [ g ] φ ( ρ 1 ( g i ) − 1 X ) , where ˆ µ is the empiricial distribution over F( X )/ G X , assigning to each element [ g ] ∈ F( X )/
G X the fraction of samples landed in [ g ] , i.e., ˆ µ [ g ] =
∣{ i ∈
[ k ]∣ g i ∈
[ g ]}∣ k − 1 .
We next present a lower bound on the probability of any particular ˆ µ that provides a good approximation ⟪ φ ⟫ F ( X ) ≈ ⟨ φ ⟩ F ( X ) .
Theorem 5.
Let F be an equivariant frame.
The probability of an arbitrary ˆ µ ∈ { ˆ µ ∣ sup φ ∈Q ∣⟨ φ ⟩ F ( X ) − ⟪ φ ⟫ F ( X )∣ ≤ (cid:15) } is bounded from below as follows,
[Equation]
where Q = { φ ∈ C( V, R ) ∣ ∣ φ ( X )∣ ≤ 1 , ∀ X ∈ V } the set of bounded, continuous functions V
→ R .
Before providing the proof let us note that this theorem provides a lower bound for each particular "good" empirical distribution ˆ µ .
The main takeoff is that for ﬁxed k and (cid:15) , the smaller m F the better the lower bound.
The counter-intuitive behaviour of this bound w.r.t.
k and (cid:15) stems from the fact that the size of the set of "good" ˆ µ , namely { φ ∈ C( V, R ) ∣ ∣ φ ( X )∣ ≤ 1 , ∀ X ∈ V } is increasing with k and (cid:15) .
Proof.
For brevity we denote m = m F .
Our setting can be formulated as follows.
We have the uniform probability distribution, denoted µ , over the discrete space [ m ] (representing the quotient set
F( X )/ G X ); µ j = 1 m , j ∈ [ m ] , and we have numbers a j = φ ( ρ 1 ( h j ) − 1 X ) , where [ h j ] ∈ F( X )/ G X , j ∈ [ m ] represent exactly one sample per orbit.
In this notation ⟨ φ ⟩ F ( X ) =
∑ mj = 1 a j m , while the approximation in equation 7 takes the form ⟪ φ ⟫ F ( X ) =
∑ mj = 1 ˆ µ j a j , where ˆ µ j = k j / k , where k j is the number of samples g i ∈
[ h j ] , i ∈
[ k ] , i.e., samples that landed in the j -th orbit.
Therefore,
[Equation]
where the latter is the total variation norm for discrete measures (Levin & Peres, 2017).
We denote
H ( ˆ µ ∣ µ ) =
∑ mj = 1 log ( ˆ µ j ) log ( ˆ µ j µ j ) the KL-divergence.
The Pinsker inequality and its inverse for discrete positive measures are (see e.g., (Polyanskiy & Wu, 2014)):
[Equation]
where α = min j ∈[ m ] µ j = m − 1 .
Therefore,
[Equation]
Now, application of Large Deviation Theory (Lemma 2.1.9 in Dembo & Zeitouni (2010)) provides that for ˆ µ so that H ( ˆ µ ∣ µ ) ≤ (cid:15) 2 m 2 :
[Equation]
A. P ROOF OF T HEOREM
Proof.
Let Ψ ∈ C
( V,W ) be an arbitrary G equivariant function, F a bounded G equivariant frame over a frame-ﬁnite domain K .
Let c > 0 be the constant from Deﬁnition 1.
For arbitrary X ∈ K ,
[Equation]
where in the ﬁrst equality we used the fact that ⟨ Ψ ⟩ F = Ψ since Ψ is already equivariant.
A. P ROOF OF P ROPOSITION
Proof.
Let us prove F is equivariant (equation 3).
Consider a transformation g =
( R , t ) ∈ G , and let (
O , s ) ∈ F( X ) , then
[Equation]
where O Λ O T is the eigen decomposition of C .
Note that the group product of these transformations is
[Equation]
We need to show ( RO , Rs + t ) ∈ F (( R , t ) X ) .
Indeed, RO ∈ O ( 3 ) and consists of eigenvectors of C = RX T ( I − 1 n 11 T ) XR T as can be veriﬁed with a direct computation.
If O , R ∈ SO ( d ) then also RO ∈ SO
( d ) .
Furthermore
[Equation]
as required.
We have shown ( R , t )F( X )
⊂ F
(( R , t ) X ) for all X
and ( R , t ) .
To show the other inclusion let X =
( R , t ) − 1 Y and get F(( R , t ) − 1 Y )
⊂ ( R , t ) − 1 F( Y ) that also holds for all Y and ( R , t ) .
In particular ( R , t )F( X ) ⊃ F(( R , t ) X ) .
The frame F is bounded since for compact
K ⊂ R n × d , the translations n − 1 X T 1 are compact and therefore uniformly bounded for X ∈ K , and orthogonal matrices always satisfy ∥ R ∥ 2 = 1 .
A. P ROOF OF P ROPOSITION
Proof.
First note that by deﬁnition S ( X ) is equivariant in rows, namely S ( ρ 1 ( g ) X )
=
ρ 1 ( g ) S ( X ) .
Therefore if g ∈ F( X ) ⊂ S n , then by deﬁnition of the frame ρ 1 ( g ) S is sorted.
Therefore, ρ 1 ( g ) S =
ρ 1 ( gh ) ρ 1 ( h ) − 1 S ( X ) = ρ 1 ( gh ) S ( ρ 1 ( h ) − 1 X ) is sorted and we get that gh ∈ F( ρ 1 ( h ) − 1 X ) .
We proved F( X ) h ⊂ F( ρ 1 ( h ) − 1 X ) for all h ∈ G and all X ∈ V .
Taking X = ρ 1 ( h ) Y for an arbitrary
Y ∈ V , and h = g − 1 for arbitrary g ∈ G , we get F( ρ 1 ( g − 1 ) Y )
⊂ F( Y ) g , for all g ∈ G and Y ∈ V .
We proved F( X ) h
= F( ρ 1 ( h ) − 1 X ) , which amounts to right action equivariance, see equation 4.
The frame is bounded since ρ 1 ( G ) is a ﬁnite set.
B E MPIRICAL FRAME ANALYSIS
[Figure]
(Satorras et al., 2021)
Repeating eigenvalues.
We empirically test the likeliness of repeating eigenvalues in the covariance matrix used from the deﬁnition of frame in Section 3.1.
We use the data from the n -body dataset (Satorras et al., 2021).
Let X ∈ R n × 3 represent the set of particles locations (each set centered around 0 and scaled to have max x i ∈
X ∥ x i ∥ 2 = 1 ) and λ 1 ≤ λ 2 ≤ λ 3 be the eigenvalues, sorted in increasing order, of the covariance matrix
C =
( X − 1 t T ) T ( X − 1 t T ) .
In order to measure the proximity of eigenvalues across the dataset we use the notion of eigenvalues spacing s i
=
λ
i + 1 − λ i ¯ s , i = 1 , 2 , where ¯ s = λ
3 − λ 1is the mean spacing.
Furthermore we deﬁne s min = min { s 1 ,s 2 } as the minimal normalized spacing, a ratio that indicates how close the spectrum is to having repeating eigenvalues.
Figure 3 presents a histogram of the minimal spacing over the training set of the n -body dataset (consists of 3000 particles sets).
The minimal spacing encountered in this experiment is of order 10 − 2 .
This empirically justiﬁes the usage of ﬁnite frames for E ( d ) equivariance.
[Figure]
Frame stability.
Here we test stability of our O ( d ) frame deﬁned in Section 3.1.
By stability we mean the magnitude of change of a frame w.r.t.
the change in the input X .
A desired attribute of the constructed frame is to be stable, that is, to exhibit small changes if the input is perturbed.
We quantify this stability by comparing the distance between our frames to frames constructed by a noisy input.
As in the previous experiment used the data from the n -body dataset (Satorras et al., 2021).
Let X ∈ R n × 3 represent the set of particles locations (each set centered around 0 and scaled to have max x i ∈
X
∥ x i ∥ 2 = 1 ) and X σ = X + Z where z i,j ∼ N( 0 ,σ ) is the noisy input sample .
We compute F( X ) and F( X σ ) and choose representatives from each set g =
( R , t ) ∈ F( X ) ,g σ =
( R σ , t σ ) ∈ F( X σ ) .
We measure the distance between the frames as a function of the representatives -
[Equation]
Notice that in the case of simple spectrum the distance is invariant to the selection of representatives.
In Figure 4 we plot distance of original and noisy frames (and its standard deviation) as function of noise level σ .
The plot validates the continuity of frames in the simple spectrum case.
C I MPLEMENTATION D ETAILS C. P OINT C LOUDS : N ORMAL E STIMATION
Here we provide implementation details for the experiment in section 5.1.
PointNet architecture.
Our backbone PointNet is based on the object part segmentation network from Qi et al. (2017a).
The network consists of layers of the form
[Equation]
where X ∈ R n × d in , W ∈ R d in × d out , b ∈ R d out are the learnable parameters, 1 ∈ R n is the vector of all ones, [⋅] is the concatenation operator, e i is the standard basis in R d in , and ν is the ReLU activation.
In this experiement, for the PointNet baseline and for Φ 3 , 3 (the backbone in FA-PointNet), we used the following architecture:
[Equation]
Note that the original PointNet network also contains two T-Net networks, applied to the input and to L 3 (the output of the third layer).
Similarly, our baseline implementation made a use of the same T-Net networks.
Note that the T-Net networks were not part of our FA-PointNet backbone architectures Φ 3 , 3 .
The FA-Local-Pointnet architecture can be seen as a composition of two parts (see Figure 5).
The ﬁrst part outputs equivariant features by applying the same MLP backbone Υ 3 k, 3 d with FA on each point’s k-nn patch.
Then, from each point’s equivariant features, the second part of the network outputs an equivariant normal estimation.
Note that the second part, ⟨ Φ ⟩ F , is an FA-Pointnet network applied to a point cloud of dimensions R n × 3 d .
For FA-Local-PointNet we made the following design choices.
Each point patch is constructed as its k nearest neighbors, with k = 20 , X i ∈ R 20 × 3 .
Then, the backbone Υ 3 ∗ 20 , 3 ∗ 42 is applied to all patches, with the following PointNet layers:
[Equation]
The second backbone, Φ 3 ∗ 42 , 3 , is built from the following PointNet layers:
[Equation]
DGCNN architecture.
Our backbone DGCNN architecture, Φ 3 , 3 , is based on the object part segmentation network from Wang et al.
It consists of EdgeConv ( n,d in ,d out ) , FC ( n,d in ,d out ) , and MaxPool layers.
EdgeConv
[Table] [Table] →
[Table] [Table] FC [Table]
Note that the DGCNN architecture incorporates T-Net network applied to the input.
Training details.
We trained our networks using the A DAM (Kingma & Ba, 2014) optimizer, setting the batch size to 32 and 16 for PointNet and DGCNN respectively.
We set a ﬁxed learning rate of 0 . 001 .
All models were trained for 250 epochs.
Training was done on a single Nvidia VGPU, using PYTORCH deep learning framework (Paszke et al., 2019).
[Equation]
We provide implementation details for the experiments in 5.2.
We used two different universal backbones, MLP and GIN equipped with identiﬁers.
The details of those architectures are presented here.
FA/GA GIN+ID architecture.
The GIN+ID backbone is based on the GIN (Xu et al., 2018a) network with an addition of identiﬁers as node features in order to increase expressiveness.
For the experiments we used a three-layer GIN with a feature dimension of size 64 and a ReLU activation function.
For the added identiﬁers we deﬁned the input node features as X 0 =
[ X 0 , I n ] , where n is the number of nodes in the graph, and to handle graphs of different sizes (EXP dataset) we padded the node features with zeros to ﬁt the size of the maximal graph.
Note that we did not apply the permutation generated by the frame on the identiﬁers.
FA/GA MLP architecture.
We used two different MLP networks for the EXP dataset and the
GRAPH8c, due to the different size of graphs in the datasets.
In the EXP dataset the maximal graph size is 64 and every node in the graph has a one dimensional binary feature, therefore the input for the MLP network is a ﬂatten representation of the graph (with additional padding according to the graph size) x ∈ R 64 2 + 64 .
Our architecture consists of layers of the form
[Equation]
where W ∈ R d out × d in , b ∈ R d out are the learnable parameters.
The ﬁnal output of the network is denoted by ( d out )
where d out is the output dimension.
The MLP network structure for the EXP-classify task:
[Equation]
with ReLU as the activation function.
For the EXP task, which has 10 -dimensional output for each graph we omitted the last layer.
The GRAPH8c is composed of all the non-isomorphic connected graphs of size 8 , hence we did not used any padding of the input here.
The nodes have no features and we just used a ﬂatten version of the adjacency matrix.
The architecture for the GRAPH8c task:
[Equation]
Training details.
We followed the protocol from (Balcilar et al., 2021) and trained our model with batch size 100 for 200 epochs.
The learning rate was set to 0 .
001 and did not change during training.
For optimization we used the A DAM optimizer.
Training was done on a single Nvidia RTXGPU, using PYTORCH deep learning framework.
Invariance Evaluation.
We quantify the permutation invariance of a model φ by comparing output of randomly permuted graphs ρ 1 ( g i )
X , g i ∈ S n , i ∈
The evaluation metric used is the invariance error , deﬁned by
[Equation]
where v = 1 m ∑ mi = 1 φ ( ρ 1 ( g i ) X ) .
The permutation invariance error of the models (FA/GA) was measured as a function of the sample size k .
We iterated over the entire GRAPH8c dataset and for every graph computed the invariance error for the FA-MLP ,GA-MLP and regular MLP models (all with the exact same backbone network).
The results presented in Figure 2 (left) are normalized by the error of the regular MLP model.
The backbone MLP we used for this experiment is of the form:
[Equation]
This section describes the implementation details for the n -body experiment from Section 5.3.
FA-GNN architecture.
We use the GNN architecture of Gilmer et al.
(2017) our base GNN layer, where for each node i ∈ [ n ] the update rule follows,
[Equation]
where N( i ) are the indices of neighbors of vertex
i , h li is the embedding of node i at layer l and a ij ∈ {− 1 , 1 } × R + are edge attributes representing attraction or repelling between pairs of particles and their distance.
h 0 ∈ R n × 6 represents the nodes input features, which in this experiment are a concatenation of the nodes initial 3 D position (rotation and translation equivariant) and velocities (rotation-equivariant and translation-invariant).
We used node feature dimension of size 60 to maintain a fair comparison with the baselines (Satorras et al., 2021).
The functions φ e and φ h are implemented as a two-layer MLP with the SiLU activation function (also chosen for consistency purposes) with an hidden dimension of 121 and 120 respectively.
To maintain fair comparison with (Satorras et al., 2021) our network is composed of 4 GNN layers. Φ ( 1 ) 6 , 3 d ′ has an additional linear embedding of the features as a preﬁx to the GNN layers while Φ ( 4 ) 3 d ′ , 3 is equipped with a two-layer MLP ( SiLU activation) as a decoder to extract ﬁnal positions.
Training details.
We followed the protocol from (Satorras et al., 2021) and trained our model with batch size 100 for 10000 epochs.
The learning rate was set to 0 .
001 and did not changed during training.
For optimization we used the A DAM optimizer.
Training was done on a single Nvidia
RTX- 6000 GPU, using PYTORCH deep learning framework.