Learning representations of molecules and materials with atomistic neural networks

TL;DR

This paper introduces neural network architectures, especially SchNet, that learn efficient, chemically intuitive representations of molecules and materials, accurately predicting properties across diverse datasets.

Contribution

The paper presents SchNet, a continuous-filter convolutional neural network that effectively models molecules and materials, demonstrating its ability to predict properties and analyze learned representations.

Findings

SchNet accurately predicts chemical properties across datasets

Learned representations align with chemical intuition

Neural networks can model complex molecular structures

Abstract

Deep Learning has been shown to learn efficient representations for structured data such as image, text or audio. In this chapter, we present neural network architectures that are able to learn efficient representations of molecules and materials. In particular, the continuous-filter convolutional network SchNet accurately predicts chemical properties across compositional and configurational space on a variety of datasets. Beyond that, we analyze the obtained representations to find evidence that their spatial and chemical properties agree with chemical intuition.

Tables

Table 1. Table 1 : Mean absolute (MAE) and root mean squared errors (RMSE) of SchNet with and without cosine cutoff for various datasets over three repetitions. For the Materials Project data, we use a smaller model ( 𝐱 i ∈ ℝ 64 subscript 𝐱 𝑖 superscript ℝ 64 \mathbf{x}_{i}\in\mathbb{R}^{64} ) and compare the cosine cutoff to a normalized filter with hard cutoff. We give the number of used reference calculations N, i.e. the size of the combined training and validation set.

Dataset	Property	Unit	$r_{\text{cut}}$ [Å]	MAE	RMSE
QM9 (N=110k)	U0	kcal mol-1	–	0.259	0.599
			$5$	0.218	0.518
	$\mu$	Debye	–	0.019	0.037
			$5$	0.017	0.033
Aspirin (N=1k)	total energy	kcal mol-1	–	0.438	0.592
			5	0.402	0.537
	atomic forces	kcal mol-1Å-1	–	1.359	1.929
			5	0.916	1.356
Aspirin (N=50k)	total energy	kcal mol-1	–	0.088	0.113
			5	0.102	0.130
	atomic forces	kcal mol-1Å-1	–	0.104	0.158
			5	0.140	0.203
Materials Project	formation energy	eV / atom	(hard) 5	0.037	0.091
(N=62k)			5	0.039	0.084

Equations

$$E$$

$$P$$

$$\mathbf{\mu }=\sum _i^N{q}_i{\mathbf{r}}_i.$$

$${\mathbf{x}}_i=f^{(1)}(Z_i)+\sum _{j\ne i}{f}^{(2)}((Z_i,{\mathbf{r}}_i),(Z_j,{\mathbf{r}}_j))+\sum _{\begin{array}{c}j,k\ne i\\ k\ne j\end{array}}{f}^{(3)}((Z_i,{\mathbf{r}}_i),(Z_j,{\mathbf{r}}_j),(Z_k,{\mathbf{r}}_k))+\dots$$

$${\mathbf{x}}_i^{(t+1)}$$

$${\mathbf{x}}_i^{(0)}={\mathbf{A}}_{Z_i}\in {\mathbb{R}}^d$$

$${\rho }^l(\mathbf{r})=\sum _{i=1}^{n_{\text{atoms}}}\delta (\mathbf{r}-{\mathbf{r}}_i){\mathbf{x}}_i^l.$$

$$(\rho \ast W)(\mathbf{r})=\underset{{\mathbf{r}}_a\in {\mathbb{R}}^3}{\int }\rho ({\mathbf{r}}_a)\circ W(\mathbf{r}-{\mathbf{r}}_a)d{\mathbf{r}}_a,$$

$$({\rho }^l\ast W)(\mathbf{r})$$

$$=\sum _{j=1}^{n_{\text{atoms}}}{\mathbf{x}}_j^l\circ W(\mathbf{r}-{\mathbf{r}}_j)$$

$${\mathbf{x}}_i^{l+1}=({\rho }^l\ast W^l)({\mathbf{r}}_i)=\sum _{j=1}^{n_{\text{atoms}}}{\mathbf{x}}_j^l\circ W({\mathbf{r}}_j-{\mathbf{r}}_i),$$

$${\mathbf{x}}_i^{(l+1)}=W^{(l)}{\mathbf{x}}_i^{(l)}+{\mathbf{b}}^{(l)}$$

$$f(x)=\ln \left(\frac{1}{2}{e}^x+\frac{1}{2}\right)$$

$$f(x)=\{\begin{array}{ll}{e}^x-1& \text{if }x<0\\ x& \text{otherwise}\end{array}$$

$${\mathbf{x}}_i^{l+1}=\sum _{j\ne i}{\mathbf{x}}_j^l\circ W({\mathbf{r}}_j-{\mathbf{r}}_i),$$

$${\hat{\mathbf{r}}}_{ij}={\left[\exp (-\gamma (r_{ij}-k\Delta \mu {)}^2)\right]}_{0\le k\le r_{\text{cut}}/\Delta \mu },$$

$$f_{\text{cut}}(r_{ij})$$

$$W_{\text{local}}(r_{ij})$$

$${\mathbf{x}}_i^{l+1}={\mathbf{x}}_{im}^{l+1}$$

$$=\sum _{j=0}^{n_{\text{atoms}}}{\mathbf{x}}_j^l\circ \underset{W}{\underbrace{\left(\sum _{b=0}^{n_{\text{cells}}}{\tilde{W}}^l({\mathbf{r}}_{jb}-{\mathbf{r}}_{ia})\right)}}.$$

$$W_{\text{normalized}}(r_{ij})=\frac{1}{n_{\text{nbh}}}W(r_{ij})$$

$${\mathbf{x}}_{\text{probe}}=({\rho }^l\ast W)({\mathbf{r}}_{\text{probe}})=\sum _{j=1}^{n_{\text{atoms}}}{\mathbf{x}}_j^l\circ W({\mathbf{r}}_{\text{probe}}-{\mathbf{r}}_j).$$

Full text

1 Introduction

In recent years, machine learning has been successfully applied to the prediction of chemical properties for molecules and materials schutt2014represent ; huo2017unified ; faber2018alchemical ; de2016comparing ; morawietz2016van ; gastegger2017machine ; faber2017prediction ; podryabinkin2017active ; brockherde2017bypassing ; bartok2017machine ; schutt2018schnet ; chmiela2018towards ; ziletti2018insightful ; dragoni2018achieving . A significant part of the research has been dedicated to engineering features that characterize global molecular similarity bartok2010gaussian ; rupp2012fast ; montavon2012learning ; hansen2013assessment ; Hansen-JCPL ; chmiela2017machine or local chemical environments behler2007generalized ; bartok2013representing ; gastegger2018wacsf based on atomic positions. Then, a non-linear regression method – such as kernel ridge regression or a neural network – is used to correlate these features with the chemical property of interest.

A common approach to model atomistic systems is to decompose them into local environments, where a chemical property is expressed by a partitioning into latent atom-wise contributions. Based on these contributions, the original property is then reconstructed via a physically motivated aggregation layer. E.g., Behler-Parinello networks behler2007generalized or the SOAP kernel bartok2013representing decompose the total energy in terms of atomistic contributions

[TABLE]

Atomic forces can be directly obtained as negative derivatives of the energy model. While this is often a suitable partitioning of extensive properties, intensive or per-atom properties can be modeled as the mean

[TABLE]

However, this might still not be a sufficient solution for global molecular properties such as HOMO-LUMO gaps or excitation energies pronobis2018capturing . To obtain a better performance, output models that incorporate property-specific prior knowledge should be used. E.g., the dipole moment can be written as

[TABLE]

such that the atomistic neural network needs to predict atomic charges $q_{i}$ gastegger2017machine ; sifain2018discovering ; yao2018tensormol ; schutt2018quantum .

The various atomistic models differ in how they obtain the energy contributions $E_{i}$, usually employing manually crafted atomistic descriptors. In contrast to such descriptor-based approaches, this chapter focuses on atomistic neural network architectures that learn efficient representations of molecules and materials end-to-end – i.e., directly from atom types $Z_{i}$ and positions $\mathbf{r}_{i}$ – while delivering accurate predictions of chemical properties across compositional and configurational space schutt2017quantum ; schutt2017schnet ; schutt2018schnet . The presented models will encode important invariances, e.g. towards rotation and translation, directly into the deep learning architecture and obtain predicted property from physically motivated output layers. Finally, we will obtain spatially and chemically resolved insights from the learned representations regarding the inner workings of the neural network as well as the underlying data.

2 The deep tensor neural network framework

In order to construct atom-centered representations $\mathbf{x}_{i}\in\mathbb{R}^{F}$, where $i$ is the index of the center atom and $F$ the number of feature dimension, a straight-forward approach is to expand the atomistic environment in terms of n-body interactions pukrittayakamee2009simultaneous ; malshe2009development , which can be written in general as

[TABLE]

However, such an approach requires to define explicit n-body models $f^{(n)}$ (e.g. using neural networks) as well as computing of a large number of higher-order terms of the atom coordinates. At the same time, all many-body networks must respect the invariances w.r.t. rotation, translation and the permutation of atoms.

An alternative approach is to incorporate higher-order interactions in a recursive fashion. Instead of explicitly modeling an $n$-body neural network, we design an interaction network $\mathbf{v}:\mathbb{R}^{F}\times\mathbb{R}\rightarrow\mathbb{R}^{F}$ that we use to model perturbations

[TABLE]

of the chemical environment $\mathbf{x}_{i}^{(t)}$ by its neighboring environments $\mathbf{x}_{j}^{(t)}$ depending on their relative position $\mathbf{r}_{ij}=\mathbf{r}_{j}-\mathbf{r}_{i}$. On this basis, we define the deep tensor neural network (DTNN) framework schutt2017quantum :

Use an embedding depending on the type of the center atom

[TABLE]

for the initial representation of local chemical environment $i$. This corresponds to the 1-body terms in Eq. 4. 2. 2.

Refine the embeddings repeatedly using the interaction networks from Eq. 5 3. 3.

Arriving at the final embedding $\mathbf{x}_{i}^{T}$ after $T$ interaction refinements, predict the desired chemical property using a property-specific output network (as described in Section 1).

The embedding matrix $A$ as well as the parameters of the interaction networks $\mathbf{v}^{(t)}$ and the output network are learned during the training procedure. This framework allows for a family of atomistic neural network models – such as the Deep Tensor Neural Network schutt2017quantum and SchNet schutt2017schnet ; schutt2018schnet – that differ in how the interactions $\mathbf{v}^{(t)}$ are modeled and the predictions are obtained from the atomistic representations $\mathbf{x}_{i}^{T}$.

3 SchNet

Building upon the principles of the previously described DTNN framework, we propose SchNet as a convolutional neural network architecture for learning representations for molecules and materials. Fig. 1 depicts an overview of the SchNet architecture as well as how the interaction refinements are modeled by interaction blocks shown on the right. In the following, we will introduce the main component of SchNet – the continuous-filter convolutional layer – before describing how these are used to construct the interaction blocks.

3.1 Continuous-filter convolutional layers

The commonly used convolutional layers lecun1989backpropagation employ discrete filter tensors since they are usually applied to data that is sampled on a grid, such as digital images, video or audio. However, such layers are not applicable for atomistic systems, since atoms can be located at arbitrary positions in space. E.g. when predicting the potential energy, the output of a convolutional layer might change rapidly when an atom moves from one grid cell to the next. Especially when we aim to predict a smooth potential energy surface, a continuous and differentiable representation is required. For this reason, we use a convolutional layer employing a continuous filter function in order to model the interactions.

Given the representations $\mathbf{x}^{l}_{i}$ of the chemical environment of atom $i$ at position $\mathbf{r}_{i}$ and layer $l$ of the neural network, the atomistic system can be described by a function

[TABLE]

In order to include the interactions of the atom-centered environments, we convolve $\rho:\mathbb{R}^{3}\rightarrow\mathbb{R}^{F}$ and a spatial filter $W:\mathbb{R}^{3}\rightarrow\mathbb{R}^{F}$ as element-wise

[TABLE]

where ”$\circ$” is the element-wise product. Here, the filter function $W$ describes the interaction of a representation $\mathbf{x}_{i}$ with an atom at the relative position $\mathbf{r}-\mathbf{r}_{i}$. The filter functions can be modeled by a filter-generating neural network similar to those used in dynamic filter networks BrabandereJTG16 . Considering the discrete location of atoms in Eq. 6, we obtain

[TABLE]

This yields a function representing how the atoms of the system act on another location in space. To obtain the rotationally-invariant interactions between atoms, we

[TABLE]

i.e., we evaluate the convolution at discrete locations in space using continuous, radial filters.

3.2 Interaction blocks

After introducing continuous-filter convolutional layers, we go on to construct the interaction blocks. Besides convolutions, we employ atom-wise, fully-connected layers

[TABLE]

that are applied separately to each atom $i$ with tied weights $W^{(l)}$.

Throughout the network, we use softplus non-linearities dugas2001incorporating that are shifted

[TABLE]

in order to conserve zero-activations: $f(0)=0$. Fig. 2 shows this activation function compared to exponential linear units (ELU) clevert2015fast :

[TABLE]

The derivatives for ELU and softplus are shown in the middle and right panel of Fig. 2, respectively. A crucial difference is that the softplus are smooth while ELUs exhibit only first-order continuity. However, the higher-order differentiability of the model, and therefore also of the employed activation functions, is crucial for the prediction of atomic forces or vibrational frequencies.

Fig. 1 (right) shows how the interaction block is assembled from these components. Since the continuous-filter convolutional layers are applied feature-wise, we achieve the mixing of feature maps by atom-wise layers before and after the convolution. This is analogous to depth-wise separable convolutional layers in Xception nets chollet2017xception which could outperform the architecturally similar InceptionV3 szegedy2016rethinking on the ImageNet dataset deng2009imagenet while having less parameters. Most importantly, feature-wise convolutional layers reduce the number of filters, which significantly reduces the computational cost. This is particularly important for continuous-filter convolutions, where each filter has to be computed by a filter-generating network.

3.3 Filter-generating networks

The architecture of the filter-generating network significantly influences the properties of the predicted filters and, consequently, the learned atomic interactions. Therefore, we can incorporate invariances or prior chemical knowledge into the filter. In the following, we describe the considerations that went into designing the SchNet filter-generating networks.

Self-interaction

In an interatomic potential, we aim to avoid self-interaction of atoms, since this is fundamentally different than the interaction with other atoms. We can encode this in the filter network by constraining the filter-network such that $W(\mathbf{r}_{i}-\mathbf{r}_{j})=0$ for $\mathbf{r}_{i}=\mathbf{r}_{j}$. Since two distinct atoms can not be at the same position, this is a unambiguous condition to exclude self-interaction. This is equivalent to modifying Eq. 9 to exclude the center atom of the environment from the sum:

[TABLE]

Rotational invariance

As the input to the filter $W:\mathbb{R}^{3}\rightarrow\mathbb{R}$ is only invariant to translations of the molecule, we additionally need to consider rotational invariance. We achieve this by using interatomic distances $r_{ij}$ as input to the filter network, resulting in radial filters $W:\mathbb{R}\rightarrow\mathbb{R}^{F}$.

Local distance regimes

In the spirit of radial basis function (RBF) networks broomhead1988multivariable ; moody1989fast , the filter-generating neural network $W(r_{ij})$ first expands the pair-wise distances

[TABLE]

with $\Delta\mu$ being the spacing of Gaussians with scale $\gamma$ on a grid ranging from 0 to the distance cutoff $r_{\text{cut}}$. This helps to decouple the various regimes of atomic interactions and allow for an easier starting point for the training procedure. On top of the RBF expansion, we apply two fully-connected layers with softplus activation functions.

As an illustrative example, Fig. 3 shows two linear models fitted to the potential energy surface of H2. Using the distance as feature directly, we obviously capture only a linear relationship. However, the expanded RBF feature space allows us to obtain a smooth and accurate fit of the potential.

From an alternative viewpoint, if we initialize a neural network with the usual weight distributions and non-linearities, the resulting function is almost linear before training as the neuron activations are close to zero. Therefore, the filter values would be strongly correlated, leading to a plateauing cost function at the beginning of training. Radial basis functions solve this issue by decoupling the various distance regimes.

Cutoffs

While in principle the size of the filter in a continuous-filter convolutional layer can be infinite, there are natural limitations on how such a filter can be trained. The interatomic distances in a dataset of molecules are upper-bound by the size of the largest molecule. More importantly, we can not consider interactions with an infinite amount of atoms in case of atomistic systems with periodic boundary conditions. Therefore, it is often beneficial or even required to restrict the filter size using a distance cutoff.

While it is certainly possible to apply a hard cutoff, this may lead to rapidly changing energies in molecular dynamics simulations. Therefore, we apply a cosine cutoff function to the filter, to obtain a local filter

[TABLE]

Periodic boundary conditions (PBC)

For materials, we have to respect the PBCs when convolving with the interactions, i.e. we have to include interactions with periodic sites of neighboring unit cells: Due to the linearity of the convolution, we can move the sum over periodic images into the filter. Given atomistic representations $\mathbf{x}_{i}=\mathbf{x}_{ia}=\mathbf{x}_{ib}$ of site $i$ for unit cells $a$ and $b$, we obtain

[TABLE]

When using a hard cutoff, we have found that the filter needs to be normalized with respect to the number of neighboring atoms $n_{\text{nbh}}$ for the training to converge:

[TABLE]

However, this is not necessary, when using a cosine cutoff function, as shown above.

4 Analysis of the representation

Having introduced the SchNet architecture, we go on to analyze the representations that have been learned by training the models on QM9 – a dataset of  130k small organic molecules with up to nine heavy atoms ramakrishnan2014quantum – as well as a molecular dynamics trajectory of aspirin chmiela2017machine . If not given otherwise, we use six interaction blocks and atomistic representations $\mathbf{x}_{i}\in\mathbb{R}^{256}$. The models have been trained using stochastic gradient descent with warm restarts loshchilov2016sgdr and the ADAM optimizer kingma2014adam .

4.1 Locality of the representation

As described above, atomistic models decompose the representation into local chemical environments. Since SchNet is able to learn a representation of such an environment, the locality of the representation may depend on whether a cutoff was used as well as the training data.

Table 1 shows the performance of SchNet models trained on various datasets with and without cutoff. We observe that the cutoff is beneficial for QM9 as well as the small aspirin training set with N=1,000 reference calculations. The cutoff function biases the model towards learning from local interactions, which helps with generalization since energy contributions from interactions at larger distances are much harder to disentangle. On the other hand, the SchNet model with trained on 50,000 aspirin reference calculation benefits from the large chemical environment when not applying a cutoff. This is because with such a large amount of training data, the model is now also able to infer less local interaction within the molecule. In case of the materials data, we observe that cosine cutoff and hard cutoff yield comparable results, where the cosine cutoff is slightly preferable since it obtains the lower root mean squared error (RMSE). Since the RMSE puts more emphasis on larger errors than MAE, this indicates that the cosine cutoff improves generalization. This may be due to the more local model which is obtained by focusing on smaller distances or by eliminating the discontinuities that are introduced by a hard cutoff.

Fig. 4 shows the atomization energies of a carbon dimer as predicted by SchNet models trained on QM9 and the aspirin trajectory of the MD17 dataset. Since the models were trained on saturated molecules, this does not reflect the real energy or atomic forces of the dimer. The reason is that the energy contribution of carbon interactions in the context of equilibrium molecules or MD trajectories, respectively, include the inferred contributions of other neighboring atoms. For instance, if we consider two carbon atoms at a distance of about 2.4Å in aspirin, they are likely to be part of the aromatic ring with other carbon at a distance of 1.4Å. In case of aspirin, we also observe a large offset since the model was not trained on molecules with a varying number of atoms. If we wanted to eliminate these model biases, we needed to train the neural networks on more diverse datasets, e.g. by explicitly including dimers with large interatomic distances. While this is necessary to obtain a general model of quantum chemistry, it might even be detrimental for the prediction of a certain subset of molecules using a give amount of training data.

Considering the above, the analysis in Fig. 4 shows how the neural network attributes interaction energies in the context of the data it was trained on. We observe that the general shape of the potential is consistent across all four models. Applying the cosine cutoff leads to constant energy contributions beyond $r_{\text{cut}}=5$Å, however, models without cutoff are nearly constant in this distance regime as well. SchNet correctly characterizes the carbon bond with its energy minimum between 1.2-1.3Å and rising energy with larger distances. For the distance regime beyond about 1.8Å, the inferred, larger environment dominates the attribution of interaction energies.

Given that the aspirin model trained on the larger dataset benefits from a larger attribution of interaction energies to larger distances, we analyze how the cutoff will affect the vibrational spectrum. Using the SchNet potentials, we have generated two molecular dynamics trajectories at 300K with a time step of 0.5fs. Fig. 5 shows the vibrational spectra of the models with and without cosine cutoff. The most distinct change is the shift of the peak at about 3600 cm*-1* to the right. This corresponds to the O-H oscillations, where the cutoff may prevent direct interactions of the hydroxyl group with the carbon ring (see Fig. 5).

4.2 Local chemical potentials

In order to further examine the spatial structure of the representation, we observe how SchNet models the influence of a molecule on a probe atom that is moved through space and acts as a test charge. This can be derived straight-forward from the definition of the continuous-filter convolutional layer in Eq. 8, which is defined for arbitrary positions in space:

[TABLE]

The remaining part of the model is left unchanged as those layers are only applied atom-wise. Finally, we visualize the predicted probe energy on a smooth isosurface around the molecule schutt2017quantum ; schutt2018quantum ; schutt2018schnet .

Fig. 6 (left) shows this for N-Formylformamide using a hydrogen probe. According to this, the probe is more likely to bond on the oxygens, as indicated by the lower probe energies. To further study this interpretation, we remove one of the hydrogens in Fig. 6 (right). In agreement with our analysis, this leads to even lower energies at the position of the missing hydrogen as well as the nearby oxygen due to the nearby unsaturated carbon.

4.3 Atom embeddings

Having examined the spatial structure of SchNet representations, we go on to analyze what the model has learned about chemical elements included in the data. As described above, SchNet encodes atom types using embeddings $\mathbf{A}_{Z}\in\mathbb{R}^{F}$ that are learned during the training process. We visualize the two leading principal components of these embeddings to examine whether they resemble chemical intuition. Since QM9 only contains five atom types (H, C, N, O, F), we perform this analysis on the more diverse Materials Project dataset as it includes 89 atom types ranging across the periodic table. Fig. 7 shows the reduced embeddings of the main group elements of the periodic table. Atoms belonging to the same group tend to form clusters.

This is especially apparent for main groups 1-7, while group 8 appears to be more scattered. Beyond that, there are partial orderings of elements according to their period within some of the groups. We observe a partial order from light to heavier elements in some groups, e.g. in group 1 (left to right: H - Li - Na - [K, Rb, Cs]), group 2 (left to right: Be - Mg - Ca - Sr - Ba) and group 5 (left to right: P-As-Sb-Bi). In some cases, the first element of the group lies further apart from the rest of the group, e.g. H, Be, C and O. These results are consistent with those we obtained from previous SchNet models trained on earlier versions of the Materials Project repository schutt2018schnet ; schutt2018quantum .

Note that these extracted chemical insights are not imposed by the SchNet architecture, but had to be inferred by the model based on the bulk systems and its energies in the training data.

5 Conclusions

We have presented the deep tensor neural network framework and its implementation SchNet, which obtains accurate predictions of chemical properties for molecules and materials. Representations of chemical environments are learned directly from atom types and position while filter-generating networks allow to incorporate invariance and prior knowledge.

In our analysis, we have found that atomic representations reflect an inferred chemical environment based on the bias of the training data. The obtained representations are dominated by local interactions and can be further localized using cosine cutoff functions that improve generalization. However, if a sufficient amount of data is available, interaction energies can be attributed reliably for larger distances which will further improve the accuracy of the model. Moreover, we have defined local chemical potentials that allow for spatially resolved chemical insights and have shown that the models learn embeddings of chemical elements that show resemblence of the structure of the periodic table.

In conclusion, SchNet presents an end-to-end atomistic neural network that we expect to facilitate further developments towards interpretable deep learning architectures to assist chemistry research.

Acknowledgements.

The authors thank Michael Gastegger for valuable discussions and feedback. This work was supported by the Federal Ministry of Education and Research (BMBF) for the Berlin Big Data Center BBDC (01IS14013A). Additional support was provided by the Institute for Information & Communications Technology Promotion and funded by the Korean government (MSIT) (No. 2017-0-00451, No. 2017-0-01779). A.T. acknowledges support from the European Research Council (ERC-CoG grant BeStMo).