Uni-Mol: A Universal 3D Molecular Representation Learning Framework
Authors: Gengmo Zhou
Zhifeng Gao
Qiankun Ding
Hang Zheng
Hongteng Xu
Zhewei Wei
Linfeng Zhang
Guolin Ke
Venue: ICLR 2023 poster (2023)
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/generation. In this paper, we propose a universal 3D MRL framework, called Uni-Mol, that significantly enlarges the representation ability and application scope of MRL schemes. Uni-Mol contains two pretrained models with the same SE(3) Transformer architecture: a molecular model pretrained by 209M molecular conformations; a pocket model pretrained by 3M candidate protein pocket data. Besides, Uni-Mol contains several finetuning strategies to apply the pretrained models to various downstream 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. The code, model, and data are made publicly available at https://github.com/dptech-corp/Uni-Mol.

Published as a conference paper at ICLR 2023
UNI-MOL: A UNIVERSAL 3D MOLECULAR REPRESENTATION LEARNING FRAMEWORK
Gengmo Zhou Zhifeng G Qiankun Ding², Hang Zheng² Hongteng Xu', Zhewei Wei', Linfeng Zhang Guuolin K RRemmin University China "DP Technology *AI for Science Institute, Beijing (zgm2015, hongtengxu, zheweil @ruc. edu. (gaozf, dinggk, zhengh, zhanglf, keql edp..tch
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 to- kens or 2D topology graphs. limiting their ability to incorporate 3D information for downstream tasks and, in particular, making it almost impossible for 3D ge- ometry prediction/generaan framework, called Uni-Mol, that significantly enlarges the representation abil- ity and application scope of MRL schemes. Uni-Mol contains two pretrained models with the same SE(3) Transformer architecture: a molecular model pre- trained by 209M molecular conformations; a pocket model pretrained by 3M candidate protein pocket data. Besides, Uni-Mol contains several finetuning strate- gies apply the pretrained models various downstream tasks. By properly incorporating 3D information, Uni-Mol outperforms SOTA in 14/15 molecular property prediction tasks. Moreover, Uni-Mol achieves superior performance 3D spatial tasks, including protein-ligand binding pose prediction, molecular con- formation generation, etc. The code, model, and data are made publicly available at htttps://gitthhi com dptpch-tcccccccccccco corp Uni Moll
In this paper, we propose universal 3D MRL
INTRODUCTION
Recently, representation learning (or pretraining, self-supervised learning) 2 has been prevailing in many applications, such as BERT [4] and GPT [5 7 in Natural Language Processing (NLP), ViT g in Computer Vision (CV), etc. These applications have a common characteristic: unlabeled data is abundant, while labeled data is limited. As solution, in typical representation learning method, one first adopts pretraining procedure to learn a good representation from large-scale unlabeled data. Then a finetuning scheme is followed to extract more information from limited supervised data.
Applications in the field of drug design share the characteristic that calls for representation learning schemes. The chemical space that a drug candidate lies is vast, while drug-related labeled data is limited. Not surprisingly, compared with traditional molecular fingerprint-basee models 10] recent molecular representation learning (MRL) models perform much better in most property prediction tasks [TT 12 133. However, to further improve the performance and extend the application scope of existing MRI. models, one is faced with a critical issue. From the perspective of life science, the properties of molecules and the effects of drugs are mostly determined by their 3D structures [14] [5]. In most current MRL methods, one starts with representing molecules as ID sequential strings, such as SMILES (16] [8] and InChI [19 20: 21), or 2D graphs 222 II 23 12 24). This may limit their ability to incorporate 3D information for downstream tasks. In particular, this makes it almost impossible for 3D geometry prediction or generation, such as, e.g., the prediction of protein- ligand binding pose 25] Even though there have been some recent attempts trying to leverage 3D
information MRL 26 the performance is less than optimal, possibly due to small size of
Equal contribution.
Corresponding authors.

---

Published as a conference paper at ICLR 2023
Pretraining
Finetune
30 Position Recovery
Masked Atom Prediction
Mo Prodiction
Backbone
molecures
Medeccde model
Molocular
Backbone
Backbone
POceeT Property Predicion
Podkect
IM candicate pocketss
Ligand Binding Complix
Figure 1: Schematic illustration of the Uni-Mol framework.
3D datasets, and 3D positions can not be used as inputs/outputs during finetuning, since they only serve as auxiliary information.
In this work, we propose Uni-Mol, to our best knowledge, the first universal 3D molecular pretraining framework, which derived from large-scale unlabeled data and able to directly take 3D positions as both inputs and outputs. In particular, Uni-Mol consists of 3 parts. 1) Backbone. A Transformer based model that can effectively capture the input 3D information, and predict 3D positions directly. 2) Pretraining. Two large-scale datasets: 209M molecular conformation dataset and 3M candidate protein pocket dataset, for pretraining 2 models on molecules and protein pockets, respectively. And two pretraining tasks: 3D position recovery and masked atom prediction, for effectively learning 3DD spatial representation. 3) Finetuning- Several finetuning strategies for various downstream tasks. For example, how use the pretrained molecular model in molecular property prediction tasks; how to combine the two pretrained models in protein binding pose prediction. We refer to Fig. for an overall schematic illustration of the Uni-Mol framework, and the details will be described in Sec. 2
To demonstrate the effectiveness of Uni-Mol, we conduct experiments on series of downstream tasks. In the molecular property prediction tasks, Uni-Mol outperforms SOTA on 14/15 datasets on the MoleculeNet benchmark. In 3D geometric tasks, Uni-Mol also achieves superior performance, For the pose prediction of protein-ligand complexes, Uni-Mol predicts 80.35% binding poses with RMSD <= 2A, 22.58% relatively better than popular docking methods, and ranks Ist in the docking power test on CASF-2016 [28) benchmark. Regarding molecular conformation generation, Uni-Mol achieves SOTA for both Coverage and Matching metrics on GEOM RMM and GEOM- -Dugs
(29] Moreover,
Uni-Mol can be successfully applied to tasks with very limited data like pocket druggability prediction.
To summarize, Uni-Mol made the following contributions: 1) To our best knowledge, Uni-Mol the first pure 3D molecular pretraining framework that can predict 3D positions, and the first molecular pretraining framework that can be directly used in 3D tasks in the field of drug design. 2) Based on extensive benchmarks, we build simple and efficient SE(3) Transformer backbone and effective 3D pretraining strategy in Uni-Mol. 3) Uni-Mol outperforms SOTA in various downstream tasks. 4) The whole Uni-Mol framework, including code, model, and data, will be made publicly available.
UNI-MOL FRAMEWORK
2.1 BACKBONE
In MRL, there are two well-known backbone models, graph neural (eNW 22 23 12 and Transformer 24 [TI. With GNN the backbone model, for efficiency purposes, locally connected graphs are often used to represent molecules. However, the locally connected graph lacks the ability to capture the long-range interactions among atoms. We believe that long-range interactions are important in MRL. Therefore, We choose Transformer as the backbone model in Uni-Mol, as it fully connects the nodes/atoms and thus can learn the possible long-range interactions.
Ithough the backhone output SEu))--quivariant positions, it based SF(3)-invariant representations.
2

---

Published as a conference paper at ICLR 2023
Protraining tasks
FFN
Picdiction
Layers Backbone Block
Iype Head
Head
Head
Softnoo
Backbone
corld paner
MotMill
Linoe
Lincer
Linour
Atomm
Cuscainnles
Atom Represertattin
Figure 2: Left: the overall pretraining architecture. Middle: the model inputs, including atom representation and pair representation. Right: details in the model block.
However, Transformer cannot handle 3D spatial data directly it was originally designed for NLP tasks. Although there are several recent works extending Transformer to 3D data (30 B1. most of them much slower than the standard Transformer due the complex additional components like Tensor Field Networks. Considering the pretraining cost in the large-scale dataset, we need an efficient backbone. To achieve that, based on the standard Transformer with Tre-LayerNorm [32] we introduce several efficient and necessary modifications, for the ability to take 3D positions as inputs and outputs.
Architecture Overview As illustrated in Fig. 2 the Uni-Mol backbone a Transformer based model. It has two inputs, atom types and atom coordinates. And two representations (atom and pair) are maintained in the model. The atom representation is initialized from atom types, by the Embedding layer; The pair representation is initialized by invariant spatial positional encoding calculated from atom coordinates. In particular, based on pair-wise Euclidean distances among atoms, the pair representation is invariant to global rotation and translation. The two representations communicate with each other in self-attention module. Details are in the following subsections.
Encode 3D positions Due its permutational invariance, Transformer cannot distinguish the posi- tions of inputs without positional encoding. Different with the discrete positions used in NLP/CV (33 34 the positions in 3D space, i.e. coordinates, ing procedure needs to be invariant under global rotation and translation. Several 3D spatial positional encodings were already proposed tackle this 33 36 37 38 and we have no interest in reinventing new one. Therefore, we benchmark existing encodings (in Appendix D.I). and use simple and effective one: Euclidean distances of atom pairs, followed by pair-type aware Gaussian kernel [39].
continuous values. Besides, the positional encod-
Furthermore, since the invariant 3D spatial positional encoding is encoded at the pair level, we also maintain a pair-level representation in Transformer, to enhance the 3D spatial representation. Specifically, the pair representation is initialized as the aforementioned spatial positional encoding. Then, to update pair representation, we use atom-to-pair communication via the result of the multi- head Query-Key Key product in self-attention. Formally, the update of ij pair representation is denoted as
d +1 Q. (K!" 1h [1,
++1 L
(I)
Vd
where d, is the pair representation of atom pair ij in 1-th layer, is the number attention heads, is the dimension of hidden representations, and Q. (k/b"t) is the Query (Key) of the i-th (j-th) atom in the I-th layer h-th head. Besides, leverage 3D information in the atom representation. we also introduce pair-to-atom communication, by using the pair representation as the bias term in self-attention. Formally, the self-attention with pair-to-aton communication denoted
Q.a(K, Vd
Attention Q/ K/h V/ =ootmax
+ 4,
V
where vith is the Value of the j-th atom the 1-th layer h-th head. As shown in the above equations, the proposed pair representation is very simple, and the extra cost of maintaining it is negligible. And our benchmark in Appendix D also demonstrates its efficiency and effectiveness.
3

---

Published as a conference paper at ICLR 2023
Algorithm 1 Corrupted Position Generation and Assignment.
Require: X E Rmxx, coordinates of m atoms, 7: noise range, reassign: re-assignment not
I: R = X + 6, where ~ Uniform( -rÅ, rA)
P Generate corrupted positions > Return directly if not need the re-assignment b Compute distance Greedy assignment based random order b Get nearest position at the i-th row D Assignment Marked the k-th column as used P Return corrupted positions
2: if (not re assign) return R
3: D = (D(s j| |X R i.j j= m)
4: for random 1aii__eerm__1_111111 m) do
5:
= argmin(Da.)
6: Y; = R
7: D inf
return
Predict 3D positions 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 coordinates, which is essential in 3D spatial tasks. To this end, we introduce SE(())-quivariant head predict the delta positions based on SEvaa-invrriant pair representation and equivariant input 2, denoted as
(e: -,jayy
Gy = ReLUU(((( R3 UWW,
where n is the number of atoms, is the number of layers 2, E R the input coordinate i-th atom, and è; E R3 is the output coordinate of i-th atom, ReLU(y) max(0, y) is Rectified Linear Unit [40]. U E TR x and W e lx are the projection matrices to convert pair representation to scalar. This head is similar to the position update procedure in EGNN [1], but much more efficient due to Uni-Mol only updating 3D positions the last layer. Besides, to be consistent with delta posi- tion prediction, Uni-Mol uses delta pair representation to update coordinates, while EGNN uses pair representation directly. Our benchmark in Appendix D.3 demonstrates the one in Uni-Mol is better.
Please note that the backbone in Uni-Mol can be replaced with any SE(3) model that can take 3D positions as inputs and outputs. However, considering the massive pretraining cost in the large-scale dataset, we favor the efficient backbone. So we only make several simple and necessary modifications to the standard Transformer model, based on the ablation benchmark results in Appendix D
2.2 PRETRAINING
Large-Scale dataset For the purpose of pretraining, we generate two large-scale datasets, one composed 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 protein-ligand structures and interactions.
The molecular pretraining dataset based on multiple public datasets (See Appendix A for more information). After normalizing and deduplicating, it contains about 19M molecules. To generate 3D conformations, we use ETKGD (2 with Merck Molecular Force Field [43] optimization in RDKit 44 to randomly generate 11 conformations for each molecule. totally 209M conformations.
The protein pocket pretraining dataset is derived from the Protein Data Bank (RCSB PDB) [45), a collection of 180K 3D structures of proteins. To extract candidate pockets, we first clean the data by adding the missing side chains and hydrogen atoms; then we use Fpocket [46] to detect possible binding pockets of the proteins; and finally, filter pockets by the number of residues. In this way, We collect dataset of 3.2M candidate pockets for pretraining.
Pretraining strategies Self-supervised 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. In Uni-Mol, we want to encourage the model to learn the 3D structural information during pretraining. To this end, design 3D position recovery self-supervised task. The main idea of the task is to recover the correct 3D positions, given the corrupted input positions. An intuitive way is mask the positions, like the token masking in BERT.
4

---

Published as a conference paper at ICLR 2023
However, the positions are continuous values, not discrete values; cannot use special value to represent the mask (like the [MASK] token in BERT).
Therefore, rather than masking, random positions used as corrupted input 3D positions, and the model is trained to predict the correct position. Nevertheless, learning the mapping from a random po- sition the ground-truth atom position is very challenging. There are two technologies to reduce the delta positions (between random and ground-truth positions), making the learning more feasible. First, re-assignment, given m atoms and mn random positions, there are mn! possible assignments. Among them, following the stationary-action principle 47 we can use the one with minimal delta positions. Due to the difficulty of finding an optimal solution, we use an efficient greedy algorithm to find a sub-optimal re- re-assignment. Second, noise range, we can limit the space of the random positions, only allowing the random positions with noise (r) around the ground-truth ones. There a tradeoff here; if r is large, the re-assignment is required to make the learning feasible; if r is small, the re-assignment may not need. We summarized the algorithm into Alg. | and benchmarked several settings, details in Appendix D.6 and found a simple and effective one: use = A, without re-assignment.
Then, with corrupted input coordinates, two additional heads are used to recover the correct positions. 1) Pair-distance prediction. Based on pair-representation, the model needs to predict the correct Fu- clidean distances of the corrupted atoms pairs. 2) Coordinate prediction. Based on SE(3)- (3(()---qivariant coordinate head, the model needs to predict the correct coordinates for the corrupted atoms.
Finally, the atom types for the corrupted atoms are masked, and head is used to predict the correct atom types. For the convenience of finetuning, similar to BERT, a special atom [CLS], whose coordinate is the center of all atoms, is used to represent the whole molecule/pocket. Both 2 pretraining models use the same self-supervised tasks described above, and Figure 2 is an illustration of the overall pretraining framework. For the detailed configurations of pretraining, please refer to Appendix CC] Similar to backbone design, we conduct an extensive benchmark for pretraining strategies (in Appendix [D.5 and D.6 and choose the above strategies based performance.
2.3 FINETUNING
To be consistent with pretraining, we use the same data prepossessing pipeline during finetuning. For molecules, multiple random conformations can be generated in short time, we can use them as data augmentation in finetuning to improve performance and robustness. 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 or ligand, we can categorize them as follow.
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 mooket- pair, we can concatenate their [CLS] representations, and then finetune with linear head.
3D prediction tasks of molecules or pockets These tasks need to predict a 3D conformation of the input, such as molecular conformation generation. Different from 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 is very time-consuming. Therefore, there are many recent works that train the model to fast generate conformations from molecular
graph [48] 49; 50 51, While in Uni-Mol, this task straightforwardly becomes a conformation optimization task: generate a new conformation based on a different input conformation. Specifically, in finetuning, the model supervised learns the mapping from Uni-Mol generated conformations to the labeled conformations. Moreover, the output conformations can be generated end-to-end by (qqiivraqi head.
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
5

---

Published as a conference paper at ICLR 2023
Table 1: Uni-Mol performance on molecular property prediction classification tasks
Classification (ROC-AUC %, higher is better f
Datasets
BBBP 2039
BACE
ClinTox
Tox21
ToxCast
SIDER 1427
HIV 41127
PCBA
MUV
Molecules
1513
1478
7831
8575
437929
93087
Tasks
1
617
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.5)
63.7(0.2)
60.6(3.2)
75.7(1.4)
80.1(1.4)
76.6(1.5)
N- A-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-Gramxca
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)
PretrainGNN
68.7(1.3)
84.5(0.7)
72.6(1.5)
78.1(.6)
65.7(0.6)
62.7(0.8)
79.9(0.7) 86.0(0.1)
81.3(2.1)
GROVER
70.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(377)
73.5(0.1)
65.3(0.5)
65.4(0.1)
68.2(1.1) 83.0(0.4)
67.3(1.8)
GvaphMVP
72.4(1.6)
81.2(009)
79.1(2.8)
75.9(0.5)
63.1(0.4)
63.9(1.2)
77.0(1.)
77.7(0.6)
MoICLR
72.2(2.1)
82.4(0.9)
91.2(3.5)
75.0(0.2)
58,9(14)
78.10.55)
79.6(1.9)
GEM
72.4(0.4)
85.6(1.1)
90.1(1.3)
78.1(0.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)
molecular ligand. Besides the conformation changes of the pocket and the molecule themselves, we also need consider how the molecule lays in the pocket, that is, the additional 6 degrees (3 rotations and 3 translations) of freedom of a rigid movement. In principle, with Uni-Mol, we can predict the complex conformation by the 3E(3)-Equivariant head in an end-to-end fashion. However, this is unstable as it is very sensitive to the initial docking positions of molecular ligands. Herein, to get rid of the initial positions, we use a 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 input of an additional 4-layer Uni-Me Mol decoder, which is finetuned to learn the pair distances of all heavy atoms in molecule and pocket. Then, with the predicted pair-distance matrix as a scoring function, directly back-propagation the loss of the current pair-distance matrix and the predicted pair-distance matrix. Thanks to the efficiency of propappop,gation, this process very fast, about 100x faster than traditional docking tools. More details can be found in Appendix C.6
first randomly place the ligand and then optimize the coordinates of its atoms by
3 EXPERIMENTS
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 C and D
3.1 MOLECULAR PROPERTY PREDICTION
a popular benchmark for molecular property prediction, including datasets focusing different molecular properties, from quantum mechanies and physical chemistry to biophysics and physiology. Following previous work GEM [3], we use scaffold splitting and report the mean and standard deviation by the results of random seeds.
Datasets and setup MoleculeNet 52
Baselines We compare Uni-Mol with multiple baselines, including supervised and pretraining baselines. D-MPNN (53] and AttentiveFP (54) are supervised GNNs methods. N-gram (55]. PretrainGNN 22), GROVER [I]] GraphMVP [26], MolCLR [12] and GEM [33 are pretraining methods. N-gram embeds the nodes in the graph and assembles them in short walks as the graph representation. Random Forest and XGBoost (56) are used as predictors for downstream tasks.
Results Table and Table show the experiment results of Uni-Mol and competitive baselines, where the best results 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), rerun it with the same

---

Published as a conference paper at ICLR 2023
Table 2: Uni-Mol performance on molecular property prediction regression tasks
Regression (lower is better L)
RMSE
MAE
Datasets
ESOL
FreeSolv
Lipo 4200
QM7 6830
QM8
QM9
# Molecules
1128
642
21786
133885 3
# Tasks
D-MPNN
(.000(0.088)
2.082(0.082)
0.68300.016)
103.5(8.6)
0.0190(0.0001)
0.00814(0.00m ((00000001)
Attentive FP
0.877(0.029)
1073(0.183)
0.721(0.001)
72.0(2.7)
0.0179(0.001)
0.00812(0) (200000001)
N-Gramar
1.774(0.07))
2.688(0.085)
0.812(0.028)
92.8(4.0)
0.0236(0. (000))
2.01037(0.0001
N- N-Gramxca
1083(00082)
5.061(0.744)
2.072(0.030)
81.9(1.9)
0.0215(0.00055
0.00964 (0.00031)
PretrainGNN
1.100(0.006)
2.764(0.002)
0.739(0.003)
113.2(0.6)
0.0200(0.0001)
O. 0.00922(0.0000)
GROVER
0.983(0.000)
2.176(0.052)
0.817(0.008)
94.5(3.8)
0.0218(0.0004)
00094400.00055)
GROVER
0.895(0.017)
1.
(.820(0.010)
92.0(0.9)
0.0224(0.0003)
0.00986(0.00025))
GvaphMVP
(.02((00033)
(0.800)
MoICLR
1.271(0.040)
2.594(0.249)
0.691(0.004)
66.8(2.3)
0.0178(0.0003)
GEM
0.798(0.029)
1887700094))
0.660(0.008)
58.9(0.8)
0.0170(0.0001)
0.00746(0.00001)
Uni-Mol
0.788(0.029)
14800000048)
0.603(0.010)
41.8(0.2)
0.0156(0.0001)
0.00467(0000004)
data split setting other baselines. From the results, we can summarize them as follows: 1) overall, Uni-Mol outperforms baselines on almost all downstream datasets. 2) In solubility (Lipo), free energy (FreeSolv), and quantum mechanical (QM7, QM8, QM9) properties prediction tasks, Uni-Mol significantly better than baselines. As 3D information is critical in these properties 14 [5] indicates that Uni-Mol can learn better 3D representation than other baselines. 3) Uni-Mol fails to beat SOTA on the SIDER dataset. After investigation, we find that Uni-Mol fails to generate 3D conformations for many molecules (like natural products and peptides) in SIDER. Therefore, due the missing 3D information, is reasonable that Uni-Mol cannot outperform others.
In summary, by better utilizing 3D information in pretraining, Uni-Mol outperforms all previous MRL models in almost all property prediction tasks.
3.2 MOLECULAR CONFORMATION GENERATION
We leave the details of molecular conformation generation to Appendix paper [57] pointed out that the current benchmark for molecular conformation generation could be wrong.
3.3 POCKET PROPERTY PREDICTION
Datasets and setup Druggability, the ability of candidate protein pocket to produce stable binding to specific molecular ligand, is one of the most critical properties of a candidate protein pocket. How- ever, this task is very challenging due to the very limited supervised data. For example, NRDLD (58), a commonly used dataset, only contains 113 data samples. Therefore, besides NRDLD, we construct regression dataset for benchmarking pocket property prediction performance. Specifically, based on Fpocket tool, we calculate Fpocket Score, Druggability Score, Total SASA, and Hydrophobicity Score for the selected 164,586 candidate pockets. The model is finetuned to predict these scores.
Baselines On the NRDLD dataset, we compare Uni-Mol with previous methods evaluated in Accuracy, recall, precision, and F1-score are used metrics for this classification task. On our created benchmark dataset, as there are no appropriate baselines, we use an additional Uni-Mol model without pretraining, denoted as Uni-Molpo pretraiedd to check the performance brought by pretraining on pocket property prediction.
Results As shown in Table 3 Uni-Mol shows the best accuracy, recall, and F F1-score on NRDLD. In our created benchmark dataset, the pretraining Uni-Mol model largely outperforms the non- pretraining one on all four scores. This indicates that pretraining candidate protein pockets indeed brings improvement pocket property prediction tasks. Unlike Molecular property prediction, due to the very limited supervised data, pocket property prediction gained much less attention. Therefore, we also release our created benchmark dataset, and hopefully, can help future research.
7

---

Published as a conference paper at ICLR 2023
Table 3: Uni-Mol performance on pocket property prediction
Classification (higher is better f>
Regression (lower is better 4)
Dataset
NRDLD
Our Created
Methods [Cavity-DnggSee DrugSoore
Volsite DrugPred PockDrug TRAPP-CNN Uni-Mol Methods
Uni Uni-Molae parriied
Uni-Mol
Accuracy
0.89
0.865
0.973
RMSTEppoter
(.1055(0.002)
0.1140(0.001)
Recall
0.957
0.913
1.000
RMS Epvaaaaiilllyyy
0.1017(0.002)
0.1001(0.001)
Precision
0.846
1,000
0.958
RMSE SASA
22.010X0.460)
20.734(0.015)
F1-score
0.898
0.955
0.979
1,41440.034)
1.28447 (0.005)
Uni-Mol
Ligand RMSD
AutodankVima
% Below Threshold f
SVinsRF20
GiokScore-Si
Methods
1.0 A
1.5 A 2.0 Á 3.0 Á
5.0 Á
DrugScoreCSD
DawaDack
Autodock Vina
44.21
57.54
64.56
73.68
84.56
GBVI WSA dBGeMEE
Vinardo
41.75
57.54
62.81
69.82
76.84
ChemPLPgBGOLD
Lg80on2@00
Smina
47.37
59.65
65.26
74.39
82.11
GietnSoomm-X0 XP
Autodock4
21.75
31.58
35.44
47.02
64.56
Top 1. Success Rate (8)
88
Uni-Mol
nee-aaiaeeeertnneed pretraidd
39.65
63.16
72.98
83.51
91.58
Uni-Mol
43.16
68.42
80.35
87.02
94.04
Figure 3: Docking power evaluation on CASF-2016 (Top 10 methods)
Table 4: Uni-Mol performance on binding pose prediction
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, then optimize the complex conformations. For the benchmark dataset, referring to the previous works [28] 660, we use CASF-2016 as the test set. For the training data used in finetuning, we use PDBbind General set v.2020 [6 (19, ,444 complexes). Notably, to examine the generalization ability, we further filter out the training complexes that are similar to the ones in the test set (CASF-2016). In particular, the complexes with both high protein sequence similarity (MMSeqs2 [62) similarity above 40%) and high molecular similarity (fingerprint similarity above 80%) are filtered out, and there are 18,404 complexes after filtering.
Two benchmarks are conducted: 1) Docking power, the default metric to benchmark the ability of a scoring function in CASF-2016. Specifically, it tests whether a scoring function can distinguish the ground truth binding pose from a set of decoys or not. CASF-2016 provides 50- 100 decoy conformations of the same ligand for each ground truth. Scoring functions are applied to rank them, and the ground truth is expected to be the top 1. 2) Binding pose accuracy. Specifically, we use the semi- aemi-flexible docking setting: keep the pocket conformation fixed, while the ligand conformation is fully flexible. We evaluate the RMSD between the prediction and the ground truth. Following previous works, we use the percentage of results below predefined RMSD thresholds as metrics.
Baselines
For the docking power benchmark, the baselines DeepDock (60 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 [63 64).
Vinardo [65], Smina [66]. and AutoDock4 [67].
Results From the docking power benchmark results shown in Figure Uni-Mol ranks the 1st, with the top success rate of 91.2%. For comparison, the previous top scoring function AutoDock Vina [63 64 achieves 90.2% of the top success rate this benchmark. From the binding pose accuracy results shown in Table 4 Uni-Mol also outperforms other baselines. Notably. Uni-Mol outperforms the second best method by relatively 22.58% 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 their interaction in the 3D space. Even without pretraining, Uni-Mol (denoted as Uni-Mol also performs very well. This demonstrates the effectiveness of Uni-Mol backbone, as it effectively learns the 3D information by limited data. In summary, by combining molecular and pocket pretraining models, Uni-Mol significantly outperforms the widely used docking tools in the protein-ligan binding tasks. We leave the efficiency benchmark and visualization for binding pose prediction to Appendix B
8

---

Published as a conference paper at ICLR 2023
4 RELATED WORK
Representation learning In recent years, representation learning 1 2 has received much attention and has been prevailing in many applications, like in NLP multi-modal
It is not doubted that representation learning becomes a default technology 4
S
68 6 77.
CV
[8] 69 70), or
[7] 72 73..
in various tasks.
Molecular representation learning Representation learning large-scale unlaheled molecules attracts much attention recently. SMILES-BERS [8] is pretrained on SMILES strings of molecules
using BERT [4. Subsequent works are mostly pretraining on 2D molecular topological graphs [23
I.I. Mo|CLR 212 applies data augmentation to molecular graphs at both node and graph levels, using a supervieed contrastive learning strategy to learn molecular representations. Further, several 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] proposes a contrastive learning GNN-based framework between 2D topology and 3D geometry. GEM 13 uses bond angles and bond length as additional edge attributes to enhance 3D information.
Equivaaqn models In many-body scenarios such as potential energy surface fitting, SE-(3)
equivariance is usually required. A series of SE(3) models proposed, such SchNet (74). tensor field networks [30], SE(3) Transformer [3I], DimmNet [75), equivariant graph neural networks (EGNN) 41, GemNet (37] and SphereNet (76). Most of these models are designed for supervised learning with energy and force.
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 utilize 3D) spatial data beyond sequential information. Early methods, such as Volsite (77), DrugPred (58), and PockDrug 788.. predict druggability based on the predefined descriptors of pockets' static structures. Later, TRAPP-CNN [59], based on 3D-CNN, proposes the analysis of proteins' conformation changes and the use of such information for druggability prediction.
Protein-ligand binding pose prediction In structure-base 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-based drug design and has been developed for the past decades. Tools such as AutoDock4 167. AutoDock Vina (44) and docking methods, such as Av RF20 [79] and DeepDock [60)] have also been developed to predict protein- ligand binding poses and assess protein-ligand binding affinity. Equibind [80 a recent graph deep learning based methods. However, Uni-Mol cannot have apple-to-apple comparison with Equibind, due to Equibind being proposed for Blind Docking. While Uni-Mol is currently designed for Targeted Docking, which follows most previous traditional tools in docking [8T]. The difference is that Blind Docking uses whole protein for docking, while Target Docking directly uses the pocket. We will extend Uni-Mol to Blind Docking tasks future work.
Smina [66] are among the most used docking programs. Also, machine learning-base
5 CONCLUSION
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.
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, is worthy of training large Uni-Mol model on a bigger dataset. 3) More high-quality benchmarks. Although there have been many applications in the field of drug design, high- gigh-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.
9

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We thank Shuqi Lu, Yuanqi Du, Zhen Wang, Yingze Wang, Junhan Chang, Xi Chen, Zhengdan Zhu and many colleagues in DP Technology for their great help in this project. We thank Yean Cheng for his discovery of bug in the pocket property task.
The work was partially done at Gaoling School of Artificial Intelligence, Peng Cheng Laboratory. Beijing Key Laboratory of Big Data Management and Analysis Methods and MOE Key Lab of Data Engineering and Knowledge Engineering. This research was supported in part by National Natural Science Foundation of China (No. U2241212, No. 61972401, No. 61932001, No. 61832017), by Beijing Natural Science Foundation (No. 4222028), by the major key project of PCL (PCL2021A12) by Beijing Outstanding Young Scientist Program No. NO.BJJWZYJHO 2019 100020098. We also wish to acknowledge the support provided by Engineering Research Center of Next-Generation Intelligent Search and Recommendation, Ministry of Education. Additionally, we acknowledge the support from Intelligent Social Governance Interdisciplinary Platform, Major Innovation & Plamning Interdisci- plinary Platform for the "Double-First Class" Initiative, Public Policy and Decision-making Research Lab, Public Computing Cloud, Renmin University of China.
3D structures in Fig. are drawn using the web service httiiieetMttttt//ee/tttee do tech).
REFERENCES
[1] Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 1798):1798-1828 2013.
[2] William L. Hamilton, Rex Ying, and Jure Leskovee. Representation learning on graphs: Methods and applications. IEEE Data Eng. Bull. 4023):22-74, 2017.
[3] Daokun Zhang, Jie Yin, Xingquan Zhu, and Chengqi Zhang. Network representation Icarning: A survey. IEEE transactions on Big Data, 6 6()):3-28, 2018.
[4] Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies. Volume I (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguisties.
[5] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understand- ing by generative pre-training. 2018.
[6] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188)9 2019.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[8] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, In International Conference Learning Representations, 2021.
[9] Qingda Zang, Kamel Mansouri, Antony J Williams, Richard S Judson, David G Allen, Warren M Casey, and Nicole Kleinstreuer. In silico prediction of physicochemiaal properties of environmental chemicals using molecular fingerprints and machine learning. Journal of chemical information and modeling. 57(1):36 49, 2017.
[10] Minjian Yang, Bingzhong Tao, Chengjuan Chen, Wenqiang Jia. Shaolei Sun. Tiantai Zhang, and Xiaojian Wang. Machine learning models based on molecular fingerprints and an extreme gradient boosting method lead to the discovery of jak2 inhibitors. Journal of Chemical Information and Modeling, 50002- 5012, 2019.
[11] Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self- supervised graph transformer on large-scale molecular data, Advances in Neural Information Processing Systems, 33, 2020.
10

---

Published as a conference paper at ICLR 2023
[12] Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence, pages 1-9, 2022.
[13] Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingho Zhou, Fan Wang. Hua Wu, and Haifeng Wang. Geometry-enhanee molecular representation learning for property prediction. Nature Machine Intelligence, pages 1-8, 2022.
[14) A Crum-Brown and TR Fraser. The connection of chemical constitution and physiological action. Trans R Soc Edinb, 25(1968- 25(1968-1969)223) 1969):257, 1865.
[15] Corwin Hansch and Toshio Fujita. p-a-it analysis. a methad for the correlation of biological activity and chemical structure. Journal of the American Chemical Society, 16(8):1616-1628 1964.
[16] David Weininger. Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of chemical information and computer sciences, 28(1):31-33) 1988.
[17] Zheng Xu, Sheng Wang. Feiyun 7hu, and Junzhou Huang. molecular embedding for drug discovery. In Proceedings of the 8th ACM international conference on bioinformatics. computational biology, and health informatics, pages 285-294, 2017.
Seq2seq fingerprint: An unsupervised deep
[18) Sheng Wang, Yuzhi Guo, Yuhong Wang. Hongmao Sun, Junzhou Huang. 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, pages 429-436, 2019.
[19] Stephen R Heller, Alan McNaught, Igor Pletnev, Stephen Stein, and Dmitrii Tchekhovskoi. Inchi, the iupac international chemical identifier. Journal of cheminformaties, 7(1):1334, 2015.
[20] Robin Winter, Floriane Montanari, Frank Noé, and Djork-Aréé Clevert. Learning continuous and data-driven molecular descriptors by translating equivalent chemical representations. Chemical science, 192-):16,,1 2019.
[21] Jennifer Handsel, Brian Matthews, Nicola Knight, and Simon J Coles. Translating the inchi: adapting neural machine translation to predict iupac names from chemical identifier. Journal of cheminformatics 13(1):1 11, 2021.
[22] Weihua Hut, Bowen Liu", Joseph Gomes, Marinka Zitiik, Percy Liang, Vijay Pande, and Jure Leskovee. Strategies for pre-training graph neural networks. In International Conference on Learning Representa- tions, 2020.
[23] Pengyong Li, Jun Wang, Yixuan Qiao, Hao Chen, Yihuan Yu, Xiaojun Yao, Peng Gao, Guotong Xie, and Sen Song. An effective self superviee framework for learning expressive molecular global representations drug discovery. Briefings Bioinformatics, 22(6): 109, 2021.
[24] 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 Neural Information Processing Systems, 34, 2021.
[25] Panagiotis Koukos, Li C Xue, and Alexandre MJJ Bonvin. Protein-ligand pose and affinity prediction: Lessons from d3r grand challenge 3. Journal of computer-aided molecular design, 33(1 8:3391, 2019.
[26] Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre-training molecular graph representation with 3d geometry. In International Conference on Learning Representa- tions, 2022.
[27] Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Gün- nemann, and Pietro Liö. 3d infomax improves gnns for molecular property prediction. arXiv preprint urXiv:2110.04122. 2021.
[28) Minyi Su, Qifan Yang, Du, Guoqin Feng, Zhihai Liu, Yan Li, and Renxiao Wang. Comparative assessment of scoring functions: the casf-2016 update. Journal of chemical information and modeling. 59(2):895-913, 2018.
[29] Andrew L Hopkins, Colin R Groom, and Alexander Alex. Ligand efficiency: a useful metric for lead selection. Drug discovery today, 910):430 431, 2004.
[30] Nathaniel Thomas. Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-cquiivvrann neural networks for 3d point clouds. arXiv preprint arXiv: o 1802.08219, 2018
11

---

Published as a conference paper at ICLR 2023
[31] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation cquivariant attention networks. Advances in Neural Information Processing Systems, 33:1970-19811 2020.
[32] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng. Chen Xing. Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. Hal Daumé III and Aarti Singh, cditors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 10524-10533. PMLR, July 2020.
[33] Guolin Ke, Di He, and Tie-Yan Liu. Rethinking positional encoding in language pre-training. In International Conference Learning Representations, 2020.
[34] Philipp Dufter, Martin Schmitt, and Hinrich Schütze. Position information in transformers: An overview. arXiv preprint arXiv:2102. 1090, 2021.
[35] Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan Hc, Shengjie Luo, Chang Liu, Di He, and Yan Liu. Benchmarking graphormer on large-scale molecular modeling datasets. arXiv preprint 11Xiv:2003.04810, 2022.
[36] Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip HS Torr, and Vladlen Koltun. Point transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 16259-16268, 2021.
[37] Johannes Kliepera, Florian Becker, and Stephan Günnemann. Gemmet: Universal directional graph neural networks for molecules. In Advances in Neural Information Processing Systems, 2021.
[38] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakoo,, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 59667887) 2021.
[39] Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C Lawrence Zitnick. Rotation invariant graph neural networks using spin convolutions. arXiv preprint 27X5::2106.0955, 2021.
[40] Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprint 2018.
[41] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E (n) equivariant graph neural networks. In International Conference on Machine Learning, pages 9323-9332. PMLR, 2021.
[42] Screina Riniker and Gregory A Landrum. Better informed distance geometry: using what we know to improve conformation generation. Journal of chemical information and modeling, 55(12):2562-25 2574, 2015.
[43] Thomas A Halgren. Merck molecular force field. basis, form, scope, parameterization, and performance of mmff94, Journal of computational chemistry, 17(5-6):490-519, 1996.
[44] Greg Landrum et al. Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling, 2013.
[45] Helen M Berman, John Westbrook, Zukang Feng, Gary Gilliland, Talapady N Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E Bourne. The protein data bank. Nucleic acids research, 28(1):33--242, 2000.
[46] Vincent Le Guilloux, Peter Schmidtke, and Pierre Tuffery. Fpocket: open source platform for ligand pocket detection. BMC bioinformatics, 10()):1-11, 2009.
[47] Richard Feynman. The Character of Physical Law, with new foreword. MIT press, 2017.
[48] Gregor Simm and Jose Miguel Hermandez-Looato A generative model for molecular distance geometry. In International Conference on Machine Learning, pages 8949-8958. PMLR, 2020.
[49] Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. In International Conference on Learning Representations, 2020.
[50] Elman Mansimov, Omar Mahmood, Scakho Kang, and Kyunghyun Cho. Molccular geometry prediction using a deep generative graph neural network. Scientific reports, 91):1113, 2019.
[51] Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen. William Green, and Tommi Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer ensembles. Advances in Neural Information Processing Systems, 34, 2021.
12

---

Published as a conference paper at ICLR 2023
[52] Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu. Karl Leswing, and Vijay Pande. Moleculenct: benchmark for molecular machine learning. Chemical science, 9()):13 223,003-330 2018.
[53] Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman-Perez Perezz Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8)33370-338 2019.
[54] Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li. Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism. Journal of medicinal chemistry, 63(16):8749 8760, 2019.
[55] Shengehao Liu, Mehmet F Demirel, and Yingyu Liang. N-gram graph: Simple unsupervised representa- tion for graphs, with applications to molecules. Advances in neural information processing systems, 32, 2019.
[56] Tianqi Chen and Carlos Guestrin. Xgboost: scalable tree boosting system. Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining, pages 785 -944794, 2016.
[57] Gengmo Zhou, Zhifeng Gao, Zhewei Wei, Hang Zheng, and Guolin Ke. Do deep learning methods really perform better in molecular conformation generation? arXiv preprint arXiv:2302.07000 077061 2023.
[58] Agata Krasowski, Daniel Muthas, Aurijit Sarkar. Stefan Schmitt, and Ruth Brenk. Drugpred: structure- based approach to predict protein druggability developed using an extensive nonredundant data set. Journal of chemical information and modeling, 51(11):2829 2842, 2011.
[59] Jui-Hung Yuan, Sungho Bosco Han. Stefan Richter. Rebecca C Wade, and Daria Kokh. Druggability assessment in trapp using machine learning approaches. Journal of Chemical Information and Modeling. 6013):1685-199) 2020.
[60] Oscar Méndez-Lucio, Mazen Ahmad, Ehecatl Antonio del Rio-Chanona, and Jörg Kurt Wegner. A geometric deep learning approach predict binding conformations of bioactive molecules. Nature Machine Intelligence, 3(12):1033 10301039-11111111111111110991111 2021.
[61] Zhihai Liu, Yan Li Han, Li, Jie Liu. Zhixiong Zhao, Wei Nie, Yuchen Liu, and Renxiao Wang. Pdb- wide collection of binding data: current status of the pdbbind database. Bioinformatics, (aaiies,,113))). 31(3):405-412. 2015.
[62] Martin Steinegger and Johannes Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature biotechnology, 10):::020112221 2017.
[63] Oleg Trott and Arthur J Olson. Autodock vina: improving the speed and accuracy of docking with new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, 31(2):455-46 461, 2010.
[64] Jerome Eberhardt, Diogo Santos-Martins, Andreas F Tillack, and Stefano Forli. Autodock vina 1.2. O: New docking methods, expanded force field, and python bindings. Journal of Chemical Information and Modeling, 61(8))3891-389, 2021.
[65] Rodrigo Quiroga and Marcos Villarreal. Vinardo: scoring function based on autodock vina improves scoring, docking, and virtual screening. PloS one, 1 115))e0155183, 2016.
[66] David Ryan Kocs, Matthew Baumgartner, and Carlos Camacho. Lessons learned in empirical scoring with smina from the esar 2011 benchmarking exercise. Journal of chemical information and modeling, 1338))18931194 2013.
[67] Garrett M Morris, Ruth Huey, William Lindstrom, Michel F Sanner, Richard K Belew, David S Goodsell, and Arthur J Olson. Autodock4 and autodocktools4: Automated docking with selective receptor flexibility. Journal of computational chemistry, 30816)27785-2711 2009.
[68] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zeulemoyer, and Veselin Stoyanov. Roberta: A robustly optmiized bert pretraining approach. arXiv preprint arXiv:1907, 1692, 2019.
[69] Hangho Bao, Li Dong, and Furu Wei. Beit: Bert pre-training of image transformers. arXiv preprint 2021.
13

---

Published as a conference paper at ICLR 2023
[70] 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. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10012- 10022, 2021.
[71] Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed EI Kholy, Faisal Ahmed, 7he Gan, Yu Cheng, and Jingjing Liu. Uniter: Universal image-text representation learning. In European conference on computer vision, pages 104-120. Springer, 2020.
[72] 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 on Machine Learning, pages PMLR, 2021.
[73] Yuqi Huo, Manli Zhang, Guangzhen Liu, Haoyu Lu, Yizhao Gao, Guoxing Yang, Jingyuan Wen, Heng Zhang, Baogui Xu, Weihao Zheng, al. Wenlan: Bridging vision and language large-scale multi-moda pre-training. arXiv preprint arXiv:2103.06561. 2021.
[74] Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmicla, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-fille convolutional neural network for modeling quantum interactions. Advances in neural information processing syxtems, 30, 2017.
[75] Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. In International Conference on Learning Representations (ICLR), 2020.
[76] 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 on Learning Representations. 2022.
[77] Jérémy Desaphy, Karima Azdimousa, Esther Kellenberger, and Didier Rognan. Comparison and drugga- bility prediction of protein-ligand binding sites from paarraccophor-annotaee cavity shapes, 2012.
[78] Alexandre Borrel, Leslic Regad, Henri Xhaard, Michel Petitjean, and Anne-Claude Camproux. Pockdrug: model for predicting pocket druggability that overcomes pocket estimation uncertainties. Journal of chemical information and modeling, 55(4):882-895, 2015.
[79] Cheng Wang and Yingkai Zhang. Improving dockkng screening powers of protein ligand scoring functions using random forest. Journal of computational chemistry, 38(3):169 177, 2017.
[80] Hannes Stärk, Octavian- Eetavian-Eugen Ganca, Lagnajit Pattanaik, Regina Barzilay, and Tommi Jaakkola. Equibind: Geometric deep learning for drug binding structure prediction, 2022.
[81] Scott LeGrand, Aaron Scheinberg, Andreas F Tillack, Mathialakan Thavappiragasam, Josh V Vermaas, Rupesh Agarwal, Jeff Larkin, Duncan Poole, Diogo Santos-Martins, Lconardo Solis- Vasquez, et al. Gpu-accelerated drug discovery with docking on the summit supercomputer. Porting, optimization, and application to covid-19 research. In Proceedings of the 11th ACM international conference on bioinformatics. computational biology and health informatics, pages 1-10, 2020.
Solis-Vasques
[82] Teague Sterling and John J Irwin. Zine 15 -igadd discovery for everyone. Journal of chemical information and modeling. 55(1 1)):234 2337, 2015.
[83] Anna Gaulton, Louisa J Bellis, A Patricia Bento, Jon Chambers, Mark Davies, Anne Hersey, Yvonne Light, Shaun McGlinchey, David Michalovich, Bissan Al- Al-Lazikani, al. Chembl: large-scale bioactivity database for drug discovery. Nucleic acids research, 40(D1):D1100 D1107, 2012.
[84] Kaitlyn M Gayvert, Neel S Madhukar, and Olivier Elemento. A data-driven approach to predicting successes and failures of clinical trials. Cell chemical biology. 129401129--133,, 2016.
[85] Matthias Rupp. Alexandre Tkatchenko, Klaus-Rober Müller, and O Anatole Von Lilienfeld. Fast and accurate modeling of molecular atomization energies with machine learning. Physical review letters. 10865), 2012.
[86] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. International Conference on Machine Learning, pages 9558 9568. PMLR, 2021.
[87] Simon Axelrod and Rafacl Bomnez--ombarllli.. Geom, energy- annotaand molecular conformations for property prediction and molecular generation. Scientific Data, 111::1114, 2022.
[88] Bharath Ramsundar, Peter Fastman, Patrick Walters, and Vijay Pande. Deep learning for the life sciences: applying deep learning to genomics, microscopy, drug discovery, and more. O' O'Reilly Media, 2019.
14

---

Published as a conference paper at ICLR 2023
[89] Paul CD Hawkins. Conformation generation: the state of the art. Journal of Chemical Information and Modeling, 57(8):1747-1756 2017.
[90] Ziyao Li. Shuwen Yang. Guojie Song, and Lingsheng Cai. Conformation -ggided molecular representation with hamiltonian neural networks. In International Conference on Learning Representations, 2020.
[91] Minkai Xu, Wujie Wang. Shitong Luo, Chence Shi, Yoshua Bengio, Rafael Gomez-Bombarelli, and Jian Tang. An end-to-end framework molecular conformation generation bilevel programming. In International Conference on Machine Learning, pages (1537-1154) PMLR, 2021.
[92] Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predicting molecular conformation via dynamic graph score matching. Advances in Neural Information Processing Systems, 34, 2021.
[93] Jinhua Zhu, Yingce Xia, Chang Liu, Lijun Wu, Shufang Xie, Tong Wang, Yusong Wang, Wengang Zhou, Tao Qin, Houqiang Li, et al. Direct molecular conformation generation. arXiv preprint 2rXiv:200.001356. 2022.
[94] Minkai Xu, Lantao Yu, Yang Song. Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[95] Dong C Liu and Jorge Nocedal. On the limited memory bigs method for large scale optimization. 1989.
Mathematical programming,
[96] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li. Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings the IEEE/CVF Conference Computer Vision and Puttern Recognition, pages 16000-16009, 2022.
[97] Xinlci Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.0422 042997, 2020.
[98] Ting Chen, Simon Kornllith, Kevin Swersky. Mohammad Norouzi, and Geoffrey E Hinton. Big self- supervised models are strong semi- superviee learners. Advances in neural information processing systems, 33:22243-222 22225,, 2020.
[99] Weitao Du, He Zhang, Yuanqi Du, Qi Meng. Wei Chen, Bin Shao, and Tie- Yan Liu. Equivariant vector field network for many-body system modeling. arXiv preprint arXiv:2/10. 10.1:111,,, 2021.
[100] Sunghwan Kim, Jie Chen, Tiejun Cheng. Asta Gindulyte, Jia Hc, Siqian Hc, Qingliang Li, Benjamin Shoemaker, Paul A Thiessen, Bo Yu, et al. Pubchem 2019 update: improved access to chemical data. Nucleic acids research, 47(DI):D1100 100--11022 2019.
15

---

Published as a conference paper at ICLR 2023
A PRETRAINING DATA
Molecular dataset The pretraining datasets use consist of two parts: one part is database collection of 12 million molecules that can be synthesized and purchased (See Table 5]. and the other part is taken from previous work [23], whose molecules are collected from the ZINC 82) and ChemBL 83 databases. After normalizing and duplicating, we obtain 19 million molecules as our pretraining dataset. To generate 3D conformations,
ETKGG 22 with Merck Molecular Force Field 43 optimization in RDKit 44 to randomly generate 11 conformations for each molecule, totally 209M conformations. In these 11 conformations of molecule, there a special flattened 3D conformation (atoms with zero z-axis coordinates) that is directly from the molecular graph. This flattened 3D conformation is used for the cases where RDKit failed generate 3D conformations, like the peptides in the SIDER task.
Candidate protein pocket dataset The pretraining dataset for candidate protein pockets is derived from the Protein Data Bank (RCSB PDB We first pre-process the data by adding missing side chains and hydrogen atoms, and then use Fpocket [46] to detect candidate binding pockets the proteins. After filtering the raw pockets by the number of residues they have contact with (10-25) and including water molecules inside the pockets, we collect pretraining dataset of 3,291,738 candidate pockets.
P 45 a collection of 180K structural data of proteins.
B DOWNSTREAM DATA SUPPLEMENTS
MOLECULAR PROPERTY PREDICTION
We conduct experiments on the MoleculeNet 522 task. MoleculeNet is a widely used benchmark for molecular property prediction. The details of the 15 datasets we used are described below.
benchmark in the molecular property prediction
BBBP Blood-brain barrier penetration (BBBP) contains the ability of small molecules to penetrate the blood-brain barrier.
BACE This dataset contains the results of small molecules as inhibitors of binding to human 8-secretase (BACE-1).
Clin'Tox This dataset contains the toxicity drug in clinical trials and status the drug for FDA approval 84).
Tox21 The dataset contains toxicity measurements of 8k molecules for 12 targets.
ToxCast This dataset is derived from toxicology data from in vitro high- throughhouuuuttt sereening and contains toxicity measurements for 8k molecules against 617 targets.
SIDER The Side Effect Resource (SIDER) contains side effects of drugs on 27 system organs. These drugs are not only small molecules but also some peptides with molecular weights over 1000.
HIV This dataset contains 40k compounds with the ability to inhibit HIV replication.
PCBA PubChem BioAssay (PCBA) is database of small molecule bioactivities generated by
tightthroughput screening. This a subset containing over 400k molecules on 128 bioassays.
MUV Maximum Unbiased Validation (MUV) is another subset of PubChem BioAssay, containing 90k molecules and 17 bioassays.
ESOL This dataset contains the water solubility of the compound and is a small dataset with 1128 molecules.
FreeSolv The dataset contains hydration free energy data for small molecules, which we use the experimental values as labels.
Lipo Lipophilicity contains the solubility of small molecules in lipids, of which we use the octanol/water distribution coefficient as the label.
hhtp://www.reshhorg/
16

---

Published as a conference paper at ICLR 2023
Table 5: Database collection of 12M purchasable molecules
Database
Molecules
Link
Targetmol
10,000
httpes:/www.arrrstttll..co
Chemdiv
1,613,931
https:/ www http:://wwwcwccwccccccccc com/
Enamine
2,734,581
Pttps://enamine... net/
Chembridge
1,557.942
https://ww..heebrriiggg com/
Life Chemical
509,975
Specs
208,670
htDp://ww..3pcc.nett/
Vitas-M
1, 1,409,339
http:://vtteemaaab.bbb7
InterBioScreen
48,627
http:///ww..baccrct COm/
Maybridge
53,352
https://ww.chhemtff55555 in/
Bionet-Key Organics
259,244
http:://ww..eyorgan..com
Asinex
530,881
https://ww..ash com/
UkrOrgSynthesis
688,952
ittps://uorsy.cm
Eximed
61,009
ittps://eximedlab..
HTS Biochemie Innovationen
58,437
httpsr//ww.tts-bloceee.. de/
Princeton BioMolecular
1,532,542
http:://pprineeecnno0..ooo
Otava
270,835
htpps//otaarcchen1cc11..
Alinda Chemical
202,332
https://ww..liidda.n
Analyticon
42,664
hww.eicooooiagnnstiiooioo dwwgnostics.co com/
QM7, QM8, QM9 The molecule in QM7 contains up heavy atoms, QM8 is 8 and QM9 is 9. These datasets provide the geometric, energetic, electronic and thermodynamic properties of the molecule, which calculated by density functional theory (DFT) 85]. QM9 contains several quantum mechanical properties of different quantitative ranges, and we select homo, lumo and gap of similar quantitative range, following the setup of the previous work
B.2 MOLECULAR CORFORMATION GENERATION
Following the settings in previous works 49 dataset in this task.
we use GEOM-QM9 and GEOM-Drugh
GEOM This dataset contains 37 million accurate conformations generated for 450,000 molecules by advanced sampling and semi-empirical density functional theory (DFT). Of these, 133,000 molecules are from QM9, and the remaining 317,000 molecules have biophysical, physiological, physical chemistry experimental data, i.c., Drugs
B.3 POCKET PROPERTY PREDICTION
NRDLD benchmark dataset for pocket druggability prediction. As NRDLD and other existing benchmark datasets are too small, we construct a regression dataset to benchmark pocket property prediction performance.
NRDLD The dataset contains 113 proteins, and predefined split is provided: 76 proteins constitute the training set and 37 proteins constitute the test set. labels 71 proteins as druggable that they noncovalently bind small drug-like ligands [59). The rest 42 proteins are labeled as less-druggable because none the ligands they cocrystallized satisfy the following requirements simultaneously:
the rule of five, clogP -2, and ligand efficiency, defined in 22]] 0.3 kcal mol heavy atom.
Our created benchmark dataset The dataset contains 164,586 candidate pockets, and four scores (Fpocket Score, Druggability Score, Total SASA, and Hydrophobicity Score) calculated by the Fpocket tool. These four scores are indicators of the druggability of candidate pockets. To avoid leaking, the selected pockets are not overlapped with the candidate protein pocket dataset used in Uni-Mol pretraining.
17

---

Published as a conference paper at ICLR 2023
Table Uni-Mol hyperparameters setup during pre-training
Hyperparameter
Molecular pretraining
Pocket pretraining
Layers
15
15
Peak learning rate
le-4
le-4
Batch size
128
128
Max training steps
1M
1M
Warmup steps
10K
10k
Attention heads
64
64
FFN dropout
0.1
Attention dropout
0.1
0.1
Emhedding dropout
0.1
0.1
Weight decay
le-4
le-4
Emhedding dim
512
512
FFN hidden
2048
2048
Gaussian kernel channels
128
128
Corrupt ratio
0.15
0.15
Activation function
GELU
GELU
Learning rate decay
Linear
Linear
Adams é
lc-6 (0.9, 0.99)
1c-6 (0.9, 0.99)
Adams (3)
Gradient clip norm Atom loss function and its weight Coordinate loss function and its weight Distance loss function and its weight Vocabulary size (atom types)
1.0 Cross entropy, 1.0
1.0 Cross entropy, 1.0
Smooth L1, 5.0
Smooth LI, 1.0
Smooth L1, 10.0
Smooth LI, 1.0 9
30
B.4 PROTEIN-LIGAND BINDING POSE PREDICTION
We use PDBbind General set v.2020 61
excluding the similar complexes in CASF-2016 [28] as the training set. And CASF-2016 used as the test set. particular, we define the pocket for each protein-ligane pair as residues the protein which have least one atom within the range of 6Å from a heavy atom in the ligand. All atoms of the selected residues are included. In addition, we draw smallest bounding box covering
the atoms in the pocket and regard the water molecules in the bounding box as a part of the pockets, too.
PDBbind General set v:2020 This dataset contains 19,443 protein-ligand complexes with binding data and processed structural files originally from the Protein Data Bank (PDB). Only complexes with experimentally determined binding affinity data are included in the general set. Notably, to examine the generalization ability, we further filter out the training complexes that are similar the ones in the test set (CASF-2011 -016).. particular, the complexes with both high protein sequence sim- ilarity (MMSeqs2 [62] similarity above 40%) and high molecular similarity (fingerprint similarity above 80%) are filtered out, and there are 18,404 complexes after filtering.
CASF-2016 CASF-2015 is the widely used benchmark for docking and scoring. This dataset, whose primary test set is known the PDBbind Core set, contains 285 protedn-ligand complexes with high-quality crystal structures and reliable binding constants from PDBbind General set. For each protein-ligand complex, CASF-2016 provides 50- 10-100 decoy molecular conformations of the same ligand for evaluation.
C EXPERIMENTS DETAILS & REPRODUCE
C.I MOLECULAR PRETRAINING SETUP
We report the detailed hyperparameters setup Uni-mol during pretraining in Table 6 Uni-Mol training loss is summed up by three components, atom (token) loss, coordinate loss, and pair-distance loss. Atoms are masked, and noise added to coordinate as described in sections 2.1 and 2.2. Since the values of the above three components differ significantly, make them have a similar influence, we enlarge the coordinate loss and distance loss. Molecular pretraining runs on 8 V100 GPUs (32GB memory, the same below), and the training time is about 20 hours.
18

---

Published as a conference paper at ICLR 2023
Table Search space for small datasets: BBBP, BACE, ClinTox, Tox21, Toxcast, SIDER, ESOL, FreeSolv, Lipo, QM7, QM8, for large datasets: PCBA, MUV, QM9, and for HIV
Hyperparameter
Small
Large
HIV
Learning rate
[5c-5, 8c-5, Ic-4, 4c-4, 5c-4]
[2c-5, lc-4]
[2e-5, 5c-5]
Batch size
[32, 64, 128, 256]
[128, 256]
[128, 256]
Epochs
[40 ,60, 80, 100]
[20, 40]
[, 10]
Pooler dropout
[0.0, 0.1. 0.2, 0.5]
[0.0. 0.1]
[0.0, 0.2]
Warmup ratio
[0.0, 0.06, 0.1]
[0.0, 0.06]
[0.0, 0.1]
C.2 POCKET PRETRAINING SETUP
The pocket Uni-Mol model is slightly different from molecule ones during pretraining: 1) We use residue-level masking strategy instead of the original atom-level, residue granularity is non- redundancy and integrity in protein. 2) Hydrogen is removed in pocket Uni-Mol pretraining, to reduce the number of used atoms and thus improve efficiency. 3) All weights of loss functions are set 1, as the residue-level masking strategy makes the 3D denoising task much harder. 4) A different vocabulary is used in pocket pretraining. In pocket data, there are amino acids, whose atoms are mostly C, N, O, S and H. While in molecule data, the atom types are more diverse, so a larger vocabulary is used, Other settings are listed in 6 Pocket pretraining runs on V100 GPUs, and the training time is about 2 days and 20 hours.
C.3 MOLECULAR PROPERTY PREDICTION
Data split In our experiments, referring previous work GEM 13 to divide the dataset into training, validation, and test sets in the ratio of 8:1:1. Scaffold splitting is more challenging than random splitting as the scaffold sets of molecules different subsets do not intersect. This splitting tests the model's generalization ability and reflects the realistic cases Since this splitting is according to the scaffold of the molecule, we find that whether or not chirality is considered when generating the scaffold using RDKit has significant impact on the division results. From the results, the splitting considering chirality makes the task harder. The original implementation of MoICLR does not consider chirality, and we reproduce the experiment by considering it. In all experiments, we choose the checkpoint with the best validation loss, and report the results on the test-set by that checkpoint.
scaffold splitting (88
Hyperparameter search space Referring to previous works, we use a grid search to find the best combination of hyperparameters for the molecular property prediction task. To reduce the time cost, we set a smaller search space for the large datasets. The specific search space is shown in Table For small datasets, we run them on single V100 GPU; for large datasets and HIV, we run them on 4 V100 GPUs.
C.4 MOLECULAR CONFORMATION GENERATION
Datasets and setup Following the settings in previous works [6). we use GEOM-QM9 and Crugs 87 dataset to perform conformation generation experiments. As described in Sec. [2.3] in this task, Uni-Mol optimizes its input generative conformations to the labeled ones. To construct the finetuning data, we leverage RDKit to generate input conformations. For each input, 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 (twice the number of labeled conformations) as previous works do. And we also use the same metrics, Coverage (COV) and Matching (MAT) as previous works. Higher COV means better diversity, while lower MAT means higher accuracy.
Metrics this task, following previous work 889 90], we use the Root of Mean Squared Deviations (RMSD) of heavy atoms to evaluate the difference between the generated conformation and the ref- erence one. Before computing RMSD, the generated conformation first aligned with the reference one, and the function & aligns conformations by applying rotations and translations to them:
RMSD R, R.
min +
| (R.)
(4)
19

---

Published as a conference paper at ICLR 2023
where R and R the generated and reference conformation, i is the i-th heavy atom, and is the number of heavy atoms.
We use Coverage (COV) and Matching (MAT) to evaluate the performance of the conformation generation model. Higher COV means better diversity, while lower MAT means higher accuracy. Formally, COV and MAT are denoted as:
R S,IRMS R, R) < R Sg
cOV(S,,
(5)
MA |S+
min RMSD R)
(6)
RES, ReS,
where Sg and S, are the set of generated and reference conformations, respectively, and 5 is given RMSD threshold. Following previous work [49 86), for GEOM-QM9, the threshold is 0.5À, and for GEOM-Drugs, the threshold value is 1.25A.
Baselines We compare Uni-Mol with 10 competitive baselines. RDKit is a traditional conformation generation method based on distance geometry. The rest baseline can be categorized into two classes. GraphDG [48]. 4GCF ConfVNE [9I]. ConfGF [86], and DGSM [92]
combine generative models with distance geometry, which first generates interatomic distance matrices and then iteratively generates atomic coordinates. CVGAE (50] GeoMol [5]]
DMCG 193), and GeoDiff [94] directly generate atomic coordinates.
Results The results are shown in Table 8 We report the mean and median of COV and MAT on GEOM-QM9 and GEOM-Drugs datasets. ConfVAE [9]] and DGSM's [92] results are from their papers, respectively. GeoMol results are from DMCG 993 paper, which DMCG reproduced after aligning the data split. We found the test sets used GeoDiff [94] and DMCG [93] are
slightly different from baselines (different filtering conditions), so we use the released model parameters of GeoDiff to reproduce the results on the same test set. For DMCG, as model parameters are not released, we use its open-source codes to reproduce from scratch. Other baseline results are from ConfGF's paper. As shown in Table g Uni-Mol exceeds existing baselines in both COV and MAT metrics on both datasets.
Experiments details We report the detailed hyperparameters setup for molecular conformation generation in Table And we run this task a single V100 GPU. Since this is 3D-related task, we only use coordinate loss and distance loss. Following AlphaFold (38), we use "recycling" to iteratively refine the output atom positions. We leverage RDKit (ETKGD) for generating inputs in molecular conformation generation tasks. Specifically, in finetuning, we randomly generate up to 2000 conformations and cluster them into conformations, as the model input. A similar pipeline is used in the inference of test data. For most baselines, as they aim to generate conformations from scratch, RDKitt-enerated conformations are not leveraged.
C.5 POCKET PROPERTY PREDICTION
The hyperparameters we search are listed in Table 10 We this task single V100 GPU. When we use the Fpocket [46] tool in our experiments, it outputs four values: Fpocket score, Druggability score, Total SASA, and Hydrophobicity Score. Specifically, the Fpocket score custom score by Fpocket; the druggability score is an empirical score calculated from evolution and homologous information. Besides, verify the effectiveness of the Fpocket tool on real-world data, we test this tool on NRDLD. Table 11 shows the performance of Fpocket tool on NRDLD dataset.
C.6 PROTEIN-LIGAND BINDING POSE PREDICTION
Data split The training set PDBbind General set v.2020 excluding the similar complexes covered by CASF-2016. We perform data preprocessing, such as adding missing atoms to both proteins and ligands and manually fixing file-loading errors, before constructing the training set. And we filter the complexes based the number of residues contained in the pockets (>= 5). Besides, filter similar complexes in PDBbind with CASF-2016 in consideration of protein sequence and ligand similarity. MMseqs2 (62] similarity above 40% and RDKit FingerprintSimilarity [44) above 80% are used as filtering conditions. The test set CASF-2016, which contains 285 protein-ligann complexes.
20

---

Published as a conference paper at ICLR 2023
Table 8: Uni-Mol performance on molecular conformation generation. Please note the previous conformation generation models were not pretrained.
Dataset
QM9
Drugs
COV(1,
M. Å)
COV(), %)
MAT(J, Á)
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
ConLVAE
80.42
85.31
0,4066
0.3891
53.14
53.98
1.2392
1.2447
ConIGF
88.49
94.13
0.2673
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
GeoDiff
92.65
95.75
0,2016
0.2006
88.45
97.09
0,8651
0.8598
DMCG
94.98
98.47
0,2365
0,2312
91.27
100.00
0.8287
0.7908
Uni-Mol no pro.ptetaieedd
97.00
100.00
0.1907
0.1754
91.68
100.00
0.8102
0.8041
Uni-Mol
97.95
100.00
0.1831
0.1659
91.91
100.00
0.7863
0.7794
Table Hyperparameters setup for molccular conformation generation
Hyperparameier
QM9
Drugs
Learning rate
20-5
5c-5
Batch size
128
128
Epochs
50
70
Warmup ratio
0.06
0.06
Recycling times
4
4
Coordinate loss function and weight
MSE, MSE, 1.0
Distance loss function and weight
MSE, 1.0
Model architecture and finetuning As shown in Figure 4 the binding pose model is encoder- decoder architecture consisting of two 15-layer Uni-Mol as encoders and a 4-layer Uni-Mol as decoder. The decoder block follows the same setting the pretraining ones. During finetuning, Uni-Mol is trained towards to cross distances (between pocket and ligand) and holo distances (pair distances in ligand itself) with MSE loss. The hyper-parameters used in fine-tuning are listed in Table [2] The Dist threshold is introduced to alleviate the large MSE loss from long-range atom pairs and thus stabilize the training. particular. the atom pairs with distances larger than Dist_threshen threshold are ignored in the loss calculation. Following AlphaFold (38, we use "recycling" to iteratively refine the output atom positions. We finetune model on V100 GPUs.
Binding pose prediction As aforementioned, Uni-Mol does not directly predict the binding pose in finetuning. We use scoring function based method to optimize the input coordinates. We first construct a scoring function by the difference between the current pair-distance matrix and the predicted pair-distance matrix. Given the coordinates of pocket atoms and ligand atoms, we can compute its pair-distance matrix. We use the mean square error to get the difference between the current pair-dite mair-iitaance matrix and the predicted one from the model. Notably, listed in Table [2] different weights are used in cross distances and holo distances calculating mean square error. Then, to get the binding pose, directly apply pack-propagttion from the scoring function to
Table 10: Search space for pocket property prediction
Hyperparameter
NRDLD
Our Created
Learning rate
[5e-5, le-4, 3e-4] [1, 4, 8, 16]
3e-4
Batch size
32
Epochs
40
20
Pooler dropout
[0, 0.1, 0.2, 0.31
0
Warmup ratio
[0.0. 0.1]
21

---

Published as a conference paper at ICLR 2023
Table 11: Performance of Fpocket tool on NRDLD
Accuracy
Recall
Precision
Fl-score
Fpocket Score
0.73
0.83 0.83
0.76 0.83
0.79
Druggability Score
0.78
0.83
Uni- Mai Block
15 Layers
repremmeemenntt representicion
BBC BI
Lancatemated storm rapresrraaaioo
Complex stor representaaioon
Parresrceenttonn
Uni-AAi Aii-MMll Mol Block
Layers
Uni Moi- Block
Iile DOM IM
ayerss
reporstaaiionn
representation
Camplee port reprrsanttaioon
Foi repeossiation
recycling
Figure 4: protein-ligand binding pose model: 1) Encoder: molecular representation and pocket representation obtained from their own pretraining Uni-Mol models; 2) Decoder: representation concatenated with atom and leir- as inputs layers Uni-Mol block learning from scratch with recycling. 3) Output: The complex representation is used project layer to learn the pair distances of molecule and pocket.
input coordinates. We choose use LBGFS 95 with learning rate of 1.0 as the optimizer for efficiency. We use 10 randomly generated molecular conformations with randomly placed positions (6 degrees of freedom) as data augmentation, and choose the one with the lowest loss after back-propagation the final result.
Exhaustiveness search for docking tools To ensure that the comparison between Uni-Mol and popular molecular docking software unbiased, we increase the exhaustiveness of the global search (roughly proportional to time) of the molecular docking software to observe the effect of computational complexity on docking power on CASF-2016 benchmark. shown Table T3] when exhaustiveness is above 16. the popular molecular docking software longer improve the performance by increasing the computational complexity.
D ABLATION STUDY
D.1 3D SPATIAL POSITIONAL ENCODINGS BENCIMARK
We investigate the performance of different 3D spatial positional encodings on the 3D molecular pretraining. In particular, benchmarked:
Gaussian kernel (GK), a simply Gaussian density function.
Gaussian kennel with pair type (GKPT) 39 Based on GK, an affine transformation according to the pair type is applied on pair distances, before applying the Gaussian kernel.
Radial Bessel basis (RBB) [75]. A Bessel based radial function.
Discretization categorical embedding (DCE). We convert the continued distances to the discrete bins, by Discretization. With binned distances, embedding-based positional encoding directly used.
Delta coordinate (DC) 36.. Following Point Transformer [6], the deltas of coordinates are directly used as pair-wise spatial relative positional encoding.
22

---

Published as a conference paper at ICLR 2023
Table 12: Hyperparameters setup for binding pose prediction
Hyperparameters for finetuning
Value
Learning
3e-4
Batch size
Epochs
50
Warmup ratio
0,06
Dropout
0.2
Dist threshold
Recycling times
Cross distance loss function and weight
MSE, 1.0
Holo distance loss function and weight
MSE, 1.0
Hyperparameters binding pose prediction
Value
Optimizer
LBFGS
Learning rate
Max iterations
100
Early stopping
5
Dist threshold
4.5
Conformation size
Cross distance weight
1.0
Holo distance weight
5.0
Table 13: Exhaustiveness study of popular docking tools on CASF-2016
Ligand RMSD
% Below
Threshold t
Methods
Exhaustiveness
0.5 Á
1.0 Á
1.5 Á
2.0 À
Autodock Vina
21.40
35.79
47.02
52.28
Autodock Vina
8
23.86
44.21
57.54
64.56
Autodock Vina
16
25.61
45.96
60.70
66.67
Autodock Vina
32
25.96
45.96
60.00
66.32
Vinardo
16.84
33.33
43.16
49.82
Vinardo
8
23.51
41.75
57.54
62.81
Vinardo
16
23.51
45.26
60.70
66.67
Vinardo
32
23.86
44.56
59.30
65.61
Smina
1
23.51
39.65
50.53
56.14
Smina
8
23.51
47.37
59.65
65.26
Smina
16
28.77
49.47
61.40
67.72
Smina
32
28.07
51.23
61.75
67.37
Autodock4
1
4.91
18.95
26.67
28.87
Autodock4
8
7.02
21.75
31.58
35.44
Autodock4
16
6.32
24.56
34.04
38.95
Autodock4
32
6.32
23.16
34.04
38.25
23

---

Published as a conference paper at ICLR 2023
Gaussian kennel with pair type and local graph (GKPTLG). Based GKPT, set up model with locally connected graphs. In particular, the cutoff radius is set to 6 Á.
We summarize the validation loss during pretraining for them in Fig. From the results, we can find:
The performance of DCE and GK are almost the same, and outperform RBB and DC. And we choose GK as the basic encoding.
Compared with GK, GKPT convergences faster. This indicates the pair type is critical the 3D spatial positional encoding.
Compared with GKPT, GKPTLG convergences slower. This indicates the locally cutoff graph is not effective for self-supervised learning, and the default fully connected graph in Transformer is more effective.
As GKPT outperforms all other encodings, we use it in the backbone model of Uni-Mol.
2.00
1.75
1.50
Valid Loss
1.25
CKPT
GKPTLG
1.00
0.75
RBB
0.50
0.25
0.00
25000
50000
75000
100000
125000
150000
175000
200000
Steps
Figure 5: Validation loss in pretraining for different 3D spatial encodings
Table 14: Ablation studies, molecular property prediction classification tasks
Classification (ROC-AUC 5;, higher is better 1)
Datasets
BBBP
BACE
ClinTox
Tox21
ToxCast
SIDER
HIV
PCBA
MUV
Uni Mol w/o pretraining
69.0(0.7)
80.9(5.4)
68.3(2.2)
75.8(0.4)
63.8(0.1)
61.9(0.5)
76.2(2.4)
86.1(0.5)
62.8(4.0)
Uni Mol wlo pair representation
71.6(1.3)
85.4(2.7)
85.5(1.7)
79.4(0.1)
69.3(0.1)
64.3(0.9)
80.2(0.2)
88.4(0.1)
71.0(7.7)
2D shortest path encoding
71.6(2.1)
85.6(1.1)
83.614.0)
79.6(0.7)
68.8(0.8)
63.7(0.1)
78.9(0.4)
88.0(0.2)
78.2(0.6)
relative positional encoding
70.3(1.))
77.8(3.7)
73.3(.77
64.9(0.2)
61.5(1.)
75.6(0.3)
77.2(1.4)
68.7(1.0)
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.))
80.8(0.3)
88.5(0.1)
82.1(1.3)
Table 15: Ablation studies, molecular property prediction regression tasks
Regression (lower is better L)
MAE
Datasets
ESOL
FreeSolv
Lipo
QM7
QM8
QM9
Uni-Mol w/o pretraining
0.924(0.037)
1.800(0.206)
0.745(0.012)
45.2(0.6)
0.0174(0.0002)
0.00653(0.000400)
Uni-Mol w/o pair representation
0.807(0.027)
1.582(0.068)
0.611(0.004)
45.2(1.0)
0.0158(0.0001)
0.00573(0.0000
2D shortest path eacoding
0.831(0.007)
1.482(0.070)
0.605(0.003)
60.6(0.2)
0.0164(0.0001)
0.00650(0.00001)
ID relative positional encoding
0.929(0.035)
2.237(0.074)
0.866(0.004)
77.5(2.7)
0.0283(0.0007)
0.02283(0.000) (0.00078)
Uni-Mol
0.788(0.029)
1.480(0.048)
0.603(0.010)
41.8(0.2)
0.0156(0.0001)
0.00467(0.00904)
24

---

B1 in Published as a conference paper at ICLR 2023
D.2 PAIR REPRESENTATION BENCHMARK
Pair representation is a new component proposed in Uni-Mol for enhancing 3D spatial positional encodings. To check its effectiveness, we conduct an ablation benchmark. We add setting that only uses spatial positional encoding as attention bias in attention layers, and the attention weights of the last layer are used for the coordinate update and distance prediction. The downstream molecular property prediction results, shown in Tables [4] and [5] demonstrate that the one without pair representation worse. Besides, according to the efficiency benchmark in Table 17, we can find the extra cost introduced by pair representation is very small. This result demonstrates the efficiency of the proposed pair representation.
D.3 SE(3) COORDINATE HEAD BENCHMARK
As aforementioned in Sec. 2.1 in SE(3) coordinate head, we use delta pair representation to calculate the coordinate updates, while previous work EGNN [41] did not. Here we conduct benchmark on molecular conformation generation. As shown in Table T6 Uni-Mol's SE(3) head outperforms the one in EGNN. The result indicates the effectiveness of the proposed SE(3) coordinate head.
Table 16; Benchmark on SE(3)ecquivariant coordinate head.
Dataset
QM9
Drugs
COv(t. %)
MAT(L, A)
COV(1. %)
MAT( A)
Methods
Mean
Median
Mean
Median
Mean
Median
Mean
Median
Uni-Mol w/ EGNN head
96.94
100.00
0.1873
0.1730
91.93
98.80
0.7952
0.7821
Uni-Mol
97.95
100.00
0.1831
0.1659
91.91
100.00
0.7863
0.7794
D.4 EFFICIENCY BENCHMARK
Due to the massive pretraining cost in the large-scale data, we favor the efficient backbone model. And the backbone model Uni-Mol designed based on the principle. To demonstrate the efficiency of our designed backbone, compared it with vanilla Transformer and terms of speed. We use the same benchmark dataset, QM9, in to test the training speed, and summarize the results in Table 17 We can that Uni-Mol backbone model is only slightly slower than the vanilla Transformer and much faster than SE(3)- Transformer(and NVIDIA optimized version). The result explicitly demonstrates that design backbone model is indeed very efficient.
NVIDIA version
Table 17: Efficiency benchmark, we test vanilla Transformer and add our components based on it. Besides, SE(3) Transformer (31) is also benchmarked. "seconds" the average cost seconds for batch. "ratio" is the slowdown ratio compared with vanilla Transformer.
Methods
Params.
seconds
ratio
vanilla Transformer
47.57M
0.062
1.00
vanilla Transformer + spatial position encoding
47.59M
0.063
1.02
vanilla Transformer spatial position encoding pair repr.
47.60M
0.064
1.03
Uni-Mol
47.6
0.066
1.06
SE(3)- TransformerpиMae
31
46.77M
14.776
238.32
SE(3)- Transssrrrofoooeeeeeddo
47.02M
0.388
6.26
D.5 PRETRAINING STRATEGIES BENCHMARK
We investigate and benchmark several pretraining strategies with Uni-Mol framework.
Masked autoencoders (MAE). Following MAE 96,, the uncorrupted atoms are used as input of the encoder, while corrupted atoms are not. Then decoder accepts both uncorrupted and corrupted atoms as input. The MAE proven effective in image-related tasks.
25

---

Published as a conference paper at ICLR 2023
Masked atom prediction only (MAP). This setting only masked atom prediction, without 3D position recovery.
MAP+3D. The setting in Uni-Mol, with MAP and 3D position recovery.
MAB=3DD+LL(Cnnrrstti Learning). A contrastive loss 197 98) is further added for the 3D conformations. In particular, two conformations are treated as "positive" they are from the same molecule; otherwise, they are treated as "negative." The loss is combined together with 3D position recovery and masked atom prediction tasks.
Besides the above strategies, we also add two additional methods. 1) Single Conformation. This method only uses one 31) conformation for each molecule, examine the performance brought by more 3D conformations. 2) Local-Connected Graph. This method uses the local-connected graph, by radius cutoff 6 Å, to examine the performance brought by the fully connected graph in Transformer.
The results are shown in Table 18 and we can find the pretraining strategies in Uni-Mol (MAP+3D) is very effective, outperforming the other settings.
From the result of Single Conformation, we can find that more 3D conformations can slightly boost the performance of the molecular property prediction tasks. As for the results of the Local-Connected Graph, we can find the performance brought by the fully connected graph in Transformer is quite significant.
Table 18: Pretraining strategies benchmark
ROC-AUC higher hetter
RMSFJ
MARI
BBBP
BACE
ClinTox
SIDER
ESOL
FreeSolv
Lipo
QM7
MAE
71.7(0.6)
86.4(0.6) 82.7(3.2) 63.8(0.)) 0.850(0.008)
1.586(0.034)
0.605(0.008)
45.7(0.7)
MAP
71.3(0.3)
86.1(0.5) 90.2(2.8) 62.5(1.1) 0.835(0.015)
1.562(0.035)
0.595(0.010)
46.6(0.2)
MAP+3D (Uni-Mol)
72.9(0.6)
85.7(0.2) 91.9(1.8) 65.9(1.3)
0.788(0.029)
1.480(0.048)
0.603(0.010)
41.8(0.2)
MAP+3D+CL
72.4(0.9)
84.2(0.6)
88.8(0.7) 65.2(0.9)
0.846(0.026)
1.499(0.020)
0.604(0.004)
44.6(2.7)
Single Conformation
73.6(0.7)
85.8(1.0)
91.3(28) 65.6(0.6)
0.796(0.022)
1.479(0.051) 0.607(0.007)
43.7(0.8)
Aucal-Connected Graph 70.2(1.1)
87.5(0.6)
87.5(388) 64.2(1.7) 0.826(0.029)
1.580(0.050)
0.593(0.011)
47.3(1.2)
D.6 CORRUPTED POSITION GENERATION AND ASSIGNMENT
To generate corrupted positions for corrupted atoms, we proposed Alg which first generates random positions based on ground-truth ones, plus with noises; then, an optional permutational re-assignment is used to further reduce the delta positions. Since there is a tradeoff between noise range and re-assignment, we benchmarked several settings, and list the result in Table I9
From the result, we can find:
"r=1A" A" is better than "r=lÅ, A, re-assign". This indicates, when is small, the re-assignment not needed.
"r=1.5A, re-assign" is better than "r=1.5A", This indicates, when r large, the re-assignment needed.
when r is large, e.g. 2Å, the final performance is not good.
"r=1Å" performs best in our benchmark, and we use as our final setting for corrupted position generation.
Table 19: Corrupted Position Generation and Assignment
ROC-AUC %. higher is better f
RMSEL
MAEL
BBBP
BACE
ClinTox
SIDER
ESOL
FreeSolv
Lipo
[Q7
2A. re-assign
70.7(1.1)
84.9(1.4) 87.4(0.8) 65.1(1.2)
1.557(0.013) 0.82700.116)
0.601(0.017)
43.1(0.5)
>=1 .A, re-assign
72.2(0.3)
85.0(1.3) 90.4(3.9) 64.6(0.5) 0.826(0.019)
1.526(0.072)
0.599(0.007)
43.8(1.1)
1= 1.SA
71.7(0.9)
85.0(1.5)
88.3(34) 64.4(1.2) 0.819(00030)
0.588(0.011) 1.884(0.072)
43.8(0.7)
7= 1. re-assign
72.1(3.2)
86.5(1.2)
88.0(0.))
64.6(2.0)
0.825(0.012)
1.613(0.092)
0.600(0.008)
44.9(0.2)
>=
72.9(0.6)
85.7(0.2)
91.9(1.8)
65.9(1.3) 0.788(0.029)
1.480(0.048)
0.603(0.010)
41.8(0.2)
26

---

Published as a conference paper at ICLR 2023
D.7
PRETRAINING OR NOT
From Tables and 13 is clear that pretraining can significantly boost the downstream performance.
D.8 ID/2D POSITIONAL ENCODINGS
We also investigate the impact of 3D spatial positional encoding on molecular property prediction tasks. Specifically, to demonstrate the effectiveness of introducing 3D information, we replace the original invariant spatial position encoding with a 2D Graphhormer-ikee [4] shortest path positional encoding and a ID BERT-like relative positional encoding on atoms. The results are summarized in the following table. Tables 14 and 13 show the results of the ablation studies. We can find that 3D spatial positional encoding largely improves the performance of molecular property prediction. is clear that 3D information indeed helps the performance of downstream tasks.
and
MORE RESULTS FOR BINDING POSE PREDICTION
Efficiency benchmark We compare Uni-Mol binding pose prediction with popular docking tools in efficiency. The baseline results taken from EquiBind [99] paper.And Uni-Mol binding pose prediction is run on a single V100 GPU, For each molecule, Uni-Mol run with 10 different initial conformations, and the total time cost is reported. As shown in Table 20 Uni-Mol is significantly faster than traditional docking tools, about 250x faster.
Table 20: Efficiency binding pose prediction. The average seconds per ligand reported.
Methods
QVINA-W
GNINA
SMINA
GLIDE (C.)
Uni-Mol
seconds
49
247
146
1405
0.2
Visualization We show protein-ligand binding pose prediction in CASF-2016 test dataset in Figure [6] Green molecules are Uni-Mol predictions while red ones are the ground truth in complexes. From the Figure, we can find that Uni-Mol can predict the accurate binding complexes, with large overlapping with ground-truth ligands.
Neljok]jneeed cone| (reeeceeeeeeeeeeeeeeeeeeeee RMSD: 0.379
PDBID: Sudi
Smilles:
CCCCIICNIICSeEHHOOCCCCCCCCC.oo. Ojeleecnnewcee
RMSD: 1.928
1:ff
Smillew www..iXCCCSABaaaapcceeeelleeeeee cleccel)e
Figure 6: Binding pose prediction visualization in CASF-2016
27

---

Published as a conference paper at ICLR 2023
Ligand conformation performance As shown in Table Uni-Mol cannot outperform popular docking tools in RMSD below 1.0A, After investigation, we found it is due to the RMSD of ligand itself is not good, as shown in Table[11 We hypothesize that the binding pose prediction in Uni- Mol is not end-to-end, and gradient descent based optimization is not physics/chemical aware, so may generate inaccurate ligand conformations. In contrast, popular docking tools are mostly physics/chemical aware, like sampling ligand conformations by enumerating rotatable bonds [65]. To tackle this, a simple workaround is to combine the Uni-Mol and docking tools: Uni-Mol mostly focuses on binding positions, while docking tools focus on ligand conformation. Considering the physics/chemical constraints and predicting binding pose end-to-end are also worthy to try. We leave the further optimizations as future work.
Table 21: Ligand-only conformation performance on CASF-2016
Ligand RMSD
% Below Threshold f
Methods
0.5 Á
1.0 Á
2.0 Á
Uni-Molyando
14.39
57.54
89.82
Uni-Mol
17.20
57.19
93.33
F TRAINING STABILITY
With Pre-LayerNorm 332] backbone and mixed-precision training, the pretraining sometimes diverges. After investigation, we found there are large numerical values in the intermediate states when divergence happens. We hypothesize that the Final-LayerNorm layer in the Pre-LayerNorm backbone results in the problem. Specifically, Final-LayerNorm is applied to the sum of all encoder layers, denoted as
o, = LayerNorm(8) 8; = o
where is the number of layers, of is the output of the i-th position in the I-th layer, and o the final output of the i-th position, after Final-LayerNorm. Therefore, due to normalization, s; can be arbitrarily large (or arbitrarily small), without affecting model results. However, too large or too small numerical value will cause the numerical unstable, especially in the mixed-precision training. To tackle this, we introduce a simple loss, to restrict the value range of 8; Formally, the loss denoted as
Cuorm meaa; max (1ss:I--
- Vd T, T.0)).
8
where d is the dimension size of 8;, T is the tolerance factor. In Uni-Mol, we set 1, and both atom-level and pair-level representations are constrained by this loss. Besides, to avoid affecting other loss functions, we set a very small loss weight (0.01) to Cnorm-
G SELF-ATTENTION MAP VISUALIZATION
For better interpretability, conduct a visualization on the self-attention map and pair distance of the molecule, shown in the Figure 7 We can easily find the average weight of all attention heads (in the right-most figures) is very similar to the distance matrix. However, when check these heads independently, there are different patterns, and are different with distance matrix. For example, most attention weights are asymmetric, while the distance matrix is symmetry. Besides, the long-range range interaction could be captured by the attention (like in Head 11 and 13). The learned long-range interaction patterns reflect our motivation for using Transformer as the backbone model.
28

---

Published as a conference paper at ICLR 2023
eesigtt head
atth weight of head
waight head
atti waiggtt haldd
welight of head
atto weight of head
attn weight of head
pitn weight head
atto wesightt of head
atto weight of head
atto weight of head 10
attn weigtt of head
Mttn weight head 12
ottn weight head
aito weight of head 14
attn weight head
Figure 7: Visualization on self-attention map of multi heads independently.
H VISUALIZATION FOR PRETRAINING TASKS.
We visualize 3D position recovery ability in Uni-Mol pretraining process. As shown in Figure 8 Uni-Mol recovers coordinates and atoms with almost no mistakes in conformational space.
noise molecule
target molecule
recovery molecule
noise target
Figure 8: Visualization of 3D position recovery in pretraining.
29

---

Published as a conference paper at ICLR 2023
Moreover, plot the pretraining masked atom predictions accuracy and the coordinate recovery loss of the validation set. As shown in Figure D during pretraining, Uni-Mol can accurately predict the atom type of the masked atoms and recover the corrupted 3D positions.
0,040
Hareed
0.2
0.000
0-88
Accuaacy
0.600
0.500
0.00
035
8.0
0.000
Stepo
Steps
(a) Accuracy of masked atom prediction
(b) Loss of coordinate recovery
Figure 9: Validation losses in pretraining.
I COMPARISON WITH GRAPHORMER AND DD-GRAPHORMER
The positional encoding (the shortest path) used in Grahpormer only handle 2D molecular graphs, not 3D positions. In the Uni-Mol backbone, we added several modifications to make the model have the ability to handle 3D inputs and outputs. Further, there is concurrent following-up work called 3D-Graphormer (35. adapting Graphormer to 3D molecules. There are several differences between 1) Both Uni-Mol and 3D-Graphormer use the pair-wise Euclidean distance and Gaussian kernel to encode 3D spatial information. 3D-Graphormer has an additional node-level centrality encoding, which the sum of spatial encodings of each node. 2) 3D-Graphormer doesn't have pair representation. 3) Our SE(3) Coordinate Head is different from the "node-level projection head" in 3D- Dr-Grophormrr. The method used 3D- -r-arrphrrmer is an attention layer for 3 axes in 3D coordinates. 4) 3D-Grraphorrer is not designed for self- superviedd pretraining.
J ONE TABLE FOR MOLECULAR PROPERTY PREDICTION
This section is one clear rable221us associate the pretraining dataset information and the corresponding results from Table and Table 2 for molecular property prediction.
NEURIPS 2022 REVIEW COMMENTS K
The Neurips 2022 review comments Uni-Mo Moi can be found /t/tttss///eeeeeeeeev.eee forum? 1 d= fF2 g10b for reference.
30

---

Published as a conference paper at ICLR 2023
Table 22: Uni-Mol performance molecular property prediction tasks, with information of pretraining datasets. "Z" "C", "p", "G", and "U" in the Data Sources column denote the used pretraining data, which are the abbreviation of ZINC15 (82). ChEMBL [83], PubChem [100) GEOM [8] and Uni-Mol (described in Appendix Aj. "*" denotes unsupervised ised pretraining on the corresponding downstream dataset. The ZINC data a subset of ZINC15. and different models may use different sample sizes. Different data sourees usually have redundancy. such as ZINC and ChemBL. and need to be de-duplicated when used.
Molemness (M) Sources
Data
Z+C
Z
Z
G
P
U
Data
# Presraining
pretrain Do prcitaii
11.00 11.00
1.000
0.05
10.00 00.00
20.00 00:00
19 0..00
0.0ON 14:0 (NN001
0,0085 1200 (0000)
0.009664(0) (0003
0.009840 0.00844000000S (005S)
0.009940 000055)
0.007-4600 0000
0.MN67 00900-))
0,023 0.0103700
0.009
QNN
MAE
0.015010.000
0.017910 00000
SIO O80S))
6.020XNO.0000) (001)
0.021 810.0004)
0.021 0000)
0.0178/0.0003
01011100000011
0.0156(8.000)
QMN
Regrossion sion (lower better J)
12
103 5)8
72 72.012.7)
81 11.9119 19
94-3 8)
94.533.5
66.8/2.3.
58.900.)
41.8(02)
QM
0.68300.166
0.721x0.000 001)
0.8 (200-028)
2.07200.030)
0,739,0.023
081770.000
0.8 700.008)
0,68 (0.000)
0.59100.0044 (0000)
0.660/0.008)
0.603000100
Lipo
2.082(0.08))
2.073/0 185)
2.68800085)
(.61(0.774) 744)
2.764
2.1 2.176(0.052)
2.1 2.17600052)
2.594 (420
L877000094
1,800
RMSE
PreeSSllv
1.50005.08
0.87700.020 029)
(.X3(0.082)
1.100
0.933000000
0.983 (9.00)
(.29(8.033))
77.710.6) 1.27100.000)
0,798(0.022)
0.738 (1.02))
ESSOL
78.611-0
76.6115 15
76.90-7)
74.8002)
81.3/.11
67.3(18)
67.311.8
79.611.9)
81.710.3)
82.1(1.3)
MUV
17
86.00.1))
60.0 2) 25.7(14) S01014
(0.) 76.52211
83.00..)
600.1)
80 (..)) (8.1)
PUBA
128
57.0/0.7) 17.100))
17-20-1)
78.)
2()))
77.0(12)
78.10..5
6009
41127
HIV
Classiication (ROC AUC %. higher is better
SIDER
65.5(0.7)
64806))
63.4(..)
63.
58.91.44
67.200.4
65.9(1.)
27
ToxClst
65.50L.)
63.70022
65.7006)
0) 74300 03:04)
65,310.5
63.10.44
69.2044
79.6(0.5) 69.600))
617
90.6/0 6) 25.900.7
76. 5)
75 ())))
76.23.72 73.50.1)
19.112.8 55.9(0.5)
75.0002)
90.1113) 18.00
Tox21
84 70.
77. 77.564.0
87 17.5(2.7)
91.213.51
91.9115)
Clin
BACE 1513
80 91.00.66 6)
78 48000022)
77.9K1 5)
79.1 19.101.33 3)
82 6/0
2)
82 4(0.9)
72.4(04) 61
(29(066) 85.70..)
71 11.00.33
64.311 8)
69. ((0.8)
10.010.1)
69.50.1)
(2
72.22111
Datusets # Molecules
D-MPNN
Attentive FP
N.Grammr
N. Gramonn
PreiranONN
GROVER...
GROVER
GraphMVP
MoCLL
Uni-Mol
# Tooks
GEM
31


Reviews:
[{"review_id": "gWgOGA0u7ki", "reviewer": "Reviewer_CPRe", "summary": "", "questions": "", "limitations": "", "rating": 10, "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 to push representation learning for molecule-related task further than before, by mixing different tasks in its per-training step, defining a representation that is generic enough to be applicable to a variety of downstream tasks, on which it performs as well or better (and several times, significantly better) than the reported SOTA.\nThe main contributions are :\n- this original idea of mixing very different pre-training tasks / going further in representation learning\n- the extensive per-training performed, in terms of the strategy: the model used is the transformer, and the strategy is to imitate the guess-a-missing-word strategy: instead, one needs to recover the corrupted position of each atom.\n- the extensive per-training performed, in terms of training data (in my perspective it seems impressive, but I haven't read the related literature)\n\nWhat is not new, as paper admits, is:\n- The transformer architecture itself\n- the various type, position, relative distances encoding (although they are well benchmarked), that are SE(3)-invariant\n- the tasks and datasets (taken individually), except for a regression task on pocket property prediction, which is new (no Baseline there).\n", "summary_of_the_review": "Aside from a few points that need further correction/clarification/discussion, the paper is sound and good, and brings a new way of pre-training 3D representations to the community, with proved efficiency.\n", "strength_and_weaknesses": "Strengths:\n- the contribution is both original (new) and strong. \n- performance is extensively review (many downstream tasks, enough detail to follow)\n- the simplicity of the ideas: mostly using and combining already existing conceptual tools, in general rather clearly explained (from someone with a good ML/DL background)\n\nWeaknesses:\n- no major conceptual novelty (may be an asset in a sense)\n- Architecture is a bit overlooked/not explained enough in the main text (page 3)\n- a few unclear parts, or parts that need a bit more discussion\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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "WoDhN98cEj", "reviewer": "Reviewer_6Ybw", "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 work proposes a novel universal molecular representation learning framework, called Uni-Mol, that incorporates 3D information. The network consisting of stack of SE(3)-equivariant transformer layers, is pretrained on two different tasks each involving molecular and protein pocket data respectively. This pretraining is established by learning to predict 3D atomic positions and masked atom representation from the noisy input. These pretrained networks when applied for finetuning various downstream tasks, it is shown to achieve superior performance on molecular / pocket property prediction, molecular conformation generation and protein-ligand binding pose prediction. ", "summary_of_the_review": "Overall, it is a good work and I am inclined towards accepting this work. It establishes two different models that can be put to use for improving accuracy in many downstream tasks. It also provides good tricks for setting up pretraining tasks on molecular / drugs data. For now, I see Uni-Mol strengths outbeats its weaknesses and its contribution may open up avenue of progress in drugs discovery. ", "strength_and_weaknesses": "Strengths:\n1. The paper is well-written and easy to follow. As per my knowledge, it is one of the first work to have shown to incorporate 3D information to improve representation capacity of molecular processing architecture.\n2. In order to preserve rotation and translation invariant properties, Uni-Mol totally operates only on input atomic representation and pairwise Euclidean distance inputs. \n3. Unlike the previously introduced SE(3) equivariant transformers that uses complex processing blocks, Uni-Mol attention layer is quite simple that transforms to and fro between pairwise and atomic representations.This is possible by relegating the 3D positions update to the final layer of the network.\n4. With efficient backbone, Uni-Mol is able to effectively pretrain on large-scale molecular and protein pockets datasets.\n\nWeaknesses:\nWhen drawing evaluation against various baselines on benchmark tasks, it is quite unfair to make comparison of models trained solely on task-specific data to Uni-Mol which has been pretrained on large datasets. For instance,\na. On molecular conformation tasks, all other baselines generates conformation using only 2D molecular graph information. On the other hand, Uni-Mol leverages the 3D position output from RDKit (that uses ETKGD with MMFF) as input for pairwise representation. This for sure provides better start-point for Uni-Mol compared to other baseline that starts with gaussian random noise as 3D position.  \nb. Moreover, pretrained Uni-Mol architecture is relatively very high capacity when compared to baselines.\n\nGiven the nature of contribution it may be difficult to establish apple to apple comparisons. However, fair comparison the efforts can be put into establishing, \na. similar capacity Uni-Mol baseline network\nb. Uni-Mol baseline that is evaluated using gaussian noise input\nc. evaluating each task using other MRL networks\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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RmRZnARRMc", "reviewer": "Reviewer_LVDL", "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": "This paper proposes a self-supervised or pretrained method for molecular representation learning. It can be summarized as follows,\n\n#### 1. A Transformer-based backbone.\n\n#### 2. Two self-supervised learning or pretraining tasks: 3D position recovery and masked atom prediction.\n\n#### 3. A few finetuning strategies are presented for fine-tuning or downstream tasks.", "summary_of_the_review": "The paper presents a  pretraining framework for molecular representation learning. I like the newly collected dataset, which may be potentially useful to the topic.  \n\nHowever, the novelty may be overclaimed. For a pretraining work, the most important thing is the self-learning design. However, predicting the property or attributes of the masked or missing parts is straightforward and has been widely used in existing works.  The experimental comparison is also questionable. ", "strength_and_weaknesses": "## Strength\n\n#### 1. This paper motivates using unlabeled data to improve molecular or protein representation learning via self-supervised learning or pertaining tasks, which is a popular topic in the research area.\n\n#### 2. The two pretraining tasks are technically sound.\n\n#### 3. Two large-scale datasets are collected for pretraining.\n\n#### 4. The paper outperforms most existing methods.\n\n#### 5. The code, model, and data will be made publicly available. \n\n\n## Weaknesses\n\n#### 1. I thought the novelty is questionable.  The authors claimed that the proposed Uni-Mol is the first pure 3D molecular pretraining framework. However, there have been already a few similar works. For example, \n\na. The Graph Multi-View Pre-training (GraphMVP)  framework leverages the correspondence and consistency between 2D topological structures and 3D geometric views. \n\nLiu et al., Pre-training Molecular Graph Representation with 3D Geometry, ICLR 2021.\n\nb. The geometry-enhanced molecular representation learning method (GEM)  proposes includes several dedicated geometry-level self-supervised learning strategies to learn molecular geometry knowledge.\n\nFang et al., Geometry-enhanced molecular representation learning for property prediction, nature machine intelligence, 2022.\n\n\nc. Guo et al. proposed a self-supervised pre-training model for learning structure embeddings from protein 3D structures.  \n\nGuo et al., Self-Supervised Pre-training for Protein Embeddings Using Tertiary Structures, AAAI 2022. \n\nd. The GeomEtry-Aware Relational Graph Neural Network (GearNet) framework uses type prediction, distance prediction and angle prediction of masked parts for pretaining. \n\nZhang et al., Protein Representation Learning by Geometric Structure Pretraining, ICML 2022 workshop. \n\n#### 2. The comparison with the SOTA methods may be unfair. The performance of the paper is based on the newly collected 209M dataset. However, the existing methods use smaller datasets. For example, GEM employs only 20M unlabeled data. Because the scale of datasets has a significant impact on the accuracy, the superior of the proposed method may be from the new large-scale datasets. \n\n#### 3. The authors claimed one of the contributions is that the proposed Uni-Mol contains a simple and efficient SE(3)-equivariant Transformer backbone. However, I thought this contribution is too weak. \n\n#### 4. The improvement is not very impressive or convincing. Although with a larger dataset for pretraining,  the improvement is a bit limited, e.g., in Table 1.  \n\n#### 5. It is not clear which part causes the main improvement: Transformer, pretraining or the larger dataset? \n\n#### 6. It could be better to show the  3D position recovery and masked atom prediction accuracy and visualize the results. \n\n#### 7. The visualization of the self-attention map and pair distance map in Appendix H is interesting. However, according to the visualization, the self-attention map is very similar to the pair distance map, as the author explained. In this case, why not directly use pair distance as attention? Or what does self-attention actually learn besides distance in the task? As self-attention is computationally expensive, is it really needed?  ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "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": "2JPRHPBD-T2", "reviewer": "Reviewer_uSFd", "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 proposes a universal 3D Molecular representation learning (MRL) framework, called Uni-Mol, that significantly enlarges the representation ability and application scope of MRL schemes. Uni-Mol contains two pretrained models with the same SE(3)-equivariant transformer architecture: a molecular model pretrained by 209M molecular conformations; a pocket model pretrained by 3M candidate protein pocket data. Besides, Uni-Mol contains several finetuning strategies to apply the pretrained models to various downstream 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.", "summary_of_the_review": "good paper.", "strength_and_weaknesses": "##########################################################################\n\nPros:\n\n- Uni-Mol is the first pure 3D molecular pretraining framework, and the first molecular pretraining framework that can be directly used in 3D tasks in the field of drug design. which is a very valuable work for 3D molecular representation learning.\n- Uni-Mol contains a simple and efficient SE(3)-equivariant Transformer backbone, and an effective 3D pretraining strategy. The ablation benchmarks demonstrate their superior performance, and Uni-Mol also outperforms SOTA in various downstream tasks.\n- The paper is well-written and the experiment section is solid.\n##########################################################################\n\nCons:\n\n- I think this paper is a solid work, a small piece of advice is since this paper focuses on 3D molecular data, in addition to quantitative experimental results, it will be more intuitive and convincing if more qualitative results can be displayed. Another concern is how much computational resources and time for 3D molecular pretraining needs to consume, which the paper does not seem to mention.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster