3D-Transformer: Molecular Representation with Transformer in 3D Space 

1 INTRODUCTION . Spatial structures are among the most crucial factors to decide molecular properties and understand their principles of action in the physical world . For example , 3D structures of proteins provide valuable information for inferring biological interventions , such as structure-based drug development and targeted mutagenesis ( Senior et al. , 2020 ; Jumper et al. , 2021 ; Baek et al. , 2021 ) . In chemistry , zeolites show obvious differences in separation properties caused by subtle changes in their 3D geometric compositions ( Chai et al. , 2020 ; Pfriem et al. , 2021 ) . Apart from that , in the pharmaceutical industry , the same compounds can have different 3D structures , resulting in different solubility ( Zhang et al. , 2017 ) . To sum up , capturing 3D spatial structures is essential to accurately forecast molecular properties . Based on these facts , researchers have studied molecular representation learning techniques ( Rao et al. , 2019 ) to include 3D spatial information ( Zhavoronkov et al. , 2019 ) . The dominant 3D molecular models are Graph Neural Networks ( GNNs ) and 3D Convolutional Neural Networks ( 3DCNNs ) ( Derevyanko et al. , 2018 ; Pagès et al. , 2019 ; Townshend et al. , 2019 ) . GNNs create edges by using either chemical bonds or finding the neighbors of each node within a distance cutoff ( Zhang et al. , 2020b ) . They encode pairwise connectivity of atoms and require running multiple hops for an atom to reach to another . 3DCNNs encode translational and permutational symmetries , but need to stack deep layers to build direct connections between distant regions , incurring significant computational costs . In contrast , Transformers rely on the self-attention mechanism to capture long-term dependencies in parallel ( Hernández & Amigó , 2021 ) . Meanwhile , Equivariant Neural Networks ( ENNs ) ( Thomas et al. , 2018 ) have emerged as a new class of methods , where geometric transformations of their inputs lead to well-defined transformations of outputs . Some ENNs adopt Transformers as the backbone but fail to surmount the intrinsic drawbacks of this architecture , including its insensibility to local patterns among non-uniformly distancing atoms and its ineffi- ciency to aggregate atom features . Some other Transformer-based methods have been proposed to fuse distance and graph neighbourhood information ( Maziarka et al. , 2020 ; 2021 ) . However , they take no consideration of employing motifs , which are frequently-occurring substructures in molecules and can be leveraged to uncover global graph properties . In this work , we present the Molformer on the basis of all preceding analysis . For the sake of injecting chemical domain knowledge , we construct a motif-template vocabulary based on functional groups and adopt trainable motif embeddings to maintain the semantic meanings of those essential substructures . Then with both motifs and atoms as input , Molformer operates on a fully-connected graph with direct connections between remote regions ( Veličković et al. , 2017 ; Joshi , 2020 ) , which reduces computational burden of multi-hop GNNs and stacked 3DCNNs . However , this characteristic limits Molformer ’ s capacity in exploiting local structures and leads to poor generalization in unseen cases ( Qi et al. , 2017 ) . Therefore , we propose a Multi-scale Self-Attention ( MSA ) module to recognize fine-grained patterns from neighborhoods . Moreover , we introduce a roto-translation invariant Convolutional Position Encoding ( CPE ) to depict position relationships among atoms and their adjacencies . After that , to retain a comprehensive representation of the entire molecule , we propose an Attentive Farthest Point Sampling ( AFPS ) module that selects important atoms with the assistance of the attention score map . To summarize , our contributions are as follows : • To the best of our knowledge , we are the foremost to incorporate motifs with knowledge of functional groups into a Transformer architecture for 3D molecular representation learning . • We propose a novel MSA to extract local patterns , a roto-translation invariant CPE method to encode relative distance at a linear computational time cost , and a simple yet effective downsampling algorithm to gather molecular representations . • We show significant improvements on several benchmarks in three domains . Code and all datasets are available at https : //github.com/smiles724/Molformer . 2 PRELIMINARIES . Problem Definition . A molecule S = ( E , P ) has N atoms and C atom classes , where E = { e1 , ... , eN } ∈ RN×C contains the one-hot atom representations and P = { p1 , ... , pN } ∈ RN×3 contains the 3D coordinates of each atom . Each one-hot ei can be converted to a dense vector xi = eiW E , with xi ∈ Rdmodel and WE ∈ RC×dmodel being the embedding matrix . The 3D coordinates of the atom i is a three-dimensional vector pi = [ pxi , p y i , p z i ] . A representation learning model f acts on S , obtaining its representation r = f ( S ) . Then r is forwarded to a prediction model g and attain the prediction of a biochemical property ŷ = g ( r ) . Self-attention Mechanism . The Transformer ( Vaswani et al. , 2017 ) has become very successful due to its core component , self-attention . Given a set of input features { xi } i=1 , ... , N , the standard dot-product attention layer is as the following : qi = fQ ( xi ) , ki = fK ( xi ) , vi = fV ( xi ) , aij = qik T j / √ dk , zi = N∑ j=1 σ ( aij ) vj ( 1 ) where { fQ , fK , fV } are embedding transformations , and { qi , ki , vi } are respectively the query , key , and value vectors with the same dimension dk . aij is the attention that the token i pays to the token j. σ denotes the Softmax function and zi is the output embedding of the token i . This formula conforms to a non-local network ( Wang et al. , 2018 ) , indicating its inability to capture fine-grained patterns in a local context . Position Encoding . Self-attention is invariant to permutation of the input ( Dufter et al. , 2021 ) , and position encoding ensures that the Transformer will reveal positional information . Position encoding methods can be either based on absolute positions or relative distances . The former takes the raw position information as input and is sensitive to spatial transformations . The latter manipulates the attention score by incorporating relative distances ( Guo et al. , 2020a ; Pan et al. , 2021 ) : aij = qik T j / √ dk + fPE ( pi−pj ) , where fPE ( · ) is the position encoding function and is translation invariant . The rotation invariance can be further accomplished by taking a L2-norm ||pi − pj ||2 ( Chen et al. , 2019b ) . 3 MOLFORMER . Molformer is based on the architecture of Transformer but adopts several significantly different and novel components ( see Figure 1 ) . First , a vocabulary of motif templates is constructed on the basis of functional groups and we extract all available motifs from each molecule . Then both atoms and motifs acquire their corresponding embeddings and are forwarded into L feature learning blocks . Each block consists of a convolutional position encoding , a multi-scale self-attention , and a feedforward network . After that , an attentive subsampling method is utilized to adaptively aggregate the molecular presentation , which is later fed into a predictor to forecast properties in a broad range of downstream tasks . 3.1 TRAINABLE MOTIF-BASED EMBEDDING . Motifs are frequently-occurring substructure patterns as well as the building blocks of complex molecular structures . They usually maintain semantic meanings and have great expressiveness of the biochemical characteristics of the whole molecule ( Zhang et al. , 2020a ) . In the chemical community , researchers have developed a set of standard criterion to recognize motifs with essential functionalities in molecules ( Milo et al. , 2002 ) . Despite that , few of prior studies directly incorporate those informative motifs into their model architectures . To fill this gap , we define a series of momentous substructures using external domain knowledge , and introduce a trainable motif embeddings method to fully exploit them in our Molformer . To begin with , all motifs are first extracted according to the motif vocabulary , which is built by functional groups . Practically , we rely on RDKit ( Landrum , 2013 ) to draw them from the SMILES ( Weininger , 1988 ) representation of each molecule . We assume M motifs { m1 , ... , mM } are detected in the molecule S , and each motif mi contains a certain number of at least two atoms . Then we regard each kind of motif as a new type of token and append them to the input . Therefore , the input for our Molformer becomes { x1 , ... , xN , xm1 , ... , xmM } , where xmi is obtained through an learnable embedding matrix WM ∈ RC′×dmodel and C ′ denotes the number of motif categories . As for the position of each motif , we adopt a weighted sum of the 3D coordinates of its component atoms as pmi = ∑ xi∈mi ( wi∑ xi∈mi wi ) · pi , where wi are the atomic weights . Our approach requires the model to automatically learn a customized embedding for each motif template through backpropagations , which follows a data-driven pattern . In some data-sufficient tasks , its greatest potential can be unlocked and those motif embeddings can be well trained . Nevertheless , in the case of few-shot learning or small datasets , each category of motif might only appear rare times . Those embeddings are not fully tuned and can be extremely biased and noisy , which will do little helps to the ultimate property prediction . 3.2 CONVOLUTIONAL POSITION ENCODING . To enable roto-translation invariance and take fully advantage of geometric information , instead of adding a term of fPE ( pi − pj ) , we propose a CPE that applies a convolutional operation to the interatomic distance D ∈ RN×N : Acov = Conv2d ( D ) A , ( 2 ) where A = [ ai , j ] i , j=1 , ···N ∈ RN×N is the attention matrix , Conv2d ( · ) denotes a 2D shallow convolutional network with a kernel size of 1× 1 , and is the element-wise product . With multi-headed self-attention , Acov is expanded in the sense that Acov ∈ RH×N×N , and Conv2d ( · ) has H output channels . The CPE method induces O ( N ) convolution operations on each atom and can drastically reduce training time when the number of atoms is very large ( Wu et al. , 2021 ) . 3.3 MULTI-SCALE SELF-ATTENTION . The self-attention mechanism in the Transformer is good at capturing global data patterns but ignores local context ( Guo et al. , 2020a ) . Exploiting local context has proven to be important for 3D spatial data such as 3D point clouds ( Qi et al. , 2017 ) . Therefore , we impose a distance-based constraint in self-attention in order to extract multi-scaled patterns from both local and global contexts . Guo et al . ( 2020b ) propose to use integer-based distance to limit attention to local word neighbors , which can not be used in molecules . This is because different types of molecules have different densities and molecules of the same type have different spatial regularity , which results in the nonuniformity of interatomic distances . Normally , small molecules have a mean interatomic distance of 1-2 Å ( Angstrom , 10−10m ) , which is denser than large molecules like proteins with approximately 5 Å on average . To address that , we design a new multi-scale methodology to robustly capture details . Specifically , we mask atoms beyond a certain distance τs ( a real number as opposed to an integer in Guo et al . ( 2020b ) ) at each scale s. We denote dij = ||pi−pj ||2 as the Euclidean distance between the i-th and j-th atom . The attention calculation is modified as : aτsij = qik T j · 1 { dij < τs } √ dk , zτsi = N∑ j=1 σ ( aτsij ) vj , ( 3 ) where 1 { dij < τs } is the indicator function . For small molecules , Equation 3 can be complementally combined with Equation 2 . Then features extracted from S different scales { τs } s=1 , ... , S as well as the informative global feature are concatenated together to form a multi-scale representation , denoted by z′i = z τ1 i ⊕ ... ⊕z τS i ⊕z global i ∈ R ( S+1 ) dk . After that , z′i is forwarded into a multi-layer perceptron to be compressed as z′′i with the original dimension dk .

Summary: This paper proposes 3D-Transformers, which is a variant of Transformer that incorporates the 3D spatial information. The authors develop a multi-scale attention module, an adaptive position encoding module, and a farthest sampling scheme for the 3D-Transformer. The authors validate the performance of 3D-Transformer on three different domains.