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: The paper presents an approach to represent 3D atomic structure in a machine learning context.   The paper distinguishes between representations for "small" and "large" molecules based on a metric that takes the spatial extension and number of atoms in the molecule into account. Structures of small molecules are represented by an interatomic distance map. Large molecules are represented by what the authors refer to as a "sinusoidal function-based absolute position encoding method". In both cases, a transformer architecture is used on top of this initial representation. The authors further introduce a subsampling procedure to select a subset of points/ atoms and aggregate information.   The authors demonstrate quantitative results on datasets relating to 1) small molecule property prediction (QM datasets), 2) protein-ligand binding (PDBbind dataset), and 3) a dataset from material science on metalorganic compounds (CoRE-MOF dataset)