marginparsep has been altered.
topmargin has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

Curvature-informed multi-task learning for graph networks

Anonymous Authors1

††footnotetext:  1Anonymous Institution, Anonymous City, Anonymous Region,
Anonymous Country. Correspondence to: Anonymous Author
<EMAIL_ADDRESS>
Preliminary work. Under review by the 2nd AI4Science Workshop at ICML 2022. Do
not distribute.

###### Abstract

Properties of interest for crystals and molecules, such as band gap,
elasticity, and solubility, are generally related to each other: they are
governed by the same underlying laws of physics. However, when state-of-the-
art graph neural networks attempt to predict multiple properties
simultaneously (the multi-task learning (MTL) setting), they frequently
underperform a suite of single property predictors. This suggests graph
networks may not be fully leveraging these underlying similarities. Here we
investigate a potential explanation for this phenomenon – the curvature of
each property’s loss surface significantly varies, leading to inefficient
learning. This difference in curvature can be assessed by looking at spectral
properties of the Hessians of each property’s loss function, which is done in
a matrix-free manner via randomized numerical linear algebra. We evaluate our
hypothesis on two benchmark datasets (Materials Project (MP) and QM8) and
consider how these findings can inform the training of novel multi-task
learning models.

##  1 Introduction

Graph networks (GNs) Battaglia et al. (2018) are considered state-of-the-art
machine learning (ML) methods for many scientific problems, including property
prediction of both inorganic crystalline materials Xie & Grossman (2018); Park
& Wolverton (2020) (hereafter “crystals”, e.g., Figure 1) and small organic
molecules Gilmer et al. (2017); Gasteiger et al. (2020) (hereafter
“molecules”, e.g., Figure 2); the same GN architectures have shown success in
both domains (e.g., Chen et al. (2019), who suggest the “crystal” vs.
“molecule” terminology and use “material” to refer to both). A common use case
might involve training a model from materials with known properties and then
screening a separate dataset lacking those properties to obtain potential
candidates to further investigate. Typically, property-prediction tasks are
formulated as single-target regression problems, where the goal is to predict
a scalar property like band gap, formation stability, elasticity, or
solubility. However, practical materials design requires optimizing multiple
properties of interest. For example, when designing a cell phone screen, a
material with both high optical transparency and hardness would be desired.

Predicting a single property can be posed as a single task for an ML model;
using a single ML model to predict multiple properties is thus a type of
multi-task learning (MTL). The most common family of MTL approaches is hard
parameter sharing, in which a single representation space is shared across
multiple tasks of interest, and task-specific networks map from that space to
output predictions Caruana (1997).

However, the latest research in novel MTL techniques largely has been focused
on classification problems in the natural language processing (NLP) and
computer vision (CV) domains – for example, the $\mathrm{MultiMNIST}$ problem
Sener & Koltun (2018), where a model must identify two separate digits in the
same picture. With some recent exceptions – e.g., Yang et al. (2021); Tan et
al. (2021); Kong et al. (2021) – there has been little work exploring the
application of MTL to multi-property GN regression problems.

Some existing papers have used hard parameter sharing for material property
prediction, in which a single model predicts a set of properties. However,
results suggest that this approach has degraded performance compared to a set
of single-property models – see, e.g., Gasteiger et al. (2020), which predicts
twelve quantum mechanical properties of molecules and finds that a single
multi-output model performs worse than a set of single-output models.

Works have explored reasons for this: Yu et al. (2020) identifies a “tragic
triad” of reasons that MTL models might perform worse compared to single-task
ones, one of which is that multi-task loss functions can have high curvature,
as characterized by the Hessian of the loss. However, Yu et al. (2020) rely on
first-order approximations to curvature and do not directly incorporate
second-order information.

Concurrently, several works have applied techniques from randomized numerical
linear algebra to directly probe the Hessians of large ML models Sagun et al.
(2017); Alain et al. (2019); Yao et al. (2019); Ghorbani et al. (2019); Papyan
(2020). These rely (1) on the fact that the product of a model’s Hessian and
another vector (a Hessian-vector product) may be efficiently obtained, without
explicitly calculating the full Hessian, and (2) that given a matrix-free
Hessian-vector product function, several key spectral properties of the
Hessian, including its eigenvalues and trace, may be estimated. However, these
papers focus only on standard image classification problems, and do not
consider the MTL setting, GNs, or regression problems.

In this paper, we formulate the application of curvature-informed techniques
to MTL models (Section 2.1), describe the application of graph networks to the
problem of multi-property prediction in the domains of materials science and
chemistry (Section 2.2 and Section 2.3), and analyze training dynamics of
multi-property prediction models using methods based on assessing local
curvature of loss function surfaces (Section 3). Our results suggest how novel
multi-property prediction models might inherently account for differences in
curvature to enhance learning efficiency.

Figure 1: Example of converting a periodic crystal lattice structure
($\text{Ba}\,\text{Ti}\,\text{O}_{3}$) to a multigraph
(a): Crystal lattice structure, created using
ase222https://wiki.fysik.dtu.dk/ase/.
(b): The corresponding multigraph representation of the crystal that is
ingested by a graph network. Figure 2: Example of converting a 2D molecule
structure ($\text{C}_{4}\,\text{H}_{6}\,\text{F}_{3}\,\text{N}$) to a graph
(a) Molecular structure as represented by rdkit
(b) The corresponding graph

##  2 Approach

Let $\mathcal{G}$ be a space of directed multigraphs corresponding to crystal
or molecule structures. A directed graph $g=(\mathcal{V},\mathcal{E})$
consists of a set of nodes $\mathcal{V}=\\{v\\}$ and edges
$\mathcal{E}=\\{(v,v^{\prime},k)\\}$. Each node $v$ has a node embedding
$x_{v}\in\mathbb{R}^{V}$, and each edge $(v,v^{\prime},k)$ has an edge
embedding $e_{v,v^{\prime},k}\in\mathbb{R}^{E}$. A pair of nodes $v$ and
$v^{\prime}$ may have multiple edges connecting them, and these multi-edges
are indexed by $k=1,\ldots,K_{v,v^{\prime}}$. Furthermore, a graph $g$ has $T$
properties $y=(y_{1},\ldots,y_{T})\in\mathcal{Y}$, here assumed to all be
scalars.

We require a model that can predict all targets $y$ for a given $g$. This
entails minimizing a loss with respect to a set of neural network weights
$\\{\theta_{sh},\theta_{1},\ldots,\theta_{T}\\}$, where $\theta_{t}$ is used
only in predicting property $t$, and $\theta_{sh}$ is shared across
predictions. The form of this problem and its neural network is discussed in
Section 2.2. First we consider the loss’s curvature.

###  2.1 Curvature assessment techniques

We view the shared parameters $\theta_{sh}$ of a graph network as a flattened
vector in $\mathbb{R}^{P}$; let $L_{t}:\mathbb{R}^{P}\to\mathbb{R}$ be a
property-specific loss function of $\theta_{sh}$, where the property-specific
parameters $\theta_{1},\ldots,\theta_{T}$ are held constant. Typically,
machine learning models consider only first-order derivative information of
$L_{t}$ when training neural networks. However, second-order information is
also useful in optimization problems – consider how a quasi-Newton method like
limited-memory L-BFGS Liu & Nocedal (1989) is more efficient than a first
order method like gradient descent. This suggests that additional insight into
the training problems of ML models can be found in using second-order
information.

The curvature of a function $L_{t}$ is characterized by its Hessian
$H_{t}=\nabla_{\theta_{sh}}^{2}L_{t}\in\mathbb{R}^{P\times P}$ for a weight
vector $\theta_{sh}\in\mathbb{R}^{P}$. For typical deep neural networks,
$P\approx 10^{6}$, and calculating $H$ directly is computationally infeasible
due to storage contraints ($H_{t}$ will have $10^{12}$ entries). However, the
product of $H_{t}$ with a vector $v$ is computable in approximately the same
time-complexity as calculating the gradient of $L_{t}$ with respect to
$\theta$ Pearlmutter (1994).333In jax, this is implemented by composing a
forward-mode Jacobian-vector product jvp function with a reverse-mode gradient
function grad. This enables calculation of properties of $H_{t}$ that can be
obtained from observing its action on a chosen vector $v$ Yao et al. (2019);
Ghorbani et al. (2019). In particular, the trace of $H_{t}$, and an
approximate probability density function of its eigenvalues can be efficiently
estimated.

The trace of $H_{t}$ is both the sum of its eigenvalues and $L_{t}$’s
Laplacian (i.e.,
$\mathrm{tr}\,H_{t}=\sum_{p}\partial_{\theta_{sh,p}\theta_{sh,p}}L_{t}(\theta)$).
Thus, the trace of $H_{t}$ can be viewed as a measurement for the complexity
of its curvature Yao et al. (2019). The sign of the trace also indicates the
sign of $L_{t}$’s local curvature. The trace of $H_{t}$ may be estimated with
the expectation

$\mathrm{tr}\,H_{t}=\mathbb{E}_{v\sim\mathbb{P}_{R}}[v^{T}H_{t}(\theta)v],$

where $\mathbb{P}_{R}$ is the distribution of random variables in
$\mathbb{R}^{P}$ with $iid$ Rademacher components Avron & Toledo (2011).

Having access to all of $H_{t}$’s eigenvalues gives us a full picture of
$L_{t}$’s curvature. These eigenvalues are often represented as a function
$\psi$ called the spectral density or density of states that is given by
$\psi(t)=\frac{1}{P}\sum_{p}\delta(t-\lambda_{p})$, where
$\lambda_{1}\leq\cdots\leq\lambda_{P}$ are the eigenvalues of $H_{t}$, and
$\delta$ is the Dirac distribution Lin et al. (2016). However, the weight
vector $\theta_{sh}$ is too high-dimensional for all of $H_{t}$’s eigenvalues
to be directly calculated in practice. Methods like power iteration can be
used to estimate the large-magnitude eigenvalues of a Hessian Yao et al.
(2019), but they scale poorly as its dimensionality increases.

A solution is to first use the Lanczos algorithm to estimate $P^{\prime}\ll P$
eigenvalues of $H_{t}$ and then smooth these estimated values into a
continuous approximation to the discrete density $\psi$ by convolving the
values with a Gaussian kernel Golub & Welsch (1969); Ghorbani et al.
(2019).444The Lanczos algorithm accurately estimates the Hessian’s most
positive and most negative eigenvalues, and the convolution approximates the
distribution of eigenvalues between these extrema. This yields a continuous
function $\hat{\psi}(t)$ that approximates $\psi$. See Appendix E for an
example of estimating $\hat{\psi}$ for a simplified loss function.

Interpreting estimated spectral densities remains a challenging task: Ghorbani
et al. (2019) and Yao et al. (2019) claim that a spectral density should be
“smooth” (i.e., concentrated within a dense region) and consider the effect of
different network architectures on density smoothness. Alain et al. (2019)
argues that having a large concentration of negative eigenvalues can lead to
inefficient training, because most optimizers fail to leverage this local
negative curvature. Papyan (2020) relate the outlier structure of a Hessian’s
eigenvalues to the number of classes in a dataset.

Compared to CV problems, graph-structured datasets have a number of
interesting attributes. For example, a graph input has a variable number of
node and edge features, which complicates the GN learning problem.
Furthermore, in contrast to the use of very deep networks in CV problems,
increasing depth in graph networks has often not been found to be effective
(see, e.g., discussion in Godwin et al. (2022)). Hessian-based information has
the potential to inform some of these comparisons.

Note that the estimations of both the trace and the spectral densities are
based on random quantities. For trace, multiple vectors $v$ must be sampled
and the quantity $v^{T}H_{t}(\theta)v$ computed, and then convergence in mean
can be assessed. In practice, we find that approximately $500$ samples are
sufficient. For spectral density, each Lanczos iterate is initialized as a
random Gaussian vector. We adapt an implementation of the Lanczos algorithm
and smoothing process developed by Ghorbani et al. (2019), in which we
reorthogonalize the Lanczos iterates during each step of estimation to promote
numerical stability. We follow Yao et al. (2019) and use $P^{\prime}=100$,
which is more than the $P^{\prime}=90$ that Ghorbani et al. (2019) validate by
comparing to an exact calculation of every eigenvalue of a smaller network.

###  2.2 Graph networks for property prediction

Let $\hat{y}:\mathcal{G}\to\mathcal{Y}$ be a neural network parameterized by
$\theta$. We follow Sener & Koltun (2018) in splitting $\theta$ into a set of
shared weights $\theta_{sh}$, and a set of property-specific weights
$\theta_{1},\ldots,\theta_{T}$. Then the predicted $t$th property is

$\hat{y}_{t}(g;\theta_{sh},\theta_{t})=f_{t}(z(g;\theta_{sh});\theta_{t}),$

where $z:\mathcal{G}\to\mathbb{R}^{D}$ is a task-independent GN Battaglia et
al. (2018), and each $f_{t}:\mathbb{R}^{D}\to\mathbb{R}$ is a task-specific
feedforward network.

We briefly outline the functionality of our graph network $z$, which maps an
input graph $g$ to a graph-level representation vector $u^{M}$. First, the
multigraph $g$ is given a global feature $u^{0}\in\mathbb{R}^{U}$ that is
initialized as a vector of ones, and, similarly to Chen et al. (2019), the
node features $x_{v}$ and edge features $e_{v,v^{\prime},k}$ are linearly
projected: $x^{0}_{v}=W_{V}\,x_{v}$ and
$e_{v,v^{\prime},k}^{0}=W_{E}\,e_{v,v^{\prime},k}$, where $W_{V}$ and $W_{E}$
are matrices learned during training. Then information is propagated across
the multigraph during $m=1,\ldots,M$ message-passing steps.

The $m$th message-passing step proceeds as follows: First, the edge features
are updated:

$e^{m}_{v,v^{\prime},k}=\phi^{E}(\mathrm{cat}(e_{v,v^{\prime},k}^{m-1},x_{v}^{m-1},x_{v^{\prime}}^{m-1},u^{m-1}))+e^{m-1}_{v,v^{\prime},k},$

where $\phi^{E}$ is a feed-forward neural network, and $\mathrm{cat}$ is the
concatenation operator. Then node updates are collected. For every node $v$,
first edge updates are calculated:

$\displaystyle h_{v,out}^{m}$ $\displaystyle=$
$\displaystyle\sum_{(v^{\prime},k):(v,v^{\prime},k)\in\mathcal{E}}e^{m}_{v,v^{\prime},k}$
$\displaystyle h_{v,in}^{m}$ $\displaystyle=$
$\displaystyle\sum_{(v^{\prime},k):(v^{\prime},v,k)\in\mathcal{E}}e^{m}_{v^{\prime},v,k},$

where the first summation is over the tuples $(v^{\prime},k)$ such that
$(v,v^{\prime},k)$ is an edge for $g$, and the second summation is defined
similarly. With these updates, the new node representation for $v$ is given by

$x^{m}_{v}=\phi^{V}(\mathrm{cat}(h_{v,out}^{m},h_{v,in}^{m},x^{m-1}_{v},u^{m-1}))+x^{m-1}_{v},$

where $\phi^{V}$ is a feed-forward neural network. Finally, the global feature
is updated:

$u^{m}=\phi^{U}\left(\mathrm{cat}\left(\sum_{v,v^{\prime},k}e^{m}_{v,v^{\prime},k},\sum_{v}x^{m}_{v},u^{m-1}\right)\right)+u^{m-1},$

where $\phi^{U}$ is a feed-forward neural network, and each feature vector
$x^{m}_{v},e^{m}_{v,v^{\prime},k},u^{m}$ is processed with a layer
normalization Ba et al. (2016) layer. After $M$ steps, the final global state
$u^{M}$ is fed into each property-specific network $f_{t}$.

Our graph networks are implemented in the
jraph555https://github.com/deepmind/jraph framework that is built on
jax666https://github.com/google/jax. We use
flax777https://github.com/google/flax to implement component neural networks
$\phi^{V}$, $\phi^{E}$, and $\phi^{U}$. Because we seek to evaluate second-
order derivative information, we ensure that our neural networks
$\phi^{V},\phi^{E},\phi^{U}$ are smooth functions with respect to their
weights by using $\tanh$ activation functions. Further details on specific
hyperparameters used to instantiate models are given in Appendix A.

We impose on each property a loss function $L_{t}$ (here assumed to be mean
squared loss for all $t$), and we collect a training dataset of multigraph
data points $\mathcal{D}=\\{(g^{n},y^{n})\\}_{n=1}^{N}$. Then we solve the
minimization problem

$\min_{\theta}L(\theta)=\sum_{t}L_{t}(\theta_{sh},\theta_{t}),$

where

$L_{t}(\theta_{sh},\theta_{t})=\sum_{n}|\hat{y}_{n}(g^{n};\theta_{sh},\theta_{t})-y^{n}_{t}|^{2}$

are task-specific losses and
$\theta=\\{\theta_{sh},\theta_{1},\ldots,\theta_{T}\\}$.

We solve this minimization with stochastic gradient descent using the
optax888https://github.com/deepmind/optax implementation of AdamW Loshchilov &
Hutter (2019) to choose stepsizes; we handle the potential difference in
scales across properties by normalizing all targets prior to training and set
$\mu_{t}=1$ for all $t$. Further details on the specifics of model training
are given in Appendix B.

###  2.3 Data sources

We evaluate our graph network curvature assessment techniques on datasets from
two scientific domains: materials science (crystals) and chemistry
(molecules). For crystals, our data featurization follows Park & Wolverton
(2020); for molecules, it follows Gilmer et al. (2017). To reduce the
complexity of our considered problem space, we choose to use simple
featurizations (one-hot encoding of atom-types as node featurizations), but
our methods are compatible with other featurizations (e.g., hand-crafted node
descriptors like those provided by rdkit). Further details about the data used
are available in Appendix C.

For materials science, we collect data from Materials Project (MP) Jain et al.
(2013), which contains results of density functional theory (DFT) calculations
for different crystals. crystal records contain compositional and structural
information, as well as some related properties. Each crystal’s structure
information is captured in a Crystallographic Information File (CIF), which we
convert into a multigraph using the VoronoiNN function from
pymatgen999https://pymatgen.org, following the approach of Park & Wolverton
(2020). As crystals are periodic structures, this process yields multiple
edges between many given pairs of nodes Xie & Grossman (2018); Park &
Wolverton (2020); Sanyal et al. (2018) (e.g., Figure 1). As targets, we use
several assessments of a crystal’s elastic properties: the Voigt-Reuss-Hill
Hill (1952) calculation for shear ($G_{VRH}$) and bulk ($K_{VRH}$) modulus,
both measured in units of GPa, as well as the isotropic Poisson ratio $\mu$, a
dimensionless quantity.

These properties are inherently coupled – $\mu$ can be calculated as a
function of $G_{VRH}$ and $K_{VRH}$. Thus, an ML model that predicts all three
serves as a test for how well it can learn underlying physical relationships.
A one-hot encoding of atom-type is used to obtain initial node features
$x_{v}$, and four bond-related properties calculated by pymatgen are used as
edge features, resulting in a dataset of 10,500 crystal structures with known
elastic properties.

For chemistry, we use the QM8 Ramakrishnan et al. (2015) dataset of organic
molecules. We use the rdkit package101010https://github.com/rdkit/rdkit to
convert SMILES strings Weininger (1988) into molecular structure graphs (see,
e.g., Figure 2). As targets, we use the first and second transition energies,
$E_{1}$ and $E_{2}$, and the first and second oscillator strengths $f_{1}$ and
$f_{2}$. QM8 contains several versions of transition energies and oscillator
strengths that have been calculated via different levels of DFT. We use the
predictions of the approximate coupled-cluster (CC2) Hättig & Weigend (2000)
method as our targets, as these are treated as the “ground truth” in
Ramakrishnan et al. (2015). A one-hot encoding of atom-type is used to obtain
initial node features $x_{v}$, and a one-hot encoding of bond-type is used to
obtain initial edge features $e_{v,v^{\prime},k}$.

Note that, in Section 2.2, we describe our input data points $g$ as directed
multigraphs. However, for our actual data points, the edges are not inherently
directed – bonds are symmetric. Thus, we duplicate each edge feature (i.e.,
$e_{v,v^{\prime},k}=e_{v^{\prime},v,k}$ for all edges) during data
preprocessing.

##  3 Results

Figure 3: Standardized MAE values for each property across train and test
splits, for the MP and QM8 prediction tasks. Despite the high variability in
curvature across training (Figure 4), training error reduces relatively
smoothly, with occasional regressions in performance for MP. Although the
training error for each property converges to rough equality, significant
inter-property variety is observed for test set error.

We demonstrate the application of our curvature-assessment techniques (Section
2.1) on trained GNs (Section 2.2) in two application domains: materials
science and chemistry (Section 2.3). Although here we focus on GNs Battaglia
et al. (2018), our assessment techniques are applicable to the broad family of
graph neural networks used in property prediction tasks.

Figure 3 shows training and test set errors for each dataset. These results do
not necessarily align with domain intuition, suggesting that the models do not
leverage scientific knowledge in their learned representations. For example,
for MP, the test set error for Poisson ratio is higher than that of $G_{VRH}$
or $K_{VRH}$, despite the fact that Poisson ratio can be calculated as a
function of $G_{VRH}$ and $K_{VRH}$. This is interesting and suggests that the
GNs are not fully leveraging known scientific knowledge. For QM8, the
oscillator strengths $f_{1}$ and $f_{2}$ have the lowest and highest,
respectively, test set errors.

Next we summarize the curvature of the loss functions using the trace of each
task’s loss’s Hessian. In Figure 4, we show that, despite the general
decreasing training error, the curvature of each property’s loss surface, as
measured by the trace of its Hessian, is highly variable. In particular, both
datasets show high heterogeneity across properties and across training epochs
in their estimated traces. Our observations here do not align with prior work
that has analyzed Hessian traces for CV problems Yao et al. (2019). In those
works, traces increased monotonically during training and were consistently
positive. Here, our estimated traces often flip between being positive and
negative.

The Hessian trace is a high-level summary of the curvature of a loss surface;
for a more granular examination, we estimate the spectral densities of each
property’s Hessian. In Figure 6(a) in D, we show that the spectral densities
are similarly variable and often feature the presence of outliers that briefly
occur during training. These outliers vary across properties. For the QM8
dataset, we zoom in on the range of eigenvalues where most density is
concentrated in Figure 5. Similar to Yao et al. (2019); Ghorbani et al.
(2019), we see that most density is concentrated near 0. This suggests that
there is a great deal of redundancy in the latent spaces learned by these
models, and their true dimensionality is likely significantly less than the
full $P$ parameters of the $\theta_{sh}$ weight vector. Unlike these works,
the spectral densities are more symmetric around 0. This matches Figure 4,
since the Hessian trace (sum of all its eigenvalues) varies between being very
positive and very negative. The exact spectral densities vary across
properties, even at the end of training.

(a) MP

(b) QM8

Figure 4: We track property-specific Hessian traces $\mathrm{tr}\,H_{t}$ while
training GNs for the MP and QM8 prediction problems. The loss surface of the
MTL problem for GNs appears to be complicated and varying across properties.
Trace signs (indicating the sign of the curvature) and trace magnitudes vary
across training and across properties, the latter by orders of magnitude. In
prior work that looked at traces of Hessians Yao et al. (2019) in CV tasks,
traces were found to be less variable. Figure 5: Snapshots of the spectral
densities of property-specific Hessian losses for the QM8 dataset. All
properties start with comparable densities closely concentrated near 0; but,
as training finishes, the densities spread out as the loss surface increases
in complexity. Each property has a different spread of eigenvalues, suggesting
that the loss surfaces do vary in curvature by property. Note that this plot
trims the range of the $x$-axis to remove outlier eigenvalues. The full
spectra are displayed in Figure 6 in Appendix D. For these plots, the $x$-axis
gives the range of eigenvalues for the loss function, and the $y$-axis gives
the density $\hat{\psi}(t)$ of eigenvalues concentrated at that point $t$ (as
described in 2.1).

##  4 Discussion

In this work, we have posited that the performance of multi-task GNs for
property prediction may be stymied by underlying variation in property-
specific loss-surface curvatures. Some existing work considers this question
Yu et al. (2020) in a simplified theoretical setting, but, to the best of our
knowledge, no one has empirically investigated the Hessian properties of
multi-task GNs. In two domains – chemistry and materials science – our results
suggest that loss surface curvature varies across each modeled property.

In order to assess curvature without calculating the full Hessian, we build on
recent work that uses matrix-free methods to estimate Hessian properties Alain
et al. (2019); Yao et al. (2019); Ghorbani et al. (2019); Papyan (2020). We
extend their results by considering a novel type of learning problem – multi-
output scalar regression – and a novel class of neural network – graph
networks. Our results echo some previous results – the majority of a Hessian’s
spectral density is concentrated near 0 – and diverges in other respects –
Hessian properties appear far noisier and more variable for GNs than for CV
tasks. We leave further investigation of this question to future work.
Potential explanations include the difference in loss functions for regression
vs. classification tasks, a higher level of noise in the datasets we examined,
and some special characteristic of GN vs. other neural network architectures.

We have here focused on a specific but representative subset of GN prediction
problems, but considerable potential variation exists. For example, many
recent GN architectures incorporate a notion of equivariance Satorras et al.
(2021); Gasteiger et al. (2020) into their feature-extraction models, and this
might impact their curvature in different ways. In addition, we have not
evaluated how the choice of optimizer (e.g., stochastic gradient descent vs.
Adam Kingma & Ba (2014) vs. AdamW) impacts the curvature properties of a
learned loss surface.

Our curvature assessment enable several future research directions. First, our
analysis here was primarily empirical. The phenomena we identify here (a
diversity of curvatures across multi-task loss functions) and the previously-
observed degradation of performance for multi-task models Gasteiger et al.
(2020) could be connected by a theoretical justification. Similarly, we find
that curvature properties of GNs appear to be much noisier and more variable
than the properties of other network architectures, and we currently lack a
theoretical justification of why. Intermediate steps might entail
investigating the curvature properties of MTL methods that, in other domains,
do out-perform single-task models (e.g., Sener & Koltun (2018)). Existing work
in analyzing curvature for computational geometry (e.g. Goldman (2005)) might
provide techniques to build upon.

## Acknowledgements

This work was supported by internal research and development funding from the
Research and Exploratory Development Mission Area of the Johns Hopkins
University Applied Physics Laboratory. Thanks to Christopher Ratto for
assistance in editing and improving the work.

## References

  * Alain et al. (2019) Alain, G., Roux, N. L., and Manzagol, P.-A.  Negative eigenvalues of the hessian in deep neural networks, 2019.  URL https://arxiv.org/abs/1902.02366.
  * Avron & Toledo (2011) Avron, H. and Toledo, S.  Randomized algorithms for estimating the trace of an implicit symmetric positive semi-definite matrix.  _J. ACM_ , 2011.  doi: 10.1145/1944345.1944349.
  * Ba et al. (2016) Ba, J. L., Kiros, J. R., and Hinton, G. E.  Layer normalization, 2016.  URL https://arxiv.org/abs/1607.06450.
  * Battaglia et al. (2018) Battaglia, P. W., Hamrick, J. B., Bapst, V., Sanchez-Gonzalez, A., Zambaldi, V., et al.  Relational inductive biases, deep learning, and graph networks, 2018.  URL https://arxiv.org/abs/1806.01261.
  * Caruana (1997) Caruana, R.  Multitask learning.  _Machine Learning_ , 1997.  doi: 10.1023/A:1007379606734.
  * Chen et al. (2019) Chen, C., Ye, W., Zuo, Y., Zheng, C., and Ong, S. P.  Graph networks as a universal machine learning framework for molecules and crystals.  _Chemistry of Materials_ , 2019.  doi: 10.1021/acs.chemmater.9b01294.
  * Gasteiger et al. (2020) Gasteiger, J., Groß, J., and Günnemann, S.  Directional message passing for molecular graphs.  In _International Conference on Learning Representations_ , 2020.
  * Ghorbani et al. (2019) Ghorbani, B., Krishnan, S., and Xiao, Y.  An investigation into neural net optimization via hessian eigenvalue density.  In _Proceedings of the 36th International Conference on Machine Learning_ , pp. 2232–2241, 09–15 Jun 2019.
  * Gilmer et al. (2017) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry.  In _Proc. of the 34th Int. Conf. on Machine Learning - Volume 70_ , 2017.
  * Godwin et al. (2022) Godwin, J., Schaarschmidt, M., Gaunt, A. L., Sanchez-Gonzalez, A., Rubanova, Y., et al.  Simple GNN regularisation for 3d molecular property prediction and beyond.  In _Int. Conf. on Learning Representations_ , 2022.
  * Goldman (2005) Goldman, R.  Curvature formulas for implicit curves and surfaces.  _Computer Aided Geometric Design_ , 2005.  doi: 10.1016/j.cagd.2005.06.005.  Geometric Modelling and Differential Geometry.
  * Golub & Welsch (1969) Golub, G. H. and Welsch, J. H.  Calculation of gauss quadrature rules.  _Mathematics of Computation_ , pp. 221–s10, 1969.
  * Hill (1952) Hill, R.  The elastic behaviour of a crystalline aggregate.  _Proc. of the Phys. Soc. Section A_ , 65(5):349–354, may 1952.  doi: 10.1088/0370-1298/65/5/307.
  * Hättig & Weigend (2000) Hättig, C. and Weigend, F.  Cc2 excitation energy calculations on large molecules using the resolution of the identity approximation.  _J. Chem. Phys._ , 2000.  doi: 10.1063/1.1290013.
  * Jain et al. (2013) Jain, A., Ong, S. P., Hautier, G., Chen, W., Richards, W. D., et al.  The Materials Project: A materials genome approach to accelerating materials innovation.  _APL Materials_ , 2013.  doi: 10.1063/1.4812323.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization, 2014.  URL https://arxiv.org/abs/1412.6980.
  * Kong et al. (2021) Kong, S., Guevarra, D., Gomes, C. P., and Gregoire, J. M.  Materials representation and transfer learning for multi-property prediction.  _Appl. Phys. Rev._ , 2021.  doi: 10.1063/5.0047066.
  * Lin et al. (2016) Lin, L., Saad, Y., and Yang, C.  Approximating spectral densities of large matrices.  _SIAM Review_ , 2016.  doi: 10.1137/130934283.
  * Liu & Nocedal (1989) Liu, D. C. and Nocedal, J.  On the limited memory bfgs method for large scale optimization.  _Mathematical Programming_ , 1989.  doi: 10.1007/BF01589116.
  * Loshchilov & Hutter (2019) Loshchilov, I. and Hutter, F.  Decoupled weight decay regularization.  In _Int. Conf. on Learning Representations_ , 2019.
  * Papyan (2020) Papyan, V.  Traces of class/cross-class structure pervade deep learning spectra.  _Journal of Machine Learning Research_ , 21(252):1–64, 2020.
  * Park & Wolverton (2020) Park, C. W. and Wolverton, C.  Developing an improved crystal graph convolutional neural network framework for accelerated materials discovery.  _Phys. Rev. Materials_ , 2020.  doi: 10.1103/PhysRevMaterials.4.063801.
  * Pearlmutter (1994) Pearlmutter, B. A.  Fast Exact Multiplication by the Hessian.  _Neural Computation_ , 1994.  doi: 10.1162/neco.1994.6.1.147.
  * Ramakrishnan et al. (2015) Ramakrishnan, R., Hartmann, M., Tapavicza, E., and von Lilienfeld, O. A.  Electronic spectra from tddft and machine learning in chemical space.  _The Journal of Chemical Physics_ , 2015.  doi: 10.1063/1.4928757.
  * Sagun et al. (2017) Sagun, L., Evci, U., Guney, V. U., Dauphin, Y., and Bottou, L.  Empirical analysis of the hessian of over-parametrized neural networks, 2017.  URL https://arxiv.org/abs/1706.04454.
  * Sanyal et al. (2018) Sanyal, S., Balachandran, J., Yadati, N., Kumar, A., Rajagopalan, P., Sanyal, S., and Talukdar, P.  Mt-cgcnn: Integrating crystal graph convolutional neural network with multitask learning for material property prediction, 2018.  URL https://arxiv.org/abs/1811.05660.
  * Satorras et al. (2021) Satorras, V. G., Hoogeboom, E., and Welling, M.  E(n) equivariant graph neural networks.  In _Proceedings of the 38th International Conference on Machine Learning_ , pp. 9323–9332, 18–24 Jul 2021.
  * Sener & Koltun (2018) Sener, O. and Koltun, V.  Multi-task learning as multi-objective optimization.  In _Proc. 32nd Int. Conf. on Neural Information Processing Systems_ , NIPS’18, 2018.
  * Tan et al. (2021) Tan, Z., Li, Y., Shi, W., and Yang, S.  A multitask approach to learn molecular properties.  _J. Chem. Inf. Mod._ , 61(8):3824–3834, 2021\.  doi: 10.1021/acs.jcim.1c00646.  PMID: 34289687.
  * Weininger (1988) Weininger, D.  Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules.  _J. Chem. Inf. Comp. Sci._ , 1988.  doi: 10.1021/ci00057a005.
  * Wu et al. (2018) Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Geniesse, C., et al.  Moleculenet: a benchmark for molecular machine learning.  _Chem. Sci._ , 9:513–530, 2018.  doi: 10.1039/C7SC02664A.
  * Xie & Grossman (2018) Xie, T. and Grossman, J. C.  Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties.  _Phys. Rev. Lett._ , 2018.  doi: 10.1103/PhysRevLett.120.145301.
  * Yang et al. (2021) Yang, A., Su, Y., Wang, Z., Jin, S., Ren, J., et al.  A multi-task deep learning neural network for predicting flammability-related properties from molecular structures.  _Green Chem._ , 2021.  doi: 10.1039/D1GC00331C.
  * Yao et al. (2019) Yao, Z., Gholami, A., Keutzer, K., and Mahoney, M.  Pyhessian: Neural networks through the lens of the hessian, 2019.  URL https://arxiv.org/abs/1912.07145.
  * Yu et al. (2020) Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C.  Gradient surgery for multi-task learning.  In _Adv. Neural Inf. Proc Sys._ , 2020.

##  Appendix A Model specifics

We use $M=5$ message-passing steps. Node, edge, and global features are
projected into a 64-dimensional feature space. $\phi^{V}$ and $\phi^{E}$ have
two layers with 256 units and a skip connection, followed by an output layer
of 64 units. $\phi^{U}$ has two layers with 192 units and a skip connection,
followed by an output layer of 64 units. All networks use $\tanh$ activations.

##  Appendix B Training specifics

Models are trained for 512 epochs using the AdamW Loshchilov & Hutter (2019)
optimizer and the default hyperparameters used in the optax implementation.
The initial rate is set as $10^{-3}$, with an exponential decay rate of 0.997
applied every epoch after the first 256 epochs.

##  Appendix C Data

We summarize our datasets in Table 1.

We scraped Materials Project for crystal records present in it as of October
2020 using the MPRester class from the pymatgen package. The Supplemental
Information contains a table of MP IDs used in this study. The raw entries of
$G_{VRH}$, $K_{VRH}$, and Poisson ratio $\mu$ contained several anomalously
small and large values, so we removed entries with values less than the 5th
percentile or greater than the 95th percentile of obtained elastic properties.
This left us with a total dataset of 10,500 crystals. Summary statistics for
the targets used are given in Table 2. Furthermore, $G_{VRH}$ and $K_{VRH}$
were log-transformed to create a more normal distribution. Roughly 70% of the
dataset (7,400 crystals) was used as training data, and all targets were
standardized prior to training.

As initial node features, we used a one-hot encoding of atom type. For initial
edge features, we used four features calculated by pymatgen: area, face_dist,
solid_angle, and volume. Following similar work in applying graph neural
networks to crystals Xie & Grossman (2018), we discretize these four features
based on deciles.

We use the version of QM8 hosted by MoleculeNet Wu et al. (2018), except that
we drop molecules with negative oscillator energies. Summary statistics for
the targets are given in Table 3. As initial node features, we use one-hot
encodings of atom type. As initial edge features, we use one-hot encodings of
bond type. Roughly 70% of the dataset (15,300 molecules) was used as training
data, and all targets were standardized prior to training.

Table 1: Summary statistics of datasets Dataset | # Data Points | # Targets | # Node Features | # Edge Features
---|---|---|---|---
MP | 10,500 | 3 | 112 | 40
QM8 | 21,725 | 4 | 9 | 4
Table 2: Summary statistics of MP targets |  $\log G_{VRH}$ (log(GPa)) |  $\log K_{VRH}$ (log(GPa)) | $\mu$
---|---|---|---
mean | 3.55068348 | 4.3697952 | 0.30674381
std | 0.73963662 | 0.67564539 | 0.06501633
min | 0.0 | 2.63905733 | 0.18
max | 4.82028157 | 5.45103845 | 0.5
Table 3: Summary statistics of QM8 targets | $E_{1}$ | $E_{2}$ | $f_{1}$ | $f_{2}$
---|---|---|---|---
mean | 0.22010506 | 0.24894142 | 0.02289806 | 0.04326408
std | 0.04382473 | 0.0345149 | 0.05336499 | 0.07312536
min | 0.06956711 | 0.10063996 | 0.0 | 0.0
max | 0.51383669 | 0.51384601 | 0.61611219 | 0.56561272

##  Appendix D Supplemental figures

(a) MP

(b) QM8

Figure 6: We visualize snapshots of property-specific estimated spectral
densities during selected training points. For both MP and QM8 models, most of
the spectral density is concentrated near 0, especially after random
initialization. Occasionally, very high-magnitude eigenvalues, both positive
and negative, spike up for a short period. Our results here elaborate on
Figure 4 – the high and varying curvature of property-specific losses is
driven by singularly large eigenvalues that vary across properties. For these
plots, the $x$-axis gives the range of eigenvalues for the loss function, and
the $y$-axis gives the density $\hat{\psi}(t)$ of eigenvalues concentrated at
that point $t$ (as described in 2.1).

##  Appendix E Example spectrum estimation

We consider a loss function defined by $L(x)=\dfrac{1}{2}x^{T}Ax$, where
$A=\dfrac{1}{2}(B+B^{T})$, for a matrix $B$ with entries sampled $iid$ from a
standard normal distribution. In this case, the Hessian is given by $A$ and so
its estimated eigenvalues can be compared to its true eigenvalues. In Figure
7, we plot the results of a sample calculation, for a $1,000\times 1,000$
matrix. We show that the estimated density has reasonable similarity to the
true distribution, and the estimated trace also matches the true trace.

Figure 7: An example showing the use of the Lanczos algorithm to estimate the
spectral density of a Hessian. The left figure is the distribution of the true
eigenvalues, and the right is the estimated spectral density. The estimated
values accurately characterize the maximal and minimal eigenvalues, and the
interior shows qualitative agreement with the known distribution.


11institutetext: University of Ljubljana, Faculty of Computer and Information
Science
Večna pot 113, Ljubljana, Slovenia
{matej.ulcar<EMAIL_ADDRESS>
# FinEst BERT and CroSloEngual BERT:
less is more in multilingual models

Matej Ulčar  Marko Robnik-Šikonja

###### Abstract

Large pretrained masked language models have become state-of-the-art solutions
for many NLP problems. The research has been mostly focused on English
language, though. While massively multilingual models exist, studies have
shown that monolingual models produce much better results. We train two
trilingual BERT-like models, one for Finnish, Estonian, and English, the other
for Croatian, Slovenian, and English. We evaluate their performance on several
downstream tasks, NER, POS-tagging, and dependency parsing, using the
multilingual BERT and XLM-R as baselines. The newly created FinEst BERT and
CroSloEngual BERT improve the results on all tasks in most monolingual and
cross-lingual situations.

###### Keywords:

contextual embeddings, BERT model, less-resourced languages, NLP

##  1 Introduction

In natural language processing (NLP), a lot of research focuses on numeric
word representations. Static pretrained word embeddings like word2vec [11] are
recently replaced by dynamic, contextual embeddings, such as ELMo [13] and
BERT [3]. These generate a word vector based on the context the word appears
in, mostly using the sentence as the context.

Large pretrained masked language models like BERT [3] and its derivatives
achieve state-of-the-art performance when fine-tuned for specific NLP tasks.
The research into these models has been mostly limited to English and a few
other well-resourced languages, such as Chinese Mandarin, French, German, and
Spanish. However, two massively multilingual masked language models have been
released: a multilingual BERT (mBERT) [3], trained on 104 languages, and newer
even larger XLM-RoBERTa (XLM-R) [2], trained on 100 languages. While both,
mBERT and XLM-R, achieve good results, it has been shown that monolingual
models significantly outperform multilingual models [20, 10]. In our work, we
reduced the number of languages in multilingual models to three, two similar
less-resourced languages from the same language family, and English. The main
reasons for this choice are to better represent each language, and keep
sensible sub-word vocabulary, as shown by Virtanen et al. [20]. We decided
against production of monolingual models, because we are interested in using
the models in multilingual sense and for cross-lingual knowledge transfer. By
including English in each of the two models, we expect to better transfer
existing prediction models from English to involved less-resourced languages.
Additional reason against purely monolingual models for less-resourced
languages is the size of training corpora, i.e. BERT-like models use
transformer architecture which is known to be data hungry.

We thus trained two multilingual BERT models: FinEst BERT was trained on
Finnish, Estonian, and English, while CroSloEngual BERT was trained on
Croatian, Slovenian, and English. In the paper, we present the creation and
evaluation of these models, which required considerable computational
resources, unavailable to most NLP researchers. We make the models which are
valuable resources for the involved less-resourced languages publicly
available111CroSloEngual BERT: http://hdl.handle.net/11356/1317
FinEst BERT: http://urn.fi/urn:nbn:fi:lb-2020061201.

##  2 Training data and preprocessing

BERT models require large quantities of monolingual data. In Section 2.1 we
first describe the corpora used, followed by a short description of their
preprocessing in Section 2.2.

###  2.1 Datasets

We trained two new BERT models from five languages: Finnish, Estonian,
Slovenian, Croatian and English. To obtain high-quality models, we used large
monolingual corpora for each language, some of them unavailable to the general
public. For English, large corpora are readily available and they are much
larger than for other languages. However, high-quality English language models
already exist and English is not the main focus of this research, we therefore
did not use all available English corpora in order to prevent English from
overwhelming the other languages in our models. Some corpora are available
online under permissive licences, others are available only for research
purposes or have limited availability. The corpora used in training are a mix
of news articles and general web crawl, which we preprocessed and
deduplicated. Details about the training set sizes are presented in Table 2,
while their description can be found in works on the involved less-resourced
languages, e,g., [18].

Table 1: The training corpora sizes in number of tokens and the ratios for each language.  Model | CroSloEngual | FinEst
---|---|---
Croatian | 31% | 0%
Slovenian | 23% | 0%
English | 47% | 63%
Estonian | 0% | 13%
Finnish | 0% | 25%
Tokens | $5.9\cdot 10^{9}$ | $3.7\cdot 10^{9}$

Table 2: The sizes of corpora subsets in millions of tokens used to create wordpiece vocabularies.  Language | FinEst | CroSloEngual
---|---|---
Croatian | / | 27
Slovenian | / | 28
English | 157 | 23
Estonian | 75 | /
Finnish | 97 | /

###  2.2 Preprocessing

Before using the corpora, we deduplicated them for each language separately,
using the Onion (ONe Instance ONly) tool222http://corpus.tools/wiki/Onion. We
applied the tool on sentence level for those corpora that did have sentences
shuffled, and on paragraph level for the rest. As parameters, we used 9-grams
with duplicate content threshold of 0.9.

BERT models are trained on subword (wordpiece) tokens. We created a wordpiece
vocabulary using bert-vocab-builder tool333https://github.com/kwonmha/bert-
vocab-builder, which is built upon tensor2tensor library [19]. We did not
process the whole corpora in creating the wordpiece vocabulary, but only a
smaller subset. To balance the language representation in vocabulary, we used
samples from each language. The sizes of corpora subsets are shown in Table 2.
The created wordpiece vocabularies contain 74,986 tokens for FinEst and 49,601
tokens for CroSloEngual model.

##  3 Architecture and training

We trained two BERT multilingual models. FinEst BERT was trained on Finnish,
Estonian, and English corpora, with altogether $3.7$ billion tokens.
CroSloEngual BERT was trained on Croatian, Slovenian, and English corpora with
together $5.9$ billion tokens.

Both models use bert-base architecture [3], which is a 12-layer bidirectional
transformer encoder with the hidden layer size of 768 and altogether 110
million parameters. We used the whole word masking for the masked language
model training task. Both models are cased, i.e. the case information was
preserved. We followed the hyper-parameters settings of Devlin et al. [3],
except for the batch size and total number of steps. We trained the models for
approximately 40 epochs with maximum sequence length of 128 tokens, followed
by approximately 4 epochs with maximum sequence length of 512 tokens. The
exact number of steps was calculated using the expression:

$s=\frac{N_{tok}\cdot E}{b\cdot\lambda}$

, where $s$ is the number of steps the models were trained for, $N_{tok}$ is
the number of tokens in the train corpora, $E$ is the desired number of epochs
(in our case 40 and 4), $b$ is the batch size, and $\lambda$ is the maximum
sequence length.

We trained FinEst BERT on a single Google Cloud TPU v3 for a total of 1.24
million steps where the first 1.13 million steps used the batch size of 1024
and sequence length 128, and the last 113 thousand steps used the batch size
256 and sequence length 512. Similarly, CroSloEngual BERT was trained on a
single Google Cloud TPU v2 for a total of 3.96 million steps, where the first
3.6 million steps used the batch size of 512 and sequence length 128, and the
last 360 thousand steps were trained with the batch size 128 and sequence
length 512. Training took approximately 2 weeks for FinEst BERT and
approximately 3 weeks for CroSloEngual BERT.

##  4 Evaluation

We evaluated the two new BERT models on three downstream evaluation tasks
available for the four involved less-resourced languages: named entity
recognition (NER), part-of-speech tagging (POS), and dependency parsing (DP).
We compared both models with BERT-base-multilingual-cased model (mBERT) on
sensible languages, i.e. FinEst BERT was compared with mBERT on Finnish,
Estonian, and English, while CroSloEngual BERT was compared with mBERT on
Croatian, Slovenian, and English.

###  4.1 Named Entity Recognition

Named entity recognition (NER) task is a sequence labeling task, which tries
to correctly identify and classify each token from an unstructured text into
one of the predefined named entity (NE) classes or, if the token is not part
of a NE, to classify it as not a named entity. Most common named entity
classes are personal names, locations and organizations. We used various
datasets, which do not cover the same set of classes. We therefore adapted the
datasets to allow a more direct comparison between languages, by reducing them
to the four labels they all have in common: PER (person), LOC (location), ORG
(organization), and O (other). All tokens, which are not named entities or
belong to any NE class other than person, location or organization, were
labeled as ’O’.

For Croatian and Slovenian, we used data from hr500k [9] and ssj500k [7],
respectively. Not all sentences in ssj500k are annotated, so we excluded those
that are not annotated. English dataset comes from CoNLL 2013 shared task
[17]. For Finnish we used Finnish News Corpus for NER [15], and for Estonian
dataset we used Nimeüksuste korpus [8]. The statistics of each dataset are
shown in Table 3.

Language | PER | LOC | ORG | Density | N
---|---|---|---|---|---
Croatian | 10241 | 7445 | 11216 | 0.057 | 506457
English | 17050 | 12316 | 14613 | 0.146 | 301418
Estonian | 8490 | 6326 | 6149 | 0.096 | 217272
Finnish | 3402 | 2173 | 11258 | 0.087 | 193742
Slovenian | 4478 | 2460 | 2667 | 0.049 | 194667
Table 3: The number of tokens labeled with each label (PER, LOC, ORG), the
density of these labels (their sum divided by the number of all tokens) and
the number of all tokens (N) for datasets in all languages.

To evaluate the performance of BERT embeddings on the NER task we trained NER
models using Huggingface’s Transformer library, basing the code on their NER
example444https://github.com/huggingface/transformers/tree/master/examples/ner.
We fine-tuned each of our BERT models with an added token classification head
for 3 epochs on the NER data. We compared the results with BERT-base-
multilingual-cased (mBERT) model, which we fine-tuned with exactly the same
parameters on the same data.

Train lang | Test lang | mBERT | CroSloEngual
---|---|---|---
Croatian | Croatian | 0.795 | 0.894
Slovenian | Slovenian | 0.903 | 0.917
English | English | 0.940 | 0.949
Croatian | English | 0.793 | 0.866
English | Croatian | 0.638 | 0.798
Slovenian | English | 0.781 | 0.833
English | Slovenian | 0.736 | 0.843
Croatian | Slovenian | 0.825 | 0.908
Slovenian | Croatian | 0.755 | 0.847
Table 4: The results of NER evaluation task on Croatian, Slovenian, and
English. The scores are average $F_{1}$ scores of the three named entity
classes. A NER model was trained on ”train language” dataset and tested on
”test language” dataset using two different BERT models for all possible
combinations of train and test languages.

We evaluated the models in a monolingual setting (training and testing on the
same language) and a crosslingual setting (training on one language, testing
on another). We present the results as macro average $F_{1}$ scores of the
three NE classes, excluding ’O’ label. Comparison between CroSloEngual BERT
and mBERT is shown in Table 4, comparison between FinEst BERT and mBERT is
shown in Table 5.

The difference in performance of each BERT on English data is negligible. In
other languages, our models outperform the multilingual BERT, the difference
is especially large in Croatian. In crosslingual setting, both FinEst BERT and
CroSloEngual BERT show a significant improvement over mBERT, especially when
one of the two languages is English. This leads us to believe that
multilingual BERT models with fewer languages are more suitable for
crosslingual knowledge transfer.

Train lang | Test lang | mBERT | FinEst
---|---|---|---
Finnish | Finnish | 0.922 | 0.959
Estonian | Estonian | 0.906 | 0.930
English | English | 0.940 | 0.942
Finnish | English | 0.692 | 0.810
English | Finnish | 0.770 | 0.901
Estonian | English | 0.765 | 0.815
English | Estonian | 0.762 | 0.839
Finnish | Estonian | 0.795 | 0.879
Estonian | Finnish | 0.839 | 0.912
Table 5: The results of NER evaluation task on Finnish, Estonian, and English.
The scores are average $F_{1}$ scores of the three named entity classes. A NER
model was trained on ”train language” dataset and tested on ”test language”
dataset using two different BERT models for all possible combinations of train
and test languages.

###  4.2 Part-of-speech tagging and dependency parsing

We evaluated BERT models on two more classification tasks: part-of-speech
(POS) tagging and dependency parsing. In the POS tagging task we attempt to
correctly classify each token within a given set of grammatical categories
(verb, adjective, punctuation, adverb, noun, etc.) Dependency parsing task
attempts to predict the tree structure, representing the syntactic relations
between words in a given sentence.

We trained classifiers on universal dependencies (UD) treebank datasets, using
universal part-of-speech (UPOS) tag set. For Croatian, we used treebank by
Agić and Ljubešić [1]. For English, we used A Gold Standard Dependency Corpus
[16]. For Estonian, we used Estonian Dependency Treebank [12], converted to
UD. Finnish treebank used is based on the Turku Dependency Treebank [5], which
was also converted to UD [14]. Slovenian treebank [4] is based on the ssj500k
corpus [7].

We used Udify tool [6] to train both POS tagger and dependency parsing
classifiers at the same time. We finetuned each BERT model for 80 epochs on
the treebank data. We kept the tool parameters at default values, except for
”warmup_steps” and ”start_step” values, which we changed to equal the number
of training batches in one epoch.

We present the results of POS tagging as UPOS accuracy score in Table 6 and
Table 7. The difference in performance between BERT models is very small on
this task. FinEst and CroSloEngual BERTs perform slightly better than mBERT on
all languages in monolingual setting, except Croatian, where mBERT and
CroSloEngual BERT are equal. The differences are more pronounced in cross-
lingual setting. When training on Slovenian, Finnish or Estonian data and
testing on English data CroSloEngual and FinEst BERT significantly outperform
mBERT. On the other hand, when training on English and testing Croatian, mBERT
outperforms CroSloEngual BERT.

Train lang. | Test lang. | mBERT | CroSloEngual
---|---|---|---
Croatian | Croatian | 0.983 | 0.983
English | English | 0.969 | 0.972
Slovenian | Slovenian | 0.987 | 0.991
English | Croatian | 0.876 | 0.869
English | Slovenian | 0.857 | 0.859
Croatian | English | 0.750 | 0.756
Croatian | Slovenian | 0.917 | 0.934
Slovenian | English | 0.686 | 0.723
Slovenian | Croatian | 0.920 | 0.935
Table 6: The embeddings quality measured on the UPOS tagging task, using UPOS accuracy score for FinEst BERT, CroSloEngual BERT and BERT-base-multilingual-cased (mBERT).  Train lang. | Test lang. | mBERT | FinEst
---|---|---|---
English | English | 0.969 | 0.970
Estonian | Estonian | 0.972 | 0.978
Finnish | Finnish | 0.970 | 0.981
English | Estonian | 0.852 | 0.878
English | Finnish | 0.847 | 0.872
Estonian | English | 0.688 | 0.808
Estonian | Finnish | 0.872 | 0.913
Finnish | English | 0.535 | 0.701
Finnish | Estonian | 0.888 | 0.919
Table 7: The embeddings quality measured on the UPOS tagging task, using UPOS
accuracy score for FinEst BERT, CroSloEngual BERT and BERT-base-multilingual-
cased (mBERT).

We present the results of dependency parsing task as unlabeled attachement
score (UAS) and labeled attachment score (LAS). In monolingual setting
CroSloEngual BERT shows improvement over mBERT on all three languages (Table
8) with the highest improvement on Slovenian and only a marginal improvement
on English. FinEst BERT outperforms mBERT on Estonian and Finnish, with the
biggest margin being on the Finnish data (Table 9). FinEst BERT and mBERT
perform equally on English data.

In crosslingual setting, the results are similar to those seen on the POS
tagging task. Major improvements of FinEst BERT and CroSloEngual BERT over
mBERT in English-Estonian, English-Finnish and English-Slovenian pairs, minor
improvements in Estonian-Finnish and Croatian-Slovenian pairs. Again, mBERT
outperformed CroSloEngual BERT when dependency parser was trained on English
data and tested on Croatian data.

Train | Test | mBERT | CroSloEngual
---|---|---|---
language | language | UAS | LAS | UAS | LAS
Croatian | Croatian | 0.930 | 0.891 | 0.940 | 0.903
English | English | 0.917 | 0.894 | 0.922 | 0.899
Slovenian | Slovenian | 0.938 | 0.922 | 0.957 | 0.947
English | Croatian | 0.824 | 0.724 | 0.822 | 0.725
English | Slovenian | 0.830 | 0.719 | 0.848 | 0.736
Croatian | English | 0.759 | 0.627 | 0.782 | 0.657
Croatian | Slovenian | 0.880 | 0.802 | 0.912 | 0.840
Slovenian | English | 0.741 | 0.578 | 0.794 | 0.648
Slovenian | Croatian | 0.861 | 0.773 | 0.891 | 0.810
Table 8: The embeddings quality measured on the dependency parsing task. Results are given as UAS and LAS for CroSloEngual BERT and BERT-base-multilingual-cased (mBERT). Train | Test | mBERT | FinEst
---|---|---|---
language | language | UAS | LAS | UAS | LAS
English | English | 0.917 | 0.894 | 0.918 | 0.895
Estonian | Estonian | 0.880 | 0.848 | 0.909 | 0.882
Finnish | Finnish | 0.898 | 0.867 | 0.933 | 0.915
English | Estonian | 0.697 | 0.531 | 0.768 | 0.591
English | Finnish | 0.706 | 0.561 | 0.781 | 0.624
Estonian | English | 0.633 | 0.492 | 0.726 | 0.567
Estonian | Finnish | 0.784 | 0.695 | 0.864 | 0.801
Finnish | English | 0.543 | 0.433 | 0.684 | 0.558
Finnish | Estonian | 0.782 | 0.691 | 0.852 | 0.778
Table 9: The embeddings quality measured on the dependency parsing task.
Results are given as UAS and LAS for FinEst BERT and BERT-base-multilingual-
cased (mBERT).

##  5 Conclusion

We built two large pretrained trilingual BERT-based masked language models,
Croatian-Slovenian-English and Finnish-Estonian-English. We showed that the
new CroSloEngual and FinEst BERTs perform substantially better than massively
multilingual mBERT on the NER task in both monolingual and cross-lingual
setting. The results on POS tagging and DP tasks show considerable improvement
of the proposed models for several monolingual and cross-lingual pairs, while
they are never worse than mBERT.

In future, we plan to investigate different combinations and proportions of
less-resourced languages in creation of pretrained BERT-like models, and use
the newly trained BERT models on the problems of news media industry.

## Acknowledgments

The work was partially supported by the Slovenian Research Agency (ARRS) core
research programme P6-0411. This paper is supported by European Union’s
Horizon 2020 research and innovation programme under grant agreement No
825153, project EMBEDDIA (Cross-Lingual Embeddings for Less-Represented
Languages in European News Media). Research was supported with Cloud TPUs from
Google’s TensorFlow Research Cloud (TFRC).

## References

  * Agić and Ljubešić [2015] Željko Agić and Nikola Ljubešić.  Universal dependencies for Croatian (that work for Serbian, too).  In _The 5th Workshop on Balto-Slavic Natural Language Processing_ , pages 1–8, 2015.
  * Conneau et al. [2019] Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov.  Unsupervised cross-lingual representation learning at scale.  _arXiv preprint arXiv:1911.02116_ , 2019.
  * Devlin et al. [2018] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dobrovoljc et al. [2017] Kaja Dobrovoljc, Tomaž Erjavec, and Simon Krek.  The universal dependencies treebank for Slovenian.  In _Proceeding of the 6th Workshop on Balto-Slavic Natural Language Processing (BSNLP 2017)_ , 2017.
  * Haverinen et al. [2013] K. Haverinen, J. Nyblom, T. Viljanen, V. Laippala, S. Kohonen, A. Missilä, S. Ojala, T. Salakoski, and F. Ginter.  Building the essential resources for Finnish: the Turku dependency treebank.  _LREC_ , 2013.
  * Kondratyuk and Straka [2019] Dan Kondratyuk and Milan Straka.  75 languages, 1 model: Parsing universal dependencies universally.  In _Proceedings of the 2019 EMNLP-IJCNLP_ , pages 2779–2795, 2019\.
  * Krek et al. [2019] Simon Krek, Kaja Dobrovoljc, Tomaž Erjavec, Sara Može, Nina Ledinek, Nanika Holz, Katja Zupan, Polona Gantar, Taja Kuzman, Jaka Čibej, Špela Arhar Holdt, Teja Kavčič, Iza Škrjanec, Dafne Marko, Lucija Jezeršek, and Anja Zajc.  Training corpus ssj500k 2.2, 2019.  Slovenian language resource repository CLARIN.SI.
  * Laur [2013] Sven Laur.  Nimeüksuste korpus.  Center of Estonian Language Resources, 2013.
  * Ljubešić et al. [2016] Nikola Ljubešić, Filip Klubička, Željko Agić, and Ivo-Pavao Jazbec.  New inflectional lexicons and training corpora for improved morphosyntactic annotation of Croatian and Serbian.  In _Proceedings of the LREC 2016_ , 2016.
  * Martin et al. [2019] Louis Martin, Benjamin Muller, Pedro Javier Ortiz Suárez, Yoann Dupont, Laurent Romary, Éric Villemonte de la Clergerie, Djamé Seddah, and Benoît Sagot.  CamemBERT: a tasty French language model.  _arXiv preprint arXiv:1911.03894_ , 2019.
  * Mikolov et al. [2013] Tomas Mikolov, Quoc V Le, and Ilya Sutskever.  Exploiting similarities among languages for machine translation.  _arXiv preprint 1309.4168_ , 2013.
  * Muischnek et al. [2016] Kadri Muischnek, Kaili Müürisep, and Tiina Puolakainen.  Estonian Dependency Treebank: from Constraint Grammar tagset to Universal Dependencies.  In _Proceedings of LREC 2016_ , 2016.
  * Peters et al. [2018] Matthew E Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer.  Deep contextualized word representations.  _arXiv preprint arXiv:1802.05365_ , 2018.
  * Pyysalo et al. [2015] Sampo Pyysalo, Jenna Kanerva, Anna Missilä, Veronika Laippala, and Filip Ginter.  Universal dependencies for Finnish.  In _Proceedings of NoDaLiDa 2015_ , 2015.
  * Ruokolainen et al. [2020] Teemu Ruokolainen, Pekka Kauppinen, Miikka Silfverberg, and Krister Lindén.  A Finnish news corpus for named entity recognition.  _Lang Resources & Evaluation_, 54(1):247–272, 2020.
  * Silveira et al. [2014] Natalia Silveira, Timothy Dozat, Marie-Catherine de Marneffe, Samuel Bowman, Miriam Connor, John Bauer, and Christopher D. Manning.  A gold standard dependency corpus for English.  In _Proceedings of LREC-2014_ , 2014.
  * Tjong Kim Sang and De Meulder [2003] Erik F. Tjong Kim Sang and Fien De Meulder.  Introduction to the CoNLL-2003 shared task: Language-independent named entity recognition.  In Walter Daelemans and Miles Osborne, editors, _Proceedings of CoNLL-2003_ , pages 142–147. Edmonton, Canada, 2003.
  * Ulčar and Robnik-Šikonja [2020] Matej Ulčar and Marko Robnik-Šikonja.  High quality elmo embeddings for seven less-resourced languages.  In _Proceedings of The 12th Language Resources and Evaluation Conference_ , pages 4731–4738, Marseille, France, May 2020. European Language Resources Association.
  * Vaswani et al. [2018] Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, et al.  Tensor2tensor for neural machine translation.  In _Proceedings of the AMT_ , pages 193–199, 2018.
  * Virtanen et al. [2019] Antti Virtanen, Jenna Kanerva, Rami Ilo, Jouni Luoma, Juhani Luotolahti, Tapio Salakoski, Filip Ginter, and Sampo Pyysalo.  Multilingual is not enough: BERT for Finnish.  _arXiv preprint arXiv:1912.07076_ , 2019.