components undergo positional changes, determining the stability and behaviour
of the system. When the force field is conservative, the forces experienced by
the atoms can be derived from the potential energy by taking the derivative
with respect to their positions,

$E=-\int F(\vec{x})d\vec{x},\quad\text{ and }\quad F=-\nabla E(\vec{x})$ (53)

By definition, these forces will always tend to minimize the potential energy,
driving the system towards a relaxed state, analogous to gradient descent
optimization.

In ML models for chemical systems, a single neural network is used to predict
the total energy of the system by summing the contributions from individual
atoms. To maintain energy conservation, the model calculates the forces on
each atom by computing the derivative of the final predicted energy with
respect to the atom’s position. This allows for accurate predictions of both
energy and forces, enabling efficient exploration and analysis of chemical
systems.

In practice, this energy conservation requirement is also incorporated into
the model’s loss function, which includes a new term calculating the
difference between predicted forces (obtained through backpropagation of the
predicted energy) and ground truth forces (energy conserving). By considering
the forces, the model adheres to the fundamental principle of energy
conservation and constrains the space of functions explored by the neural
network.

However, it is worth noting that this energy conserving requirement may hamper
model performance. In certain cases, breaking free from this constraint may be
beneficial. For example, in the Open Catalyst Project [Chanussot et al.,
2021], non-energy conserving models525252These models predict forces from
final atom representations using a separate neural network and include a loss
term about force predictions. outperformed energy-conserving ones [Gasteiger
et al., 2021, Shuaibi et al., 2021, Duval et al., 2022]. Whether to strictly
enforce the energy conservation principle remains an active area of research.
Furthermore, energy conservation itself does not necessarily yield accurate
approximation of forces given that modeling function does not guarantee
effectively modeling its gradient. As such further research is needed to
explore the effects of energy and force conservation in atomic systems.

###  A.7 GNN architectural details

####  A.7.1 Message Passing

Message passing is often used to describe the functioning of the family of
Graph Neural Networks (GNN) models. Why? Because updating the representation
of each node $i$ can be seen as passing a message from neighbouring nodes
($j\in\mathcal{N}_{i}$) to the node of interest. In its simplest form, the
message $\mathbf{{m}}_{ij}$ is computed via a learnable message function
$\underline{f_{1}}$ of the neighbour’s representation ${\bm{s}}_{j}$ and
${\bm{s}}_{i}$. The updated representation at layer (t+1) is obtained by
applying a learnable update function $\underline{f_{2}}$ on the aggregated
messages $\oplus_{j\in\mathcal{N}_{i}}\mathbf{{m}}_{ij}$ coming from
neighbouring nodes and the existing representation $\mathbf{{s}}_{i}$.
$\oplus$ is not learnable and often denotes the sum or mean operator.

$\displaystyle\mathbf{{m}}_{ij}$
$\displaystyle=\underline{f_{1}}(\mathbf{{s}}_{i}^{(t)},\mathbf{{s}}_{j}^{(t)})$
$\displaystyle\mathbf{{s}}_{i}^{(t+1)}$
$\displaystyle=\underline{f_{2}}(\mathbf{{s}}_{i}^{(t)},\oplus_{j\in\mathcal{N}_{i}}\mathbf{{m}}_{ij})$

Message Passing has proven very successful so far but it also comes with its
set of own limitations. Among them we find over-smoothing, where node features
become too similar after multiple message passing layers, loosing
discriminative power due to excessive aggregation, and over-squashing, which
denotes an excessive information compression through bottleneck edges, likely
leading to information loss [Giraldo et al., 2023].

Going further → A Gentle Introduction to Graph Neural Networks (Distil, 2021)
blogarticle blog video blog blog blog book blog technical blog visual blog []
visualarticle visual video visual blog visual book visual technical visual
visual visual [] → Graph Representation Learning Book (Hamilton, 2020)
bookarticle book video book blog book book book technical book visual book []

####  A.7.2 Activation function

As for any deep learning model, the choice of non-linear activation function
plays a central role in modeling complex non-linearities of atomic
interactions. The Rectified Linear Unit (ReLU) activation [Glorot et al.,
2011] is widely used in many deep learning models [Sanchez-Gonzalez et al.,
2020, Bapst et al., 2020]. However, given the inherent nature of forces, ReLU
may not be ideal to model atomic forces because its output is modeled as
piece-wise linear hyperplanes with sharp boundaries. As a result, a wide array
of alternatives have been explored: Tanh, Leaky-ReLU [Maas et al., 2013],
SoftPlus [Dugas et al., 2001], Shifted SoftPlus [Schütt et al., 2018], and
Swish [Ramachandran et al., 2017].

While the choice of activation function ultimately depends on the specific
problem domain and dataset characteristics, Geometric GNNs most commonly use
the Swish activation ${\rm swish}(x)=x\cdot{\rm sigmoid}(x)$ as it offers
several advantages. It provides a smoother output landscape and has non-zero
activation for negative inputs. This smoothness is essential when dealing with
atomic systems, ensuring that small perturbations in the input space result in
gradual changes in the output space, which promotes stability and avoids
abrupt changes in predictions. Additionally, Swish has a smooth gradient,
making backpropagation more stable and efficient, mitigating issues like
vanishing or exploding gradients in deep networks. Finally, its non-linear
behavior allows Geometric GNNs to model complex molecular interactions and
spatial relationships effectively.

Opinion The swish activation works well but the community should not stop
looking for better alternatives, endowed with similar desirable properties.

####  A.7.3 Basis functions

In Geometric GNNs, the choice of basis function
$\psi:\mathbb{R}^{d}\to\mathbb{R}^{a}$ is essential to transform geometric
information, i.e. atom relative positions $\vec{x}_{ij}$, into discriminative
representations. In different words, basis functions encode the spatial
relationship between atoms, enabling Geometric GNNs to effectively model the
structural properties of the system. The dimension $a$ of the encoded
geometric information is a hyperparameter depending on the choice of the basis
function. It is usually chosen to be significantly larger than the default
dimensionality to capture fine-grained distinctions in atom positions. Below,
we describe the most widespread basis functions:

  1. 1.

Identity: $\psi_{\rm id}(\vec{x}_{ij})=\mathbf{{x}}_{ij}$. The standard basis
function uses the “scalarised” geometric information for each edge $(i,j)$,
basically keeping the same values but loosing its geometric aspect
($\vec{x}\in\mathbb{R}^{d}$ to $\mathbf{{x}}\in\mathbb{R}^{d}$).

  2. 2.

Radial Basis Function (RBF): $\psi_{\rm
rbf}(\vec{x}_{ij})=[\psi_{1},\ldots,\psi_{a}]$, where $\psi_{k}$ is the output
of the $k$-th basis function $\psi_{k}(\vec{x}_{ij})={\rm
e}^{\left(\|\vec{x}_{ij}\|-\mu_{k}\right)^{2}/(2\cdot\sigma^{2})}$
representing the distance between atoms $(i,j)$, encoded using Gaussian
functions that decay with distance. By considering pairwise distances, the RBF
basis function can encode the varying degrees of influence that atoms exert on
each other based on their spatial proximity. The Gaussian means are evenly
distributed on $[0,1]$, i.e., $\mu_{k}=k/(a-1)$ and the standard deviation is
$\sigma=1/(a-1)$. Values of $\vec{x}_{ij}$ are often normalized to lie in
$[0,1]$. $\psi_{\rm rbf}$ results in a $a$-dimensional vector.

  3. 3.

Sine: $\psi_{\rm sin}(\hat{x}_{ij})=[\psi_{1},\ldots,\psi_{b}]$, where
$\psi_{k}$ is the output of the $k$-th basis function
$\psi_{k}(\hat{x}_{ij})=\sin(1.1^{k}\hat{x}_{ij})$. This design is based on
function approximation using the Fourier series. $\psi_{\rm sin}$ is applied
to each dimension of the unit vector $\hat{x}_{ij}$, resulting in an
$a=b\times d$ vector.

  4. 4.

Spherical Harmonics: $\psi_{\rm
sph}(\hat{x}_{ij})=[\vec{\mathbf{Y}}^{(0)}(\theta,\phi),\ldots,\vec{\mathbf{Y}}^{(L)}(\theta,\phi)]$,
where the polar and azimuthal angles $\theta$ and $\phi$ are directly computed
from the unit directional vector
$\hat{x}_{ij}\in{\mathbb{R}}^{3},\|\vec{x}_{ij}\|=1$. The vector-valued
function $\vec{\mathbf{Y}}^{(l)}:S^{2}\rightarrow\mathbb{R}^{(2l+1)}$ is the
list of Laplace’s spherical harmonics
$\vec{Y}^{(l)}_{m}:S^{2}\rightarrow\mathbb{R}$ with order $l\geq 0$ and degree
$m\in\\{-l,\ldots,l\\}$ [MacRobert, 1947]. The spherical harmonics are special
function which encode angular information from the surface of the sphere,
forming an orthonormal basis for Fourier transformations of functions on the
sphere, like sine waves on $\mathbb{R}$. They are equivariant in SO$(3)$,
transforming in a predictable manner when the input is rotated by
$\mathbf{{R}}\in\text{SO}(3)$:

$Y^{(l)}_{m}(\mathbf{{R}}\cdot\hat{x}_{ij})=\sum_{m^{\prime}}\vec{D}^{(l)}_{mm^{\prime}}(\mathbf{{R}})\vec{Y}^{(l)}_{m^{\prime}}(\hat{x}_{ij})$

where $\vec{D}^{(l)}_{mm^{\prime}}$ are the entries of the Wigner-D matrix
$\vec{\mathbf{D}}^{(l)}\in\mathbb{R}^{(2l+1)\times(2l+1)}$. For Geometric
GNNs, spherical harmonics are often used as convolutional filter (or simple
input feature) as they enable to capture the spatial orientation and angular
relationships between atoms. The list of spherical harmonics is displayed here
and an implementation is available here.

  5. 5.

MLP: $\psi_{\rm
MLP}(\vec{x}_{ij})=\sigma(\underline{\bm{W}}\cdot\vec{x}_{ij}+\underline{\bm{b}})$,
where $\sigma(\cdot)$ is a non-linear activation function, and
$\underline{\bm{W}}$ and $\underline{\bm{b}}$ are learnable parameters. This
method is flexible because it can be applied to any geometric quantity (atom
relative position, concatenation of scalar distance and angles, [Duval et al.,
2023]); and powerful because two-layers MLP with enough hidden layers are
universal approximators. However, the MLP basis function cannot preserve
equivariance due to the presence of non-linearities.

Going further → Radial Basis Functions, RBF Kernels, & RBF Networks Explained
Simply (Ye, 2020) blogarticle blog video blog blog blog book blog technical
blog visual blog [] visualarticle visual video visual blog visual book visual
technical visual visual visual [] → Achieving Rotational Invariance with
Bessel-Convolutional Neural Networks (Delchevalerie, 2021) articlearticle
article video article blog article book article technical article visual
article [] videoarticle video video video blog video book video technical
video visual video [] → Spherical Harmonic blogarticle blog video blog blog
blog book blog technical blog visual blog []

####  A.7.4 Examples of architecture

In this subsection, we display the architecture of GemNet [Gasteiger et al.,
2021] to demonstrate that best-performing invariant methods often come at the
cost of a complex functioning. We also describe below the well known SchNet
[Schütt et al., 2017] architecture, having as main objective to depict the
widely used continuous convolution.

SchNet introduced the use of continuous filters to handle unevenly spaced
data, e.g. atoms at arbitrary positions. Given 3D atom input positions
$\vec{\mathbf{x}}\in\mathbb{R}^{n\times 3}$, the continuous-filter
convolutional layer $t$ requires a filter-generating function

$\underline{\psi}^{t}:\mathbb{R}^{3}\rightarrow\mathbb{R}^{a},$

that maps positions to the corresponding discrete convolution filter values.
This learnable filter generating function is modeled with a neural network
where the (invariant scalar) output $\mathbf{{s}}_{i}^{t+1}$ of the continuous
convolutional layer at position $\vec{x}_{i}$ is given by

$\mathbf{{s}}_{i}^{t+1}=\sum_{j\in\mathcal{N}_{i}}\mathbf{{s}}^{t}_{j}\odot\underline{\psi}^{t}(\vec{x}_{ij})=\sum_{j\in\mathcal{N}_{i}}\mathbf{{s}}^{t}_{j}\odot\underline{f}^{t}(\psi(d_{ij}))$
(54)

where "$\odot$" represents the element-wise multiplication, $\underline{f}$
represents a two-layer MLP with softplus activation function and $\psi$ a
radial basis function✜ . These feature-wise convolutions are applied for
computational efficiency. The interactions between feature maps are handled by
separate object-wise or, specifically, atom-wise layers in SchNet.

Figure 25: The full GemNet architecture. $\square$ denotes the layer’s input,
$\|$ concatenation, $\sigma$ the SiLU non-linearity, and orange a layer with
weights shared across interaction blocks. Differences between two-hop message
passing (Q-MP) and one-hop message passing (T-MP) are denoted by dashed lines.
Numbers next to connecting lines denote embedding sizes. Taken from the
original paper [Gasteiger et al., 2021].

Going further → Graph ML in 2023: The State of Affairs, (Galkin, 2022)
blogarticle blog video blog blog blog book blog technical blog visual blog []
visualarticle visual video visual blog visual book visual technical visual
visual visual []

###  A.8 Expressive power of Geometric GNN

Figure 26: The Geometric Weisfeiler-Leman Test, an upper bound on the
expressive power of equivariant GNNs [Joshi et al., 2023a]. GWL distinguishes
non-isomorphic geometric graphs $\mathcal{G}_{1}$ and $\mathcal{G}_{2}$ by
injectively assigning colours to distinct neighbourhood patterns, up to global
symmetries (here $O(d)$). Each iteration expands the neighbourhood from which
geometric information can be gathered (shaded for node $i$). Example inspired
by Schütt et al. [2021b].

The graph isomorphism problem and the Weisfeiler-Leman (WL) test for
distinguishing non-isomorphic graphs have become a powerful tool for analysing
the expressive power of traditional GNNs [Jegelka, 2022]. It was shown by Xu
et al. [2019], Morris et al. [2019] that message passing GNNs are at most as
powerful as WL at distinguishing non-isomorphic graphs and suffer from the
same failure modes as WL. The WL framework has since become a major driver of
progress in designing more expressive GNNs [Dwivedi et al., 2023, Bodnar et
al., 2021].

However, WL does not directly apply to geometric graphs as they exhibit a
stronger notion of geometric isomorphism that must accounts for spatial
symmetries. Two geometric graphs can only be _geometrically isomorphic_ if the
underlying graphs are isomorphic _and_ the geometric attributes are
equivalent, up to global group actions like rotations and reflections.

Joshi et al. [2023a] recently proposed the Geometric WL (GWL) framework for
characterising the expressivity of Geometric GNNs by their ability to solve
geometric graph isomorphism, i.e. to provide distinct representations for any
two different geometric graphs, up to group actions. In addition to their
theoretical contributions, they proposed several synthetic experiments to test
new Geometric GNNs’ expressivity in a practical manner. For instance, one task
asks models to distinguish between several counterexamples from Pozdnyakov and
Ceriotti [2022]. These edge cases consists of pairs of configurations that are
indistinguishable when comparing their set of k-body scalars, illustrated in
Figure 27. In other terms, any model for atom-centred properties that uses 2,
3 or 4-body order features will incorrectly give identical results for the
different configurations. Hence, Geometric GNNs that do not process enough
information to uniquely distinguish two atomic systems, such as SchNet or
DimeNet, automatically fails this experiment.

Figure 27: Counterexamples from [Pozdnyakov and Ceriotti, 2022] where pairs
are distinct but cannot be discriminated by the unordered list of distances or
distances and angles between atoms. They were created to demonstrate that some
Geometric GNNs cannot distinguish between these pairs of structures using
k-body scalarisation. (a) The two tetrahedra share the same list of pairwise
distances, as per color coding. (b) The two structure share the same list of
pairwise distances, and in addition the list of distances of each atom
relative to its neighbors. (c) The two environments share the same list of
distances and angles relative to the central (gray) atom. (d) The two
environments share the same list of distances, angles and tetrahedra around
the central (gray) atom.

Opinion Under the GWL framework, Geometric GNNs require injective aggregation,
update and readout functions to be maximally expressive or universal. While
some current architectures have claimed to be universal or _complete_ under
specific conditions [Gasteiger et al., 2022, Wang et al., 2022a], we believe
that a provably universal, equivariant GNN on sparse graphs with finite tensor
and body order remains an open question.

Going further → On the Expressive Power of Geometric Graph Neural Networks.
articlearticle article video article blog article book article technical
article visual article [] → Geometric GNN Dojo: Reference implementations and
synthetic experiments to evaluate Geometric GNN expressivity in practice.
visualarticle visual video visual blog visual book visual technical visual
visual visual []

###  A.9 On inductive biases

A major research area in deep learning is to find ways to express preferences
over the kinds of functions we would like to solve our problems. For instance
with Maximum A Posteriori (MAP) estimation one could use a prior distribution
over a model’s parameters to express a preference, a prior belief, over the
possible values of the function’s parameters. Those preferences are called
inductive biases because they are baked-in the algorithms in ways that
voluntarily restrict the accessible function space during learning, or favour
specific regions. Those restrictions are most often designed to improve data
efficiency (the number of data points required to get to some level of
performance) and/or generalization (the ability to keep performing well on new
data). Typical inductive biases include: parameter regularization (as in the
MAP example) [Bishop, 1995, Srivastava et al., 2014, Nusrat and Jang, 2018],
enforcing invariance or equivariance to certain input transformations [Lee et
al., 2019, Rowley et al., 1998, Bronstein et al., 2017] (as is the case for
Geometric GNNs) or into the very architecture of the neural network [He et
al., , Yu and Koltun, 2016, Battaglia et al., 2018, Sherstinsky, 2020, Wang et
al., 2023c]. Discovering, leveraging and evaluating inductive biases is
therefore of paramount importance when developing learning approaches to solve
real-world problems, especially when data (or compute) is scarce.

Going further → Inductive Biases for Deep Learning of Higher-Level Cognition
(Goyal & Bengio, 2020) articlearticle article video article blog article book
article technical article visual article [] technicalarticle technical video
technical blog technical book technical technical technical visual technical
[] → Useful Inductive Biases for Deep Learning in Molecular Science (Welling,
2022) videoarticle video video video blog video book video technical video
visual video []

##  Appendix B An opinionated history of methods

Similarly to other field, the first prediction models for molecular energy and
forces relied on hand-crafted representations [Behler, 2016, Halgren, 1996,
Chmiela et al., 2018] built on physical properties. This was the case until
recently, where research moved to end-to-end machine learning models based on
Graph Neural Networks [Gori et al., 2005].

Multiple work came expanding the application of ML techniques to a broad set
of materials modeling tasks ranging from solid-state [Zitnick et al., 2020,
Miret et al., 2023, Lee et al., 2023b] to molecular [Hoja et al., 2021,
Ramakrishnan et al., 2014] structures. Most existing GNN architectures have
tried to incorporate physics-informed 3D symmetries, either directly in the
model architecture, making model predictions explicitly invariant or
equivariant to the desired transformations, or via the data. There are
Geometric GNNs.

A first line of Geometric GNNs were constructed to be E(3)-invariant by
extracting scalar representations from atomic relative positions [Unke and
Meuwly, 2019, Klicpera et al., 2020, Liu et al., 2021, Shuaibi et al., 2021,
Ying et al., 2021, Adams et al., 2022]. The evolution was steady as we went
from SchNet [Schütt et al., 2017] using pair-wise distances, to DimeNet
[Klicpera et al., 2020] adding bond angles, to SphereNet [Liu et al., 2022]
and GemNet [Gasteiger et al., 2021, 2022] additionally incorporating torsion
angles (i.e. quadruplet of atoms). The scalarisation of geometric information
enables to apply traditional message passing schemes (with any non-linearity)
and the use of additional information enables models to distinguish a larger
set of atomic systems. Nevertheless, this comes at a greater computational
cost since GemNet must look at 3-hop neighbourhoods to compute torsion
information for each update step, at both training and inference time.
Besides, these models extract a set of pre-defined scalars representing
geometric information and cannot represent equivariant properties directly,
due to their invariant nature.

In parallel, Equivariant GNNs with Spherical Coordinates [Thomas et al., 2018,
Anderson et al., 2019, Fuchs et al., 2020, Brandstetter et al., 2021, Batatia
et al., 2022b, Frank et al., 2022], also called Spherical EGNN, focused on
enforcing equivariance using irreducible representations of the SO(3) group.
Such models build upon the concepts of steerability and equivariance
introduced by Cohen and Welling [Cohen and Welling, 2016]. They use spherical
tensors as node embeddings and ensure equivariance to SO(3) by placing
constraints on the operations that can be performed [Kondor et al., 2018].
Specifically, they compute linear operations with a generalized tensor product
between the atom embeddings and edges’ directions. Tensor Field Networks
[Thomas et al., 2018], NequIP [Batzner et al., 2022], SEGNN [Brandstetter et
al., 2021], MACE [Batatia et al., 2022b], Allegro [Musaelian et al., 2022],
Equiformer [Liao and Smidt, 2023], and various others lie in this category.
They are commonly referred to as e3nn networks. While these methods are
expressive and generalize well, they can be hard to implement, constraint a
lot the functional space and optimization landscape, and very computationally
expensive at training and inference (due to the Clebsch-Gordan tensor product
with spherical representations).

For this reason, Equivariant GNNs with Cartesian Coordinates [Schütt et al.,
2021a, Satorras et al., 2021b, Thölke and De Fabritiis, 2022], also called
Cartesian EGNNs, started to model equivariant interactions in the Cartesian
space, updating both scalar and vector representations. These GNNs achieve
good performance while being relatively fast by avoiding the expensive
equivariant operations of Spherical EGNNs. The authors manually design two
separate sets of functions to deal with each type of representation and often
use complex operations to mix their information, restricting the set of
possible operations (e.g. vector element-wise dot product) to preserve
equivariance.. As downsides, in addition to rendering the global architecture
hard to understand, the decomposition of Cartesian tensors into spherical
tensors offers many nice properties that Cartesian EGNNs lack. For instance,
they lack structured, hierarchical and compact representation of geometric
information, which improves model efficiency and memory utilization. They also
don’t capture information across different angular directions, which makes the
model more sensitive to angular variations and dependencies present in atomic
interactions.

Overall, we have witnessed in recent years incredible progress in terms of
model architectures across invariant and equivariant GNNs. However, all above
approaches restrict the model learning space, reducing the number of possible
operations inside the model architecture. Although this is theoretically
desirable, in practice, it may hamper the learning capacity of the model. For
this reason, we have seen the appearance of unconstrained Geometric GNNs.
Unconstrained GNNs do not enforce symmetries via the model architecture.
Instead, they attempt to implicitly learn them through simple data
augmentation [Hu et al., 2021]; they relax them [Zitnick et al., 2022] or they
enforce them by mapping input data to a unique canonical space of all
euclidean representations [Duval et al., 2023].

In this work, we attempt to bridge an existing gap in the literature by
providing a holistic and opinionated overview of the field of Geometric GNNs,
encompassing all important aspects. While we are the first work of such kind,
we acknowledge the presence of some great recent works also attempting to
bridge this gap [Wang et al., 2022c, Han et al., 2022, Liu et al., 2023, Atz
et al., 2021]. Wang et al. [2022c] focuses on the the molecular applications
of GNNs; Liu et al. [2023] proposes a benchmarking platform for various GNNs
and molecular datasets; Han et al. [2022], Atz et al. [2021] offer a short
summary of the field. Overall, they all provide a descriptive and relatively
general overview of the field, not always specific to Geometric GNNs. Our
distinguishing contributions can be listed as follows: (1) a thorough
description of all steps and variations in the Geometric GNNs modeling
pipeline; (2) a novel taxonomy of methods containing a clear mathematical
relation between them; (3) a concise description of how Geometric GNNs power
different applications, with associated datasets; (4) a detailed list of
promising future research directions, plus opinions on many ongoing
discussions that divide the field; (5) the inclusion of almost all contextual
material needed to understand the field; (6) a new notation scheme; (7) an
exhaustive list of approaches and datasets that need to be updated by the
community. We hope that our work will be useful to the whole community of
researchers, helping experienced ones to efficiently navigate the field and
guiding newcomers to integrate it.

Going further → Towards Geometric Deep Learning (Bronstein, 2023) blogarticle
blog video blog blog blog book blog technical blog visual blog []

##  Appendix C Data

###  C.1 Data splits

  1. 1.

Random split: ensures training, validation, and test data are sampled from the
same underlying probability distribution.

    1. –

Pros: Simple to implement.

    2. –

Cons: May not preserve the underlying distribution of data. May result in
variability in performance due to random variations in the data split. May not
capture specific challenges or biases in the data.

    3. –

When to use? They can provide a good baseline evaluation, especially when the
dataset is well-balanced and representative.

  2. 2.

Stratified split: The dataset is divided while maintaining a similar class
distribution across sets (e.g., balanced representation of active and inactive
compounds).It ensures each of the training, validation, and test sets to cover
the full range of provided labels.

    1. –

Pros: Helps ensure that each split contains a representative distribution of
different classes or properties, reducing the risk of biased evaluations.

    2. –

Cons: Requires class or property information for stratification, which may not
always be available or applicable. Small or imbalanced datasets may still pose
challenges. Not optimal to measure generalisation ability.

    3. –

When to use ? When there is a class or property imbalance in the dataset.

  3. 3.

Extrapolation split: The dataset is divided such that some targeted molecules
are placed in the test set without being included in the training set. This is
often referred to as out-of-distribution (OOD).

    1. –

Pros: Aims to assess a model’s ability to generalize to unseen chemical
environments. It challenges the model to predict properties based on general
chemical principles rather than memorizing specific training instances,
yielding more robust models.

    2. –

Cons: Such split requires careful curation of OOD examples that maintain
chemical relevance while representing unseen combinations. The design of OOD
examples may inadvertently introduce biases or artifacts (e.g. low density
region), impacting the fairness of model evaluations.

    3. –

When to use ? When we care about ood generalisation of the model. For e.g., in
materials or drug discovery applications.

  4. 4.

Time Split: The dataset is divided based on a chronological order, such as
using earlier time points for training and later time points for validation or
testing.

    1. –

Pros: Reflects real-world scenarios where models are trained on historical
data and tested on future data. Allows evaluation of model performance under
temporal variations.

    2. –

Cons: Assumes that data collected at different time points are representative
of the same underlying distribution. May not be suitable for all datasets or
tasks.

    3. –

When to use? In presence of meaningful temporal data.

  5. 5.

Group split: The dataset is divided based on specific groups or categories
present in the data (e.g., different targets, protein families, chemical
series).

    1. –

Pros: Enables evaluation of model performance on specific subsets of the data
that may have distinct characteristics or challenges. Can provide insights
into target or group-specific performance.

    2. –

Cons: Requires prior knowledge or information about the groups or categories.
May introduce biases if the groups are not representative or if data in
different groups have varying characteristics.

    3. –

When to use? When the data can be grouped based on specific characteristics or
challenges.

###  C.2 Examples of predicted properties

  * •

band gap: determines a material’s electronic behavior and is relevant in areas
such as semiconductor design and solar cell applications.

  * •

dielectric constant: measures the ability of a substance or material to store
electrical energy. It is an expression of the extent to which a material holds
or concentrates electric flux. Dieletric constant has important applications
in energy storage devices and electrical substation equipments.

  * •

Refractive index is a measure of how light propagates through a material. It
is an important property in optics and photonics applications, as it
determines the material’s ability to manipulate light.

  * •

glass: is a classification property indicating if a material is a glass former
or not. Glass-forming materials are important in the field of materials
science, as their amorphous structure offers unique properties and
applications in fields like optics, electronics, and energy storage.

  * •

jdft2d: Exfoliation energy represents the energy required to separate or
exfoliate a layered material into individual layers. This property is relevant
in the field of 2D materials, where exfoliation plays a crucial role in
obtaining thin layers with desired properties.

  * •

Formation energy of a material provides insights into its stability and the
energy involved in its formation. It is useful for understanding the
feasibility of synthesizing new materials and their thermodynamic properties

  * •

Phonon energy: the frequency of the highest frequency optical phonon mode peak
provides information about the lattice vibrations in a material. It is
relevant for understanding thermal conductivity, phonon transport, and thermal
properties of materials.

  * •

Forces: the forces acting on each atom, commonly used to perform molecular
dynamics simulations.

  * •

Relaxed energy: refers to the minimum energy state of a molecule or system,
obtained through computational methods. It provides insights into the
energetics of the molecule and aids in understanding its structural properties
and interactions.

  * •

Dipole Moment: is a measure of the separation of positive and negative charges
within a molecule. It quantifies the molecule’s polarity and can be important
for predicting its behavior in certain chemical reactions or interactions.

  * •

HOMO (Highest Occupied Molecular Orbital): refers to the highest energy level
in a molecule that is occupied by electrons. It is an electronic property used
to describe the reactivity and stability of molecules.

  * •

Polarizability: is a property that describes the ability of a molecule to be
deformed by an external electric field. It reflects the molecule’s response to
changes in the electric field and is important in studying intermolecular
forces and interactions.

  * •

Heat Capacity: measures the amount of heat energy required to raise the
temperature of a substance. It can be used to understand and predict how a
molecule will respond to changes in temperature, providing insights into its
thermodynamic properties.

  * •

Toxicity: refers to the degree to which a substance can cause harm or adverse
effects to living organisms. Machine learning models trained to predict
toxicity can help identify potentially toxic compounds and aid in drug
development and safety assessment.

  * •

Solubility: is the property of a substance to dissolve in a solvent to form a
homogeneous solution. Solubility prediction involves determining the
likelihood of a compound to dissolve in a particular solvent or under specific
conditions. Accurate solubility predictions are valuable in environmental
studies, drug discovery and material science.

###  C.3 Atom-type rescaling

In this data pre-processing approach, the target values (e.g., molecule
energy) are shifted or corrected based on a learned or calculated energy
contribution from individual atom types. The idea is to factor out the
contributions from different atom types and create a more interpretable target
variable that represents the interaction energy between atoms. This can help
the model focus on learning the residual energy variations after accounting
for the atomic contributions.

To calculate the shift values (e.g. energy contributions) of each atom type
based on the training dataset, one must first construct a matrix ${\bm{A}}$ of
shape $(n,z_{max}+1)$, where $n$ is the number of dataset samples and
$z_{max}$ represents the largest atomic number. This matrix counts the
occurrences of each atomic type for each structure.

Next, one must solve a linear equation ${\bm{A}}{\bm{x}}={\bm{q}}$ to
calculate the per atom shifts, where ${\bm{x}}$ is a vector of shape
$(z_{max}+1)$ representing the unknown shifts for each atom type, and
${\bm{q}}$ is the target quantity of shape $(n)$, which can be the ground
truth energy. The solution ${\bm{x}}$ provides the shifts for each atom type.

Finally, to obtain the shifted energies $y^{\prime}$, the sum of the shifts
corresponding to every atom in the atomic system is subtracted from the ground
truth energy $y$. This ensures that the energy values are adjusted according
to the atom-type specific shifts, effectively balancing the contributions of
different atom types to the overall potential energy.

$y^{\prime}=y-\sum_{i\in\mathcal{{G}}_{j}}x_{i}$

Target value shifting is particularly common in the field of computational
chemistry. It’s used to improve the prediction accuracy of ML models by
factoring out known atomic contributions from the target values, allowing to
focus on the finer details of atomic interactions. In different terms, the
objective is to obtain rescaled target values exhibiting a desirable
distribution and reduced variance, which enhances the quality of the training
data and improves the stability/convergence of the GNN model during training.
While it requires accurate calculation of energy contributions, this approach
can help capture subtleties that might be challenging to learn directly from
raw data. For instance, it is commonly applied on the QM7-X dataset [Hoja et
al., 2021], as described in this work.

Finally, note that although slightly less common on materials and molecules,
atom type shift can be applied to input features instead of target variables.
This is the case when ML models use atomic properties having very different
scales, ultimately making it easier for the model to learn patterns across
different atom types.

Opinion If utilised, this pre-processing step must be reported clearly in the
paper for reproducibility, with corresponding shift values.