Equivariant Message Passing Neural Network for Crystal Material
  Discovery
Automatic material discovery with desired properties is a fundamental
challenge for material sciences. Considerable attention has recently been
devoted to generating stable crystal structures. While existing work has shown
impressive success on supervised tasks such as property prediction, the
progress on unsupervised tasks such as material generation is still hampered by
the limited extent to which the equivalent geometric representations of the
same crystal are considered. To address this challenge, we propose EMPNN a
periodic equivariant message-passing neural network that learns crystal lattice
deformation in an unsupervised fashion. Our model equivalently acts on lattice
according to the deformation action that must be performed, making it suitable
for crystal generation, relaxation and optimisation. We present experimental
evaluations that demonstrate the effectiveness of our approach.
Astrid Klipfel, Olivier Peltre, Najwa Harrati, Yaël Fregier, Adlane Sayede, Zied Bouraoui
# Equivariant Message Passing Neural Network for Crystal Material Discovery

###### Abstract

Automatic material discovery with desired properties is a fundamental challenge for material sciences. Considerable attention has recently been devoted to generating stable crystal structures. While existing work has shown impressive success on supervised tasks such as property prediction, the progress on unsupervised tasks such as material generation is still hampered by the limited extent to which the equivalent geometric representations of the same crystal are considered. To address this challenge, we propose EMPNN a periodic equivariant message-passing neural network that learns crystal lattice deformation in an unsupervised fashion. Our model equivalently acts on lattice according to the deformation action that must be performed, making it suitable for crystal generation, relaxation and optimisation. We present experimental evaluations that demonstrate the effectiveness of our approach.

## 1 Introduction

Discovering thermodynamic stable materials with desired properties is a fundamental challenge for material sciences. Considerable attention has recently been devoted to crystalline (crystal) material generation. Crystals are involved everywhere in our modern society from metal alloys to semiconductors. Contrarily to organic molecules which are mostly composed of wide carbon chains with a limited variety of atoms, crystals are three-dimensional periodic structures composed of a wider variety of chemical bonds and atoms. The periodic structure is often represented as a parallelepiped tiling, a.k.a crystal lattice or unit cell.

Within the broad aim of automated stable (crystal) material discovery, various strategies mainly based on simulation or Machine Learning (ML) can be explored. Simulation allows the properties of a given structure to be predicted by applying physics laws while ML consists of modelling and predicting the physical properties. Notice that simulation can also be used for material relaxation, i.e. modifying a structure to improve its stability. The success of ML has led to a paradigm shift in materials science. In particular, ML techniques are used for performing molecule design, modelling physical properties or at the early stage of material discovery. Recently, several works have been introduced to manipulate crystal structures, e.g. [12, 13]. Most notably, models based on geometrically equivariant ML techniques such as Message Passing Neural Networks (MPNNs) have shown good performance in theoretical chemistry, in particular, on supervised tasks such as properly predictions on both organic and crystalline structures, e.g. [11, 12]. However, the majority of existing models are not fully equivariant, making them unsuitable for unsupervised tasks such as generation or representation learning. For example, the method from [11] is only equivariant to SO(3) (rotation group), making it not suitable for crystal lattice deformation where the shape of the structure is unknown in advance. To this end, some methods have been proposed to approximate Density Functional Theory (DFT) simulation using MPNNs for unsupervised tasks, e.g. [11, 12]. They rely on self-simulations to gather information about the interaction forces of a few specific structures to perform generation. However, discovering new materials requires a consequent amount of data to obtain out-of-distribution generalization, i.e. knowledge needed to generalise to unknown structures and perform arbitrary lattice deformation.

We propose EMPNN an equivariant MPNN that acts on crystal lattice without any label from the interaction forces and stress tensors. Previous works already showed the advantage of using MPNN acting on atomic position for both organic molecules and crystals. But acting on crystal lattices without explicit stress tensors remains a challenging problem. Our model enforces a structuring bias adapted to crystals using group actions incorporated by the equivariance property of MPNN layers. To illustrate intuition, given a pair of atoms, if we know their interaction force in a given state, we can generalize this interaction to any other orientation as long as the state and the relative distance remain the same. Hence, we can take advantage of this property, and the equivariant representation to enhance the generalisation capability. This allows our model to equivalently act on crystal lattice according to the deformation action that needs to be performed. We consider equivariance with respect to the Euclidean group \(Euc(3)\) and \(\text{SL}_{3}(\mathbb{Z})\) group. To the best of our knowledge, our model is the first general framework that
formulates an equivariant MPNN on the periodic structures. To demonstrate the effectiveness of our model, we propose a number of evaluation tasks to compare multiple equivariant MPNNs and losses.

## 2 Related works

Within the area of automatic stable material discovery, We can identify three classes of related work according to the molecular descriptors used to represent data.

**Fingerprint.** This class of methods uses handcrafted features of the materials. They are based on fingerprint representation that includes atomic positions and lattice parameters [14]. Additional information such as electronegativity, atomic radius or interatomic distances can also be incorporated, e.g. [13, 15]. Those works mainly rely on Feed-forward Neural Network (FFN) architectures to build Variational Autoencoder [12] or Generative Adversarial Networks [11] to achieve generation or optimization tasks. However, fingerprints do not satisfy the uniqueness property, i.e. the same crystal can have different representations. As FFNs are not equivariant to permutation, alternative representations of the same material can be processed differently. The same observation can be made for other group actions. Finally, existing models don't take into account periodicity.

**Voxel.** Offering a convenient way to represent data in 3-dimensional space, voxels allow encoding lattice parameters and atomic positions [13, 14, 15, 16]. However, voxel-based representation is limited since input data are by nature sparse and discontinuous in the space. Moreover, voxels do not take into account periodicity, which can lead to an edge effect. Finally, the aforementioned methods are not equivariant. As shown in section 4, there are multiple equivalent representations of a given material. Therefore, a set of equivalent representations may lead to inconsistent results. This is a clear limitation of voxel-based representation models.

**Graph-based Representation.** Graph representation of materials can represent the local environment of each atom and structure periodicity. Recent works suggested using Graph Neural Networks (GNN) for materials [10]. MPNNs allow to process sparse data and can be designed to be invariant or equivariant to many group actions. Most of the existing works are equivariant to \(\text{SO}(3)\)[17] thanks to a spherical basis that allows us to predict lattice properties and perform simulations. However, these methods are not able to deform crystal lattices where the shape of the lattice is unknown in advance. In addition, these works are equivariant to subgroups of the Euclidian group but do not consider other group actions such as \(\text{SL}_{3}(\mathbb{Z})\). Several methods have been proposed to approximate DFT simulation with GNN. These methods work by learning interaction forces and stress tensors to lower the total energy of a structure with methods analogue to DFT calculation [18, 19, 20, 21]. These equivariant methods require a lot of additional information about interaction forces, which are not always available. They mainly use self-simulations to gather data, but only for a few specific structures. To discover new materials, we need a lot of data and cannot rely on randomly generated structures, as they lead in general to unstable structures.

## 3 Problem Setting

Crystalline materials can be defined as infinite point clouds. A periodic structure can be represented as a network where a group of points is repeated by a discrete translation, which is is equivalent to parallelepiped tiling containing a cloud of atoms as illustrated in Figure 1. A crystal can be described as atomic positions \(x_{i}\in[0,1[^{3}\) with an associated feature space \(F\) representing the chemical information of each atom \(z_{i}\in F\) and a lattice \(\rho\in\text{GL}_{3}(\mathbb{R})\) representing the material periodicity. The infinite point cloud generated by this representation can be defined as follows:

\[\left\{\left(\rho(x_{i}+\tau),\;z_{i}\right)|\;\tau\in\mathbb{Z}^{3},\;1\leq i \leq n\right\}\ \subseteq\ \mathbb{R}^{3}\times F \tag{1}\]

Where \(\tau\) acts as a \(\mathbb{Z}^{3}\) vector that translates the point cloud. Equation 1 defines the space in which the atoms are located as a torus. In fact, when atoms leave by one side of the lattice they enter by the opposite side with the same orientation. \(\text{GL}_{3}(\mathbb{R})\) defines the shape of the lattice, i.e the periodicity. \(F\) is the feature space that can encode chemical information such as atomic number or charge. For crystal generation, we need to define a model capable to deform the geometry of a structure in order to minimize the total energy and hence obtain a stable structure. Such actions are performed on the material lattice \(\rho\) resulting in the updated lattice \(\rho^{\prime}\) and on atomic positions \(x_{i}\) resulting in the updated positions \(x_{i}^{\prime}\).

\[\begin{cases}\rho^{\prime}=h\rho\\ x_{i}^{\prime}=[x_{i}+h_{i}]\end{cases}. \tag{2}\]

We aim to predict the action \(h\in\text{GL}_{3}(\mathbb{R})\) on the lattice and the actions \(h_{i}\in\mathbb{R}^{3}\) on the atomic position. The atomic positions are brought back into the crystal lattice by truncation. In the following, we introduce our model that learns arbitrary deformations on crystal lattices. We first explain, in Section 4, why group actions are needed for materials, recall the notion of equivariance, and define our group actions on crystals while providing their properties. Finally, Section 5 gives an explicit description of our model along with equivariance results. Proofs and additional materials are provided in an online ArXiv appendix.

Figure 1: Periodic structure represented as a lattice (in dotted lines). The multi-graph associated with a material (blue arrow) can overlap on the adjacent repetition of the lattice and a pair of nodes can have multiple connections.
Equivariance and Group Actions

Crystals materials can be seen as an infinite cloud of atoms as \(\langle m\rangle\subseteq\mathbb{R}^{d}\times F\). As such, equivalences between materials are defined by isometries, i.e. by the group action of \(\mathrm{Euc}(d)\) regardless of lattice generators1. As a crystal lattice can have multiple space-tiling representations resulting in an identical infinite atomic cloud, the \(SL_{d}(\mathbb{Z})\) group action is needed for paving. Consequently, the group \(G=\mathrm{Euc}(d)\times SL_{d}(\mathbb{Z})\times\mathfrak{S}_{n}\) acts on the lattice without affecting its properties. \(\mathfrak{S}_{n}\) is the permutation group that acts by changing the numbering of atoms, where \(n\) is the number of atoms. Please note that atoms are always in the same place, but not with the same index. As chirality has an impact on the properties of a chemical structure, the reflection action should be excluded. The special Euclidean group that doesn't include reflection should be then considered. However, in this work, we consider \(\mathrm{Euc}(d)\) that acts on the chirality assuming that this limitation will not be problematic with inorganic material. We consider crystals described by an infinite cloud of atoms that is invariant under a discrete subgroup \(L\subseteq\mathbb{R}^{d}\) of maximal rank. For any choice of generators \((\tau_{1},\ldots,\tau_{d})\in L\), we consider the unique automorphism \(\rho\in GL_{d}(\mathbb{R})\) that maps the canonical basis of \(\mathbb{R}^{d}\) to the generating basis of \(L\) to represent \(L\).

Footnote 1: A generator is a lattice property that defines pattern repetition) and atom indices

**Definition 1**.: _The representation space of featured materials \(\mathcal{M}^{F}\) is the disjoint union \(\coprod_{n\in\mathbb{N}}\mathcal{M}^{F}_{n}\) where:_

\[\mathcal{M}^{F}_{n}=\left\{(\rho,x,z)\,|\,\rho\in GL_{d}(\mathbb{R}),\;x\in[0,1[^{n\times d},\;z\in F^{n}\right\}\]

_Chemical materials are represented in \(\mathcal{M}=\mathcal{M}^{\mathbb{N}}\), with atomic numbers as feature sequence \(z\)._

\(\mathcal{M}^{F}_{n}\) is an infinite set of triplet \(\rho\), \(x\), \(z\) that represent all possible materials with \(n\) atoms. The atomic number has a chemistry reference, e.g. 1 for hydrogen or 6 for carbon.

**Definition 2**.: _The infinite point cloud \(\langle M\rangle\) associated to a material \(M=(\rho,x,z)\) in \(\mathcal{M}^{F}_{n}\) is defined as:_

\[\langle M\rangle=\left\{\left(\rho\cdot(x_{i}+\tau),\,z_{i}\right)|\,\tau\in \mathbb{Z}^{d},\,1\leq i\leq n\right\}\ \subseteq\ \mathbb{R}^{d}\times F\]

_The cloud \(\langle M\rangle\) is invariant under the action of the lattice \(L=\rho\cdot\mathbb{Z}^{d}\subseteq\mathbb{R}^{d}\)._

The \(\mathrm{Euc}(d)\) group acts naturally on subsets of \(\mathbb{R}^{d}\) and two materials \(M\) and \(M^{\prime}\) should be considered physically identical if they span isometric point clouds. Let us write \(M\sim M^{\prime}\) if there exists an isometry \(g\in\mathrm{Euc}(d)\) such that \(\langle M^{\prime}\rangle=g\cdot\langle M\rangle\). Let \(\langle\mathcal{M}^{F}\rangle\) be the image of \(\mathcal{M}^{F}\) in \(\mathcal{P}(\mathbb{R}^{d}\times F)\) under \(\langle-\rangle\). The quotient space \(\mathcal{M}/\sim\) of equivalent materials is defined by the following universal diagram:

Infinite point clouds can only be represented by non-intrinsic representatives \(M\in\mathcal{M}^{F}\). In the following, we describe how the relation \(\sim\) is related to group actions on \(\mathcal{M}^{F}\). The following proposition introduces the group actions that don't change the properties of materials, i.e. actions that lead to producing equivalent materials.

**Proposition 1**.: _The following four actions on \(\mathcal{M}^{F}_{n}\) preserve the equivalence class of material:_

* \(\mathfrak{S}_{n}\)_permutation group, acting by_ \(\sigma\cdot(\rho,x,z)=(\rho,x\circ\sigma^{-1},z\circ\sigma^{-1})\)__
* \(O(d)\) _orthogonal group, acting by_ \(g\cdot(\rho,x,z)=(g\cdot\rho,x,z)\)__
* _Ex translation group_2_, acting by_ \(v\cdot(\rho,x,z)=(\rho,\left(x+\rho^{-1}v\right],z)\)__

Footnote 2: The actions of \(E\) and \(\mathbb{R}^{d}\) are equivalent, being simply intertwined by the isomorphism \(\rho:\mathbb{R}^{d}\to E\). The action of \(E\) is more natural, extending the action of \(O(d)\) to \(\mathrm{Euc}(d)\) but the action of \(\mathbb{R}^{d}\) is more convenient in our representation space.

Footnote 3: Permutations acting trivially on \(\langle\mathcal{M}\rangle\).

_These actions are free and proper on \(\mathcal{M}^{F}_{n}\). The point cloud map \(\langle-\rangle\) commutes with these actions4._

Footnote 4: The actions of \(E\) and \(\mathbb{R}^{d}\) are equivalent, being simply intertwined by the isomorphism \(\rho:\mathbb{R}^{d}\to E\). The action of \(E\) is more natural, extending the action of \(O(d)\) to \(\mathrm{Euc}(d)\) but the action of \(\mathbb{R}^{d}\) is more convenient in our representation space.

Footnote 4: Permutations acting trivially on \(\langle\mathcal{M}\rangle\).

Performing modification by permutations and isometries is not enough to get a faithful representation of \(\mathcal{M}^{F}/\sim\). Different choices of lattice \(L\subseteq\mathbb{R}^{d}\) lead to different primitive point clouds in \([0,1[^{d}\). The action of \(SL_{d}(\mathbb{Z})\) on \(GL_{d}(\mathbb{R})\) describes all the possible choices of generators for \(L\). However, \(SL_{d}(\mathbb{Z})\) cannot simply act by left multiplication on \(\mathcal{M}^{F}\) like \(\mathrm{Euc}(d)\) without distorting the relative positions of atoms in the primitive cell \(\rho\cdot[0,1[^{d}\). We complete Proposition 1 by specifying how to repave the space while being equivalent to the structure we start with.

**Proposition 2**.: _The group \(SL_{d}(\mathbb{Z})\) acts on \(\mathcal{M}^{F}\) by letting for every change of lattice generators \(g\):_

\[g\cdot(\rho,x,z)=(\rho\cdot g^{-1},[gx],z)\]

_where \([gx]_{i}\) denotes the unique element of \([0,1[^{d}\) in the orbit of \(gx_{i}\) under \(\mathbb{Z}^{d}\). Identifying the reference cell \([0,1[^{d}\) with the torus \(\mathbb{T}^{d}\), the action of \(SL_{d}(\mathbb{Z})\) on \(\mathcal{M}^{F}\simeq GL_{d}(\mathbb{R})\times(\mathbb{T}^{d})^{n}\times F^{n}\) is free and proper. The point cloud map is invariant under the action of \(SL_{d}(\mathbb{Z})\)._

The reference cell is the base cell we use to pave the space with \(L\). It is a parallelepiped of atoms and \(L\) is the translation that allows the parallelepiped moving to pave the space.

Figure 2: In definition 2 the point cloud is a space tiling (top left corner). The actions from \(\mathrm{Euc}(2)\) and \(SL_{2}(\mathbb{Z})\) groups commute and do not affect interatomic distances.
**Proposition 3**.: _The actions of \(\mathrm{Euc}(d)\), \(\mathfrak{S}_{n}\) and \(SL_{d}(\mathbb{Z})\) on \(\mathcal{M}_{n}^{F}\) commute as shown in Figure 2._

Let \(G\) be the product of \(\mathrm{Euc}(d)\times\mathfrak{S}_{n}\times SL_{d}(\mathbb{Z})\). Propositions 1 and 2 imply that the quotient of \(\mathcal{M}^{F}\) under the action of \(G\) is a well-formed topological space. This quotient is not the space \(\mathcal{M}^{F}/\sim\) of equivalent materials, because the lattice associated with a material representation \(M\in\mathcal{M}^{F}\) is not always a maximal symmetry subgroup of its point cloud.

Graph equivarianceInternal forces acting on a crystal structure are equivariant to the aforementioned group actions. As the properties of a crystal depend on interatomic interaction, equivariance could be then considered as the solution to obtain generalization capability. In this work, we take advantage of the equivariance of the graph representation of materials under \(G\) the product of \(\mathrm{Euc}(d)\times\mathfrak{S}_{n}\times SL_{d}(\mathbb{Z})\).

**Definition 3**.: _A neural network \(f_{\theta}:\mathcal{M}^{F}\rightarrow\mathbb{R}^{k}\) is said invariant under \(G\) if for all \(g\in G\):_

\[f_{\theta}(g\cdot M)=f_{\theta}(M)\]

**Definition 4**.: _A neural network \(\varphi_{\theta}:\mathcal{M}^{F}\rightarrow\mathcal{M}^{F^{\prime}}\) is said equivariant under \(G\) if for all \(g\in G\):_

\[\varphi_{\theta}(g\cdot M)=g\cdot\varphi_{\theta}(M)\]

## 5 Equivariant GNN for Materials

We now introduce our MPNN that performs arbitrary deformation by reasoning on relative atomic distances and angles. A spatial equivariance is enforced by the MPNN. We first associate a graph with a material and then take advantage of the local invariance (input quantities are themselves invariant: distance, angle, etc.) and equivariance of the graph to define equivariant actions on crystal lattices.

**Definition 5**.: _We call directed 2-graph \(\Gamma=(\Gamma_{0},\Gamma_{1},\Gamma_{2})\) a triplet of sets together with applications:_

* \(\pi_{1}:\Gamma_{1}\rightarrow\Gamma_{0}\times\Gamma_{0}\)_, written_ \(\pi_{1}(\gamma)=(\mathrm{src}(\gamma),\mathrm{tgt}(\gamma))\)__
* \(\pi_{2}:\Gamma_{2}\rightarrow\Gamma_{0}\times\Gamma_{0}\times\Gamma_{0}\)__

_We call \(\Gamma\) a directed 1-graph when \(\Gamma_{2}=\varnothing\)._

The aforementioned graphs are often called "multi"-graphs. Recall that \(\pi_{1}\) and \(\pi_{2}\) may not be injective. They are called "hyper"-graphs as well, because they generalise 1-graphs to dimensions \(\geq 1\) and "directed" because we do not assume any symmetry on \(\Gamma\) w.r.t vertice permutations.

**Definition 6**.: _Let \(M=(\rho,x,z)\) in \(\mathcal{M}_{n}^{F}\) be a material and \(c_{i}>0\) for \(1\leq i\leq n\) denotes cutoff distances. We define a directed 2-graph \(\Gamma=\Gamma_{M,c}\) by the graded components:_

* \(\Gamma_{0}=\{1,\ldots,n\}\)__
* \(\Gamma_{1}=\left\{(i,j,\tau)\in\Gamma_{0}\times\Gamma_{0}\times\mathbb{Z}^{d }\,\big{|}\,||\rho(x_{j}-x_{i}+\tau)||<c_{i}\right\}\)__
* \(\Gamma_{2}=\left\{(\gamma,\gamma^{\prime})\in\Gamma_{1}\times\Gamma_{1}\, \big{|}\,\mathrm{tgt}(\gamma)=\mathrm{src}(\gamma^{\prime})\right\}\)__

_with obvious projections, i.e. with \(\pi_{1}:(i,j,\tau)\mapsto(i,j)\) and \(\pi_{2}:(\gamma,\gamma^{\prime})\mapsto(\mathrm{src}(\gamma),\mathrm{tgt}( \gamma),\mathrm{tgt}(\gamma^{\prime}))\)._

This graph construction includes many definitions of material graphs, making it versatile and usable in most contexts since a material graph is built from the local environment of atoms. This definition includes a graph built from a constant cutoff distance (i.e. \(c_{i}\) is constant) and a graph built from \(k\) nearest neighbour or built from chemical properties as the covalent radii. Definition 6 generalizes to most of the graphs defined in previous works [1, 1, 1, 10]. The key feature of this construction is the invariance of edges and triplets. As interatomic distances and unoriented angles are invariants to \(\mathrm{Euc}(d)\) and \(SL_{d}(\mathbb{Z})\) groups, any graph constructed from the local environment of the atoms will be invariant. More details about graph construction are in the appendix. We now introduce notations needed to define our model.

**Definition 7**.: _Let consider \(M=(\rho,x,z)\in\mathcal{M}^{F}\) and \(\Gamma=\Gamma_{M,c}\) we introduce the following notations:_

* \(e_{ij}^{\tau}=(x_{j}-x_{i}+\tau)\) _for edge vector in lattice coordinates,_
* \(v_{ij}^{\tau}=\rho(e_{ij}^{\tau})\) _for the edge vector in physical space,_
* \(r_{ij}^{\tau}=||v_{ij}^{\tau}||\) _for the physical edge length,_
* \(\theta r_{ijk}^{\tau\tau^{\prime}}\) _as the unoriented angle between_ \(v_{ij}^{\tau}\) _and_ \(v_{jk}^{\tau^{\prime}}\)__
* \(\mathcal{A}_{ijk}^{\tau\tau^{\prime}}\) _as the area of the triangle_ \(x_{i}\)_,_ \(x_{j}+\tau\) _and_ \(x_{k}+\tau^{\prime}\)__

_Let us also write \(e_{\gamma},v_{\gamma},r_{\tau},\theta_{\gamma\gamma^{\prime}},\mathcal{A}_{ijk }^{\tau\tau^{\prime}}\) for the same quantities when we do not need to make vertices explicit. Note that \(r_{\gamma}\), \(\theta_{\gamma\gamma^{\prime}}\) and \(\mathcal{A}_{ijk}^{\tau\tau^{\prime}}\) are natural Euclid invariants._

### Gradient of the invariant geometry

To build vector fields of our equivariant MPNN, we take advantage of the gradient of the invariant geometry of crystal graphs. For 0-chains, i.e. vertices \(i\in\Gamma_{0}\), the Euclid group acts transitively on spatial coordinates such that \(I_{i}\) is trivial (a point) and \(r_{i}\) is a constant. For 1-chains, i.e. directed edges \(\gamma\in\Gamma_{1}\), the only Euclid invariant is the length of the associated vector. For \(I_{\gamma}=\mathbb{R}\) and for \(\gamma:i\xrightarrow{\tau}j\), we let:

\[r_{\gamma}(x_{\gamma})=r_{ij}^{\tau} \tag{3}\]

For 2-chains \(\bar{\gamma}=i\xrightarrow{\tau}j\xrightarrow{\tau^{\prime}}k\), we find more convenient to define invariants as two vector lengths and the angle at their common point, i.e. \(I_{\bar{\gamma}}=\mathbb{R}^{3}\) with:

\[r_{\bar{\gamma}}=\left(\theta_{ijk}^{\tau\tau^{\prime}},\,r_{ij}^{\tau},\,r_{ jk}^{\tau^{\prime}}\right) \tag{4}\]

For a tangent vector at \(\rho\in GL_{d}(\mathbb{R})\), we have:

\[\frac{\partial v_{ij}^{\tau}}{\partial\rho}=\rho\cdot(x_{j}-x_{i}+\tau)=\rho \cdot e_{ij}^{\tau} \tag{5}\]

The differential edge distances with respect to \(\rho\) projects on the source and image edge vectors \(e_{ij}^{\tau}\) and \(u_{ij}^{\tau}\) respectively. It is equal to 1 on the rank 1 linear map \(|u_{ij}^{\tau}\rangle\langle e_{ij}^{\tau}|\). \(u_{ij}^{\tau}\rangle\) denotes the normalized vector \(v_{ij}^{\tau}\) such as \(u_{ij}^{\tau}=v_{ij}^{\tau}/\tau_{ij}^{\tau}\).

\[\frac{\partial r_{ij}^{\tau}}{\partial\rho}=\langle u_{ij}^{\tau},\,\rho\cdot e_{ ij}^{\tau}\rangle \tag{6}\]

The angle differentials with respect to \(\rho\) are computed by assuming that the middle point is fixed (it is true up to a translation in the target space, which does not alter the angle). \(\omega_{ijk}^{\tau\tau^{\prime}}\) denotes the unit normal vector to \((v_{ij}^{\tau},v_{jk}^{\tau})\)

\[\frac{\partial\theta_{ijk}^{\tau\tau^{\prime}}}{\partial\rho}=\langle\omega_{ijk}^{ \tau\tau^{\prime}}\times u_{jk}^{\tau^{\prime}}\rho\cdot e_{jk}^{\tau^{\prime}} \rangle-\langle\omega_{ijk}^{\tau\tau^{\prime}}\times u_{ij}^{\tau}\rho\cdot e_{ij}^{\tau}\rangle \tag{7}\]

The mixed product coincides with the determinant and is invariant under cyclic permutations.
### Equivariant Message Passing Neural Network

We now introduce a general definition of our equivariant MPNN based on vector fields. We formally define \(\lambda\) as the vector field used in Equation10. It allows for defining how the GNN acts on the crystal lattice.

**Definition 8**.: _To every edge \(\gamma\in\Gamma_{1}\) and every 2-region \(\gamma\gamma^{\prime}\in\Gamma_{2}\) we associate the infinitesimal lattice deformations \(\lambda_{\bar{\gamma}}:\mathcal{M}_{\bar{\gamma}}\to\mathfrak{gl}_{d}\) defined by:_

* \(\lambda_{\gamma\gamma}(M_{\bar{\gamma}})=\left|\,u_{\gamma}\,\right\rangle \left\langle\,u_{\gamma}\,\right|\)__
* \(\lambda_{\gamma\gamma^{\prime}}(M_{\bar{\gamma}})=\left|\,u_{\gamma}\,\right\rangle \left\langle\,u_{\gamma^{\prime}}\,\right|+\left|\,u_{\gamma^{\prime}}\,\right\rangle \left\langle\,u_{\gamma}\,\right|\)__

_The \(\left|\,-\,\right\rangle\)\(\left\langle\,\,\,\,-\,\right\rangle\) is a notation in quantum physics to denote the matrix obtained as the product of a column vector (\(\left|\,(\,\right\rangle\)) \(N\)) is \(V\) seen as a column) and a line vector (\(\left\langle\,W\,\right|\) is W seen as line vector). In our case \(\left|\,-\,\right\rangle\)\(\left\langle\,\,\,\,-\,\right\rangle\) with two vectors \(u,v\in\mathbb{R}^{d}\) we have \(\left|\,u\,\right\rangle\left\langle\,v\,\right|=uv^{\intercal}\). Alternatively, we can directly use gradients of the geometric invariant such as:_

* \(\lambda_{\gamma}(M_{\bar{\gamma}})=\frac{\partial r_{ij}^{\tau^{\prime}}}{ \partial\rho}\)__
* \(\lambda_{\gamma\gamma^{\prime}}(M_{\bar{\gamma}})=\frac{\partial r_{ij}^{\tau^ {\prime}}}{\partial\rho}\) _or_ \(\frac{\partial r_{ij}^{\tau^{\prime}}}{\partial\rho}\) _or_ \(\frac{\partial\theta r_{ij}^{\tau^{\prime}}}{\partial\rho}\) _or_ \(\frac{\partial\mathcal{A}_{ij}^{\tau^{\prime}}}{\partial\rho}\)__

_To ensure transversality with \(\mathfrak{so}_{d}\), \(\lambda_{\bar{\gamma}}\) for all \(\bar{\gamma}\in\Gamma\) is symmetric as equivariance means that the lattice is searched among an equivalence class in \(GL_{d}(\mathbb{R})/SO_{d}\)._

An equivariant GNN that acts on materials is as follows:

**Proposition 4**.: _A neural network \(\varphi_{\theta}:\mathcal{M}_{n}^{F}\to\mathcal{M}_{n}^{F^{\prime}}\), written \(\varphi_{\theta}:(\rho,x,z)\mapsto\rho^{\prime}\) is decomposed as follows:_

_The generation of messages from the edges and the triplets of the graph such as \(\varphi_{\theta}^{m^{(k)}}:\mathbb{R}^{f^{(k)}\times\Gamma_{k}}\to\mathbb{R} ^{h^{(k)}\times\Gamma_{k}}\)_

\[m_{ij\tau}= \varphi_{\theta}^{m^{(1)}}(z_{i},z_{j},||v_{ij\tau}||) \tag{8a}\] \[m_{\gamma,\gamma^{\prime}}= \varphi_{\theta}^{m^{(2)}}(z_{i},z_{j},z_{k},||v_{\gamma}||,|v_{ \gamma^{\prime}}||,\theta_{\gamma,\gamma^{\prime}}) \tag{8b}\]

_The aggregation and update of the messages at each node is \(\varphi_{\theta}^{\varphi^{(k)}}:\mathbb{R}^{h^{(k)}\times\Gamma_{k}}\to \mathbb{R}^{h^{\prime(k)}\times\Gamma_{k}}\) and \(\varphi_{\theta}^{u}:\mathbb{R}^{z\times\Gamma_{0}}\times\mathbb{R}^{h^{\prime (k)}\times\Gamma_{1}}\times\mathbb{R}^{h^{\prime(2)}\times\Gamma_{2}}\to \mathbb{R}^{z\times\Gamma_{0}}\)_

\[z_{i}^{\prime}= \varphi_{\theta}^{u}(z_{i},\sum_{\gamma\in\Gamma_{1}(i)}\varphi _{\theta}^{z^{(1)}}(m_{\gamma}),\sum_{(\gamma,\gamma^{\prime})\in\Gamma_{2}(i )}\varphi_{\theta}^{z^{(2)}}(m_{\gamma\gamma^{\prime}})) \tag{9}\]

\(\varphi_{\theta}^{\rho^{(k)}}\) _is the weight of a vector field \(\lambda_{\bar{\gamma}}\) such as \(\varphi_{\theta}^{\rho^{(k)}}:\mathbb{R}^{f^{(k)}\times\Gamma_{k}}\to\mathbb{R }^{\Gamma_{k}}\)_

\[\rho^{\prime}=\exp\left(\frac{1}{|\Gamma_{1}|}\sum_{\gamma\in \Gamma_{1}}\varphi_{\theta}^{\rho^{(1)}}(m_{\gamma})\cdot\lambda_{\gamma} \right)\cdot\rho \tag{10a}\] \[\rho^{\prime}=\exp\left(\frac{1}{|\Gamma_{2}|}\sum_{(\gamma, \gamma^{\prime})\in\Gamma_{2}}\varphi_{\theta}^{\rho^{(2)}}(m_{\gamma},m_{ \gamma^{\prime}},\theta_{\gamma\gamma^{\prime}})\cdot\lambda_{\gamma\gamma^{ \prime}}\right)\cdot\rho \tag{10b}\]

\(\varphi_{\theta}\) _is equivariant under \(G=\mathrm{Euc}(d)\times\mathfrak{S}_{n}\times SL_{d}(\mathbb{Z})\) if the vector field \(\lambda_{\bar{\gamma}}\) is invariant to \(SL_{d}(\mathbb{Z})\) and equivariant to \(\mathrm{Euc}(d)\) such as \(\lambda_{\bar{\gamma}}(g\cdot M)=g\lambda_{\bar{\gamma}}(M)g^{-1}\) for all \(g\in O(d)\), as the translation doesn't act on the crystal lattice._

From proposition 4, a GNN architecture acting on crystal material that satisfies Equations 8-10 is equivariant.

### EMPNN for Crystal Lattice Deformation

To empirically evaluate our approach, we defined EMPNN as a simple but effective GNN model that fits with Proposition 4. We chose to keep our model simple to facilitate the comparison between multiple vector fields. The architecture is illustrated in Figure 3. We slightly adapted equation 10 by adding a first-order approximation of the matrix exponential to both vector fields over the edges and the triplets. Further details are given in Section B.2 of the appendix.

Loss functionsThe goal of a loss function is to reproduce the shape and volume of the target crystal, i.e. \(\varphi(\tilde{\rho}\cdot h^{-1})=g\cdot\rho\cdot h^{-1},g\in O(3)\) and \(h\in SL_{3}(\mathbb{Z})\) (as \(\mathrm{Euc}(3)\) acts on \(\rho\) as \(O(3)\)). There exist multiple ways to define loss functions, but all the definitions will have implicit bias. To evaluate this bias, we use a classical loss function over the normalized lattice parameters. Another approach is to compute a matrix distance between the metric tensors. Both the lattice parameters and metric tensor losses are invariant to the euclidean group but equivariant to \(SL_{3}(\mathbb{Z})\). We tested the mean absolute error (MAE) \(\mathcal{L}_{\text{max}}^{\text{Param}}\) and the mean squared error (MSE) \(\mathcal{L}_{\text{max}}^{\text{Param}}\) of the normalized lattice parameters. We have also tested the MAE \(\mathcal{L}_{\text{max}}^{\rho}\), the MSE \(\mathcal{L}_{\text{mse}}^{\rho}\) and the invariant Riemannian metric \(\mathcal{L}_{\text{Riemann}}^{\text{Param}}\) of the metric tensors. The loss expressions are available in the appendix.

## 6 Experiments

Our main goal is to show the capability of our EMPNN to perform arbitrary crystal lattice deformation by improving the total energy of crystal structures, i.e. the thermodynamic stability. We rely on denoising of the crystal lattice as evaluation task4. We considered datasets of stable crystals where each structure is in local minima of formation energy. Applying a small random deformation to a structure leads to a less stable one with a high energy level (as the energy increases in all directions locally). We can then generate pairs of stable and less stable structures that we used to teach our model how to deform the less stable structure to obtain a stable one. In general, denoising tasks are more insightful than generative tasks as they show how a model acts on a

Figure 3: (a) The EMPNN model comprises an embedding layer, standard MPNN layers and EMPNN layers to perform deformation. (b) A EMPNN layer is composed of an MPNN with vector fields deforming the lattice \(\rho\).
crystal lattice. More specifically, external bias can be better controlled when performing denoising. The chemical composition and atomic positions have an important impact on the outcome. For example, binary and ternary compounds with a light element are known to be significantly easier to generate than ternary compounds without light elements or quaternary compounds. Consequently, a generator may tend to produce simple stable materials instead of a representative sample. In this case, an improvement of the metrics may not reflect lattice improvement. The quality of a crystal is also more difficult to evaluate. Namely, if a generative model can not produce some specific lattice shapes, quantitative metrics will struggle to measure the bias. Therefore, the performance of a generator is not a good measure to evaluate the performance of our model on arbitrary lattice deformation.

Evaluation metricsWe introduce three evaluation metrics defined as the average improvement of lattice parameters and the total energy. Let us denote the lattice parameters by \(abc\in\mathbb{R}^{3}\) and \(\alpha\beta\gamma\in\mathbb{R}^{3}\) and the total energy by \(E\in\mathbb{R}\). Given a parameter \(y\), let \(\tilde{y}\) be the noisy parameter and \(y^{\prime}\) the denoised parameter. The metrics are defined as follows:

\[\text{length}= \frac{1}{3N}\sum_{k=1}^{N}l1(\widetilde{abc}_{k},abc_{k})-l1(abc_{k }^{\prime},abc_{k}) \tag{11}\] \[\text{angle}= \frac{1}{3N}\sum_{k=1}^{N}l1(\widetilde{\alpha\beta\gamma}_{k}, \alpha\beta\gamma_{k})-l1(\alpha\beta\gamma_{k}^{\prime},\alpha\beta\gamma_{k})\] (12) \[\text{energy}= \frac{1}{N}\sum_{k=1}^{N}E_{k}^{\prime}-\tilde{E}_{k} \tag{13}\]

Namely, the improvement could be geometrical, i.e. based on lattice parameters or chemical, i.e. lowering of the formation energy. Evaluating the formation energy is computationally expressive and only done on a small subset of the test set.

Experimental setting and datasetsWe considered three datasets of stable crystals for which we perform denoising: Perov-5 [10, 11], Carbon-24 [24] and Mp-20 [12]. Perov-5 contains perovskite (cubic) structures that have highly uniform shapes but with different chemical compositions between structures. Carbon-24 is composed of carbon atoms having a large variety of shapes. This dataset is used to evaluate the performance of our EMPNN without negative bias in case of a poor chemical encoding of atoms. Mp-20 is a subset of the material project proposed in [10] that has a large sample of shapes and chemical compositions. It is the most representative of ordinary structures. We used the same training, validation and test splits as [10]. To train our model, we apply random deformations on the lattices \(\rho\) as \(\tilde{\rho}=\exp(A)\rho\) with \(A\sim\mathcal{N}(0,\sigma)\). All the conducted experiments, use grid search on hyperparameters. More information about the experiments is given in the supplementary materials. We conducted three experiments to evaluate (1) the loss functions of Section 5.3, (2) the vector fields and (3) the reconstruction capability of our model.

Loss functions evaluationTable 1 shows the relationship between geometrical and chemical metrics. The best energy improvements are generally associated with the best lattice error improvement. Based on lattice parameters comparison, we obtain better performance of the loss functions. However, we may expect that this evaluation is biased. But since energy-based metrics show similar results to geometry based metrics, we conclude that the bias is negligible.

Force field evaluationWe evaluated force field configurations acting on the lattice. We first considered the edge information: \(\{|\gamma\rangle\langle\gamma|\subseteq\Gamma_{1}\}\) and \(\{r_{\gamma}\subseteq\Gamma_{1}\}\). Second, we consider triplets information without angle and area: \(\{|\gamma\rangle\langle\gamma|\subseteq\Gamma_{2}\}\) and \(\{r_{\gamma},r_{\gamma^{\prime}}\subseteq\Gamma_{2}\}\). As geometrical information such as angles can determine crystal properties, we include triplets information as unoriented angles and area: \(\{|\gamma\rangle\langle\gamma^{\prime}|\subseteq\Gamma_{2}\}\) and \(\{|\gamma\rangle\langle\gamma|,|\gamma\rangle\langle\gamma^{\prime}|_{\text{ sym}}\subseteq\Gamma_{2}\}\). \(\cup\), represents the union of several vector fields and \(\bullet\) denotes a wildcard that takes all vector fields into account for a given n-graph. We also evaluate the benefits of symmetric matrices on the lattice as suggested in Definition 8. Any matrix in \(GL_{d}(\mathbb{R})\) can be seen as the composition of a rotation and a symmetric matrix, i.e. polar decomposition such as \(M=RS\) with \(R\in\mathit{SO}_{d}\) and \(S\in GL_{d}(\mathbb{R})/SO_{d}\). As rotation doesn't act on material properties, then acting on the lattice with \(M\) is equivalent to acting on the lattice with a symmetric matrix \(S\). Forcing this action to be a symmetric matrix may then lead to interesting results. We conduct experiments with relaxed symmetry constraint "sym" when the symmetric vector fields are used.

As baselines, we first consider the (Feed-forward FF) method proposed in [10] which is an invariant

\begin{table}
\begin{tabular}{c|c c c|c c c|c c c} loss & \multicolumn{3}{c|}{Carbon-24} & \multicolumn{3}{c|}{Mp-20} & \multicolumn{3}{c}{Perov-5} \\  & lengths & angle & energy & lengths & angle & energy & lengths & angle & energy \\ \hline \(\mathcal{L}_{\text{me}}^{\text{Param}}\) & **0.696** & **8.390** & **-0.655** (62.5) & **0.785** & **5.093** & **3.124** (51.7) & **0.967** & **15.227** & -3.426 (93.8) \\ \(\mathcal{L}_{\text{me}}^{\text{Param}}\) & 0.677 & 8.148 & -0.413 (65.6) & 0.710 & 4.752 & 5.485 (44.8) & 0.983 & 15.437 & -3.634 (93.8) \\ \(\mathcal{L}_{\text{me}}^{\rho}\) & 0.599 & 4.306 & 0.526 (62.5) & 0.540 & 1.674 & 11.268 (40.7) & 0.964 & 15.074 & -1.518 (90.6) \\ \(\mathcal{L}_{\text{me}}^{\rho}\) & 0.655 & 5.563 & 2.432 (40.6) & 0.683 & 2.645 & 10.964 (18.5) & 0.974 & 15.047 & **-3.741** (93.8) \\ \(\mathcal{L}_{\text{Riemann}}^{\rho}\) & 0.637 & 5.352 & 0.864 (43.8) & 0.729 & 3.777 & 6.859 (51.7) & 0.967 & 15.367 & -3.088 (93.8) \\ \end{tabular}
\end{table}
Table 1: Metrics are defined as the average improvement of lattice parameters and the average improvement of total energy. The metrics are calculated between noisy structure and denoised structure. The lengths are given in Å (angström), the angle in degree (higher is better) and the energy in eV/atom (lower is better). The value between parenthesis is the percentage of structure with lower energy. Energy is calculated with VASP[10, 10, 11] on a subset of 32 structures because of the high computational budget of DFT calculation.
method aiming to predict lattice parameters (distances and angles) using an invariant encoder with a simple FF. This allows us to compare the performance of our model with an invariant model. The second baseline (DFT) is a DFT calculation that evaluates the stress tensor of the crystal and optimizes its geometry. The configuration of the DFT calculation is given in the supplementary materials. DFT is not based on ML, as such, it is computationally heavy compared to EMPNN. DFT is unsuited for generating crystals without optimization techniques. So it cannot really be compared with ML models (baselines and our model), but we chose to use it to provide insight into the metrics.

Table 2 shows an enhanced denoising capability of our model for most of the proposed variants. Including triplets information improves the results when vector fields are defined from the gradient of invariant geometry (Section 5.1). However, vector fields defined from edges information achieved more consistent results than those defined from triplets, especially on ket-bra. Our model outperforms FF on Carbon-24 and Mp-20 with a significant improvement of the lattice parameters but not on Perov-5 (although the performance is very close). This suggests the importance of equivariance. The FF is not capable to achieve fine-grained deformation contrary based on vector fields. In fact, FF converges much faster during the first training steps but cannot improve the loss above a certain threshold. The only case where FF outperforms our model is when the crystal shape is extremely uniform, which is the case of Perov-5 where all the structures are cubic. In Perov-5, the angle improvement is not relevant as FF uses normalized lattice parameters. A random model or a constant parameter will produce similar results. Regarding DFT, it improves the lattice parameters on Carbon-24 and Mp-20 but not on Perov-5. This suggests that the crystals before random deformation probably remain close to local minima of the formation energy after deformation on Carbon-24 and Mp-20, but not on Perov-5. Our methods can take advantage of the biased distribution on Perov-5 while DFT is not capable of. Finally, comparing multiple configurations of vector fields shows that ket-bra works better on 1-graph while gradient-based vector fields work better on 2-graph. Triplets vector fields obtain better results with the area and angle information.

**Reconstruction task evaluation.** The reconstruction is close to the generative task and aims to build a crystal lattice from scratch. This cannot be performed with chemical simulation techniques such as DFT. We start from the point cloud as if it was in a cubic lattice of one A on a side. From this cubic lattice, the EMPNN performs the reconstruction. The main hypothesis is that there is a single stable cell which corresponds to the starting atomic positions. Our model consistently outperforms the FF model as shown in Table 3.

## 7 Conclusion

We proposed a general equivariant MPNN framework for material science by taking into consideration \(SL_{3}(\mathbb{Z})\) group action on crystal materials. In particular, our model uses multiple vector fields to act on crystal lattices. We showed the benefits of our model compared to equivariant baselines that do not consider \(SL_{3}(\mathbb{Z})\). We also compared different loss functions and results with DFT calculation to give insight into methods based on lattice reconstruction such as those using auto-encoder.

\begin{table}
\begin{tabular}{c c|c c|c c|c c}  & \multicolumn{2}{c|}{carbon-24} & \multicolumn{2}{c|}{Mp-20} & \multicolumn{2}{c}{Perov-5} \\  & method & lengths & angle & lengths & angle & lengths & angle \\ \hline \multirow{8}{*}{\(|\bullet\rangle\langle\bullet|\subseteq\Gamma_{1}\rangle\)} & **0.084** & **1.266** & **0.115** & **1.437** & 0.290 & 5.487 \\  & \(\{|\gamma\rangle\langle\gamma|\subseteq\Gamma_{2}\rangle\)} & 0.056 & 0.596 & 0.053 & 0.283 & 0.287 & 5.209 \\  & \(\{|\bullet\rangle\langle\bullet|\subseteq\Gamma_{1}\rangle\cup\{|\gamma\rangle \langle\gamma|\subseteq\Gamma_{2}\}\)} & 0.063 & 0.454 & 0.063 & 0.270 & **0.296** & 5.733 \\  & \(\{|\bullet\rangle\langle\bullet|\subseteq\Gamma_{1}\rangle\cup\{|\gamma\rangle \langle\gamma|,|\gamma\rangle\langle\gamma^{\prime}|_{\text{sym}}\subseteq \Gamma_{2}\}\)} & 0.065 & 0.670 & 0.066 & 0.353 & **0.296** & 5.733 \\  & \(\{|\bullet\rangle\langle\bullet|\subseteq\Gamma_{1}\rangle\cup\{|\bullet \rangle\langle\bullet|\subseteq\Gamma_{2}\}\)} & 0.065 & 0.725 & 0.066 & 0.420 & **0.296** & **5.765** \\ \hline \multirow{8}{*}{\(\nabla\)} & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\) & 0.075 & 1.183 & 0.102 & **1.479** & 0.259 & 4.654 \\  & \(\{r_{\gamma},r_{\gamma^{\prime}}\subseteq\Gamma_{2}\}\) & 0.060 & 0.488 & 0.085 & 0.391 & 0.289 & **5.560** \\  & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\cup\{r_{\gamma},r_{\gamma^{\prime}}\subseteq \Gamma_{2}\}\) & 0.101 & 1.232 & 0.101 & 0.541 & 0.292 & 5.514 \\  & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\cup\{r_{\gamma},r_{\gamma^{\prime}},\mathcal{A} _{\gamma\gamma^{\prime}}\subseteq\Gamma_{2}\}\) & 0.087 & 1.093 & **0.106** & 0.717 & 0.265 & 4.990 \\  & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\cup\{r_{\gamma},r_{\gamma^{\prime}},\theta_{ \gamma\gamma^{\prime}}\subseteq\Gamma_{2}\}\) & **0.107** & **1.283** & 0.088 & 0.617 & **0.293** & 5.550 \\ \hline \multirow{8}{*}{\(\nabla_{\text{sym}}\)} & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\) & 0.083 & 1.307 & 0.064 & 0.816 & 0.281 & 5.134 \\  & \(\{r_{\gamma},r_{\gamma^{\prime}}\subseteq\Gamma_{2}\}\) & **0.100** & 1.188 & 0.101 & 0.503 & 0.281 & 4.959 \\  & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\cup\{r_{\gamma},r_{\gamma^{\prime}}\subseteq \Gamma_{2}\}\) & 0.097 & **1.375** & 0.098 & 0.672 & 0.226 & 3.188 \\  & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\cup\{r_{\gamma},r_{\gamma^{\prime}},\mathcal{A} _{\gamma\gamma^{\prime}}\subseteq\Gamma_{2}\}\) & 0.099 & 1.328 & **0.124** & **1.160** & 0.285 & 5.457 \\  & \(\{r_{\gamma}\subseteq\Gamma_{1}\}\cup\{r_{\gamma},r_{\gamma^{\prime}},\theta_{ \gamma\gamma^{\prime}}\subseteq\Gamma_{2}\}\) & **0.100** & 1.289 & -0.001 & -0.007 & **0.291** & **5.617** \\ \hline \multicolumn{2}{l|}{feed forward (FF)} & -0.191 & -5.277 & -0.304 & -3.304 & 0.303 & 6.438 \\ \hline \multicolumn{2}{l|}{DFT} & 0.164 & 5.442 & 0.345 & 5.648 & 0.150 & -1.446 \\ \end{tabular}
\end{table}
Table 2: Metrics are defined as the average improvement of lattice parameters. The experiment is split into five categories of vector fields: from the ket-bra \(|\bullet\rangle\langle\bullet|\), from the gradient of invariant geometric without symmetric action \(\nabla\), the gradient with symmetric action \(\nabla_{\text{sym}}\), lattice predicted by a FF readout function and lattice obtained from a DFT calculation with VASP.

\begin{table}
\begin{tabular}{c|c c|c c} model & \multicolumn{2}{c|}{carbon-24} & \multicolumn{2}{c}{mp-20} \\  & length & angles & length & angles \\ \hline EMPNN & **0.200** & **3.199** & **0.174** & **1.965** \\ baseline & 0.469 & 13.693 & 0.534 & 6.324 \\ \end{tabular}
\end{table}
Table 3: MAE between lattice parameters of the original cell and the reconstructed cell (Å and degree).
## 8 Acknowledgments

This work has been supported by ANR-22-CE23-0002 ERI-ANA, ANR-20-THIA-0004 and by HPC resources from GENCI-IDRIS (Grant 2022-[AD011013338]).
