Algorithmic Differentiation for Automated Modeling of Machine Learned Force Fields

Abstract

Reconstructing force fields (FFs) from atomistic simulation data is a challenge since accurate data can be highly expensive. Here, machine learning (ML) models can help to be data economic as they can be successfully constrained using the underlying symmetry and conservation laws of physics. However, so far, every descriptor newly proposed for an ML model has required a cumbersome and mathematically tedious remodeling. We therefore propose using modern techniques from algorithmic differentiation within the ML modeling process, effectively enabling the usage of novel descriptors or models fully automatically at an order of magnitude higher computational efficiency. This paradigmatic approach enables not only a versatile usage of novel representations and the efficient computation of larger systems—all of high value to the FF community—but also the simple inclusion of further physical knowledge, such as higher-order information (e.g., Hessians, more complex partial differential equations constraints etc.), even beyond the presented FF domain.

Most physical quantities are represented by differential equations (DEs)

1

where is a finite linear combination of differential operators of order n.

Examples of eq 1 are found in electrodynamics, fluid dynamics, quantum mechanics, etc., always imposing strong constraints on the form of a physical solution f(x). For several years, machine learning (ML) has started to broadly contribute to physical modeling across its disciplines, for instance in particle physics,¹ atomistic simulations and force fields (e.g., see refs (2−6)), fluid dynamics (e.g., see refs (7 and 8)) or lattice gauge theories.⁹⁻¹² Note, however, that the standard procedure of ML models has been so far to learn from empirical data only. Recently, ML modeling has also started to incorporate physical constraints, be it as regularization terms (e.g., see refs (13 and 14)) or by explicitly taking into account conservation laws, DEs such as eq 1, symmetries (e.g., see refs (3, 15, and 16)), various operator response properties^17,18 or empirical correction terms, e.g. for long-ranged electrostatics.¹⁹⁻²² Notably, the direct usage of DEs allows the method to greatly simplify modeling and to be highly data efficient since the known laws of physics do not need to be learned from empirical data anymore (e.g., see refs (3, 15, and 16)).

Let us study the popular example of eq 1 force fields (FFs), namely, the negative gradient , where u(x) is a scalar energy potential mapped to a vector field of forces f(x). The unknown force field (FF) is sampled at several locations and f(x_i) are obtained from ab initio atomistic reference data.^5,15,23 The constraint

2

which is equivalent to both energy conservation and curl-free vector fields, enables the estimation of a faithful representation of true energy-conserving FFs as opposed to the estimation of general (unphysical) vector fields.^15,24,25

This approach unfolds its usefulness well beyond the FF domain for higher-order differential operators and compositions thereof,²⁶⁻³⁰ which characterize complex properties across all physics domains, e.g., Hessians or Laplacian operators.³¹⁻³³ Prior studies^{15−17,34−46} have demonstrated the effectiveness of explicit DE constraints, but their adoption still remains low due to an inefficiency that hampered model development and training so far: (1) Differential operator transforms inflate model complexity, making their full algebraic derivation mathematically tedious and error-prone. (2) Training and evaluation of a model using DE constraints is often associated with significant time and memory costs. While there exist prior works with focus on efficiency improvement, they rely on manual derivations and are limited to first-order gradients and primitive kernels without physical descriptors.²⁸⁻³⁰ On the other hand, recent existing work that focuses on automating the derivation of (first-order) constraints is not able to avoid a costly dense instantiation of the full model.⁴⁶

In this work, we address both issues simultaneously using methods from algorithmic differentiation (AD) (e.g., refs (47−50)). While AD is now routinely used to automate the computation of derivatives, another benefit, which turns out to be essential for this work, is that it can simplify models that have a DE structure. As we will see, having a DE structure allows the collapse of certain portions of the computational graph early during evaluation, by preaccumulating derivatives according to the chain rule. In summary, there are three key aspects that make this work useful for a broad set of applications in physical modeling:

• It enables the systematic construction of empirical models subject to complex differential constraints.

• It is easy to use and versatile, even for elaborate neural network descriptors.

• It affords orders of magnitude gains in computational speed.

We demonstrate this for a number of popular FF models³ based on Gaussian process (GP) estimators under linear operator transformations. Such models can then be decomposed into tensor products of operators due to unique properties of the kernel function (see Figure 1).

Figure 1.

(A and B) More efficient kernel-vector products. Kernel functions subject to linear differential operator constraints have a tensor-product structure (purple), which can be leveraged by managing the contraction order when computing kernel-vector products. Gaussian processes can benefit enormously from this improvement, since the full instantiation of the kernel matrix can be avoided during inference and training. (C) The effectiveness of differential equation constraints. A Gaussian process (RBF kernel, length scale σ = 1) is used to reconstruct the two-dimensional Rosenbrock function u(x₁, x₂) = (white contour lines) from just two training points. Function value samples alone are insufficient to recover any surface features (left), but gradient constraints already enable the approximate recovery of extrema, despite not having been directly sampled (middle). Hessian constraints further refine the local curvature of the reconstruction (right).

As we will show, AD can yield efficient and automatic construction of GPs. Previously, a full instantiation of the DE constrained models was necessary, which due to the laborious manual derivations required, warranted a separate publication for each new constrained model (e.g. see refs (15−17, 25, and 51)). Using AD, we show that GPs can be created ad hoc and automatically (avoiding the above-discussed manual derivation steps) for arbitrary DE constraints, as we demonstrate by reimplementing a broad selection of popular classic as well as new ML-FFs; all can now be studied within our novel framework.

This newfound efficiency also enables us to easily recombine promising concepts from existing models into new, even more powerful physics models.

Specifically, AD enables us to break down differential expressions into sets of linear primitive operations, which are evaluated one-by-one to avoid a full instantiation of all intermediates of whenever the response is needed only for a specific transformation of x. Due to the significant internal structure of derivatives, the individual terms that comprise are often low-dimensional and can thus be evaluated more efficiently by optimally managing the contraction order without having to undergo an unnecessary inflation.

Following this general idea, it turns out that most operator constraints can be applied with a surprisingly low overhead, often without even increasing the asymptotic computational complexity class of the original model. This fact is rather unexpected as it is unintuitive, since a constrained model often involves significantly larger terms, which are expensive if evaluated naively. AD draws its efficiency from just two fundamental operations:⁴⁸ Given a differentiable function with total evaluation time cost C_u, AD guarantees that for , , and

• 1.Jacobian-vector product J_u(x)v has cost ,
• 2.Jacobian-transpose-vector product has cost .

Since any linear differential operator can be composed using these two simple rules, thus often at similar low cost (see Figure 1 and Tables SI and SII), e.g., Hessian-vector-products in instead of quadratic complexity. These improvements are possible because the full Jacobian never has to be calculated or stored.

Such inexpensive Jacobian products can be leveraged in any differentiable model, but they are particularly useful when applying DE constraints to GPs (for FFs) as we will now discuss

3

is fully defined by a mean and a covariance function and is readily learned from data.^2,24,52−54 For instance, many FFs, with n = 3N atomic degrees of freedom, are modeled by GPs from ab initio reference calculations (cf. refs (3−5 and 55)).

A natural benefit of GPs for DE-based regression is that they are closed under linear operators, leading to the form

4

where subscripts indicate the argument of the kernel for each operator to act on.

As a concrete example, consider observing function values, gradients, and Hessians, i.e. within the same model (cf. Figure 1). The corresponding kernel will take the form

5

where each differential constraint appears as a cross-covariance combination with all other terms (e.g., the Hessian-Hessian covariance part being a fourth-order derivative). We leverage AD to construct this exceedingly complex matrix without the need to instantiate the corresponding analytical expressions. Rather, each term is constructed efficiently by contracting local Jacobian-vector products on-the-fly (cf. Figure 1).

Likewise, for third-order operators, the kernel uses a sixth-order derivative, and so on. At this complexity growth rate, a manual algebraic derivation becomes increasingly hopeless.

But with the help of AD, a mere definition of k(x, x′) is sufficient, avoiding taking manual derivatives. Moreover, computationally the action of the operator avoids an intermediate generation of the full matrix expression, which gives rise to high computational efficiency.

We would like to discuss further the efficiency aspect of our AD framework, which also goes beyond standard AD usage. While the standard computation of the GP equations, for example, for the predictive mean are f(x) = , where coefficients α_i are obtained by solving a linear system (see eqs S2 and S3 and refs (2) and (54)). We note that this computation is usually dense and therefore slow.

Due to linearity, we can alternatively compute the same at roughly an order more efficiently (see Figure S1 and Table SI) by

6

as the operator tensor product is resolved at first into a scalar operator which can take full advantage of the automatic contraction rules inherent to AD. For an intuitive explanation see Figure 1, and for a more detailed derivation see the Supporting Information.

We will now apply the AD framework discussed above to the atomistic simulations domain. For this we proceed with the construction of potential energy surfaces (PESs) for small molecules from the well-established MD17 benchmark data set,¹⁵ using various GP-based ML-FFs that we recreate within our novel framework. First, we consider the Symmetric Gradient-domain Machine Learning (sGDML)^16,56 model, which uses derivative constraints (eq 2) to simultaneously reconstruct conservative molecular FFs and their corresponding PESs. This model employs a twice differentiable kernel function from the parametric Matérn family,⁵⁷⁻⁵⁹ which is symmetrized to be invariant with respect to the relevant rigid space group, as well as dynamic nonrigid symmetries of the system at hand. It is then combined with a descriptor that enumerates all unique pairwise inverse distances between atoms. This composition of functions, if implemented in the standard manner,⁶⁰ leads to an increased cost since all atomic degrees of freedom of the molecule enter the model as separate constraints. Our AD framework resolves this complexity and avoids unnecessary instantiation of operator tensor products, yielding an order 3N improvement (where N is the number of atoms) in the case of gradient operators.

Furthermore, we have reimplemented other FF GP-based models, including FCHL19.^25,61 Within our AD framework, we could do so by specifying the respective kernel functions without needing to manually implement derivatives. These models are compared to (s)GDML models trained on the exact same data set splits. To verify that our own implementations are correct, we have compared our test errors with the respective original publications.

Analyzing running times of the highly optimized FCHL19²⁵ reference implementation with our own AD-based one, we can already see the high intrinsic optimization abilities and computational advantages for our AD framework, which is able to generate predictions up to 2 orders of magnitude faster, while using only a single GPU instead of a CPU cluster node (cf. Table 2 in ref (25) and Table 3). We note, however, as a word of caution to this comparison that different compute architectures and programming languages are being used.

Table 3. Benchmarked Force Prediction Times (s) for Different Kernels^a.

	ethanol (N = 9)			aspirin (N = 21)
model	dense	contraction	speedup	dense	contraction	speedup
sGDML	0.0780	0.0018	×43.3	0.2832	0.0087	×32.5
FCHL19GPR	57.9439	0.9815	×59.0	414.0216	10.1438	×40.8
sGDML[RBF]	0.0800	0.0015	×53.3	0.2811	0.0074	×37.9
global-FCHL19	57.5194	0.9653	×59.5	458.5465	10.0728	×45.5
sFCHL19	60.3951	1.0704	×56.4	419.3778	10.3336	×40.5

^aEach model was trained using 1000 points and evaluated for a batch of 10 points. All timings are averaged over 10 runs (excluding an initial run for just-in-time compilation). Our approach (contraction) yields consistent speedups by a up to two orders of magnitude over the direct (dense) implementation of the constrained models. All measurements are done on a single Nvidia Titan RTX 24 GB GPU.

Table 3 contains a running time comparison of all considered models within our AD framework for the largest (aspirin) and smallest (ethanol) molecule in the MD17 data set, both with and without relying on early contraction of intermediates. Since AD alleviates the need for laborious manual derivations, we can easily replicate further ML-FFs within the same code to arrive at fair running time comparisons. This type of analysis could not be done accurately before.

The ultimate advantage of our AD framework is the freedom to recombine promising concepts from existing approaches effortlessly on a broad scale in order to discover even more powerful task-specific models. So far, there exists no single universally best ML-FF; the field is rather characterized by specialized solutions. We demonstrate this new modeling flexibility by creating several better variants of the GP-based models mentioned above. sGDML[RBF] is an sGDML variant that uses the RBF kernel instead of the Matérn kernel. global-FCHL19 is a global variant of FCHL19, which parametrizes each atom interaction individually at the cost of permutational invariance. And finally, sFCHL19 is a symmetrized global variant using the symmetries recovered by the sGDML model. Table 1 compares these variants to the respective original models and demonstrates that highly significant advances in prediction accuracy become possible for all MD17 data sets, simply by systematically combining existing ideas. Even linearly combining kernels shows consistent improvements (see Table SIV). Furthermore, AD allows an effortless gradient-based hyperparameter optimization, as we demonstrate in section D in the Supporting Information.

Table 1. Test Errors (MAE) for Force Learning on MD17 with 1000 Reference Points Using Various New GP-FF Variants^a.

model	aspirin	benzene	ethanol	malonaldehyde	naphthalene	salicylic acid	toluene	uracil
sGDML	0.702	0.163	0.341	0.410	0.115	0.286	0.145	0.242
FCHL19GPR	0.628	0.179	0.180	0.302	0.185	0.277	0.246	0.147
sGDML[RBF]	0.506	0.147	0.186	0.237	0.068	0.167	0.096	0.140
global-FCHL19	0.748	0.235	0.319	0.429	0.292	0.216	0.370	0.120
sFCHL19	0.430	0.158	0.160	0.274	0.094	0.172	0.111	0.130

^aAll errors are in kcal mol^–1 Å^–1. Best results are in bold. The automation of AD allowed us to recombine existing ideas from different ML-based FFs on a broad scale and find better-performing model variants for all MD17 datasets. The best result for each dataset is marked in bold.

We can even go a step further and use the representations learned by modern deep architectures⁶²⁻⁶⁴ as pretrained descriptors D(x) within GDML-type models. Each descriptor then yields a new composite kernel

7

Our numerical results show that this construction interestingly yields models that are more accurate than their (pretrained) ingredients. Tables 2 and 3 summarize the force prediction performance, showing that significant improvements are possible following this simple strategy. We reiterate that this advance was enabled by the unprecedented flexibility to remix different models enabled by our framework.

Table 2. Using Molecular Representations Generated by Various Deep Neural Network Architectures as Descriptor for GDML-Type Models^a.

model	descriptor	aspirin	ethanol	malonaldehyde	naphthalene	salicylic acid	toluene	uracil
SchNet		0.824	0.225	0.428	0.342	0.490	0.325	0.307
sGDML[RBF]	SchNet	0.930	0.176	0.359	0.294	0.459	0.281	0.248
PaiNN		0.389	0.220	0.336	0.103	0.232	0.122	0.176
sGDML[RBF]	PaiNN (scalar features)	0.422	0.186	0.321	0.102	0.241	0.118	0.182

^aAll models have been trained on the same 1000 reference points for each molecule in the MD17 dataset. All (force) test errors (MAE) are in kcal mol^–1 Å^–1. The best result for each dataset is marked in bold.

In order to demonstrate that the AD framework can also be applied beyond gradient observations (constraint from eq 2), we will now illustrate the effect of employing more complex higher-order DE constraints, namely, Hessians for the learning model.

Using the Hessian-kernel in eq 5, we reconstruct a toy surface from just two gradient/Hessian observations (see Figure 1). Notably, this example shows that constraints with higher derivatives are more informative and aid regression models significantly. In our example, the reconstruction error improves by 1 order of magnitude as a result of including second-order measurements. With that, even a small sample size becomes sufficient to identify the correct model. For further illustrative examples using other DE constraints and ML-based DE solving, see Table SI.

In summary, an overwhelming number of physical phenomena are governed by linear DEs that can be used as highly effective physical knowledge in data-driven estimation problems. By excluding physically infeasible solutions, such constraints play a crucial role in obtaining data-efficient and robust models. We have used methods from AD to address two key challenges that have so far hindered widespread adoption of this approach: (1) the full algebraic instantiation of the constrained model is expensive (or even unfeasible) to train and evaluate, and (2) the application of DEs to GPs by hand is tedious and error-prone for a modeler. Due to these obstacles, the construction of DE-constrained models did not leave much room for explorative and swift model creation in the past. To change that, we have contributed how our AD framework can be used to automate the model construction process and how regularities in the differential structure can be leveraged to gain efficiency in practice. This general framework has enabled us to construct and integrate various differential operator constraints that constitute the basic building blocks of most physical laws into ML-FFs. Following the same principle, more complex operators can be constructed and turned into constrained ML models; we used mainly GPs to show this point. For further examples, see the Supporting Information.

Finally, we have demonstrated in a series of numerical experiments how our framework can be used to readily replicate some state-of-the-art DE-constrained GP-based FFs by simply changing a few lines of code without the need to go through tedious manual derivations. Going one step further, we were able to demonstrate how competitive new kernelized variants of existing deep learning-based FFs can be developed and combined.

Supporting Information Available

The Supporting Information is available free of charge at https://pubs.acs.org/doi/10.1021/acs.jpclett.2c02632.

• Further details about computational settings and benchmarks, toy examples, derivations of computational complexity considerations, Tables SI–SIV and Figures S1 and S2, as well as source code for all experiments (PDF)

• Transparent Peer Review report available (PDF)

The authors declare no competing financial interest.