Data efficiency and extrapolation trends in neural network interatomic
  potentials
Over the last few years, key architectural advances have been proposed for
neural network interatomic potentials (NNIPs), such as incorporating
message-passing networks, equivariance, or many-body expansion terms. Although
modern NNIP models exhibit small differences in energy/forces errors,
improvements in accuracy are still considered the main target when developing
new NNIP architectures. In this work, we show how architectural and
optimization choices influence the generalization of NNIPs, revealing trends in
molecular dynamics (MD) stability, data efficiency, and loss landscapes. Using
the 3BPA dataset, we show that test errors in NNIP follow a scaling relation
and can be robust to noise, but cannot predict MD stability in the
high-accuracy regime. To circumvent this problem, we propose the use of loss
landscape visualizations and a metric of loss entropy for predicting the
generalization power of NNIPs. With a large-scale study on NequIP and MACE, we
show that the loss entropy predicts out-of-distribution error and MD stability
despite being computed only on the training set. Using this probe, we
demonstrate how the choice of optimizers, loss function weighting, data
normalization, and other architectural decisions influence the extrapolation
behavior of NNIPs. Finally, we relate loss entropy to data efficiency,
demonstrating that flatter landscapes also predict learning curve slopes. Our
work provides a deep learning justification for the extrapolation performance
of many common NNIPs, and introduces tools beyond accuracy metrics that can be
used to inform the development of next-generation models.
Joshua A. Vita, Daniel Schwalbe-Koda
# Data efficiency and extrapolation trends in neural network interatomic potentials

###### Abstract

Over the last few years, key architectural advances have been proposed for neural network interatomic potentials (NNIPs), such as incorporating message-passing networks, equivariance, or many-body expansion terms. Although modern NNIP models exhibit small differences in energy/forces errors, improvements in accuracy are still considered the main target when developing new NNIP architectures. In this work, we show how architectural and optimization choices influence the generalization of NNIPs, revealing trends in molecular dynamics (MD) stability, data efficiency, and loss landscapes. Using the 3BPA dataset, we show that test errors in NNIP follow a scaling relation and can be robust to noise, but cannot predict MD stability in the high-accuracy regime. To circumvent this problem, we propose the use of loss landscape visualizations and a metric of loss entropy for predicting the generalization power of NNIPs. With a large-scale study on NequIP and MACE, we show that the loss entropy predicts out-of-distribution error and MD stability despite being computed only on the training set. Using this probe, we demonstrate how the choice of optimizers, loss function weighting, data normalization, and other architectural decisions influence the extrapolation behavior of NNIPs. Finally, we relate loss entropy to data efficiency, demonstrating that flatter landscapes also predict learning curve slopes. Our work provides a deep learning justification for the extrapolation performance of many common NNIPs, and introduces tools beyond accuracy metrics that can be used to inform the development of next-generation models.

## 1 Introduction

Machine learning (ML) has proven extremely valuable in the materials and chemical sciences as a tool for data analysis and generation [1, 2, 3, 4]. Particularly in atomistic simulations, ML-based models offer a compelling balance between high-accuracy, high-cost quantum chemistry calculations and low-accuracy, low-cost classical force fields [5, 6, 7]. Whereas several models based on kernel regression or Gaussian processes have been proposed [8, 9, 10, 11, 12, 13, 6], recent developments in neural network (NN) interatomic potentials (IPs) have shown promise due to their low inference time, scalability to large datasets, and high accuracy in predicting potential energy surfaces (PESes) [14, 7, 9]. These methods have been used for a variety of applications, including molecular simulation, excited-state dynamics, phase transitions, chemical reactions, and more [15, 16, 17, 18, 9].

Despite their successes, NNIPs still struggle with data efficiency and robust generalization. Over the last few years, several different model architectures were proposed to reduce errors in PES fitting, decrease the amount of data required to train the models, and improve predictions for configurations beyond the training domain. In particular, NN architectures incorporating physics concepts such as directional representations and equivariance [19, 20, 21, 22, 23] or many-body interactions [24, 25, 26] have gained popularity due to higher accuracy and data efficiency. Nevertheless, recent works show that accuracy metrics over datasets are insufficient to quantify the models' quality in production simulations and motivate the use of alternative metrics such as computational speed or simulation stability [27, 28, 29, 30, 31, 32]. Different NNIP models may have similar test accuracy, but completely different extrapolation ability [33, 34]. This begs the question: **which metrics can distinguish between NNIPs with similar test error but different extrapolation behavior?**
In this work, we investigate trends in robustness and extrapolation behavior of NNIPs, and propose a metric to predict their stability in production simulations and data efficiency. In particular, we provide the following contributions:

* Using literature data and two state-of-the-art NNIP models (NequIP and MACE), we show that extrapolation trends can be obtained from error metrics in the 3BPA dataset. For example, the NNIPs are able to recover the underlying PES despite being trained to noisy labels, and have extrapolation errors that follow scaling relations against in-domain accuracy metrics. Although these trends hold across different architectures, we show that this scaling relation is not recovered in the high-accuracy regime, and thus cannot be used to downselect robust models.
* To circumvent the limitations above, we propose that loss landscapes (LLs) can provide evaluation strategies for NNIPs beyond accuracy metrics. Qualitative inspection of LLs explains some heuristic training regimes for NNIPs, such as the use of higher weights for forces loss, weight cycling, or separating learning rates for different parts of the architecture, thus providing theoretical justification for certain hyperparameters.
* Using a large-scale study with NequIP and MACE, we show that flatness of LLs correlates with model robustness. In particular, we quantify the loss entropy around the optimized models and relate them to errors in the extrapolation regime. Furthermore, we show that molecular dynamics simulations of models with flatter LLs exhibit less unphysical behavior than their sharper counterparts.

This combination of theoretical justification, benchmarking, and a new metric can aid the development of newer NNIP architectures with higher accuracy, trainability, and robustness.

## 2 Background

**NNIPs:** NN-based force fields have been first proposed using feedforward NNs and symmetry-based representations [14]. In these systems, improvements are achieved by designing new representations to better capture the atomic environment [35, 36, 37, 38, 39, 40, 41]. More recently, message-passing neural networks (MPNNs) [42] showed remarkable ability to fit PESes using learned representations. In this area, most works compare models according to their accuracy with respect to standardized datasets, such as QM9 [43, 44], MD17 [8], and others. MPNN-based potentials often vary according to their interaction blocks, handling of symmetry operations, and general architectural choices.

**NN representation capacity and generalization:** although NNs have a large number of parameters, obtaining NN-based models with low generalization error is not uncommon. Preventing overfitting may require regularization techniques such as changing the loss function, e.g. with weight regularization or decay, augmenting the dataset, or using training protocols such as adaptive learning rates, dropout, and more [45, 46, 47]. However, these are not requirements to controlling the generalization error [48]. For example, NNs with good generalization capacity have been shown to overfit to random labels and input data [48]. This suggests that, given enough parameters, even architecturally regularized NNs exhibit wide representation capacity for arbitrary data sets. In some cases, however, training to noisy labels can only be overcome by adapting architectures, regularization techniques, and correcting the loss function [45].

**Loss landscapes (LLs):** the shape of the LL is strongly correlated with the trainability of NN architectures. The minimization of the empirical risk is easier when the (non-convex) LL is smoother, as it yields more predictive gradients [49, 50]. Furthermore, LLs with several local minima exhibit lower trainability than their smoother counterparts [51, 47]. Some works proposed that flatter LLs are also related to lower NN generalization error [51, 52, 53]. Although this relationship is often complicated by factors such as batch size [46], normalization [54], or weight decay, the sharpness of the LL is the most predictive of generalization error in NNs [55]. Whereas this sharpness/flatness can be quantified using the Hessian of the loss [46, 47] or assumptions of a prior on the weights [56], visualizations of the LL have proven useful to illustrate these minima with respect to the weight space [57, 58, 59, 60] without the cost associated to calculating the full Hessian. Analogies with potential energy landscapes have also been performed to explore the cost functions of other machine learning methods and explain the shape of these optimization landscapes [61, 62].

## 3 Methods

### Visualizing loss landscapes

The loss landscape \(\ell\) of a neural network can be plotted by evaluating the loss function \(\mathcal{L}\) along a trajectory between two parameter sets \(\mathbf{\theta}\) and \(\mathbf{\theta}^{\prime}\). The simplest approach is to linearly interpolate the weights [57], choosing a scalar \(t\in[0,1]\) such that \(\mathbf{\theta}(t)=(1-t)\mathbf{\theta}+t\mathbf{\theta}^{\prime}\). Then, the loss landscape \(\ell\) for a model becomes

\[\ell(t)=\mathcal{L}\left(\mathbf{\theta}(t)\right)=\mathcal{L}\left((1-t)\mathbf{ \theta}+t\mathbf{\theta}^{\prime}\right). \tag{1}\]
In this work, the loss \(\mathcal{L}\) is evaluated on the training set of the model with parameters \(\mathbf{\theta}\).

In the absence of the reference weights \(\mathbf{\theta}^{\prime}\), the LL can be constructed by sampling a random vector \(\mathbf{\delta}\) in the parameter space and plotting the LL around \(\mathbf{\theta}\) as

\[\ell(t;\mathbf{\delta})=\mathcal{L}\left(\mathbf{\theta}+t\mathbf{\delta}\right), \tag{2}\]

where the domain of \(t\) is appropriately chosen to span a neighborhood of \(\mathbf{\theta}\). This notion can be extended to 2D LLs by taking two orthogonal vectors \(\mathbf{\delta}_{1},\mathbf{\delta}_{2}\) such that

\[\ell(t_{1},t_{2};\mathbf{\delta}_{1},\mathbf{\delta}_{2})=\mathcal{L}\left(\mathbf{\theta} +t_{1}\mathbf{\delta}_{1}+t_{2}\mathbf{\delta}_{2}\right), \tag{3}\]

with scalars \(t_{1},t_{2}\) chosen to span a (two-dimensional) neighborhood of \(\mathbf{\theta}\). These approaches have been used to study the LLs of NN classifiers in image datasets, interpolate between sets of classifiers, and explore the loss function around degenerate minima [63, 64, 65, 46].

One challenge when analyzing LLs is comparing different models according to their parameters. Because activation functions such as ReLU allow for scale-invariance of NN weights, especially when coupled with batch normalization techniques, the magnitude of the vector \(\mathbf{\delta}\) is not transferable from model to model. This prevents a fair comparison of LLs, curvatures, and sharpness metrics. To account for this, we use the filter normalization technique proposed by Li et al. [58]. Therein, each random vector \(\mathbf{\delta}\) is normalized by the scale of each filter \(i\), in each layer \(j\) of the NN, i.e.

\[\mathbf{\delta}^{(i,j)}=\frac{\mathbf{\delta}^{(i,j)}}{\|\mathbf{\delta}^{(i,j)}\|}\|\mathbf{ \theta}^{(i,j)}\|, \tag{4}\]

where \(\|.\|\) is the Frobenius norm. Then, the LL is plotted according to the filter-normalized vector \(\mathbf{\bar{\delta}}\),

\[\ell(t;\mathbf{\bar{\delta}})=\mathcal{L}\left(\mathbf{\theta}+t\mathbf{\bar{\delta}} \right), \tag{5}\]

and analogously for 2D LLs.

Although informative, sampling LLs can have cost comparable to training the NN, since evaluating the loss for each interpolated weight in each direction is equivalent to one training epoch. Depending on the number of parameters and directions \(\mathbf{\delta}_{n}\) under analysis, the loss may be evaluated over the entire training dataset a large number of times.

### Quantitative comparison of loss landscapes

Despite the usefulness of visualizing loss landscapes, comparing them beyond qualitative insights requires metrics to differentiate them. The most commonly used metric in the field is the curvature of the LL [55], which is related to the magnitude of the eigenvalues of the loss Hessian. Although informative, the Hessian is a local property and cannot fully capture "valley-like" degeneracies in loss landscapes. Furthermore, computing the full Hessian for a system with millions of training parameters would be intractable. Thus, alternative metrics to compare LLs are required.

In the case of NNNPs, which are often trained both to energies and forces, two LLs are obtained per model, one for each target value being trained to. Although they cannot be completely disentangled, both LLs should be compared simultaneously to assess model performance. To derive one of such metrics, we propose the use of "loss entropy" [66, 51] to quantify loss flatness around optimized minima. Although variations of this quantity he been proposed, we use the formula

\[S(T)=k\log\left[\sum_{t}\exp\left(\frac{-\bar{\ell}\left(t\right)}{kT}\right) \right], \tag{6}\]

where \(S\) is the entropy of the loss landscape \(\bar{\ell}(t)\) computed with respect to energy or forces and averaged over \(N\) orthogonal, filter-normalized weight displacements,

\[\bar{\ell}\left(t\right)=\frac{1}{N}\sum_{n}\ell(t;\mathbf{\bar{\delta}}_{n}). \tag{7}\]
and \(kT\) is a weighting parameter that quantifies the "flatness" of the LL with respect to a certain acceptable threshold of training loss. This is similar to distributions of microstates accessible in a given "temperature," and ensures that low-loss states contribute to a much larger entropy than high-loss states.

As the entropy of the LL of a NNIP takes into account both energy and forces losses, we adopt a strategy similar to the one used during training of the NNs by balancing energy and forces losses with a weighted sum,

\[S=\alpha S_{E}(T_{E})+(1-\alpha)S_{F}(T_{F}), \tag{8}\]

where \(\alpha\) is a dimensionless parameter between 0 and 1 that weights the entropy of the energy loss (\(S_{E}\)) and forces loss (\(S_{F}\)). As these losses have different units, they have to be computed with different normalization parameters \(k_{E}T_{E}\) and \(k_{F}T_{F}\), each of which takes into account "thermal randomness" of the training loss. For simplicity, in this work, we adopt \(k=1\) as a dimensionless parameter, and assign the adequate units to the "thermal error" for better interpretability of the loss entropy.

### NNIPs and dataset

**NequIP**[23] is an equivariant NNIP that uses Clebsch-Gordon transformations and spherical harmonics to incorporate equivariance in the model. NequIP demonstrates state-of-the-art accuracy in several datasets, data efficiency, and has been employed to simulate a variety of organic and inorganic systems.

**MACE**[25] is an equivariant NNIP that uses higher-order messages to efficiently embed information beyond two-body interactions in traditional MPNNs. The model demonstrates state-of-the-art performance in a variety of benchmarks, faster learning, and competitive computational cost.

**Other NNIPs** proposed recently and not benchmarked in this study include Allegro [26], GemNet [67], DimeNet [68], HIP-NN [69], NewtonNet [70], BOTNet [24], SchNet [71], PaiNN [72], and others [21, 37, 39, 73].

Training details are provided in Appendix A.

**The 3BPA dataset**[28] under study in this work was chosen due to its previous use in the literature for benchmarking NNIP models in extrapolation behavior [24, 25, 26]. The benchmark involves training models on low-temperature samples and evaluating their performance on held-out samples from high-temperature simulations. Distributions of energies and forces of the 3BPA dataset are shown in Appendix D.1, Fig. S2. Additional results using the ANI-Al dataset [74] be found in Appendix D.2.

## 4 Results and Discussion

### Robustness to noise and extrapolation trends in NNIPs

When comparing NNIPs, metrics of interest typically include errors in predicting forces and energies of a test dataset, and are appropriately used as baselines for assessing model quality. However, accuracy metrics are not necessarily predictive of extrapolation power [29]. A first hypothesis to consider when analyzing NNIP extrapolation is whether NNIPs can learn a PES despite being trained on noisy data. Although this test does not measure extrapolation to out-of-distribution data, it verifies whether models are expected to overfit to corrupted training data, which would lower their robust generalization power. For example, NN-based classifiers can overfit to random labels in image datasets or to completely random inputs [48], even in architectures with good generalization error that are designed to prevent overfitting. Most NNIPs have enough parameters to memorize the training data, but standard regularization and architectural choices can curb overfitting in NNIPs, leading to lower generalization errors.

To test this hypothesis, we trained four different NNIPs to the 3BPA dataset and analyzed their _training_ error. Following the lead from the deep learning literature [48], we then gradually corrupted the labels of the training set by randomly adding a sample from \(\mathcal{N}(0,\sigma\cdot\sigma_{\text{DFT}})\) to the true forces, where \(\sigma_{\text{DFT}}\) is the standard deviation in the forces predicted by density functional theory (DFT), and \(\sigma\) is a scalar ranging from 0.0 to 0.1 (see Appendix D.1). In principle, NN regressors with arbitrary levels of expressivity (or absent regularization) could achieve low training error even in these noisy PESes. Figure 0(a) shows the error for NNIPs trained to the 3BPA dataset with corrupted forces, and tested on the original, un-corrupted data. When the forces are not corrupted, models exhibit reasonable training errors lower than 40 meV/A, as expected by their nominal performances in energy prediction [23, 25]. However, even small amounts of noise in the forces prevent the noisy dataset from being memorized with high accuracy, with the training loss plateauing instead of tending to zero. The ability of the models to predict the noisy forces saturates at the limit of the noise, indicating that these NNIPs do not memorize the high-frequency labels. On the other hand, when the test error of the
models trained with corrupted labels is computed with respect to the uncorrupted dataset, the error is substantially smaller than the noise baseline (see the "original" panel of Fig. 1a). Thus, the NNIPs under analysis are able to learn the underlying PES in the 3BPA dataset despite the added noise.

Contrary to the overfitting hypothesis, these results suggest that data redundancy in the training set may help NNIPs to "denoise" the data. To illustrate this effect, we show in Appendix C how data redundancy downplays the effect of external noise in a toy system. Fig. S1 shows how a large number of training points can counterbalance the effect of external noise when predicting the original, non-noisy data for the case of linear regressor models. In this case, the model averages out the noise and recovers the true function even at high levels of added noise. On the other hand, at the low-data regime, the regression model is unable to recover the true function, and its error quickly grows. Although the results from the linear model may not directly translate to the case of NNs, Fig. 1a shows that, to an extent, NNIPs are able to "denoise" energies from the dataset due to architectural, training, or implicit data regularization.

To verify if this behavior was specific to the 3BPA dataset or could generalize to other datasets, we corrupted the energies or forces of the ANI-Al dataset and trained different models on the noisy values (see Appendix D.2), obtaining similar results even when only energies were considered. As the ANI-Al example exhibited even better results than the 3BPA, we tested whether a drastic increase in the injected noise could still be denoised by the NNIPs under study. Following the results of Fig. S1, we trained the two NNIP architectures under study to PESes with energy noises up to twenty times higher than those in Fig. 1 (see Fig. S12), and up to twice the standard deviation of the original dataset distribution of the ANI-Al dataset. Although the distribution of per-atom energies shows that the noisy PES is completely different than the original one (Fig. S12b,c), all models succeeded in modeling the underlying PES below the error baseline (Fig. S12). As also illustrated by the toy example, the performance of the models degrades as extremely large amounts of noise are added. Nonetheless, errors with respect to the non-noisy dataset are remarkably low considering the corruption baselines.

The toy example from Appendix C can be considered an upper bound of the "dataset denoising" ability, given that a functional form of the inputs is known and the model could, in principle, fit perfectly to the data. As the trends of NNIPs trained on the 3BPA or ANI-Al potential approach this behavior, it can be concluded that the two NNIP architectures are able to "denoise" the external noise added to the datasets, possibly due to data redundancy. As generalization tests

Figure 1: **a**, Force RMSE values (eV/Å) for models trained to only the forces of increasingly noisy versions of the 3BPA dataset. The models trained on noisy data are then evaluated on their noisy training set (“noisy” panel) or on the original, uncorrupted dataset (“original”). The dashed lines correspond to the amount of noise that was added to the DFT forces in units of eV/Å. The legend denotes the maximum body order of the MACE models (\(v\)), the number of interaction layers in the NequIP models (\(n\)), or the maximum order of the model irreps (\(L\)). All MACE models in **a** used only 2 interaction layers. **b**, Relationship between extrapolation ability, quantified using the slope of the test errors on the 3BPA dataset with increasing temperature (“extrapolation slope”), and the test error at 300 K. Slopes were computed by performing a linear fit to values taken from the literature for MACE [25], BOTNet [24], and others [26] (see Fig. S3). The dashed line is a visual guide. **c**, Forces RMSE on the 300 K test set of several NequIp and MACE models with different architectures, but similar testing errors on the 300 K split. (top) Correlation between test errors on all splits (300, 600, and 1200 K) and test error on the 300 K split for models with similar test errors. (bottom) Average time-to-failure of MD simulation for each of the models with similar test errors.
assume the model is being tested on unseen data, it is not clear whether the accuracy reflects the quality of the model predictions or simply their ability to reproduce local environments existing in the training data. This is particularly important in the high-data regime, when the test dataset may be correlated to the train dataset in non-obvious ways.

To bypass this problem, alternative strategies were proposed to measure generalization power, including separating train-test splits according to sampling temperature, as is the case of the 3BPA dataset [28]. While testing a model on held-out samples from high-temperature simulations is typically considered an independent evaluation of its performance, we have found that extrapolation errors are correlated with low-temperature test errors across various model architectures. This is performed by fitting a linear model to the errors on the 3BPA testing sets at 300, 600, and 1200 K from the literature [24, 25, 26] (see Fig. S3), then using the slope of the fitted line as associated metric. Fig. 1b shows that all the models that were trained only to 3BPA frames at 300 K follow an approximate linear scaling relation between the extrapolation slope and the log of the low-temperature errors. This correlation between the low- and high-temperature data does not strictly preclude the use of the 3BPA dataset for assessing a fitted model's extrapolation abilities, as the extrapolation slope represents how much the accuracy of a given model degrades as the sampling temperature increases. Nevertheless, it suggests that data and model regularization effects may be enforcing extrapolation trends in wide error ranges. For example, a model like ACE [24, 41] is known to provide functional forms that aid extrapolation beyond the training data [28]. Despite this, the ACE model for 3BPA follows the same trends of more complex message-passing NNs. In fact, the generalization slopes are correlated to their low-temperature error regardless of significant architectural differences. This indicates that, for this benchmark, the root mean squared errors (RMSEs) at higher temperatures can be estimated from the test error at 300 K even without evaluating the model at the other test sets.

The exceptions to this rule are the two ANI models ("ANI-2x" and "ANI-pretrained"), which were pre-trained to the 8.9 million configurations from the ANI-2x dataset [75]. As seen in other fields of deep learning [76], the pre-trained models extrapolate better than all other models, though fine-tuning on 3BPA ("ANI-pretrained") leads to slightly worse extrapolation slope. These results suggest that: (1) more diverse datasets may be required for assessing the extrapolation and generalization capacity of a model; and (2) pre-training on large datasets may be required to create universal NNIPs [77], given that pre-trained ANI models were able to escape the scaling relation seen in Fig. 1b.

Although the scaling relation can estimate the extrapolation slope within reasonable bounds, it is unable to recover trends within the same architecture in the low-error regime. As best-performing models often have differences of force errors smaller than 5 meV/A in the 300 K test set, it is not clear whether their extrapolation behavior can be accurately recovered from the scaling relation. Indeed, Fig. 1c shows the correlation between the forces RMSE computed on the 300 K split and all splits (300, 600, and 1200 K, "high T") for several NequIP and MACE models with different hyperparameters and training methods (see Tables S3 and S6), but similar testing error at the 300 K split. Although there is a positive correlation between the two RMSE metrics, the dispersion of points indicates that models with nearly the same RMSE for the 300 K data split show discrepancies in error above 10 meV/A when all splits are taken into account.

This scenario is aggravated when molecular dynamics (MD) simulations are used to test the extrapolation power of the NNIPs. When measuring the average simulation times for each of the models (see Appendix B for details on the MD simulations), no clear correlation is obtained between the error on the 300 K split of 3BPA and the average (physically meaningful) simulation length (Fig. 1c). This motivates the creation of a metric that captures robustness trends in NNIPs.

### Training insights derived from loss landscapes

Beyond data regularization, architectural and training choices strongly influence generalization ability of NN models. Relevant aspects include model initialization, hyperparameter optimization, choice of optimizer, batch sizes, and many others. Although the process of training NNs strongly affects their extrapolation behavior, good NN models systematically outperform their counterparts by optimizing to better local minima in the LL [54]. Thus, using these insights, we propose that _loss landscapes of NNIP models can predict their generalization ability towards unseen data despite using only the training data_. Correlations between robust generalization and loss sharpness have been observed for other NN models in the literature [55, 51], but not yet explored in the context of NNIPs.

To first verify if qualitative insights could be derived from LLs in the context of NNIPs, we investigated the behavior of the loss function around the optimized minima of the NNIP models trained to the non-noisy 3BPA dataset. To ensure the LL visualizations were statistically meaningful, we sampled 20 different orthogonal directions for each set of parameters and models, and interpolated them using the filter-normalized method described in Sec. 3.1 (see also Fig. 2a). Then, we compared the LLs of the NNIP models using 2D visualizations, as often done for NN classifiers [58]. Qualitative inspection of the 2D LLs in Fig. 2b reveals the presence of weight degeneracies in the prediction of energies in models (see also Fig. S14 for another example in the ANI-Al dataset). This "valley-like" landscape represents a
subspace of weights leading to similar accuracy in energy [78], and reflects the interplay between energy and forces during training. These results agree with the literature regarding LLs of over-parameterized models [79], as well as the notion that physical systems often result in so-called "sloppy" models [80; 81], and can improve trainability and interpolation [82].

Qualitative analysis of LLs also explain other factors typically found as heuristics of NNIP training. For example, the energy/forces coefficients in the final loss are often defined from hyperparameter optimization [23; 83] and have ratios varying from 1:10 to 1:100 for energy:forces RMSE. Nevertheless, the success of these higher ratios can be justified from the perspective of the LL. In Fig. 2b, we show how a higher weight on force losses leads to LLs with less saddle points around the optimized minima, thus favoring training. Although energy and forces are related and completely disentangling their effects is not achievable, the interpolation in Fig. 2b shows how mixing both can lead to better optimization landscapes for NNIPs. Based on these results, an effective training regimen would start with a relatively large weight for forces loss for a fast (and smoother) optimization of forces. Once the force errors are reasonably converged, the weight can be decreased until a desired threshold in energy errors is achieved. Similar strategies with weight cycling and scheduling -- e.g., starting with an energy:force loss weighting of 1:10, but then switch to 1000:1 in the later stages of training -- can also be effective.

One limitation of visualizations of the high-dimensional LL is that different directions in weight perturbation may lead to similar losses. This results in only slight variations in landscape topography with different random directions, as noted in Fig. 2. This may be related to the dominance of specific layers of the model in the filter normalization technique. Figure S13 shows how the distribution of weights is non-uniform, suggesting some layers have higher sensitivity to weight perturbation than others. As the filter normalization displaces the model weights along random directions with magnitude proportional to the norm of each filter and layer, parameters with higher weights may influence certain regions of the loss landscapes. An exception to this point is when the parameters are intertwined with functions embedded in the architecture, such as the trainable Bessel functions in NequIP. As shown in Fig. S6, freezing certain high-magnitude weights when generating the LLs can help flatten the landscape and remove spurious minima, emphasizing the importance of proper regularization and training regimen taking these effects into account (e.g., separate learning rates for certain layers in the model).

### Loss landscapes predict extrapolation trends in NequIP

In the context of NNIPs, robust generalization has important ramifications in two distinct areas: the stability of the model in production (e.g., MD simulations), and its data efficiency during training. To probe the first of these features, we trained NequIP models using various choices of model architecture and optimization techniques (see Table S1) in the low-temperature split of the 3BPA dataset. Then, we performed MD simulations in the NVT ensemble with a high temperature of 1600 K (Appendix B) to ensure all models were extrapolating well beyond their training split. Furthermore, as the 3BPA benchmark already provides geometries sampled from simulations up to 1200 K, the MD

Figure 2: **a**, A schematic describing the process of generating loss landscapes. Starting at the origin with a trained model, landscapes are constructed by performing a grid-sampling along a random filter-normalized direction [58] in parameter space up to a chosen maximum distance. In order to ensure that that results were consistent regardless of the choice of random sampling direction, landscapes were averaged over multiple random directions, \(\delta_{1}\dots\delta_{n}\), with \(n=20\) usually showing only slight variations in landscape topography. **b**, 2D LLs for a MACE model trained to the 3BPA dataset, generated by sampling a plane defined by two random directions. The model was trained using \(E_{w}=1\) and \(F_{w}=1000.0\), but the final landscape was re-weighted for the purpose of this figure using a fixed value of \(E_{w}=1.0\) and \(F_{w}\) increasing from 0.0 to 10.0 to demonstrate the effects of different weights on optimization.
trajectories at 1600 K could provide information beyond the 1200 K splits provided by the dataset. For each model in the study, 30 MD trajectories were simulated to obtain reasonable statistics on the model behavior, with simulation timelengths of up to 6 ps and a timestep of 1 fs. In principle, a model's ability to extrapolate beyond the training set can be assessed by the trajectories that behave in a physical way [34], as well as the average trajectory length until the model fails to extrapolate (Appendix B).

Using the results from MD simulations, we investigate whether LLs can predict the extrapolation behavior and trajectory stability of models. While the test performance at low-temperature samples is unable to perform such predictions (Sec. 4.1), the test error at high-temperature samples is still expected to be a reasonable predictor of such stability. Nevertheless, whereas high-temperature samples are available in the 3BPA dataset, out-of-domain data is often not available for an arbitrary material system. In addition to being expensive to generate using ground-truth methods, sampling new configurations is rarely performed in an exhaustive way. Thus, the major advantage of extrapolation tests based on LLs is their reliance only on the training dataset. To quantify the sharpness of the LLs [55], we compute the loss entropy as described in Sec. 3.2, with \(T_{E}\) and \(T_{F}\) taken to be reasonable "room temperature" values of the energy/force RMSEs, respectively (\(T_{E}=4\) meV/atom, \(T_{F}=40\) meV/A). The weight \(\alpha\) was set to \(0.2\) to resemble the higher force weighting used during training without ignoring the contributions of energy LLs in model stability. Nevertheless, a sensitivity analysis of 5 with respect to these parameters showed that results remained consistent around these error ranges, as long as unreasonable temperature values were not used (see Fig. S4).

The relationship between the model/training parameters, the MD stability, and LLs entropy are shown in Fig. 3 (see also Table S2). The analysis is grouped into four experiments (columns in Fig. 3) to isolate the effects of specific portions of the training procedure and model architecture, thus uncovering useful trends for NNIP practitioners. For example, considering only the distributions of time to failure from MD simulations in Fig. 3b, it can be seen that rescaling the energies predicted by the model (models "no rescaling" vs. "rescaling") greatly improves the stability of the model, especially when the Bessel weights used by the radial basis are trainable. This behavior is reflected in the energy and forces LLs (Fig. 3a), with the LL of the model without rescaling showing considerable sharpness compared to the models using rescaled energies. Quantitative trends are also obtained when the loss entropy and test RMSE at all splits in the 3BPA dataset are computed (Fig. 3c, first column, and Table S3). Models with rescaling showed higher average simulation time with physical trajectories, as well as increased entropy and lower forces RMSE.

When comparing how the training regime can change the extrapolation behavior of NequIP, the results show that the AMSGrad variant of the Adam optimizer [84] leads to consistent improvements in the model extrapolation both for the 2-layer and the 5-layer model (Fig. 3b, second and third column) compared to the baseline, which does not use AMSGrad. The improvements of simulation quality are particularly pronounced in the 5-layer models, where the model trained using AMSGrad shows significant improvements in simulation stability compared to the baseline despite not using trainable Bessel functions (Table S1). On the other hand, the exponential moving average (EMA) appears to degrade the extrapolation performance of the NNIIPs, often leading to sharper LLs if used in isolation and with constant learning rates. These trends are reflected in the extrapolation metrics (Fig. 3c). Although the loss entropy underperforms compared to the forces RMSE in the case of the 2-layer models, the trend is correctly captured for the 5-layer model, where the model trained with AMSGrad showed both higher loss entropy and higher MD stability.

Finally, NequIP models can be compared according to the number of message-passing layers. Using a fixed training regime, a higher number of layers showed pronounced improvement both in the stable simulation time (Fig. 3b), as well as in the forces RMSE and loss entropy (Fig. 3c), with the trend remaining consistent across these models. Interestingly, the improvement in loss entropy for the 5-layer model is reflected mostly in the energy LL rather than the force LL (Fig. 3a). Although the energy is not required when integrating the equations of motion in the NVT ensemble during an MD simulation, this improvement in LL may reflect the model's higher generalization capacity, and thus wider minima. This phenomenon can be explained by recognizing that the out-of-domain data often results in a horizontally shifted version of the loss landscape [46]. To explain this effect, we computed the LL of NequIP models with respect to the test set instead of the training set (see Fig. S5). As expected, the test LL shows higher errors compared to the train LL, as seen in the vertical shift of the curves. However, while the forces LL only undergo a vertical shift, test energy LLs also undergo a horizontal shift, indicating that the optimal model for the training set is not necessarily optimal for predicting energies of the testing set. Furthermore, as the forces loss is often derived from the energy in NNIIPs, the propagation of these mismatches may be responsible for degradation in extrapolation performance. Thus, in general, the results in Fig. 3 show that architectures and optimization strategies which result in loss landscapes with higher entropy (i.e., flatter landscapes) tend to demonstrate improved stability in MD simulations that sample out-of-domain configurations.

### Loss landscapes predict extrapolation trends in MACE

To confirm that the results from Sec. 4.3 could be extended beyond NequIP, we performed a similar study using the MACE framework. Given the differences in architecture and available code, MACE has different hyperparameters and
optimizer choices than NequIP (see Table S4). Nevertheless, a similar study for MACE reproduces the trends seen for NequIP, as shown in Fig. 4.

As seen in Sec. 4.3, model stability can be greatly improved by using rescaling and AMSGrad. Also for MACE, EMA appears to lower the time-to-failure if used in isolation, with similar trends in out-of-domain RMSE, loss entropy, and MD stability seen in the NequIP case (Fig. 4). In addition, we also explored the effects of stochastic weight averaging (SWA) [85] and weight cycling (WC) (column 2 of Fig. 4) implemented in the MACE code. Consistent with results from the previous studies, the effect of the stochastic weight averaging on the MSE is also reduced by the factor \(\sim 10^{-10}\).

Figure 3: Analysis of LLs and MD simulation stability results for NequIP models while varying: (column 1) rescaling and Bessel weights, (column 2) optimizer settings using a model with two interaction blocks or (column 3) five interaction blocks, and (column 4) the number of interaction blocks. **a**, LLs for each model and **b**, distributions of time to failure computed over 30 MD simulations. **c**, Time to failure compared to (top) the entropy of the loss landscape (as defined in Sec. 3.2) or the forces RMSE on the high-temperature test sets. The horizontal line, box, and whiskers in **b** depict the median, interquartile range, and points within 150% of the interquartile range of the distribution. Isolated diamonds are points outside of this range.
from the use of SWA for image classification tasks [85], SWA + WC is shown to improve model generalizability, leading to higher simulation times and flatter loss landscapes. Interestingly, the model trained with SWA + WC exhibits similar force RMSE in high-temperature samples compared to the baseline, but a higher loss entropy and higher MD stability (Fig. 4c). As SWA flattens the loss landscape by design [85], implementing this strategy in other NNIP models may be a low-cost modification for improving the generalization capacity of different architectures. The use of WC can also help the optimization process (Sec. 4.2) and lead to better energy landscapes overall.

Figure 4: Analysis of loss landscapes and MD simulation stability results for MACE models while varying: (column 1) the use of rescaling; (column 2) optimizer settings using a model with a fixed body order, \(v=2\), and symmetry order, \(L=3\); (column 3) the symmetry order of the edge features with a fixed body order of 2; and (column 4) the body order with a fixed symmetry order of 3. **a**, Loss landscapes for each model. **b**, distributions of time-to-failure computed over 30 MD simulations. **c**, Time-to-failure compared to (top) the entropy of the loss landscape (as defined in Sec. 3.2) or the forces RMSE on all test sets of the 3BPA dataset. The horizontal line, box, and whiskers in **b** depict the median, interquartile range, and points within 150% of the interquartile range of the distribution. Isolated diamonds are points outside of this range.
For the MACE study, we also analyzed the relationship between the bond order of the model (column 4) and their stability in production MD simulations. In this case, the loss entropy recovers the trend in MD stability better than the forces RMSE (Fig. 4c, column 4). Interestingly, despite the higher errors of the \(v=1,L=3\) model compared to its \(v=2,L=3\) counterpart, the former still exhibits a higher average time to failure. This observation can be traced back to the flatter energy LL (Fig. 4a, column 4), as also seen in the case of NequIP.

Finally, the models were compared according to the symmetry order of the edge features (column 3). In general, increasing the value of \(L\) from \(L=0\) to \(L=2\) led to more stable simulations, although the model with \(L=3\) shows a degradation of the production behavior. Nevertheless, the LLs do not follow the expected trends explained so far in this work. We believe the filter normalization technique discussed in Sec. 3.1 cannot provide comparable LLs upon changes of \(L\) in MACE. As this parameter affects the tensor order of some model components [25], the number of parameters in these particular filters grows with \(\mathcal{O}(N^{L})\), which may be affecting the filter normalization. In contrast, adding message-passing layers with similar numbers of parameters, as in NequIP, allows the filter normalization technique to create comparable loss landscapes. This result shows some limitations of the loss entropy when comparing models with different architectures or hyperparameters, and may require different strategies for computing them in the future.

### Loss landscapes and data efficiency

As extrapolation power and data efficiency are conceptually related, a natural extension of Secs. 4.3 and 4.4 is to verify whether the loss landscape entropy can also be used to predict the data efficiency of a model during training. As with many other applications of deep learning, generating training data is often one of the most costly steps in the development process. Especially considering the "denoising" effects of NNIPs demonstrated in Sec. 4.1, identifying training techniques and model architectures that lead to more data-efficient training has the potential to greatly reduce the computational cost of building new NNIPs.

To test whether loss entropy can predict the data efficiency of models, we computed the learning curves for NequIP and MACE models with the hyperparameters selected in Sections 4.3 and 4.4. This is achieved by training models with the specified training parameters and architectures to the same subsets of the 3BPA training set using only 25, 125, 250, or 500 samples from the 300 K split (Tables S3 and S6). Following previous work [23], the slopes of the learning curves were then computed by fitting the line \(\log n=m\log\varepsilon+b\) to the number of training samples \(n\) and the force RMSEs \(\varepsilon\) calculated using all test splits (300, 600, and 1200 K), and comparing the slopes \(m\). Consistent with the results from Sections 4.3 and 4.4 (combined in Figures 5a,b), the high-temperature force RMSEs (i.e., the points from the learning curves using \(n=500\)) are good predictors of the learning curve slope. In addition, Figs. 5c,d show that the loss entropy also has a high correlation with the slope of the learning curve despite being computed using only the training set. As discussed in Sec. 4.4, these correlations are less explicit for the case of MACE models, and thus show that none of the metrics is truly universal for predicting stability in simulations and learning curve slopes. Nevertheless, these results further demonstrate that loss landscapes provide information on the extrapolation behavior of NNIPs despite being derived only from the training set.

## 5 Conclusion

In this work, we motivate the need for additional metrics beyond RMSE by showing that in-domain errors fail to predict model stability in the high-accuracy regime despite large-scale trends in extrapolation behavior and NNIP robustness to noise. We propose the use of loss entropy as a metric for quantifying the extrapolation behavior, and demonstrate that it correlates well with out-of-distribution error and stability in production simulations. Using large studies with NequIP and MACE, we show how models containing flatter loss landscapes exhibit better extrapolation behavior, and how different training parameters can be used to achieve them. For example, rescaling, AMSGrad, and SWA were shown to increase the loss entropy and MD stability, and may be important tools when training NNIP models. Similarly, models with similar test error can be distinguished by their energy loss landscape, with models displaying broader minima performing better in extrapolation tasks. Future studies can address shortcomings of loss landscape visualizations in NNIP by analyzing how filter-normalization can be made more suitable for NNIP architectures, or to better isolate the effects of key architectural changes in the loss. Nevertheless, these results can better inform the development of new model architecture and optimization strategies in NNIPs, facilitating their use in general materials simulation.

## Code Availability

The package ip_explorer and the additional code/data used to reproduce the results of this paper will be made available after internal review at LLNL. This preprint will be updated with the relevant links. Loss landscape calculations were performed using the the code from the public package [https://github.com/marcellodebernardi/](https://github.com/marcellodebernardi/)
loss-landscapes. Training codes for NequIP and MACE are available from their original authors as described in Appendix A.

## Data Availability

The datasets used to train the models in this work were obtained directly from their original sources. For convenience, we provide the links to each source here: [https://github.com/davkovacs/BOTNet-datasets/tree/main/dataset_3BPA](https://github.com/davkovacs/BOTNet-datasets/tree/main/dataset_3BPA) (3BPA) and [https://github.com/atomistic-ml/ani-al](https://github.com/atomistic-ml/ani-al) (ANI-AI).

## Author Contributions

**Joshua Vita:** Methodology, Software, Validation, Investigation, Data Curation, Writing - Original Draft, Writing - Review & Editing, Visualization. **Daniel Schwalbe-Koda:** Conceptualization, Methodology, Software, Validation, Investigation, Data Curation, Writing - Original Draft, Writing - Review & Editing, Visualization, Supervision.

## Acknowledgments

This work was performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344, funded by the Laboratory Directed Research and Development (LDRD) Program at LLNL under project tracking code 22-ERD-055. The authors thank Vincenzo Lordi, the Quantum Simulations Group at LLNL, and Simon Batzner for the discussions. We also thank Ilyes Batatia for the support with the MACE code.

Manuscript released as LLNL-JRNL-845001-DRAFT.

Figure 5: Loss landscape entropy and high-temperature force test RMSEs compared to **a,b** MD time-to-failure or **c,d** learning curve slope for **a,c** NequIP and **b,d** MACE models. The points. Learning curve slopes were computed as described in Sec. 4.5. Colors for points in **a** and **b** were chosen to match the colors in Fig. 3 and Fig. 4 respectively. Lines in **c** and **d** represent linear fits to the data.