Atomic-scale representation and statistical learning of tensorial properties

TL;DR

This paper presents a framework for incorporating three-dimensional symmetries into Gaussian process regression models to accurately predict tensorial properties of atomic structures, enhancing the modeling of molecular polarizability and electron density.

Contribution

It introduces a symmetry-adapted Gaussian process regression approach with specialized representations and kernels for tensorial properties, extending existing methods to account for geometric covariance.

Findings

Effective learning of molecular polarizability.

Accurate prediction of ground-state electron density.

Framework generalizes smooth overlap representations.

Abstract

This chapter discusses the importance of incorporating three-dimensional symmetries in the context of statistical learning models geared towards the interpolation of the tensorial properties of atomic-scale structures. We focus on Gaussian process regression, and in particular on the construction of structural representations, and the associated kernel functions, that are endowed with the geometric covariance properties compatible with those of the learning targets. We summarize the general formulation of such a symmetry-adapted Gaussian process regression model, and how it can be implemented based on a scheme that generalizes the popular smooth overlap of atomic positions representation. We give examples of the performance of this framework when learning the polarizability and the ground-state electron density of a molecule.

Equations

$$y(\mathcal{X})=⟨w||\mathcal{X}⟩.$$

$$y(\mathcal{X})\sim \mathcal{N}(⟨w||\mathcal{X}⟩,{\sigma }^2).$$

$$L(w)=\sum _{n=1}^N{\sigma }^{-2}(y_n-⟨w||{\mathcal{X}}_n⟩{)}^2+{\alpha }^2⟨w||w⟩,$$

$$|w⟩={\hat{C}}^{-1}\sum _{n=1}^N|{\mathcal{X}}_n⟩y_n,$$

$$\hat{C}=\sum _{n=1}^N|{\mathcal{X}}_n⟩⟨{\mathcal{X}}_n|+{\eta }^2\hat{I},$$

$$k(\mathcal{X},{\mathcal{X}}^{\prime })=⟨\mathcal{X}||{\mathcal{X}}^{\prime }⟩.$$

$$\text{Cov}[y(\mathcal{X}),y({\mathcal{X}}^{\prime })]=k(\mathcal{X},{\mathcal{X}}^{\prime })+{\sigma }^2{\delta }_{\mathcal{X}{\mathcal{X}}^{\prime }}.$$

$$y(\mathcal{X})=\mathbf{k}(\mathcal{X}{)}^T{\left[\mathbf{K}+{\sigma }^2\mathbf{I}\right]}^{-1}\mathbf{y}=\mathbf{k}(\mathcal{X}{)}^T\mathbf{x}(\{\mathbf{y}\}),$$

$$y_{\alpha \beta ...}(\hat{t}\mathcal{X})=y_{\alpha \beta ...}(\mathcal{X}).$$

$$y_{\alpha \beta ...}(\hat{R}\mathcal{X})=R_{\alpha {\alpha }^{\prime }}{R}_{\beta {\beta }^{\prime }}\times ...\times y_{{\alpha }^{\prime }{\beta }^{\prime }...}(\mathcal{X}).$$

$$y_{\alpha \beta ...}(\hat{Q}\mathcal{X})=Q_{\alpha {\alpha }^{\prime }}{Q}_{\beta {\beta }^{\prime }}\times ...\times y_{{\alpha }^{\prime }{\beta }^{\prime }...}(\mathcal{X}).$$

$$|\mathcal{X}\alpha \beta ...⟩=|\mathcal{X}⟩\otimes |\alpha ⟩\otimes |\beta ⟩\otimes ...,$$

$$y_{\alpha \beta ...}(\mathcal{X})=⟨w||\mathcal{X}\alpha \beta ...⟩,$$

$$y_{\alpha \beta ...}(\mathcal{X})=⟨w_{\alpha \beta ...}||\mathcal{X}\alpha \beta ...⟩.$$

$$\left[\hat{I}\otimes \hat{R}\otimes \hat{R}\otimes ...\right]{|\left(\hat{R}\mathcal{X}\right)\alpha \beta ...⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}={|\mathcal{X}\alpha \beta ...⟩}_{\mathcal{S}\mathcal{O}\left(3\right)},$$

$${|\mathcal{X}\alpha \beta ...⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}\equiv \left[{\hat{R}}_{\mathcal{X}\to }\otimes {\hat{R}}_{\mathcal{X}\to }\otimes ...\right]|\left({\hat{R}}_{\mathcal{X}\to }\mathcal{X}\right)\alpha \beta ...⟩$$

$${|\mathcal{X}\alpha \beta ...⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}\equiv \int \text{d}\hat{R}\left[\hat{I}\otimes \hat{R}\otimes ...\right]|\left(\hat{R}\mathcal{X}\right)\alpha \beta ...⟩.$$

$${|\mathcal{X}\alpha \beta ...⟩}_{\mathcal{O}(3)}={|\mathcal{X}\alpha \beta ...⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}+\left[\hat{I}\otimes \hat{Q}\otimes ...\right]{|\left(\hat{Q}\mathcal{X}\right)\alpha \beta ...⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}$$

$$|w⟩=\sum _{n=1}^N\sum _{\alpha \beta ...}{\hat{C}}^{-1}{|{\mathcal{X}}_n\alpha \beta ...⟩}_{\mathcal{O}(3)}{y}_{\alpha \beta ...}(n),$$

$$\hat{C}=\sum _{n=1}^N\sum _{\alpha \beta ...}{|{\mathcal{X}}_n\alpha \beta ...⟩}_{\mathcal{O}(3)}{⟨{\mathcal{X}}_n\alpha \beta ...|}_{\mathcal{O}(3)}+\eta \hat{I}.$$

$$k_{\alpha \beta ...}^{{\alpha }^{\prime }{\beta }^{\prime }...}(\mathcal{X},{\mathcal{X}}^{\prime })=⟨{\mathcal{X}}^{\prime }{\alpha }^{\prime }{\beta }^{\prime }...|{|\mathcal{X}\alpha \beta ...⟩}_{\mathcal{O}(3)}.$$

$$\int \text{d}\hat{R}\int \text{d}{R}^{\prime }⟨\hat{R}\mathcal{X}||{\hat{R}}^{\prime }{\mathcal{X}}^{\prime }⟩(R{R}^{\prime }{)}_{\alpha {\alpha }^{\prime }}(R{R}^{\prime }{)}_{\beta {\beta }^{\prime }}...$$

$$|\hat{R}\mathcal{X}⟩=\hat{U}[\hat{R}]|\mathcal{X}⟩,$$

$$k_{\alpha \beta ...}^{{\alpha }^{\prime }{\beta }^{\prime }...}(\mathcal{X},{\mathcal{X}}^{\prime })=\int \text{d}\hat{R}k(\mathcal{X},\hat{R}{\mathcal{X}}^{\prime })R_{\alpha {\alpha }^{\prime }}{R}_{\beta {\beta }^{\prime }}...,$$

$${\mathbf{k}}_{\mathcal{S}\mathcal{O}\left(3\right)}(\mathcal{X},{\mathcal{X}}^{\prime })=⟨{\mathcal{X}}^{\prime }|{\hat{R}}_{{\mathcal{X}}^{\prime }\to \mathcal{X}}|\mathcal{X}⟩{\mathbf{R}}_{{\mathcal{X}}^{\prime }\to \mathcal{X}},$$

$$\ket{\mathcal{X}\lambda\mu}=\ket{\mathcal{X}}\otimes\ket{\lambda\mu},\$$

$${|\mathcal{X}\lambda \mu ⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}=\int \text{d}\hat{R}|\hat{R}\mathcal{X}⟩\otimes \hat{R}|\lambda \mu ⟩.$$

$${|\mathcal{X}\lambda \mu ⟩}_{\mathcal{O}(3)}={|\mathcal{X}\lambda \mu ⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}+(-1{)}^{\lambda }{|\left(\hat{i}\mathcal{X}\right)\lambda \mu ⟩}_{\mathcal{S}\mathcal{O}\left(3\right)}.$$

$${\mathbf{k}}_{\mathcal{S}\mathcal{O}\left(3\right)}^{\lambda }(\mathcal{X},{\mathcal{X}}^{\prime })=\int \text{d}\hat{R}⟨\mathcal{X}||\hat{R}{\mathcal{X}}^{\prime }⟩{\mathbf{D}}^{\lambda }(\hat{R})$$

$${\psi }_{{\mathcal{X}}_j}^{\alpha }(\mathbf{r})=\sum _{i\in {\mathcal{X}}_j,\alpha }\exp \left\{-\frac{|\mathbf{r}-\left({\mathbf{r}}_i-{\mathbf{r}}_j\right){|}^2}{2{\sigma }^2}\right\}$$

Full text

††preprint: APS/123-QED

I Introduction

The purpose of a statistical learning model is the prediction of regression targets by means of simple and easily accessible input parameters mlbook . In chemistry, physics and materials science, regression targets are usually scalars or tensors, including electronic energies bartok2010 ; jain2013 ; calderon2015 ; ward2017 , quantum-mechanical forces Li ; Glielmo2017 ; Glielmo2018 , electronic multipoles Yuan2014 ; Bereau2015 ; Bereau2017 , response functions liang2017 ; grisafi2018 ; Wilkins2019 ; Christensen2019 and scalar fields like the electron density burke2017 ; Alred2018 ; grisafi2019 . For ground-state properties, the regression input usually consists of all the information connected with the atomic structure at a given point of the Born-Oppenheimer surface, e.g. nuclear charges and atomic positions. A more or less complex manipulation of these primitive inputs leads to what is usually called a structural descriptor, or representation (Fig. 1).

It is widely recognized that an essential ingredient for maximizing the efficiency of machine learning models is to use representations that mirror the properties one wants to predict. Here we discuss an effective approach to build linear regression models for tensors. The notion that the representation should mirror the property means when a symmetry operation is applied to an atomic structure, the associated representation should transform in a way that mimics the transformation of the properties of the structure. It should be stressed that it is completely possible to build a machine learning model that does not incorporate such transformation properties. The universal symmetries of the property must then be learned by the model through exposure to data in the training set, making the training process less efficient. A crucial focus of this chapter is the creation of symmetry-adapted representations.

Once one has a symmetry-adapted representation at hand, the linear regression model is bound to fulfill the symmetry requirements imposed by the property braa-bowm09irpc ; behler2007 ; bartok2013 ; Shapeev2015 ; Zhang2018 . There is, however, another important consideration when building a model for tensors, expressed in terms of a Cartesian reference system. It is well known that any tensor can be decomposed into a set of spherical components that transform independently under rotations Weinert1980 ; Stone . Particularly for high-order tensors, the irreducible spherical decomposition of a tensor simplifies greatly the learning task, compared to the Cartesian representation, as we will discuss later on.

The process of symmetry-adapting a representation is general but rather abstract, and for it to be practical one must choose the initial representation with care. For this purpose we use the Smooth Overlap of Atomic Positions (SOAP) framework, which is based on the representation of atom-centered environments constructed from a smooth atom density built up using Gaussians centered on each neighbor of the central atom. This density-based representation can be adapted to incorporate correlations between atoms to any order. It has been applied successfully to a vast number of machine learning investigations for physical properties of atomic structures Sandip2016 ; musi+18cs ; bartok2017 . After summarizing the derivation and efficient implementation of an extension to SOAP, called $\lambda$-SOAP, which is particularly well-suited to the learning of tensorial properties, we present a few examples to demonstrate its effectiveness for this task.

II Linear Regression

Suppose one wanted to build a linear regression model to predict a scalar property $y(\mathcal{X})$ for an input $\mathcal{X}$,

[TABLE]

In this equation $\ket{w}$ represents the weight vector we wish to learn and $\ket{\mathcal{X}}$ is a representation of the input. The usual approach for learning the weight vector is to suppose the properties are independently and normally distributed, i.e.

[TABLE]

One then maximizes the log likelihood of a set of $N$ observations $\{y_{n}\}$ with respect to the weight vector. The log likelihood (loss or cost function) is given by

[TABLE]

where the regularizer $\alpha^{2}\bra{w}\ket{w}$ appears if one introduces a Gaussian prior on $w$ with variance $\alpha^{-2}$. $L(w)$ attains its maximum at

[TABLE]

where the covariance $\hat{C}$ is

[TABLE]

and $\eta=\alpha/\sigma$.

The preceding linear regression scheme in which one handles the representation $\ket{\mathcal{X}}$ explicitly is often called the primal formulation. There is in fact another, complementary formulation called the dual (Kernel Ridge Regression (KRR) or Gaussian Process Regression (GPR)) in which the equations take a slightly different form. In the dual, one does not handle the representation explicitly but rather introduces a kernel function which – roughly speaking – measures the similarity between two inputs. The link between the primal and dual lies in the observation that a positive-definite kernel $k(\mathcal{X},\mathcal{X}^{\prime})$ can always be written as an inner product mlbook ,

[TABLE]

This means that given a kernel one can always construct a representation and vice versa. From the perspective of GPR, the kernel is interpreted as the covariance between the properties of its two arguments,

[TABLE]

The properties are assumed to be normally distributed, which means one can straightforwardly find the conditional distribution of the property $y(\mathcal{X})$ given a set of observations in a training set $\{y_{n}\}$. The mean of this distribution is given by

[TABLE]

where the $j$th component of $\mathbf{k}(\mathcal{X})$ is $k(\mathcal{X},\mathcal{X}_{j})$, $K_{jk}=k(\mathcal{X}_{j},\mathcal{X}_{k})$ and $\mathbf{y}$ is a vector formed from $\{y_{n}\}$.

When the feature space associated with a kernel is known explicitly, and finite-dimensional, the primal and dual formulations are formally equivalent, and the choice of which to use is an important but purely practical question. Constructing a primal model requires inversion of the covariance matrix, while the dual requires inversion of the kernel matrix $\bm{K}$. If the feature space, i.e. the space occupied by the representation, is larger than the training set then the GPR approach is more convenient. Of course, the real utility of the kernel trick becomes apparent when the kernel is a complex, non-linear function for which the feature space is unknown and/or infinite-dimensional. In these circumstances, working in the dual makes it possible to formulate regression as a linear problem, where reference configurations (or a sparse set of representative states) is used to define a basis for the target, as in the r.h.s. of Eq. (8). As such, all the complexity of the input space representation is contained in the definition of the kernel function.

III Tensors, symmetries and correlations

The previous discussion defines the general architecture of regression models which can be used to predict any scalar quantity associated with the molecular geometry. We now discuss the implications of learning tensors, or, similarly, any quantity that is not invariant under a rigid rotation or reflection of the atomic structure. In so doing, we will introduce a formalism which is general enough to encompass both proper Cartesian tensors, such as molecular polarizabilities, and three-dimensional scalar fields that can be conveniently decomposed in atom-centered contributions, such as the ground-state charge density of a molecule.

Let us start by considering the prototypical case of a Cartesian tensor $y\equiv y_{\alpha\beta...}$ of rank $r$, with the combination of indices $\{\alpha\beta...\}$ running over a number of Cartesian components equal to $3^{r}$. Given any arbitrary distorted atomic structure with no particular internal symmetry, we are interested in characterizing the transformations of the tensor under only three families of symmetry operations, viz. translations, rotations and reflections. Since these symmetry operations do not affect the internal geometry of an atomic structure, we can think equivalently in terms of active transformations, in which the system undergoes the symmetry operation and the reference frame remains fixed, or in terms of passive transformations, in which the reference frame undergoes the symmetry operation and the system remains fixed. In the following, we summarize the symmetry operations by adopting an active picture and assume the system is not subjected to an external field.

Translations. Any physical property of an atomic structure $\mathcal{X}$ remains unchanged under a rigid translation $\hat{t}$ of atomic positions, i.e.,

[TABLE]

Rotations. Under the application of a rigid rotation $\hat{R}$ to an atomic structure $\mathcal{X}$, we assume that each Cartesian component of the tensor undergoes a covariant linear transformation. Using Einstein notation for convenience, and representing by $\bm{R}$ the rotation matrix corresponding to $\hat{R}$, the rotated tensor is

[TABLE]

Reflections. Applying a reflection operator $\hat{Q}$ to an atomic structure $\mathcal{X}$ through any mirror plane leads to the following reflected tensor,

[TABLE]

III.1 Covariant descriptors

In general terms, a primitive representation that mirrors a tensor of a given rank $r$ could formally be built by considering

[TABLE]

where $\ket{\mathcal{X}}$ is an arbitrary description of the system, while $\ket{\alpha}$ represents a set of Cartesian axes which is rigidly attached to the system. When using this primitive representation in a linear regression model, the tensor component corresponding to $\alpha\beta...$ would be

[TABLE]

or

[TABLE]

After maximizing the log likelihood, the former possibility leads to a model that predicts every component to be the same, while the latter ignores the known correlations between the components and is therefore likely to overfit. For example, consider a training set in which only one of the tensor components is non-zero. All but one of the regression weights $\{\ket{w_{\alpha\beta...}}\}$ would be driven towards zero to maximize the log likelihood, so the trained model would only predict a finite value for the component it had been explicitly exposed to in the training set. The model would therefore incorrectly predict the tensor components for a structure differing only by a rigid rotation from one in the training set.

To address these problems, one should adapt the primitive descriptor so that it fulfills each of the symmetries detailed in Eqs. (9-11). Since the Cartesian basis vectors are invariant under translations, Eq. (9) implies the core representation should itself be invariant under translations. Using Haar integration one can construct a core representation that is invariant under translations by integrating an arbitrary representation over the translation operator $\hat{t}$ willatt2018arxiv . One can then proceed to consider covariance under $\mathcal{SO}\left(3\right)$ group operations. Eq. (10) implies that a covariant representation for $\ket{\mathcal{X}\alpha\beta...}$ should satisfy the invariance relationship

[TABLE]

for any rotation $\hat{R}$. Starting from the primitive definition of Eq. (12), there are a variety of ways to enforce this invariance relationship. One possibility is to use

[TABLE]

where the operator $\hat{R}_{\mathcal{X}\to}$ is defined to rotate $\mathcal{X}$ into a specified orientation which is common to all the molecules of the dataset (Fig. 2).

This works under the assumption that it is always possible to define a unique (and therefore unambiguous) internal reference frame to rotate $\mathcal{X}$ into a specified orientation, which might be possible when the system involved has a particularly rigid internal structure. A more general strategy, which does not require any assumption on the molecular geometry to be made, consists in considering the covariant integration over the operator $\hat{R}$ (Haar integration),

[TABLE]

On the top of this definition, the requirement that a representation be covariant in $\mathcal{O}(3)$, including the reflection symmetry of the tensor as in Eq. (11), means that a simple linear combination of the $\mathcal{SO}\left(3\right)$ descriptor with its reflected counterpart with respect to any arbitrary mirror plane of the system must be considered; that is,

[TABLE]

for any arbitrary choice of $\hat{Q}$. Any other reflection operation can be automatically included by having made the descriptor covariant under rotations.

III.2 Covariant regression

Having shown how to build a symmetry-adapted representation of the system, let us see the implications of this procedure for linear regression. Using a symmetry-adapted representation in a linear regression model leads to the following solution for the regression weight,

[TABLE]

where the covariance is

[TABLE]

Note that the solution for the linear regression weight does not change when the training structures and corresponding tensors simultaneously undergo a symmetry operation that the representation has been adapted to. In other words, the same model results regardless of the arbitrary orientation of structures in the training set.

When moving to the dual, we find the kernel to be

[TABLE]

This result corresponds to

[TABLE]

As stressed earlier, performing the linear regression in the dual using this kernel leads to a formally-equivalent model to that resulting from the primal formulation described above, yet this kernel appears to be more complicated than a symmetry-adapted descriptor since it involves two integrations over rotations. If, however, we assume the core representation $\ket{\mathcal{X}}$ undergoes a unitary transformation when the system is rotated,

[TABLE]

the kernel reduces to

[TABLE]

where $k(\mathcal{X},\mathcal{X}^{\prime})=\bra{\mathcal{X}}\ket{\mathcal{X}^{\prime}}$ is the kernel corresponding to the core representation. The requirement that the core representation should undergo a unitary transformation when the system is rotated is reasonable since, if it were not true, the autocorrelation $k(\mathcal{X},\mathcal{X})$ would depend on the absolute orientation of $\mathcal{X}$, which is unphysical given our assumption of the absence of external fields.

Note that upon defining a collective tensorial index $\{\alpha\beta...\}$, a kernel matrix of size $3^{r}N\times 3^{r}N$ can be constructed by stacking together each of the $3^{r}\times 3^{r}$ vector-valued correlation functions Alvarez2012 . Then, a covariant tensorial prediction of the property of interest can eventually be carried out according to the GPR prescription of Eq. (8).

It is instructive to compare the symmetry-adapted kernel definition of Eq. (24) to the kernel that one gets from the aligned descriptors of Eq. (16). In this case, building a kernel function on the top of this descriptor effectively means carrying out the structural comparison in a common reference frame where the two molecules are mutually aligned. One can then conveniently learn the tensor of interest component-by-component through a much simpler scalar regression framework. For the simple case of rank-1 tensors, for instance, we would get,

[TABLE]

where we have defined the best alignment operator as $\hat{R}_{\mathcal{X}^{\prime}\to\mathcal{X}}=\hat{R}_{\mathcal{X}^{\prime}\to}\hat{R}^{T}_{\mathcal{X}\to}$. This strategy has been successfully used in the learning of electronic multipoles of organic molecules Bereau2015 as well as for predicting optical response functions of water molecules in their liquid environments liang2017 . For the latter example, a representation of the best-alignment structural comparison is reported in Fig. 3.

This method for tensor learning has the clear drawback of relying on the definition of a rigid molecular geometry, for which an internal reference frame can be effectively used to perform the procedure of best alignment. Following this line of thought, the availability of a covariant kernel function allows us to implicitly carry out both the structural comparison and the geometric alignment of two molecules simultaneously, neglecting any prior consideration about the internal structure of the molecule at hand.

III.3 Spherical representation

The family of Cartesian symmetry-adapted descriptors previously introduced can be effectively used, in principle, to predict any Cartesian tensor of arbitrary order. However, we should notice that having a tensor product for each additional Cartesian axis makes the cost of the regression scale unfavorably with the tensor order, producing a global kernel matrix of dimension $(3^{r})^{2}$. In fact, it is well established that a more natural representation of Cartesian tensors is given by their irreducible spherical components (ISC) Stone . As described in Ref. Stone , the transformation matrix from Cartesian to spherical tensors can be found recursively, starting from the known transformation for rank-2 tensors.

Upon trivial manipulations, the components of a spherical tensor transform separately as the irreducible representations of the $\mathcal{SO}\left(3\right)$ group. Each $\lambda$-component of the tensor spans an orthogonal subspace of dimension $2\lambda+1$. For instance, the 9 components of a rank-2 tensor separate out into a term (proportional to the trace) that transforms like a scalar, three terms that transform like $Y^{1}_{m}$, and five terms that transform like $Y^{2}_{m}$. When using a spherical representation, the kernel matrix is block diagonal, which greatly reduces the number of non-zero entries, and makes it possible to learn separately the different components. An additional advantage is that the possible symmetry of the tensor can be naturally incorporated by retaining only the spherical components $\lambda$ that have the same parity as the tensor rank $r$. For instance, the $\lambda=1$ component of a symmetric rank-2 tensor vanishes identically, meaning that only the 6 surviving elements of the tensor need to be considered when doing the regression. Especially for high rank tensors, this property means that the number of components can be cut down significantly.

In light of the discussion carried out for Cartesian tensors, it is straightforward to realize how a symmetry-adapted descriptor that transforms covariantly with spherical harmonics of order $\lambda$ should look. Since each ISC is effectively a vector of dimension $2\lambda+1$, we can first write a primitive spherical harmonic representation as

[TABLE]

where $\ket{\lambda\mu}$ is an angular momentum state of order $\lambda$, such that $\bra{\hat{\mathbf{r}}}\ket{\lambda\mu}=Y^{\lambda}_{\mu}(\hat{\mathbf{r}})$. Its symmetry-adapted counterpart, which is covariant in $\mathcal{SO}\left(3\right)$, is

[TABLE]

Finally, since odd spherical harmonics are anti-symmetric with respect to the inversion operator $\hat{i}$, a spherical tensor descriptor that is covariant in $\mathcal{O}(3)$ can be obtained by considering

[TABLE]

Note that a tensorial kernel function built on the top of this descriptor would transform under rotations as the Wigner-$D$ matrix of order $\lambda$, $D_{\mu\mu^{\prime}}^{\lambda}=\langle\lambda\mu|\hat{R}|\lambda\mu^{\prime}\rangle$:

[TABLE]

In addition to being the most natural strategy to perform the regression of Cartesian tensors, using a representation like that of Eq. (28) comes in handy when building regression models for the many physical properties that can be decomposed in a basis of atom-centered spherical harmonics. In the following sections, we will give an example of this kind by predicting the ground-state electronic charge density of molecular systems.

IV SOAP representation

We now proceed to characterize the exact functional form of a symmetry-adapted representation of order $\lambda$ which can be used to carry out a covariant prediction of any property that transforms as a spherical harmonic. In Sec. II it was pointed out that, within a framework of linear regression, both the primal and the dual formulation can be adopted to actually implement the interpolation of a given tensorial property. In what follows, however, we will focus our attention on the dual formulation, discussing in parallel the feature vector associated with the $\lambda$-SOAP representation and the corresponding kernel function. This choice is justified by the greater flexibility of the kernel formulation, allowing a non-linear extension of the framework as discussed below.

An atom-centered environment $\mathcal{X}_{j}$ describes the set of atoms that are included within a spherical cutoff $r_{\text{cut}}$ around the central atom $j$. We will label as $\ket{\mathcal{X}_{j}}$ the abstract vector which describes the local structure. A convenient definition of $\ket{\mathcal{X}_{j}}$ in real space can be obtained by writing a smooth probability amplitude, for each atomic species $\alpha$, as a superposition of Gaussians with spread $\sigma$ that are centered on the positions $\{\mathbf{r}_{i}\}$ of the atoms that surround the central atom $j$:

[TABLE]

This definition descends naturally from the requirement of translational invariance of a representation of the entire structure willatt2018arxiv and corresponds to the construction that is used in Ref. bartok2013 to define the Smooth Overlap of Atomic Positions (SOAP) kernel. Formally, one can then write

[TABLE]

with the ket $\ket{\alpha}$ tagging the identity of each species. Even though it might be convenient to use a lower-dimensional chemical space will+18pccp , particularly when building models for dataset containing many elements, in what follows we will assume that each element is associated with an orthogonal subspace, i.e. $\bra{\alpha}\ket{\beta}=\delta_{\alpha\beta}$. This implies that, when using this representation to define a scalar-product kernel, only the density distributions of the same atomic type are overlapped,

[TABLE]

With this choice, the two adjustable parameters $r_{\text{cut}}$ and $\sigma$ determine respectively the range and the resolution of the representation. To simplify the notation, we will omit the $\alpha$ labels, assuming that a single element is present. The extension to the case with multiple chemical species follows straightforwardly.

IV.1 $\lambda$-SOAP(1) representation

To the first order in structural correlations, including the environmental state $\ket{\mathcal{X}_{j}}$ in the definition of a local symmetry-adapted descriptor of order $\lambda$ reads

[TABLE]

The real space representation of $\ket{\mathcal{X}^{(1)}_{j}\lambda\mu}$ can be understood as a rotational average of the environmental density which is rigidly attached to a spherical harmonic of order $\lambda$,

[TABLE]

A more concise, and easily-computed version of this representation results from projecting $\ket{\mathcal{X}^{(1)}_{j}\lambda\mu}_{\mathcal{SO}\left(3\right)}$ on a basis of spherical harmonics, in which the integral over rotations can be performed analytically,

[TABLE]

It is clear that many of the indices in this representation are redundant, and would have no effect when taking an inner product between two such representations. The most concise form that produces the same scalar product kernel as Eq. (34) is

[TABLE]

where we introduced the spherical density component

[TABLE]

This ket corresponds to the kernel

[TABLE]

which is straightforward to calculate using e.g. quadrature in $r$ or an expansion on a radial basis.

It is insightful to consider the explicit expression for Eq. (36) in terms of the atom density. Taking for instance the case of $\lambda=1,\mu=0$, for which $Y^{1}_{0}(\hat{\mathbf{r}})\propto x/r$:

[TABLE]

One sees that the 2-body $\lambda$-SOAP representation corresponds to moments of the smooth atom density, resolved over different shells around the central atom. A linear model built on these features can respond to changes in the atomic density at different distances, simultaneously adapting the magnitude and geometric orientation of the target property.

IV.2 $\lambda$-SOAP(2) representation

Describing an atomic environment in a way that goes beyond the two-body structural correlations ($\nu>1$) is of fundamental importance, because information on distances alone is not sufficient to uniquely determine an atomic structure. Building on the definition of Eq. (33), and on the symmetrized-atom-density framework of Ref. willatt2018arxiv , this can be achieved by introducing an additional tensor product in the environmental state $\ket{\mathcal{X}_{j}}$ within the rotational average,

[TABLE]

By projecting on a real-space basis, the representation becomes

[TABLE]

Similarly to the $\nu=1$ case, one can compute the ket without an explicit rotational average by projecting on a basis of spherical harmonics,

[TABLE]

where the parentheses denote a Wigner 3j symbol. Just as for the $\lambda$-SOAP(1) case considered earlier, it is clear that many of the indices in this expression are redundant. When taking an inner product between two such representations, one can use orthogonality of Wigner 3j symbols to simplify to an inner product between two objects with the following form,

[TABLE]

The Clebsch-Gordan coefficient $\bra{lk,lk^{\prime}}\ket{\lambda\mu}$ has the role of combining two angular momentum components of the atomic environment $\mathcal{X}_{j}$ to be compatible with the spherical tensor order $\lambda$. This contains all the essential information of the abstract representation $\ket{\mathcal{X}^{(2)}_{j}\lambda\mu}_{SO(3)}$ that is needed for $\lambda$-SOAP(2) linear regression. Note that $\bra{lk,lk^{\prime}}\ket{\lambda\mu}$ is zero unless $k+k^{\prime}=\mu$, that the indices $l$, $l^{\prime}$ and $\lambda$ must satisfy the inequality $|l-l^{\prime}|\leq\lambda\leq l+l^{\prime}$ and that the representation is invariant under transposition of $r$ and $r^{\prime}$.

Let us see how the representation changes under inversion. Given the parity of the spherical harmonics,

[TABLE]

it follows that

[TABLE]

This condition implies that a representation that is covariant in $\mathcal{O}(3)$ can be easily obtained by retaining only the components of the feature vectors for which $l+l^{\prime}+\lambda$ is even. Generalization of this procedure to higher orders of $\lambda$-SOAP is tedious but straightforward using well-known formulae for integrals of products of Wigner-$D$ matrices over rotations.

IV.3 Non-linearity

As already mentioned in the introduction, a crucial aspect to improve regression performance is to incorporate non-linearities in the construction of the representation. For instance, tensor products of the scalar representation introduce higher body order correlations, in a way that can be easily implemented in a kernel framework by raising the kernel to an integer power willatt2018arxiv . When working with tensorial representations, however, one has to be careful to avoid breaking the covariant transformation properties of the feature vector. Taking products of $\ket{\mathcal{X}_{j}^{(\nu)}\lambda\mu}_{\mathcal{O}(3)}$ kets would require re-projecting the product onto the irreducible representations of the group, which would be as cumbersome as increasing the body order exponent $\nu$. One obvious solution to this problem is to multiply the spherical kernel of order $\lambda$ by its scalar and rotationally invariant counterpart, which can then be raised to an integer power $\zeta$ without breaking the tensorial nature of the kernel. For any generic order $\nu$ and $\nu^{\prime}$ in structural correlations, this procedure consists in considering the tensor product

[TABLE]

which leads to the kernel definition

[TABLE]

For $\zeta=1$, one recovers the original tensorial kernel, while a non-linear behavior is introduced for $\zeta>1$. A considerable improvement of the learning power is usually obtained when using $\zeta=2$, while negligible further improvement is observed for $\zeta>2$.

These considerations also apply to the use of fully non-linear machine-learning models like a neural network. To guarantee that the prediction of the model is consistent with the group covariances, the tensorial $\lambda$-SOAP features must enter the network at the last layer, and all the previous non-linear layers can only contribute to different linear combinations of the tensorial features, e.g.

[TABLE]

where each of the $f_{rr^{\prime}ll^{\prime}}$ can be an arbitrary non-linear combination of the scalar SOAP features (see Fig. 4). Similar ideas have already been implemented in the context of generalizing the construction of spherical convolutional neural networks kondor2018arxiv .

IV.4 Implementation

In the previous discussion it was pointed out that beyond the formal definition of the structural descriptor in real space, the kernel evaluation eventually requires the computation of the SOAP density power spectrum $\bra{rr^{\prime}ll^{\prime}}\ket{\mathcal{X}_{j}^{(2)}\lambda\mu}$. In turn, computing this quantity requires the evaluation of the density expansion coefficients $\bra{rlm}\ket{\mathcal{X}_{j}}$. In practice, the continuous variable $r$ can be replaced by an expansion over a discrete set of orthogonal radial functions $R_{n}(r)$ that are defined within the spherical cutoff $r_{\text{cut}}$. For this reason, we will refer, from now on, to the density expansion coefficients as $\bra{nlm}\ket{\mathcal{X}_{j}}$.

Having represented the environmental density distribution as a superposition of Gaussian functions centered on each atom, the spherical harmonics projection can be carried out analytically Kaufmann , leading to:

[TABLE]

where the sum over $i$ runs over the neighboring atoms of a given chemical element, and $\iota_{l}$ represents a modified spherical Bessel function of the first kind. Under suitable choices of the functions $R_{n}(r)$, the radial integration can also be carried out analytically, too.

One possibility is to start with non-orthogonal Gaussian type functions, $\tilde{R}_{k}(r)$, reminiscent of Gaussian-type orbitals commonly used in quantum chemistry:

[TABLE]

where $\mathcal{N}_{k}$ is a normalization factor, such that $\int_{0}^{\infty}drr^{2}\tilde{R}^{2}_{k}(r)=1$. The set of Gaussian widths $\{\sigma_{k}\}$ can be chosen to effectively span the radial interval involved in the environment definition. For instance, one can take $\sigma_{k}=r_{\text{cut}}\max(\sqrt{k},1)/n_{\text{max}}$, obtaining functions that have equally-spaced peaks between 0 and $r_{\text{cut}}$. The explicit functional form of the primitive radial integrals is

[TABLE]

where $\Gamma$ is the Gamma function, while ${}_{1}F_{1}$ is the confluent hypergeometric function of the first kind. These primitive integrals can be finally orthogonalized by applying the orthogonalization matrix $\bm{S}^{-1/2}$, with $\bm{S}$ representing the overlap matrix between primitive functions,

[TABLE]

for which well-known analytical expressions exist Gradshteyn2007 .

V Examples

In this section, the effectiveness of a kernel ridge regression model that is adapted to the fundamental physical symmetries of the target is demonstrated, considering two very different quantities as examples. The first example involves the prediction of the first order dipole polarizability $\bm{\alpha}$ for a broad dataset of organic molecules, when training a $\lambda$-SOAP(2) regression model on the irreducible spherical components of the tensor $\bm{\alpha}$. In the second example, we show how to predict the charge density $\rho(\mathbf{r})$ of a small, yet flexible, hydrocarbon molecule like butane, by decomposing $\rho(\mathbf{r})$ into atom-centered spherical harmonics. In both cases, a comparison of the prediction performance is carried out between $\lambda$-SOAP(2) descriptors that are covariant in $\mathcal{SO}\left(3\right)$ - which were used in previous work - and those that have been made fully $\mathcal{O}(3)$ compliant by symmetrization over $\hat{i}$.

V.1 Molecular Polarizabilities

The polarizability of a molecule is a rank-2 tensor describing the second-order response of the molecular energy $U$ with respect to an applied electric field $\mathbf{E}$, with components $\alpha_{ij}=\partial^{2}U/\partial E_{i}\partial E_{j}$. By construction this tensor is symmetric, meaning that it can be decomposed into two components: a $\lambda=0$ (scalar) component,

[TABLE]

Fig. 5 shows learning curves for the prediction of the $\mathbf{\alpha}^{(2)}$ component for polarizabilities of molecules in the QM7b database mont+13njp , which contains 7,211 small molecules. The results are shown for two types of kernel: the first of these is an $\mathcal{SO}\left(3\right)$ kernel as used in Ref. Wilkins2019 and the second of this is an $\mathcal{O}(3)$ kernel accounting for inversion symmetry. All kernels were built using 8 radial functions, maximum spherical harmonic cutoff $l_{\text{ max}}=6$, $\sigma=0.3~{}\text{\AA}$ and $\zeta=2$, and kernels with different values of the radial cutoff $r_{\text{cut}}$ were combined to form a multiscale kernelbartok2017 as $\bm{k}^{\lambda}(\mathcal{X},\mathcal{X}^{\prime})=\sum_{i}c_{i}\bm{k}^{\lambda}_{i}(\mathcal{X},\mathcal{X}^{\prime})$, where $\bm{k}^{\lambda}_{i}(\mathcal{X},\mathcal{X}^{\prime})$ is the kernel with $r_{\text{ cut}}/\text{\AA}=i$ and $c_{2}=0.04053$, $c_{3}=0.00997$, $c_{4}=0.02250$, $c_{5}=0.01560$. The use of a $\mathcal{O}(3)$ reduces approximately by half the size of the feature vector, and at the same time it leads to a small - but consistent - reduction in the test error.

V.2 Electronic charge densities

Another learning task that can benefit from a symmetry-adapted regression scheme involves the learning of scalar fields such as the electron charge density. Machine-learning models for the charge density have been proposed based on the coefficients in a plane wave basis burke2017 – this is convenient due to orthogonality, but leads to poor transferability when considering flexible molecules, or learning across different molecular species – or based on direct prediction of the density on a real-space grid Alred2018 ; Chandrasekaran2019 . By expanding the density on an atom-centred basis set, composed of radial functions multiplied by spherical harmonics,

[TABLE]

one obtains a model that is localized and transferable, concise, and easily integrated with the many electronic structure codes that are based on atom-centered basis functions. The coefficients in the expansion transform under rotations like spherical harmonics, and can therefore be learned efficiently using a SA-GPR model,

[TABLE]

where the sum runs over a set of reference environments $Z_{i}$ centered around atoms of the same kind as $i$, and the weights are computed by a regression procedure that is complicated by the fact that the basis set is not orthogonal grisafi2019 .

In Fig. 6 we report the result obtained for a dataset of butane molecules (C4H10), for which 1000 reference pseudo-valence densities have been computed at the DFT/PBE level. The dimensionality of the regression problem is defined by considering the 100 most diverse atomic environments, out of a total of 14000, selected by farthest point samplingceri+13jctc through the 0-SOAP(2) distance metric. Given that in our previous work the learning performance was essentially limited by the basis set expansion error for the density, we decided to compare the optimized basis set used in Ref. 18 with a resolution of the identity (RI) basis set, usually adopted in the context of avoiding the computation of the four-center Hartree integral in electronic structure theory Hattig2005 . When considering in particular the RI-cc-pV5Z basis, which accounts for basis functions up to $l$=4 of angular momentum, we find that the basis set decomposition error is almost halved ($\sim$0.6%) with respect to Ref. 18, as shown by the asymptotic convergence in Fig. 6. The figure also compares, in the case of the RI basis, the learning performances associated with $\lambda$-SOAP(2) descriptors that have been made covariant in $\mathcal{SO}\left(3\right)$ and $\mathcal{O}$(3) respectively. As seen for the case of polarizability, the $\mathcal{O}$(3) features improve – although only slightly – the prediction accuracy. The improvement is more substantial at the smallest training set size, where the incorporation of prior knowledge on the symmetries of the system can make up for the scarcity of data.

VI Outlook

The previous examples show how statistical learning of a tensorial quantity across the configurational space of atomic coordinates and composition represents a challenging methodological task which requires considerable modifications to the architecture of more familiar scalar learning models. The efficiency of a regression model benefits greatly from the incorporation of symmetry, as it effectively reduces the dimensionality of the space in which the algorithm is asked to interpolate the values of the target property. Symmetry of tensorial quantities should be included in two distinct ways. First, one should decompose the tensor into irreducible spherical components, so as to minimize the amount of information that is needed to account for geometric covariance. Particularly for high-rank Cartesian tensors, the matrix of correlations between tensor elements can be made block diagonal, which reflects on the size and complexity of the associated kernel matrices. Second, by constructing representations of the molecular structure that are made isomorphic with the tensor of interest, one can obtain a linear basis that satisfies the expected covariant transformations. An important aspect to consider is that, in order to preserve the properties of the symmetry-matched basis, non-linearities have to be treated with care. We discuss how it is possible to do so in the context of kernel ridge regression models, and how one should proceed to design a covariant neural network that can be used to efficiently accomplish a symmetry-adapted regression task.

We discuss a practical implementation of these ideas within the framework of the smooth overlap of atomic positions representations, that uses a spherical-harmonics representation of the atom density and is therefore particularly well-suited to incorporate $\mathcal{SO}\left(3\right)$ covariance. We discuss an extension, that we refer to as $\lambda$-SOAP, that provides a natural linear basis to regress quantities that transform like spherical harmonics, and can be made to represent arbitrarily high body-order correlations between atomic coordinates. As an original result of this work, we also discuss how to satisfy the inversion symmetry of the tensor, showing that representations that incorporate the full $\mathcal{O}(3)$ covariances improve the performance of the machine-learning model, particularly in the limit of a small training set. We also show an example of the use of $\lambda$-SOAP representations to learn a scalar field in three-dimension as a sum of atom-centered contributions, choosing the electron density as a physically relevant example. We believe that this strategy – although more complex than alternatives that use orthogonal basis functions or a real-space grid – has the best promise to be transferable across different systems, and to be combined with standard electronic structure packages.

Acknowledgements.

The Authors acknowledge support by the European Research Council under the European Union’s Horizon 2020 research and innovation programme (Grant Agreement No. 677013-HBMAP).