Uni-Mol: A Universal 3D Molecular Representation Learning Framework
Authors: Gengmo Zhou
Zhifeng Gao
Qiankun Ding
Hang Zheng
Hongteng Xu
Zhewei Wei
Guolin Ke
Linfeng Zhang
Venue: NeurIPS 2022 Submitted (2022)
Keywords: Representation Learning
Large-Scale 3D Molecular Pretraining
Molecular Property
Protein-Ligand Complex

Abstract: Molecular representation learning (MRL) has gained tremendous attention due to its critical role in learning from limited supervised data for applications like drug design. In most MRL methods, molecules are treated as 1D sequential tokens or 2D topology graphs, limiting their ability to incorporate 3D information for downstream tasks and, in particular, making it almost impossible for 3D geometry prediction or generation. Herein, we propose Uni-Mol, a universal MRL framework that significantly enlarges the representation ability and application scope of MRL schemes. Uni-Mol is composed of two models with the same SE(3)-equivariant transformer architecture: a molecular pretraining model trained by 209M molecular conformations; a pocket pretraining model trained by 3M candidate protein pocket data. The two models are used independently for separate tasks, and are combined when used in protein-ligand binding tasks. By properly incorporating 3D information, Uni-Mol outperforms SOTA in 14/15 molecular property prediction tasks. Moreover, Uni-Mol achieves superior performance in 3D spatial tasks, including protein-ligand binding pose prediction, molecular conformation generation, etc. Finally, we show that Uni-Mol can be successfully applied to the tasks with few-shot data like pocket druggability prediction. 

Uni-Mol: A Universal 3D Molecular Representation Learning Framework
Anonymous Author(s) Affiliation Address
email
Abstract
Molecular representation learning (MRL) has gained tremendous attention due its critical role in learning from limited supervised data for applications like drug design. In most MRL methods, molecules are treated as 1D sequential tokens 2D topology graphs, limiting their ability to incorporate 3D information for downstream tasks and, in particular, making it almost impossible for 3D geometry prediction or generation. Herein, we propose Uni-Mol, a universal MRL framework that significantly enlarges the representation ability and application scope of MRL schemes. Uni-Mol composed of two models with the same SE(()-equivariant transformer architecture: a molecular pretraining model trained by 209M molecular conformations; a pocket pretraining model trained by 3M candidate protein pocket data. The two models are used independently for separate tasks, and are combined when used in protein-ligand binding tasks. By properly incorporating 3D infor- mation, Uni-Mol outperforms SOTA in 14/15 molecular property prediction tasks. Moreover, Uni-Mol achieves superior performance in 3D spatial tasks, including protein-ligand binding pose prediction, molecular conformation generation, etc. Finally, we show that Uni-Mol can be successfully applied to the tasks with few-shot data like pocket druggability prediction. The model and data will be made publicly available htgp:::///tttbb
3
5
10
11
12
13
14
15
16
18
corp/Uni-Mor
Introduction
19
Recently, representation learning (or pretraining, self-supervised learning) I has been prevailing 21 in many applications, such as BERT 4 and GPT [5] in Natural Language Processing (NLP), VIT [8] in Computer Vision (CV), etc. These applications have common characteristic: unlabeled 29 data is abundant, while labeled data is limited. As solution, in a typical representation learning method, first adopts pretraining procedure to learn good representation from large-scale unlabeled data, and then a finetuning scheme is followed to extract more information from limited supervised data.
25
26
27
Applications in the field drug design share the characteristic that calls for representation learning schemes. The chemical space that drug candidate lies in is vast, while drug-related labeled data is limited. Not surprisingly, compared with traditional molecular fingerprint based models [9] 10]. recent molecular representation learning (MRL) models perform much better in most property prediction 12 13). However, to further improve the performance and extend the application scope of existing MRL models, one is faced with critical issue. From the perspective of life science, the properties of molecules and the effects drugs are mostly determined by their 3D structures [14] 15 In most current MRL methods, one starts with representing molecules ID sequential strings,
29
31
tasks
39
34
35
such as SMILES [16] 17 18 and InChl [19 20 21). or 2D graphs [22 I1 23 12 24). This may limit their ability incorporate 3D information for downstream tasks. particular, this makes
36
37
almost impossible for 3D geometry prediction or generation, such as, e.g., the prediction of protein-
Submitted to 36th Conference on Neural Information Processing Systems (NeurIPS 2022). Do not distribute.

---

atom level
Milecclorr Property
209M 3D molecules
pretrain
level
Molenular Conformation Gonieraiin
Conformations
Pretraining
Representation
Downstream tasks
3M candidate pockets
Pocket pretrain
level
Proporty
Figure 1: Schematic illustration of the Uni-Mol framework. Uni-Mol is composed of two models: molecular pretraining model trained by 209M molecular 3D conformations; pocket pretraining model trained by 3M candidate protein pocket data. The two models are used independently for separate tasks, and are combined when used in protein-ligand binding tasks.
38 ligand binding pose 25 Even though there have been some recent attempts trying to leverage 3D
the performance is less than optimal, possibly due to the small size of 40 3D datasets, and 3D positions can not be used as inputs/outputs during finetuning, since they only serve auxiliary information.
information in MRL 26 27
42 In this work, we propose Uni-Mol, to our best knowledge, the first universal molecular pretraining framework, which is derived from large-scale unlabeled data and is able to directly take 3D positions 44 as both inputs and outputs. Uni-Mol consists of 3 parts. 1) Backbone. Based on Transformer, the invariant spatial positional encoding and pair level representation are added to better capture the 3D information. Moreover, an equivariant head is used to directly predict 3D positions. 2) Pretraining. 46
We create two large-scale datasets, 209M molecular conformation dataset and 3M candidate protein pocket dataset, for pretraining 2 models on molecules and protein pockets, respectively. For the pretraining tasks, besides masked atom prediction, 3D position denoising task is used for learning 3D spatial representation. 3) Finetuning. According to specific downstream tasks, the used pretraining models are different. For example, in molecular property prediction tasks, only the molecular pretraining model is used; in protein-ligand binding pose prediction, both two pretraining models used. We refer to Fig. for an overall schematic illustration of the Uni-Mol framework.
48
50
52
53
54 To demonstrate the effectiveness of Uni-Mol, we conduct experiments on a series of downstream tasks. In the molecular property prediction tasks, Uni-Mol outperforms SOTA 14/15 datasets the MoleculeNet benchmark. In 3D geometric tasks, Uni-Mol also achieves superior performance. For the pose prediction of protein- ligand complexes, Uni-Mol predicts 88.07% binding poses with
56
RMSD <= 2A, 22.8 1% more than popular docking methods, and ranks 1st in the docking power test 59 on CASF-2016 [28) benchmark. Regarding molecular conformation generation, Uni-Mol achieves
SOTA for both Coverage and Matching metrics GEOM-QM? and GEOM-Drug (29). Moreover, Uni-Mol can be successfully applied to tasks with very limited data like pocket druggability prediction.
61
2
Uni-Mol Framework
In this section, we introduce the Uni-Mol framework by showing the details of the backbone, the 85 pretraining scheme, and the finetuning scheme. We refer Fig. for a schematic illustration of the model architecture

---

+
Pretraining tasks
Masked Abom Prodiction
FFN
3D Position Danoting
Layers
Uni-Mol Block
MatMul
Atom
Coord
Pain-dis
Softmee
Uni-Mol
Pai-Type
Scale
Represontation
oxrd petene
Unear
Buciiiaan Distance
Atoms
Coordinates
Atam Reprisentation
Par Represertation
modocalc
Figure 2: Left: the overall pretraining architecture. Middle: the model inputs, including atoms and spatial positional encoding created by pair Euclidean distance. Right: pair representation and its update process.
67
2.1 Backbone
Transformer (30] is widely used as backbone model in representation learning. However, Trans- former was originally designed for NLP tasks and cannot handle 3D spatial data directly. To tackle this, based the standard Transformer with Pre-LayerNore [3] backbone, introduce several modifications.
89
71
72
Invariant spatial positional encoding Due to its permutationally invariant property. Transformer cannot distinguish the positions of inputs without positional encoding. Different with the discrete (ordinal) positions used NLP/CV [32 33). the positions 3D space, i.e. coordinates, are continuous values. Besides, the positional encoding procedure needs to be invariant under global rotation and translation. To achieve that, similar to the relative positional encoding, we simply use Fuclidean distances of all atom pairs, as well as pair-type aware Gaussian kernels (34) Formally, the D-channel positional encoding of atom pair ij denoted as
73
74
75
76
77
78
Psy (G(A(dis tejs a, b), uh, " DI, A(d,r;a, = a.d +b,
2) is Gaussian density function with parameters js and a, dij the 80 Euclidean distance of atom pair ij. and fıj is the pair-type of atom pair ij. Please note the pair-type
79
where ((d. 14, == ov2*
81 here is not the chemical bond, and it is determined by the atom types pair ij. A(day layi a, b) is the affine transformation with parameters a and b, it affines di corresponding to its pair-type lig- Except dj and Lij. all remaining parameters trainable and randomly initialized.
82
83
84 Pair representation By default, Transformer maintains the token(atom) level representation, which later used in finetuning downstream tasks. Nevertheless, as the spatial positions are encoded at levr- we also maintain the pair-level representation, better learn the 3D spatial representation. Specifically, the pair representation is initialized as the aforementioned spatial positional encoding, the atom-to- to-pair communication via the multi-head Query- Key product results self-attention. Formally, the update of ij pair representation denoted as
85
87
Then, update pair representation,
89
1+1 - di +1 Q (K
U Pi;M,
1h [1, H],
Vd
90 where g/y is the pair representation of atom pair ij 1-th layer, is the number attention heads, 91 is the dimension of hidden representations, Q.Ptt (KC).") ) the Query (Key) of the i-th (j-th) atom 82 in the /-th layer h-th head, and M E DI the projection matrix to make the representation the 93 same shape multi-head Query-Key product results.
34 Besides, to leverage 3D information the atom representation, we also introduce the pair-to-atom communication, by using the pair representation as the bias term in self attentttnn Formally, the

---

96
self-attention with pair-to-atonn communication is denoted as
(KA
Attention ahahh K viah)) softmax(
+ Piş M
(3)
Vd
where v is the Value of the j-th atom in the 1-th layer h-th head. The pair representation and atom-pair communication are firstly proposed in the Evoformer in AlphaFold [35]. but the cost of
97
98
99
Evoformer is extremely large. In Uni-Mol, as we keep them as simple as possible, the extra cost of maintaining pair representation negligible.
100
coordinate head With 3D spatial positional encoding and pair representation, the model can learn a good 3D representation. However, it still lacks the ability to directly output co- ordinates, which is essential in 3D spatial tasks. To this end, we add a simple -qqiivaquie head
(B(33)Equivaring
102
103
104
to Uni-Mol. Following the idea of EGNN (36). the design of SE(3) -E(3)-equvarriane head is denoted as
Djjay n2.
Gi = ReLU UU)
(4)
where n is the number of total atoms, L is the number of layers in model, æ; R is the input coordinate of i-th atom, and d; E IR is the output coordinate of i-th atom, ReLU(y) = max (0, y)
105
106
Rectified Linear Unit (37]. U R and W " are the projection matrices convert
107
108
pair representation to scalar.
109
2.2
Pretraining
110
For the purpose of pretraining, we generate two large-scale datasets, one composed of 3D structures of organic molecules, and another composed of 3D structures of candidate protein pockets. Then, two models are pretrained using these two datasets, respectively. As pockets are directly involved in many drug design tasks, intuitively. the pretraining on candidate protein pockets can boost the performance of tasks related to protein-ligand structures and interactions.
111
112
113
114
115
The molecular pretraining dataset is based on multiple public datasets (See Appendix ?? for more information). After normalizing and deduplicating. it contains about 19M molecules. To generate 3D conformations, we use ETKGD (38) with Merck Molecular Force Field (39) optimization in RDKit [40] to randomly generate 10 conformations for each molecule. We also generate an additional 2D conformation (based on the molecular graph), to avoid some rare cases that fail to
117
118
119
120
generate 3D conformations.
121
The protein pocket pretraining dataset is derived from the Protein Data Bank (RCSB PDB [4] collection of 180K 3D structures of proteins. To extract candidate pockets, first clean the data by adding the missing side chains and hydrogen atoms; then we use Fpocket (42) to detect possible binding pockets of the proteins; and finally, we filter pockets by the number of residues in contact with and retains water molecules in the pocket. In this way, We collect a dataset composed of 3.2M candidate pockets for pretraining.
123
125
126
127
Self- supervuper task is vitally important for effective learning from large-scale unlabeled data. For example, the masked token prediction task in BERT (4) encourages the model to learn the contextual information. Similar to BERT, the masked atom prediction task is used in Uni-Mol. For each molccule/pocket, we add special atom [CLS], whose coordinate is the center of all atoms, to represent the whole molccule/pocket. However, as 3D spatial positional encoding leaks chemical bonds, atom types could be inferred easily, and therefore, the masked atom prediction cannot encourage the model to learn useful information. To tackle this, as well as learning from 3D information, we design a 3D position denoising task. Particularly, uniform noises of [-1 Á, Å are added to the random 15% atom coordinates, then the spatial positional encoding is calculated based on corrupted coordinates. In this way, the masked atom prediction task becomes non con-trivial. Besides, two additional heads are used to recover the correct spatial positions. 1) Pair-distance prediction. Based on pair-represnntation, the model needs to predict the correct Fuclidean distances of the atoms pairs with corrupted coordinates. 2) Coordinate prediction. Based SE(3)- Equivarana coordinate head, the model needs to predict the correct coordinates for the atoms with corrupted coordinates.
129
131
132
133
134
136
138
140
httt://www.resb.orgg

---

141
Both 2 pretraining models use the same self- sele--upe tasks described above, and Figure i the
142
illustration of the overall pretraining framework. For the detailed configurations of pretraining, please refer to Appendix ??.
143
2.3
Finetuning
145
To be consistent with pretraining, we use the same data prepossessing pipeline during finetuning. For molecules, as multiple random conformations can be generated a short time, we can use them as data augmentation in finetuning to improve performance and robustness. Some molecules may fail to generate 3D conformations, and we use their molecular graph as 2D conformation. For tasks that provide atom coordinates, we use them directly and skip the 3D conformation generation process. As there are 2 pretraining models and several types of downstream tasks, we should properly use them in the finetuning stage. According to the task types, and the involvement of protein ligand, we can categorize them as follow.
146
148
150
152
153
Non-3D prediction tasks These tasks do not need to output 3D conformations, Examples include molecular property prediction, molecule similarity, pocket druggability prediction, protein-ligand binding affinity prediction, etc. Similar to NLP/CV, we can simply use the representation of [CLS] which represents the whole molecule/pocket, or the mean representation of all atoms, with linear head to finetune on downstream tasks. In the tasks with mocket-molecule pair, we can concatenate their [CLS] representations, and then finetune with linear head.
154
155
156
157
158
159
3D prediction tasks of molecules or pockets These tasks need to predict a 3D conformation of the input, such as molecular conformation generation. Different with the fast conformation generation method used in Uni-Mol, molecular conformation generation task usually requires running advanced sampling and semi-empirical density functional theory (DFT) to account for the ensemble of 3D conformers that are accessible to molecule, and this very time-consuming. Therefore, there are many recent works that train the model to fast generate conformations from molecular graph [43] 44 45 46). While in Uni-Mol, this task straightforwardly becomes a conformation optimization task: generate new conformation based on different input conformation. Specifically, in finetuning, the model supervised learns the mapping from Uni-Mol generated conformations to the labeled conformations. Moreover, the optimized conformations can be generated end-to-end by SE(3) Eqqivvriine coordinate head.
160
161
165
165
167
168
169
170
3D prediction tasks of protein-ligand pairs This is one of the most important tasks in structure- based drug design. The task is to predict the complex structure of a protein binding site and molecular ligand. Besides the conformation changes of the pocket and the molecule themselves, we also need to consider how the molecule lays in the pocket, that is, the additional 6 degrees (3 rotations and 3 translations) of freedom of rigid movement. In principle, with Uni-Mol, we can predict the complex conformation by the Squiv-rqqi coordinate head in an end-to-end fashion. However, this is unstable as it is very sensitive to the initial docking positions of molecular ligand. Herein, to get rid of the initial positions, we use scoring function based optimization method this paper. In particular, the molecular representation and pocket representation are firstly obtained from their own pretraining models by their own conformations; then, their representations are concatenated as the
171
172
173
174
175
178
177
178
179
180
input of an additional layer Uni-Mol decoder, which is finetuned to learn the pair distances of all atoms in molecule and pocket. With the predicted pair-distance matrix as the scoring function, we
181
182
use a simple differential evolution algorithm (47 to sample and optimize the complex conformations.
183
More details can be found in Appendix ??
3
Experiments
184
185
To verify the effectiveness of our proposed Uni-Mol model, we conduct extensive experiments on multiple downstream tasks, including molecular property prediction, molecular conformation generation, pocket property prediction, and protein-ligand binding pose prediction. Besides, we also conduct several ablation studies. Due to space restrictions, we leave the detailed experimental settings and ablation studies to Appendix ??.
186
187
188
190
3.1
Molecular property prediction
191
Datasets and setup MoleculeNet (48) is a widely used benchmark for molecular property prediction, including datasets focusing different levels of properties of molecules, from quantum
192
5

---

Table 1: Uni-Mol performance on molecular property prediction classification tasks
Classification (ROC-AUC %; higher is better f)
Dutasets
BBBP 2039
BACE
ClinTox
Tox21
ToxCast
SIDER
HIV 41127
PCBA
MUV
# Molecules
1513
1478
7831
8575
1427
437929
93087
Tasks
1
1
2
12
617
27
1
128
17
D-MPNN
71.0(0.3)
80.9(0.6)
90.6(0.6)
75.9(0.7)
65.5(0.3)
57.0(0.7)
77.1(0.5)
86.2(0.1)
78.6(1.4)
Attentive FP
64.3(1.8)
78.4(0.022)
84.7(0.3)
76.1(0
63.7(0.2)
60.6(3.2)
75.7(1.4)
80.1(1.4)
76.6(1.5)
N-Gramar
69.7(0.6)
77.9(1.5)
77.5(4.0)
74.3(0.4)
66.8(0.7)
77.2(0.1)
76.9(0.7)
N- Grammyxc
69.1(0.8)
79.1(1.3)
87.5(2.7)
75.8(0.9)
65.5(0.7)
78.7(0.4)
74.8(0.2)
PretrainGNNN
68.7(1.3)
84.5(0.7)
72.6(1.5)
78.1
65.7(0.6)
62.7(0.8)
79.9(0.7) 86.0(0.1)
81.3(.1)
GROVER
10.0(0.1)
82.6(0.7)
81.2(3.0)
74.3(0.1)
65.4(0.4)
64.8(0.6)
62.5(0.9) 76.5(2.1)
67.3(1.8)
GROVER
69.5(0.1)
81.0(1.4)
76.2(3.7)
73.5(0.1)
65.3(0.5)
65.4(0.1)
68.2(1.1)
83.0(0.4)
67.3(1.8)
Gr
72.4(1.6) 81.2(0.9)
79.1(2.8
75.9(0.5)
63.1(0.4)
63.9(1.2)
77.0(1.2)
77.7(0.6)
MolCLR
72.2(2.1)
82.4(0.9)
91.2(3.5)
75.0(0.2)
58.9(1.4)
78.1(0.5)
79.6(1.9)
GFM
72.4(0.4)
85.6(1.1)
90.1(1.3)
10.1( 69.2(0.4)
67.2(0.4)
80.6(0.9)
86.6(0.1)
81.7(0.5)
Uni-Mol
72.9(0.6)
85.7(0.2)
91.9(1.8
79.6(0.5) 69.6(0.1) 65.9(1.3)
80.8(0.3)
88.5(0.1) 82.1(1.3)
Table 2: Uni-Mol performance on molecular property prediction regression tasks
Regression (lower is hetter
RMSE
MAE
Datasets
ESOL 1128
FreeSolv 642
Lipo 4200
QM7 6830
QM8
QM9
# Molecules
21786
133885 3
# Tasks
12
D-MPNN
1.050(0.008)
2.082(0.082)
0.683(0.016)
103.5(8.6)
0.0190(0.000)
0.008 (400 1(0.000)))
Attentive FP
0.877(0.029)
2.073(0.183
0.721(0.001)
72.0(27)
0.0179(0.001)
0.00812 1.00812(0.00) (200.00001)
N-GrameF
1.074(0.107
2.688(0.085)
0.812(0.028)
92.8(4.0)
0.0236(0.0006
40.037(0.00016
N-Gramxea
1083((.082))
5.061(0.744)
2.072(0.030)
81.9(19))
0.0215(0.0005)
100964(0000033)
PretrainGNN
L.100(0.006)
2.764(0.002)
0.739(0.003)
113.2(0.6)
(0000)))0001)
100922(0.00000
GROVERase
0.983(0.090) 2.176(0.052)
0.817(0.008)
94.5(3.8)
0.0218(0.0004)
0.00984(0.00055))
GROVER
0.895(0.017)
2.272(0.051)
0.823(00100)
92.0(0.9)
0.0024
0. 0.00986(0.00025))
GraphMVP
1.029(0.033)
0.681(0.010)
MolCLR
(0.04().040)
2.594(0.249)
(.691().004)
66.8(2.3)
0.0178(0.0003)
GEM
0.798(0.029)
1887700099)
0. 0.660000088
58.9(0.8)
0.0171(0.0001)
0.00746(0.00001)
Uni-Mol
0.788(0.029)
1.620(0.035)
0.603(0.010)
41.8(0.2)
0.0156(0.0001)
0.0046700.0004))
193
mechanics and physical chemistry to biophysics and physiology. Following previous work GEM 13]
194
we use scaffold splitting for the dataset and report the mean and standard deviation of the results for three random seeds.
195
We compare Uni-Mol with multiple basclines, including supervised and pretraining basclines. D-MPNN [49] and AttentiveFP [50) are supervised GNNs methods. N-gram (5I].
Baselines
197
198
PretrainGNN 22) GROVER 11), GraphMVP [26). MolCLR [12), and GEM [13] pretraining
199
methods. N-gram embeds the nodes in the graph and assembles them in short walks as the graph
200
representation. Random Forest and XGBoost (52) used the predictor for downstream tasks.
201
Table and Table show the experiment results of Uni-Mol and competitive baselines, where the best results are marked in bold. Most baseline results are from the paper of GEM, except for the recent works GraphMVP and MoICLR. The results of GraphMVP are from its paper. As MolCLR uses different data split setting (without considering chirality). we rerun it with the same data split setting as other basclines. From the results, we can summarize them as follows: 1) overall, Uni-Mol outperforms basclines on almost all downstream datasets. 2) In solubility (ESOL, Lipo), free energy (FreeSolv), and quantum mechanical (QM7, QM8, QM9) properties prediction tasks, Uni-Mol is significantly better than baselines. As 3D information is critical in these properties, it indicates that
Results
203
205
207
209
Uni-Mol can learn a better 3D representation than other baselines. 3) Uni-Mol fails to beat SOTA on the SIDER dataset. After investigation, we find Uni-Mol fails to generate 3D conformations (and rollbacks to 2D graphs) for many molecules (like natural products and peptides) in SIDER. Therefore,
211
212
due to the missing 3D information, it is reasonable that Uni-Mol cannot outperform others.
213 In summary, by better utilizing 3D information in pretraining, Uni-Mol outperforms all previous
MRL models in almost all property prediction tasks.

---

traditional confor- Table Uni-Mol performance on molecular conformation generation
Dataset
QM9
Drugs
Cov(t. %)
MA A)
COV(1, %)
MAT( A)
Methods
Mean
Median
Mean
Median
Mean
Median
Mean
Median
RDKit
83,26
90.78
0.3447
0.2935
60.91
65.70
1.2026
1.1252
CVGAE
0.09
0.00
1.6713
1.6088
0.00
0.00
3.0702
2.9937
GraphDG
73.33
84.21
0.4245
0.3973
8.27
0.00
1.9722
1.9845
CGCF
78.05
82.48
0.4219
0.3900
53.96
57.06
1,2487
1.2247
ConfVAR
80.42
85.31
0.4066
0.3891
53.14
53.98
1.2392
1.2447
ConfGF
88.49
94.13
0.26773
0.2685
62.15
70.93
1.1629
1.1596
GeoMol
71.26
72.00
0.3731
0.3731
67.16
71.71
1.0875
1.0586
DGSM
91.49
95.92
0,2139
0.2137
78.73
94.39
1.0154
0.9980
DMCG
96.34
99.53
0,2065
0,2003
96.69
100,00
0.7223
0.7236
GcoDiff
90.07
93.39
0,2090
0.1988
89.13
97.88
0.8629
0.8529
Uni-Mol
98.68
100.00
0.1806
0.1510
92.69
100.00
0.6596
0.6215
215
3.2
Molecular conformation generation
216
Datasets and setup Following the settings in previous works [44 533. we use GEOM-QM9 and GEOM-Drugs (54) dataset to perform conformation generation experiments. As described in Sec. 2.3 in this task, Uni-Mol optimizes its generative conformations to the labeled ones. To construct the finetuning data, we first randomly generate 10 conformations. Then, for each of them, we calculate the RMSD between it and labeled conformations, and choose the one with minimal RMSD as its optimizing target. For the inference in the test set, we generate the same number of conformations
217
218
219
220
221
222
(twice the number of labeled conformations) as previous works do. And we use the same metrics, Coverage (COV) and Matching (MAT). Higher COV means better diversity, while lower MAT means
223
224
higher accuracy.
225
Baselines We compare Uni-Mol with 10 competitive baselines. RDKit 38
226
mation generation method based on distance geometry. The rest baseline can be categorized into two classes. GraphDG [43). CGCF(44), ConfVAE [55], ConfGF [53). and DGSM (56) combine gener- ative models with distance geometry, which first generates interatomic distance matrices and then iteratively generates atomic coordinates. CVGAE (45), GeoMol (46). DMCG (57] and GeoDiff (58)
228
229
230
directly generate atomic coordinates.
231
The results are shown in Table We report the mean and median of COV and MAT on GEOM-QM9 and GEOM-Drugs datasets. ConfVAE (55). GeoMol 44). DGSM (56). DMCG (57). GeoDiff's s (58) results are from their papers, respectively. Other baseline results are from ConfGi's paper. As shown in Table 3 Uni-Mol exceeds existing baselines in both COV and MAT metrics both datasets. Although Uni-Mol outperforms SOTA, we suspect that the above benchmark cannot satisfy the real-world demand of conformation generation tasks in the field of drug design. Since the ensemble of molecular conformations in biological systems is different from that in vacuum or general solution environment, the ensemble of bioactive conformation must be considered in order to
Results
232
233
235
236
237
238
239
apply the conformation generation model in the context of drug design, while the GEOM dataset just ignores this. Establishing a reasonable benchmark will be crucial in this research direction.
240
241
3.3
Pocket property prediction
242
Datasets and setup Druggability, the ability of a candidate protein pocket to produce stable binding to specific molecular ligand, is one of the most critical properties of a candidate protein pocket. However, this task is very challenging due to the very limited supervised data. For example, NRDLD [59]. commonly used dataset, only contains 113 data samples. Therefore, besides NRDLD, we construct a regression dataset for benchmarking pocket property prediction performance. Specifically, based on Fpocket tool, calculate Fpocket Score, Druggability Score. Total SASA, and Hydrophobicity Score for the selected 164,586 candidate pockets. Model is trained to predict these scores. To avoid leaking, the selected pockets are not overlapped with the candidate protein pocket dataset used Uni-Mol pretraining.
244
246
248
250
On the NRDLD dataset, we compare Uni-Mol with 6 previous methods evaluated in 60) Accuracy recall, precision, and F1-score are used as metries for this classification task. On our created benchmark dataset, there are no appropriate basclines, we use an additional Uni-Mol model
Baselines
252
7

---

Table 4: Uni-Mol performance on pocket property prediction
Classification (higher better f)
Regression (lower hetter
Dataset
NRDLD
Fpacket Scores
Methods
Cavity- DrugScore Vulsite DrugPted PockDrug
TRAPP-CNN Uni-Mol
Methods
Uni-Molasson
Uni-Mol
Accuracy
0.82
0.89
0.89
0.865
0.946
0.946
MSFipecter
(0.22(0
(..50()
Recall
0.957
0.913
1,000
MSErugahily
0.60100.02)
0.49910.007)
Precision
0.846
1.000
0,920
MSE BABA
0.197(0.008) 0.129(0.005)
F1-score
0.898
0.955
0.958
0.035700017)) 0.0127(0.0005)
without pretraining, denoted Uni-Mol to check the performance brought by pretraining on pocket property prediction. MSE (mean square error) is used as the metric.
255
Results As shown in Table 4 Uni-Mol shows the best accuracy, recall, and Fl-score on NRDLD, the few-show dataset. In our created benchmark dataset, the pretraining Uni-Mol model largely outperforms the penpprrraiinnng one on all four scores. This indicates that pretraining on candidate protein pockets indeed brings improvement in pocket property prediction tasks.
257
259
280
Unlike Molecular property prediction, due to the very limited supervised data, pocket property prediction gained much less attention. Therefore, we also plan to release our created benchmark dataset, and hopefully, it can help future research.
261
262
263
3.4
Protein-ligand binding pose prediction
Datasets and setup As mentioned above, protein-ligand binding pose prediction is one of the most important tasks drug design. And Uni-Mol combines both the molecular and pocket pretraining models to learn distance matrix based scoring function, and then sample and optimize the complex conformations. For the benchmark dataset, referring to the previous works [8 61 we use CASF.
265
267
2016 as the test set. For the training data used in finetuning, (19,443 443 protein-ligand complexes), excluding complexes that already exist in the CASF-2016.
PDBbind General set v:2020 [62]
289
270
Two benchmarks are conducted: 1) Docking power, the default metric to benchmark the ability of scoring function CASF-2016. Specifically, it tests whether a scoring function can distinguish the ground truth binding pose from set of decoys or not. For each ground truth, CASF-2016 provides
271
272
273
50 100 decoy conformations of the same ligand. Scoring functions are applied to rank them, and the use the
ground truth binding pose is expected to be the top 1. 2) Binding pose accuracy. Specifically, the semi-flexible docking setting: keep the pocket conformation fixed, while the conformation ligand is fully flexible. We evaluate the RMSD between the predicted binding pose and the ground truth. Following previous works, we use the percentage of results that are below predefined RMSD thresholds metries.
275
277
279
For docking power benchmark, the baselines are DeepDock [61] and the top 10 scoring functions reported in (28), including both conventional scoring functions and machine learning- based ones. For the binding pose accuracy, the baselines are Autodock Vina 663 64), Vinardo (65).
Baselines
280
281
Smina (66). and AutoDock4 (67].
Results From the docking power benchmark results shown in Figure 3 Uni-Mol ranks the 1st,
284
with the top success rate of 91.6%. For comparison, the previous top scoring function AutoDock Vina (63 64) achieves 90.2% of the top success rate in this benchmark. From the binding pose accuracy results shown in Table 5] Uni-Mol also surpasses all other baselines. Notably, Uni-Mol
285
286
outperforms the second best method by 22.81 % under the threshold of 2Á. This result indicates that Uni-Mol can effectively learn the 3D information from both molecules and pockets, as well as the interaction in 3D space of them. Even without pretraining, Uni-Mol (demoted as Uni-Mol is
287
288
289
290
also better than other baselines. This demonstrates the effectiveness of Uni-Mol backbone, as effectively learns the 3D information by limited data.
281
292
In summary, by combining molecular and pocket pretraining models, Uni-Mol significantly outper- forms the widely used docking tools in the prggedn-ligand binding tasks.
293
4 Related work
295
Molecular representation learning Representation learning on large-scale unlabeled molecules
attracts much attention recently. SMILES- BERT [18 is pretrained on SMILES strings of molecules
297
using BERT [.. Subsequent works are mostly pretraining on 2D molecular topological graphs [23]
II
MolCLR
12
applies data augmentation to molecular graphs at both node and graph levels, using
8

---

Uni-Mol
Ligand RMSD
AuttxaackVinu
% Below Threshold 1
dVeiulF20
Giddeccor--S SP
Methods
0.5
1.0
1.5 A 2.0 Á 3.0 Á
5.0 A
DrugScoreCSD
DeesDock
Autodock Vina
23.86
44.21
57.54
64.56
73.68
84.56
GBVL WWS--GGGMMM
Vinardo
23.51
41.75
57.54
62.81
69.82
76.84
ThemPLPg6GOL
Smina
23,51
47.37
59.65
65.26
74.39
82.11
GlideBoore
Autodack4
7.02
21.75
31.58
35,44
47.02
64.56
00
Top Success
Uni-Molladoo
14.04
49.47
65.26
75.44
87.02
98.60
Figure 3: Docking power evaluation on CASF-2016 (Top 10 methods)
Uni-Mol
24.91
70.53
84.21
88.07
94.74
98.95
Table 5: Uni-Mol performance on binding pose prediction
299
self-supervised contrastive learning strategy learn molecular representations. Further, several
300
recent works try to leverage the 3D spatial information of molecules, and focus on contrastive or transfer learning between 2D topology and 3D geometry of molecules. For example, GraphMVP [26)
301
proposes a contrastive learning GNN-based framework between 2D topology and 3D geometry.
302
303
GEM [13) uses bond angles and bond length as additional edge attributes to enhance 3D information. As aforementioned, due to the inability of handling 3D information, most previous representation learning models cannot be used in the important 3D prediction tasks.
304
305
308
Equivarann models In many-body scenarios such as potential energy surface fitting, SE-(3) cquivariance is usually required. A series of SE(3) models are proposed, such as SchNet (68). tensor
307
field networks [69]. SE(3) Transformer (70., DimmNet (71). equivariant graph neural networks
308
309
(EGNN) [30). GemNet [72] and SphereNet (73). Most of these models are used in supervised
310
learning with energy and force. In Uni-Mol, based on the standard Transformer, we introduce several minor changes to make the model (5(3))Equivariant
312
Pocket druggability prediction Druggability prediction of protein binding pockets is crucial for drug discovery as druggable pockets need to be identified at the beginning. Since proteins undergo conformation changes that might alter the druggability of pockets, it is necessary to utilize 3D spatial data beyond sequential information. Early methods, such as Volsite (74), DrugPred (59), and PockDrug (75). predict druggability based on the predefined descriptors of pockets' static structures.
313
314
315
316
317
Later, TRAPP-CNN 660) based on 3D-CNN, proposes the analysis of proteins' conformation changes
and the use of such information for druggability prediction.
318
319
Protein-ligand binding pose prediction In structure-based drug design, it is crucial to understand the interactions between protein targets and ligands. The in vitro estimation of the binding pose and affinity, such as docking, allows for lead identification and guides molecular optimization. In particular, docking is one of the most important approaches in structure-baser drug design and has
320
321
322
323
been developed for the past decades. Tools such as AutoDock4 (67), AutoDock Vina [63] (64), and Smina (66) are among the most used docking programs. Also, machine learning- docking methods, such as
324
(76), DeepDock (61] and Equibind (77), have also been developed to predict protein-ligand binding poses and assess protein-ligand binding affinity.
326
5 Conclusion
327
328
In this paper, to enlarge the application scope and representation ability of molecular representation learning (MRL), we propose Uni-Mol, the first universal large-scale 3D MRL framework. Uni-Mol consists of 3 parts: a Transformer based backbone to handle 3D data; two large-scale pretraining models to learn molecular and pocket representations respectively; finetuning strategies for all kinds of downstream tasks. Experiments demonstrate that Uni-Mol can outperform existing SOTA in various downstream tasks, especially in 3D spatial tasks.
329
331
332
333
334
There are 3 potential future directions. 1) Better interaction mechanisms for finetuning two pretraining models together. As the interaction between the pretraining pocket model and the pretraining molecular model is simple in the current version of Uni-Mol, we believe there is a large room for further improvement. 2) Large Uni-Mol models. As larger pretraining models often perform better,
335
336
337
338
is worthy of training large Uni-Mol model on a bigger dataset. 3) More -ughltyultty benchmarks.
339
Although there have been many applications the field of drug design, high-quality public datasets
have been lacking. Many public datasets cannot satisfy real-world demand due to the low data quality.
We believe the high-quality benchmarks will be the lighthouse of the entire field, and will significantly accelerate the development of drug design.
341
9

---

References
343
[1]
Yoshua Bengio, Aaron Courville. and Pascal Vincent. "Representation learning: A review and new perspectives". In: IEEE transactions on pattern andlysis and machine intelligence 35.8 (2013), pp. 1798- 1828.
345
346
347
[2]
William L. Hamilton. Rex Ying, and Jure Leskovee. "Representation Learning on Graphs: Methods and Applications".. In: IEEE Data Eng. Bull. 40.3 (2017), pp. 52-74. URL: lattp://sites , computer org/
348
349
debull/A17sept:// pdff
350
[3]
Daokun Zhang et al. "Network representation learning: A survey". In: IEEE transactions on Big Data
351
6.1 (2018), pp. 3-28.
352
[4]
Jacob Devlin ct al. "BERT: Pre-training of Deep Bidircctional Transformers for Language Under- standing". In: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Longuage Technologies. Volume I (Long and Short Papers).
353
355
Minneapolis, Minnesota: Association for Computational Linguistics, June 2019, pp. 4171-4186, DOE:
356
10. N0.18633/v1/H199 1423 URL: https: htttps//aclanthhen http://ccaaatolooyyoooooo org//19 Baarg/199-4444
357
[5]
Alec Radford et al. "Improving language understanding by generative pre- pre-training". In: (2018).
358
[6]
Alec Radford et al. "Language models are unsupervised multitask learners". In: OpenAl blog 1.8 (2019).
359
p. 9.
360
Tom Brown et al. "Language models are few-shot learners.. In: Advances in neural information process- ing systems 33 (2020), pp. 1877-1901.
381
362
[8]
Alexey Dosovitskiy et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale". In: International Conference on Learning Representations. 2021. URL:: https: //ppr//veeeeevee net net/forumiid-ee d--YebbdTty
363
364
385
19 Qingda Zang et al. "In silico prediction of physicochemical properties of environmental chemicals using molecular fingerprints and machine learning". In: Journal of chemical information and modeling 57.1 (2017), pp. 36-49.
[10]
Minjian Yang et al. "Machine learning models based on molecular fingerprints and an extreme gradient boosting method lead the discovery of JAK2 inhibitors". In: Journal of Chemical Information and Modeling 59.12 (2019), pp. 5002-50011
389
370
371 Yu Rong et al. Supervised Graph Transformer on Large-Scale Molecular Data". In: Advances in Neural Information Processing Systems 33 (2020).
372
373
112]
Yuyang Wang et al. "Molecular contrastive learning of representations via graph neural networks". In:
374
Nature Machine Intelligence (2022), pp. 1-9. Dor: 10. 1038/ /4225642266662 - 022 00447-x
375
13
Xiaomin Fang et al. "Geometry enhhancer--eehnneedd molecular representation learning for property prediction". In:
378
Nature Machine Intelligence (2022), pp. 1-8. DOI: 10. 10w11033//220 021 02000500-21-00008---10
377
[14]
A Crum-Brown and TR Fraser. "The connection of chemical constitution and physiological action". In: Trans R Soc Edinb 25. 18.1968-199 (1865), p. 257.
378
379
115]
Corwin Hansch and Toshio Fujita. "p-o-T Analysis. A Method for the Correlation of Biological Activity and Chemical Structure". In: Journal of the American Chemical Society 86.8 (1964), pp. 1616-1626.
380
381
[16]
David Weininger. "SMILES, a chemical language and information system. 1. Introduction to methodology and encoding rules". In: Journal of chemical information and computer sciences 28.1 (1988), pp. 31-36.
383
117]
Zheng Xu al. "Seq2seq fingerprint: An unsupervised deep molecular embedding for drug discovery". In: Proceedings of the 8th ACM international conference bioinformatics, computational biology, and health informatics. 2017. pp. 285-294.
385
386
[18]
Sheng Wang et al. "Smiles-bert: large scale unsupervised pre-training for molecular property prediction". In: Proceedings of the 10th ACM international conference on bioinformatics, computational biology and health informatics. 2019, pp. 429-436.
387
388
389
199
Stephen R Heller et al. "InChI, the IUPAC international chemical identifier". In: Journal of cheminfor- matics 7.1 (2015), pp. 1-34.
390
391
[20]
Robin Winter et al. "Learning continuous and data-driven molecular descriptors by translating equivalent chemical representations".. In: Chemical science 10.6 (2019), pp. 1692-1701.
392
393
[1
Jennifer Handsel et al. "Translating the InChl: adapting neural machine translation to predict IUPAC names from a chemical identifier". In: Journal of cheminformattic 13.1 (2021). pp. 1-11.
394
385
[22]
Weihua Hut et al. "Strategies for Pre-training Graph Neural Networks". In: International Conference on http:/ooormmmmddd1.....1.ama
396
Learning Representations. 2020. URI.: httpp::/opperrviieee
387
[23]
Pengyong Li et al. "An effective self-supervised framework for learning expressive molecular global representations drug discovery"". In: Briefings in Bioinformatics 22.6 (2021). bbab109.
389
[24]
Chengxuan Ying et al. "Do Transformers Really Perform Badly for Graph Representation?" In: Advances in Neural Information Processing Systems 34 (2021).
400
10

---

401
[25]
Panagiotis I Koukos, Li C Xue. and Alexandre MIJ Bonvin. "Protein-ligand pose and affinity prediction: Lessons from D3R Grand Challenge 3". In: Journal of computer-uided molecular dexign 33.1 (2019).
402
403
pp. 83-91.
[26]
Shengehao Liu et al. "Pre-training Molecular Graph Representation with 3D Geometry". In: International Conference on Learning Representations. 2022. URL: hntpss : //penreviis net / forummm ? id=
405
406
xQUe1pOKPam
407
[27]
Hannes Stärk et al. "3D Infomax improves GNNs for Molecular Property Prediction". In: arXiv preprint 0rXiv:2110044122 (2021).
408
409
[28]
Minyi Su et al. "Comparative assessment of scoring functions: the CASF-2016 update". In: Journal of chemical information and modeling 59.2 (2018), pp. 895-913.
410
411
[29]
Andrew L Hopkins, Colin R Groom, and Alexander Alex. "Ligand efficiency: a useful metric for lead selection." In: Drug discovery today 9.10 (2004), pp. 430-431.
412
415
[30]
Ashish Vaswani et al. " "Attention is all you need". In: Advances in neural information processing systems
414
30 (2017).
415
[3]]
Ruibin Xiong et al. "On Layer Normalization in the Transformer Archttecuure.. In: Proceedings of the 37th International Conference on Machine Learning. Ed. by Hal Daumé III and Aarti Singh. Vol. 119, Proceedings of Machine Learning Research. PMLR, July 2020, pp. 10524-10533.
416
418
[32]
Guolin Ke, Di He, and Tie-Yan Liu. "Rethinking Positional Encoding in Language Pre-training". In: International Conference Learning Representations. 2020.
420
[33
Philipp Dufter, Martin Schmitt, and Hinrich Schüitze. "Position information in transformers: An overview""
In: arXiv preprint arXiv:2002.1 11090 (2021).
422
[34]
Muhammed Shuaibi et al. "Rotation invariant graph neural networks using spin convolutions". In: arXiv
423
preprint arXiv:2 106. 09575 (2021).
424
[35]
John Jumper et al. "Highly accurate protein structure prediction with AlphaFold". In: Nature 596.7873 (2021), pp. 583-589.
425
36
Victor Garcia Satoraa,, Emiel Hoogeboom, and Max Welling. "B (n) equivariant graph neural networks". In: International Conference on Machine Learning. PMLR. 2021, pp. 9323-9332.
427
428
337
Abien Fred Agarap. "Deep learning using rectified linear units (relu)". In: arXiv preprint
429
arXiv:1803.088 08035 (2018).
[38]
Sereina Riniker and Gregory A Landrum. "Better informed distance geometry: using what know to improve conformation generation". In: Journal of chemical information and modeling 55.12 (2015). pp. 2562-2574.
431
432
433
[39
Thomas Halgren. "Merek molecular force field. Basis, form, scope, parameterization, and perfor- mance of MMFF94", In: Journal of computational chemistry 17.5-6 (1996), pp. 490-519.
434
435
[40]
Greg Landrum et al. RDKit; software suite for cheminformatics, computaational chemistry, and predictive modeling. 2013.
436
437
[41]
Helen M Berman et al. "The protein data bank". In: Nucleic acids research 28.1 (2000), pp. 235-242. Vincent Le Guilloux, Peter Schmidtke, and Pierre Tuffery. "Fpocket: an open source platform for ligand pocket detection". " In: BMC bioinformotics 10.1 (2009), pp. 1-11.
438
[42]
439
410
[43]
Gregor Simm and Jose Miguel Hernandez-Lobar "A Generative Model for Molecular Distance Geome- try". In: International Conference on Machine Learning. PMLR. 2020, pp. 8949-8958.
441
442
[44]
Minkai Xu al. "Learning Neural Generative Dynamies for Molecular Conformation Generation", In: International Conference on Learning Representations. 2020.
443
444
[45]
Elman Mansimov et al. "Molecular geometry prediction using a deep generative graph neural network". In: Scientific reports 9.1 (2019), pp. 1-13.
445
448
[46]
Octavian Ganea et al. "Geomol: Torsionll geometric generation molecular 3d conformer ensembles".
447
In: Advances in Neural Information Processing Systems 34 (2021).
448
[47]
Rainer Storn and Kenneth Price. "Differential evolution-a simple and efficient heuristic for global optimization over continuous spaces". In: Journal of global optimization 11.4 (1997), pp. 341-359, Zhenqin Wu et "MoleculeNet: a benchmark for molecular machine learning". ". In: Chemical science
449
450
[48]
451
9.2 (2018), pp. 513-530.
452
[49]
Kevin Yang et al. "Analyziig learned molecular representations for property prediction". In: Journal of chemical information and modeling 59.8 (2019), pp. 3370-3388.
453
454
50
Zhaoping Xiong et al. "Pushing the boundarics of molecular representation for drug discovery with the graph attention mechanism". In: Journal of medicinal chemistry 63.16 (2019), pp. 8749-8760. 8760.
455
456
[1
Shengchao Liu, Mehmet F Demirel, and Yingyu Liang. "N-gram graph: Simple unsupervised representa- tion for graphs, with applications to molecules". In: Advances in neural information processing systems 32 (2019).
457
459
[52]
Tianqi Chen and Carlos Guestrin. "Xgboost: A scalable tree boosting system".. In: Proceedings of the 22nd acm sigkdd international conference knowledge discovery and data mining. 2016, pp. 784-794.
11

---

461
[53]
Chence Shi et al. "Learning gradient fields for molecular conformation generation". In: International Conference on Machine Learning. PMLR. 2021, pp. 9558-9568.
462
463
[54]
Simon Axelrod and Rafael Gomez-Bombaree property prediction and molecular generation". In: Scientific Data 9.1 (2022), pp- 1-14.
"GEOM, energy-amnotated molecular conformations for
464
[55]
Minkai Xu et al. " "An end-to ro-end framework for molecular conformation generation via bilevel program- ming". In: International Conference on Machine Learning. PMLR. 2021, pp. 11537-11547.
468
[56]
Shitong Luo al. "Predicting Molecular Conformation via Dynamic Graph Score Matching". In: Advances in Neural Information Processing Syxtems 34 (2021).
468
469
157]
Jinhua Zhu al. "Direct molecular conformation generation". In: arXiv preprint arXiv:2202.0/3356 (2022).
470
471
[58]
Minkai Xu al. "GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation", In: International Conference on Learning Representations. 2022.
472
473
[59]
Agata Krasowski et al. "DrugPred: structure-based approach to predict protein druggability developed using an extensive nonredundant data set". In: Journal of chemical information and modeling 51.11 (2011), pp. 2829-2842,
474
475
476
[60]
Jui-Hung Yuan et al. "Druggability assessment in TRAPP using machine learning approaches". In:
477
Journal of Chemical Information and Modeling 60.3 (2020), pp. 1685-1699.
478
[61]
Oscar Méndez-Lucio et al. "A geometric deep learning approach to predict binding conformations of bioactive molecules", In: Nature Machine Intelligence 3.12 (2021), pp. 1033-1035
480
622
Zhihai Liu et al. "PDB-wii -iid collection of binding data: current status of the PDBbind database". In: Bioinformatics 31.3 (2015), pp. 405-412.
482
663
Oleg Trott and Arthur J Olson. "AutoDoek Vina: improving the speed and accuracy of docking with
483
new scoring function, efficient optimization, and multithreading". In: Journal of computational chemistry 31.2 (2010), pp. 455-461.
484
485
644
Jerome Eberhardt et al. "AutoDock Vina 1.2. O: New docking methods, expanded force field, and python hindings". In: Journal of Chemical Information and Modeling 61.8 (2021), pp. 3891-3898.
486
487
665
Rodrigo Quiroga and Marcos A Villarreal. "Vinardo: A scoring function based on autodock vina improves scoring, docking, and virtual screening". In: PloS one 11.5 (2016), c0155183.
488
489
[66]
David Ryan Koes, Matthew P Baumgartner, and Carlos Camacho. "Lessons learned in empirical scoring with smina from the CSAR 2011 I benchmarking exercise". In: Journal of chemical information and modeling 53.8 (2013), pp. 1893-1904.
490
491
492
[67]
Garrett M Morris et al. "AutoDock4 and AutoDockTools4: Automated docking with selective receptor ficxibility". In: Journal of computational chemistry 30.16 (2009), pp. 2785-2791.
494
[68]
Kristof Schütt et al. "Schnet: A continuous- filter convolutional neural network for modeling quantum interactions". In: Advances in neural information processing systems 30 (2017).
496
[69]
Nathaniel Thomas et al. "Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds". In: arXiv preprint arXiv: 1802.08219 (2018).
770
Fabian Fuchs et al. "Sc transfooooor 3d roto- trontranslation equivariant attention networks". In: Advances in Neural Information Processing Systems 33 (2020), pp. 1970-1981.
71]
Johannes Gasteiger. Janek GroB, and Stephan Günnemann. "Directional Message Passing for Molecular Graphs". In: International Conference on Learning Representations (ICLR). 2020.
501
502
[72]
Johannes Kliepera, Florian Becker, and Stephan Günnemann. "GemNet: Universal Directional Graph Neural Networks for Molecules", In: Advances Neural Information Processing Systems. 2021.
503
[73
Yi Liu al. "Spherical Message Passing for 3D Molecular Graphs". In: International Conference on Learning Representations. 2022. URL.: httpps:/ /astpa::/openreiiiie
505
508
[74]
Jérémy Desaphy et al. Comparison and druggability prediction of protein ligand binding sites from pharmacophore-awwan annootate- cavity shapes. 2012.
507
508
[75]
Alexandre Borrel et al. "PockDrug: A model for predicting pocket druggability that overcomes pocket estimation uncertainties". In: Journal of chemical information and modeling 55.4 (2015). pp. 882-895. Cheng Wang and Yingkai Zhang. "Improving boriingdocking-seee screeninn powers of protein-ligand scoring functions using random forest". In: Journal of computational chemistry 38.3 (2017), pp. 169- 177.
509
510
[76]
511
512
513
[77]
Hannes Stärk et al. EquiBind: Geometric Deep Learning for Drug Binding Structure Prediction. 2022.
Checklist
515
1. For all authors...
516
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
12

---

518
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
520
(d) Have you read the ethics review guidelines and ensured that your paper conforms them?
521
[Yes]
622
2. If you are including theoretical results...
523
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
524
(b) Did you include complete proofs of all theoretical results? [N/A]
525
3. If you ran experiments...
528
(a) Did you include the code. data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as URL)? [Yes] The data we used are all from public databases and details in data processing are explained in Appendix. The data, code, and instructions will be made public upon the acceptance of the paper.
527
528
529
530
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] We report all the training details for the experiemnt in Appendix.
532
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] We report the mean and std for different runs of experiments in Table I
533
534
Table and Table 4
535
(d) Did you include the total amount of compute and the type of resources used (c.g.. type of GPUs, internal cluster, or cloud provider)? [Yes] We report the detailed computing resources used for the experiment in Appendix.
537
538
you are using existing assets (e.g., code, data, models) or curating/releaing new assets... (a) If your work uses existing assets, did you cite the creators? [Yes] We discuss all the used datasets in the experiment section 3 datasets and setup part.
539
(b) Did you mention the license of the assets? [Yes] We mention the license for the datasets used
542
in Appendix.
(c) Did you include any new assets either in the supplemental material or as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you 're using/curating? [N/A]
545
548
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
547
5. If you used crowdsourcing or conducted research with human subjects...
549
Did you include the full text of instructions given to participants and screenshots, if applica-
550
ble? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
552
(IRB) approvals, if applicable? [N/A]
553
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "pD2kvQRCZ8O", "reviewer": "Reviewer_Ms4D", "summary": "The authors proposed the same SE(3)-equivariant transformer variant, called Uni-Mol, for both molecular pretraining and protein pocket pretraining. To make the model SE(3)-equivariant, they devise an invariant spatial positional encoding for the transformer variant.  Specifically, a affine transformation is used to fuse euclidean distance of atom pairs and types of atom pairs, then a Gaussian density function is used to get continuous representations of the atom pairs. After that, the self-attention mechanism takes the pair representation as a bias term in softmax calculation to better leverage 3D information of atoms. Following the idea of EGNN, the authors also propose a SE(3)-equivariant coordinate head to to directly predict coordinates of atoms. According to different downstream tasks, the two models are used independently or jointly.", "questions": "1、For molecular conformation generation task, what is the input of the model for fine-tuning? That is, does the model need some kind of ‘random atomic coordinates’ as input for molecular conformation optimization? Moreover, how to randomly generate 10 conformations as the authors mention in Sec. 3.2? I guess ETKGD is utilized again to generate coordinates as inputs of Uni-Mol. If so, do other models need 3d coordinates as input, i.e., using ETKGD in a similar way? Do the authors check whether any molecules exist in both pretraining data set and test set?\n2、For pocket property prediction task, the authors construct a regression dataset based on Fpocket tool. What is the difference between Fpocket Score and Druggability Score? Moreover, it seems only Fpocket Score is conducted in Table 4. Besides, how well Fpocket tool can perform in reality? Can the authors test this tool on NRDLD ( if possible )?\n3、For 3D prediction tasks of protein-ligand pairs, the authors claim that they use a simple differential evolution algorithm to sample and optimize the complex conformations. More details should be revealed.\n4、The authors should investigate the impact of the pair Euclidean distance matrix in ablation studies.\n5、The authors should demonstrate the effectiveness of pretraining process. That is, how much improvement can pretraining bring to model performance?\n6、The authors mention that 3D spatial positional encoding leaks chemicl bonds in Sec. 2.2 which is contradictory to some content in Sec. 2.1 (third to last sentence in invariant spatial positional encoding part).\n7、In Table 2 in supplementary material, max number of atoms is 256 for protein pocket. Is this value big enough for a protein? Moreover, why vocabulary size is different between molecules and proteins (30 for molecule and 9 for pocket )?\n\n", "limitations": "1、 The proposed SE(3)-equivariant transformer architecture has a limitation on the length of input molecules. In another word, it can not handle with molecules or protein pockets with length bigger than the max atom numbers (in this paper, the number is 256).", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This manuscript describes a universal pretrained model for both molecular and protein pocket representation. Specifically, it shows an interesting way of integrating attention mechanism with pair Euclidean distance matrix and atom pair type matrix, keeping the transformer variant SE(3)-equivariant. The authors conduct many downstream tasks to evaluate the performance of the pretrained model and the results are mostly promising.  However, it is a little pity that some aspects of the approach are not clearly communicated. ", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "jZWLM_S3Kq1", "reviewer": "Reviewer_byGu", "summary": "This paper proposes a method to incorporate 3D information into molecular representation learning, Uni-Mol. Particularly, Uni-Mol includes three parts: \n\n1. Two models (a molecular model and a pocket model) built upon the SE(3)-equivariant transformer architecture; \n2. Two large-scale datasets (a 209M molecular conformation dataset and a 3M candidate protein pocket dataset) and corresponding pretraining strategies; \n3. Finetuning strategies for various downstream tasks. \n\nIn the experiments, Uni-Mol outperforms SOTA in most of the tested tasks and demonstrates its ability in few-shot learning settings. ", "questions": "* In line 114 and line 119, it is mentioned that the 3D conformations and pocket binding positions are obtained by using simulation or optimization toolkits. Will these toolkits introduce bias to the data? If yes, do you have any idea for alleviating such bias?\n* In line 130, it states that for the 3D position denoising task, the noise is uniformly sampled from [-1 A, 1 A]. How is the sampling strategy determined? Have you tried other sampling strategies like using a Gaussian distribution?\n* In line 204, it is concluded that 3D information helps the model to learn better representations. Are there any ablation study results that can justify this more convincingly? \n* Will the datasets be released for academic use?", "limitations": "The limitations are not highlighted or summarized in the paper. \n\nTo me, my most concern is about the data used to pretrain the model. As mentioned in the Questions section, since the 3D information in the data mostly comes from computational results instead of ground-truth annotations, how do we know such bias in data won't influence the model performance? If bias is inevitable, what benefits can be guaranteed by using such data?", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\n* (originality) The proposed method, Uni-Mol, is original and has reasonable designs: (1) For the backbone model architecture, modifications (including invariant spatial positional encoding, pair representation, and a coordinate head) have been made to the standard Transformer to adapt to 3D molecular representation; (2) Two new large-scale datasets as well as the proposed 3D position denoising task for pretraining; (3) Detailed guidances of how to fine-tune Uni-Mol on downstream tasks. \n* (quality) The experiments were conducted thoroughly and touched four aspects of molecular representation learning: molecular property prediction, molecular conformation generation, pocket property prediction, and protein-ligand binding pose prediction. For each task, the datasets and baselines are chosen carefully. \n* (significance) The proposed method Uni-Mol outperforms previous work in almost every tested task, which is impressive. In some tasks, like molecular conformation generation and pocket property prediction, the improvements are significant. \n* (clarity) The paper is presented clearly and easy to follow. Contributions are properly emphasized. \n\n**Weaknesses**:\n\n* (originality) The design of Uni-Mol is largely based on previous work. For the backbone model architecture, it is not new to use Transformers in molecular representation (e.g., GROVER), and most modifications can also be found in previous work (e.g., pairwise Euclidean distance matrix representation, the SE(3)-equivariant head as in EGNN, the masked atom prediction task). This limits the novelty of the proposed method. \n* (significance) Some experiment results didn't show significant improvement over previous methods (e.g., molecular property prediction tasks and ESOL). But this is acceptable since tasks like molecular property prediction have already been studied for long, and the baseline models are very competitive. \n* The related work focuses more on molecular representation methods, while the literature on pretraining has been ignored. ", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "BJHsVZ0g2ot", "reviewer": "Reviewer_hrSe", "summary": "This work provides representation learning for 3D molecules (like BERT in NLP). The authors provide two pretrained models using the same model architecture. The first one is trained on 209M molecular conformations, and the second one is trained on 3M protein pocket data. The model architecture is a modified Transformer, and the pretraining tasks include masked atom prediction and 3D position denoising. The pretrained models can be used for various tasks, such as molecular property prediction, molecular conformation generation, protein-ligand binding pose prediction, and pocket druggability prediction. The performance is very good.", "questions": "1. What is the difference between the proposed model architecture and existing Transformer architectures like Graphormer?\n2. I suggest the authors compare (or discuss) the newly-proposed graph DL-based methods on the protein-ligand binding pose prediction task.\n3. Can we extend protein pockets to whole proteins since in some cases, we don’t know the pocket? If not, what is the challenge?", "limitations": "The authors have addressed the limitations of their work.", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. This work is well-written and easy to follow.\n2. It provides two good pretrained models for small molecules and protein pockets. The models can be finetuned for various downstream tasks.\n3. The experiment results show that the provided models can outperform previous methods on various tasks.\n\nWeakness: \n1. About the model architecture: There are many existing Transformer architectures like [1] for graph representations. The authors should discuss the differences with these methods.\n[1] Ying, Chengxuan, et al. \"Do transformers really perform badly for graph representation?.\" \n2. About baseline results: For the protein-ligand binding pose prediction task, there are several graph DL-based methods like [2], I suggest the authors include such papers.\n[2] Stärk, Hannes, et al. \"Equibind: Geometric deep learning for drug binding structure prediction.\" \n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3djj7SBxF8H", "reviewer": "Reviewer_Ue6N", "summary": "The paper proposes a pretraining / self-supervised learning method for 3D molecular representation learning. The backbone is Transformer based. The pretraining is adding noise to atom coordinates. The pretraining  Experiments on molecular property prediction, molecular conformation generation, pocket property prediction and protein-ligand binding pose prediction show the effectiveness of the proposed method. \n", "questions": "1. The novelty of the proposed method seems not that significant. \n- One difference or advantage that the paper claims is 3D encoding. However, the Euclidean distance map (e.g.,[1]) or Gaussian functions for distance processing [2,3]  has been widely used in protein 3D structure learning. \n\n[1] Highly accurate protein structure prediction with AlphaFold\n\n[2] GraphQA: protein model quality assessment using graph convolutional networks \n\n[3] Learning from Protein Structure with Geometric Vector Perceptrons\n\n- There have been a few methods for the atom-level molecule or protein 3D structure modeling (e.g., [4]). The proposed method should be compared with them via experiments.\n\n[4] Intrinsic-Extrinsic Convolution and Pooling for Learning on 3D Protein Structures\n\n- The pretraining method seems very simple. It is more like a trick. \n- (Minor) There are a few existing pretraining / self-supervised learning methods for protein representation learning. The authors might want to add a discussion about them.\n\n[5] Structure-aware Protein Self-supervised Learning\n\n[6] Contrastive Representation Learning for 3D Protein Structures\n\n[7] Protein Structure Representation Learning by Geometric Pretraining\n\n\n2. The motivation to employ Transformer is not clear. The explanation that \"Transformer is the default backbone in representation learning\" is not convincing or even incorrect. Moreover, as a Transformer-related work, it could be better to visualize the self-attention map to verify the motivation for employing Transformers. \n\n\n3. Besides adding the pair representation, is there any other difference with the original Transformer?  Moreover, using Transformers for 3D modeling has been explored in 3D vision. Some important competitors should be compared, such as [8], to show the effectiveness of the proposed method. \n\n[8] Point Transformer.\n\n\n\n4. Some important experiments or ablation studies are missing. \n\na) Accuracy w/ and w/o pertaining. \n\nb) Effect of invariant spatial positional encoding.\n\nc) Effect of pair representation.\n\nd) Effect of pair representation.\n\ne) Effect of SE(3)-equivariance coordinate head.\n", "limitations": "NA", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1. Among the early efforts, the paper proposes a pretraining / self-supervised learning method for universal 3D molecular representation\nlearning framework. \n\n2. The proposed method achieves SOTAs on several tasks and datasets.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Reject