GPU-accelerated approximate kernel method for quantum machine learning

Data-driven approximate machine learning (ML) methods have become increasingly prominent in theoretical chemistry in recent years.^1–11 In particular, supervised learning can be used to augment accurate but computationally prohibitive electronic structure calculations. For tasks such as ab initio molecular dynamics (AIMD), approximate variants of the electronic Schrödinger equation are solved for every coordinate update. Consequently, surrogate ML models that can partially substitute for quantum calculations are extremely beneficial for reducing the overall computational burden (and carbon footprint). In general, these ML models first transform atomic coordinates into an intermediate symmetry-preserving representation, which is then passed on to a nonlinear machine learning model, usually one based on neural network (NNs)¹² or Gaussian process (GP) methods.¹³ 

While NNs have received significant focus to enable them to scale performantly on graphics processing unit (GPU) architectures through GPU-accelerated packages, such as PyTorch¹⁴ or TensorFlow,¹⁵ there has been much less focus in providing similar accelerations and tooling to GP methods, as the vanilla variant of these tend to scale unfavorably on GPU devices.

However, there are improved time complexity GP models relying sparsification, whereby a subset of the training data is used to enable training on the full dataset.^16,17 In particular, the training step time complexity is linearized with respect to the amount of sparse training data and the inference time of these models is favorable. This strategy has enabled sparse GP models to be trained in a “Big Data” context while retaining sufficient accuracy with respect to vanilla GP.

Alternatively, the mapping provided by the kernel function central to GP models can be approximated to yield a desired linearization. In this work, we briefly discuss the seminal work on random Fourier features (RFF)¹⁸ in which the kernel mapping is approximated via a low-dimensional feature function. We also summarize a significantly more computationally efficient variant, termed structured orthogonal features (SORF),¹⁹ which further reduces the computational cost.

We provide a software package to perform training and prediction of Quantum Machine Learning (QML) models, termed QML-Lightning. It includes GPU implementations of both sparse GP, RFF, and SORF models as well as a GPU implementation of FCHL19,²⁰ an accurate atom-centered representation. QML-Lightning is built upon the PyTorch software package, with additional CUDA C implementations for critical components to improve its computational throughput. A thorough benchmark of the predictive accuracy of QML-Lightning has been performed on several established datasets of chemical compounds from literature, comparing against existing kernel- and neural network based models. To assess the models’ performance across chemical compound space, we have benchmarked the model against the QM9 dataset,²¹ which contains 134k small organic molecules with elements C, H, O, N, and F. For dynamics and structural relaxation applications, we have benchmarked the model against both the MD17^22–24 and rectified rMD17²⁵ datasets, which contain trajectories of 10 small organic molecules. To infer the models’ extrapolative performance, QML-Lightning has been benchmarked against the challenging 3PBA dataset,²⁶ which contains three sets of MD trajectories at temperatures 300, 600, and 1200 K. We note that we have focused on energetic properties in this work as these are the most critical for AIMD applications; however, other QM properties can also be used within the QML-Lightning framework straightforwardly. Finally, we detail the wall time to perform 1 K MD steps times for a variety of systems including periodic systems with up to ∼14k atoms, using an OpenMM²⁷ plugin that is provided as part of this work.

Note that while finalizing work on this paper, Dhaliwal et al.²⁸ most recently published RFF-based interatomic potentials for performing molecular dynamics with promising results for central processing unit (CPU) based applications. Here, however, we focus on a GPU implementation of the more efficient and accurate SORF model.

The QML-Lightning software is provided under an MIT license at https://github.com/nickjbrowning/qml-lightning. We additionally provide an OpenMM wrapper for QML-Lightning under the same license at https://github.com/nickjbrowning/openmm-qml-lightning.

This section first summarizes a subset of kernel methods to learn quantum mechanical properties. First, Gaussian process regression (GP)^29,30 is introduced to learn both energies and forces of chemical compounds. Then, sparse GP^17,30 is discussed, and finally the approximate kernel methods random Fourier features (RFF)¹⁸ and structured orthogonal random features (SORF)¹⁹ are introduced.

In the following, upper case indices denote the index of chemical compounds in a dataset, while lower-case indices denote the index of atomic centers in each chemical compound.

In GP kernel models of quantum mechanical properties, one constructs a linear model using a basis of kernel functions k(p_J, ⋅). For example, to learn potential energies, a suitable functional form would be

where ρ_I and ρ_J are the atomic representations of atom I and J, respectively, and N_t is the number of training molecules. We note that these atomic representations are functions of the set of local atomic charges and coordinates {z, r}, which has been omitted for brevity. The sets $CI$ and $CJ$ contain the set of atoms for training and query molecules i and j, respectively. The coefficients α_i are obtained from the regularized minimization problem given by

which has the following solution in matrix form:

The hyperparameter λ is a small number in order to regularize and ensure numerical stability upon kernel inversion.³¹ For learning potential energies and atomic forces simultaneously, one can construct an expression for the potential energy as follows:

where matrix notation is introduced for simplicity and $∂∂r⃗$ is used as a shorthand to stack the following derivatives:

where l indices the coordinate components from the query atom J. The Hessian in Eq. (1) has the following form:

where k indices the coordinate components from training atom I. The dimension of the full GP kernel is (3MN_t + N_t) × (3MN_t + N_t), where M is the average number of atoms per molecule in the entire training set. In particular, the Hessian term in Eq. (1) has a compute time scaling as $O(36Nt2M4)$⁠, which will severely limit applicability with respect to both large training set sizes and large systems.

To alleviate the complexity of the above, which will be referred to as “full GP” in the following, Bartok et al.^16,29 introduced sparse GPs for modeling molecular systems. For a thorough review of GP and sparsification, please see Ref. 30. In sparse GP, rather than using all available training points, one selects a subset of representations as the basis. For example, a subset of the kernels centered on atoms can be used as the sparse basis such that the expression for the potential energy is as follows:

where N_s is the number of sparse points selected. The sparse points I are typically selected from the training database using uniform random sampling, CUR matrix decomposition,³² or farthest point sampling (FPS).¹⁷ The key difference from full GP presented above is that there is the use of one coefficient α_I per sparse point, rather than one coefficient per training molecule. For potential energy models, this formalism reduces the cost of building the kernel and predictions to $O(Ns2NtMd)$ and O(N_sMd), respectively, where M in this context is the (average) number of atoms in the (training) query molecule and d is the size of the representation. The model coefficients can be obtained by solving the following loss function in matrix form:

which yields the solution

Additionally, atomic forces can be included in the training scheme resulting in the following model:

where the matrix $−∂∂r⃗*K∈R3NtM×Ns$ stacks the 3N_tM partial Cartesian derivatives with all sparse points. In the following, we have dropped the sGP label for simplicity, and we also introduce the notation ∂K as a synonym for the matrix $−∂∂r⃗*K$⁠. This yields the loss function,

where the notation [U, F] is used to indicate the concatenation of the matrices U and F. The coefficients are then obtained from the following equation:

We stress that including derivatives in the above formulation does not increase the size of basis in which the total energy is expanded, so these derivatives only appear in the loss function. To compute these model coefficients, the dominant term is now the formation of the ∂K^T∂K matrix product, which scales as $O(3NtMNs2)$⁠.

Alternatively to sparsifying over the training data to reduce the computational complexity of training and inference, it is also possible to approximate the mapping provided by the kernel function such that linear models result. In the seminal work of Rahimi and Recht,¹⁸ a lower-dimensional lifting function was introduced to approximate the inner product synonymous with the kernel method,

As a consequence of Bochner’s theorem, the Fourier transform of a shift-invariant kernel k(ρ_I, ρ_J) = k(ρ_I − ρ_J) is a proper probability distribution. Consequently, one can readily define an explicit feature map that approximates the kernel via Monte Carlo integral estimation,

where N_f is the number of independent vectors w drawn from the probability distribution p(w) and i is the imaginary unit. For different kernels, the distribution p(w) takes different forms; however, for Gaussian kernels used here, p(w) is also Gaussian. After applying Euler’s formula to z(ρ) and after simplification, the following low-dimensional feature map is obtained:

where b is sampled from a uniform distribution on [0, 2π]. Since the potential energies are extensive, one can partition them into atomic energy contributions ϵ_I, and the representation of the atomic environment is passed into this low-dimensional feature mapping,

where $Z∈RNt×Nf$ is the feature matrix corresponding to N_t training observations. The weights α are the solution to the following regularized normal equation:

where the coefficients are obtained first by an LU decomposition of $ZTZ+λI$⁠. The time complexity of constructing the Z^TZ matrix and performing inference scale as $O(Nf2NtMd)$ and O(N_fMd), respectively. To include forces in the training scheme, the derivatives of the feature vectors $∂zIl∂rik$ are computed, where l and k are the feature and coordinate component indices, respectively, and stored in a derivatives feature matrix $∂Z∈R3NtM×Nf$⁠. The following regularized normal equation is then solved:

where the notation $Z,∂Z$ indicates the concatenation of the feature matrix Z with the derivative features ∂Z. The dominant term in constructing the normal equations is the ∂Z^T∂Z matrix product, which scales as $O(3NtMNf2)$⁠.

The above formulation revolves around computing the affine transformation $W[ρI,1]T$⁠, which is applied to ρ_I prior to the cosine operation $(z(ρI)=cos(W[ρI,1]T))$⁠. Storing and computing this linear transformation has O(N_fd) space and time complexity. To reduce this, Yu et al.¹⁹ introduced structured orthogonal random features (SORF). In this method, the matrix W is replaced by a special structured matrix consisting of products of random binary diagonal matrices and Walsh–Hadamard matrices. The resulting linear transformation has O(N_f log d) time complexity and O(d) or O(1) space complexity, depending on its implementation. Briefly, in the case of Gaussian kernel approximation, one can replace the transformation

where $G∈RNf×d$ is a random Gaussian matrix, with the following transformation:

where Q is a uniformly distributed random orthogonal matrix (e.g., obtained via QR decomposition of G) and S is a diagonal matrix with entries sampled i.i.d from the χ-distribution with d degrees of freedom. The resulting matrix SQ is an unbiased estimator of the Gaussian kernel with low variance.¹⁹ While this construction still has O(N_fd) time complexity as well as the additional cost of computing the QR decomposition in a preprocessing step, one can further approximate this transformation as

where S has first been replaced by a scalar $d$ and the random orthogonal matrix Q has been replaced by a special type of structured matrix. The matrices $Di∈Rd×d$ are diagonal sign-flipping matrices, where each entry is sampled from a Rademacher distribution, and H is the Walsh–Hadamard matrix,

Note that when the number of features N_f > d, the operation is simply repeated $Nfd$ times, with the resulting vectors concatenated into a length N_f vector. Crucially, the product W_SORFρ_I now has time complexity O(N_transformN_f log d), since the multiplication with H can be efficiently implemented via the fast Hadamard transform using in-place operations in O(d log d) time. Finally, since the Walsh–Hadamard matrix is only defined in $R2n×2n$⁠, ρ_I must also be projected into 2ⁿ dimensions. This is achieved via an SVD decomposition on a subset of the atomic environment representations for each element e, concatenated into the matrix $Z̃e$ given by

and the atomic representations are projected into a lower dimension via the following matrix-vector product,

where only the first N_PCA columns from the matrix P_e are used. The subscript e indicates that the matrix $Z̃e$ is built using only atomic representations of atom type e; hence, each element has its own projection matrix $PeNPCA$⁠. Here, we have found N_PCA = 128 or 256 to be sufficient. In the rest of the text, d refers to the dimension of the atomic representation after projection.

We note that there are a number of other approximate kernel methods that aim to reduce computational complexity, including other RFF-type approximations^33–35 as well as those based on the Nyström method,³⁶ which relies on low-rank structure in the kernel matrix. Here, however, we have opted to use SORF due to its simplicity, computational efficiency, and accuracy.³⁵ 

Finally, for RFF and SORF models there are two user-defined hyperparameters: σ, λ. In all results presented in this work, the kernel width parameter (σ) and regularization parameter (λ) are selected from a grid using fivefold cross-validation; however, we note that suitable σ and λ values can be found within the range of 2–12 and 10⁻⁵ to 10⁻⁸, respectively.

In this work, we use FCHL19²⁰ as the permutationally and rotationally invariant atomic environment featurization layer. FCHL19 is an atom-centered representation consisting of two- and three-body elemental bins, similar in construction to the atom-centered symmetry functions (ACSFs) used in the work of Behler.^12,37 The functional form is briefly summarized here. For every unique combination of two elements $X,Y$⁠, the representation for each atom I is constructed as follows:

where $Z,RX,Y$ refers to the set of atomic charges and coordinates that have either element $X$ or $Y$⁠. The two-body function is given by

where R_s are the n₂ radial grid centers linearly distributed between 0 and r_cut and μ(r_IJ) and σ(r_IJ) are the parameters of the log normal distribution,

where η₂ is a hyperparameter. The cutoff function to smoothly decay the representation to zero at r_cut is defined as

The three-body term $GX,Y3-body$ is given by the following function:

The radial term $Gradial3-body$ is given by the following expression:

where η₃ is a parameter that controls the width of the n₃ radial distribution functions, located at R_s grid points. The three-body scaling function ξ₃ is the Axilrod–Teller–Muto term^38,39 with modified exponents,⁴⁰ 

where θ_KIJ is the angle between atoms I, J, K, with I at the center, c₃ is a weight term, and N₃ is a three-body scaling factor. Finally, the angular term is given by a Fourier expansion as follows:

where the cosine and sine terms are given by

where ζ is a parameter describing the width of the angular Gaussian function and n > 0 is the expansion order. Similar to previous work, only the two n = 1 cosine and sine terms are used.

The optimal parameters to generate the FCHL19 representation differ here from those used in the original implementation.²⁰ We have found improved parameters using a lower cutoff of r_cut = 6.0 Å. To fit these parameters, we employed a subset of 576 distorted geometries of small molecules with up to five atoms of the type CNO, saturated with hydrogen atoms, for which forces and energies have been obtained from density functional theory (DFT) calculations.^20,41 This dataset is identical to that used in the original FCHL19 publication. This dataset is randomly divided into a training set of 384 geometries and a test set of 192 geometries. Models are fitted to the training set and predictions on the test set are used to minimize the following cost function with respect to the parameters:

where U_i is the energy of molecule i and F_i and n_i are the forces and number of atoms for the same molecule, respectively. A greedy Monte Carlo optimization was used to perform this optimization, where real-type parameters are optimized by multiplying with a factor randomly chosen from a normal distribution centered on 1 with a variance of 0.05 and integer-type parameters by randomly adding +1 or −1. The final parameters that were found to work best are listed in Table I.

The FCHL19 representation is constructed by assigning each atom I to each block in a batch. One block is launched for each atom in the system. For each block, a total of 256 threads are used in a two-dimensional thread grid. The first dimension of this grid contains 16 threads and enumerates over all two-body interactions with central atom I to construct $GI2$⁠, while the second dimension contains eight threads, which enumerates the third index in the three-body interaction $GI3$⁠. For the forwards pass of the model, the reduction of all $GI2$ and $GI3$ scalar elements is performed global memory. For the backward pass, however, the atomic force components are stored and summed in (local, on-chip) memory, therefore significantly increasing throughput. A simple tiled neighbor-list is used to linearize the cost of FCHL19 with respect to increasing number of atoms in the local environments; therefore, only atoms within the cutoff radius are considered when constructing the representation. Once the atomic representation has been constructed, it is projected to a lower dimension of size N_PCA using a matrix obtained from an SVD of a randomly selected subset of atomic representations from the training set. The size of this lower-dimensional vector is constrained to be a power of 2 for the purposes of the Hadamard transform. Each element within the training database has its own projection matrix and is used to project down each atomic representation separately.

For the SORF forward pass, each block handles all $Nfd$ Hadamard transforms for a single atom to produce a feature vector of the desired length N_f. The Hadamard transform itself operates on the projected FCHL19 representation (dimension d = 2ⁿ, n ≥ 1), after multiplication with the diagonal sign-flipping matrix, using a shared-memory butterfly operation that has O(d log d) complexity. This operation is performed iteratively N_transform times. In this work, either d = 128 or d = 256 is used and N_transform is set to 1. For the backward pass, the gradients are stored and reduced in shared memory.

We begin by briefly comparing the performance of RFF-type and SORF-type approximations with those of a sparse GP method. We have implemented a sparse GP model in QML-Lightning that uses a Gaussian kernel where the sparse points for each element are uniformly randomly selected from the available training data. The left column of Fig. 1 shows the convergence of the out-of-sample mean absolute errors (MAEs) for energy with respect to increasing N_f used to approximate the kernel as well as the wall time to compute 1k configurations. For the sparse GP model, we set the total number of sparse points N_s to be equal to N_f. Here, we have used the aspirin trajectory from the unrectified MD17 database.^22–24 The amount of training data remains fixed, using 1k configurations and training on energies only. The right column of Fig. 1 displays the aforementioned energy MAEs but additionally against an estimate of the number of operations performed per atom for SORF, RFF, and sparse GP models. Since we are only concerned with relative differences, the horizontal axis has been normalized to 1.

For a fixed number of features, the SORF and RFF models outperform the sparse GP model, with SORF having the lowest error. Furthermore, SORF models are faster than both RFF and sparse GP models for all model sizes. We note that, while RFF and sparse GP models have the same formal cost, warp matrix-multiply-accumulate (WMMA) operations are accelerated on TensorCore-enabled GPUs such as the RTX 3080 used here, thereby making the calculation of matrix-vector product $W[ρI,1]T$ significantly faster. However, increasing N_f incurs a numerical error in this matrix-vector product for the RFF model if performed in single-point precision. To remedy this controlled-error, low-precision matrix-vector products can be computed with 3 WMMA operations, which has not been implemented in this work. While this would increase the cost of RFF, it would not significantly change the ordering of results presented here. We note that this does not appear to be an issue for the SORF model since the product $W[ρI,1]T$ is computed via a series of fast Hadamard transforms requiring only simple additions.

In Fig. 2, the predictive accuracy of several explicit kernel models and SORF models for atomization energy of molecules in the QM9 dataset²¹ are compared. These models include FCHL18⁴⁰ and FCHL19²⁰ using the OQML²⁰ regressor, which is equivalent to a sparse GP model in which all training atoms are used in the kernel basis. For the SORF models, learning curves using both N_f = 16 384 and N_f = 32 768 are shown. We find that the SORF models with FCHL19 perform similar to OQML with FCHL19: The MAE for OQML/FCHL19 and SORF-32k/FCHL19 at 75 000 training samples are 11 and 12 meV, respectively. We note that there is a small deviation away from linearity in the learning curve at $≈50000$ training samples, indicating that more features may be required.

Figure 3 reports the energy and force MAEs as a function of number of training samples using three molecules from the MD-17dataset.^22–24 To be consistent with previous literature,²⁰ we use the unrectified MD-17 dataset, which is known to contain significant noise on the energy values.²⁵ The learning curves for SORF/FCHL19 with both N_f = 16 384, 32 768 are reported. We compare against OQML models based on FCHL19²⁰ as well as GDML²² and sGDML⁴² models. Additionally, SchNet⁴³ and state-of-the-art NequIP⁴⁴ neural networks have been included.

For energy learning, the SORF/FCHL19 models in general display similar accuracies to the OQML/FCHL19 model. For toluene, naphthalene, and salicylic acid, the OQML/FCHL19 model slightly outperforms the SORF/FCHL19 models for both N_f = 16 384 and N_f = 32 768. However, in all other cases, the SORF models display similar or better accuracy. Both sGDML and GDML perform worse than OQML and SORF models in general; however, for toluene and naphthalene specifically, sGDML has the lowest error among the kernel models.

For force learning, the SORF-based models are as accurate or better than OQML/FCHL19 for N_f = 16 384, and they are reasonably more accurate than OQML for N_f = 32 768. These models outperform SchNet in all cases, while NequIP outperforms the SORF models in all cases. We note that the training times for SORF/FCHL19 are on the order of seconds, while OQML/FCHL19 and sGDML models take several minutes to train. Furthermore, GDML models take several hours, and SchNet and NequIP are trained over hours to days. We additionally provide a comparative benchmark of the revised MD-17 dataset,²⁵ a recomputed version of the original MD-17 dataset with tighter SCF convergence criteria. Table II lists the out-of-sample MAEs for the largest SORF model constructed in this work with FCHL19, GP with FCHL19,²⁵ sGDML, ACE,⁴⁵ and NequIP, with rotation orders l = 0 and l = 3. We note that FCHL19 is a comparatively simplistic atomic featurization layer; consequently, it is expected that it does not perform as well as state-of-the-art equivariant many-body neural networks^46–48 such as NequIP. For a more reasonable comparison, the l = 0 channel NequIP model, which contains at most 3-body terms similar to FCHL19, has been included here. While the force errors are similar to the MD-17 results, the energy errors are significantly lower across all models, which is consistent with the observation that the noise floor on the original MD-17 data is higher on the energies.

The 3PBA dataset^26,45 contains both ambient and high temperature configurational samples for the small drug-like molecule, 3-(bezyloxy)pyridin-2-amine (3BPA). This molecule has 3 central rotatable dihedral angles (α, β, γ), leading to a complex dihedral potential energy surface with many local minima, which can be challenging for both classical and ML-based potentials.⁴⁹ In particular, at ambient temperatures (300 K) the accessible phase space is small; however, the dataset contains both 600 and 1200 K configurational samples, which have increasingly large phase space. Therefore, a crucial test on whether an ML model can extrapolate well is if the model, when trained on the 300 K samples, can accurately predict the 600 and 1200 K samples. Consequently, we have trained the SORF/FCHL19 model on the 300 K subset of the 3PBA dataset in order to compare against results obtained from ACE, two kernel models sGDML and GAP²⁹ using SOAP,^50,51 as well as two related neural network architectures ANI^52,53 and ANI-2x.⁵⁴ We note that these results have been summarized directly from the ACE publication.⁴⁵ 

Table III lists the energy and force root-mean-squared errors (RMSEs) of a variety of different models listed in the ACE paper,⁴⁵ including the SORF model with both 16 384 and 32 768 features. For the 300 K test dataset, the SORF model with N_f = 32 768 achieves very low out-of-sample errors, having more accurate forces than all models except ACE. For the 600 K dataset, the SORF models do not fare as well as the ACE and ANI models; however, they significantly outperform both kernel based models sGDML and GAP with SOAP features. For the 1200 K dataset, the SORF model has lower force errors than ACE; however, the ANI models have a reasonably lower error. We note that ACE contains up to five-body terms in its cluster expansion, while FCHL19 contains only up to three-body terms. Therefore, one would expect ACE to outperform the FCHL19/SORF models in this setting.

Table IV shows the total time spent in seconds when training various CPU and GPU-accelerated models using both energies and forces, given a fixed amount of training data (N_t = 1k) from the MD-17 database. The device used to train these models has been listed for reference. Three kernel models are shown: OQML and sGDML. We note that OQML and GP models use the same representation as the GPU variant implemented here. We additionally compare against NequIP; however, the training times shown are an upper bound as models with good predictive accuracy can likely be obtained with early stopping. The timings for the SORF model with both 16 384 and 32 768 features are shown using a consumer RTX 3080 GPU and scientific A100 GPU. It should be noted that there is a significant cumulative round-off error when constructing the Z^TZ matrix in floating-point precision (FP32), which leads to ill-conditioning. However, Kahan’s summation⁵⁵ has been implemented here to minimize this error, allowing the normal equations to be constructed in FP32 format to yield a threefold reduction in training time compared to FP64 precision. For the A100 GPU, only FP64 performance is shown as we observed minimal gain in speed when using FP32. Using the RTX 3080 GPU, SORF models are found to be several orders of magnitude faster to train than both GP and neural networks. Furthermore, they are on average ≈3 and ≈210 times faster to train than OQML and sGDML models, respectively, across the dataset shown when using FP64 matrix multiplication, while they are ≈10 and ≈700 times faster to train when using FP32 with Kahan’s summation. On the A100 GPU, trained models can be obtained in seconds. Note that this work focuses on models that can fit into GPU global memory, i.e., both the Z^TZ matrix and the normal equations are constructed and solved on the GPU. However, QML-Lightning also supports an out-of-core approach, where the Z^TZ matrix is tiled, with each tile being copied to the host device and summed. Finally, the normal equations are then solved on the CPU. This second variant is necessary when the number of features yields matrices that exceed the memory required to both construct the Z^TZ matrix as well as solve the normal equations. The choice of CPU(s) or out-of-core CPU/GPU implementation will significantly determine the model training time in these cases. However, this aspect has not been investigated in this work.

In Fig. 4, we show the force accuracy of the SORF model with FCHL19 as a function of the number of features, and the corresponding wall time to complete 1k MD steps for azobenzene using the OpenMM QML-Lightning interface. For additional context, we compare against Performant ACE⁵⁶ (PACE), which is an MPI-parallelized and optimized version of ACE.⁴⁵ For the SORF model, an RTX 3080 was used, whereas results for PACE have been provided by the first author of the ACE paper, where 8 MPI-tasks were found to achieve best time performance on a 16-core Intel(R) Xeon(R) Gold 5218 CPU @ 2.30 GHz processor. The reader should be cautioned that there are a number of challenges in directly comparing GPU and CPU codes and that constructing a benchmark that is 100% fair is very challenging. PACE yields significantly lower wall times for smaller models, being approximately an order of magnitude faster than the SORF model when using 4k basis functions and features, respectively. Small molecules represent worst-case systems for GPU codes since there are a number of idle streaming multiprocessors, and those that are active are computing relatively little amount of work. However, it is encouraging to see that even for small systems the SORF model achieves a similar error and wall time as PACE at 16k features. Doubling the number of features to 33k reasonably improves the error while increasing the wall time cost by only 1 s; however, the error begins to saturate here and we do not expect a significant improvement in error at 65k features. Conversely, the linearization achieved in PACE is over an increasing number of higher order many-body correlations; therefore, the model’s accuracy improves as the number of basis functions increases, albeit with a larger computational cost.

Table V lists the total wall time in seconds for completing 1000 MD steps for a variety of systems using the OpenMM QML-Lightning interface. Additionally, per atom per MD step timings and the average and maximum number of neighbors per atom are provided. For representative small-molecule timings, azobenzene is listed. For larger nonperiodic systems, a brass cluster^57,58 and a vitamin D receptor (5OWD)⁵⁹ with interstitial water molecules are included. For larger periodic systems, two metal organic frameworks (MOFs) obtained from the materials project⁶⁰ and liquid water are included. For the two monoclinic MOF systems, a 3 × 3 × 3 supercell was used. For the Zn₄H₄₄C₅₁N₈O₁₈ system (qmof-0b5e989),^61,62 the lattice constants and angles are [a, b, c] = [16.02, 21.41, 29.93] Å and [α, β, γ] = [90.0, 90.0, 97.8]°, respectively. For the Nd₂C₃₆H₄₂N₁₂O₁₈S₃ system (qmof-bf74329),^61,62 the lattice constants and angles are [a, b, c] = [12.14, 16.95, 25.66] Å and [α, β, γ] = [102.2, 90.0, 90.0]°, respectively. For liquid water, a 2 × 2 × 2 supercell was used with a cubic unit cell with a lattice constant of a = 25.0 Å, generated using OpenMMTools.⁶³ We note that since the our optimization strategy parallelizes the computation over atoms, with each atom being assigned to a different block in round-robin fashion, smaller systems yield an underutilized GPU as each atom is mapped to a different SM, of which there are 68 on an RTX 3080. Furthermore, the small number of neighbors yields a limited amount of thread-parallel work to perform when constructing FCHL19 within each SM. For the brass cluster, the timings per atom per step are significantly improved since the SMs are fully utilized. However, approximately only four atoms are computed per SM; therefore, there is limited work, which yields comparatively poor timings with respect to the larger systems. For large systems, near-microsecond per atom per MD step timings can be observed; for example, forces are evaluated in 5.5 µs/(atom ⋅step) for the Zn₄H₄₄C₅₁N₈O₁₈ system and 8.56 µs/(atom ⋅step) for the Nd₂C₃₆H₄₂N₁₂O₁₈S₃ system. We expect that sub-microsecond timings could be obtained using low-precision computation.

For a component breakdown of the computational cost, Fig. 5 shows the approximate on-device time for key components in computing energies and forces for azobenzene (N_atoms = 24), a brass cluster (N_atoms = 245), a vitamin D receptor (N_atoms = 4369), and the periodic MOF system Nd₂C₃₆H₄₂N₁₂O₁₈S₃ (N_atoms = 12 204) using N_f = 16 384. Profiling PyTorch GPU codes yields an additional cost to each node on the autograd graph; therefore, we display the results in terms of percentage of the total device time. For azobenzene, the forward and backward passes of FCHL19 and SORF require only 28% of the total device time, while indexing operations, initialization, and intermediary operations require ∼70%. We note that many of these operations could be optimized away through combining them into single kernel calls. In particular, it is likely that the indexing costs, which are mostly associated with sub-selecting the atomic representations for projection with the relevant matrices $PeNPCA$⁠, could be removed through using customized CUDA C functions to generate these indices for use directly in the FCHL19 and SORF modules. For the larger brass cluster, forward and backward passes of FCHL19 and SORF occupy 47% of the total device time; however, indexing and other costs still occupy a large percentage of the total device time. Here, the percentage cost of the forward and backward passes of SORF compared to azobenzene are reduced from 23% to 18%. For the vitamin D receptor, the forward and backward passes of FCHL19 begin to dominate the cost, accounting for 24.5% and 45.7%, respectively, while the cost of the SORF transforms is reduced to 10% of the total device time. Finally, for the periodic MOF system, the forward and backward passes of FCHL19 occupy 80.5% of the total cost, while the SORF transform only requires 7.6% of the total device time.

In this paper, we have introduced a PyTorch-based library, termed QML-Lightning, which contains approximate kernel models and an efficient representation designed for learning quantum mechanical properties. We have provided a low-cost, PyTorch-wrapped CUDA C implementation of structured orthogonal random features, a variant of the well-known random Fourier features, as well as a computationally efficient implementation of FCHL19, an accurate atom-centered representation.

The combination of structured orthogonal features and FCHL19 has been benchmarked against existing datasets, yielding not only similar or better accuracy than explicit kernel models with FCHL19 but also competitive accuracy with contemporary models, with significantly reduced training time and performant prediction time.

This research was supported by the NCCR MARVEL, a National Center of Competence in Research, funded by the Swiss National Science Foundation (Grant No. 182892). Calculations were performed at sciCORE (http://scicore.unibas.ch/) scientific computing center at University of Basel.

We would also like to thank David Kovacs, Ilyes Batatia, and Gábor Csányi for fruitful discussions and David Kovacs for providing PACE timings.

The authors have no conflicts to disclose.

Nicholas J. Browning: Conceptualization (lead); Data curation (lead); Formal analysis (equal); Investigation (lead); Methodology (lead); Software (lead); Visualization (equal); Writing – original draft (lead); Writing – review & editing (equal). Felix A. Faber: Formal analysis (equal); Validation (lead); Visualization (equal); Writing – original draft (equal); Writing – review & editing (equal). O. Anatole von Lilienfeld: Funding acquisition (lead); Project administration (lead); Writing – review & editing (equal).

The data that support the findings of this study are available within the article.

$C$

set of atoms belonging to a compound

d

size of atomic representation.

F

matrix containing force components

G

atom density function

H

Walsh-Hadamard matrix

I

identity matrix

I, J

indices denoting atoms I and J, respectively

K

Kernel matrix

k(⋅, ⋅)

Kernel function

L

loss function

$PeNPCA$

atomic projection matrix for element e

N_f

number of features for the RFF approximation

N_s

number of sparse points in sparse GP basis

N_t

number of training instances

U

matrix containing total formation energies

U

total formation energy of a compound

w

angular frequency of a Fourier transform

Z

matrix containing approximate feature mappings

z(⋅)

RFF approximation of the feature mapping

α

regression weights

ρ

atomic representations

ϕ(⋅)

feature mapping of a kernel function

ϵ

atomic energy contributions