# Graph Neural Network with Local Frame for Molecular Potential Energy Surface

Xiyuan Wang1 Muhan Zhang1,2,
1Institute for Artificial Intelligence, Peking University
2Beijing Institute for General Artificial Intelligence
wangxiyuan<EMAIL_ADDRESS>Corresponding author: Muhan Zhang
(muhan@pku.edu.cn).

###### Abstract

Modeling molecular potential energy surface is of pivotal importance in
science. Graph Neural Networks have shown great success in this field,
especially those using rotation-equivariant representations. However, they
either suffer from a complex mathematical form or lack theoretical support and
design principle. To avoid using equivariant representations, we introduce a
novel local frame method to molecule representation learning and analyze its
expressive power. With a frame and the projection of equivariant vectors on
the frame, GNNs can map the local environment of an atom to a scalar
representation injectively. Messages can also be passed across local
environments with frames’ projection on frames. We further analyze when and
how we can build such local frames. We prove that local frames always exist
when the local environments have no symmetry, as is often the case in
molecular dynamics simulations. For symmetric molecules, though only
degenerate frames can be built, we find that the local frame method may still
achieve high expressive power in some common cases due to the reduced degrees
of freedom. Using only scalar representations allows us to adopt existing
simple and powerful GNN architectures. Our model outperforms a range of state-
of-the-art baselines in experiments. Simpler architectures also lead to higher
scalability. Our model only takes about $30\%$ inference time compared with
the fastest baseline.

##  1 Introduction

Prediction of molecular properties is widely used in fields such as material
searching, drug designing, and understanding chemical reactions (Schmitz et
al., 2019). Among properties, potential energy surface (PES) (Lewars, 2016),
the relationship between the energy of a molecule and its geometry, is of
pivotal importance as it can determine the dynamics of molecular systems and
many other properties. Many computational chemistry methods have been
developed for the prediction, but few of them can achieve both high precision
and scalability.

In recent years, machine learning (ML) methods, which are both accurate and
efficient, have emerged. Among these ML methods, Graph Neural Networks (GNNs)
are especially promising. In the past few years, they have improved
continuously (Schutt et al., 2017; Schütt et al., 2017; Klicpera et al., 2020;
Anderson et al., 2019; Schütt et al., 2021; Gasteiger et al., 2021; Batzner et
al., 2021; Thölke and Fabritiis, 2022) and achieved state-of-the-art
performance on many benchmark datasets. Compared with popular GNNs used in
other graph tasks (Kipf and Welling, 2017), these GNNs need a lot of special
designs, as molecules are more than a graph composed of merely nodes and
edges. Atoms are in the continuous three-dimensional space, and the prediction
targets like energy are sensitive to the position of atoms. Therefore, to
model molecules, GNNs must include geometric information. Moreover, for
generalization, symmetry of the predicted properties must be taken into
consideration. For example, if an isolated molecule rotates, the energy will
not change. Thus, the prediction of a model should not change either.

Various theories and models have been proposed to capture the target’s
symmetry. Equivariant representations are widely used in these approaches.
Among them, some models (Thomas et al., 2018; Batzner et al., 2021; Anderson
et al., 2019) use irreducible representations of the SO$(3)$ group and suffer
from complex mathematical form and low scalability. Other models (Schütt et
al., 2021; Thölke and Fabritiis, 2022) manually design equivariant functions
for vector and scalar representations and achieve outstanding performance.
However, to deal with two kinds of representations, functions for these two
kinds must be defined separately, and complex operation is needed to mix the
information contained in the two kinds of representations. Furthermore, the
design is often complicated and ad hoc due to the lack of design principle and
theoretical support, hampering these models’ scalability and
comprehensibility. In contrast, models using only invariant representations
(Schütt et al., 2017; Klicpera et al., 2020; Gasteiger et al., 2021) can use
the elegant architectures and theories of ordinary GNNs, though they are
usually less expressive than those using equivariant representations.
Therefore, we wonder whether equivariant representations can be replaced with
invariant representations while keeping the high expressive power.

In this work, we propose a simple, scalable, yet expressive GNN for molecular
property prediction. As shown in Figure 1, our model first produces a local
frame for each atom based on its local environment. Then it projects the
relative positions and frames of neighbor atoms on the frame to create a graph
with only scalar features. Finally, a normal GNN can work on the graph to
produce a prediction. Our model outperforms all baselines on a majority of the
targets in the MD17 and QM9 datasets. It also achieves three times higher
inference speed than the fastest baseline on the PES task.

Figure 1: An illustration of our model. A local frame is generated for each
atom, which is used to encode the geometric information of its neighborhood
into scalars. Then an ordinary GNN is applied.

##  2 Related work

We can classify existing ML models for PES into two classes: manual
descriptors and GNNs.

Manual Descriptor. These models first use manually designed functions with few
learnable parameters to convert one molecule to a descriptor vector and then
feed the vector into some ordinary ML models like kernel regression (Rupp et
al., 2012; Chmiela et al., 2019; Christensen et al., 2020) and neural networks
(Behler and Parrinello, 2007; Smith et al., 2017; Zhang et al., 2018) to
produce the prediction. Compared with GNNs, these methods are more scalable
and data-efficient. However, they are less accurate and cannot process
variable-size molecules or different kinds of atoms due to the hard-coded
descriptors. Our model outperforms one state-of-the-art model of them
(Christensen et al., 2020) on the MD17 dataset with the same data split while
using a more flexible GNN architecture.

GNN. These GNNs mainly differ in the way to incorporate geometric information.
Schutt et al. (2017) and Schütt et al. (2017) only consider the distance
between atoms. Klicpera et al. (2020) introduce angular features to describe
the direction of interaction, and Gasteiger et al. (2021) further use dihedral
angles. All these models use rotation-invariant geometric features only, so we
call them invariant models. Though it is theoretically expressive, GemNet
(Gasteiger et al., 2021) has high time complexity of up to $O(Nn^{3})$, where
$n$ is the number of atoms in the neighborhood of one atom and $N$ is the
number of atoms in the molecule. Recent works have also utilized rotation-
equivariant features. Some (Anderson et al., 2019; Thomas et al., 2018;
Batzner et al., 2021) are based on irreducible representations of the SO$(3)$
group. Though having certain theoretical guarantees, these models have complex
formulations in their full forms and are hard to implement. Their scalability
also suffers from special functions used for group representations. Other
works (Schütt et al., 2021; Thölke and Fabritiis, 2022) model equivariant
interactions in Cartesian space using both invariant representations and
vector representations, which are sequences of 3-vectors rotating with the
molecule. They are much simpler and achieve state-of-the-art performance with
low time complexity of only $O(Nn)$. However, these models have no theoretical
guarantees. Their architectures are also complex due to different sets of
functions that must be designed separately for different input and output
types. Our work avoids both irreducible representations and vector
representations. We introduce an O$(3)$-equivariant frame and project all
equivariant features on the frame injectively, thus deserting vector
representations while keeping the expressive power. The proposed method is
simple and scalable. Though it has the same time complexity $O(Nn)$ as
existing state-of-the-art models, our model only needs to process scalar and
has no computation overhead for transforming vector or mixing scalar and
vector.

Besides potential energy surface models, methods designed for other molecular
tasks and point cloud also use GNNs to capture geometric information. Some of
them (Du et al., 2022; Satorras et al., 2021; Luo et al., 2021; Huang et al.,
2022) also use the term "frame" or some designs similar to frame. However,
these methods are quite different from ours. The most notable difference is
that even using frame, existing models still keep vector features, while our
model completely deserts vector features and only operates on a vanilla scalar
graph after projection. This can greatly simplify the architecture without
losing expressive power. The comparison is detailed in Appendix I. Moreover,
these methods can be categorized into models using both invariant
representations and vector representations. Our model can beat the
representative methods (Schütt et al., 2021; Thölke and Fabritiis, 2022) of
this kind in both accuracy and scalability on potential energy surface
prediction tasks.

##  3 Preliminaries

Modeling PES. PES is the relationship between molecular energy and geometry.
Given a molecule with $N$ atoms, our model takes the kinds of atoms
$z\in{\mathbb{Z}}^{N}$ and the three-dimensional coordinates of atoms
$\vec{r}\in{\mathbb{R}}^{N\times 3}$ as input to predict the energy
$\hat{E}\in{\mathbb{R}}$ of this molecule. It can also predict the force on
atoms $\hat{\vec{F}}\in{\mathbb{R}}^{N\times 3}=-\nabla_{\vec{r}}\hat{E}$.
With PES, many properties can be calculated by molecular dynamics (MD)
simulations, which has been well studied in the computational chemistry field
and is not our focus. Nevertheless, our model can also predict other molecular
properties directly.

Equivariance. We define equivariant and invariant functions the same as in
(Villar et al., 2021).

######  Definition 3.1.

Given a function $h:{\mathbb{X}}\to{\mathbb{Y}}$ and a group $G$ acting on
${\mathbb{X}}$ and ${\mathbb{Y}}$ as $\star$. We say that $h$ is

$\displaystyle G\text{-invariant:}~{}~{}~{}~{}~{}~{}~{}~{}\text{if}~{}h(g\star
x)=h(x),~{}~{}\forall x\in{\mathbb{X}},g\in G$ (1) $\displaystyle
G\text{-equivariant:}~{}~{}~{}~{}\text{if}~{}h(g\star x)=g\star
h(x),~{}~{}\forall x\in{\mathbb{X}},g\in G$ (2)

The energy is invariant to the permutation of atoms, coordinates’
translations, and coordinates’ orthogonal transformations (rotations and
reflections). The architecture of GNN naturally keeps the permutation
invariance. As the relative position $\vec{r}_{ij}=\vec{r}_{i}-\vec{r}_{j}$,
which is invariant to translation, is usually used as the input to GNNs, the
translation invariance can also be ensured easily. So we focus on orthogonal
transformations. Orthogonal transformations of coordinates form the group
$\text{O}(3)=\\{Q\in{\mathbb{R}}^{3\times 3}~{}|~{}QQ^{T}=I\\}$, where $I$ is
the identity matrix. Representations are considered as functions of $z$ and
$\vec{r}$. Below we define scalar representations and vector representations,
respectively.

######  Definition 3.2.

Representation $s$ is called a scalar representation if
$s(z,\vec{r})=s(z,\vec{r}o^{T}),\forall
o\in\text{O}(3),z\in{\mathbb{Z}}^{N},\vec{r}\in{\mathbb{R}}^{N\times 3}$.
Representation $\vec{v}$ is called a vector representation if
$\vec{v}(z,\vec{r})o^{T}=\vec{v}(z,\vec{r}o^{T}),\forall
o\in\text{O}(3),z\in{\mathbb{Z}}^{N},\vec{r}\in{\mathbb{R}}^{N\times 3}$.

Scalar and vector representations are also called invariant and equivariant
representations in some previous work (Thölke and Fabritiis, 2022). Sometimes
we simply call them scalars and vectors. Note that the scalars in our
definition are not restricted to one number, and the vectors may not only be
of shape $(1,3)$. In Lemma 3.1, we introduce some basic operations of scalars
and vectors.

######  Lemma 3.1.

  * •

(Function of scalar) Any function of scalar $s$ will produce a scalar.

  * •

(Scale of vector) Let $s$ denote a scalar of shape $(F)$, $\vec{v}$ denote a
vector of shape $(F,3)$. We define $s\circ\vec{v}$ to output a matrix of shape
$(F,3)$, whose $(i,j)$th element is $s_{i}\vec{v}_{ij}$. When $\vec{v}$ is of
shape $(1,3)$, we first broadcast it along the first dimension. Then the
output is also a vector representation.

  * •

(Projection) Let $\vec{v}_{1}$, $\vec{v}_{2}$ denote two vectors of shape
$(F_{1},3)$ and $(F_{2},3)$, respectively. The projection of $\vec{v}_{1}$ to
$\vec{v}_{2}$, denoted as
$P(\vec{v}_{1},\vec{v}_{2}):=\vec{v}_{1}\vec{v}_{2}^{T}$, is a scalar of shape
$(F_{1},F_{2})$.

  * •

(Inner product) Let $\vec{v}_{1}$, $\vec{v}_{2}$ denote two vectors of shape
$(F,3)$. The inner product
$\langle\vec{v}_{1},\vec{v}_{2}\rangle:=\text{diag}(P(\vec{v}_{1},\vec{v}_{2}))$
is a scalar of shape $(F)$.

Local Environment. Most message passing frameworks for PES encode the local
environment of each atom in each message passing layer as defined in
Definition 3.3.

######  Definition 3.3.

Let $r_{ij}$ denote $||\vec{r}_{ij}||$. The local environment of atom $i$ is
$\text{LE}_{i}=\\{(s_{j},\vec{r}_{ij})|r_{ij}<r_{c}\\}$, the set of scalar
atom features $s_{j}$ and relative positions $\vec{r}_{ij}$ of atoms $j$
within the sphere centered at $i$ with cutoff distance $r_{c}$, where $r_{c}$
is usually a hyperparameter.

We just let the scalar feature $s_{j}$ be initialized with the kind of atom
$z_{j}$, but it is not restricted to that. In this work, orthogonal
transformation of a set/sequence means transforming each element in the
set/sequence, respectively. For example, an orthogonal transformation $o$ will
map $\\{(s_{j},\vec{r}_{ij})|r_{ij}<r_{c}\\}$ to
$\\{(s_{j},\vec{r}_{ij}o^{T})|r_{ij}<r_{c}\\}$.

##  4 Local frame

In general, our model projects all vector representations of neighbors on the
local frame of each atom to enable using ordinary GNNs on a graph with scalars
only. This section first illustrates how local frames capture geometric
information and then explains how to build local frames. With a frame, we can
easily convert geometric features into scalars without information loss and
achieve high expressive power with ordinary GNNs.

###  4.1 Rethinking vector representations

Though vector representations have been used for a long time, it is still
unclear how to ideally transform them. Existing methods (Schütt et al., 2021;
Thölke and Fabritiis, 2022; Deng et al., 2021; Villar et al., 2021) either
have no theoretical guarantee or tend to use too many parameters, resulting in
possible redundancy of computation. In this study, we ask a fundamental
question, can we totally give up vector representations during the message
passing? The motivation is that, given scalars only, we have mature ordinary
GNN architectures with guaranteed expressive power. It also eliminates the
complex conversions between scalars and vectors.

Assume there exist three linearly independent vector representations in
${\mathbb{R}}^{1\times 3}$ forming a frame $\vec{E}\in{\mathbb{R}}^{3\times
3}$, whose rows are equivariant to orthogonal transformations. We will discuss
how to generate the frame in Section 4.3. We define the projection function on
the frame $\vec{E}$ as follows:

$P_{\vec{E}}(\vec{x}):=P(\vec{x},\vec{E})=\vec{x}\vec{E}^{T}.$ (3)

Given the frame $\vec{E}$, the projection will contain all the information of
the input vector $\vec{x}$ as shown in Proposition 4.1.

######  Proposition 4.1.

$P_{\vec{E}}$ is a bijective function, which maps a vector shape $(F,3)$ to a
scalar of shape $(F,3)$. Its inverse function $P^{-1}_{\vec{E}}$ maps a scalar
of shape $(F,3)$ to a vector of shape $(F,3)$.

Therefore, we can use $P_{\vec{E}}$ and $P_{\vec{E}}^{-1}$ to map between
scalar representations and vector representations. Moreover, it allows us to
express any equivariant function acting on vectors with an invariant function
acting on scalars as shown in the following proposition.

######  Proposition 4.2.

Given any equivariant function $g$, there exists a function
$\tilde{g}=P_{\vec{E}}\cdot g\cdot P_{\vec{E}}^{-1}$ which takes scalar
representations as input and outputs scalar representations, where $\cdot$
means function composition. And $g=P_{\vec{E}}^{-1}\cdot\tilde{g}\cdot
P_{\vec{E}}$.

We can use an MLP to approximate the function $\tilde{g}$ and thus achieving
universal approximation for all $\text{O}(3)$-equivariant functions. This
demonstrates one advantage of the frame method—we can transform vectors to
scalars in the beginning, fully operate on the scalar space without losing any
expressive power. We can also transform scalars back to vector prediction if
needed.

###  4.2 How frames boost the expressive power

Encoding local environment. On the one hand, most GNNs for PES only pass
messages to the center atom from atoms within the sphere of radius $r_{c}$.
Therefore, encoding local environments is equivalent to one message passing
layer, and its expressive power is a part of the expressive power of the whole
GNN for PES. On the other hand, the interaction of two atoms will converge to
zero if the distance approaches infinity. Therefore, the state of an atom is
usually only closely related to the atoms close to it. Neglecting distant
atoms and encoding local environments can capture the physical image better.
Therefore, we study how well frames can be used to encode local environments.

The expressive power of encoding local environment reaches maximum when local
environments can be mapped injectively to center atom representations.
However, some popular models are under-expressive. For example, SchNet (Schütt
et al., 2017) only considers the distance between atoms and neglects the
angular information, leading to inability to differentiate some simple local
environments as shown in Figure 2(a). Moreover, Figure 2(b) illustrates that
though DimeNet (Klicpera et al., 2020) adds angular information to message
passing, its expressive power is still limited, which may be attributed to the
loss of high-order geometric information like dihedral angles. In contrast,
Proposition 4.3 shows that there exists a message passing scheme encoding
local environments injectively with a local frame.

######  Proposition 4.3.

Given a frame $\vec{E}$, we can build a function $\phi$ as follows, which
encodes local environments injectively, with suitable functions $\rho$ and
$\varphi$.

$\phi(\\{(s_{j},\vec{r}_{ij})|r_{ij}<r_{c}\\})=\rho(\sum_{r_{ij}<r_{c}}\varphi(\text{Concatenate}(P_{\vec{E}}(\vec{r}_{ij}),s_{j})).$
(4)

(a)

(b)

(c)

Figure 2: The green balls in the figure are the center atoms. We use balls
with different colors to represent different kinds of atoms. (a) SchNet cannot
distinguish two local environments due to the inability to capture angle. (b)
DimeNet cannot distinguish two local environments with the same set of angular
embeddings. Blue lines form a regular icosahedron and help visualization. The
center atom is at the symmetrical center of the icosahedron. (c) Some
invariant GNN models fail to pass the orientation information. However,
including projection of frame vectors can solve this problem (right). We only
show one vector (red) to represent the frame for simplicity. As a part of the
molecule rotates, the $\text{O}(3)$-equivariant local frame will also rotate,
leading to a different projection.

Passing messages across local environments. Though an atom can only collect
information from the local environment in one message passing step, messages
can be passed across local environments along a path of atoms. However, most
invariant models’ expressive power of multi-step message passing is still low.
They lose the orientation information with only invariant messages passed
across local environments. When a part of the molecule rotates, the center
atom representations will be invariant, leading to the same energy prediction
produced by them, while the true energy can change.

We can solve this problem by frame-frame projection, i.e., projecting the
frame of atom $j$ on the frame of atom $i$ and adding the projection into the
message from $j$ to $i$. As shown in Figure 2(c), frame-frame projection
elimilates the information loss. A more formalized discussion is in Appendix
G.

###  4.3 How to build a frame?

Usually, a frame should contain three linearly independent vectors. A
straightforward idea is to use the atom feature, like kind of atom, to produce
the frame. However, systems with central symmetry (like a single atom) cannot
produce non-zero vectors as a frame. We formally characterize this phenomenon
in the following proposition.

######  Proposition 4.4.

Assume that vector $\vec{v}$ is a function of a set of atoms
$\\{(s_{i},\vec{r}_{i})|i=1,2,...\\}$. If
$\\{(s_{i},\vec{r}_{i})|i=1,2,...\\}=\\{(s_{i},-\vec{r}_{i})|i=1,2,...\\}$,
then $\vec{v}=\vec{0}$.

Therefore, we consider producing the frame from the local environment of each
atom. Intuitively, if an atom $i$ has three non-symmetric neighbors not in the
same plane, the neighbors’ relative positions $\vec{r}_{ij}$ (sorted by their
norm) can naturally build a frame for the center atom which is
O$(3)$-equivariant. However, for atoms with highly symmetric local
environments, things can be different. Nevertheless, our main theorem shown
below states that we can always find an O$(3)$-equivariant frame for an atom
from its local environment no matter how the local environment varies.
Furthermore, the frame has full rank if the local environment has no symmetry.

######  Theorem 4.1.

There exists an O$(3)$-equivariant function $g$ mapping the local environment
$\text{LE}_{i}$ to a frame $\vec{E}_{i}\in{\mathbb{R}}^{3\times 3}$. The frame
has full rank iff there does not exist $o\in\text{O}(3)$, $o\neq I$,
$o(\text{LE}_{i})=\text{LE}_{i}$.

Theorem 4.1 guarantees that we can always build a full-rank frame for an atom
if its local environment has no symmmetry.

In general, building a frame for symmetric molecules remains a problem but
will not seriously hamper our model. Firstly, our model can produce reasonable
output even with symmetric input and is provably more expressive than a widely
used model SchNet (Schütt et al., 2017) (see Appendix J). Secondly, symmetric
molecules are rare, which form a zero-measure set. In our two representative
real-world datasets, less than $0.01\%$ of molecules (about ten molecules in
the whole datasets of several hundred thousand molecules) are symmetric.
Thirdly, symmetric geometry may be captured with a degenerate frame. As shown
in Figure 3(a), water is a symmetric molecule. We can use a frame with one
vector to describe its geometry. Last, a solution to frame degeneration is in
Appendix K.

(a)

(b)

Figure 3: (a) The left part shows the symmetry of the water molecule in
equilibrium, which has two planes of symmetry and a rotation axis. Its
equivariant vectors must be parallel to the rotation axis. However, with a
frame composed of only one vector, its geometry can be described. The right
part shows that with the projection of $\vec{r}_{ij}$ on the frame and the
distance between two atoms, the angle $\theta$ and the position of $j$ atom
can be determined. (b) The left part is a molecule with central symmetry. All
its vector features, including the global frame, will be zero. However, when
selected as the center (green), the atom’s environment has no central
symmetry.

A message passing process for frame generation. There exists a universal
framework for approximating O$(3)$-equivariant functions (Villar et al., 2021)
which can be used to implement the function $g$ in Theorem 4.1. For
scalability, we use a simplified form of that framework as follows.

$\vec{E}_{i}=\sum_{j\neq
i,r_{ij}<r_{c}}g^{\prime}(r_{ij},s_{j})\circ\frac{\vec{r}_{ij}}{r_{ij}},$ (5)

where $g^{\prime}$ is some function mapping scalar features and distance to
scalar weights and the entire framework reduces to a message passing process.
The derivation is detailed in Appendix B.

Local frame vs global frame. With the message passing framework, a frame is
produced for each atom. However, these local frames can also be summed to
produce a global frame.

$\vec{E}=\sum_{i=1}^{n}\vec{E}_{i}.$ (6)

Global frame can replace local frame and keep the invariance of energy
prediction. Moreover, if all atoms use the same frame, passing messages across
local environments will be straightforward, as the frame-frame projection
become trivial. However, the global frame is more likely to degenerate due to
the symmetry of molecules. As shown in Figure 3(b), the benzene molecule has
central symmetry and produces a zero global frame. However, when choosing each
atom as the center, the central symmetry is broken and a meaningful local
frame can be produced. We further formalize this intuition and prove that
global frame is more likely to degenerate in Appendix L. We also discuss why
we do not use the Gram-Schmidt process to make our frame vectors orthonormal
in Appendix C.

##  5 GNN with local frame

Based on our previous analysis, we formally introduce our GNN with local frame
(GNN-LF) model. The full architecture is detailed in Appendix D.

We first convert the input features of a molecule
$\vec{r}\in{\mathbb{R}}^{N\times 3}$ and $z\in{\mathbb{N}}^{N}$ to a graph.
The initial node feature $s_{i}^{(0)}\in{\mathbb{R}}^{F}$ is an embedding of
the kind of atom $z_{i}$. For edge $ij$, we construct two features: the edge
weight $w^{(e)}_{ij}=\text{cutoff}(r_{ij})$ (where cutoff means the cutoff
function), and the radial basis expansion of the distance
$s^{(e)}_{ij}=\text{rbf}(r_{ij})$. These special functions are detailed in
Appendix D.

Scalar message function. The scalar message $m_{ij}$ of edge $ij$ is a
function of the atom representations of $i$ and $j$ and the scalar features of
edge $ij$. Concatenating them is enough to keep the expressive power. However,
following Schütt et al. (2017), we use edge scalar features to produce a
filter $f_{ij}$ to filter the representation of atom $j$. Moreover, to make
the message smooth to the cutoff ratio, edge weight $w^{(e)}_{ij}$ is used to
scale the message. The message $m_{ij}$ is computed as follows.

$m_{ij}=w^{(e)}_{ij}(f_{ij}\odot s_{j}),$ (7)

where $\odot$ means element-wise product. Initially, $f_{ij}=f(s^{(e)}_{ij})$,
where $f$ is an invariant function outputting an $F$-dimensional scalar. This
message function empirically works well, which may be attributed to its
similarity to a quantum chemistry method (Tsubaki and Mizoguchi, 2020).

Producing frame. Combining Equation 5 and the form of scalar message in
Equation 7, we get the form of the message passing scheme for producing the
local frame (denoted by $\vec{E}_{i}$) as follows.

$\vec{E}_{i}=\sum_{j\neq i,r_{ij}<r_{c}}w^{(e)}_{ij}(f_{ij}\odot
s_{j})\circ\frac{\vec{r}_{ij}}{r_{ij}}.$ (8)

Note that $\vec{E}_{i}\in{\mathbb{R}}^{F\times 3}$ is not restricted to have
three vectors, but grows with the embedding dimension $F$. This allows
dimension consistency across the network and reduces the risk of degeneracy.

Projection of vector feature. With local frame, we can project the vector
features on the frame and add them to the scalar edge features. For edge
direction:

$d^{1}_{ij}=P(\frac{\vec{r}_{ij}}{r_{ij}},\vec{E}_{i}),~{}d^{2}_{ij}=P(\frac{\vec{r}_{ij}}{r_{ij}},\vec{E}_{j}).$
(9)

Note that $d^{1}$ is the edge direction projection on the center node, while
$d^{2}$ is that on the neighbor node. Though $d^{1}$ is enough in our
theoretical analysis, adding $d^{2}$ can bring a slight improvement in some of
our experiments. Thus, we treat whether to use $d^{2}$ as a hyperparameter.

Directly using frame-frame projection leads to a large matrix. Therefore, we
only use the diagonal elements of the projection. To keep the expressive
power, we transform the frame with a linear layer before the inner product
operation.

$d^{3}_{ij}=\langle W_{1}\vec{v}_{j},W_{2}\vec{v}_{i}\rangle.$ (10)

After all the projections, we update the filter function in (7) by including
the additional scalars:

$f_{ij}=f(s^{(e)}_{ij},d^{1}_{ij},d^{2}_{ij},d^{3}_{ij}),$ (11)

which is used in the later scalar message passing steps.

Decomposing filter. The filter function is further decomposed. The distance
information is easier to learn as it has been expanded with a set of bases.
Thus a linear layer $g_{1}$ is enough. In contrast, the directional
information is the inner product of vectors, which may need a more expressive
MLP $g_{2}$. Therefore, we generate two filters separately.

$f(s^{(e)}_{ij},d^{1}_{ij},d^{2}_{ij},d^{3}_{ij})=g_{1}(s^{(e)}_{ij})\odot
g_{2}(d^{1}_{ij},d^{2}_{ij},d^{3}_{ij}).$ (12)

Sharing filters. Generating different filters for each message passing layer
is time-consuming, as $O(Nn)$ filters need to be computed in each layer.
Therefore, we share filters between different layers. Experimental results
show that sharing filters only leads to minor performance loss.

##  6 Experiment

In this section, we compare GNN-LF with existing models on different tasks and
do ablation analysis. We report the mean absolute error (MAE) on test set (the
lower the better) for all tasks. All our results are averaged over three
random dataset splits. Baselines’ results are from their original papers. The
best results are shown in bold and the second best results are shown with
underlines in tables. Our experimental settings are detailed in Appendix E.

###  6.1 Modeling PES

We first evaluate our models for modeling PES on the MD17 dataset (Chmiela et
al., 2017), which consists of MD trajectories of small organic molecules. We
compare our model with a manual descriptor model: FCHL (Christensen et al.,
2020) , invariant models: SchNet (Schütt et al., 2017), DimeNet (Klicpera et
al., 2020), GemNet (Gasteiger et al., 2021), a model using irreducible
representations of SO$(3)$ group: NequIP (Batzner et al., 2021), and models
using vector representations: PaiNN (Schütt et al., 2021) and TorchMD (Thölke
and Fabritiis, 2022). The results are shown in Table 1.

Table 1: Results on the MD17 dataset. Units: energy (E) (kcal/mol) and forces (F) (kcal/mol/Å).  |  | FCHL | SchNet | DimeNet | GemNet | PaiNN | NequIP | TorchMD | GNN-LF
---|---|---|---|---|---|---|---|---|---
Aspirin | E | 0.182 | 0.37 | 0.204 | - | 0.167 | - | 0.124 | 0.1342
F | 0.478 | 1.35 | 0.499 | 0.2168 | 0.338 | 0.348 | 0.255 | 0.2018
Benzene | E | - | 0.08 | 0.078 | - | - | - | 0.056 | 0.0686
F | - | 0.31 | 0.187 | 0.1453 | - | 0.187 | 0.201 | 0.1506
Ethanol | E | 0.054 | 0.08 | 0.064 | - | 0.064 | - | 0.054 | 0.0522
F | 0.136 | 0.39 | 0.230 | 0.0853 | 0.224 | 0.208 | 0.116 | 0.0817
Malonaldehyde | E | 0.081 | 0.13 | 0.104 | - | 0.091 | - | 0.079 | 0.0755
F | 0.245 | 0.66 | 0.383 | 0.1545 | 0.319 | 0.337 | 0.176 | 0.1239
Naphthalene | E | 0.117 | 0.16 | 0.122 | - | 0.166 | - | 0.085 | 0.1136
F | 0.151 | 0.58 | 0.215 | 0.0553 | 0.077 | 0.097 | 0.060 | 0.0550
Salicylic acid | E | 0.114 | 0.20 | 0.134 | - | 0.166 | - | 0.094 | 0.1074
F | 0.221 | 0.85 | 0.374 | 0.1268 | 0.195 | 0.238 | 0.135 | 0.1005
Toluence | E | 0.098 | 0.12 | 0.102 | - | 0.095 | - | 0.074 | 0.0950
F | 0.203 | 0.57 | 0.216 | 0.0600 | 0.094 | 0.101 | 0.066 | 0.0529
Uracil | E | 0.104 | 0.14 | 0.115 | - | 0.106 | - | 0.096 | 0.1037
F | 0.105 | 0.56 | 0.301 | 0.0969 | 0.139 | 0.173 | 0.094 | 0.0751
average | rank | 3.93 | 6.63 | 5.38 | 2.00 | 4.36 | 5.25 | 2.25 | 1.75

Our model outperforms all the baselines on nine out of sixteen targets and
achieves the second-best performance on all other seven targets, which is
better than any other model. Compared with GemNet, the model with the second
highest average rank, our model leads to $10\%$ lower loss on average. The
outstanding performance verifies the effectiveness of the local frame method
for modeling PES. Moreover, our model also uses fewer parameters and less than
$30\%$ computation time compared with the baselines as shown in Appendix F.

###  6.2 Ablation study

We perform an ablation study to verify our model designs. The results are
shown in Table 2. On average, ablation of $d^{3}$ term in Equation (10)
(NoDir3) leads to $20\%$ higher MAE, which verifies the necessity of frame-
frame projection. In contrast, ablation of the $d^{2}$ term in Equation (9)
(NoDir2) only results in $1\%$ higher MAE on three datasets, which illustrates
that, as shown in our theoretical analysis, projection of direction on the
target node and frame is enough for high expressive power. The column Global
replaces the local frames with the pooled global frame, resulting in $100\%$
higher loss, which verifies the advantages of local frames. Ablation of filter
decomposition (NoDecomp) leads to $9\%$ higher loss, indicating the advantage
of separately processing distance and projections. Finally, we find although
using different filters for each message passing layer (NoShare) uses much
more computation time (1.67$\times$) and parameters (3.55$\times$), it only
leads to $0.01\%$ lower loss on average, illustrating that sharing filters
does little harm to the expressive power.

Table 2: Ablation results on the MD17 dataset. Units: energy (E) (kcal/mol) and forces (F) (kcal/mol/Å). GNN-LF does not use $d^{2}$ for some molecules, so the NoDir2 column is empty.  |  | GNN-LF | NoDir2 | NoDir3 | Global | NoDecomp | GNN-LF | Noshare
---|---|---|---|---|---|---|---|---
Aspirin | E | 0.1342 | - | 0.1435 | 0.2280 | 0.1411 | 0.1342 | 0.1364
F | 0.2018 | - | 0.2799 | 0.6894 | 0.2622 | 0.2018 | 0.1979
Benzene | E | 0.0686 | - | 0.0716 | 0.0972 | 0.0688 | 0.0686 | 0.0713
F | 0.1506 | - | 0.1583 | 0.3520 | 0.1499 | 0.1506 | 0.1507
Ethanol | E | 0.0520 | - | 0.0532 | 0.0556 | 0.0518 | 0.0520 | 0.0514
F | 0.0814 | - | 0.0930 | 0.1465 | 0.0847 | 0.0814 | 0.0751
Malonaldehyde | E | 0.0755 | 0.0764 | 0.0776 | 0.0923 | 0.0762 | 0.0755 | 0.0790
F | 0.1239 | 0.1259 | 0.1466 | 0.3194 | 0.1321 | 0.1239 | 0.1210
Naphthalene | E | 0.1136 | - | 0.1152 | 0.1276 | 0.1254 | 0.1136 | 0.1168
F | 0.0550 | - | 0.0834 | 0.2069 | 0.0553 | 0.0550 | 0.0547
Salicylic acid | E | 0.1074 | 0.1081 | 0.1087 | 0.1224 | 0.1123 | 0.1074 | 0.1091
F | 0.1005 | 0.1048 | 0.1328 | 0.2890 | 0.1399 | 0.1005 | 0.1012
Toluence | E | 0.0950 | 0.0930 | 0.0942 | 0.1000 | 0.0932 | 0.0950 | 0.0942
F | 0.0529 | 0.0543 | 0.0770 | 0.1659 | 0.0695 | 0.0529 | 0.0519
Uracil | E | 0.1037 | - | 0.1069 | 0.1075 | 0.1053 | 0.1037 | 0.1042
F | 0.0751 | - | 0.0964 | 0.1901 | 0.0825 | 0.0751 | 0.0754

###  6.3 Other chemical properties

Though designed for PES, our model can also predict other properties directly.
The QM9 dataset (Wu et al., 2018) consists of 134k stable small organic
molecules. The task is to use the kind of atoms $z$ and coordinates of atoms
$\vec{r}$ to predict $12$ properties of these molecules. We compare our model
with invariant models: SchNet (Schütt et al., 2017), DimeNet++ (Gasteiger et
al., 2020), a model using irreducible representations: Cormorant (Anderson et
al., 2019), and models using vector representations: PaiNN (Schütt et al.,
2021) and TorchMD (Thölke and Fabritiis, 2022). The results are shown in Table
3. Our model outperforms all other models on seven out of twelve tasks and
achieves the second-best performance on four out of the five left tasks, which
illustrates that the local frame method has the potential to be applied to
other fields.

Table 3: Results on the QM9 dataset. Target | Unit | SchNet | DimeNet++ | Cormorant | PaiNN | Torchmd | GNN-LF
---|---|---|---|---|---|---|---
$\mu$ | D | 0.033 | 0.0297 | 0.038 | 0.012 | 0.002 | 0.013
$\alpha$ | $a_{0}^{3}$ | 0.235 | 0.0435 | 0.085 | 0.045 | 0.01 | 0.0353
$\epsilon_{HOMO}$ | meV | 41 | 24.6 | 34 | 27.6 | 21.2 | 23.5
$\epsilon_{LUMO}$ | meV | 34 | 19.5 | 38 | 20.4 | 17.8 | 17.0
$\Delta\epsilon$ | meV | 63 | 32.6 | 61 | 45.7 | 38 | 37.1
$\langle R^{2}\rangle$ | $a_{0}^{2}$ | 0.073 | 0.331 | 0.961 | 0.066 | 0.015 | 0.037
ZPVE | meV | 1.7 | 1.21 | 2.027 | 1.28 | 2.12 | 1.19
$U_{0}$ | meV | 14 | 6.32 | 22 | 5.85 | 6.24 | 5.30
$U$ | meV | 19 | 6.28 | 21 | 5.83 | 6.30 | 5.24
$H$ | meV | 14 | 6.53 | 21 | 5.98 | 6.48 | 5.48
$G$ | meV | 14 | 7.56 | 20 | 7.35 | 7.64 | 6.84
$C_{v}$ | cal/mol/K | 0.033 | 0.023 | 0.026 | 0.024 | 0.026 | 0.022

##  7 Conclusion

This paper proposes a local frame method, which avoids passing vector
representations during message passing while keeping high expressive power. We
show that local frames enable GNNs to encode local environments around atoms
injectively and pass orientation information across local environments.
Furthermore, we give conditions for the existence of local frames and propose
ways to construct them. In the experiments, our model outperforms all
baselines on modeling PES and predicting other properties, while using fewer
parameters and less inference time, which verifies the simplicity,
effectiveness and scalability of the local frame method.

## References

  * Schmitz et al. (2019) Gunnar Schmitz, Ian Heide Godtliebsen, and Ove Christiansen.  Machine learning for potential energy surfaces: An extensive database and assessment of methods.  _The Journal of Chemical Physics_ , 150(24):244113, 2019.
  * Lewars (2016) Errol G. Lewars.  _The Concept of the Potential Energy Surface_ , pages 9–49.  Springer International Publishing, 2016.
  * Schutt et al. (2017) K. T. Schutt, F. Arbabzadah, S. Chmiela, K.-R. Muller, and A. Tkatchenko.  Quantum-chemical insights from deep tensor neural networks.  _Nature Communications_ , 8:13890, 2017.
  * Schütt et al. (2017) Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller.  Schnet: A continuous-filter convolutional neural network for modeling quantum interactions.  In _Advances in Neural Information Processing Systems 30_ , pages 991–1001, 2017.
  * Klicpera et al. (2020) Johannes Klicpera, Janek Groß, and Stephan Günnemann.  Directional message passing for molecular graphs.  In _International Conference on Learning Representations_ , 2020.
  * Anderson et al. (2019) Brandon M. Anderson, Truong-Son Hy, and Risi Kondor.  Cormorant: Covariant molecular neural networks.  In _Advances in Neural Information Processing Systems_ , pages 14510–14519, 2019.
  * Schütt et al. (2021) Kristof Schütt, Oliver T. Unke, and Michael Gastegger.  Equivariant message passing for the prediction of tensorial properties and molecular spectra.  In _International Conference on Machine Learning_ , volume 139, pages 9377–9388, 2021.
  * Gasteiger et al. (2021) Johannes Gasteiger, Florian Becker, and Stephan Günnemann.  Gemnet: Universal directional graph neural networks for molecules.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Batzner et al. (2021) Simon L. Batzner, Tess E. Smidt, Lixin Sun, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, and Boris Kozinsky.  Se(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials.  _CoRR_ , abs/2101.03164, 2021.
  * Thölke and Fabritiis (2022) Philipp Thölke and Gianni De Fabritiis.  Equivariant transformers for neural network based molecular potentials.  In _International Conference on Learning Representations_ , 2022.
  * Kipf and Welling (2017) Thomas N. Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  _International Conference on Learning Representations_ , 2017.
  * Thomas et al. (2018) Nathaniel Thomas, Tess E. Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley.  Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds.  _CoRR_ , abs/1802.08219, 2018.
  * Rupp et al. (2012) Matthias Rupp, Alexandre Tkatchenko, Klaus-Robert Müller, and O. Anatole von Lilienfeld.  Fast and accurate modeling of molecular atomization energies with machine learning.  _Physical Review Letter_ , 108:058301, Jan 2012.
  * Chmiela et al. (2019) Stefan Chmiela, Huziel E. Sauceda, Igor Poltavsky, Klaus-Robert Müller, and Alexandre Tkatchenko.  sgdml: Constructing accurate and data efficient molecular force fields using machine learning.  _Computer Physics Communications_ , 240:38–45, 2019.
  * Christensen et al. (2020) Anders S Christensen, Lars A Bratholm, Felix A Faber, and O Anatole von Lilienfeld.  Fchl revisited: Faster and more accurate quantum machine learning.  _The Journal of chemical physics_ , 152(4):044107, 2020.
  * Behler and Parrinello (2007) Jörg Behler and Michele Parrinello.  Generalized neural-network representation of high-dimensional potential-energy surfaces.  _Physical Review Letter_ , 98:146401, 2007.
  * Smith et al. (2017) J. S. Smith, O. Isayev, and A. E. Roitberg.  Ani-1: an extensible neural network potential with dft accuracy at force field computational cost.  _Chem. Sci._ , 8:3192–3203, 2017.
  * Zhang et al. (2018) Linfeng Zhang, Jiequn Han, Han Wang, Wissam Saidi, Roberto Car, and Weinan E.  End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems.  In _Advances in Neural Information Processing Systems_ , 2018.
  * Du et al. (2022) Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie-Yan Liu.  SE(3) equivariant graph neural networks with complete local frames.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, _International Conference on Machine Learning_ , 2022.
  * Satorras et al. (2021) Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling.  E(n) equivariant graph neural networks.  In _International Conference on Machine Learning_ , 2021.
  * Luo et al. (2021) Shitong Luo, Jiahan Li, Jiaqi Guan, Yufeng Su, Chaoran Cheng, Jian Peng, and Jianzhu Ma.  Equivariant point cloud analysis via learning orientations for message passing.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , pages 16296–16305, 2021.
  * Huang et al. (2022) Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang.  Equivariant graph mechanics networks with constraints.  _International Conference on Learning Representations_ , 2022.
  * Villar et al. (2021) Soledad Villar, David W Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith.  Scalars are universal: Equivariant machine learning, structured like classical physics.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Deng et al. (2021) Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J. Guibas.  Vector neurons: A general framework for so(3)-equivariant networks.  In _International Conference on Computer Vision_ , pages 12180–12189. IEEE, 2021.
  * Tsubaki and Mizoguchi (2020) Masashi Tsubaki and Teruyasu Mizoguchi.  On the equivalence of molecular graph convolution and molecular wave function with poor basis set.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 1982–1993. Curran Associates, Inc., 2020.
  * Chmiela et al. (2017) Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller.  Machine learning of accurate energy-conserving molecular force fields.  _Science Advances_ , 3(5):e1603015, 2017.
  * Wu et al. (2018) Zhenqin Wu, Bharath Ramsundar, Evan N. Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S. Pappu, Karl Leswing, and Vijay Pande.  Moleculenet: a benchmark for molecular machine learning.  _Chemical Science_ , 9:513–530, 2018.
  * Gasteiger et al. (2020) Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann.  Fast and uncertainty-aware directional message passing for non-equilibrium molecules.  In _Machine Learning for Molecules Workshop, NeurIPS_ , 2020.
  * Zaheer et al. (2017) Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alexander J. Smola.  Deep sets.  In _Advances in Neural Information Processing Systems_ , pages 3391–3401, 2017.
  * Kravtsiv and Zagorodnyuk (2012) Viktoriia Kravtsiv and Andriy Zagorodnyuk.  On algebraic bases of algebras of block-symmetric polynomials on banach spaces.  _Matematychni Studiï_ , 37, 03 2012.
  * Unke and Meuwly (2019) Oliver Unke and Markus Meuwly.  Physnet: A neural network for predicting energies, forces, dipole moments and partial charges.  _J Chem Theory Comput._ , 6:3678–3693, 02 2019.
  * Segol and Lipman (2020) Nimrod Segol and Yaron Lipman.  On universal equivariant set networks.  In _International Conference on Learning Representations_ , 2020.
  * Dym and Maron (2021) Nadav Dym and Haggai Maron.  On the universality of rotation equivariant point cloud networks.  In _International Conference on Learning Representations_ , 2021.
  * Puny et al. (2022) Omri Puny, Matan Atzmon, Heli Ben-Hamu, Edward J. Smith, Ishan Misra, Aditya Grover, and Yaron Lipman.  Frame averaging for invariant and equivariant network design.  _International Conference on Learning Representations_ , 2022.

## Checklist

The checklist follows the references. Please read the checklist guidelines
carefully for information on how to answer these questions. For each question,
change the default [TODO] to [Yes] , [No] , or [N/A] . You are strongly
encouraged to include a justification to your answer, either by referencing
the appropriate section of your paper or providing a brief inline description.
For example:

  * •

Did you include the license to the code and datasets? [No] The code and the
data are proprietary.

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes]

    3. (c)

Did you discuss any potential negative societal impacts of your work? [N/A]

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]

    2. (b)

Did you include complete proofs of all theoretical results? [Yes]

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes]

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [No] Existing works seldom report error bars.

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We use
internal cluster and clarify our GPU resource.

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes]

    2. (b)

Did you mention the license of the assets? [Yes]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [No]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A]

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A]

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

##  Appendix A Proofs

Due to repulsive force, atoms cannot be too close to each other in stable
molecules. Therefore, we assume that there exist an upper bound $N$ of the
number of neighbor atoms.

###  A.1 Proof of Lemma 3.1

###### Proof.

For all function $g$, scalar $s$, transformation $o\in\text{O}(3)$,
$g(s(z,\vec{r}o^{T}))=g(s)$. Therefore, $g(s)$ is a scalar.

For all scalar $s$, vector $\vec{v}$, and transformation $o\in\text{O}(3)$,
$s(z,\vec{r}o^{T})\circ\vec{v}(z,\vec{r}o^{T})=s(z,\vec{r})\circ(\vec{v}(z,\vec{r})o^{T})=(s(z,\vec{r})\circ\vec{v}(z,\vec{r}))o^{T}$.
Therefore, $s\circ\vec{v}$ is a vector.

For all vectors $\vec{v}_{1},\vec{v}_{2}$,
$\vec{v}_{1}(z,\vec{r}o^{T})\vec{v}_{2}(z,\vec{r}o^{T})^{T}=\vec{v}_{1}(z,\vec{r})o^{T}o\vec{v}_{2}(z,\vec{r})^{T}=\vec{v}_{1}(z,\vec{r})\vec{v}_{2}(z,\vec{r})^{T}$.

As $P(\vec{v}_{1},\vec{v}_{2})$ is a scalar,
$\langle\vec{v}_{1},\vec{v}_{2}\rangle=\text{diag}(P(\vec{v}_{1},\vec{v}_{2}))$
must be a scalar. ∎

###  A.2 Proof of Proposition 4.1

###### Proof.

Let $\vec{v}$ denote a vector. $P_{\vec{E}}(\vec{v})$ is scalar because

$\displaystyle\forall o\in\text{O}(3),P_{\vec{E}}(\vec{v})(z,\vec{r}o^{T})$
$\displaystyle=\vec{v}(z,\vec{r}o^{T})\vec{E}(z,\vec{r}o^{T})^{T}$ (13)
$\displaystyle=\vec{v}(z,\vec{r})o^{T}o\vec{E}(z,\vec{r})^{T}$ (14)
$\displaystyle=P_{\vec{E}}(\vec{v})(z,\vec{r}).$ (15)

$P_{\vec{E}}$ is bijective as
$\vec{x}=P_{\vec{E}}(\vec{x})(\vec{E}^{-1})^{T}$. Obviously,
$P_{\vec{E}}^{-1}(s)=s(\vec{E}^{-1})^{T}$, $s\in{\mathbb{R}}^{F\times 3}$, $s$
is a scalar representation. To prove that $P^{-1}(s)$ is a vector
representation, we note that $P^{-1}(s)$ is itself a function of
$(z,\vec{r})$, which can be written as $P^{-1}(s)(z,\vec{r})$. We have

$\displaystyle\forall o\in\text{O}(3),P_{\vec{E}}^{-1}(s)(z,\vec{r})o^{T}$
$\displaystyle=s(z,\vec{r})[(\vec{E}(z,\vec{r})o^{T})^{-1}]^{T}$ (16)
$\displaystyle=s(z,\vec{r}o^{T})(\vec{E}(z,\vec{r}o^{T})^{-1})^{T}$ (17)
$\displaystyle=P_{\vec{E}}^{-1}(s)(z,\vec{r}o^{T}).$ (18)

Therefore, $P_{\vec{E}}^{-1}(s)$ is a vector representation. ∎

###  A.3 Proof of Proposition 4.2

###### Proof.

Assume that $s$ is a scalar representation.

$\displaystyle\tilde{g}(s)$
$\displaystyle=P_{\vec{E}}(g(P_{\vec{E}}^{-1}(s)))$ (19)
$\displaystyle=g(s(\vec{E}(z,\vec{r})^{-1})^{T}))\vec{E}(z,\vec{r})^{T}.$ (20)

The representation $\tilde{g}(s)$ can be written as a function of
$(z,\vec{r})$. Then, we have

$\displaystyle\forall o\in\text{O}(3),\tilde{g}(s)(z,\vec{r}o^{T})$
$\displaystyle=g(s(\vec{E}(z,\vec{r}o^{T})^{-1})^{T}))\vec{E}(z,\vec{r}o^{T})^{T}$
(21)
$\displaystyle=g(s(\vec{E}(z,\vec{r})^{-1})^{T})o^{T})o\vec{E}(z,\vec{r})^{T}$
(22)
$\displaystyle=g(s(\vec{E}(z,\vec{r})^{-1})^{T}))o^{T}o\vec{E}(z,\vec{r})^{T}$
(23) $\displaystyle=g(s(\vec{E}(z,\vec{r})^{-1})^{T}))\vec{E}(z,\vec{r})^{T}$
(24) $\displaystyle=\tilde{g}(s)(z,\vec{r}).$ (25)

Therefore, $\tilde{g}(s)$ is also a scalar representation. ∎

###  A.4 Proof of Proposition 4.3

Let $\gamma$ denote a function mapping local environments to a set of scalars
with a frame.

$\gamma(\\{(s_{j},\vec{r}_{ij})|r_{ij}<r_{c}\\})=\\{\text{Concatenate}(s_{j},P_{\vec{E}}(\vec{r}_{ij}))|r_{ij}<r_{c}\\}.$
(26)

As it has an inverse function, $\gamma$ is injective.

$\gamma^{-1}(\\{s_{j}|j=1,2,..\\})=\\{(s_{j}[:-3],P^{-1}_{\vec{E}}(s_{j}[-3:])|j=1,2,...)\\}.$
(27)

Therefore, we only need an expressive scalar aggregation function. Zaheer et
al. [2017] proved the existence of such function for scalar in $[0,1]$.
However, our scalar feature $\in{\mathbb{R}}^{F}$, where $F$ is a
hyperparameter usually $\geq 4$, so we extend the conclusion of Zaheer et al.
[2017].

######  Lemma A.1.

There exists a function $\psi$ mapping a scalar set of size $\leq N$ to a
scalar injectively.

###### Proof.

For a set of scalar $S=\\{s_{i}|i=1,2,...,n,s^{(i)}\in{\mathbb{R}}^{F}\\}$, we
define a polynomial as follows.

$p^{(S)}(x_{1},...,x_{F})=\prod_{i=1}^{n}[\sum_{j=1}^{F}(x_{i}-s_{ij})^{2}],$
(28)

where $s_{ij}$ is the $j^{\text{th}}$ element of $s_{i}$.
$p^{(S)}(x_{1},...,x_{F})$ is a $2nF$-order polynomial. Let $C_{t}$ donate the
coefficient of the $t$-th term by some order. $C_{t}$ is a polynomial of
$s^{(i)}_{j},i=1,2,...,n,j=1,2,...,F$ and can be formalized as follows.

$C_{t}=\sum_{D}\alpha^{(t,D)}\prod_{i=1}^{n}\prod_{j=1}^{F}s_{ij}^{d_{ij}},$
(29)

where
$D\in\\{[d_{11},d_{12},...,d_{ij},...,d_{nF}]|d_{ij}\in{\mathbb{N}},\sum_{ij}d_{ij}\leq
2nF\\}$. Note that $C_{t}$ is symmetric to the permutation of $i$. Such
polynomials are called block-symmetric polynomial [Kravtsiv and Zagorodnyuk,
2012].

As shown in [Kravtsiv and Zagorodnyuk, 2012], the generating elements of
block-symmetric polynomial are as follows.

$H_{\kappa}^{k_{1},k_{2},...,k_{F}}=\sum_{i=1}^{n}\prod_{j=1}^{F}s_{ij}^{k_{i}},\sum_{j=1}^{F}k_{j}=\kappa.$
(30)

Because degree of $C_{t}$ is less than $2nF+1$, $C_{t}$ is a function of
$[H_{\kappa}^{k_{1},k_{2},...,k_{F}}|\sum_{j=1}^{F}k_{j}=\kappa,\kappa=1,2,...,2nF]$.

Let
$\varphi(s_{i})=[\prod_{j=1}^{F}s_{ij}^{k_{i}}|\sum_{j=1}^{F}k_{j}=\kappa,\kappa=1,2,...,2nF]$.
$\sum_{i}\varphi(s_{i})=[H_{\kappa}^{k_{1},k_{2},...,k_{F}}|\sum_{j=1}^{F}k_{j}=\kappa,\kappa=1,2,...,2nF]$
determines $C_{t}$ and the polynomial $p^{(S)}(x_{1},...,x_{F})$.

Assuming that $S\neq S^{\prime}$ but
$\sum_{i}\varphi(s_{i})\neq\sum_{i}\varphi(s^{\prime}_{i})$. Therefore,
$p^{(S)}(x_{1},...,x_{F})=p^{(S^{\prime})}(x_{1},...,x_{F})$. Therefore, for
all $s_{i}\in S$,
$p^{(S^{\prime})}(s_{i1},s_{i2},...,s_{iF})=p^{(S)}(s_{i1},s_{i2},...,s_{iF})=0$,
$s\in S^{\prime}$. Therefore, $S\subseteq S^{\prime}$. As
$|S|=|S^{\prime}|=n$, $S=S^{\prime}$, leading to a contradiction.

Therefore, given two set of scalars
$S=\\{s_{i}|s_{i}\in{\mathbb{R}}^{F},i=1,2,...,n\\}$,
$S^{\prime}=\\{s^{\prime}_{i}|s^{\prime}_{i}\in{\mathbb{R}}^{F},i=1,2,...,n\\}$,
$S\neq
S^{\prime}\longleftarrow\sum_{i}\varphi(s_{i})\neq\sum_{i}\varphi(s^{\prime}_{i})$.
There also exists a function $\psi_{n}$ map sets of $n$ elements from
${\mathbb{R}}^{F}$ injectively.

$\psi_{n}(\\{s_{i}|s_{i}\in{\mathbb{R}}^{F},i=1,2,...,n\\})=\rho_{n}(\sum_{i=1}^{n}\phi_{n}(x_{i})),$
(31)

where $\rho_{n}$, $\phi_{n}$ are suitable functions.

Let $\phi(s)$ denote the sequence $[1,\phi_{0}(s),...,\phi_{N}(s)]$. Therefore
$\sum_{i=1}^{n}\phi(s_{i})$ is a sequence
$[n,\sum_{i=1}^{n}\phi_{0}(s_{i}),...,\sum_{i=1}^{n}\phi_{N}(s_{i})]$. Let
$\rho(n,\sum_{i=1}^{n}\phi_{0}(s_{i}),...,\sum_{i=1}^{n}\phi_{N}(s_{i}))$
produce a output pair $[n,\rho_{n}(\sum_{i=1}^{n}\phi_{n}(s_{i}))]$. Let
$\psi(S)=\tilde{\rho}(\sum_{s\in S}\tilde{\phi}(s))$. Therefore, $\psi$ can
map feature sets of size $\leq N$ injectively. ∎

###  A.5 Proof of Proposition 4.4

$\displaystyle\vec{v}(\\{(s_{i},-\vec{r}_{i})\\})$
$\displaystyle=\vec{v}(\\{(s_{i},\vec{r}_{i}\begin{bmatrix}-1&0&0\\\ 0&-1&0\\\
0&0&-1\end{bmatrix})\\})$ (32)
$\displaystyle=\vec{v}(\\{(s_{i},\vec{r}_{i})\\})\begin{bmatrix}-1&0&0\\\
0&-1&0\\\ 0&0&-1\end{bmatrix}$ (33)
$\displaystyle=-\vec{v}(\\{(s_{i},\vec{r}_{i})\\}).$ (34)

Therefore, if $\\{(s_{i},\vec{r}_{i})\\}=\\{(s_{i},-\vec{r}_{i})\\}$,
$\vec{v}(\\{(s_{i},\vec{r}_{i})\\})=-\vec{v}(\\{(s_{i},\vec{r}_{i})\\})=0$.

###  A.6 Proof of Theorem 4.1

###### Proof.

(1) We first prove there exists an O(3)-equivariant function $g$ mapping the
local environment $\text{LE}_{i}$ to a frame
$\vec{E}_{i}\in{\mathbb{R}}^{3\times 3}$. The frame has full rank if there
does not exist $o\in\text{O}(3),o\neq I,o(LE_{i})=LE_{i}$.

Let $\gamma$ denote a function mapping local environments to sets of vectors.

$\gamma(\\{(\vec{r}_{ij},s_{j})|r_{ij}<r_{c}\\})=\\{\text{Concatenate}(s_{j}\vec{r}_{ij},\vec{r}_{ij})|r_{ij}<r_{c}\\},$
(35)

in which $s_{j}$ is reshaped as $(F,1)$, $\vec{r}_{ij}$ is of shape $(1,3)$.
$\gamma$ is $\text{O}(3)$-equivariant. Therefore, we discuss the aggregation
function on a set of vector, denoted as
$\\{\vec{u}_{i}|i=1,2,...,n,\vec{u}_{i}\in\mathbb{R}^{F\times 3}\\}$.

Assume that $V=\\{\\{\vec{u}_{i}|i=1,2,...,n,\vec{u}_{i}\in\mathbb{R}^{F\times
3}\\}|n=1,2,...,N\\}$, where $N$ is the upper bound of the size of local
environment, is the set of sets of vector messages in local environment.

An equivalence relation can be defined on $V$: $v_{1}\in V,v_{2}\in
V,v_{1}\sim v_{2}$ iff there exists $o\in\text{O}(3),o(v_{1})=v_{2}$. Let
$\tilde{V}=V/\sim$ denote the quotient set. For each equivalence class $[v]$
with no symmetry, a representative $v$ can be selected. We can define a
function $r:\tilde{V}-\\{[v]|[v]\in\tilde{V},\exists
v\in[v],o\in\text{O}(3),o\neq I,o(v)=v\\}\to V$ as $r([v])=v$ mapping each
equivalence class with no symmetry to its representative. For a message set
with no symmetry, the transformation from its representative to it is also
unique. Let $h:V-\\{v|v\in V,\exists o\in\text{O}(3),o\neq
I,o(v)=v\\}\to\text{O}(3)$. $h(v)=o,o(r([v]))=v$.

Therefore, the function $g$ can take the form as follows.

$g(v)=\begin{cases}\left[\begin{matrix}0&0&0\\\ 0&0&0\\\
0&0&0\end{matrix}\right]&\text{if there exists }o\in\text{O(3)},o\neq
I,ov=v\\\ h(v)^{T}&\text{otherwise}\end{cases}$

Therefore, $g\cdot\gamma$ is the required function.

We further detail how to select the representative elements: We first define a
linear order relation $\leq_{l}$ in $V$. If $v_{1},v_{2}\in
V,|v_{1}|<|v_{2}|$, $v_{1}<_{l}v_{2}$. So we only consider the order relation
between two sets of the same size $n$.

We first define a function $\psi$ mapping message set to a sequence
injectively.

$\psi(\\{u_{i}|i=1,2,...,n,u_{i}\in{\mathbb{R}}^{F\times
3}\\})=[\text{flatten}(u_{\pi(i)})|i=1,2,...,n,\\\ \pi\text{ is a permutation
sorting }u_{i}\text{ by lexicographical order}].$ (36)

Forall $v_{1},v_{2}\in V,|v_{1}|=|v_{2}|$, $v_{1}\leq_{l}v_{2}$ iff
$\psi(v_{1})\leq\psi(v_{2})$ by lexicographical order. As the size of local
environent is bounded, the sequence is also of a finite length. Therefore, the
lexicographical order and thus the linear order relation $\leq_{l}$ are well-
defined.

All permutations of $\\{1,2,...,n\\}$ form a permutation set $\Pi_{n}$.

For all $[v]\in\tilde{V}$, let
$r([v])=\operatorname*{arg\,min}_{v^{\prime}\in[v]}\psi(v^{\prime})$. To
illustrate the existence of such minimal sequence, we reform it.

$\displaystyle min_{v^{\prime}\in[v]}\psi(v^{\prime})$
$\displaystyle=\min_{\pi\in\Pi_{n},o\in\text{O}(3)}S(o,\pi)$ (37)
$\displaystyle=\min_{\pi\in\Pi_{n}}\min_{o\in\text{O}(3)}S(o,\pi),$ (38)

where $S(o,\pi)=[\text{flatten}(u_{\pi(i)}o^{T})|i=1,2,...,n]$. Each element
of this sequence is continuous to $o$.

We first fix the $\pi$. As $\text{O}(3)$ is a compact group,
$\operatorname*{arg\,min}_{o\in\text{O}(3)}S(o,\pi)_{1}$ exists. Let
$L_{1}=\\{o|o\in\text{O}(3),S(o,\pi)_{1}=\min_{o^{\prime}\in\text{O}(3)}S(o^{\prime},\pi)_{1}\\}$
is still a compact set. Therefore, $\operatorname*{arg\,min}_{o\in
L_{1}}S(o,\pi)_{1}$ exists. Let $L_{2}=\\{o|o\in
L_{1},S(o,\pi)_{2}=\min_{o^{\prime}\in L_{1}}S(o^{\prime},\pi)_{2}\\}$.
Similarly, $L_{3},L_{4},...,L_{3Fn}$ can be defined and they are non-empty
set. Forall $o_{1},o_{2}\in L_{3Fn}$, as $S(o_{1},\pi)\leq S(o_{2},\pi)$ and
$S(o_{2},\pi)\leq S(o_{1},\pi)$ by lexicographical order,
$S(o_{2},\pi)=S(o_{1},\pi)$ and thus $o_{1}(v)=o_{2}(v)$. If v has no
symmetry, $\forall o\in\text{O}(3),o\neq I,o(v)\neq v$,
$o_{1}(v)=o_{2}(v)\Longrightarrow o_{1}=o_{2}$. Therefore, $L_{3Fn}$ contains
a unique element $o^{(0)}_{v}$ and $\min_{o\in\text{O}(3)}S(o,\pi)$ is unique.

As $\Pi_{n}$ is a finite set, if $\min_{o\in\text{O}(3)}S(o,\pi)$ exist for
all $\pi\in\Pi_{n}$, $\min_{\pi\in\Pi_{n}}\min_{o\in\text{O}(3)}S(o,\pi)$ must
exist. Therefore the minimal sequence exists. As $\leq_{l}$ is a linear order,
the minimal sequence is unique. With the unique sequence, the unique
representative can be determined.

(2) Then we prove there does not exist $o\in\text{O}(3),o\neq
I,o(LE_{i})=LE_{i}$ if the frame has full rank.

The frame $\vec{E}$ is a function of local environment. If there exists

$o\in\text{O}(3),o(LE_{i})=LE_{i}.$

Then $\vec{E}(o(LE_{i}))=\vec{E}(LE_{i})o^{T}=\vec{E}(LE_{i})$.

As $\vec{E}$ is a invertible matrix, $o=I$. Therefore, $LE_{i}$ has no
symmetry.

∎

##  Appendix B Derivation of the message passing section for frame

The framework proposed by Villar et al. [2021] is

$h_{n}(\\{\vec{m}_{i1},\vec{m}_{i2},\vec{m}_{i2},...,\vec{m}_{in}\\})=\sum_{j=1}^{n}g(\vec{m}_{ij},\\{\vec{m}_{i1},...,\vec{m}_{in}\\}-\\{\vec{m}_{ij}\\})\circ\vec{m}_{ij},$
(39)

where $h_{n}$ is the aggregation function for $n$ messages. $g$ is a
$\text{O}(3)$-invariant functions. We can further reform it.

$h_{n}(\\{\\{\vec{m}_{i1},\vec{m}_{i2},...,\vec{m}_{in}\\}\\})=\sum_{j=1}^{n}g_{1}^{(n)}(g_{2}^{(n)}(\vec{m}_{ij}),h_{n-1}(\\{\vec{m}_{i1},\vec{m}_{i2},...,\vec{m}_{i,n}\\}-\\{\vec{m}_{ij}\\}))\vec{m}_{ij},$
(40)

where $g_{1}^{(n)},g_{2}^{(n-1)}$ are two $\text{O}(3)$-invariant functions.
With this equation, we can recursively build $n$ messages aggregation function
$h_{n}$ with $h_{n-1}$. Its universal approximation power has been proved in
[Villar et al., 2021].

However, as they can have varied numbers of neighbors, different nodes have to
use different aggregation functions, which is hard to implement. Therefore, we
assume that the aggregation function keeps the same form for different sizes
of the message multiset and desert the recursive term $h_{n-1}$.

$h_{n}(\\{\\{\vec{m}_{i1},\vec{m}_{i2},...,\vec{m}_{in}\\}\\})=\sum_{j=1}^{n}g(\vec{m}_{ij})\vec{m}_{ij}.$
(41)

The message $\vec{m}_{ij}$ can have the form
$\text{Concatenate}(1,s)\circ\vec{r}_{ij}$ in Proposition 4.1. As $g$ is a
invariant function, we can further simplify Equation 41.

$h_{n}(\\{\\{\vec{m}_{i1},\vec{m}_{i2},...,\vec{m}_{in}\\}\\})=\sum_{j=1}^{n}g^{\prime}(r_{ij},s_{j})\frac{\vec{r}_{ij}}{r_{ij}},$
(42)

where $g^{\prime}$ is a function mapping scalars to scalars.

##  Appendix C Are orthonormal vectors necessary?

Usually, frame means a set of orthonormal vectors, while $\vec{E}_{i}$ is just
some vectors. A straightforward method is to use the Gram-Schmidt process.
However, it leads to exploding gradients, which can be explained as follows.

$\displaystyle v$ $\displaystyle=LQ,$ (43) $\displaystyle\frac{\partial
l}{\partial v}$ $\displaystyle=(L^{-1})^{T}[\frac{\partial l}{\partial
Q}-\text{copyltu}(\frac{\partial l}{\partial Q}Q^{T})Q],$ (44)

where $Q\in{\mathbb{R}}^{3\times 3}$ is the orthonormal vectors,
$L\in{\mathbb{R}}^{3\times 3}$ is the coefficients of the Gram-Schmidt
process, $l$ is the loss function, copyltu means generates a symmetric matrix
by copying the lower triangle to the upper triangle. Note that $L$ can be a
singular matrix, and $(L^{-1})^{T}$ can make the gradient of loss explode.
Moreover, the analysis before does not require the orthonormality of the
frame. Therefore, we use $\vec{E}_{i}$ as the local frame.

Figure 4: The architecture of GNN-LF. (a) The full architecture of GNN-LF
contains four parts: an embedding block, a geo2filter block, message passing
(MP) layers, and an output module. Embedding block consists of an embedding
layer converting kind of atoms to learnable tensors and a neighborhood
embedding block proposed by Thölke and Fabritiis [2022]. (b) The geo2filter
block builds a graph with the coordinates of atoms, passes messages to produce
local frames, projects vector features onto the frames, and uses edge
invariant features to produce edge filters. (c) A message passing layer
filters atom representations with edge filters to produce messages and
aggregates these messages to update atom embeddings. (d) The projection block
produces $d^{1},d^{2},d^{3}$ and concatenates them.

##  Appendix D Archtecture of GNN-LF

Following Thölke and Fabritiis [2022], we also use a neighborhood embedding
block which aggregates neighborhood information as the initial atom feature.

$s_{i}^{(0)}=\text{Emb}_{1}(z_{i})+\sum_{r_{ij}<r_{c}}Emb_{2}(z_{j})\odot
f(s^{(e)}_{ij}).$ (45)

where $\text{Emb}_{1}$ and $\text{Emb}_{2}$ are two embedding layers and $f$
is the filter function.

These special functions are proposed by previous methods [Behler and
Parrinello, 2007, Unke and Meuwly, 2019].

$\text{cutoff}(r)=\begin{cases}\frac{1}{2}(1+\cos{\frac{\pi
r}{r_{c}}}),r<r_{c}\\\ 0,r>r_{c}\end{cases}$ (46)

$\text{rbf}_{k}(r_{ij})=e^{-\beta_{k}(\exp(-r_{ij})-\mu_{k})^{2}},$ (47)

where $\beta_{k}$, $\mu_{k}$ are coefficients of the $k^{\text{th}}$ basis.

The full architecture is shown in Figure 4. For PES tasks, the output module
is a sum pooling and a linear layer. Other scalar prediction tasks can also
use this module. However, on the QM9 dataset, we design special output modules
for two properties. For dipole moment $\mu$, given node representations
$[s_{i}|i=1,2,...,N]$ and atom coordinates $[\vec{r}_{i}|i=1,2,...,N]$, our
prediction is as follows.

$\hat{\mu}=|\sum_{i}(q_{i}-\text{average}_{j}(q_{j}))\vec{r}_{i}|,$ (48)

where $q_{i}\in{\mathbb{R}}$, the prediction of charge, is function of
$s_{i}$. We use a linear layer to convert $s_{i}$ to $q_{i}$. As the whole
molecule is electroneutral, we use $q_{i}-\text{average}_{j}(q_{j})$.

For electronic spatial extent $\langle R^{2}\rangle$, we make use of atom mass
(known constants) $[m_{i}|i=1,2,...,N]$. The output module is as follows.

$\displaystyle\vec{r}_{c}$
$\displaystyle=\frac{\sum_{i}m_{i}\vec{r}_{i}}{\sum_{i}m_{i}}$ (49)
$\displaystyle\hat{\langle R^{2}\rangle}$
$\displaystyle=|\sum_{i}x_{i}|\vec{r}_{i}-\vec{r}_{c}|^{2}|,$ (50)

where $x_{i}\in{\mathbb{R}}$ is a scalar feature of node $i$. We also use a
linear layer to convert $s_{i}$ to $x_{i}$.

Table 4: The inference time and the GPU memory consumption of f random batches of 32 molecules (16 molecules for GemNet) from the MD17 dataset. The format is inference time in ms/GPU memory comsumption in MB. | DimeNet | GemNet | torchmd | GNN-LF | NoShare
---|---|---|---|---|---
number of parameters | 2.1 M | 2.2 M | 1.3 M | 0.8M$\sim$1.3M | 2.4M$\sim$5.3M
aspirin | 133/5790 | 612/15980 | 32/2065 | 10/279 | 22/883
benzene |  94/1831 | 393/3761 | 33/918 |  8/95 | 11/213
ethanol |  95/784 | 344/1565 | 32/532 |  8/54 | 11/115
malonaldehyde |  88/784 | 355/1565 | 32/532 |  7/68 | 10/127
naphthalene | 112/4470 | 498/11661 | 32/1694 |  9/175 | 15/491
salicylic_acid |  92/3489 | 430/8182 | 34/1418 |  9/176 | 15/424
toluene | 113/3148 | 423/7153 | 45/1322 |  8/176 | 15/458
uracil | 107/1782 | 354/3735 | 32/907 |  8/99 | 14/302
average | 104/2760 | 426/6700 | 34/1174 |  9/140 | 14/377

##  Appendix E Experiment settings

Computing infrastructure. We leverage Pytorch for model development.
Hyperparameter searching and model training are conducted on an Nvidia A100
GPU. Inference times are calculated on an Nvidia RTX 3090 GPU.

Training process. For MD17/QM9 dataset, we set an upper bound (6000/1000) for
the number of epoches and use an early stop strategy which finishes training
if the validation score does not increase after 500/50 epoches. We utilize
Adam optimizer and ReduceLROnPlateau learning rate scheduler to optimize
models.

Model hyperparameter tuning. Hyperparameters were selected to minimize l1 loss
on the validation sets. The best hyperparameters selected for each model can
be found in our code in the supplement materials. For MD17/QM9, we fix the
initial lr to $1e-3/3e-4$, batch size to $16/64$, hidden dimension to $256$.
The cutoff radius is selected from $[4,12]$. The number of message passing
layer is selected from $[4,8]$. The dimension of rbf is selected from
$[32,96]$. Please refer to our code for the detailed settings.

Dataset split. We randomly split the molecule set into train/validation/test
sets. For MD17, the size of the train and validation set are $950,50$
respectively. All remaining data is used for test. For QM9: The sizes of
randomly splited train/validation/test sets are $11000,10000,10831$
respectively.

##  Appendix F Scalability

To assess the scalability of our model, we show the inference time of random
MD17 batches of $32$ molecules on an NVIDIA RTX 3090 GPU. Note that GemNet
consumes too much memory, and only batches of $16$ molecules can fit in the
GPU. Our model only takes $30\%$ time and $12\%$ space compared with the
fastest baseline. Moreover, NoShare use $260\%$ more space and $67\%$ more
computation time than GNN-LF with filter sharing technique.

##  Appendix G Passing messages across local environments

If the cutoff radius is large enough, an injective message passing layer can
encode the geometry of the whole molecule. However, for computation
efficiency, the local environment is usually used which contains only a
proportion of atoms in the molecule. Therefore, models need several message
passing layers to pass messages along a chain of neighboring atoms and across
local environments to model long-range interactions. For example, in Figure
5(a), though atom $5$ is not in the local environment of atom $3$, GNN can
pass messages from atom $5$ to atom $3$ with two message passing layers: The
first layer passes messages of atom $5$ to atom $4$, and the second layer send
messages of atom $4$, which contains information of atom $5$, to atom $3$.

(a)

(b)

Figure 5: (a) Messages can be passed across local environments. (b) Frame-
frame projection can avoid information loss in a multi-step message passing
process.

We have proved that with a frame, a single message passing layer proposed in
Proposition 4.3 can injectively encode local environments. However, some
information may get lost when the message is passed along a chain of nodes by
several message passing layers. For example, the two molecules are different
in Figure 5(b) as the orientation of atom $5$ changes, but GNN cannot
differentiate them even if each message passing layer has the form in
Proposition 4.3. It seems strange because, intuitively, atom $3$ can receive
message from atom $4$ in the second message passing process, reconstruct the
local environment of atom $4$, and thus find the position change of atom $5$.
However, such reconstruction of the local environment is only possible when
atom $3$ has information about atom $4$’s frame—the message sent from $4$ to
$3$ is invariant before and after atom $5$ changes its position, and the
position change actually alters atom $4$’s frame instead of its encoding.

Therefore, our way to solve the above information loss is very
straightforward: collecting the information of neighboring nodes’ frames in
the message passing process as in Equation 51.

$\phi(\\{(s_{j},\vec{r}_{ij},\vec{E}_{j})|r_{ij}<r_{c}\\})=\rho(\sum_{r_{ij}<r_{c}}\varphi(\text{Concatenate}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j})),$
(51)

where $\phi(\\{(s_{j},\vec{r}_{ij},\vec{E}_{j})|r_{ij}<r_{c}\\})$ aggregates
information from atom $i$’s local environment to atom $i$, $\vec{E}_{i}$ means
the frame of atom $i$, other symbols have the same meaning as in Proposition
4.3. $P_{\vec{E}_{i}}(\vec{E}_{j})$ is the frame-frame projection.

Note that a single layer defined in Equation 51 can still encode the local
environment, as extra frame-frame projection cannot lower the expressive
power.

######  Proposition G.1.

Given a frame $\vec{E}$, with suitable functions $\rho$ and $\varphi$, $\phi$
defined in Equation 51 encodes the local environment injectively.

###### Proof.

According to proposition 4.3, there exists $\rho^{\prime}$ and
$\varphi^{\prime}$ so that
$\rho(\sum_{r_{ij}<r_{c}}\varphi(\text{Concatenate}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j}))$
can encode local environment injectively. Let $\rho$ equals to
$\rho^{\prime}$,
$\varphi(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j}))=\varphi^{\prime}(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),s_{j}))$.
Then
$\phi(\\{(s_{j},\vec{r}_{ij},\vec{E}_{j})|r_{ij}<r_{c}\\})=\rho^{\prime}(\sum_{r_{ij}<r_{c}}\varphi^{\prime}(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),s_{j}))$
encodes local environment injectively. ∎

Moreover, this form can solve the information loss problem in the multi-step
message passing process. Given a molecule with atom coordinates
$\vec{r}\in{\mathbb{R}}^{N\times 3}$ and atomic features (like embedding of
atomic number) $s\in{\mathbb{R}}^{N\times d}$, let ${\mathcal{G}}$ denote the
undirected graph corresponding to the molecule. Node $i$ in ${\mathcal{G}}$
represents the atom $i$ in the molecule. ${\mathcal{G}}$ has edge $(i,j)$ iff
$r_{ij}<r_{c}$, where $r_{c}$ is the cutoff radius. Let $d({\mathcal{G}},i,j)$
denote the shortest path distance between node $i$ and $j$ in graph
${\mathcal{G}}$. The following theorem shows that multi-step message passing
can encode the full molecule injectively.

######  Theorem G.1.

If ${\mathcal{G}}$ is a connected graph and its diameter is $L$, GNN with $L$
message passing layers defined in Equation 51 can encode the
$\\{(s_{j},\vec{r}_{ij})|j\in\\{1,2,...,n\\}\\}$ injectively into the
embedding of node $i$.

###### Proof.

We use cat to represent concatenate throughout the proof. Let $N(i)_{l}$
denote $\\{j|d({\mathcal{G}},i,j)\leq l\\}$. The $l^{\text{th}}$ message
passing layer has the following form.

$s_{i}^{(l)}=\rho_{l}(\sum_{j\in
N_{1}(i)}\varphi_{l}(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j}^{(l-1)})),$
(52)

where $s_{i}^{(0)}=s_{i}$.

By enumeration on $l$, we prove that there exist $\rho_{l}$, $\varphi_{l}$ so
that $s_{i}^{(l)}=\psi(\\{\text{cat}(s_{j},P_{\vec{E}_{i}}(\vec{r}_{ij}))|j\in
N_{l}(i)\\})$.

We first define some auxiliary functions.

According to [Segol and Lipman, 2020], there exists a multiset function $\psi$
mapping a multiset of scalar to a scalar injectively. $\psi$ can have the
following form

$\psi(\\{x_{i}|i\in{\mathbb{I}}\\})=\sum_{i}\varphi(x_{i}),$ (53)

where ${\mathbb{I}}$ is some finite index set. As $\psi$ is injective, it has
an inverse function.

We define function $m,m^{\prime},m^{\prime\prime}$ to extract scalar from
concatenated node features.

$m(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j}^{(0)}))=\text{cat}(s_{j}^{(0)},P_{\vec{E}_{i}}(\vec{r}_{ij})).$
(54) $m^{\prime}(\text{cat}(s_{j},P_{\vec{E}_{i}}(\vec{r}_{ij})))=s_{j}.$ (55)
$m^{\prime\prime}(\text{cat}(s_{j},P_{\vec{E}_{i}}(\vec{r}_{ij})))=P_{\vec{E}_{i}}(\vec{r}_{ij}).$
(56)

Last but not least, there exist a function $T$ transforming coordinate
projections from one frame to another frame.

$T(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),P_{\vec{E}_{j}}(\vec{r}_{jk}))=P_{\vec{E}_{i}}(\vec{r}_{ij})+P_{\vec{E}_{j}}(\vec{r}_{jk})P_{\vec{E}_{i}}(\vec{E}_{j})=P_{\vec{E}_{i}}(\vec{r}_{ik})$
(57)

$l=1$: let $\varphi_{1}=\varphi\circ m$ , $\rho_{1}$ is identity mapping.

$l>1$: Assume for all $l^{\prime}<l$,
$s_{i}^{(l^{\prime})}=\psi(\\{\text{cat}(s_{j},P_{\vec{E}_{i}}(\vec{r}_{ij}))|j\in
N_{l^{\prime}}(i)\\})$.

$\varphi_{l}$ has the following form.

$\varphi_{l}(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j}^{(l-1)}))=\varphi(\psi(\\\
\\{\text{cat}(m^{\prime}(x),T(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),m^{\prime\prime}(x))|x\in\psi^{-1}(s_{j}^{(l-1)})\\})).$
(58)

Therefore

$\varphi_{l}(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j}^{(l-1)}))=\varphi(\psi(\\{\text{cat}(s_{k},P_{\vec{E}_{i}}(\vec{r}_{ik}))|k\in
N_{l-1}(j)\\})).$ (59)

Note $\varphi_{l}$ transforms coordinate projection from an old frame to a new
frame.

Therefore, the input of $\rho_{l}$, namely $a_{i}^{(l)}$, has the following
form.

$\displaystyle a_{i}^{(l)}$ $\displaystyle=\sum_{j\in
N(i)}\varphi_{l}(\text{cat}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j}^{(l-1)})$
(60) $\displaystyle=\sum_{j\in
N(i)}\varphi(\psi(\\{\text{cat}(s_{k},P_{\vec{E}_{i}}(\vec{r}_{ik}))|k\in
N_{l-1}(j)\\}))$ (61)
$\displaystyle=\psi(\\{\psi(\\{\text{cat}(s_{k},P_{\vec{E}_{i}}(\vec{r}_{ik}))|k\in
N_{l-1}(j)\\})|j\in N(i)\\})$ (62)

We can transform $a_{i}^{(l)}$ to a set of set of scalars with the following
function.

$\eta(a_{i}^{(l)})=\\{\psi^{-1}(s)|s\in\psi^{-1}(a_{i}^{(l)})\\}.$ (63)

Therefore,
$\eta(a_{i}^{(l)})=\\{\\{\text{cat}(s_{k},P_{\vec{E}_{i}}(\vec{r}_{ik}))|k\in
N_{l-1}(j)\\}|j\in N(i)\\}$

We can use another function $\iota$ unions scalar sets in set ${\mathbb{S}}$
to a set of scalar.

$\iota({\mathbb{S}})=\bigcup_{s\in{\mathbb{S}}}s.$ (64)

$\rho_{l}$ has the following form.

$\rho_{l}(a_{i}^{(l)})=\psi\circ\iota\circ\eta(a_{i}^{(l)}).$ (65)

Therefore, the output is

$\rho_{l}(a_{i}^{(l)})=\psi(\\{\text{cat}(s_{k},P_{\vec{E}_{i}}(\vec{r}_{ik}))|k\in
N_{l}(i)\\}\\})$ (66)

Therefore, $\forall l\in{\mathbb{N}}$, there exists $\rho_{l}$, $\varphi_{l}$
so that
$s_{i}^{(l)}=\psi(\\{\text{cat}(s_{j},P_{\vec{E}_{i}}(\vec{r}_{ij}))|j\in
N_{l}(i)\\})$.

As $L$ is the diameter of ${\mathcal{G}}$,
$s_{i}^{(L)}=s_{i}^{(l)}=\psi(\\{\text{cat}(s_{j},P_{\vec{E}_{i}}(\vec{r}_{ij}))|j\in
N_{L}(i)\\})=\psi(\\{\text{cat}(s_{j},P_{\vec{E}_{i}}(\vec{r}_{ij}))|j\in\\{1,2,...,n\\}\\})$.
As $\psi$ is an injective function, GNN with $L$ message passing layers
defined in Equation 51 can encode the
$\\{(s_{i},\vec{r}_{ij})|j\in\\{1,2,...,n\\}\\}$ injectively to $s_{i}^{(L)}$.
∎

##  Appendix H Full ablation study of $d^{2}$

Table 5: Ablation results on the MD17 dataset. Units: energy (E) (kcal/mol) and forces (F) (kcal/mol/Å). GNN-LF does not use $d^{2}$ for some molecules, so the corresponding NoDir2 column is empty.  Molecule |  | Dir2 | NoDir2
---|---|---|---
Aspirin | E | 0.1385 | 0.1342
F | 0.1967 | 0.2018
Benzene | E | 0.0708 | 0.0686
F | 0.1503 | 0.1506
Ethanol | E | 0.0516 | 0.0522
F | 0.0798 | 0.0817
Malonaldehyde | E | 0.0773 | 0.0764
F | 0.1260 | 0.1259
Naphthalene | E | 0.1182 | 0.1136
F | 0.0522 | 0.0550
Salicylic acid | E | 0.1080 | 0.1081
F | 0.1005 | 0.1048
Toluene | E | 0.0931 | 0.0930
F | 0.0530 | 0.0543
Uracil | E | 0.1060 | 0.1037
F | 0.0758 | 0.0751

The comparison between GNN-LF with the $d^{2}$ term in Equation 9 (Dir2) and
GNN-LF without $d^{2}$ term (NoDir2) is shown in Table 5. $d^{2}$ term boosts
accuracy on $8/16$ targets and leads to only $0.2\%$ lower loss on average.
Such negligible accuracy improvement verifies that projection of direction on
the target node and frame is enough for high expressive power. This result is
an indicator of the success of our theory.

##  Appendix I Existing methods using frame

Though some previous works [Du et al., 2022, Satorras et al., 2021, Huang et
al., 2022, Luo et al., 2021] also use the term "frame" or designs similar to
"frame", they are very different methods from ours.

The primary motivation of our work is to get rid of vector (a.k.a.
equivariant) representation for higher and provable expressive power, simpler
architecture, and better scalability. We only use vectors in the frame
generation and projection process. After projection, all the remaining parts
of our model only operates on scalars. In contrast, existing works [Du et al.,
2022, Satorras et al., 2021, Huang et al., 2022, Luo et al., 2021] still use
both vector and scalar representations, resulting in extra complexity even
after using frame. For example, functions for vectors and scalars still need
to be defined separately, and complex operation is needed to mix the
information contained in the two kinds of representations. In addition, our
model can beat the representative methods of this kind in both accuracy and
scalability on potential energy surface prediction tasks.

Other than the different primary motivation, our model has an entirely
different architecture from existing ones.

  1. 1.

Generating frame: ClofNet [Du et al., 2022] produces a frame for each pair of
nodes and use cross product to produce the frame. Both EGNN [Satorras et al.,
2021] and GMN [Huang et al., 2022] use coordinate embeddings which are
initialized with coordinates. Luo et al. [2021] initializes the frame with
zero. Then these models use some schemes to update the frame. Our model uses a
novel message passing scheme to produce frames and will not update it, leading
to simpler architecture and low computation overhead.

  2. 2.

Scalarization: Existing models [Du et al., 2022, Satorras et al., 2021, Huang
et al., 2022, Luo et al., 2021] only project vector features onto the frame,
while we also use frame-frame projection, which is verified to be critical
both experimentally and theoretically.

  3. 3.

Message passing layer: Existing models [Du et al., 2022, Satorras et al.,
2021, Huang et al., 2022, Luo et al., 2021] all use both scalar and vector
features and pass both scalar and vector messages, which needs to mix scalars
and vectors, update scalars, and update vectors, while our model only uses
scalars, resulting in an entirely different and much simpler design with
significantly higher scalability.

  4. 4.

Our designing tricks, including: message passing scheme to produce frame,
filter decomposition, and filter sharing, are not used in [Du et al., 2022,
Satorras et al., 2021, Huang et al., 2022, Luo et al., 2021]. Our experiments
and ablation study verified their effectiveness.

Furthermore, existing models use different groups to describe symmetry. Du et
al. [2022], Luo et al. [2021] designs $\text{SO}(3)$-equivariant model, while
our model is $\text{O}(3)$-equivariant. We emphasize that this is not a
contraint of our model but a requirement of the task. As most target
properties of molecules are $\text{O}(3)$-equivariant (including energy and
force we aim to predict), our model can fully describe the symmetry.

Our theoretical analysis is also novel. Satorras et al. [2021], Huang et al.
[2022], Luo et al. [2021] have no theoretical analysis of expressive power. Du
et al. [2022]’s analysis is primarily based on the framework of Dym and Maron
[2021], which is further based on the polynomial approximation and the group
representation theory. The conclusion is that a model needs many message
passing layers to approximate high-order polynomials and achieve universality.
Our theoretical analysis gets rid of polynomial and group representation and
provides a much simpler analysis. We also prove that one message passing layer
proposed in our paper are enough to be universal.

In summary, although also using “frame”, our work is significantly different
to any existing work in either method, theory, or task.

##  Appendix J Expressive power with symmetric input

We use the symbol in Equation 51 SchNet’s message can be formalized as
follows.

$\phi(\\{(s_{j},\vec{r}_{ij},\vec{E}_{j})|r_{ij}<r_{c}\\})=\rho(\sum_{r_{ij}<r_{c}}\varphi(\text{Concatenate}(s_{j},r_{ij})).$
(67)

In implementation, GNN-LF has the following form.

$\phi(\\{(s_{j},\vec{r}_{ij},\vec{E}_{j})|r_{ij}<r_{c}\\})=\rho(\sum_{r_{ij}<r_{c}}\varphi(\text{Concatenate}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j},r_{ij})).$
(68)

Therefore, for all input molecules, GNN-LF is at least as expressive as
SchNet.

######  Theorem J.1.

$\forall L\in{\mathbb{N}}^{+}$, for all $L$ layer SchNet, there exists a $L$
layer GNN-LF produce the same output for all input molecule.

###### Proof.

Let the following equation denote the $l^{\text{th}}$ layer SchNet.

$\phi^{\prime}_{l}(\\{(s_{j}^{(l)},\vec{r}_{ij},\vec{E}_{j})|r_{ij}<r_{c}\\})=\rho^{\prime}_{l}(\sum_{r_{ij}<r_{c}}\varphi^{\prime}_{l}(\text{Concatenate}(s_{j}^{(l-1)},r_{ij})).$
(69)

Let $\rho_{l}=\rho^{\prime}_{l}$,
$\varphi_{l}(\text{Concatenate}(P_{\vec{E}_{i}}(\vec{r}_{ij}),P_{\vec{E}_{i}}(\vec{E}_{j}),s_{j},r_{ij}))=\varphi^{\prime}_{l}(\text{Concatenate}(s_{j},r_{ij}))$,
which neglects the projection input.

Let the $l^{\text{th}}$ layer of GNN-LF have the following form.

$\phi_{l}(\\{(s_{j}^{(l)},\vec{r}_{ij},\vec{E}_{j})|r_{ij}<r_{c}\\})=\rho_{l}(\sum_{r_{ij}<r_{c}}\varphi_{l}(\text{Concatenate}(s_{j}^{(l-1)},r_{ij})).$
(70)

This GNN-LF produces the same output as the SchNet.

∎

##  Appendix K How to overcome the frame degeneration problem.

As shown in Theorem 4.1, if the frame is $\text{O}(3)$-equivariant, no matter
what scheme is used, the frame will degenerate when the input molecule is
symmetric. In other words, the projection degeneration problem roots in the
symmetry of molecule. Therefore, we try to break the symmetry by assigning
node identity features $s^{\prime}$ to atoms. The $i^{\text{th}}$ row of
$s^{\prime}$ is $i$. We concatenate $s$ and $s^{\prime}$ as the new node
feature $\tilde{s}\in{\mathbb{R}}^{N\times(d+1)}$. Let $\eta$ denote a
function concatenating node feature $s$ and node identity features
$s^{\prime}$, $\eta(s)=\tilde{s}$. Its inverse function removes the node
identity $\eta^{-1}(\tilde{s})=s$.

In this section, we assume that the cutoff radius is large enough so that
local environments cover the whole molecule. Let $[n]$ denote the sequence
$1,2,...,n$. Let $s\in{\mathbb{R}}^{N\times d}$ denote the scalar atomic
features, $\vec{r}\in{\mathbb{R}}^{N\times 3}$ denote the 3D coordinates of
atoms, and $\vec{r}_{i}$, the $i^{\text{th}}$ row of $\vec{r}$, denote the
coordinates of atom $i$. Let $\vec{r}-\vec{r}_{i}$ denote an $N\times 3$
matrix whose $j^{\text{th}}$ row is $\vec{r}_{j}-\vec{r}_{i}$. We assume that
$N>1$ throughout this section.

Now each atom in the molecule has a different feature. The frame generation is
much simpler now.

######  Proposition K.1.

With node identity features, there exists an O(3)-equivariant function mapping
the local environment $LE_{i}=\\{(\tilde{s}_{i},\vec{r}_{ij})|j\in[N]\\}$ to a
frame $\vec{E}_{i}\in{\mathbb{R}}^{3\times 3}$, and the first
$\text{rank}(\vec{E}_{i})$ rows of $\vec{E}_{i}$ form an orthonormal basis of
$\text{span}(\\{\vec{r}_{ij}|j\in[N]\\})$ while other rows are zero.

###### Proof.

For node $i$, we can use the following procedure to produce a frame.

Initialize $\vec{E}_{i}$ as an empty matrix. For j in $[1,2,3,...,n]$, if
$\vec{r}_{ij}$ is linearly independent to row vectors in $\vec{E}_{i}$, add
$\vec{r}_{ij}$ as a row vector of $\vec{E}_{i}$.

Therefore, when the procedure finishes, the row vectors of $\vec{E}$ form a
maximal linearly independent system of $\\{\vec{r}_{ij}|j\in[N]\\}$.

Then, we use the Gram-Schmidt process to orthonormalize the non-empty row
vectors in $\vec{E}$, and then use zero to fill the empty rows in $\vec{E}$ to
form a ${3\times 3}$ matrix. Therefore, the first $\text{rank}(\vec{E}_{i})$
rows of $\vec{E}_{i}$ are orthonormal, and can linearly express all vector in
$\\{\vec{r}_{ij}|j\in[N]\\}$. In other words, the first
$\text{rank}(\vec{E}_{i})$ rows of $\vec{E}_{i}$ form an orthonormal basis of
$\text{span}(\\{\vec{r}_{ij}|j\in[N]\\})$.

Note that $\vec{r}_{ij}$, $\vec{0}$ are $\text{O}(3)$-equivariant vectors.
Therefore, the frame produced with this scheme is $\text{O}(3)$-equivariant. ∎

With the frame, GNN-LF has the universal approximation property.

######  Proposition K.2.

Assuming that the node identity features are used, and the frame is produced
by the method in Proposition K.1. For all $\text{O}(3)$-invariant and
translation-invariant functions $f(s,\vec{r})$, $f$ can be written as a
function of the embeddings of node $1$ produced by one message passing layer
proposed in Proposition 4.3.

###### Proof.

Let $eye_{r}\in{\mathbb{R}}^{3\times 3}$ denote a diagonal matrix whose first
$r$ diagonal elements are $1$ and others are $0$.

With node identity features and method in Proposition K.1, the first
$\text{rank}(\vec{E}_{i})$ rows of $\vec{E}_{i}$ form an orthonormal basis of
$\text{span}(\\{\vec{r}_{ij}|j\in[N]\\})$ while other rows are zero. Especial,
all vectors in $\\{\vec{r}_{1j}|j\in[N]\\}$ can be written as linear
combination of rows in $\vec{E}_{1}$, $\vec{r}_{1j}=w_{1j}\vec{E}_{j}$.
Therefore, the projection operation
$P_{\vec{E}_{1}}:\\{\vec{r}_{1j}|j\in[N]\\}\to\\{(\vec{r}_{1j})\vec{E}_{1}^{T}|j\in[N]\\}$
is injective, as
$P_{\vec{E}_{1}}(\vec{r}_{1j})\vec{E}_{1}=w_{1j}\vec{E}_{1}\vec{E}_{1}^{T}\vec{E}_{1}=w_{1j}eye_{\text{rank}(\vec{E}_{1})}\vec{E}_{1}=w_{1j}\vec{E}_{1}=\vec{r}_{1j}$.

According to the proof of Proposition 4.3, there exists injective function
$\phi$ so that node embeddings
$z_{1}=\phi(\\{\tilde{s}_{j},P_{\vec{E}_{1}}(\vec{r}_{1j})|j\in[N]\\})$. Note
that both $\vec{E}_{1}$ and $z_{1}$ are functions of $LE_{1}$.

Let $\tau$ denote a function
$(z_{1},\vec{E}_{1})=\tau(\\{(\tilde{s}_{j},\vec{r}_{1j})|j\in[N]\\})$.
$\forall o\in\text{O}(3)$,
$(z_{1},\vec{E}_{1}o^{T})=\tau(\\{(\tilde{s}_{j},\vec{r}_{1j}o^{T})|j\in[N]\\})$.
Moreover, $\tau$ is also a invertible function because
$\\{(\tilde{s}_{j},\vec{r}_{1j})|j\in[N]\\}=\\{(s,p\vec{E}_{1})|(s,p)\in\phi^{-1}(z_{1})\\}$.

Because the last column of $\tilde{s}$ is the node identity feature, there
exists an bijective function $\psi$ converting set of features to matrix of
features.
$\psi(\\{(\tilde{s}_{j},\vec{r}_{1j})|j\in[N]\\})=(s,-(\vec{r}-\vec{r}_{1}))$.
Intuitively, it puts the features of atom with node identity $i$ to the
$i^{\text{th}}$ row of feature matrix. Similarly,
$\psi^{\prime}(\\{(\tilde{s}_{j},P_{\vec{E}_{1}}(\vec{r}_{1j}))|j\in[N]\\})=(s,-(\vec{r}-\vec{r}_{1})\vec{E}_{1}^{T})$.

As $f$ is a translation- and O(3)-invariant function,
$f(s,\vec{r})=f(s,\vec{r}-\vec{r}_{1})=f(s,-(\vec{r}-\vec{r}_{1}))=f(\psi(\\{(\tilde{s}_{j},\vec{r}_{1j})|j\in[N]\\}))$.
Let $g=f\circ\psi\circ\tau^{-1}$, $g(\vec{E}_{1},z_{1})=f(s,\vec{r})$.
Moreover, $\forall o\in\text{O}(3)$,
$g(\vec{E}_{1}o^{T},z_{1})=f(s,\vec{r}o^{T})=f(s,\vec{r})$.

Let $\text{extend}(E_{i})\in O(3)$ denote any matrix whose first
$\text{rank}(E_{i})$ rows equals to $E_{i}$’s first rows. Therefore,
$f(s,\vec{r})=f(s,\vec{r}\text{extend}(\vec{E}_{1})^{T})=g(\vec{E}_{1}\text{extend}(\vec{E}_{1})^{T},z_{1})=g(eye_{\text{rank}(\vec{E}_{1})},z_{1})=g^{\prime}(\text{rank}(\vec{E}_{1}),z_{1})$.

Note that
$\text{rank}(\vec{E}_{1})=\text{rank}(\vec{r}-\vec{r}_{1})=\text{rank}(P_{\vec{E}_{1}}(\vec{r}-\vec{r}_{1}))=\text{rank}(\iota\circ\psi^{\prime}\circ\phi^{-1}(z_{i}))$,
where $\iota$ is a selection function:
$\iota(z,-(\vec{r}-\vec{r}_{0})\vec{E}_{1}^{T})=-(\vec{r}-\vec{r}_{0})\vec{E}_{1}^{T}$.
Therefore,
$f(s,\vec{r})=g^{\prime}(\text{rank}(\iota\circ\phi^{-1}(z_{1})),z_{1})=g^{\prime\prime}(z_{1})$.
∎

For simplicity, let function $\varphi$ denote GNN-LF with node identity
features (including adding node identity feature, generating frame, and a
message passing layer proposed in Proposition 4.3), $\varphi(z,\vec{r})$ is
the embeddings of node $1$.

Node identity features help avoiding expressive power loss caused by frame
degeneration. However, GNN-LF’s output is no longer permutation invariant.
Therefore, we use the relational pooling method [Puny et al., 2022], which
introduces extra computation overhead and keeps the permutation invariance.

To illustrate this method, we first define some concepts. Function
$\pi:[n]\to[n]$ is a permutation iff it is bijective. All permutation on $[n]$
forms the permutation group $S_{n}$. We compute the output of all possible
atom permutations and average them, in order to keep permutation invariance.
We define the permutation of matrix here: for all matrix of shape $N\times m$,
$\forall\pi\in S_{N}$, the $i^{\text{th}}$ row of $\pi(M)$ equals to the
$(\pi^{-1}(i))^{\text{th}}$ row of $M$.

######  Proposition K.3.

For all $\text{O}(3)$-invariant, permutation-invariant and translation
invariant function $f(s,\vec{r})$, there exists GNN-LF $\varphi$ and some
function $g$, with which $\frac{1}{N!}\sum_{\pi\in
S_{N}}g(\varphi(\pi(s),\pi(\vec{r})))$ is permutation invariant and equals to
$f(s,\vec{r})$.

###### Proof.

Define a "frame" (defined in Definition 1 in [Puny et al., 2022]) $F:V\to
2^{S_{n}}$, where $V$ is the embedding space. $\forall v\in V$, $F(v)=S_{n}$.
So the relational pooling of GNN-LF with node identity features $\langle
g\circ\varphi\rangle_{F}(s,\vec{r})=\frac{1}{N!}\sum_{\pi\in
S_{N}}g(\varphi(\pi(s)),\pi(\vec{r}))$. Note that the permutation operation
$\pi$ and O(3) operation $o$ commute: $\pi(\vec{r}o^{T})=\pi(\vec{r})o^{T}$.
According to Theorem 2 in [Puny et al., 2022], $\langle
g\circ\varphi\rangle_{F}$ is permutation invariant.

According to Theorem 4 in [Puny et al., 2022], if there exist function
$g^{\prime}$ and $\varphi^{\prime}$ so that
$g^{\prime}\circ\varphi^{\prime}=f$ (the existence is shown in Proposition
K.2), there will also exist GNN-LF $\varphi$ and function $g$, so that
$\langle g\circ\varphi\rangle_{F}(s,\vec{r})=f(s,\vec{r})$. ∎

Therefore, we can completely solve the frame degeneration problem with the
relational pooling trick and node identity features. However, the time
complexity is up to $O(N!N^{2})$, so we only analyze this method
theoretically.

##  Appendix L Why is global frame more likely to degenerate than local frame?

Let $[N]$ denote the sequence $1,2,...,N$. $N$ is the number of atoms in the
molecule.

We first consider when local frame degenerates. As shown in Theorem 4.1, the
degeneration happens if and only if the local environment is symmetric under
some orthogonal transformations.
$\text{rank}(\vec{E}_{i})<3\Leftrightarrow\exists o\in\text{O}(3),o\neq
I,\\{(s_{i},\vec{r}_{ij}o^{T})|r_{ij}<r_{c}\\}=\\{(s_{i},\vec{r}_{ij})|r_{ij}<r_{c}\\}$.

The global frame has the following form,

$\vec{E}=\sum_{i=1}^{N}\vec{E}_{i}.$ (71)

We first prove some properties of $\vec{E}$ function.

######  Proposition L.1.

$\vec{E}$ is an $\text{O}(3)$-equivariant, translation-invariant, and
permutation-invariant function.

###### Proof.

$\text{O}(3)$-equivariance: $\forall o\in O(3)$,
$\vec{E}_{i}(s,\vec{r}o^{T})=\vec{E}_{i}(s,\vec{r})o^{T}$. Therefore,

$\vec{E}(s,\vec{r}o^{T})=\sum_{i=1}^{N}\vec{E}_{i}(s,\vec{r}o^{T})=(\sum_{i=1}^{N}\vec{E}_{i}(s,\vec{r}))o^{T}=\vec{E}(s,\vec{r})o^{T}.$
(72)

Translation-invariance: For all translation $\vec{t}\in{\mathbb{R}}^{3}$, let
$\vec{r}+\vec{t}$ denote a matrix of shape $(N,3)$ whose $i^{\text{th}}$ row
is $\vec{r}_{i}+\vec{t}$. As $\vec{E}_{i}$ is a function of
$\vec{r}_{i}-\vec{r}_{j}=\vec{r}_{i}+\vec{t}-(\vec{r}_{j}+\vec{t})$,
$\vec{E}_{i}(z,\vec{r}+t)=\vec{E}_{i}(z,\vec{r})$. Therefore,

$\vec{E}(s,\vec{r}+t)=\sum_{i=1}^{N}\vec{E}_{i}(s,\vec{r}+t)=\sum_{i=1}^{N}\vec{E}_{i}(s,\vec{r})=\vec{E}(s,\vec{r}).$
(73)

Permutation-invariance: for all permutation $\pi\in S_{n}$,
$\pi(\vec{r})_{i}=\pi(\vec{r})_{\pi^{-1}(i)}$.

$\displaystyle\vec{E}(\pi(s),\pi(\vec{r}))$
$\displaystyle=\sum_{i=1}^{N}\vec{E}_{i}(\pi(s),\pi(\vec{r}))$ (74)
$\displaystyle=\sum_{i=1}^{N}\vec{E}_{i}(\\{(\pi(s)_{j},\pi(\vec{r})_{i}-\pi(\vec{r})_{j}|~{}|\pi(\vec{r})_{i}-\pi(\vec{r})_{j}|<r_{c}\\})$
(75)
$\displaystyle=\sum_{i=1}^{N}\vec{E}_{i}(\\{(s_{\pi^{-1}(j)},\vec{r}_{{\pi^{-1}(i)}}-\vec{r}_{\pi^{-1}(j)})|r_{\pi^{-1}(i)\pi^{-1}(j)}<r_{c}\\})$
(76)
$\displaystyle=\sum_{i=1}^{N}\vec{E}_{i}(\\{(s_{j},\vec{r}_{i}-\vec{r}_{j}|r_{ij}<r_{c}\\})$
(77) $\displaystyle=\vec{E}(s,\vec{r}).$ (78)

∎

Then we prove a sufficient condition for global frame degeneration.

######  Proposition L.2.

$\text{rank}(\vec{E}<3)$ if there exists $\vec{t}\in{\mathbb{R}}^{3}$ and
$o\in\text{O}(3),o\neq I$ such that
$\\{(s_{i},\vec{r}_{i}-\vec{t})|i\in[N]\\}=\\{(s_{i},(\vec{r}_{i}-\vec{t})o^{T})|i\in[N]\\}$.

###### Proof.

As $\vec{E}$ is a permutation invariant function,
$\vec{E}=\vec{E}(\\{(s_{i},\vec{r}_{i})|i\in[N]\\})$.

As $\vec{E}$ is a translation-invariant and $\text{O}(3)$-equivariant
function,

$\vec{E}(\\{(s_{i},(\vec{r}_{i}-\vec{t})o^{T})|i\in[N]\\})=\vec{E}(\\{(s_{i},(\vec{r}_{i}-\vec{t}))|i\in[N]\\})o^{T}\\\
=\vec{E}(\\{(s_{i},\vec{r}_{i}))|i\in[N]\\})o^{T}.$ (79)

Therefore, under the condition
$\\{(s_{i},\vec{r}_{i}-\vec{t})|i\in[N]\\}=\\{(s_{i},(\vec{r}_{i}-\vec{t})o^{T})|i\in[N]\\}$,
we have

$\displaystyle\vec{E}(\\{(s_{i},\vec{r}_{i}))|i\in[N]\\})o^{T}=\vec{E}(\\{(s_{i},\vec{r}_{i}))|i\in[N]\\}),$
(80) $\displaystyle\implies$
$\displaystyle\vec{E}(\\{(s_{i},\vec{r}_{i}))|i\in[N]\\})(I-o^{T})=0.$ (81)

Therefore, $\text{rank}(\vec{E})+\text{rank}(I-o^{T})-3\leq 0$. As $I\neq
o^{T}$, $\text{rank}(I-o^{T})>0$, $\text{rank}(\vec{E})<3$. ∎

The main difference between the degeneration conditions is the choice of
origin. The local frame of atom $i$ degenerates when the molecule is symmetric
with atom $i$ as the origin point, while the global frame degenerates if the
molecule is symmetric with any origin point. Therefore, the global frame is
more likely to degenerate.

######  Corollary L.1.

Assume the cutoff radius is large enough so that local environments contain
all atoms. If there exists $i$, $\text{rank}(\vec{E}_{i})<3$, then
$\text{rank}(\vec{E})<3$.

###### Proof.

As $\text{rank}(\vec{E}_{i})<3$, $\exists o\in\text{O}(3),o\neq
I,\\{(s_{j},(\vec{r}_{i}-\vec{r}_{j})o^{T})|j\in[N]\\}=\\{(s_{j},(\vec{r}_{i}-\vec{r}_{j}))|j\in[N]\\}$.

Therefore,

$\displaystyle\\{(s_{j},-(\vec{r}_{i}-\vec{r}_{j})o^{T})|j\in[N]\\}=\\{(s_{j},-(\vec{r}_{i}-\vec{r}_{j}))|j\in[N]\\}$
(82) $\displaystyle\implies$
$\displaystyle\\{(s_{j},(\vec{r}_{j}-\vec{r}_{i})o^{T})|j\in[N]\\}=\\{(s_{j},\vec{r}_{j}-\vec{r}_{i}))|j\in[N]\\}$
(83)

Let $\vec{t}=\vec{r}_{i}$, according to Proposition L.2,
$\text{rank}(\vec{E})<3$. ∎

Therefore, when the cutoff radius is large enough, the global frame will also
degenerate if some local frame degenerates.