Atomistic Graph Neural Networks for metals: Application to bcc iron

The prediction of the atomistic structure and properties of crystals including defects based on ab-initio accurate simulations is essential for unraveling the nano-scale mechanisms that control the micromechanical and macroscopic behaviour of metals. Density functional theory (DFT) can enable the quantum-accurate prediction of some of these properties, however at high computational costs and thus limited to systems of ∼ 1,000 atoms. In order to predict with quantum-accuracy the mechanical behaviour of nanoscale structures involving from thousands to several millions of atoms, machine learning interatomic potentials have been recently developed. Here, we explore the performance of a new class of interatomic potentials based on Graph Neural Networks (GNNs), a recent field of research in Deep Learning. Two state-of-the-art GNN models are considered, SchNet and DimeNet, and trained on an extensive DFT database of ferromagnetic bcc iron. We find that the DimeNet GNN Fe potential including three-body terms can reproduce with DFT accuracy the equation of state and the Bain path, as well as defected configurations (vacancy and surfaces). To the best of our knowledge, this is the first demonstration of the capability of GNN of reproducing the energetics of defects in bcc iron. We provide an open-source implementation of DimeNet that can be used to train other metallic systems for further exploration of the GNN capabilities. Introduction Recently, Graph Neural Networks (GNNs) have become one of the most active research fields in Artificial Intelligence1. GNNs are a class of Deep Learning methods introduced to analyze data which display a graph structure. Graphs represent the topology of a great variety of data structures in which objects (nodes) are connected with each other by some kind of relation (edges). Due to the very general nature of graphs, applications of GNNs are found in very different contexts, such as computer vision2–5, natural language processing6–9, social sciences10, 11, and natural sciences including biology12, particle physics13 and astrophysics14. The topology of a graph can also reflect that of atomistic crystal structures: indeed, a graph can be generated by connecting each atom (nodes) with its neighbors (edges), within a specified cutoff radius15. The Message Passing Neural Networks framework (MPNN)16 has been introduced as a common GNN paradigm for atomistic structures in quantum chemistry applications. Within an atomistic GNN, the atoms and their connections are associated with numerical lists of “features”, also named embedding vectors. Features are updated by the Message Passing framework, which is a two-step process. In the first step, each atom receives a message that is an aggregate of its neighbour’s embeddings. In the second step, an updated embedding of the atom is evaluated, by means of a function that depends on the message and on the current atomic embedding. By iterating this scheme n times, each atom will receive messages from atoms that are distant up to n connections, thus accounting for long-range interactions. A GNN model for atomistic graphs is therefore determined: (i) by the nature and the size of the embeddings, which convey the informative content of the specific atomistic system and (ii) by the operations it executes on them, i.e. the procedure used to aggregate and update embeddings. Once the aforementioned characteristics are defined, the model can be trained to predict the system potential energy surface (PES). A number of GNN models have been proposed in the recent years to model atomistic systems. Most of them were first introduced in molecular research and further applied to crystals. Deep Tensor Neural Networks (DTNN)17 and PhysNet18 *These authors contributed equally: L. Cian and G. Lancioni. ar X iv :2 10 9. 14 01 2v 1 [ co nd -m at . m tr lsc i] 2 8 Se p 20 21 aggregate the atomic embeddings by means of filters that ensure that the resulting message changes smoothly with respect to small changes of the interatomic distances. The main difference between DTNN and PhysNet lies in how distances are represented and how messages are aggregated. Crystal Graph Convolutional Neural Networks (CGCNN)19 were explicitly developed to deal with materials displaying a crystal structure, such as metals. Unlike DTNN and PhysNet, CGCNN considers both atomic (node) and edge embeddings; however, distances are not regularized with continuous functions: the range of the distances is partitioned in ten equally spaced segments, and interatomic distances are encoded within a single vector in which all components are zero but the one associated with the matching segment. Thus, this model lacks the ability to smoothly change the embeddings with respect to small displacements of the atomic positions. SchNet20 is based on DTNN and introduces continuous filter convolutions: distances are used as input of a neural layer that generates a continuous mapping to an embedding space. In an updated version21, periodic boundary conditions (PBCs) have been introduced, and the model has been applied to the prediction of formation energies of bulk crystals. Also inspired by SchNet and sharing its overall architecture, the “Neural message passing with edge updates”22 uses both node and edge embeddings in the form of a concatenation of the two connected atoms embeddings. This makes edge embeddings directional as they depend on the order of the concatenated elements. MatErials Graph Network (MEGNet)23 leverages a similar scheme, by incorporating both directional edge and node updates, while also introducing a global state vector which stores the molecule/crystal-level or state attributes, e.g. the temperature of the system. Updates of atoms, bonds and global state vector are performed in a sequence. All these approaches employ filters that rely only on the distance between pairs of atoms to aggregate and update the atomic embeddings. It is well-established that classical, empirical interatomic potentials24 that rely on pairwise interactions often fail to reproduce structural changes25 and some crucial properties of dislocations in metals26. In the case of phase transitions, the addition of directionality, i.e. angular dependence of the interatomic potential, as well as second nearest-neighbor interactions, has lead to the improved qualitative reproduction of quantum-mechanical PES25. Within the context of GNNs, there is a remarkable shortage of approaches that rely also on the angle between edges connecting atomic pairs. Embeddings of edges connecting triplets of atoms convey the angular information, and once they are updated via the message passing scheme, they can be used to update the atomic embeddings. With this aim, DimeNet27 also leverages the Directional Message (hence the name) by considering the direction of the pairwise connections and by introducing the angle between two edges connected within atomic triplets. DimeNet employs a continuous filter convolution by expanding both distances and angles in a Bessel-Fourier basis. However, to date, DimeNet has been applied merely to isolated molecules and has not been investigated to model crystals such as metals. Although GNNs have been scarcely explored in the context of interatomic potentials for metals, they introduce a number of advantages with respect to other ML methods15. First, interactions among neighbouring atoms are straightforwardly modeled as pair-wise connections. Previous ML approaches need to introduce specific geometrical descriptors of the environment around atoms (within a cut-off radius), such as atom-centered symmetry functions in the Neural Networks Potentials28, or bispectrum components and then smooth overlap of atomic positions (SOAP)29 in Gaussian Approximation Potentials (GAP)30. Second, iterating the process makes the model able to consider the contributions of distant atoms, so as to mimic the influence of long-range interactions beyond the cut-off distance that limits pairwise interactions. This can be easily achieved by stacking message-passing layers in the network. Previous ML approaches either lack these long-range contributions or account for them by adding extra long-range terms to the total energy, e.g. for electrostatic interactions31. Third, the GNN approach guarantees scalability of the system, as the pair-wise nature of the connections means that complex clusters of atoms can be modeled by simply increasing the number of iterations, at a limited computational cost. Finally, since the approach is only dependent on the relative positions of the atoms which determine the connections inside the cut-off radius, it is also invariant with respect to isometric transformations, i.e. reflections, translations, rotations, and combinations of those, and to permutation of atoms. Here, we use GNNs to explore their ability to reproduce with quantum-accuracy the potential energy surface (PES) of metals, by taking as a reference the challenging and technologically crucial example of ferromagnetic body-centered-cubic (BCC) iron. We consider SchNet as a prototypical GNN framework that is based on the distance of atomic pairs, and we consider DimeNet to assess the performance of a GNN scheme that also includes angular (three-body) interactions. To this purpose, we have implemented periodic boundary conditions (PBCs) and made the new DimeNet implementation that includes PBCs available at https://github.com/AilabUdineGit/GNN_atomistics/. In order to machine-learn the GNN interatomic potential, we use an existing database32 that was previously trained to develop a Gaussian Approximation Potential (GAP)33. The remainder of this paper is organized as follows. Section Results is divided in two main subsections: “Implementation and training” reports a summary of computational details, together with some performance metrics; “Testing the GNN interatomic potentials for bcc Fe” shows a comparative analysis of the networks based on their predi