# Improving Molecular Modeling with Geometric GNNs: an Empirical Study

Ali Ramlaoui  Théo Saulus  Basile Terver  Victor Schmidt  David Rolnick
Fragkiskos D. Malliaros  Alexandre Duval

###### Abstract

Rapid advancements in machine learning (ML) are transforming materials science
by significantly speeding up material property calculations. However, the
proliferation of ML approaches has made it challenging for scientists to keep
up with the most promising techniques. This paper presents an empirical study
on Geometric Graph Neural Networks for 3D atomic systems, focusing on the
impact on performance, scalability, and symmetry enforcement of different (1)
canonicalization methods, (2) graph creation strategies, and (3) auxiliary
tasks. Our findings and insights aim to guide researchers in selecting optimal
modeling components for molecular modeling tasks. Our code is available at
https://github.com/RolnickLab/ocp.

Machine Learning, ICML

##  1 Introduction

The field of computational materials science has witnessed an increasing
interest in recent years, with the explosion of machine learning approaches to
model material properties at the quantum scale. This is possible thanks to the
release of large-scale datasets such as OC20 (Chanussot et al., 2020) and
QM7-X (Hoja et al., 2021), which contain millions of molecular structures
along with various properties (forces, energy, band gap) computed using
quantum mechanical simulations involving Density Functional Theory (DFT) (Kohn
et al., 1996).

ML models have been trained to approximate DFT, thus constituting an even
faster proxy to the Schrödinger equation, modeling atomic interactions and
systems’ behavior. They enable the calculation of material properties in
seconds instead of hours or days, offering the potential to accelerate
scientific discovery via high-throughput screening of novel materials (Batatia
et al., 2023a; Merchant et al., 2023). However, the explosion of ML approaches
in recent years makes it hard to keep up with promising techniques for
scientists in the field. While some surveys have attempted to structure the
different categories of ML approaches (Han et al., 2022; Duval et al., 2024),
these do not focus on empirical evaluation.

In this paper, we propose an empirical study of some key modeling aspects of
Geometric GNNs for 3D atomic systems. Specifically, we investigate the impact
of 1) the recent canonicalization methods used to enforce or approximate
Euclidean symmetries, 2) the graph creation step when modeling a 3D atomic
system, and 3) adding several auxiliary tasks. We focus on the OC20 dataset
modeling the relaxed adsorption energy of an adsorbate-catalyst system. We
hope that the conclusions and insights drawn from our experiments will benefit
the community, making it possible to quickly choose the right modeling
component.

##  2 Choice of Canonicalization

A function $f:\mathcal{X}\rightarrow\mathcal{X}$ is said to be equivariant
with respect to a transformation $t$ if $\forall
x\in\mathcal{X},f(t(x))=t(f(x))$. In particular, $f$ is $E(3)$-equivariant if
it is equivariant to rotations, translations, and reflections.
$E(3)$-equivariance is a desirable property in molecular modeling to learn
representations that are best suited for physically meaningful tasks, such as
force predictions on atoms (e.g., S2EF task of OC20). This can be enforced in
the architecture of the model during the message passing steps by using
equivariant features of the input’s representation (Schütt et al., 2021;
Batatia et al., 2023b; Liao & Smidt, 2023), which comes at the cost of
expensive feature computations. A recent alternative to these equivariant
architectures lies in unconstrained GNNs, which do not enforce
$E(3)$-equivariance by model design but instead with a coordinate-
preprocessing step referred to as canonicalization (Hu et al., 2021; Duval et
al., 2023; Pozdnyakov & Ceriotti, 2024). This process grants unconstrained
GNNs with greater flexibility, scalability, and often expressivity (Duval et
al., 2024). Commonly, it involves projecting (e.g., with PCA or an equivariant
network) the input atomic system into a canonical space such that every
Euclidean transformation of the same system gets mapped to the same canonical
representation, i.e. handling symmetries in the data pre-processing.

Since these canonicalization methods are all recent and no comparison has yet
been drawn, we benchmark the proposed approaches on QM9 and OC20 tasks,
evaluating their impact on performance and symmetry enforcement.

###  2.1 Canonicalization for direct predictions

In this subsection, we evaluate several canonicalization procedures with the
FAENet backbone architecture (Duval et al., 2023), a powerful unconstrained
GNN which is not equivariant by itself. We assume familiarity of the reader
with these approaches but provide a description in Appendix A.

  * •

Vector Neurons Network (VNN) (Deng et al., 2021) using the VN re-
implementation of PointNet (Qi et al., 2017) and DGCNN (Wang et al., 2019).
This class of canonicalization networks is $SO(3)$-equivariant by design and
is applied to obtain a canonical representation of the data following the
method of Kaba et al. (2023).

  * •

Stochastic Frame Averaging (SFA) (Duval et al., 2023), designed to avoid
averaging the predictions over $8$ elements of the frame as required by Frame
Averaging (Puny et al., 2021). Instead, we sample one canonical orientation at
random at each epoch, similar to data augmentation on a small and complete
set.

  * •

A novel method denoted SFA+SignNet, where we propose to map the several
projection matrices of SFA to a single one using a sign-equivariant network
proposed in SignNet (Lim et al., 2023b, a). The rationale is to handle the
sign ambiguity problem of PCA that exists in SFA with a small dedicated
network to get a unique canonical representation of Euclidean transformations.
We propose two implementations of SignNet, either using VNNs to have a perfect
$E(3)$-equivariance when combined with SFA or using MLPs without theoretical
guarantees.

For VNNs and SFA+SignNet methods, we test both training and freezing the
weights of the Canonicalization Network (CN). We report performance on the
OC20 IS2RE, OC20 S2EF, and QM9 in Tables 1, 5, 6, 7, 8, and 9.

In our OC20 IS2RE experiment, we found there are almost no differences between
the various canonicalization methods, with MAE of 594, 598, and 593 for SFA,
VN-PointNet, and VN-DGCNN. Specially, non-exact canonicalizations (SFA and
SFA+SignNet) demonstrate equal or better MAE than perfectly equivariant
methods (e.g., VN-based). This suggests that heuristics approximation of
equivariance should be sufficient in some practical applications like OC20.
This is aligned with what Duval et al. (2023) suggested when showing that SFA
outperforms exact Frame Averaging in terms of downstream performance.

In terms of symmetry enforcement, non-exact methods are surprisingly nearly as
effective as fully invariant methods, suggesting that the FAENet backbone
implicitly learns to handle symmetries.

Regarding exact canonicalization methods, we observe that training or not the
network and swapping one method for another has little impact on model
performance. This tends to indicate that the canonical networks’s ability to
introduce equivariance is more critical than the choice of the
canonicalization method.

Canonicalization | Cano. trained | avg. MAE | EwT (ID) | 3D Rotation
---|---|---|---|---
| parameters | (meV) $\downarrow$ | (%) $\uparrow$ | Invariance $\downarrow$
SFA | 0 | 594 | 4.40 | $1.30\cdot 10^{-2}$
(U) SFA+MLP-SignNet | 0 | 580 | 4.48 | $9.71\cdot 10^{-2}$
(T) SFA+MLP-SignNet | 454 | 583 | 4.46 | $4.00\cdot 10^{-2}$
(U) SFA+VN-SignNet | 0 | 592 | 4.69 | $7.58\cdot 10^{-3}$
(T) SFA+VN-SignNet | 2,620 | 599 | 4.25 | $2.57\cdot 10^{-2}$
(U) VN-Pointnet | 0 | 605 | 4.09 | $4.62\cdot 10^{-3}$
(T) VN-Pointnet | 1,310 | 598 | 4.12 | $\mathbf{3.80\cdot 10^{-3}}$
(U) VN-DGCNN | 0 | 600 | 4.31 | $3.11\cdot 10^{-2}$
(T) VN-DGCNN | 663,804 | 593 | 4.42 | $9.10\cdot 10^{-3}$

Table 1: Invariance comparison of canonicalization methods on OC20 IS2RE
dataset. (U) (resp. (T)) indicates an untrained (resp. trained)
canonicalization network. FAENet backbone has 4,147,731 parameters. More
details in Tables 5 and 6.

###  2.2 Canonicalization for relaxed IS2RE

Previous work showed that solving the IS2RE task yields better results by
performing relaxed energy predictions rather than direct energy estimation
(Liao et al., 2023). Here, we evaluate whether canonicalization methods also
perform well at relaxing a trajectory. Table 2 reports the performances of
FAENet with multiple symmetry-preserving methods, with the invariant SchNet
model and direct IS2RE acting as baseline. Our findings suggest that relaxed
IS2RE predictions are competitive with direct IS2RE predictions and are
interesting directions to explore for improving molecular property predictions
with these architectures. Moreover, a potential explanation as to why the
relaxations do not yield significant improvements over direct IS2RE may
involve the approximate equivariance or the lack of continuity in some of
these canonicalization methods, as pointed out in Dym et al. (2024), which may
hamper accurate and smooth relaxations trajectories. This is mainly true for
SFA, where a frame is randomly picked at each step of the relaxation, meaning
that the canonical inputs can also be far from each other. Lastly, note that
exact equivariant methods for relaxed IS2RE predictions such as the untrained
VN-PointNet (UTPN) implementation from (Kaba et al., 2023) yields better MAE
than the approximate equivariance module SFA. Still, during our experiments we
found that with further focus on the accuracy of S2EF models, canonicalization
methods used to enforce symmetries can become more appealing using iterative
relaxation methods.

| IS2RE | IS2RS
---|---|---
Model | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$ | DwT (%) $\uparrow$ | Pos. MAE $\downarrow$
FAENet (Direct) | 4.05 | 0.551 | - | -
FAENet (SFA) | 4.92 | 0.587 | 31.1 | 0.390
FAENet (UTPN) | 5.64 | 0.560 | 33.7 | 0.381
SchNet | 1.89 | 0.912 | 15.0 | 0.461

Table 2: Results on the IS2RE and IS2RS tasks for the VAL-ID validation
dataset of OC20 using iterative relaxations. The S2EF models’ results are
reported in Appendix A.5. Note that these results can only be obtained by
keeping the tag 0 atoms of the subsurface as they are important for
relaxations.

##  3 Graph Creation Study

For large molecular structures, electrostatic long-range interactions are non-
negligible components of the system’s dynamics (Gasteiger et al., 2019). While
various methods have tried to model long-range interactions between far-away
atoms, they often suffer from the over-smoothing effect when increasing the
number of interaction layers (Liao & Smidt, 2023). As a result, models where
the geometric graph is modelled using a cutoff distance between neighbour
atoms have been shown to work best. Since properly handling these interactions
is essential to accurately simulate the system, we explore in this section the
impact of this creation step and the rewiring strategies.

###  3.1 Graph cutoff and rewiring

First, we vary the cutoff distance111When representing the 3D point cloud with
a graph, we create an edge between any two atoms if their are within a fixed
cutoff distance, and no edge otherwise. used during the creation of the graph
to check whether linking more atoms with each other helps in learning these
couplings. We report the results for the FAENet model on the OC20 IS2RE task
in Table 3.

A small cutoff of 1.0 Å leads to the weakest performance, which makes sense
since the associated graph is almost empty and nodes are isolated, i.e.
messages can not pass correctly. A large cutoff or a complete graph also leads
to poor performance despite every atomic interaction being modeled. Within an
intermediary range of cutoff values, the model achieves optimal learning with
computational efficiency. This is in accordance with past observations where
the locality bias of GNN models seemed to fit really well with atomic system
property prediction tasks. Thus, although the creation of the graph through a
well-chosen cutoff is important, the margin for fine-tuning this parameter is
large enough.

The fact that better graphs are adapted to GNN functioning rather than the
precise modeling of the situation, where all atoms interact with each other,
echoes what Duval et al. (2022) have stated. Indeed, they showed that to fit
the GNN message passing scheme, removing repeating subsurface atoms of the
adsorbate did not affect the model performance on IS2RE tasks. Similarly to
having a moderate cutoff, the performance improvements of such a strategy are
decisive for scalability. As an empirical example, we tried to run an
EquiformerV2 (Liao et al., 2023) model for IS2RE with the remove-tag-0 method
on an 80GB A100 GPU, leading to a 5$\times$ speed-up with no performance loss
and proving the relevance of this technique even on very large models.

| ID
---|---
Model | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$
Cutoff 30 - Max neighbours 40 | 2.65 | 0.697
Cutoff 20 - Max. neighbours 40 | 3.08 | 0.673
Cutoff 20 - Max. neighbours 10 | 2.25 | 0.768
Cutoff 10 - Max. neighbours 50 | 4.17 | 0.553
Cutoff 10 - Max. neighbours 10 | 4.49 | 0.553
Cutoff 6 - Max. neighbours 40 | 4.31 | 0.553
Cutoff 1 - Max. neighbours 40 | 1.35 | 1.069

Table 3: Impact of the cutoff on the performance of FAENet on the OC20 IS2RE
task. Full Table in Appendix B.1.

###  3.2 Ewald-based long range message passing

Since a small value for the cutoff seems to be the most interesting one, we
want to model the long interactions differently than adding links between all
atoms. Ewald-based Message Passing (EMP) (Kosmala et al., 2023) is introduced
in this perspective. It incorporates a physics-based prior in the architecture
to model the long-range electrostatic potential via a nonlocal Fourier space
scheme, drawing edges based on a cutoff on frequency instead of distances.

Our experiments, given in Table 4, show that EMP is interesting for an
invariant method like SchNet (Schütt et al., 2017), which limits its geometric
information to atom pairwise distances. However, EMP does not benefit more
advanced GNNs like FAENet. To understand why, we plot in Figure 1 the cosine
similarity between the embeddings throughout interaction layers. We observe
that SchNet and FAENet learn very different representations. Indeed, while the
representation of each atom in SchNet tends to be similar to nearby atoms,
FAENet is able to give very different embeddings for them. A potential
explanation could be that because FAENet is a more expressive model, the
propagated messages are less constrained and thus can lead to very diverse
atom representations without Ewald. On the other hand, embeddings of SchNet
only become diversified using Ewald, maybe because incorporating longer
messages helps create more distinct representations. More plots can be found
in Appendix B.2. One interesting takeaway is that EMP helps take into account
long-range interactions for simple models and GNNs with symmetry-constrained
layers (Kosmala et al., 2023) but is less efficient on expressive models.

| ID
---|---
Model | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$
FAENet (Graph Rewiring) | 4.05 | 0.551
FAENet (Graph Rewiring) + Ewald | 4.12 | 0.562
FAENet (No Graph Rewiring) | 4.54 | 0.544
FAENet (No Graph Rewiring) + Ewald | 4.11 | 0.556
SchNet (Graph Rewiring) | 3.18 | 0.641
SchNet (Graph Rewiring) + Ewald | 3.54 | 0.604
SchNet (No Graph Rewiring) | 2.93 | 0.654
SchNet (No Graph Rewiring) + Ewald | 3.48 | 0.597

Table 4: Comparison of the performances of FAENet and SchNet with and without
Ewald Message Passing on the IS2RE task. Full table and results for the QM9
dataset in Appendix B.2.

(a) SchNet without Ewald

(b) FAENet without Ewald

(c) SchNet with Ewald

(d) FAENet with Ewald

Figure 1: Similarity matrix of the embeddings of the atoms of a randomly
picked system for different interaction blocks from the training set of OC20.
The same system is used every time to be able to compare the different
results.

##  4 Auxiliary Tasks

In this section, we study how to leverage other tasks to improve the
performance of FAENet on IS2RE.

###  4.1 Noisy nodes IS2RS auxiliary task

First, we recall that increasing the number of interaction blocks above 8 does
not yield better performance for most geometric GNNs, with a dramatic loss of
information after 14 layers in the classical FAENet IS2RE setup, as
illustrated in Table 14. To address the oversmoothing issue, Godwin et al.
(2022) propose to use noisy regularisation, introducing their method called
Noisy Nodes. It consists of adding an auxiliary node-level denoising task that
encourages diversity in the latent representations of the nodes (more details
in C.2).

Here, we implement Noisy Nodes for the IS2RE task, which is done by adding a
position decoding head running in parallel to the original energy prediction
head. Our implementation takes inspiration from other state-of-the-art models
such as EquiformerV1 (Liao & Smidt, 2023), which benefit from using Noisy
Nodes on the IS2RE task (more details in C.2).

To understand the reason for the performance drop for the classical FAENet
IS2RE setup of (Duval et al., 2023) when adding interaction layers in Table
14, we plot the Mean Average Distance (MAD) (Chen et al., 2019), averaged over
50 input adsorbate-catalyst pairs, of their embeddings throughout the
interaction blocks in Figure 2. The classical FAENet models (”FAENet top”)
with 14 interaction layers or more see their latent node representations all
collapse to almost the same vector, since the MAD goes to almost zero as we go
deeper in the model’s interaction layers, which is a manifestation of
oversmoothing. Figure 2 shows that the models trained with Noisy Nodes IS2RS
auxiliary task (”FAENet aux”) do not suffer from oversmoothing (i.e. MAD going
to zero) even when going as deep as 28 interaction layers.

Figure 2: MAD values of the graph embeddings (averaged over 50 randomly
sampled graphs of the train set) throughout the interaction layers for various
models. “FAENet top” models are trained with the top configs of the classical
FAENet model (Duval et al., 2023) but with more epochs and lower batch size
(128). “FAENet aux” models are our models trained on IS2RE with IS2RS
auxiliary task. A model having xx interaction layers is indicated as “XXi”.

Then, we compare the performances of our new FAENet models trained on IS2RE
Noisy Nodes, varying the number of interaction layers, as summarized in Table
15. First, we observe a clear correlation between the performances and the
number of blocks. Second, the best model trained with Noisy Nodes, which has
26 interaction blocks, outperforms the best models of (Duval et al., 2023),
allowing the average MAE to decrease from 568 meV to 525 and the average EwT
percentage to increase from 3.78% to 4.43%. Details about our experimental
setup are given in Appendix C.1. We observe that the throughput at inference
is divided by two between the smallest (5 interaction layers) and biggest
model (28 layers). The best tradeoff seems to be around 16 interaction layers,
similar to what (Liao & Smidt, 2023) choose for their models trained with
IS2RS auxiliary task.

The current most general state-of-the-art models for atomic property
prediction, such as (Shoghi et al., 2024) only make use of backbone GNNs (such
as GemNet-OC (Gasteiger et al., 2022)) with not more than 6 interaction
blocks. Noisy Nodes proves to be a regularisation technique that can be used
across models, having equivariant (Liao & Smidt, 2023) or non-equivariant
(Duval et al., 2023) features, to leverage the power of more interaction
layers and improve performance. On direct IS2RE, our experiments only show a
slightly increasing performance with the number of blocks. Yet, unlocking the
potential of deeper GNNs with this simple regularisation allows greater
freedom in training or pretraining on larger datasets or for more complex
tasks, such as relaxed IS2RE 2.2. Our focus for FAENet was on IS2RE, but we
expect similar results on S2EF, IS2RS, and other datasets such as QM9 or
QM7-X, as obtained by (Godwin et al., 2022) with other GNNs such as GNS
(Sanchez-Gonzalez et al., 2020).

###  4.2 Learning equivariance with more interaction layers

As previously stated, equivariant networks are preferred when handling 3D
point clouds. However, in the field of protein structure prediction, we were
caught by the AlphaFold3 architecture (Abramson et al., 2024) that chose non-
equivariant networks in its diffusion module, contrary to AlphaFold2 (Jumper
et al., 2021). Thus, we tested whether a deep GNN with no equivariance
enforced could learn the equivariance from the data and match performances
with canonicalized models using the FAENet backbone, as in Section 2.

We compare our models trained on IS2RE with IS2RS auxiliary task in two
settings: with SE(3)-SFA and without any canonicalization method (No-FA). The
results are displayed in Table 16. They reveal that, even with many additional
interaction blocks (up to 26), both the MAE and the equivariance property do
not improve compared to using SE(3)-SFA, meaning that imposing equivariance is
still a beneficial inductive bias. Unfortunately, in our setting, having a
deeper network does not allow to learn invariance and equivariance more
effectively. Further analysis would be needed with larger datasets and other
methods to reinforce this argument or find its limits (if any).

###  4.3 Pre-training on different tasks

With the release of larger datasets, the community has increasingly shifted
towards pre-training and transfer learning approaches (Batatia et al., 2023a;
Shoghi et al., 2024; Deng et al., 2023). We make a step in this direction in
this subsection by pre-training our model on the S2EF, which contains roughly
two orders of magnitude more data points, and by fine-tuning it on IS2RE,
hoping to transfer some knowledge of atomic interactions. More precisely, we
leverage the extensive S2EF dataset to utilize all trajectories while
maintaining high inference throughput by performing direct energy predictions.
Although training on S2EF is more time-consuming than on IS2RE, separating the
learning of relaxation and molecular interactions may yield better results
than training IS2RE directly from scratch.

Figure 7 in Appendix C.6 shows that with this approach, the energy MAE starts
at a better value during training but evolves much slower and converges to a
slightly better result (0.53 eV vs. 0.55 eV on the validation ID split). This
shows that while the model starts with a good representation of molecular
dynamics, the differences are not significant because dynamics are not taken
into account by the S2EF training process. Learning molecular interactions
through auxiliary tasks with or without dynamics is helpful to achieve better
performances, but this pre- or joint training needs to be correctly
incorporated into the architecture so as not to overwrite the learned
information during the downstream task. This opens the way for new
architectures designed to leverage materials design knowledge between tasks
and datasets (Shoghi et al., 2024). Whether by training with auxiliary tasks
or on other datasets, there seems to be transfer learning and generalization
capabilities in atomic property prediction, as in NLP. Hence, we recommend to
further explore this promising area of research.

##  5 Conclusion

In this study, we explored various techniques aimed at enhancing the
performance of geometric GNNs for molecular modeling. Our empirical study
covered several aspects, with the main observations being summarized below.

##### Canonicalization methods.

While exact methods provide the best theoretical guarantees for equivariance,
approximative heuristics such as SFA seem to yield better performance. This
opens questions about how to design canonicalizations that are the most
effective in practice, beyond theoretical guarantees, and leaves a broader set
of possibilities for the model design, too.

##### Graph creation.

Although accurate graph construction is important, many viable options can be
considered without significant differences in performance. Furthermore,
physics-inspired modules such as Ewald-based message passing demonstrate
improved performance for symmetry-constrained models such as SchNet but do not
provide any benefits to more expressive models such as FAENet.

##### Auxiliary tasks.

Implementing Noisy Nodes as an auxiliary task significantly enhances the
performance of FAENet by leveraging the benefits of much deeper GNNs. As with
pretraining on different tasks such as S2EF, there is evidence of transfer
learning for atomic property prediction, and we recommend more exploration of
this path in the flavor of (Shoghi et al., 2024).

Future work could focus on refining these techniques, exploring their
applications across a wider spectrum of datasets, and developing new methods
to combine the strengths of various approaches.

## Acknowledgement

The authors thank Sékou-Oumar Kaba, Derek Lim, Joshua David Robinson, and
Hannah Lawrence for their insightful comments and discussions, as well as the
anonymous reviewers for their suggestions and feedback. Supported in part by
ANR (French National Research Agency) under the JCJC project GraphIA
(ANR-20-CE23-0009-01) and the Canada CIFAR AI Chairs program. This research
was enabled in part by computing resources provided by Mila (mila.quebec) and
material support from NVIDIA Corporation in the form of computational
resources.

## References

  * Abramson et al. (2024) Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., et al.  Accurate structure prediction of biomolecular interactions with alphafold 3.  _Nature_ , pp. 1–3, 2024.
  * Batatia et al. (2023a) Batatia, I., Benner, P., Chiang, Y., Elena, A. M., Kovács, D. P., Riebesell, J., Advincula, X. R., Asta, M., Baldwin, W. J., Bernstein, N., et al.  A foundation model for atomistic materials chemistry.  _arXiv preprint arXiv:2401.00096_ , 2023a.
  * Batatia et al. (2023b) Batatia, I., Kovács, D. P., Simm, G. N. C., Ortner, C., and Csányi, G.  Mace: Higher order equivariant message passing neural networks for fast and accurate force fields, 2023b.
  * Chanussot et al. (2020) Chanussot, L., Das, A., Goyal, S., Lavril, T., Shuaibi, M., Rivière, M., Tran, K., Heras-Domingo, J., Ho, C., Hu, W., Palizhati, A., Sriram, A., Wood, B., Yoon, J., Parikh, D., Zitnick, C. L., and Ulissi, Z. W.  The open catalyst 2020 (oc20) dataset and community challenges.  _ACS Catalysis_ , 2020.  doi: 10.1021/acscatal.0c04525.
  * Chen et al. (2019) Chen, D., Lin, Y., Li, W., Li, P., Zhou, J., and Sun, X.  Measuring and relieving the over-smoothing problem for graph neural networks from the topological view, 2019.
  * Deng et al. (2023) Deng, B., Zhong, P., Jun, K., Riebesell, J., Han, K., Bartel, C. J., and Ceder, G.  Chgnet: Pretrained universal neural network potential for charge-informed atomistic modeling.  _arXiv preprint arXiv: 2302.14231_ , 2023.
  * Deng et al. (2021) Deng, C., Litany, O., Duan, Y., Poulenard, A., Tagliasacchi, A., and Guibas, L. J.  Vector neurons: A general framework for so (3)-equivariant networks.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 12200–12209, 2021.
  * Duval et al. (2022) Duval, A., Schmidt, V., Miret, S., Bengio, Y., Hernández-García, A., and Rolnick, D.  Phast: Physics-aware, scalable, and task-specific gnns for accelerated catalyst design.  _arXiv preprint arXiv: 2211.12020_ , 2022.
  * Duval et al. (2023) Duval, A., Schmidt, V., Garcia, A., Miret, S., Malliaros, F. D., Bengio, Y., and Rolnick, D.  Faenet: Frame averaging equivariant gnn for materials modeling.  _International Conference on Machine Learning_ , 2023.  doi: 10.48550/arXiv.2305.05577.
  * Duval et al. (2024) Duval, A., Mathis, S. V., Joshi, C. K., Schmidt, V., Miret, S., Malliaros, F. D., Cohen, T., Liò, P., Bengio, Y., and Bronstein, M.  A hitchhiker’s guide to geometric gnns for 3d atomic systems.  2024\.
  * Dym et al. (2024) Dym, N., Lawrence, H., and Siegel, J. W.  Equivariant frames and the impossibility of continuous canonicalization, 2024.
  * Gasteiger et al. (2019) Gasteiger, J., Groß, J., and Günnemann, S.  Directional message passing for molecular graphs.  In _International Conference on Learning Representations_ , 2019.
  * Gasteiger et al. (2022) Gasteiger, J., Shuaibi, M., Sriram, A., Günnemann, S., Ulissi, Z., Zitnick, C. L., and Das, A.  Gemnet-oc: Developing graph neural networks for large and diverse molecular simulation datasets, 2022.
  * Godwin et al. (2022) Godwin, J., Schaarschmidt, M., Gaunt, A., Sanchez-Gonzalez, A., Rubanova, Y., Veličković, P., Kirkpatrick, J., and Battaglia, P.  Simple gnn regularisation for 3d molecular property prediction & beyond, 2022.
  * Han et al. (2022) Han, J., Rong, Y., Xu, T., and Huang, W.  Geometrically equivariant graph neural networks: A survey.  _arXiv preprint arXiv:2202.07230_ , 2022.
  * Hoja et al. (2021) Hoja, J., Sandonas, L. M., Ernst, B. G., Vázquez-Mayagoitia, A., Jr., R. A. D., and Tkatchenko, A.  Qm7-x: A comprehensive dataset of quantum-mechanical properties spanning the chemical space of small organic molecules.  _Scientific Data_ , 2021.  doi: 10.1038/s41597-021-00812-2.
  * Hu et al. (2021) Hu, W., Shuaibi, M., Das, A., Goyal, S., Sriram, A., Leskovec, J., Parikh, D., and Zitnick, C. L.  Forcenet: A graph neural network for large-scale quantum calculations.  _arXiv preprint arXiv: 2103.01436_ , 2021.
  * 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.
  * Kaba et al. (2023) Kaba, S.-O., Mondal, A. K., Zhang, Y., Bengio, Y., and Ravanbakhsh, S.  Equivariance with learned canonicalization functions.  In _International Conference on Machine Learning_ , pp. 15546–15566. PMLR, 2023.
  * Kohn et al. (1996) Kohn, W., Becke, A. D., and Parr, R. G.  Density functional theory of electronic structure.  _The Journal of Physical Chemistry_ , 100(31):12974–12980, 1996.
  * Kosmala et al. (2023) Kosmala, A., Gasteiger, J., Gao, N., and Günnemann, S.  Ewald-based long-range message passing for molecular graphs.  _arXiv preprint arXiv:2303.04791_ , 2023.
  * Liao & Smidt (2023) Liao, Y.-L. and Smidt, T.  Equiformer: Equivariant graph attention transformer for 3d atomistic graphs, 2023.
  * Liao et al. (2023) Liao, Y.-L., Wood, B., Das, A., and Smidt, T.  Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations, 2023.
  * Lim et al. (2023a) Lim, D., Robinson, J., Jegelka, S., and Maron, H.  Expressive sign equivariant networks for spectral geometric learning, 2023a.
  * Lim et al. (2023b) Lim, D., Robinson, J. D., Zhao, L., Smidt, T., Sra, S., Maron, H., and Jegelka, S.  Sign and basis invariant networks for spectral graph representation learning.  In _The Eleventh International Conference on Learning Representations_ , 2023b.  URL https://openreview.net/forum?id=Q-UHqMorzil.
  * Merchant et al. (2023) Merchant, A., Batzner, S., Schoenholz, S. S., Aykol, M., Cheon, G., and Cubuk, E. D.  Scaling deep learning for materials discovery.  _Nature_ , 624(7990):80–85, 2023.
  * Pozdnyakov & Ceriotti (2024) Pozdnyakov, S. and Ceriotti, M.  Smooth, exact rotational symmetrization for deep learning on point clouds.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Puny et al. (2021) Puny, O., Atzmon, M., Smith, E. J., Misra, I., Grover, A., Ben-Hamu, H., and Lipman, Y.  Frame averaging for invariant and equivariant network design.  In _International Conference on Learning Representations_ , 2021.
  * Qi et al. (2017) Qi, C. R., Su, H., Mo, K., and Guibas, L. J.  Pointnet: Deep learning on point sets for 3d classification and segmentation, 2017.
  * Sanchez-Gonzalez et al. (2020) Sanchez-Gonzalez, A., Godwin, J., Pfaff, T., Ying, R., Leskovec, J., and Battaglia, P.  Learning to simulate complex physics with graph networks.  In III, H. D. and Singh, A. (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 8459–8468. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/sanchez-gonzalez20a.html.
  * Schreiner et al. (2022) Schreiner, M., Bhowmik, A., Vegge, T., Busk, J., and Winther, O.  Transition1x – a dataset for building generalizable reactive machine learning potentials, 2022.
  * Schütt et al. (2017) Schütt, K., Kindermans, P.-J., Sauceda Felix, H. E., Chmiela, S., Tkatchenko, A., and Müller, K.-R.  Schnet: A continuous-filter convolutional neural network for modeling quantum interactions.  _Advances in neural information processing systems_ , 30, 2017.
  * Schütt et al. (2021) Schütt, K. T., Unke, O. T., and Gastegger, M.  Equivariant message passing for the prediction of tensorial properties and molecular spectra, 2021.
  * Shoghi et al. (2024) Shoghi, N., Kolluru, A., Kitchin, J. R., Ulissi, Z. W., Zitnick, C. L., and Wood, B. M.  From molecules to materials: Pre-training large generalizable models for atomic property prediction, 2024.
  * Smith et al. (2020) Smith, J. S., Zubatyuk, R., Nebgen, B., Lubbers, N., Barros, K., Roitberg, A. E., Isayev, O., and Tretiak, S.  The ani-1ccx and ani-1x data sets, coupled-cluster and density functional theory properties for molecules.  _Scientific data, 7(1):1–10_ , 2020.
  * Tran et al. (2023) Tran, R., Lan, J., Shuaibi, M., Wood, B. M., Goyal, S., Das, A., Heras-Domingo, J., Kolluru, A., Rizvi, A., Shoghi, N., Sriram, A., Therrien, F., Abed, J., Voznyy, O., Sargent, E. H., Ulissi, Z., and Zitnick, C. L.  The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysts.  _ACS Catalysis_ , 13(5):3066–3084, February 2023\.  ISSN 2155-5435.  doi: 10.1021/acscatal.2c05426.  URL http://dx.doi.org/10.1021/acscatal.2c05426.
  * Wang et al. (2019) Wang, Y., Sun, Y., Liu, Z., Sarma, S. E., Bronstein, M. M., and Solomon, J. M.  Dynamic graph cnn for learning on point clouds, 2019.
  * Zaidi et al. (2022) Zaidi, S., Schaarschmidt, M., Martens, J., Kim, H., Teh, Y. W., Sanchez-Gonzalez, A., Battaglia, P., Pascanu, R., and Godwin, J.  Pre-training via denoising for molecular property prediction, 2022.

##  Appendix A Canonicalization

###  A.1 Stochastic Frame Averaging

##### Frame Averaging

We recall the idea of Frame Averaging introduced by Puny et al. (2021), in
which $\mathcal{X}$ and $\mathcal{Y}$ denote normed linear space with
respective representations $\rho_{1}$ and $\rho_{2}$ of a group $G$. In our
case, the group of interest is $E(3)$. A frame is defined as a function
$\mathcal{F}:\mathcal{X}\to 2^{G}$ taking values in a non-empty subset of the
group $G$ such that it is $G$-equivariant and bounded. Under such conditions,
and for every $\Phi:\mathcal{X}\to\mathcal{Y}$, the function
$\langle\Phi\rangle_{\mathcal{F}}:\mathcal{X}\to\mathcal{Y}$, called average
over the frame $\mathcal{F}$, and defined as

${\langle\Phi\rangle}_{\mathcal{F}}:x\mapsto\frac{1}{|\mathcal{F}(x)|}\sum_{g\in\mathcal{F}(x)}\rho_{2}(g)\Phi(\rho_{1}(g)^{-1}x),$
(1)

is $G$-equivariant. This allows to create an arbitrary neural network, and
guarantee the symmetry by averaging the outputs over a well-chosen frame.

##### Choosing the frame.

As in Duval et al. (2023), the neural network takes as input $X$ the atoms
position, $Z$ is their atomic numbers, and has outputs in $\mathcal{Y}$, where
$\mathcal{Y}=\mathbb{R}^{n\times 3}$ in the case of force predictions, and
$\mathcal{Y}=\mathbb{R}$ in the case of energy prediction. The group $E(3)$
only acts on $X$, and not $Z$. A Principal Component Analysis (PCA) on the
atomic structure allows to decompose the covariance matrix of the points cloud
$\Sigma=U^{T}\Lambda U$ derived from the centroid of the positions
$t=\frac{1}{n}X^{T}\bm{1}$, with $\Lambda$ the diagonal matrix containing the
three eigenvalues $\lambda_{1}>\lambda_{2}>\lambda_{3}$ (assumed distinct
because we consider non-planar structures). The frame is taken as

$\mathcal{F}(X)=\\{(U,t)\mid U=[\pm u_{1},\pm u_{2},\pm u_{3}]\\},$ (2)

which is a subset of $E(3)$. The authors prove that the frame $\mathcal{F}$
defined as such is $G$-equivariant and bounded.

##### Stochastic Frame Averaging.

This process requires to average the predictions of the neural network over
$|\mathcal{F}(X)|=2^{3}=8$ elements of the frame. In order to make the
computations faster, Duval et al. (2023) sample only one element from the
frame instead of performing the average. Although Stochastic Frame-Averaging
(SFA) does not have theoretical guarantees, it has been experimentally shown
to learn almost perfect equivariance.

###  A.2 SFA+SignNet

Initially proposed to help spectral graph representation learning, SignNet
(Lim et al., 2023b) is a network which outputs are invariant to sign flips.
The authors state that a continuous function
$\eta:\mathbb{R}^{n}\to\mathbb{R}^{d}$ is sign-invariant if and only if
$\eta(x)=\kappa(x)+\kappa(-x)$ for some (freely chosen) continuous function
$\kappa:\mathbb{R}^{n}\to\mathbb{R}^{d}$. $\text{SignNet}:\mathbb{R}^{n\times
k}\to\mathbb{R}^{n\times k}$ is then defined as:

$\text{SignNet}(x_{1},\dots,x_{k})=\mu\left([\kappa(x_{i})+\kappa(-x_{i})]_{i=1}^{k}\right),$
(3)

where $\mu$ and $\kappa$ are neural networks chosen freely.

We propose to apply such a network on the sampled element of the frame $U$ and
parametrize $\mu$ and $\kappa$ either with MLPs or with VN-PointNets. In order
to constrain the output, we further orthonormalize the output with a Gram-
Schmidt process:

$U^{\prime}=\text{Gram-Schmidt}(\text{SignNet}(U)).$ (4)

When parametrizing SignNet with MLPs, since we apply non-linearities on the
orthogonal matrix $U$, there is no theoretical guarantee for the whole process
to be $E(3)$-equivariant a priori. This is coherent with empirical
observations when using this method, although training $\mu$ and $\kappa$
helps better enforce the equivariance.

When parametrizing SignNet with VN-PointNets, the whole network is made
exactly $E(3)$-equivariant. As explained by Lim et al. (2023a) (section 2.2),
the matrix $U$ is orthogonal and unique up to sign changes, while the SignNet
function is both sign invariant by design and $O(3)$-equivariant thanks to the
use of VN-PointNets.

###  A.3 Vector Neurons Network

VNNs are a class of $SO(3)$-equivariant models, where usual neurons are
replaced with so-called Vector Neurons: for a given layer, non-linearities
output a matrix of size $h\times 3$ instead of a vector of length $h$.

Using this framework, Deng et al. (2021) re-implement classic operations such
as linear layers, non-linearities, pooling operations, and normalization
layers. They also prove that those layers are all $SO(3)$-equivariants, which
allows to re-implement classical networks into their VN variant. In
particular, they re-implement the VN variants of two well-known networks from
the point clouds literature: PointNet (Qi et al., 2017) and DGCNN (Wang et
al., 2019), and test them on classification, segmentation, and reconstruction
tasks. They show that accuracy increases compared to the classic
implementations and that their equivariance property is indeed (almost)
perfectly enforced.

To obtain an $O(3)$ (and then $E(3)$) equivariance, the output of the VNN has
to be further orthonormalized with a Gram-Schmidt process to canonicalize the
representation in $O(3)$, as described by Kaba et al. (2023).

The following Vector Neurons Networks (VNNs) are used:

  * •

VN-Pointnet with a varying number of VNLinearLeakyReLU layers (between 1 and
3), with the implementation of Deng et al. (2021).

  * •

VN-DGCNN, with the implementation is the one of Deng et al. (2021).

To summarize, we use VNNs to learn the transformation $U$ from the positions
$X$:

$\displaystyle\text{VNN}:\mathbb{R}^{n\times 3}$
$\displaystyle\to\mathbb{R}^{3\times 3}$ $\displaystyle X$
$\displaystyle\mapsto U.$

$U$ is then orthonormalized with a Gram-Schmidt process.

###  A.4 Experimental comparisons

OC20 IS2RE: Tables 5 and 6

OC20 S2EF: Tables 7 and 8

QM9: Table 9

In all OC20 experiments of this section, for a fair assessment, SFA is used in
3D mode, i.e. without the computational trick to force the $z$ axis to remain
fixed during canonicalization, which is specific to OC20. The same goes for
the methods derived from SFA. As a consequence, the reported performances are
lower than reported in other sections.

Canonicalization | Cano. trainable | 2D Rotation $\downarrow$ | 3D Rotation $\downarrow$ | Reflection $\downarrow$
---|---|---|---|---
| parameters | Invariance | Invariance | Invariance
SFA | 0 | $1.29\cdot 10^{-2}$ | $1.32\cdot 10^{-2}$ | $1.30\cdot 10^{-2}$
Untrained SFA+MLP-SignNet | 0 | $1.01\cdot 10^{-1}$ | $1.00\cdot 10^{-1}$ | $9.71\cdot 10^{-2}$
Trained SFA+MLP-SignNet | 454 | $4.21\cdot 10^{-2}$ | $7.74\cdot 10^{-2}$ | $4.00\cdot 10^{-2}$
Untrained SFA+VN-SignNet | 0 | $6.89\cdot 10^{-3}$ | $7.58\cdot 10^{-3}$ | $7.37\cdot 10^{-3}$
Trained SFA+VN-SignNet | 2,620 | $2.45\cdot 10^{-2}$ | $2.66\cdot 10^{-2}$ | $2.43\cdot 10^{-2}$
Untrained VN-Pointnet (2 hid.) | 0 | $4.61\cdot 10^{-3}$ | $4.62\cdot 10^{-3}$ | $4.62\cdot 10^{-3}$
Trained VN-Pointnet (2 hid.) | 1,310 | $\mathbf{3.63\cdot 10^{-3}}$ | $\mathbf{3.72\cdot 10^{-3}}$ | $\mathbf{3.80\cdot 10^{-3}}$
Untrained VN-Pointnet (1 hid.) | 0 | $4.28\cdot 10^{-3}$ | $4.28\cdot 10^{-3}$ | $4.29\cdot 10^{-3}$
Untrained VN-Pointnet (0 hid.) | 0 | $2.76\cdot 10^{-2}$ | $2.76\cdot 10^{-2}$ | $2.79\cdot 10^{-2}$
Trained VN-Pointnet (0 hid.) | 24 | $1.86\cdot 10^{-2}$ | $2.31\cdot 10^{-2}$ | $2.36\cdot 10^{-2}$
Untrained VN-DGCNN | 0 | $3.03\cdot 10^{-2}$ | $3.08\cdot 10^{-2}$ | $3.11\cdot 10^{-2}$
Trained VN-DGCNN | 663,804 | $9.89\cdot 10^{-3}$ | $2.49\cdot 10^{-2}$ | $9.10\cdot 10^{-3}$

Table 5: Invariance comparison of canonicalization methods on OC20 IS2RE
dataset. The FAENet backbone for this task and dataset has 4,147,731
parameters (5 interaction blocks). We measure the rotation invariance and
reflection invariance property as the difference in prediction between every
samples D1 (of the ID val split) and D2 defined as a SO(3) transformation of
D1, in eV.

| ID | OOD-CAT | OOD-ADS | OOD-BOTH
---|---|---|---|---
Canonicalization | EwT (%) $\uparrow$ | MAE (meV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (meV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (meV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (meV) $\downarrow$
SFA | 4.40 | 566 | 4.12 | 563 | 2.56 | 652 | 2.77 | 594
Untrained SFA+MLP-SignNet | 4.48 | 554 | 4.46 | 552 | 2.75 | 637 | 2.88 | 576
Trained SFA+MLP-SignNet | 4.46 | 554 | 4.51 | 551 | 2.67 | 642 | 2.78 | 586
Untrained SFA+VN-SignNet | 4.69 | 563 | 4.60 | 558 | 2.62 | 651 | 2.59 | 595
Trained SFA+VN-SignNet | 4.25 | 572 | 4.27 | 568 | 2.92 | 658 | 2.97 | 596
Untrained VN-Pointnet (2 hid.) | 4.09 | 567 | 4.66 | 565 | 2.60 | 673 | 2.85 | 615
Trained VN-Pointnet (2 hid.) | 4.12 | 568 | 4.33 | 563 | 2.77 | 658 | 2.75 | 604
Untrained VN-Pointnet (1 hid.) | 4.37 | 565 | 4.20 | 561 | 2.64 | 666 | 2.73 | 614
Untrained VN-Pointnet (0 hid.) | 4.01 | 581 | 3.92 | 571 | 2.75 | 660 | 2.64 | 615
Trained VN-Pointnet (0 hid.) | 4.14 | 567 | 4.36 | 563 | 2.56 | 675 | 2.88 | 614
Untrained VN-DGCNN | 4.31 | 567 | 4.14 | 562 | 2.58 | 660 | 2.72 | 610
Trained VN-DGCNN | 4.42 | 560 | 4.40 | 556 | 2.78 | 656 | 2.81 | 601

Table 6: Performance comparison of canonicalization methods on OC20 IS2RE
dataset. All models were trained for 12 epochs using Duval et al. (2023)
config.

| Cano. trainable | Energy invariance | Forces equivariance
---|---|---|---
Canonicalization | parameters | 3D Rotation $\downarrow$ | Reflection $\downarrow$ | 3D Rotation $\downarrow$ | Reflection $\downarrow$
SFA | 0 | $1.88\cdot 10^{-2}$ | $1.88\cdot 10^{-2}$ | $7.17\cdot 10^{-2}$ | $8.34\cdot 10^{-3}$
Untrained SFA+MLP-SignNet | 0 | $7.81\cdot 10^{-2}$ | $7.61\cdot 10^{-2}$ | $7.44\cdot 10^{-2}$ | $2.04\cdot 10^{-2}$
Trained SFA+MLP-SignNet | 454 | $6.57\cdot 10^{-2}$ | $3.57\cdot 10^{-2}$ | $7.35\cdot 10^{-2}$ | $1.17\cdot 10^{-2}$
Untrained SFA+VN-SignNet | 0 | $2.07\cdot 10^{-2}$ | $2.04\cdot 10^{-2}$ | $6.86\cdot 10^{-2}$ | $9.42\cdot 10^{-3}$
Trained SFA+VN-SignNet | 2,620 | $1.92\cdot 10^{-2}$ | $1.89\cdot 10^{-2}$ | $6.55\cdot 10^{-2}$ | $8.72\cdot 10^{-3}$
Untrained VN-Pointnet (2 hid.) | 0 | $1.80\cdot 10^{-2}$ | $1.80\cdot 10^{-2}$ | $6.92\cdot 10^{-2}$ | $8.78\cdot 10^{-3}$
Trained VN-Pointnet (2 hid.) | 1,310 | $1.67\cdot 10^{-2}$ | $1.67\cdot 10^{-2}$ | $6.89\cdot 10^{-2}$ | $8.77\cdot 10^{-3}$
Untrained VN-Pointnet (0 hid.) | 0 | $3.50\cdot 10^{-2}$ | $3.49\cdot 10^{-2}$ | $6.96\cdot 10^{-2}$ | $1.08\cdot 10^{-2}$
Trained VN-Pointnet (0 hid.) | 24 | $3.31\cdot 10^{-2}$ | $3.34\cdot 10^{-2}$ | $7.00\cdot 10^{-2}$ | $1.05\cdot 10^{-2}$
Untrained VN-DGCNN | 0 | $\mathbf{1.50\cdot 10^{-2}}$ | $\mathbf{1.50\cdot 10^{-2}}$ | $\mathbf{6.83\cdot 10^{-2}}$ | $\mathbf{3.58\cdot 10^{-3}}$
Trained VN-DGCNN | 663,804 | $2.02\cdot 10^{-2}$ | $1.50\cdot 10^{-2}$ | $6.91\cdot 10^{-2}$ | $8.09\cdot 10^{-3}$

Table 7: Equivariance comparison of canonicalization methods on OC20 S2EF
dataset. The FAENet backbone for this task and dataset has 5,675,410
parameters (7 interaction blocks). We measure the energy rotation invariance,
energy reflection invariance, force rotation equivariance, and force
reflection equivariance properties as the difference in prediction between
every sample D1 (of the ID val split) and D2 defined as a SO(3) transformation
of D1, in eV.

| Energy MAE (mEV) $\downarrow$ | Force MAE (mEV) $\downarrow$
---|---|---
Canonicalization | ID | OOD Cat | OOD Ads | OOD Both | ID | OOD Cat | OOD Ads | OOD Both
SFA | 424 | 445 | 579 | 680 | 55.6 | 55.2 | 63.2 | 74.6
Untrained SFA+MLP-SignNet | 420 | 444 | 515 | 631 | 54.0 | 53.8 | 61.4 | 72.4
Trained SFA+MLP-SignNet | 422 | 446 | 558 | 666 | 54.2 | 53.9 | 62.7 | 73.8
Untrained SFA+VN-SignNet | 439 | 458 | 565 | 673 | 56.5 | 56.0 | 65.3 | 76.7
Trained SFA+VN-SignNet | 442 | 464 | 590 | 701 | 58.0 | 57.5 | 65.2 | 76.9
Untrained VN-Pointnet (2 hid.) | 435 | 455 | 596 | 697 | 56.0 | 55.6 | 66.9 | 77.5
Trained VN-Pointnet (2 hid.) | 435 | 453 | 585 | 696 | 56.1 | 55.8 | 64.2 | 75.6
Untrained VN-Pointnet (0 hid.) | 440 | 459 | 597 | 705 | 56.0 | 55.6 | 64.7 | 75.8
Trained VN-Pointnet (0 hid.) | 440 | 459 | 572 | 671 | 55.8 | 55.4 | 63.7 | 74.9
Untrained VN-DGCNN | 456 | 474 | 593 | 763 | 55.7 | 55.5 | 65.8 | 76.9
Trained VN-DGCNN | 432 | 453 | 662 | 762 | 55.5 | 55.2 | 71.2 | 80.7

Table 8: Performance comparison of canonicalization methods on OC20 S2EF
dataset. All models were trained for 12 epochs using Duval et al. (2023)
config.

| Cano. trainable | MAE (meV) $\downarrow$ | Energy invariance (eV)
---|---|---|---
Canonicalization | parameters | ID | Test | 3D Rotation $\downarrow$ | Reflection $\downarrow$
SFA | 0 | 9.20 | 9.06 | $1.65\cdot 10^{-3}$ | $1.76\cdot 10^{-3}$
Untrained SFA+MLP-SignNet | 0 | 11.3 | 11.2 | $2.12\cdot 10^{-3}$ | $2.20\cdot 10^{-3}$
Trained SFA+MLP-SignNet | 454 | 10.5 | 10.7 | $1.60\cdot 10^{-3}$ | $1.66\cdot 10^{-3}$
Untrained SFA+VN-SignNet | 0 | 9.41 | 9.40 | $1.28\cdot 10^{-3}$ | $1.37\cdot 10^{-3}$
Trained SFA+VN-SignNet | 2,620 | 10.1 | 10.2 | $1.33\cdot 10^{-3}$ | $1.40\cdot 10^{-3}$
Untrained VN-Pointnet (2 hid.) | 0 | 10.4 | 10.3 | $\mathbf{1.19}\cdot 10^{-3}$ | $1.30\cdot 10^{-3}$
Trained VN-Pointnet (2 hid.) | 1,310 | 10.1 | 9.85 | $1.30\cdot 10^{-3}$ | $1.44\cdot 10^{-3}$
Untrained VN-Pointnet (0 hid.) | 0 | 9.51 | 9.49 | $1.21\cdot 10^{-3}$ | $\mathbf{1.24}\cdot 10^{-3}$
Trained VN-Pointnet (0 hid.) | 24 | 11.4 | 11.5 | $1.64\cdot 10^{-3}$ | $1.74\cdot 10^{-3}$
Untrained VN-DGCNN | 0 | 9.92 | 9.94 | $1.32\cdot 10^{-3}$ | $1.46\cdot 10^{-3}$
Trained VN-DGCNN | 663,804 | 9.34 | 9.25 | $1.79\cdot 10^{-3}$ | $1.79\cdot 10^{-3}$

Table 9: Equivariance and performance comparison of canonicalization methods
on QM9 dataset for the target property $U_{0}$ (internal energy at 0 Kelvin).
The FAENet backbone for this task has 6,495,127 parameters (5 interaction
blocks). We measure the energy rotation invariance as the difference in
prediction between every samples D1 (of the ID val split) and D2 defined as a
SO(3) transformation of D1, in eV. All models were trained for 300 epochs
using Duval et al. (2023) config.

###  A.5 Relaxations from S2EF model for IS2RE

The models used to run the experiments with the relaxation methods were
trained on the 2M train split of the S2EF dataset from OC20. This dataset has
been shown to converge to similar performances as the complete dataset, which
is way larger and takes too long to train on (Gasteiger et al., 2022). We
report in Table 10 the performances of these models, which might help
interpret some of the results for the relaxation.

Model | EwT $\downarrow$ | Force MAE $\downarrow$ | Forces $\cos\quad$$\uparrow$
---|---|---|---
FAENet with SFA | 10.7 | 0.044 | 0.32
FAENet with Untrained PointNet | 10.5 | 0.0043 | 0.33
FAENet without SFA | 10.0 | 0.042 | 0.34
SchNet Base | 5.1 | 0.061 | 0.07

Table 10: Performance comparison of models on OC20 S2EF dataset on the VAL-ID
split. The energy within threshold, Force MAE, and $\cos$ similarity are
reported for these S2EF models that are then used for relaxations. Note that
this method yields way longer training and inference times when compared to
direct IS2RE as reported.

##  Appendix B Graph creation study

###  B.1 Cutoff

The cutoff defines the distance within which a link between two atoms is
created. All atoms that are at a distance smaller than this cutoff will be
linked. However, in order to avoid cluttering the created graphs, most methods
impose a maximum number of neighbors for every atom. This parameter is usually
taken around 40 neighbors.

| ID | OOD-ADS | OOD-CAT | OOD-BOTH
---|---|---|---|---
Model | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$
Cutoff 30 - Max neighbours 40 | 2.65 | 0.697 | 1.45 | 0.906 | 2.86 | 0.691 | 1.53 | 0.846
Cutoff 20 - Max. neighbours 40 | 3.08 | 0.673 | 1.85 | 0.808 | 2.86 | 0.669 | 1.86 | 0.757
Cutoff 20 - Max. neighbours 10 | 2.25 | 0.768 | 1.51 | 0.988 | 2.52 | 0.754 | 1.38 | 0.928
Cutoff 10 - Max. neighbours 50 | 4.17 | 0.553 | 2.81 | 0.640 | 4.12 | 0.551 | 3.02 | 0.585
Cutoff 10 - Max. neighbours 40 | 4.29 | 0.555 | 2.95 | 0.631 | 4.33 | 0.553 | 2.71 | 0.587
Cutoff 10 - Max. neighbours 30 | 4.43 | 0.551 | 2.65 | 0.655 | 4.51 | 0.552 | 2.51 | 0.611
Cutoff 10 - Max. neighbours 20 | 4.38 | 0.551 | 2.46 | 0.676 | 4.45 | 0.551 | 2.55 | 0.621
Cutoff 10 - Max. neighbours 10 | 4.49 | 0.553 | 2.84 | 0.627 | 4.34 | 0.549 | 3.01 | 0.582
Cutoff 6 - Max. neighbours 40 | 4.31 | 0.553 | 3.00 | 0.626 | 4.39 | 0.554 | 2.81 | 0.577
Cutoff 1 - Max. neighbours 40 | 1.35 | 1.069 | 1.32 | 1.112 | 1.33 | 1.051 | 1.37 | 1.018

Table 11: Impact of the cutoff on the performances of FAENet on the OC20 IS2RE
task. Full table on all validation splits.

###  B.2 Ewald-based Long-Range Message Passing

The main idea behind Ewald summation used in section 3.2 is to decompose the
electrostatic interaction potential with a given charge into a short-range
interaction and a long-range interaction term. The short-range contribution
can be computed with real spatial features and the long-range contribution is
computed using a Fourier transform. This principle is illustrated in Figure 3.
This allows for computational methods in electrostatics to converge faster and
with higher accuracy because the long-range interaction becomes more
tractable.

Figure 3: Ewald Summation interactions. The interaction term (left) is the
result of the short-range interaction (middle) and the long-term interaction
(right) which are both computed using cutoffs on respectively the real and the
Fourier space. Adapted from (Kosmala et al., 2023)

In the case of GNNs, the short-range interaction is already computed by the
currently implemented interaction blocks using a distance cutoff between the
atoms, which omits the negligible parts of this interaction for further atoms.
However, the heavy-tail of long-range interactions is then reported to a new
term in real space, which doesn’t diverge anymore for closer atoms. It can
then be computed using the same cutoff idea but in the Fourier space where a
spatial frequency cutoff is used to make it tractable as introduced in Kosmala
et al. (2023).

##### Periodic case.

In the case where there exists a spatial periodic tiling of materials (OC20
for example), it is possible to define the set of periodic cells localization
$\Lambda=\\{\lambda_{1}\mathbf{v_{1}},\lambda_{2}\mathbf{v_{2}},\lambda_{3}\mathbf{v_{3}}\mid(\lambda_{1},\lambda_{2},\lambda_{3})\in\mathbb{Z}^{3}\\}$,
where $\mathbf{v_{1}},\mathbf{v_{2}},\mathbf{v_{3}}$ define the periodic cell
lattice, similarly to the periodic interval in the 1D case. In the real space,
the long-range interaction component would be written as a sum over all of the
elements on the infinite tiling, which can be decomposed as a Fourier series
expansion using the reciprocal lattice $\Lambda^{\prime}$. This reciprocal
lattice would be similar to the $2\pi$ in the 1D case of the Fourier
transform. It is the periodic space of all the wavevectors of the Fourier
series. This results in the proposed expansion for Ewald message passing:

$M^{lr}(x_{i})=\sum_{\mathbf{k}\in\Lambda^{\prime}}\exp{(i\mathbf{k}^{T}x_{i})}\cdot\sum_{j\in\mathcal{S}}h_{j}\exp{(-i\mathbf{k}^{T}x_{j})}\cdot\hat{\Phi}^{lr}(||\mathbf{k}||),$
(5)

where $M^{lr}(x_{i})$ corresponds to the long-range message computed at node
$i$ from all of the nodes in the system $S$, and $\hat{\Phi}^{lr}$ is a
learned Fourier coefficient of a radial basis function representing the
interaction. The cutoff in the Fourier basis $c_{k}$ is then set to make the
sum finite over the set $\\{k\in\Lambda^{\prime},||k||\leq c_{k}\\}$. Since
the number of wavevectors used for the computation is finite,
$\hat{\Phi}(\|\mathbf{k}\|)$ is learned for every $\mathbf{k}$. Since Ewald
summation applies to periodic structures, the authors of Ewald message passing
Kosmala et al. (2023) propose tricks to deal with the aperiodic case by
assuming an infinite tiling.

(a) SchNet with Ewald - Interaction Blocks

(b) FAENet with Ewald - Interaction Blocks

(c) SchNet with Ewald - Ewald Blocks

(d) FAENet with Ewald - Ewald Blocks

Figure 4: Similarity matrix of the embeddings of the atoms of a system for
different interaction blocks on FAENet and SchNet with Ewald message passing
on the models. The visualized layers here are the standard interaction blocks
in the first row and the Ewald interaction blocks in the second row. The two
outputs are summed to get the final representation for Ewald shown in Figure
1.

(a) SchNet without Ewald

(b) FAENet without Ewald

(c) SchNet with Ewald

(d) FAENet with Ewald

Figure 5: Same plots as Figure 1 but with a second randomly picked system from
the OC20 train split.

(a) SchNet with Ewald - Interaction Blocks

(b) FAENet with Ewald - Interaction Blocks

(c) SchNet with Ewald - Ewald Blocks

(d) FAENet with Ewald - Ewald Blocks

Figure 6: Same plots as Figure 4 but with a second randomly picked system from
the OC20 train split.

| ID | OOD-ADS | OOD-CAT | OOD-BOTH
---|---|---|---|---
Model | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$ | EwT (%) $\uparrow$ | MAE (eV) $\downarrow$
FAENet (Graph Rewiring) | 4.05 | 0.551 | 2.65 | 0.650 | 4.29 | 0.550 | 2.76 | 0.601
FAENet (Graph Rewiring) + Ewald | 4.12 | 0.562 | 2.68 | 0.648 | 4.14 | 0.563 | 2.83 | 0.597
FAENet (No Graph Rewiring) | 4.54 | 0.544 | 2.59 | 0.657 | 4.66 | 0.539 | 2.65 | 0.601
FAENet (No Graph Rewiring) + Ewald | 4.11 | 0.556 | 2.75 | 0.626 | 4.13 | 0.553 | 2.85 | 0.569
SchNet (Graph Rewiring) | 3.18 | 0.641 | 2.53 | 0.720 | 3.00 | 0.638 | 2.59 | 0.642
SchNet (Graph Rewiring) + Ewald | 3.54 | 0.604 | 2.53 | 0.665 | 3.53 | 0.599 | 2.67 | 0.608
SchNet (No Graph Rewiring) | 2.93 | 0.654 | 2.22 | 0.700 | 3.04 | 0.646 | 2.54 | 0.656
SchNet (No Graph Rewiring) + Ewald | 3.48 | 0.597 | 2.76 | 0.647 | 3.56 | 0.599 | 2.73 | 0.612

Table 12: Energy prediction errors with and without Ewald Message Passing.
Graph Rewiring refers to removing the subsurface atoms from the system (Duval
et al., 2022). In this table, FAENet is taken with 5 interaction layers (top
config), while SchNet uses 3 interaction layers.

Model | MAE (meV) $\downarrow$ | MSE ((meV)2) $\downarrow$
---|---|---
FAENet | 8.44 | 0.600
FAENet + Ewald | 8.34 | 0.574
SchNet | 16.0 | 1.18
SchNet + Ewald | 11.5 | 0.73

Table 13: Energy prediction errors with and without Ewald Message Passing on
the test split of the QM9 dataset for the target property $U_{0}$ (Internal
energy at 0K).

##  Appendix C Noisy Nodes

###  C.1 Experimental setup

Runs are done on a single Nvidia Quadro RTX 8000 GPU with 48 GB memory or, if
indicated, on a single Nvidia A100 GPU with 80 GB memory.

###  C.2 Noisy Nodes implementation

In practice, we perturb the input node positions of the graph $G$ with a noise
$\sigma$ and train the model with two loss terms, a Noisy Nodes loss term and
the primary loss (associated with the main task) term

$\mathcal{L}=\lambda\cdot\mathcal{L}_{NNodes}(\hat{G}^{\prime},V^{\prime})+\mathcal{L}_{Primary}(\hat{G}^{\prime},V^{\prime}),$
(6)

where $\lambda$ is the weight we assign to the auxiliary denoising task,
$\hat{G}^{\prime}=\tilde{\Phi}(\tilde{G})$ is the output of the model
$\tilde{\Phi}$, $\tilde{G}$ is the noised graph, and $V^{\prime}$ can either
be the target nodes features (e.g. atom positions at equilibrium, that is the
IS2RS task) or the initial nodes positions (see next subsection on denoising
pre-training).

For the input, we first interpolate between initial structure and relaxed
structure and then add Gaussian noise, that is for each node $i$, the input
positions of the input ”Noisy Nodes graph” $\tilde{x}^{i}$ are defined as

$\tilde{x}^{i}=\begin{cases}\gamma(x^{i}_{rel}-x^{i}_{init})+Z^{i}\textrm{
with probability 0.5}\\\ x^{i}_{init}\textrm{ with probability
0.5},\end{cases}$ (7)

with random interpolation factor $\gamma\sim U[0,1]$ independent between
graphs and iid Gaussian noises $Z^{i}\sim\mathcal{N}(0,\sigma)$ with
$\sigma=0.3$. The Noisy Nodes target is
$\Delta_{pos}^{i}=x_{rel}^{i}-\tilde{x}^{i}$ and therefore the auxiliary loss
term is $\|\Delta_{pos}^{i}-\Phi(\tilde{x}^{i})\|_{1}$ for the model $\Phi$.
Our total loss is the sum of the energy MAE loss and the auxiliary loss
weighted by a number $\lambda$, as in Equation 6. Both the primary loss and
Noisy Nodes Loss (before multiplication by $\lambda$) typically have the same
value between 1 and 2.5.

For the IS2RE training, we add a position decoding head to the preexisting
energy prediction head. The position decoding head, as with the force
prediction head introduced in (Duval et al., 2023), is a 2-layer MLP with
Batchnorm.

###  C.3 Related work

Since Noisy Nodes performs denoising as an auxiliary task during training, the
representation learning benefits of denoising are limited to the downstream
dataset on which the model is trained. Zaidi et al. (2022) propose to rather
perform denoising as a pre-training objective on a large, unlabelled dataset
of atomic structures.

Shoghi et al. (2024) introduce Joint Multi-domain Pre-training (JMP), a
supervised pre-training strategy that simultaneously trains on various
datasets from different chemical domains (OC20 (Chanussot et al., 2020), OC22
(Tran et al., 2023), ANI-1x (Smith et al., 2020), and Transition-1x (Schreiner
et al., 2022)), treating each dataset as a unique pre-training task within a
multi-task framework.

###  C.4 Training hyperparameters

The most obvious changes to the training hyperparameters that should
theoretically allow to leverage the denoising auxiliary task are to increase
the depth of the network and the number of epochs. First, we observe that we
reach convergence on the validation set for the energy MAE more slowly when
using the auxiliary task because of a more complex loss to minimize and a
higher number of model weights (due to the supplementary interaction blocks).
Hence, unless otherwise stated, the number of epochs is 50 in the IS2RE with
IS2RS auxiliary task experiments.

(Liao & Smidt, 2023) use a linearly decaying weight from 15 to 1 for the
auxiliary IS2RS loss to encourage the model to learn more from the auxiliary
task in the beginning but focus on the primary task at the end of training. We
also tested a constant auxiliary weight of 1, a weight decaying from 30, and
cosine annealing with a linear warmup scheduler, but this yielded worse or
equivalent results. Thus, unless otherwise stated, we always use as an
auxiliary weight scheduler the one of Equiformer in the following IS2RE with
IS2RS auxiliary task experiments.

Since we always use the MAE as energy loss for the IS2RE main task, it proves
essential to use the MAE loss for the auxiliary position loss to leverage the
benefits of Noisy Nodes. Indeed, our experiments using the MSE position loss
led to a collapse of the node embeddings at the last interaction layers, that
could be observed by plotting the MAD throughout the interaction layers.

###  C.5 Results

####  C.5.1 Comparing Depth for classical FAENet

In the results of Table 14, the number of warmup steps is 6000 for about 180k
steps. The batch size is 128, hence the lower throughput than the one with the
same configs of (Duval et al., 2023). In these runs, the hyperparameters are
the top configs of (Duval et al., 2023) except for slight differences: the
number of hidden channels in the embedding blocks is a bit lower. Also,
contrary to the top configs of (Duval et al., 2023), we train for 50 epochs
(instead of 12) with no early stopping to be in a comparable setup to the
IS2RE with IS2RS auxiliary task models.

Interaction | Parameters | Time | Energy MAE (meV) $\downarrow$ | EwT (%) $\uparrow$
---|---|---|---|---
blocks | (millions) | Train $\downarrow$ | Infer. $\uparrow$ | ID | OOD Ads | OOD Cat | OOD Both | Average | ID | OOD Ads | OOD Cat | OOD Both | Average
5 | 5.9 | 19min | 786 | 556 | 685 | 552 | 636 | 607 | 4.41 | 2.26 | 4.51 | 2.47 | 3.41
8 | 9.2 | 25min | 676 | 554 | 643 | 558 | 596 | 588 | 4.31 | 2.66 | 4.40 | 2.75 | 3.53
10 | 11.4 | 29min | 621 | 552 | 649 | 554 | 603 | 590 | 4.56 | 2.77 | 4.20 | 2.72 | 3.56
12 | 13.7 | 34min | 527 | 551 | 661 | 556 | 609 | 594 | 4.01 | 2.57 | 4.21 | 2.53 | 3.33
14 | 15.9 | 37min | 498 | 621 | 749 | 608 | 693 | 668 | 3.42 | 2.06 | 3.27 | 2.31 | 2.77
16 | 18.2 | 40min | 457 | 590 | 704 | 592 | 647 | 633 | 4.05 | 2.51 | 3.90 | 2.53 | 3.25

Table 14: Comparison varying the number of interaction blocks for FAENet
without auxiliary task with the top configs of (Duval et al., 2023) except for
slight differences. Scalability is measured with training time for one epoch
(Train, in minutes) and inference throughput (Infer., number of samples
processed in a second). The best score is in bold and the second-best score is
underlined. Performances are slightly worse than the best performances
presented in (Duval et al., 2023) because of different hyperparameters, but
the main conclusion is that the performances worsen when we add interaction
layers after 12.

####  C.5.2 Comparing depth for FAENet with Noisy Nodes IS2RS auxiliary task

In Table 15, we observe a clear positive correlation between the depth of the
model and its performances in terms of energy MAE and EwT. The gain in
performance is very clear between 5 and 16 interaction layers, then increases
much more slowly up to 28 layers.

Interaction | Parameters | Time | Energy MAE (meV) $\downarrow$ | EwT (%) $\uparrow$
---|---|---|---|---
blocks | (millions) | Train $\downarrow$ | Infer.$\uparrow$ | ID | OOD Ads | OOD Cat | OOD Both | Average | ID | OOD Ads | OOD Cat | OOD Both | Average
5 | 4.2 | 17min | 724 | 523 | 592 | 525 | 545 | 546 | 4.59 | 2.92 | 4.40 | 3.15 | 3.76
8 | 6.6 | 22min | 679 | 518 | 600 | 526 | 561 | 551 | 5.13 | 2.88 | 5.21 | 2.74 | 3.99
10 | 8.1 | 27min | 584 | 513 | 606 | 521 | 561 | 550 | 5.17 | 2.82 | 5.03 | 2.81 | 3.96
12 (A100) | 9.7 | 30min | 504 | 517 | 589 | 523 | 546 | 544 | 5.10 | 3.17 | 4.92 | 3.04 | 4.06
14 | 11.2 | 33min | 525 | 507 | 580 | 519 | 541 | 537 | 5.20 | 3.13 | 5.21 | 3.25 | 4.20
16 | 12.8 | 35min | 465 | 505 | 566 | 516 | 527 | 529 | 5.17 | 3.63 | 5.00 | 3.69 | 4.37
18 | 14.3 | 38min | 446 | 508 | 596 | 518 | 554 | 544 | 5.31 | 3.22 | 5.58 | 2.97 | 4.27
20 | 15.9 | 42min | 423 | 508 | 593 | 513 | 553 | 542 | 5.39 | 3.16 | 5.18 | 2.88 | 4.15
22 | 17.4 | 45min | 401 | 507 | 569 | 515 | 528 | 530 | 5.24 | 3.26 | 5.06 | 3.32 | 4.22
24 | 19.0 | 48min | 375 | 500 | 574 | 510 | 534 | 529 | 5.70 | 3.43 | 5.43 | 3.31 | 4.47
26 | 20.6 | 50min | 335 | 504 | 562 | 512 | 521 | 525 | 5.15 | 3.68 | 5.20 | 3.67 | 4.43
28 | 22.1 | 56min | 349 | 505 | 567 | 513 | 517 | 525 | 5.20 | 3.64 | 5.27 | 3.28 | 4.35

Table 15: Comparison of number of interaction blocks for IS2RE with auxiliary
IS2RS. Scalability is measured witht raining time for one epoch (Train, in
minutes) and inference throughput (Infer., number of samples processed in a
second). Here the canonicalization technique is SE(3)-SFA, the number of
warmup steps is 192000 (about half total number of steps), and the batch size
is 64. Best score is in bold, second best score is underlined.

####  C.5.3 Canonicalization comparison

We use SE(3)-SFA because it is less equivariant to reflections than Stochastic
FA (which samples a frame in E(3)), but it has to learn data symmetries from
fewer frames, which helps training. For both SE(3)-SFA and No-FA models, we
observe in Table 16 that the equivariance-invariance do not seem to be
correlated to the number of interaction layers. Moreover, we do not observe a
clear correlation between the Average energy MAE and the equivariance-
invariance.

When using SE(3)-SFA, we see a clear correlation between the increase in the
number of interaction layers and the performances. In the No-FA case, the
impossibility to learn equivariance and invariance even when adding more
layers to the model might account for the less clear positive correlation
between the performances (in terms of energy MAE) and the number of
interaction layers than when we were applying SE(3)-SFA.

Canonicalization | Interactions | 2D E-RI $\downarrow$ | 2D E-Refl-I $\downarrow$ | 2D Pos-RI $\downarrow$ | 2D Pos-Refl-I $\downarrow$ | Average E-MAE (meV) $\downarrow$
---|---|---|---|---|---|---
SE(3)-SFA | 5 | 22.7 | 34.7 | 55.0 | 80.3 | 546
SE(3)-SFA | 14 | 29.7 | 43.2 | 61.7 | 90.4 | 537
SE(3)-SFA | 16 | 21.4 | 33.6 | 56.0 | 83.4 | 529
SE(3)-SFA | 18 | 19.0 | 30.2 | 53.9 | 78.4 | 544
SE(3)-SFA | 20 | 22.6 | 34.7 | 56.2 | 82.8 | 542
SE(3)-SFA | 22 | 20.8 | 32.6 | 56.3 | 81.7 | 530
SE(3)-SFA | 28 | 25.5 | 37.3 | 60.8 | 86.5 | 525
SE(3)-SFA no aux | 5 | 6.9 | 8.9 |  |  | 569
No-FA | 8 | 111 | 107 | 253 | 238 | 561
No-FA | 10 | 110 | 107 | 251 | 233 | 577
No-FA | 14 | 121 | 117 | 275 | 254 | 562
No-FA | 18 | 115 | 111 | 273 | 257 | 608
No-FA | 22 | 116 | 111 | 260 | 240 | 554
No-FA | 26 | 120 | 117 | 288 | 271 | 579

Table 16: Comparison of using SE(3)-SFA on both the energy and position
prediction heads to No FA and to the top configs of FAENet with no auxiliary
task from (Duval et al., 2023). The symmetry-preservation metrics are in meV
for the energy and milli-Angstroms for the positions. The best model for each
of the 2 categories is in bold and the second best is underlined. We do not
seem to have a correlation between the quality of the invariance and
equivariance with the number of interaction layers in both categories.

###  C.6 Pre-tranining on different tasks

Figure 7: Validation curve during training for both a model trained directly
from scratch for IS2RE and an S2EF model fine-tuned on IS2RE