# LCAONet: Message-passing with physically optimized atomic basis functions

Kento Nishio
School of Engineering
The University of Tokyo
Tokyo, 113-8656
<EMAIL_ADDRESS>
Kiyou Shibata, Teruyasu Mizoguchi
Institute of Industrial Science
The University of Tokyo
Tokyo, 153-8505
<EMAIL_ADDRESS>
http://www.edge.iis.u-tokyo.ac.jp

###### Abstract

A Model capable of handling various elemental species and substances is
essential for discovering new materials in the vast phase and compound space.
Message-passing neural networks (MPNNs) are promising as such models, in which
various vector operations model the atomic interaction with its neighbors.
However, conventional MPNNs tend to overlook the importance of physicochemical
information for each node atom, relying solely on the geometric features of
the material graph. We propose the new three-body MPNN architecture with a
message-passing layer that utilizes optimized basis functions based on the
electronic structure of the node elemental species. This enables conveying the
message that includes physical information and better represents the
interaction for each elemental species. Inspired by the LCAO (linear
combination of atomic orbitals) method, a classical method for calculating the
orbital interactions, the linear combination of atomic basis constructed based
on the wave function of hydrogen-like atoms is used for the present message-
passing. Our model achieved higher prediction accuracy with smaller parameters
than the state-of-the-art models on the challenging crystalline dataset
containing many elemental species, including heavy elements. Ablation studies
have also shown that the proposed method is effective in improving accuracy.
Our implementation is available online.

##  1 Introduction

Machine learning (ML) has recently attracted attention in computational
materials as an alternative to the high computational cost of first-principles
methods such as density functional theory (DFT). By using ML method, we can
accurately simulate large systems for which first-principles calculations are
not applicable, and thus gain insight that is closer to real materials that
have vacancies, substitution sites, and so on. Many ML methods have been
developed using features (descriptors) that encode geometrical, topological,
physical, and chemical information about the local environment of the atoms.
The most important aspect of descriptor construction is its uniqueness in
representing a material [1, 2]. In addition to the topology, physical and
chemical information such as orbital-orbital interactions, atomic charges, and
electric dipoles are also crucial for the functional manifestation of a
material. Therefore, developing descriptors that adequately represent these
properties is essential for accurate modeling, but it is not easy [3]. The
development and selection of descriptors require a lot of manual coordination
and years of experience and background knowledge about the material science.

Against this background, a number of deep learning methods based on Message-
Passing neural networks (MPNNs), which can extract the local environment of
atoms without requiring manual adjustments through message-passing
architecture, have been developed and shown promising results in predicting
quantum mechanical properties of molecular and crystalline systems. In the
message-passing layer, geometric information of a material, such as
interatomic distances and bond angles, are converted into vectors with
physically motivated basis functions, and the graph embedding vectors are
updated by vector operations. In most cases, the functional form of the basis
functions for vectorizing the geometry is the same regardless of the node
elemental species. However, in actual materials, even if the structure is the
same, the atomic interaction should be different depending on the type of
neighboring atoms. In addition, such message-passing methods can capture
geometric information but not physical and chemical properties of elemental
species. Therefore, extracting physically rich representations such as
descriptors through message-passing layers is currently difficult.

In this work, we propose an improved method of three-body MPNN that optimize
the form and the number of basis functions according to the electronic
structure of node atoms. Based on a new set of basis functions developed by
the physical formulation of atomic interactions, we develop a physically
motivated weighted message-passing that can propagate interactions between
orbitals inspired by the LCAO-MO method, a classical computational method for
modeling orbital interactions (LCAONet). We also develop a method to generate
chemically richer initial node vectors by embedding the number of electrons in
the ground-state occupied orbitals in addition to the atomic number. The
proposed improved methods are investigated by ablation studies, and it is
shown that each method contributes significantly to the accuracy improvements.
Modeling that incorporates physicochemically rich representations such as
descriptors without human intervention via a message-passing architecture
outperforms existing ML methods. It is suggested that in order to develop a
universal model that can handle a variety of elemental species, it may be
useful to include more inductive bias in the model according to node atoms.

##  2 Related work

#### Descriptor based models.

ML methods for predicting material properties have traditionally used
descriptors, which contain geometrical, physical and chemical information.
SOAP [4, 5], for example, encodes the local environment by expanding the
atomic density based on the spherical harmonics and radial basis functions.
OFM [6] encodes the interaction between valence orbitals of the neighboring
atoms based on the electron configuration. A concept common to many of these
descriptors is the encoding of information about the system based on
"nearsightedness" [7]. After atomic numbers and atomic positions of the system
is transformed using manually adjusted descriptors, it is input into the
Gaussian process [8, 9, 10] or neural networks [11, 12, 13, 14, 15]. Although
descriptors have a sophisticated expression, their hyperparameters need to be
well tuned to the system to develop an accurate model.

#### MPNNs for molecules.

Graph neural networks (GNNs) were first proposed in the 90s [16, 17] and 00s
[18, 19, 20] and then applied to molecular graphs as message-passing neural
networks [21, 22]. When the GNN is applied to material graphs, the geometric
features of the graph are vectorized by physically motivated basis functions,
which are used to model interactions with neighboring atoms. The geometric
features such as the distance between two nodes [23, 24, 25, 26], the angle
between two edges consisting of three nodes [27, 28, 29], and also the
dihedral angle generated between four nodes [30, 31] are used. In this study,
from the viewpoint of computational cost, message-passing was performed
including the three-body interaction. Message-passing with rotational
equivariant features has also been developed to efficiently incorporate
higher-order geometric features [32, 33, 34, 35]. Although these studies have
demonstrated excellent accuracy for molecular and crystalline dataset, the
basis function form is constant regardless of the node elemental species. In
this study, we provide a new physically motivated message-passing framework to
select basis functions based on the electronic structure of the node atom.

#### MPNNs with physicochemical information.

There have been efforts to incorporate physical information such as orbital
interactions into MPNNs. In OGCNN [36], OFM is incorporated during the
generation of graph embeddings, and message-passing incorporates orbital
interactions into the node embedding vectors. However, OGCNN only deals with
valence orbital information, and this study incorporates richer orbital
information, including inner-shell orbitals, and higher-order geometric
features up to three-body interactions. In OrbNet [37], orbital features are
generated in advance by the semi-empirical DFT calculation, and the graph-
attention model is constructed using pre-computed features. In this study,
since the goal is to learn the topology and physical interactions from the
graph only, no additional computed numerical features were added to the model.
However, it is possible to add such precomputed features to this model. In
SpookyNet [38], the spin state of the molecular system and all electrons are
added as inputs and long-range interactions are taken into account in the
message-passing layer to achieve high accuracy for molecular systems. In this
study, following SpooKyNet, the electronic structure information of each atom
is used as the embedding. Furthermore, the difference in messages depending on
the spin state is included in the basis function.

##  3 LCAO-MO method

The Linear Combination of Atomic Orbital - Molecular Orbital (LCAO-MO) method
is a calculation approach that approximates the molecular orbital or, more
broadly, the electronic state of a system as a linear combination of atomic
orbitals that are strongly bounded for each atom. Due to its simplicity, the
LCAO-MO method has long been used in the analysis of many molecular systems
[39].

Here, we apply the LCAO-MO method to the simplest system, hydrogen molecular
ion ($\mathrm{H_{2}^{+}}$), and outline the methodology. Hydrogen molecular
ion is the simplest system consisting of two atomic nucleus with one nucleon
number and one electron. Consider approximating the molecular orbitals (MO) of
the entire system by the LCAO-MO method. Note that in hydrogen molecular ion,
the Schrödinger equation can be solved annalytically under the Born-
Oppenheimer approximation [40] but here we obtain an approximate solution to
understand the LCAO-MO method.

Let $\phi$ be the 1s orbital of the one-electron wave function bounded around
each nucleus, then the MO $\psi$ of $\mathrm{H_{2}^{+}}$ approximated by the
LCAO-MO is as follows:

$\psi=c_{1}\phi_{A}+c_{2}\phi_{B}.$ (1)

Each coefficient can be obtained from the variational principle; Determining
the coefficients so that the energy $E$ obtained from $\psi$ and the
Schrödinger equation $H\psi=E\psi$ is minimized. Energy is expressed using MO
as follows:

$E=\frac{\Braket{\psi}{H}{\psi}}{\Braket{\psi}{\psi}}=\frac{\sum_{i}\sum_{j}c_{i}^{*}c_{j}H_{ij}}{\sum_{i}\sum_{j}c_{i}^{*}c_{j}S_{ij}},$
(2)

where $S_{ij}=\Braket{\psi_{i}}{\psi_{j}}$,
$H_{ij}=\Braket{\psi_{i}}{H}{\psi_{j}}$, and $c_{i}$ is determined from the
minimum energy condition $\frac{\partial E}{\partial c_{i}}=0$. Determine the
coefficients by solving the equations $(\bm{H}-E\bm{S})\bm{c}=0$ that is
derived from minimization conditions. Then the coefficients are determined, a
relationship, $c_{2}/c_{1}=\pm 1$, is derived between the two coefficients.
This means that the contribution of each nucleus is equal, which is a
physically correct depiction. The specific values of the coefficients can be
obtained from the MO normalization conditions.

Thus, the LCAO-MO method is a very simple yet powerful method for gaining
physical insight.

##  4 Optimized basis functions with node elemental species

Figure 1: The message-passing scheme of this study, and the radial basis
functions developed from the Schrödinger equation for particles moving in the
potential of the hydrogen-like atom. The radial basis function has more nodes
as the value of $l$ increases (more nodes in $p$ and $d$ orbitals than in $s$
orbitals). The figure on the left shows how the number and type of basis
functions are restricted according to elemental species. When considering
message-passing for the central atom $i$, the hydrogen atom uses only the $1s$
orbital for message-passing. On the other hand, the Na atom uses the four
basis functions $1s,2s,2p$ and $3s$. They all correspond to occupied orbitals
in the ground state.

#### Graph neural networks for materials.

Material graphs can be uniquely represented by only pairs of atomic numbers
$\bm{Z}=\\{z_{1},\ldots,z_{i},\ldots,z_{N}\\}$ and atomic coordinates
$\bm{X}=\\{\bm{x}_{1},\ldots,\bm{x}_{i},\ldots,\bm{x}_{N}\\}$. $N$ represents
the number of atoms constituting the system, and in the case of a periodic
system, the number of atoms in the supercell is considered. GNNs input atomic
numbers and atomic coordinates to predict quantum mechanical properties such
as formation energy, highest occupied molecular orbital (HOMO) energy, dipole
moment and so on. In this study, we predict the scalar regression target
$t\in\mathbb{R}$, i.e., our GNN model is defined using the parameters $\theta$
via $f_{\theta}:\\{\bm{Z},\bm{X}\\}\rightarrow t$ as mentioned in DimeNet
[27].

#### Geometric representation.

To represent the interaction between nodes, we use the pairwise interatomic
distance $d_{ij}=||\bm{x}_{i}-\bm{x}_{j}||_{2}$, and the three-body angle
$\alpha_{(kj,ji)}=\angle\bm{x}_{k}\bm{x}_{j}\bm{x}_{i}$ as geometric features.
Then, a geometric feature vector $\bm{e}_{RBF}^{(ij)},\bm{e}_{SBF}^{(ijk)}$ is
created using basis functions and used for message-passing. As the radial
basis function for vectorizing $\bm{e}_{RBF}^{(ij)}$, previous studies have
used a set of Gaussian type radial basis functions [23, 24, 25, 26], but a set
of spherical Bessel functions has been developed as an orthogonal basis more
suitable for quantum systems and incorporated into many models [27, 29, 34].
In this study, we propose a radial basis function that can more accurately
represent physical interactions based on the same spirit as the orthogonal
basis proposed in DimeNet [27]. Furthermore, by varying the number and type of
functions used depending on the electronic structure of the node elemental
species, we improve parameter efficiency and bring useful inductive bias to
the model. As the angular basis function for $\bm{e}_{SBF}^{(ijk)}$,
normalized spherical harmonic function are used to vectorize as in the
previous studies [27]. This follows naturally from the physical formulation
that expresses the elemental interaction.

#### Physically motivated basis function.

To construct the basis function based on the physical formulation, we consider
a one-electron wave function bounded by a spherically symmetric potential, as
in DimeNet. In view of the fact that the model is to be constructed to
reproduce the DFT results that calculate the electronic state using the
electron density $\rho(\bm{r})=||\Psi(\bm{r})||^{2}$, it is desirable to
construct the basis function based on the electron wave function
$\Psi(\bm{r})$. The electron wave function is defined by the time-independent
Schrödinger equation
$\left(-\frac{\hbar^{2}}{2m}\nabla^{2}+V(\bm{r})\right)\Psi(\bm{r})=E\Psi(\bm{r})$,
with constant mass $m$ and energy $E$. Since it is difficult to know in
advance the potential $V(\bm{r})$, DimeNet derived the electron wave function
by assuming a spherically symmetric potential that is 0 inside the cutoff
sphere and $\infty$ outside the cutoff sphere. The spherical Bessel function
derived by the variable separation method was then used as the radial basis
function, and the spherical harmonic function was used as the angular basis
function. However, the actual electrons are thought to be bound by the
potential from the nucleus. In other words, it is more natural to assume that
the Coulomb potential
$V(\bm{r})=-\frac{z_{i}e^{2}}{4\pi\epsilon_{0}}\frac{1}{||\bm{r}||}$ is
working against the electrons in the cutoff sphere. The form of the one-
electron wave function in motion under the Coulomb potential of a nucleus is
well known and can be obtained by the variable separation method in a polar
coordinates $(r,\theta,\phi)$ [41].

$\Psi_{nlm}(r,\theta,\phi)=R_{nl}(r)Y_{m}^{l}(\theta,\phi),$ (3)

where $n$, $l$, and $m$ are the principal, azimuthal, and magnetic quantum
number, respectively. Each quantum number is an integer satisfying the
condition, $n=1,2,3\dots$, $l=0,1,2\dots,n-1$ and $||m||\leq l$, that
determine the shape and orientation of the orbit. For example, $n=1$, $l=0$
corresponds to the 1s orbital, while $n=2$, $l=1$ corresponds to the $2p$
orbital. $Y_{m}^{l}$ is the normalized spherical harmonic function, a
spherically symmetric function, and $R_{nl}$ is a radial wave function with

$R_{nl}(r)=C_{nl}\left(\frac{2z_{i}}{na}\right)\left(\frac{2z_{i}r}{na}\right)^{l}\mathrm{exp}\left(-\frac{z_{i}r}{na}\right)L_{n+l}^{2l+1}\left(\frac{2z_{i}r}{na}\right),$
(4)

where $L_{n+l}^{2l+1}$ is associated Laguerre polynomials and $C_{nl}$ is the
normalization factor, both functions of $n$ and $l$. In this study, the
$C_{nl}$ were determined by normalizing within the cutoff radius for
crystalline systems and by normalizing over the whole space for molecular
systems. $r$ represents the distance from the center of the orbit i.e. atomic
nucleus. As described above, the spherical harmonic and the radial wave
functions are derived from the Schrödinger equation of a particle moving
around an hydrogen-like atomic potential, and are used as basis functions for
vectorizing the graph. In the next section, we will discuss how to optimize
the basis functions for each node atom.

#### Optimization of basis functions according to elemental species.

The number and type of basis functions are then optimized for each elemental
species by using as the basis function only those radial wave functions
corresponding to the orbitals occupied by the ground state, i.e., the orbitals
in which electrons are present in the ground state (see Figure 1). For
example, for a carbon atom (C), the ground state electron configuration is
$(1s)^{2}(2s)^{2}(2p)^{2}$. Message-passing is performed using only the basis
corresponding to $1s,2s,2p$ and no other basis (e.g., $n=3$ or more). Note
that such optimization restricts the basis function of, for instance, a
hydrogen atom to only the $1s$ orbital, so that only messages that decay from
the neighborhood are propagated. In actual molecules and crystals, the spread
of interactions is expected to be different for each elemental species, so the
optimization of basis functions can more clearly represent the message
differences among elemental species. Since information on occupied orbitals is
explicitly included in the basis, it is expected that information on the
element’s electronic structure and inter-orbital interactions can be included
in the message passing.

##  5 LCAONet

Figure 2:  The LCAONet architecture. $\Box$ denotes the layer’s input,
$\parallel$ denotes concatenation, $\cdot$ denotes vector-scalar
multiplication and $\odot$ denotes element-wise multiplication. The edge
distance $d_{ij}$ and $d_{kj}$ is vectorized by the radial wave function of
the hydrogen atom, and the three-body angle $\alpha_{(kj,ji)}$ is vectorized
by the spherical harmonic function. In the embedding block, an initial
embedding vector $\bm{h}^{(0)}_{i}$ is generated based on the atomic number
and electronic structure information of the node atom $i$. In multiple
interaction blocks, the message is propagated by the neighboring embeddings
$\bm{h}_{j}$, the basis $\bm{e}_{\mathrm{RBF}}^{(ij)}$ optimized for each
elemental species, and the basis coefficient vector $\bm{C}_{ji}^{nl}$, which
updates the node embedding $\bm{h}_{i}$. In the output block, the embedding
vectors are transformed to calculate node-wise scalar values, and finally they
are summed up to output the predictions.

###  5.1 Embedding block

#### Electron embedding in the ground state.

In this study, the number of electrons in the ground state is embedded in
addition to the atomic number. The number of electrons in each orbital is
represented by a learnable, randomly initialized electron number weight
vectors $\bm{a}^{nl}_{e_{z}}\in\mathbb{R}^{F}$ shared across the materials.
$n,l$ is the quantum number of the orbital and $e^{nl}$ is the number of
electrons in the orbital of quantum number $n,l$ in the ground state, meaning
that a different embedding is generated for each orbital. (For example, the
embedding corresponding to the case where the 1s orbital has one occupied
electron number is different from the embedding corresponding to the case
where the 2s orbital has one occupied electron number.)

#### Embedding node and coefficient vectors.

Atomic embedding $\bm{h}_{i}$ are initialized for each node, and the
coefficient vectors $\bm{C}_{ji}^{nl}$ for adding the atomic bases are
initialized for each edge. These embedding vectors are initialized using the
atomic number $z$ and the number of electrons of each orbital $e^{nl}$ as
described above. Atomic numbers are similarly embedded with a weight vector
$\bm{a}_{z}\in\mathbb{R}^{F}$, which is learnable and randomly initialized.
This weight vector is also shared among materials. We generate the initial
node embedding vector $\bm{h}^{(0)}_{i}\in\mathbb{R}^{F_{n}}$ by combining
orbital electron weights $\bm{a}_{e_{z}}^{nl}$ with atomic number weights
$\bm{a}_{z}$ via

$\bm{h}^{(0)}_{i}=f^{\mathrm{enc}}\left(\bm{a}_{z_{i}}\parallel\sum^{nl}{\bm{a}^{nl}_{e_{z_{i}}}}\right),$
(5)

where $\parallel$ is the concatenation and $f^{\mathrm{enc}}$ is the encoding
function. In this study, a two-layer neural network was used to enable
learning.

Similarly, the coefficient vector $C^{nl}_{ji}\in\mathbb{R}^{F}$ is also
initialized using orbital electron weights
$\bm{\tilde{a}}_{e_{z}}^{nl}\in\mathbb{R}^{F}$ and atomic number weights
$\bm{\tilde{a}}_{z}\in\mathbb{R}^{F}$. Here, the tilde on the atomic number
and electron number weights indicates that a different parameter is used than
the weights used for the node embedding. In this study, a linear combination
of each orbit is used for messages from the neighborhood. $C^{nl}_{ji}$ is
used as the coefficient in the linear combination, but we want the coefficient
vector of the orbit corresponding to the $n,l$ of the unoccupied orbit to be a
zero vector. To satisfy this conditions, initialize the coefficient vectors in
three steps as follows. First, the electron number weights
$\bm{\tilde{a}}_{e_{z}}^{nl}$ are initialized with a zero vector when the
number of the electron $e$ is zero (unoccupied orbital). Even when learning
parameters by back-propagation, the atomic basis used can be restricted,
depending on the elemental species, as always a zero vector by eliminating
gradient propagation. Second, the weight of the orbital electron number is
transformed by a certain function $f^{\mathrm{e}}$ to make
$\bm{x}_{e_{z}}^{nl}\in\mathbb{R}^{F}$ i.e.
$\bm{x}_{e_{z}}^{nl}=f^{\mathrm{e}}(\bm{\tilde{a}}_{e_{z}}^{nl})$. In this
process, the embedding $\bm{x}_{e_{z}}^{nl}$ is generated so that the vector
of non-occupied orbitals becomes a zero vector. To achieve this,
$f^{\mathrm{e}}$ is a two-layer neural network without a bias term. Finally,
using an electron number embedding $\bm{x}_{e_{z}}^{nl}$ where only the
vectors corresponding to the $(n,l)$ of the occupied orbitals have values, the
coefficient vectors are initialized via

$\bm{C}^{nl}_{ji}=\left\\{\bm{x}_{e_{z_{j}}}^{nl}+\bm{x}_{e_{z_{j}}}^{nl}\odot
f^{\mathrm{z}}\left(\bm{\tilde{a}}_{z_{i}}\parallel\bm{\tilde{a}}_{z_{j}}\right)\right\\},$
(6)

where $f^{\mathrm{z}}$ is a function that converts concatenated atomic number
weights with learnable parameters. $\odot$ represents the element-wise
multiplication of the vectors, and the atomic number embedding is multiplied
by the electron embedding $\bm{x}_{e_{z_{j}}}^{nl}$ to keep the unoccupied
orbitals of neighboring $j$ atom in the zero vector. $e_{z_{j}}$ represents
the number of electrons in the neighboring atoms. The reason for using the
number of electrons in the $j$ atom to embed is that the information on the
electronic structure and orbitals of the neighboring atoms is aggregated in
the central atom $i$. That is, $\bm{C}^{nl}_{ji}$ is modeled as how the $nl$
orbitals of the neighboring atom $j$ interact with the central atom $i$.

###  5.2 Interaction block

#### Message with electronic structure and orbital information.

As described in 4, the hydrogen-like wave functions, developed to more
accurately model the atomic interactions, are used as atomic basis functions
to vectorize the geometric features of the graph. Then, inspired by the LCAO
method, these linear combinations are used for message-passing, which includes
electronic structure and orbital information. Furthermore, by initializing the
coefficient vectors of the linear combination as described in 5.1, only the
basis corresponding to the occupied orbitals is used for message-passing, and
the number and type of basis functions are optimized according to the
neighboring node elemental species. The goal is thereby to capture inductive
bias in the message according to atomic species.

#### Three body orbital interaction.

First, directional message-passing[27] is performed based on the angle
information consisting of the three-body nodes. As in previous studies [27,
28], the geometric information is vectorized by a spherical harmonic basis
with $m=0$. First, we create a linear combination of orbitals between the
three-body to create the orbital message of $\mathcal{L}$ th layer
$\Psi_{ijk}^{\mathcal{L}}\in\mathbb{R}^{C}$ via

$\Psi_{ijk}^{(\mathcal{L})}=\sum_{nl}R_{nl}(d_{kj})Y_{l}^{0}\left(\alpha_{(kj,ji)}\right)\cdot
f_{kj}^{\mathcal{L}}\left(\bm{C}^{nl}_{kj}\right),$ (7)

where $\cdot$ denotes vector-scalar multiplication and $f_{kj}^{\mathcal{L}}$
is a two-layer neural network, and as in 5.1, the bias term is eliminated for
the embedding of unoccupied orbitals is kept to $0$ vector. Note that only the
coefficient vectors of the occupied orbitals have non-zero values, so the
basis used depends on the electronic structure of the $k$ atom. For example,
if the $k$ atom has only 1s orbitals, then only the basis corresponding to
$n=1,l=0$ is used for the message. This allows the message to include an
appropriate spatial spread for each atom. Since the three-body directional
message-passing expresses the relative positional relationship between edges,
in this study, this message aggregated with a coefficient vector that
expresses the interaction of the edges:

$\tilde{\bm{C}}_{ji}^{nl}=\bm{C}_{ji}^{nl}+\bm{C}_{ji}^{nl}\odot
f_{\mathrm{three}}^{\mathcal{L}}\left(\sum_{k\in\mathcal{N}_{j}\setminus
i}\sigma\left(\bm{W}_{k}\bm{h}^{(\mathcal{L})}_{k}+\bm{b}_{k}\right)\odot\Psi_{ijk}^{(\mathcal{L})}\right).$
(8)

Again, an element-wise product is taken with the three-body message and
coefficient vectors and a residual network is used so that the unoccupied
orbital coefficient vector is kept to 0 vector.

#### Edge orbital interaction.

For nodes between two-body, the linear combination of the radial basis
functions $\Psi_{ji}^{(\mathcal{L})}=\sum_{nl}R_{nl}(d_{ij})\cdot
f_{ji}^{\mathcal{L}}\left(\bm{C}^{nl}_{ji}\right)\in\mathbb{R}^{C}$ is also
used for message-passing. We update node embedding with orbital messages
containing electronic structure information via

$\bm{h}^{(\mathcal{L}+1)}_{i}=\bm{h}^{(\mathcal{L})}_{i}+\bm{W}_{\mathrm{out}}\left[\sum_{j\in\mathcal{N}_{i}}f_{\mathrm{node}}^{\mathcal{L}}\left\\{\left(\bm{W}_{i}\bm{h}^{(\mathcal{L})}_{i}+\bm{b}_{i}\right)\parallel\left(\bm{W}_{j}\bm{h}^{(\mathcal{L})}_{j}+\bm{b}_{j}\right)\right\\}\circ\bm{W}_{\Psi}\Psi_{ji}^{(\mathcal{L})}\right].$
(9)

###  5.3 Output block

The node embedding vector $\bm{h}_{i}^{(L)}$, updated through the $L$-layer
interaction blocks, is finally passed to the output block. In the output
block, the node embedding is transformed by a neural network to obtain a
atomic scalar value $t_{i}$. These outputs are summed or averaged across the
graph to obtain the final predicted scalar value $t$, depending on whether the
physical property value is extensive or intensive.

##  6 Experiments and results

Table 1: MAE on Perovskite for the test set. The results are averaged over five randomly created splits. SchNet was trained using SchNetPack [42], and M3GNet was also trained using publicly available code [43]. All models were trained from scratch. Ablation studies were also performed on LCAONet to verify the use of spherical Bessel functions developed on DimeNet as radial basis functions, the use of the same basis functions for all elemental species without optimizing the atomic basis functions, and the use of no electron configuration information for node embedding. When the interaction layers are aligned to 4 layers and the graph embedding dimension to 128 dimensions, the number of parameters is 1,147,174 for M3GNet, while 869,185 for LCAONet, achieving the better accuracy with about 75% of the number of parameters.  model | MAE (meV/atom)
---|---
OGCNN result from [36] | 50
SchNet | 51.8 $\pm$ 1.2
M3GNet | 38.5 $\pm$ 0.5
LCAONet |  37.5 $\pm$ 1.3
LCAONet with spherical Bessel function as $\bm{e}_{RBF}^{(ij)}$ | 41.3 $\pm$ 1.3
LCAONet No optimized basis | 37.9 $\pm$ 1.6
LCAONet No electronic info into $\bm{h}^{(0)}$ | 38.9 $\pm$ 2.0

The accuracy of the model was verified for the periodic systems with a various
kind of elemental species. The model hyperparameters and training settings see
the supplemental material.

#### Models.

For comparison, we used the state-of-the-art model for the crystalline dataset
such as SchNet[44], M3GNet[29] and OGCNN[36]. M3GNet has been developed as a
universal model that can handle a wide variety of crystalline materials with
various elemental species on the periodic table. Comparisons were also made
with OGCNN, which combines the orbital representation with MPNN.

#### Cubic-Perovskite - periodic dataset.

Cubic-Perovskite (space group $Pm\bar{3}m$) dataset[45] is created to search
for efficient photoelectrochemical cells (PEC), consisting of $18,928$
different structures and formation energy by calculated with DFT, including
hypothetical materials. This dataset holds 56 elemental species, including
heavy elements with atomic numbers up to 83, and LCAONet’s performance was
tested in terms of its ability to handle a wide variety of elemental species
for the periodic system. Training was split into $80\%,10\%,10\%$ each for
training, validation, and test set, following prior research [36], and SchNet,
M3GNet and LCAONet were learned from scratch with using open source codes [42,
43]. Table 1 reports the mean absolute error (MAE) for each model, with
LCAONet achieving state-of-the-art accuracy. Accurate prediction of heavy
elements and hypothetical structures is an extremely important property for
models to have in the search for new materials, and it is clear that LCAONet
is capable of handling a large number of elemental species. Note that M3GNet
is being developed as a model covering the periodic table. In addition, While
the OGCNN achieved improved accuracy by explicitly incorporating the valence
orbital interaction into the MPNN, it is believed that the model in this study
was able to incorporate orbital information more efficiently through message
passing. When the hyperparameter is aligned with 4 interaction layers and 128
embedded dimensions of the graph, the number of parameters is 1,147,174 for
M3GNet while 869,185 for LCAONet, which is better accuracy with about 75% of
the number of parameters, indicating that the developed basis functions and
message-passing are superior in parameter efficiency. Note that M3GNet uses
spherical Bessel functions as the radial basis functions.

#### Ablation study.

Individual ablation studies were performed on the Perovskite dataset to verify
whether the developed basis functions, the optimization of basis functions for
each elemental species, and the embedding of electronic structure information
contributed to the improved accuracy. Using the spherical Bessel basis
functions proposed in DimeNet as the radial basis function
$\bm{e}_{RBF}^{(ij)}$, the error increased by 10% over the original LCAONet,
suggesting that the basis functions developed in this study are effective in
improving accuracy. The accuracy also deteriorated when the radial basis
functions were not optimized and the same functional form was used for all
elemental species. Note that optimizing the basis functions according to
elemental species limits the geometric representation. The fact that
optimization of basis functions leads to improved accuracy despite these
limitations means that it is important to properly represent the geometric
spread for each elemental species. The LCAONet error increased by about 2%
when electronic structure information was not used in the generation of the
initial node embedding vectors. These results indicate that the optimization
of basis functions and embedding of electronic structure information can
provide an effective inductive bias to the model.

##  7 Conclusion

In this work, we developed physically motivated atomic radial basis functions
and introduced electronic structure message-passing, which encompasses orbital
information, in order to make accurate predictions for a wide variety of
elemental species on the periodic table. The geometrical features of the graph
are vectorized by the radial basis functions of hydrogen-like atoms for the
inter-atomic distances and by the spherical harmonic basis for the directional
information. It was shown that by optimizing the form and number of basis
functions according to the electron configuration of the ground state of the
atoms, it was possible to propagate a message that expresses an appropriate
spatial extent depending on the elemental species. Furthermore, we showed that
message-passing can include physicochemically rich representations by
initializing the embedding vectors of nodes and coefficients using electronic
structure information. We constructed LCAONnet, an MPNN that encapsulates
these improved methods, and achieved state-of-the-art accuracy for the Cubic
perovskite dataset, which contain various elemental species. Predictive models
that can handle arbitrary elemental species are essential in the search for
new materials. In the future, more detailed orbital information, such as spin
degeneracy and crystal field splitting, can be included in the message-passing
to cover a wider variety of materials. In addition, our model uses only $E(3)$
invariant geometrical features, and we believe that extending the equivariant
MPNN to include such physicochemical information would also be useful.

## Acknowledgments and Disclosure of Funding

This research was supported by MEXT (Nos 19H00818, and 19H05787) and JST-CREST
(JPMJCR1993). The authors of this work take full responsibilities for its
content.

## References

  * [1] Felix Faber, Alexander Lindmaa, O Anatole von Lilienfeld, and Rickard Armiento.  Crystal structure representations for machine learning models of formation energies.  Int. J. Quantum Chem., 115(16):1094–1101, August 2015.
  * [2] Haoyan Huo and Matthias Rupp.  Unified representation of molecules and crystals for machine learning.  Mach. Learn.: Sci. Technol., 3(4):045017, November 2022.
  * [3] Luca M Ghiringhelli, Jan Vybiral, Sergey V Levchenko, Claudia Draxl, and Matthias Scheffler.  Big data of materials science: critical role of the descriptor.  Phys. Rev. Lett., 114(10):105503, March 2015.
  * [4] Albert P Bartók, Risi Kondor, and Gábor Csányi.  On representing chemical environments.  Phys. Rev. B Condens. Matter, 87(18):184115, May 2013.
  * [5] Sandip De, Albert P Bartók, Gábor Csányi, and Michele Ceriotti.  Comparing molecules and solids across structural and alchemical space.  Phys. Chem. Chem. Phys., 18(20):13754–13769, May 2016.
  * [6] Tien Lam Pham, Hiori Kino, Kiyoyuki Terakura, Takashi Miyake, Koji Tsuda, Ichigaku Takigawa, and Hieu Chi Dam.  Machine learning reveals orbital interaction in materials.  Sci. Technol. Adv. Mater., 18(1):756–765, October 2017.
  * [7] E Prodan and W Kohn.  Nearsightedness of electronic matter.  Proc. Natl. Acad. Sci. U. S. A., 102(33):11635–11638, August 2005\.
  * [8] 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.  Phys. Rev. Lett., 104(13):136403, April 2010.
  * [9] Matthias Rupp, Alexandre Tkatchenko, Klaus-Robert Müller, and O Anatole von Lilienfeld.  Fast and accurate modeling of molecular atomization energies with machine learning.  Phys. Rev. Lett., 108(5):058301, February 2012.
  * [10] Albert P Bartók, Sandip De, Carl Poelking, Noam Bernstein, James R Kermode, Gábor Csányi, and Michele Ceriotti.  Machine learning unifies the modeling of materials and molecules.  Sci Adv, 3(12):e1701816, December 2017.
  * [11] Jörg Behler and Michele Parrinello.  Generalized neural-network representation of high-dimensional potential-energy surfaces.  Phys. Rev. Lett., 98(14):146401, April 2007.
  * [12] Jörg Behler.  Atom-centered symmetry functions for constructing high-dimensional neural network potentials.  J. Chem. Phys., 134(7):074106, February 2011.
  * [13] J S Smith, O Isayev, and A E Roitberg.  ANI-1: an extensible neural network potential with DFT accuracy at force field computational cost.  Chem. Sci., 8(4):3192–3203, April 2017.
  * [14] M Gastegger, L Schwiedrzik, M Bittermann, F Berzsenyi, and P Marquetand.  wACSF—Weighted atom-centered symmetry functions as descriptors in machine learning potentials.  J. Chem. Phys., 148(24):241709, June 2018.
  * [15] Tsz Wai Ko, Jonas A Finkler, Stefan Goedecker, and Jörg Behler.  A fourth-generation high-dimensional neural network potential with accurate electrostatics including non-local charge transfer.  Nat. Commun., 12(1):398, January 2021.
  * [16] Igor I Baskin, Vladimir A Palyulin, and Nikolai S Zefirov.  A neural device for searching direct correlations between structures and properties of chemical compounds.  J. Chem. Inf. Comput. Sci., 37(4):715–721, July 1997.
  * [17] A Sperduti and A Starita.  Supervised neural networks for the classification of structures.  IEEE Trans. Neural Netw., 8(3):714–735, 1997.
  * [18] M Gori, G Monfardini, and F Scarselli.  A new model for learning in graph domains.  In Proceedings. 2005 IEEE International Joint Conference on Neural Networks, 2005., volume 2, pages 729–734 vol. 2, July 2005.
  * [19] F Scarselli, Sweah Liang Yong, M Gori, M Hagenbuchner, Ah Chung Tsoi, and M Maggini.  Graph neural networks for ranking web pages.  In The 2005 IEEE/WIC/ACM International Conference on Web Intelligence (WI’05). IEEE, 2005.
  * [20] Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini.  Computational capabilities of graph neural networks.  IEEE Trans. Neural Netw., 20(1):81–102, January 2009.
  * [21] David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, Alan Aspuru-Guzik, and Ryan P Adams.  Convolutional networks on graphs for learning molecular fingerprints.  In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
  * [22] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl.  Neural message passing for quantum chemistry.  In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1263–1272. PMLR, 06–11 Aug 2017\.
  * [23] Tian Xie and Jeffrey C Grossman.  Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties.  Phys. Rev. Lett., 120(14):145301, April 2018.
  * [24] K T Schütt, H E Sauceda, P-J Kindermans, A Tkatchenko, and K-R Müller.  SchNet - a deep learning architecture for molecules and materials.  J. Chem. Phys., 148(24):241722, June 2018.
  * [25] Chi Chen, Weike Ye, Yunxing Zuo, Chen Zheng, and Shyue Ping Ong.  Graph networks as a universal machine learning framework for molecules and crystals.  Chem. Mater., 31(9):3564–3572, May 2019.
  * [26] Oliver T Unke and Markus Meuwly.  PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges.  J. Chem. Theory Comput., 15(6):3678–3693, June 2019.
  * [27] Johannes Gasteiger, Janek Groß, and Stephan Günnemann.  Directional message passing for molecular graphs.  In International Conference on Learning Representations (ICLR), 2020\.
  * [28] 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, 2020.
  * [29] Chi Chen and Shyue Ping Ong.  A universal graph deep learning interatomic potential for the periodic table.  Nature Computational Science, 2(11):718–728, November 2022.
  * [30] Johannes Gasteiger, Florian Becker, and Stephan Günnemann.  Gemnet: Universal directional graph neural networks for molecules.  In Conference on Neural Information Processing Systems (NeurIPS), 2021.
  * [31] Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ward Ulissi, C. Lawrence Zitnick, and Abhishek Das.  Gemnet-OC: Developing graph neural networks for large and diverse molecular simulation datasets.  Transactions on Machine Learning Research, 2022.
  * [32] Nathaniel Thomas, T Smidt, S Kearnes, Lusann Yang, L Li, Kai Kohlhoff, and Patrick F Riley.  Tensor field networks: Rotation- and Translation-Equivariant neural networks for 3D point clouds.  ArXiv, 2018.
  * [33] Brandon Anderson, Truong-Son Hy, and Risi Kondor.  Cormorant: covariant molecular neural networks.  In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pages 14537–14546. Curran Associates Inc., Red Hook, NY, USA, December 2019.
  * [34] Kristof Schütt, Oliver Unke, and Michael Gastegger.  Equivariant message passing for the prediction of tensorial properties and molecular spectra.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9377–9388. PMLR, 18–24 Jul 2021.
  * [35] Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji.  Spherical message passing for 3d molecular graphs.  In International Conference on Learning Representations, 2021.
  * [36] Mohammadreza Karamad, Rishikesh Magar, Yuting Shi, Samira Siahrostami, Ian D Gates, and Amir Barati Farimani.  Orbital graph convolutional neural network for material property prediction.  Phys. Rev. Mater., 4(9):093801, September 2020.
  * [37] Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R Manby, and Thomas F Miller, 3rd.  OrbNet: Deep learning for quantum chemistry using symmetry-adapted atomic-orbital features.  J. Chem. Phys., 153(12):124111, September 2020.
  * [38] Oliver T Unke, Stefan Chmiela, Michael Gastegger, Kristof T Schütt, Huziel E Sauceda, and Klaus-Robert Müller.  SpookyNet: Learning force fields with electronic degrees of freedom and nonlocal effects.  Nat. Commun., 12(1):7273, December 2021.
  * [39] D.A. McQuarrie and J.D. Simon.  Physical Chemistry: A Molecular Approach.  Physical Chemistry: A Molecular Approach. University Science Books, 1997\.
  * [40] M. Born and R. Oppenheimer.  Zur quantentheorie der molekeln.  Annalen der Physik, 389(20):457–484, 1927.
  * [41] David J Griffiths and Darrell F Schroeter.  Introduction to Quantum Mechanics.  Cambridge University Press, 3rd edition, August 2018.
  * [42] Kristof T. Schütt, Stefaan S. P. Hessmann, Niklas W. A. Gebauer, Jonas Lederer, and Michael Gastegger.  SchNetPack 2.0: A neural network toolbox for atomistic machine learning.  The Journal of Chemical Physics, 158(14):144801, 04 2023.
  * [43] Code of M3GNet.  https://github.com/materialsvirtuallab/m3gnet.
  * [44] K T Schütt, H E Sauceda, P-J Kindermans, A Tkatchenko, and K-R Müller.  SchNet - a deep learning architecture for molecules and materials.  J. Chem. Phys., 148(24):241722, June 2018.
  * [45] Ivano E Castelli, David D Landis, Kristian S Thygesen, Søren Dahl, Ib Chorkendorff, Thomas F Jaramillo, and Karsten W Jacobsen.  New cubic perovskites for one- and two-photon water splitting using the computational materials repository.  Energy Environ. Sci., 5(10):9034–9043, September 2012.