Relative Molecule Self-Attention Transformer

1 INTRODUCTION . Predicting molecular properties is central to applications such as drug discovery or material design . Without accurate prediction of properties such as toxicity , a promising drug candidate is likely to fail clinical trials . Many molecular properties can not be feasibly computed ( simulated ) from first principles and instead have to be extrapolated , from an often small experimental dataset ( Chan et al. , 2019 ; Bender & Cortés-Ciriano , 2021 ) . The prevailing approach is to train a machine learning model such a random forest ( Korotcov et al. , 2017 ) or a graph neural network ( Gilmer et al. , 2017 ) from scratch to predict the desired property for a new molecule . Machine learning is moving away from training models purely from scratch . In natural language processing ( NLP ) , advances in large-scale pretraining ( Devlin et al. , 2018 ; Howard & Ruder , 2018 ) and the development of the Transformer ( Vaswani et al. , 2017 ) architecture have culminated in large gains in data efficiency across multiple tasks ( Wang et al. , 2019a ) . Instead of training models purely from scratch , the models in NLP are commonly first pretrained on a large unsupervised corpora . The chemistry domain might be at the brink of an analogous revolution , which could be especially transformative due to the high cost of obtaining large experimental datasets . In particular , recent work has proposed Molecule Attention Transformer ( MAT ) , a Transformer-based architecture adapted to processing molecular data ( Maziarka et al. , 2020 ) and pretrained using self-supervised learning for graphs ( Hu et al. , 2020 ) . Several works have shown further gains by improving network architecture or the pretraining tasks ( Chithrananda et al. , 2020 ; Fabian et al. , 2020 ; Rong et al. , 2020 ) . However , pretraining has not yet led to such transformative data-efficiency gains in molecular property prediction . For instance , non-pretrained models with extensive handcrafted featurization tend to achieve very competitive results ( Yang et al. , 2019a ) . We reason that architecture might be a key bottleneck . In particular , most Transformers for molecules do not encode the three dimensional structure of the molecule ( Chithrananda et al. , 2020 ; Rong et al. , 2020 ) , which is a key factor determining many molecular properties . On the other hand , performance has been significantly boosted in other domains by enriching the Transformer architecture with proper inductive biases ( Dosovitskiy et al. , 2021 ; Shaw et al. , 2018 ; Dai et al. , 2019 ; Ingraham et al. , 2021 ; Huang et al. , 2020 ; Romero & Cordonnier , 2021 ; Khan et al. , 2021 ; Ke et al. , 2021 ) . Motivated by this perspective , we methodologically explore the design space of the self-attention layer , a key computational primitive of the Transformer architecture , for molecular property prediction . In particular , we explore variants of relative self-attention , which has been shown to be effective in various domains such as protein design and NLP ( Shaw et al. , 2018 ; Ingraham et al. , 2021 ) Our main contribution is a new self-attention layer for molecular graphs . We tackle the aforementioned issues with Relative Molecule Attention Transformer ( R-MAT ) , our pre-trained transformer-based model , shown in Figure 1 . We propose Relative Molecule Self-Attention , a novel variant of relative self-attention , which allows us to effectively fuse distance and graph neighbourhood information ( see Figure 2 ) . Our model achieves state-of-the-art or very competitive performance across a wide range of tasks . Satisfyingly , R-MAT outperforms more specialized models without using extensive handcrafted featurization or adapting the architecture specifically to perform well on quantum prediction benchmarks . The importance of representing effectively distance and other relationships in the attention layer is evidenced by large performance gains compared to MAT . An important inspiration behind this work was to unlock the potential of large pretrained models for the field , as they offer unique long-term benefits such as simplifying machine learning pipelines . We show that R-MAT can be trained to state-of-the-art performance with only tuning the learning rate . We also open-source weights and code as part of the Huggingmolecules ( Gaiński et al. , 2021 ) package . 2 RELATED WORK Pretraining coupled with the efficient Transformer architecture unlocked state-of-the-art performance in molecule property prediction ( Maziarka et al. , 2020 ; Chithrananda et al. , 2020 ; Fabian et al. , 2020 ; Rong et al. , 2020 ; Wang et al. , 2019b ; Honda et al. , 2019 ) . First applications of deep learning did not offer large improvements over more standard methods such as random forests ( Wu et al. , 2018 ; Jiang et al. , 2021 ; Robinson et al. , 2020 ) . Consistent improvements were enabled by more efficient architectures adapted to this domain ( Mayr et al. , 2018 ; Yang et al. , 2019a ; Klicpera et al. , 2020 ) . In this spirit , our goal is to further advance modeling for any chemical task by redesigning self-attention for molecular data . Encoding efficiently the relation between tokens in selfattention has been shown to substantially boost performance of Transformers in vision , language , music and biology ( Shaw et al. , 2018 ; Dai et al. , 2019 ; Ingraham et al. , 2021 ; Huang et al. , 2020 ; Romero & Cordonnier , 2021 ; Khan et al. , 2021 ; Ke et al. , 2021 ) . The vanilla selfattention includes absolute encoding of position , which can hinder learning when the absolute position in the sentence is not informative.1 Relative positional encoding featurizes the relative distance between each pair of tokens , which led to substantial gains in the language and music domains ( Shang et al. , 2018 ; Huang et al. , 2020 ) . However , most Transformers for the chemical domain predominantly used no positional encoding in the self- attention layer ( Chithrananda et al. , 2020 ; Fabian et al. , 2020 ; Rong et al. , 2020 ; Wang et al. , 2019b ; Honda et al. , 2019 ; Schwaller et al. , 2019 ) , which gives rise to similar issues with representing relations between atoms . We directly compare to ( Maziarka et al. , 2020 ) , who introduced first self-attention module tailored to molecular data , and show large improvements across different tasks . Our work is also closely related to ( Ingraham et al. , 2021 ) that used relative self-attention fusing three dimensional structure with positional and graph based embedding , in the context of protein design . 1This arises for example when input is an arbitrary chunks of the text ( Huang et al. , 2020 ) ( e.g . in the next sentence prediction task used in BERT pretraining ) . 3 RELATIVE MOLECULE SELF-ATTENTION . 3.1 MOLECULAR SELF-ATTENTIONS . We first give a short background on how prior work on adapting self-attention for molecular data , and point out their potential shortcomings . Text Transformers Multiple works have applied directly the Transformer to molecules encoded as text using the SMILES representation ( Chithrananda et al. , 2020 ; Fabian et al. , 2020 ; Wang et al. , 2019b ; Honda et al. , 2019 ; Schwaller et al. , 2019 ) . SMILES is a linear encoding of a molecule into a string of characters according to a deterministic ordering algorithm ( Weininger , 1988 ; Jastrzębski et al. , 2016 ) . For example , the SMILES encoding of carbon dioxide is C ( =O ) =O . Adding a single atom can completely change the ordering of atoms in the SMILES encoding . Hence , the relative positions of individual characters are not easily related to their proximity in the graph or space . This is in contrast to natural language processing , where the distance between two words in the sentence can be highly informative ( Shaw et al. , 2018 ; Huang et al. , 2020 ; Ke et al. , 2021 ) . We suspect this makes the use of self-attention in SMILES models less effective . Another readily visible shortcoming is that the graph structure and distances between atoms of the molecule are either completely encoded or entirely thrown out . Graph Transformers Several works have proposed Transformers that operate directly on a graph ( Maziarka et al. , 2020 ; Rong et al. , 2020 ; Nguyen et al. , 2019 ) . The GROVER and the U2GNN models take as input a molecule encoded as a graph ( Rong et al. , 2020 ; Nguyen et al. , 2019 ) . In both of them , the self-attention layer does not have a direct access to the information about the graph . Instead , the information about the relations between atoms ( existence of a bond or distance in the graph ) is indirectly encoded by a graph convolutional layer that is run in GROVER within each layer , and in U2GNN only at the beginning . Similarly to Text Transformers , Graph Transformers also do not take into account the distances between atoms . Structured Transformer introduced in ( Ingraham et al. , 2021 ) uses relative self-attention that operates on amino-acids in the task of protein design . Self-attention proposed by ( Ingraham et al. , 2021 ) , similarly to our work , provides the model with information about the three dimensional structure of the molecule . As R-MAT encodes the relative distances between pairs of atoms , Structured Transformer also uses relative distances between modeled amino-acids . However , it encodes them in a slightly different way . We incorporate their ideas , and extend them to enable processing of molecular data . Molecule Attention Transformer Our work is closely related to Molecule Attention Transformer ( MAT ) , a transformer-based model with self-attention tailored to processing molecular data ( Maziarka et al. , 2020 ) . In contrast to most of the aforementioned models , MAT incorporates the distance information in its self-attention module . MAT stacks N Molecule Self-Attention blocks followed by a mean pooling and a prediction layer . For a D-dimensional state x ∈ RD , the standard , vanilla self-attention operation is defined as A ( x ) = Softmax ( QKT√ dk ) V , ( 1 ) where Q = xWQ , K = xWK , and V = xWV . Molecule Self-Attention extends Equation ( 1 ) to include additional information about bonds and distances between atoms in the molecule as A ( x ) = ( λa Softmax ( QKT√ dk ) + λd g ( D ) + λgA ) V , ( 2 ) where λa , λd , λg are the weights given to individual parts of the attention module , g is a function given by either a softmax , or an element-wise g ( d ) = exp ( −d ) , A is the adjacency matrix ( with A ( i , j ) = 1 if there exists a bond between atoms i and j and 0 otherwise ) and D is the distance matrix , where D ( i , j ) represents the distance between the atoms i and j in the 3D space . Self-attention can relate input elements in a highly flexible manner . In contrast , there is little flexibility in how Molecule Self-Attention can use the information about the distance between two atoms . The strength of the attention between two atoms depends monotonically on their relative distance . However , molecular properties can depend in a highly nonlinear way on the distance between atoms . This has motivated works such as ( Klicpera et al. , 2020 ) to explicitly model the interactions between atoms , using higher-order terms .

Summary: This paper proposes a relative self-attention layer for the Transformer model. The relative self-attention of two atoms consists of their relative distance, their shortest path distance in the molecular graph, and their physiochemical. The proposed relative molecule attention Transformer can be first pretrained with a contextual property prediction task and then a graph-level prediction task. The pretrained Transformer can be finetuned on downstream molecular property prediction tasks and achieves excellent performance.