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: The paper proposes a new transformer network architecture to pre-train the molecule datasets. Based on the Molecule Attention Transformer, the proposed model, R-MAT, incorporates a few handcrafted features into the self-attention layer of a transformer architecture. The features can incorporate distances between atoms from multiple perspectives. The experimental results show that the pre-trained model is useful to predict various properties of molecules.