[1]Alex Morehead [1]Jianlin Cheng

[1]Electrical Engineering & Computer Science, University of Missouri-Columbia,
W1024 Lafferre Hall, Columbia, 65211, Missouri, United States of America

# Geometry-Complete Perceptron Networks for
3D Molecular Graphs

<EMAIL_ADDRESS><EMAIL_ADDRESS>*

###### Abstract

The field of geometric deep learning has had a profound impact on the
development of innovative and powerful graph neural network architectures.
Disciplines such as computer vision and computational biology have benefited
significantly from such methodological advances, which has led to
breakthroughs in scientific domains such as protein structure prediction and
design. In this work, we introduce GCPNet, a new geometry-complete,
SE(3)-equivariant graph neural network designed for 3D molecular graph
representation learning. Rigorous experiments across four distinct geometric
tasks demonstrate that GCPNet’s predictions (1) for protein-ligand binding
affinity achieve a statistically significant correlation of 0.608, more than
5% greater than current state-of-the-art methods; (2) for protein structure
ranking achieve statistically significant target-local and dataset-global
correlations of 0.616 and 0.871, respectively; (3) for Newtownian many-body
systems modeling achieve a task-averaged mean squared error less than 0.01,
more than 15% better than current methods; and (4) for molecular chirality
recognition achieve a state-of-the-art prediction accuracy of 98.7%, better
than any other machine learning method to date. The source code, data, and
instructions to train new models or reproduce our results are freely available
at https://github.com/BioinfoMachineLearning/GCPNet.

###### keywords:

Geometric deep learning, equivariance, graph neural networks, computational
biology, molecules, proteins

##  1 Introduction

(a)

PDB ID: 8EXU

(b)

PDB ID: 8GU4

(c)

PDB ID: 8HDG

Figure 1: Mol* [1] visualizations of various kinds of representative 3D
molecular structures studied in this work including (a) small molecules (i.e.,
ligands), (b) proteins, and (c) protein-ligand complexes.

Over the last several years, the field of deep learning has pioneered many new
methods designed to process graph-structured inputs. Being a ubiquitous form
of information, graph-structured data arises from numerous sources such as the
fields of physics and chemistry, as shown in Figure 1. Moreover, the
relational nature of graph-structured data allows one to identify and
characterize topological associations between entities in large real-world
networks (e.g., social networks).

In particular, 3D data often emerges in domains such as computer vision and
can be readily described as graph-structured inputs [2]. Studies such as those
of [3, 4, 5] have demonstrated the utility of this approach to modeling 3D
data as graphs. Additionally, to process and analyze such 3D information in a
meaningful, powerful, and concise way, one must also carefully consider the
symmetries present in such data to reduce the geometric redundancies they
might present to a machine learning model [6].

Many disciplines present types of data for which 3D geometric information can
be carefully analyzed to produce meaningful and reliable predictions about the
system at hand. For example, in protein biology, knowing the 3D structure of a
protein macromolecule is a key step towards developing a deeper understanding
of its molecular function in living organisms [7]. In a context-specific
manner, similar geometric insights have been proposed in fields such as
neurobiology [8] and materials design [9]. In light of such insights, the
field of deep learning has grown to account for the importance of geometry in
the representation learning of real-world objects.

In the case of being presented with geometric (non-Euclidean domain) data,
machine learning systems have been developed to process arbitrarily-structured
inputs such as graphs and meshes [10, 11]. A subfield of geometric machine
learning, geometric deep learning, has recently received much attention from
researchers for its various success cases in using deep neural networks to
faithfully model geometric data in the form of 3D graphs and manifolds [12,
13].

Previous works in geometric deep learning [14] have explored the use of neural
networks for modeling physical systems [15, 16, 17]. Some of the earliest
neural networks applied to physical systems include convolutional networks
(CNNs) [18, 19, 20, 21], graph neural networks (GNNs) [22, 23, 24, 25], and
point cloud neural networks [26, 27, 28]. Likewise, recurrent neural networks
(RNNs) have been used to classify sequential real-world data such as speech
activity [29, 30] and have been adopted as case studies for how to develop
next-generation wave-based electronics hardware [31].

In scientific domains such as computational biology and chemistry, graphs are
often used to represent the 3D structures of molecules [32, 33], chemical
compounds [34], and even large biomolecules such as proteins [35, 36]. Graphs
have even been used in fields such as computational physics to model complex
particle physics simulations [37] as well as in real-world traffic systems to
predict travel times and delays [38]. Underlying many of these successful
examples of graph representations are GNNs, a class of machine learning
algorithms specialized in processing irregularly-structured input data such as
graphs. Careful applications of graph neural networks in scientific domains
have considered the physical symmetries present in many scientific data such
as molecular state symmetries [39] or physical dynamics constraints [40] and
have leveraged such symmetries to design new attention-based neural network
architectures [41, 42].

Throughout their development, geometric deep learning methods have expanded to
incorporate within them equivariance to various geometric symmetry groups to
enhance their generalization capabilities and adversarial robustness. Methods
such as group-equivariant CNNs [43], Tensor Field Networks [44],
SE(3)-Transformers [45], and equivariant GNNs [45, 46, 47, 48, 49, 50, 51, 52,
53, 54, 55] have paved the way for the development of future deep learning
models that respect physical symmetries present in 3D data (e.g., rotation
equivariance with respect to input data symmetries). Concurrently to these
efforts, self-supervised learning methods have begun to facilitate automatic
detection and enforcement of the symmetries present in input data within the
network’s representations for such inputs [56]. Nonetheless, deciding how to
optimize self-supervised learning algorithms for one’s desired level of
equivariance has proven to be a challenging task [57].

In total, in this work, we make connections between geometric graph neural
networks, equivariance, and geometric information-completeness guarantees that
provide one with a rich foundation on which to build new graph neural network
architectures. In particular, we introduce a new graph neural network model,
GCPNet, that is equivariant to the group of 3D rotations and translations
(i.e., SE(3), the special Euclidean group) and guarantees geometric
information completeness following graph message-passing on 3D point clouds.
We demonstrate its expressiveness and flexibility for modeling physical
systems through rigorous experiments for distinct molecular-geometric tasks.
In detail, we provide the following contributions:

  * •

In contrast to prior geometric networks for molecules that are insensitive to
their chemical chirality, cannot detect global physical forces acting upon
each atom, or do not directly learn geometric features, we present the first
geometric graph neural network architecture with the following desirable
properties for learning from 3D molecules: (1) the ability to directly predict
translation and rotation-invariant scalar properties and rotation-equivariant
vector-valued quantities for nodes and edges, respectively; (2) a rotation and
translation-equivariant method for iteratively updating node positions in 3D
space; (3) sensitivity to molecular chirality; and (4) a means by which to
learn from and account for the global forces acting upon the atoms within its
inputs.

  * •

We establish new state-of-the-art results for four distinct molecular-
geometric representation learning tasks - molecular chirality recognition,
protein-ligand binding affinity prediction, protein structure ranking, and
Newtonian many-body-systems modeling - where model predictions vary from
analyzing individual nodes to summarizing entire graph inputs. GCPNet’s
performance for these tasks is statistically significant and surpasses that of
previous state-of-the-art machine learning methods for 3D molecules.

##  2 Results

In this work, we consider four distinct modeling tasks comprised of seven
datasets in total, where implementation details are discussed in Appendix C.
We note that additional experiments are included in Appendix B for interested
readers.

Assessing model sensitivity to molecular chirality. Molecular chirality is an
essential geometric property of 3D molecules for models to consider when
making predictions for downstream tasks. Simply put, this property describes
the ”handedness” of 3D molecules, in that, certain molecules cannot be
geometrically superimposed upon a mirror reflection of themselves using only
3D rotation and translation operations. This subsequently poses a key
challenge for machine learning models: Can such predictive models effectively
sensitize their predictions to the effects of molecular chirality such that,
under 3D reflections, their molecular feature representations change
accordingly? To answer this question using modern machine learning methods, we
adopt the rectus/sinister (RS) 3D molecular dataset of [58] to evaluate the
ability of state-of-the-art machine learning methods to distinguish between
right-handed and left-handed versions of a 3D molecule. In addition, we
carefully follow their experimental setup including dataset splitting and
evaluation criteria, where we evaluate each method’s classification accuracy
in distinguishing between right and left-handed versions of a molecule.
Baseline methods for this task include state-of-the-art invariant neural
networks (INNs) and equivariant neural networks (ENNs), where we list each
method’s latest results for this task as reported in [59].

Table 1: Comparison of GCPNet with baseline methods for the RS task. The results are averaged over three independent runs. The top-1 (best) results for this task are in bold, and the second-best results are underlined. Type | Method | R/S Accuracy (%) $\uparrow$
---|---|---
INN | ChIRo [2023] | 98.5
| SchNet [2023] | 54.4
| DimeNet++ [2023] | 65.7
| SphereNet [2023] | 98.2
ENN | EGNN [2023] | 50.4
| SEGNN [2023] | 83.4
Ours |  GCPNet w/o Frames | 50.2 $\pm$ 0.6
| GCPNet |  98.7 $\pm$ 0.1

Contribution of frame embeddings for chirality sensitivity. Table 1 shows that
GCPNet is more accurately able to detect the effects of molecular chirality
compared to all other baseline methods, even without performing any
hyperparameter tuning. In particular, GCPNet outperforms ChIRo [58], a GNN
specifically designed to detect different forms of chirality in 3D molecules.
Moreover, when we ablate GCPNet’s embeddings of local geometric frames, we
find that this E(3)-equivariant (i.e., 3D rotation and reflection-equivariant)
version of GCPNet is no longer able to solve this important molecular
recognition task, resulting in prediction accuracies at parity with random
guessing. These two previous observations highlight that (1) GCPNet’s local
frame embeddings are critical components of the model’s sensitivity to
molecular chirality and that, (2) using such frame embeddings, GCPNet can
flexibly learn representations of 3D molecules that are more predictive of
chemical chirality compared to hand-crafted methods for such tasks.

Table 2: Comparison of GCPNet with baseline methods for the LBA task. The results are averaged over three independent runs. The top-1 (best) results for this task are in bold, and the second-best results are underlined. Type | Method | RMSE $\downarrow$ |  $p$ $\uparrow$ |  $Sp$ $\uparrow$
---|---|---|---|---
CNN | 3DCNN [2020] | 1.416 $\pm$ 0.021 | 0.550 | 0.553
| DeepDTA [2018] | 1.866 $\pm$ 0.080 | 0.472 | 0.471
| DeepAffinity [2019] | 1.893 $\pm$ 0.650 | 0.415 | 0.426
RNN | Bepler and Berger [2019] | 1.985 $\pm$ 0.006 | 0.165 | 0.152
| TAPE [2019] | 1.890 $\pm$ 0.035 | 0.338 | 0.286
| ProtTrans [2021] | 1.544 $\pm$ 0.015 | 0.438 | 0.434
GNN | GCN [2020] | 1.601 $\pm$ 0.048 | 0.545 | 0.533
| DGAT [2021] | 1.719 $\pm$ 0.047 | 0.464 | 0.472
| DGIN [2021] | 1.765 $\pm$ 0.076 | 0.426 | 0.432
| DGAT-GCN [2021] | 1.550 $\pm$ 0.017 | 0.498 | 0.496
| MaSIF [2020] | 1.484 $\pm$ 0.018 | 0.467 | 0.455
| IEConv [2021] | 1.554 $\pm$ 0.016 | 0.414 | 0.428
| Holoprot-Full Surface [2021] | 1.464 $\pm$ 0.006 | 0.509 | 0.500
| Holoprot-Superpixel [2021] | 1.491 $\pm$ 0.004 | 0.491 | 0.482
| ProNet-Amino-Acid [2023] | 1.455 $\pm$ 0.009 | 0.536 | 0.526
| ProNet-Backbone [2023] | 1.458 $\pm$ 0.003 | 0.546 | 0.550
| ProNet-All-Atom [2023] | 1.463 $\pm$ 0.001 | 0.551 | 0.551
| GeoSSL-DDM [2023] | 1.451 $\pm$ 0.030 | 0.577 | 0.572
ENN | Cormorant [2019] | 1.568 $\pm$ 0.012 | 0.389 | 0.408
| PaiNN [2021] | 1.698 $\pm$ 0.050 | 0.366 | 0.358
| ET [2022] | 1.490 $\pm$ 0.019 | 0.564 | 0.532
| GVP [2021] | 1.594 $\pm$ 0.073 | 0.434 | 0.432
| GBP [2022] |  1.405 $\pm$ 0.009 | 0.561 | 0.557
Ours |  GCPNet w/o Frames | 1.485 $\pm$ 0.015 | 0.521 | 0.504
|  GCPNet w/o ResGCP | 1.514 $\pm$ 0.008 | 0.471 | 0.468
|  GCPNet w/o Scalars | 1.685 $\pm$ 0.000 | 0.050 | 0.000
|  GCPNet w/o Vectors | 1.727 $\pm$ 0.005 | 0.270 | 0.304
| GCPNet |  1.352 $\pm$ 0.003 | 0.608 | 0.607

Evaluating predictions of protein-ligand binding affinity. Protein-ligand
binding affinity (LBA) prediction challenges methods to estimate the binding
affinity of a protein-ligand complex as a single scalar value [60]. Accurately
estimating such values in a matter of seconds using a machine learning model
can provide invaluable and timely information in the typical drug discovery
pipeline [73]. The corresponding dataset for this SE(3)-invariant task is
derived from the ATOM3D dataset [60] and is comprised of 4,463 nonredundant
protein-ligand complexes, where cross-validation splits are derived using a
strict 30% sequence identity cutoff. Results are reported in terms of the root
mean squared error (RMSE), Pearson’s correlation ($p$), and Spearman’s
correlation ($Sp$) between a method’s predictions on the test dataset and the
corresponding ground-truth binding affinity values represented as
$pK=-\log_{10}(K)$, where $K$ is the binding affinity measured in Molar units.
Baseline comparison methods for this task include a variety of state-of-the-
art CNNs, recurrent neural networks (RNNs), GNNs, and ENNs.

The results shown in Table 2 reveal that, in operating on atom-level protein-
ligand graph representations, GCPNet achieves the best performance for
predicting protein-ligand binding affinity by a significant margin, notably
improving performance across all metrics by 7% on average. Here, to the best
of our knowledge, GCPNet is also the first method capable of achieving Pearson
and Spearman binding affinity correlations greater than 0.6 on the PDBBind
dataset [74] when employing a strict 30% sequence identity cutoff. Moreover,
we find that these correlations are highly statistically significant (i.e.,
Pearson’s p-value of $2e-50$, Spearman’s p-value of $2e-49$, and Kendall’s tau
correlation of 0.432 with a p-value of $3e-45$).

Ablating network components reveals impact of model design. Denoted as ”GCPNet
w/o …” in Table 2, our ablation studies with GCPNet for the LBA task
demonstrate the contribution of each component in its model design. In
particular, our proposed local frame embeddings improve GCPNet’s performance
by more than 15% across all metrics (GCPNet w/o Frames), where we hypothesize
these performance improvements come from using these frame embeddings to
enhance the model’s sensitivity to molecular chirality. Similarly, our
proposed residual GCP module (i.e., ResGCP) improves GCPNet’s performance by
23% on average.

Specifically of interest is the observation that independent removal of scalar
and vector-valued features within GCPNet appears to severely decrease GCPNet’s
performance for LBA prediction. Notably, removing the model’s access to
scalar-valued features degrades performance by 70% on average, while not
allowing the model to access vector-valued features reduces performance by 42%
on average. One possible explanation for these observations is that both types
of feature representations the baseline GCPNet model learns (i.e., scalars and
vectors) are useful for understanding protein-ligand interactions. In
addition, our ablation results in Table 2 suggest that our proposed frame
embeddings and ResGCP module are complementary to these scalar and vector-
valued features in the context of predicting the binding affinity of a
protein-ligand complex.

Table 3: Comparison of GCPNet with baseline methods for the PSR task. Local metrics are averaged across target-aggregated metrics. The best results for this task are in bold, and the second-best results are underlined. N/A denotes a metric that could not be computed. |  | Local |  |  | Global |
---|---|---|---|---|---|---
Method |  $p$ $\uparrow$ |  $Sp$ $\uparrow$ |  $K$ $\uparrow$ |  $p$ $\uparrow$ |  $Sp$ $\uparrow$ |  $K$ $\uparrow$
3DCNN [2020] | 0.557 | 0.431 | 0.308 | 0.780 | 0.789 | 0.592
GCN [2020] | 0.500 | 0.411 | 0.289 | 0.747 | 0.750 | 0.547
ProQ3D [2017] | 0.444 | 0.432 | 0.304 | 0.796 | 0.772 | 0.594
VoroMQA [2017] | 0.412 | 0.419 | 0.291 | 0.688 | 0.651 | 0.505
RWplus [2010] | 0.192 | 0.167 | 0.137 | 0.033 | 0.056 | 0.011
SBROD [2019] | 0.431 | 0.413 | 0.291 | 0.551 | 0.569 | 0.393
Ornate [2019] | 0.393 | 0.371 | 0.256 | 0.625 | 0.669 | 0.481
DimeNet [2020] | 0.302 | 0.351 | 0.285 | 0.614 | 0.625 | 0.431
GraphQA [2021] | 0.357 | 0.379 | 0.251 | 0.821 | 0.820 | 0.618
PaiNN [2021] | 0.518 | 0.444 | 0.315 | 0.773 | 0.813 | 0.611
ET [2021] | 0.564 | 0.466 | 0.330 | 0.813 | 0.814 | 0.611
GVP [2021] | 0.581 | 0.462 | 0.331 | 0.805 | 0.811 | 0.616
GBP [2022] | 0.612 | 0.517 | 0.372 | 0.856 | 0.853 | 0.656
GCPNet w/o Frames | 0.588 | 0.512 | 0.367 | 0.854 | 0.851 | 0.657
GCPNet w/o ResGCP | 0.576 | 0.509 | 0.365 | 0.852 | 0.847 | 0.648
GCPNet w/o Scalars | N/A | N/A | N/A | N/A | N/A | N/A
GCPNet w/o Vectors | 0.571 | 0.497 | 0.356 | 0.802 | 0.804 | 0.608
GCPNet | 0.616 | 0.534 | 0.385 | 0.871 | 0.869 | 0.676

Evaluating ranking predictions for protein structure decoys. Protein structure
ranking (PSR) requires methods to predict the overall quality of a 3D protein
structure when comparing it to a reference (i.e., native) protein structure
[60]. The quality of a protein structure is reported as a single scalar value
representing a method’s predicted global distance test (GDT_TS) score [82]
between the provided decoy structure and the native structure. Such
information is crucial in drug discovery efforts when one is tasked with
designing a drug (e.g., ligand) that should bind to a particular protein
target, notably when such targets have not yet had their 3D structures
experimentally determined and have rather had them predicted computationally
using methods such as AlphaFold 2 [42]. The respective dataset for this
SE(3)-invariant task is also derived from the ATOM3D dataset [60] and is
comprised of 40,950 decoy structures corresponding to 649 total targets, where
cross-validation splits are created according to a target’s release year in
the Critical Assessment of Techniques for Protein Structure Prediction (CASP)
competition [83]. Results are reported in terms of the Pearson’s correlation
($p$), Spearman’s correlation ($Sp$), and Kendall’s tau correlation ($K$)
between a method’s predictions on the test dataset and the corresponding
ground-truth GDT_TS values, where local results are averaged across
predictions for individual targets and global results are averaged directly
across all targets. Baseline comparison methods for this task include a
composition of state-of-the-art CNNs, GNNs, and ENNs, as well as previous
statistics-based methods.

Conveying a similar message to that in Table 2, the results in Table 3
demonstrate that, in operating on atom-level protein graphs, GCPNet performs
best against all other state-of-the-art models for the task of estimating a 3D
protein structure’s quality (i.e., PSR). In this setting, GCPNet outperforms
all other methods across all local and global metrics by 2.5% on average. Once
again, GCPNet’s predictions are highly statistically significant, this time
with Pearson, Spearman, and Kendall tau p-values all below $1e-50$,
respectively.

Identifying components for effective protein structure ranking. Our ablation
studies with GCPNet, in the context of PSR, once more reveal that the design
of our local frames, ResGCP module, and scalar and vector feature channels are
all beneficial for enhancing GCPNet’s ability to analyze a given 3D graph
input. Here, in sensitizing the model to chemical chirality, our local frame
embeddings improve GCPNet’s performance for PSR by 4% on average. Similarly,
our ResGCP module improves the model’s performance by 5%. Interestingly,
without access to scalar-valued node and edge features, GCPNet is unable to
produce valid predictions for the PSR test dataset due to what appears to be a
phenomenon of vector-wise latent variable collapse [84]. This finding suggests
that, for the PSR task, the baseline GCPNet model relies strongly on the
scalar-valued representations it produces. Lastly, including vector-valued
features within GCPNet improves the model’s performance for the PSR task by
9%.

Table 4: Comparison of GCPNet with baseline methods for the NMS task. Results are reported in terms of the MSE for future position prediction over four test datasets of increasing modeling difficulty, graph sizes, and composed force field complexities. The final column reports each method’s MSE averaged across all four test datasets. The best results for this task are in bold, and the second-best results are underlined. N/A denotes an experiment that could not be performed due to a method’s numerical instability. Method | ES(5) | ES(20) | G+ES(20) | L+ES(20) | Average
---|---|---|---|---|---
GNN [2022] | 0.0131 | 0.0720 | 0.0721 | 0.0908 | 0.0620
TFN [2022] | 0.0236 | 0.0794 | 0.0845 | 0.1243 | 0.0780
SE(3)-Transformer [2022] | 0.0329 | 0.1349 | 0.1000 | 0.1438 | 0.1029
Radial Field [2022] | 0.0207 | 0.0377 | 0.0399 | 0.0779 | 0.0441
PaiNN [2021] | 0.0158 | N/A | N/A | N/A | N/A
ET [2022] | 0.1653 | 0.1788 | 0.2122 | 0.2989 | 0.2138
EGNN [2022] | 0.0079 | 0.0128 | 0.0118 | 0.0368 | 0.0173
ClofNet [2022] | 0.0065 | 0.0073 | 0.0072 | 0.0251 | 0.0115
GCPNet w/o Frames | 0.0067 | 0.0074 | 0.0074 | 0.0200 | 0.0103
GCPNet w/o ResGCP | 0.0090 | 0.0135 | 0.0099 | 0.0278 | 0.0150
GCPNet w/o Scalars | 0.0119 | 0.0173 | 0.0170 | 0.0437 | 0.0225
GCPNet | 0.0070 | 0.0071 | 0.0073 | 0.0173 | 0.0097

Evaluating trajectory predictions for Newtonian many-body systems. Newtonian
many-body systems modeling (NMS) asks methods to forecast the future positions
of particles in many-body systems of various sizes [53], bridging the gap
between the domains of machine learning and physics. In our experimental
results for the NMS task, the four systems (i.e., datasets) on which we
evaluate each method are comprised of increasingly more nodes and are
influenced by force fields of increasingly complex directional origins for
which to model, namely electrostatic force fields for 5-body (ES(5)) and
20-body (ES(20)) systems as well as for 20-body systems under the influence of
an additional gravity field (G+ES(20)) and Lorentz-like force field
(L+ES(20)), respectively. The four datasets for this SE(3)-equivariant task
were generated using the descriptions and source code of [53], where each
dataset is comprised of 7,000 total trajectories. Results are reported in
terms of the mean squared error (MSE) between a method’s node position
predictions on the test dataset and the corresponding ground-truth node
positions after 1,000 timesteps. Baseline comparison methods for this task
include a collection of state-of-the-art GNNs, ENNs, and Transformers.

The results in Table 4 show that GCPNet achieves the lowest MSE averaged
across all four NMS datasets, improving upon the state-of-the-art MSE for
trajectory predictions in this task by 19% on average. In particular, GCPNet
achieves the best results for two of the four NMS datasets considered in this
work, where these two datasets are respectively the first and third most
difficult NMS datasets for methods to model. On the two remaining datasets,
GCPNet matches the performance of prior state-of-the-art methods. Moreover,
across all four datasets, GCPNet’s trajectory predictions yield an RMSE of
0.0963 and achieve Pearson, Spearman, and Kendall’s tau correlations of 0.999,
0.999, and 0.981, respectively, where all such correlation values are highly
statistically significant (i.e., p-values $<1e-50$). Note that, to calculate
these correlation values, we score GCPNet’s vector-valued predictions
independently for each coordinate axis and then average the resulting metrics.

Analyzing components for successful trajectory forecasting. Once again, our
ablation studies with GCPNet demonstrate the importance of GCPNet’s local
frame embeddings, scalar information, and ResGCP module. Here, we note that we
were not able to include an ablation study on GCPNet’s vector-valued features
since they are directly used to predict node position displacements for
trajectory forecasting. Table 4 shows that each model component
synergistically enables GCPNet to achieve new state-of-the-art results for the
NMS task. In enabling the model to detect global forces, our proposed local
frame embeddings improve GCPNet’s ability to learn many-body system dynamics
by 6% on average across all dataset contexts. Specifically interesting to note
is that these local frame embeddings improve the model’s trajectory
predictions within the most complex dataset context (i.e., L+ES(20)) by 14%,
suggesting that such frame embeddings improve GCPNet’s ability to learn many-
body system dynamics even in the presence of complex global force fields.
Furthermore, GCPNet’s ResGCP module and scalar-valued features improve the
model’s performance for modeling many-body systems by 35% and 57%,
respectively.

Across all tasks studied in this work, GCPNet improves upon the overall
performance of all previous methods. Our experiments demonstrate this for both
node-level (e.g., NMS) and graph-level (e.g., LBA) prediction tasks, verifying
GCPNet’s ability to encode useful information for both scales of granularity.
Furthermore, we have demonstrated the importance of each model component
within GCPNet, showing how these components are complementary to each other in
the context of representation learning over 3D molecular data. We will now
proceed to describe the design and operations of our proposed GCPNet model
architecture.

##  3 Methods

###  3.1 Preliminaries

####  3.1.1 Overview of the Problem Setting

We represent a 3D molecular structure as a 3D $k$-nearest neighbors ($k$-NN)
graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ with $\mathcal{V}$ and
$\mathcal{E}$ representing the graph’s set of nodes and set of edges,
respectively, and $N=|\mathcal{V}|$ and $E=|\mathcal{E}|$ representing the
number of nodes and the number of edges in the graph, respectively. In
addition, $\mathbf{X}\in\mathbb{R}^{N\times 3}$ represents the respective
Cartesian coordinates for each node. We then design E(3)-invariant (i.e., 3D
rotation, reflection, and translation-invariant) node features
$\mathbf{H}\in\mathbb{R}^{N\times h}$ and edge features
$\mathbf{E}\in\mathbb{R}^{E\times e}$ as well as O(3)-equivariant (3D rotation
and reflection-equivariant) node features
$\bm{\chi}\in\mathbb{R}^{N\times(m\times 3)}$ and edge features
$\bm{\xi}\in\mathbb{R}^{E\times(x\times 3)}$, respectively.

Upon constructing such features, we apply several layers of graph message-
passing using a neural network $\bm{\Phi}$ (which later on we refer to as
GCPNet) that updates node and edge features using invariant and equivariant
representations for the corresponding feature types. Importantly, $\bm{\Phi}$
guarantees, by design, SE(3) equivariance with respect to its vector-valued
input coordinates and features (i.e., $x_{i}\in\mathbf{X}$,
$\chi_{i}\in\bm{\chi}$, and $\xi_{ij}\in\bm{\xi}$) and SE(3)-invariance
regarding its scalar features (i.e., $h_{i}\in\mathbf{H}$ and
$e_{ij}\in\mathbf{E}$). In addition, $\bm{\Phi}$’s scalar graph
representations achieve geometric self-consistency for the 3D structure of the
input molecular graph $\mathcal{G}$, sensitizing them to the effects of
molecular chirality while making them uniquely identifiable under 3D
rotations. Lastly, geometric completeness requires methods that accept 3D
molecular graph inputs to be able to discern the local geometric environment
of a given atom with no directional ambiguities. This enables geometry-
complete methods such as $\bm{\Phi}$ to detect the presence and influence of
global force fields acting on the graph inputs. We formalize these
equivariance, geometric self-consistency, and geometric completeness
constraints using the three following definitions, where $\Box^{\prime}$
represents a feature that has been updated by our neural network.

######  Definition 1.

(SE(3) Equivariance).

Given
$(\mathbf{H}^{\prime},\mathbf{E}^{\prime},\mathbf{X}^{\prime},\bm{\chi}^{\prime},\bm{\xi}^{\prime})=\bm{\Phi}(\mathbf{H},\mathbf{E},\mathbf{X},\bm{\chi},\bm{\xi})$,
we have
$(\mathbf{H}^{\prime},\mathbf{E}^{\prime},\mathbf{Q}\mathbf{X}^{\prime^{\mathrm{T}}}+\mathbf{g},\mathbf{Q}\bm{\chi}^{\prime^{\mathrm{T}}},\mathbf{Q}\bm{\xi}^{\prime^{\mathrm{T}}})=\bm{\Phi}(\mathbf{H},\mathbf{E},\mathbf{Q}\mathbf{X}^{\mathrm{T}}+\mathbf{g},\mathbf{Q}\bm{\chi}^{\mathrm{T}},\mathbf{Q}\bm{\xi}^{\mathrm{T}})$,
$\quad\forall\mathbf{Q}\in\textit{SO}(3),\forall\mathbf{g}\in\mathbb{R}^{3\times
1}$.

######  Definition 2.

(Geometric Self-Consistency).

Given a pair of molecular graphs $\mathcal{G}_{1}$ and $\mathcal{G}_{2}$,
with $\mathbf{X}^{1}=\\{\mathbf{x}_{i}^{1}\\}_{i=1,...,N}$ and
$\mathbf{X}^{2}=\\{\mathbf{x}_{i}^{2}\\}_{i=1,...,N}$, respectively,
a geometric representation
$\bm{\Phi}(\mathbf{H},\mathbf{E})=\bm{\Phi}(\mathcal{G})$ is considered
geometrically self-consistent if
$\bm{\Phi}(\mathcal{G}^{1})=\bm{\Phi}(\mathcal{G}^{2})\Longleftrightarrow\exists\mathbf{Q}\in\textit{SO}(3),\exists\mathbf{g}\in\mathbb{R}^{3\times
1}$,
for
$i=1,...,n,\mathbf{X}_{i}^{1^{\mathrm{T}}}=\mathbf{Q}\mathbf{X}_{i}^{2^{\mathrm{T}}}+\mathbf{g}$
[85].

######  Definition 3.

(Geometric Completeness).

Given a positional pair of nodes $(x_{i}^{t},x_{j}^{t})$ in a 3D graph
$\mathcal{G}$,
with vectors $a_{ij}^{t}\in\mathbb{R}^{1\times 3}$,
$b_{ij}^{t}\in\mathbb{R}^{1\times 3}$, and $c_{ij}^{t}\in\mathbb{R}^{1\times
3}$ derived from $(x_{i}^{t},x_{j}^{t})$,
a local geometric representation
$\bm{\mathcal{F}}_{ij}^{t}=(a_{ij}^{t},b_{ij}^{t},c_{ij}^{t})\in\mathbb{R}^{3\times
3}$ is considered
geometrically complete if $\bm{\mathcal{F}}_{ij}^{t}$ is non-degenerate,
thereby forming a
local orthonormal basis located at the tangent space of $x_{i}^{t}$ [53].

####  3.1.2 SE(3)-equivariant complete representations

Representation learning on 3D molecular structures is a challenging task for a
variety of reasons: (1) an expressive representation learning model should be
able to predict arbitrary vector-valued quantities for each atom and atom pair
in the molecular structure (e.g., using $\bm{\chi}^{\prime}$ and
$\bm{\xi}^{\prime}$ to predict side-chain atom positions and atom-atom
displacements for each residue in a 3D protein graph); (2) arbitrary rotations
or translations to a 3D molecular structure should affect only the vector-
valued representations a model assigns to a molecular graph’s nodes or edges,
whereas such 3D transformations of the molecular structure should not affect
the model’s scalar representations for nodes and edges [53]; (3) the
geometrically invariant properties of a molecule’s 3D structure should be
uniquely identifiable by a model; and (4) in a geometry-complete manner,
scalar and vector-valued representations should mutually exchange information
between nodes and edges during a model’s forward pass for a 3D input graph, as
these information types can be correlatively related (e.g., a scalar feature
such as the $L_{2}$ norm of a vector $v$ can be associated with the vector of
origin $v$) [54, 41].

In line with this reasoning, we need to ensure that the coordinates our model
predicts for the node positions in a molecular graph $\mathcal{G}$ transform
according to SE(3) transformations of the input positions. This runs in
contrast to previous methods that remain strictly E(3)-equivariant or
E(3)-invariant to 3D transformations of the input $\mathcal{G}$ and
consequently ignore the important effects of molecular chirality. At the same
time, the model should jointly update the scalar and vector-valued features of
$\mathcal{G}$ according to their respective molecular symmetry groups to
increase the model’s expressiveness in approximating geometric and physical
quantities [86]. To increase its generalization capabilities, the model should
also disambiguate any geometric directions within its local node environments
and should maintain SE(3)-invariance of its scalar representations when the
input graph is transformed in 3D space. Following [70], this helps prevent the
model from losing important geometric or chiral information (i.e., becoming
geometrically self-inconsistent) during graph message-passing. One way to do
this is to introduce a new type of message-passing neural network.

Figure 2: A framework overview for our proposed Geometry-Complete Perceptron
Network (GCPNet). Our framework consists of (i.) a graph (topology) definition
process, (ii.) a GCPNet-based graph neural network for 3D molecular
representation learning, and (iii.) demonstrated application areas for GCPNet.
Zoom in for the best viewing experience.

###  3.2 GCPNet Model Architecture

Towards this end, we introduce our architecture for $\bm{\Phi}$ satisfying
Defs. (1), (2), and (3) which we refer to as the Geometry-Complete
SE(3)-Equivariant Perceptron Network (GCPNet). We illustrate the GCPNet
algorithm in Figure 2 and outline it in Algorithm 1. Subsequently, we expand
on our definition for $\mathbf{GCP}$ and $\mathbf{GCPConv}$ in Sections 3.2.1
and 3.3.1, respectively, while further illustrating $\mathbf{GCP}$ in Figure
3.

Figure 3: An overview of our proposed Geometry-Complete Perceptron
($\mathbf{GCP}$) module. The $\mathbf{GCP}$ module introduces node and edge-
centric encodings of 3D frames as input features that are used to directly
update both scalar and vector-valued features with geometric information-
completeness guarantees as well as chirality sensitivity.

We can then prove the following three propositions (see Appendices A.1 through
A.3 for a more detailed description of the GCPNet algorithm and its
equivariant properties).

  * •

Proposition 1. GCPNets are SE(3)-equivariant
$\longrightarrow$ Def. (1).

  * •

Proposition 2. GCPNets are geometry self-consistent
$\longrightarrow$ Def. (2).

  * •

Proposition 3. GCPNets are geometry-complete
$\longrightarrow$ Def. (3).

####  3.2.1 Geometry-Complete Perceptron Module

As illustrated in Figure 3, GCPNet represents the features for nodes within an
input graph as a tuple $(h,\chi)$ to distinguish scalar features
$(h\in\mathbb{R}^{h})$ from vector-valued features
$(\chi\in\mathbb{R}^{m\times 3})$. Similarly, GCPNet represents an input
graph’s edge features as a tuple $(e,\xi)$ to differentiate scalar features
$(e\in\mathbb{R}^{e})$ from vector-valued features $(\xi\in\mathbb{R}^{x\times
3})$. For conciseness, we will subsequently refer to both node and edge
feature tuples as $(s,V)$. We then define
$\mathbf{GCP}_{\mathcal{F}_{ij},\lambda}(\cdot)$ to represent the
$\mathbf{GCP}$ encoding process, where $\lambda$ represents a downscaling
hyperparameter (e.g., 3) and $\bm{\mathcal{F}}_{ij}\in\mathbb{R}^{3\times 3}$
denotes the SO(3)-equivariant (i.e., 3D rotation-equivariant) frames
constructed using the $\mathbf{Localize}$ operation (i.e., the
$\mathbf{EquiFrame}$ operation of [53]) in Algorithm 1. Specifically, the
frame encodings are defined as
$\bm{\mathcal{F}}_{ij}^{t}=(a_{ij}^{t},b_{ij}^{t},c_{ij}^{t})$, with
$a_{ij}^{t}=\frac{x_{i}^{t}-x_{j}^{t}}{\lVert
x_{i}^{t}-x_{j}^{t}\rVert},b_{ij}^{t}=\frac{x_{i}^{t}\times x_{j}^{t}}{\lVert
x_{i}^{t}\times x_{j}^{t}\rVert},$ and $c_{ij}^{t}=a_{ij}^{t}\times
b_{ij}^{t}$, respectively. In Appendix A.3, we discuss how these frame
encodings are direction information-complete for edges, allowing networks
incorporating them to effectively detect and leverage for downstream tasks the
force fields present within real-world many-body systems such as small
molecules and proteins.

Expressing Vector Representations with $V$. The $\mathbf{GCP}$ module then
expresses vector representations $V$ as follows. The features $V$ with
representation depth $r$ are downscaled by $\lambda$.

$z=\\{v\bm{w}_{d_{z}}|\bm{w}_{d_{z}}\in\mathbb{R}^{r\times(r/\lambda)}\\}$ (1)

Additionally, $V$ is separately downscaled in preparation to be subsequently
embedded as direction-sensitive edge scalar features.

$V_{s}=\\{v\bm{w}_{d_{s}}|\bm{w}_{d_{s}}\in\mathbb{R}^{r\times(3\times 3)}\\}$
(2)

Deriving Scalar Representations $s^{\prime}$. To update scalar
representations, the $\mathbf{GCP}$ module, in the following manner, derives
two invariant sources of information from $V$ and combines them with $s$:

$q_{ij}=(V_{s}\cdot\mathcal{F}_{ij})\in\mathbb{R}^{9}$ (3) $\displaystyle
q=\begin{cases}\frac{1}{\lvert\mathcal{N}(i)\rvert}\sum_{j\in\mathcal{N}(i)}q_{ij}&\text{if
$V_{s}$ represents nodes}\\\ q_{ij}&\text{if $V_{s}$ represents
edges}\end{cases}$ (4) $s_{(s,q,z)}=s\cup q\cup\lVert z\rVert_{2}$ (5)

where $\cdot$ denotes the inner product, $\mathcal{N}(\cdot)$ represents the
neighbors of a node, and $\lVert\cdot\rVert_{2}$ denotes the $L_{2}$ norm.
Then, denote $t$ as the representation depth of $s$, and let
$s_{(s,q,z)}\in\mathbb{R}^{t+9+(r/\lambda)}$ with representation depth
$(t+9+(r/\lambda))$ be projected to $s^{\prime}$ with representation depth
$t^{\prime}$:

$s_{v}=\\{s_{(s,q,z)}\bm{w}_{s}+\bm{b}_{s}|\bm{w}_{s}\in\mathbb{R}^{(t+9+(r/\lambda))\times
t^{\prime}}\\}$ (6) $s^{\prime}=\sigma_{s}(s_{v})$ (7)

Note that embedding geometric frames $\mathcal{F}_{ij}$ as $q_{ij}$ in
Equation 3 ultimately enables GCPNet to iteratively learn chirality-sensitive
and global force-aware representations of each 3D network input. Moreover,
Equation 4 allows GCPNet to encode local geometric substructures for each
node, where the theoretical importance of such network behavior is discussed
in detail by [87].

Deriving Vector Representations $V^{\prime}$. The $\mathbf{GCP}$ module then
concludes by updating vector representations as follows:

$V_{u}=\\{z\bm{w}_{u_{z}}|\bm{w}_{u_{z}}\in\mathbb{R}^{(r/\lambda)\times
r^{\prime}}\\}$ (8)
$V^{\prime}=\\{V_{u}\odot\sigma_{g}(\sigma^{+}(s_{v})\bm{w}_{g}+\bm{b}_{g})|\bm{w}_{g}\in\mathbb{R}^{t^{\prime}\times
r^{\prime}}\\}$ (9)

where $\odot$ represents element-wise multiplication and the gating function
$\sigma_{g}$ is applied row-wise to preserve SO(3) equivariance within
$V^{\prime}$.

Conceptually, the $\mathbf{GCP}$ module is autoregressively applied to tuples
$(s,V)$ a total of $\omega$ times to derive rich scalar and vector-valued
features. The module does so by blending both feature types iteratively with
the 3D direction and information completeness guarantees provided by geometric
frame encodings $\bm{\mathcal{F}}_{ij}$. We note that this model design runs
in contrast with prior graph neural networks for physical systems such as GVP-
GNNs [53] and ClofNet [46], which are either insensitive to chemical chirality
and global atomic forces or do not directly learn geometric features for
downstream prediction tasks, making the proposed $\mathbf{GCP}$ module well
suited for learning directly from 3D molecular graphs.

###  3.3 Learning from 3D Graphs with GCPNet

In this section, we propose a flexible manner in which to perform 3D graph
convolution with our proposed $\mathbf{GCP}$ module, as illustrated in Figure
2 and employed in Algorithm 1.

####  3.3.1 Geometry-Complete Graph Convolution.

Let $\mathcal{N}$(i) denote the neighbors of node $n_{i}$, selected using a
distance-based metric such as k-nearest neighbors or a radial distance cutoff.
Subsequently, we define a single layer $l$ of geometry-complete graph
convolution as

$n_{i}^{l}=\phi^{l}(n_{i}^{l-1},\mathcal{A}_{\forall
j\in\mathcal{N}(i)}\Omega_{\omega}^{l}(n_{i}^{l-1},n_{j}^{l-1},e_{ij},\mathcal{F}_{ij})),$
(10)

where $n_{i}^{l}=(h_{i}^{l},\chi_{i}^{l})$;
$e_{ij}=(e_{ij}^{0},\xi_{ij}^{0})$; $\bm{\Phi}$ is a trainable function
denoted as $\mathbf{GCPConv}$; $l$ signifies the representation depth of the
network; $\mathcal{A}$ is a permutation-invariant aggregation function; and
$\Omega_{\omega}$ represents a message-passing function corresponding to the
$\omega$-th $\mathbf{GCP}$ message-passing layer. We proceed to expand on the
operations of each graph convolution layer as follows.

To start, messages between source nodes $i$ and neighboring nodes $j$ are
first constructed as

$m_{ij}^{0}=\mathbf{GCP}(n_{i}^{0}\cup n_{j}^{0}\cup e_{ij},\mathcal{F}_{ij})$
(11)

where $\cup$ denotes a concatenation operation. Then, up to the $\omega$-th
iteration, each message is updated by the $m$-th message update layer using
residual connections as

$\displaystyle\Omega_{\omega}^{l}$
$\displaystyle=\mathbf{ResGCP}_{\omega}^{l}(m_{ij}^{l-1},\mathcal{F}_{ij}),$
(12) $\displaystyle\mathbf{ResGCP}_{\eta}^{l}(z_{i}^{l-1},\mathcal{F}_{ij})$
$\displaystyle=z_{i}^{l-1}+\mathbf{GCP}_{\eta}^{l}(z_{i}^{l-1},\mathcal{F}_{ij}),$
(13)

where we empirically find such residual connections between message
representations to reduce oversmoothing within GCPNet by mitigating the
problem of vanishing gradients.

Updated node features $\hat{n}^{l}$ are then derived residually using an
aggregation of generated messages as

$\hat{n}^{l}=n^{l-1}+f(\\{\Omega_{\omega,v_{i}}^{l}|v_{i}\in\mathcal{V}\\}),$
(14)

where $f$ represents an aggregation function such as a summation or mean that
is invariant to permutations of node ordering. The residual connection between
$\hat{n}^{l}$ and $n^{l}$ is established here to encourage the network to
update the representation space of node features in a layer-asynchronous
manner.

To encourage GCPNet to make its node feature representations independent of
the size of each input graph, we then employ a node-centric feed-forward
network to update node representations. Specifically, we apply to
$\hat{n}^{l}$ a linear $\mathbf{GCP}$ function with shared weights $\phi_{f}$
followed by $r$ $\mathbf{ResGCP}$ modules, operations concisely portrayed as

$\displaystyle\tilde{n}_{r-1}^{l}$ $\displaystyle=\phi_{f}^{l}(\hat{n}^{l})$
(15) $\displaystyle n^{l}$
$\displaystyle=\mathbf{ResGCP}_{r}^{l}(\tilde{n}_{r-1}^{l}).$ (16)

Lastly, if one desires to update the positions of each node in $\mathcal{G}$
(e.g., as we do for tasks involving position-related predictions such as NMS),
we propose a flexible, SE(3)-equivariant method to do so using a dedicated
$\mathbf{GCP}$ module as follows:

$\displaystyle(h_{p_{i}}^{l},\chi_{p_{i}}^{l})$
$\displaystyle=\mathbf{GCP}_{p}^{l}(n_{i}^{l},\mathcal{F}_{ij})$ (17)
$\displaystyle x_{i}^{l}$
$\displaystyle=x_{i}^{l-1}+\chi_{p_{i}}^{l},\mbox{where
}\chi_{p_{i}}^{l}\in\mathbb{R}^{1\times 3}.$ (18)

####  3.3.2 The GCPNet Algorithm

Algorithm 1 GCPNet

1: Input: $(h_{i}\in\mathbf{H}$, $\chi_{i}\in\bm{\chi})$,
$(e_{ij}\in\mathbf{E}$, $\xi_{ij}\in\bm{\xi})$, $x_{i}\in\mathbf{X}$, graph
$\mathcal{G}$

2: Initialize
$\mathbf{X}^{0}=\mathbf{X}^{C}\leftarrow\mathbf{Centralize}(\mathbf{X})$

3:
$\bm{\mathcal{F}}_{ij}=\mathbf{Localize}(x_{i}\in\mathbf{X}^{0},x_{j}\in\mathbf{X}^{0})$

4: Project $(h_{i}^{0}$, $\chi_{i}^{0})$, $(e_{ij}^{0}$,
$\xi_{ij}^{0})\leftarrow\mathbf{GCP}_{e}((h_{i}$, $\chi_{i})$, $(e_{ij}$,
$\xi_{ij})$, $\mathcal{F}_{ij})$

5: for $l=1$ to $L$ do

6: $(h_{i}^{l}$, $\chi_{i}^{l})$,
$x_{i}^{l}=\mathbf{GCPConv}^{l}((h_{i}^{l-1}$, $\chi_{i}^{l-1})$,
$(e_{ij}^{0}$, $\xi_{ij}^{0})$, $x_{i}^{l-1}$, $\mathcal{F}_{ij})$

7: end for

8: if Updating Node Positions then

9: $\bm{\mathcal{F}}_{ij}^{L}=\mathbf{Localize}(x_{i}\in\mathbf{X}^{l}$,
$x_{j}\in\mathbf{X}^{l})$

10: Finalize $(\mathbf{X}^{L})\leftarrow\mathbf{Decentralize}(\mathbf{X}^{l})$

11: else

12: $x_{i}^{L}=x_{i}^{0}$

13: end if

14: Project $(h_{i}^{L}$, $\chi_{i}^{L})$, $(e_{ij}^{L}$,
$\xi_{ij}^{L})\leftarrow\mathbf{GCP}_{p}((h_{i}^{l}$, $\chi_{i}^{l})$,
$(e_{ij}^{0}$, $\xi_{ij}^{0})$, $\mathcal{F}_{ij}^{L})$

15: Output: $(h_{i}^{L}$, $\chi_{i}^{L})$, $(e_{ij}^{L}$, $\xi_{ij}^{L})$,
$x_{i}^{L}$

In this section, we describe our overall learning algorithm driven by GCPNet
(Algorithm 1). We also discuss the rationale behind our design decisions for
GCPNet and provide examples of use cases in which one might apply GCPNet for
specific learning tasks.

On Line 2 of Algorithm 1, the $\mathbf{Centralize}$ operation removes the
center of mass from each node position in the input graph to ensure that such
positions are subsequently 3D translation-invariant.

Thereafter, following [53], the $\mathbf{Localize}$ operation on Line 3 crafts
translation-invariant and SO(3)-equivariant frame encodings
$\bm{\mathcal{F}}_{ij}^{t}=(a_{ij}^{t},b_{ij}^{t},c_{ij}^{t})$. As described
in more detail in Appendix A.3, these frame encodings are chirality-sensitive
and direction information-complete for edges, imbuing networks that
incorporate them with the ability to more easily detect force field
interactions present in many real-world atomic systems, as we demonstrate
through corresponding experiments in Section 2.

Before applying any geometry-complete graph convolution layers, on Line 4 we
use $\mathbf{GCP}_{e}$ to embed our input node and edge features into scalar
and vector-valued values, respectively, while incorporating geometric frame
information. Subsequently, in Lines 5-6, each layer of geometry-complete graph
convolution is performed autoregressively via $\mathbf{GCPConv}^{l}$ starting
from these initial node and edge feature embeddings, all while maintaining
information flow originating from the geometric frames
$\bm{\mathcal{F}}_{ij}$.

On Lines 8 through 12, we finalize our procedure with which to update in an
SE(3)-equivariant manner the position of each node in an input 3D graph. In
particular, we update node positions by residually adding learned vector-
valued node features ($\chi_{v_{i}}^{l}$) to the node positions produced by
the previous $\mathbf{GCPConv}$ layer ($l-1$). As shown in Appendix A.1, such
updates are initially SO(3)-equivariant, and on Line 10 we ensure these
updates also become 3D translation-equivariant by adding back to each node
position the input graph’s original center of mass via the
$\mathbf{Decentralize}$ operation. In total, this procedure produces
SE(3)-equivariant updates to node positions. Additionally, for models that
update node positions, we note that Line 9 updates frame encodings
$\bm{\mathcal{F}}_{ij}$ using the model’s final predictions for node positions
to provide more information-rich feature projections on Line 14 via
$\mathbf{GCP}_{p}$ to conclude the forward pass of GCPNet.

####  3.3.3 Network Utilities.

In summary, GCPNet receives an input 3D graph $\mathcal{G}$ with node
positions $\bm{x}$, scalar node and edge features, $h$ and $e$, as well as
vector-valued node and edge features, $\chi$ and $\xi$. The model is then
capable of e.g., (1) predicting scalar node, edge, or graph-level properties
while maintaining SE(3) invariance; (2) estimating vector-valued node, edge,
or graph-level properties while ensuring SE(3) equivariance; or (3) updating
node positions in an SE(3)-equivariant manner.

##  4 Discussion

In light of our impressive results with GCPNet, future work on the model could
involve researching more computationally-efficient variations of GCPNet that
require fewer GCP message-passing layers within each GCP convolution layer or
that embed geometric frames sparsely rather than in each GCP layer.
Interestingly, [87] recently showed that GCPNet, in particular, is
theoretically robust in terms of its geometric expressiveness of local
geometric substructures and global geometric interaction terms. Nonetheless,
[87] also highlight future directions for improving such methods’
expressiveness for learning equivariant-valued features including
incorporating higher-order equivariant tensors into one’s message-passing
procedure. Enhancing its geometric expressiveness to thereby increase its
runtime efficiency would allow GCPNet to be used increasingly in new
scientific and deep learning applications requiring high computational
throughput (e.g., virtual screening of new drugs).

##  5 Conclusion

In this work, we introduced GCPNet, a state-of-the-art GNN for 3D molecular
graph representation learning. We have demonstrated its utility through
several benchmark studies. In future work, we aim to develop extensions of
GCPNet that increase its geometric expressiveness as well as explore
applications of GCPNet for generative modeling of molecular structures.

Code Availability The source code for GCPNet is available at
https://github.com/BioinfoMachineLearning/GCPNet.

Data Availability The data required to train new GCPNet models or reproduce
our results for the NMS task are available under a Creative Commons
Attribution 4.0 International Public License at
https://zenodo.org/record/7293186. The data required to train new GCPNet
models or reproduce our results for the RS task are available under an MIT
License at https://figshare.com/s/e23be65a884ce7fc8543. All other data
required to train new GCPNet models or reproduce our results for the remaining
tasks are available to download using our model training scripts available at
https://github.com/BioinfoMachineLearning/GCPNet.

Acknowledgments This work is partially supported by two NSF grants (DBI1759934
and IIS1763246), two NIH grants (R01GM093123 and R01GM146340), three DOE
grants (DE-AR0001213, DE-SC0020400, and DE-SC0021303), and the computing
allocation on the Summit compute cluster provided by the Oak Ridge Leadership
Computing Facility.

Author Contributions Statement AM and JC conceived the project. AM designed
the experiments. AM performed the experiments and collected the data. AM
analyzed the data. AM and JC wrote the manuscript. AM and JC edited the
manuscript.

Competing Interests Statement The authors declare no competing interests.

## References

  *   * Sehnal et al. [2021] Sehnal, D., Bittrich, S., Deshpande, M., Svobodová, R., Berka, K., Bazgier, V., Velankar, S., Burley, S.K., Koča, J., Rose, A.S.: Mol* viewer: modern web app for 3d visualization and analysis of large biomolecular structures. Nucleic Acids Research 49(W1), 431–437 (2021)
  * Valsesia et al. [2018] Valsesia, D., Fracastoro, G., Magli, E.: Learning localized generative models for 3d point clouds via graph convolution. In: International Conference on Learning Representations (2018)
  * Qi et al. [2017] Qi, X., Liao, R., Jia, J., Fidler, S., Urtasun, R.: 3d graph neural networks for rgbd semantic segmentation. In: Proceedings of the IEEE International Conference on Computer Vision, pp. 5199–5208 (2017)
  * Zhang and Rabbat [2018] Zhang, Y., Rabbat, M.: A graph-cnn for 3d point cloud classification. In: 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 6279–6283 (2018). IEEE
  * Zhou et al. [2021] Zhou, H., Feng, Y., Fang, M., Wei, M., Qin, J., Lu, T.: Adaptive graph convolution for point cloud analysis. In: Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4965–4974 (2021)
  * Esteves [2020] Esteves, C.: Theoretical aspects of group equivariant neural networks. arXiv preprint arXiv:2004.05154 (2020)
  * Hegyi and Gerstein [1999] Hegyi, H., Gerstein, M.: The relationship between protein structure and function: a comprehensive survey with application to the yeast genome. Journal of molecular biology 288(1), 147–164 (1999)
  * Umulis and Othmer [2012] Umulis, D.M., Othmer, H.G.: The importance of geometry in mathematical models of developing systems. Current opinion in genetics & development 22(6), 547–552 (2012)
  * Matter and Niederberger [2022] Matter, F., Niederberger, M.: The importance of the macroscopic geometry in gas-phase photocatalysis. Advanced Science 9(13), 2105363 (2022)
  * Hamilton [2020] Hamilton, W.L.: Graph representation learning. Synthesis Lectures on Artifical Intelligence and Machine Learning 14(3), 1–159 (2020)
  * Cao et al. [2022] Cao, W., Zheng, C., Yan, Z., He, Z., Xie, W.: Geometric machine learning: research and applications. Multimedia Tools and Applications, 1–53 (2022)
  * Masci et al. [2016] Masci, J., Rodolà, E., Boscaini, D., Bronstein, M.M., Li, H.: Geometric deep learning. In: SIGGRAPH ASIA 2016 Courses. SA ’16. Association for Computing Machinery, New York, NY, USA (2016). https://doi.org/10.1145/2988458.2988485 . https://doi.org/10.1145/2988458.2988485
  * Bronstein et al. [2017] Bronstein, M.M., Bruna, J., LeCun, Y., Szlam, A., Vandergheynst, P.: Geometric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine 34(4), 18–42 (2017)
  * Bronstein et al. [2021] Bronstein, M.M., Bruna, J., Cohen, T., Veličković, P.: Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478 (2021)
  * Cao et al. [2020] Cao, W., Yan, Z., He, Z., He, Z.: A comprehensive survey on geometric deep learning. IEEE Access 8, 35929–35949 (2020)
  * Kiarashinejad et al. [2020] Kiarashinejad, Y., Zandehshahvar, M., Abdollahramezani, S., Hemmatyar, O., Pourabolghasem, R., Adibi, A.: Knowledge discovery in nanophotonics using geometric deep learning. Advanced Intelligent Systems 2(2), 1900132 (2020)
  * James [2021] James, J.: Citywide traffic speed prediction: A geometric deep learning approach. Knowledge-Based Systems 212, 106592 (2021)
  * LeCun et al. [1995] LeCun, Y., Bengio, Y., et al.: Convolutional networks for images, speech, and time series. The handbook of brain theory and neural networks 3361(10), 1995 (1995)
  * He et al. [2016] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 770–778 (2016)
  * He et al. [2017] He, K., Gkioxari, G., Dollár, P., Girshick, R.: Mask r-cnn. In: Proceedings of the IEEE International Conference on Computer Vision, pp. 2961–2969 (2017)
  * Krizhevsky et al. [2017] Krizhevsky, A., Sutskever, I., Hinton, G.E.: Imagenet classification with deep convolutional neural networks. Communications of the ACM 60(6), 84–90 (2017)
  * Kipf and Welling [2016] Kipf, T.N., Welling, M.: Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907 (2016)
  * Gilmer et al. [2017] Gilmer, J., Schoenholz, S.S., Riley, P.F., Vinyals, O., Dahl, G.E.: Neural message passing for quantum chemistry. In: International Conference on Machine Learning, pp. 1263–1272 (2017). PMLR
  * Veličković et al. [2017] Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., Bengio, Y.: Graph attention networks. arXiv preprint arXiv:1710.10903 (2017)
  * Dwivedi and Bresson [2020] Dwivedi, V.P., Bresson, X.: A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09699 (2020)
  * Qi et al. [2017] Qi, C.R., Su, H., Mo, K., Guibas, L.J.: Pointnet: Deep learning on point sets for 3d classification and segmentation. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 652–660 (2017)
  * Liu et al. [2019] Liu, Y., Fan, B., Xiang, S., Pan, C.: Relation-shape convolutional neural network for point cloud analysis. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8895–8904 (2019)
  * Zhang et al. [2019] Zhang, Z., Hua, B.-S., Yeung, S.-K.: Shellnet: Efficient point cloud convolutional neural networks using concentric shells statistics. In: Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1607–1616 (2019)
  * Graves et al. [2013] Graves, A., Mohamed, A.-r., Hinton, G.: Speech recognition with deep recurrent neural networks. In: 2013 IEEE International Conference on Acoustics, Speech and Signal Processing, pp. 6645–6649 (2013). Ieee
  * Graves and Jaitly [2014] Graves, A., Jaitly, N.: Towards end-to-end speech recognition with recurrent neural networks. In: International Conference on Machine Learning, pp. 1764–1772 (2014). PMLR
  * Hughes et al. [2019] Hughes, T.W., Williamson, I.A., Minkov, M., Fan, S.: Wave physics as an analog recurrent neural network. Science advances 5(12), 6946 (2019)
  * Duvenaud et al. [2015] Duvenaud, D.K., Maclaurin, D., Iparraguirre, J., Bombarell, R., Hirzel, T., Aspuru-Guzik, A., Adams, R.P.: Convolutional networks on graphs for learning molecular fingerprints. In: Cortes, C., Lawrence, N., Lee, D., Sugiyama, M., Garnett, R. (eds.) Advances in Neural Information Processing Systems, vol. 28. Curran Associates, Inc., ??? (2015). https://proceedings.neurips.cc/paper/2015/file/f9be311e65d81a9ad8150a60844bb94c-Paper.pdf
  * Liu et al. [2021] Liu, Y., Wang, L., Liu, M., Zhang, X., Oztekin, B., Ji, S.: Spherical message passing for 3d graph networks. arXiv preprint arXiv:2102.05013 (2021)
  * Akutsu and Nagamochi [2013] Akutsu, T., Nagamochi, H.: Comparison and enumeration of chemical graphs. Computational and structural biotechnology journal 5(6), 201302004 (2013)
  * Xia and Ku [2021] Xia, T., Ku, W.-S.: Geometric graph representation learning on protein structure prediction. In: Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pp. 1873–1883 (2021)
  * Morehead et al. [2022] Morehead, A., Chen, X., Wu, T., Liu, J., Cheng, J.: Egr: Equivariant graph refinement and assessment of 3d protein complex structures. arXiv preprint arXiv:2205.10390 (2022)
  * Shlomi et al. [2020] Shlomi, J., Battaglia, P., Vlimant, J.-R.: Graph neural networks in particle physics. Machine Learning: Science and Technology 2(2), 021001 (2020)
  * Derrow-Pinion et al. [2021] Derrow-Pinion, A., She, J., Wong, D., Lange, O., Hester, T., Perez, L., Nunkesser, M., Lee, S., Guo, X., Wiltshire, B., et al.: Eta prediction with graph neural networks in google maps. In: Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 3767–3776 (2021)
  * Ye et al. [2020] Ye, S., Liang, J., Liu, R., Zhu, X.: Symmetrical graph neural network for quantum chemistry with dual real and momenta space. The Journal of Physical Chemistry A 124(34), 6945–6953 (2020)
  * Han et al. [2022] Han, J., Huang, W., Ma, H., Li, J., Tenenbaum, J.B., Gan, C.: Learning physical dynamics with subequivariant graph neural networks. arXiv preprint arXiv:2210.06876 (2022)
  * Morehead et al. [2022] Morehead, A., Chen, C., Cheng, J.: Geometric transformers for protein interface contact prediction. In: International Conference on Learning Representations (2022). https://openreview.net/forum?id=CS4463zx6Hi
  * Jumper et al. [2021] Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Žídek, A., Potapenko, A., et al.: Highly accurate protein structure prediction with alphafold. Nature 596(7873), 583–589 (2021)
  * Cohen and Welling [2016] Cohen, T., Welling, M.: Group equivariant convolutional networks. In: International Conference on Machine Learning, pp. 2990–2999 (2016). PMLR
  * Thomas et al. [2018] Thomas, N., Smidt, T.E., Kearnes, S.M., Yang, L., Li, L., Kohlhoff, K., Riley, P.F.: Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds. ArXiv abs/1802.08219 (2018)
  * Fuchs et al. [2020] Fuchs, F., Worrall, D., Fischer, V., Welling, M.: Se(3)-transformers: 3d roto-translation equivariant attention networks. Advances in Neural Information Processing Systems 33, 1970–1981 (2020)
  * Jing et al. [2020] Jing, B., Eismann, S., Suriana, P., Townshend, R.J., Dror, R.: Learning from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.01411 (2020)
  * Jing et al. [2021] Jing, B., Eismann, S., Soni, P.N., Dror, R.O.: Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint arXiv:2106.03843 (2021)
  * Kofinas et al. [2021] Kofinas, M., Nagaraja, N., Gavves, E.: Roto-translated local coordinate frames for interacting dynamical systems. Advances in Neural Information Processing Systems 34, 6417–6429 (2021)
  * Gasteiger et al. [2021] Gasteiger, J., Becker, F., Günnemann, S.: Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems 34, 6790–6802 (2021)
  * Schütt et al. [2021] Schütt, K., Unke, O., Gastegger, M.: Equivariant message passing for the prediction of tensorial properties and molecular spectra. In: International Conference on Machine Learning, pp. 9377–9388 (2021). PMLR
  * Huang et al. [2022] Huang, W., Han, J., Rong, Y., Xu, T., Sun, F., Huang, J.: Equivariant graph mechanics networks with constraints. arXiv preprint arXiv:2203.06442 (2022)
  * Thölke and De Fabritiis [2022] Thölke, P., De Fabritiis, G.: Equivariant transformers for neural network based molecular potentials. In: International Conference on Learning Representations (2022)
  * Du et al. [2022] Du, W., Zhang, H., Du, Y., Meng, Q., Chen, W., Zheng, N., Shao, B., Liu, T.-Y.: SE(3) equivariant graph neural networks with complete local frames. In: Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., Sabato, S. (eds.) Proceedings of the 39th International Conference on Machine Learning. Proceedings of Machine Learning Research, vol. 162, pp. 5583–5608. PMLR, ??? (2022). https://proceedings.mlr.press/v162/du22e.html
  * Aykent and Xia [2022] Aykent, S., Xia, T.: Gbpnet: Universal geometric representation learning on protein structures. In: Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. KDD ’22, pp. 4–14. Association for Computing Machinery, New York, NY, USA (2022). https://doi.org/10.1145/3534678.3539441 . https://doi.org/10.1145/3534678.3539441
  * Batzner et al. [2022] Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J.P., Kornbluth, M., Molinari, N., Smidt, T.E., Kozinsky, B.: E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications 13(1), 2453 (2022)
  * Dangovski et al. [2021] Dangovski, R., Jing, L., Loh, C., Han, S., Srivastava, A., Cheung, B., Agrawal, P., Soljacic, M.: Equivariant self-supervised learning: Encouraging equivariance in representations. In: International Conference on Learning Representations (2021)
  * Xie et al. [2022] Xie, Y., Wen, J., Lau, K.W., Rehman, Y.A.U., Shen, J.: What should be equivariant in self-supervised learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4111–4120 (2022)
  * Adams et al. [2021] Adams, K., Pattanaik, L., Coley, C.W.: Learning 3d representations of molecular chirality with invariance to bond rotations. arXiv preprint arXiv:2110.04383 (2021)
  * Schneuing et al. [2023] Schneuing, A., Du, Y., Harris, C., Jamasb, A.R., Igashov, I., Du, Blundell, T.L., Lio, P., Gomes, C.P., Welling, M., Bronstein, M.M., Correia, B.: Structure-based Drug Design with Equivariant Diffusion Models (2023). https://openreview.net/forum?id=uKmuzIuVl8z
  * Townshend et al. [2020] Townshend, R.J., Vögele, M., Suriana, P., Derry, A., Powers, A., Laloudakis, Y., Balachandar, S., Jing, B., Anderson, B., Eismann, S., et al.: Atom3d: Tasks on molecules in three dimensions. arXiv preprint arXiv:2012.04035 (2020)
  * Öztürk et al. [2018] Öztürk, H., Özgür, A., Ozkirimli, E.: Deepdta: deep drug–target binding affinity prediction. Bioinformatics 34(17), 821–829 (2018)
  * Karimi et al. [2019] Karimi, M., Wu, D., Wang, Z., Shen, Y.: Deepaffinity: interpretable deep learning of compound–protein affinity through unified recurrent and convolutional neural networks. Bioinformatics 35(18), 3329–3338 (2019)
  * Bepler and Berger [2019] Bepler, T., Berger, B.: Learning protein sequence embeddings using information from structure. In: International Conference on Learning Representations (2019). https://openreview.net/forum?id=SygLehCqtm
  * Rao et al. [2019] Rao, R., Bhattacharya, N., Thomas, N., Duan, Y., Chen, P., Canny, J., Abbeel, P., Song, Y.: Evaluating protein transfer learning with tape. In: Wallach, H., Larochelle, H., Beygelzimer, A., Alché-Buc, F., Fox, E., Garnett, R. (eds.) Advances in Neural Information Processing Systems, vol. 32. Curran Associates, Inc., ??? (2019). https://proceedings.neurips.cc/paper/2019/file/37f65c068b7723cd7809ee2d31d7861c-Paper.pdf
  * Elnaggar et al. [2021] Elnaggar, A., Heinzinger, M., Dallago, C., Rehawi, G., Wang, Y., Jones, L., Gibbs, T., Feher, T., Angerer, C., Steinegger, M., et al.: Prottrans: Toward understanding the language of life through self-supervised learning. IEEE transactions on pattern analysis and machine intelligence 44(10), 7112–7127 (2021)
  * Nguyen et al. [2021] Nguyen, T., Le, H., Quinn, T.P., Nguyen, T., Le, T.D., Venkatesh, S.: Graphdta: Predicting drug–target binding affinity with graph neural networks. Bioinformatics 37(8), 1140–1147 (2021)
  * Gainza et al. [2020] Gainza, P., Sverrisson, F., Monti, F., Rodola, E., Boscaini, D., Bronstein, M., Correia, B.: Deciphering interaction fingerprints from protein molecular surfaces using geometric deep learning. Nature Methods 17(2), 184–192 (2020)
  * Hermosilla et al. [2021] Hermosilla, P., Schäfer, M., Lang, M., Fackelmann, G., Vázquez, P.-P., Kozlikova, B., Krone, M., Ritschel, T., Ropinski, T.: Intrinsic-extrinsic convolution and pooling for learning on 3d protein structures. In: International Conference on Learning Representations (2021). https://openreview.net/forum?id=l0mSUROpwY
  * Somnath et al. [2021] Somnath, V.R., Bunne, C., Krause, A.: Multi-scale representation learning on proteins. In: Beygelzimer, A., Dauphin, Y., Liang, P., Vaughan, J.W. (eds.) Advances in Neural Information Processing Systems (2021). https://openreview.net/forum?id=-xEk43f_EO6
  * Wang et al. [2023] Wang, L., Liu, H., Liu, Y., Kurtin, J., Ji, S.: Learning hierarchical protein representations via complete 3d graph networks. In: The Eleventh International Conference on Learning Representations (2023). https://openreview.net/forum?id=9X-hgLDLYkQ
  * Liu et al. [2023] Liu, S., Guo, H., Tang, J.: Molecular geometry pretraining with SE(3)-invariant denoising distance matching. In: The Eleventh International Conference on Learning Representations (2023). https://openreview.net/forum?id=CjTHVo1dvR
  * Anderson et al. [2019] Anderson, B., Hy, T.S., Kondor, R.: Cormorant: Covariant molecular neural networks. Advances in neural information processing systems 32 (2019)
  * Rezaei et al. [2020] Rezaei, M.A., Li, Y., Wu, D., Li, X., Li, C.: Deep learning in drug design: protein-ligand binding affinity prediction. IEEE/ACM Transactions on Computational Biology and Bioinformatics (2020)
  * Wang et al. [2005] Wang, R., Fang, X., Lu, Y., Yang, C.-Y., Wang, S.: The pdbbind database: methodologies and updates. Journal of medicinal chemistry 48(12), 4111–4119 (2005)
  * Uziela et al. [2017] Uziela, K., Menéndez Hurtado, D., Shu, N., Wallner, B., Elofsson, A.: Proq3d: improved model quality assessments using deep learning. Bioinformatics 33(10), 1578–1580 (2017)
  * Olechnovič and Venclovas [2017] Olechnovič, K., Venclovas, Č.: Voromqa: Assessment of protein structure quality using interatomic contact areas. Proteins: Structure, Function, and Bioinformatics 85(6), 1131–1145 (2017)
  * Zhang and Zhang [2010] Zhang, J., Zhang, Y.: A novel side-chain orientation dependent potential derived from random-walk reference state for protein fold selection and structure prediction. PloS one 5(10), 15386 (2010)
  * Karasikov et al. [2019] Karasikov, M., Pagès, G., Grudinin, S.: Smooth orientation-dependent scoring function for coarse-grained protein quality assessment. Bioinformatics 35(16), 2801–2808 (2019)
  * Pagès et al. [2019] Pagès, G., Charmettant, B., Grudinin, S.: Protein model quality assessment using 3d oriented convolutional neural networks. Bioinformatics 35(18), 3313–3319 (2019)
  * Klicpera et al. [2020] Klicpera, J., Groß, J., Günnemann, S.: Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123 (2020)
  * Baldassarre et al. [2021] Baldassarre, F., Menéndez Hurtado, D., Elofsson, A., Azizpour, H.: Graphqa: protein model quality assessment using graph convolutional networks. Bioinformatics 37(3), 360–366 (2021)
  * Zemla [2003] Zemla, A.: Lga: a method for finding 3d similarities in protein structures. Nucleic acids research 31(13), 3370–3374 (2003)
  * Kryshtafovych et al. [2021] Kryshtafovych, A., Schwede, T., Topf, M., Fidelis, K., Moult, J.: Critical assessment of methods of protein structure prediction (casp)—round xiv. Proteins: Structure, Function, and Bioinformatics 89(12), 1607–1617 (2021)
  * Dieng et al. [2019] Dieng, A.B., Kim, Y., Rush, A.M., Blei, D.M.: Avoiding latent variable collapse with generative skip models. In: Chaudhuri, K., Sugiyama, M. (eds.) Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics. Proceedings of Machine Learning Research, vol. 89, pp. 2397–2405. PMLR, ??? (2019). https://proceedings.mlr.press/v89/dieng19a.html
  * Wang et al. [2022] Wang, L., Liu, Y., Lin, Y., Liu, H., Ji, S.: ComENet: Towards complete and efficient message passing for 3d molecular graphs. In: Oh, A.H., Agarwal, A., Belgrave, D., Cho, K. (eds.) Advances in Neural Information Processing Systems (2022). https://openreview.net/forum?id=mCzMqeWSFJ
  * Brandstetter et al. [2021] Brandstetter, J., Hesselink, R., Pol, E., Bekkers, E.J., Welling, M.: Geometric and physical quantities improve e (3) equivariant message passing. arXiv preprint arXiv:2110.02905 (2021)
  * Du et al. [2023] Du, W., Du, Y., Wang, L., Feng, D., Wang, G., Ji, S., Gomes, C., Ma, Z.-M.: A new perspective on building efficient and expressive 3d equivariant graph neural networks. arXiv preprint arXiv:2304.04757 (2023)
  * Zaheer et al. [2017] Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R.R., Smola, A.J.: Deep sets. Advances in neural information processing systems 30 (2017)
  * Qiao et al. [2022] Qiao, Z., Nie, W., Vahdat, A., Miller III, T.F., Anandkumar, A.: Dynamic-backbone protein-ligand structure prediction with multiscale generative diffusion models. arXiv preprint arXiv:2209.15171 (2022)
  * Ingraham et al. [2019] Ingraham, J., Garg, V., Barzilay, R., Jaakkola, T.: Generative models for graph-based protein design. Advances in neural information processing systems 32 (2019)
  * Kuhlman and Baker [2000] Kuhlman, B., Baker, D.: Native protein sequences are close to optimal for their structures. Proceedings of the National Academy of Sciences 97(19), 10383–10388 (2000)
  * Gao et al. [2022] Gao, Z., Tan, C., Li, S.Z.: Pifold: Toward effective and efficient protein inverse folding. arXiv preprint arXiv:2209.12643 (2022)
  * Van Rossum and Drake [2009] Van Rossum, G., Drake, F.L.: Python 3 Reference Manual. CreateSpace, Scotts Valley, CA (2009)
  * Paszke et al. [2019] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., Chintala, S.: Pytorch: An imperative style, high-performance deep learning library. In: Wallach, H., Larochelle, H., Beygelzimer, A., Alché-Buc, F., Fox, E., Garnett, R. (eds.) Advances in Neural Information Processing Systems 32, pp. 8024–8035. Curran Associates, Inc., ??? (2019)
  * Falcon [2019] Falcon, e.a. WA: Pytorch lightning. https://github.com/PyTorchLightning/pytorch-lightning 3 (2019)
  * Fey and Lenssen [2019] Fey, M., Lenssen, J.E.: Fast graph representation learning with pytorch geometric. ArXiv abs/1903.02428 (2019)

##  Appendix A Proofs.

###  A.1 Proof of Proposition 1.

Proof. Suppose the vector-valued features given to the corresponding
$\mathbf{GCPConv}$ layers in GCPNet are node features $\chi_{i}$ and edge
features $\xi_{ij}$ that are O(3)-equivariant (i.e., 3D rotation and
reflection-equivariant) by way of their construction. Additionally, suppose
the scalar-valued features given to the respective $\mathbf{GCPConv}$ layers
in GCPNet are E(3)-invariant (i.e., 3D rotation, reflection, and translation-
invariant) node features $h_{i}$ and edge features $e_{ij}$.

Translation equivariance. In line with [53], the $\mathbf{Centralize}$
operation on Line 2 of Algorithm 1 first ensures that $\mathbf{X}^{0}$ becomes
3D translation invariant by the following procedure. Let
$\mathbf{X}(t)=(\mathbf{x}_{1}(t),...,\mathbf{x}_{n}(t))$ represent a many-
body system at time $t$, where the centroid of the system is defined as

$C(t)=\frac{\mathbf{x}_{1}(t)+...+\mathbf{x}_{n}(t)}{n}.$ (19)

Note that in uniformly translating the position of the system by a vector
$\mathbf{v}$, we have $\mathbf{X}(t)+\mathbf{v}\longrightarrow
C(t)+\mathbf{v}$, meaning that the centroid of the system translates in the
same manner as the system itself. However, note that if at time $t=0$ we
recenter the origin of $\mathbf{X}$ to its centroid, we have

$\displaystyle\mathbf{X}(t)-C(0)\xrightarrow{\mbox{translation by
}\mathbf{v}\mbox{ at }t=0}\mathbf{X}(t)-C(0)$

which implies the system $\mathbf{X}$ is translation-invariant under the
centralized reference $\mathbf{X}(t)-C(0)$ when the translation vector
$\mathbf{v}$ is applied to $\mathbf{X}$ at time $t=0$. Concretely, in the case
of translation-invariant tasks such as predicting molecular properties or
classifying point clouds, here we have successfully achieved 3D translation
invariance. Moreover, for translation-equivariant tasks such as forecasting
the positions of a many-body system, we can achieve translation equivariance
by simply adding $C(0)$ back to the predicted positions. Therefore, using the
above methodology, GCPNets are translation equivariant.

Permutation equivariance. Succinctly, we note that since GCPNet operates on
graph-structured input data, permutation equivariance is guaranteed by design.
For further discussion of why our proposed method as well as why other graph-
based algorithms proposed previously are inherently permutation-equivariant,
we refer readers to [88]. Therefore, GCPNets are permutation-equivariant.

SO(3)-equivariant frames. On Line 3 of Algorithm 1, the $\mathbf{Localize}$
operation constructs SO(3)-equivariant (i.e., 3D rotation-equivariant) frames
$\bm{\mathcal{F}}_{ij}$ in the following manner.

Define our frame encodings as

$\bm{\mathcal{F}}_{ij}^{t}=(a_{ij}^{t},b_{ij}^{t},c_{ij}^{t}),$ (20)

where we have

$a_{ij}^{t}=\frac{x_{i}^{t}-x_{j}^{t}}{\lVert
x_{i}^{t}-x_{j}^{t}\rVert},b_{ij}^{t}=\frac{x_{i}^{t}\times x_{j}^{t}}{\lVert
x_{i}^{t}\times x_{j}^{t}\rVert},c_{ij}^{t}=a_{ij}^{t}\times b_{ij}^{t}.$ (21)

The proof that $\bm{\mathcal{F}}_{ij}^{t}$ is equivariant under SO(3)
transformations of its input space is included in [53]. However, for
completeness, we include a version of it here.

Let $g\in SO(3)$ be an action under which the positions in $X$ transform
equivariantly, and $\bm{\mathcal{F}}_{ij}^{t}$ be defined as we have it in
Equation 20 above. That is, we have

$\displaystyle(\mathbf{x}_{1}(t),...,\mathbf{x}_{n}(t))\xrightarrow{g}(g\mathbf{x}_{1}(t),...,g\mathbf{x}_{n}(t)),$

where from the definition of $a_{ij}^{t}$ in Equation 21 we have

$\displaystyle a_{ij}^{t}\xrightarrow{g}ga_{ij}^{t}.$

Considering $b_{ij}^{t}$, from Equation 21 we have

$\displaystyle(gx_{i}(t))\times(gx_{j}(t))$
$\displaystyle=\det(g)(g^{T})^{-1}(x_{i}(t)\times x_{j}(t))$
$\displaystyle=g(x_{i}(t)\times x_{j}(t)),$ (22)

where using $g^{-1}=g^{T}$ for the orthogonal matrix $g$ gives us Equation
A.1. Consequently, $b_{ij}^{t}\xrightarrow{g}gb_{ij}^{t}$. Lastly, by applying
Equation A.1 once again, we have that $c_{ij}^{t}\xrightarrow{g}gc_{ij}^{t}$.

Moreover, note that under reflections of $x$, we have $R:x\rightarrow-x$ which
gives us $a_{ij}^{t}\rightarrow-a_{ij}^{t}$. Thereafter, by the right-hand
rule, the cross product of two equivariant vectors gives us a pseudo-vector
$b_{ij}^{t}=x_{i}^{t}\times x_{j}^{t}\rightarrow b_{ij}^{t}$, where
subsequently it is implied that $c_{ij}^{t}\rightarrow-c_{ij}^{t}$.
Consequently, we have $\det(-a_{ij}^{t},b_{ij}^{t},-c_{ij}^{t})=1$, informing
us that the frame encodings $\bm{\mathcal{F}}_{ij}^{t}$ are not reflection-
equivariant (a symmetry that is important to not enforce when learning
representations of chiral molecules such as proteins). Therefore, the frame
encodings within GCPNet are SO(3)-equivariant.

Note, after the construction of these frames, that they are used on Line 4 of
Algorithm 1 to embed all node and edge features (i.e., $h_{i}$, $e_{ij}$,
$\chi_{i}$, and $\xi_{ij}$) using a single $\mathbf{GCP}$ module as well as in
all subsequent $\mathbf{GCP}$ modules. We will now prove that the feature
updates each $\mathbf{GCP}$ module makes with the frame encodings
$\bm{\mathcal{F}}_{ij}^{t}$ defined in Equation 20 are SO(3)-equivariant.

SO(3)-equivariant GCP module. The operations of a $\mathbf{GCP}$ module are
illustrated in Figure 3 and derived in Section 3.2.1. Their SO(3) invariance
for scalar feature updates and SO(3) equivariance for vector-valued feature
updates is proven as follows.

Following the proof of O(3) equivariance for the GVP module in [46], the proof
of SO(3) equivariance within the $\mathbf{GCP}$ module is similar, with the
following modifications. Within the $\mathbf{GCP}$ module, the vector-valued
features (processed separately for nodes and edges) are fed not only through a
bottleneck block comprised of downward and upward projection matrices
$\mathbf{D_{z}}$ and $\mathbf{U_{z}}$ but are also fed into a dedicated
downward projection matrix $\mathbf{D_{\mathcal{S}}}$. The output of matrix
multiplication between O(3)-equivariant vector features and
$\mathbf{D_{\mathcal{S}}}$ yields O(3)-equivariant vector features
$v_{i_{\mathcal{S}}}$ that are used as unique inputs for an SO(3)-invariant
scalarization operation. In particular, the following demonstrates the
invariance of our design for matrix multiplication with our $\mathbf{GCP}$
module’s projection matrices (e.g., $\mathbf{D_{\mathcal{S}}}$). Suppose
$\mathbf{W}_{h}\in\mathbb{R}^{h\times v}$, $\mathbf{V}\in\mathbb{R}^{v\times
3}$, and $\mathbf{Q}\in SO(3)\in\mathbb{R}^{3\times 3}$. In line with [46],
observe for
$\mathbf{D}=(\mathbf{Q}\mathbf{V}^{\mathrm{T}})\in\mathbb{R}^{3\times v}$ that

$\displaystyle\lVert\mathbf{W}_{h}\mathbf{D}^{\mathrm{T}}\rVert_{2}=\lVert\mathbf{W}_{h}(\mathbf{V}^{\mathrm{T}})^{\mathrm{T}}\rVert_{2}=\lVert\mathbf{W}_{h}\mathbf{V}\rVert_{2}.$

Specifically, our SO(3)-invariant scalarization operation is defined as

$q_{ij}=(v_{i_{\mathcal{S}}}\cdot a_{ij}^{t},v_{i_{\mathcal{S}}}\cdot
b_{ij}^{t},v_{i_{\mathcal{S}}}\cdot c_{ij}^{t}),$ (23)

where $\bm{\mathcal{F}}_{ij}^{t}=(a_{ij}^{t},b_{ij}^{t},c_{ij}^{t})$ denotes
the SO(3)-equivariant frame encodings defined in Equations 20 and 21.

To prove that Equation 23 yields SO(3)-invariant scalar features, let $g\in
SO(3)$ be an arbitrary orthogonal transformation. Then we have
$v_{i_{\mathcal{S}}}\rightarrow gv_{i_{\mathcal{S}}}$, and similarly
$\bm{\mathcal{F}}_{ij}^{t}=(a_{ij}^{t},b_{ij}^{t},c_{ij}^{t})\rightarrow(ga_{ij}^{t},gb_{ij}^{t},gc_{ij}^{t})$.
Now, similar to [53], we can derive that Equation 23 becomes

$\displaystyle(v_{i_{\mathcal{S}}}\cdot a_{ij}^{t},v_{i_{\mathcal{S}}}\cdot
b_{ij}^{t},v_{i_{\mathcal{S}}}\cdot c_{ij}^{t})$
$\displaystyle\rightarrow((v_{i_{\mathcal{S}}})^{T}g^{T}ga_{ij}^{t},(v$
${}_{i_{\mathcal{S}}})^{T}g^{T}gb_{ij}^{t},(v_{i_{\mathcal{S}}})^{T}g^{T}gc_{ij}^{t})$
$\displaystyle=(v_{i_{\mathcal{S}}}\cdot a_{ij}^{t},v_{i_{\mathcal{S}}}\cdot
b_{ij}^{t}$ $\displaystyle,v_{i_{\mathcal{S}}}\cdot c_{ij}^{t}),$ (24)

where we used the fact that $g^{T}g=I$ due to the orthogonality of $g$ (with
$I$ being the identity matrix). Therefore, the scalarization operation
proposed in Equation 23, and previously in Equation 3 (in an alternative
form), yields SO(3)-invariant scalars, which is in line with the results of
[89].

The output of Equation 23, $q_{ij}$, is then aggregated in Equation 4 and
concatenated in Equation 5 with the $\mathbf{GCP}$ module’s remaining
O(3)-invariant scalar features (i.e., $L_{2}$ vector norm features). Note that
introducing SO(3)-invariant scalar information into the $\mathbf{GCP}$ module
in this way breaks the 3D reflection symmetry that previous geometric graph
convolution modules enforced [46], now giving rise within the $\mathbf{GCP}$
module to SO(3)-invariant and SO(3)-equivariant updates to scalar and vector-
valued features, respectively. Therefore, scalar and vector-valued feature
updates for nodes and edges within the $\mathbf{GCP}$ module are
SO(3)-invariant and SO(3)-equivariant, respectively.

As in Section 3.3.1, we now turn to discuss the operations within a single
$\mathbf{GCPConv}$ layer, in particular proving that they maintain the
respective SO(3) invariance and SO(3) equivariance for scalar and vector-
valued features that the $\mathbf{GCP}$ module provides.

SO(3)-equivariant GCPConv layer. Via the corresponding proof in [46], by way
of induction all such operations in Equations 10-15 are respectively
SE(3)-invariant and SO(3)-equivariant for features
$m_{ij}^{l}=(m_{e_{ij}}^{l},m_{\xi_{ij}}^{l})$. Thereby, so are features
$n_{i}^{l}=(h_{i}^{l},\chi_{i}^{l})$, given that the proof of equivariance for
the equivariant $\mathbf{LayerNorm}$ and $\mathbf{Dropout}$ operations
employed within each $\mathbf{GCPConv}$ has previously been concretized by
[46]. Equation 17 concludes the operations of a single $\mathbf{GCPConv}$
layer by, as desired, updating the positions of each node $i$ in the 3D input
graph. To do so, $\mathbf{GCPConv}$ residually updates current node positions
$x_{i}^{l-1}$ using SO(3)-equivariant vector-valued features
$\chi_{p_{i}}^{l}$. Therefore, $\mathbf{GCPConv}$ layers are SO(3)-invariant
for scalar feature updates and SO(3)-equivariant for vector-valued node
position and feature updates.

SE(3)-equivariant GCPNet. Lastly, as desired, Line 10 of Algorithm 1 adds
$C(0)$ back to the predicted node positions $\mathbf{X}^{l}$ as provided by
each $\mathbf{GCPConv}$ layer, ultimately imbuing position updates within
$\mathbf{X}^{l}$ with SE(3) equivariance. Line 14 then concludes GCPNet by
using the latest frame encodings $\bm{\mathcal{F}}_{ij}^{t}$ to perform, as
desired, a final SO(3)-invariant and SO(3)-equivariant projection for scalar
and vector-valued features, respectively. Therefore, as desired, GCPNets are
SE(3)-invariant for scalar feature updates, SE(3)-equivariant for vector-
valued node position and feature updates, and, as a consequence, satisfy the
constraint proposed in Def. 1.

∎

###  A.2 Proof of Proposition 2.

Proof. The proof of SE(3) invariance for scalar node and edge features,
$h_{i}$ and $e_{ij}$, follows as a corollary of Appendix A.1
(SE(3)-equivariant GCPNet). Therefore, GCPNets are SE(3)-invariant concerning
their predicted scalar node and edge features and, as a consequence, are
geometrically self-consistent according to the constraint in Def. 2.

∎

###  A.3 Proof of Proposition 3.

Proof. Suppose that GCPNet designates its local geometric representation for
layer $t$ to be
$\bm{\mathcal{F}}_{ij}^{t}=(a_{ij}^{t},b_{ij}^{t},c_{ij}^{t})$, where
$a_{ij}^{t}=\frac{x_{i}^{t}-x_{j}^{t}}{\lVert
x_{i}^{t}-x_{j}^{t}\rVert},b_{ij}^{t}=\frac{x_{i}^{t}\times x_{j}^{t}}{\lVert
x_{i}^{t}\times x_{j}^{t}\rVert},$ and $c_{ij}^{t}=a_{ij}^{t}\times
b_{ij}^{t}$, respectively. As in [53], this formulation of
$\bm{\mathcal{F}}_{ij}^{t}$ is proven in Appendix A.1 (SO(3)-equivariant
frames) to be an SO(3)-equivariant local orthonormal basis at the tangent
space of $x_{i}^{t}$ and is thereby geometrically complete. Note this implies
that GCPNet permits no loss of geometric information as discussed in Appendix
A.5 of [53]. Therefore, GCPNets are geometry-complete and satisfy the
constraint proposed in Def. 3.

∎

##  Appendix B Additional Experiments and Results.

In this section, we explore an additional modeling task, computational protein
design, with its implementation details being discussed in Appendix C.

CPD, Node Classification. Computational protein design (CPD) investigates a
method’s ability to design native-like protein sequences. In our CPD
experiments, we explore fixed-backbone sequence design, where methods are
provided with the 3D backbone structure of a protein and asked to generate a
corresponding sequence. We train and evaluate each CPD method on the CATH 4.2
dataset created by [90]. This dataset contains 18,204, 608, and 1,120
training, validation, and test proteins, respectively, where all available
protein structures with 40% nonredundancy are partitioned by their CATH
(class, architecture, topology/fold, homologous superfamily) classification.
Baseline comparison methods for this task include a mixture of state-of-the-
art Transformers, GNNs, and ENNs.

Under the assumption that native sequences are optimized for their structures
[91], the metrics with which we evaluate each method measure how well a method
can distinguish a native-like sequence from a non-native sequence. In
particular, following [90], we adopt model perplexity as a measure of how well
a method can model the language of native protein sequences. Similarly, we
employ native sequence recovery (i.e., amino acid recovery) rates as a way of
evaluating, on average, how well each method can design sequences that
resemble native protein sequences.

Table 5: Comparison of GCPNet with baseline methods for the CPD task. Results are reported in terms of the perplexity and amino acid recovery rates of each method for fixed-backbone sequence design. The best results for this task are in bold, and the second-best results are underlined. |  | Perplexity $\downarrow$ |  |  | Recovery $\uparrow$ |
---|---|---|---|---|---|---
Method | Short | Single | All | Short | Single | All
STran* [2019] | 8.54 | 9.03 | 6.85 | 28.30 | 27.60 | 36.40
SGNN* [2020] | 8.31 | 8.88 | 6.55 | 28.40 | 28.10 | 37.30
GVP* [2021] | 7.10 | 7.44 | 5.29 | 32.10 | 32.00 | 40.20
GBP* [2022] | 6.14 | 6.46 | 5.03 | 33.22 | 33.22 | 42.70
GCPNet w/o Frames | 7.71 | 8.18 | 5.87 | 31.82 | 31.72 | 41.18
GCPNet w/o ResGCP | 9.63 | 9.88 | 7.09 | 27.27 | 27.02 | 35.27
GCPNet w/o Scalars | 18.51 | 18.37 | 18.17 | 8.70 | 8.55 | 8.62
GCPNet w/o Vectors | 10.41 | 10.53 | 8.87 | 26.42 | 26.23 | 28.99
GCPNet | 8.22 | 8.60 | 6.06 | 33.33 | 32.86 | 40.32

Table 5 shows that, in representing proteins as amino acid residue-level
graphs, GCPNet matches or exceeds the performance of several state-of-the-art
prediction methods for CPD. In particular, GCPNet improves upon state-of-the-
art short sequence recovery rates of previous methods by 0.5% on average while
maintaining competitive performance against other methods in all other
metrics. We note that all CPD methods marked with * perform model inference
autoregressively, introducing a significant computational bottleneck for real-
world applications of these models. Inference with GCPNet, in contrast, is
designed for direct prediction of amino acid sequences corresponding to a 3D
protein structure, thereby decreasing inference runtime by more than a factor
of two compared to other methods. While being a simple direct-shot prediction
method for CPD, GCPNet is still able to achieve competitive results in terms
of amino acid recovery rates for sequence generation, with reasonable results
in terms of perplexity as well.

Interestingly, in the context of CPD, an ablation of our equivariant local
frames $\bm{\mathcal{F}}_{ij}$ reveals that such frames are not useful for
increasing GCPNet’s confidence in its structural understanding of the language
of proteins (i.e., its perplexity). This suggests that future work could
involve exploring alternative geometric encoding schemes for residue-based
graphs when approaching the CPD task [92] with GCPNet. This finding highlights
the fact that the local frames $\bm{\mathcal{F}}_{ij}$ appear to be most
useful in the context of representation learning on atomic graphs where lower-
level molecular motifs are likely to appear, implying that future work towards
improving CPD results with GCPNet could involve developing novel atom-level
encoding schemes for residue-based graph predictions to leverage the promising
results GCPNet yields in other dataset contexts. Nonetheless, our remaining
ablations demonstrate that other design characteristics of GCPNet such as the
$\mathbf{ResGCP}$ module and scalar and vector-valued feature representations
enable GCPNet to better decode sequence-based information from 3D protein
structures.

Table 6: Summary of GCPNet’s node and edge features for 3D input graphs derived for the LBA and PSR tasks. Here, $N$ and $E$ denote the number of nodes and edges in $\mathcal{G}$, respectively. | Feature | Type | Shape
---|---|---|---
Node Features ($h$) | One-hot encoding of atom type | Categorical (Scalar) | $N\times 9$
Node Features ($\chi$) | Directional encoding of orientation | Numeric (Vector) | $N\times 2$
Edge Features ($e$) | Radial basis distance embedding | Numeric (Scalar) | $E\times 16$
Edge Features ($\xi$) | Pairwise atom position displacement | Numeric (Vector) | $E\times 1$
Total | Node features |  | $N\times 11$
| Edge features |  | $E\times 17$
Table 7: Summary of GCPNet’s node and edge features for 3D input graphs derived for the CPD task. | Feature | Type | Shape
---|---|---|---
Node Features ($h$) | Dihedral angle encoding | Numeric (Scalar) | $N\times 6$
Node Features ($\chi$) | Orientation and sidechain encoding | Numeric (Vector) | $N\times 3$
Edge Features ($e$) | Distance and positional embedding | Numeric (Scalar) | $E\times 32$
Edge Features ($\xi$) | Pairwise atom position displacement | Numeric (Vector) | $E\times 1$
Total | Node features |  | $N\times 9$
| Edge features |  | $E\times 33$
Table 8: Summary of GCPNet’s node and edge features for 3D input graphs derived for the NMS task. | Feature | Type | Shape
---|---|---|---
Node Features ($h$) | Invariant velocity encoding | Numeric (Scalar) | $N\times 1$
Node Features ($\chi$) | Velocity and orientation encoding | Numeric (Vector) | $N\times 3$
Edge Features ($e$) | Edge and distance embedding | Numeric (Scalar) | $E\times 17$
Edge Features ($\xi$) | Pairwise atom position displacement | Numeric (Vector) | $E\times 1$
Total | Node features |  | $N\times 4$
| Edge features |  | $E\times 18$

##  Appendix C Implementation Details.

Featurization. As shown in Table 6, for the LBA and PSR tasks, in each 3D
input graph, we include as a scalar node feature an atom’s type using a
9-dimensional one-hot encoding vector for each atom. As vector-valued node
features, we include forward and reverse unit vectors in the direction of
$x_{i+1}-x_{i}$ and $x_{i-1}-x_{i}$, respectively (i.e., the node’s 3D
orientation). For the input 3D graphs’ scalar edge features, we encode the
distance $\lVert x_{i}-x_{j}\rVert_{2}$ using Gaussian radial basis functions,
where we use 16 radial basis functions with centers evenly distributed between
0 and 20 units (e.g., Angstrom). For the graphs’ vector-valued edge features,
we encode the unit vector in the direction of $x_{i}-x_{j}$ (i.e., pairwise
atom position displacements).

As displayed in Table 7, for the CPD task, in each 3D input graph, we include
as scalar node features an encoding of each amino acid residue’s dihedral
angles $\\{\sin,\cos\\}\circ\\{\phi,\psi,\omega\\}$, where $\bm{\Phi}$,
$\psi$, and $\omega$ are the dihedral angles computed from the corresponding
protein’s $C_{i-1}$, $N_{i}$, $C_{i}$, and $N_{i+1}$ backbone atoms. We then
include as vector-valued node features each node’s 3D orientation. For edge
features, we use Gaussian radial basis function distance encodings as scalar
edge features and pairwise atom position displacements as vector-valued edge
features.

As illustrated in Table 8, for the NMS task, in each 3D input graph, we
include as a scalar node feature an invariant encoding of each node’s velocity
vector, namely $\sqrt{v_{i}^{2}}$. Each node’s velocity and orientation are
encoded as vector-valued node features. Scalar edge features are represented
as Gaussian radial basis distance encodings as well as the product of the
charges in each node pair (i.e., $c_{i}c_{j}$). Lastly, vector-valued edge
features are represented as pairwise atom position displacements.

Hardware Used. The Oak Ridge Leadership Facility (OLCF) at the Oak Ridge
National Laboratory (ORNL) is an open science computing facility that supports
HPC research. The OLCF houses the Summit compute cluster. Summit, launched in
2018, delivers 8 times the computational performance of Titan’s 18,688 nodes,
using only 4,608 nodes. Like Titan, Summit has a hybrid architecture, and each
node contains multiple IBM POWER9 CPUs and NVIDIA Volta GPUs all connected
with NVIDIA’s high-speed NVLink. Each node has over half a terabyte of
coherent memory (high bandwidth memory + DDR4) addressable by all CPUs and
GPUs plus 800GB of non-volatile RAM that can be used as a burst buffer or as
extended memory. To provide a high rate of I/O throughput, the nodes are
connected in a non-blocking fat-tree using a dual-rail Mellanox EDR InfiniBand
interconnect. We used the Summit compute cluster to train all our models. For
the LBA and NMS tasks, we used 16GB NVIDIA Tesla V100 GPUs for model training,
whereas for the memory-intensive PSR and CPD tasks, we used 32GB V100 GPUs
instead.

Software Used. We used Python 3.8.12 [93], PyTorch 1.10.2 [94], PyTorch
Lightning 1.7.7 [95], and PyTorch Geometric 2.1.0post0 [96] to run our deep
learning experiments. For each model trained, PyTorch Lightning was used to
facilitate model checkpointing, metrics reporting, and distributed data
parallelism across 6 V100 GPUs. A more in-depth description of the software
environment used to train and run inference with our models is available at
https://github.com/BioinfoMachineLearning/GCPNet.

Hyperparameters. As shown in Tables 9, 10, 11, and 12, we use a learning rate
of $10^{-4}$ for all GCPNet models. The learning rate is kept constant
throughout each model’s training. For the NMS task, each model is trained for
a minimum of 100 epochs and a maximum of 12,000 epochs. For all other tasks,
each model is trained for a minimum of 100 epochs and a maximum of 1,000
epochs. For a given task, models with the best loss on the corresponding
validation data split are then tested on the test split for the respective
task. Note that, for the RS task, we do not perform any model hyperparameter
tuning, following previous conventions from [59].

Table 9: Hyperparameter search space for all GCPNet models through which we searched to obtain strong performance on the LBA task’s validation split. The final parameters for the standard GCPNet model for the LBA task are in bold. Hyperparameter | Search Space
---|---
Number of GCPNet Layers | 7, 8
Number of $\mathbf{GCP}$ Message-Passing Layers | 8
$\chi$ Hidden Dimensionality |  16, 32
Learning Rate |  0.0001, 0.0003
Weight Decay Rate | 0
$\mathbf{GCP}$ Dropout Rate |  0.1, 0.25
Dense Layer Dropout Rate |  0.1, 0.25
Table 10: Hyperparameter search space for all GCPNet models through which we searched to obtain strong performance on the PSR task’s validation split. The final parameters for the standard GCPNet model for the PSR task are in bold. Hyperparameter | Search Space
---|---
Number of GCPNet Layers | 5
Number of $\mathbf{GCP}$ Message-Passing Layers | 8
$\chi$ Hidden Dimensionality |  16, 32
Learning Rate |  0.0001, 0.0003
Weight Decay Rate |  0, 0.0001
$\mathbf{GCP}$ Dropout Rate |  0.1, 0.25
Dense Layer Dropout Rate |  0.1, 0.25
Table 11: Hyperparameter search space for all GCPNet models through which we searched to obtain strong performance on the CPD task’s validation split. The final parameters for the standard GCPNet model for the CPD task are in bold. Hyperparameter | Search Space
---|---
Number of GCPNet Encoder Layers | 9
Number of GCPNet Decoder Layers | 3
Number of $\mathbf{GCP}$ Message-Passing Layers | 8
$\chi$ Hidden Dimensionality |  16, 32
Learning Rate | 0.0001
Weight Decay Rate | 0.0, $\mathbf{10}^{\mathbf{-8}}$, 0.0001
$\mathbf{GCP}$ Dropout Rate | 0.1, 0.2, 0.25, 0.4
Decoder Residual Updates | False, True
Table 12: Hyperparameter search space for all GCPNet models through which we searched to obtain strong performance on the NMS task’s validation split. The final parameters for the standard GCPNet model for the NMS task are in bold. Hyperparameter | Search Space
---|---
Number of GCPNet Layers |  4, 7
Number of $\mathbf{GCP}$ Message-Passing Layers | 8
$\chi$ Hidden Dimensionality | 16
Learning Rate |  0.0001, 0.0003
Weight Decay Rate | 0
$\mathbf{GCP}$ Dropout Rate | 0.0, 0.1