Complete Neural Networks for Complete Euclidean Graphs
Neural networks for point clouds, which respect their natural invariance to
permutation and rigid motion, have enjoyed recent success in modeling geometric
phenomena, from molecular dynamics to recommender systems. Yet, to date, no
model with polynomial complexity is known to be complete, that is, able to
distinguish between any pair of non-isomorphic point clouds. We fill this
theoretical gap by showing that point clouds can be completely determined, up
to permutation and rigid motion, by applying the 3-WL graph isomorphism test to
the point cloud's centralized Gram matrix. Moreover, we formulate an Euclidean
variant of the 2-WL test and show that it is also sufficient to achieve
completeness. We then show how our complete Euclidean WL tests can be simulated
by an Euclidean graph neural network of moderate size and demonstrate their
separation capability on highly symmetrical point clouds.
Snir Hordan, Tal Amir, Steven J. Gortler, Nadav Dym
# Complete Neural Networks for Euclidean Graphs

###### Abstract

We propose a \(2\)-WL-like geometric graph isomorphism test and prove it is complete when applied to Euclidean Graphs in \(\mathbb{R}^{3}\). We then use recent results on multiset embeddings to devise an efficient geometric GNN model with equivalent separation power. We verify empirically that our GNN model is able to separate particularly challenging synthetic examples, and demonstrate its usefulness for a chemical property prediction problem.

Equivariant machine-learning models are models that respect data symmetries. Notable examples include Convolutional Neural Networks, which respect translation symmetries of imgaes, and Graph Neural Networks (GNNs), which respect the symmetry of graphs to permutations of their vertices.

In this paper we focus on equivariant networks for point clouds (which are often also called Euclidean or geometric graphs). Point clouds are sets of \(n\) points in \(\mathbb{R}^{d}\), whose symmetries include permutation of the \(n\) points, as well as translation, rotation and possibly also reflection. We denote the group of permutations by \(S_{n}\), the group of translations and rotations by \(SE(d)\), and the group obtained by also including reflections by \(E(d)\). Our interest is thus in functions on \(\mathbb{R}^{d\times n}\) that are invariant or equivariant to the action of \(E(d)\times S_{n}\) or \(SE(d)\times S_{n}\). In the past few years many works have focused on symmetry-preserving networks for point clouds in \(\mathbb{R}^{3}\), and on their applications for 3D computer vision and graphics (Deng et al., 2021), chemistry (Gasteiger et al., 2020) and physics simulation (Kondor, 2018). There are also applications for \(d>3\) for graph generation (Victor Garcia Satorras, 2021) and processing of Laplacian eigendecompositions (Lim et al., 2022).

The search for equivariant networks with good empirical performance is complemented by the theoretical study of these networks and their approximation power. These typically focus on two strongly related concepts: (i) _separation_ - the ability of a given invariant architecture to distinguish between two objects that are not related by a group symmetry, and (ii) _universality_ - the ability of the architecture to approximate any continuous equivariant function. These two concepts are intimately related and one typically implies the other, as discussed in Section 4 and in (Chen et al., 2019).

Recent results (Pozdnyakov and Ceriotti, 2022) show that distance-based message passing networks cannot separate point clouds better than a geometric variant of the 1-WL test (1-Geo), and that this test cannot separate all point clouds. On the other extreme, several works describe equivariant architectures that are universal, but these rely on high-dimensional representations of \(SO(d)\)(Dym and Maron, 2020; Finkelshtein et al., 2022; Gasteiger et al., 2021) or \(S_{n}\)(Lim et al., 2022). Thus, there is a large gap between the architectures needed for universality and those used in practice. For example, (Lim et al., 2022) requires hidden tensors of dimension \(n^{n}\) for universality, but uses tensors of dimension \(n^{2}\) in practice.

### Main results

In this paper, we make a substantial step towards closing this gap, by showing that complete separation can be obtained using efficient architectures of practical size.

We begin by expanding upon the notion of _geometric graph isomorphism_ tests, recently presented in (Pozdnyakov and Ceriotti, 2022; Anonymous, 2023). We show that while 1-Geo is not complete, it does separate _almost all_ distinct pairs. We then build on ideas from (Kurlin, 2022) to propose a _geometric \(2\)-WL_ test (2-Geo), which separates _any_ pair of 3D point clouds. Similarly, for general \(d\), we achieve separation using a geometric \(d-1\)-WL test. These results and some variations are discussed in Section 2.

In Section 3 we explain how to construct invariant architectures whose separation power is equivalent to 2-Geo (or 1-Geo), and thus are separating. This problem has been addressed successfully for graphs with discrete labels (see discussion in Section 7), but is more challenging for geometric graphs and other graphs with continuous labels. The main difficulty in this construction is the ability to construct
efficient continuous injective multiset-valued mappings. We will show that the complexity of computing standard injective multiset mappings is very high, and show how this complexity can be considerably improved using recent results from (Dym and Gortler, 2022). As a result we obtain \(\mathcal{SO}(3)\) and \(\mathcal{O}[3,n]\) separating invariant architectures with a computational complexity of \(O(n^{4}\log(n))\) and embedding dimension of \(6n+1\), which is approximately \(n^{2}\) times lower than what can be obtained with standard approaches. This advantage is even more pronounced when considering 'deep' 1-Geo tests (see Figure 2).

In Section 4 we use our separation results to prove the universality of our models, which is obtained by appropriate postprocessing steps to our separating architectures.

To empirically validate our findings, we present a dataset of point-cloud pairs that are difficult to separate, based on examples from (Pozdnyakov and Ceriotti, 2022; Pozdnyakov et al., 2020) and new challenging examples we construct. We verify that our architectures can separate all of these examples, and evaluate the performance of some competing architectures as well. We also show that our architectures achieve improved results on a benchmark chemical property regression task, in comparison to similar non-universal architectures. These results are described in Section 5.

## 1 Mathematical notation

A (finite) _multiset_\(\{\!\!\{y_{1},\ldots,y_{N}\}\!\!\}\) is an unordered collection of elements where repetitions are allowed.

Let \(\mathcal{G}\) be a group acting on a set \(\mathcal{X}\) and \(f:\mathcal{X}\to\mathcal{Y}\) a function. We say that \(f\) is _invariant_ if \(f(gx)=f(x)\) for all \(x\in X,g\in G\), and we say that \(f\) is _equivariant_ if \(\mathcal{Y}\) is also endowed with some action of \(G\) and \(f(gx)=gf(x)\) for all \(x\in\mathcal{X},g\in\mathcal{G}\).

A separating invariant mapping is an invariant mapping that is injective, up to group equivalence. Formally, we denote \(X\underset{\mathcal{G}}{=}Y\) if \(X\) and \(Y\) are related by a group transformation from \(\mathcal{G}\), and we define

**Definition 1.1** (Separating Invariant).: Let \(\mathcal{G}\) be a group acting on a set \(\mathcal{X}\). We say \(F:\mathcal{X}\to\mathbb{R}^{K}\) is a _\(\mathcal{G}\)-separating invariant_ if for all \(X,Y\in\mathcal{X}\),

1. **(Invariance)**\(X\underset{\mathcal{G}}{=}Y\Rightarrow F(X)=F(Y)\)
2. **(Separation)**\(F(X)=F(Y)\Rightarrow X\underset{\mathcal{G}}{=}Y\).

We call \(K\) the _embedding dimension_ of \(F\).

We focus on the case where \(\mathcal{X}\) is some Euclidean domain and require the separating mapping to be continuous and differentiable almost everywhere, so that it can be incorporated in deep learning models -- which typically require this type of regularity for gradient-descent-based learning.

The natural symmetry group of point clouds \((x_{1},\ldots,x_{n})\in\mathbb{R}^{d\times n}\) is generated by a translation vector \(t\in\mathbb{R}^{d}\), a rotation matrix \(R\in\mathcal{SO}(d)\), and a permutation \(\sigma\in S_{n}\). These act on a point cloud by

\[(R,t,\sigma)_{*}(x_{1},\ldots,x_{n})=(Rx_{\sigma^{-1}(1)}{+}t,\ldots,Rx_{ \sigma^{-1}(n)}{+}t).\]

We denote this group by \(\mathcal{SO}[d,n]\). In some instances, reflections \(R\in\mathcal{O}(d)\) are also permitted, leading to a slightly larger symmetry group, which we denote by \(\mathcal{O}[d,n]\).

For simplicity of notation, throughout this paper we focus on the case \(d=3\). In Appendix D we explain how our constructions and theorems can be generalized to \(d>3\).

## 2 Geometric Graph isomorphism tests

In this section we discuss geometric graph isomorphism tests, namely, tests for checking whether two given point clouds \(X,Y\in\mathbb{R}^{3\times n}\) are related by a permutation, rotation and translation (and possibly also reflection). Given two point clouds \(X,Y\), these tests typically compute some feature \(F(X),F(Y)\) and check whether \(F(X)=F(Y)\). This feature is \(\mathcal{G}\)-invariant, with \(\mathcal{G}\) denoting our symmetry group of choice, so that \(F(X)\neq F(Y)\) automatically implies that \(X\underset{\mathcal{G}}{\neq}Y\). Ideally, we would like to have _complete_ tests, meaning that \(X\underset{\mathcal{G}}{\neq}Y\) implies that \(F(X)\neq F(Y)\). Typically, these require more computational resources than _incomplete tests_.

### Incomplete geometric graph isomorphism test

Perhaps the most well-known graph isomorphism test is 1-WL. Based on this test, (Pozdnyakov and Ceriotti, 2022) formulated the following test, which we refer to as the 1-Geo test: Given two point clouds \(X=(x_{1},\ldots,x_{n})\) and \(Y=(y_{1},\ldots,y_{n})\) in \(\mathbb{R}^{3\times n}\), this test iteratively computes for each point \(x_{i}\) an \(\mathcal{O}(3)\)-invariant feature \(h_{i}^{t}\) via

\[h_{i}^{t}=\mathbf{Embed}^{(t)}\left(h_{i}^{t-1},\{\!\!\{h_{j}^{t-1},\|x_{i}-x_ {j}\|\}\,,j\neq i\!\!\}\right), \tag{1}\]

using an arbitrary initialization \(h_{i}^{t}\) for \(t=0\). This process is repeated \(T\) times, and then a final global feature for the point cloud is computed via

\[F^{\text{1-Geo}}(X)=\mathbf{Embed}^{(T+1)}\{\!\!\{h_{i}^{T}\mid i=1,\ldots,n \!\!\}.\]

A similar computation is performed on \(Y\) to obtain \(F^{\text{1-Geo}}(Y)\). In this test \(\mathbf{Embed}^{(t)}\) are hash functions, namely, they are discrete mappings of multisets to vectors, defined such that they assign distinct values to the finite number of multisets encountered during the computation of \(F^{\text{1-Geo}}\) for \(X\) and \(Y\). Note that by this construction, these functions are defined differently for different pairs \(X,Y\).
The motivation in (Pozdnyakov & Ceriotti, 2022) in considering this test is that many distance-based symmetry-preserving networks for point clouds are in fact a realization of this test, though they use **Embed** functions that are continuous, defined globally on \(\mathbb{R}^{3\times n}\), and in general may assign the same value to different multisets. Consequently, the separation power of these architectures is at most that of \(F^{\text{1-Geo}}\) with discrete hash functions. We note that continuous multiset functions _can_ be used to construct architectures with the separation power of geometric isomorphism tests. This will be discussed in Section 3.

The separation power of \(F^{\text{1-Geo}}\) is closely linked to the notion of _geometric degree_: For a point cloud \(X=(x_{1},\ldots,x_{n})\), we define the geometric degree \(d(i,X)\) to be the multiset

\[d(i,X)=\{\hskip-1.422638pt\{\|x_{1}-x_{i}\|,\ldots,\|x_{n}-x_{i}\|\}\hskip-1.422638pt \},\quad i\in[d],\]

and the geometric degree histogram \(d_{H}(X)\) to be

\[d_{H}(X)=\{\hskip-1.422638pt\{d(1,X),\ldots,d(n,X)\}\hskip-1.422638pt\}.\]

It is not difficult to see that if \(d_{H}(X)\neq d_{H}(Y)\) then \(X\) and \(Y\) can be separated by \(F^{\text{1-Geo}}\) even with a single iteration \(T=1\). With \(T=2\), as we show in the following theorem, \(F^{\text{1-Geo}}\) can do even more, and separate \(X\) and \(Y\) even if \(d_{H}(X)=d_{H}(Y)\), provided that the values in their histograms are all distinct, namely that \(X\) and \(Y\) belong to

\[\mathbb{R}^{3\times n}_{distinct}=\{X\in\mathbb{R}^{3\times n}|\,d(i,X)\neq d (j,X)\;\;\forall i\neq j\}.\]

Figure 1 depicts the distance matrices of two point clouds \(A,B\) that belong to \(\mathbb{R}^{3\times n}_{distinct}\) while having the same degree histogram.

**Theorem 2.1**.: _Suppose that \(X,Y\in\mathbb{R}^{3\times n}_{distinct}\), and \(\textbf{Embed}^{(t)},t=1,2,3\), are multiset-to-vector functions that assign distinct values to the finite number of multi-sets encountered when computing \(F^{\text{1-Geo}}(X)\) and \(F^{\text{1-Geo}}(Y)\). Then \(F^{\text{1-Geo}}(X)=F^{\text{1-Geo}}(Y)\) if and only if \(X\underset{\mathcal{O}[3,n]}{=}Y\)._

While _almost any_ pair of point clouds \(X,Y\) (in the Lebesgue sense) belongs to \(\mathbb{R}^{3\times n}_{distinct}\), and thus can be separated by \(F^{\text{1-Geo}}\), this test is not complete in general. This was shown in (Pozdnyakov & Ceriotti, 2022), by providing an example of point clouds \(C,D\in\mathbb{R}^{3\times 6}\) (see Figure 1)) that cannot be distinguished by \(F^{\text{1-Geo}}\).

### \(\mathcal{SO}[3,n]\)-isomorphism test

We now describe a \(2\)-WL-like geometric graph isomorphism test for \(\mathcal{SO}[d,n]\), which we name 2-Geo. Unlike 1-Geo, this test is _complete_. It is inspired by a similar test described in (Kurlin, 2022). The relationship between this work and ours is discussed in Section 7.

Figure 1: Distance matrices (Left) and degree histograms \(d_{H}\) (Right) of three pairs of point clouds \((A,B)\), \((C,D)\), \((E,F)\). These pairs are hard to separate by distance-based methods, as they have the same degree histogram. Nonetheless, \((A,B)\) can be separated by two iterations of 1-Geo, since \(A,B\in\mathbb{R}^{3\times n}_{distinct}\). Each of \(C\), \(D\) is comprised of three pairs of points, each of which share the same degree. While it was shown in (Pozdnyakov & Ceriotti, 2022) that 1-Geo cannot separate \(C\) from \(D\), our 2-Geo can separate _any_ distinct pair of 3D point clouds. \(E\) and \(F\) are especially challenging 6-dimensional point clouds in which all points have the same geometric degree.
As a first step, we eliminate the translation symmetry by centering the point clouds \(X\) and \(Y\). The centering of \(X=(x_{1},\dots,x_{n})\) is the point cloud \((x_{1}^{c},\dots,x_{n}^{c})\) defined by \(x_{i}^{c}=x_{i}-\frac{1}{n}\sum_{j=1}^{n}x_{j}\). It is known (Dym & Gortler, 2022) that the original \(X\) and \(Y\) are related by a symmetry in \(\mathcal{SO}[3,n]\) if and only if the centralized point clouds are related by a rotation and permutation.

Now, let us make two simplifying assumptions that we shall later dispose of: (a) the first two points of \(X\) and \(Y\) are in correspondence- meaning that if \(X\) and \(Y\) can be aligned by an \(\mathcal{SO}[3,n]\) transformation, then the permutation component assigns \(x_{1}\) to \(y_{1}\) and \(x_{2}\) to \(y_{2}\), and (b) the first two points in each point cloud are linearly independent.

Under these assumptions, we can define bases for \(\mathbb{R}^{3}\) by \(x_{1},x_{2},x_{1}\times x_{2}\) and \(y_{1},y_{2},y_{1}\times y_{2}\). These two bases are related by a rotation if and only if the Gram matrices of these two bases are identical. If indeed they are, this still only implies that the first two points are related by a rotation. To check whether the remaining points are related by a rotation and permutation, it suffices to check that the unordered collection of inner products of the remaining points with the basis we defined are identical.

Formally, we define for \((i,j)=1,2\) or any other pair of indices \((i,j)\)

\[X_{[i,j]} =\big{[}x_{i}^{c},x_{j}^{c},x_{i}^{c}\times x_{j}^{c}\big{]}\in \mathbb{R}^{3\times 3} \tag{2}\] \[P_{[i,j,k]} =X_{[i,j]}^{T}x_{k}^{c}\] (3) \[G_{[i,j]}(X) =X_{[i,j]}^{T}X_{[i,j]}\] (4) \[h_{[i,j]}(X) =\textbf{Embed}^{(1)}\big{\{}\mathbb{R}P_{[i,j,k]}\;\mid\;k=3, \dots,n\big{\}}\] (5) \[m_{[i,j]}(X) =\big{(}G_{[i,j]}(X),h_{[i,j]}(X)\big{)} \tag{6}\]

and we define \(m_{[i,j]}(Y)\) a similar manner, where **Embed** is some multiset-valued hash functions. The above construction guarantees that if \(X,Y\) satisfy the simplifying assumptions (a)-(b), then \(X\) and \(Y\) are related by a symmetry in \(\mathcal{SO}[3,n]\) if and only if \(m_{[1,2]}(X)=m_{[1,2]}(Y)\).

Let us now remove the simplifying assumption (a)-(b). Since we no longer know the correspondence, instead of just considering \(m_{[1,2]}\), we consider the multiset of all possible \(m_{[i,j]}\) and define

\[F^{2\text{-Geo}}(X)=\textbf{Embed}^{(2)}\big{\{}m_{[i,j]}(X)\;\mid\;1\leq i \neq j\leq n\big{\}}. \tag{7}\]

We define \(F^{2\text{-Geo}}(Y)\) similarly. In the appendix we prove

**Theorem 2.2**.: _Let \(X,Y\in\mathbb{R}^{3\times n}\), and let \(\textbf{Embed}^{(1)},\textbf{Embed}^{(2)}\) be multiset-to-vector functions that assign distinct values to the finite number of multisets encountered when computing \(F^{2\text{-Geo}}(X)\) and \(F^{2\text{-Geo}}(Y)\). Then \(F^{2\text{-Geo}}(X)=F^{2\text{-Geo}}(Y)\) if and only if \(\underset{\mathcal{SO}[3,n]}{=}Y\)._

Proof ideaIf the centralized point cloud \(X^{c}\) has rank \(\geq 2\), there are some \(i,j\) such that \(x_{i}^{c},x_{j}^{c}\) are linearly independent. If \(F^{2\text{-Geo}}(X)=F^{2\text{-Geo}}(Y)\) then there are some \(s,t\) such that \(m_{[i,j]}(X)=m_{[s,t]}(Y)\), and the argument above then shows that \(\underset{\mathcal{SO}[3,n]}{=}Y\). The full proof (which does not assume any rank assumptions on \(X^{c}\)) is given in the appendix.

### \(\mathcal{O}[3,n]\)-isomorphism test

The 2-Geo test described above can be modified to the scenario where reflections are also considered symmetries of the point cloud, so we would like the test to distinguish point clouds up to \(\mathcal{O}[3,n]\) symmetries. A simple way to achieve this is to consider for each pair \(x_{i},x_{j}\) both orientations of the vector product

\[X_{[i,j]}^{pos} =X_{i,j}=[x_{i},x_{j},x_{i}\times x_{j}]\] \[X_{[i,j]}^{neg} =[x_{i},x_{j},-x_{i}\times x_{j}].\]

The details of this construction are given in Appendix B. Ultimately this leads to a complete \(\mathcal{O}[3,n]\) test with twice the time and space complexity of the 2-Geo test for \(\mathcal{SO}[3,n]\).

An interesting but less efficient alternative to the above test is to use the standard \(3\)-WL graph isomorphism test, with the initial label for each triplet of indices taken to be the Gram matrix corresponding to those indices. The details of this construction are described in Appendix B as well.

## 3 Separating Architectures

In the previous section we described incomplete and complete geometric graph isomorphism tests for \(\mathcal{SO}[3,n]\) and \(\mathcal{O}[3,n]\) symmetries. Our goal now is to build separating architectures based on these tests. Let us first focus on the 2-Geo test for \(\mathcal{SO}[3,n]\).

To construct a separating \(\mathcal{SO}[3,n]\)-invariant architecture based on our 2-Geo test, we first need to choose a realization for the multiset-to-vector functions \(\textbf{Embed}^{(1)}\) and \(\textbf{Embed}^{(2)}\). To this end, we use parametric functions \(\textbf{Embed}_{\alpha}:\mathbb{R}^{3\times(n-2)}\rightarrow\mathbb{R}^{K_{1}}\) and \(\textbf{Embed}_{\beta}:\mathbb{R}^{K_{1}\times(n^{2}-n)}\rightarrow\mathbb{R}^{ K_{2}}\) which are invariant to permutation of their second coordinate. Note that this also renders \(F^{2\text{-Geo}}\) parametric \(F^{2\text{-Geo}}_{\alpha,\beta}\). We will also want them to be continuous and piecewise differentiable with respect to \(X,\alpha\) and \(\beta\).

The main challenge is of course to guarantee that for some parameters \(\alpha,\beta\), the function \(F^{2\text{-Geo}}_{\alpha,\beta}\) is \(\mathcal{SO}[3,n]\) separating. A standard way (see (Anonymous, 2023; Maron et al., 2019)) to achieve this is to require that for some \(\alpha,\beta\), the functions \(\textbf{Embed}_{\alpha}\) and \(\textbf{Embed}_{\beta}\) are injective as functions of multisets, or equivalently that they are permutation-invariant and separating. Note that by Theorem 2.2, this requirement will certainly suffice to guarantee that \(F^{2\text{-Geo}}\) is an \(\mathcal{SO}[3,n]\)-separating invariant. Our next step is therefore
to choose permutation invariant and separating \(\mathbf{Embed}_{\alpha}\), \(\mathbf{Embed}_{\beta}\). Naturally, we will want to choose these so that the dimensions \(K_{1},K_{2}\) and the complexity of computing these mappings are as small as possible.

### \(S_{n}\) separating invariants

We now consider the problem of finding separating invariants for the action of \(S_{N}\) on \(\mathbb{R}^{D\times N}\). Let us begin with the scalar case \(D=1\). Two well-known separating invariant mappings in this setting are the power-sum polynomials \(\Psi_{pow}\) and the sort mapping \(\Psi_{sort}\), defined by

\[\Psi_{sort}(s_{1},\dots,s_{N}) =\mathrm{sort}(s_{1},\dots,s_{N})\] \[\Psi_{pow}(s_{1},\dots,s_{N}) =\left(\sum_{j=1}^{N}s_{j},\sum_{j=1}^{N}s_{j}^{2},\dots,\sum_{j= 1}^{N}s_{j}^{N}\right).\]

It is clear that \(\mathrm{sort}\) is permutation-invariant and separating. It is also continuous and piecewise linear, and thus meets the regularity conditions we set out for separating invariant mappings. The power-sum polynomials are clearly smooth. Their separation can be obtained from the separation of the elementary symmetric polynomials, as discussed e.g., in (Zaheer et al., 2017).

We now turn to the case \(D>1\), which is our case of interest. One natural idea is to use lexicographical sorting. However, for \(D>1\), this sorting is not continuous. The power-sum polynomials can be generalized to multi-dimensional input, and these were used in the invariant learning literature (Maron et al., 2019). However, a key disadvantage is that to achieve separation, they require an extremely high embedding dimension \(K=\binom{N+D}{D}\).

A more efficient approach was recently proposed in (Dym and Gortler, 2022). This method initially applies linear projections to obtain \(N\) scalars and then applies a continuous \(1\times N\)-separating mapping \(\Psi=\Psi_{pow}\) or \(\Psi=\Psi_{sort}\), namely, one-dimensional power-sum polynomials or sorting. In more detail, for some natural \(K\), the function \(\mathbf{Embed}_{\theta}:\mathbb{R}^{D\times N}\rightarrow\mathbb{R}^{K}\) is determined by a vector \(\theta=(a_{1},\dots,a_{K},b_{1},\dots,b_{K})\in\mathbb{R}^{K(D+N)}\) where each \(a_{i}\) and \(b_{i}\) are \(D\)- and \(N\)-dimensional respectively, and

\[\mathbf{Embed}_{\theta}(X)=\langle b_{j},\ \Psi\left(a_{j}^{T}X\right) \rangle,\ j=1,\dots,K. \tag{8}\]

The following theorem shows that this mapping is permutation invariant and separating.

**Theorem 3.1** ((Dym and Gortler, 2022)).: _Let \(\mathcal{X}\) be an \(S_{N}\)-invariant semi-algebraic subset of \(\mathbb{R}^{D\times N}\) of dimension \(D_{\mathcal{X}}\). Denote \(K=2D_{\mathcal{X}}+1\). Then for Lebesgue almost every \(\theta\in\mathbb{R}^{K(D+N)}\) the mapping \(\mathbf{Embed}_{\theta}:\mathcal{X}\rightarrow\mathbb{R}^{K}\) is \(S_{N}\) invariant and separating._

When choosing \(\mathcal{X}=\mathbb{R}^{D\times N}\) we get that \(D_{\mathcal{X}}=N\cdot D\). The embedding dimension of \(\mathbf{Embed}_{\theta}\) would then be \(2N\cdot D+1\). This already is a significant improvement over the cardinality of the power-sum polynomials. Another important point that we will soon use is that if \(\mathcal{X}\) is a strict subset of \(\mathbb{R}^{D\times N}\) the number of separators will depend linearly on the intrinsic dimension \(D_{\mathcal{X}}\), and not on the ambient dimension \(ND\).

To conclude this subsection, we note that sort-based permutation invariants such as those we obtain when choosing \(\Psi=\Psi_{sort}\) are common in the invariant learning literature (Zhang et al., 2018, 2019). In contrast, polynomial-based choices such as \(\Psi=\Psi_{pow}\) are not so popular. However, this choice does provide us with the following corollary.

**Corollary 3.2**.: _Under the assumptions of the previous theorem, there exists a smooth parametric function \(q_{\theta}:\mathbb{R}^{3}\rightarrow\mathbb{R}^{2D_{\mathcal{X}}+1}\) such that the separating permutation-invariant mapping \(\mathbf{Embed}_{\theta}:\mathbb{R}^{3\times N}\rightarrow\mathbb{R}^{2D_{ \mathcal{X}}+1}\) defined using \(\Psi=\Psi_{pow}\) is given by_

\[\mathbf{Embed}_{\theta}(x_{1},\dots,x_{N})=\sum_{i=1}^{N}q_{\theta}(x_{i})\in \mathbb{R}^{2D_{\mathcal{X}}+1}.\]

Accordingly, in all approximation results which we will get based on the embedding \(\Psi_{pow}\) we can approximate \(\mathbf{Embed}_{\theta}\) with a function of the form \(\sum_{i=1}^{n}\mathcal{N}(x_{i})\) where \(\mathcal{N}\) is a neural networks whose input and output dimension are the same as those of \(\mathbf{Embed}_{\theta}\).

### Dimensionality of separation

From the discussion above we see that we can choose \(\mathbf{Embed}_{\alpha}\) to be a separating invariant mapping on \(\mathbb{R}^{3\times(n-2)}\), with an embedding dimension of \(K_{1}=6n-11\). It would then seem natural to choose the embedding dimension of \(\mathbf{Embed}_{\beta}\) so that separation on \(\mathbb{R}^{K_{1}\times(n^{2}-n)}\) is guaranteed. This would require a rather high embedding dimension of \(\sim n^{3}\). We note that any permutation invariant and separating mapping on \(\mathbb{R}^{D\times N}\) with reasonable regularity will have embedding dimension of at least \(N\cdot D\) (see (Anonymous, 2023)), and as a result we will always obtain an embedding dimension of \(n^{3}\) when requiring \(\mathbf{Embed}_{\alpha},\mathbf{Embed}_{\beta}\) to be separating on all of their (ambient) domain.

Significant savings can be obtained by the following observation: while the ambient dimension of the domain of \(\mathbf{Embed}_{\beta}\) is large, we only need injectivity for multisets in this domain which were obtained from some point cloud in \(\mathbb{R}^{3\times n}\). Accordingly, (once \(\alpha\) is fixed) \(\mathbf{Embed}_{\beta}\) needs only to be injective on a subset \(\mathcal{X}\) of the domain whose intrinsic dimension is at most \(3n\). Using Theorem 3.1 (see details in the proof of Theorem 3.4 stated below) we can take the embedding dimension of \(\mathbf{Embed}_{\beta}\) to be \(K_{2}=6n+1\). This idea is visualized in Figure 2(a).

We note that the advantage of the intrinsic separation technique presented here is even more pronounced when considering the implementation of \(F^{1\text{-Geo}}\) with \(T\) large. If we
require the mappings \(\textbf{Embed}^{(t)}\) to be separating and invariant on their (ambient) domain, the embedding dimension \(K_{t}\) of the \(t\)-th mapping is at least \(n+1\) times larger than the previous embedding dimension \(K_{t-1}\), so that the final embedding dimension is roughly \(\sim n^{T+1}\). In contrast, since the intrinsic dimension at each step is \(3n\), we get a constant embedding dimension of \(\sim 6n\) for all \(t\), by using a variation of Theorem 3.1 for vector-multiset pairs. See Appendix E for a full explanation, and Figure 2(b) for an illustration.

### Separation by feed-forward Neural Network Architectures

To summarize our discussion, the \(\mathcal{SO}[3,n]\) geometric graph isomorphism discussed in Theorem 2.2 can be realized as a separating invariant architecture by replacing \(\textbf{Embed}^{(1)}\) and \(\textbf{Embed}^{(2)}\) with the parametric functions \(\textbf{Embed}_{\alpha}\) and \(\textbf{Embed}_{\beta}\) respectively as in (8), with embedding dimension \(K_{1}=6n-11\) and \(K_{2}=6n+1\) as discussed above, and \(\Psi=\Psi_{sort}\) or \(\Psi=\Psi_{pow}\). This leads to the architecture described in Algorithm 1.

``` Input:\(X=(x_{1},...,x_{n})\in\mathbb{R}^{3\times n}\) \(X_{[i,j]}\leftarrow\big{[}x_{i}^{c},x_{j}^{c},x_{i}^{c}\times x_{j}^{c}\big{]}\) \(h_{[i,j]}(X)\leftarrow\textbf{Embed}_{\alpha}\big{\{}\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!
### Universality for Rotation-Equivariant models

Separating invariants are useful not only for proving universality of _invariant_ models, but also for _equivariant_ models. In our context, we use a result from (Villar et al., 2021) which showed that a permutation invariant and equivariant functions can be written as combinations of general invariant functions and simple equivariant functions. Since this result requires \(\mathcal{O}(3)\) invariance, we use a modification of Algorithm 2 to \(\mathcal{O}(3)\) invariance which is described formally in the appendix in Algorithm 3. We then obtain the following equivariant universality result:

**Theorem 4.2** (Equivariant Universality).: _Let \(f:\mathbb{R}^{3\times n}\rightarrow\mathbb{R}^{3}\) be continuous, \(\mathcal{O}(3)\)-equivariant and translation and permutation invariant. Then for any compact \(M\subset\mathbb{R}^{3\times n}\) and any \(\varepsilon>0\), \(f\) can be approximated to \(\epsilon\)-accuracy uniformly on \(M\) by functions of the form_

\[\tilde{f}(X)=\sum_{k=1}^{n}\mathcal{N}(h_{k},h_{global})x_{k}^{c},\]

_where \(h_{k}(X),h_{global}(X)\) are the output of Algorithm 3 and \(\mathcal{N}\) is a fully connected neural network._

## 5 Experiments

### Separation experiment

To evaluate the separation power of different architectures, we constructed a dataset consisting of pairs of point clouds that are particularly difficult to separate. This dataset will be made available for public use.

Each pair of point clouds \(X_{1},X_{2}\) is used as a prototype, from which we generate data samples for a binary classification task. Samples are generated by randomly choosing one of \(X_{1},X_{2}\), applying a random rotation and permutation to it and adding noise. The task is to determine whether each sample originates from \(X_{1}\) or \(X_{2}\).

We used the following challenging pairs: (i) **Hard1-Hard3**: Three pairs of 3D point clouds from (Pozdnyakov et al., 2020). In each pair, both clouds have the same degree histogram, but are members of the set \(\mathbb{R}^{3\times n}_{distinct}\) -- which can be separated by 1-Geo according to Theorem 2.1. The distance matrices for one such pair are visualized in \(A,B\) of Figure 1. (ii) **Harder**: A pair of 3D point clouds from (Pozdnyakov and Ceriotti, 2022) that are not in \(\mathbb{R}^{3\times n}_{distinct}\), and provably cannot be separated by 1-Geo. These are \(C,D\) in Figure 1. (iii) **Cholesky dim=d**: Pairs \(X_{1},X_{2}\) of \(d\) points in \(\mathbb{R}^{d}\), with \(d=6,8,12\). All points in \(X_{1}\), \(X_{2}\) have the

\begin{table}
\begin{tabular}{l c c c c c c c} \hline \hline Point Clouds & GramNet & GeoEGNN & EGNN & LinearEGNN & MACE & TFN & DimeNet & GVPGNN \\ \hline Hard1[2] & 1.0 & 0.998 & 0.5 & 1.0 & 1.0 & 0.5 & 1.0 & 1.0 \\ Hard2 [2] & 1.0 & 0.97 & 0.5 & 1.0 & 1.0 & 0.5 & 1.0 & 1.0 \\ Hard3 [2] & 1.0 & 0.85 & 0.5 & 1.0 & 1.0 & 0.55 & 1.0 & 1.0 \\ Harder [1] & 1.0 & 0.899 & 0.5 & 0.5 & 1.0 & 0.5 & 1.0 & 1.0 \\ Cholesky dim=6 & 1.0 & Irrelevant & 0.5 & 0.5 & 1.0 & Irrelevant & Irrelevant & Irrelevant \\ Cholesky dim=8 & 1.0 & Irrelevant & 0.5 & 0.5 & 1.0 & Irrelevant & Irrelevant & Irrelevant \\ Cholesky dim=12 & N/A & Irrelevant & 0.5 & 0.5 & 0.5 & Irrelevant & Irrelevant & Irrelevant \\ \hline \hline \end{tabular}
\end{table}
Table 1: Results of our models on challenging point clouds [1](Pozdnyakov and Ceriotti, 2022)(Pozdnyakov et al., 2020) Fig. S4

Figure 2: The standard method for constructing architectures that are equivalent to isomorphism tests uses injective multiset mappings on the ambient domain. The embedding dimension of such mapping increases exponentially with depth. In contrast, Theorem 3.1 allows for injective multiset functions whose dimensionality is twice the intrinsic dimension of the features which is always \(3n\). The figure shows the implementation of these two approaches for the computation of \(F^{\text{2-Geo}}\) and \(F^{\text{1-Geo}}\).
same degree histogram. The point clouds \(E,F\) for \(d=6\) appear in Figure 1. Further details appear in Appendix A.

The results appear in Table 1. First, as our theory predicts, \(\mathbf{GramNet}\) achieves perfect separation on all examples, while in the \(0.1\) noise level we use \(\mathbf{GeoEGNN}\) to achieve good but not perfect separation.

Next, note that EGNN (Victor Garcia Satorras, 2021) fails for all examples. Surprisingly, replacing the neural networks in EGNN with simple linear functions (LinearEGNN) does yield successful separation of examples in \(\mathbb{R}^{3\times n}_{distinct}\), as predicted in Theorem 2.1.

Finally, note that Tensor Field Networks (Thomas et al., 2018) does not separate our examples well, while MACE (Batatia et al.,), DimeNet (Gasteiger et al., 2020) and GVPGNN (Jing et al.) do. Of these methods, only MACE is applicable to problems with \(d>3\), and it has failed to separate for the \(d=12\) case. **GramNet** was unable to run for \(d=12\) due to memory constraints.

### Invariant regression on QM9

We evaluated our architectures on the QM9 Dataset for molecular property prediction (Ramakrishnan et al., 2014). To implement \(\mathbf{GeoEGNN}\) we used the original implementation of **EGNN**(Victor Garcia Satorras, 2021) augmented by the addition of our \(h_{[i,j]}\) of Equation (5) as edge features. As shown in Table 2, this minor modification of EGNN typically leads to improved results. In contrast, despite its excellent separation properties \(\mathbf{GramNet}\) is not competitive on this task.

## 6 Conclusion

We presented fully separating architectures whose embedding dimension depends linearly on the point cloud's dimension, in stark contrast to contemporary approaches with an exponential dependence. Our implementation of these architectures achieves good separation results in practice and yields improved results in several tasks on the QM9 benchmark. We believe these results will open the door to further improvements, both in the theoretical understanding of what is necessary for separation and in the development of separating architectures with good practical performance.

## 7 Related Work

WL equivalence.The relationship between the \(k\)-WL test and GNNs is very well-studied. Proofs of equivalence of GNNs and \(k\)-WL often assume a countable domain (Xu et al., 2018) or require separation only for a single graph (Morris et al., 2018). To the best of our knowledge, our separation result is the first one in which, for a fixed parameter vector (and in fact for almost all such vectors), separation is guaranteed for _all_ graphs with features in a _continuous_ domain. This type of separation could be obtained using the power-sum methodology in (Maron et al., 2019), but the complexity of this construction is exponentially worse than ours (see Subsections 3.1 and 3.2).

Complete Invariants and universality.As mentioned earlier, several works describe \(\mathcal{SO}[3,n]\) and \(\mathcal{O}[3,n]\) equivariant point-cloud architectures that are universal. However, these rely on high-dimensional representations of \(SO(d)\)(Dym and Maron, 2020; Finkelshtein et al., 2022; Gasteiger et al., 2021) or \(S_{n}\)(Lim et al., 2022).

In the planar case \(d=2\), universality using low-dimensional features was achieved in (Bokman et al., 2022). For general \(d\), a complete test similar to our 2-Geo was proposed in (Kurlin, 2022). However, it uses Gram-Schmidt orthogonalization, which leads to discontinuities at point clouds with linearly-dependent points. Moreover, the complete invariant features defined there are not vectors, but rather sets of sets. As a result, measuring invariant distances for \(d=3\) requires \(O(n^{7.5}+n^{3.5}log^{3}(n))\) arithmetic operations, whereas using GramNet invariant features only requires \(O(n^{4}log(n))\) operations. Finally, we note that more efficient tests for equivalence of geometric graphs were suggested in (Brass and Knauer, 2000), but there does not seem to be a straightforward way to modify these constructions to efficiently compute a complete, continuous invariant feature.

Weaker notions of universalityWe showed that 1-Geo is complete on the subset \(\mathbb{R}^{3\times n}_{distinct}\). Similar results for a simpler algorithm, and with additional restrictions, were obtained in (Widdowson and Kurlin, 2022). Efficient separation/universality can also be obtained for point clouds with distinct principal axes (Puny et al., 2021; Kurlin, 2022), or when only considering permutation (Qi et al., 2017) or rigid (Wang et al., 2022) symmetries, rather than considering both symmetries simultaneously.

\begin{table}
\begin{tabular}{c c c c c c c c c c c c} \hline \hline Property & \(\alpha\) & \(\varepsilon_{HOMO}\) & \(H\) & \(\varepsilon_{LUMO}\) & \(\Delta\varepsilon\) & \(\mu\) & \(C_{\nu}\) & \(G\) & \(R^{2}\) & \(U\) & \(U_{0}\) & \(ZPVE\) \\ \([units]\) & \(bohr^{3}\) & meV & meV & meV & meV & D & cal/mol K & meV & \(bohr^{3}\) & meV & meV & meV \\ \hline
**EGNN** & 0.073 & 29 & 12 & **25** & 48 & **0.029** & **0.031** & 12 & 0.106 & 12 & 11 & **1.55** \\
**GeoEGNN** & **0.068** & **27.9** & **11.6** & 38.3 & **45.8** & 0.032 & **0.031** & **10.75** & **0.1004** & **11.5** & **10.5** & 1.61 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Results on QM9 Dataset
AcknowledgementsN.D. acknowledges the support of the Horev Fellowship.
