NeurIPS 2024

SubGDiff: A Subgraph Diffusion Model to Improve Molecular Representation Learning

Jiying Zhang, Zijing Liu, Yu Wang, Yu Li

Main figure for SubGDiff: A Subgraph Diffusion Model to Improve Molecular Representation Learning

Abstract

Molecular representation learning has shown great success in advancing AI-based drug discovery. The core of many recent works is based on the fact that the 3D geometric structure of molecules provides essential information about their physical and chemical characteristics. Recently, denoising diffusion probabilistic models have achieved impressive performance in 3D molecular representation learning. However, most existing molecular diffusion models treat each atom as an independent entity, overlooking the dependency among atoms within the molecular substructures. This paper introduces a novel approach that enhances molecular representation learning by incorporating substructural information within the diffusion process. We propose a novel diffusion model termed SubGDiff for involving the molecular subgraph information in diffusion. Specifically, SubGDiff adopts three vital techniques: i) subgraph prediction, ii) expectation state, and iii) k-step same subgraph diffusion, to enhance the perception of molecular substructure in the denoising network. Experimentally, extensive downstream tasks demonstrate the superior performance of our approach. The code is available at https://github.com/youjibiying/SubGDiff.

Proposed method: SubGDiff

Benchmark Results

Tested on 7 datasets 18 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
46
ContextPred
actual #46 vs paper #259 213 places worse
0.8060 ±0.0120
best reported
actual #46 vs paper #259 213 places worse
paper 0.7420 ±0.0006
paper 0.7420 ±0.0006 Δ at pub 0.0640
48
3D Infomax
actual #48 vs paper #179 131 places worse
0.8050 ±0.0320
best reported
actual #48 vs paper #179 131 places worse
paper 0.7610 ±0.0018

Difference within reported uncertainty; not counted as underperformance.

paper 0.7610 ±0.0018 Δ at pub 0.0440
61
GraphMVP
actual #61 vs paper #179 118 places worse
0.7990 ±0.0180
best reported
actual #61 vs paper #179 118 places worse
paper 0.7620 ±0.0038
paper 0.7620 ±0.0038 Δ at pub 0.0370
67
MolCLR
actual #67 vs paper #282 215 places worse
0.7980 ±0.0070
best reported
actual #67 vs paper #282 215 places worse
paper 0.7300 ±0.0016
81
InfoGraph
actual #81 vs paper #279 198 places worse
0.7930
best reported
actual #81 vs paper #279 198 places worse
paper 0.7320 ±0.0043
paper 0.7320 ±0.0043 Δ at pub 0.0610
145
SubGDiff
0.7720 ±0.0039
This paper
153
MoleculeSDE (VP)
0.7700 ±0.0012
This paper
214
random init
0.7530 ±0.0022
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.0025

Difference within reported uncertainty; not counted as underperformance.

paper 0.7480 ±0.0025 Δ at pub 0.0030
#ModelROC-AUC ↑Paper
1
ProtoMol
0.9140 ±0.0030
2
MMSG
0.9080 ±0.0050
3
MemGNN
0.9070
9
MolCLR
actual #9 vs paper #270 261 places worse
0.8900 ±0.0030
best reported
actual #9 vs paper #270 261 places worse
paper 0.7150 ±0.0317
43
GIN
0.8646 ±0.0081
129
SubGDiff
0.8210 ±0.0096
This paper
142
MoleculeSDE (VP)
0.8140 ±0.0066
This paper
149
GraphMVP
actual #149 vs paper #179 30 places worse
0.8120 ±0.0090
best reported
actual #149 vs paper #179 30 places worse
paper 0.7980 ±0.0074

Difference within reported uncertainty; not counted as underperformance.

paper 0.7980 ±0.0074 Δ at pub 0.0140 Δ today 0.0140
183
3D Infomax
0.7970 ±0.0154
This paper
182
AttrMask
0.7970 ±0.0033
This paper
186
ContextPred
actual #186 vs paper #242 56 places worse
0.7960 ±0.0120
best reported
actual #186 vs paper #242 56 places worse
paper 0.7600 ±0.0208
paper 0.7600 ±0.0208 Δ at pub 0.0360 Δ today 0.0360
204
InfoGraph
actual #204 vs paper #223 19 places worse
0.7890 ±0.0110
best reported
actual #204 vs paper #223 19 places worse
paper 0.7780 ±0.0088

Difference within reported uncertainty; not counted as underperformance.

paper 0.7780 ±0.0088 Δ at pub 0.0110
234
random init
0.7670 ±0.0251
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
41
3D Infomax
actual #41 vs paper #196 155 places worse
0.8770 ±0.0140
best reported
actual #41 vs paper #196 155 places worse
paper 0.6830 ±0.0112
paper 0.6830 ±0.0112 Δ at pub 0.1940
40
ContextPred
actual #40 vs paper #259 219 places worse
0.8770 ±0.0260
best reported
actual #40 vs paper #259 219 places worse
paper 0.6570 ±0.0062
paper 0.6570 ±0.0062 Δ at pub 0.2200
49
GraphMVP
actual #49 vs paper #167 118 places worse
0.8600 ±0.0340
best reported
actual #49 vs paper #167 118 places worse
paper 0.6940 ±0.0021
paper 0.6940 ±0.0021 Δ at pub 0.1660
78
MolCLR
actual #78 vs paper #241 163 places worse
0.7360 ±0.0050
best reported
actual #78 vs paper #241 163 places worse
paper 0.6660 ±0.0189
142
SubGDiff
0.7020 ±0.0223
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.0236
paper 0.6500 ±0.0236 Δ at pub 0.0520 Δ today 0.0520
145
MoleculeSDE (VP)
0.7010 ±0.0135
This paper
173
InfoGraph
actual #173 vs paper #219 46 places worse
0.6920 ±0.0080
best reported
actual #173 vs paper #219 46 places worse
paper 0.6750 ±0.0011
paper 0.6750 ±0.0011 Δ at pub 0.0170 Δ today 0.0170
200
random init
0.6810 ±0.0059
This paper
#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
12
MolCLR
actual #12 vs paper #50 38 places worse
0.9320 ±0.0170
best reported
actual #12 vs paper #50 38 places worse
paper 0.8610 ±0.0095
19
GCN
0.9130 ±0.0170
35
SubGDiff
0.8820 ±0.0157
This paper
40
AttrMask
0.8770 ±0.0119
This paper
49
GraphMVP
0.8650 ±0.0170
This paper
63
random init
0.8370 ±0.0293
This paper
65
MoleculeSDE (VP)
0.8260 ±0.0364
This paper
89
3D Infomax
0.7990 ±0.0349
This paper
106
InfoGraph
actual #106 vs paper #117 11 places worse
0.7810
best reported
actual #106 vs paper #117 11 places worse
paper 0.7650 ±0.0107

Difference within reported uncertainty; not counted as underperformance.

paper 0.7650 ±0.0107 Δ at pub 0.0160
115
ContextPred
0.7720 ±0.0088
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
14
MolCLR
actual #14 vs paper #219 205 places worse
0.6800 ±0.0110
best reported
actual #14 vs paper #219 205 places worse
paper 0.5750 ±0.0177
66
GIN
0.6403 ±0.0104
70
GraphMVP
actual #70 vs paper #162 92 places worse
0.6390 ±0.0120
best reported
actual #70 vs paper #162 92 places worse
paper 0.6050 ±0.0025
104
SubGDiff
0.6220 ±0.0097
This paper
103
ContextPred
0.6220 ±0.0059
This paper
134
AttrMask
0.6120 ±0.0012
This paper
147
MoleculeSDE (VP)
0.6080 ±0.0104
This paper
155
3D Infomax
0.6060 ±0.0078
This paper
162
InfoGraph
actual #162 vs paper #183 21 places worse
0.6054 ±0.0090
best reported
actual #162 vs paper #183 21 places worse
paper 0.5990 ±0.0030

Difference within reported uncertainty; not counted as underperformance.

paper 0.5990 ±0.0030 Δ at pub 0.0064
224
random init
0.5700 ±0.0133
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 #110 85 places worse
0.7160 ±0.0130
best reported
actual #25 vs paper #110 85 places worse
paper 0.6480 ±0.0025
paper 0.6480 ±0.0025 Δ at pub 0.0680
28
ContextPred
actual #28 vs paper #177 149 places worse
0.7140 ±0.0180
best reported
actual #28 vs paper #177 149 places worse
paper 0.6250 ±0.0031
paper 0.6250 ±0.0031 Δ at pub 0.0890
35
InfoGraph
actual #35 vs paper #135 100 places worse
0.7050
best reported
actual #35 vs paper #135 100 places worse
paper 0.6370 ±0.0050
paper 0.6370 ±0.0050 Δ at pub 0.0680
50
GraphMVP
actual #50 vs paper #113 63 places worse
0.6890 ±0.0100
best reported
actual #50 vs paper #113 63 places worse
paper 0.6450 ±0.0020
paper 0.6450 ±0.0020 Δ at pub 0.0440
107
SubGDiff
0.6500 ±0.0048
This paper
112
MoleculeSDE (VE)
0.6470 ±0.0006
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.0011

Difference within reported uncertainty; not counted as underperformance.

paper 0.6290 ±0.0011 Δ at pub 0.0040
166
MolCLR
0.6290 ±0.0038
This paper
183
random init
0.6210 ±0.0019
This paper

Quantum Chemistry

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 today — OK at publication; today's leaderboard best is elsewhere