Geometric Transformer for End-to-End Molecule Properties Prediction

1 Introduction

Properties of chemical compounds can generally be estimated using methods such as density functional theory (DFT) or ab initio quantum chemistry [Jensen, 2017]. However, these can be computationally expensive and therefore have a limited applicability, especially for larger systems. In recent years, many approaches have started leveraging machine learning to reduce the computational complexity required for efficiently predicting molecular properties.

In this vein, many contributions have focused on the creation of handcrafted representations at the atomic or molecular level [Christensen and others, 2020; Huang and others, 2016] as input for various machine learning methods. Schr¨odinger's equation indicates that the system variables that define the ground-state properties of a given molecule are a function of the inter-atomic distances and the nuclear charges solely. [Jensen, 2017]. Based on this observation, several recent methods predict molecular properties in an end-to-end fashion where the input is defined by the atoms' type and spatial position. Such methods often incorporate quantum chemistry knowledge and rely on extensive hyperparameter tuning.

Since atomic interactions are challenging to simulate, many recent works make use of graph neural networks as a natural tool to model molecules [Gilmer et al., 2017]. Related to graph neural networks, Transformers [Vaswani et al., 2017] have recently become extremely popular in numerous application domains.

In this paper we extend the ubiquitous Transformer to chemical compounds data in order to predict their groundstate properties. Our work does not employ extended domain knowledge, and is based solely on the simple distance relevance assumption, namely, that the bigger the distance between atomic elements, the lower the interaction.

Contrary to other works, the framework does not assume any extended input, e.g. quantummechanical properties[Qiao and others, 2020], complex geometric constants such as bending or torsion angles [Klicpera et al., 2020b], additional solvers, e.g. fast DFT as residual solvers (i.e. delta learning) [Unke and others, 2019; Qiao and others, 2020]), or even knowledge adaptation from quantum chemistry into the machine learning model design [Sch¨utt et al., 2018a; Klicpera et al., 2020b].

The Transformer we design is endowed with an adapted positional encoder, and with learned inter-atomic geometric embedding at the different levels of the model, allowing increased representational power, while maintaining the molecule invariance to rigid transformations and permutation. For better regularization, we suggest to augment the training set by merging pairs of molecules positioned far apart.

The experimental results demonstrate the representational power of the end to end model, potentially allowing removal of undesirable inductive bias [Goyal and Bengio, 2020] such as handcrafted envelope functions. Also, by allowing the model to fit the data while minimizing the domainknowledge, scientific insights can emerge, such as effective atomic cut-off radius, or most relevant molecular interactions for a given property.

Related Work Classical molecule property prediction methods combined handcrafted features generally based on force field methods at the atom level [Christensen and others, 2020] or molecular level [Huang and others, 2016], integrated into various machine learning models such as Kernel methods [Christensen and others, 2020] , Gaussian processes [Bartok et al., 2010] or Neural networks [Sch¨utt et al., 2018a].

These methods have recently been superseded by end-toend neural networks, alleviating the need for handcrafted signatures. The most popular and powerful models are based on graph neural networks, which allow a natural representation of the molecular graph [Sch¨utt et al., 2018a; Unke and others, 2019; Anderson et al., 2019; Klicpera et al., 2020b; Qiao and others, 2020]. These architectures are generally designed based on the message-passing mechanism, by aggregating features obtained from atom types, geometric invariants such as pairwise inter-atomic distances [Sch¨utt et al., 2018a; Unke and others, 2019], bending or torsion angles [Klicpera et al., 2020b], or handcrafted atomic features derived from quantum mechanics [Qiao and others, 2020]. Existing methods generally involve a deep understanding and a cautious adaptation of the underlying physics in order to provide better preconditioning [Klicpera et al., 2020b; Qiao and others, 2020].

Transformer neural networks were originally introduced for machine translation [Vaswani et al., 2017] and they now dominate most applications in the field of Natural Language Processing. Transformer encoders primarily rely on the selfattention operation in conjunction with feed-forward layers, allowing manipulation of variable-size sequences and learning of long-range dependencies. Many works have augmented the self-attention mechanism using domain-specific knowledge [Chen et al., 2017; Bello et al., 2019].

Recently, a transformer architecture called MAT has been proposed for chemical molecules [Maziarka et al., 2020; Maziarka et al., 2021]. MAT modifies the self-attention module by summing the inverse exponent of the pairwise distance matrix to the self-attention tensor. As one contribution of our work, we show that the MAT self-attention architecture is sub-optimal in modelling interactions, and we propose a better self-attention module capable of capturing the connectivity of the graph. Our method also outperforms concurrent efforts [Wu et al., 2021; Kwak et al., 2021] which integrate mollifiers from the literature [Sch¨utt et al., 2018a;

Klicpera et al., 2020b] to the key element of the self-attention.

2 Method

A molecule is defined by the atomic numbers z = {z1, . . . , zN} ∈ Z+, which serve to identify each type of atom, and the three dimensional positions X = {x1, . . . , xN} ∈ R3 of the N atoms composing it. Molecular predictions must satisfy fundamental symmetries and invariance of physical laws such as invariance to rigid spatial transformation (rotation and translation) and permutation (atoms of the same type are indistinguishable). Therefore, the positional input is transformed to interatomic Euclidean distances D = {dij}Ni,j=1 where dij = ∥xi -xj∥2 for rigid transform invariance, while permutation invariance is obtained via equal initial atomic representations of identical particles.

In this work, we design a parameterized deep neural network fθ for scalar regression of properties p ∈ R such that fθ : {z,D} → R. We do not include any further auxiliary features in the input (e.g. adjacency, bonds type, hybridization [Gilmer et al., 2017; Maziarka et al., 2020], DFT solver [Unke and others, 2019; Qiao and others, 2020]), neither any kind of knowledge adaptation from quantum physics/chemistry [Unke and others, 2019; Klicpera et al., 2020b; Klicpera et al., 2020a; Qiao and others, 2020]. Figure 1 depicts the proposed architecture. The positional encoding provides an initial geometry aware embedding of the atoms while the self-attention mechanism enables the accurate learning of the molecule geometry as well as the determination of the complex geometric interactions that are modeled in order to perform the regression task.

Transformer Transformer was introduced by [Vaswani et al., 2017] as a novel, attention-based building block for machine translation. The input sequence is first embedded into a high-dimensional space, coupled with positional embedding for each element. The embeddings are then propagated through multiple normalized self-attention and feed-forward blocks.

The self-attention mechanism introduced by Transformers is based on a trainable associative memory with (key, value) vector pairs where a query vector q ∈ Rd is matched against a set of k key vectors using scaled inner products as follows

where Q ∈ RN×d , K ∈ Rk×d and V ∈ Rk×d represent the packed N queries, k keys and values tensors respectively. Keys, queries and values are obtained using linear transformations of the sequence' elements. A multi-head self-attention layer is defined by extending the self-attention using h attention heads, i.e. h self-attention functions applied to the input, reprojected to values via a dh ×D linear layer.

Geometric Positional Encoding The initial embedding of the molecule is based solely on the atoms' type and thus is unable to differentiate similar atoms since the molecule geometry is omitted. The original Transformer's positional encoding module aims to transfer a measure of proximity of the sequence elements to the initial embedding. In our case, since the input is defined as a set rather than a sequence, the positional encoder needs to be adapted in order to provide a geometry-aware initial embedding. Here we propose to use the pairwise inter-atomic distance matrix in order to bring positional information to each atom. For each atom we first embed its pairwise distance vector to a single scalar in order to keep the module invariant to the size of the molecule, and then project it to the initial embedding space. Formally, denoting the atom embedding Emb(zi) : Z+ → Rd and the pairwise distance matrix D ∈ R+n×n such that (D)ij := (Di)j = dij , we have

Here, yi denotes the obtained initial positioned embedding of atom zi, fpos : R+ → R denotes the mapping of the pairwise distance parameterized as a shallow neural network , and W ∈ Rd denotes the projection matrix of the atomic onedimensional embedding onto the embedding space. Figure 2 depicts the impact of the positional encoder on the initial embedding. As can be observed, the initial embedding does not differentiate between atoms of the same type, while positional encoding brings information about the geometry. The convergence plot in Figure 3 demonstrates that the proposed positional encoding allows the transformer to learn the molecular geometry and thus be able to predict properties of the molecule. We also compare our method with the positional encoder of [Dwivedi and Bresson, 2020] where, since in our setting no graph is given, we propose to compute the 15 (half the biggest molecule size) first Laplacian eigenmaps [Belkin and Niyogi, 2003] instead of the combinatorial Laplacian. Our approach allows to encode the global geometry of the molecule with respect to each atom, while the permutationinvariant aggregation part in Eq. (2) maintains the symmetry invariance of the embedding. This preconditioning enables faster convergence and slightly better performance.

Geometric Self-Attention In order to maintain the invariant properties of molecules, we propose to augment the initial positional encoding layer and import the pairwise information directly into the self-attention layer. This is a natural choice since the self-attention layer already computes pairwise similarity of the atoms' representations via the normalized inner product. The self-attention layer (Eq. 1) is then extended to

where D denote the pairwise distance matrix, and φω a parameterized, potentially learned, similarity function. Several formulations can be conceived, however every extension should satisfy the distance relevance assumption such that φω is a vanishing mapping with respect to distance Dij , such that for a given positive δ scalar for all Dij > δ we would have φω(Qi, Kj , Dij ) → 0.

These assumptions are certainly reminiscent of cut-off distance (radii) and of many inter-atomic formulations such as the Van der Waals force or the Axilrod-Teller-Muto potential [Axilrod and Teller, 1943; Muto, 1943].

Assuming a vanishing mapping ψω(D) is given and applied element-wise, several approaches have been proposed in order to extend the self-attention mechanism. One popular extension is performed inside the softmax function as in [Wang and others, 2020] such that

This extension has the ability to fulfill the relevance assumption requirement if the function ψω(D), because of the softmax function, assigns negative values to distant atoms (at the limit limd→∞ψω(d) = -∞). However, such a requirement can be hard to design or even to learn with parameterized representation. Another option is to extend the transformer outside of the normalization such that

In this vein and concurrent to our work, [Maziarka et al., 2020] suggested MAT, a transformer architecture where the self-attention mechanism is defined as ψω(D) = ωexp(-D), with hyper-parameter ω ∈ R. This approach clearly fails in fulfilling the assumption presented above since distant atoms still impact the self-attention via the softmax similarity.

We propose to directly multiply the distance relation by the similarity tensor as follows

with ⊙ denoting the Hadamard product. This way, the interatomic distance has a direct gating impact on the pairwise atomic contribution obtained from the embedding. Such representation is especially important for the augmentation scheme we present later.

Other multiplicative alternatives are possible, such as transferring the Hadamard product inside the softmax [Wu et al., 2021; Fuchs et al., 2020; Kwak et al., 2021]. However there is a computationally demanding and hard to train need to couple the input of ψ with the similarity tensor values in order to ensure the desired values of the softmax function, i.e., in that case ϕω(Qi,Kj,Dij) → 0 for large Dij implies that ψω(Dij) → ±∞ and that QTi Kj and ψω(Dij) are of opposite signs.)

Learning the Graph Geometry Many methods have struggled to model the interaction function ψω. From force field methods to the recent learning-based approach, there is a need to empirically redefine the Euclidean pairwise distance in order to satisfy physical experimentation and/or performance.

Many handcrafted methods adopt molecular mechanics approximations of un/bonded interactions (e.g. stretch, bending, electrostatic or Van der Waals energies [Christensen and others, 2020]) where molecular properties are obtained via the modification of the interatomic distance, generally using exponential mapping of the distance.

Neural network-based methods adopt a similar approach, where atomic embedding is obtained by modifying the pairwise distance using various learned handcrafted mollifiers, Gaussian radial basis functions, or complex basis in the corresponding function space [Sch¨utt et al., 2018a; Klicpera et al., 2020b; Qiao and others, 2020].

One of the most critical hyper-parameters present in every method is the cut-off distance parameter, connecting only atoms which lie within the cut-off sphere. This hyperparameter may also change according to the property to be predicted [Schutt et al., 2018b].

Here we propose to learn the pairwise metric and the molecule connectivity at each level of the Transformer. The transformation of the Euclidean distance coupled with the self-attention mechanism suggested above allows us to directly optimize the inter-atomic representation as well as the cut-off distance according to the prediction objective, removing cumbersome hyper-parameterization, and allowing soft and differentiable construction of the adjacency matrix learned in a self-adaptive fashion. The similarity function is now simply given by

where ψω : R → R is an element-wise learnable function.

We parameterize ψω as a shallow, fully connected neural network, and we further enforce the positiveness of the new similarity map by squaring the filtered distances. Transforming the (element-wise) inverse of the distance D-1 instead of D speeds up the training since ψω transforms an already vanishing function (i.e. the multiplicative inverse function).

In contrast with existing methods that use envelopes [Unke and others, 2019; Klicpera et al., 2020b] or distance mollifiers [Sch¨utt et al., 2018a], as well as cut-off parameters, all inducing a handcrafted graph connectivity, we optimally unify the learning of the pairwise metric and of the graph adjacency.

We present the learned metrics for several of the Transformer blocks in Figure 4. As can be seen, the metric obtained is both more complex and more abstract than monotonic or Gaussian functions used in previous works [Sch¨utt et al., 2018a; Unke and others, 2019; Maziarka et al., 2020]. It is interesting to notice that some of the obtained cut-off distances lie around values empirically set in other works (2 - 6 ˚A). The diversified filter bank demonstrates the dynamic graph connectivity the network learns via the induced masking of the similarity map.

Regularization via Molecule Augmentation Transformers are generally extremely large and over-parameterized models. Dropout layers commonly used in Transformers in order to avoid overfitting cannot be used in our setting, because of the permutation-invariance requirement as described in [Lee et al., 2019]. One of the most efficient techniques for reducing overfitting is data augmentation. However it is not straightforward to augment molecular data (elements of a set), especially not for regression tasks since modification of one atom type or its spatial positions has unpredictable effects on molecular properties.

While many augmentation methods modify each datum, Mixup strategies [Zhang and others, 2017] intend to create new data samples from pairs (or more) of data. Here, following our initial distance relevance assumption, we propose to extend the Mixup idea to molecules where a new data sample is obtained by creating a new system of two molecules positioned far apart.

In our mixup scheme, we constrain the property of the new system to be the sum of the properties of the two submolecules, even if the predicted property is intensive, i.e., unlike extensive properties it is not physically additive. The sum of the contribution of each atom in the output module allows the model to learn to disentangle the two distanced sub-systems and reduce overfitting.

Formally, for a given property p and given two centered molecules Mi,Mj such that Mk = {zk , Xk } we have

with a rotation matrix R ∈ SO(3), and T = t · 1, t ∈ R a spatial translation vector ensuring large enough distance between the molecules so that the interaction is null or negligible (e.g. t > 103 ˚A). Thus, based on inter-atomic distances, we want our model to be able to differentiate the two systems such that the target properties p(Mi) and p(Mj) sum, thus

An illustration of the augmentation scheme is given in Figure 5. We note that methods relying on cut-off distances cannot use this augmentation strategy in a straightforward manner since the cut-off parameters would automatically separate the two molecules. Also, this approach can be extended to more than just pairs of atoms. The main drawback of this method is an increase in training time since the created molecules can be up to twice as large as the largest molecule of the dataset. However, this augmentation method enables a significant improvement of generalization, even on relatively small datasets, as demonstrated in the Discussion section.

3 Experiments

The experimental setup including the architecture details and the training procedure is provided in the Appendix 1 .

QM9 The popular QM9 dataset [Ramakrishnan et al., 2014] contains 130, 831 molecules with up to 9 atoms of the type C,N,O, and F saturated with hydrogen atoms in their equilibrium geometries, with chemical properties computed with DFT solvers. Following previous work, we split the dataset to 110, 000, 10, 000 and 10, 831 molecules for the training, validation and testing sets respectively. We use the atomization energy for U0, U, H, and G.

1 https://arxiv.org/pdf/2110.13721.pdf

In Table 1 (left) we report the mean absolute error (MAE) on all QM9 targets and compare it to the state-of-the-art models SchNet [Sch¨utt et al., 2018a], PhysNet [Unke and others, 2019], MGCN [Lu et al., 2019], Cormorant [Anderson et al., 2019], and DimeNet [Klicpera et al., 2020b]. We also compare with concurrent molecular data Transformers, RMAT[Maziarka et al., 2021], 3D-T [Wu et al., 2021], SE(3)-T [Fuchs et al., 2020] and GeoT [Kwak et al., 2021].

The method outperforms or is similar to state-of-the-art methods for most of the properties and surpasses Dimenet by 5.22% average performance ratio, and other Transformers methods by large margins. In contrast to other works [Schutt et al., 2018b; Klicpera et al., 2020a], the same model and training procedure were applied for all properties. The proposed framework does not require tuning of physical hyperparameters or chemical approximations, making it a true endto-end framework.

The recent SOTA work DimeNet++ [Klicpera et al., 2020a] substantially outperforms DimeNet (by 9% in average) with careful initialization and architectural modifications, and outperforms our framework by 0.78% (average performance ratio). We believe that a similar thorough architecture search may have similar impact on the proposed approach, depending on available computational resources since Transformers are computationally intensive. It is also important to notice that the proposed method does not take into account computationally heavy bond angles between triplets of atoms (i.e. potentially inducing cubic complexity) as in many recent frameworks.

MD17 We use MD17 [Chmiela et al., 2017] to test model performance in molecular dynamics simulations. The goal of this benchmark is to predict, for eight small organic molecules, the Cartesian atomic forces acting on each atom due to the overall potential energy. A separate model is to be trained for each molecule, in order to provide accurate individual predictions. We test our model in the challenging 1000 training samples setting [Chmiela et al., 2018]. The original training objective is extended to molecular dynamics predictions by backpropagating to the atom coordinates X as follows

where F denotes the nuclear three-dimensional Cartesian forces to be predicted and ρ is the forces' loss coefficient. In our experiments ρ is set to 103 , the augmentation scheme is extended straightforwardly to forces (i.e. concatenation) due to their translation invariance, and ReLU activations are replaced with GELU non-linearities to enforce twice continuous differentiability. As shown in Table 1 (right) our framework sets or reaches SOTA performances even when the training size remains extremely small, a challenging setting for Transformer models. It demonstrates our method's flexibility and its ability to generalize to other tasks and datasets.

Transformer based architectures are to the right of the vertical line.

3.1 Discussion

Geometric Self-attention Analysis Figure 6, presents typical impact of the proposed self-attention mechanism on the similarity map. As can be seen, different filters applied to the pairwise distance have a major impact on the similarity matrix and redefine the adjacency matrix at each level of the

network. One can observe the drastic impact (color flipping)

of the learned metric on the similarity map.

Comparison and Ablation Studies Our ablation study compares typical impact of the different self-attention modules. We present the convergence curves of our method from Eq. (7), the concurrent Transformers based MAT method [Maziarka et al., 2020] from Eq. (5), the MAT method with a learned metric ψω, and the sum self-attention from Eq. (4) [Wang and others, 2020] with learned metric; we refer to the last method as SUM SA. The compared networks and the training procedure are exactly the same, except for the aforementioned self-attention equation itself and the distance mapping module. The results are presented in Figure 7 (left). As can be seen, the MAT architecture presents the worse convergence, while the metric learning module significantly improves the performance. The proposed self-attention mechanism greatly surpasses all other architectures.

Finally, we present the typical effect of the data augmentation procedure on the generalization of the network during training. Figure 7 (right) presents the impact of the augmentation on the MAE convergence of the model for both an extensive and intensive property, namely U0 and µ. As can be seen, in both cases, when applying augmentation, the generalization gap is extremely reduced while the validation loss is much lower. The gain for the property µ is 175% in terms of testing MAE and 75% for U0.

Figure 7: Left: Comparison of validation MAE losses between MAT, MAT with the proposed learned metric (MAT+LM), the summed self-attention mechanism with learned metric (SUM SA +LM), and our method. Right: The impact of data augmentation on the MAE convergence and generalization of the network (val=validation) for U0 and µ. Continuous lines are consistently below their corresponding dashed ones.

4 Conclusion

We introduce a new Transformer architecture and a training scheme for molecular predictions. The proposed model allows effective representation of interactions based solely on pairwise distances. The graph geometry and connectivity can be learned in a soft fashion by the network via the geometric self-attention module. We further propose a molecular data augmentation procedure based on mixing strategies, which leads to a clear improvement in generalization for the proposed scheme. Our results indicate that our method is the first Transformer, as far as we can ascertain, that is able to model molecular data successfully without requiring any assumptions on the underlying physical model or involving complex geometric priors. We believe the advent of new datasets and new transformer architectures will allow the development of more efficient models also less prone to overfitting, making it a tool of predilection for the analysis of molecular data.

Acknowledgments

This project has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant ERC CoG 725974). The contribution of the first author is part of a PhD thesis research conducted at Tel Aviv University.