Unified representation of molecules and crystals for machine learning

The computational study of atomistic systems such as molecules and crystals requires accurate treatment of interactions at the atomic and electronic scale. Accurate first-principles methods, however, are limited by their high computational cost. In settings that require many calculations, such as dynamics simulations, phase diagrams, or high-throughput searches, machine learning (ML) [1, 2] can reduce overall costs by orders of magnitude via accurate interpolation between reference calculations [3–5]. For this, the problem of repeatedly solving a complex equation such as Schrödinger’s equation for many related inputs is mapped onto a non-linear regression problem: instead of numerically solving new systems, they are statistically estimated based on a reference set of known solutions [6, 7]. This ansatz enables, among other applications, screening larger databases of molecules and materials [5, 8], running longer dynamics simulations [9], investigating larger systems [10], and increasing the accuracy of calculations [5, 11].

Kernel-based ML models [12–15] for data-efficient accurate prediction of ab initio properties require a single space in which regression is carried out. Representations [16] are functions that map atomistic systems to elements in such spaces, either directly or via a kernel [17]. Representations should be (a) invariant against transformations preserving the predicted property, in particular translations, rotations, and nuclear permutations of same elements, as learning these invariances from data would require many reference calculations; non-scalar properties can require equivariance instead of invariance; (b) unique, that is, variant against transformations changing the property, as systems with identical representation that differ in property would introduce errors [18]; (c) continuous, and ideally differentiable, as discontinuities work against the smoothness assumption of the ML model and model gradients are often useful; (d) general in the sense of being able to encode any atomistic system, including finite and periodic systems; (e) fast to compute, as the goal is to reduce computational cost; (f) data-efficient in the sense of requiring few reference calculations to reach a given target error. Constant size is an advantage, [19] as is the ability to encode the whole system as well as local atomic environments. Requirements (e) and (f) are in practice determined empirically. See [16, 20–22] for details on these and further requirements.

Some representations fulfill these requirements only partially, such as the Coulomb matrix (CM) [6] and bag of bonds (BoB) [23] discussed below. State-of-the-art representations often fulfill these requirements in some limit, such as infinite expansion order. See [16] for a comprehensive and detailed discussion. The descriptors used in cheminformatics, [24] and sometimes in materials informatics, often violate (b) and (c), in particular if they do not include atomic coordinate information or rely on cutoff-based definitions of chemical bonds. Such descriptors serve the different purpose of directly predicting derived properties that are not functions of a single conformation, such as solubility or binding affinity to a macromolecule.

We introduce a many-body tensor representation (MBTR) derived from CM/BoB and concepts of many-body expansions. It is related [16] to Behler–Parrinello symmetry functions [25] and histograms of distances, angles, and dihedral angles [26]. MBTR fulfills the above requirements in the limit, is interpretable, allows visualization (figure 1), and describes finite and periodic systems. State-of-the-art empirical performance is demonstrated by us for organic molecules and inorganic crystals, as well as applicability to phase diagrams of Pt-group / transition metal binary systems, and by others for predicting and optimizing various molecular [27–33] and crystalline properties [34–37]. Implementations of MBTR are publicly available (see Code Availability section at the end).

Zoom In Zoom Out Reset image size

We start from the CM [6, 10, 38], which represents a molecule $\mathcal{M}$ as a symmetric atom-by-atom matrix

$$\begin{equation} \boldsymbol{M}_{i,j} = \begin{cases} \frac{1}{2} Z_i^{2.4} & i = j \\[3pt] \frac{Z_i Z_j}{d_{i,j}} & i \neq j \end{cases} , \end{equation} \tag{ 1 }$$

where Z_i are atomic numbers and $d_{i,j} = ||\boldsymbol{R_i}-\boldsymbol{R_j}||$ is Euclidean distance between atoms i and j. To avoid dependence on atom ordering (in the input), which would violate (a), M is either diagonalized, loosing information which violates (b) [18], or sorted, causing discontinuities that violate (c). Another shortcoming is the use of Z, which is not well suited for interpolation [39] as it overly decorrelates chemical elements from the same column of the periodic table.

The related BoB [23] representation uses the same terms, but arranges them differently. For each pair of chemical elements, corresponding CM terms are stored in sorted order, which can be viewed as an $N_\mathrm{e} \times N_\mathrm{e} \times d$ tensor, or an $N_\mathrm{e} \times (N_\mathrm{e} + 1)/2 \times d$ tensor if symmetry is taken into account, where $N_\mathrm{e}$ is number of elements and d is sufficiently large. Unlike the CM, it can not distinguish homometric molecules [21], which might distort its feature space [40]. While the BoB tensor itself does not suffer from discontinuities, its derivative does.

To derive MBTR, we retain stratification by elements, but avoid the sorting by arranging distances on a real-space axis:

$$\begin{equation} f_{\textrm{BoB}}\bigl(x,z_1,z_2\bigr) = \sum_{i,j = 1}^{N_\mathrm{a}} \delta\bigl(x-d_{i,j}^{\,-1}\bigr) \delta(z_1,Z_{i}) \delta(z_2,Z_{j}) , \end{equation} \tag{ 2 }$$

where x is a real number, $z_1,z_2$ are atomic numbers, $N_\mathrm{a}$ is number of atoms, $\delta(\cdot)$ is Dirac’s delta, and $\delta(\cdot,\cdot)$ is Kronecker’s delta. f_BoB has mixed continuous-discrete domain and encodes all (inverse) distances between atoms with elements z₁ and z₂. For a smoother measure, we replace Dirac’s δ with another probability distribution $\mathcal{D}$, ‘broadening’ or ‘smearing’ it [20, 41]. In this work, we use the normal distribution. Other distributions can be used, in particular symmetric and short-tailed ones, for example, the Laplace distribution or the uniform distribution. We did not observe significant differences in performance, however. Adding a weighting function w₂ and replacing the Kronecker δ functions by an element correlation matrix $C \in \boldsymbol{R}^{N_\mathrm{e} \times N_\mathrm{e}}$ yields

$$\begin{equation} f_2\bigl(x,z_1,z_2\bigr) \! = \!\! \sum_{i,j = 1}^{N_\mathrm{a}} \!w_2(i,j) \, \mathcal{D}\bigl(x,g_2(i,j)\bigr) C_{z_1,Z_{i}} C_{z_2,Z_{j}} \end{equation} \tag{ 3 }$$

of which equation (2) is a special case. In general, g₂ describes a relation between atoms i and j, such as their inverse distance, $\mathcal{D}$ broadens the result of g₂, and w₂ allows to weight down contributions, for example, from far-away atoms. Building on the idea of many-body expansions, [42, 43] we generalize from f₂ in equation (3), which encodes two-body terms, to the MBTR equation

$$\begin{equation} f_k(x,\boldsymbol{z}) = \sum_{\boldsymbol{i} = 1}^{N_\mathrm{a}} w_k(\boldsymbol{i})\mathcal{D}\bigl(x, g_k(\boldsymbol{i})\bigr) \prod_{j = 1}^k C_{z_{j},Z_{i_j}} , \end{equation} \tag{ 4 }$$

where $\boldsymbol{z} \in \boldsymbol{N}^k$ are atomic numbers, $\boldsymbol{i} = (i_1,\ldots,i_k) \in \{1,\ldots,N_\mathrm{a}\}^k$ are index tuples, and w_k , g_k assign a scalar to k atoms in $\mathcal{M}$ [44]. Canonical choices of g_k for $k \! = \! 1,2,3,4$ are atom counts, (inverse) distances, angles, and dihedral angles. The element correlation matrices C allow exploitation of similarities between chemical element species (‘alchemical learning’), for example, within the same column of the periodic table [45–47].

We measure the similarity of two atomistic systems $\mathcal{M}$ and $\mathcal{M}^{^{\prime}}$ as the Euclidean distance between their representations,

$$\begin{equation} d_k^{2}(\mathcal{M},\mathcal{M}^{^{\prime}}) = \sum_{\boldsymbol{z}} \int \bigl( f_k(x,\boldsymbol{z}) - f_k^{^{\prime}} (x,\boldsymbol{z}) \bigr)^2 \mathrm{d} x.\end{equation} \tag{ 5 }$$

In practice, we adjust equation (4) for symmetries. Discretizing the continuous axis as $\left(x_{\min},x_{\min}+\right.$ $\left.\Delta x,\ldots,x_{\max}\right)$ results in a rank $k\!+\!1$ tensor of dimensions $N_\mathrm{e} \times \cdots \times N_\mathrm{e} \times N_x$ with $N_x = (x_{\max}-x_{\min})/\Delta x$, where $x_{\min}$ and $x_{\max}$ are the smallest and largest values for which $f_k(x,\boldsymbol{z}) \neq 0$ for all z and $\mathcal{M}$. Linearizing element ranks yields $N_\mathrm{e}^k \times N_x$ matrices, allowing for visualization (figure 1) and efficient numerical implementation via linear algebra routines. For systems with many element species, discretization can lead to large matrices, requiring substantial amounts of memory. In such settings, memory-efficient implementation via sparse matrix formats or on-the-fly calculation of distances and inner products (see, e.g. [45]) of MBTR matrices might be preferable.

Periodic systems, used to model bulk crystals and surfaces, can be viewed as unit cells surrounded by infinitely many translated images of themselves. For such systems, $N_\mathrm{a} = \infty$ and the sum in equation (4) diverges. We prevent this by requiring one index of i to be in the (same) primitive unit cell [48]. This accounts for translational symmetry and prevents double-counting. Use of weighting functions w_k such as exponentially decaying weights [49] then ensures convergence of the sum. Figure 1 (right) presents the resulting distributions of angles for face-centered cubic (FCC) NaCl as an example. Note that the k-body terms g_k do not depend on choice of unit cell geometry (lattice vectors). This ensures unique representation of Bravais lattices where the choice of basis vectors is not unique, for example 2D hexagonal lattices where the angle between lattice vectors can be $\frac{1}{3}\pi$ or $\frac{2}{3}\pi$.

Many applications, including dynamics simulations and structural relaxation, require forces, the negative gradient of the energy with respect to atomic coordinates. The gradient of equation (4) is given by

$$\begin{align} \nabla f_k(x,\boldsymbol{z}) & = \sum_{\boldsymbol{i} = 1}^{N_\mathrm{a}} \Biggl( \mathcal{D}\bigl(x, g_k(\boldsymbol{i})\bigr) \nabla w_k(\boldsymbol{i}) \nonumber\\ &\quad+ w_k(\boldsymbol{i}) \frac{\partial\,\mathcal{D}\bigl(x, g_k\bigr)}{\partial g_k} \nabla g_k(\boldsymbol{i}) \Biggr) \prod_{j = 1}^k C_{z_{j},Z_{i_j}}. \end{align} \tag{ 6 }$$

The gradient $\nabla f_k(x,\boldsymbol{z}$) can be derived analytically if this is possible for $\nabla w_k$, $\nabla g_k$, and $\nabla \mathcal{D}$. Alternatively, automatic differentiation [50, 51] can be used, removing the need for manual derivation. Figure 1 visualizes MBTR gradients.

To validate MBTR, we demonstrate accurate predictions for properties of molecules and crystals. Focusing on the representation, we employ plain kernel ridge regression (KRR) models [7] unless stated otherwise.

3.1. Changes in molecular structure

3.2. Changes in crystal chemistry

Interpolation across changes in the chemistry of crystalline solids is demonstrated for a dataset of 11k elpasolite structures (ABC₂D₆, AlNaK₂F₆ prototype) [56, 57] composed of 12 different elements, with geometries and energies computed at DFT/PBE level of theory. Predicting formation energies with MBTR yields a root-mean-squared-error (RMSE) of 8.1 meV/atom and MAE of 4.7 meV/atom (figure 2) for a training set of 9k crystals.

Zoom In Zoom Out Reset image size

Adding chemical elements should increase the intrinsic dimensionality of the learning problem, and thus prediction errors. To verify this, we created a dataset of 4611 ABC₂ ternary alloys containing 22 non-radioactive elements from groups 1, 2, 13–15, spanning five rows and columns of the periodic table. Structures were taken from the Open Quantum Materials Database (OQMD) [58, 59], with geometries and properties also computed via DFT/PBE. As expected, energy predictions exhibit larger errors (RMSE 31 meV/atom, MAE 23 meV/atom) compared to an elpasolite model of same training set size (RMSE 23 meV/atom, MAE 15 meV/atom).

3.3. Changes in molecular geometry

For interpolation of changes in molecular geometry, we employ a benchmark dataset [13, 62] of ab initio molecular dynamics trajectories of eight organic molecules. Each molecule was simulated at a temperature of 500 K for between 150k and 1M time steps of 0.5 fs, with energies and forces computed at the DFT/PBE level of theory and the Tkatchenko–Scheffler model [63] for van der Waals interactions. Table 2 presents results for models trained and evaluated only on energies (right-hand side) and only on forces (left-hand side).

Table 2.  Energy and force prediction errors for changes in geometry of organic molecules. Shown are prediction errors for total energies (kcal mol⁻¹) and atomic forces (kcal mol⁻¹ Å⁻¹). MAE = mean absolute error, RMSE = root mean squared error, PaiNN = polarizable atom interaction neural network [60], FCHL19 = Faber–Christensen–Huang–Lilienfeld representation [45, 47], sGDML = symmetric gradient domain machine learning [13], sMatérn = Matérn kernel augmented with symmetric permutations for sGDML [61], CM^md = Coulomb matrix variant, MBTR = many-body tensor representation.

	Trained only on forces, 1k references								Trained only on energies, 10k references
	PaiNN		FCHL19		sGDML/CMmd		sGDML/MBTR		CMmd	MBTR	MBTR
Kernel	—		Gaussian		sMatérn		sMatérn		Gaussian	linear	Gaussian
	Energy	Force	Energy	Force	Energy	Force	Energy	Force	Energy	Energy	Energy
Molecule	MAE	MAE	MAE	MAE	MAE	MAE	MAE	MAE	MAE	MAE	MAE
Benzene	—	—	—	—	0.07 a	0.06 a	0.07	0.15	0.03	0.03	0.03
Uracil	0.11	0.13	0.10	0.10	0.11	0.24	0.11	0.17	0.05	0.10	0.03
Naphthalene	0.12	0.08	0.12	0.15	0.12	0.11	0.11	0.09	0.12	0.10	0.07
Aspirin	0.17	0.34	0.17	0.50	0.19	0.68	0.17	0.48	0.36	0.21	0.25
Salicylic acid	0.12	0.20	0.12	0.22	0.12	0.28	0.11	0.18	0.11	0.13	0.07
Malonaldehyde	0.10	0.34	0.08	0.25	0.10	0.41	0.09	0.36	0.18	0.21	0.10
Ethanol	0.06	0.22	0.05	0.14	0.07	0.33	0.06	0.26	0.17	0.17	0.06
Toluene	0.10	0.09	0.10	0.20	0.10	0.14	0.09	0.13	0.16	0.11	0.10

^a We observed higher noise in predictions of benzene, whose reported prediction errors are also inconsistent in different publications. To make results comparable, we retrained the sGDML/CM^md model (originally reported values are 0.10 and 0.06).

Energy-only models were trained on 10k configurations and validated on 2k other ones, employing MBTR (parametrized for dynamics data, see supplement) and a similarly modified CM (CM^md, see supplement). Non-linear MBTR regression performs best overall, with the linear kernel again being competitive.

On the one hand, differentiating energy-based ML potentials can introduce errors, for example, from small oscillations between training samples due to insufficient regularization, and from insufficient model constraints in directions not covered by the training data [64]. On the other hand, electronic structure calculations often provide reference forces at not additional cost. It is therefore beneficial to include these in model training. This often reduces the required number of reference calculations by an order of magnitude [9, 13, 61, 65].

Force-only models require an adaptation of plain KRR. To accommodate forces into training and to demonstrate use of MBTR with other regression approaches, we employ MBTR in the framework of symmetrized gradient-domain machine-learning (sGDML) [61]. This approach uses the Matérn kernel, augmentation by symmetric molecular permutations, and reference forces for training, while providing energy and force predictions.

Table 2 (left-hand side) compares performance of the original sGDML approach (based on the CM^md representation [13, 61]) and of sGDML based on a two-body MBTR representation. Both models were trained on 1000 configurations, leading to kernel matrices with dimensionalities between 27k and 63k. For reference, we also present results for the Faber–Christensen–Huang–Lilienfeld (FCHL) representation [45, 47] and the polarizable atom interaction neural network (PaiNN) [60]. sGDML/MBTR performs as good or better than sGDML/CM^md for energy and force predictions. Compared to PaiNN, sGDML/MBTR performs better for energy predictions, but worse for forces. For a more fine-grained comparison between the sGDML models, figure 3 presents learning curves of relative MAE ratios of sGDML/MBTR over sGDML/CM^md, together with standard deviations over five runs starting from different random seeds. sGDML/MBTR consistently outperforms sGDML/CM^md with error reductions up to 50%–60%, especially when less than 100 training samples are used. For more symmetric molecules such as benzene, malonaldehyde, and ethanol, use of MBTR is less beneficial but still an improvement.

Zoom In Zoom Out Reset image size

3.4. Phase diagrams

We demonstrate applicability by identifying the convex hull of the phase diagram for Pt-group/transition metal binary alloys, relevant for industrial applications [66]. For a given dataset of candidate structures, we predict the energy of each structure and identify those with the lowest energy, which form the convex hull in a phase diagram. Compositions that lie on or slightly above the convex hull correspond to stable and meta-stable alloys, respectively.

To demonstrate this, we use a dataset [66] of 153 alloys computed at the DFT/PBE level of theory. This dataset contains at most a few hundred structures for each alloy. Due to this small amount of data direct application of ML models results in errors in predicted energies that are large enough to lead to wrong convex hulls. We address this by employing a simple active learning [67] scheme.

Starting with a few randomly selected structures, we iteratively train ML models on these and predict energies of candidate structures. In each iteration, we calculate (look up) DFT energies only for structures predicted to be low in energy and include these in the training dataset of the next iteration. This procedure prevents computationally expensive DFT calculations for high-energy structures that lie above the convex hull, saving up to 48% of all DFT calculations while still identifying the correct convex hull.

Figure 4 presents results for AgPt. The active learning model selected 357 DFT calculations for training and predicted energies of 331 (48%) other structures, with a MAE of 39 meV/atom. The trade-off between the number of saved calculations and the probability of failing to identify the correct convex hull can be explicitly controlled by adjusting the energy threshold below which DFT calculations are requested. In this simple demonstration, structures are given and not derived from composition by relaxation. While structural relaxation is possible with ML, it brings its own challenges [68–76].

Zoom In Zoom Out Reset image size

3.5. Other uses

MBTR is a general representation (numerical description, feature set) of atomistic systems for fast accurate interpolation between quantum-mechanical calculations via ML. It is based on distributions of k-atom terms stratified by chemical elements. Despite, or because of, this simple principle, it is connected to many other representations, including CM [6], BoB [23], histograms of distances, angles and dihedral angles [26], atom-centered symmetry functions [25], partial radial distribution functions [81], FCHL representation [45, 47], as well as cluster expansion [82]. See [16] for further details on these and other relationships.

MBTR represents whole molecules and crystals. With increasing number of atoms, and thus degrees of freedom, this approach is likely to degrade, and exploitation of locality via prediction of additive atomic energy contributions becomes appealing [9, 83]. This requires representing local chemical environments [20], for which MBTR can be modified [34, 84, 85].

We note in passing that problems in the training of ML models, such as outliers, can often be traced back to problems in the underlying reference calculations, such as unconverged fast Fourier transform grids or inconsistent settings (violating the assumption that a single function is being fitted), a phenomenon also observed by others [86]. This suggests that automated identification of errors in big datasets of electronic structure calculations via parametrization of ML models might be a general approach for validation of such datasets. We rationalize this hypothesis by ML models identifying regularity (correlations) in data, and faulty calculations deviating in some way from correct ones.

Continuing advances in electronic structure codes and increasing availability of large-scale computing resources have led to large collections of ab initio calculations, such as Materials Project [87], AFLOWlib [88], OQMD [59], and Novel Materials Discovery Laboratory [89]. Representations such as MBTR are key to combine quantum mechanics with machine learning (QM/ML) for fast, accurate and precise interpolation in these settings.

M R and H H thank Matthias Scheffler for helpful discussions and support. M R thanks the Institute for Pure and Applied Mathematics (IPAM) for hospitality and support, participants of its program on Understanding Many-Particle Systems with Machine Learning for feedback, and Albert Bartók-Pártay, Gábor Csányi, Alexander Shapeev, Alexandre Tkatchenko and the late Alessandro de Vita for extended discussions. M R acknowledges funding from EU Horizon 2020 program Grant 676580, The Novel Materials Discovery (NOMAD) Laboratory, a European Center of Excellence.

This article is an extended and updated version of the original arXiv preprint 1704.06439, whose formal publication was delayed for non-technical reasons.

All datasets used in this study are publicly available. Implementations of MBTR are available as part of the DScribe [85] and qmmlpack [90] libraries. Code to reproduce results of reported experiments is available at: https://github.com/hhaoyan/mbtr.

The data that support the findings of this study are openly available at the following URL/DOI: https://doi.org/10.6084/m9.figshare.19567324.v1.