Relative Molecule Self-Attention Transformer
Authors: Lukasz Maziarka
Dawid Majchrowski
Tomasz Danel
Piotr Gaiński
Jacek Tabor
Igor T. Podolak
Pawel Morkisz
Stanislaw Kamil Jastrzebski
Venue: ICLR 2022 Submitted (2022)
Keywords: molecular property prediction
transformer-based methods
graph neural networks
self-supervised learing

Abstract: Self-supervised learning holds promise to revolutionize molecule property prediction - a central task to drug discovery and many more industries - by enabling data efficient learning from scarce experimental data. Despite significant progress, non-pretrained methods can be still competitive in certain settings. We reason that architecture might be a key bottleneck. In particular, enriching the backbone architecture with domain-specific inductive biases has been key for the success of self-supervised learning in other domains. In this spirit, we methodologically explore the design space of the self-attention mechanism tailored to molecular data. We identify a novel variant of self-attention adapted to processing molecules, inspired by the relative self-attention layer, which involves fusing embedded graph and distance relationships between atoms. Our main contribution is Relative Molecule Attention Transformer (R-MAT): a novel Transformer-based model based on the developed self-attention layer that achieves state-of-the-art or very competitive results across a~wide range of molecule property prediction tasks. 

Under review as a conference paper at ICLR 2022
RELATIVE FORMER
MOLECULE
SELF-ATTENTION
TRANS-
Anonymous authors Paper under double-blind review
ABSTRACT
Pretraining using self-supervised learning holds promise to revolutionize molecule property prediction a central task to drug discovery and many more industries by enabling data efficient learning from scarce experimental data. However, despite significant progress, non- con-prerrained methods can be still competitive certain settings. We reason that architecture might be a key bottleneck. In particular, enriching the backbone architecture with domain-speciff inductive biases has been key for the success of self-supervised learning in other domains. Inspired by this, we methodologically explore the design space of the self-attention mechanism for molecular data. Our main contribution is Relative Molecule Attention Transformer (R-MAT): a novel Transformer-basee model that achieves state-of-the-art or very competitive results across wide range of molecule property prediction tasks. Relative Molecule Attention Transformer uses a novel self-attention variant that can be readily incorporated into future models for processing molecular data.
INTRODUCTION
Predicting molecular properties is central to applications such drug discovery material design. Without accurate prediction of properties such as toxicity, a promising drug candidate is likely to fail clinical trials. Many molecular properties cannot 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- CCriinn (021). The prevailing approach is to train machine learning model such a random forest (Korotcov et al., 2017) or graph neural network (Gilmer et al. 201 7) from scratch to predict the desired property for 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. [018] Howard Ruder (0188) 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 training models purely from scratch, the models in NLP are commonly first pretrained
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 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 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 fcaturization tend to achieve very competitive results (Yang et al., 2019a). We reason that architecture might be a key bot- tleneck. In particular, most Transformers for molecules not encode the three dimensional structure of the molecule (Chithrananda et al.) [020) Rong et al.) (020). which is a key factor determining many molecular properties. On the other hand, performance has been significantly boosted other domains by enriching the Transformer architecture with proper inductive biases (Dosoviskkiy et al. |021 Shaw
let al.) 2018, Dai et al. 2019 Ingraham et al.] 2021 Huang et al. 2020; Romero & Cordonnier 2021 Khan al. 202 Ke et al. 2021). Motivated by this perspective,
methodologically explore the design space of the self-attention layer, a key computational primitive of the Transformer architecture, for molecular property prediction. In particular, explore variants of relative self-ate self-attention, which

---

Under review as a conference paper at ICLR 2022
has been shown to be effective in various domains such as protein design and NLP Ingraham et al. 2021)
(Shaw et al.
2018
Our main contribution is 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 We propose Relative Molecule Self-Attention, 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 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 ffte-offthe-art performance with only tuning the learning rate. We also open-source weights and code as part of the Huggingmolecules (Gaiñski et al. 2021 package.
RELATED WORK
Prodiction
Fully Connected Layer
Pretraining coupled with the efficient Transformer archi- tecture unlocked state-of-the-an performance in molecule property prediction (Maziarka et al. 2020 Chithrananda
Astencion Pooling
Layer Norm
al. 2020 Fabian al. 2020 Rong 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., 202T Robinson et al.) 2020). Consistent im- provements 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 re- designing self-attention for molecular data.
Position Feed Forward
Laper Narm
Reletive Medeenlar Self- Attention
distance emboddinns
bond embeekdings
Encoding efficiently the relation between tokens in self- attention has been shown to substantially boost perfor- mance of Transformers in vision, language, music and biology (Shaw al., 2018 Dai al.] 2019 Ingraham et al. 2021]] Huang et al.] 2020 Romero & Cordonnier] 2021 Khan al. | [202]] Ke et al. 2021). The vanilla self- attention includes absolute encoding of position, which can hinder learning when the absolute position in the sen- tence is not informative Relative positional encoding featurizes the relative distance between each pair of to- kens, 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-
an inigibouriooo calldddings
Laper Narm
Fanbedding g Layer
Impat
Figure 1: Relative Molecule Attention Transformer uses a novel relative self- attention block tailored to molecule prop- erty prediction. It fuses three types of features: distance embedding, bond em- bedding and neighbourhood embedding.
attention layer (Chithrananda et al. 20200 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.) (02T) that used relative self-attention fusing three dimensional structure with positional and graph based embedding, in the context of protein design.
'This arises for example when input is arbitrary chunks of the text (Huang al. 2020 (e.g. in the next sentence prediction task used in BERT pretraining).
2

---

Under review as a conference paper at ICLR 2022
Dummy Node
Conjugated
1.5
15
1.5
2
3
1.5
1.5
2
2
H
Aromatic Ring
1.5
-II
3
15
R
F2
63
...
(a) neighbors
(b) bonds
(c) distances
The Relative Molecule Self-Attention layer is based on the following features: (a) neigh- bourhood embedding one-hot encodes graph distances (neighhourhood order) from the source node marked with arrow; (b) bond embedding one-hot encodes the bond order (numbers next to the graph edges) and other bond features for neighbouring nodes; (c) distance embedding uses radial basis functions to encode pairwise distances in the 3D space. These features are fused according to Equation
Figure 2:
RELATIVE MOLECULE SELF- SELF-ATTENTION
3.1 MOLECULAR SELF-ATTENTIONS
We first give 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.)
2019h Honda et al. 2019 Schwaller al. 2019). SMILES is linear encoding of molecule into
a string of characters according to a deterministic ordering algorithm [Weininger] 1988 Jastrzębski
al. 2016). For example, the SMILES encoding of carbon dioxide C(=0)=O.
Adding 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 al., 2020 Ke al.) (02T). 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 entirely thrown out.
Graph Transformers Several works have proposed Transformers that operate directly on graph (Maziarka et al. 2020 Rong al.) 2020 Nguyen et al.) 2019). The GROVER and the U2GNN models take as input a molecule encoded as graph (Rong et al. 2020 Nguyen et al.) 2019). In both of them, the self-attention layer does not have direct access to the information about the graph. Instead, the information about the relations between atoms (existence of bond or distance in the graph) indirectly encoded by a graph convolutional layer that is rum GROVER within each layer, and in U2GNN only at the beginning. Similarly to Text Transformers, Graph Transformers also do 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 atoms, Structured Transformer also uses relative distances between modeled amino-acids. However, it encodes them a slightly different way. We incorporate their ideas, and extend them to enable processing of molecular data.
3

---

Under review as a conference paper at ICLR 2022
Molecule Attention Transformer Our work is closely related to Molecule Attention Transformer (MAT), a transformer-bat model with self 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 mean pooling and prediction layer. For D-dimensional state E TR,, the standard, vanilla self-attention operation is defined as
QK? Vdz
A(x) = Softmax
V
(I)
where Q x W K x xW and V - x W Molecule Self-Attention extends Equation include additional information about bonds and distances between atoms in the molecule as
!
QK? vdx
A(x) Aa Softmax
Aag(D) + +>gA) v,
where Ao Ads 0g the weights given individual parts of the attention module, y is function given by either a softmax, or an element-wise g(d) exp(-d), A is the adjacency matrix (with (i.)) if there exists bond between atoms and j and otherwise) and D is the distance matrix, where D (i,)) represents the distance between the atoms i and in the 3D space.
Self-attention can relate input elements in highly flexible manner. In contrast, there is little flexibility in how Molecule Self- 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 highly nonlinear way the distance between atoms. This has motivated works such as (Klicpera al. 2020) to explicitly model interactions between atoms, using higher-order terms.
3.2 RELATIVE POSITIONAL. ENCODING
In natural language processing, vanilla self-attention layer does not take into account the positional information of the input tokens (i.e. if we permute the layer input, the output will stay the same). In order to add the positional information into the input data, vanilla transformer enriches it with encoding of the absolute position. On the other hand, relative positional encoding (Shaw et al. 2018) adds the relative distance between each pair of tokens, which leads to substantial gains in the learned task. In our work, we use relative self-attention to encode the information about the relative neighbourhood, distances and physicochemical features between all pairs of atoms in the input molecule (See Figure
3.3
ATOM RELATION EMBEDDING
Our core idea to improve Molecule Self-A Attentitt is to add flexibility in how it processes graph and dis- tance information. Specifically, we adapt positional relative encoding processing molecules (Shaw et al. 2018 Dai et al. 2019 Huang et al. 2020) Ke et al. 2021), which we note was already hinted (Shaw et al. 2018) high-level future direction. The key idea in these works enrich the self-attention block to efficiently represent information about relative positions of items in the input sequence.
What reflects the relative position of two atoms in molecule? Similarly to MAT, we delineate three inter-related factors: (1) their relative distance, (2) their distance in the molecular graph, and (3) their physiochemical relationship (e.g. are they within the same aromatic ring).
In the next step, we depart from Molecule Belf-Attention (Maziarka al. 2020) and introduce new factors to the relation embedding. Given two atoms, represented by vectors 2;, 2 E RP, we encode their relation using atom relation embedding bij E B This embedding will then be used in the self-attention module after projection layer. Next, we describe three components that are concatenated to form the embedding bij.
4

---

Under review as a conference paper at ICLR 2022
Neighbourhood embeddings First, encode the neighbourhood order between atoms as a dimensional one hot encoding, with information about how many other vertices are between nodes and j in the original molecular graph (see Figure 2 and Table from Appendix A
Distance embeddings As discussed earlier, we hypothesize that much more flexible represen- tation of the distance information should be facilitated in MAT. To achieve this, we use radial basis distance encoding proposed by (Klicpera et al. 2020
2 sin N:T d) en (d) = 2
d
where d is the distance between atoms, the predefined cutoff distance, E fI,.... Nemb) and Nemb is the total number of radial basis functions that we use. Then obtained numbers are passed to the polynomial envelope function m(d) 1 (p+1)(p+2) d + plp + 2)d p(p-1) 1) ge+2, with p in order to get the final distance embedding.
Bond embeddings Finally, we featurize each bond to reflect the physical relation between pairs of atoms that might arise from, for example, being part of the same aromatic structure in the molecule. Molecular bonds are embedded in as 7 dimensional vector following (Coley et al. 2017 (see Table F from Appendix [A). When the
atoms not connected by molecular bond, all dimensions are set to zeros. We note that while these features can be easily learned in pretraining, we hypothesize that this featurization might be highly useful for training R-MAT smaller datasets.
3.4 RELATIVE MOLECULE SELF-ATTENTION
Equipped with the embedding bij for each pair of atoms in the molecule, we now use it to define novel self-attention layer that we refer to Relative Molecule Self-Attention.
First, mirroring the key-query-valut design in the vanilla self-attention (c.f. Equation (T). we transform by into key and value specific vectors by. bK using two neural networks Φv and Φx- Each neural network consists two layers. A hidden layer, shared between all attention heads and output layer, that create separate relative embedding for different attention heads.
Consider Equation in index notation:
T (x,WV),
A(x); = Σ Softmax
Eij Vd=
where the unnormalized attention is Cig = (x; WKyT. By analogy, in Relative Molecule Self-Attentio,, we compute Eij as
eij (z,W (a, (2, + (z,
+ (
(Z,
(3)
vanilla self attention
content dependent positional bios for query
content dependent positional bins for key
globul content bias
global positional blas
where u. D' are trainable vectors. We then define Relative Molecule Self-Attention operation:
A Σ Softmax
Cij Vd.
WV + +b))
In other words, we enrich the self-attention layer with atom relations embedding, In the phase of attention weights calculation, we add content-dependent positional bias, global context bias and global positional bias DDai et al. 2019 Huang et al. 2020) (that are calculated based 6 the layer. Then, during calculation of the attention weighted average, we also include the information about the other embedding b/g-
5

---

Under review as a conference paper at ICLR 2022
3.5 RELATIVE MOLECULE ATTENTION TRANSFORMER
Finally, we use Relative Molecule Self-Attention to construct Relative Molecule Attention Trans- former (R-MAT). The key changes compared MAT are: (1) the use of Relative Molecule Self- Attention, (2) extended atom featurization, and (3) extended pretraining procedure. Figure Iillustrates the R-MAT architecture.
The input is embedded as a matrix of size N Naton x 36 where each atom of the input is embedded following (Coley al. 2017 Pocha et al. 2020), see Ta Appendix We process the input using N stacked Relative Molecule Self-Attention attention layers. Each attention layer is followed by position-wise feed-forward Network (similar as the classical transformer model ( (Vaswani et al. 2017 which consists of linear layers with a CeakyReLUU nonlinearity between them.
After processing the input using attention layers, we pool the representation into a constant-sizes vector. We replace simple mean pooling with attention- based pooling layer. After applying N self-attention layers, we use the following self-attention pooling (Lin et al. 2017) in order to get the graph-level embedding of the molecule:
P Softmax(W,, (anh(W,H")) Flatten(PH),
g
where H is the hidden state obtained from self-attention layers, Wi E PxD and W2 E RSP are pooling attention weights, with equal to the pooling hidden dimension and S equal to the number of pooling attention heads. Finally, the graph embedding g is then passed to the two layer MLP, with leaky-ReLU activation in order to make the prediction.
Pretraining We used two-step pretraining procedure. In the first step, network is trained with the contextual property prediction task proposed by (Rong et al.) 2020), where we mask not only selected atoms, but also their neighbours. The goal of the task is to predict the whole atom context. This task is much more demanding for the network than the classical masking approach presented by (Maziarkal let al,, (2020) since the network has to encode more specific information about the masked atom neighbourhood. Furthermore, the size of the context vocabulary is much bigger than the size of the atoms vocabulary in the MAT pretraining approach. The second task is graph-level prediction proposed by (Fabian et al. 2020) in which the goal to predict a set of real-valued descriptors of physicochemical properties. For more detailed information about the pretraining procedure and ablations, see Appendix B
Other details Similarly to (Maziarka al. 2020), we add an artificial dummy node to the input molecule. The distance of the dummy node to any other atom in the molecule is set to the maximal cutoff distance, and the edge connecting the dummy node with any other atom has its unique index (see index 5 in Tf Appendix [A). Moreover, the dummy node has its own index in the input atom embedding. We calculate distance information in the similar manner as (Maziarka et al. 2020). The 3D molecular conformations that are used to obtain distance matrices UFFOPTIMIZE
calculated using
function from the RDKit package (Landrum] 2016) with the default parameters. Finally, consider variant of the model extended with 200 rdkit features as in (Rong) et al. 2020). The features are concatenated to the final embedding and processed using a prediction MLP.
4
EXPERIMENTS
4.1
SMALL HYPERPARAMETER BUDGET
The industrial drug discovery pipelines focus on fast iterations of compound screenings and adjusting the models to new data incoming from the laboratory. We start by comparing this setting R-MAT to DMPNN (Yang al. 2019a). MAT (Maziarka et al. 2020) and GROVER (Rong et al. 2020). representative state-of-the thee---rrhherr art models on popular molecular property prediction tasks. We followed the evaluation in (Maziarka et al. 2020), where the only changeable hyperparameter is the learning rate, which was checked with 7 different values.
The BBBP and Estrogen-8 3 datasets scaffold splits, while all the other datasets use random splits. Splits were proposed by (Maziarka et al. 2020). For every dataset we calculate scores based on different splits, we report the mean test score based the hyperparameters that obtained the best.

---

Under review as a conference paper at ICLR 2022
validation score, in parentheses include the standard deviation. In this and the next experiments, we denote models extended with additional rdkit features (see Section (35) as GROVER and R-MAT adkät- More information about the models and datasets used in this benchmark are given in Appendix C.4.
Table 1 shows that R-MAT outperforms other methods in 3 out of 6 tasks. For comparison, also cite representative results of other methods from (Maziarka et al.) 2020). Satisfyingly, we observe marked improvement on the solubility prediction tasks (ESOL and FreeSolv). Understanding solubility depends to a large degree on a detailed understanding of spatial relationships between atoms. This suggests that the improvement performance might be related better utilization of the distance or graph information.
Table 1: Results on molecule property prediction benchmark from (Maziarka et al. 2020). We only tune the learning rate for models in the first group. First two datasets are regression tasks (lower is better). other datasets are classification tasks (higher is better). For reference, we include results for non-pretrained baselines (SVM, RF, GCN (Duvenaud et al., 2015), and DMPNN (Yang et al.) 2019a)) from (Maziarka et al. 2020 We also include SVMaak and RF two baseline methods with added rdkit features. Rank-plot for these experiments is in Appendix D.I
ESOL
FreeSolv
BBBP
Estrogen-3
MetStab
MetStabrigh
MAT
(.20)
(665)
(.09)
(773
.862
884(.030)
GROVER
(303
(770)
(26
758
(.92(.031)
(.19))
GROVERnIkä
(.08)
(308.
.726(.003)
.788(
.873
888
R-MAT
(.30)
2332
.745(.010)
(788)
.887 (.028)
(.27)
R-MAT rdkit
-246(0024)
.239
(306))
.746
.91(,010)
(88(.032)
(886(.031)
SVM
(779(005))
(61(7077)
(.30(.00)
7772
(.93()030)
.990
SVM,dka
.279(.02))
.285
(.01.
.781(0001)
.895
(.84()03))
RF
(373)
5524(0098
(721(.003)
.791(.012)
.892(.026)
(.80)
RF
289((000)
.337 (.26)
(743).002)
(80(.0003)
.903
(886(.028)
GCN
(.99))
(.68)
(695
(300
(883)
(875
DMPNN
(.446.
(252).044)
(.09(.001)
.776(.006)
(885(.026)
(.19)))
4.2
LARGE HYPERPARAMETER HYPERPARAMETTE BUDGET
In contrast to the previous setting, we test R-MAT against similar set of models but using a large-scale hyperparameter search (300 different hyperparameter combinations). This setting has been proposed in (Rong et al.] 2020). For comparison, we include results under small (7 different learning rates) hyperparameter budget. All datasets use scaffold split. Scores are calculated based on 3 different data splits. While the ESOL and FreeSolv datasets are the same as in the previous paragraph, here they scaffold split and labels are normalized (unlike in the previous paragraph). Additional information about the models and datasets used in this benchmark are given in Appendix C.5
Table summmarizes the experiment. Results show that for large hyperparameters budget R-MAT outperforms other methods tasks and along with GROVER are the best in one more task. In overall in this setting our method achieves comparable results to GROVER, having the same median rank and being slightly worse in terms of mean rank (see left side of Figure 7 On the other hand, for small hyperparameters budget R-MAT achieves the best results, both in terms of the mean and the median ranks (see right side of Figure
4.3 LARGE-SCALE EXPERIMENTS
Finally, to better understand how R-MAT performs in a setting where pretraining likely to less influence results, we include results on QM9 dataset (Ramakrishnan et al. 2014). QM9 is a quantum mechanics benchmark that encompasses prediction 12 simulated properties across around 130k small molecules with at most heavy (non-hydrogen) atoms. The molecules are provided with their atomic 3D positions for which the quantum properties were initially calculated. For these experiments, used learning rate equal to 0.015 (we selected this learning rate value as it returned the best results for a dataset among 4 different learning rates that we tested: {0.005,0.01 11155,001,0.00... 00000.01,,0002), Additional information about the dataset and models used in this benchmark are given in Appendix C.6
7

---

Under review as a conference paper at ICLR 2022
Table 2: Results on the benchmark from (Rong al. 2020). Models are fine-tuned under large hyperparameters budget. Additionally, models fine- tuned with only tuning the learning rate are presented in the last group. The last two datasets are classification tasks (higher better), the remaining datasets non-pretrained baselines (GraphConv (Kipf Welling, 2016), Weave (Kearnes et al. 2016) and DMPNN (Yang et al. 2019a)) from (Rong et 2000). We also include RFrakit baseline method with added rdkit features. Rank-plot for these experiments is in Appendix D.2] We bold the best scores over all models and underline the best scores for learning rate tuned models only.
regression tasks (lower is better). For reference, we include results for
ESOL
FreeSolv
Lipo
QM7
BACE
BBBP
RF RF
(96)
2.625
.739 (008))
124.3(3.5)
(,400)
.928( (.025)
GraphConv
1.068(.050)
2.905
..199.049)
118.9
884(.011)
.877
Weave
1.158 (.055)
(.250)
(13(.012)
(2.7
(.00))
1837
DMPNN
(.50)
2.977(9914)
.653 (.366
(046)
105.8
882((.55)
.999 (.30)
GROVER
(.10)
1.592
563, (000)
72.5(5.9)
(878(.01))
.936
R-MAT rdkit
(786(.133)
2.044(.662)
577(.028)
68. (1.699(11122))
.871 (7028)
MAT
(.53(.159)
(.125
(608(.07)
102.8(2.94)
(46(,025)
920(.039)
GROVER
(.100.
2.262
604(.015)
82.623
(.2(
.908 (053))
GROVER
(924(.129)
2.096
(.29)
84. (4.625(4.174))
.873(.031)
.931 (.21))
R-MAT
.801 8801(.132) (.132)
1.912
(.29)
77 77.248 819)
.811.0011
.931(.010) (.016)
R-MAT rdkit
(149.
2.057 (.434)
(.099
70.929 (3.568)
.858
.200
Figure compares R-MAT performance with various models. More detailed results could be find in Table TO from Appendix D.3 R-MAT achieves highly competitive results, with state-of-the-an performance on 4 out of the T2 tasks. We attribute higher variability of performance to the limited small hyperparameter scarch we performed.
4.4 EXPLORING THE DESIGN SPACE OF SELF-ATTENTION LAYER
Achieving strong empirical results hinged on a methodologically explo- ration the design space of different variants of the self-attention layer. We document here this exploration and relevant ablations. Due to space lim- itations, we defer most results to the Appendix E We perform all exper- iments on the ESOL, FreeSolv and BBBP datasets with different seaf- fold splits. We did not use any pre- training for these experiments. We follow the same fine-tuning methodol- ogy as in Section 4.1
NMP
Schnet
Cormorant
Model
L1Net
UeConv
EGNN
Rank
Figure Rank plot of scores obtained on the QM9 bench- mark, which consists of 12 different quantum property pre- diction tasks.
Importance of different sources of information in self-attention The self-attention module in R-MAT incorporates three auxiliary sources of information: (1) distance information, (2) graph information (encoded using neighbourhood order), and (3) bond features. Table 3 (Left), we show the effect on performance of ablating each of this elements. Importantly, find that each component is important to R-MAT performance, including the distance matrix.
Maximum neighbourhood order We take closer look at how we encode the molecular graph. (Maziarka et al. 2020% used simple binary adjacency matrix encode the edges. We enriched this representation by adding one-hot encoding of the neighbourhood order. For example, the order of for pair of atoms means that there atoms. In R-MAT we used 4 as the maximum order neighbourhood distance. That is, we encoded separate features if two atoms are 1, 2, 3 or 4 hops away in the molecular graph. In Table (Right)
other vertices the shortest path between this pair of
8

---

Under review as a conference paper at ICLR 2022
Table 3: Ablations of Relative Molecule Self-Attention; other ablations are included in the Appendix.
(a) Test performances of R-MAT for different rela- (b) Test set performances of R-MAT for different tive attention features.
choices of maximum neighbourhood order.
BBBP
ESOL
FreeSolv
BBBP
ESOL
FreeSolv
R-MAT
(00{.0a))
.77((027)
(338
R-MAT
.908
.37
.188
distance
858(.064)
(.38))
(.8(.
Max order
(.18 (887.
.461 (,09)
neighbourhood
(043)
(.20).
(545(.023)
Max order 2 .900..88)
(.4((000)
,519,038)
bond features
(60
(955,,02))
(5366003)
Max order =
(T33
.455
,065)
we ablate this choice. The result suggests that R-MAT performance benefits from including separate feature for all the considered orders.
Adjacency
Distances
Closer comparison to Molecule Atten- tion Transformer Our main motivation for improving self-attention MAT was to make it easier to represent attention pat- terns that depend in a more complex way on the distance and graph information. We qualitatively explore here whether R-MAT achieves this goal, comparing its attention patterns to that MAT. From the Figure that indeed R-MAT seems pable of learning more complex attention patterns than MAT. We add more detailed comparison, with more visualised attention heads in Appendix D.4
(a) Molecule
head
head
head
(b) Relative Molecule Attention Transformer
head
head
head
head
5 CONCLUSIONS
Transformer has been successfuly adapted to various domain by incorporating into its architecture minimal of inductive biases. In similar spirit, we methodolog- ically explored the design space of the self- attention layer, and identified highly ef- fective Relative Molecule Self-Attentic
(c) Molecule Attention Transformer
Figure 4: Visualization of the learned self-attention for each of the first attention heads the second layer of pretrained R-MAT (middle) and the first 4 atten- tion heads in pretrained MAT (bottom), for a molecule from the ESOL dataset. The top Figure visualizes the molecule and its adjacency and distance matrices. The self-attention pattern in MAT dominated by the adja- cency and distance matrix, while R-MAT seems capable of learning more complex attention patterns.
Relative Molecule Attention Transformer, a model based on Relative Molecule Self- Attention, achieves state-of-the-arr competitive results across a wide range molecular property prediction tasks. R-MAT highly versatile model, showing state-of-the-art results in both quantum property prediction tasks, as well as on biological datasets. We also show that R-MAT is easy train and requires tuning only the learning rate achieve competitive results, which together with open-sourced weight and code, makes it is highly accessible.
very
Relative Molecule Self-Attention encodes an inductive bias consider relationships between atoms that are commonly relevant to a chemist, but on the other hand leaves flexibility to unlearn them if needed. Relatedly, Vision Transformers learn global processing in early layers despite being equipped with locality inductive bias (Dosovitskiy et al.) (021). Our empirical results show in new context that picking the right set inductive biases key for self-supervised learning to work well. also how that Relative Molecule Self-Ar Attf-Attnntion will help improve other models for molecular property prediction.
Learning useful representations for molecular property prediction is far from solved. Achieving state- of- results, while less dependent on them, still relied on using certain large sets of handcrafted features both in fine-tuning and pretraining. At the same time, these features are beyond doubt learnable from data. Developing methods that will push representation learning towards discovering these and better features automatically from data is exciting challenge for the future.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Brandon Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. arXiv preprint arXiv: 04015, 2019.
Andreas Bender and Isidro Portés-Ciriano. Artificial intelligence in drug discovery: what is re- alistic. what are illusions? part 1: Ways to make an impact, and why we are not there yet. Drug Discovery Today, 26(2):511 1--5244, 2021. ISSN 1359-6446. doi: https///do..org/10..00 1016/ j.drudis 12.0. 12.009. URL https:/ httpp/::///ww/////// www. www.sccieieeeeeeeeooco.cc.cceeee wttpeeecoo/..eeeeeeetttttteeeeee
Gino Brumner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Wattenhofer. On identifiability in transformers. arXiv preprint (eXiv:1908.042100 2019.
H.C. Stephen Chan, Hanbin Shan, Thamani Dahoun, Horst Vogel, and Shuguang Yuan. Advancing drug discovery via artificial intelligence. Trends in Pharmacological Sciences, 40(8):592-600 604, 2019. ISSN 0163-6147. doi: http:///oi..og//1.00011..t.....p 20..0060000.000..00.0...0 URL https: www sci ecceedddi i rect com poom/cciin /ttps///e/////ttctc/////////////
35x Special Issue:
Rise of Machines in Medicine.
Seyone Chithrananda, Gabriel Grand, and Bharath Ramsundar. Chemberta: Large-scale self- supervised pretraining for molecular property prediction, 2020.
Kevin Clark, Minh- Minh-Thang Luong, Quoe V Le, and Christopher D Manning. Electra: Pre-training text encoders as discriminators rather than generators. arXiv preprint arXiv:2003. 10555, 2020.
Connor W. Coley, Regina Barzilay, William H. Green, Tommi S. Jaakkola, and Klavs F. Jensen. Convolutional embedding of attributed molecular graphs for physical property prediction. Journal of Chemical Information and Modeling, 2017.
Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. 2018. 8.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL net. fort/ d= cbFdNTTy!
David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarelli, Timothy Hirzel, Alan Aspuru-Guzik, and Ryan P Adams. Convolutional networks on graphs for learning molecular fingerprints. In Advances in Neural Information Processing Systems 28. 2015.
Benedek Fabian, Thomas Edlich, Héléna Gaspar, Marwin Segler, Joshua Meyers, Marco Fiscato, and Mohamed Ahmed. Molecular representation learning with language models and domain-relevant auxiliary tasks, 2020.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning, pp. 3176-3176. PMLR, 2020.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (r)-tsfsfformr:s: 3d
roto-translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
Piotr Gaiński, Fukasz Maziarka, Tomasz Danel, and Stanislaw Jastrzębski. Huggingmolecules: an open-source library for transformer-basee molecular property prediction. 2021.
Anna Gaulton, Louisa J. Bellis, A. Patricia Bento, Jon Chambers, Mark Davies, Anne Hersey, Yvonne Light, Shaun McGlinchey. David Michalovich, Bissan Al-Lazikani, and John P. Overington. ChEMBL: a large-scale bioactivity database for drug discovery. Nucleic Acids Research, 40(D1), 09 2011.
10

---

Under review as a conference paper at ICLR 2022
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Machine Learning, 2017.
Shion Honda, Shoi Shi, and Hiroki R. Ueda. SMILES transformer: Pre-trained molecular fingerprint for low data drug discovery, 2019.
Howard and Sebastian Ruder. Fine-tuned language models for text classification. 2018.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for Pre-training Graph Neural Networkss. In International Conference on Learning Representations, 2020.
Zhiheng Huang, Davis Liang, Peng Xu, and Bing Xiang. Improve transformer models with better relative position embeddings. 2020.
John Ingraham, Vikas Kamur Garg, Regina Barzilay, and Tommi S Jaakkola. Generative models for graph-based protein design. 2021.
Stanislaw Jastrzębski, Damian Leániak, and Wojciech Marian Czamecki. Learning to smile (s). arXiv preprint 1rX.v:660.006280 2016.
Dejun Jiang, Zhenxing Wu, Chang-Ya Yu Hsich, Guangyong Chen, Ben Liao, Zhe Wang, Chao Shen, Dongsheng Cao, Jian Wu, and Tingjun Hou. Could graph neural networks learn better molecular representation for drug discovery? comparison study of descriptor-based and graph-based models. Journal of cheminformatics, 13(1):1 23, 2021.
Guolin Ke, Di He, and Tie- Yan Liu. Rethinking positional encoding in language pre-training. In International Conference on Learning Representations, 2021. URL hattps:/
net/ funt/forum?iln :iterummid-099-- 528y2F
Steven Kearnes, Kevin McCloskey, Mare Berdll, Vijay Pande, and Patrick Riley. Molecular graph convolutions: Moving beyond fingerprints. Journal of Computer-Aided Molecular Design, 2016.
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers in vision: A survey, 2021.
Thomas N Kipf and Max Welling, Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv: wot.::::9000000 2016
Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. International Conference on Learning Representations, 2020.
Alexandru Korotcov, Valery Tkachenko, Daniel P. Russo, and Sean Ekins. Comparison of deep learning with multiple machine learning methods and metrics using diverse drug discovery data sets. Molecular Pharmaceutics, 2017.
Greg Landrum. Rdkit: Open-source cheminformatics software. 2016.
Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentive sentence embedding. 2017.
Christopher A Lipinski, Franco Lombardo. Beryl W Dominy, and Paul J Feeney. Experimental and computational approaches to estimate solubility and permeability in drug discovery and development settings. Advanced drug delivery reviews, 23(1-3):3-25, 1997.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. 2019.
Andreas Mayr, Günter Klambauer, Thomas Unterthiner, Marvin Steijaert, Jörg K. Wegner, Hugo Ceulemans, Djork-Amé Clevert, and Sepp Hochreiter. Large-scale comparison of machine learning methods for drug target prediction chembl. Chem. Sei., 2018.
11

---

Under review as a conference paper at ICLR 2022
Lukasz Maziarka, Tomasz Danel, Slawomir Mucha, Krzysztof Rataj. Jacek Tabor, and Stanistaw Jastrzębski. Molecule attention transformer. 2020.
Benjamin Kurt Miller, Mario Geiger, Tess E Smidt, and Frank Noé. Relevance of rotationally equivariant convolutions for predicting molecular properties. arXiv preprint arXiv:2008 08461, 2020.
Dai Quoe Nguyen, Tu Dinh Nguyen, and Dinh Phung. Unsupervised universal selfatttetion network for graph classification. CoRR, abs/1909. 1855, 2019. hKl/sttp://araii..rr//bbb/
11855
Agnieszka Pocha, Tomasz Danel, and Lukasz Maziarka. Comparison of atom representations in graph neural networks for molecular property prediction. 2020.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum
chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
MC Robinson, RC Glen, and AA Lee. Validating the validation: reanalyzing a large-scale comparison of deep learning and machine learning models for bioactivity prediction. Journal of Computer- Aided Molecular Design, 34:717-730, 2020. doi: 10.
URL http:
7/dx doii wrw/oo 1 1007 01101000002-0010000------0--
David Rogers and Mathew Hahn. Extended--comecciityy fingerprints. Journal of chemical information and modeling, 2010.
David W. Romero and -an-ttap Cordonnier. Group equivariant stand-alone self-attention
for vision. In International Conference Learning Representations, 2021. URL https
//openreview net (tte://foommmddd/k//mdmo
Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer large-scale molecular data. Advances Neural Information Processing Systems, 33, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, (n) equivariant graph neural networks arXiv preprint arXiv:2102.09844 2021.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmicla, Alexandre Tkatchenko, and Klaus-Rober Müller. SchNet: A continuous-filtt filler convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems 30. 2017.
Philippe Schwaller, Teodoro Laino, Theophile Gaudin, Peter Bolgar, Christopher Hunter, Costas chemical
Bekas, and Alpha A Lee. Molecular transformer: A model for reaction prediction. ACS central science, 2019.
Alexander Sergeev and Mike Del Balso. Horovod: fast and easy distributed deep learning in tensorflow. arXiv preprint arXiv: / 18820055999 2018.
Chao Shang. Qinqing Liu, Ko-Shin Chen, Jiangwen Sun, Jin Lu, Jinfeng Yi, and Jinbo Bi. Edge Attention- based Multi-Relational Graph Convolutional Networks. 2018.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics; Human Language Technologies, Volume 2 (Short Papers), 2018.
Teague Sterling and John J Irwin. Zinc 15-ligand discovery for everyone. Journal of chemical information and modeling. 2015.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivaat equivariant neural networks for 3d point clouds. 08219, 2018
arXiv preprint
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 2017.
12

---

Under review as a conference paper at ICLR 2022
Alex Wang, Yada Pruksachatkum, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: stickier benchmark for general-purpose language understanding systems. In Advances in Neural Information Processing Systems, 2019a.
Sheng Wang, Yuzhi Guo, Yuhong Wang, Hongmao Sun, and Junzhou Huang. SMILES-BERTT Large scale unsupervised pre-training for molecular property prediction. In Proceedings of the 10th ACM International Conference on Bioinformatics. Computational Biology and Health Informatics, BCB '19 2019b.
David Weininger. Smiles, a chemical language and information system. 1. introduction to methodol- ogy and encoding rules. Chem. Inf. Comput. Sci., 28(1):31-36, 1988.
Zhenqin Wu, Bharath Ramsundar, Evan N. Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: benchmark for molecular machine learning. Chem. Sci., 2018.
Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman- Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 2019a.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 2019b.
13

---

Under review as a conference paper at ICLR 2022
A R-MAT NODE AND EDGE FEATURES
In the following section, we present the node and edge features used by R-MAT.
A.1 EDGE FEATURES
In R-MAT, all atoms are connected with an edge. The vector representation of every edge contains information about atoms neighbourhood, distances between them and physicochemical features of bond exists (see Figure
Neighbourhood embeddings The neighbourhood information of atom pair is represented by 6-dimensional one-hot encoded vector, with features presented in Table 4 Every neighbourhood embedding contains the information about how many other vertices are between nodes and the original molecular graph.
Table 4: Featurization used to embed neighbourhood order in R-MAT.
Indices
Description
0
1
Atoms and j are connected with bond In the shortest path between atoms and j there is one atom
2
In the shortest path between atoms and j there are two atoms
3
4 In the shortest path between atoms and j there are three or more atoms 5 Any of the atoms i j is dummy node
Bond embeddings Molecular bonds are embedded in 7-dimensional vector following (Coley et al. 2017), with features specified in Table 5 When the two atoms are not connected by true molecular bond, all 7 dimensions are set to zeros.
Table 5: Featurization used to embed molecular bonds in R-MAT.
Indices
Description
0
3
Bond order one-hot hot vector of 1, 1.5, 2, aromatic
5
Is conjugated Is in ring
A.2 NODE FEATURES
The input molecule is embedded as matrix of size Natom x 36 where each atom of the input is embedded following (Coley al. 2017 Pocha et al. 2020). All features are presented in Table 6
B PRETRAINING
We extend the pretraining procedure (Maziarka et al.) 2020), who used a masking task based on (Devlin et al. 2018: Hu et al. 2020):: they masked types of some of the graph atoms and treat them as the label, that should be predicted by the neural network. Such approach works well in NLP where models pretrained with the masking task create the state-of-the-arr representation (Devlin et al.) 2018 Liu et al., 2019 Yang et al., (2019)). However in chemistry, otherwise than NLP, the size of atoms vocabulary is much smaller. Moreover, usually only type of atom fits given place and thus the
14

---

Under review as a conference paper at ICLR 2022
Table 6: Featurization used to embed atoms in R-MAT.
Indices
Description
Atomic identity one-hot vector of B, N, C, CI, Br, Dummy, other
0 11
Number of heavy neighbors as one-hot
12 17
vector of 0, 1, 2, 3, 4,5
Number of hydrogen atoms as
18 22
one-hot vector of 0, 1, 2, 3,
Formal charge as one-hot vector of -5, -4, 4, 5
23
33
34
ring Is aromatic
35
representation trained with the masking task has problems with encoding meaningful information in chemistry.
B. CONTEXTUAL PRETRAINING
Instead of atom masking, we used two-step pretraining that combines the procedures proposed by (Rong et al. 2020, Fabian al. 2020). the first step, the network is trained with the contextual property prediction task (Rong et al. (020), where we mask not only the selected atoms, but also their neighbours. The task is then predict the whole atom context, e.g. if the selected atom's 's type is carbon connected with nitrogen with double bond and with oxygen with single bond, we encode the atom neighbourhood as C_N-DOUBLEI SEILLEEO-SINGGIE O-SINGLEI (we list all the node-edge counts terms in the alphabetical order), then the network has to predict the specific type of the masked neighbourhood for every masked atom. This task is much more demanding for the network than the classical masking approach presented by (Maziarka et al., 2020) as the network has to encode more specific information about the masked atom's 's neighbourhood. Furthermore, the size of context vocabulary is much bigger than the size of atoms vocabulary in the MAT pretraining approach (2925 for R-MAT vs 35 for MAT).
B.2 GRAPI-LEVEL PRETRAINING
The second task is the graph-level property prediction proposed by (Fabian et al., 2020). In this pre- training procedure, the task is to predict 200 real-valued descriptors of physicochemical characteristics of every given molecule.
The list of all 200 descriptors from RDKit is as follows:
BalabanJ, BertzCI, Chi0, Chion, Chiov, Chil, Chiln, Chilv, Chi2n, Chi2v, Chi3n, Chi3v, Chi4n, Chi4v, rState_vsAl, EState_VSA10, EState_VSAl1, EState_VSA2, EState_vsA3, EState_VSA4, EState_VSA5, EState_VSAG, EState_VSAT rppensityMorganl, FppensityMorgan2,, FppensityMorgan3, FractioncsP3,
Setatt_VVAS, eState_VSAt ExactKolWt,
HallKierAlpha, HeavyAtomCount, HeavyлtоmMolwt, Ipc, Kappal, Kappa2,
Kappa3, LabuteasA, MaxAbsestate MaxibsPartiallharreee MaxtStateindex,
MaxFartialCharge, MinAosEStateIndex. MinAbsPartiialCaare. XinEStateIndex, MinFartialCharge, MolLogP, MoLMR, MolNe, NHоHCount, NCCount, AmnklippatticCaatooyyassss NumAl humalipppaaicteeeercccc..... NumAliphaticRing NumbAromatticGarbbccyeless Numaromatiicheterrceyyee
NumAromaticRings,
NumHAcccptorrs, NumHDonors, NumHeterooatoms, NunReddiallEloctrron.. NumRotatabbeeBondds Munsstturatcddaabbcyyele. Munsattraateddeetrroyyyet MumSaturateeRRings, JumVallenceElectrrnss PEGE PEOE_VSAL, PEOE _ese_,_ PEOE VSA11, PEOE VSA12, PEOE VSA13, PECE VSA14, PEOE VSA2, PECE VSA3, PEOE VSA4, PEOE_VSAS, PEOE_ VSA6, PEOE PECE_VSA7, PEOE VEAS, PECE_ VSA9, Ri ngCount, SMR_VSAI SMR_VSAI VSA1 SMR_, SMR_,SA3, SMR_ [MR_ SMR_VSA5,
15

---

Under review as a conference paper at ICLR 2022
SMR VSAG, SMR VSA7, SMR VSAB, SMR VSA9, SlogP VSAI, SlogP VSA10, CLogP_VSAlL SlogP VSA12, SlogP VSA2, SlogP VSA3, Slog? VSA4, SlogP _Loge_VSè,, VSA6, LlogP_VSAL SlogP_VSA SiogP_VSA9, VSA9, TPSA, VSA EStatel,
VSA_KState10, vsA_astate?, VSA_EState3, VSA_RStated, VSA_kState5, VSA_PState6, VSA_ RStatee, V VSA_RState9, fr_A1_CO0, fr_A1_OH, fr_Al_oH_noTert, fr_Ar,, fr_Ar_cOO, fr fr_Ar_,, fr_Ar_NH, fr fr_Ar_CH, fr_coo, fr02, fr___o, fr_c___neCOO, fr_c_s, fr_HOCCN, fr_Imine, fr fr_NHO, fr_ fH_NHl, fr NH2, fr fr_N_o, fr "r_NNealkkllationl, fr fr_Ndealkylation., fr #t_Nhpyrrele, fr_SH, fr sr_aldehyde, Er_alky__carrbamate,, fr fr_elkyl fr_allylic_oxid, fr_amide, fr fr_amidine, fr_aniline, fr_aryl_methyl, fr_azide, fr_azo, fr_barbitur, fr_benzene, fr fr_benzodiazeping, fr fr_bicyclic, fr_diazo, fr fr_d@hhddrpyrriddne, fr_epoxide, fr_ester, fr_ether, fr fr_furan, ir fr_guanido, fr fr_halogen, ir fr_hdrzine, fr fr_hdrzone, fr fr_imidazole, fr fr_imide, fr fr_isocyan, fr fr_isothiccyan, fr_ket,ne, fr t=_ketone_Toplie Coroins_Topiss, fr fr_lactam, fr_lactone, fr fr_methoxy, fr_morpholine, fr_nitrile, fr_nitro, fr fr_nitro_aronn arem, fr_oxime, fr_para_hydroxylatioon fr_phenol, fr fr_phenol_noOrtholbond fr_phos_acid, fr_phos_ester, fr_piperdine, fr_piperzine, fr_priamide, fr prisulfonamd, fr pyridine, fr ir_quatN, fr_sulfide, fr sr_sulfonamd, fr sulfone, =r_term _ acetylene, ir ir_tetrazole, fr tr_thiazele, fr #r_thiocyan, fr #r_thiophene, fr_unbrch. alkane, ir fr_urea, ged
fr_nitroso, fr_oxazole,
C EXPERIMENTAL SETTING
C.I MODEL HYPERRAARMETERS
R-MAT model consists of 10 layers with 12 attention heads in each, dmodet 768. The distance layer consists of 32 radial functions (Nemb 32) and cutoff distance is set to 20 Å. Attention pooling consists of 4 pooling heads and pooling hidden dimension set to 128. Prediction MI.P consists of one hidden layer with dimension set to 1024 and dropout 0.1. R-MAT uses Reaky-ReLU with slope 0.1 as a non-lincarity in all experiments. This set of hyperparameters defines model with 48 millions of parameters, which is equal the number of parameters of GROVER base and is slightly higher than the number of parameters of MAT (42M).
C.2 3D CONFORMATIONS
The 3D molecular conformations that are used to obtain distance matrices were calculated using UFFOPTIMIZEMOLECUUI function from the RDKit package (Landrum) 2016) with the default parameters (MAXITERS=200, "DWTHRESH=1D CONFID= IGNOREINTEEFRGIITER.CC TIONs=True).
One vinadvantage of this approach is the costly calculation of the distance matrix. Which is not burdensome for small datasets with small molecules (e.g FreeSolv), however, they can be a problem with larger ones (e.g. dataset for pretraining). In Table 7 we present dataset sizes and molecular statistics for three different datasets (FreeSolv, ESOL, BBBP), as well as time needed for molecular conformations calculation. In Figure 5 present these times for every single molecule from these datasets. Based on these results one can see, that the larger conformation calculation time. Moreover even for BBBP, that is not rather large dataset, calculating the conformations could take almost minutes. We leave exploring other conformation methods (more accurate faster) interesting topic for the future.
molecule, the longer the
Table 7: Time needed for molecular conformations calculation for different datasets.
Dataset size
Average number of atoms
Calculation time (s)
Average calculation time (s)
FreeSoly
622
8.72
5.225
0.008
ESOL
1128
13.28
27 27.865
0.024
BBBP
2037
24.02
267. 267.084
0.131
16

---

Under review as a conference paper at ICLR 2022
TreeSolv
ESOL
BBBP
(4)
Time
501
Number Atoms
Atoms
Numbbe
Figure 5: Time needed for conformations calculation for every single molecule for different datasets.
C.3 PRETRAINING
We pretrained R-MAT on 4 millions of unlabelled molecules. Molecules were obtained from the ZINCI5 (Sterling& & Irwin 2015) and ChEMBL (Gaulton et al. 2011) datasets, by first taking sample of 10M molecules and then filtering them using the Lipinski's 's rule of five (Lipinski et al.) 1997). We have split the data into training and validation datasets, where validation dataset consists of 5% of the data. We pretrained R-MAT for 150 epochs. We used dropout equal to 0.1, learning rate 0.001, the Noam optimizer (Vaswani al. (017) trained with 20000 warm-up steps and batch size 256. We pretrained R-MAT with 8 Nvidia AT00 GPUs, using the Horovod package (Sergeev &
Del Balso 2018).
Overlap between pre-trained and task datasets In Table g we provide the information about number of atoms that overlaps between pre-trained dataset and all tasks datasets used this paper. One can see that usually the overlap is at most few percent of all dataset molecules. This overlap does not affect the test performance, as pre-training consists semi-supervised task of context prediction and other task, associated with the prediction of simple physico-chemin characteristics, loosely related to the fine-tuning tasks, which are additionally added to the final graph embedding during the prediction.
Table 8: Overlap between pre-trained dataset and different tasks datasets.
dataset
overlap molecules
overlap percentage
bbbp
34
1.6%
bace
0.0%
esol
83
7.3%
estrogen-beta
45
2.2%
freesoly
0.3%
lipo
128
3.0%
mesta-high
608
28.5%
mesta-low
608
28.5%
qm7
0.0%
qm9
0.0%
C.4 SMALL HYPERPARAMETER BUDGET
Models We compared R-MAT with four models trained from scratch: Support Vector Machine with RBF kernel (SVM) and Random Forest (RF) that both works on ECFP fingerprints (Rogers & Hahn 2010), Graph Convolutional Network (Duvenaud et al. 2015) (GCN) and Directed Mcssage
Passing Neural Network (Yang et al. 2019a) (DMPNN).
The comparison also includes two different pretrained models: MAT Mazziaka et al. 2020 and GROVER (Rong et al. 2020).
Datasets The benchmark is based on important molecule property tasks in the drug discovery domain. The first datasets are ESOL and FreeSoly, in which the task is to predict the solubility of
17

---

Under review as a conference paper at ICLR 2022
molecule in water key property of any drug and the error is measured using RMSE. The goal in BBBP and Estrogen-3 3 is to classify correctly whether a given molecule is active against a biological target. For details other tasks please (Maziarka et al.) 2020). BBBP and Estrogen- 3 used scaffold split, rest datasets used random split method. For every dataset 6 different splits were created. Labels of regression datasets (ESOL and FreeSolv) were normalized before training. We did not include the Estrogen-a dataset that was also presented by (Maziarka et al., 2020), due to the GPU memory limitations (as the biggest molecule from this dataset consists of over 500 atoms).
Training perraaaraeetess We fine-tune R-MAT the target tasks for 100 epochs, with batch size equal to 32 and Noam optimizer with warm-up equal to 30% of all steps. The only hyper- parameter that tune is the learning rate, which is selected from the set of 7 possible options: (le-3, 5,,4, 1e-4, 5e-5, le-5, 5e-6, 1e-6]. This small budget for hyperparameter selection reflects the long-term goal of this paper of developing easy use models for molecule property prediction. The fine-tuning was conducted using Nvidia V100 GPU.
C.5 LARGE IYPERPARAMETER BUDGET
Models For large hyperparameters budget we compared R-MAT with three models trained from
scratch: GraphConv (Kipt & Welling 2016), Weave (Kearnes et al. 2016) and DMPNN Yangg
et al. 2019a) and two ppee-rained models: MAT (Maziarka et al. (020) and GROVER (Rong et al. 2020). For small hyperparameters budget we compared R-MAT to MAT and GROVER. We note that R-MAT, MAT and GROVER use different pretraining methods. MAT was pretrained with 2M molecules from ZINC database and GROVER was pretrained with 10M molecules from the ZINC and ChEMBL databases.
Datasets All datasets were splitted using scaffold split. The resulting splits are different than in the MAT benchmark. For every dataset, 3 different splits were created. In our comparison we included only the subset of single-task datasets from original GROVER work (Rong et al. 2020). This is the reason why we use smaller number of datasets. The obtained regression scores for ESOL differs significantly from the small hyperparameters budget benchmark because this time labels are not normalized.
Training hy hyperparrmeters For learning rate tuning we used the same hyperparameters settings as in the MAT benchmark (see Appendix C.4).
For large hyperparameters budget we run random search with the hyperparameters listed Table
Table 9: Relative Molecule Attention Transformer large grid hyperparameters ranges
parameters
warmup
0.05, 0.1, 0.2, 0.3 0.005, 0.001, 0.0005, 0.0001, 0.00005, 0.00001, 0.000005, 0.000001 100
learning rate
epochs pooling hidden dimension
64, 128, 256, 512, 1024
pooling attention heads
2, 8 8
prediction MLP layers
1,
prediction MLP dim
256, 512, 1024, 2048
prediction MLP dropout
0.0, 0.1. 0.2
C.6 LARGE-SCALE EXPERIMENTS
We compared our R-MAT with 8 different models: NMP Gilmer et al. Schnet (Schütt et al.) 2017), Cormorant (Anderson et al.] 2019). LINet Miller al. LieConv (Finzi et al. 2020), TEN (Thomas et al. 2018), SE(3)- Tr. (Fuchs et al. 2020). EGNN ras et al. 2021).
Models
201 b.
2020
(Sator
18

---

Under review as a conference paper at ICLR 2022
Datasets The QM9 dataset (Ramakrishnan et al. 2014) is dataset that consists of molecules, with up to heavy atoms (H, C. N, O, F) and up to 29 atoms overall per molecule. Each atom from this dataset is additionally associated with 3D position. The dataset consists of 12 different regression tasks named: a, Ae, θномо, éLUO. 14, Cv. G. H, R², U. Uo. ZPVE, for which the mean absolute emror is standard metric. The dataset has 130k molecules. We data splits proposed by (Anderson et al. 2019). which gives us 100k trainin molecules, 18k molecules for validation and 13k molecules for testing.
Training hy perraarameeer We trained R-MAT for 1000 epochs, with batch size equal to 256 and learning rate equal to 0.015. We report the test set MAE for the epoch with the lowest validation MAE. We selected this learning rate value as it returned the best results for among 4 different learning rates that we tested: {0.005, 0.01, 0.015, 0.02}.
C.7 ABLATIONS
Datasets For the ablations section, used the BBBP, ESOL and FreeSolv datasets, splitted using a scaffold split, with different splits. Labels of the regression datasets (ESOL and FreeSolv) were normalized before training. Scores obtained in this section differ significantly from the previous benchmarks due to the different data splits, different model hyperaarameters and no pretraining used.
Training hyperparameters Similarly
for our main benchmarks, we tuned only the learning rate, which was selected from the set of 7 possible options: le- 1,,,1e,,,530 1e-5, 5e-6, le-6). We used batch size equal to 32 and Noam optimizer with warm-up equal to 20% of all steps. Moreover we use single layer, instead of two-layers MLP as our classification part.
D ADDITIONAL EXPERIMENTAL RESULTS
D.I SMALL HYPERPARAMETER BUDGET
In Figure O one can find rank plots for results from Table R-MAT and R-MAT dkkä obtained the best median rank among compared models. The performance gain of Rakr over RF and SVM over SVM is interesting. Overall these two baseline models are worse only than R-MAT.
SVM
SVM rdkit
RF
RF rdkit
GCN
Model
DMPNN
MAT
GROVER
GROVER rdkit
RMAT
RMAT rdkit
10
Rank
Figure 6: Rank plot for small hyperparameter budget experiments.
D.2 LARGE HYPERPARAMETE BUDGET
In Figure 7 one can find rank plots for results from Table 2 trained with the large grid search (Left) and for models with only learning rate tuning (Right).
We present separate plots for models
19

---

Under review as a conference paper at ICLR 2022
AF nise
GrachCorv
Pboee
CMPN
RMAT
GROVEn
ANNT nik
1.5
2.0
5.0
Pank
Pank
Large grid-scarch
(b) Learning only grid-search
Figure 7: Rank plot for large hyperparameter budget (Left) as well as for models trained with only the learning rate tunning (Right).
D.3 LARGE-SCALE EXPERIMENTS
In Table I0 one can find detailed results of comparison R-MAT performance with other various models. R-MAT achieves highly competitive results, with state-of-the-art performance 4 out of the 12 tasks, which proves how universal this model is.
Table 10: Mean absolute error QM9, benchmark including various quantum prediction tasks. Results are cited from literature.
Task Units
As meV
STОМо meV
ELUMO meVV
C
C
H
R bohr
U meV
Ug meV
ZPVE
b
D
cal/mol K
meV V
me
meV
NMP
.092
69
43
38
.030
.040
19
17
.180
20
20
1.50
Schnet
.235
63
41
34
.033
.033
14
14
.073
19
14
1.70
Cormorant
0885
61
34
38
.038
.026
20
21
.961
21
22
2.03
LINet
.88
68
46
35
.043
.031
14
14
.354
14
13
1.56
LicConv
49
30
.032
22
.800
19
19
2.28
58
.064
SF
.142
53
.051
.054
EGNN
.071
48
29
25
.029
.031
12
12
106
12
11
1.55
R-MAT aIa
.082
48
.110
.676
10
12
2.23
D.4
CLOSER COMPARISON TO MOLECULE ATTENTION TRANSFORMER
Our main motivation for improving self-attention in MAT was to make it easier to represent attention patterns that depend in a more complex way on the distance and graph information. We qualitatively explore here whether R-MAT achieves this goal, comparing its attention patterns to that of MAT.
For this purpose we compared attention patterns learned by pretrained MAT (weights taken from (Maziarka al., 20200) and R-MAT for selected molecule from the ESOL dataset. Fig- & shows that different heads of Relative Molecule Self-Attention are focusing on different atoms the input atom
in the input molecule. We can that self-attention strength is concentrated (head 5), on the closest neighbours (heads 0 and 11), on the second order neighbours (head 7), on the dummy node (head 1) or on some substructure that occurs the molecule (heads 6 and 10 concentrated on atoms and 2). In contrast, self-attention in MAT focuses mainly on the input atoms and its closest neighbours, the information from other regions of the molecule is not strongly propagated. This likely happens due to the construction of the Molecule Self-Attent Equation (2). where the output atom representation is calculated from equally weighted messages based on the adjacency matrix, distance matrix and self-attention. Due to its construction, it is more challenging for MAT than for R-MAT to learn to attend to a distant neighbour.
in MAT (c.f.
E EXPLORING THE DESIGN SPACE OF MOLECULAR
Identifying the Relative Molecule Self-Attention layer required large-scale and methodological exploration of the self-attention design space. In this section, present experimental data that
20

---

Under review as a conference paper at ICLR 2022
Adjacency
Distancees
Molecule
54 head
5A head
SA hcad 2
SA head
SA head
SA head
SA hcad
SA head
St headid
SA head 5
S4 head 4
54 head
SA headid
SA head
54 Mhhed
St head
SA head 0
54 head
54 head
54 head
St head 11
head
SA 20
head
54 head 12
SA 1o
54 head 14
head
(b) Relative Molecule Attention Transformer
(c) Molecule Attention Transformer
Figure 8: Visualization of the learned self-attention for each of all attention heads in the second layer of pretrained R-MAT (left) and all attention heads in pretrained MAT (right), for a molecule from the ESOL dataset. The top Figure visualizes the molecule and its adjacency and distance matrices. The self-attention pattern in MAT is dominated by the adjacency and distance matrix, while R-MAT seems capable of learning more complex attention patterns.
informed our choices. We also hope it will inform future efforts designing attention mechanism for molecular data. We follow here the same evaluation protocol as in Section 4.4 and show how different natural variants compare against R-MAT.
E. SELF-ATTENTION VARIANTS
Relative Molecule Self-Attention is designed to better incorporate the relative spatial position of atoms in the molecule. The first step is embedding each pair of atoms. Then, the embedding used to re-weight self-attention. To achieve this, Relative Molecule Self-Attention combines ideas from natural language processing (Shaw et al. 2018 Dai et al.) 2019)] Huang al. 2020). These works focus encoding better relative positions of tokens in the input.
compare to three specific variants from these works that can written using our previously introduced notation as:
1. Relative self attention
Shaw et al. 2018):
Gij = (z, W (x) WKyT + (a, (1,
2. Relative self attention with attentive bias Dai et al. 2019):
Eij = (I, (a WKJT - (I. (I,W
uu (zy K o² bi
3. Improved relative self-attention Huang et al. 2020):
Eij = (z, ( (x, w
21

---

Under review as a conference paper at ICLR 2022
Table 11 shows that the attention operation used in R-MAT outperforms variants and across all three tasks and variant on two tasks, being comparable on the third one. This might be expected given that Relative Molecule Self-Attention combines these ideas (c.f. Equation (3)
Table 11: Test set performances of R-MAT for different choices of the relative self self-attention.
BBBP
ESOL
FreeSolv
R-MAT
.908(.03)
(.7(()07)
(.366)3))
Relative type = .859(.05)
.37 (37((.01)
.509 (.028)
Relative type = (566)
.24((014)
.472 (.07))
Relative type 3 (851)
.389 (.9(0000)
.441 (.021)
E.2 ENRICHING BOND FEATURES WITH ATOM FEATURES
In Relative Molecule Self-Attenton, we use a small number of bond features to construct the atom pair embedding. We investigate here the effect of extending bond fcaturization.
Inspired by (Shang et al. 2018), we added information about the atoms that an edge connects. We tried three different variants. In the first one, extend the bond representation with concatenated input features of atoms that the bond connects. In the second one, instead of raw atoms' features, we
tried the on-hoo-encoding of the type of the bond connection (i.e. when the bond connects atoms C and N, we encode it as a bond 'C_N" and take the one-hot-encoding of this information). Finally, we combined these two approaches together.
Table 12: Test set performances of R-MAT for different choices of bond featurization.
BBBP
FreeSolv .388
ESOL
R-MAT
.908(.039)
.488
Connected atoms features
(866(.773)
.406,.088
.489 (.46)
Connection type one-hot
(.12(
.411 (.028)
.510
Both
(873..04)
.390(.u0)
.502
The results are shown in Table Surprisingly, find that adding this type of information to the bond features negatively affects performance of R-MAT. This suggests that R-MAT can already access these features efficiently from the input (which featurize using the same set of features). This could also happen due to the fact that after a few layers, the attention is not calculated over the input atoms anymore. Instead, it works over hidden embeddings, which themselves be mixed representations of multiple atom embeddings (Brunner et al. 2019). where the proposed additional representation contains only information about the input features.
E.3 DISTANCE ENCODING VARIANTS
R-MAT uses a specific radial base distance encoding proposed by (Klicpera et al. 22200) followed by the envelope function, with Nemb 32. We compare here several other natural choices.
We tested the following distance encoding variants (1) removal of the envelope function, (2) increasing the number of distance radial functions 128, (3) using distance embedding from the popular MAT model (Maziarka et al. 2020). (4) using distance embedding from the popular SchNet model (Schütt et al. 2017). The distance in MAT is encoded e(d) exp( -d). The distance in SchNet is encoded e,, exp(- (ppl-lldd 11,, for y 10Å and OÅ p.. 30Å divided into Nemb equal sections, with Nemb set to 32 or 128.
The results shown in Table These results corroborate that a proper representation of distance information a key in adapting self-attention to molecular data. We observe that all variants under- perform compared to the radial base encoding used in Relative Molecule Self-Attentinn. Noteworthy
22

---

Under review as a conference paper at ICLR 2022
is the fact that encoding distances in MAT way caused lot of instability during the training of R-MAT.
Table 13: Test set performances of R-MAT' for different choices of distance modeling.
BBBP
ESOL
FreeSolv
R-MAT
.78((027)
.438(.036)
Nemb 128
(.00(.10)
.417(.025)
..16).
no envelope
(.025)
.397
.473(.01)
Nemb 128, no envelope
.901( (.030)
(416(.014)
(,08)
MAT dist
886(.024)
(40((.045)
(.12))
SchNet dist
(165)
.398 (,433)
(.00)
Nemb 128, SchNet dist
(.85.)
.400
.45(.010)
F ADDITIONAL COMPARISON OF GRAPH PRETRAINING
F.1 PRETRAINING METHODS BENCHMARK
As pretraining is nowadays the main component of big Transformer architectures (Devlin et al. 2018 Liu al. 2019 Clark et al., 2020), decided devote more attention this issue. For this purpose we compared various graph pretraining methods to identify the best one and use it in the final R-MAT model.
Pretraining methods We used various pretraining methods proposed in the molecular property prediction literature (Hu et al. 2020; Maziarka et al. 2020) Rong et al. [0200] Fabian et al. 2020). To be more specific, we tried R-MAT with the following pretraining procedures:
No pretraining - as baseline we include reuslts for R-MAT fine-tuned from scratch, without any pretraining.
Masking masked pretraining used in (Maziarka et al. 2020). This is an adaptation of standard MLM pretraining used NLP (Devlin et al. 20188) to the graphical data. this
approach, mask features of 15% of all atoms in the molecule and then pass it throught the model. The goal is to predict what was the masked features.
Contextual - contextual pretraining method proposed by (Rong et al. 2020). We described it further in Appendix B
Graph-motift graph-level motif prediction method proposed by (Rong et al. 2020), where for every molecule we obtain the fingerprint with information whether specified molecular functional groups present in our molecule. The network's 's task is multi-label classification, where has to predict, whether every predefined functional group in the given molecule.
Physicochemical - graph-level prediction method proposed by (Fabian et al. 2020 We
described it further Appendix B
GROVER pretraining used by authors of GROVER (Rong et al. 2020). Combination of two pretraining methods: contextual and graph-motifs.
R-MAT- pretraining used in this paper. Combination of two pretraining methods: contextual and physicochemical.
Experimental setting We pretrained every model using described earlier dataset with 4M molecules. For every pretraining choice, we trained the model for 50 epochs, using the same training settings for our standard pretraining (see Appendix C.3
Model hyperparameters and fine-tuning settings were the same as for MAT and GROVER benchmarks. We used the BBBP, ESOL and FreeSoly datasets, splitted using scaffold split, with 3 different data splits. Data splits are different than in previous experiments, and this causes the fact that the results differs from the results from other paragraphs'.
23

---

Under review as a conference paper at ICLR 2022
Results Results of this benchmark are presented in Table 14 We draw some interesting conclusions from them. One can see, that using any kind of pretraining helps in obtaining better results than for the model trained from scratch. Using physicochemical features for graph-level training gives better results than graph-motifs. Therefore R-MAT pretraining (contextual + physicochmmical better than GROVER pretraining (contextual + graph-motfis). Moreover combination of two tasks in pretraining usually gives better results than pretraining using only one task. Interestingly, both node-level pretraining methods (masking and contextual pretraining), returns similar results.
Table 14: Test set performances of R-MAT for different choices of pretraining used.
BBBP
ESOL
FreeSolv
No pretraining
(855.0081)
(221)
(.06))
Masking
(67(.006)
.376
(.044)
Contextual
.901 (.09))
(.84)
.133 (,047)
Graph-motif's
(866,005))
(389
.473
Physiochemical
(.0(2))))
.406(,072)
.400
GROVER
.977
(7287)277)
(56(.
R-MAT
(,31.)15)
(600.)12)
.402
F.2 PRETRAINING LEARNING CURVES
In Figure we present learning curves for R-MAT pretrained with our procedure (contextual + physiochemical). Left-side image shows that training loss flattens out quite quickly, however slowly decreases until the end of the training. Moreover one can see that contextual task harder for the network than graph-level property prediction, their losses vary by several orders of magnitude. Left and middle images shows that R-MAT predictions for validation datasets are of good quality, moreover these curves also presents that our model learns all the time, reaching the best values the end of training
0.010
<190
Loss
Contestual lose
0.000
0.30
0.90
loss
0.00e
0.004
035
690
0.002
0.000
25
105
100
Fpach
Epoch
Train step
(b) Validation RMSE for graph- level property prediction task
(c) Validation ROC AUC for con- textual task
(a) Train loss
Figure 9: Learning curves for R-MAT pretraining. On the left-side figure one can see train losses for both contextual and graph-level property prediction tasks. On the middle figure one can see the RMSE for graph-level property prediction for validation dataset obtained by R-MAT during pretraining, On the right-side figure one can see ROC AUC for the classification task of contextual prediction for validation dataset obtained by R-MAT during pretraining.
We also tested, whether longer pretraining allows R-MAT to get better results during the fine-tuning process. In Figure 10 we present fine-tuning scores obtained by models pretrained with different number of pretraining epochs, for FreeSolv and ESOL datasets. Interestingly, longer training does not always results better fine-tuning scores. This is indeed the case for the FreeSoly dataset (left-side image), however for ESOL (right-side image) we cannot draw such conclusion.
24

---

Under review as a conference paper at ICLR 2022
FreeSolv score
ESOL score
042
040
040
0:38
0.35
60
100
120
140
20
40
80
300
120
140
Pretraining epoch
Pretraining epoch
(a) FreeSolv score for R-MAT pretrained for given numbers pretraining epochs
ESOL score for R-MAT pretrained for given numbers pretraining epochs
Figure 10: Fine-tuning scores obtained by R-MAT pretrained with different number of pretraining epochs.
25


Reviews:
[{"review_id": "vAM7fNWvuYI", "reviewer": "Reviewer_GZ33", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Pros\n    - The experimental results are comprehensive.\n    - The paper is easy to follow. Well written.\n\n- Cons\n    - Lack of intuitive understanding of proposed features w.r.t. tasks to solve.\n    - Theoretical justification on why it is impossible to learn these features with the vanilla Transformer or how difficult to learn such features.\n    - Difficult to reproduce the result due to the large scale experiments\n\n- Questions\n    - How many overlaps between the task datasets and pre-trained datasets? Does this overlap influence the test performance?\n    - Increasing the maximum neighborhood order does not improve the performance.\n        - The results seem inconsistent. How do we interpret the result?\n    - What makes EGNN perform the best with the QM9 dataset (Figure 3) and why the proposed method cannot achieve a similar result?", "summary_of_the_paper": "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.", "summary_of_the_review": "Based on the pros and cons written above, I think the paper is marginally above the acceptance threshold.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Vme8EQOQZlU", "reviewer": "Reviewer_Lgqf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Adding distance information directly to the learned \"distance\" before computing the softmax as well as additional information as in MAT is a very good idea (similar to ALiBi [1]). Results in Table 3 are supportive for adding distance-, bond- and a neighborhood-embedding in the attention matrix.\nThere are datasets the support a performace increase of the proposed method over baselines and established methods, however\n 1. Distance Matrix in MAT vs. In R-MAT: is that it is first embedded, then a learn projection; rather than taking the distance matrix directly times an e.g. a learned scalar; is there a clear advantage of one over the other?\n 2. Table 1 reports the best performance for ESOL and FreeSolv, given that R-MAT was pretrained on 200 rdkit-descriptors and in one model-variant uses those descriptors in the last layer, please make sure to also include a baseline variant of e.g. RF with the 200 rdkit-descriptors (a few of whom describe solubility), (additional to using fingerprints)\n 3. Please state why between the Experiments (1-2) splits were changed from random to scaffold, as well as the normalization of the labels\n 4. In Table 9: \"R-MAT outperforms other variants across the three tasks\" however in ESOL type 1 performs best/not significantly differently\n 5. One disadvantage is the costly calculation of the distance matrix, which should be mentioned at least once in the paper\n 6. Reproducibility: Code to reproduce the results is not provided\n\n \nMinor Issues:\n- \"we report the mean test score\" but the values in brackets for table 1, 2, .., are they std, var, or SD?\n- How were the fixed-hyperparameters for Experiment 1, small-hyperparameter budget selected?\n- 8 nVidia A100 GPUs --> Nvidia\n- Table 2: QM7 should be changed to QM9\n- Equation (1): by adding to vanilla self-attention e_{ij} , doesn't have a variance of sqrt(d_z), so the denominator in the softmax might be adjusted\n- The distance matrix might change dependent on calculation, any improvements from running multiple runs?; distance is not static; rotatable bonds\n\n[1] Press, O., Smith, N.A., & Lewis, M. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ArXiv, abs/2108.12409.", "summary_of_the_paper": "The paper introduces Relative Molecule Attention Transformer, adding distance-, bond- and a neighborhood-embedding  in the attention matrix, and pretrain the model on 4M molecules via context pretraining and on 200 rdkit-descriptors.", "summary_of_the_review": "Several issues have been discussed: Experimental results should be compared to additional baselines, changes in the experimental setup are not justified and seem unnecessary, drawbacks are not adequately discussed.  In the current form, the paper is marginally below the acceptance threshold but the reviewer is inclined to change the rating, in light of new supportive evidence.\n\n**update**: The authors successfully tackled the issues I raised with the manuscript, therefore I am happy to raise my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "50F9Edao8nt", "reviewer": "Reviewer_DtcY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nThe proposed relative self-attention is carefully designed to consider three factors: relative distance, shortest path distance in the molecular graph, and physiochemical. Inspired by the positional bias and content bias from NLP methods, a novel relative molecule self-attention is proposed and helps the Transformer model achieve good experimental performance. \n\nWeaknesses:\n\nSome concerns about the evaluation. In Table 1, why do only two datasets (BBBP and Estrogen-$\\beta$) use scaffold splits, and the other four datasets use the random split? The baseline method GROVER evaluates on multiple graph classification datasets, why only BBBP is evaluated in Table 1? While another dataset BACE is only evaluated in Table 2? Could the authors include results on more datasets like GROVER? Evaluating only on selected datasets makes the readers suspicious of the generalization of the proposed method.\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper proposes a novel relative molecule self-attention layer for the Transformer model. The model can be first pretrained and then finetuned on downstream molecule property prediction tasks. Experimental results demonstrate the effectiveness of the proposed method. However, there are some concerns about the evaluation as described in the Weaknesses section.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents an extension to the MAT model by using relative attention that considers graph-level distance, geometric distance, and bond type between nodes during the attention computation.  This is shown to lead to improved performance on several benchmark datasets against MAT and GROVER.  The inclusion of additional information into the attention computation is a sensible and natural choice for transformer with the success of relative positional embedding in the NLP and vision domains. But it is somewhat a straightforward extension of the existing ideas from other domains to MAT, hence the novelty is somewhat limited.  It is also worth noting that the proposed method should be considered in the context of the larger body of research on 3D GNNs. The authors drew inspiration from DimeNet's design in the encoding of geometric distances but do not consider it in the empirical comparisons. Instead, it focused exclusively on transformer based models. This limits the scope of conclusions that we can draw from these experiments and makes it difficult to gauge the practical impact of RMAT in comparison to many other GNN methods that uses 3D geometries of the molecule.
Decision: Reject