ICML 2024

UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning

Shikun Feng, Yuyan Ni, Minghao Li, Yanwen Huang, Zhiming Ma, Weiying Ma, Yanyan Lan

Main figure for UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning

Abstract

Recently, a noticeable trend has emerged in developing pre-trained foundation models in the domains of CV and NLP. However, for molecular pre-training, there lacks a universal model capable of effectively applying to various categories of molecular tasks, since existing prevalent pre-training methods exhibit effectiveness for specific types of downstream tasks. Furthermore, the lack of profound understanding of existing pre-training methods, including 2D graph masking, 2D-3D contrastive learning, and 3D denoising, hampers the advancement of molecular foundation models. In this work, we provide a unified comprehension of existing pre-training methods through the lens of contrastive learning. Thus their distinctions lie in clustering different views of molecules, which is shown beneficial to specific downstream tasks. To achieve a complete and general-purpose molecular representation, we propose a novel pre-training framework, named UniCorn, that inherits the merits of the three methods, depicting molecular views in three different levels. SOTA performance across quantum, physicochemical, and biological tasks, along with comprehensive ablation study, validate the universality and effectiveness of UniCorn.

Benchmark Results

Tested on 9 datasets 14 baseline values reported too low

Tables show presented results in global context (best 3 papers + GCN + MLP).

MoleculeNet

#ModelROC-AUC ↑Paper
1
HIMP
0.8740 ±0.0050
2
GINE w/ VN
0.8720 ±0.0030
3
NaiveGINe+ K=3 w/ VN
0.8700 ±0.0040
14
GCN
0.8400 ±0.0040
15
GROVER
actual #15 vs paper #253 238 places worse
0.8380 ±0.0170
best reported
actual #15 vs paper #253 238 places worse
paper 0.7430 ±0.0010
paper 0.7430 ±0.0010 Δ at pub 0.0950
48
3D Infomax
actual #48 vs paper #246 198 places worse
0.8050 ±0.0320
best reported
actual #48 vs paper #246 198 places worse
paper 0.7450 ±0.0070
paper 0.7450 ±0.0070 Δ at pub 0.0600
61
GraphMVP
actual #61 vs paper #246 185 places worse
0.7990 ±0.0180
best reported
actual #61 vs paper #246 185 places worse
paper 0.7450 ±0.0040
paper 0.7450 ±0.0040 Δ at pub 0.0540
82
UniCorn
0.7930 ±0.0050
This paper
125
MoleBlend
0.7780 ±0.0080
This paper
150
Mole-BERT
actual #150 vs paper #160 10 places worse
0.7710 ±0.0040
best reported
actual #150 vs paper #160 10 places worse
paper 0.7680 ±0.0050

Difference within reported uncertainty; not counted as underperformance.

paper 0.7680 ±0.0050 Δ at pub 0.0030
161
MoleculeSDE
0.7680 ±0.0030
This paper
202
GraphMAE
0.7550 ±0.0060
This paper
229
AttrMask
actual #229 vs paper #238 9 places worse
0.7510 ±0.0090
best reported
actual #229 vs paper #238 9 places worse
paper 0.7480 ±0.0020

Difference within reported uncertainty; not counted as underperformance.

paper 0.7480 ±0.0020 Δ at pub 0.0030
#ModelROC-AUC ↑Paper
1
ProtoMol
0.9140 ±0.0030
2
MMSG
0.9080 ±0.0050
3
MemGNN
0.9070
43
GIN
0.8646 ±0.0081
59
UniCorn
0.8580 ±0.0120
This paper
96
MoleBlend
0.8370 ±0.0140
This paper
105
GraphMAE
0.8310 ±0.0090
This paper
118
GROVER
actual #118 vs paper #118
0.8260 ±0.0070
actual #118 vs paper #118
paper 0.8260 ±0.0070
paper 0.8260 ±0.0070 Δ at pub 0.0310
143
Mole-BERT
actual #143 vs paper #164 21 places worse
0.8140 ±0.0100
best reported
actual #143 vs paper #164 21 places worse
paper 0.8080 ±0.0140

Difference within reported uncertainty; not counted as underperformance.

paper 0.8080 ±0.0140 Δ at pub 0.0060
149
GraphMVP
actual #149 vs paper #232 83 places worse
0.8120 ±0.0090
best reported
actual #149 vs paper #232 83 places worse
paper 0.7680 ±0.0110
paper 0.7680 ±0.0110 Δ at pub 0.0440 Δ today 0.0440
170
MoleculeSDE
actual #170 vs paper #187 17 places worse
0.8040 ±0.0092
best reported
actual #170 vs paper #187 17 places worse
paper 0.7950 ±0.0210

Difference within reported uncertainty; not counted as underperformance.

paper 0.7950 ±0.0210 Δ at pub 0.0090
182
AttrMask
0.7970 ±0.0030
This paper
183
3D Infomax
0.7970 ±0.0150
This paper
#ModelROC-AUC ↑Paper
1
MAGPrompt+
0.9432 ±0.0238
2
CamS-LLaMA
0.9420
3
SE(3)-I MPPH
0.9400 ±0.0210
34
GIN
0.8875 ±0.0049
37
GROVER
actual #37 vs paper #150 113 places worse
0.8870 ±0.0060
best reported
actual #37 vs paper #150 113 places worse
paper 0.7000 ±0.0010
paper 0.7000 ±0.0010 Δ at pub 0.1870
41
3D Infomax
actual #41 vs paper #174 133 places worse
0.8770 ±0.0140
best reported
actual #41 vs paper #174 133 places worse
paper 0.6910 ±0.0100
paper 0.6910 ±0.0100 Δ at pub 0.1860
49
GraphMVP
actual #49 vs paper #190 141 places worse
0.8600 ±0.0340
best reported
actual #49 vs paper #190 141 places worse
paper 0.6850 ±0.0020
paper 0.6850 ±0.0020 Δ at pub 0.1750
71
UniCorn
0.7420 ±0.0110
This paper
80
MoleculeSDE
actual #80 vs paper #107 27 places worse
0.7320 ±0.0048
best reported
actual #80 vs paper #107 27 places worse
paper 0.7180 ±0.0070
paper 0.7180 ±0.0070 Δ at pub 0.0140
83
MoleBlend
0.7300 ±0.0080
This paper
96
Mole-BERT
actual #96 vs paper #106 10 places worse
0.7230 ±0.0070
best reported
actual #96 vs paper #106 10 places worse
paper 0.7190 ±0.0160

Difference within reported uncertainty; not counted as underperformance.

paper 0.7190 ±0.0160 Δ at pub 0.0040
101
GraphMAE
0.7200 ±0.0060
This paper
141
AttrMask
actual #141 vs paper #269 128 places worse
0.7020 ±0.0050
best reported
actual #141 vs paper #269 128 places worse
paper 0.6500 ±0.0230
paper 0.6500 ±0.0230 Δ at pub 0.0520 Δ today 0.0520
#ModelROC-AUC ↑Paper
1
SE(3)-I MPPH
0.9930 ±0.0040
2
KA-GCN
0.9920 ±0.0050
3
KA-GAT
0.9910 ±0.0050
15
UniCorn
0.9210 ±0.0040
This paper
19
GCN
0.9130 ±0.0170
40
AttrMask
0.8770 ±0.0110
This paper
42
MoleBlend
0.8760 ±0.0070
This paper
46
MoleculeSDE
0.8700 ±0.0050
This paper
49
GraphMVP
actual #49 vs paper #96 47 places worse
0.8650 ±0.0170
best reported
actual #49 vs paper #96 47 places worse
paper 0.7900 ±0.0250
71
GraphMAE
0.8230 ±0.0120
This paper
76
GROVER
actual #76 vs paper #79 3 places worse
0.8170 ±0.0250
best reported
actual #76 vs paper #79 3 places worse
paper 0.8120 ±0.0300

Difference within reported uncertainty; not counted as underperformance.

paper 0.8120 ±0.0300 Δ at pub 0.0050
88
Mole-BERT
actual #88 vs paper #96 8 places worse
0.8010 ±0.0360
best reported
actual #88 vs paper #96 8 places worse
paper 0.7890 ±0.0300

Difference within reported uncertainty; not counted as underperformance.

paper 0.7890 ±0.0300 Δ at pub 0.0120
89
3D Infomax
0.7990 ±0.0340
This paper
#ModelROC-AUC ↑Paper
1
KA-GAT
0.8470 ±0.0020
2
KA-GCN
0.8420 ±0.0010
3
GraphKAN
0.8370 ±0.0010
50
MoleBlend
0.6490 ±0.0030
This paper
52
GROVER
actual #52 vs paper #52
0.6480 ±0.0060
actual #52 vs paper #52
paper 0.6480 ±0.0060
paper 0.6480 ±0.0060 Δ at pub 0.0100
66
GIN
0.6403 ±0.0104
69
UniCorn
0.6400 ±0.0180
This paper
70
GraphMVP
actual #70 vs paper #99 29 places worse
0.6390 ±0.0120
best reported
actual #70 vs paper #99 29 places worse
paper 0.6230 ±0.0160

Difference within reported uncertainty; not counted as underperformance.

paper 0.6230 ±0.0160 Δ at pub 0.0160 Δ today 0.0160
90
Mole-BERT
0.6280 ±0.0110
This paper
134
AttrMask
0.6120 ±0.0010
This paper
148
MoleculeSDE
0.6080 ±0.0030
This paper
155
3D Infomax
0.6060 ±0.0070
This paper
175
GraphMAE
0.6030 ±0.0110
This paper
#ModelROC-AUC ↑Paper
1
GeomGCL
0.7630
2
LiGhT DeepGPT
0.7570 ±0.0110
3
CoMPT
0.7550
14
GCN
0.7350 ±0.0020
25
3D Infomax
actual #25 vs paper #117 92 places worse
0.7160 ±0.0130
best reported
actual #25 vs paper #117 92 places worse
paper 0.6440 ±0.0080
paper 0.6440 ±0.0080 Δ at pub 0.0720
34
GROVER
actual #34 vs paper #99 65 places worse
0.7080 ±0.0150
best reported
actual #34 vs paper #99 65 places worse
paper 0.6540 ±0.0040
paper 0.6540 ±0.0040 Δ at pub 0.0540
43
UniCorn
0.6940 ±0.0110
This paper
50
GraphMVP
actual #50 vs paper #168 118 places worse
0.6890 ±0.0100
best reported
actual #50 vs paper #168 118 places worse
paper 0.6270 ±0.0010
paper 0.6270 ±0.0010 Δ at pub 0.0620
87
MoleBlend
0.6610
This paper
103
MoleculeSDE
actual #103 vs paper #106 3 places worse
0.6520 ±0.0031
best reported
actual #103 vs paper #106 3 places worse
paper 0.6500 ±0.0020

Difference within reported uncertainty; not counted as underperformance.

paper 0.6500 ±0.0020 Δ at pub 0.0020
114
Mole-BERT
actual #114 vs paper #121 7 places worse
0.6450 ±0.0040
best reported
actual #114 vs paper #121 7 places worse
paper 0.6430 ±0.0020

Difference within reported uncertainty; not counted as underperformance.

paper 0.6430 ±0.0020 Δ at pub 0.0020
128
GraphMAE
0.6410 ±0.0030
This paper
150
AttrMask
actual #150 vs paper #165 15 places worse
0.6330 ±0.0060
best reported
actual #150 vs paper #165 15 places worse
paper 0.6290 ±0.0010

Difference within reported uncertainty; not counted as underperformance.

paper 0.6290 ±0.0010 Δ at pub 0.0040
#ModelRMSE ↓Paper
1
Graphormer-SPIS
0.4840 ±0.0050
2
Graphormer-SPD
0.4920 ±0.0040
3
conan-FGW
0.5140 ±0.0190
7
GCN
0.5640 ±0.0150
49
UniCorn
0.8170 ±0.0340
This paper
#ModelRMSE ↓Paper
1
SE(3)-I MPPH
0.3540 ±0.0530
2
ECFP-4 + RF
0.5600 ±0.0660
3
GPSE + GPS
0.6820
28
UniCorn(DFT + RDKit)
1.5550 ±0.0750
This paper
69
GIN
2.3460 ±0.1220

Quantum Chemistry

#ModelMAE ↓Paper
1
QHNet
0.0000
2
QHNetV2
0.0000
3
PhiSNet
0.0000
42
Frad
actual #42 vs paper #58 16 places worse
0.0431
best reported
actual #42 vs paper #58 16 places worse
paper 0.0530
43
UniCorn
0.0460
This paper
58
Coord
actual #58 vs paper #58
0.0529
best reported
actual #58 vs paper #58
paper 0.0530
paper 0.0530 Δ at pub 0.0001
86
SE(3)-DDM
0.1220
This paper
104
MoleculeSDE
actual #104 vs paper #104
0.2820
actual #104 vs paper #104
paper 0.2820
106
MoleculeJAE
0.3450
This paper
SOTA performance (not significantly different from the best) Proposed method (this paper) Baseline / comparison (this paper) Global context (other papers) Best at pub — reported result worse than the best known when published (best reported in green) Best at pub — reported value improved over the best known when published Best today — OK at publication; today's leaderboard best is elsewhere