One Transformer Can Understand Both 2D & 3D Molecular Data
Authors: Shengjie Luo
Tianlang Chen
Yixian Xu
Shuxin Zheng
Tie-Yan Liu
Liwei Wang
Di He
Venue: ICLR 2023 poster (2023)
Keywords: Transformer
general-purpose molecular model
2D molecular representation
3D molecular representation

Abstract: Unlike vision and language data which usually has a unique format, molecules can naturally be characterized using different chemical formulations. One can view a molecule as a 2D graph or define it as a collection of atoms located in a 3D space. For molecular representation learning, most previous works designed neural networks only for a particular data format, making the learned models likely to fail for other data formats. We believe a general-purpose neural network model for chemistry should be able to handle molecular tasks across data modalities. To achieve this goal, in this work, we develop a novel Transformer-based Molecular model called Transformer-M, which can take molecular data of 2D or 3D formats as input and generate meaningful semantic representations. Using the standard Transformer as the backbone architecture, Transformer-M develops two separated channels to encode 2D and 3D structural information and incorporate them with the atom features in the network modules. When the input data is in a particular format, the corresponding channel will be activated, and the other will be disabled. By training on 2D and 3D molecular data with properly designed supervised signals, Transformer-M automatically learns to leverage knowledge from different data modalities and correctly capture the representations. We conducted extensive experiments for Transformer-M. All empirical results show that Transformer-M can simultaneously achieve strong performance on 2D and 3D tasks, suggesting its broad applicability. The code and models will be made publicly available at https://github.com/lsj2408/Transformer-M.

Published as a conference paper at ICLR 2023
ONE TRANSFORMER CAN UNDERSTAND BOTH 2D & 3D MOLECULAR DATA
Tianlang Chen Shuxin Zheng', Tin Liu Nattioaal Key Laboratory of General Artificial Intelligence. School Intelligence Science
Shengjie Luo',
Yixian Xu
Liwei Wang He
and Technology, Peking University School of EECS, Peking University *Microsoft Research
*Center for Data Science, Peking University
"Shanghai Artificial Intelligence Laboratory aku..duu.e. on, tichen@pku.c edu. cn, xyx050@stu. ppu... pwp.ppu.u..o.opo edu.cn (shuz, tyliu) Emi wrwsoft.ccom (wanglw, dihe (ph.)ppk.edd...... .ewt..educccm
1 uus jesseesuu. pku. edu.
ABSTRACT
Unlike vision and language data which usually has a unique format, molecules can naturally be characterized using different chemical formulations. One view a molecule as a 2D graph or define it as a collection of atoms located in a 3D space. For molecular representation learning, most previous works designed neural networks only for a particular data format, making the learned models likely to fail for other data formats. We believe a general-purpose neural network model for chemistry should be able to handle molecular tasks across data modalities. To achieve this goal, in this work, we develop a novel Transformer- based Molecular model called Transformer-M, which can take molecular data of 2D or 3D formats as input and generate meaningful semantic representations. Using the standard Transformer as the backbone architecture, Transformer-M M develops two separated channels to encode 2D and 3D) structural information and incorporate them with the atom features in the network modules. When the input data is in particular format, the corresponding channel will activated, and the other will be disabled. By training 2D and 3D molecular data with properly designed supervised signals, Transformer-N automatically learns to leverage knowledge from different data modalities and correctly capture the representations. We conducted extensive experiments for Transformer-M. All empirical results show that Transformer-M simultaneously achieve strong performance on 2D) and 3D tasks, suggesting its broad applicability. The code and models will be made publicly available at https: / gi thub. cwm 1 1sj24s4884 4084/T/rara.. Traansaanssesa forser-----r-- M.
INTRODUCTION
Deep learning approaches have revolutionized many domains, including computer vision (He et al., 2016), natural language processing (Devlin et al., 2019; Brown et al., 2020), and games (Mnih et al., 2013; Silver et al., 2016). Recently, researchers have started investigating whether the power of neural networks could help solve important scientific problems in chemistry, e.g., predicting the property of molecules and simulating the molecular dynamics from large-scale training data (Hu et al., 2020a; 2021: Zhang et al., 2018: Chanussot et al., 2020).
One key difference between chemistry and conventional domains such as vision and language is the multimodality of data. In vision and language, data instance is usually characterized in particular form. For example, an image defined as RGB values a pixel grid, while a sentence is defined tokens in sequence. In contrast, molecules naturally have different chemical formulations. A molecule can be represented as a sequence (Weininger, 1988), 2D graph (Wiswesser, 1985). or collection of atoms located in 3D) space. 2D and 3D structures are the most popularly used formulations as many valuable properties and statistics can be obtained from them (Chmiela et al., 2017; Stokes al., 2020). However, as far as we know, most previous works focus on designing neural network models for either 2D or 3D structures, making the model learned in one form fail to be applied in tasks of the other form.
We argue that general-purposs purpose neural network model in chemistry should at least be able to handle molecular tasks across data modalities. In this paper, we take the first step toward this goal by
"These two authors contributed equally to this project Corresponnence Di He <di he@pku edu. en> and Liwei Wang <wangl w@pku. edu. cnz.

---

Published as a conference paper at ICLR 2023
developing Transformer-M. a versatile transformrrbassed Molecular model that performs well for both 2D and 3D molecular representation learning. Note that for molecule, its 2D and 3D forms describe the same collection of atoms but use different characterizations of the structure. Therefore, the key challenge is to design a model expressive and compatible in capturing structural knowledge in different formulations and train the parameters to learn from both information. Transformer is more favorable than other architectures as it can explicitly plug structural signals in the model as bias terms (e.g., positional encodings (Vaswani et al., 2017; Raffel et al., 2020)). We can conveniently set 2D and 3D structural information as different bias terms through separated channels and incorporate them with the atom features in the attention layers.
Architecture. The backbone network of our Transformer-N M is composed of standard Transformer blocks. We develop two separate channels to encode 2D and 3D structural information. The 21) channel uses degree encoding, shortest path distance encoding, and edge encoding extracted from the 2D graph structure, following Ying et al. (2021a). The shortest path distance encoding and edge encoding reflect the spatial relations and bond features of pair of atoms and are used as bias terms in the softmax attention. The degree encoding is added to the atom features in the input layer. For the 3D channel, we follow Shi et al. (2022) to use the 3D distance encoding to encode the spatial distance between atoms in the 3D geometric structure. Each atom pair's Euclidean distance is encoded via the Gaussian Basis Kernel function (Scholkopf et al., 1997) and will be used bias term in the softmax attention. For each atom, we sum up the 3D distance encodings between it and all other atoms, and add atom features in the input layer. See Figure/ 1 for an illustration.
Training. Except for the parameters in the two structural channels, all other parameters in Transformer-M (e.g., self-attention and feed-forward networks) are shared for different data modali- ties. We design joint-training approach for Transformer-M to learn its parameters. During training, when the instances in a batch are only associated with 2D graph structures, the 2D channel will be activated, and the 3D channel will be disabled. Similarly, when the instances in a batch use 3D geometric structures, the 3D channel will be activated, and the 2D channel will be disabled. When both 2D and 3D information are given, both channels will be activated. In such a way, we can collect 2D and 3D data from separate databases and train Transformer-M with different training objectives, making the training process more flexible. We expect single model to learn to identify and incorporate information from different modalities and efficiently utilize the parameters, leading to better generalization performance.
Experimental Results. We use the PCQM4Mv2 dataset in the OGB Large-Scale Challenge (OGB- LSC) (Hu et al., 2021) train our Transformer-M, which consists of 3.4 million molecules of both 2D and 3D forms. The model is trained to predict the pre-computed HOMO-LUMO gap of each data instance in different formats with a pre-text 3D denoising task specifically for 3D data. With the pre-trained model, we directly use or fine-tune the parameters for various molecular tasks of different data formats. First, we show that on the validation set of the PCQM4Mv2 task, which only contains 2D molecular graphs, our Transformer-M M surpasses all previous works by a large margin. The improvement is credited to the joint training, which effectively mitigates the overfitting problem. Second, On PDBBind (Wang et al., 2004; 2005b) (2D&3D), the fine-tuned Transformer-M achieves state-of-the-ar performance compared to strong baselines. Lastly, on QM9 (Ramakrishnan et al., 2014) (3D) benchmark, the fine-tuned Transformer-M models achieve competitive performance compared to recent methods. All results show that our Transformer-N M has the potential to be used as general-purpose model in a broad range of applications in chemistry.
RELATED WORKS
Neural networks for learning 2D molecular representations. Graph Neural Network (GNN) is popularly used in molecular graph representation learning (Kipf & Welling, 2016: Hamilton et al., 2017: Gilmer et al., 2017: Xu et al., 2019; Veličković et al., 2018). GNN learns node and graph representations by recursively aggregating (i.e., message passing) and updating the node representations from neighbor representations. Different architectures are developed by using different aggregation and update strategies. We refer the readers to Wu et al. (2020) for a comprehensive survey, Recently, many works extended the Transformer model to graph tasks (Dwivedi & Bresson, 2020; Kreuzer al., 2021; Ying et al., 2021a; Luo et al., 2022; Kim et al., 2022; Rampášek et al., 2022; Park et al., 2022; Hussain et al., 2022; Zhang et al., 2023). Seminal works include Graphormer (Ying et al., 2021a), which developed graph structural encodings and used them in standard Transformer model.
2

---

Published as a conference paper at ICLR 2023
Neural networks for learning 3D molecular representations. Learning molecular representations with 3D geometric information is essential in many applications, such as molecular dynamics simulation. Recently, researchers have designed architectures to preserve invariant and equivariant properties for several necessary transformations like rotation and translation. Schütr et al. (2017) used continuous-filte convolutional layers to model quantum interactions in molecules. Thomas et al. (2018) used filters built from spherical harmonics to construct a rotation- and translation- equivariant neural network. Kliepera et al. (2020) proposed directional message passing, which ensured their embeddings to be rotationally equivariant. Liu et al. (2022); Wang et al. (2022) use spherical coordinates to capture geometrie information and achieve equivariance. Hutchinson et al. (2021); Thölke & De Fabritiis (2021) built Transformer models preserving equivariant properties. Shi et al. (2022) extended Ying et al. (2021a) to a 3D Transformer model which attains better results on large-scale molecular modeling challenges (Chanussot et al., 2020).
Multi-view learning for molecules. The 2D graph structure and 3D geometric structure can be con- sidered as different views of the same molecule. Inspired by the contrastive pre-training approach in vision (Chen et al., 2020; He et al., 2020; Radford et al., 2021), many works studied pre-training meth- ods for molecules by jointly using the 2D and 3D information. Stärk et al. (2022) used two encoders to encode the 2D and 3D molecular information separately while maximizing the mutual information between the representations. Liu et al. (2021a) derived the GraphMVP framework, which uses contrastive learning and reconstruction to pre-train a 2D encoder and a 3D encoder. Zhu et al. (2022) unified the 2D and 3D pre-training methods above and proposed 2D GNN model that can be en- hanced by 3D geometric features. Different from these works, we aim to develop single model which compatible with both 2D and 3D molecular tasks. Furthermore, all the above works train models using paired 2D and 3D data, while such paired data is not a strong requirement to train our model.
General-purpose models. Building a single agent that works for multiple tasks, even across modalities, is a recent discovery in deep learning. In the early years, researchers found that a single multilingual translation model can translate tens languages using the same weights and perform better than a bilingual translation model for rare languages (Lample & Conneau, 2019; Conneau et al., 2019; Xue et al., 2020; Liu et al., 2020). Large-scale language model (Devlin et al., 2019; Brown et al., 2020) is another example that can be applied to different downstream tasks using in-context learning or fine-tuning. Reed et al. (2022) further pushed the boundary by building a single generalist agent, Gato. This agent uses the same network with the same weights but can play Atari, caption images, and make conversations like a human. Our work also lies in this direction. We focus on developing general-purpose model in chemistry, which can take molecules in different formats as input and perform well on various molecular tasks with a small number of additional training data.
3 TRANSFORMER-N
In this section, we introduce Transformer-M M, versatile Transformer serving as a general architecture for 2D and 3D molecular representation learning. First, we introduce notations and recap the prelimi- naries in the backbone Transformer architecture (Section 3. After that, we present the proposed Transformer-M model with two structural channels for different data modalities (Section 3.2).
3.1 NOTATIONS AND THE BACKBONE TRANSFORMER
A molecule M is made up of a collection of atoms held together by attractive forces. We denote X E Ruaxd as the atoms with features, where n is the number of atoms, and d the feature dimension. The structure of M can be represented in different formulations, such as 2D graph structure and 3D geometric structure. For the 2D graph structure, atoms are explicitly connected by chemical bonds, and we define MD = (X, E), where (i.j) E denotes the edge feature (i.e., the type of the bond) between atom and j if the edge exists. For the 3D geometric structure, for each atom i, its position r; in the Cartesian coordinate system is provided. We define MD (X, R), where R (r,,.., ..., ra) and r, E R Our goal is to design a parametric model which can take either MdD or M (or both of them) input, obtain contextual representations, and make predictions on downstream tasks.
Transformer layer. The backbone architecture we use in this work is the Transformer model (Vaswani al., 2017). A Transformer is composed of stacked Transformer blocks. A Transformer block consists of two layers: a self-attention layer followed by a feed-forward layer, with both layers having normalization (e.g., LayerNorm (Ba et al., 2016)) and skip connections (He al., 20. Denote X (4) as the input to the (1 + 1)-th block and define (0) X. For an input x the (1 1)-th block works follows:
3

---

Published as a conference paper at ICLR 2023
w WX (x"'
(²("=) = softmax
(1)
Vd
4 x x=
A x² W/ W.
(2)
A=1
(1+1 x GELU(X wi)
(3)
where bhth E Rddudd xd W W w E Rdxds, W E R E H is the number of attention heads, dæ is the dimension of each head, and r is the dimension of the hidden layer. A (X) is usually referred to the attention matrix.
Positional encoding. Another essential component in the Transformer is positional encoding. Note that the self-attention layer and the feed-forward layer do not make use of the order of input elements (e.g., word tokens), making the model impossible to capture the structural information. The original paper (Vaswani et al., 2017) developed effective positional encodings to encode the sentence structural information and explicitly integrate them as bias terms into the model. Shortly, many works realized that positional encoding plays crucial role in extending standard Transformer more complicated data structures beyond language. By carefully designing structural encoding using domain knowledge, Transformer has successfully been applied the image and graph domain and achieved impressive performance (Dosovitskiy et al., 2020; Liu et al., 2021b; Ying et al., 2021a).
3.2 TRANSFORMER-M AND TRAINING STRATEGY
As we can see, the two molecular formulations defined in Section 3.1 use the same atom feature space but different characterizations of the structure (graph structure E v.s. geometric structure It). Therefore, the key challenge is to design a compatible architecture that can utilize either structural information in E or R (or both) and incorporate them with the atom features in a principled way.
The Transformer is suitable backbone achieve the goal as bias terms and properly plug them into different modules. Furthermore, with Transformer, we can treat E and R in unified way by decomposing the struetural information into pair-wise and atom-wise encodings. Without loss of generality, we choose to use the encoding strategies in the graph and geometric Transformers proposed by Ying et al. (2021a); Shi et al. (2022). For the sake of completeness, we briefly introduce those structural encodings and show how to leverage them in Transformer-M M. Note that our design methodology also works with other encoding strategies (Hussain et al., 2022; Park et al., 2022: Thölke & De Fabritiis, 2021). See Appendix B.5 for the detailed results.
encode structural information
Encoding pair-wise relations in E. We use two terms to encode the structural relations between any atom pairs in the graph. First, we encode the shortest path distance (SPD) between two atoms to reflect their spatial relation. Let 4SSD denote the SPD encoding between atom i and j. which is a learnable scalar determined by the distance of the shortest path between i and j. Second, we encode the edge features (e.g., the chemical bond types) along the shortest path between and j to reflect the bond information. For most molecules, there exists only one distinct shortest path between any two atoms. Denote the edges in the shortest path from to j as SP = (e1,e2, e2, .... eN). and the edge Σ. e., ew,w),)² where w,, are learnable vectors
encoding between i and is defined as & Edge of the same dimension as the edge feature. Denote $SSD and Eddge as the matrix form of the SPD encoding and edge encoding, both of which are of shape n X n.
ij N
Encoding pair-wise relations in R. We encode the Euclidean distance to reflect the spatial relation between any pair of atoms in the 3D space. For each atom pair (i, we first process their Euclidean distance with the Gaussian Basis Kernel function (Scholkopf et al., 1997), oik (KJ)
V2R exp -
1(166)) (66.) (76.)) r--ry +1-ry-+-
1 where K is the number of Gaussian
[64]
Basis kernels. Then the 3D Distance encoding T. Distance is obtained according to TD3D Distance
GELU ((a,)) W3., where (b(,)) = (w...) v (1.)) Wbb R W3 E R are learnable parameters. ((.j)): P(i.j) are learnable scalars indexed by the pair of atom types, and μ* o are learnable kernel center and learnable scaling factor of the k-th Gaussian Basis Kernel. Denote T.DD Distance as the matrix form of the 3D distance encoding, whose shape is n. x
4

---

Published as a conference paper at ICLR 2023
SPD Encodings
Edge Encodings
Degree Encodings
CH2OH
Atom Feature
ATT
FFN
ATT
FFN
Encodings
Encodings
Figure 1: An illustration of our Transformer- M model architecture. We build two channels on the backbone Transformer. The red channel is activated for data with 2D graph structures to incorporate degree, shortest path distance, and edge information. The purple channel is activated for data with 3D geometric structures to leverage Euclidean distance information. Different encodings are located in appropriate modules.
Integrating $SDD, Edge and 3D Distance in Transformer-ML All pair-wise encodings defined above capture the interatomic information, which is in similar spirit to the relative positional encoding for sequential tasks (Raffel et al., 2020). Therefore, we similarly locate those pair-wise signals in the self-attention module to provide complementary information the dot-produe term XW XWo(X (XX(W For simplicity, we omit the index of attention head and layer and the modified attention matrix is defined
XW (WO(XWR) (XWR)
A(X) softmax
SPD S Edge
3D Distance
vd
2D pair-wise channel
3D pair-wise channel
Encoding atom-wise structural information in E. For atom i, Eqn. (4) computes the normalized weights according to the semantic (first term) and spatial relation (last three terms) between and other atoms. However, the information still not sufficient. For example, the importance (i.e., centrality) of
each atom is missing in the attention. For each atom i, we use its degree as the centrality information. Formally, let v Degree denote the degree encoding of the atom i, which is d-dimensional learnable vector determined by the degree of the atom. Denote Degree [T Degree centrality encoding of all the atoms, which is shape n x d.
Degree Degree
V Degree the
Encoding atom-wise structural information in R. Similar to the 2D atom-wise centrality encoding, for geometric data, we encode the centrality of each atom in the 3D space. For each atom i, we sum up the 3D Distance encodings between and all other atoms. Let of Distance denote the centrality encoding of atom have Su Distance
Eje(n (1.)) W3 where W3 R is a learnable weight matrix. Similarly, we define Sum of 3D Distance as the encoding of all atoms, whose shape d.
Integrating Degree and y Sum of 3D Distance in Transformer-M. We add the atom-wise encodings of 2D and 3D structures to the atom features in the input layer. Formally, the input X (0) is modified as:
x (0) X
Degree
T Sum of 3D Distance
2D channel
atom-wise channel
Through this simple way, the structural information of molecules in both 2D and 3D formats is integrated into one Transformer model. is easy to check that Transformer-N preserves equivariant properties for both data formats.

---

Published as a conference paper at ICLR 2023
Training. The next step is to learn the parameters in Transformer-M to capture meaningful represen- tations from each data format. To achieve this goal, we develop a simple and flexible joint training method to learn Transformer-M. We first collect datasets in different formats (2D/3D) and define snpeervssed/ell--uervviii
tasks (e.g., energy regression) on each format, and train the model on all the data toward each objective, respectively. To be concrete, during training, if data instance comes from a dataset in the 2D format, the 2D channel is activated, and the 3D channel is disabled. The model parameter will be optimized minimize the corresponding (i.e., 2D) objective. When data instance comes from dataset in the 3D format, only the 3D channel is activated, and the model will learn to minimize the 3D objective. Both channels are activated the model takes molecules in both 2D and 3D formats as input. Compared with the multi-view learning approaches, we can train Transformer-M using unpaired 2D and 3D data, making the training process more flexible.
The Transformer-M may generalize better due to the joint training. Several previous works (Liu et al., 2021a) observed that 21) graph structure and 3D geometric structure contain complementary chemical knowledge. For example, the 2D graph structure only contains bonds with bond type, while the 3D geometric structure contains fine-gr ggaineraind information such as lengths and angles. As another example, the 3D geometric structures are usually obtained from computational simulations like Density Functional Theory (DFT) (Burke, 2012), which could have approximation errors. The 2D graphs are constructed by domain experts, which some extent, provide references the 3D structure. By jointly training using 2D and 3D data with parameter sharing, our model can learn more chemical knowledge instead of overfitting to data noise and perform better on both 2D and 3D tasks.
Future Directions. As an initial attempt, our Transformer-M opens up a way to develop general- purpose molecular models to handle diverse chemical tasks in different data formats. We believe is a starting point with more possibilities to explore in the future. For example, in this work, we use a simple way and linearly combine the structural information of 2D and 3D structures, and we believe there should be other efficient ways to fuse such encodings. Our model can also be combined with previous multi-view contrastive learning approaches. worth investigating how to pre-train our model using those methods.
4 EXPERIMENTS
In this section, we empirically study the performance of Transformer-M. First, we pre-train our model on the PCQM4Mv2 training set from OGB Large-Scale Challenge (Hu et al., 2021) (Section 4. I). With the pre-trained model, we conduct experiments molecular tasks in different data formats and evaluate the versatility and effectiveness of our Transformer-M Due to space limitation, we study three representative tasks, PCQM4Mv2 (2D, Section 4.2), PDBBind (2D & 3D, Section 4.3) and QM9 (3D, Section 4,4). Ablation studies are presented in Section 4.5. All codes are implemented based on the official codebase of Graphormer (Ying et al., 2021a) in PyTorch (Paszke al., 2019).
4.1 LARGE-SCALE PRE-TRAINING
Our Transformer-M is pre-trained on the training set of PCQM4Mv2 from the OGB Large- Scale Challenge (Hu et al., 2021). The total number of training samples is 3.37 million. Each molecule is associated with its 2D graph structure and 3D) geometric structure. The HOMO-LUMO energy gap of each molecule is provided as its label, which is obtained by DFT-based methods (Burke, 2012).
We follow Ying et al. (2021a) and employ a 12-layer Transformer-M model. The dimension of hidden layers and feed-forward layers is set to 768. The number of attention heads is set to 32. The number of Gaussian Basis kernels is set to 128. To train Transformer-M, we provide three modes for each data instance: (1) activate the 2D channels and disable the 3D channels (2D mode); (2) activate the 3D channels and disable the 2D channels (3D mode); (3) activate both channels (2D+3D mode). The mode of each data instance during training is randomly drawn on the fly according to a pre-defined distribution, implemented similarly to Dropout (Srivastava et al., 2014). In this work, we use two training objectives. The first one is supervised learning objective, which aims to predict the HOMO-LUMO energy gap of each molecule. Besides, we also use a self-supervised learning objective called 3D Position Denoising (Godwin et al., 2022; Zaidi et al., 2022), which is particularly effective. During training, if a data instance is in the 3D mode, we add Gaussian noise to the position of each atom and require the model to predict the noise from the noisy input. The model is optimized to minimize a linear combination of the two objectives above. Details of settings are in Appendix I.

---

Published as a conference paper at ICLR 2023
Table Results on PCQM4Mv2 validation set in OGB Large-Scale Challenge (Hu et al., 2021). The evaluation metric is the Mean Absolute Error (MAE) [eV]. We report the official results of baselines from OGB and use * to indicate our implemented results. Bold values indicate the best performance.
method
#param.
Valid MAE
MLP--Fingerprin (Hu et al., 2021)
16.1M
0.1753
GCN (Kipf Welling, 2016)
2.0M
0.1379
GIN (Xu et al., 2019)
3.8M
0. 0.1195
GINE-VN VN (Brossard et al., 2020; Gilmer et al., 2017)
13.2M
0.1167*
GCN-VN (Kipf & Welling. 2016; Gilmer et al., 2017)
4.9M
0.1153
GIN-vN (Xu et al., 2019; Gilmer et al., 2017)
6.7M
0.1083
DeeperGCN-vN (Li al., 2020; Gilmer et al., 2017)
25.5M
0.1021
GraphGPSssAL (Rampášek et al., 2022)
6.2M
0.0938
CoAIGIN (Cui, 2022)
5.2M
0.0933
TokenGT (Kim et al., 2022)
48.5M
0.0910
GRPEBASE (Park et al., 2022)
46.2M
0.0890
EGT (Hussain et al., 2022)
89.3M
0.0869
GRPELARGE (Park ct al., 2022)
46.2M
0.0867
Graphormer (Ying et al., la2 Shi et al., 2022)
47.1M
0.0864
GraphGPSBAss (Rampášck ct al., 2022)
19.4M
0.0858
Transformer-M (ours)
47.1M
0.0787
4.2 PCQM4MV2 PERFORMANCE (2D)
After the model is pre-trained, we evaluate our Transformer-M on the validation set of PCQM4Mv2. Note that the validation set of PCQM4My2 consists of molecules in the 2D format only. Therefore, it to evaluate how well Transformer-M performs 2D molecular data. The goal of the task is to predict the HOMU-LUMO energy gap, and the evaluation metric is the Mean Absolute Error (MAE). As our training objectives include the HOMO-LUMO gap prediction task, we didn't fine-tune the model parameters on any data. During inference, only the 2D channels are activated. We choose several strong baselines covering message passing neural network (MPNN) variants and Graph Transformers. Detailed descriptions of baselines are presented in Appendix B.2.
The results are shown in Table 1. It can be easily seen that our Transformer-M surpasses all baselines by a large margin, c.g., 8.2% relative MAE reduction compared to the previous best model (Rampášek al., 2022), establishing new sfate-of-the-ard on PCQM4Mv2 dataset. Note that our general architecture is the same as the Graphormer model (Ying et al., 2021a). The only difference between Transformer-M M and the Graphormer baseline is that Graphormer is trained on 2D data only, while Transformer-M is trained using both 2D and 3D structural information. Therefore, we can conclude that Transformer-M performs well on 2D molecular data, and the 2D-3D joint training with shared parameters indeed helps the model learn more chemical knowledge.
4.3 PDBBIND PERFORMANCE (2D& & 3D)
To verify the compatibility of our Transformer-M, we further fine-tune model the PDBBind dataset (version 2016, Wang et al. (2004; 2005b)), one of the most widely used datasets for structure- based virtual screening (Jiménez et al., 2018; Drepniewska-Dziubinns 2019). PDBBind dataset consists of protein-ligane complexes as data instances, which are obtained in bioassay experiments associated with the pKa (or - log Ka, log K,) affinity values. For each data instance, the 3D geometric structures are provided and the 2D graph structures are constructed via pre-defined rules. The task requires models to predict the binding affinity of protein-ligand complexes, which is extremely vital for drug discovery. After pre-trained on the PCQM4Mv2 training set, our Transformer-M model is fine-tuned and evaluated on the core set of the PDBBind dataset.
al., 2018; Zheng al.,
We compare our model with competitive baselines covering classical methods, CNN-based methods, and GNNs. All experiments are repeated five times with different seeds. Average performance is reported. Due space limitation, present the details of baselines and experiment settings in Appendix B.3. The results are presented in Table Our Transformer-N consistently outperforms all the baselines on all evaluation metrics by large margin, e.g., 3.3% absolute improvement on Pearson's 's correlation coefficient (R). It is worth noting that data instances of the PDBBind dataset are protein-ligant complexes, while our model is pre-trained on simple molecules, demonstrating the transferability of Transformer-M.
7

---

Published as a conference paper at ICLR 2023
Table 2: Results on PDBBind core set (version 2016) (Wang al., 2004; 2005b). The evaluation metrics include Pearson's s correlation coefficient (R), Mean Absolute Error (MAE), Root-Mean Squared Error (RMSE), and Standard Deviation (SD). We report the official results of baselines from Li et al. (2021). Bold values indicate the best performance.
PDBBind core set
Method
R
MAE J.
RMSE
SD
T.R
20.00
1.67510.008 10.000
1.612.0.00
SVR
1.264410.00 0.027
1.555
0.000
1..90
RF-Score (Ballester et al., 2010)
0.789
1.161 10.007
1.446
1.335
Pafnucy (Stepniewska-Dziubbll Dziubinska al., 2018)
0.695
1.204
1.585
1.563 10.022
OnionNet (Zheng et al., 2019)
0.768
1.078
4.002
40.0744..04
1.391 0:0.38
GraphDTAGcs (Nguyen et al., 2020)
0.613
1.343+0.087
10.05
1.719
GraphDTAcxt (Nguyen et al., 2020)
0.601#0.0z
+..04
1.765
1.740
GraphDTAGux (Nguyen et al., 2020)
0.667
1.261
4-0.044
1.640
1.621 +:0.03
GraphDTAGxr.GNN (Nguyen et al., 2020)
0.697
+0.016
1.562
1.562+40.022
1.558 +0.088
GNN-DTI (Lim et al., 2019)
0.736
1.192+0.082
1.492.2.002
[.471200051
DMPNN (Yang al., 2019)
0.729
1.188+0.00e
1.493
1.489 +0.044
SGCN (Danel et al., 2020)
0.68644.015
1.250
+..82
MAT (Maziarka et al., 2020)
0.747
1.154
40.
1.445 00.033
DimeNet (Klicpera et al., 2020)
0.752
1.138:00.026
1,453 022
1.434
CMPNN (Song et al., 2020)
0.765 0.00
1.117
1,408 028
1.399:10.025
SIGN (Li al., 2021)
0.797
1.027
1.316
1.312
Transformer-M (ours)
0.830
0.940
1.232.10.013
1.207 100007
Table 3: Results on QM9 (Ramakrishnan et al., 2014). The evaluation metric is the Mean Absolute Error (MAE). We report the official results of baselines from Thölke & De Fabritiis (2021); Godwin et al. (2022); Jiao et al. (2022). Bold values indicate the best performance.
mothod
Au
R
ZPVE
Ug
U
H
G
C,
EdgePred (Hamilton et al., 2017)
0.039
0.086
31.9
58.2
0.1
14.7
14.8
14.5
0.038
AurMask (Ilu et 2019)
0.020
0.072
31.3
37.8
50.0
0.423
1.90
10.7
10.8
11.4
11.2
0.062
IntoGraph (Sun et 2019)
0.041
0,099
48.1
38.1
72.2
0.114
1.69
16.4
14.9
14.5
16.5
0.030
GraphCL. et 2020)
0.027
0,066
22.9
0.095
9.7
9.6
10.2
GPT-GNN (Hu al., 2020b)
0.039
0.103
35.7
28.8
54.1
0.158
1.75
12.0
24,8
14,8
12.2
0.032
GmphMVP (Ting et al., 2021)
0.031
0.070
28.5
26.3
46.9
0.082
1.63
10.2
10.3
10.4
11.2
0.033
GEM (Fang et al., 2021)
0.034
27.7
52.1
0.089
13.3
13.2
3D Infomax (Stärk al., 2022)
0.034
0.075
29.8
25.7
48.8
0.122
1.67
12.7
12.5
12.4
13.0
0.033
PosPred (Jiao ct al., 2022)
0.024
0.067
25.1
20.9
40.6
0.115
1.46
10.2
10.3
10.2
10.9
0.035
3D-MGP (Jiao
2022)
0.020
0.057
37.1
0.092
8.6
8.7
9.3
0.026
Schnel (Schilt al., 2017)
0.033
0.235
41
34
63
0.073
14
19
14
14
0.033
PhysNet (Unke & Mcuwly, 2019)
0.0529
0.0615
32.9
24.7
42.5
0.765
1.39
8.15
8.34
8.42
9.4
0.028
Cormorant (Anderson al. 2019)
0.038
0,085
0.961
21
0.026
DimeNet++ (Kliepera al. 2020)
0,0297
0.0435
24.6
19.5
32.6
0.331
1.21
6.32
6,28
6.53
7.56
0.023
PaiNN (Schitt et al., 2021)
0.012
0.045
27.6
20.4
45.7
0.066
1.28
5.85
5.83
5.98
7.35
0.024
LieTF (Hutchinson el al.,, 2021)
0,082
0.448
17
TorchMDDNeet (Thölke & De Fabritiis, 2021)
0.011
0.059
20.3
17.5
36,1
0.033
1.84
6.15
6.38
6.16
7.62
0.026
EGNN (Satorras al., 2021)
0.029
0.071
29
25
48
0.106
1.55
11
12
12
12
0.031
NoisyNodo (Godwin al., 2022)
0.025
0.052
20.4
18.6
28.6
0,70
1.16
7,30
7.57
7.43
8.30
0.025
Transtormer-M Ml (ours)
0.037
0,041
17.5
16.2
27.4
0.075
1.18
9.37
9.41
9.39
9.63
0.022
4.4 QM9 PERFORMANCE (3D)
We use the QM9 dataset (Ramakrishnan et al., 2014) to evaluate our Transformer-M on molecular tasks in the 3D data format. QM9 is a quantum chemistry benchmark consisting of 134k stable small organic molecules. These molecules correspond to the subset of all 133,885 species out of the GDB-17 chemical universe of 166 billion organic molecules. Each molecule associated with 12 targets covering its energetic, electronic, and thermodynamic properties. The 3D geometric structure of the molecule is used as input. Following Thölke & De Fabritiis (2021), we randomly choose 10,000 and 10,831 molecules for validation and test evaluation, respectively. The remaining molecules are used to fine-tune our Transformer-M model. We observed that several previous works used different data splitting ratios didn't describe the evaluation details. For fair comparison, choose baselines that use similar splitting ratios in the original papers. The details of baselines and experiment settings presented in Appendix B.4.
8

---

Published as a conference paper at ICLR 2023
The results are presented in Table 3. It can be seen that our Transformer-N M achieves competitive performance compared to those baselines, suggesting that the model is compatible with 3D molecular data. In particular, Transformer-M performs best on HUMO, LUMO, and HUMO-LUMO predictions. This indicates that the knowledge learned the pre-training task transfers better to similar tasks. Note that the model doesn't perform quite well some other tasks. We believe the Transformer-M can be improved in several aspects, including employing a carefully designed output layer (Thölke & De Fabritis,, 2021) or pre-training with more self-supervised training signals.
4.5 ABLATION STUDY
In this subsection, we conduct a series of experiments to investigate the key designs of our Transformer-M In this paper, we use two training objectives to train the model, and we will ablate the effect of the two training objectives. Besides, we use three modes activate different channels with pre-defined distribution, and we will study the impact of the distribution on the final performance. Due to space limitation, we present more analysis on our Transformer-M model in Appendix B.5.
Table 4: Impact of pre-training tasks and mode distribution on Transformer-M. All other hyperparam- eters are kept the same for a fair comparison.
Tasks
PCQM4My2
QM9
Mode Distribution
2D-3D joint Pre-training
3D Position Denoising
(pin: PID pmna.))
Valid MAE
ELUMO
A
x
0.0878
26.5
23.8
42.3
1:2:1
0.0811
20.1
19.3
31.2
1:2:1
0.0787
17.5
16.2
27.4
1:1:1
0.0796
18.1
16.9
27.9
1:2:2
0.0789
17.8
16.6
27.5
Impact of the pre-training tasks. As stated in Section 4.1, our Transformer-M M model is pre-trained on the PCQM4Mv2 training set via two tasks: (1) predicting the HOMO-LUMO gap of molecules in both 2D and formats. (2) 3D position denoising. We conduct ablation studies on both PCQM4Mv2 and QM9 datasets to check whether both objectives benefit downstream tasks. In detail, we conduct two additional experiments. The first experiment training Transformer-M models from seratch PCQM4Mv2 using its 2D graph data and QM9 using its 3D geometric data to check the benefit of the overall pre-training method. The second experiment is pre-training Transformer-M without using the 3D denoising task to study the effectiveness of the proposed 2D-3D joint pre-training approach. The results are shown in Table 4. It can be seen that the joint pre-training significantly boosts the performance both PCQM4Mv2 and QM9 datasets. Besides, the 3D Position Denoising task is also beneficial, especially on the QM9 dataset in the 3D format.
Impact of mode distribution. Denote (P2D- PSD. P2D&3D) as the probability the modes mentioned in Section 4.1. We conduct experiments to investigate the influence of different distributions on the model performance. We select three distribution with (PID. P.D. PID&ED) being: 1:1:1, 1:2:2, and 1:2:1. The results are presented in Table 4. We obtain consistent conclusions on both PCQM4Mv2 and QM9 datasets: 1) for all three configurations, our Transformer-M model achieves strong performance, which shows that our joint training is robust to hyperparameter selection; 2) Using slightly larger probability on the 3D mode achieves the best results.
5 CONCLUSION
In this work, we take the first step toward -eneral-purpose molecular models. The proposed Transformer-M offers promising way to handle molecular tasks 2D and formats. We use two separate channels to encode 2D and 3D structural information and integrate them into the backbone Transformer. When the input data in a particular format, the corresponding channel will be acti- vated, and the other will be disabled. Through simple training tasks on 2D and 3D molecular data, our model automatically learns to leverage chemical knowledge from different data formats and correctly capture the representations. Extensive experiments conducted, and all empirical results show that our Transformer-M M can achieve strong performance on 2D and 3D tasks simultancously. The potential of our Transformer-M can be further explored in broad range of applications in chemistry.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We thank Shanda Li for the helpful discussions. We also thank all the anonymous reviewers for the very careful and detailed reviews as well as the valuable suggestions. Their help has further enhanced our work. This work is supported by National Key R&D Program of China (2022ZD01 14900) and National Science Foundation of China (NSFC62276005), This work partially supported by the Shanghai Committee of Science and Technology (Grant No. 21DZ1 100100).
REFERENCES
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprini arXiv: 1607.06500 2016.
Ballester, P. J., Mitchell, J., and B. A machine learning approach to predicting protein- binding affinity with applications to molecular docking. Bioinformatics Oxford, 2010.
Rémy Brossard, Oriel Frigo, and David Dehaene. Graph convolutions that can finally model local structure. arXiv preprint urXiv:201 1506, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler. Mateusz Litwin, Scott Gray, Ben- jamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL https: http:///cocceece...c.ce/. / htwpooeeeeeinnn..eee............ Paper cc/paper/ (atper//020/flle/// httpee/111cccddcfcaeaaaa 967 774 8bfb8acclllccoooceee 42f644 Paper .dff
Kieron Burke. Perspective on density functional theory. The Journal of chemical physics, 136(15): 150901, 2012.
Chanussot, A Da, S Goyal, Lavril, M Shuaibi, Riviere, Tran, J dleras-Domingo, C Ho, W Hu, et al. The open catalyst 2020 (oc20) dataset and community challenges. arxiv. arXiv, 2020.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5)e1603015, 2017.
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Fran- cisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Unsupervised cross-lingual representation learning at scale. arXiv preprint arXiv: 1XX.:1911.00216.. 2019.
Xuefeng Cui. Coatgin: Marrying convolution and attention for graph-based molecule property
prediction. bioRxiv, 2022.
Tomasz Danel, Przemyslaw Spurek, Jacek Tabor, Marek Smieja, Lukasz Struski, Agnieszka Slowik, and Lukasz Maziarka. Spatial graph convolutional networks. In International Conference on Neural Information Processing, pp. 668-675. Springer, 2020.
Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4186, 2019.
10

---

Published as a conference paper at ICLR 2023
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Pavlo O Dral, Xin Wu, Lasse Spörkel, Axel Koslowski, and Walter Thiel. Semiempirical quantum- chemical wrthogonalizaaion-coreee methods: benchmarks for ground-state properties. Journal of Chemical Theory and Computation, 12(3):1097-1129, 2016.
Vijay Prakash Dwivedi and Xavier Bresson. generalization of transformer networks graphs.
arXiv preprint arXiv:2012.09699 2020.
Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, and Haifeng Wang. Chemrl-gem Geometry enhanced molecular representation learning for property prediction. arXiv preprint 10X.1::0066061300 2021.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp. 1263-1272. PMLR, 2017.
Jonathan Godwin, Michael Schaarschmidt, Alexander L Gaunt, Alvaro
Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia. Simple GNN regularisation for 3d molecular property prediction and beyond. In International Conference on Learning Representations, 2022. URL https: / haapp:://opeeeee /oppnreview....c///o//e wwrriiewnee//eee net/ wwwoomooooooliii.o wwe.ddwwwwwWweoll.
William L Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In NIPS, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9738, 2020.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovee. Strategies for pre-training graph neural networks. arXiv preprint 1905. 12265, 2019.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint 0rXiv:2005.00633 2020a.
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Isc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.09433 2021.
Ziniu Hu, Yuxiao Dong, Kuansan Wang. Kei-Wei Chang, and Yizhou Sun. Gpt-gnn: Generative pre-training of graph neural networks. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1857-1867, 2020b.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European conference on computer vision, pp. 646-661. Springer, 2016.
Md Shamim Hussain, Mohammed J Zaki, and Dharmashankar Subramanian. Global self-attention as a replacement for graph convolution. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 655 6655666,410 2022.
Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lictransformer: Equivariant self-attention for lie groups. In International Conference on Machine Learning, pp. 4533 PMLR, 2021.
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. 3d equivariant molecular graph pretraining, arXiv preprint arXiv:2207.08822 0.8844.. 2022.

---

Published as a conference paper at ICLR 2023
José Jiménez, Miha Skalic, Gerard Martinez-Rosell and Gianni De Fabritiis. K deep: protein-ligan absolute binding affinity prediction via 3d-convolutional neural networks. Journal of chemical information and modeling, 58(2)2887-29, 2018.
Baoyu Jing, Yuejia Xiang, Xi Chen, Yu Chen, and Hanghang Tong. Graph-mvp: Multi-view prototypical contrastive learning for multiplex graphs. arXiv preprint arXiv:2109.03560, 2021.
Jinwoo Kim, Tien Dat Nguyen, Sconwoo Min, Sungjun Cho, Moontae Lee, Honglak Lec, and Scunghoon Hong. Pure transformers are powerful graph learners. arXiv preprint arXiv:2207.02505, 2022.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1412. 6980, 2014.
Thomas N Kipf and Max Welling, Semi-supervised classification with graph convolutional networks. arXiv preprint 1609. 02900, 2016.
Johannes Klicpera, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03/23, 2020.
Kezhi Kong, Guohao Li, Mucong Ding, Zuxuan Wu, Chen Zhu, Bernard Ghanem, Gavin Taylor, and Tom Goldstein. Flag: Adversarial data augmentation for graph neural networks. arXiv preprint arXiv:2010.09891, 2020.
Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Létourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 01221618-2122) 2021.
Guillaume Lample and Alexis Conncau. Cross-lingual language model pretraining. arXiv preprint arXiv:1901.07291, 2019.
Greg Landrum. Rdkit: Open-source cheminformatics software. Github, 2016. URL https:
Guohao Li, Chenxin Xiong, Ali Thabet, and Bernard Ghanem. Deepergen: All you need to train deeper gens, arXiv preprint arXiv:2006.07739, 2020.
Shuangli Li, Jingbo Zhou, Tong Xu, Liang Huang, Fan Wang, Haoyi Xiong, Weili Huang, Dejing Dou, and Hui Xiong. Structure-aware interactive graph neural networks for the prediction of protein- ligand binding affinity. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining. pp. 975, 2021.
J. Lim, S. Ryu, K. Park, J. Choe, and W. Y. Kim. Predicting drug-target interaction using novel graph neural network with 3d structure- aruueuue-embeddded graph representation. Journal of Chemical Information and Modeling, 59(9), 2019.
Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre- training molecular graph representation with 3d geometry. arXiv preprint 110.0:2110.07728. 2021a.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference Learning Representations (ICLR), 2022.
Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, and Luke Zettlemoyer. Multilingual denoising pre-training for neural machine translation. Transactions of the Association for Computational Linguistics, 8:726-742, 2020.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.1 14030, 2021b.
12

---

Published as a conference paper at ICLR 2023
Shengjie Luo, Shanda Li, Shuxin Zheng, Tie- Liu, Liwei Wang, and Di He. Your transformer may not be as powerful as you expect. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: wwwpeeeeeeeeeeeeeeeeeeeeeeeeeeee ew net / forum? d d=NgFFNdsOGD.
Nakata Maho. The pubchemqe project: A large chemical database from the first principle calculations. In AIP conference proceedings, volume 1702, pp. 090058. AIP Publishing L.C., 2015.
Lukasz Maziarka, Tomasz Danel, Slawomir Mucha, Krzysztof Rataj. Jacek Tabor, and Stanislaw Jastrzebski. Molecule attention transformer. arXiv preprint Stt::2002.082664 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, arXiv preprint arXiv:1312.5602, 2013.
Maho Nakata and Tomomi Shimazaki. Pubchemqe project: a large-scale first-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 57()): 1300-1308, 2017.
T. Nguyen, H. Le, T. P. Quinn, T. Nguyen, and S. Venkatesh. Graphdta: Predicting binding affinity with graph neural networks. Bioinformatics, 2020.
Wonpyo Park, Woong-Gi Chang, Donggeon Lee, Juntae Kim, et al. Grpe: Relative positional encoding for graph transformer. In ICLR2022 Machine Learning for Drug Discovery, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference Machine Learning, pp. 8748-8763. PMLR, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL
http://mmmlooro//.....ee www.oooppp.p../oaeeooooo peneers/vv1/ /v21/20- 0t4.html.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatolc Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Ladislav Rampášek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. arXiv preprint arXiv:2205.12455 2022.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez. Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint 2rXiv:2205.06173 2022.
Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying. Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 8459-8468 PMLR, 13-18 Jul 2020. URL https: ppwceedeed i ngs ml wwm.m..pp... presss / vv1v 9 wwwanhee-gozzznoo e2 html.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. (n) equivariant graph neural networks. In International conference on machine learning, pp. 9323-9332, PMLR, 2021.
Bernhard Scholkopf. Kah-Kay Kay Sung, Christopher JC Burges, Federico Girosi, Partha Niyogi, Tomaso Poggio, and Vladimir Vapnik. Comparing support vector machines with gaussian kernlls to radial basis function classifiers. IEEE transactions on Signal Processing, 45(1 45(11):2758-276) 1997.
13

---

Published as a conference paper at ICLR 2023
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robers Müller. Schnet: continuous-fin filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference Machine Learning, pp. 9377-9388. PMLR, 2021.
Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie- Yan Liu. Benchmarking graphormer large-scale molecular modeling datasets. arXiv preprint arXiv:2203.04810 2022.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
Y. Song. S. Zheng, Z. Niu, Z. H. Fu, and Y. Yang. Communicative representation learning on attributed molecular graphs. In Twenty-Ninth International Joint Conference on Artificial Intelligence and Seventeenth Pacific Rim International Conference on Artificial Intelligence IJCAI-PRICAI-20, 2020.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ily Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1955 2014.
Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning, pp.
PMLR, 2022.
Marta M Seapniwwska-Dziubinne Piotr Zielenkiewicz, and Pawel Siedlecki. Development and evaluation of deep learning model for protein-ligan binding affinity prediction. Bioinformatics, 34(21):36663674, 2018.
Jonathan M Stokes, Kevin Yang, Kyle Swanson, Wengong Jin, Andres Cubillos-Ruiz, Nina M Donghia, Craig R MacNair, Shawn French, Lindsey Carfrae, Zohar Bloom-Ackermann, et al. A deep learning approach antibiotic discovery. Cell, 180(4):688-702, 2020.
Fan-Yun Sun, Jordan Hoffmann, Vikas Verma, and Jian Tang. Infograph: Unsupervised and semi- supervised graph-level representation learning via mutual information maximization. arXiv preprint arXiv:1908.00 01000, 2019.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. International Conference on Learning Representations, 2021.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translationeequivarian neural networks for 3d point clouds. arXiv preprini 1102.08800.0822 2018
Oliver T Unke and Markus Meuwly. Physnet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of chemical theory and computation, 1576))3678-369, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. ICLR, 2018.
Limei Wang, Yi Liu, Yuchao Lin, Haoran Liu, and Shuiwang Ji. Comenct: Towards complete and efficient message passing for 3d molecular graphs. arXiv preprint arXiv:2206.085/ 2022.
R. Wang. X. Fang, Lu, C. Y. Yang, and S. Wang. The pdbbind database:? methodologies and updates. Journal of Medicinal Chemistry, 48(12):4111-411) 2005a.
14

---

Published as a conference paper at ICLR 2023
Renxiao Wang, Xueliang Fang. Yipin Lu, and Shaomeng Wang. The pdbbind database: Collection of binding affinities for protein- ligand complexes with known dice-dimensional structures. Journal of medicinal chemistry, 47(12):2977-299) 2004.
Renxiao Wang. Xueliang Fang, Yipin Lu, Chao-Yie Yang, and Shaomeng Wang. The pdbbind database: methodologies and updates. Journal of medicinal chemistry, (8(12):4111-4111 2005b.
David Weininger. Smiles, a chemical language and information system. 1. introduction to methodol- ogy and encoding rules. Journal of chemical information and computer sciences, 28(1):31-38) 1988.
William J Wiswesser. Historic development of chemical notations. Journal of chemical information and computer sciences, 25(3):258-28) 263, 1985.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 2020.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2019. URL https: //openrevieww.oo//////r/. net htt////ormmmiiiieeeeeeeeeeeemmee
1otp:/ew.nn////ommmtidorr-- masaeiddrrr6i115S0 61A5Km.
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: A massively multilingual pre-trained text-to-text transformer. arXiv preprint arXiv:2010.1 1934, 2020.
K. Yang. K. Swanson, W. Jin, C. W. Coley, and Barzilay. Analyzing learned molecular rep- resentations for property prediction. Journal of Chemical Information and Modeling. 59(8). 2019.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877-28885 2021a.
Chengxuan Ying, Mingqi Yang, Shuxin Zheng, Guolin Ke, Shengjie Luo, Tianle Cai, Chenglin Wu, Yuxin Wang, Yanming Shen, and Di He. First place solution of kdd cup 2021 ogb large-scale challenge graph-level track. arXiv preprint arXiv:2106.08273 2021b.
Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. Advances in Neural Information Processing Systems, 33: 5812-5823, 2020.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez- Gonzzlz. Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.0011 00133, 2022.
Bohang Zhang, Shengjie Luo, Liwei Wang, and Di He. Rethinking the expressive power of GNNs via graph biconnectivity. In International Conference on Learning Representations, 2023. URL
httpps://pperreview net. forum? fidammiidrrrh// 76K
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and EJPRL Weinan. Deep potential molccular dynamics: scalable model with the accuracy of quantum mechanies. Physical review letters, 120 1430:143001, 2018.
I. Zheng, J. Fan, and Y. Mu. Onionnet: a multiple-layer inter-molecular contact based convolutional neural network for protein-ligand binding affinity prediction. ACS Omega, 4(14). 2019.
Chen Zhu, Yu Cheng, Zhe Gan, Sigi Sun, Tom Goldstein, and Jingjing Liu. Freelb: Enhanced adversarial training for natural language understanding. In ICLR, 2020.
Jinhua Zhu, Yingee Xia, Lijun Wu, Shufang Xie, Tao Qin, Wengang Zhou, Houqiang Li, and Yan Liu. Unified 2d and 3d pre- pre-training of molecular representations. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 2626-2636. 2022.
15

---

Published as a conference paper at ICLR 2023
A IMPLEMENTATION DETAILS OF TRANSFORMER-N
2D-3D joint Pre-training. To effectively utilize molecular data in both 2D and 3D formats, use simple strategy. During training, each data instance in a batch has three modes: (1) activate the 2D channels, and disable the 3D channels (2D mode): (2) activate the 3D channels, and disable the 2D channels (3D mode); (3) activate both the 2D and 3D channels (2D+3D mode). The mode of each data instance during training is randomly drawn on the fly according to a pre-defined distribution (PED. P3D. PED&3D) implemented similarly to Dropout (Srivastava al., 2014). For each data instance, the model is required to predict its HOMO-LUMO energy gap across the above three modes.
Prediction head for position output. (3D Position Denoising) Following Godwin et al. (2022); Zaidi et al. (2022), we further adopt the 3D Position Denoising task as a self-supervised learning objective. During training, if data instance is in the 3D mode, add Gaussian noise to the positions of each atom. The model is required to predict the noise from the noisy input. Formally, let ry E R denote the atom positions of molecule, and the noisy version of the atom positions is denoted by R {ri + σc1, r2 + oe,, Pn + ac,); where a is the scaling factor of noise and ~ N (0, I). The prediction of the model is denoted by (ćI.... ) Following Shi et al. (2022), we use an SE(3) equivariant attention layer as the prediction head:
R (rı, r2,..,,
Ea W W k= 0,
(6)
BBEV
where x is the output of the last Transformer block, aiy the attention score between atom and j calculated by Eqn.4, A is the k-th element of the directional vector ||r--, between atom and
j. and W E Rdxd W E R are learnable weight matrices. The denoising loss of a batch of molecules & = [vi ..... V is calculated by the cosine similarity between the predicted noises and the ground-truth noises as:
Cpos
Σ V'EE JEV
1 [3
Σ
- (1 (1-cos cos (E), 2j))
Prediction head for scalar output. We follow Ying et al. (2021a) to add a pseudo atom to each molecule and make connections between the pseudo atom and other atoms individually. The bias terms in the self-attention layer between this pseudo atom and other atoms are set to the same learnable scalar. For molecule-level prediction, we can simply use the representation of this pseudo atom to predict targets.
B EXPERIMENTAL DETAILS
B. LARGE-SCALE PRE-TRAINING
Dataset. Our Transformer-M model is pre-trained on the training set of PCQM4Mv2 from the OGB Large-Scale Challenge (Hu al., 2021). PCQM4Mv2 is quantum chemistry dataset originally curated under the PubChemQC project (Maho, 2015; Nakata & Shimazaki, 2017). The total number of training samples is 3.37 million. Each molecule in the training is associated with both 2D graph structures and 3D geometric structures. The HOMO-LUMO energy gap of each molecule is provided, which is obtained by DFT-based geometry optimization (Burke, 2012). According to the OGB-LSC (Hu et al., 2021), the HOMO-I.UMO energy gap is one of the most practically-relevamt quantum chemical properties of molecules since it is related to reactivity, photoexcitation, and charge transport. Being the largest publicly available dataset for molecular property prediction, PCQM4Mv2 is considered to be a challenging benchmark for molecular models.
Settings. Our Transformer-M M model consists of 12 layers. The dimension of hidden layers and feed-forward layers is set to 768. The number of attention heads is set to 32. The number of Gaussian Basis kernels is set to 128. We use AdamW (Kingma & Ba, 2014) the optimizer and set its hyperparameter e to 1e-8 and (3) 82) to (0.9,0.999). The gradient clip norm is set to 5.0. The peak learning rate is set to 2e-4. The batch size is set to 1024. The model trained for 1.5 million steps
16

---

Published as a conference paper at ICLR 2023
with 90k-step warm-up stage. After the warm-up stage, the learning rate decays linearly to zero. The dropout ratios for the imput embeddings, attention matrices, and hidden representations are set to 0.0, 0.1, and 0.0 respectively. The weight decay is set to 0.0. We also employ the stochastic depth (Huang et al., 2016) and set the probability to 0.2. The probability (PaD, PaD, (DDD2)))))) of each data instance entering the three modes mentioned in Section 4.1 is set to (0.2, 0.5, 0.3). The scaling factor a of added noise in the 3D Position Denoising task is set to 0.2. The ratio of the supervised loss the denoising loss set to 1:1, All models are trained on 4 NVIDIA Tesla A100 GPUs.
B.2 PCQM4MV2
Baselines. We compare our Transformer-M with several competitive baselines. These models fall into two categories: message passing neural network (MPNN) variants and Graph Transformers.
For MPNN variants, we include two widely used models, GCN (Kipf & Welling, 2016) and GIN (Xu et al., 2019), and their variants with virtual node (VN) (Gilmer et al., 2017; Hu al.. 2020a). Additionally, we compare GINE-vN (Brossard et al., 2020) and DeeperGCN-vN (Li et al., 2020). GINE is the multi-hop version of GIN. DeeperGCN a 12-layer GNN model with carefully designed aggregators. The result of MLP-Fingerprint (Hu et al., 2021) is also reported.
We also compare several Graph Transformer models. Graphormer (Ying et al., 2021a) developed graph structural encodings and integrated them into standard Transformer model. It achieved impressive performance across several world competitions (Ying et al., 2021b; Shi et al., 2022). CoAt- GIN (Cui, 2022) is a hybrid architecture combining both Convolution and Attention. TokenGT (Kim et al., 2022) adopted the standard Transformer architecture without graph-specific modifications. GraphGPS (Rampášek et al., 2022) proposed framework to integrate the positional and structural encodings, local paassigg- mechanism, and global attention mechanism into the Transformer model. GRPE (Park et al., 2022) proposed graph-specific relative positional encoding and considerd both node-spatial and node-edge relations. EGT (Hussain et al., 2022) exclusively used global self-attention as an aggregation mechanism rather than static localized convolutional aggregation, and utilized edge channels to capture structural information.
B.3 PDBBIND
Dataset. PDBBind is a well-known dataset that provides comprehensive collection of experimentally measured binding affinity data for biomolecular complexes deposited in the Protein Data Bank (PDB) (Wang et al., 2005a). The task requires models to predict the binding affinity value pK, (or - log Ka, - log K2) of protein-ligand complexes, which is extremely vital for drug discovery. In our experiment, we use the PDBBind v2016 dataset, which is widely used in recent works (Li et al., 2021). The PDBBind dataset includes three overlapped subsets called the general, refined, and core sets. The general set contains all 13,283 poteinnllgaadi complexes, while the 4,057 complexes in the refined set selected out of the general set with better quality. Moreover, the core set serves the highest quality benchmark for testing. To avoid data leakage, we remove the data instances in the core set from the refined set. After training, we evaluate our model on the core set. The evaluation metrics include Pearson's is correlation coefficient (R), Mean Absolute Error (MAE), Root-Mean Squared Error (RMSE), and Standard Deviation (SD).
Baselines. We compare our Transformer- MM with several competitive baselines. These models mainly fall into three categories: classic Machine Learning methods, Convolution Neural Network (CNN) based methods, Graph Neural Network (GNN) based methods.
First, we report the results of LR, SVR, and RF-Score (Ballester al., 2010), which employed traditional machine learning approaches to predict the binding affinities. Second, inspired by the success of CNNs in computer vision, Stepniewska-Dziubeee model that represents the complexes via 3D grid and utilizes 3D convolution to produce feature maps. Zheng et al. (2019) introduced OnionNet, which also used CNNs to extract features based on rotation-free clement-pair specific contacts between atoms of proteins and ligands.
al. (2018) proposed the Pafnucy
There are also several works that leverage GNNs to improve the performance of the PDBBind dataset. GraphDTA (Nguyen et al., 2020) represented protein-ligand complexes as 2D graphs and used GNN models to predict the affinity score. GNN-DTI (Lim et al., 2019) incorporated the 3D structures of ligand complexes into GNNs. DMPNN (Yang et al., 2019) operated over hybrid representation that combines convolutions and descriptors. SGCN (Danel et al., 2020) is a GCN-inspired architecture that leverages node positions. MAT (Maziarka et al., 2020) augmented the
17

---

Published as a conference paper at ICLR 2023
attention mechanism in the standard Transformer model with inter-atomic distances and molecular graph structures. DimeNet (Kliepera et al., 2020) developed the atom-pair embeddings and utilized directional information between atoms. CMPNN (Song et al., 2020) introduced communicative kernel and a message booster module to strengthen the message passing between atoms. SIGN (Li al., 2021)) )) proposed polar-inspired graph attention layers and pairwise interactive pooling layers to utilize the biomolecular structural information.
Settings. We fine-tune the pre-trained Transformer-M on the PDBBind dataset. We use AdamW W (Kingma & Ba, 2014) as the optimizer and set its hyperparameter e to 1e-8 and (8), to (0.9,0.999). The gradient clip norm is set 5.0. The peak learning rate set to 2c-4. The total number of epochs is set to 30. The ratio of the warm-up steps to the total steps is set to 0.06. The batch size is set to 16. The dropout ratios for the input embeddings, attention matrices, and hidden representations are set to 0.0, 0.1, and 0.0 respectively. The weight decay is set to 0.0. Following ( (Ying et al., 2021a), We use FLAG (Kong et al., 2020) with minor modifications for graph data augmentation. In particular, except for the step size and the number of adversarial attack steps m, we also employ a projection step in Zhu et al. (2020) with maximum perturbation y. These hyperparaters are set to the following configurations: 0.01, = 4, 7 0.01. All models trained on 2 NVIDIA Tesla V100 GPUs.
B.4 QM9
Dataset. QM9 (Ramakrishnan et al., 2014) a quantum chemistry benchmark consisting of 134k stable small organic molecules. These molecules correspond to the subset of all 133,885 species out of the GDB-17 chemical universe of 166 billion organic molecules. Each molecule is associated with 12 targets covering its energetic, electronic, and thermodynamic properties. The 3D geometric structure of the molecule is used as input. Following Thölke & De Fabritiss (2021), we randomly choose 10,000 and 10,831 molecules for validation and test evaluation, respectively, The remaining molecules are used to fine-tune our Transformer-N -M model.
Baselines. We comprehensively compare our Transformer-M M with both pre-training methods and 3D molecular models. First, we follow Jiao et al. (2022) to compare several pre-training methods. Hu et al. (2019) proposed strategy pre-train GNNs via both node-level and graph-level tasks. Sun et al. (2019) maximized the mutual information between graph-level representations and substructure representations the pre-training tasks. You al. (2020) instead used contrastive learning to pre-train GNNs. There are also several works that utilize 3D geometric structures during pre- pre-training. Jing et al. (2021) maximized the mutual information between 2D and 3D representations. Fang et al. (2021) proposed strategy to learn spatial information by utilizing both local and global 3D structures. Stärk et al. (2022) used two encoders to capture 21) and 3DD structural information separately while maximizing the mutual information between 2D and 3D representations. Jiao et al. (2022) adopted an cquivariant energy-basee model and developed a node-level pretraining loss for force prediction. We report the results of these methods from (Jiao et al., 2022) for comparison.
Second, we follow Thölke& & De Fabritiis (2021) to compare 3D molecular models. Schütt et al. (2017) used continuous-filee convolution layers to model quantum interactions in molecules. An- derson et al. (2019) developed a GNN model equipped with activation functions being covariant to rotations. Klicpera et al. (2020) proposed directional message passing, which uses atom-pair embeddings and utilizes directional information between atoms. Schütt et al. (2021) proposed the polarizable atom interaction neural network (PaiNN) that uses an equivariant message passing mech- anism. Hutchinson et al. (2021) built upon the Transformer model consisting of attention layers that are equivariant arbitrary Lie groups and their discrete subgroups. Thölke & De Fabritiis (2021) also developed a Transformer variant with layers designed by prior physical and chemical knowledge. Satorras et al. (2021) proposed the EGNN model which does not require computationally expensive higher-order representations in immediate layers to keep equivariance, and can be easily scaled to higher-dimeesional spaces. Godwin et al. (2022) proposed the 3D position denoising task and verified it the Graph Network-based Simulator (GNS) model Sanchez-Gonzalee et al., 2020).
Settings. We fine-tune the pre-trained Transformer-M on the QM9 dataset. Following Thölke & De Fabritiis (2021), we adopt the Mean Squared Error (MSE) loss during training and use the Mean Absolute Error (MAE) loss function during evaluation. We also adopt label standardization for stable training. We use AdamW as the optimizer, and set the hyper-parameter (8,, B2) to (0.9.0.999). The gradient clip norm is set to 5.0. The peak learning rate is set to 7c-5. The batch size is set 128. The dropout ratios for the input embeddings, attention matrices, and hidden
to 1e-8 and
18

---

Published as a conference paper at ICLR 2023
representations set to 0.0, 0.1, and 0.0 respectively. The weight decay is set to 0.0. The model is fine-tuned for 600k steps with 60k-step warm-up stage. After the warm-up stage, the learning rate decays linearly to zero. All models are trained on NVIDIA A100 GPU.
B.5 MORE ANALYSIS
Investigation on the generality of the design methodology of Transformer-M. In this work, we develop our Transformer-M model based on the Transformer backbone and integrate separate 2D and 3D channels (implemented by encoding methods) to encode the structural information of 2D general design methodology for handling
and 3D molecular data. As stated in Section 3.2. molecular data in different forms, which works well with different structural encoding instantiations. To demonstrate its generality and effectiveness, we further conduct experiments with other structural encodings from GRPE (Park et al., 2022) and EGT (Hussain et al., 2022), which are competitive baselines on the PCQM4Mv2 benchmark as shown in Table All the hyperparameters are kept the same as the settings in Appendix B.I for fair comparison. The results are presented in Table 5. It can be easily seen that our Transformer- .M model equipped with different encoding methods can consistently obtain significantly better performance than the corresponding vanilla 2D models, which indeed verifies the generality and effectiveness of the design methodology of Transformer-M.
Table 5: Performance of Transformer-M with different structural encodings on PCQM4Mv2 validation set. The evaluation metric is the Mean Absolute Error (MAE) [eV]. For fair comparison, all hyperparameters are kept the same as the settings in Appendix B.I.
Training in
Training in the Transformer-M framework
Structural Encodings
A
the vanilla setting (2D only)
Graphormer (Y (Ying et al., 2021a)
0.0878
0.0787
10.4%
GRPE (Park et al., 2022)
0,0884
0.0798
9.7%
EGGTHHssaii et al., 2022)
0.0876
0.0794
9.4%
Table 6: Ablation study on the impact of 3D conformers calculated by different methods. Experiments conducted the PCQM4Mv2 dataset. The evaluation metric is the Mean Absolute Error (MAE) [eV]. For a fair comparison, all hyperparameters are kept the same as the settings in Appendix B.
3D Conformer [by RDKit (Landrum, 2016)] [by DFT, from PCQM4Mv2 (Hu 2021)]
3D Conformer
2D
Valid MAE
0.0878
0.0872
0.0792
0.0787
Investigation the impact of 3D conformers calculated by different methods. Besides the versatility to handle molecules in different formats, our Transformer-M further achieves strong performance on various challenging molecular tasks, shown in Section 4. On PCQM4My2 validation set (2D only), our Transformer-M establishes a new state-of-the-art, which mainly credits the newly introduced 2D-3D joint training strategy in Section 3.2. The chemical knowledge in the 3D geometric structure can be leveraged during joint training and boost the performance of 2D tasks. Since the benefits of the 3D geometric structure are observed, it is natural ask how the quality of calculated 3D conformers influences the performance of Transformer-M.
To investigate this question, we additionally use the RDKit (Landrum, 2016) generate one 3D conformed for each molecule in the training set of PCQM4Mv2. Compared to officially provided DFT-optimized geometric structures, the structures are less costly obtain by using RDKit while also being less accurate. Thus, each molecule has its 2D molecular graph, 3D conformer calculated by DFT, and 3D conformer calculated by RDKit. Based on such a dataset, we conduct three additional experiments. Firstly, we train our Transformer-M model only using 2D molecular graphs. In this experiment, only the 2D channels are activated. Secondly, we train our Transformer-M model using both 2D molecular graphs (encoded by 2D channels) and 3D conformers generated by RDKit (encoded by 3D channels). Thirdly, we train our Transformer-N M model using 2D molecular graphs, 3D conformers generated by RDKit, and 3D conformers calculated by DFT. In this experiment,
19

---

Published as a conference paper at ICLR 2023
use two sets of 3D channels to separately encode structural information of 3D RDKit conformers and 3D DFT conformers. During training. when a data instance enters 3D or 2D+3D modes, both sets of 3D channels are activated and integrated. For all three experiments, the hyperparameters. of Transformer-M are kept the same as the settings in Appendix B.1. The results are presented in Table 6.
We that the quality of the 3D conformer matters in the final performance: leveraging 3D con- formers generated by RDKit (second line) brings minor gains compared to using 2D molecular graphs only (first line). On the contrary, when leveraging 3D conformers calculated by DFT, the improvement is significant (the last two lines). From the practical view, it will be interesting to investigate the influence of 3D conformers calculated by methods that are more accurate than RDKit while more efficient than DFT, e.g., semiempirical methods (Dral et al., 2016), which we leave as future work.
Investigation on the effectiveness of Transformer-M pre-training. We provide additional results on the effectiveness of our model on both PDBBind (2D+3D) and QM9 (3D) downstream datasets.
Firstly, to verify the effectiveness of Transformer-M MM pre-training on the PDBBind dataset, we further pre-train the Graphormer model (Ying et al., 2021a) on the same PCQM4Mv2 dataset as competitive pre-trained baseline. Since the Graphormer model can only handle graph data, we only use the 2D molecular graph of each data instance. All hyperparameters are kept the same as the settings in Appendix B.1. The results are presented in Table
Table 7: Investigation on the effectiveness of Transformer-N pre-training on PDBBind core set (version 2016). The evaluation metrics include Pearson's correlation coefficient (R), Mean Absolute Error (MAE), Root-Mean Squared Error (RMSE), and Standard Deviation (SD).
PDBBind core set
Method
R 1
MAE 1
RMSE +
SD J
SIGN (I.i et al., 2021) (the hest baseline in Table 2) Graphormer (Ying et al., 2021a)
0.79710.012
1.027 20.020.023
1.316
1.312 035
0.804
0.998
1.285 20.000
1.271 t000
Transformer-h (ours)
0.830
00.00
1.23220.018
1.207=0.007
We can draw the following conclusions: (1) pre-training is helpful (e.g., 0.797 (R of SIGN model, the best baseline) -> 0.804 (R of pre-trained Graphormer model)); (2) our pre-training method is more significant (e.g.. 0.804 -> 0.830), which indeed demonstrate the effectiveness of our framework.
Secondly, we demonstrate that our pre-training strategy helps learn better Transformer model on downstream QM9 dataset. We conduct two additional experiments the QM9 dataset. In the first ex- periment, we train the 3D geometric Transformer model (Transformer-M with 3D channel only) from scratch. In the second experiment, the 3D geometric Transformer on the PCQM4Mv2 and fine-tune the pre-trained checkpoint on QM9, Due to the time limits and constrained resources, selected six QM9 targets for comparison. All the hyperparameters pre-training and fine-tuning are kept the same. The results are presented in Table 8.
use the 3D Position Denoising task the objective pre-train
Table 8: Investigation on the effectiveness of Transformer-M pre-training on QM9 (Ramakrishnan al., 2014). The evaluation metric the Mean Absolute Error (MAB).
Method
Éномо
ELUMO
Ug
U
H
G
No Pre-training
26.5
23.8
14.78
14.56
14.82
14.95
3D Position Denoising only
18.8
17.6
10.26
10.18
10.33
10.27
2D-3D joint Pre-training + 3D Position Denosing
17.5
16.2
9.37
9.41
9.39
9.63
can be easily seen that our pre-training methods consistently and significantly improve the down- stream performance all six tasks, which indeed demonstrates the effectiveness of our general framework for 3D molecular data. We are aware that we achieve competitive rather than SOTA performance compared with baselines (5 best performance out of 12 targets, see Table 3). For UO, U, H. and G, there still exists a performance gap between our Transformer-M and some latest baselines, which use pretty complicated neural architectures. We believe that exploring more model alternatives and leveraging the wisdom in those networks into our Transformer-N M will further improve the performance, which we will keep working on.
20


Reviews:
[{"review_id": "pQ3xPfpnw3", "reviewer": "Reviewer_yHnH", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper documents a simple method for combining the input representations of a small molecule in 2D and 3D in a single neural network model. The work demonstrates how one can train models using subsets of 2D, 3D, or mixed representations, by modifying the attention head of a transformer architecture to explicitly add a sum over channels of pairwise features generated from (2D) small molecule graphs and channels of pairwise features generated from the small molecule coordinates in 3D. The authors combine a small number of concepts from previous state-of-the art models for small molecules in selecting these features and the unsupervised pre-training tasks, and they manage to show that their new architecture achieves peak performance in the opengraph large-scale challenge for small molecules, and decent performance throughout other tasks.\n", "summary_of_the_review": "This paper ranks top in a rather important small-molecule benchmark; the idea is simple and powerful and the implementation is straightforward.  I think it deserves publication at ICLR2023.\n", "strength_and_weaknesses": "A core strength of this paper is that the high-level idea is simple and clean, yet the results show that it works well: Transformer-M ranks top on the small-molecule OpenGraph large-scale challenge, a world-wide public competition challenge.\n\nThis work is rather strong.  The paper starts to explain the rationale for why the method works as well as it does by performing an ablation study.  I'd have liked to see a couple additional ablations, or other ways to explain why the method works so well: can one keep 3D position denoising for the 3D entries, but drop the joint pre-trainig? How strongly do the various specific features in the pair-wise channels influence the final result?  Do we expect the 2D and 3D features to be \"aligned\" for a given molecule after pre-training the model? Finally, and only as a minor curiosity, what is the impact, if any, of the mode distribution on a downstream task like the QM9 (and if you had that info, could you possibly combine tables 4 and 5 to a single one?)\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "__v-WfygIoi", "reviewer": "Reviewer_JC8f", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes Transformer-M, a Transformer-based model that can take both 2D and 3D molecular formats as input. It adopts various positional encoding techniques to unify 2D and 3D information into transformer as attention bias terms. This model is claimed to be a general-purpose model for molecular tasks. Experiments on several 2D and 3D molecular tasks have been conducted to evaluate the developed method.", "summary_of_the_review": "Overall, I think this work is a good start for developing a general-purpose model for molecular tasks. The empirical results can be valuable to the community. However, the novelty of this work is below the standard of ICLR.\n\n#####Questions#####\n\nHow does Transformer-M deal with the molecules where there might have multiple shortest path between two atoms?\n\n\n#####After rebuttal#####\n\nSince most of my original concerns are well addressed, I increased the score from 5 to 6. I highly recommend the authors to include the additional experiments in the revision and to consider the technical details of the method, as in my response below.", "strength_and_weaknesses": "\n#####Strength\n\n1. The motivation of developing a general-purpose model for molecular tasks is great and might be inspiring to the community.\n\n2. The empirical result of the proposed method is competitive, especially on PCQM4Mv2, where the task is the same as the pretraining objective.\n\n3. This manuscript is well organized and easy to follow.\n\n#####Weaknesses\n\n1. The main concern for this paper is that the novelty is not enough to me. To be specifically, both the 2D and 3D branches are proposed by previous work. The positional encodings are verified to be effective by existing work. This work combines these two branches as a single model, which is not novel in terms of technical contribution.\n\n2. Also, the empirical result on PCQM4Mv2 is good. However, the effectiveness of pretraining Transformer-M has not been demonstrated well, empirically. Specifically, the methods used for comparison in Table 2 (PDBBind) and Table 3 (QM9) do not use the same extra data for pretraining. In this sense, the comparison is not so rigorous and fair to me. Even though, the performance on QM9 is not good, except for the energy-related tasks that are consistent with the pretraining tasks. This cannot support the claim that the proposed Transformer-M is a general-purpose model.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Nnt65getf5", "reviewer": "Reviewer_8vsr", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Molecules can be represented in a variety of chemical data modes, such as a 2D graph or a collection of atoms in 3D space. Most previous work in molecular representation learning has designed networks for a specific data mode, with the risk of failing to learn from other modes. The authors argue that a neural network that is chemically generalized should be able to handle molecule-related tasks across data modes. To accomplish this, the authors created Transformer-M, which is based on the Transformer and can be fed with 2D or 3D molecular data. The results of experiments indicate that Transformer-M performs well in 2D, 3D, and 2D&3D tasks.", "summary_of_the_review": "In summary, this paper is reader-friendly, but novelty is not enough, and some descriptions of the experiments are not clear enough.", "strength_and_weaknesses": "Strengths\n\n• The paper is generally well-written and easy to follow. \n\n• Transformer-M can encode 2D or 3D structural information as bias terms and add to the attention matrix. It also encodes atomic centrality and adds to the atom features. Based on this, the model can obtain molecular representations from different data modes.\n\n• During pretraining, the authors label the data with different data modes for joint training. This training strategy may improve the performance of Transformer-M on downstream tasks from the results.\n\nWeaknesses\n\n• 2D and 3D information encoding are from previous work. This work just simply combines them, and the model architecture is the same as Graphormer[1]. The novelty is not enough.\n\n• Supervised pretraining based on the prediction of homo-lumo gap may lead to negative transfer. For example, on QM9 in downstream experiments, Transformer-M performs poorly on most tasks other than homo, lumo, and gap. This may be contradictory to the description \"general-purpose neural network model\" claimed in this paper.\n\n• Lack of description of PDBBind data processing and splitting in downstream tasks.\n\n• Absence of some ablation experiments: ①(p2D, p3D, p2D&3D) = (1:0:0) / (0:1:0) / (0:0:1); ②Only using the 3D position denoising task while pretraining.\n\nOther questions：\n\n• Do the authors consider encoding 1D molecular data mode, e.g., SMILES,  simultaneously?\n\n• What do the authors think about the possibility of negative transfer on downstream tasks due to the supervised signal introduced during pretraining?\n\n• Whether there is data leakage during finetuning on PDBBind, because we know that the general, refined, and core sets have overlapping parts.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "6oCV4fS5zs", "reviewer": "Reviewer_do3S", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a simple modifications over the Graphormer 3D [1] , which previously encoded node, graph level information through additional biases added to the self-attention modules, namely node-encoding via in/out degree embeddings, path encoding via shortest path and edge embedding aggregation along this path. To encode the spatial and centrality information for 3D molecules they embed these using a GBF assuming a fully connected graph as opposed to using cutoffs. \n\nNow, the authors simply introduce a switching mechanism in this model, where the corresponding 2D/3D encodings are switched on and off.  This is done as it may be the case that 2D information is available and 3D information is not on a particular input sample, so to allow the model to distill \"3D\" level information to enrich the 2D representation, for example it is the case for the PCQM4Mv2 dataset which includes a highly optimized conformer for the train-set but not the rest.\n\n\n\n[1] Shi, Yu, et al. \"Benchmarking graphormer on large-scale molecular modeling datasets.\" arXiv preprint arXiv:2203.04810 (2022).", "summary_of_the_review": "To summarize it's a simple idea that works, and the empirical results are good.", "strength_and_weaknesses": "The modifications are only a few lines of code over the original Graphormer 3D proposed in [1]. Still, this simple modification results in a new STOTA on PCQM4Mv2, namely introducing a mask in the input based on whether the network should operate in 2D/3D/Both.\n\nThe experiments on the task probabilities, which is the primary modification, are limited; changing this ratio can drastically change the model's performance. Most of the gains are attributed to the pretraining task and 3D encoding in the bias term, which was previously leveraged in [1].\n\nI would have liked to see more experiments on trying to further leverage 3D information on 2D inputs that don't have the position information, for example, adding conformers as well. Would the network operated in 3D/2D mode perform the best over simply 2D? \nWhat happens if you generate multiple conformers and switch these modes on and off? Would there be a gain in performance? The authors could have done more experiments on this front.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster