A neural network potential with self-trained atomic fingerprints: a test
  with the mW water potential
We present a neural network (NN) potential based on a new set of atomic
fingerprints built upon two- and three-body contributions that probe distances
and local orientational order respectively. Compared to existing NN potentials,
the atomic fingerprints depend on a small set of tuneable parameters which are
trained together with the neural network weights. To tackle the simultaneous
training of the atomic fingerprint parameters and neural network weights we
adopt an annealing protocol that progressively cycles the learning rate,
significantly improving the accuracy of the NN potential. We test the
performance of the network potential against the mW model of water, which is a
classical three-body potential that well captures the anomalies of the liquid
phase. Trained on just three state points, the NN potential is able to
reproduce the mW model in a very wide range of densities and temperatures, from
negative pressures to several GPa, capturing the transition from an open random
tetrahedral network to a dense interpenetrated network. The NN potential also
reproduces very well properties for which it was not explicitly trained, such
as dynamical properties and the structure of the stable crystalline phases of
mW.
Francesco Guidarelli Mattioli, Francesco Sciortino, John Russo
# A neural network potential with self-trained atomic fingerprints:

###### Abstract

We present a neural network (NN) potential based on a new set of atomic fingerprints built upon two- and three-body contributions that probe distances and local orientational order respectively. Compared to existing NN potentials, the atomic fingerprints depend on a small set of tuneable parameters which are trained together with the neural network weights. To tackle the simultaneous training of the atomic fingerprint parameters and neural network weights we adopt an annealing protocol that progressively cycles the learning rate, significantly improving the accuracy of the NN potential. We test the performance of the network potential against the mW model of water, which is a classical three-body potential that well captures the anomalies of the liquid phase. Trained on just three state points, the NN potential is able to reproduce the mW model in a very wide range of densities and temperatures, from negative pressures to several \(GPa\), capturing the transition from an open random tetrahedral network to a dense interpenetrated network. The NN potential also reproduces very well properties for which it was not explicitly trained, such as dynamical properties and the structure of the stable crystalline phases of mW.

## I Introduction

Machine learning (ML) potentials represent one of the emerging trends in condensed matter physics and are revolutionising the landscape of computational research. Nowadays, different methods to derive ML potentials have been proposed, providing a powerful methodology to model liquids and solid phases in a large variety of molecular systems [1; 2; 3; 4; 5; 6; 7; 8; 9; 10; 11; 12; 13; 14; 15; 16; 17]. Among these methods, probably the most successful representation of a ML potential so far is given by Neural Network (NN) potentials, where the potential energy surface is the output of a feed-forward neural network [18; 19; 20; 21; 22; 23; 24; 25; 26; 27; 28; 29; 30; 31; 32; 33; 34; 35].

In short, the idea underlying NN potentials construction is to train a neural network to represent the potential energy surface of a target system. The model is initially trained on a set of configurations generated ad-hoc, for which total energies and forces are known, by minimizing a suitable defined loss-function based on the error in the energy and force predictions. If the training set is sufficiently broad and representative, the model can then be used to evaluate the total energy and forces of any related atomic configuration with an accuracy comparable to the original potential. Typically the original potential will include additional degrees of freedom, such as the electron density for DFT calculations, or solvent atoms in protein simulations, which make the full computation very expensive. By training the network only on a subset of the original degrees of freedom one obtaines a coarse-grained representation that can be simulated at a much reduced computational cost. NN potentials thus combine the best of two worlds, retaining the accuracy of the underlying potential model, at the much lower cost of coarse-grained classical molecular dynamics simulations. The accuracy of the NN potential depends crucially on how local atomic positions are encoded in the input of the neural network, which needs to retain the symmetries of the underlying Hamiltonian, i.e. rotational, translational, and index permutation invariance. Several methods have been proposed in the literature [36; 12], such as the approaches based on the Behler-Parrinello (BP) symmetry functions [18], the Smooth Overlap of Atomic Positions (SOAP) [37], N-body iterative contraction of equivariants (NICE) [38] and polynomial symmetry functions [39], or frameworks like the DeepMD [23], SchNet [22] and RuNser [18]. In all cases, atomic positions are transformed into atomic fingerprints (AFs). The choice of the AFs is particularly relevant, as it greatly affects the accuracy and generality of the resulting NN potential.

We develop here a fully learnable NN potential in which the AFs, while retaining the simplicity of typical local fingerprints, do not need to be fixed beforehand but instead are learned during the training procedure. The coupled training of the atomic fingerprint parameters and of the network weights makes the NN training process more efficient since the NN representation is spontaneously built on a variable atomic fingerprint representation. To tackle the combined minimization of the AF parameters and of the network weights we adopt an efficient annealing procedure, that periodically cycles the learning rate, i.e. the step size of the minimization algorithm, resulting in a fast and accurate training process.

We validate the NN potential on the mW model of water [40], which is a one-site classical potential that has found widespread adoption to study water's anomalies [41; 42] and crystallization phenomena [43; 44]. Since the first pioneering MD simulations [45; 46], water is often chosen as a prototypical case study, as the large number of distinct local structures that are compatible with its tetrahedral coordination make it the molecule with the most complex thermodynamic behavior [47], for ex
ample displaying a liquid-liquid critical point at supercooled conditions [48; 49; 50; 51; 52]. NN potentials for water have been developed starting from density functional calculations, with different levels of accuracy [53; 54; 55; 56; 57; 58; 59; 60]. NN potentials have also been proposed to parametrise accurate classical models for water with the aim of speeding up the calculations when multi-body interactions are included [61], as in the MBpol model [62; 63; 64] or for testing the relevance of the long range interactions, as for the SPC/E model [65]. We choose the mW potential as our benchmark system because its explicit three-body potential term offers a challenge to the NN representation that is not found in molecular models built from pair-wise interactions. We stress that we train the NN-potential against data which can be generated easily and for which structural and dynamic properties are well known (or can be evaluated with small numerical errors) in a wide range of temperatures and densities. In this way, we can perform a quantitative accurate comparison between the original mW model and the hereby proposed NN model.

Our results show that training the NN potential at even just one density-temperature state point provides an accurate description of the mW model in a surrounding phase space region that is approximately a hundred kelvins wide. A training based on three different state points extends the convergence window extensively, accurately reproducing state points at extreme conditions, i.e. large negative and (crushingly) positive pressures. We will show that the NN reproduces thermodynamic, structural and dynamical properties of the mW liquid state, as well as structural properties of all the stable crystalline phases of mW water.

The paper is organized as follows. In Section II we describe the new atomic fingerprints and the details about the Neural Network potential implementation, including the _warm restart_ procedure used to train the weights and the fingerprints at the same time. In Section III we present the results, which include the accuracy of the models built from training sets that include one or three state points, and a comparison of the thermodynamic, structural and dynamic properties with those of the original mW model. We conclude in Section IV.

## II The Neural Network Model

The most important step in the design of a feed-forward neural network potential is the choice on how to define the first and the last layers of the network, respectively named the _input_ and _output_ layers. We start with the output layer, as it determines the NN potential architecture to be constructed. Here we follow the Behler Parrinello NN potential architecture [18], in which the total energy of the system is decomposed as the sum of local fields (\(E_{i}\)), each one representing the contribution of a local environment centered around atom \(i\). Being this a many-body contribution, it is important to note that \(E_{i}\) is not the energy of the single atom \(i\), but of all its environment (see also the Appendix A). With this choice, the total energy of the system is simply the sum over all atoms, \(E=\sum E_{i}\), and the force \(\vec{f}_{i}\) acting on atom \(i\) is the negative gradient of the _total_ energy with respect to the coordinates \(\nu\) of atom \(i\), e.g. \(f_{i\nu}=\partial E/\partial x_{i\nu}\). We have to point out that a NN potential is differentiable and hence it is possible to evaluate the gradient of the energy analytically. This allows to compute forces of the NN potential in the same way of other force fields, e.g. by the negative gradient of the total potential energy.

The input layer is built from two-body (distances) and three-body (angles) descriptors of the local environment, \(\vec{D}^{(i)}\) and \(\vec{T}^{(i)}\) respectively, ensuring translational and rotational invariance. The first layer of the neural network is the _Atomic Fingerprint Constructor_ (_AFC_), as shown in Fig. 1, which applies an exponential weighting on the atomic descriptors, restoring the invariance under permutations of atomic indexes. The outputs of this first layer are the atomic fingerprints (AFs) and in turn these are given to the first _hidden layer_. We will show how this organization of the AFC layer allows for the internal parameters of the exponential weighting to be trained together with the weights in the hidden layers of the network. In the following we describe in detail the construction of the inputs and the calculation flow in the first layers.

### The atomic fingerprints

The choice of input layer presents considerably more freedom, and it is here that we deviate from previous NN potentials. The data in this layer should retain all the information needed to properly evaluate forces and energies of the particles in the system, possibly exploiting the internal symmetries of the Hamiltonian (which in isotropic fluids are the rotational, translational and permutational invariance) to reduce the number of degenerate inputs. Given that the output was chosen as \(E_{i}\), the energy of the atomic environment surrounding atom \(i\), the input uses an atom-centered representation of the local environment of atom \(i\).

In the input layer, we define an atom-centered representation of the local environment of atom \(i\), considering both the distances \(r_{ij}\) with the nearest neighbours \(j\) within a spatial cut-off \(R_{c}\), and the angles \(\theta_{jik}\) between atom \(i\) and the pair of neighbours \(jk\) that are within a cut-off \({R_{c}}^{\prime}\). More precisely, for each atom \(j\) within \(R_{c}\) from \(i\) we calculate the following descriptors

\[D_{j}^{(i)}(r_{ij};R_{c})=\begin{cases}\frac{1}{2}\left[1+\cos\left(\pi\frac{r _{ij}}{R_{c}}\right)\right]&r_{ij}\leq R_{c}\\ 0&r_{ij}>R_{c}\end{cases} \tag{1}\]

and, for each triplet \(j-i-k\) within \({R_{c}}^{\prime}\) from \(i\),

\[T_{jk}^{(i)}(r_{ij},r_{ik},\theta_{jik})= \tag{2}\] \[\frac{1}{2}\left[1+\cos\left(\theta_{jik}\right)\right]\ D_{j}^{(i )}(r_{ij};{R_{c}}^{\prime})\ D_{k}^{(i)}(r_{ik};{R_{c}}^{\prime})\]
Here \(i\) indicates the label of i-th particle while index \(j\) and \(k\) run over all other particles in the system. In Eq. 1, \(D_{j}^{(i)}(r_{ij};R_{c})\) is a function that goes continuously to zero at the cut-off (including its derivatives). The choice of this functional form guarantees that \(D_{j}^{(i)}\) is able to express contributions even from neighbours close to the cut-off. Other choices, based on polynomials or other non-linear functions, have been tested in the past [31]. For example, we tested a parabolic cutoff function which produced considerably worse results than the cutoff function in Eq. 1. The function \(T_{jk}^{(i)}(r_{ij},r_{ik},\theta_{jik})\) is also continuous at the triplet cutoff \(R_{c}\). The angular function \(\frac{1}{2}\left[1+\cos\left(\theta_{jik}\right)\right]\) guarantees that \(0\leq T_{jk}^{(i)}(r_{ij},r_{ik},\theta_{jik})\leq 1\). We note that the use of relative distances and angles in Eq. 1-2 guarantees translational and rotational invariance.

The pairs and triplets descriptors are then fed to the _AFC layer_ to compute the atomic fingerprints, AFs. These are computed by projecting the \(D_{j}^{(i)}\) and \(T_{jk}^{(i)}\) descriptors on a exponential set of functions defined by

Figure 1: Schematic representation of the Neural Network Potential flow. (A) Starting from the relative distances and the triplets angles between neighbouring atoms, the input layer evaluates the atomic descriptors \(\vec{D}^{(i)}=\{D_{j}^{(i)}\) (Eq. 1) and \(\mathbf{T}^{(i)}=\{T_{jk}^{(i)}\}\) (Eq. 2). (B) The first layer is the Atomic Fingerprint Constructed (AFC) and it combines the atomic descriptors into atomic fingerprints, weighting them with an exponential function. The red nodes perform the calculation of Eq. 5, where from the two-body descriptors a weighting vector \(\vec{D}_{w}^{(i)}(\alpha)=\{e^{\alpha D_{j}^{(i)}}\}\) is calculated (square with \(\alpha\)) and then the scalar product \(\vec{D}^{(i)}\cdot\vec{D}_{w}^{(i)}(\alpha)\) is computed (square with point) and finally a logarithm is applied (circle). The blue nodes perform the calculation of Eq. 7, where two weighting vectors are calculated from the two-body descriptors namely \(\vec{D}_{w}^{(i)}(\gamma)\) and \(\vec{D}_{w}^{(i)}(\delta)\) and one weighting matrix from the three-body descriptors \(\mathbf{T}_{w}^{(i)}(\beta)=\{e^{\vec{\sigma}_{jk}^{(i)}}/2\}\). Finally in the compression unit (Eq 6) values are combined as \(0.5[\vec{D}^{(i)}\circ\vec{D}_{w}^{(i)}(\gamma)]^{T}[\mathbf{T}^{(i)}\circ\mathbf{T}_ {w}^{(i)}(\beta)][\vec{D}^{(i)}\circ\vec{D}_{w}^{(i)}(\delta)]\) where we use the circle symbol for the element-wise multiplication. The output value of the compression unit is given to the logarithm function (circle). The complete network (D) is made of ten AFC units and two hidden layers with 25 nodes per layer and here is depicted 2.5 times smaller.
\[\overline{D}^{(i)}(\alpha) = \ln\left[\sum_{j\neq i}D_{j}^{(i)}e^{\alpha D_{j}^{(i)}}+\epsilon \right]-Z_{\alpha} \tag{3}\] \[\overline{T}^{(i)}(\beta,\gamma,\delta) = \ln\left[\sum_{j\neq k\neq i}\frac{T_{jk}^{(i)}e^{\beta T_{jk}^{(i )}}e^{\gamma D_{j}^{(i)}}e^{\delta D_{k}^{(i)}}}{2}+\epsilon\right]\] (4) \[-Z_{\beta\gamma\delta}\]

These AFs are built summing over all pairs and all triplets involving particle \(i\), making them invariant under permutations, and multiplying each descriptor by an exponential filter whose parameters are called \(\alpha\) for distance AFs, and \(\beta,\ \gamma,\ \delta\) for the triplet AFs. These parameters play the role of feature selectors, i.e. by choosing an appropriate list of \(\alpha,\ \beta,\ \gamma,\ \delta\) the AFs can extract the necessary information from the atomic descriptors. The best choice of \(\alpha,\ \beta,\ \gamma,\ \delta\) will emerge automatically during the training stage. In Eqs. 3-4, the number \(\epsilon\) is set to \(10^{-3}\) and fixes the value of energy in the rare event that no neighbors are found inside the cutoff. Parameters \(Z_{\alpha}\) and \(Z_{\beta\gamma\delta}\) are optimized during the training process, shifting the AFs towards positive or negative values, and act as normalization factors that improve the representation of the NN.

The definitions in equations 3-4 can be reformulated in terms of product between vectors and matrices in the following way. The descriptors in equations 1-2 for particle i can be represented as a vector \(\vec{D}^{(i)}=\{D_{j}^{(i)}\}\) and a matrix \(\mathbf{T}^{(i)}=\{T_{jk}^{(i)}\}\) respectively. Given a choice of \(\alpha\), \(\beta\), \(\gamma\) and \(\delta\), three weighting vector \(\vec{D}_{w}^{(i)}(\alpha)=\{e^{\alpha D_{j}^{(i)}}\}\), \(\vec{D}_{w}^{(i)}(\gamma)=\{e^{\gamma D_{j}^{(i)}}\}\) and \(\vec{D}_{w}^{(i)}(\delta)=\{e^{\delta D_{j}^{(i)}}\}\) and one weighting matrix \(\mathbf{T}_{w}^{(i)}(\beta)=\{e^{\beta T_{jk}^{(i)}}/2\}\) are calculated from \(\vec{D}^{(i)}\) and \(\mathbf{T}^{(i)}\). The 2-body atomic fingerprint (Eq. 3) is finally computed as

\[\overline{D}^{(i)}(\alpha)=\ln\left[\vec{D}^{(i)}\cdot\vec{D}_{w}^{(i)}( \alpha)+\epsilon\right]-Z_{\alpha} \tag{5}\]

The 3-body atomic fingerprint (Eq. 4) is computed first by what we call _compression_ step in Fig. 1 as

\[\overline{T}_{c}^{(i)}=\frac{[\vec{D}^{(i)}\circ\vec{D}_{w}^{(i)}(\gamma)]^{T }[\mathbf{T}^{(i)}\circ\mathbf{T}_{w}^{(i)}(\beta)][\vec{D}^{(i)}\circ\vec{D}_{w}^{(i )}(\delta)]}{2} \tag{6}\]

and finally by

\[\overline{T}^{(i)}(\beta,\gamma,\delta)=\ln\left[\overline{T}_{c}^{(i)}(\beta, \gamma,\delta)+\epsilon\right]-Z_{\beta\gamma\delta} \tag{7}\]

where we use the circle symbol for the element-wise multiplication. The NN potential flow is depicted in Figure 1 following the vectorial representation.

In summary, our AFs select the local descriptors useful for the reconstruction of the potential by weighting them with an exponential factor tuned with exponents \(\alpha\), \(\beta\), \(\gamma\), \(\delta\). A similar weighting procedure has been showed to be extremely powerful in the selection of complex patterns and is widely applied in the so-called _attention layer_ first introduced by Google Brain [66]. However the AFC layer imposes additionally physically motivated constraints on the neural network representation.

We note that the expression for the system energy is a sum over the fields \(E_{i}\), but the local fields \(E_{i}\) are not additive energies, involving all the pair distances and triplets angles within the cut-off sphere centered on particle \(i\). This non-additive feature favours the NN ability to capture higher order correlations (multi-body contribution to the energy), and has been shown to outperform additive models in complex datasets [67]. The NN non-additivity requires the derivative of the whole energy \(E\) (as opposed to \(E_{i}\)) to estimate the force on a particle \(i\). In this way, contributions to the force on particle \(i\) come not only from the descriptors of \(i\) but also from the descriptors of all particles who have \(i\) as a neighbour, de facto enlarging the effective region in space where interaction between particles are included. This allows the network to include contributions from length-scales larger than the cutoffs that define the atomic descriptors. The Appendix A provides further information on this point.

### Hidden layers

We employ a standard feed-forward fully-connected neural network composed of two hidden layers with 25 nodes per layer and using the hyperbolic tangent (tanh) as the activation function. The nodes of the first hidden layer are fully connected to the ones in the second layer, and these connections have associated weights \(W\) which are optimized during the training stage.

The input of the first hidden layer is given by the AFC layer where we used five nodes for the two-body AFs (Eq. 3) and five nodes for the three-body AFs (Eq. 4) for a total of 10 AFs for each atom. We explore the performance of some combinations for the number of two-body and three-body AF in Appendix D and we find that the choice of five and five is the more efficient.

The output is the local field \(E_{i}\), for each atomic environment \(i\), whose sum \(E=\sum_{i=1}^{N}E_{i}\) represents the NN estimate of the potential energy \(E\) of the whole system.

### Loss function and training strategy

To train the NN-potential we minimize a loss function computed over \(n_{f}\) frames, i.e. the number of independent configurations extracted from an equilibrium simulation of the liquid phase of the target potential (in our case the mW potential). The loss function is the sum of two contributions.

The first contribution, \(H[\{\Delta\epsilon^{k},\Delta f_{i\nu}^{k}\}]\), expresses the difference in each frame \(k\) between the NN estimates and the target values for both the total potential energy (normalized by total number of atoms) \(\epsilon^{k}\) and the atomic
forces \(f_{i\nu}^{k}\) acting in direction \(\nu\) on atom \(i\). The \(n_{f}\) energy \(\epsilon^{k}\) values and \(3Nn_{f}\) force \(f_{i\nu}^{k}\) values are combined in the following expression

\[H[\{\Delta\epsilon^{k},\Delta f_{i\nu}^{k}\}]=\frac{p_{e}}{n_{f}} \sum_{k=1}^{n_{f}}h_{\text{Huber}}(\Delta\epsilon^{k})+\] \[\frac{p_{f}}{3Nn_{f}}\sum_{k=1}^{n_{f}}\sum_{i=1}^{N}\sum_{\nu=1} ^{3}h_{\text{Huber}}(\Delta f_{i\nu}^{k}) \tag{8}\]

where \(p_{e}=0.1\) and \(p_{f}=1\) control the relative contribution of the energy and the forces to the loss function, and \(h_{\text{Huber}}(x)\) is the so-called Huber function

\[h_{\text{Huber}}(x)=\begin{cases}0.5x^{2}\;\text{ if }|x|\leq 1\\ 0.5+(|x|-1)\;\text{ if }|x|>1\end{cases} \tag{9}\]

\(p_{e}\) and \(p_{f}\) are hyper-parameters of the model, and we selected them with some preliminary tests that found those values to be near the optimal ones. The Huber function [68] is an optimal choice whenever the exploration of the loss function goes through large errors caused by outliers, i.e. data points that differ significantly from previous inputs. Indeed when a large deviation between the model and data occur, a mean square error minimization may gives rise to an anomalous trajectory in parameters space, largely affecting the stability of the training procedure. This may happen especially in the first part of the training procedure when the parameter optimization, relaxing both on the energy and forces error surfaces may experience some instabilities.

The second contribution to the loss function is a regularization function, \(R[\{\alpha^{l},\beta^{m},\gamma^{m},\delta^{m}\}]\), that serves to limit the range of positive values of \(\alpha^{l}\) and of the triplets \(\beta^{m},\gamma^{m},\delta^{m}\) (where the indexes \(l\) and \(m\) run over the five different values of \(\alpha\) and five different triplets of values for \(\beta\), \(\gamma\) and \(\delta\)) in the window \(-\infty\) to \(5\). To this aim we select the commonly used relu function

\[r_{\text{relu}}(x)=\begin{cases}x-5\;\text{ if }x>5\\ 0\;\text{ if }x\leq 5\end{cases} \tag{10}\]

and write

\[R[\{\alpha^{l},\beta^{m},\gamma^{m},\delta^{m}\}]=\sum_{l=1}^{5} r_{\text{relu}}(\alpha^{l})+\] \[\sum_{m=1}^{5}[r_{\text{relu}}(\beta^{m})+r_{\text{relu}}( \gamma^{m})+r_{\text{relu}}(\delta^{m})] \tag{12}\]

Thus, the \(R\) function is activated whenever one parameters of the _AFC layer_ becomes, during the minimization, larger than \(5\).

To summarize, the global loss function \(\mathcal{L}\) used in the training of the NN is

\[\mathcal{L}[\epsilon,f]=H[\{\Delta\epsilon^{k},\Delta f_{i\nu}^{k}\}]+p_{b}R [\{\alpha^{l},\beta^{m},\gamma^{m},\delta^{m}\}] \tag{13}\]

where \(p_{b}=1\) weights the relative contribution of \(R\) compared to \(H\).

Compared to a standard NN-potential, we train not only the network weights \(W\) but also the AFs parameters \(\Sigma\equiv\{\alpha^{l},\beta^{m},\gamma^{m},\delta^{m}\}\) at the same time. The simultaneous optimization of the weights W and AFs \(\Sigma\) prevents possible bottleneck in the optimisation of W at fixed representation of \(\Sigma\). Other NN potential approaches implement a separate initial procedure to optimise the \(\Sigma\) parameters followed by the optimisation of W at fixed \(\Sigma\)[69]. The two-step procedure not only requires a specific methodological choice for optimising \(\Sigma\), but also may not result in the optimal values, compared to a search in the full parameter space (i.e. both \(\Sigma\) and W). Since the complexity of the loss function has increased, we have investigated in some detail some efficient strategies that lead to a fast and accurate training. Firstly, we initialize the parameters \(W\) via the Xavier algorithm, in which the weights are extracted from a random uniform distribution [70]. To initialize the \(\Sigma\) parameters we used a uniform distribution in interval \([-5,5]\). We then minimize the loss function using the _warm restart procedure_ proposed in reference [71]. In this procedure, the learning rate \(\eta\) is reinitialized at every cycle \(l\) and inside each cycle it decays as a function of the number of training steps \(t\) following

\[\eta^{(l)}(t)=A_{l}\left\{\frac{(1-\xi_{f})}{2}\left[1+\cos\left( \frac{\pi t}{T_{l}}\right)\right]+\xi_{f}\right\} \tag{14}\] \[0\leq t\leq T_{l}\]

where \(\xi_{f}=10^{-7}\), \(A_{l}=\eta_{0}\xi_{0}^{l}\) is the initial learning rate of the \(l\)-th cycle with \(\eta_{0}=0.01\) and \(\xi_{0}=0.9\), \(T_{l}=b\tau^{l}\) is the period of the \(l\)-th cycle with \(\tau=1.4\) and \(b=40\). The absolute number of training steps \(n\) during cycle \(l\) can be calculated summing over the length of all previous cycles as \(n=\tau+\sum_{m=0}^{l-1}T_{m}\).

We also select to evaluate the loss function for groups of four frames (mini-batch) and we randomly select \(200\) frames \(n_{f}=200\) for a system of \(1000\) atoms and hence we split this dataset in \(160\) frames (\(\%80\)) for the training set and the \(40\) frames (\(\%20\)) for the test set.

In Fig. 2(A) we represent the typical decay of the learning rate of the warm restart procedure, which will be compared to the standard exponential decay protocol in the Results section.

### The Target Model

To test the quality of the proposed novel NN we train the NN with data produced with the mW [40] model
of water. This potential, a re-parametrization of the Stillinger-Weber model for silicon [72], uses a combination of pairwise functions complemented with an additive three-body potential term

\[E=\sum_{i}\sum_{j>i}U_{2}(r_{ij})+\lambda\sum_{i}\sum_{j\neq i}\sum_{j>k}U_{3} \left(r_{ij},r_{ik},\theta_{jik}\right) \tag{15}\]

where the two body contribution between two particles \(i\) and \(j\) at relative distance \(r_{ij}\) is a generalized Lennard-Jones potential

\[U_{2}\left(r_{ij}\right)=A\epsilon\left[B\left(\frac{\sigma}{r_{ij}}\right)^{ p}-\left(\frac{\sigma}{r_{ij}}\right)^{q}\right]exp\left(\frac{\sigma}{r_{ij}-a \sigma}\right) \tag{16}\]

where the \(p=12\) and \(q=6\) powers are substituted by \(q=0\) and \(p=4\), multiplied by an exponential cut-off that brings the potential to zero at \(a\sigma\), with \(a=1.8\) and \(\sigma=2.3925\) A. \(A\epsilon\) (with \(A=7.049556277\) and \(\epsilon=6.189\) kcal mol\({}^{-1}\)) controls the strength of the two body part. B controls the two-body repulsion (with \(B=0.6022245584\)).

The three body contribution is computed from all possible ordered triplets formed by the central particle with the interacting neighbors (with the same cut-off \(a\sigma\) as the two-body term) and favours the tetrahedral coordination of the atoms via the following functional form

\[U_{3}\left(r_{ij},r_{ik},\theta_{jik}\right)=\epsilon\left[\cos \left(\theta_{jik}\right)-\cos\left(\theta_{0}\right)\right]^{2}\times\\ \exp\left(\frac{\gamma\sigma}{r_{ij}-a\sigma}\right)\exp\left( \frac{\gamma\sigma}{r_{ik}-a\sigma}\right) \tag{17}\]

where \(\theta_{jik}\) is the angle formed in the triplet \(jik\) and \(\gamma=1.2\) controls the smoothness of the cut-off function on approaching the cut-off. Finally, \(\theta_{0}=109.47^{\circ}\) and \(\lambda=23.15\) controls the strength of the angular part of the potential.

Figure 3: Comparison of the root mean square error calculated on the validation set for 60 replicas differing in the initial seed of the training procedure using both an exponential decay of the learning rate (points) and the warm restart method (squares), for the energy (panel A) and for the forces (panel B). For the forces, a significant improvement both in the average error and in its variance is found for the warm restart schedule.

Figure 2: Model convergence properties: (A) Learning rate schedule (Eq. 14) as a function of the absolute training step \(n\) (one step is defined as an update of the network parameters). (B) The training and validation loss (see \(\mathcal{L}[\epsilon,f]\) in Eq. 13) evolution during the training procedure, reported as a function of the number of epoch \(n_{e}\) (an epoch is defined as a complete evaluation of the training dataset). Root mean square (RMS) error of the total potential energy per particle (C) and of the force cartesian components (D) during the training evaluated in the test dataset. Data in panels B-C-D refers to the NN3 model and the green point shows the best model location.
The mW model, with its three-body terms centered around a specific angle and non-monotonic radial interactions, is based on a functional form which is quite different from the radial and angular descriptors selected in the NN model. The NN is thus agnostic with respect to the functional form that describes the physical system (the mW in this case). But having a reference model with explicit three body contributions offers a more challenging target for the NN potential compared to potential models built entirely from pairwise interactions. The mW model is thus an excellent candidate to test the performance of the proposed NN potential.

## III Results

### Training

We study two different NN models, indicated with the labels NN1 and NN3, differing in the number of state points included in the training set. These two models are built with a cut-off of \(R_{c}=4.545\) A for the two-body atomic descriptors and a cut-off of \(R_{c}^{\prime}=4.306\) A for the three-body atomic descriptors. \(R_{c}^{\prime}\) is the same as the mW cutoff while \(R_{c}\) was made slightly larger to mitigate the suppression of information at the boundaries by the cutoff functions. The NN1 model uses only training information based on mW equilibrium configurations from one state point at \(\rho_{1}=1.07\) g cm\({}^{-3}\), \(T_{1}=270.9\) K where the stable phase is the liquid. The NN3 model uses training information based on mW liquid configurations in three different state points, two state points at \(\rho_{1}=0.92\) g cm\({}^{-3}\), \(T_{1}=221.1\) K and \(\rho_{2}=0.92\) g cm\({}^{-3}\), \(T_{2}=270.9\) K where the stable solid phase is the clathrate Si34/Si136 [73] and one state point at \(\rho_{3}=1.15\) g cm\({}^{-3}\), \(T_{2}=270.9\) K.

This choice of points in the phase diagram is aimed to improve agreement with the low temperature-low density as well as high density regions of the phase diagram. Importantly, all configurations come from either stable or metastable liquid state configurations. Indeed, the point at \(\rho_{2}=0.92\) g cm\({}^{-3}\), \(T_{2}=270.9\) K is quite close to the limit of stability (respect to cavitation) of the liquid state.

To generate the training set, we simulate a system of \(N=1000\) mW particles with a standard molecular dynamics code in the NVT ensemble, where we use a time step of 4 fs and run \(10^{7}\) steps for each state point. From these trajectory, we randomly select 200 configurations (frames) to create a dataset of positions, total energies and forces. We then split the dataset in the _training_ and in the _test_ data sets, the first one containing 80% of the data. We then run the training for 4000 epochs with a minibatch of 4 frames. At the end of every epoch, we check if the validation loss is improved and we save the model parameters. In Fig. 2 we plot the loss function for the training and test datasets (B), the root mean square error of the total energy per particle (C), and of the force (D) for the NN3 model. The results show that the learning rate schedule of Eq. 14 is very effective in reducing both the loss and error functions.

Interestingly, the neural network seems to avoid overfitting (i.e. the validation loss is decreasing at the same rate as the loss on the training data), and the best model (deepest local minimum explored), in a given window of training steps, is always found at the end of that window, which also indicates that the accuracy could be further improved by running more training steps. Indeed we found that by increasing the number of training steps by one order of magnitude the error in the forces decreases by a further 30%. Similar accuracy of the training stage is obtained also for the NN1 model (not shown).

The training procedure always terminates with an error on the test set equal or less than \(\Delta\epsilon\simeq 0.01\) kcal mol\({}^{-1}\) (0.43 meV) for the energy, and of \(\Delta f\simeq 1.55\) kcal mol\({}^{-1}\) nm\({}^{-1}\) (6.72 meV A\({}^{-1}\)) for the forces. These values are comparable to the state-of-the-art NN potentials [54; 55; 61; 23], and within the typical accuracy of DFT calculations [74].

We can compare the precision of our model with that of alternative NN potentials trained on a range of water models. An alternative mW neural network potential has been trained on a dataset made of 1991 configurations of 128 particles system at different pressure and temperature (including both liquid and ice structures) with Behler-Parinello symmetry functions [24]. The training of this model (which uses more atomic fingerprints and a larger cutoff radius) converged to an error in energy of \(\Delta\epsilon\simeq 0.0062\) kcal mol\({}^{-1}\) (0.27 meV), and \(\Delta f\simeq 3.46\) kcal mol\({}^{-1}\) nm\({}^{-1}\) (15.70 meV A\({}^{-1}\)) for the forces. In a recent work searching for liquid-liquid transition signatures in an ab-initio water NN model [55], a dataset of configurations spanning a temperature range of \(0-600\) K and a pressure range of \(0-50\) GPa was selected. For a system of 192 particles, the training converged to an error in energy of \(\Delta\epsilon\simeq 0.010\) kcal mol\({}^{-1}\) (0.46 meV), and \(\Delta f\simeq 9.96\) kcal mol\({}^{-1}\) nm\({}^{-1}\) (43.2 meV A\({}^{-1}\)) for the forces. In the NN model of MB-POL [61], a dataset spanning a temperature range from 198 K to 368 K at ambient pressure was selected. In this case, for a system of 256 water molecule, an accuracy of \(\Delta\epsilon\simeq 0.01\) kcal mol\({}^{-1}\) (0.43 meV) and \(\Delta f\simeq 10\) kcal mol\({}^{-1}\) nm\({}^{-1}\) (43.36 meV A\({}^{-1}\)) was reached. Finally, the NN for water at \(T=~{}300\) K used in Ref. [54], reached precisions of \(\Delta\epsilon\simeq 0.046\) kcal mol\({}^{-1}\) (2 meV) and \(\Delta f\simeq 25.36\) kcal mol\({}^{-1}\) nm\({}^{-1}\) (110 meV A\({}^{-1}\)).

While a direct comparison between NN potentials trained on different reference potentials is not a valid test to rank the respective accuracies, the comparisons above show that our NN potential reaches a similar precision in energies, and possibly an improved error in the force estimation.

The accuracy of the NN potential could be further improved by extending the size of the dataset and the choice of the state points. In fact, while the datasets in Ref. [54; 55; 61] have been built with optimized proce
dures, the dataset used in this study was prepared by sampling just one (NN1) or three (NN3) state-points. Also the size of the datasets used in the present work is smaller or comparable to the ones of Ref. [54; 55; 61].

In Fig. 3 we compare the error in the energies (A) and the forces (B) between sixty independent training runs using the standard exponential decay of the learning rate (points) and the warm restart protocol (squares). The figure shows that while the errors in the energy computations are comparable between the two methods, the warm restart protocol allows the forces to be computed with higher accuracy. Moreover we found that the warm restart procedure is less dependent on the initial seed and that it reaches deeper basins than the standard exponential cooling rate.

### Comparing NN1 with NN3

The NN potential model was implemented in a custom MD code that makes use of the tensorflow C API [75]. We adopted the same time step (4 fs), the same number of particles (\(N=1000\)) and the same number of steps (\(10^{7}\)) as for the simulations in the mW model.

As described in the Training Section, we compare the accuracy of two different training strategies: NN1 which was trained on a single state point, and NN3 which is instead trained on three different state point. In Fig. 4 we plot the energy error (\(\Delta\epsilon\)) between the NN potential and the mW model with both NN1 (panel A) and NN3 (panel B). Starting from NN1, we see that the model already provides an excellent accuracy for a large range of temperatures and for densities close to the training density. The biggest shortcoming of the NN1 model is at densities lower than the trained density, where the NN potential model cavitates and does not retain the long-lived metastable liquid state displayed by the mW model. We speculate that this behaviour is due to the absence of low density configurations in the training set, which prevents the NN potential model from correctly reproducing the attractive tails of the mW potential.

To overcome this limitation we have included two additional state points at low density in the NN3 model. In this case, Fig. 4B shows that NN3 provides a quite accurate reproduction of the energy in the entire explored density and temperature window (despite being trained only with data at \(\rho=0.92\) g cm\({}^{-3}\) and \(\rho=1.15\) g cm\({}^{-3}\)).

We can also compare the accuracy obtained during production runs against the accuracy reached during training, which was \(\Delta\epsilon\simeq 0.01\) kcal mol\({}^{-1}\). Fig. 4B shows the error is of the order of 0.032 kcal mol\({}^{-1}\) (1.3 meV), for density above the training set density. But in the density region between 0.92 and 1.15, the error is even smaller, around 0.017 kcal mol\({}^{-1}\) (0.7 meV) at the lowest density boundary.

We can thus conclude that the NN3 model, which adds to the NN1 model information at lower density and temperature, in the region where tetrahedality in the water structure is enhanced, is indeed capable to represent, with only three state points, a quite large region of the phase space, encompassing dense and stretched liquid states. This suggests that a training based on few state points at the boundary of the density/temperature region which needs to be studied is sufficient to produce a high quality NN model. In the following we focus entirely on the NN3 model.

### Comparison of thermodynamic, structural and dynamical quantities

In Fig. 5 we present a comparison of thermodynamic data between the mW model (squares) and its NN potential representation (points) across a wide range of state points. Fig. 5A plots the energy as function of density for temperatures ranging from melting to deeply supercooled conditions. Perhaps the most interesting result is that the NN potential is able to capture the energy minimum, also called the _optimal network forming density_, which is a distinctive anomalous property of water and other _empty liquids_[76].

Fig. 5(B) shows the pressure as a function of the temperature for different densities, comparing the mW with the NN3 model. Also the pressure shows a good agreement between the two models in the region of densities between \(\rho=0.92\) g cm\({}^{-3}\) and \(\rho=1.15\) g cm\({}^{-3}\), which, as for the energy, tends to deteriorate at \(\rho=1.22\) g cm\({}^{-3}\).

In the large density region explored, the structure of the liquid changes considerably. On increasing density, a transition from tetrahedral coordinated local structure, prevalent at low \(T\) and low \(\rho\), towards denser local envi

Figure 4: Comparison between the mW total energy and the NN1 model (A) and NN3 model (B) for different temperatures and densities. While the NN3 model is able to reproduce the mW total energy with a good agreement in a wide region of densities and temperatures, the NN1 provide a good representation only in a limited region of density and temperature values. Blue squares represent the state points used for building the NN models.
ronments with interstitial molecules included in the first coordination shell takes place. This structural change is well displayed in the radial distribution function, shown for different densities at fixed temperature in Fig. 6. Fig. 6 also shows the progressive onset of a peak around 3.5 A  developing on increasing pressure, which signals the growth of interstitial molecules, coexisting with open tetrahedral local structures [77; 78]. At the highest density, the tetrahedral peak completely merges with the interstitial peak. The NN3 model reproduces quite accurately all features of the radial distribution functions, maxima and minima positions and their relative amplitudes, at all densities, from the tetrahedral-dominated to the interstitial-dominated limits. In general, NN3 model reproduces quite well the mW potential in energies, pressure and structures and it appreciably deviates from mW pressures and energies quantities only at densities (above 1.15 g/cm\({}^{3}\)) which are outside of the training region.

To assess the ability of NN potential to correctly describe also the crystal phases of the mW potential, we compare in Fig. 7 the \(g(r)\) of mW with the \(g(r)\) of the NN3 model for four different stable solid phases [73]: hexagonal and cubic ice (\(\rho=1.00\) g cm\({}^{-3}\) and \(T=246\) K), the dense crystal SC16 (\(\rho=1.20\) g cm\({}^{-3}\) and \(T=234\) K) and the clathrate phase Si136 (\(\rho=0.80\) g cm\({}^{-3}\) and \(T=221\) K). The results, shown in Fig. 7, show that, despite no crystal configurations have been included in the training set, a quite accurate representation of the crystal structure at finite temperature is provided by the NN3 model for all distinct sampled lattices.

Finally, we compare in Fig. 8 the diffusion coefficient (evaluated from the long time limit of the mean square displacement) for the mW and the NN3 model, in a wide range of temperatures and densities, where water displays a diffusion anomaly. Fig. 8 shows again that, also for

Figure 5: Comparison between the mW total energy and the NN3 total energy as a function of density along different isotherm (A) and comparison between the mW pressure and the NN3 pressure as a function of temperature along different isochores (B). The relative error of the NN vs the mW potential grows with density, but remains within 3% even for densities larger than the densities used in the training set.

Figure 6: Comparison between the mW radial distribution functions \(g(r)\) and the NN3 \(g(r)\) at \(T=270.9\) K for four different densities. The tetrahedral structure (signalled by the peak at 4.54 Å ) progressively weakens in favour of an interstitial peak progressively growing at \(3.5-3.8\) Å. Different \(g(r)\) have been progressively shifted by two to improve clarity.

Figure 7: Comparison between the mW radial distribution functions \(g(r)\) and the NN3 \(g(r)\) for four different lattices: (A) hexagonal diamond (the oxygen positions of the ice \(I_{h}\)); (B) cubic diamond (the oxygen positions of the ice \(I_{c}\); (C) the SC16 crystal (the dense crystal form stable at large pressures in the mW model) and (D) the Si136 clathrate structure, which is stable at negative pressures in the mW model. Different \(g(r)\) have been progressively shifted by four to improve clarity.
dynamical quantities, the NN potential offers an excellent representation of the mW potential, despite the fact that no dynamical quantity was included in the training set. A comparison between fluctuations of energy and pressure of mW and NN3 potential is reported in Appendix B.

## IV Conclusions

In this work we have presented a novel neural network (NN) potential based on a new set of atomic fingerprints (AFs) built from two- and three-body local descriptors that are combined in a permutation-invariant way through an exponential filter (see Eq. 3-4). One of the distinctive advantages of our scheme is that the AF's parameters are optimized during the training procedure, making the present algorithm a self-training network that automatically selects the best AFs for the potential of interest.

We have shown that the added complexity in the concurrent training of the AFs and of the NN weights can be overcome with an annealing procedure based on the warm restart method [71], where the learning rate goes through damped oscillatory ramps. This strategy not only gives better accuracy compared to the commonly implemented exponential learning rate decay, but also allows the training procedure to converge rapidly independently from the initialisation strategies of the model's parameters.

Moreover we show in Appendix C that the potential hyper-surface of the NN model has the same smoothness as the target model, as confirmed by (i) the possibility to use the same timestep in the NN and in the target model when integrating the equation of motion and (ii) by the possibility of simulate the NN model even in the NVE ensemble with proper energy conservation.

We test the novel NN on the mW model [40], a one-component model system commonly used to describe water in classical simulations. This model, a re-parametrization of the Stillinger-Weber model for silicon [72], while treating the water molecule as a simple point, is able to reproduce the characteristic tetrahedral local structure of water (and its distortion on increasing density) via the use of three-body interactions. Indeed water changes from a liquid of tetrahedrally coordinated molecules to a denser liquid, in which a relevant fraction of interstitial molecules are present in the first nearest-neighbour shell. The complexity of the mW model, both due to its functional form as well as to the variety of different local structures which characterise water, makes it an ideal benchmark system to test our NN potential.

We find that a training based on configurations extracted by three different state points is able to provide a quite accurate representation of the mW potential hyper-surface, when the densities and temperatures of the training state points delimit the region of in which the NN potential is expected to work. We also find that the error in the NN estimate of the total energy is low, always smaller than \(0.03\) kcal mol\({}^{-1}\), with a mean error of \(0.013\) kcal mol\({}^{-1}\). The NN model reproduces very well not only the thermodynamic properties but also structural properties, as quantified by the radial distribution function, and the dynamic properties, as expressed by the diffusion coefficient, in the extended density interval from \(\rho=0.92\) g cm\({}^{-3}\) to \(\rho=1.22\) g cm\({}^{-3}\).

Interestingly, we find that the NN model, trained only on disordered configurations, is also able to properly describe the radial distribution of the ordered lattices which characterise the mW phase diagram, encompassing the cubic and hexagonal ices, the SC16 and the Si136 clathrate structure [73]. In this respect, the ability of the NN model to properly represent crystal states suggests that, in the case of the mW, and as such probably in the case of water, the geometrical information relevant to the ordered structures is contained in the sampling of phase space typical of the disordered liquid phase. These findings have been recently discussed in reference [80] where it has been demonstrated that liquid water contains all the building blocks of diverse ice phases.

We conclude by noticing that the present approach can be generalized to multicomponent systems, following the same strategy implemented by previous approaches [23; 18]. Work in this direction is underway.

###### Acknowledgements.

 FGM and JR acknowledge support from the European Research Council Grant DLV-759187 and CINECA grant ISCAB NNPROT.

Figure 8: Comparison between the mW diffusion coefficient \(D\) and the NN3 corresponding quantity for different temperatures and densities, in the interval \(221-271\) K. In this dynamic quantity, the relative error is, for all temperatures, around \(8\%\). Note also that in this \(T\) window the diffusion coefficient shows a clear maximum, reproducing one of the well-know diffusion anomaly of water. Diffusion coefficients have been calculated in the NVT ensemble using the same Andersen thermostat algorithm [79] for mW and NN3 potential.