Cormorant: Covariant Molecular Neural Networks

Introduction

In principle, quantum mechanics provides a perfect description of the forces governing the behavior of atoms, molecules and crystalline materials such as metals. However, for systems larger than a few dozen atoms, solving the Schrödinger equation explicitly at every timestep is not a feasible proposition on present day computers. Even Density Functional Theory (DFT) (Hohenberg and Kohn 1964), a widely used approximation to the equations of quantum mechanics, has trouble scaling to more than a few hundred atoms.

Consequently, the majority of practical work in molecular dynamics today falls back on fundamentally classical models, where the atoms are essentially treated as solid balls and the forces between them are given by pre-defined formulae called atomic force fields or empirical potentials, such as the CHARMM family of models (Brooks et al. 1983; Brooks et al. 2009). There has been a widespread realization that this approach has inherent limitations, so in recent years a burgeoning community has formed around trying to use machine learning to learn more descriptive force fields directly from DFT computations (Behler and Parrinello 2007; Bartók et al. 2010; Rupp et al. 2012; Shapeev 2015; Chmiela et al. 2016; Zhang et al. 2018; Schütt et al. 2017; Hirn et al. 2017). More broadly, there is considerable interest in using ML methods not just for learning force fields, but also for predicting many other physical/chemical properties of atomic systems across different branches of materials science, chemistry and pharmacology (Montavon et al. 2013; Gilmer et al. 2017b; Smith et al. 2017; Yao et al. 2018).

At the same time, there have been significant advances in our understanding of the equivariance and covariance properties of neural networks, starting with (Cohen and Welling 2016a; Cohen and Welling 2016b) in the context of traditional convolutional neural nets (CNNs). Similar ideas underly generalizations of CNNs to manifolds (Masci et al. 2015; Monti et al. 2016; Bronstein et al. 2017) and graphs (Bruna et al. 2014; Henaff et al. 2015). In the context of CNNs on the sphere, Cohen et al. 2018 realized the advantage of using “Fourier space” activations, i.e., expressing the activations of neurons in a basis defined by the irreducible representations of the underlying symmetry group (see also (Esteves et al. 2017)), and these ideas were later generalized to the entire $\textrm{SE}(3)$ group (Weiler et al. 2018). Kondor and Trivedi 2018 gave a complete characterization of what operations are allowable in Fourier space neural networks to preserve covariance, and Cohen et al generalized the framework even further to arbitrary gauge fields (Cohen et al. 2019). There have also been some recent works where even the nonlinear part of the neural network’s operation is performed in Fourier space: independently of each other (Thomas et al. 2018) and (Kondor 2018) were to first to use the Clebsch–Gordan transform inside rotationally covariant neural networks for learning physical systems, while (Kondor et al. 2018) showed that in spherical CNNs the Clebsch–Gordan transform is sufficient to serve as the sole source of nonlinearity.

The Cormorant neural network architecture proposed in the present paper combines some of the insights gained from the various force field and potential learning efforts with the emerging theory of Fourier space covariant/equivariant neural networks. The important point that we stress in the following pages is that by setting up the network in such a way that each neuron corresponds to an actual set of physical atoms, and that each activation is covariant to symmetries (rotation and translation), we get a network in which the “laws” that individual neurons learn resemble known physical interactions. Our experiments show that this generality pays off in terms of performance on standard benchmark datasets.

The nature of physical interactions in molecules

Ultimately interactions in molecular systems arise from the quantum structure of electron clouds around constituent atoms. However, from a chemical point of view, effective atom-atom interactions break down into a few simple classes based upon symmetry. Here we review a few of these classes in the context of the multipole expansion, whose structure will inform the design of our neural network.

The simplest type of physical interaction is that between two particles that are pointlike and have no internal directional degrees of freedom, such as spin or dipole moments. A classical example is the electrostatic attraction/repulsion between two charges described by the Coulomb energy

Here $q_{A}$ and $q_{B}$ are the charges of the two particles, \mbox{\boldmathr}_{\!A} and \mbox{\boldmathr}_{\hskip-0.81949ptB} are their position vectors, \mbox{\boldmathr}_{\!AB}=\mbox{\boldmathr}_{\!A}\hskip-1.00006pt-\hskip-1.00006pt\mbox{\boldmathr}_{\hskip-0.81949ptB}, and $\epsilon_{0}$ is a universal constant. Note that this equation already reflects symmetries: the fact that (1) only depends on the length of \mbox{\boldmathr}_{\!AB} and not its direction or the position vectors individually guarantees that the potential is invariant under both translations and rotations.

One step up from the scalar case is the interaction between two dipoles. In general, the electrostatic dipole moment of a set of $N$ charged particles relative to their center of mass $r$ is just the first moment of their position vectors weighted by their charges:

The dipole/dipole contribution to the electrostatic potential energy between two sets of particles $A$ and $B$ separated by a vector \mbox{\boldmathr}_{\!AB} is then given by

One reason why dipole/dipole interactions are indispensible for capturing the energetics of molecules is that most chemical bonds are polarized. However, dipole/dipole interactions also occur in other contexts, such as the interaction between the magnetic spins of electrons.

One more step up the multipole hierarchy is the interaction between quadropole moments. In the electrostatic case, the quadropole moment is the second moment of the charge density (corrected to remove the trace), described by the matrix

Quadropole/quadropole interactions appear for example when describing the interaction between benzene rings, but the general formula for the corresponding potential is quite complicated. As a simplification, let us only consider the special case when in some coordinate system aligned with the structure of $A$, and at polar angle $(\theta_{A},\phi_{A})$ relative to the vector \mbox{\boldmathr}_{\!AB} connecting $A$ and $B$, \mbox{\boldmath\Theta}_{A} can be transformed into a form that is diagonal, with $[\Theta_{\hskip-0.81949ptA}]_{zz}\hskip-1.00006pt=\hskip-1.00006pt\vartheta_{A}$ and $\smash{[\Theta_{\hskip-0.81949ptA}]_{xx}\hskip-1.00006pt=\hskip-1.00006pt[\Theta_{\hskip-0.81949ptA}]_{yy}\hskip-1.00006pt=\hskip-1.00006pt-\vartheta_{A}/2}$ (Stone 1997). We make a similar assumption about the quadropole moment of $B$. In this case the interaction energy becomes

Higher order interactions involve moment tensors of order 3,4,5, and so on. One can appreciate that the corresponding formulae, especially when considering not just electrostatics but other types of interactions as well (dispersion, exchange interaction, etc), quickly become very involved.

Spherical tensors and representation theory

Fortunately, there is an alternative formalism for expressing molecular interactions, that of spherical tensors, which makes the general form of physically allowable interactions more transparent. This formalism also forms the basis of the our Cormorant networks described in the next section.

The key to spherical tensors is understanding how physical quantities transform under rotations. Specifically, in our case, under a rotation $R$:

The above imply that there is a fixed unitary transformation matrix $C^{(k)}$ which reduces the $k$’th order rotation operator into a direct sum of irreducible representations:

is a scalar, and hence it is a candidate for being a term in the potential energy. Note the similarity of this expression to the bispectrum (Kakarala 1992; Bendory et al. 2018), which is an already established tool in the force field learning literature (Bartók et al. 2013).

Almost any rotation invariant interaction potential can be expressed in terms of iterated Clebsch–Gordan products between spherical tensors. In particular, the full electrostatic energy between two sets of charges $A$ and $B$ separated by a vector \mbox{\boldmathr}=(r,\theta,\phi) expressed in multipole form (Jackson 1999) is

We emphasize that our discussion of electrostatics is only intended to illustrate the algebraic structure of interatomic interactions of any type, and is not restricted to electrostatics. In what follows, we will not explicitly specify what interactions the network will learn. Nevertheless, there are physical constraints on the interactions arising from symmetries, which we explicitly impose in our design of Cormorant.

CORMORANT: COvaRiant MOleculaR Artificial Neural neTworks

The goal of using ML in molecular problems is not to encode known physical laws, but to provide a platform for learning interactions from data that cannot easily be captured in a simple formula. Nonetheless, the mathematical structure of known physical laws, like those discussed in the previous sections, give strong hints about how to represent physical interactions in algorithms. In particular, when using machine learning to learn molecular potentials or similar rotation and translation invariant physical quantities, it is essential to make sure that the algorithm respects these invariances.

The second important feature of our architecture is that each neuron corresponds to either a single atom or a set of atoms forming a physically meaningful subset of the system at hand, for example all atoms in a ball of a given radius. This condition helps encourage the network to learn physically meaningful and interpretable interactions. The high level definition of Cormorant nets is as follows.

Let $\mathcal{S}$ be a molecule or other physical system consisting of $N$ atoms. A “Cormorant” covariant molecular neural network for $\mathcal{S}$ is a feed forward neural network consisting of $m$ neurons $\mathfrak{n}_{1},\ldots,\mathfrak{n}_{m}$, such that

Every neuron $\mathfrak{n}_{i}$ corresponds to some subset $\mathcal{S}_{i}$ of the atoms. In particular, each input neuron corresponds to a single atom. Each output neuron corresponds to the entire system $\mathcal{S}$.

The type of each output neuron is \smash{\mbox{\boldmath\tau}_{\!\textrm{out}}\hskip-1.00006pt=\hskip-1.00006pt(1)}, i.e., a scalar. Cormorant can learn data of arbitrary SO(3)-vector outputs. We restrict to scalars here to simplify the exposition.

Condition (C3) guarantees that whatever function a Cormorant network learns will be invariant to global rotations. Translation invariance is easier to enforce simply by making sure that the interactions represented by individual neurons only involve relative distances.

Here and in the following $\oplus$ denotes the appropriate concatenation of vectors and matrices. In Cormorant, however, as a slight departure from (7), to reduce the quadratic blow-up in the number of columns, we always have $n_{1}\hskip-1.00006pt=\hskip-1.00006ptn_{2}$ and use the restricted “channel-wise” CG-product,

2 One-body and two-body interactions

As stated in Definition 2, the covariant neurons in a Cormorant net correspond to different subsets of the atoms making up the physical system to be modeled. For simplicty in our present architecture there are only two types of neurons: those that correspond to individual atoms and those that correspond to pairs. For a molecule consisting of $N$ atoms, each layer $s=0,1,\ldots,S$ of the covariant part of the network has $N$ neurons corresponding to the atoms and $N^{2}$ neurons corresponding to the $(i,j)$ atom pairs. By loose analogy with graph neural networks, we call the corresponding $F_{i}^{s}$ and $g^{s}_{i,j}$ activations vertex and edge activations, respectively.

The actual form of the vertex activations captures “one-body interactions” propagating information from the previous layer related to the same atom and (indirectly, via the edge activations) “two-body interactions” capturing interactions between pairs of atoms:

3 Overall structure and comparison with other architectures

In addition to the covariant neurons described above, our network also needs neurons to compute the input featurization and the the final output after the covariant layers. Thus, in total, a Cormorant networks consists of three distinct parts:

We leave the details of the input and output featurization to the Supplement.

A key difference between Cormorant and other recent covariant networks (Tensor Field Networks (Thomas et al. 2018) and $\textrm{SE}(3)$-equivariant networks (Weiler et al. 2018)) is the use of Clebsch-Gordan non-linearities. The Clebsch-Gordan non-linearity results in a complete interaction of every degree of freedom in an activation. This comes at the cost of increased difficulty in training, as discussed in the Supplement. We further note that $\textrm{SE}(3)$-equivariant networks use a three-dimensional grid of points to represent data, and ensure both translational and rotational covariance (equivariance) of each layer. Cormorant on the other hand uses activations that are covariant to rotations, and strictly invariant to translations.

Experiments

We present experimental results on two datasets of interest to the computational chemistry community: MD-17 for learning molecular force fields and potential energy surfaces, and QM-9 for learning the ground state properties of a set of molecules. The supplement provides a detailed summary of all hyperparameters, our training algorithm, and the details of the input/output levels used in both cases. Our code is available at https://github.com/risilab/cormorant.

MD-17 (Chmiela et al. 2016) is a dataset of eight small organic molecules (see Table 1(b)) containing up to 17 total atoms composed of the atoms H, C, N, O, F. For each molecule, an ab initio molecular dynamics simulation was run using DFT to calculate the ground state energy and forces. At intermittent timesteps, the energy, forces, and configuration (positions of each atom) were recorded. For each molecule we use a train/validation/test split of 50k/10k/10k atoms respectively. The results of these experiments are presented in Table 1(b), where the mean-average error (MAE) is plotted on the test set for each of molecules. (All units are in kcal/mol, as consistent with the dataset and the literature.) To the best of our knowledge, the current state-of-the art algorithms on this dataset are DeepMD (Zhang et al. 2018), DTNN (Schütt et al. 2017), SchNet (Schütt et al. 2017), GDML (Chmiela et al. 2016), and sGDML (Chmiela et al. 2018). Since training and testing set sizes were not consistent, we used a training set of 50k molecules to compare with all neural network based approaches. As can be seen from the table, our Cormorant network outperforms all competitors.

Conclusions

This project was supported by DARPA “Physics of AI” grant number HR0011837139, and used computational resources acquired through NSF MRI 1828629.

We thank E. Thiede for helpful discussion and comments on the manuscript.

References

Architecture

As discussed in the main text, our Cormorant architecture is constructed from three basic building blocks: (1) an input featurization that takes $(Z_{i},\mathbf{r}_{i})$ and outputs a scalar, (2) a set of covariant CG layers that update $F_{i}^{s}$ to $F_{i}^{s+1}$, (3) a layer that takes the set of covariant activations $F_{i}^{s}$, and construct a permutation and rotation invariant regression target.

See Table for a more complete table of symbols used in the supplement and main text.

2 Overall structure

networks are constructed from three basic units:

This design is organized in a modular way to separate the input featurization, the covariant SO(3)-vector layers, and the output regression tasks. Importantly, the ${\rm INPUT}$ and ${\rm OUTPUT}$ networks are different for GDB9 and MD17. However, the covariant SO(3)-vector layers ${\rm CGNet}$ were identical in design and hyperparameter choice. We include these designs and choices below.

3 Input featurization

We found for MD-17, a complex input featurization network was not significantly beneficial, and that this input parametrization was sufficiently expressive.

3.2 QM-9

4 Covariant S​O​(3)SO(3)-vector layers

For both datasets, the central covariant $SO(3)$-vector layers of our Cormorant are identical. In both cases, we used $S=4$ layers with $L=3$, followed by a single $SO(3)$-vector layer with $L=0$. The number of channels of the input tensors at each level is fixed to $n_{c}=16$, and similarly the set of weights $W$ reduce the number of channels of each irreducible representation back to $n_{c}=16$.

The algorithm can be implemented as iterating over the function

The function $\left(g_{ij}^{s+1},F_{i}^{s+1}\right)\leftarrow{\rm CGLayer}\left(g_{ij}^{s},F_{i}^{s},\mathbf{r}_{i}\right)$ is itself constructed in the following way:

$g_{ij}^{s+1}\leftarrow{\rm EdgeNetwork}\left(g_{ij}^{s},\mathbf{r}_{ij},F_{i}^{s}\right)$

$F_{i}^{s+1}\leftarrow{\rm VertexNetwork}\left(F^{s+1}_{ij},F_{i}^{s}\right)$

${\rm VertexNetwork}\left(G^{s+1}_{ij},F_{i}^{s}\right):(V^{s})^{N\times N}\times(V^{s})^{N}\rightarrow(V^{s+1})^{N}$ updates the vertex SO(3)-vector activations by combining a “Clebsch-Gordan aggregation”, a CG non-linearity, a skip connection, and a linear mixing layer.

4.2 Edge networks

Our edge network is an extension of the “edge networks” in Message Passing Neural Networks Gilmer et al. 2017a. The $\rm EdgeNetwork$ function takes three different types of pair features, concatenates them, and then mixes them. We express write the edge network (Eq. (9)) in the main text) with all indices explicitly included:

where $\sigma\left(x\right)$ is the sigmoid activation, $r_{c,{\rm soft}}^{s}$ is a soft cutoff that drops off with width $w_{c}^{s}$.

4.3 From edge scalar representations to S​O​(3)SO(3)-vector

4.4 Vertex networks

The function ${\rm VertexNetwork}$ is found by concatenating three operations:

$F_{i}^{s+1,{\rm ag}}=\sum_{j\in N\left(i\right)}G^{s+1}_{ij}\otimes_{\rm cg}F_{j}^{s}$ is a CG-aggregation step and $G^{s}_{ij}$ is the set of edge representations calculated by ${\rm Edge2Vertex}$.

$F_{i}^{s+1,{\rm nl}}=F_{i}^{s}\otimes_{\rm cg}F_{i}^{s}$ is a CG non-linearity.

$F_{i}^{s+1,{\rm id}}=F_{i}^{s}$ is just the identity function, or equivalently a skip connection.

5 Output featurization

The output featurization of the network starts with the construction of a set of scalar invariants from the set of activations $F_{i}^{s}$ for all atoms $i$ and all levels $s=0\ldots S$. We extract three scalar invariants from each activation $F$ (dropping the $i$ and $s$ indices):

The output for the MD-17 network is straightforward. The scalars $x_{i}$ are summed over, and then a single linear layer is applied: $y=A\left(\sum_{i}x_{i}\right)+b$.

5.2 QM-9

6 Weight initialization

We chose the gain to ensure that the activations at each level were order unity when the network is initialized. We found that if the gain was too low, the CG products in higher levels would not significantly contribute to training, and information would only flow through linear (one-body) operations. This would result in convergence to poor training error. On the other hand, if the gain is set too high, the CG non-linearities dominate at initialization and would increase the change of the instabilities discussed above.

Experimental details

We trained our network using the AMSGrad [j.2018on] optimizer with a constant learning rate of $5\times 10^{-4}$ and a mini-batch size of $25$. We trained for 512 and 256 epoch respectively for MD-17 and QM-9. For each molecule in MD-17, we uniformly sampled 50k/10k/10k data points in the training/validation/test splits respectively. In QM-9 the dataset was randomly split to 100k molecules in the train set, with $10\%$ in the test set, and the remaining in the validation set. We removed the 3054 molecules that failed consistency requirements [Ramakrishnan et al. 2014], and also subtracted the thermochemical energy [Gilmer et al. 2017a] for the targets $C_{v}$, $U_{0}$, $U$, $G$, $H$, ZPVE.

For both datasets, we used, $S=4$ CGLayers with $L=3$ and we used $N_{c}=16$ channels at the output of each CGLayer. This gave networks with 299808 and 154241 parameter respectively for QM-9 and MD-17. Training time for QM-9 is takes roughly 48 hours on a NVidia 2080 Ti GPU. Training time for MD-17 varies based upon the molecule being trained, but typically ranges between 26 and 30 hours.

Training our Cormorant had several subtleties that we both believe are related to the nature of the CG non-linearity. We found a poor choice of weight initialization or optimization algorithm will frequently result in either: (1) an instability resulting in very large training loss ($>10^{6}$), from which the network will never recover, or (2) convergence to weights where the activation of CG non-linearities in higher layers turn off, and the resulting training error is poor.

We believe these difficulties are a result of the CG non-linearity, which is quadratic and unbounded. In fact, our network is just a high-order polynomial function of learnable parameters. This is true for MD-17, although for QM-9, the presence of non-linearities in the fully-connected MLPs adds a more conventional non-linearity. For the hyperparameters used in our experiments, the prediction at the top is a sixteenth order polynomial of our network’s parameters. As a result, in certain regions of parameter space small gradient updates can result in rapid growth of the output amplitude or a rapid drop in the importance of some channels.

These issues were more significant when we used Adam [Kingma2015AdamAM] then AMSGrad, and when the network’s parameters were not initialized in a narrow range. Using the weight initialization scheme discussed in Sec. 7.6, we were able to consistently converge to low training and validation error, provided we were limited to at most four CG layers.