Original: R ELATIVE
Revised: R ELATIVE

Original: M
Revised: M

Original: OLECULE S ELF -A TTENTION T RANS -
Revised: OLECULE S ELF -A TTENTION T RANS -

Original: FORMER
Revised: FORMER

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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 efﬁcient learning from scarce experimental data.
Revised: 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 efﬁcient learning from scarce experimental data.

Original: However, despite signiﬁcant progress, non-pretrained methods can be still competitive in certain settings.
Revised: However, despite signiﬁcant progress, non-pretrained methods can be still competitive in certain settings.

Original: We reason that architecture might be a key bottleneck.
Revised: We reason that architecture might be a key bottleneck.

Original: In particular, enriching the backbone architecture with domain-speciﬁc inductive biases has been key for the success of self-supervised learning in other domains.
Revised: In particular, enriching the backbone architecture with domain-speciﬁc inductive biases has been key for the success of self-supervised learning in other domains.

Original: Inspired by this, we methodologically explore the design space of the self-attention mechanism for molecular data.
Revised: Inspired by this, we methodologically explore the design space of the self-attention mechanism for molecular data.

Original: Our main contribution is Relative Molecule Attention Transformer (R-MAT): a novel Transformer-based model that achieves state-of-the-art or very competitive results across a wide range of molecule property prediction tasks.
Revised: Our main contribution is Relative Molecule Attention Transformer (R-MAT): a novel Transformer-based model that achieves state-of-the-art or very competitive results across a wide range of molecule property prediction tasks.

Original: Relative Molecule Attention Transformer uses a novel self-attention variant that can be readily incorporated into future models for processing molecular data.
Revised: Relative Molecule Attention Transformer uses a novel self-attention variant that can be readily incorporated into future models for processing molecular data.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Predicting molecular properties is central to applications such as drug discovery or material design.
Revised: Predicting molecular properties is central to applications such as drug discovery or material design.

Original: Without accurate prediction of properties such as toxicity, a promising drug candidate is likely to fail clinical trials.
Revised: Without accurate prediction of properties such as toxicity, a promising drug candidate is likely to fail clinical trials.

Original: Many molecular properties cannot be feasibly computed (simulated) from ﬁrst principles and instead have to be extrapolated, from an often small experimental dataset (Chan et al., 2019; Bender & Cortés-Ciriano, 2021).
Revised: Many molecular properties cannot be feasibly computed (simulated) from ﬁrst principles and instead have to be extrapolated, from an often small experimental dataset (Chan et al., 2019; Bender & Cortés-Ciriano, 2021).

Original: 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.
Revised: 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.

Original: Machine learning is moving away from training models purely from scratch.
Revised: Machine learning is moving away from training models purely from scratch.

Original: 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 efﬁciency across multiple tasks (Wang et al., 2019a).
Revised: 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 efﬁciency across multiple tasks (Wang et al., 2019a).

Original: Instead of training models purely from scratch, the models in NLP are commonly ﬁrst pretrained on a large unsupervised corpora.
Revised: Instead of training models purely from scratch, the models in NLP are commonly ﬁrst pretrained on a large unsupervised corpora.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Several works have shown further gains by improving network architecture or the pretraining tasks (Chithrananda et al., 2020; Fabian et al., 2020;
Revised: Several works have shown further gains by improving network architecture or the pretraining tasks (Chithrananda et al., 2020; Fabian et al., 2020;

Original: Rong et al., 2020).
Revised: Rong et al., 2020).

Original: However, pretraining has not yet led to such transformative data-efﬁciency gains in molecular property prediction.
Revised: However, pretraining has not yet led to such transformative data-efﬁciency gains in molecular property prediction.

Original: For instance, non-pretrained models with extensive handcrafted featurization tend to achieve very competitive results (Yang et al., 2019a).
Revised: For instance, non-pretrained models with extensive handcrafted featurization tend to achieve very competitive results (Yang et al., 2019a).

Original: We reason that architecture might be a key bottleneck.
Revised: We reason that architecture might be a key bottleneck.

Original: 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.
Revised: 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.

Original: On the other hand, performance has been signiﬁcantly boosted in other domains by enriching the Transformer architecture with proper inductive biases (Dosovitskiy et al., 2021;
Revised: On the other hand, performance has been signiﬁcantly boosted in other domains by enriching the Transformer architecture with proper inductive biases (Dosovitskiy et al., 2021;

Original: Shaw et al., 2018; Dai et al., 2019; Ingraham et al., 2021; Huang et al., 2020; Romero & Cordonnier, 2021;
Revised: Shaw et al., 2018; Dai et al., 2019; Ingraham et al., 2021; Huang et al., 2020; Romero & Cordonnier, 2021;

Original: Khan et al., 2021; Ke et al., 2021).
Revised: Khan et al., 2021; Ke et al., 2021).

Original: 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.
Revised: 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.

Original: 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;
Revised: 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;

Original: Ingraham et al., 2021)
Revised: Ingraham et al., 2021)

Original: Our main contribution is a new self-attention layer for molecular graphs.
Revised: Our main contribution is a new self-attention layer for molecular graphs.

Original: We tackle the aforementioned issues with Relative Molecule Attention Transformer (R-MAT), our pre-trained transformer-based model, shown in Figure 1.
Revised: We tackle the aforementioned issues with Relative Molecule Attention Transformer (R-MAT), our pre-trained transformer-based model, shown in Figure 1.

Original: 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).
Revised: 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).

Original: Our model achieves state-of-the-art or very competitive performance across a wide range of tasks.
Revised: Our model achieves state-of-the-art or very competitive performance across a wide range of tasks.

Original: Satisfyingly, R-MAT outperforms more specialized models without using extensive handcrafted featurization or adapting the architecture speciﬁcally to perform well on quantum prediction benchmarks.
Revised: Satisfyingly, R-MAT outperforms more specialized models without using extensive handcrafted featurization or adapting the architecture speciﬁcally to perform well on quantum prediction benchmarks.

Original: The importance of representing effectively distance and other relationships in the attention layer is evidenced by large performance gains compared to MAT.
Revised: The importance of representing effectively distance and other relationships in the attention layer is evidenced by large performance gains compared to MAT.

Original: An important inspiration behind this work was to unlock the potential of large pretrained models for the ﬁeld, as they offer unique long-term beneﬁts such as simplifying machine learning pipelines.
Revised: An important inspiration behind this work was to unlock the potential of large pretrained models for the ﬁeld, as they offer unique long-term beneﬁts such as simplifying machine learning pipelines.

Original: We show that R-MAT can be trained to state-of-the-art performance with only tuning the learning rate.
Revised: We show that R-MAT can be trained to state-of-the-art performance with only tuning the learning rate.

Original: We also open-source weights and code as part of the Huggingmolecules (Gai´nski et al., 2021) package.
Revised: We also open-source weights and code as part of the Huggingmolecules (Gai´nski et al., 2021) package.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: [Figure]
Revised: [Figure]

Original: Transformer uses a novel relative selfattention block tailored to molecule property prediction.
Revised: Transformer uses a novel relative selfattention block tailored to molecule property prediction.

Original: It fuses three types of features: distance embedding, bond embedding and neighbourhood embedding.
Revised: It fuses three types of features: distance embedding, bond embedding and neighbourhood embedding.

Original: Pretraining coupled with the efﬁcient Transformer architecture unlocked state-of-the-art performance in molecule property prediction (Maziarka et al., 2020; Chithrananda et al., 2020; Fabian et al., 2020;
Revised: Pretraining coupled with the efﬁcient Transformer architecture unlocked state-of-the-art performance in molecule property prediction (Maziarka et al., 2020; Chithrananda et al., 2020; Fabian et al., 2020;

Original: Rong et al., 2020; Wang et al., 2019b; Honda et al., 2019).
Revised: Rong et al., 2020; Wang et al., 2019b; Honda et al., 2019).

Original: First applications of deep learning did not offer large improvements over more standard methods such as random forests (Wu et al., 2018;
Revised: First applications of deep learning did not offer large improvements over more standard methods such as random forests (Wu et al., 2018;

Original: Jiang et al., 2021; Robinson et al., 2020).
Revised: Jiang et al., 2021; Robinson et al., 2020).

Original: Consistent improvements were enabled by more efﬁcient architectures adapted to this domain (Mayr et al., 2018;
Revised: Consistent improvements were enabled by more efﬁcient architectures adapted to this domain (Mayr et al., 2018;

Original: Yang et al., 2019a; Klicpera et al., 2020).
Revised: Yang et al., 2019a; Klicpera et al., 2020).

Original: In this spirit, our goal is to further advance modeling for any chemical task by redesigning self-attention for molecular data.
Revised: In this spirit, our goal is to further advance modeling for any chemical task by redesigning self-attention for molecular data.

Original: Encoding efﬁciently 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;
Revised: Encoding efﬁciently 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;

Original: Dai et al., 2019; Ingraham et al., 2021; Huang et al., 2020; Romero & Cordonnier, 2021; Khan et al., 2021; Ke et al., 2021).
Revised: Dai et al., 2019; Ingraham et al., 2021; Huang et al., 2020; Romero & Cordonnier, 2021; Khan et al., 2021; Ke et al., 2021).

Original: The vanilla selfattention includes absolute encoding of position, which can hinder learning when the absolute position in the sentence is not informative.
Revised: The vanilla selfattention includes absolute encoding of position, which can hinder learning when the absolute position in the sentence is not informative.

Original: 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).
Revised: 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).

Original: However, most Transformers for the chemical domain predominantly used no positional encoding in the selfattention layer (Chithrananda et al., 2020; Fabian et al., 2020;
Revised: However, most Transformers for the chemical domain predominantly used no positional encoding in the selfattention layer (Chithrananda et al., 2020; Fabian et al., 2020;

Original: Rong et al., 2020; Wang et al., 2019b;
Revised: Rong et al., 2020; Wang et al., 2019b;

Original: Honda et al., 2019; Schwaller et al., 2019), which gives rise to similar issues with representing relations between atoms.
Revised: Honda et al., 2019; Schwaller et al., 2019), which gives rise to similar issues with representing relations between atoms.

Original: We directly compare to (Maziarka et al., 2020), who introduced ﬁrst self-attention module tailored to molecular data, and show large improvements across different tasks.
Revised: We directly compare to (Maziarka et al., 2020), who introduced ﬁrst self-attention module tailored to molecular data, and show large improvements across different tasks.

Original: 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.
Revised: 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.

Original: N
Revised: N

Original: S
Revised: S

Original: N H
Revised: N H

Original: H        3 4 +
Revised: H        3 4 +

Original: Dummy Node 0 1 2 3 4 + D (a) neighbors
Revised: Dummy Node 0 1 2 3 4 + D (a) neighbors

Original: N
Revised: N

Original: S
Revised: S

Original: N H
Revised: N H

Original: H 1.
Revised: H 1.

Original: AromaticRing 1 1.5 2 3 A C R 1.
Revised: AromaticRing 1 1.5 2 3 A C R 1.

Original: Conjugated (b) bonds
Revised: Conjugated (b) bonds

Original: N
Revised: N

Original: S
Revised: S

Original: N H
Revised: N H

Original: H e 1 e 2 ee N emb n= n= n=3 N e n (c) distances
Revised: H e 1 e 2 ee N emb n= n= n=3 N e n (c) distances

Original:  R ELATIVE M OLECULE S ELF -A TTENTION 3. M OLECULAR SELF - ATTENTIONS
Revised:  R ELATIVE M OLECULE S ELF -A TTENTION 3. M OLECULAR SELF - ATTENTIONS

Original: We ﬁrst give a short background on how prior work on adapting self-attention for molecular data, and point out their potential shortcomings.
Revised: We ﬁrst give a short background on how prior work on adapting self-attention for molecular data, and point out their potential shortcomings.

Original: 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).
Revised: 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).

Original: SMILES is a linear encoding of a molecule into a string of characters according to a deterministic ordering algorithm
Revised: SMILES is a linear encoding of a molecule into a string of characters according to a deterministic ordering algorithm

Original: (Weininger, 1988; Jastrz˛ebski et al., 2016).
Revised: (Weininger, 1988; Jastrz˛ebski et al., 2016).

Original: For example, the SMILES encoding of carbon dioxide is C(=O)=O.
Revised: For example, the SMILES encoding of carbon dioxide is C(=O)=O.

Original: Adding a single atom can completely change the ordering of atoms in the SMILES encoding.
Revised: Adding a single atom can completely change the ordering of atoms in the SMILES encoding.

Original: Hence, the relative positions of individual characters are not easily related to their proximity in the graph or space.
Revised: Hence, the relative positions of individual characters are not easily related to their proximity in the graph or space.

Original: 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).
Revised: 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).

Original: We suspect this makes the use of self-attention in SMILES models less effective.
Revised: We suspect this makes the use of self-attention in SMILES models less effective.

Original: Another readily visible shortcoming is that the graph structure and distances between atoms of the molecule are either completely encoded or entirely thrown out.
Revised: Another readily visible shortcoming is that the graph structure and distances between atoms of the molecule are either completely encoded or entirely thrown out.

Original: Graph Transformers Several works have proposed Transformers that operate directly on a graph (Maziarka et al., 2020;
Revised: Graph Transformers Several works have proposed Transformers that operate directly on a graph (Maziarka et al., 2020;

Original: Rong et al., 2020; Nguyen et al., 2019).
Revised: Rong et al., 2020; Nguyen et al., 2019).

Original: The GROVER and the
Revised: The GROVER and the

Original: U2GNN models take as input a molecule encoded as a graph (Rong et al., 2020; Nguyen et al., 2019).
Revised: U2GNN models take as input a molecule encoded as a graph (Rong et al., 2020; Nguyen et al., 2019).

Original: In both of them, the self-attention layer does not have a direct access to the information about the graph.
Revised: In both of them, the self-attention layer does not have a direct access to the information about the graph.

Original: 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.
Revised: 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.

Original: Similarly to Text Transformers, Graph Transformers also do not take into account the distances between atoms.
Revised: Similarly to Text Transformers, Graph Transformers also do not take into account the distances between atoms.

Original: Structured Transformer introduced in (Ingraham et al., 2021) uses relative self-attention that operates on amino-acids in the task of protein design.
Revised: Structured Transformer introduced in (Ingraham et al., 2021) uses relative self-attention that operates on amino-acids in the task of protein design.

Original: 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.
Revised: 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.

Original: As R-MAT encodes the relative distances between pairs of atoms, Structured
Revised: As R-MAT encodes the relative distances between pairs of atoms, Structured

Original: Transformer also uses relative distances between modeled amino-acids.
Revised: Transformer also uses relative distances between modeled amino-acids.

Original: However, it encodes them in a slightly different way.
Revised: However, it encodes them in a slightly different way.

Original: We incorporate their ideas, and extend them to enable processing of molecular data.
Revised: We incorporate their ideas, and extend them to enable processing of molecular data.

Original: Molecule Attention Transformer
Revised: Molecule Attention Transformer

Original: 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).
Revised: 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).

Original: In contrast to most of the aforementioned models, MAT incorporates the distance information in its self-attention module.
Revised: In contrast to most of the aforementioned models, MAT incorporates the distance information in its self-attention module.

Original: MAT stacks N Molecule Self-Attention blocks followed by a mean pooling and a prediction layer.
Revised: MAT stacks N Molecule Self-Attention blocks followed by a mean pooling and a prediction layer.

Original: For a D -dimensional state x ∈ R D , the standard, vanilla self-attention operation is deﬁned as
Revised: For a D -dimensional state x ∈ R D , the standard, vanilla self-attention operation is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where Q = x W Q , K = x W K , and V = x W V .
Revised: where Q = x W Q , K = x W K , and V = x W V .

Original: Molecule Self-Attention extends Equation (1) to include additional information about bonds and distances between atoms in the molecule as
Revised: Molecule Self-Attention extends Equation (1) to include additional information about bonds and distances between atoms in the molecule as

Original: [Equation]
Revised: [Equation]

Original: 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 ) =
Revised: 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 ) =

Original: exp( − d ) , A is the adjacency matrix (with
Revised: exp( − d ) , A is the adjacency matrix (with

Original: 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.
Revised: 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.

Original: Self-attention can relate input elements in a highly ﬂexible manner.
Revised: Self-attention can relate input elements in a highly ﬂexible manner.

Original: In contrast, there is little ﬂexibility in how Molecule Self-Attention can use the information about the distance between two atoms.
Revised: In contrast, there is little ﬂexibility in how Molecule Self-Attention can use the information about the distance between two atoms.

Original: The strength of the attention between two atoms depends monotonically on their relative distance.
Revised: The strength of the attention between two atoms depends monotonically on their relative distance.

Original: However, molecular properties can depend in a highly nonlinear way on the distance between atoms.
Revised: However, molecular properties can depend in a highly nonlinear way on the distance between atoms.

Original: This has motivated works such as (Klicpera et al., 2020) to explicitly model the interactions between atoms, using higher-order terms.
Revised: This has motivated works such as (Klicpera et al., 2020) to explicitly model the interactions between atoms, using higher-order terms.

Original: 3. R ELATIVE POSITIONAL ENCODING
Revised: 3. R ELATIVE POSITIONAL ENCODING

Original: In natural language processing, a 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).
Revised: In natural language processing, a 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).

Original: In order to add the positional information into the input data, a vanilla transformer enriches it with encoding of the absolute position.
Revised: In order to add the positional information into the input data, a vanilla transformer enriches it with encoding of the absolute position.

Original: 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.
Revised: 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.

Original: 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 2).
Revised: 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 2).

Original: 3. A TOM RELATION EMBEDDING
Revised: 3. A TOM RELATION EMBEDDING

Original: Our core idea to improve Molecule Self-Attention is to add ﬂexibility in how it processes graph and distance information.
Revised: Our core idea to improve Molecule Self-Attention is to add ﬂexibility in how it processes graph and distance information.

Original: Speciﬁcally, we adapt positional relative encoding to processing molecules (Shaw et al., 2018; Dai et al., 2019; Huang et al., 2020; Ke et al., 2021), which we note was already hinted at in (Shaw et al., 2018) as a high-level future direction.
Revised: Speciﬁcally, we adapt positional relative encoding to processing molecules (Shaw et al., 2018; Dai et al., 2019; Huang et al., 2020; Ke et al., 2021), which we note was already hinted at in (Shaw et al., 2018) as a high-level future direction.

Original: The key idea in these works is to enrich the self-attention block to efﬁciently represent information about relative positions of items in the input sequence.
Revised: The key idea in these works is to enrich the self-attention block to efﬁciently represent information about relative positions of items in the input sequence.

Original: What reﬂects the relative position of two atoms in a molecule?
Revised: What reﬂects the relative position of two atoms in a molecule?

Original: 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).
Revised: 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).

Original: In the next step, we depart from Molecule Self-Attention (Maziarka et al., 2020) and introduce new factors to the relation embedding.
Revised: In the next step, we depart from Molecule Self-Attention (Maziarka et al., 2020) and introduce new factors to the relation embedding.

Original: Given two atoms, represented by vectors x i , x j ∈ R D , we encode their relation using an atom relation embedding b ij ∈ R D (cid:48) .
Revised: Given two atoms, represented by vectors x i , x j ∈ R D , we encode their relation using an atom relation embedding b ij ∈ R D (cid:48) .

Original: This embedding will then be used in the self-attention module after a projection layer.
Revised: This embedding will then be used in the self-attention module after a projection layer.

Original: Next, we describe three components that are concatenated to form the embedding b ij .
Revised: Next, we describe three components that are concatenated to form the embedding b ij .

Original: Neighbourhood embeddings First, we encode the neighbourhood order between two atoms as a 6 dimensional one hot encoding, with information about how many other vertices are between nodes i and j in the original molecular graph (see Figure 2 and Table 4 from Appendix A).
Revised: Neighbourhood embeddings First, we encode the neighbourhood order between two atoms as a 6 dimensional one hot encoding, with information about how many other vertices are between nodes i and j in the original molecular graph (see Figure 2 and Table 4 from Appendix A).

Original: Distance embeddings As we discussed earlier, we hypothesize that a much more ﬂexible representation of the distance information should be facilitated in MAT.
Revised: Distance embeddings As we discussed earlier, we hypothesize that a much more ﬂexible representation of the distance information should be facilitated in MAT.

Original: To achieve this, we use a radial basis distance encoding proposed by (Klicpera et al., 2020):
Revised: To achieve this, we use a radial basis distance encoding proposed by (Klicpera et al., 2020):

Original: [Equation]
Revised: [Equation]

Original: where d is the distance between two atoms, c is the predeﬁned cutoff distance, n ∈ { 1 , . .
Revised: where d is the distance between two atoms, c is the predeﬁned cutoff distance, n ∈ { 1 , . .

Original: . , N emb } and N emb is the total number of radial basis functions that we use.
Revised: . , N emb } and N emb is the total number of radial basis functions that we use.

Original: Then obtained numbers are passed to the polynomial envelope function u ( d ) = 1 − ( p +1)
Revised: Then obtained numbers are passed to the polynomial envelope function u ( d ) = 1 − ( p +1)

Original: ( p +2)
Revised: ( p +2)

Original: [Equation]
Revised: [Equation]

Original: d p +2 , with p = 6 , in order to get the ﬁnal distance embedding.
Revised: d p +2 , with p = 6 , in order to get the ﬁnal distance embedding.

Original: Bond embeddings Finally, we featurize each bond to reﬂect the physical relation between pairs of atoms that might arise from, for example, being part of the same aromatic structure in the molecule.
Revised: Bond embeddings Finally, we featurize each bond to reﬂect the physical relation between pairs of atoms that might arise from, for example, being part of the same aromatic structure in the molecule.

Original: Molecular bonds are embedded in as a 7 dimensional vector following (Coley et al., 2017) (see
Revised: Molecular bonds are embedded in as a 7 dimensional vector following (Coley et al., 2017) (see

Original: Table 5 from Appendix A).
Revised: Table 5 from Appendix A).

Original: When the two atoms are not connected by a true molecular bond, alldimensions are set to zeros.
Revised: When the two atoms are not connected by a true molecular bond, alldimensions are set to zeros.

Original: 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 on smaller datasets.
Revised: 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 on smaller datasets.

Original: 3. R ELATIVE M OLECULE S ELF -A TTENTION
Revised: 3. R ELATIVE M OLECULE S ELF -A TTENTION

Original: Equipped with the embedding b ij for each pair of atoms in the molecule, we now use it to deﬁne a novel self-attention layer that we refer to as Relative Molecule Self-Attention.
Revised: Equipped with the embedding b ij for each pair of atoms in the molecule, we now use it to deﬁne a novel self-attention layer that we refer to as Relative Molecule Self-Attention.

Original: First, mirroring the key-query-value design in the vanilla self-attention (c.f. Equation (1)), we transform b ij into a key and value speciﬁc vectors b Vij , b Kij using two neural networks φ V and φ K .
Revised: First, mirroring the key-query-value design in the vanilla self-attention (c.f. Equation (1)), we transform b ij into a key and value speciﬁc vectors b Vij , b Kij using two neural networks φ V and φ K .

Original: Each neural network consists of two layers.
Revised: Each neural network consists of two layers.

Original: A hidden layer, shared between all attention heads and output layer, that create a separate relative embedding for different attention heads.
Revised: A hidden layer, shared between all attention heads and output layer, that create a separate relative embedding for different attention heads.

Original: Consider Equation (1) in index notation:
Revised: Consider Equation (1) in index notation:

Original: [Equation]
Revised: [Equation]

Original: where the unnormalized attention is e ij =
Revised: where the unnormalized attention is e ij =

Original: ( x i W Q )( x j W K ) T .
Revised: ( x i W Q )( x j W K ) T .

Original: By analogy, in Relative Molecule
Revised: By analogy, in Relative Molecule

Original: Self-Attention, we compute e ij as
Revised: Self-Attention, we compute e ij as

Original: [Equation]
Revised: [Equation]

Original: where u , v ∈ R D (cid:48) are trainable vectors.
Revised: where u , v ∈ R D (cid:48) are trainable vectors.

Original: We then deﬁne Relative Molecule Self-Attention operation:
Revised: We then deﬁne Relative Molecule Self-Attention operation:

Original: [Equation]
Revised: [Equation]

Original: In other words, we enrich the self-attention layer with atom relations embedding.
Revised: In other words, we enrich the self-attention layer with atom relations embedding.

Original: In the phase of attention weights calculation, we add content-dependent positional bias, global context bias and global positional bias (Dai et al., 2019; Huang et al., 2020) (that are calculated based on b Kij ) to the layer.
Revised: In the phase of attention weights calculation, we add content-dependent positional bias, global context bias and global positional bias (Dai et al., 2019; Huang et al., 2020) (that are calculated based on b Kij ) to the layer.

Original: Then, during calculation of the attention weighted average, we also include the information about the other embedding b V ij .
Revised: Then, during calculation of the attention weighted average, we also include the information about the other embedding b V ij .

Original: 3. R ELATIVE M OLECULE A TTENTION T RANSFORMER
Revised: 3. R ELATIVE M OLECULE A TTENTION T RANSFORMER

Original: Finally, we use Relative Molecule Self-Attention to construct Relative Molecule Attention Transformer (R-MAT).
Revised: Finally, we use Relative Molecule Self-Attention to construct Relative Molecule Attention Transformer (R-MAT).

Original: The key changes compared to MAT are: (1) the use of Relative Molecule SelfAttention, (2) extended atom featurization, and (3) extended pretraining procedure.
Revised: The key changes compared to MAT are: (1) the use of Relative Molecule SelfAttention, (2) extended atom featurization, and (3) extended pretraining procedure.

Original: Figure 1 illustrates the R-MAT architecture.
Revised: Figure 1 illustrates the R-MAT architecture.

Original: The input is embedded as a matrix of size N atom × 36 where each atom of the input is embedded following (Coley et al., 2017; Pocha et al., 2020), see Table 6 of Appendix A.
Revised: The input is embedded as a matrix of size N atom × 36 where each atom of the input is embedded following (Coley et al., 2017; Pocha et al., 2020), see Table 6 of Appendix A.

Original: We process the input using N stacked Relative Molecule Self-Attention attention layers.
Revised: We process the input using N stacked Relative Molecule Self-Attention attention layers.

Original: Each attention layer is followed by position-wise feed-forward Network (similar as in the classical transformer model (Vaswani et al., 2017)), which consists of 2 linear layers with a leaky-ReLU nonlinearity between them.
Revised: Each attention layer is followed by position-wise feed-forward Network (similar as in the classical transformer model (Vaswani et al., 2017)), which consists of 2 linear layers with a leaky-ReLU nonlinearity between them.

Original: After processing the input using attention layers, we pool the representation into a constant-sized vector.
Revised: After processing the input using attention layers, we pool the representation into a constant-sized vector.

Original: We replace simple mean pooling with an attention-based pooling layer.
Revised: We replace simple mean pooling with an attention-based pooling layer.

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: where H is the hidden state obtained from self-attention layers, W 1 ∈ R P × D and W 2 ∈ R S × P are pooling attention weights, with P equal to the pooling hidden dimension and S equal to the number of pooling attention heads.
Revised: where H is the hidden state obtained from self-attention layers, W 1 ∈ R P × D and W 2 ∈ R S × P are pooling attention weights, with P equal to the pooling hidden dimension and S equal to the number of pooling attention heads.

Original: Finally, the graph embedding g is then passed to the two layer MLP, with leaky-ReLU activation in order to make the prediction.
Revised: Finally, the graph embedding g is then passed to the two layer MLP, with leaky-ReLU activation in order to make the prediction.

Original: Pretraining We used two-step pretraining procedure.
Revised: Pretraining We used two-step pretraining procedure.

Original: In the ﬁrst 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.
Revised: In the ﬁrst 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.

Original: The goal of the task is to predict the whole atom context.
Revised: The goal of the task is to predict the whole atom context.

Original: This task is much more demanding for the network than the classical masking approach presented by (Maziarka et al., 2020) since the network has to encode more speciﬁc information about the masked atom neighbourhood.
Revised: This task is much more demanding for the network than the classical masking approach presented by (Maziarka et al., 2020) since the network has to encode more speciﬁc information about the masked atom neighbourhood.

Original: Furthermore, the size of the context vocabulary is much bigger than the size of the atoms vocabulary in the MAT pretraining approach.
Revised: Furthermore, the size of the context vocabulary is much bigger than the size of the atoms vocabulary in the MAT pretraining approach.

Original: The second task is a graph-level prediction proposed by (Fabian et al., 2020) in which the goal is to predict a set of real-valued descriptors of physicochemical properties.
Revised: The second task is a graph-level prediction proposed by (Fabian et al., 2020) in which the goal is to predict a set of real-valued descriptors of physicochemical properties.

Original: For more detailed information about the pretraining procedure and ablations, see Appendix B.
Revised: For more detailed information about the pretraining procedure and ablations, see Appendix B.

Original: Other details Similarly to (Maziarka et al., 2020), we add an artiﬁcial dummy node to the input molecule.
Revised: Other details Similarly to (Maziarka et al., 2020), we add an artiﬁcial dummy node to the input molecule.

Original: 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 Table 4 of Appendix A).
Revised: 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 Table 4 of Appendix A).

Original: Moreover, the dummy node has its own index in the input atom embedding.
Revised: Moreover, the dummy node has its own index in the input atom embedding.

Original: We calculate distance information in the similar manner as (Maziarka et al., 2020).
Revised: We calculate distance information in the similar manner as (Maziarka et al., 2020).

Original: The 3D molecular conformations that are used to obtain distance matrices are calculated using
Revised: The 3D molecular conformations that are used to obtain distance matrices are calculated using

Original: UFFO PTIMIZE M OLECULE function from the RDKit package (Landrum, 2016) with the default parameters.
Revised: UFFO PTIMIZE M OLECULE function from the RDKit package (Landrum, 2016) with the default parameters.

Original: Finally, we consider a variant of the model extended with 200 rdkit features as in (Rong et al., 2020).
Revised: Finally, we consider a variant of the model extended with 200 rdkit features as in (Rong et al., 2020).

Original: The features are concatenated to the ﬁnal embedding g and processed using a prediction
Revised: The features are concatenated to the ﬁnal embedding g and processed using a prediction

Original: MLP.
Revised: MLP.

Original:  E XPERIMENTS 4. S MALL HYPERPARAMETER BUDGET
Revised:  E XPERIMENTS 4. S MALL HYPERPARAMETER BUDGET

Original: The industrial drug discovery pipelines focus on fast iterations of compound screenings and adjusting the models to new data incoming from the laboratory.
Revised: The industrial drug discovery pipelines focus on fast iterations of compound screenings and adjusting the models to new data incoming from the laboratory.

Original: We start by comparing in this setting R-MAT to DMPNN (Yang et al., 2019a), MAT (Maziarka et al., 2020) and GROVER (Rong et al., 2020), representative state-of-the-art models on popular molecular property prediction tasks.
Revised: We start by comparing in this setting R-MAT to DMPNN (Yang et al., 2019a), MAT (Maziarka et al., 2020) and GROVER (Rong et al., 2020), representative state-of-the-art models on popular molecular property prediction tasks.

Original: 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.
Revised: 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.

Original: The BBBP and Estrogen- β datasets use scaffold splits, while all the other datasets use random splits.
Revised: The BBBP and Estrogen- β datasets use scaffold splits, while all the other datasets use random splits.

Original: Splits were proposed by (Maziarka et al., 2020).
Revised: Splits were proposed by (Maziarka et al., 2020).

Original: For every dataset we calculate scores based on different splits, we report the mean test score based on the hyperparameters that obtained the best validation score, in parentheses we include the standard deviation.
Revised: For every dataset we calculate scores based ondifferent splits, we report the mean test score based on the hyperparameters that obtained the best validation score, in parentheses we include the standard deviation.

Original: In this and the next experiments, we denote models extended with additional rdkit features (see Section 3.5) as GROVER rdkit and
Revised: In this and the next experiments, we denote models extended with additional rdkit features (see Section 3.5) as GROVER rdkit and

Original: R-MAT rdkit .
Revised: R-MAT rdkit .

Original: More information about the models and datasets used in this benchmark are given in
Revised: More information about the models and datasets used in this benchmark are given in

Original: Appendix C.4.
Revised: Appendix C.4.

Original: Table 1 shows that R-MAT outperforms other methods in 3 out of 6 tasks.
Revised: Table 1 shows that R-MAT outperforms other methods in 3 out of 6 tasks.

Original: For comparison, we also cite representative results of other methods from (Maziarka et al., 2020).
Revised: For comparison, we also cite representative results of other methods from (Maziarka et al., 2020).

Original: Satisfyingly, we observe a marked improvement on the solubility prediction tasks (ESOL and FreeSolv).
Revised: Satisfyingly, we observe a marked improvement on the solubility prediction tasks (ESOL and FreeSolv).

Original: Understanding solubility depends to a large degree on a detailed understanding of spatial relationships between atoms.
Revised: Understanding solubility depends to a large degree on a detailed understanding of spatial relationships between atoms.

Original: This suggests that the improvement in performance might be related to better utilization of the distance or graph information
Revised: This suggests that the improvement in performance might be related to better utilization of the distance or graph information

Original: [Table] [Table]
Revised: [Table] [Table]

Original: 4. L ARGE HYPERPARAMETER BUDGET
Revised: 4. L ARGE HYPERPARAMETER BUDGET

Original: In contrast to the previous setting, we test R-MAT against a similar set of models but using a large-scale hyperparameter search (300 different hyperparameter combinations).
Revised: In contrast to the previous setting [Table] we test R-MAT against a similar set of models but using a large-scale hyperparameter search (300 different hyperparameter combinations).

Original: This setting has been proposed in (Rong et al., 2020).
Revised: This setting has been proposed in (Rong et al., 2020).

Original: For comparison, we include results under small (7 different learning rates) hyperparameter budget.
Revised: For comparison, we include results under small (7 different learning rates) hyperparameter budget.

Original: All datasets use a scaffold split.
Revised: All datasets use a scaffold split.

Original: Scores are calculated based on 3 different data splits.
Revised: Scores are calculated based on 3 different data splits.

Original: While the ESOL and FreeSolv datasets are the same as in the previous paragraph, here they use a scaffold split and labels are not normalized (unlike in the previous paragraph).
Revised: While the ESOL and FreeSolv datasets are the same as in the previous paragraph, here they use a scaffold split and labels are not normalized (unlike in the previous paragraph).

Original: Additional information about the models and datasets used in this benchmark are given in Appendix C.5.
Revised: Additional information about the models and datasets used in this benchmark are given in Appendix C.5.

Original: Table 2 summmarizes the experiment.
Revised: Table 2 summmarizes the experiment.

Original: Results show that  R-MAT outperforms other methods in 3 out of 6 tasks  in large grid mode.
Revised: Results show that for large hyperparameters budget R-MAT outperforms other methods in 2 tasks and along with GROVER are the best in one more task.

Original: And in 4 out of 6 for tuning only the learning rate mode.
Revised: 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 7).

Original: 4. L ARGE - SCALE EXPERIMENTS
Revised: 4. L ARGE - SCALE EXPERIMENTS

Original: Finally, to better understand how R-MAT performs in a setting where pretraining is likely to less inﬂuence results, we include results on QM9 dataset (Ramakrishnan et al., 2014).
Revised: Finally, to better understand how R-MAT performs in a setting where pretraining is likely to less inﬂuence results, we include results on QM9 dataset (Ramakrishnan et al., 2014).

Original: QM9 is a quantum mechanics benchmark that encompasses prediction of 12 simulated properties across around 130k small molecules with at most 9 heavy (non-hydrogen) atoms.
Revised: QM9 is a quantum mechanics benchmark that encompasses prediction of 12 simulated properties across around 130k small molecules with at most 9 heavy (non-hydrogen) atoms.

Original: The molecules are provided with their atomic 3D positions for which the quantum properties were initially calculated.
Revised: The molecules are provided with their atomic 3D positions for which the quantum properties were initially calculated.

Original: For these experiments, we used learning rate equal to 0.015 (we selected this learning rate value as it returned the best results for α dataset among 4 different learning rates that we tested: {0.005,0.01,0.015,0.02}).
Revised: For these experiments, we used learning rate equal to 0.015 (we selected this learning rate value as it returned the best results for α dataset among 4 different learning rates that we tested: {0.005,0.01,0.015,0.02}).

Original: Additional information about the dataset and models used in this benchmark are given in Appendix C.
Revised: Additional information about the dataset and models used in this benchmark are given in Appendix C.

Original: Figure 3 compares R-MAT performance with various models.

Original: More detailed results could be ﬁnd in Table 10 from Appendix D.3.

Original: R-MAT achieves highly competitive results, with state-of-the-art

Original: [Table]
Revised: [Table]

Original: 936 ( . 008 ) R-MAT rdkit .
Revised: 936 ( . 008 ) R-MAT rdkit .

Original: MAT .
Revised: MAT .

Original: GROVER .
Revised: GROVER .

Original: 908 ( . 053) GROVER rdkit .
Revised: 908 ( . 053) GROVER rdkit .

Original: R-MAT .
Revised: R-MAT .

Original: R-MAT rdkit .
Revised: R-MAT rdkit .

Original: 920 ( . 021) performance on 4 out of the 12 tasks.
Revised: Figure 3 compares R-MAT performance with various models. More detailed results could be ﬁnd in Table 10 from Appendix D.3. R-MAT achieves highly competitive results, with state-of-the-art performance on 4 out of the 12 tasks.

Original: We attribute higher variability of performance to the limited small hyperparameter search we performed.
Revised: We attribute higher variability of performance to the limited small hyperparameter search we performed.

Original: 4. E XPLORING THE DESIGN SPACE OF SELF - ATTENTION LAYER
Revised: 4. E XPLORING THE DESIGN SPACE OF SELF - ATTENTION LAYER

Original: Achieving strong empirical results hinged on a methodologically exploration the design space of different variants of the self-attention layer.
Revised: Achieving strong empirical results hinged on a methodologically exploration the design space of different variants of the self-attention layer.

Original: We document here this exploration and relevant ablations.
Revised: We document here this exploration and relevant ablations.

Original: Due to space limitations, we defer most results to the
Revised: Due to space limitations, we defer most results to the

Original: Appendix E. We perform all experiments on the ESOL, FreeSolv and
Revised: Appendix E. We perform all experiments on the ESOL, FreeSolv and

Original: BBBP datasets with 3 different scaffold splits.
Revised: BBBP datasets with 3 different scaffold splits.

Original: We did not use any pretraining for these experiments.
Revised: We did not use any pretraining for these experiments.

Original: We follow the same ﬁne-tuning methodology as in Section 4.1.
Revised: We follow the same ﬁne-tuning methodology as in Section 4.1.

Original:  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.
Revised: 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.

Original: In Table 3 (Left), we show the effect on performance of ablating each of this elements.
Revised: In Table 3 (Left), we show the effect on performance of ablating each of this elements.

Original: Importantly, we ﬁnd that each component is important to R-MAT performance, including the distance matrix.
Revised: Importantly, we ﬁnd that each component is important to R-MAT performance, including the distance matrix.

Original: Maximum neighbourhood order We take a closer look at how we encode the molecular graph.
Revised: Maximum neighbourhood order We take a closer look at how we encode the molecular graph.

Original: (Maziarka et al., 2020) used a simple binary adjacency matrix to encode the edges.
Revised: (Maziarka et al., 2020) used a simple binary adjacency matrix to encode the edges.

Original: We enriched this representation by adding one-hot encoding of the neighbourhood order.
Revised: We enriched this representation by adding one-hot encoding of the neighbourhood order.

Original: For example, the order offor a pair of atoms means that there are two other vertices on the shortest path between this pair of atoms.
Revised: For example, the order offor a pair of atoms means that there are two other vertices on the shortest path between this pair of atoms.

Original: In R-MAT we used 4 as the maximum order of neighbourhood distance.
Revised: In R-MAT we used 4 as the maximum order of neighbourhood distance.

Original: That is, we encoded as separate features if two atoms are 1, 2, 3 or 4 hops away in the molecular graph.
Revised: That is, we encoded as separate features if two atoms are 1, 2, 3 or 4 hops away in the molecular graph.

Original: In Table 3 (Right) we ablate this choice.
Revised: In Table 3 (Right) [Table] [Table] [Table]

Revised: we ablate this choice.

Original: The result suggests that R-MAT performance beneﬁts from including separate feature for all the consideredorders.
Revised: The result suggests that R-MAT performance beneﬁts from including separate feature for all the considered

Original: [Table]
Revised: [Table]

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original:  patterns.
Revised: attention pattern in MAT is dominated by the adjacency and distance matrix, while R-MAT seems capable of learning more complex attention patterns.

Original: 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.
Revised: 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.

Original: We qualitatively explore here whether R-MAT achieves this goal, comparing its attention patterns to that of MAT.
Revised: We qualitatively explore here whether R-MAT achieves this goal, comparing its attention patterns to that of MAT.

Original: From the Figureone can see that indeed R-MAT seems capable of learning more complex attention patterns than MAT.
Revised: From the Figureone can see that indeed R-MAT seems capable of learning more complex attention patterns than MAT.

Original: We add a more detailed comparison, with more visualised attention heads in Appendix D.4.
Revised: We add a more detailed comparison, with more visualised attention heads in Appendix D.4.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: Transformer has been successfuly adapted to various domain by incorporating into its architecture a minimal set of inductive biases.
Revised: Transformer has been successfuly adapted to various domain by incorporating into its architecture a minimal set of inductive biases.

Original: In a similar spirit, we methodologically explored the design space of the selfattention layer, and identiﬁed a highly effective Relative Molecule Self-Attention.
Revised: In a similar spirit, we methodologically explored the design space of the selfattention layer, and identiﬁed a highly effective Relative Molecule Self-Attention.

Original: Relative Molecule Attention Transformer, a model based on Relative Molecule SelfAttention, achieves state-of-the-art or very competitive results across a wide range of molecular property prediction tasks.
Revised: Relative Molecule Attention Transformer, a model based on Relative Molecule SelfAttention, achieves state-of-the-art or very competitive results across a wide range of molecular property prediction tasks.

Original: R-MAT is a highly versatile model, showing state-of-the-art results in both quantum property prediction tasks, as well as on biological datasets.
Revised: R-MAT is a highly versatile model, showing state-of-the-art results in both quantum property prediction tasks, as well as on biological datasets.

Original: We also show that R-MAT is easy to train and requires tuning only the learning rate to achieve competitive results, which together with open-sourced weight and code, makes it is highly accessible.
Revised: We also show that R-MAT is easy to train and requires tuning only the learning rate to achieve competitive results, which together with open-sourced weight and code, makes it is highly accessible.

Original: Relative Molecule Self-Attention encodes an inductive bias to consider relationships between atoms that are commonly relevant to a chemist, but on the other hand leaves ﬂexibility to unlearn them if needed.
Revised: Relative Molecule Self-Attention encodes an inductive bias to consider relationships between atoms that are commonly relevant to a chemist, but on the other hand leaves ﬂexibility to unlearn them if needed.

Original: Relatedly, Vision Transformers learn global processing in early layers despite being equipped with a locality inductive bias (Dosovitskiy et al., 2021).
Revised: Relatedly, Vision Transformers learn global processing in early layers despite being equipped with a locality inductive bias (Dosovitskiy et al., 2021).

Original: Our empirical results show in a new context that picking the right set of inductive biases is key for self-supervised learning to work well.
Revised: Our empirical results show in a new context that picking the right set of inductive biases is key for self-supervised learning to work well.

Original: We also how that Relative Molecule Self-Attention will help improve other models for molecular property prediction.
Revised: We also how that Relative Molecule Self-Attention will help improve other models for molecular property prediction.

Original: Learning useful representations for molecular property prediction is far from solved.
Revised: Learning useful representations for molecular property prediction is far from solved.

Original: Achieving stateof-the-art results, while less dependent on them, still relied on using certain large sets of handcrafted features both in ﬁne-tuning and pretraining.
Revised: Achieving stateof-the-art results, while less dependent on them, still relied on using certain large sets of handcrafted features both in ﬁne-tuning and pretraining.

Original: At the same time, these features are beyond doubt learnable from data.
Revised: At the same time, these features are beyond doubt learnable from data.

Original: Developing methods that will push representation learning towards discovering these and better features automatically from data is an exciting challenge for the future.
Revised: Developing methods that will push representation learning towards discovering these and better features automatically from data is an exciting challenge for the future.

Original: A R-MAT NODE AND EDGE FEATURES
Revised: A R-MAT NODE AND EDGE FEATURES

Original: In the following section
Revised: In the following section

Original: [Table]
Revised: [Table]

Original: we present the node and edge features used by R-MAT.
Revised: we present the node and edge features used by R-MAT.

Original: A. E DGE FEATURES
Revised: A. E DGE FEATURES

Original: In R-MAT, all atoms are connected with an edge.
Revised: In R-MAT, all atoms are connected with an edge.

Original: The vector representation of every edge contains information about atoms neighbourhood, distances between them and physicochemical features of a bond if it exists (see Figure 2)
Revised: The vector representation of every edge contains information about atoms neighbourhood, distances between them and physicochemical features of a bond if it exists (see Figure 2)

Original: Neighbourhood embeddings The neighbourhood information of an atom pair is represented by a 6 -dimensional one-hot encoded vector, with features presented in Table 4.
Revised: Neighbourhood embeddings The neighbourhood information of an atom pair is represented by a 6 -dimensional one-hot encoded vector, with features presented in Table 4.

Original: Every neighbourhood embedding contains the information about how many other vertices are between nodes i and j in the original molecular graph.
Revised: Every neighbourhood embedding contains the information about how many other vertices are between nodes i and j in the original molecular graph.

Original: [Table]
Revised: [Table]

Original: Bond embeddings Molecular bonds are embedded in a 7 -dimensional vector following (Coley et al., 2017), with features speciﬁed in Table 5.
Revised: Bond embeddings Molecular bonds are embedded in a 7 -dimensional vector following (Coley et al., 2017), with features speciﬁed in Table 5.

Original: When the two atoms are not connected by a true molecular bond, all 7 dimensions are set to zeros.
Revised: When the two atoms are not connected by a true molecular bond, all 7 dimensions are set to zeros.

Original: [Table]
Revised: [Table]

Original: A. N ODE FEATURES
Revised: A. N ODE FEATURES

Original: The input molecule is embedded as a matrix of size N atom × 36 where each atom of the input is embedded following (Coley et al., 2017; Pocha et al., 2020).
Revised: The input molecule is embedded as a matrix of size N atom × 36 where each atom of the input is embedded following (Coley et al., 2017; Pocha et al., 2020).

Original: All features are presented in Table
Revised: All features are presented in Table

Original: [Table]
Revised: [Table]

Original: B P RETRAINING
Revised: B P RETRAINING

Original: We extend the pretraining procedure of (Maziarka et al., 2020), who used a masking task based on (Devlin et al., 2018;
Revised: We extend the pretraining procedure of (Maziarka et al., 2020), who used a masking task based on (Devlin et al., 2018;

Original: 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.
Revised: 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.

Original: Such approach works well in NLP where models pretrained with the masking task create the state-of-the-art representation (Devlin et al., 2018;
Revised: Such approach works well in NLP where models pretrained with the masking task create the state-of-the-art representation (Devlin et al., 2018;

Original: Liu et al., 2019; Yang et al., 2019b).
Revised: Liu et al., 2019; Yang et al., 2019b).

Original: However in chemistry, otherwise than in NLP, the size of atoms vocabulary is much smaller.
Revised: However in chemistry, otherwise than in NLP, the size of atoms vocabulary is much smaller.

Original: Moreover, usually only one type of atom ﬁts a given place and thus the
Revised: Moreover, usually only one type of atom ﬁts a given place and thus the

Original: [Table]
Revised: [Table]

Original: representation trained with the masking task has problems with encoding meaningful information in chemistry.
Revised: representation trained with the masking task has problems with encoding meaningful information in chemistry.

Original: B. C ONTEXTUAL PRETRAINING
Revised: B. C ONTEXTUAL PRETRAINING

Original: Instead of atom masking, we used a two-step pretraining that combines the procedures proposed by (Rong et al., 2020; Fabian et al., 2020).
Revised: Instead of atom masking, we used a two-step pretraining that combines the procedures proposed by (Rong et al., 2020; Fabian et al., 2020).

Original: In the ﬁrst step, the network is trained with the contextual property prediction task (Rong et al., 2020), where we mask not only the selected atoms, but also their neighbours.
Revised: In the ﬁrst step, the network is trained with the contextual property prediction task (Rong et al., 2020), where we mask not only the selected atoms, but also their neighbours.

Original: The task is then to predict the whole atom context, e.g. if the selected atom’s type is carbon connected with a nitrogen with a double bond and with an oxygen with a single bond, we encode the atom neighbourhood as C_N-DOUBLE1_O-SINGLE1 (we list all the node-edge counts terms in the alphabetical order), then the network has to predict the speciﬁc type of the masked neighbourhood for every masked atom.
Revised: The task is then to predict the whole atom context, e.g. if the selected atom’s type is carbon connected with a nitrogen with a double bond and with an oxygen with a single bond, we encode the atom neighbourhood as C_N-DOUBLE1_O-SINGLE1 (we list all the node-edge counts terms in the alphabetical order), then the network has to predict the speciﬁc type of the masked neighbourhood for every masked atom.

Original: 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 speciﬁc information about the masked atom’s neighbourhood.
Revised: 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 speciﬁc information about the masked atom’s neighbourhood.

Original: Furthermore, the size of context vocabulary is much bigger than the size of atoms vocabulary in the MAT pretraining approach ( for R-MAT vs 35 for MAT).
Revised: Furthermore, the size of context vocabulary is much bigger than the size of atoms vocabulary in the MAT pretraining approach ( for R-MAT vs 35 for MAT).

Original: B. G RAPH - LEVEL PRETRAINING
Revised: B. G RAPH - LEVEL PRETRAINING

Original: The second task is the graph-level property prediction proposed by (Fabian et al., 2020).
Revised: The second task is the graph-level property prediction proposed by (Fabian et al., 2020).

Original: In this pretraining procedure, the task is to predict 200 real-valued descriptors of physicochemical characteristics of every given molecule.
Revised: In this pretraining procedure, the task is to predict 200 real-valued descriptors of physicochemical characteristics of every given molecule.

Original: The list of all 200 descriptors from RDKit is as follows:
Revised: The list of all 200 descriptors from RDKit is as follows:

Original: BalabanJ, BertzCT, Chi0, Chi0n, Chi0v, Chi1, Chi1n, Chi1v, Chi2n,
Revised: BalabanJ, BertzCT, Chi0, Chi0n, Chi0v, Chi1, Chi1n, Chi1v, Chi2n,

Original: Chi2v, Chi3n, Chi3v, Chi4n, Chi4v, EState_VSA1, EState_VSA10,
Revised: Chi2v, Chi3n, Chi3v, Chi4n, Chi4v, EState_VSA1, EState_VSA10,

Original: EState_VSA11, EState_VSA2, EState_VSA3, EState_VSA4, EState_VSA5,
Revised: EState_VSA11, EState_VSA2, EState_VSA3, EState_VSA4, EState_VSA5,

Original: EState_VSA6, EState_VSA7, EState_VSA8, EState_VSA9, ExactMolWt,
Revised: EState_VSA6, EState_VSA7, EState_VSA8, EState_VSA9, ExactMolWt,

Original: FpDensityMorgan1, FpDensityMorgan2, FpDensityMorgan3, FractionCSP3,
Revised: FpDensityMorgan1, FpDensityMorgan2, FpDensityMorgan3, FractionCSP3,

Original: HallKierAlpha, HeavyAtomCount, HeavyAtomMolWt, Ipc, Kappa1, Kappa2,
Revised: HallKierAlpha, HeavyAtomCount, HeavyAtomMolWt, Ipc, Kappa1, Kappa2,

Original: Kappa3, LabuteASA, MaxAbsEStateIndex, MaxAbsPartialCharge, MaxEStateIndex,
Revised: Kappa3, LabuteASA, MaxAbsEStateIndex, MaxAbsPartialCharge, MaxEStateIndex,

Original: MaxPartialCharge, MinAbsEStateIndex, MinAbsPartialCharge, MinEStateIndex,
Revised: MaxPartialCharge, MinAbsEStateIndex, MinAbsPartialCharge, MinEStateIndex,

Original: MinPartialCharge, MolLogP, MolMR, MolWt, NHOHCount, NOCount,
Revised: MinPartialCharge, MolLogP, MolMR, MolWt, NHOHCount, NOCount,

Original: NumAliphaticCarbocycles, NumAliphaticHeterocycles, NumAliphaticRings,
Revised: NumAliphaticCarbocycles, NumAliphaticHeterocycles, NumAliphaticRings,

Original: NumAromaticCarbocycles, NumAromaticHeterocycles, NumAromaticRings,
Revised: NumAromaticCarbocycles, NumAromaticHeterocycles, NumAromaticRings,

Original: NumHAcceptors, NumHDonors, NumHeteroatoms, NumRadicalElectrons,
Revised: NumHAcceptors, NumHDonors, NumHeteroatoms, NumRadicalElectrons,

Original: NumRotatableBonds, NumSaturatedCarbocycles, NumSaturatedHeterocycles,
Revised: NumRotatableBonds, NumSaturatedCarbocycles, NumSaturatedHeterocycles,

Original: NumSaturatedRings, NumValenceElectrons, PEOE_VSA1, PEOE_VSA10,
Revised: NumSaturatedRings, NumValenceElectrons, PEOE_VSA1, PEOE_VSA10,

Original: PEOE_VSA11, PEOE_VSA12, PEOE_VSA13, PEOE_VSA14, PEOE_VSA2, PEOE_VSA3,
Revised: PEOE_VSA11, PEOE_VSA12, PEOE_VSA13, PEOE_VSA14, PEOE_VSA2, PEOE_VSA3,

Original: PEOE_VSA4, PEOE_VSA5, PEOE_VSA6, PEOE
Revised: PEOE_VSA4, PEOE_VSA5, PEOE_VSA6, PEOE

Original: [Table]
Revised: [Table]

Original: SMR_VSA6, SMR_VSA7, SMR_VSA8, SMR_VSA9, SlogP_VSA1, SlogP_VSA10,
Revised: SMR_VSA6, SMR_VSA7, SMR_VSA8, SMR_VSA9, SlogP_VSA1, SlogP_VSA10,

Original: SlogP_VSA11, SlogP_VSA12, SlogP_VSA2, SlogP_VSA3, SlogP_VSA4, SlogP_VSA5,
Revised: SlogP_VSA11, SlogP_VSA12, SlogP_VSA2, SlogP_VSA3, SlogP_VSA4, SlogP_VSA5,

Original: SlogP_VSA6, SlogP_VSA7, SlogP_VSA8, SlogP_VSA9, TPSA, VSA_EState1,
Revised: SlogP_VSA6, SlogP_VSA7, SlogP_VSA8, SlogP_VSA9, TPSA, VSA_EState1,

Original: VSA_EState10, VSA_EState2, VSA_EState3, VSA_EState4, VSA_EState5,
Revised: VSA_EState10, VSA_EState2, VSA_EState3, VSA_EState4, VSA_EState5,

Original: VSA_EState6, VSA_EState7, VSA_EState8, VSA_EState9, fr_Al_COO, fr_Al_OH, fr_Al_OH_noTert, fr_ArN, fr_Ar_COO, fr_Ar_N, fr_Ar_NH, fr_Ar_OH, fr_COO, fr_COO2, fr_C_O, fr_C_O_noCOO, fr_C_S, fr_HOCCN, fr_Imine, fr_NH0, fr_NH1, fr_NH2, fr_N_O, fr_Ndealkylation1, fr_Ndealkylation2, fr_Nhpyrrole, fr_SH, fr_aldehyde, fr_alkyl_carbamate, fr_alkyl_halide, fr_allylic_oxid, fr_amide, fr_amidine, fr_aniline, fr_aryl_methyl, fr_azide, fr_azo, fr_barbitur, fr_benzene, fr_benzodiazepine, fr_bicyclic, fr_diazo, fr_dihydropyridine, fr_epoxide, fr_ester, fr_ether, fr_furan, fr_guanido, fr_halogen, fr_hdrzine, fr_hdrzone, fr_imidazole, fr_imide, fr_isocyan, fr_isothiocyan, fr_ketone, fr_ketone_Topliss, fr_lactam, fr_lactone, fr_methoxy, fr_morpholine, fr_nitrile, fr_nitro, fr_nitro_arom, fr_nitro_arom_nonortho, fr_nitroso, fr_oxazole, fr_oxime, fr_para_hydroxylation, fr_phenol, fr_phenol_noOrthoHbond, fr_phos_acid, fr_phos_ester, fr_piperdine, fr_piperzine, fr_priamide, fr_prisulfonamd, fr_pyridine, fr_quatN, fr_sulfide, fr_sulfonamd, fr_sulfone, fr_term_acetylene, fr_tetrazole, fr_thiazole, fr_thiocyan, fr_thiophene, fr_unbrch_alkane, fr_urea, qed
Revised: VSA_EState6, VSA_EState7, VSA_EState8, VSA_EState9, fr_Al_COO, fr_Al_OH, fr_Al_OH_noTert, fr_ArN, fr_Ar_COO, fr_Ar_N, fr_Ar_NH, fr_Ar_OH, fr_COO, fr_COO2, fr_C_O, fr_C_O_noCOO, fr_C_S, fr_HOCCN, fr_Imine, fr_NH0, fr_NH1, fr_NH2, fr_N_O, fr_Ndealkylation1, fr_Ndealkylation2, fr_Nhpyrrole, fr_SH, fr_aldehyde, fr_alkyl_carbamate, fr_alkyl_halide, fr_allylic_oxid, fr_amide, fr_amidine, fr_aniline, fr_aryl_methyl, fr_azide, fr_azo, fr_barbitur, fr_benzene, fr_benzodiazepine, fr_bicyclic, fr_diazo, fr_dihydropyridine, fr_epoxide, fr_ester, fr_ether, fr_furan, fr_guanido, fr_halogen, fr_hdrzine, fr_hdrzone, fr_imidazole, fr_imide, fr_isocyan, fr_isothiocyan, fr_ketone, fr_ketone_Topliss, fr_lactam, fr_lactone, fr_methoxy, fr_morpholine, fr_nitrile, fr_nitro, fr_nitro_arom, fr_nitro_arom_nonortho, fr_nitroso, fr_oxazole, fr_oxime, fr_para_hydroxylation, fr_phenol, fr_phenol_noOrthoHbond, fr_phos_acid, fr_phos_ester, fr_piperdine, fr_piperzine, fr_priamide, fr_prisulfonamd, fr_pyridine, fr_quatN, fr_sulfide, fr_sulfonamd, fr_sulfone, fr_term_acetylene, fr_tetrazole, fr_thiazole, fr_thiocyan, fr_thiophene, fr_unbrch_alkane, fr_urea, qed

Original: C E XPERIMENTAL SETTING C. M ODEL HYPERPARAMETERS
Revised: C E XPERIMENTAL SETTING C. M ODEL HYPERPARAMETERS

Original: R-MAT model consists of 10 layers with 12 attention heads in each, d model = 768 .
Revised: R-MAT model consists of 10 layers with 12 attention heads in each, d model = 768 .

Original: The distance layer consists of 32 radial functions (N emb = 32 ) and cutoff distance c is set to 20 Å. Attention pooling consists of 4 pooling heads and pooling hidden dimension is set to 128 .
Revised: The distance layer consists of 32 radial functions (N emb = 32 ) and cutoff distance c is set to 20 Å. Attention pooling consists of 4 pooling heads and pooling hidden dimension is set to 128 .

Original: Prediction MLP consists of one hidden layer with dimension set to 1024 and dropout 0 .
Revised: Prediction MLP consists of one hidden layer with dimension set to 1024 and dropout 0 .

Original: R-MAT uses leaky-ReLU with slope 0 .
Revised: R-MAT uses leaky-ReLU with slope 0 .

Original: 1 as a non-linearity in all experiments.
Revised: 1 as a non-linearity in all experiments.

Original: This set of hyperparameters deﬁnes a model with 48 millions of parameters, which is equal to the number of parameters of GROVER base and is slightly higher than the number of parameters of MAT ( 42 M).
Revised: This set of hyperparameters deﬁnes a model with 48 millions of parameters, which is equal to the number of parameters of GROVER base and is slightly higher than the number of parameters of MAT ( 42 M).

Original: C. 3D CONFORMATIONS
Revised: C. 3D CONFORMATIONS

Original: The 3D molecular conformations that are used to obtain distance matrices were calculated using
Revised: The 3D molecular conformations that are used to obtain distance matrices were calculated using

Original: UFFO PTIMIZE M OLECULE function from the RDKit package (Landrum, 2016) with the default parameters ( MAX I TERS = 200 , VDW T HRESH = 10 . 0 , CONF I D = − 1 , IGNORE I NTERFRAG I NTERACTIONS =
Revised: UFFO PTIMIZE M OLECULE function from the RDKit package (Landrum, 2016) with the default parameters ( MAX I TERS = 200 , VDW T HRESH = 10 . 0 , CONF I D = − 1 , IGNORE I NTERFRAG I NTERACTIONS =

Original: True).
Revised: True).

Original: One disadvantage of this approach is the costly calculation of the distance matrix.
Revised: One disadvantage of this approach is the costly calculation of the distance matrix.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: In Figure 5 we present these times for every single molecule from these datasets.
Revised: In Figure 5 we present these times for every single molecule from these datasets.

Original: Based on these results one can see, that the larger the molecule, the longer the conformation calculation time.
Revised: Based on these results one can see, that the larger the molecule, the longer the conformation calculation time.

Original: Moreover even for BBBP, that is not rather a large dataset, calculating the conformations could take almost 5 minutes.
Revised: Moreover even for BBBP, that is not rather a large dataset, calculating the conformations could take almost 5 minutes.

Original: We leave exploring other conformation methods (more accurate or faster) as an interesting topic for the future.
Revised: We leave exploring other conformation methods (more accurate or faster) as an interesting topic for the future.

Original: [Table]
Revised: [Table]

Original: C. P RETRAINING
Revised: C. P RETRAINING

Original: We pretrained R-MAT on 4 millions of unlabelled molecules.
Revised: We pretrained R-MAT on 4 millions of unlabelled molecules.

Original: Molecules were obtained from the
Revised: Molecules were obtained from the

Original: ZINC15 (Sterling & Irwin, 2015) and ChEMBL (Gaulton et al., 2011) datasets, by ﬁrst taking a sample of 10 M molecules and then ﬁltering them using the Lipinski’s rule of ﬁve (Lipinski et al., 1997).
Revised: ZINC15 (Sterling & Irwin, 2015) and ChEMBL (Gaulton et al., 2011) datasets, by ﬁrst taking a sample of 10 M molecules and then ﬁltering them using the Lipinski’s rule of ﬁve (Lipinski et al., 1997).

Original: We have split the data into training and validation datasets, where validation dataset consists of 5% of the data.
Revised: We have split the data into training and validation datasets, where validation dataset consists of 5% of the data.

Original: We pretrained R-MAT for 150 epochs.
Revised: We pretrained R-MAT for 150 epochs.

Original: We used dropout equal to 0 .
Revised: We used dropout equal to 0 .

Original: 1 , learning rate 0 .
Revised: 1 , learning rate 0 .

Original: 001 , the Noam optimizer (Vaswani et al., 2017) trained with 20000 warm-up steps and batch size 256 .
Revised: 001 , the Noam optimizer (Vaswani et al., 2017) trained with 20000 warm-up steps and batch size 256 .

Original: We pretrained R-MAT with 8 Nvidia A100 GPUs, using the Horovod package (Sergeev &
Revised: We pretrained R-MAT with 8 Nvidia A100 GPUs, using the Horovod package (Sergeev &

Original: Del Balso, 2018).
Revised: Del Balso, 2018).

Original: Overlap between pre-trained and task datasets In Table 8 we provide the information about number of atoms that overlaps between pre-trained dataset and all tasks datasets used in this paper.
Revised: Overlap between pre-trained and task datasets In Table 8 we provide the information about number of atoms that overlaps between pre-trained dataset and all tasks datasets used in this paper.

Original: One can see that usually the overlap is at most few percent of all dataset molecules.
Revised: One can see that usually the overlap is at most few percent of all dataset molecules.

Original: This overlap does not affect the test performance, as pre-training consists of semi-supervised task of context prediction and other task, associated with the prediction of simple physico-chemical characteristics, loosely related to the ﬁne-tuning tasks, which are additionally added to the ﬁnal graph embedding during the prediction.
Revised: This overlap does not affect the test performance, as pre-training consists of semi-supervised task of context prediction and other task, associated with the prediction of simple physico-chemical characteristics, loosely related to the ﬁne-tuning tasks, which are additionally added to the ﬁnal graph embedding during the prediction.

Original: [Table]
Revised: [Table]

Original: C. S MALL HYPERPARAMETER BUDGET
Revised: C. S MALL HYPERPARAMETER BUDGET

Original: Models We compared R-MAT with four models trained from scratch:
Revised: Models We compared R-MAT with four models trained from scratch:

Original: Support Vector Machine with RBF kernel (SVM) and Random Forest (RF) that both works on ECFP ﬁngerprints (Rogers &
Revised: Support Vector Machine with RBF kernel (SVM) and Random Forest (RF) that both works on ECFP ﬁngerprints (Rogers &

Original: Hahn, 2010), Graph Convolutional Network (Duvenaud et al., 2015) (GCN) and Directed Message
Revised: Hahn, 2010), Graph Convolutional Network (Duvenaud et al., 2015) (GCN) and Directed Message

Original: Passing Neural Network (Yang et al., 2019a) (DMPNN).
Revised: Passing Neural Network (Yang et al., 2019a) (DMPNN).

Original: The comparison also includes two different pretrained models: MAT (Maziarka et al., 2020) and
Revised: The comparison also includes two different pretrained models: MAT (Maziarka et al., 2020) and

Original: GROVER (Rong et al., 2020).
Revised: GROVER (Rong et al., 2020).

Original: Datasets The benchmark is based on important molecule property tasks in the drug discovery domain.
Revised: Datasets The benchmark is based on important molecule property tasks in the drug discovery domain.

Original: The ﬁrst two datasets are ESOL and FreeSolv, in which the task is to predict the solubility of a molecule in water – a key property of any drug – and the error is measured using RMSE.
Revised: The ﬁrst two datasets are ESOL and FreeSolv, in which the task is to predict the solubility of a molecule in water – a key property of any drug – and the error is measured using RMSE.

Original: The goal in
Revised: The goal in

Original: BBBP and Estrogen − β is to classify correctly whether a given molecule is active against a biological target.
Revised: BBBP and Estrogen − β is to classify correctly whether a given molecule is active against a biological target.

Original: For details on other tasks please see (Maziarka et al., 2020).
Revised: For details on other tasks please see (Maziarka et al., 2020).

Original: BBBP and Estrogen − β used scaffold split, rest datasets used random split method.
Revised: BBBP and Estrogen − β used scaffold split, rest datasets used random split method.

Original: For every dataset 6 different splits were created.
Revised: For every dataset 6 different splits were created.

Original: Labels of regression datasets (ESOL and FreeSolv) were normalized before training.
Revised: Labels of regression datasets (ESOL and FreeSolv) were normalized before training.

Original: We did not include the Estrogen − α 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).
Revised: We did not include the Estrogen − α 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).

Original: Training hyperparameters We ﬁne-tune R-MAT on the target tasks for 100 epochs, with batch size equal to 32 and Noam optimizer with warm-up equal to 30% of all steps.
Revised: Training hyperparameters We ﬁne-tune R-MAT on the target tasks for 100 epochs, with batch size equal to 32 and Noam optimizer with warm-up equal to 30% of all steps.

Original: The only hyperparameter that we tune is the learning rate, which is selected from the set of 7 possible options: { 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 , 5 e − 6 , 1 e − 6 } .
Revised: The only hyperparameter that we tune is the learning rate, which is selected from the set of 7 possible options: { 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 , 5 e − 6 , 1 e − 6 } .

Original: This small budget for hyperparameter selection reﬂects the long-term goal of this paper of developing easy to use models for molecule property prediction.
Revised: This small budget for hyperparameter selection reﬂects the long-term goal of this paper of developing easy to use models for molecule property prediction.

Original: The ﬁne-tuning was conducted using Nvidia V100 GPU.
Revised: The ﬁne-tuning was conducted using Nvidia V100 GPU.

Original: C. L ARGE HYPERPARAMETER BUDGET
Revised: C. L ARGE HYPERPARAMETER BUDGET

Original: Models For large hyperparameters budget we compared R-MAT with three models trained from scratch: GraphConv (Kipf & Welling, 2016), Weave (Kearnes et al., 2016) and DMPNN (Yang et al., 2019a) and two pre-trained models: MAT (Maziarka et al., 2020) and GROVER (Rong et al., 2020).
Revised: Models For large hyperparameters budget we compared R-MAT with three models trained from scratch: GraphConv (Kipf & Welling, 2016), Weave (Kearnes et al., 2016) and DMPNN (Yang et al., 2019a) and two pre-trained models: MAT (Maziarka et al., 2020) and GROVER (Rong et al., 2020).

Original: For small hyperparameters budget we compared R-MAT to MAT and GROVER.
Revised: For small hyperparameters budget we compared R-MAT to MAT and GROVER.

Original: We note that R-MAT, MAT and GROVER use different pretraining methods.
Revised: We note that R-MAT, MAT and GROVER use different pretraining methods.

Original: MAT was pretrained with 2M molecules from ZINC database and GROVER was pretrained with 10M molecules from the ZINC and ChEMBL databases.
Revised: MAT was pretrained with 2M molecules from ZINC database and GROVER was pretrained with 10M molecules from the ZINC and ChEMBL databases.

Original: Datasets All datasets were splitted using a scaffold split.
Revised: Datasets All datasets were splitted using a scaffold split.

Original: The resulting splits are different than in the MAT benchmark.
Revised: The resulting splits are different than in the MAT benchmark.

Original: For every dataset, 3 different splits were created.
Revised: For every dataset, 3 different splits were created.

Original: In our comparison we included only the subset of single-task datasets from original GROVER work (Rong et al., 2020).
Revised: In our comparison we included only the subset of single-task datasets from original GROVER work (Rong et al., 2020).

Original: This is the reason why we use a smaller number of datasets.
Revised: This is the reason why we use a smaller number of datasets.

Original: The obtained regression scores for ESOL differs signiﬁcantly from the small hyperparameters budget benchmark because this time labels are not normalized.
Revised: The obtained regression scores for ESOL differs signiﬁcantly from the small hyperparameters budget benchmark because this time labels are not normalized.

Original: Training hyperparameters For learning rate tuning we used the same hyperparameters settings as in the MAT benchmark (see Appendix C.4).
Revised: Training hyperparameters For learning rate tuning we used the same hyperparameters settings as in the MAT benchmark (see Appendix C.4).

Original: For large hyperparameters budget we run random search with the hyperparameters listed in a Table 9.
Revised: For large hyperparameters budget we run random search with the hyperparameters listed in a Table 9.

Original: [Table]
Revised: [Table]

Original: C. L ARGE - SCALE EXPERIMENTS
Revised: C. L ARGE - SCALE EXPERIMENTS

Original: Models We compared our R-MAT with 8 different models: NMP (Gilmer et al., 2017),
Revised: Models We compared our R-MAT with 8 different models: NMP (Gilmer et al., 2017),

Original: Schnet (Schütt et al., 2017), Cormorant (Anderson et al., 2019), L1Net (Miller et al., 2020),
Revised: Schnet (Schütt et al., 2017), Cormorant (Anderson et al., 2019), L1Net (Miller et al., 2020),

Original: LieConv (Finzi et al., 2020), TFN (Thomas et al., 2018), SE(3)-Tr.
Revised: LieConv (Finzi et al., 2020), TFN (Thomas et al., 2018), SE(3)-Tr.

Original: (Fuchs et al., 2020), EGNN (Satorras et al., 2021).
Revised: (Fuchs et al., 2020), EGNN (Satorras et al., 2021).

Original: Datasets The QM9 dataset (Ramakrishnan et al., 2014) is a dataset that consists of molecules, with up to 9 heavy atoms (H, C, N, O, F) and up to 29 atoms overall per molecule.
Revised: Datasets The QM9 dataset (Ramakrishnan et al., 2014) is a dataset that consists of molecules, with up to 9 heavy atoms (H, C, N, O, F) and up to 29 atoms overall per molecule.

Original: Each atom from this dataset is additionally associated with 3D position.
Revised: Each atom from this dataset is additionally associated with 3D position.

Original: The dataset consists of 12 different regression tasks named: α , ∆ ε , ε HOMO , ε LUMO , µ , C ν , G , H , R 2 , U , U 0 , ZPVE, for which the mean absolute error is a standard metric.
Revised: The dataset consists of 12 different regression tasks named: α , ∆ ε , ε HOMO , ε LUMO , µ , C ν , G , H , R 2 , U , U 0 , ZPVE, for which the mean absolute error is a standard metric.

Original: The dataset has over 130k molecules.
Revised: The dataset has over 130k molecules.

Original: We use data splits proposed by (Anderson et al., 2019), which gives us 100k trainin molecules, 18k molecules for validation and 13k molecules for testing.
Revised: We use data splits proposed by (Anderson et al., 2019), which gives us 100k trainin molecules, 18k molecules for validation and 13k molecules for testing.

Original: Training hyperparameters We trained R-MAT for 1000 epochs, with batch size equal to 256 and learning rate equal to 0 .
Revised: Training hyperparameters We trained R-MAT for 1000 epochs, with batch size equal to 256 and learning rate equal to 0 .

Original: We report the test set MAE for the epoch with the lowest validation
Revised: We report the test set MAE for the epoch with the lowest validation

Original: MAE.
Revised: MAE.

Original: 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 } .
Revised: 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 } .

Original: C. A BLATIONS
Revised: C. A BLATIONS

Original: Datasets For the ablations section, we used the BBBP, ESOL and FreeSolv datasets, splitted using a scaffold split, with 3 different splits.
Revised: Datasets For the ablations section, we used the BBBP, ESOL and FreeSolv datasets, splitted using a scaffold split, with 3 different splits.

Original: Labels of the regression datasets (ESOL and FreeSolv) were normalized before training.
Revised: Labels of the regression datasets (ESOL and FreeSolv) were normalized before training.

Original: Scores obtained in this section differ signiﬁcantly from the previous benchmarks due to the different data splits, different model hyperparameters and no pretraining used.
Revised: Scores obtained in this section differ signiﬁcantly from the previous benchmarks due to the different data splits, different model hyperparameters and no pretraining used.

Original: Training
Revised: Training

Original: hyperparameters Similarly as for our main benchmarks, we tuned only the learning rate, which was selected from the set of 7 possible options: { 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 , 5 e − 6 , 1 e − 6 } .
Revised: hyperparameters Similarly as for our main benchmarks, we tuned only the learning rate, which was selected from the set of 7 possible options: { 1 e − 3 , 5 e − 4 , 1 e − 4 , 5 e − 5 , 1 e − 5 , 5 e − 6 , 1 e − 6 } .

Original: We used batch size equal to 32 and Noam optimizer with warm-up equal to 20% of all steps.
Revised: We used batch size equal to 32 and Noam optimizer with warm-up equal to 20% of all steps.

Original: Moreover we use single layer, instead of two-layers MLP as our classiﬁcation part.
Revised: Moreover we use single layer, instead of two-layers MLP as our classiﬁcation part.

Original: D A DDITIONAL EXPERIMENTAL RESULTS D. S MALL HYPERPARAMETER BUDGET
Revised: D A DDITIONAL EXPERIMENTAL RESULTS D. S MALL HYPERPARAMETER BUDGET

Original: In Figure 6 one can ﬁnd rank plots for results from Table 1.
Revised: In Figure 6 one can ﬁnd rank plots for results from Table 1.

Original: R-MAT and R-MAT rdkit obtained the best median rank among all compared models.
Revised: R-MAT and R-MAT rdkit obtained the best median rank among all compared models.

Original: The performance gain of RF rdkit over RF and SVM rdkit over SVM is interesting.
Revised: The performance gain of RF rdkit over RF and SVM rdkit over SVM is interesting.

Original: Overall these two baseline models are worse only than R-MAT.
Revised: Overall these two baseline models are worse only than R-MAT.

Original: D.
Revised: D.

Original: L ARGE HYPERPARAMETER BUDGET
Revised: L ARGE HYPERPARAMETER BUDGET

Original: In Figure 7 one can ﬁnd rank plots for results from Table 2.
Revised: In Figure 7 one can ﬁnd rank plots for results from Table 2.

Original: We present separate plots for models trained with the large grid search (Left) and for models with only learning rate tuning (Right).
Revised: We present separate plots for models trained with the large grid search (Left) and for models with only learning rate tuning (Right).

Original: [Figure]
Revised: [Figure]

Original: D. L ARGE - SCALE EXPERIMENTS
Revised: D. L ARGE - SCALE EXPERIMENTS

Original: In Table 10 one can ﬁnd detailed results of comparison R-MAT performance with other various models.
Revised: In Table 10 one can ﬁnd detailed results of comparison R-MAT performance with other various models.

Original: R-MAT achieves highly competitive results, with state-of-the-art performance on 4 out of the 12 tasks, which proves how universal this model is.
Revised: R-MAT achieves highly competitive results, with state-of-the-art performance on 4 out of the 12 tasks, which proves how universal this model is.

Original: [Table]
Revised: [Table]

Original: D. C LOSER COMPARISON TO M OLECULE A TTENTION T RANSFORMER
Revised: D. C LOSER COMPARISON TO M OLECULE A TTENTION T RANSFORMER

Original: 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.
Revised: 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.

Original: We qualitatively explore here whether R-MAT achieves this goal, comparing its attention patterns to that of MAT.
Revised: We qualitatively explore here whether R-MAT achieves this goal, comparing its attention patterns to that of MAT.

Original: For this purpose we compared attention patterns learned by pretrained MAT (weights taken from (Maziarka et al., 2020)) and R-MAT for a selected molecule from the ESOL dataset.
Revised: For this purpose we compared attention patterns learned by pretrained MAT (weights taken from (Maziarka et al., 2020)) and R-MAT for a selected molecule from the ESOL dataset.

Original: Figure 8 shows that different heads of Relative Molecule Self-Attention are focusing on different atoms in the input molecule.
Revised: Figure 8 shows that different heads of Relative Molecule Self-Attention are focusing on different atoms in the input molecule.

Original: We can see that self-attention strength is concentrated on the input atom (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 in the molecule (heads 6 and 10 are concentrated on atoms 1 and 2).
Revised: We can see that self-attention strength is concentrated on the input atom (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 in the molecule (heads 6 and 10 are concentrated on atoms 1 and 2).

Original: 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.
Revised: 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.

Original: This likely happens due to the construction of the Molecule Self-Attention in MAT (c.f.
Revised: This likely happens due to the construction of the Molecule Self-Attention in MAT (c.f.

Original: Equation (2)), where the output atom representation is calculated from equally weighted messages based on the adjacency matrix, distance matrix and self-attention.
Revised: Equation (2)), where the output atom representation is calculated from equally weighted messages based on the adjacency matrix, distance matrix and self-attention.

Original: Due to its construction, it is more challenging for MAT than for R-MAT to learn to attend to a distant neighbour.
Revised: Due to its construction, it is more challenging for MAT than for R-MAT to learn to attend to a distant neighbour.

Original: E E XPLORING THE DESIGN SPACE OF MOLECULAR SELF - ATTENTION
Revised: E E XPLORING THE DESIGN SPACE OF MOLECULAR SELF - ATTENTION

Original: Identifying the Relative Molecule Self-Attention layer required a large-scale and methodological exploration of the self-attention design space.
Revised: Identifying the Relative Molecule Self-Attention layer required a large-scale and methodological exploration of the self-attention design space.

Original: In this section, we present experimental data that
Revised: In this section, we present experimental data that

Original: [Table]
Revised: [Table]

Original: [Table] 
Revised: [Table] [Table] [Table] SA

Original: [Table]
Revised: [Table]

Revised: SA

Original: [Table]
Revised: [Table]

Original: SA
Revised: SA

Original: [Table]
Revised: [Table]

Original: SA
Revised: SA

Revised: [Table]

Revised: SA

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: SA
Revised: SA

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: SA
Revised: [Table] [Table] [Table]

Original: [Table]
Revised: [Table]

Original: SA
Revised: SA

Original: [Table]
Revised: [Table]

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: SA

Original: [Table]
Revised: [Table]

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: SA

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: informed our choices.
Revised: informed our choices.

Original: We also hope it will inform future efforts in designing attention mechanism for molecular data.
Revised: We also hope it will inform future efforts in designing attention mechanism for molecular data.

Original: We follow here the same evaluation protocol as in Section 4.4 and show how different natural variants compare against R-MAT.
Revised: We follow here the same evaluation protocol as in Section 4.4 and show how different natural variants compare against R-MAT.

Original: E. S ELF - ATTENTION VARIANTS
Revised: E. S ELF - ATTENTION VARIANTS

Original: Relative Molecule Self-Attention is designed to better incorporate the relative spatial position of atoms in the molecule.
Revised: Relative Molecule Self-Attention is designed to better incorporate the relative spatial position of atoms in the molecule.

Original: The ﬁrst step is embedding each pair of atoms.
Revised: The ﬁrst step is embedding each pair of atoms.

Original: Then, the embedding is used to re-weight self-attention.
Revised: Then, the embedding is used to re-weight self-attention.

Original: To achieve this, Relative Molecule Self-Attention combines ideas from natural language processing (Shaw et al., 2018; Dai et al., 2019; Huang et al., 2020).
Revised: To achieve this, Relative Molecule Self-Attention combines ideas from natural language processing (Shaw et al., 2018; Dai et al., 2019; Huang et al., 2020).

Original: These works focus on encoding better relative positions of tokens in the input.
Revised: These works focus on encoding better relative positions of tokens in the input.

Original: We compare to three speciﬁc variants from these works that can be written using our previously introduced notation as:
Revised: We compare to three speciﬁc variants from these works that can be written using our previously introduced notation as:

Original: Relative self attention (Shaw et al., 2018):
Revised: Relative self attention (Shaw et al., 2018):

Original: [Equation]
Revised: [Equation]

Original: Relative self attention with attentive bias (Dai et al., 2019):
Revised: Relative self attention with attentive bias (Dai et al., 2019):

Original: [Equation]
Revised: [Equation]

Original: Improved relative self-attention (Huang et al
Revised: Improved relative self-attention (Huang et al

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Table 11 shows that the attention operation used in R-MAT outperforms variants 2 and 3 across all three tasks and variant 1 on two tasks, being comparable on the third one.
Revised: Table 11 shows that the attention operation used in R-MAT outperforms variants 2 and 3 across all three tasks and variant 1 on two tasks, being comparable on the third one.

Original: This might be expected given that Relative Molecule Self-Attention combines these ideas (c.f. Equation (3)).
Revised: This might be expected given that Relative Molecule Self-Attention combines these ideas (c.f. Equation (3)).

Original: [Table]
Revised: [Table]

Original: E. E NRICHING BOND FEATURES WITH ATOM FEATURES
Revised: E. E NRICHING BOND FEATURES WITH ATOM FEATURES

Original: In Relative Molecule Self-Attention
Revised: In Relative Molecule Self-Attention

Original: [Table]
Revised: [Table]

Original: we use a small number of bond features to construct the atom pair embedding.
Revised: we use a small number of bond features to construct the atom pair embedding.

Original: We investigate here the effect of extending bond featurization.
Revised: We investigate here the effect of extending bond featurization.

Original: Inspired by (Shang et al., 2018)
Revised: Inspired by (Shang et al., 2018)

Original: [Table]
Revised: [Table]

Original: we added information about the atoms that an edge connects.
Revised: we added information about the atoms that an edge connects.

Original: We tried three different variants.
Revised: We tried three different variants.

Original: In the ﬁrst one
Revised: In the ﬁrst one

Original: [Table]
Revised: [Table]

Original: we extend the bond representation with concatenated input features of atoms that the bond connects.
Revised: we extend the bond representation with concatenated input features of atoms that the bond connects.

Original: In the second one
Revised: In the second one

Original: [Table]
Revised: [Table]

Original: instead of raw atoms’ features
Revised: instead of raw atoms’ features

Original: [Table]
Revised: [Table]

Original: we tried the one-hot-encoding of the type of the bond connection (i
Revised: we tried the one-hot-encoding of the type of the bond connection (i

Original: [Table]
Revised: [Table]

Original: e. when the bond connects atoms C and N
Revised: e. when the bond connects atoms C and N

Original: [Table]
Revised: [Table]

Original: we encode it as a bond ’C_N’ and take the one-hot-encoding of this information).
Revised: we encode it as a bond ’C_N’ and take the one-hot-encoding of this information).

Original: Finally, we combined these two approaches together.
Revised: Finally, we combined these two approaches together.

Original: [Table]
Revised: [Table]

Original: type one-hot .
Revised: type one-hot .

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Both .
Revised: Both .

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: The results are shown in Table 12.
Revised: The results are shown in Table 12.

Original: Surprisingly, we ﬁnd that adding this type of information to the bond features negatively affects performance of R-MAT.
Revised: Surprisingly, we ﬁnd that adding this type of information to the bond features negatively affects performance of R-MAT.

Original: This suggests that R-MAT can already access these features efﬁciently from the input (which we featurize using the same set of features).
Revised: This suggests that R-MAT can already access these features efﬁciently from the input (which we featurize using the same set of features).

Original: This could also happen due to the fact that after a few layers, the attention is not calculated over the input atoms anymore.
Revised: This could also happen due to the fact that after a few layers, the attention is not calculated over the input atoms anymore.

Original: Instead, it works over hidden embeddings, which themselves can be mixed representations of multiple atom embeddings (Brunner et al., 2019), where the proposed additional representation contains only information about the input features.
Revised: Instead, it works over hidden embeddings, which themselves can be mixed representations of multiple atom embeddings (Brunner et al., 2019), where the proposed additional representation contains only information about the input features.

Original: E. D ISTANCE ENCODING VARIANTS
Revised: E. D ISTANCE ENCODING VARIANTS

Original: R-MAT uses a speciﬁc radial base distance encoding proposed by (Klicpera et al., 2020), followed by the envelope function, with N emb = 32 .
Revised: R-MAT uses a speciﬁc radial base distance encoding proposed by (Klicpera et al., 2020), followed by the envelope function, with N emb = 32 .

Original: We compare here to several other natural choices.
Revised: We compare here to several other natural choices.

Original: We tested the following distance encoding variants : (1) removal of the envelope function, (2) increasing the number of distance radial functions to 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).
Revised: We tested the following distance encoding variants : (1) removal of the envelope function, (2) increasing the number of distance radial functions to 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).

Original: The distance in MAT is encoded as e ( d ) =
Revised: The distance in MAT is encoded as e ( d ) =

Original: exp( − d ) .
Revised: exp( − d ) .

Original: The distance in
Revised: The distance in

Original: SchNet is encoded as e n ( d ) = exp( − γ (cid:107) d − µ n (cid:107) 2 ) , for γ = 10 Å and 0 Å ≤ µ n ≤ 30 Å divided into
Revised: SchNet is encoded as e n ( d ) = exp( − γ (cid:107) d − µ n (cid:107) 2 ) , for γ = 10 Å and 0 Å ≤ µ n ≤ 30 Å divided into

Original: N emb equal sections, with N emb set to 32 or 128 .
Revised: N emb equal sections, with N emb set to 32 or 128 .

Original: The results are shown in Table 13.
Revised: The results are shown in Table 13.

Original: These results corroborate that a proper representation of distance information is a key in adapting self-attention to molecular data.
Revised: These results corroborate that a proper representation of distance information is a key in adapting self-attention to molecular data.

Original: We observe that all variants underperform compared to the radial base encoding used in Relative Molecule Self-Attention.
Revised: We observe that all variants underperform compared to the radial base encoding used in Relative Molecule Self-Attention.

Original: Noteworthy is the fact that encoding distances in MAT way caused a lot of instability during the training of
Revised: Noteworthy is the fact that encoding distances in MAT way caused a lot of instability during the training of

Original: R-MAT.
Revised: R-MAT.

Original: [Table]
Revised: [Table]

Original: F A DDITIONAL COMPARISON OF GRAPH PRETRAINING F. P RETRAINING METHODS BENCHMARK
Revised: F A DDITIONAL COMPARISON OF GRAPH PRETRAINING F. P RETRAINING METHODS BENCHMARK

Original: As pretraining is nowadays the main component of big Transformer architectures (Devlin et al., 2018;
Revised: As pretraining is nowadays the main component of big Transformer architectures (Devlin et al., 2018;

Original: Liu et al., 2019; Clark et al., 2020)
Revised: Liu et al., 2019; Clark et al., 2020)

Original: [Table]
Revised: [Table]

Original: we decided to devote more attention to this issue.
Revised: we decided to devote more attention to this issue.

Original: For this purpose we compared various graph pretraining methods to identify the best one and use it in the ﬁnal R-MAT model.
Revised: For this purpose we compared various graph pretraining methods to identify the best one and use it in the ﬁnal R-MAT model.

Original: Pretraining methods We used various pretraining methods proposed in the molecular property prediction literature (Hu et al., 2020; Maziarka et al., 2020;
Revised: Pretraining methods We used various pretraining methods proposed in the molecular property prediction literature (Hu et al., 2020; Maziarka et al., 2020;

Original: Rong et al., 2020; Fabian et al., 2020).
Revised: Rong et al., 2020; Fabian et al., 2020).

Original: To be more speciﬁc
Revised: To be more speciﬁc

Original: [Table]
Revised: [Table]

Original: we tried R-MAT with the following pretraining procedures:
Revised: we tried R-MAT with the following pretraining procedures:

Original: • No pretraining – as a baseline we include reuslts for R-MAT ﬁne-tuned from scratch, without any pretraining.
Revised: • No pretraining – as a baseline we include reuslts for R-MAT ﬁne-tuned from scratch, without any pretraining.

Original: Masking – masked pretraining used in (Maziarka et al., 2020).
Revised: Masking – masked pretraining used in (Maziarka et al., 2020).

Original: This is an adaptation of standard MLM pretraining used in NLP (Devlin et al., 2018) to the graphical data.
Revised: This is an adaptation of standard MLM pretraining used in NLP (Devlin et al., 2018) to the graphical data.

Original: In this approach, we mask features of 15% of all atoms in the molecule and then pass it throught the model.
Revised: In this approach, we mask features of 15% of all atoms in the molecule and then pass it throught the model.

Original: The goal is to predict what was the masked features.
Revised: The goal is to predict what was the masked features.

Original: • Contextual – contextual pretraining method proposed by (Rong et al., 2020).
Revised: • Contextual – contextual pretraining method proposed by (Rong et al., 2020).

Original: We described it further in Appendix B. • Graph-motifs – graph-level motif prediction method proposed by (Rong et al., 2020), where for every molecule we obtain the ﬁngerprint with information whether speciﬁed molecular functional groups are present in our molecule.
Revised: We described it further in Appendix B. • Graph-motifs – graph-level motif prediction method proposed by (Rong et al., 2020), where for every molecule we obtain the ﬁngerprint with information whether speciﬁed molecular functional groups are present in our molecule.

Original: The network’s task is multi-label classiﬁcation, where it has to predict, whether every predeﬁned functional group is in the given molecule.
Revised: The network’s task is multi-label classiﬁcation, where it has to predict, whether every predeﬁned functional group is in the given molecule.

Original: • Physicochemical – graph-level prediction method proposed by (Fabian et al., 2020).
Revised: • Physicochemical – graph-level prediction method proposed by (Fabian et al., 2020).

Original: We described it further in Appendix B. • GROVER – pretraining used by authors of GROVER (Rong et al., 2020).
Revised: We described it further in Appendix B. • GROVER – pretraining used by authors of GROVER (Rong et al., 2020).

Original: Combination of two pretraining methods: contextual and graph-motifs.
Revised: Combination of two pretraining methods: contextual and graph-motifs.

Original: • R-MAT– pretraining used in this paper.
Revised: • R-MAT– pretraining used in this paper.

Original: Combination of two pretraining methods: contextual and physicochemical.
Revised: Combination of two pretraining methods: contextual and physicochemical.

Original: Experimental setting We pretrained every model using described earlier dataset with 4M molecules.
Revised: Experimental setting We pretrained every model using described earlier dataset with 4M molecules.

Original: For every pretraining choice, we trained the model for 50 epochs, using the same training settings as for our standard pretraining (see Appendix C.3).
Revised: For every pretraining choice, we trained the model for 50 epochs, using the same training settings as for our standard pretraining (see Appendix C.3).

Original: Model hyperparameters and ﬁne-tuning settings were the same as for MAT and GROVER benchmarks.
Revised: Model hyperparameters and ﬁne-tuning settings were the same as for MAT and GROVER benchmarks.

Original: We used the BBBP, ESOL and FreeSolv datasets, splitted using a scaffold split, with 3 different data splits.
Revised: We used the BBBP, ESOL and FreeSolv datasets, splitted using a scaffold split, with 3 different data splits.

Original: Data splits are different than in previous experiments, and this causes the fact that the results are differs from the results from other paragraphs’.
Revised: Data splits are different than in previous experiments, and this causes the fact that the results are differs from the results from other paragraphs’.

Original: Results Results of this benchmark are presented in Table 14.
Revised: Results Results of this benchmark are presented in Table 14.

Original: We can draw some interesting conclusions from them.
Revised: We can draw some interesting conclusions from them.

Original: One can see, that using any kind of pretraining helps in obtaining better results than for the model trained from scratch.
Revised: One can see, that using any kind of pretraining helps in obtaining better results than for the model trained from scratch.

Original: Using physicochemical features for graph-level training gives better results than graph-motifs.
Revised: Using physicochemical features for graph-level training gives better results than graph-motifs.

Original: Therefore R-MAT pretraining (contextual + physicochemical) is better than GROVER pretraining (contextual + graph-motﬁs).
Revised: Therefore R-MAT pretraining (contextual + physicochemical) is better than GROVER pretraining (contextual + graph-motﬁs).

Original: Moreover combination of two tasks in pretraining usually gives better results than pretraining using only one task.
Revised: Moreover combination of two tasks in pretraining usually gives better results than pretraining using only one task.

Original: Interestingly, both node-level pretraining methods (masking and contextual pretraining), returns similar results.
Revised: Interestingly, both node-level pretraining methods (masking and contextual pretraining), returns similar results.

Original: [Table]
Revised: [Table]

Original: F. P RETRAINING LEARNING CURVES
Revised: F. P RETRAINING LEARNING CURVES

Original: In Figure 9 we present learning curves for R-MAT pretrained with our procedure (contextual + physiochemical).
Revised: In Figure 9 we present learning curves for R-MAT pretrained with our procedure (contextual + physiochemical).

Original: Left-side image shows that training loss ﬂattens out quite quickly, however it slowly decreases until the end of the training.
Revised: Left-side image shows that training loss ﬂattens out quite quickly, however it slowly decreases until the end of the training.

Original: Moreover one can see that contextual task is harder for the network than graph-level property prediction, as their losses vary by several orders of magnitude.
Revised: Moreover one can see that contextual task is harder for the network than graph-level property prediction, as their losses vary by several orders of magnitude.

Original: 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 at the end of training
Revised: 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 at the end of training

Original: [Figure]
Revised: [Figure]

Original: We also tested, whether longer pretraining allows R-MAT to get better results during the ﬁne-tuning process.
Revised: We also tested, whether longer pretraining allows R-MAT to get better results during the ﬁne-tuning process.

Original: In Figure 10 we present ﬁne-tuning scores obtained by models pretrained with different number of pretraining epochs, for FreeSolv and ESOL datasets.
Revised: In Figure 10 we present ﬁne-tuning scores obtained by models pretrained with different number of pretraining epochs, for FreeSolv and ESOL datasets.

Original: Interestingly, longer training does not always results in better ﬁne-tuning scores.
Revised: Interestingly, longer training does not always results in better ﬁne-tuning scores.

Original: This is indeed the case for the FreeSolv dataset (left-side image), however for ESOL (right-side image) we cannot draw such a conclusion.
Revised: This is indeed the case for the FreeSolv dataset (left-side image), however for ESOL (right-side image) we cannot draw such a conclusion.

Original: [Figure]
Revised: [Figure]