Effectively predicting molecular interactions has the potential to accelerate molecular dynamics by multiple orders of magnitude and thus revolutionize chemical simulations. Graph neural networks (GNNs) have recently shown great successes for this task, overtaking classical methods based on fixed molecular kernels. However, they still appear very limited from a theoretical perspective, since regular GNNs cannot distinguish certain types of graphs. In this work we close this gap between theory and practice. We show that GNNs with directed edge embeddings and two-hop message passing are indeed universal approximators for predictions that are invariant to translation, and equivariant to permutation and rotation. We then leverage these insights and multiple structural improvements to propose the geometric message passing neural network (GemNet). We demonstrate the benefits of the proposed changes in multiple ablation studies. GemNet outperforms previous models on the COLL, MD17, and OC20 datasets by 34, 41, and 20, respectively, and performs especially well on the most challenging molecules. Our implementation is available online. [<https://www.daml.in.tum.de/gemnet>]

§ INTRODUCTION

Graph neural networks (GNNs) have shown great promise for predicting the energy and other quantum mechanical properties of molecules. They can predict these properties orders of magnitudes faster than methods from quantum chemistry – at comparable accuracy. GNNs can thus enable the accurate simulation of systems that are orders of magnitude larger. However, they still exhibit severe theoretical and practical limitations. Regular GNNs are only as powerful as the 1-Weisfeiler Lehman test of isomorphism and thus cannot distinguish between certain molecules [Morris et al., 2019, Xu et al., 2019]. Moreover, they require a large number of training samples to achieve good accuracy.

In this work we first resolve the questionable expressiveness of GNNs by proving sufficient conditions for universality in the case of invariance to translations and rotations and equivariance to permutations; and then extending this result to rotationally equivariant predictions. Simply using the full geometric information (e.g. all pairwise atomic distances) in a layer does not ensure universal approximation. For example, if our model uses a rotationally invariant layer we lose the relative information between components. Such a model thus cannot distinguish between features that are rotated differently. This issue is commonly known as the “Picasso problem”: An image model with rotationally invariant layers cannot detect whether a person's eyes are rotated correctly. Instead, we need a model that preserves relative rotational information and is only invariant to global rotations. To prove universality in the rotationally invariant case we extend a recent universality result based on point cloud models that use representations of the rotation group $\SO3$ [Dym & Maron, 2021]. We prove that spherical representations are actually sufficient; full $\SO3$ representations are not necessary. We then generalize this to rotationally equivariant predictions by leveraging a recent result on extending invariant to equivariant predictions [Villar et al., 2021]. We then discretize spherical representations by selecting points on the sphere based on the directions to neighboring atoms. We can connect this model to GNNs by interpreting these directions as directed edge embeddings. For example, the embedding direction of atom $a$ would be defined by atom $c$, resulting in the edge embedding $e_{ca}$. Updating the spherical representation of atom $a$ based on atom $b$ then corresponds to two-hop message passing between the edges $e_{ca}$ and $e_{db}$ via $e_{ba}$, with atoms $c$ and $d$ defining the embedding directions. This message passing formalism naturally allows us to obtain the molecule's full geometrical information (distances, angles, and dihedral angles), and the direct correspondence proves the model's universality.

We call this edge-based two-hop message passing scheme geometric message passing, and propose multiple structural enhancements to improve the practical performance of this formalism. Based on these changes we develop the highly accurate and sample-efficient geometric message passing neural network (GemNet). We furthermore show that stabilizing the variance of GemNet's activations with predetermined scaling factors yields significant improvements over regular normalization layers.

We investigate the proposed improvements in a range of ablation studies, and show that each of them significantly reduces the model error. These changes introduce little to no computational overhead over two-hop message passing. Altogether, GemNet outperforms previous models for force predictions on COLL by 34, on MD17 by 41, and on OC20 by 20 on average. We observe the largest improvements for the most challenging molecules, which exhibit dynamic, non-planar geometries.
In summary, our contributions are:

* Showing the universality of spherical representations and two-hop message passing with directed edge embeddings for rotationally equivariant predictions.

* Geometric message passing: Symmetric message passing enhanced by geometric information.

* Incorporating all proposed improvements in the Geometric Message Passing Neural Network (GemNet), which significantly outperforms previous methods for molecular dynamics prediction.

§ RELATED WORK

Machine learning potentials. Research on predicting a molecule's energy and forces (so-called machine learning potentials) started with hand-fitted analytical functions and then gradually moved towards fully learned models. Arguably, classical force fields are their very first instances. They use analytical functions with coefficients that were hand-tuned based on experimental data. A popular example for these is the Merck Molecular Force Field (MMFF94) [Halgren, 1996]. The next wave of methods used kernel ridge regression based on fixed, hand-crafted molecular representations [Bartók et al., 2010, Chmiela et al., 2017, Faber et al., 2018]. Finally, modern research mostly focusses on fully end-to-end learnable models based on GNNs [Gilmer et al., 2017, Schütt et al., 2017]. These models can also be combined with molecular features from quantum mechanical calculations to improve performance [Qiao et al., 2020]. We consider this combination as orthogonal research.

Directional GNNs. We can also achieve equivariance and invariance to rotations without relying on group representations. Directional GNNs achieve this by representing directional information explicitly [Schütt et al., 2021] or in the form of angles [Gasteiger et al., 2020] and dihedral angles [Flam-Shepherd et al., 2021, Liu et al., 2021]. Our work is focused on this class of models, proving their universality and proposing an improved variant, GemNet.

Expressiveness of GNNs. A large part of GNN research has been focused on their (limited) expressiveness. Xu et al., 2019, Morris et al., 2019 first proved that they are only as expressive as the Weisfeiler-Lehman test of isomorphism and Garg et al., 2020 showed the limitations of basic directional message passing. Maron et al., 2019, Morris et al., 2019, Morris et al., 2020, Kondor et al., 2019 then investigated higher-order representations to circumvent this issue. Finally, Maron et al., 2019, Azizian & Lelarge, 2020 showed that so-called folklore GNNs are the most expressive GNNs for a given tensor order.

Equivariant neural networks. Equivariance and invariance have recently emerged as one of the foundational principles of modern neural networks [Cohen & Welling, 2016, Cohen et al., 2019]. This is especially relevant for models in physics, for which we often know the symmetries a priori. Equivariant models for the $\SO3$ group were first investigated in the context of spherical convolutions by Cohen et al., 2018, Esteves et al., 2018, Kondor et al., 2018. These methods leverage group representations to achieve full equivariance. They were then transferred to the context of 3D point clouds and molecules by Weiler et al., 2018, Thomas et al., 2018, Anderson et al., 2019, and further developed by Fuchs et al., 2020, Batzner et al., 2021, Finzi et al., 2020. Importantly, Yarotsky, 2021 proved the universality of 2D convolutional networks, and Bogatskiy et al., 2020 extended this result to general groups. Maron et al., 2020 proved universality for models invariant to $S_n$ and equivariant to an additional symmetry. Dym & Maron, 2021 combined these results to prove universality for the joined group of translations, rotations, and permutations. Apart from reflections this is the exact group relevant for general molecules.

§ UNIVERSALITY OF SPHERICAL REPRESENTATIONS

GNNs for molecules typically incorporate directional information in one of two ways: Via $\SO3$ representations [Thomas et al., 2018, Anderson et al., 2019] or by using directions in real space [Gasteiger et al., 2020, Schütt et al., 2021]. Directions in real space are associated with the three-dimensional $S^2$ sphere, while the $\SO3$ group is double covered by the four-dimensional $S^3$ sphere. Directional representations thus use one degree of freedom less than $\SO3$ representations, making them significantly cheaper. And, as we will prove in this section, directional representations actually provide the same expressivity as $\SO3$ representations for predictions in $\R^3$. We achieve this by showing that the $\SO3$-based tensor field network (TFN) [Thomas et al., 2018] variant used by Dym & Maron, 2021 is equivalent to a similar model based on spherical representations, in the case of rotationally invariant predictions. We then generalize a recent result by Villar et al., 2021, which lets us extend our theorem to the rotationally equivariant case. Afterwards, we relate this universality to directional GNNs by interpreting them as a discretization of spherical representations.

Preliminaries. We consider a point cloud with $n$ points (atoms), each associated with a position and a set of rotationally invariant features (e.g. atom types), defined as $\mX \in \R^{3 \times n}$ and $\mH_{\text{in}} \in \R^{h \times n}$. In this section we define model classes by sets of functions $\gF$. As a first step, we are interested in proving that the set $\gF$ defining our model is equal to the full set of functions $\gG'$ that are invariant to the group of translations $\sT^3$ and rotations $\SO3$, and equivariant to the group of permutations $S_n$. We denote the codomain of functions in $\gG'$ as $W_{\text{T}}^n$, where $W_{\text{T}}$ is some representation of $\SO3$. We denote a vector's norm by $x = \| \vx \|_2$, its direction by $\hat{\vx} = \vx / x$, and the relative position by $\vx_{ba} = \vx_b - \vx_a$. Proofs are deferred to the appendix. Note that this section is not intended as an introduction to the $\SO3$ group. For a concise introduction in the context of machine learning see e.g. <cit.> or Kondor et al., 2018.

Tensor field network. In order to show the equivalence of the TFN to spherical representations, we first need to define this model. Following Dym & Maron, 2021, we split the model into two parts: Embedding functions $\gF_{\text{feat}}$ that lifts the input into an equivariant representation, and pooling functions $\gF_{\text{pool}}$ that aggregate the results of multiple embedding functions on each point and computes the model output. The overall model is then defined as the set of functions

\begin{equation}
\gF_{K(D), D}^{\text{TFN}} = \{ f \mid f(\mX, \mH_{\text{in}}) = \sum_{k=1}^K f_{\text{pool}}^{(k)*}(f_{\text{feat}}^{(k)}(\mX, \mH_{\text{in}})), f_{\text{pool}}^{(k)} \in \gF_{\text{pool}}^{\text{TFN}}(D), f_{\text{feat}}^{(k)} \in \gF_{\text{feat}}^{\text{TFN}}(D) \},
\end{equation}

where $D \in \sN$ denotes the function's maximum polynomial degree, $K(D) \in \sN$ is chosen such that <ref> is fulfilled (Dym & Maron, 2021 only prove the existence of this function), and $f^*$ denotes elementwise application of $f$ on all points. We then define the set $\gF_{\text{pool}}^{\text{TFN}}$ as all rotationally equivariant linear functions on the $\SO3$ group, i.e. all $\SO3$ convolutions [Kostelec & Rockmore, 2008]. Note that these are more expressive than the self-interaction layers used originally [Thomas et al., 2018].
The embedding functions $\gF_{\text{feat}}^{\text{TFN}}(D) = \{ \pi_2 \circ f^{(2D)} \circ \dots \circ f^{(1)} \mid f^{(i)} \in \gF_{\text{prod}}^{\text{TFN}} \}$ consist of an auxiliary function $\pi_2(\mX, \mH) = \mH$ and a series of tensor product functions (called convolution by Dym & Maron, 2021) $\gF_{\text{prod}}^{\text{TFN}} = \{ f \mid f(\mX, \mH) = (\mX, \tilde{\mH}^{\text{TFN}}(\mX, \mH)) \}$. The intermediate representations are $\mH \in W_{\text{feat}}^n$, where $W_{\text{feat}}$ is a representation of $\SO3$ indexed by the degree $l$ and the order $m$. For $\mH_{\text{in}}$ we have $l\!=\!m\!=\!0$. The main update is defined by

\begin{equation}
\tilde{\mH}^{\text{TFN}(l_o)}_{am_o}(\mX, \mH) = \theta \mH^{(l_o)}_{am_o} + \sum_{l_f, m_f} \sum_{l_i, m_i} C_{(l_f, m_f), (l_i, m_i)}^{(l_o, m_o)} \sum_{b \in \gN_a} F^{(l_f)}_{\text{TFN},m_f}(\vx_b - \vx_a) \mH^{(l_i)}_{bm_i},
\label{eq:tfn}
\end{equation}

where $\theta$ is a (learned) scalar and $\gN_a$ are the neighbors of point $a$. The Clebsch-Gordan coefficients $C_{(l_f, m_f), (l_i, m_i)}^{(l_o, m_o)}$ arise from decomposing the tensor product of two input $\SO3$ representations (the filter and input representations) into a sum of output representations. Their exact values are not relevant for this discussion. We index the output with degree $l_o$ and order $m_o$, the learned filter with $l_f$ and $m_f$, and the input with $l_i$ and $m_i$. $F^{(l)}_{\text{TFN},m}(\vx) = R^{(l)}(x) Y_{lm}(\hat{\vx})$ is a rotationally equivariant filter, with a (learned) radial part $R$, which is any polynomial of degree $\le D$, and the real spherical harmonics $Y_{lm}$ with degree $l$ and order $m$. The spherical harmonics are the basis for the Fourier transformation of functions on the sphere, analogously to sine waves for functions on $\R$. We can prove universality for TFNs by using the universality of polynomial regression and showing that TFNs can fit any polynomial (see Dym & Maron, 2021 for details), resulting in:

Consider the set of functions $\gG$ mapping $\R^{3 \times n + h \times n} \to W_{\text{\upshape T}}^n$ that are equivariant to rotations and permutations and invariant to translations. For all $n \in \sN$,

* For $D \in \sN_0$, every polynomial $p \in \gG$ of degree $D$ is in $\gF_{K(D), D}^{\text{\upshape TFN}}$.

* Every continuous function $f \in \gG$ can be approximated uniformly on compact sets by functions in $\bigcup_{D \in \sN_0} \gF_{K(D), D}^{\text{\upshape TFN}}$.

Spherical networks. Instead of intermediate $\SO3$ representations we now switch to spherical representations, which are functions on the sphere $\mH: S^2 \to \R$. We define the set of functions $\gF_{K(D), D}^{\text{sphere}}$ analogously to $\gF_{K(D), D}^{\text{TFN}}$. However, for $\gF_{\text{feat}}^{\text{sphere}}(D)$ we use

\begin{equation}
\tilde{\mH}^{\text{sphere}}_a(\mX, \mH)(\hat{\vr}) = \theta \mH_a(\hat{\vr}) + \sum_{b \in \gN_a} F_{\text{sphere}}(\vx_b - \vx_a, \hat{\vr}) \mH_b(\hat{\vr}),
\label{eq:sphere}
\end{equation}

with the filter function $F_{\text{sphere}}(\vx, \hat{\vr}) = \sum_{l, m} R^{(l)}(x) \Re[Y_m^{(l)*}(\hat{\vx}) Y_{m}^{(l)}(\hat{\vr})]$, using the real part $\Re$ of the complex spherical harmonics $Y_m^{(l)}$. The set of pooling functions for invariant predictions is

\begin{equation}
\gF_{\text{pool}}^{\text{sphere}} = \{ f \mid f(\mH) = \theta_{\text{pool}} \int_{S^2} \mH(\hat{\vr}) \,\text{d}\hat{\vr} \},
\end{equation}

with the learnable parameter $\theta_{\text{pool}}$.
We obtain the universality theorem by showing the equivalence between this model and TFN for rotationally invariant functions. The proof is based on the connection between spherical harmonics and the Clebsch-Gordan coefficients <cit.> (see <ref>).

Consider the set of functions $\gG'$ mapping $\R^{3 \times n + h \times n} \to W_{\text{\upshape T}}^n$ that are equivariant to permutations and invariant to translations and rotations. For all $n \in \sN$,

* For $D \in \sN_0$, every polynomial $p \in \gG'$ of degree $D$ is in $\gF_{K(D), D}^{\text{\upshape sphere}}$.

* Every continuous function $f \in \gG'$ can be approximated uniformly on compact sets by functions in $\bigcup_{D \in \sN_0} \gF_{K(D), D}^{\text{\upshape sphere}}$.

Next, we extend <ref> to rotationally equivariant functions. We do this by generalizing a recent result by Villar et al., 2021 to obtain (see <ref>):

Let $h \colon \R^{d \times n + h \times n} \to \R^{d \times n}$ be any function that is equivariant to permutations and rotations and invariant to translations. For all $a \in [1, n]$, let the set of relative vectors $\{\vx_{ca} \mid c \in [1, n]\}$ not span a $(d-1)$-dimensional space. Then there are $n - 1$ functions $f^{(c)} \colon \R^{d \times n + h \times n} \to \R^{n}$ such that

\begin{equation}
\vh_a(\mX, \mH) = \sum_{\substack{c=1\\c \neq a}}^n f_a^{(c)}(\mX, \mH) \vx_{ca},
\label{eq:vec_pred}
\end{equation}

where $f^{(c)}$ is equivariant to permutations, but invariant to rotations and translations.

This theorem lets us extend a rotationally invariant model to an equivariant one, while providing universality guarantees. Together, <ref> and <ref> (with $d = 3$) thus show that we can approximate any rotationally equivariant function using only representations on the $S^2$ sphere. We thus do not need $\SO3$ representations, spin-weighted spherical harmonics [Esteves et al., 2020], triplet embeddings, or complex-valued functions. This result puts theory back in line with practice, where the best results are currently achieved without relying on these more expensive representations [Schütt et al., 2021].

§ FROM SPHERICAL REPRESENTATIONS TO DIRECTIONAL MESSAGE PASSING

Directional representations. To use spherical representations in a model we first need to find a tractable description. Instead of using spherical harmonics, we propose to sample the representations in specific directions $\hat{\vr}_i$. If we look at recent models, we see that they implicitly use the directions to each atom's neighbors for this purpose, i.e. they embed the edges in the molecule's graph. These directions define an equivariant mesh that circumvents the aliasing effects that would arise from fixed grids [Kondor et al., 2018]. Schütt et al., 2021 flexibly define the directional mesh in each layer by aggregating directions, while Gasteiger et al., 2020 and others use a fixed mesh for each atom. We can refine this mesh of directions e.g. by using more neighbors or by interpolating between directions. The approximation error of this directional mesh is related to the spherical harmonic expansion via the mesh norm and the separating distance between directions [Jetter et al., 1999, Keiner et al., 2007]. Note that depending on the discretization scheme the resulting mesh might not provide a universal approximation guarantee.

<ref> only defines the relationship for a fixed direction, while models commonly use different directional meshes for the input and output. To incorporate this we add a convolution with a learned filter $F_2$, which can only improve the model's expressiveness. Since the input and output are spherical functions, the used filter $F_2$ has to be zonal, i.e. it can only depend on one angle. This can be expressed as [Driscoll & Healy, 1994]

\begin{equation}
\begin{aligned}
\tilde{\mH}^{\text{dir}}_a(\mX, \mH)(\hat{\vr}_o) &= \theta \mH_a(\hat{\vr}_o) + \int_{\SO3} \sum_{b \in \gN_a} F_{\text{sphere}}(\vx_{ba}, \mR \hat{\vn}) \sum_{i \in \gR_b} \mH_{bi} \delta(\mR \hat{\vn} - \hat{\vr}_i) F_2(\mR^{-1} \hat{\vr}_o) \,\text{d}\mR\\
&= \theta \mH_a(\hat{\vr}_o) + \sum_{b \in \gN_a} \sum_{i \in \gR_b} F_{\text{sphere}}(\vx_{ba}, \hat{\vr}_i) \mH_{bi} F_2(\measuredangle \hat{\vr}_o \hat{\vr}_i),
\end{aligned}
\end{equation}

where $\gR_b$ denotes the directional mesh of atom $b$ with mesh directions denoted by $\hat{\vr}_i$, and $\hat{\vr}_o$ specifies the output direction. The integral vanishes due to the Dirac delta $\delta$.

General filters. To see the relationship to GNNs we furthermore need to generalize the filter $F_{\text{sphere}}(\vx_{ba}, \hat{\vr}_i)$. This filter only depends on the angle $\measuredangle \hat{\vr}_i \hat{\vx}_{ba}$ since it is rotationally invariant:

$F_{\text{\upshape sphere}}(\mR \vx, \mR \hat{\vr}) = F_{\text{\upshape sphere}}(\vx, \hat{\vr})$ for any rotation matrix $\mR$.

We can therefore substitute $F_{\text{sphere}}$ with a general learnable filter $F_1$ that is parametrized by this relative angle. Since $F_{\text{sphere}}$ arises as a special case we do not lose expressivity. We thus obtain

\begin{equation}
\tilde{\mH}^{\text{gem}}_a(\mX, \mH)(\hat{\vr}_o) = \theta \mH_a(\hat{\vr}_o) + \sum_{b \in \gN_a} \sum_{i \in \gR_b} F_1(x_{ba}, \measuredangle \hat{\vr}_i \hat{\vx}_{ba}) F_2(\measuredangle \hat{\vr}_o \hat{\vr}_i) \mH_{bi}.
\label{eq:gem}
\end{equation}

We have now arrived at a message passing scheme that has universal approximation guarantees and is only based on relative directional information. To see the connection to GNNs we interpret these discretized spherical representations as edge embeddings pointing towards $\hat{\vr}_o$ and $\hat{\vr}_i$. <ref> then corresponds to two-hop message passing between the edge embeddings of $\hat{\vr}_o$ and $\hat{\vr}_i$ via the edge $\hat{\vx}_{ba}$. Interestingly, the central learnable part of <ref> is the product of the filters $F_1(x_{ba}, \measuredangle \hat{\vr}_i \hat{\vx}_{ba})$ and $F_2(\measuredangle \hat{\vr}_o \hat{\vr}_i)$ with the input representation, which is strikingly similar to the Hadamard product used in modern GNNs [Schütt et al., 2017, Gasteiger et al., 2020] – except that these only use one-hop message passing.

§ GEOMETRIC MESSAGE PASSING

Angles used in geometric message passing. The dihedral angle $\theta_{cabd}$ becomes visible when rotating the molecule so that atoms $a$ and $b$ lie on top of each other (right).

Geometric representation. We now develop a specific two-hop message passing scheme based on <ref>. We use embeddings based on interatomic directions, and embed all atom pairs with distance $x_{ca} \le c_{\text{emb}}$. $\hat{\vr}_o$ and $\hat{\vr}_i$ are thus instantiated as the interatomic directions $\hat{\vx}_{ca}$ and $\hat{\vx}_{db}$. We denote directional embeddings as $\vm_{ca} = \mH_a(\hat{\vx}_{ca})$. Message passing is thus based on quadruplets of atoms – two atoms are interacting ($a$ and $b$) and two atoms define the directions ($c$ and $d$). We denote the angle between directions by $\varphi_{abd} = \measuredangle \hat{\vx}_{ab} \hat{\vx}_{db}$. To improve empirical performance we additionally use the dihedral angle $\theta_{cabd} = \measuredangle \hat{\vx}_{ca} \hat{\vx}_{db} \perp \hat{\vx}_{ba}$ and substitute $\measuredangle \hat{\vr}_o \hat{\vr}_i = \measuredangle \hat{\vx}_{ca} \hat{\vx}_{db}$ with $\varphi_{cab}$. <ref> illustrates the three angles $\varphi_{cab}$, $\varphi_{abd}$, and $\theta_{cabd}$ we use for updating the embedding $\vm_{ca}$ based on $\vm_{db}$. To ensure that all angles are well-defined we exclude overlapping atom quadruplets, i.e. $a\!\neq\!b\!\neq\!c\!\neq\!d$. We represent the relative directional information using spherical Fourier-Bessel bases with polynomial radial envelopes to ensure smoothly differentiable predictions, as proposed by Gasteiger et al., 2020. We split the basis into three parts to incorporate all available geometric information. Before the envelope, these are:

\begin{align}
\tilde{\ve}_{\text{RBF}, n}(x_{db}) &= \sqrt{\frac{2}{c_{\text{emb}}}} \frac{\sin(\frac{n \pi}{c_{\text{emb}}} x_{db})}{x_{db}},\\
\tilde{\ve}_{\text{CBF}, ln}(x_{ba}, \varphi_{abd}) &= \sqrt{\frac{2}{c_{\text{int}}^3 j_{l+1}^2(z_{ln})}} j_l(\frac{z_{ln}}{c_{\text{int}}} x_{ba}) Y_{l0}(\varphi_{abd}),\\
\tilde{\ve}_{\text{SBF}, lmn}(x_{ca}, \varphi_{cab}, \theta_{cabd}) &= \sqrt{\frac{2}{c_{\text{emb}}^3 j_{l+1}^2(z_{ln})}} j_l(\frac{z_{ln}}{c_{\text{emb}}} x_{ca}) Y_{lm}(\varphi_{cab}, \theta_{cabd}),
\end{align}

with the interaction cutoff $c_{\text{int}}$, the spherical Bessel functions $j_l$, and the  root of the  Bessel function $z_{ln}$. Note that Gasteiger et al., 2020 only used the first two parts $\ve_{\text{RBF}}$ and $\ve_{\text{CBF}}$. These representations are then transformed using two linear layers to obtain the filter $F$. In order to maintain a smoothly differentiable cutoff we cannot use a bias in this transformation. Altogether, the core geometric message passing scheme is

\begin{equation}
\begin{aligned}
\tilde{\vm}_{ca} = \sum_{\substack{b \in \gN_a^{\text{int}} \setminus \{c\}, \\ d \in \gN_b^{\text{emb}} \setminus \{a,c\}}} \Big( &(\mW_{\text{SBF}1} \ve_{\text{SBF}}(x_{ca}, \varphi_{cab}, \theta_{cabd}))^T \tW ((\mW_{\text{CBF}2} \mW_{\text{CBF}1} \ve_{\text{CBF}}(x_{ba}, \varphi_{abd}))\\
&\odot (\mW_{\text{RBF}2} \mW_{\text{RBF}1} \ve_{\text{RBF}}(x_{db})) \odot \vm_{db} )\Big),
\end{aligned}
\label{eq:core_geom}
\end{equation}

where $\mW$ denotes a weight matrix, $\tW$ denotes a weight tensor. The first weight matrix of each representation part has a small output dimension. This causes a bottleneck that improves generalization.

Symmetric message passing. Whenever we have a directional embedding $\vm_{ca}$, we also have the opposing embedding $\vm_{ac}$, since both are based on the same cutoff $c_{\text{emb}}$. Whether we associate the embedding $\vm_{ca}$ or $\vm_{ac}$ with atom $a$ is arbitrary. A more principled approach is to jointly interpret both embeddings as a representation of the atom pair $a$ and $c$. In this view, an update to $\vm_{ca}$ should also influence $\vm_{ac}$. This would normally require executing the above message passing scheme twice, once for updating $\vm_{ca}$ based on $\vm_{db}$, and once for updating $\vm_{ac}$ based on $\vm_{db}$. We propose to circumvent this double execution by calculating the update (<ref>) only once and then using it for both $\vm_{ca}$ and $\vm_{ac}$. To preserve the distinction between the two directions and ensure that $\vm_{ca} \neq \vm_{ac}$, we transform the two updates using two separate learnable weight matrices. One single message passing update thus carries information for both embeddings, which is then dissected by the two weight matrices. In practice, this only requires a simple re-indexing operation that maps the edge $ca$ to $ac$.

Efficient bilinear layer. The whole message passing scheme, i.e. basis transformation, neighbor aggregation, and bilinear layer, only use linear functions. We can therefore freely optimize the order of summation without changing the result, as proposed by Wu et al., 2019 (see <ref> for details). Doing so can provide a faster and more memory-efficient model, reducing memory usage by 50 even for Hadamard products. Moreover, since everything is based on efficient matrix products, this allows us to use the bilinear layer at practically no additional cost compared to a Hadamard product. Note that this requires using padded matrices instead of the usual gather-scatter operations to prevent excessively large intermediate results.

§ GEMNET: GEOMETRIC MESSAGE PASSING NEURAL NETWORK

The GemNet architecture (comprehensive version in <ref>). $\square$ denotes the layer's input, $\|$ concatenation, and $\sigma$ a non-linearity. Directional embeddings $\vm_{ca}$ are updated using three forms of interaction: Two-hop geometric message passing (Q-MP), one-hop geometric message passing (T-MP), and atom self-interactions. Differences between Q-MP and T-MP are denoted by colors and dashed lines.

GemNet. The geometric message passing neural network (GemNet) is a significantly refined architecture based on DimeNet$^{++}$ <cit.>. GemNet predicts the molecular energy $E$ and forces $\mF \in \R^{3 \times n}$ based on the atomic positions $\mX \in \R^{3 \times n}$ and the atomic numbers $\vz \in \sN^n$. The architecture is illustrated in <ref>. A comprehensive version with low-level layers and hyperparameters is described in <ref>. GemNet was developed on the COLL dataset, but generalizes to other datasets such as MD17 without architectural changes. Every change we propose either improves model performance or reduces model complexity. For example, GemNet uses no biases since we found them to be irrelevant or even detrimental to accuracy. We show the impact of the most relevant changes via ablation studies in <ref>. =-1

Interactions. GemNet incorporates three forms of interactions. The first is geometric message passing, as described in <ref>. The second is a one-hop form of geometric message passing. This interaction uses a single cutoff $c = c_{\text{emb}}$ and passes messages only between directional embeddings pointing towards the same atom, similarly to DimeNet [Gasteiger et al., 2020]. This provides both angle-based pair interactions and atom self-interactions, thanks to the symmetric message passing scheme described in <ref>. The third interaction is a pure atom self-interaction based on atom embeddings. We first aggregate the directional embeddings of one atom to obtain an atom embedding. We then use this atom embedding to update all directional embeddings. We found all three interaction forms to be beneficial, and show this in our ablation studies.

Stabilizing activation variance. The variance of activations in a model is usually stabilized using normalization methods, which has various positive effects on training [De & Smith, 2020, Santurkar et al., 2018, Luo et al., 2019]. However, they also have multiple undesirable side effects, especially in the context of molecular regression. Batch normalization introduces correlations between separate molecules and atoms. Layer normalization forces all activation scales to be constant, while atomic interactions actually cover a large range of scales – directly bonding atoms have a substantially stronger interaction than atoms at a long range. To circumvent these issues, we stabilize GemNet's variance by introducing constant scaling factors, as suggested by Brock et al., 2021. We found that the activation variance is primarily impacted by four components: Skip connections, non-linearities, message aggregation, and Hadamard/bilinear layers. The two summands in a skip connection $y = x + f(x)$ have no covariance at initialization due to random weight matrices. We can thus remove its impact by scaling the output by $1/\sqrt{2}$. We remove the non-linearity's impact by scaling its output with a gain of $\gamma = 1/0.6$ for SiLU, similarly to [Klambauer et al., 2017]. Note that we do not center SiLU's output but instead choose a slightly lower $\gamma$ to account for mean shift. Additionally, we standardize the weight matrices to have exactly zero mean and $1/\text{fan-in}$ variance. The sum aggregation and Hadamard/bilinear layers have a more complex impact on the variance, which we cannot determine a priori (see <ref> for details). We therefore estimate the variance after these layers based on random batches of data. We then rescale their output accordingly to obtain roughly the variance of the layer input at initialization. These simple empirical scaling factors are sufficient to keep the activation variance roughly constant (see <ref>). We found that other measures such as adaptive gradient clipping [Brock et al., 2021], scaled weight standardization [Brock et al., 2021], or weighting the residual block with zero at initialization [De & Smith, 2020] are not beneficial for model accuracy.

GemNet-Q and GemNet-T. Geometric message passing is comparatively expensive since it is based on quadruplets of atoms. Its runtime thus scales with $\gO(nk_{\text{int}}k_{\text{emb}}^2)$, where $k_{\text{int}}$ is the number of interacting neighbors, and $k_{\text{emb}}$ is the number of embedded directions. For this reason we investigate two message passing models in our experiments – one with two-hop geometric message passing (GemNet-Q) and one using only the two cheaper forms of interaction (GemNet-T). Their complexities are $\gO(nk_{\text{int}}k_{\text{emb}}^2)$ and $\gO(nk_{\text{emb}}^2)$, respectively. Note that GemNet-T is thus a direct ablation of the two-hop message passing scheme implied by our theoretical results.

Direct force predictions. GemNet predicts forces by calculating $\mF_a = - \partial E / \partial \vx_a$ via backpropagation. This form of calculation guarantees a conservative force field, which is important for the stability of simulations. However, by using <ref> we can also directly predict forces and other vector quantities. This essentially means predicting a magnitude for each directional embedding and then summing up over the vectors defined by this magnitude and the embedding's associated direction, similarly to Park et al., 2021. We denote this variant as GemNet-dQ and GemNet-dT. Interestingly, GemNet is thus able to generate rotationally equivariant predictions despite only using invariant representations. Direct predictions substantially accelerate the model, especially for training. For most datasets, the resulting accuracy is on par with most previous models, but significantly worse than GemNet's accuracy via backpropagation. However, this is not true for OC20, where we found GemNet-dT to converge faster and perform on par with GemNet-T.

Limitations. GemNet is focused on one specific, important task: Predictions for molecular simulations. We do not make any statements regarding its performance beyond this scope. The GemNet architecture might seem more complex than some previous models, due to its larger variety of interactions and blocks. However, its number of parameters and training or inference time is actually on par with previous models. Two-hop message passing introduces significant computational overhead. We mitigate this effect with a down-projection layer and additionally introduce the ablated GemNet-T model. This model performs surprisingly well on MD17, but not on COLL. This suggests that one-hop message passing is expressive enough for some practical use cases, but two-hop message passing gives an advantage for the more challenging task of fitting multiple molecules at once.

Societal impacts. Accelerating molecular simulations can have positive effects in a wide range of applications in physics and chemistry. At the same time, however, this can be used for malicious purposes such as developing chemical agents or weapons. To the best of our knowledge, this work does not promote these use cases more than regular chemistry research does. To somewhat mitigate negative effects we will publish our source code under the Hippocratic license [Ehmke, 2020].

§ EXPERIMENTS

MAE on COLL, in [per-mode=symbol] and . GemNet is 34 more accurate for forces. The higher energy error is due to its lower loss weight.

Force MAE for MD17@CCSD in [per-mode=symbol]. GemNet outperforms previous methods by 44 on average.

Force MAE for MD17 in [per-mode=symbol]. GemNet outperforms all previous methods by a wide margin, on average by 41.

Results for the three tasks of the open catalyst dataset (OC20), averaged across its four test sets. GemNet outperforms all previous methods in all measures, on average by 20.

*DimeNet$^{++}$-large uses separate models for energy and force prediction for IS2RE.

Experimental setup. We evaluate our model on four molecular dynamics datasets. COLL <cit.> consists of configurations taken from molecular collisions of different small organic molecules. MD17 [Chmiela et al., 2017] consists of configurations of multiple separate, thermalized molecules, considering only one molecule at a time. MD17@CCSD [Chmiela et al., 2018] uses the same setup, but calculates the forces using the more accurate and expensive CCSD or CCSD(T) method. The open catalyst (OC20)  dataset <cit.> consists of energy relaxation trajectories of solid catalysts with adsorbate molecules. This dataset is split into three tasks: (1) Structure to energy and forces (S2EF), which is the same task as used by the COLL and MD17 datasets, (2) initial structure to relaxed structure (IS2RS), where an energy optimization is carried out based on the model's predictions and we measure how close the final structure is to the true relaxed structure (average distance within threshold, ADwT) and whether the final forces are close to zero (average forces below threshold, AFbT), and (3) initial structure to relaxed energy (IS2RE), where we predict the energy of the relaxed structure, based on an energy optimization starting at the initial structure. All presented OC20 models are trained on the S2EF data.
Following the setup of Batzner et al., 2021, we use 1000 training and validation configurations for MD17, and 950 training and 50 validation configurations for MD17@CCSD. We focus on force predictions and use a high force loss weight since they determine the accuracy of molecular simulations. We measure the mean absolute error (MAE), averaged over all samples, atoms, and components. We compare with the results reported by several state-of-the-art models: sGDML [Chmiela et al., 2018], FCHL19 [Christensen et al., 2020], SchNet [Schütt et al., 2017], DimeNet [Gasteiger et al., 2020], DimeNet$^{++}$ [Gasteiger et al., 2020], SphereNet [Liu et al., 2021], NequIP [Batzner et al., 2021], PaiNN [Schütt et al., 2021], ForceNet [Hu et al., 2021], and SpinConv [Shuaibi et al., 2021]. For further details see <ref>.

Results. <ref> show that GemNet-T and GemNet-Q consistently perform best on all molecular dynamics datasets investigated – and by a large margin. This is true both in comparison to previous GNNs and for kernel methods – despite the latter typically being more sample efficient. The improvements are largest for chain-like molecules, such as ethanol and malonaldehyde. These molecules are the most challenging since they exhibit a wide range of movement.
GemNet even performs better than some previous models that were trained with 50x more training samples. For example, it performs better than SchNet with 50000 training samples on six out of eight MD17 molecules (see <ref>).
Interestingly, the two-hop message passing scheme implied by our theoretical results (GemNet-Q) yields significant improvements on COLL, but performs approximately on par with the ablated GemNet-T on MD17. To investigate this disagreement we trained GemNet on a combined dataset of all MD17 molecules. <ref> shows that GemNet-Q again performs better than GemNet-T in this setting. These results suggest that regular MD17 is too simple to show the benefits of two-hop message passing. It seems to be particularly important in more difficult settings that cover a large variety of configurations and molecules.

Computational aspects. GemNet-Q is roughly two times slower than GemNet-T (see <ref>). Thanks to the efficient aggregation, GemNet with bilinear layers is as fast as with regular Hadamard products. Efficient aggregation also reduces the memory usage for regular Hadamard products by around 50 (from 4.1GB to 2.2GB for a batch of 32 Toluene molecules). Note that GemNet has not been optimized for runtime and can likely be accelerated substantially. GemNet-Q uses 2.2M and GemNet-T 1.9M parameters, which is comparable to previous models such as DimeNet$^{++}$, which uses 1.9M parameters. See <ref> for further details.

Direct force prediction. Directly predicting the forces accelerates training by four times on average and inference by 1.6 times on average in our experiments (see <ref>), while reducing memory consumption by roughly a factor of two. While using direct predictions instead of backpropagation increases the MAE by 44 on COLL and by 48 on MD17 (see <ref>), they actually perform better on the S2EF task on OC20. This is likely due to OC20 being orders of magnitude larger than COLL and MD17. Whether to use direct predictions thus depends on the dataset and the application's computational requirements.

Ablation studies on COLL. Force MAE in [per-mode=symbol] after 500000 training steps. All proposed components yield significant improvements.

Ablation studies. We investigate the proposed architectural improvements on COLL in <ref>. The proposed symmetric message passing scheme yields significant accuracy improvements, as does using a bilinear layers instead of a Hadamard product. We also see that removing any of the three interaction forms described in <ref> increases the error, showing that this combination is indeed beneficial. The proposed scaling factors also yield decent improvements, while regular layer normalization actually increases the error. Two-hop message passing yields the largest single improvement. <ref> shows that our architectural improvements yield similar benefits for DimeNet$^{++}$. Overall, the error improvements are quite evenly distributed. This suggests that GemNet's improved performance is not due to one single change, but rather due to the full range of improvements proposed in this work.

§ CONCLUSION

In this work we proved the universality for GNNs using directional embeddings. We proposed geometric message passing based on these insights, and improved this scheme with symmetric message passing and efficient bilinear layers. We incorporated these improvements in the GemNet architecture, which substantially improves the error on various molecular dynamics datasets. We showed that all of the proposed enhancements yield significant performance improvements. Most of our proposed improvements are of independent interest for other molecular GNNs.

We would like to thank Soledad Villar for help with proving <ref>, as well as Nicholas Gao and Aleksandar Bojchevski for their invaluable feedback and comments.

This research was supported by the Deutsche Forschungsgemeinschaft (DFG) through the Emmy Noether grant GU 1409/2-1 and the TUM International Graduate School of Science and Engineering (IGSSE), GSC 81.

[Anderson et al., 2019]
Brandon M. Anderson, Truong-Son Hy, and Risi Kondor.
Cormorant: Covariant Molecular Neural Networks.
In NeurIPS, 2019.

[Azizian & Lelarge, 2020]
Waiss Azizian and Marc Lelarge.
Expressive Power of Invariant and Equivariant Graph Neural
In ICLR, 2020.

[Bartók et al., 2010]
Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi.
Gaussian Approximation Potentials: The Accuracy of Quantum
Mechanics, without the Electrons.
Physical Review Letters, 1040 (13):0 136403,

[Batzner et al., 2021]
Simon Batzner, Tess E. Smidt, Lixin Sun, Jonathan P. Mailoa, Mordechai
Kornbluth, Nicola Molinari, and Boris Kozinsky.
SE(3)-Equivariant Graph Neural Networks for
Data-Efficient and Accurate Interatomic Potentials.
arXiv, 2101.03164, 2021.

[Bogatskiy et al., 2020]
Alexander Bogatskiy, Brandon Anderson, Jan Offermann, Marwah Roussi, David
Miller, and Risi Kondor.
Lorentz Group Equivariant Neural Network for Particle
In ICML, 2020.

[Brock et al., 2021]
Andrew Brock, Soham De, and Samuel L. Smith.
Characterizing signal propagation to close the performance gap in
unnormalized ResNets.
In ICLR, 2021.

[Brock et al., 2021]
Andrew Brock, Soham De, Samuel L. Smith, and Karen Simonyan.
High-Performance Large-Scale Image Recognition Without
arXiv, 2102.06171, 2021b.

[Chanussot et al., 2021]
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed
Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua
Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh,
C. Lawrence Zitnick, and Zachary Ulissi.
Open Catalyst 2020 (OC20) Dataset and Community Challenges.
ACS Catalysis, 110 (10):0 6059–6072, 2021.

[Chmiela et al., 2017]
Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky,
Kristof T. Schütt, and Klaus-Robert Müller.
Machine learning of accurate energy-conserving molecular force
Science Advances, 30 (5):0 e1603015, 2017.

[Chmiela et al., 2018]
Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre
Towards exact molecular dynamics simulations with machine-learned
force fields.
Nature Communications, 90 (1):0 1–10, 2018.

[Christensen & Lilienfeld, 2020]
Anders S. Christensen and O. Anatole von Lilienfeld.
On the role of gradients for machine learning of molecular energies
and forces.
Machine Learning: Science and Technology, 10
(4):0 045018, 2020.

[Christensen et al., 2020]
Anders S. Christensen, Lars A. Bratholm, Felix A. Faber, and O. Anatole von
FCHL revisited: Faster and more accurate quantum machine
The Journal of Chemical Physics, 1520 (4):0
044107, 2020.

[Cohen & Welling, 2016]
Taco Cohen and Max Welling.
Group Equivariant Convolutional Networks.
In ICML, 2016.

[Cohen et al., 2018]
Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling.
Spherical CNNs.
In ICLR, 2018.

[Cohen et al., 2019]
Taco S Cohen, Mario Geiger, and Maurice Weiler.
A General Theory of Equivariant CNNs on Homogeneous
In NeurIPS, 2019.

[De & Smith, 2020]
Soham De and Sam Smith.
Batch Normalization Biases Residual Blocks Towards the
Identity Function in Deep Networks.
In NeurIPS, 2020.

[Driscoll & Healy, 1994]
J. R. Driscoll and D. M. Healy.
Computing Fourier Transforms and Convolutions on the
Advances in Applied Mathematics, 150 (2):0
202–250, 1994.

[Dym & Maron, 2021]
Nadav Dym and Haggai Maron.
On the Universality of Rotation Equivariant Point Cloud
In ICLR, 2021.

[Ehmke, 2020]
Coraline Ada Ehmke.
The Hippocratic License 2.1: An Ethical License for Open
Source., 2020.
URL <https://firstdonoharm.dev>.

[Elfwing et al., 2018]
Stefan Elfwing, Eiji Uchibe, and Kenji Doya.
Sigmoid-weighted linear units for neural network function
approximation in reinforcement learning.
Neural Networks, 107:0 3–11, 2018.

[Esteves et al., 2018]
Carlos Esteves, Christine Allen-Blanchette, Ameesh Makadia, and Kostas
Learning SO(3) Equivariant Representations with Spherical
In ECCV, 2018.

[Esteves et al., 2020]
Carlos Esteves, Ameesh Makadia, and Kostas Daniilidis.
Spin-Weighted Spherical CNNs.
In NeurIPS, 2020.

[Faber et al., 2018]
Felix A. Faber, Anders S. Christensen, Bing Huang, and O. Anatole von
Alchemical and structural distribution based representation for
universal quantum machine learning.
The Journal of Chemical Physics, 1480 (24):0
241717, 2018.

[Finzi et al., 2020]
Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson.
Generalizing Convolutional Neural Networks for Equivariance
to Lie Groups on Arbitrary Continuous Data.
In ICML, 2020.

[Flam-Shepherd et al., 2021]
Daniel Flam-Shepherd, Tony C. Wu, Pascal Friederich, and Alan Aspuru-Guzik.
Neural Message Passing on High Order Paths.
Machine Learning: Science and Technology, 2021.

[Fuchs et al., 2020]
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling.
SE(3)-Transformers: 3D Roto-Translation Equivariant
Attention Networks.
In NeurIPS, 2020.

[Garg et al., 2020]
Vikas Garg, Stefanie Jegelka, and Tommi Jaakkola.
Generalization and Representational Limits of Graph Neural
In ICML, 2020.

[Gasteiger et al., 2020]
Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann.
Fast and Uncertainty-Aware Directional Message Passing for
Non-Equilibrium Molecules.
In Machine Learning for Molecules Workshop, NeurIPS,

[Gasteiger et al., 2020]
Johannes Gasteiger, Janek Groß, and Stephan Günnemann.
Directional Message Passing for Molecular Graphs.
In ICLR, 2020.

[Gilmer et al., 2017]
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and
George E. Dahl.
Neural Message Passing for Quantum Chemistry.
In ICML, 2017.

[Halgren, 1996]
Thomas A. Halgren.
Merck molecular force field. I. Basis, form, scope,
parameterization, and performance of MMFF94.
Journal of Computational Chemistry, 170
(5-6):0 490–519, 1996.

[Hu et al., 2021]
Weihua Hu, Muhammed Shuaibi, Abhishek Das, Siddharth Goyal, Anuroop Sriram,
Jure Leskovec, Devi Parikh, and C. Lawrence Zitnick.
ForceNet: A Graph Neural Network for Large-Scale
Quantum Calculations.
arXiv, 2103.01436, 2021.

[Jetter et al., 1999]
Kurt Jetter, Joachim Stöckler, and Joseph Ward.
Error estimates for scattered data interpolation on spheres.
Mathematics of Computation, 680 (226):0
733–747, 1999.

[Keiner et al., 2007]
Jens Keiner, Stefan Kunis, and Daniel Potts.
Efficient Reconstruction of Functions on the Sphere from
Scattered Data.
Journal of Fourier Analysis and Applications, 130
(4):0 435–458, 2007.

[Klambauer et al., 2017]
Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter.
Self-normalizing neural networks.
In NeurIPS, 2017.

[Kondor et al., 2018]
Risi Kondor, Zhen Lin, and Shubhendu Trivedi.
Clebsch-Gordan Nets: a Fully Fourier Space Spherical
Convolutional Neural Network.
In NeurIPS, 2018.

[Kondor et al., 2019]
Risi Kondor, Truong Son Hy, Horace Pan, Brandon M. Anderson, and Shubhendu
Covariant Compositional Networks For Learning Graphs.
In International Workshop on Mining and Learning with
Graphs (MLG), KDD, 2019.

[Kostelec & Rockmore, 2008]
Peter J. Kostelec and Daniel N. Rockmore.
FFTs on the Rotation Group.
Journal of Fourier Analysis and Applications, 140
(2):0 145–179, 2008.

[Liu et al., 2021]
Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji.
Spherical Message Passing for 3D Graph Networks.
arXiv, 2102.05013, 2021.

[Loshchilov & Hutter, 2018]
Ilya Loshchilov and Frank Hutter.
Decoupled Weight Decay Regularization.
In ICLR, 2018.

[Luo et al., 2019]
Ping Luo, Xinjiang Wang, Wenqi Shao, and Zhanglin Peng.
Towards Understanding Regularization in Batch Normalization.
In ICLR, 2019.

[Maron et al., 2019]
Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman.
Provably Powerful Graph Networks.
In NeurIPS, 2019.

[Maron et al., 2019]
Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman.
On the Universality of Invariant Networks.
In ICML, 2019.

[Maron et al., 2020]
Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya.
On Learning Sets of Symmetric Elements.
In ICML, 2020.

[Morris et al., 2019]
Christopher Morris, Martin Ritzert, Matthias Fey, William L. Hamilton, Jan Eric
Lenssen, Gaurav Rattan, and Martin Grohe.
Weisfeiler and Leman Go Neural: Higher-Order Graph
Neural Networks.
In AAAI, 2019.

[Morris et al., 2020]
Christopher Morris, Gaurav Rattan, and Petra Mutzel.
Weisfeiler and Leman go sparse: Towards scalable higher-order
graph embeddings.
In NeurIPS, 2020.

[Park et al., 2021]
Cheol Woo Park, Mordechai Kornbluth, Jonathan Vandermause, Chris Wolverton,
Boris Kozinsky, and Jonathan P. Mailoa.
Accurate and scalable graph neural network force field and molecular
dynamics with direct force architecture.
npj Computational Materials, 70 (1):0 1–9,

[Qiao et al., 2020]
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and
Thomas F. Miller.
OrbNet: Deep learning for quantum chemistry using
symmetry-adapted atomic-orbital features.
The Journal of Chemical Physics, 1530 (12):0
124111, 2020.

[Qiao et al., 2021]
Zhuoran Qiao, Anders S. Christensen, Matthew Welborn, Frederick R. Manby, Anima
Anandkumar, and Thomas F. Miller III.
UNiTE: Unitary N-body Tensor Equivariant Network with
Applications to Quantum Chemistry.
arXiv, 2105.14655, 2021.

[Reddi et al., 2018]
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar.
On the Convergence of Adam and Beyond.
In ICLR, 2018.

[Sakurai & Tuan, 1993]
J. J. Sakurai and San Fu Tuan.
Modern Quantum Mechanics.
Revised, subsequent edition edition, 1993.

[Santurkar et al., 2018]
Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Mądry.
How does batch normalization help optimization?
In NeurIPS, 2018.

[Schütt et al., 2017]
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan
Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller.
SchNet: A continuous-filter convolutional neural network for
modeling quantum interactions.
In NeurIPS, 2017.

[Schütt et al., 2021]
Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger.
Equivariant message passing for the prediction of tensorial
properties and molecular spectra.
In ICML, 2021.

[Shuaibi et al., 2021]
Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram,
Zachary Ulissi, and C. Lawrence Zitnick.
Rotation Invariant Graph Neural Networks using Spin
arXiv, 2106.09575, 2021.

[Thomas et al., 2018]
Nathaniel Thomas, Tess Smidt, Steven M. Kearnes, Lusann Yang, Li Li, Kai
Kohlhoff, and Patrick Riley.
Tensor Field Networks: Rotation- and
Translation-Equivariant Neural Networks for 3D Point Clouds.
arXiv, 1802.08219, 2018.

[Villar et al., 2021]
Soledad Villar, David W. Hogg, Kate Storey-Fisher, Weichi Yao, and Ben
Scalars are universal: Equivariant machine learning, structured
like classical physics.
In NeurIPS, 2021.

[Weiler et al., 2018]
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen.
3D Steerable CNNs: Learning Rotationally Equivariant
Features in Volumetric Data.
In NeurIPS, 2018.

[Wu et al., 2019]
Wenxuan Wu, Zhongang Qi, and Li Fuxin.
PointConv: Deep Convolutional Networks on 3D Point
In CVPR, 2019.

[Xu et al., 2019]
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka.
How Powerful are Graph Neural Networks?
In ICLR, 2019.

[Yarotsky, 2021]
Dmitry Yarotsky.
Universal Approximations of Invariant Maps by Neural
Constructive Approximation, 2021.


# The Impact of Data Corruption on Named Entity Recognition for Low-resourced
Languages

Manuel Fokam Michael Beukman
School of Computer Science and Applied Mathematics,
University of the Witwatersrand, Johannesburg, South Africa
<EMAIL_ADDRESS>
###### Abstract

Data availability and quality are major challenges in natural language
processing for low-resourced languages. In particular, there is significantly
less data available than for higher-resourced languages. This data is also
often of low quality, rife with errors, invalid text or incorrect annotations.
Many prior works focus on dealing with these problems, either by generating
synthetic data, or filtering out low-quality parts of datasets. We instead
investigate these factors more deeply, by systematically measuring the effect
of data quantity and quality on the performance of pre-trained language models
in a low-resourced setting. Our results show that having fewer completely-
labelled sentences is significantly better than having more sentences with
missing labels; and that models can perform remarkably well with only 10% of
the training data. Importantly, these results are consistent across ten low-
resource languages, English, and four pre-trained models.

##  1 Introduction

Natural Language Processing (NLP) is a rapidly growing field that has been
applied to a wide range of tasks and domains (Vaswani et al., 2017; Conneau et
al., 2020). However, much of the focus in NLP has been on high-resource
languages such as English (Vaswani et al., 2017; Radford et al., 2018, 2019).
While this has led to notable advancements for these languages, low-resourced
languages have not received as much attention, resulting in a significant
performance gap between high- and low-resourced languages. This has prompted
an increasing number of studies focused exclusively on low-resourced
languages, resulting in the development of models (Ogueji et al., 2021; Alabi
et al., 2022) and the introduction of datasets (Adelani et al., 2021, 2022a,
2022b).

Despite this impressive progress, data remains a limiting factor for low-
resourced NLP (Adelani et al., 2022a, b). In particular, the two main problems
are the availability and quality of data. First, the datasets available for
low-resourced languages tend to be smaller than those for high-resourced
languages, and for many languages, no data exists at all (Martinus and Abbott,
2019; Adelani et al., 2021). Secondly, the available datasets are often of
questionable quality, containing invalid text or incorrect annotations
(Kreutzer et al., 2022), which has detrimental effects on the models trained
on these datasets (Abdul-Rauf et al., 2012; Alabi et al., 2019).

This means that many existing datasets in low-resourced NLP are either small
or of low quality. This observation has led to research that investigates the
tradeoff between the amount and quality of data (Gascó et al., 2012; Alabi et
al., 2019; de Gibert Bonet et al., 2022). This line of work has provided
valuable insights that allow NLP practitioners to make informed decisions when
faced with a choice of which dataset should be used to train a model. However,
many of these works focus on comparing different datasets, often from
different domains, without clearly quantifiable tradeoffs between data
quantity and quality (Alabi et al., 2019). While these approaches can be
useful, a more comprehensive and precise approach is needed to fully
understand the tradeoff between data quantity and quality in low-resourced
NLP.

To address this, we take a different perspective and focus on systematically
and quantifiably reducing the quality of datasets and examining the effects of
this on the performance of NLP models. Additionally, by altering the amount of
data used to train our models, we can compare the tradeoffs between quality
and quantity. We do this by devising various controllable corruption
strategies, and training models on different levels of corrupted data. Our
focus is on a named-entity recognition task due to its prevalence in many NLP
systems and the availability of a few high-quality datasets in low-resourced
languages. We fine-tune existing pre-trained language models, as this is a
common and high-performing approach, especially for low-resourced languages
(Ogueji et al., 2021; Adelani et al., 2021; Alabi et al., 2022).

We provide systematic evidence to support prior findings that the quality of
data, in general, is strongly preferred over quantity. Furthermore, our
findings are consistent across eleven different languages and four pre-trained
models, suggesting that our conclusions hold true in a general sense.

##  2 Background and Related Work

###  2.1 Named Entity Recognition

Named entity recognition (NER) is a token classification task, where the goal
is to classify each token or word in a text as an Organisation, Location,
Person, Date, or indicate that the token does not correspond to a named entity
by giving it the label of “Other”. NER as a field has many impactful
applications in NLP pipelines and use-cases(Sang and Meulder, 2003; Lample and
Chaplot, 2017; Adelani et al., 2021). A typical NER dataset consists of
multiple sentences, with each sentence containing both the words and their
associated labels.

The prevailing approach to train NER models is to use a pre-trained large
language model (such as BERT (Devlin et al., 2019), XLM-Roberta (Conneau et
al., 2020), etc.) and fine-tune it on a small amount of NER data (Conneau et
al., 2020; Adelani et al., 2021). These models were pre-trained on a large
corpus of unlabelled text, and resulted in improved downstream performance
after fine-tuning compared to training on NER data from scratch. The overall
classification F1 score, calculated as the harmonic mean of precision and
recall, is generally used as the main metric of performance in NER (Sang and
Meulder, 2003; Adelani et al., 2021).

###  2.2 Data Collection and Annotation

Since the lack of data has traditionally been a major limiting factor for low-
resourced NLP research, multiple different approaches have developed to
effectively collect data in resource-constrained settings. In particular,
community involvement has played a large part in this (Nekoto et al., 2020,
2022), where native speakers annotate or create datasets to be used in
research. This has led to the creation of many different datasets (Adelani et
al., 2021; Nekoto et al., 2022), but it relies on community members instead of
trained annotators, which may result in some aspects of the annotation being
less accurate. Furthermore, while this approach can successfully develop
datasets for low-resourced languages, due to logistic challenges and a limited
amount of unlabelled text, these datasets are often significantly smaller than
high-resourced datasets (Conneau et al., 2020; Adelani et al., 2021).

###  2.3 Analysis of Quality vs Quantity in Low-resourced Languages

While there has been significant progress in recent years, datasets for low-
resourced language are often quite small and limited, or exhibit low quality.
Both of these factors can lead to poorly-performing models. For instance,
Kreutzer et al. (2022) perform a large-scale audit of several web-scale and
automatically extracted multilingual datasets, and find that the quality is
often poor, with non-linguistic or otherwise invalid text being commonplace.

This lack of quality can have great effects on the performance of models.
Alabi et al. (2019) show that for certain low-resourced African languages,
using a significantly smaller, but curated dataset outperforms training a
model on a large, but noisy dataset. Abdulmumin et al. (2022) find similar
results, where training on filtered data of higher quality improved the
performance of translation models for low-resourced languages. Many of these
works consider one or two completely different datasets, and compare the
relative quality and quantity. This, however, lacks a systematic approach that
controls for other factors such as the domain of the data. In addition, the
filtering-based approaches often use a learned model as a filter and select
only sentences that have a predicted quality value above a certain threshold
(Abdulmumin et al., 2022; de Gibert Bonet et al., 2022). While this does
provide a quantifiable level of quality, it may not be comparable across
datasets or different filtering models. Additionally, datasets may exhibit
different levels of certain problems, e.g. some datasets may have many tokens
corresponding to punctuation whereas others may have sentences in a different
language to the rest of the data. These problems may make it hard to
accurately compare the results of these studies and use their conclusions in
practice.

##  3 Methodology

Our aim is to analyse and quantify the impact of data corruption on the
performance of pre-trained language models. This understanding would enable us
to make more informed decisions about the relative importance of data quality
and quantity, ultimately leading to improved data creation processes and
selection of NLP training data for practitioners.

While we can corrupt NER datasets in various ways, we choose corruptions that
simulate a mislabelling scenario during the annotation process, e.g.
mislabelling a person in a sentence as an organisation. There are two main
reasons for this choice. Firstly, many NER datasets are formed by taking an
existing text source, which is usually of high quality, such as news data
(Adelani et al., 2021) and annotating each word; thus, errors are more likely
to appear during the annotation process. Secondly, it is challenging to
corrupt the base sentences in a reasonable, quantifiable and incremental way,
as sentences encompass meaning which is often hard to change atomically.

Thus, we focus on corrupting only the labels, using different strategies
detailed in Section 3.1. For each corruption strategy, we uniformly vary the
amount of corruption and train our models using the new, corrupted dataset.
This process allows us to evaluate how each corruption strategy affects the
model as we adjust the degree of corruption. As an additional experiment, we
vary the size of the data available to the model by using only a subset of the
sentences without corrupting any labels, allowing us to determine the effect
of varying the amount of data on performance.

We discuss the data used in this study in Section 3.2.

Figure 1: An illustration of the different corruption strategies we use. When
(Left) capping labels, we effectively remove a certain fraction of labels,
replacing them with O. When (Right) swapping labels, we instead randomly
replace a label with an incorrect one. This figure is just illustrative, and
in our experiments we have a percentage threshold; for instance, corrupting
30% of the labels across the entire dataset.

###  3.1 Different Corruption Strategies

This section contains descriptions of the corruption strategies that we use,
with Figure 1 providing a visual illustration of our quality-related
corruption strategies. We only change the training data while leaving the
evaluation data unchanged to obtain an objective comparison of different
corruption strategies.

####  3.1.1 Sentence Capping

Dataset annotation is generally expensive and logistically challenging when
multiple participants are involved. As a result, low-resourced NLP datasets
are often not particularly large (Adelani et al., 2021). Due to this
observation, we first evaluate the effect of varying the size of the data
available to our NER models. In this strategy, we randomly remove sentences
from the original dataset to create sub-datasets with fewer sentences than the
original dataset. This process allows us to measure the model performance as a
function of data quantity. We choose to represent quantity as a function of
the number of sentences because removing words can alter the meaning of a
sentence in ways we cannot control.

####  3.1.2 Entity Label Capping

A rich NER dataset would be a dataset with a high annotation density, i.e., a
high number of annotated entities per sentence. This strategy aims at
inhibiting the model by thresholding the number of entity annotations allowed
in the dataset. In the real world, this would be equivalent to a situation
where an annotator failed to label a particular token or span of tokens as one
of the entities PER, LOC, ORG, DATE, instead giving it the default entity type
O, which generally means not relevant.

Here we globally corrupt the data by choosing a certain percentage of labels
to keep across the entire dataset. For example, $50\%$ would mean that we
randomly remove half of all entity labels (replacing them with O), which may
leave some sentences unmodified and others entirely without annotations.

For this and the following corruption strategy, we consider the atomic element
to be a single named entity, even if this consists of multiple words. As a
result, we change the entire span of an entity label instead of just a part
thereof.

####  3.1.3 Entity Label Swapping

Another scenario that could happen during the annotation procedure would be
the mislabelling of a span of tokens with the wrong entity. For example, A
person named Christian Dior mistakenly labelled as an organisation due to some
bias in the knowledge of the annotator. These mistakes may create datasets
with contradictory labels, with the same tokens being used in very similar
contexts but labelled differently. Therefore, our goal behind this corruption
strategy is to determine how robust large pre-trained language models are to
such mistakes. Here we again choose a global percentage, randomly selecting
labelled entities according to this percentage and swapping their labels with
incorrect ones.

###  3.2 Data

We use the MasakhaNER dataset (Adelani et al., 2021), a high-quality NER
dataset for ten low-resourced African languages licensed under CC-BY-4.0-NC.
We specifically focus on low-resourced languages, as these languages often
suffer from the aforementioned problems. Furthermore, this dataset is of high
quality, which allows us to evaluate the full spectrum of quality, from gold-
standard to completely corrupted.

As a baseline, we also use the CONLL NER dataset, which is a staple NER
dataset in English (Sang and Meulder, 2003), with many more sentences than any
of the MasakhaNER languages. Table 1 contains information about the number of
sentences and entities for each NER corpus.

Table 1: Information about the data. Entity Density refers to the fraction of
tokens that are entities. All languages use the Latin script, except Amharic,
which uses the Fidel script.

Code Code # Sentences # Tokens # Entities Entities per % Entities Sentence  in
Tokens Amharic amh 1,750 25,829 3,995 2.3 15.5 Luganda lug 1,428 33,003 5,039
3.5 15.3 Luo luo 644 18,577 2,704 4.2 14.6 English en 14,042 203,621 29,450
2.1 14.5 Nigerian Pidgin pcm 2,124 52,604 7,392 3.5 14.1 Kinyarwanda kin 2,116
47,912 6,104 2.9 12.7 Swahili swa 2,109 56,599 7,161 3.4 12.7 Hausa hau 1,912
55,010 6,836 3.6 12.4 Igbo ibo 2,235 42,719 5,294 2.4 12.4 Yorùbá yor 2,171
56,274 6,324 2.9 11.2 Wolof wol 1,871 36,805 2,157 1.2 5.9

Table 2: Information about the different pre-trained language models we use.
In the MasakhaNER Languages column, we list only the languages the model pre-
trained on that are included in the MasakhaNER dataset.

Name Model Version Source Parameters MasakhaNER Languages AfriBERTa afriberta-
large Ogueji et al. (2021) 126M amh, hau, ibo, kin, pcm, swa, yor Afro-XLM-R
afro-xlmr-base Alabi et al. (2022) 270M amh, hau, ibo, kin, pcm, swa, yor XLM
Roberta xlm-roberta-base Conneau et al. (2020) 270M amh, hau, swa Multilingual
BERT bert-base-multilingual-cased Devlin et al. (2019) 110M swa, yor

##  4 Experiments

Having described our corruption strategies, we now perform our experiments and
showcase our results. We consider the three corruption strategies described
above, and use four different pre-trained language models, described in
Section 4.1. Each run consists of fine-tuning a single pre-trained model on a
single language’s dataset, either the original one or a corrupted version. We
fine-tune models for 25 epochs, as the results were similar to 50 epochs
(which Adelani et al. (2021) used) and trained much faster. We use a learning
rate of $5e-5$, a batch size of $64$, and a sequence length of $200$. We run
all experiments over three seeds and average the results. The compute nodes
used to run our experiments are equipped with NVIDIA RTX 3090 GPUs.

We specifically investigate the effect of progressively corrupting data on the
performance of each model, measured by the overall F1 score. This simulates
the effect of having low data quality (for instance, due to incorrect
annotations), but allows us to study this in a controlled setting. We do not
modify the test datasets at all.

Then, to normalise results across models and languages, we divide each F1
score by the value obtained when training the same model on the full,
uncorrupted dataset. This effectively measures what fraction of performance is
lost when corrupting data and allows us to transform all of the metrics to
fall between 0 and 1, resulting in the metrics being comparable across
languages and models.

In the cap sentences strategy, where we train models on a subset of data, we
specifically remove a certain percentage of the data and train the model on
the remaining sentences. Since the specific fraction we keep may have an
effect, we run this experiment three times, each time with different random
selections of data. We average over these permutations and find that the
results are very similar across them.

###  4.1 Different Pre-trained Language Models

We use four different pre-trained language models. We first consider two
models developed specifically for low-resourced African languages, AfriBERTa
and Afro-XLM-R. The other two models are traditional multilingual models, with
the majority of the training datasets consisting of high-resourced languages,
XLM-R and mBERT (multilingual BERT). AfriBERTa (Ogueji et al., 2021) was pre-
trained on less than 1GB of African language text. Afro-XLM-R (Alabi et al.,
2022) used language adaptive fine-tuning, where a pre-trained language model
is fine-tuned on unlabelled data using the same objective that was used during
pre-training. Afro-XLM-R performed this process on 20 languages, 17 of them
from Africa, starting from XLM Roberta. XLM Roberta (Conneau et al., 2020) is
a high-performing model that was pre-trained on 100 languages. Finally, mBERT
(Devlin et al., 2019) used the standard BERT training process on 104
languages, using data from Wikipedia.

We choose the specific model versions to be roughly comparable in terms of the
number of parameters. More information about the models is included in Table
2.

###  4.2 Initial Results

In Table 3, we show the results when each model trains on the entire training
dataset. Overall, most models perform well on most languages, with Afro-XLM-R
performing the best on average. mBERT, on the other hand, performs the worst
overall, with an F1 score of 0 on Amharic, as it was not pre-trained on data
containing this script (Adelani et al., 2021).

Table 3: The F1 score when fine-tuning each model on unaltered training data.
Bold indicates the best performance per language.

Model AfriBERTa Afro-XLM-R XLM-R mBERT Average Language amh 72.1 (0.9) 75.9
(1.9) 71.9 (0.9) 0.0 (0.0) 55.0 en 88.5 (0.3) 92.8 (0.1) 92.7 (0.2) 92.6 (0.2)
91.7 hau 90.0 (0.5) 90.8 (0.4) 89.8 (0.2) 87.2 (0.5) 89.5 ibo 87.1 (0.3) 87.0
(0.6) 83.2 (0.2) 84.7 (0.5) 85.5 kin 74.1 (0.7) 78.1 (0.2) 72.5 (1.3) 70.7
(0.5) 73.8 lug 78.7 (0.2) 81.3 (0.2) 77.7 (0.4) 79.6 (0.7) 79.3 luo 68.1 (0.9)
69.2 (4.9) 69.4 (2.2) 71.7 (0.9) 69.6 pcm 85.5 (0.6) 89.2 (0.3) 86.2 (1.5)
88.0 (0.1) 87.2 swa 87.5 (0.6) 88.3 (0.2) 87.5 (0.6) 86.0 (0.7) 87.3 wol 61.4
(1.4) 66.1 (1.6) 63.9 (0.8) 63.4 (0.9) 63.7 yor 79.3 (0.6) 80.9 (1.0) 76.5
(1.1) 78.7 (0.7) 78.8 Average 79.3 81.8 79.2 73.0 78.3

###  4.3 How Corruption Affects Performance

Here we compare the three different corruption strategies: (1) deleting a
certain fraction of sentences, (2) deleting (i.e., setting to O) a certain
fraction of labels, and (3) swapping (i.e., replacing with another, but
incorrect entity) a certain fraction of labels. The results are shown in
Figure 2. The first conclusion we can draw from this is that the number of
sentences is far less important than the quality of annotations. In
particular, when deleting 90% of the sentences (leaving us with about 10% of
the labels of the original dataset), we can still recover around 75% of the
performance of training on the entire dataset. When we set 90% of the labels
to O, however, the performance is much worse, just above 10% of training on
the original dataset. Thus, even though the number of labels is roughly equal
for each case, having incorrectly labelled data affects the models much more
than having fewer sentences that are completely labelled.111We do verify that
when keeping only $X\%$ of the sentences, the fraction of labels we are left
with is very close to $X\%$ compared to the original dataset, confirming that
labels are mostly distributed uniformly throughout the sentences. When we swap
labels with incorrect entities, the models perform similarly, but slightly
worse compared to replacing these labels with O. This suggests that when an
annotator is uncertain, it is better to leave an entity out compared to
labelling it incorrectly.

Overall, this experiment shows that the fraction of correct annotations is an
important factor in NER. This means that, for every amount of labels, having
fewer, but completely labelled sentences is significantly better than having
the same number of sentences, but with incomplete or incorrect labels.

Figure 2: Comparing the result of (blue) deleting sentences, (orange) deleting
labels and (green) swapping labels. The X-axis represents the level of
quality: $1.0$ is the original dataset, whereas $0.1$ means that we keep
$10\%$ of the sentences or $10\%$ of the labels and corrupt/delete the other
90%. The mean here is shown, with the standard deviation across seeds,
languages and models shaded.

###  4.4 Performance Across Models and Languages

Having demonstrated that, on average, the quality of data is much more
important than the quantity of data, we now investigate how these results
differ across pre-trained models and languages. In the top row of Figure 3 we
plot the results of each corruption strategy, showing the performance of each
model separately.

Overall, all models perform similarly, with AfriBERTa and mBERT being slightly
more robust to being trained on small amounts of data than the larger XLM-R
models. We study the effect of language on our results in Figure 3 (bottom),
averaging over the random seeds and pre-trained models. We find that, on
average, each language performs similarly. The only exception here is when we
reduce the number of sentences the models were trained on, English performs
better than average, whereas Luo performs worse. One reason for these
observations is that the CONLL dataset is significantly larger than the other
languages’ datasets, making it perform well when given only a small fraction
of the data, as this still corresponds to a large absolute number of
sentences. Similarly, Luo has the smallest dataset out of all eleven
languages, making it more susceptible to having even less data. This also
suggests that having fewer, higher-quality sentences is preferred, but that
having too few sentences can drastically reduce performance.

Figure 3: Showcasing the effect of each corruption when isolating each (top)
pre-trained model or (bottom) language. We plot the mean and standard
deviation over random seeds and the other feature (the language in the top row
and the model in the bottom row).

###  4.5 The Tradeoff Between Quantity and Quality

The results in the previous sections indicate that prioritising data quality
and the correctness of annotations is important, more so than the number of
sentences we train on. We investigate this further by looking at the
relationship between quantity and quality, combining the removal of sentences
with the deletion of labels. Since we have shown that our results are mostly
consistent across languages and models, we consider only Afro-XLM-R and mBERT
as well as three languages, English, Swahili and Luo. We choose these two
models due to their differences in number of parameters and pre-training
languages. Most languages exhibit roughly similar performance, but English and
Luo had slightly different behaviour when deleting sentences, due to their
different dataset sizes. Swahili is used as a baseline as it had roughly
average behaviour. The results are shown in Figure 4 and they confirm that
deleting labels is more detrimental to performance than removing a similar
percentage of sentences. For example, for Afro-XLM-R and Swahili, keeping 25%
of the sentences but not removing any labels results in close to optimal
performance, at 99%. Having 25% of the labels but keeping all of the sentences
gives much worse performance, at 37%, even though the overall number of
correctly labelled entities is similar. Having 50% of the sentences and 50% of
the labels, the results are in-between, just below 70%. However, as can be
seen when the fraction of remaining sentences becomes too small (e.g. having
10% or fewer sentences for Luo), the models’ performance suffers drastically,
regardless of quality. Thus, while having correctly annotated labels is a
priority, if we do not have enough data, even perfect-quality annotations will
result in poor performance.222mBERT shows a similar trend to Afro-XLM-R but,
as before, it is slightly more robust to being fine-tuned on small amounts of
data.

Figure 4: Showing the effect of varying both the number of sentences and the
fraction of deleted labels on (top) mBERT and (bottom) Afro-XLM-R for en, swa
and luo. The x-axis shows the percentage of sentences remaining whereas the
y-axis lists the percentage of labels remaining. For example, at (50%, 50%),
we first remove half of the sentences and then delete half of the labels in
the remaining sentences. Each cell contains the fraction of F1 obtained when
training on this data compared to training on the original dataset, averaged
over 3 seeds.

##  5 Discussion & Future Work

Our work follows a recent trend of questioning whether more data is always
better in NLP, even at the cost of quality (Alabi et al., 2019; Abdulmumin et
al., 2022). In contrast to many of these works, we systematically and
quantifiably investigate the effects of reducing the quantity and quality of
data. Firstly, our findings demonstrate that the quantity of data used for
training does not greatly impact final performance, as we were able to achieve
around 80% performance with just 10-20% of the original dataset’s sentences.
However, removing entity labels, which simulates a reduction in annotation
quality, had a significant impact on performance. This supports prior findings
that we do not need a large amount of data to perform well when leveraging
pre-trained language models (Adelani et al., 2022a). We do note, however, that
when the number of sentences falls below a certain threshold (roughly between
200 and 400 sentences), performance drops significantly, indicating that we do
need a minimum amount of data to perform well. Furthermore, while quality has
anecdotally been shown to be more important than quantity of data (Alabi et
al., 2019; Abdulmumin et al., 2022), here we quantify this effect in NER. Our
results imply that when we have the budget to label $N$ entities, using fewer
fully-labelled sentences is better than using more sentences that are only
partially or incorrectly labelled. Our results also suggest that even modest
data-collection and annotation efforts should be able to result in datasets
that are large enough to obtain decent performance. Quality, however, is of
great importance and should be prioritised in the data creation process.

The second overarching observation we can make is that, in most cases, all
models exhibit roughly equal behaviour, in terms of the dropoff in
performance, as the level of corruption increases. This ranges from the
African language-centric models (AfriBERTa and Afro-XLM-R) to the
predominantly high-resourced models (XLM-R and mBERT). This suggests that the
behaviour we see here is quite general, as opposed to being specific to just a
particular model. The variation across the eleven languages is also remarkably
low, again highlighting the consistency of our results. This is notable as we
considered languages that were included in the pre-trained models’ training
data as well as some that were not.

There are numerous avenues for future work. One option would be to expand our
work into other NLP tasks such as machine translation or question answering,
to determine whether the same trend holds. Developing new corruption
strategies that are applicable to other tasks and cover other aspects of
quality would also be necessary to achieve this. Investigating the effect of
the corruptions on data characteristics, such as the prevalence of rare
tokens, would also be valuable. Furthermore, while we do not directly consider
methods to address the problems introduced by corrupted data, our work
motivates research into developing better methods for dealing with corrupted
datasets, mitigating some of the negative effects of training on low-quality
data. One promising option would be to use active learning, to choose which
sentences should be labelled, or verified by a human annotator (Gal et al.,
2017).

Finally, we randomly remove sentences and find that this can still result in
high performance. Recent work, however, has demonstrated that carefully
choosing the data points to train on can result in significantly more data
efficiency, requiring vastly less data while achieving comparable performance
(Mindermann et al., 2022; Sorscher et al., 2022). This line of work is
promising and has great potential in NLP for low-resourced languages.

##  6 Conclusion

In this paper, we present a systematic analysis of the impact of data quality
and quantity on the performance of pre-trained models in a named entity
recognition task for low-resourced languages. By designing multiple corruption
strategies and fine-tuning models on datasets with varying degrees of
corruption, we are able to provide useful insights into the relationship
between data quality and model performance.

Our results, which are consistent across pre-trained models and languages,
demonstrate that pre-trained models can perform effectively with minimal data
and that missing or incorrect annotations have a much greater negative impact
than having fewer fully-labelled sentences. The findings of this study have
the potential to inform future NER dataset creation efforts and aid NLP
practitioners in selecting appropriate datasets for fine-tuning.

## References

  * Abdul-Rauf et al. (2012) Sadaf Abdul-Rauf, Mark Fishel, Patrik Lambert, Sandra Noubours, and Rico Sennrich. 2012.  Extrinsic evaluation of sentence alignment systems.  _Workshop on Creating Cross-language Resources for Disconnected Languages and Styles_.
  * Abdulmumin et al. (2022) Idris Abdulmumin, Michael Beukman, Jesujoba O. Alabi, Chris Emezue, Everlyn Asiko, Tosin Adewumi, Shamsuddeen Hassan Muhammad, Mofetoluwa Adeyemi, Oreen Yousuf, Sahib Singh, and Tajuddeen Rabiu Gwadabe. 2022.  Separating grains from the chaff: Using data filtering to improve multilingual translation for low-resourced african languages.  In _Proceedings of the Seventh Conference on Machine Translation: Shared Task Papers_.
  * Adelani et al. (2021) David Ifeoluwa Adelani, Jade Z. Abbott, Graham Neubig, Daniel D’souza, Julia Kreutzer, Constantine Lignos, Chester Palen-Michel, Happy Buzaaba, Shruti Rijhwani, et al. 2021.  Masakhaner: Named entity recognition for african languages.  _Trans. Assoc. Comput. Linguistics_ , 9:1116–1131.
  * Adelani et al. (2022a) David Ifeoluwa Adelani, Jesujoba O. Alabi, Angela Fan, Julia Kreutzer, Xiaoyu Shen, Machel Reid, Dana Ruiter, Dietrich Klakow, Peter Nabende, et al. 2022a.  A few thousand translations go a long way! leveraging pre-trained models for african news translation.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL 2022, Seattle, WA, United States, July 10-15, 2022_ , pages 3053–3070. Association for Computational Linguistics.
  * Adelani et al. (2022b) David Ifeoluwa Adelani, Graham Neubig, Sebastian Ruder, Shruti Rijhwani, Michael Beukman, Chester Palen-Michel, Constantine Lignos, Jesujoba O. Alabi, Shamsuddeen Hassan Muhammad, et al. 2022b.  Masakhaner 2.0: Africa-centric transfer learning for named entity recognition.  _CoRR_ , abs/2210.12391.
  * Alabi et al. (2022) Jesujoba O Alabi, David Ifeoluwa Adelani, Marius Mosbach, and Dietrich Klakow. 2022\.  Multilingual language model adaptive fine-tuning: A study on african languages.  _CoRR_ , abs/2204.06487.
  * Alabi et al. (2019) Jesujoba O. Alabi, Kwabena Amponsah-Kaakyire, David Ifeoluwa Adelani, and Cristina España-Bonet. 2019.  Massive vs. curated word embeddings for low-resourced languages. the case of yorùbá and twi.  _CoRR_ , abs/1912.02481.
  * Conneau et al. (2020) Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. 2020.  Unsupervised cross-lingual representation learning at scale.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020_ , pages 8440–8451. Association for Computational Linguistics.
  * de Gibert Bonet et al. (2022) Ona de Gibert Bonet, Ksenia Kharitonova, Blanca Calvo Figueras, Jordi Armengol-Estapé, and Maite Melero. 2022.  Quality versus quantity: Building Catalan-English MT resources.  In _Proceedings of the 1st Annual Meeting of the ELRA/ISCA Special Interest Group on Under-Resourced Languages_ , pages 59–69, Marseille, France. European Language Resources Association.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pages 4171–4186. Association for Computational Linguistics.
  * Gal et al. (2017) Yarin Gal, Riashat Islam, and Zoubin Ghahramani. 2017.  Deep bayesian active learning with image data.  In _ICML_ , volume 70 of _Proceedings of Machine Learning Research_ , pages 1183–1192. PMLR.
  * Gascó et al. (2012) Guillem Gascó, Martha-Alicia Rocha, Germán Sanchis-Trilles, Jesús Andrés-Ferrer, and Francisco Casacuberta. 2012.  Does more data always yield better translations? In _Proceedings of the 13th Conference of the European Chapter of the Association for Computational Linguistics_ , pages 152–161, Avignon, France. Association for Computational Linguistics.
  * Kreutzer et al. (2022) Julia Kreutzer, Isaac Caswell, Lisa Wang, Ahsan Wahab, Daan van Esch, Nasanbayar Ulzii-Orshikh, Allahsera Tapo, Nishant Subramani, Artem Sokolov, et al. 2022.  Quality at a glance: An audit of web-crawled multilingual datasets.  _Trans. Assoc. Comput. Linguistics_ , 10:50–72.
  * Lample and Chaplot (2017) Guillaume Lample and Devendra Singh Chaplot. 2017.  Playing FPS games with deep reinforcement learning.  In _Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence, February 4-9, 2017, San Francisco, California, USA_ , pages 2140–2146. AAAI Press.
  * Martinus and Abbott (2019) Laura Martinus and Jade Z. Abbott. 2019.  A focus on neural machine translation for african languages.  _CoRR_ , abs/1906.05685.
  * Mindermann et al. (2022) Sören Mindermann, Jan Markus Brauner, Muhammed Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N. Gomez, Adrien Morisot, Sebastian Farquhar, and Yarin Gal. 2022.  Prioritized training on points that are learnable, worth learning, and not yet learnt.  In _International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 15630–15649. PMLR.
  * Nekoto et al. (2022) Wilhelmina Nekoto, Julia Kreutzer, Jenalea Rajab, Millicent Ochieng, and Jade Abbott. 2022.  Participatory translations of oshiwambo: Towards sustainable culture preservation with language technology.  In _3rd Workshop on African Natural Language Processing_.
  * Nekoto et al. (2020) Wilhelmina Nekoto, Vukosi Marivate, Tshinondiwa Matsila, Timi E. Fasubaa, Taiwo Fagbohungbe, Solomon Oluwole Akinola, Shamsuddeen Hassan Muhammad, Salomon Kabongo Kabenamualu, Salomey Osei, et al. 2020.  Participatory research for low-resourced machine translation: A case study in african languages.  In _Findings of the Association for Computational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020_ , volume EMNLP 2020 of _Findings of ACL_ , pages 2144–2160. Association for Computational Linguistics.
  * Ogueji et al. (2021) Kelechi Ogueji, Yuxin Zhu, and Jimmy Lin. 2021.  Small data? no problem! exploring the viability of pretrained multilingual language models for low-resourced languages.  In _Proceedings of the 1st Workshop on Multilingual Representation Learning_ , pages 116–126, Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. 2018.  Improving language understanding by generative pre-training.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9.
  * Sang and Meulder (2003) Erik F. Tjong Kim Sang and Fien De Meulder. 2003.  Introduction to the conll-2003 shared task: Language-independent named entity recognition.  In _Proceedings of the Seventh Conference on Natural Language Learning, CoNLL 2003, Held in cooperation with HLT-NAACL 2003, Edmonton, Canada, May 31 - June 1, 2003_ , pages 142–147. ACL.
  * Sorscher et al. (2022) Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari S. Morcos. 2022.  Beyond neural scaling laws: beating power law scaling via data pruning.  _CoRR_ , abs/2206.14486.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pages 5998–6008.