SubgDiff: A Subgraph Diffusion Model to Improve Molecular Representation Learning
Authors: Jiying Zhang
Zijing Liu
Yu Wang
Bin Feng
Yu Li
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion model
Subgraph
Molecular representation learning
Primary area: diffusion_based_models

Abstract: Molecular representation learning has shown great success in advancing AI-based drug discovery. A key insight of many recent works is that the 3D geometric structure of molecules provides essential information about their physicochemical properties. Recently, denoising diffusion probabilistic models have achieved impressive performance in molecular 3D conformation generation.  However, most existing molecular diffusion models treat each atom as an independent entity, overlooking the dependency among atoms within the substructures. This paper introduces a novel approach that enhances molecular representation learning by incorporating substructural information in the diffusion model framework. 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. Experiments on extensive downstream tasks, especially the molecular force predictions, demonstrate the superior performance of our approach.

SubgDiff: A Subgraph Diffusion Model to Improve Molecular Representation Learning
Jiying Zhang, Zijing Liu", Yu Wang, Bin Feng, Yu Li" International Digital Economy Academy (IDEA)
(ennbiin (zhangjiying, ,,i,g,,,uui,, liluiiijiann jingiiinggnnn fenggiiinn liyu) edu. en
https: //pp:://iia trub. com
com/IDEA- IDEA XL/SubgDiffff
Abstract
Molecular representation learning has shown great success in advancing Al-based drug discovery. key insight of many recent works is that the 3D geometric structure of molecules provides essential information about their physicochemical properties. Recently, denoising diffusion probabilistic models have achieved im- pressive performance in molecular 3D conformation generation. However, most existing molecular diffusion models treat each atom as an independent entity, overlooking the dependency among atoms within the substructures. This paper introduces novel approach that enhances molecular representation learning by incorporating substructural information in the diffusion model framework. We propose novel diffusion model termed SubgDiff for involving the molecular subgraph information in diffusion. Specifically, Subg adopts three vital tech- niques: i) subgraph prediction, ii) expectation state, and iii) k-step same subgraph diffusion, to enhance the perception of molecular substructure in the denoising network. Experiments on extensive downstream tasks, especially the molecular force predictions, demonstrate the superior performance of our approach.
Introduction
Molecular representation learning (MRL) has at- tracted tremendous attention due significant role in learning from limited labeled data for applications
Solution
0.8
Adsorbed
Surface
like Al-based drug discovery [37, 3, 4, 25] and ma- terial science [33]. From the perspective of physical chemistry, the 3D molecular conformation can largely determine the properties of molecules and the activi- ties of drugs [6, 7]. Thus, numerous geometric neural network architectures and self-supervised learning strategies have been proposed explore 3D molecu- lar structures to improve performance on downstream molecular property prediction tasks [35, 56, 23, 58].
0.6
Solid
0.2
0
Meanwhile, diffusion probabilistic models (DPMs) have shown remarkable power to generate realistic samples, especially in synthesizing high-qualit im- ages and videos [40, 10]. By modeling the generation as reverse diffusion process, DPMs transform ran- dom noise into sample in the target distribution. Recently, diffusion models/flow models have also
Figure 1: Equilibrium probability of the six dif- ferent 3D structures (c1-c6) of the same molecule ibuprofen (C13H1802) in four different conditions. (Adapted with permission from [27). Copyright 2018 American Chemical Society.)
"Corresponding author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

DDPM
6
Ege Gaussian noise;
R
R
tomic poondinates 3D molecule
SUBGDIFF
E2
E3
Sy
R
R²
Subgraph selected by mask vector's 5t
Figure 2: Comparison of forward process between DDPM [10] and subgraph diffusion. For each step, DDPM adds noise into all atomic coordinates, while subgraph diffusion selects subset of the atoms to diffuse.
demonstrated strong capabilities of molecular 3D conformation generation [52, 16, 59, 43, 61, 59]. The training process of a DPM for conformation generation can be viewed as the reconstruction of the original conformation from noisy version, where the noise is modulated by different time steps. Consequently, the denoising objective in the diffusion model can naturally be used as a self-supervised representation learning technique [32]. Inspired by this intuition, several works have used this technique for molecule pretraining [24, 56]. Despite considerable progress. the potential of DPMs for molecular representation learning has not been fully explored. We therefore raise the question: Can we effectively enhance MRL with the denoising nelwork (noise predictor) of DPM? If yes, how to achieve it?
To answer this question, we first analyze the gap between the current DPMs and the characteristics of molecular structures. Most diffusion models on molecules propose to independently inject continuous Gaussian noise into the every node feature [11] atomic coordinates of 3D molecular geometry [52, 56). However, this approach treats each atom as an individual particle, overlooking the substructure within molecules, which pivotal in molecular representation learning [55, 46, 28]. As shown in Figure I, the substructures are mostly invariant in different 3D molecular conformations, which contains crucial information about the properties, such as the equilibrium distribution, crystallization and solubility [27]. As result, uniformly adding same-scale Gaussian noise to all atoms makes difficult for the denoising network to learn the joint distribution of global structure and local substructure, hindering the downstream prediction performance on molecular properties closely related to 3D conformation (e.g. most physicochemical properties). So here we try to tackle the previous question by designing a DPM involving the knowledge of substructures.
Toward this goal, we propose novel diffusion model termed SubgDiff, adding distinct Gaussian noise to different substructures of 3D molecular conformation. Specifically, instead of adding the same Gaussian noise to every atomic coordinate, SubgDiff only adds noise to randomly selected subgraph at each time step in the diffusion process (Figure 2). In the training phase of SubgDiff, subgraph prediction task is integrated into the training objective, which explicitly directs the denoising network to capture substructure information from the molecules. Additionally, we propose two techniques: expectation state diffusion and ststep same-subgraph diffusion, to train the model effectively and enhance its sampling capability.
With the ability capture the substructure information from the noisy 3D molecule, the denoising networks tend to gain more representation power. The experiments on various 2D and 3D molecular property prediction tasks demonstrate the superior performance of our approach. To summarize, our contributions are as follows: (1) we incorporate the substructure information into diffusion models to improve molecular representation learning; (2) we propose a new diffusion model Subg that adopts subgraph prediction, expectation state and k-step subgruh diffusion to improve its sampling and training: (3) the proposed representation learning method achieves superior performance on various downstream tasks, especially molecular forces prediction.

---

2 Related work
Diffusion models graphs. The diffusion models graphs be mainly divided into two
categories: continuous diffusion and discrete diffusion. Continuous diffusion applies a Gaussian noise process on each node edge [14, 30J, including GeoDiff [52], EDM [11]. SubDiff [54]. Meanwhile, discrete diffusion constructs the Markov chain on discrete space, including Digress [9] and GraphARM [18]. However, it remains open to exploring fusing the discrete characteristic into the continuous Gaussian on graph learning, although a closely related work has been proposed for images and cannot be used for generation [32]. Our work, SubgDiff, is the first diffusion model fusing subgraph, combining discrete characteristics and the continuous Gaussian.
Conformation generation. Various deep generative models have been proposed for conformation generation, including CVGAR [26], GraphDG [39]. CGCF [50], ConfVAE [51], ConfGF [38] and GeoMol [8]. Recently, diffusion-based methods have shown competitive performance. Torsional Diffusion [16] raises diffusion process on the hypertorus defined by torsion angles, However, it is not suitable as a representation learning technique due the lack of local information (length and angle of bonds). GeoDiff [52] generates molecular conformation with a diffusion model on atomic coordinates. However, it views the atoms as separate particles, without considering the dependence between atoms from the substructure.
SSL for molecular property prediction. There exist several works leveraging the 3D molecular conformation to boost the representation learning, including GraphMVP [22], GeoSSL. [24], the denoising pretraining approach raised by Zaidi al. [56] and MoleculeSDE [23], etc. However, those studies have not considered the molecular substructure in the pertaining. this paper, we concentrate how to boost the perception of molecular substructure in the denoising networks through the diffusion model.
The discussion with more related works (e.g. MDM [32], MDSM [19] and SSSD [1]) can be found in Appendix B.I.
3 Preliminaries
Notations. We use I to denote the identity matrix with dimensionality implied by context, @ to represent the element product, and diag(s) denote the diagonal matrix with diagonal elements of the vector s. If not specified, both and represent noise sampled from the standard Gaussian distribution N( I). The topological molecular graph can be denoted as G(V, &, X) where the set of nodes, & is the set of edges, X is the node feature matrix, and its corresponding 3D Conformational Molecular Graph is represented as GG R), where R = [R1, ... Ryyl E g3 is the set of 3D coordinates of atoms.
DDPM. Denoising diffusion probabilistic models (DDPM) [10] is a typical diffusion model [40] which consists of diffusion (aka forward) and reverse process. In the setting of molecular conformation generation, the diffusion model adds noise on the 3D molecular coordinates R [52].
Forward and reverse process. Given the fixed variance schedule 31, B2, Gr, the posterior distribution q( (R. (R9) that is fixed to a Markov chain can be written as T
(R RD)
II (Q(RR' (R-))" q('(RR R'-)) A((R; /1 3RR' B.I..
(1)
t-1
To simplify notation, we consider the diffusion on single atom coordinate R, and omit the subscript
e to get the general notion A throughout the paper. Let a; = 11 3r, ā. = [1 B.), and then the sampling of R² at any time step has the closed form: q(' (To) = A Vo, (1 (1- I)).
The reverse process of DDPM is defined as Markov chain starting from Gaussian distribution p(R*) = NR";; I)
T
(ooo)) - p(R") II (R
( - -N((T" ( (o( ( .t),
where 0: = 1-ō.-1 3e denote time denendennddent constant. In DDPM, (o RR, t) is parameterized as 1-n;
Vi- C t)) and co, i.e. the denoising network, is parameterized by neural network where the inputs are R² and time step t.
Ho(R', (R', t)= = 6, (R'
B.
3

---

Training and sampling, The training objective of DDPM is:
simmple (0) = - o v1 f)1l]
After training, samples are generated through the reverse process po po( Specifically, R² first sampled from W(O, I), and I2 in each step is predicted as follows,
R
(R'
ae
(R + ~
(4)
VI
4 SubgDiff
Directly using DDPM on atomic coordinates 3D molecules means each atom is viewed as an independent single data point. However, the substructures play important role in molecular generation [15] and representation learning [57]. Ignoring the inherent interactions among atoms within substructures may hinder the denoising network from learning features for molecular prop- erties related to 3D conformation. In this paper, we propose to involve a mask operation in each diffusion step, leading new diffusion SubgDiff for molecular representation learning. Each mask corresponds to subgraph in the molecular graph, aligning with the substructure in the 3D molecule. Furthermore, we incorporate a subgraph predictor and reset the state of the Markov Chain to the expectation of atomic coordinates, thereby enhancing the effectiveness of SubgDiff in sampling. Additionally, we also propose k-step subgraph diffusion for training to effectively capture the substructure information.
4.1 Involving subgraph into diffusion process
In the forward process of DDPM, we have R. VT R. + /Btt-1, Yu E V. in which the Gaussian noise ét-1 is injected every atom. Moreover, the training objective in Equation 3 shows that the denoising networks would always predict Gaussian noise for all atoms. Neither the forward nor reverse process of DDPM takes into account the substructure of the molecule. Instead, in SubgDiff, mask vector s, = [st,:" ... stj T (0. 1] is introduced to determine which atoms will be added noise at step t. The mask vector St is sampled from discrete distribution (S G) to select a subset of the atoms. In molecular graphs, the discrete mask distribution Ps, (S G) is equivalent to the subgraph distribution, defined over predefined sample space Further, the distribution Ps, (S g) should keep the selected connected subgraph cohere with the molecular substructures. Here, we adopt a Torsional prrsed decomposition method [16] (Details in Appendix A.2). With the mask vector latent variables $11,, the state transition of the forward process can be formulated as (Figure 3):
0
St
R
R
Figure 3: The Markov Chain of SubgDiff is lazy Markov Chain.
(Grnb
1. fi-1. where each sample is a connected subgraph extracted from G.
[V 3. + Vicu-1 R
st. if St. O,
R
(5)
which can be rewritten as R /I st, B,R, + st, BIE
The posterior distribution q( [20, s1.T) can be expressed as matrix form:
QR 1: R
Stor) = II RR'RR (RR' s,,) NRR T B,diag(s,)RR' (,ddiag(s,) ",, (6)
To simplify the notation, consider the diffusion on single node v and omit the subscript in R and st, to get the notion Rt and st. By defining % := 1 st,r, 71 : I[ s,B,), the closed form of sampling R' given Rdth is R
N I)..
4.2 Reverse process learning
The reverse process is decomposed as follows:

---

T (o.o(R .11T) p((1" II Po - '[2t, s)po(s(
Po.o
(8)
where po( RR' 84) and po(s; R²) are both learnable models. In the context of molecular learning, the model can be regarded as first predicting which subgraph 8, should be denoised and then using the noise prediction network pe po( [R', .s)) to denoise the node position in the subgraph.
However, it is tricky to generate 3D structure by adopting the typical training and sampling method used in Ho et al. [10]. Specifically, following Ho et al. [10], the reverse process can be optimized by maximizing the variational lower bound (VLB) of log p(R)) as follows.
(s( [ 9((x)
log p( R 2
Beg wellt)) log
+ Egure (11) llg (R [R))
reconstruction tenn
subgraph prodiction term
T
Egtors) DKalg(R2 R S1T) I(RR))
E
Becre Bae Dia (g( R² sux) Po(R stl)].
masching
denoising matching
Details of the derivation are provided in the Appendix D.2. The subgraph predictor (s R²) in the first term can be parameterized by node classifier sg. For the denoising matching term that is closely related to sampling, by Bayes rule, the posterior (R"'1|A"
|R', R .1::) can be written as:
q( (RR R², s1.) xx WVR' (R', R², S14, ()))
(10)
Mgg (R' Rº Fo) v1 Best
(Tt
81 1-11) B. E0):
Bis
where (t) is the standard deviation and s1:-1 are contained in 1- Following DDPM and parameterizing po (R--1 R se) (N Mg RR,, R 81--1, (0 R² co(g, R², (())). the training objective
Laimpte (0, 0) = all.
co R E)D ABCE s. (GG ()),
(11)
where BCE(S,, sy) is the binary cross entropy loss, is the weight used for the trade-off, and sæ is the subgraph predictor implemented as node classifier with Gzn(G. RR) input and shares molecule encoder with cg. The BCE loss employed here uses the subgraph selected at time-step as the target, thereby explicitly compelling the denoising network to capture substructure information from molecules. Eventually, the sy can be used to infer the mask vector š. = so (G. R',t) during sampling. Thus, the sampling process is:
n²= - = (T, nR SL: RI,, co, R²,)) (t)z).
(12)
However, using Equation 11 and Equation 12 directly for training and sampling faces two issues. First, the inability to access $1:1-1 in My during the sampling process hinders the step-wise denoising procedure, posing a challenge to the utilization of conventional sampling methods in this context. Inferring $1:-1 solely from Ht using another model po (s |R', st) is also difficult due to the intricate modulation of noise introduced in Rs through multi-step Gaussian noising. Second, training the subgraph predictor with Equation 11 is challenging. To be specific, the subgraph predictor should be capable of perceiving the sensible noise change between time steps t 1 and f. However, the noise scale 3, is relatively small when t small, especially if the diffusion step is large (e.g. 1000). As result, it is difficult to precisely predict the subgraph.
Next, effectively tackle the above issues, we design two techniques: expectation state diffusion and k-step same subgraph diffusion.
4.3 Expectation state diffusion.
We first devise new way calculate the denoising term. To eliminate the effect of mask series s and improve the training of subgraph prediction loss, we use a new lower bound the denoising term follows:
Ey a no) [Dxc(q(R'- R 814) Il RPo[R''1 RR', st]] < P.. (Ka((((t" 1 R² 81a) || (T -22-11 ' +:)]]
5

---

where (( Rrr,,,,,, s1.) is defined as (( E2 11 12-11
1..... _R .9. (e)a 111-1 12 110)
It is an approxi- mated posterior that only relies on the expectation of R2 and St. This lower bound defines new forward process, in which, state 0 to state t - use the Ean. ,R'1 and state t remains as Equation 6. Assume each node V E V. St, Bern(p) (i.idd. w.r.t. t). Formally, we have
q(R' R .t) - v1 StBB (s.8))),
(13)
E ((ER--11R s NNER VouR P II a50n,
1-1
(14)
where 0r :=
p)² and ã := I Og are general form of aj and 0jj DDPM
(p 1), respectively. Intuitively, this process equivalent to using mean state ER' to replace R- during the forward process. This estimation is reasonable since the expectation = R- is
like cluster center of R which can represent R properly. Thus, the approximated posterior becomes
q(
x ( " " RP, S13 6o). of(1),
(15)
where
Po(R", 810) =
(R'
9,
v1 B.a
V + =n Dimi Sent
0c-1 /./( + +1 s. o300)= = 8, Σ
0e-1 8.).
We parameterize (o Rt, st) as N p RP, 814-1, Sa .),, co ((')))) t)), og)(),)()) and adopt the same training objective as Equation 11. By employing the sampling method R! Mag( RV, Ro, -1 Re, 1), eo R2', ()) ())+ +,((((), we observe that the proposed expectation state enables step-by-step execution of the sampling process. Moreover, using expectation is beneficial to reduce the complexity of R² for predicting the mask & during training. This will improve the denoising network to perceive the substructure when we the diffusion model for self-supervised learning.
4.4 k-step same-subgraph diffusion.
To reduce the complexity of the mask series (8), S2. ,8)) and accumulate more noise on the same subgraph for facilitating the convergence of the subgraph prediction loss, we generalize the one-step subgraph diffusion to k-step same subgraph diffusion (Figure 8 in Appendix), in which the selected subgraph will be continuously diffused & steps. After that, the difference between the selected and unselected parts will be distinct enough to help the subgraph predictor perceive it. The forward process of k-step same subgraph diffusion can be written as (t > ki k E N):
( R " ( N R²,
(1 18)) i-t-
II
(16)
where of = -1- -1 -11311
4.5
Training and sampling of SubgDiff
By combining the expectation state and k-step same-subgraph diffusion, SubgDiff first divides the entire diffusion step T into T/k diffusion intervals. In each interval [ki, kii + 1)], the mask vectors k(f+1)
{s, f=kiz are equal to Ski+l- SubgDiff then adopts the expectation state at the split time step {ikli 1... } to eliminate the effect of {s i .)., that is, gets the expectation of ER at step ik w.r.t. S Overall, the diffusion process of SubgDiff is two-phase diffusion process. In the first phase, the state to state ( 1)/k] use the expectation state diffusion, while in the second phase, state k(l(t 1) //]] + 1 to state t use the k-step same subgraph diffusion (see Figure 4), With m == [[t - 1)/k:,, the two phases can be formulated as follows.
am (1 a D441 B.)) Phase I: Step 0 - km: Esg mm Rkm = Van Rd + PV E
where aj (Pv 1122 (1 -1-18 (1)))) +1 p)² is a general forms of az in Equation 14 (in which

---

Expectation state diffusion
kmn-step same- subgraph diffusion
km+
S Pence1
Elik
Exedi2
f(k- Daosiimm
(5, - 5, Jewa
(5, 5-g21liie
(5) Aen-m--12eee
Eegs R²
R
Esute R
4i mask random variable; Gy: Gamesian moise; E.R. the oeperiation of winls
subgraph selected. by mask sample SuNvE
Figure 4: The forward process of SubgDiff. The state 0 to km uses the expectation state and the mask variables are the same in the interval [kii, ki + kj, 0, 1, m 1. The state km + to applies the same subgraph diffusion.
Reverse Process
S2te1 33x
CLuk
E2k+1:lk
ST
Syal
S
1:2k
R
R
Spm+
Figure 5: The reverse process Sufg The selected subgraph is the same in the interval [ki, min(ki + k, T)). 0,.., m.
case k 1) and ã; = II}=1 aj. In the rest of the paper, aj denotes the general version without a special statement. Actually, Een.o Rku only calculate the expectation w.r.t. random variables (si++1/i 1,2,.....
Phase II: Step km + + t: The phase a km))-ttp same mask diffusion. R VIli-k++1 (1 3.8kmm Eargem Rkm + Ili (1 B18km+1)Esm
Let 7/2 = 318km+11 7/t I Ni and Br = IT (1 - Bc). We can drive the single-step state transition: q(121 NRR:, T. (1 T) and
a(D0-1 (7)) = N(R=
1am ,6); From
5t-1
Oan 5
(1
But
Re-1
(17)
B
Feom
Then we reuse the training objective in Equation 11 as the objective of SubgDiff: Lasggte (0, 0) E, s... deress..eldddaaae (l))ii co 1)) ABCE(s,, (oG.R'
(18)
where R is calculated by Equation 17.
Sampling. Although the forward process uses the expectation state w.r.t. s, we can only update the because sampling only needs to get a subgraph from the distribution in the k-step interval. Eventually, adopting the 8 defined in Equation 17, the sampling process (Figure 5) shown below,
mask &, at ik; = 1, ... ...
R²
1 Vee
(R'-
Vni 8em
1km+1.
Ro .f)+
Shant 1
2 Tem
Vró 8km=
61 E
5am
2,
(19)
where N 1), [(t 1)/k] and 8km+1 s (9 Rkm++ km .mm+ kk). The subgraph selected by Skm+1 will be generated in from the steps km + & to km. The mask predictor can be viewed as discriminator of important subgraphs, indicating the optimal subgraph should be recovered in the next & steps. After one subgraph (substructure) is generated properly, the model can gently fine-tune the other parts the molecule (c.f. the video supplementary material). This subgraph diffusion would intuitively increase the robustness and generalization of the generation process, which is also verified by the experiments Section 5.2. The training and sampling algorithms of SubgDiff are summarized in Algorithm and Algorithm 2.
7

---

Algorithm 1: Training SubgDiff
Input: molecular graph GAD, k for same mask diffusion, m = [(| - 1)/k]
Sample ~ ((,,...,,))
10,
Sample Skm+1 Psxm+1 (5|
Sample subgraph
R' + (R*{R))
Equation 17
L1 BCE 86 (9, R', .)) f))
Subgraph prediction loss
L2 = ||diag (Sag+()++1))) ( (9, R ())
> Denoising loss
optimizer. step R sv.e (AL1 + fa))
b Optimize parameters 0, 0
Algorithm 2: Sampling from SubgDiff is the same as training, for k-step same-subgraph diffusion; Sample R. ~ I),
> Random noise Initialization
for /= = 7 to I do
2 ~~~ (~, if 1, clse z 0
>> Random noise > Subgraph prediction > Posterior sampling
If t%k = 0 or = T: T 89 R
R'-1 <- Equation 19
end
return R
5
Experiments
We conduct experiments to address the following two questions: 1) Can substructures improve the representation ability of the denoising network when using diffusion self-supervised learning? 2) How does the proposed subgraph diffusion affect the generative ability of the diffusion models? For the first question, we employ SubgDiff' as denoising pretraining task and evaluate the performances of the denoising network on various downstream tasks. For the second one, we compare SubgDiff with the vanilla diffusion model GeoDiff [52] the task of molecular conformation generation.
5.1 SubgDiff improves molecular representation learning
To verify the introduced substructure in the diffusion can enhance the denoising network for represen- tation learning, we pretrain with SubgDiff objective and finetune on various downstream tasks.
Dataset and settings. For pretraining, we follow [23] and use PCQM4MM2 dataset [12]. It's sub- dataset of PubChemQC [29] with 3.4 million molecules with 3D geometric conformations. We use various molecular property prediction datasets downstream tasks. For tasks with 3D conformations, we consider the dataset MD17 and follow the literature [35, 36, 24] of using 1K for training and IK for validation, while the test set (from 48K to 991K) is much larger. For downstream tasks with only 2D molecule graphs, we use eight molecular property prediction tasks from MoleculeNet [48].
Pretraining framework. To explore the potential of the proposed method for representation learning, we consider MoleculeSDE [23], a SOTA pretraining framework, to be the training backbone, where SubgDiff is used for the 2D - 3D model and the mask operation is extended to the node feature and graph adjacency for the 3D - 2D model. The details can be found in Appendix A.4.2.
Baselines. For 3D tasks, we incorporate two self self-supervised methods [TTyee Prediction, Angle Pre- diction], and three contrastive methods [InfoNCE [31] and EBM-NCE [22] and 3D InfoGraph [24]]. Two denoising baselines are also included [GeoSSL [24], Denoising [56] and MolcculeSDE). For 2D tasks, the baselines are AttrMask, ContexPred [13], InfoGraph [45], MolCLR [47], 3D InfoMax [44], GraphMVP [22] and MoleculeSDE.
Results. As shown in Table and Table 2, SubgDiff outperforms MoleculeSDE in both 2D and 3D downstream tasks. Particularly, the significant improvement on the MD17 force prediction task, which is closely related to 3D molecular conformation, demonstrates that the introduced subgraph diffusion helps the perception of molecular 3D structure in the denoising network during pretraining. Further, Diff achieves SOTA performance compared all the baselines. This also indicates that the proposed SufgDiff objective is promising for molecular representation learning due to the
8

---

Table I: Results (mean absolute error) on MD17 force prediction. The hest and second best results are marked in bold and underlined.
Pretruining
Aspirin
Benzene
Ethanol
Mulonaldehyde
Naphthulene
Salicylic 4
Toluene
Uracil v
(random init)
1.203
0.380
0.386
0.794
0.587
0.826
0.568
0.773
Type Prediction
1.383
0.402
0,450
0.879
0.622
1.028
0.662
0.840
Angle Prediction
1.542
11.447
0.669
1.022
0.680
1.032
0.623
0.768
3D InfoGraphh
1,610
0.415
0.788
0.768
InfoNCE
1.132
0.395
0,466
0.888
0.542
0.831
0.554
0.664
HBM-NCR
1.251
0.373
0.457
0.824
0.512
0.990
0.560
0.742
Denoising
1.364
0,391
0.432
0.830
0.599
0.817
0.628
0.607
GeoSSL
1.107
0.360
0.357
0.737
0,568
0,902
0.484
0.502
SDD (VE)
1.112
0.304
0,282
0,520
0,455
0.725
0.515
0.447
MoD (VP)
1.244
0.315
0,338
0.488
0.432
0.712
0.478
0.468
Ours
0.880
0.252
0.258
0.459
0.325
0.572
0.362
0.420
Table 2: Results for MoleculeNet (with 2D topology only). We report the mean (and standard deviation) ROC-AUC of three random seeds with scaffold splitting for each task. The backbone is GIN. The hest and second best results marked bold and underlined, respectively.
Pre-training
BBBP
Tox21 f
ToxCast f
Sider
ClinTox
MUV f
1
Bace f
Avg 1
(random init)
68.140.59
75.3+0.22
62.14-0.19
83.74-2.3
74.64+2.35
75.2400.0
76.742251
AttMaask
65.0:22.36
74.8::.25
62.9:10111
61.2=0.12
17.7
73.442202
76.81:0.33
79.7#0.33
72.68
ContextPred
6..7+0062
74.2+0.06
62.5++0.11
62.2-0.59
77.2:1088
75.3++1.57
77.1+-0.6
76.0
71.28
InfoGraph
67.5:0.11
73.240.43
63.71.0.00
59.9=0.30
76.51:007
74.110.74
75.11:.99
77.820.8
MolCLR
6..6:11.89
73.0::.16
62.9:10.38
57.5=1.77
86.1:
72.54:2.38
76.24:1.51
71.5
70.79
InfoMax
68.3:1.12
76.1=0.18
64.81:0.55
60.6=0.78
79.9:3.49
74.44:245
75.94:0.99
79.7
72.47
GrPphMVP
69.4:0.21
76.2+0.38
64.5+0020
60.5=0.25
86.5+1170
76.2422.8
76.24+0.11
79.840.74
73.66
MolvculeSDE/VE)
68.3::0.25
76.9-10.33
64.71:0.06
60.2=0.29
80.8:1.53
76.8+17.1
77.011.68
19.9
73.15
7..1+1.35
77.0=1012
64.0:10.07
60.8=1004
82.6++364
76.6::3.25
77.3411.1
0-0066
73.73
70.2422.33
77.240.39
65.0-10.48
62.2+
1852
77.3+-117
77.6:1051
involvement of the knowledge of substructures during training. More results on the QM9 dataset [35] on the quantum mechanics property prediction be found in Appendix A.4.3.
5.2 SubgDiff benefits conformation generation
We have proposed a new diffusion model to enhance molecular representing learning, where the base diffusion model (GeoDiff) is initially designed for conformation generation. To evaluate the effects of SubgDiff on the generative ability of diffusion models, we assess its generation performance and generalization ability. Following prior works [52], we utilize the GEOM-QN MMM [34] and GEOM- Drugs [2] datasets. The detailed description of the dataset splitting, metrics and model architecture can be found in Appendix A.5,
Conformation generation.
The comparison with GeoDiff on the GEOM-QM9 dataset is reported in that SubgDiff significantly outperforms the GeoDiff baseline
Table 3. From the results, it is easy on both metrics (COV-R and MAT-R) across different sampling steps. It indicates that by training with the substructure information, SufgDiff has positive effect on the conformation generation task. Moreover, SubgDiff with 500 steps achieves much better performance than GeoDiff with 5000 steps on 5 out of 8 metrics, which implies our method accelerate the sampling efficiency (10x).
Table 3: Results for conformation generation on GEOM-QM9 dataset with different diffusion timesteps. DDPM [10] is the sampling method used in GcoDiff. Our proposed sampling method (Algorithm can be viewed DDPM variant.
denotes SubgDiff oupprffommiiundepprfonnsss GeoDiff.
COV-R (%)
MAT-R (A) 1
COV-P (S)
MAT (AT-P((A) (Á) J
Models
Timesteps
Sampling method
Mean
Median
Mean
Median
Mean
Median
Mean
Median
GeoDiff
5000
DDPM
80.36
83.82
0.2820
0.2799
53.66
50.85
0.6673
0.4214
SuhgDii
5000
DDPM (ours)
90.91
95.59
0.2461
0.2351
50.16
48.01
0.6114
0.4791
GeoDiff
500
DDPM
80.20
83.59
0.3617
0.3412
45.49
45.45
1.1518
0.5087
SubgDiff
500
DDPM (ours)
89.78
94.17
0.2417
0.2449
50.03
48.31
0.5571
0.4921
GeoDiff
200
DDPM
69.90
72.04
0.4222
0.4272
36.71
33.51
0.8532
0.5554
Diff
200
DDPM (ours)
85.53
88.99
0.2994
0.3033
47.76
45.89
0.6971
0.5118

---

Domain generalization. To further illustrate the benefits of SubgDiff, we design two cross- domain tasks: (1) Training on QM9 (small molecular with up to 9 heavy atoms) and testing on Drugs (medium-sized organic com- pounds); (2) Training on Drugs and testing QM9. The results (Ta- ble and Appendix Table 13) show that SubgDiff consistently outper- forms GeoDiff and other models trained the in-domain dataset, demonstrating the introduced sub-
Table 4: Results the GEOM-QM9 dataset for domain generalization. Except GeoDiff and SubgDiff, the other methods are trained with in-domain data.
COV-R (5) f
MAT-R (A)
4 Median
Models
Train data
Mean
Median
Mean
CVGAE [26]
QM9
0.09
0.00
1.6713
1.6088
GraphDG [39]
QM9
73.33
84.21
0.4245
0.3973
CGCF [50]
QM9
78.05
82.48
0.4219
0.3900
ConfVAE [51]
QM9
77.84
88.20
0.4154
0.3739
GeoMol [8]
QM9
71.26
72.00
0.3731
0.3731
GcoDiff
Drugs
74.94
79.15
0.3492
0.3392
SubgDiff
Drugs
83.50
88.70
0.3116
0.3075
structure effectively enhances the robustness and generalization of the diffusion model.
6 Conclusion
We present novel diffusion model SubgDiff, which involves the subgraph constraint in the diffusion model by introducing a mask vector to the forward process. Benefiting from the expectation state and k-step same-subgraph diffusion, SubgDiff effectively boosts the perception of molecular substructure in the denoising network, thereby achieving state-of-the-art performance at various downstream property prediction tasks. There are several exciting avenues for future work. The mask distribution can be made flexible such that more chemical prior knowledge may be incorporated into efficient subgraph sampling. Besides, the proposed SubgDiff can be generalized to proteins such that the denoising network can learn meaningful secondary structures.
Acknowledgement
This project was supported by Shenzhen Hetao Shenzhen-Hong Kong Science and Technology Innovation Cooperation Zone, under Grant No. HTHZZOSW-KKCYDD-00300
References
[1] Juan Lopez Alcaraz and Nils Strodthoff. Diffusion-based time series imputation and forecasting with structured state space models. Transactions on Machine Learning Research, 2022.
[2] Simon Axelrod and Rafael Homez-Bombarelli GEOM, energy-annotated molecular conforma- tions for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
[3] He Cao, Zijing Liu, Xingyu Lu, Yuan Yao, and Yu Li. Instructmol: Multi-modal integration for building a versatile and reliable molecular assistant in drug discovery. arXiv preprint arXiv:2311.16200 2023.
[4] He Cao, Yanjun Shao, Zhiyuan Liu, Zijing Liu, Xiangru Tang, Yuan Yao, and Yu Li. PRESTO: Progressive pretraining enhances synthetic chemistry outcomes. arXiv preprint arXiv:2406.1311 13193, 2024.
[5] Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 1603015, 2017.
[6] Aurora Cruz-Cabeza and Joel Bernstein. Conformational polymorphism. Chemical reviews, 114(4)22100-2191 2014.
[7] Bin Feng. Zequn Liu, Nanlan Huang, Zhiping Xiao, Haomiao Zhang, Srbuhi Mirzoyan, Hanwen Xu, Jiaran Hao, Yinghui Xu, Ming Zhang, et al. bioactivity foundation model using pairwise meta-learning, Nature Machine Intelligence, 6(8)9962-974, 2024.
10

---

[8] 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: 13757-13776,, 2021.
[9] Kilian Konstantin Haefeli, Karolis Martinkus, Nathanaël Perraudin, and Roger Wattenhofer. Diffusion models for graphs benefit from discrete state spaces. In The First Learning Graphs Conference, 2022.
[10] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840 6851, 2020.
[11] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. International conference on machine learning, pp- 8867-8887, 2022.
[12] Weihua Hu, Matthias Fey, Marinka Zitmik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint arXiv:2005.000 00680, 2020.
[13] Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovee. Strategies for pre-training graph neural networks. International Conference on Learning Representations, 2020.
[14] John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph-based protein design. Advances in neural information processing systems, 32, 2019.
[15] Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Hierarchical generation of molecular graphs using structural motifs. In International conference on machine learning, pp. 2020.
[16] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi S Jaakkola. Torsional diffusion for molecular conformer generation. In Advances Neural Information Processing Systems, 2022.
[17] Wolfgang Kabsch. solution for the best rotation to relate two sets of vectors. Acta Crystallo- graphica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallography, 32 (5):922 109922-923,111 1976.
[18) Lingkai Kong, Jiaming Cui, Haotian Sun, Yuchen Zhuang, B. Aditya Prakash, and Chao Zhang. Autoregressive diffusion model for graph generation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, volume 202 of Proceedings of Machine Leurning Research, pp. 17391-17408, 23-29 Jul 2023.
[19] Jiachen Lei, Peng Cheng, Zhongjie Ba, and Kui Ren. Masked diffusion models are fast learners. arXiv preprint arXiv:2306. 11363, 2023.
[20] Haitao Lin, Yufei Huang, Odin Zhang, Yunfan Liu, Lirong Wu, Siyuan Li, Zhiyuan Chen, and Stan Z. Li. Functional-groupp-esss diffusion for pocket-specific molecule generation and elaboration. Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[21] Shengchao Liu, Mehmet F Demirel, and Yingyu Liang. N-gram graph: Simple unsupervised representation for graphs, with applications to molecules. Advances in neural information processing systems, 32, 2019.
[22] 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 Representations, 2022.
[23] Shengchao Liu, Weitao Du, Zhi-Ming Ma, Hongyu Guo, and Jian Tang. A group symmetric stochastic differential equation model for molecule multi-modal pretraining. In International Conference on Machine Learning, pp. 21497-21526, 2023.

---

[24] Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)- invariant denoising distance matching, The Eleventh International Conference on Learning Representations, 2023.
[25] Xingyu Lu, He Cao, Zijing Liu, Shengyuan Bai, Leqing Chen, Yuan Yao, Hai-Tao Zheng, and Yu Li. MoleculeQA: dataset evaluate factual accuracy in molecular comprehension. arXiv preprint S0.iv::200308/922 2024.
[26] Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry prediction using deep generative graph neural network. Scientific reports, 911:20338, 2019.
[27] Veselina Marinova, Geoffrey PF Wood, Ivan Marziano, and Matteo Salvalaglio. Dynamics and thermodynamies of ibuprofen conformational isomerism at the crystal/solution interface. Journal of chemical theory and computation, (4(12):6444-4499) 2018.
[28] Siqi Miao, Mia Liu, and Pan Li. Interpretable and generalizable graph learning via stochastic attention mechanism. In International Conference on Machine Learning, pp. 15524-15544 2022.
[29] Maho Nakata and Tomomi Shimazaki. Pubchemqe project: a large-scale first fist-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 1706):130--1308 2017.
[30] Chenhao NIu, Yang Song, Jiaming Song. Shengjia Zhao, Aditya Grover, and Stefano Ermon. Permutation invariant graph generation via score-based generative modeling. In International Conference on Artificial Intelligence and Statistics, pp. 4474 2020.
[31] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 2018.
[32] Zixuan Pan, Jianxu Chen, and Yiyu Shi. Masked diffusion as self-supervised representation learner. arXiv preprint arXiv:2308. 05695, 2023.
[33] Robert Pollice, Gabriel dos Passos Gomes, Matteo Aldeghi, Riley J Hickman, Mario Krenn, Cyrille Lavigne, Michael Lindner-D's Addario, AkshatKumar Nigam, Cher Tian Ser, Zhenpeng Yao, et al. Data-driven strategies for accelerated materials design. Accounts of Chemical Research, 54(4):849-860, 2021.
[34] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[35] Kristof Schütt, Picter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmicla, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-filte convolutional neural network for modeling quantum interactions. Advances in newral information processing systems, 30, 2017.
[36] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377-9388, 2021.
[37] Jie Shen and Christos A Nicolaou. Molecular property prediction: recent trends in the era of artificial intelligence. Drug Discovery Today: Technologies, 32:29-36, 2019.
[38] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International conference on machine learning, pp. 9558-9568, 2021.
[39] Gregor Simm and Jose Miguel Hernandez-Lobato. A generative model for molecular distance geometry. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119, pp. 8949-89588 2020.
[40] Jascha Sohl- Sohk-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. International conference on machine learning, pp. 2256-2265. 2015.
12

---

[41] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems. pp. 11918-11930, 2019.
[42] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Leurning Representations, 2020.
[43] Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei- Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36, 2023.
[44] Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3D infomax improves GNNs for molecular property prediction. In International Conference on Machine Learning, pp. 20479-20508 2022.
[45] Fan- Yun Sun, Jordan Hoffmann, Vikas Verma, and Jian Tang, Infograph: Unsupervised and semi-supervised graph-level representation learning via mutual information maximization. In International Conference on Learning Representations (ICLR), 2020.
[46] Yuyang Wang, Rishikesh Magar, Chen Liang, and Amir Barati Farimani. Improving molecular contrastive learning via faulty negative mitigation and decomposed fragment contrast. Journal of Chemical Information and Modeling, (((11):27132277) 2022.
[47] Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence, 4(3):299- 287, 2022.
[48] Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. MoleculeNet: a benchmark for molecular machine learning. Chemical science, 9()):513 53000, 2017.
[49] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2018.
[50] 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, 2021.
[51] Minkai Xu, Wujie Wang. Shitong Luo, Chence Shi. Yoshua Bengio, Rafael and Jian Tang. An end-to-end framework for molecular conformation generation via bilevel programming. In International Conference on Machine Learning, pp. 11537-11547,0021
[52] 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.
[53] Zhao Xu, Youzhi Luo, Xuan Zhang, Xinyi Xu, Yaochen Xie, Meng Liu, Kaleb Dickerson, Cheng Deng, Maho Nakata, and Shuiwang Ji. Molecule3d: A benchmark for predicting 3d geometries from molecular graphs arXiv preprint arXhv:2110.0171 2021.
[54] Kuo Yang, Zhengyang Zhou, Limin Li, Pengkun Wang, Xu Wang, and Yang Wang. Subdiff: Subgraph latent diffusion model, 2024.
[55] Zhaoning Yu and Hongyang Gao. Molecular representation learning via heterogeneous motif graph neural networks. In International Conference on Machine Learning, pp. 25581-25594, 2022.
[56] Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Gonzalz, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. In The Eleventh International Conference on Learning Representations, 2023.
[57] Xuan Zang, Xianbing Zhao, and Buzhou Tang. Hierarchical molecular graph self-supervised learning for property prediction. Communications Chemistry, 6(1):34, 2023.
13

---

[58] Jiying Zhang, Xi Xiao, Long-Koa Kai Huang, Yu Rong, and Yatao Bian. Fine-tuning graph neural networks via graph topology induced optimal transport. In Lud De Raedt (ed.), Proceedings of the Thirty-Firsi International Joint Conference on Artificial Intelligence, IJCAI-22, pp. 3730- 3736. International Joint Conferences on Artificial Intelligence Organization, 2022. doi: 10.24%63/jjci..22255 2022/518 8. Main Track.
[59] Jiying Zhang, Zijing Liu, Shengyuan Bai, He Cao, Yu Li, and Lei Zhang. Efficient antibody structure refinement using emergy-guided se (3) flow matching. In IEEE International Conference Bioinformatics and Biomedicine (BIBM). IEEE, 2024.
[60] Yangtian Zhang, Zuobai Zhang, Bozitao Zhong, Sanchit Misra, and Jian Tang. DiftPack:
A torsional diffusion model for autoregressive protein side-chain packing. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
[61] Qi Zhengyang, Liu Zijing. Zhang Jiying. Cao He, and Li Yu. Controlmol: Adding substruture control to molecule diffusion models. In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2025.
14

---

Appendix
Contents
Experiment details and more results
15
A.1 Visualization of representations
15
A.2 Mask distribution
16
A.3 Hyperparameters
16
A.4 Settings and more results on molecular representation learning task
17
A.5 Settings and more results on conformation generation task
19
A.6 Sentivity analysis of in k-step sume-subgraph diffusion
21
B More discussions
22
B.1 More related works
22
B.2 Limitations
22
C An important lemma for diffusion model
23
D Derivations of training objectives
26
D.1 SubgDiff (1-same step and without expectation state)
................................
26
D.2 ELBO of SufgDiff
................................
28
D.3 Single-step subgraph diffusion
30
E Expectation state distribution
33
F The derivation of SubgDiff
34
A Experiment details and more results
A.1 Visualization of representations
We conduct an alignment analysis validate that method capture subgraphs during pretraining. Specifically, we employ t-distributed stochastic neighbor embedding (t-SNE) to visually represent molecules with various scaffolds. The purpose is to investigate whether molecules sharing the same scaffold exhibit similar representations, which are extracted by the pretrained molecular encoder. scaffold, which is the core structure of molecule, holds significant importance in the field of chemistry as it serves as a foundation for systematic exploration of molecular cores and building blocks. is usually represented by substructure of molecule and can be regarded as the subgraph in our SufgDiff.
In our analysis, select the nine most prevalent scaffolds from each dataset (BBBP, Sider, Clin and Bacе) and assign each molecule to a cluster according to scaffold. To quantify the molecule embedding, compute the Silhouette index of the embeddings for each dataset.
As shown Table 5, SubgDiff enables the generation more distinctive representations molecules with different scaffolds. This implies that SubgDiff enables the denoising network (molecular encoder) to better capture the subgraph (scaffold) information. We also provide the t-SNE visualizations Figure 6.
15

---

BBBP
Sider
Silbowette: 0,2632
Silhoustte: 0.1739
XCO
go
Qg--
+
dgp-
80
ono
0
o
Clintox
BACE
Silhouctie: 0.1780
Silhourtte: 0,2571
on
a
8
XHO
CO
8
orggo
Figure 6: T-distributed stochastic neighbor embedding (t-SNE) visualization of the learned molecules representa- tions, colored by the scaffolds of the molecules.
Table 5: Silhouette index (higher is better) of the molecule embeddings on Moleculenet dataset (with 2D topology only)
BBBP f
ToxCast
Sider 1
ClinTox
Bace 1
MoleculeSDE SubgDifff (ours)
0.2344
0.0611 0.0650
0.1664 0.1739
0.1394
0.1860
0.2632
0.1780
0.2571
A.2 Mask distribution
In this paper, pre-define the mask distribution be discrete distribution, with sample space = { and (S = G )mb,) 1/N,t ( 1, where G is the subgraph split by the Torsional- based decomposition methods [16]. The decomposition approach will cut off one torsional edge in 3D molecule to make the molecule into two components, each of which contains at least two atoms. 1). Thus
The two components are represented as complementary mask vectors (i.e. torsional edges in G will generate 2n subgraphs. Finally, for each atom the st, ~ (0r5) ie. p 0.5 SubgDiff.
A.3 Hyperparameters
All models are trained with SGD using the ADAM optimizer. Pre-training takes around 24 hours with a single Nvidia A6000 GPU of 48GB RAM. The hyperparameters can be seen in Table 6 and Table 7.
Table 6: Additional hyperparameters of our SubgDiff.
Task
31
3T
scheduler
T'
k (k-same mask)
Batch Size
Train Iter.
QM9
1e-7
2e-3
sigmoid
5000
250
10Å
64
2M
Drugs
1e-7
2e-3
sigmoid
5000
250
10Â
32
6M
16

---

Table 7: Additional hyperparamelers of SubyDiff with different timesteps.
Task
31
Br
3 scheduler
k (k-same mask)
Batch Size
Train Iter.
500-step QM9
le-7
2e-2
sigmoid
500
25
10Å
64
2M
200-step QM9
1c-7
5e-2
sigmoid
200
10
10Å
64
2M
500-step Drugs
lc-7
2c-2
sigmoid
500
25
10Å
32
4M
1000-step Drugs
le-7
9e-3
sigmoid
500
50
10A
32
4M
A.4
Settings and more results on molecular representation learning task
A.4.1 Model architecture
We use the pretraining framework MoleculeSDE proposed by [23] and extend our SubgDiff to multi- modality pertaining. The two key components MoleculeSDE are two SDEs(stochastic differential equations Song et al. [42]]]: from 3D conformation to 2D topology (3D + 2D). In practice, these two SDEs can be replaced by discrete diffusion models. In this paper, we use the proposed SubgDiff to replace the SDEs.
SDE from 2D topology to 3D conformation (2D + 3D) and an SDE
2D topoloical molecular graph. topological molecular graph is denoted as JrD g(V. X), where X is the atom attribute matrix and X is the bond attribute matrix. The 2D graph representation with graph neural network (GNN)
H2D GIN(s = GIN(G, X), where GIN is the powerful 2D graph neural network [49] and H2D (1 12 ... where hip is the i-th node representation.
(20)
3D conformational molecular graph. The molecular conformation is denoted gaD G3D R). The conformational representations are obtained by a 3D GNN SchNet [35]:
y A H3D SchNet(gsD) SchNet(G, R), where HaD [b%. .... and hpp is the i-th node representation.
(21)
An Equivaraut Conformation Generation The first objective is the conditional generation from topology to conformation, p implemented as SubgDifff. The denoising network adopt is the SE(3)--quivariance network (S (S2D-3D) used MoleculeSDE. The details of the network architecture refer [23].
Therefore, the training objective from 2D topology graph to 3D confirmation is:
LID = R.,. FR.
(22)
[|| diag(s:)(e S 320 (x, Re, t)) BCE(s,, a 20-330 (ax, R., t))]
where >> (D. (æ. R,,() gets the invariant feature from Sg and introduces mask head (MLP) to read out the mask prediction.
An SE(3)-Invariant Topology Generation. The second objective is to reconstruct the 2D topology from 3D conformation, i.c., p( We also use the SE(3)-invariant score network S S:DD-2D proposed by MoleculeSDE. The details of the network architecture refer to [23]. For modeling S:D-+2 it needs to satisfy the SE(()-invariance symmetry property. The inputs are 3D conformational representation y. the noised 2D information æ; time I, and time I. The output of S0D+22 is the
Gaussian noise, as (eX (e) The diffused 2D information contains two parts: H; (X,, Ex). For
node feature X, the training objective is
FEx,y Fr,,,
(23)
diag(s,)(c S x. ///
X,,)
BCE(s, " X.,)))
(24)
For edge feature E, we define mask matrix S from mask vector Siy 1 if or s, otherwise, Sty = 0. Eventually, the ojective can be written as:
Lí EE,, Ex, R
(25)
S, 530 (v. F,, F...))) BCF(Ss, 2p T0-20 F..1))]
(26)
17

---

Then the score network S3 is also decomposed into two parts for the atoms and bonds: S (zi) and Sg (at). Similarly, the mask predictor 2 > is also decomposed into two parts for the atoms and bonds: 8 (at) and 8 (az4).
Similar to the topology conformation generation procedure, the s)p (2z R4, t) gets the invariant feature from - and introduces mask head (MLP) read out the mask prediction.
Learning. Following MoleculeSDE, we incorporate contrastive loss called EBM-NCE [22]. EBM- NCE provides an alternative approach to estimate the mutual information T T(X;Y) and is anticipated to complement the generative self-supervised learning (SSI.) method. As result, the ultimate objective is:
Laverill Conitastivee + a2 ozs (SLDDDD2D + L L3D-2D)
(27)
where 01, 022 ary are three coefficient hyperparameters.
A.4.2 Dataset and settings
Dataset. For pretraining, following MoleculeSDE, we use PCQM4Mv2 [12]. It's a sub-dataset of PubChemQC [29] with 3.4 million molecules with both the topological graph and geometric conformations. For finetuning, in addition to QM9 [34], we also include MD17. To be specific, small organic molecules. These
MD17 comprises eight molecular dynamics simulations focused datasets were initially presented by Chmiela et al. [5] for the development of mhergy-conserving force fields using GDML. Each dataset features the trajectory of an individual molecule, encompassing broad spectrum of conformations. The objective is to predict energies and forces for each trajectory by employing single model.
Baselines for 3D property prediction We begin by incorporating three coordinate-MI-unaee SSL methods: (1) Type Prediction, which aims to predict the atom type of masked atoms; (2) Angle Prediction, which focuses predicting the angle among triplet atoms, specifically the bond angle prediction; (3) 3D InfoGraph, which adopts the contrastive learning paradigm by considering the node-graph pair from the same molecule geometry as positive and negative otherwise. Next, in accordance with the work of [24], we include two contrastive baselines: (4) GeoSSI-InfoNCE [31] and (5) GeOSSL-EBM-NCE [22]. We also incorporate generative SSL baseline named (6) GeoSSL-RR (RR for Representation Reconstruction). The above baselines are pre- pre- pre-trained on a subset of IM molecules with 3D geometries from Molecule3D [53] and we reuse the results reported by [24] with SchNet as backbone.
Baselines for 2D topology pretraining. We pick up the most promising ones as follows. Attr- Mask [13, 21], ContexPred [13], InfoGraph [45], and MoICLR [47].
Baselines for 2D and 3D multi-modality pretraining. We include MoleculeSDDE)233 Variance Exploding (VE) and Variance Preserving (VP) as erucial baseline to verify the effectiveness of our methods due to the same pertaining framework. We reproduce the results from the released Code.
A.4.3 3D molecular property prediction Results on QM9.
By adopting the pertaining setting in Appendix A.4.2, we also take the QM9 dataset for finetuning and follow the literature [35, 36, 23], using 1 110K for training, 10K for validation and 11k for testing. In addition, the QM9 dataset encompasses 12 tasks that pertain to quantum properties, which are commonly used for evaluating representation learning tasks [35, 24]. The experimental results can be seen in Table 8. The results also suggest the superior performance of our method.
A.4.4 Compared with GeoDiff.
We directly reuse the pre-trained model of the molecular conformation generation in Section 5.2 for fine-tuning, to compare our method with GefDiff from naive denoising pretraining perspective [56]. The results are shown in Table 9.
18

---

Table 8: Results 12 quantum mechanics prediction tasks from QM9. We take I 110K for training, 10K for validation, and 11K for testing. The evaluation is mean absolute error (MAE), and the best and the second best results are marked bold and underlined. respectively. The backbone SchNet.
Pretraining
Alpha 1
Gap 1
HOMOL
LUMO
Mu 4
Cv
G298 1
H298
R2
U298 4
U0. 4
Zpve
0.070
50.59
33
14.85
14.70
14.44
1.698
Supervised
0.070
51.34
32.62
27.61
0.030
0.032
14.08
14,09
0.141
14.13
13.25
1.727
Type Prediction
0.084
56.07
34.55
30.65
0,040
0.034
18.79
19.39
0.201
19.29
18.86
2.001
Angle Prediction
0.084
57.01
15.81
16.41
3D InaoGraph
0.076
53.33
33.92
28.55
0.030
0.032
15.97
16.28
0.117
16.17
15.96
1.666
GeossL-RR
0.073
52.57
34.44
28.41
0.033
0.038
15.74
16.11
0.194
15.58
14.76
1.804
GeeosL-InfoNCCC
0.075
53.00
1.675
GeossI -EBM-CCE
0.073
$2.86
33.74
28. 07
0.031
0.032
14.02
13.65
0.121
13.70
13.45
1.677
MoleculeSDE
0.062
47.74
28.02
24.60
0,028
0.029
13.25
12.70
0.120
12.68
12.93
1.643
Ours
0.054
44.88
25.45
23.75
0.027
0.028
12.03
11.46
0.110
11.32
11.25
1.568
Table 9: Results on 12 quantum mechanics prediction tasks from QM9. We take 1 110K for training, 10K for validation, and 11K for testing. The evaluation is mean absolute error (MAE), and the best and the second hest results marked bold and underlined, respectively. The hackhone is SchNet.
Pretraining
Alpha 4
Gap 1
HOMO]
LUMO
Mu L
Cv
G298 1
H298 4
R2
U298
00.1
Zpve
GeoDiff'
28.29
0.028
0.035
15.76
15.24
1.869
SuhgDiff (ours)
0.076
50.80
31,15v
26.62
0.025
0.032
14.92A
12.86A
0.129A
14.74
14.53a
1.710%
A.5
Settings and more results conformation generation task
A.5.1
Dataset and network.
Following prior works [52], we utilize the GEOM-QM9 [34] and GEOM Drugs [2] datasets. The former dataset comprises small molecules of up to 9 heavy atoms, while the latter contains larger drug- like compounds. We reuse the data split provided by Xu et al. [52]. For both datasets, the training dataset comprises 40, 000 molecules, each with 5 conformations, resulting in 200, 000 conformations in total. The test split includes 200 distinctive molecules, with 14, 324 conformations for Drugs and 22, 408 conformations for QM9.
We adopt the graph field network (GFN) from [52] as the GNN encoder for extracting the 3D molecular information. In the 1-th layer, the GFN receives node embeddings h' E Ra*b (where b represents the feature dimension) and corresponding coordinate embeddings x RRx3 as input. then produces the output h and x according to the following process:
m'y = T (h/, h/,, ||x
x²) 1/11 Giy;
(28)
hlt1 4% hn/, Σ mlyi0h)
(29)
JEN(()
x
Σ (EN(()
(R, R2) 20 (m'ri 0m)
(30)
des
where & are implemented as feed-forward networks and diy denotes interatomic distances. The initial embedding h" is composed of atom embedding and time step embedding while x" represents atomic coordinates. (iii) is the neighborhood of jth node, consisting of connected atoms and other ones within a radius threshold т, helping the model capture long-range interactions explicitly and support disconnected molecular graphs.
Eventually, the Gaussian noise and mask can be predicted as follows (C.f. Figure 7):
& = x
(31)
& MLP(h.)
(32)
where é is equivalent and &; is invariant.
Settings. For GeoDiff [52] with 5000 steps, we use the checkpoints released in public GitHub to reproduce the results. For 200 and 500 steps. we retrain it and do the DDPM sampling.
A.5.2 Evaluation metries for conformation generation.
To compare the generated and ground truth conformer ensembles, we employ the same evaluation metrics in a prior study [8]: Average Minimum RMSD (AMR) and Coverage. These metrics
19

---

= (G..R.
F R,O)
RCR(S, &,); L2 [|diag(s)(( f
Adding boise
Mask bead
Noise head
Rd
R
3D Molecules
GNN Encoder & Mask prediction &y: Noise prediction
Figure The model architecture for denoising SubgDiff.
enable us to assess the quality of the generated conformers from two perspectives: Recall (R) and Precision (P). Recall measures the extent to which the generated ensemble covers the ground-truth ensemble, while Precision evaluates the accuracy of the generated conformers.
The four metries built upon moo-mean--quure deviation (RMSD), which is defined the normalized Frobenius norm of two atomic coordinates matrices, after aligmment by Kabsch algorithm [17]. Formally, let Sg and S, denote the sets of generated and reference conformers respectively, then the Coverage and Matching metrics [50] can be defined as:
COV-R(S,, S,) [S, {ce E, RM)D(,,,,,) S
(33)
Σ min RMSD(C,C), MAT-R(S,, S,) |s,| CES.
(34)
CES,
where ò a threshold. The other two metrics COV-P and MAT-P can be defined similarly but with the generated sets Sy and reference sets S, exchanged.
COV-P(S,, Sy))
{ce E S, RMSD(C.C) s,cc &.C.. [
(35)
MAT-P(S,, |S ĆES,
min RMSD(C, CES,
(36)
In practice, Sg is set as twice of the size of S, for each molecule.
A.5.3 Comparison with GeoDiff' using Langevin Dynamics sampling method.
In order to verify that proposed diffusion process can bring benefits to other sampling methods, we conduct the experiments to compare our proposed diffusion model with GeoDiff by adopting a typical sampling method Langevin dynamics (L.) sampling)[41]:
R e--1 = R + outo(G, R²,e) 20,
(37)
where ~ (O, I) and ho? h is the hyper-parameter referring to step size and o, is the noise schedule in the forward process. We use various stme-step to evaluate the generalization and robustness of the proposed method, and the results shown in Table 10 indicate that our method significantly outperforms GeoDiff, especially when the time-step relatively small (200,500), which implies that our training method can effectively improve the efficiency of denoising. Similar results are also observed on the GEOM-drugs dataset (Table 11).
Table 10: Results GEOM-QM9 dataset with different time steps. Langevin dynamics [41] is typical sampling method used in DPM. a denotes SubgDiff outperforms GeoDiff. The threshold 5 = 0.5À.
COV-R (%)
MAT-R (A) 1
COV-P (56
MAT-P (A)
Steps
Sampling method
Models
Mean
Median
Mean
Median
Mean
Median
Mean
Median
500
Langevin dynamics
GooDiff
87.80
93.66
0.3179
0.3216
46.25
45.02
0.6173
0,5112
Langevin dynamics
SubgDiff
91.40a
95.39
0.2543
0.2601
51.71
48.50A
0.5035
0.4734A
200
Langevin dynamics
GeoDiff
86.60
0.3532
0.3574
42.98
42.60
0.5563
0.5367
200
Langevin dynamics
SubgDiff
9
9
0.3064
0,3098
4
46.46
0,5540
0.5082
20

---

Table 11: Results the GEOM-Drugs dataset under different diffusion timesteps. DDPM [10] is the sam- pling method used GeoDiff and Langevin dynamics [41] a typical sampling method used in DPM. Our proposed sampling method (Algorithm 2) can be viewed as a DDPM variant. A/v denotes SubgDiff outper-
forms/underperforns GeoDiff. The threshold 1.25A.
COV-R (%) f
MAT-R (A) 4
Models
Timesteps
Sampling method
Mean
Median
Mean
Median
GeoDiff
500
DDPM
50.25
48.18
1.3101
1.2967
SubgDiff
500
DDPM (ours)
76.16
86.43
1.0463
1.0264
GeoDiff
500
LD
64.12
75.56
1.1444
1.1246
SubgDiff
500
L. (ours)
74.30
77.87 A
1.0003
0.9905
Table 12: Results GEOM-QM9 dataset. The threshold 8 = 0.5Å.
COV-R (%) f
MAT-R (A) 4
COV-P (%) f
MAT-P (Å) 4
Models
Mean
Median
Mean
Median
Mean
Median
Mean
Median
CVGAE
0.09
0.00
1.6713
1.6088
GraphDG
73.33
84.21
0.4245
0.3973
43.90
35.33
0.5809
0.5823
CGCF
78.05
82,48
0,4219
0.3900
36.49
33.57
0.6615
0.6427
ConfVAE
77.84
88.20
0,4154
0.3739
38.02
34.67
0.6215
0.6091
GeoMol
71.26
72.00
0.3731
0.3731
ConfGF
88.49
94.31
0,2673
0,2685
46.43
43,41
0.5224
0.5124
GeoDiff
80.36
83.82
0.2820
0.2799
53.66
50.85
0.6673
0.4214
SubgDiff 90.91
95.59
0.2460
0.2351
50.16
48.01
0.6114
0.4791
A.5.4
Comparison with SOTAs.
We compare SubgDiff with 7 state-of-the-ar baselines: CVGAB [26]. GraphDG [39], CGCF [50], ConfVAE [51]. ConfGF [38] and GcoDiff [52]. For the above baselines, we reuse the experimental results reported by [52]. For GeoDiff [52], we use the checkpoints released in public GitHub to reproduce the results. ii) Results: The results the GEOM-QM9 dataset reported in Table 12. From the results, we get the following observation: SubgDiff significantly outperforms the baselines COV-R, indicating the SubgDiff tends to explore more possible conformations. This implicitly demonstrates the subgraph will help fine-tune the generated conformation to be a potential conformation.
i) Baselines:
A.5.5 Domain generelizaion
The results of Training on QM9 (small molecular with up to 9 heavy atoms) and testing on Drugs (medium-sized organic compounds) can be found in table 13.
Table 13: Results on GEOM-Drugs dataset. The threshold 1. 1.25A
Train data
COV-R (%)
MAT-R (Å) 4
Models
Mean
Median
Mean
Median
CVGAE
Drugs
0.00
0.00
3.0702
2.9937
GraphDG
Drugs
8.27
0.00
1.9722
1,9845
GeoDiff
QM9
7.99
0.00
2.7704
2.3297
SubgDiff
QM9
24.01
9.93
1.6128
1.5819
A.6 Sentivity analysis of in k-step same-subgraph diffusion
The results from Table 14 indicate that k 25 is more likely to give the best performance when the diffusion step N is 500. From our experience, and the number of the diffusion step N can maintain certain ratio, e.g. N k = 20 in experiments, and the model can achieve better performance.
21

---

Table 14: The sensitivity analysis for different &: in k-step same subgraph diffusion on conformation generation
k:
lld sampling"
10
25
50
COV-R(Mean) t
89.70
88.06
89.78
89.02
COV-R(Median)
93.96
93.26
94.17
93.21
COV-P(Mean) (%)
49.90
47.38
50.03
48.63
COV-P(Median) (%) f
47.00
47.06
48.31
46.77
MAT-R(Mean) (A)
0.5235
0.2623
0.2417
0.2706
MAT-R(Median) (A) L
0.2710
0.2597
0.2449
0.2709
MAT-P((Men)) (A)
4.7816
4.2922
0.5571
0.7512
MAT-P(MMedan) (A)
0.5378
0.5615
0.4921
0.4995
'The subgruph predictor cannot predict the correct subgraphs
Langevin Dynamics sampling rather than DDPM
B
More discussions
B.1 More related works
Masks on diffusion models. Previous works also share a similar idea of subgraph (mask) diffusion, such as MDM [32], MDSM [19] and SSSD [1]. However, the difference between our SubgDiff and theirs mainly lies in the following two aspects: i) Usage: the mask matrix/vector in SSSD and MDSM is fixed in all training steps, which means some segments of the data (time series or images) will never be diffused. But our method samples the S, Pa, (S) at each time step. hence suitable discrete distribution p(S) can ensure that almost all nodes can be added noise. ii) Purpose: MDSM and MDM concentrate on self- self-supervised pre-training, while SubgDiff serves as potent generative model and supervupe pre-training algorithm. Notably, when S; = So: Vt. SubgDiff can recover to MDSM.
Graph generation models. D3FG [20]: D3FG adopts three different diffusion models (D3PM, DDPM, and SO(3) Diffusion) to generate three different parts of molecules(linker types, center atom position, and functional group orientations), respectively. In general, these three parts can also be viewed as three subgraphs(subset). DifiPACK[60] is Autoregressive generative method that predicts the torsional angle (i = 1,2, 4) of protein side-chains with the condition X1 Xi....d-1. where Xi is predefined subset of atoms. It uses a resionall-based diffusion model to approximate the distribution (ixx/xi......111 1) in which every subset Xs needs a separate score network to estimate. Essentially, both D3FG and DiffPACK can be viewed selecting subset first and then only adding noise on the fixed subset during the entire diffusion process. In contrast, our method proposes to randomly sample a subset from mask distribution p(S) in each time-step during the forward process. [18) proposes an autoregressive diffusion model named GraphARM, which absorbs one node in each time step by masking it along with its connecting edges during the forward process. Differently from GraphARM, our SubgDiff selects a subgraph in each time step to inject the Gaussian noise, which is equivalent to masking several nodes during the forward process. In addition, the number of steps in GraphARM must be the same as the number of nodes due to the usage of the absorbing state, while our method can set any time-step during diffusion theoretically since we use the real-value Gaussian noise. Concurrently, SubDiff [54] is proposed to use subgraphs as minimum units to train latent diffusion model, while our method directly involves the subgraph during the forward process, which is a new type of diffusion model.
B.2 Limitations
One limitation of our method is that it does not currently align with the Boltzmann distribution over the conformations and it cannot accurately estimate the likelihood of sampled conformation. Currently, we generate the subgraph by dividing the molecule with the rotatable bonds. More chemical prior knowledge can be used to guide the mask of the subgraph. Additionally, there is requirement for the number of steps in sampling to be the same as that training, which can be improved with the recent progress of diffusion models such as DDIM.
22

---

C An important lemma for diffusion model
According to [40, 10), the diffusion model is trained by optimizing the variational bound on the negative log- log-likelihood - llog Pa RR), in which the tricky terms are Lt-1 DKr ( R², (
T 2 i t> 1. Here we provide a lemma that tells us the posterior distribution q( RR' It)) used in the training and sampling algorithms of the diffusion model can be determined by ( R R)), q RR' R"). Formally, we have
Lemma C.I Assume the forward and reverse processes of the diffusion model are both Markov chains. Given the forward Gaussian distribution ( R N μn R' o
q(Rt R (NR U2 R o22) and ~ NO, I), the distribution q R R is
g Q R X (R 1
of
o
(R'
P1
V +0
1902 + oj
Parameterizing ) in the reverse process as
NR 1 - (R'
of
0.00 1),
co (R),
f11
VMloz +
11/02 o1
the training objective of the DPPM can be written as
of
((0)
Ro
||e colMINER
Vugo3 + off.
2ufo3
and the sampling process is R²-
of
111
R¹
v o + of
c0 (o(((,
+
VIţ0ž + oo 2,
0102
(38)
where: z 1),
Once we get the variables (/11, 01, (12, 02), we can directly obtain the training objective and sampling process via lemma C.I, which will help the design of new diffusion models.
Proof: Given the forward Gaussian distribution (( RR) = 11 and q( R (RR) (/ P2 o2I), we have
((11)10)) = ,,")) 1122) = -N' R +1
(39)
From the DDPM, we know training diffusion model should optimize the ELBO of the data
(R) 2 E E log PRRO:T qR 1.7 (o))
(40)
R [log
Dia RD) |PRR) ))
( (Daa(((R RR, Rd) Po
R
RR))
reconstruction DETOO
pror manching 1200
denoising mutching tenm
(41)
23

---

To compute the KL divergence DKL(O(R'-(RE) RR) || Po RR)), we first rewrite (R' [R' RR) by Bayes rule
(R' r²)) q (R R'
R DRR RRO)
(42)
(RRP
N pn R of P2 og.1
(43)
NRR; M1 {of ((02)1) 1)
RR* 14 R 192 +
(R²
p2
(R¹ - 12
X exp
20
203
2(0 1 m903)
(44)
(1l P1 11
(1t-
p2 )2
(R' 11
exp
(45)
2
of
o2
of + pi
1
(-2p1 RRRR-11 + pP(R-119
((R'-1)+ - 2422R'1 R
exp
C(R', R²)
2
of
02
(46)
1
2μr R R'R
17(R1-112
(RL-1)2
2 R' R
X exp
(47)
2
of
of
0²
1
1
1
Mi R'
P2 R
exp
(R--1)2
-2
Rd
(48)
2
0
di
of
1
of
+ 13
R
1/2 12RR0
exp
(Rl 1R 2
Rd
(49)
2
0f
of
0
Mo R R
exp
o
+ 1 2
(R' 2
of
+
03
R
(50)
0
of 1902 0703
14 Rd
R
0 2
exp
(R' 1R1-110 -2
of + u
R
(51)
(R' 11211122221 299-21103 1 Rd + uz0 R
exp
off3
of 10
R
(52)
oftuio]
X A(R 1, 1103 Rd + P20fR
0902 1)
(53)
of + 1i
of + 1102
(R
E,(()
We can rewrite our variance equation as (t) = =t](I,, where:
oio3
20(1) of + 13
(54)
From (39), we have the relationship between R6 and R2D:
R R² V of + mNO3E
(55)
11/21
Substituting this into (R,' R2)), we can get
m1oRR Rd + 120
Pg (R))
(56)
of + 1ÃOz
Vottnjoje
1103R + p207
(57)
of + 1i0z
110RR' + a RD Vo
H1
B1
(58)
of 1ojo
1 _RR
of
R²
(59)
M1 of + 1/02
24

---

Thus,
of
0f03
q( (R' RO) x NN 1,
M1 (R²
I)
(60)
Voi + pioz
o + uÍ02
My(R),) (4)
E.(()
Parameterizing
the reverse co Co t)), pfastoj I) and the corresponding optimization problem becomes:
Po PR RR*) of
process
as
(R'
arg DKI (q ( "R',, R2)) (P((T""((R"))
0
mig (N (R! Mag Eg (t)) N (R' ( ; Eg (t)))
(61)
of
0f
arg min
(Rt
(62)
8
(1)
P1 o + 1100
111 v of + 1102
of
arg min 0
(co Co t))
(63)
(t)
M1 of + 1102
2
of 11 o 1503
|1co0 ( 10112]
arg min (t) of
(64)
||co ( |
arg min
(65)
203m3
Therefore, the training objective of the DPM can written as
of E col + V VIo + ofe,
(()) R,
(66)
212033
During the reverse process, we sample R²-1 ~ po (1 process is
Formally, the sampling (reverse)
Rd
1 RC 1n
v 1 + to (T((T( +
of
0102
~~01
(67)
V + o
25

---

This lemma can be easily extended to the conditional version:
Lemma C.2 Assume the forward and reverse processes of the diffusion model both Markov chains. {R') Pezo are the states and y1, y2 are the given conditions. Given the forward Gaussian distribution
119. v1) = (N( q(R²~1RR² 322 = NN u2ll? ,0?.)
; 111 ,,91);;
and ED ~ (0,1), we have the distribution g( , 91, = NN 11 (oî + 1(03))) Thus, the posterior distribution q( RR'-11R3 RR) is
1 (R'
of
0 1).
q(R' R x NN
cD),
P1
V + of
11 + 0
Parameterizing Po(R |R , y1, 8222 in the reverse process
NRR 1 -R*
of
0102 (R',
f11
Vnio oo
1903 + 0i
the training objective of the DPM can be written as
of
(()) - ,e,
||e - c0 + V /noo + off, 0)13
2p402
and the sampling process is
R²11
v V1 of
of
co (o(B,, +
Vužož o 2,
0102
(68)
where: Z ~ I).
D Derivations of training objectives
D.1 SubgDiff (1-same step and without expectation state)
Here, we utilize the binary characteristic the mask vector to derive the ELBO for SubgDiff, and we also provide general proof in Section D.2:
pl 11 S1-T)
log p(
Ey(RET 189)
log
(69)
R HD, (1(11)
p( RTT I Po ( (2*)
log
(70)
I1 a( R R ()+))())))
Rd IT- po (R - )
=Eggne
log
(71)
IT. q( (RR R (()))(s))
p0(se R2))
+ log p(R²" IT.-_
Po R (R*)
log
(72)
[T ((8x)
I al RR' R 1,8))
T
Poo pa (R*)
log p( (RT)) IT.
Po (1' (1²)
Eygrer S
> log
((s))
I a( R ((
mask prediction ter
(73)
(74)
The first term is mask prediction while the second term similar the ELBO of the classical diffusion model. The only difference is the st in (( st). According to Bayes rule, we can rewrite each transition as:
Q(R" RR Rd (9(RR(1())))
if 8: 1
q(té)) ,I",, 8()
a1ft 189)
(75)
p
(R.).
if 8 = 0
26

---

where (a()): := 6( a) is Dirac delta function, that is, (z(z) = 0 * and I_ (z) = Without loss of generality, assume that $1 and sT both equal 1. Armed with this new equation, we drive the second term:
log P(R²)
Po (R ( R²)
(76)
I @ Rs Rt-1 88))
pl R² po A R I Po R IRt)
log
(77)
(R RD) I1 9 R$ R 8t)
( po (IP R2-) IT Po (R*
log
(78)
(R 11 R2 R- R st)
log Po [TTT po (R
po It' (2())
= -E (P))
log
(79)
( [T))
( 1Q-11 1, 12, 8t 1=2
log p(l po [ (2)) (R) [R)
(H (tt)
Pe [tt
EgcReT RR))
+ log
log
R(R' 1 n (R
[C 0]
(80)
p(R²" )pol (po(Te) (Te R²) ( )10)
1R log II
Po ( (1²)
log
+ log
II
8 (1T1)
9( 1 _i BBL---"" R
(E((m=0))
16(()=11)
(81)
Po( 12.) (
p( yoo ( R2
+ log (1T4
po (It" (I2' (2
R log
log
+ log
q(H²" ( (10)
( 12 120)
1E((--1)
(82)
log Po (
log P(R² ( (R))
Po Po (
E.gRRUTTRR) )o)
Σ
log
(R²)
(70)
q(R' R2', R?)
1)
(83)
log po (It (17))
Σ
[log Bog (((R))))
(84)
8 - (Rt)
(R²)
=){1]s==0)
P(RT)
po Po(
E (.
log "( (eo)
Σ
E no)
log q(I2-11 12', 120)
(85)
log Po [T' 1172))
Σ
( [log po (It" ( [e')]
(86)
d (RC)
P(R") log q(RR) R9)
Pou log P (
+E +EqiReieen
+
Σ
(87)
q( 1 RD)
=1(1)a==11
log po ( I(8
Equi (o
(88)
5 R2))
=))
term
decay term
DKLl R2") ((TT)
Σ 16(()=11)
B Dxi(g(n' nt, ne') po (o (It' ((2'))
(89)
prior matching term
denoising matching term
27

---

Here, the decay term represents the terms with 8, 0, which are unnecessary to minimize when we set po (R' |R')) 5 (R)). Eventually, the ELOB can rewritten follows:
E BIRRTITY log Po(sil ((sc)
logp(R") R*)
B Beuiin [log po (RR) R*)
nocossiructio
mask prediction -DR RR) ll ((
( [Dala((** RR) Po Doo(R ((²)))
prior matching term
te(1),111
denoising matching term
(90)
The mask prediction term can be implemented by node classifier and the denoising matching term can be calculated via Lemma C.1. In detail,
( R²) = N V1 R -1 (B)s))
(91)
a( a( = N( - (1 - -I))
(92)
Thus, the training objective of SubgDiff is:
[(0, 0) = R. 2(1 s, 20-1
||c col (1,,,), (o
(93)
To recover the existing work, we omit the mask prediction term (i.e. Let pa(s; R*) = g(s))) of SubgDiff in the main text.
D.2 ELBO of SubgDiff
Here, can derive the ELBO for SubgDiff:
/ p( sur )dr d31:T
log p(tt)) = log
(94)
// p (1.r) ((( R s1 RR)
log
AR daver
(95)
q( (RR.T B R?)) { [/puio (1 (IRL n1P #1:T (1(T)
log
dR dser TRser
(96)
g (R.T, s1. (1))
pl p(10O:T s1:r))
log ( Begree ssrr)
(97)
(H S1.T H()
p(RO: ,1.7)
2 (r p19 s1rr)
log B (sur)) q RD (1.T))((187))
(98)
P(R") I (R' (
RR)
log
IL a(1 (1),)) (()((()))
(99)
P(R") Po (po(s)))))
B .aim Jom R2)
log
(100)
11 Po(s) log p(H I ( ( 1 si)
B .air Ro)
log
(101)
II g(xt)
9 11 12--1,00 1.
log Po(s: ( g(s()
log p(R²) 1 11 po (R' 1
warr sur RD)
11: 9( RM -1
mask prediction term
(102)
(103)
According to Bayes rule, we can rewrite each transition as:
q ( R², s1.)) q(Rt R RO s1.1 )q( q (R' R s14)
(104)
q 11.0 812-1)
28

---

Armed with this new equation, we drive the second term:
p(R)
II Po [
1 [(',,,) st)
log
(105)
Il 0 R Rd 88)
p(R") po ( (
11== Po (o' |R', |R', se)
= -
log
(106)
q( [t 81) 1- q | 1 8t)
log p( po (IT) R$ I1=
Po ( - (8))
- (
(107)
( R2 81 I1t ( R8 ( R2
Ro, 81:0)
pl RD po RR)) R1)s))
pol R R² 8x) 9 [R-11 11, s1:1)
log
log
(108)
(R [10ss)) .1)
(=2
p( ) (o (PPlRt s1)
T
pel R Rt st)
-E
log
+ log
(109)
qR s1)
g ne R2 $n. l R I .1s1)
t-2
(R 1 .11
pl (R Po (R |it 81)
Po 1 17
-E an 189)
log
log
(110)
q [ 1t,, $1)
q 1 |2 R .6 )q( a
p( (2TTT) po R2 [t $1)
g
Po P0 11',
ExcRer (89)
log
log
(( s1:r) + log
gf |
| si
(R 1 Rt R s14)
(111)
p(R )po R R 81)
T
Pe R 100 8t)
log q([eT (1:T)
log
(112)
q(I0=-1 1 17 81:1)
E 189)
[logpo (It") R ,s1]]
(113)
)(H²")
T
log (o( (H' (1'tss))))) si)
+B.gRuoT (89))
log 2RR' R s1rr)
E gR G
(114)
qRR IR,,, R s1,)
-F (899)) [logro( R R². s1)]
(115)
p(RT))
1 11,,, 8x)
+ pe
log ( [7,, s1TT)
E Eg w w²-1 - saal
log 9(12'-1 12', 17.0 s11))
(116)
-E [R s1)]
(117)
rcconstruction term
T
-Batara) Pxlggi HP, S p((?))
B (R ard
(Kala(Hf * Hº s1;) Po (o
"1 s,])
prior matching
denoising matching term
(118)
Eventually, the ELOB can be rewritten as follows:
(o(x
log p(I())
E a log ((x))
Egens in 188) (o nt' s1)]
(119)
neconstruction torm
predicction
Ey Dx (g(" IA s1T) p((e))
E DKala(Hf " RP s1:)) |l po (1 '128 ss))]
prior matching
denoising matching term
(120)
The mask prediction term can be implemented by node classifier 8y- For the denoising matching term, by Bayes rule, the q( RR' R2 $1:)) can be written as:
q 1R" R² s1x) g(R* q R
Rd,
s1a ( R $18-1)
(121)
R8 18,, 81a)
29

---

For the naive SufgDiff, we have
Q R' R s1=) N= v1 1 B, (3)
(122)
1 R², 81:--1 WR vT-
--1)
(123)
Then the denoising matching term can also be calculated via Lemma C.1 (let R R RR) := (
q(t [R--,, R RR the training objective of SubgDiff is:
$1:t) qft2 ( == q 1 $1x-1) and Po Po (R) = Pa ()( Thus,
8,
C( 0) = Ro 2(1 5,) 3o
||e ea co(v + v(1 t,,,, ABCE 80 s0 (9.
(124)
D.2.1 Expectation of 81:7
The denoising matching term in (120) can be calculated by only sampling (R', 86) instead of (R', 81a). Specifically, substitute (121) into the denoising matching term:
log pa 1 | 8t)
(125)
q( R³, R $1:t)
pe [1 1AL
Rt . St)
log R 1!- 18 elg( FC
(126)
11
9 12 .010s)
log Po R² 'R 8c)
=Eqcst, (a
log logg(11' s10-1)
(127)
9(12 18- v( RC Re...))
2 E (1
(Eylera) log Po (R'-1
R . ( s1:1)]
(128)
log Becan ,11 log B 1) 12 12 .1:)
(129)
i=(EE 110)
== 11-1 1100 (2)
Po IR', 8x)
Bgione) log
Ey(ae) log (E, 1( E.coe) ( R
(130)
Pol I MI
Bginnd log
(131)
R
otre 17 (010))
Pe [R' [
log 9(8 Rt- R a On, (82 (12 11e (yin)
(132)
denoising matching term
Pe (1 8t).
EygRy, ne- (e 18)))
log @(H'-1 [ RD, $14)
(133)
Eag "Iil "1tt, Rt s1a) Il po (Ht "1 st)
(134)
denoising matching
Thus, should focus calculating the distribution
E.R E. R st (t)q(E,
1 RP, s1a)
(135)
((R) 129 81:2)
By lemma C.I, if we can gain the expression of q(' E. R st) and R', RR9), we can get the training objective and sampling process.
D.3 Single-step subgraph diffusion
D.3.1 Training
Step 0 to Step (Ro - R(-1);: The state space of the mask diffusion should be the mean of the random state.
30

---

E.Rt ~N(E,, V1 I.-,, ,811)
(136)
q =1t)= .v.lt", -.))
(137)
Form (137), we have:
R' /1 Rt + /
(138)
ER' = (pvI N 31 + 1 -EE) + PV/
(139)
(pVVI 8 (DV 1
(pv -111)
3r + p)py 3 16 + PV B
(140)
= CPVI - p) -1 -
(pV 1 Bi + 1 B + [pV B1
(141)
(142)
(PV 1 Bg 11 - +
(PVI - Bz+ + p)pv 32) + tco
j=2
3=3
(143)
(PV 1 - B, + 1 1) +
E 1=
(pVI- - -1-
(144)
VaR +
Vappy B.1
(145)
j=it1
0 R +
Σ II
333ico
(146)
1=1 j=
Ro +p,
Σ
02
3EED
(147)
0
1=1
(148)
where 0; :0:=(10 (=ppVI-P 1 p)
and ā. IT: 0;.
((BR'RR))= = NRR:; Vo
BTT)
(149)
II: Step 1 to Step t (R' + R³): We build the step t 1 + t is a discrete transition from
(R' (R!-1 RD), with
t-1
( RL. (RD) NRR
oo Rd "
II
(150)
t=1
1=f+1
qR' RRE,F DEERRRL (,)) = -N( / B, -
(151)
31

---

R = /1 B + V
(152)
/I - 3t Vāi- +P
0t11 BrE0
+ v 11
(153)
a;
V
8e 201//111112 RR + +PV1 Sti1
0it-1
Br00 V8rBar0-1
(154)
āi
V BeBev Bollooon110 1R + p2 (221-1880 81.3)) Σ
0h-1 1
Be + 8e3.50
(155)
a
Step 0 to Step (R + R²)):
q( R*) = /
(156)
s,Bi (0'(1((011110 - -t,B)) ā 3, + st
(157)
Thus, from subsection 1.2.1, the training objective of 1-step SubgDiff is:
Lsimpte (0, 0) = E, s [selle - co 1)1 (CE 80 (a (R))
(158)
where (8,, 80) = 8; log 8 R'' + +1 8t) log (1 so (F ()) is Binary Cross Entropy loss. However, training the SufgDiff is not trivial. The challenges come from two aspects: 1) the mask predictor should be capable of perceiving the sensible noise change between (t - 1)-th and t-th step. However, the noise scale 3, is relatively small when / is small, especially if the diffusion step is larger than a thousand, thereby mask predictor cannot precisely predict. 2) The accumulated noise for each node at (1 1) step would be mainly affected by the mask sampling from 1 to L 1 step. which heavily increases the difficulty of predicting the noise added between (t s)-step to t-step.
D.3.2 Sampling
Finally, the sampling can be written as:
(1 (1-
(1BssB))p'>> Et=
1-1 āe Eis a,
B, + stBe R² (srB1y p2(1 31
Di=l 2:-1
3 + s1 te(R(t)
R!-1
+ 0
VI S, + (1 111-1,,111 B,B,)p
مo- Gr-13,)
i==
dr.
(159)
s5 / 12 s Σ i=1
8e B, + S
R²
60 (R', ()) + 0,z (160)
V1 - srBr
vI 31 (8 (1 B22
Σイ=1
Ge-1
Bi)
1 R¹
srB1
co t) + 022 (161)
V1 s
v1 - scfe v 81,8 + (1 st
S,Bx)p''
أeーt 3 a,
-=
(162)
where 81 = SR' t) and
so (R', t)p
Σ1
d. 3
Tt
,o( 2² So (R', ())&k)
مو- B
(163)
32

---

Fonsi
Crari Szmel
Cuest Sur
Crse Funii
Ponet
R
Rehe-1
R
Figure &: An example of k-step same subgraph diffusion, where the mask vectors same as 8km+1 from step km (m + 1)k, m E N+
E Expectation state distribution
The state space of the mask diffusion should be the mean of the random state.
E., ~ N(ER'; / R R ,c11)
(164)
Form Equation 137, we have:
R' = VI 8, + v B1BBE-111
(165)
ER = (PV1 - 31 +1 D)ER + +PV 3
(166)
= =PPVI 3 +1 P)(PV 3i-1 + 1 P)ERR*
(167)
+ 3r + 1)D 3r- 12 + /v//v//// B
(168)
= 3: + 1
P
(169)
+ +VI(PVI - 3r +1 + |PV Ec--
(170)
(171)
[I -
(172)
- +1- 31
[pVI- Br 1
p)pv
B21
teo (173)
j=3
II (PV 1- 8, + 1-
Σ
1 +1-
(174)
Ro R
Vaspy BiPED
(175)
j=1+1
v Ro R +p
Σ
II
3,53
(176)
-
!
R Rd +p,
Σ
3,50
(177)
σi
i=1
(178)
where o= = (pv1 11 and C = IT 04-
Finally, the Expectation state distribution is:
(ERR(RR) = (VEE II
R P
Σ II a,BiT)
(179)
==1
j=it
33

---

The derivation of SubgDiff
When is an integer multiple of k
t/k
ER
I
(1-3))+1 D)
(180)
f=1
111)+1++1
2
k/k
t/k
kj
+
(p
II 1-- p/P
1
(1 - 3:)
cD
(181)
1=1
1=++1
1J8+1
13
t/k
t/4 Σ
II
Vaj +p
II ay(1 j=1+1
II (=(1-))++
(1 -3
(182)
1/k
Vorkk Rd +p
Σ
04/k
II B,))
(183)
ii
(=(1-1)8+11
where 03- (PV/ / [1 118+1 (1 B,) + + 1-)
When E N, let = [(t 1)/kj. and we have
R¹
(1 8(8km+1) BR 1
II
(1 B.8km- mmck
(184)
kmm++1
kmm++1
m
kl
II (1 Bi8km+1)
Vam Rd +
II
(1 B:))eo
(185)
kmm+1
(=(1-))++1
II (1 Biskmt1)em
(186)
II
RivaR
(187)
+m
II ii
2
m
II
(1- 3.))
1
%
E0
(188)
i=
1)k+1
1=knn+1
where 7 = Biskm+1-
((11119)) = N ((etm
II
Rm
(189)
i= kmm+1
p
a'm
II
(1-
+1
1) (190)
i=
(1-1)8411
i=mm+1
Let 7. = II Yt. and Br = 11 - (1-31)
R(R* N(AA
V Ft
RamR"
Zt
p²
m
Om (1
BR )
7t )m))) (191)
7km
Fkom
α
B11111
Fkm
34

---

F.0.1 Sampling
[41
V1 8km Be:
(192)
0 = 5km+1 1t
(193)
7t-1 V0m
122
(194)
Tkm
02
2
---1
m
Omm
II
(1- +
2--1
(195)
Tkm
On
7km
1=1
im
(1-1)
According to the Lemma C.I, we have
of
Rº- Ht' V 11903 + of
co((e())
0102
(196)
P1
VMÃ03 + o
(11"
(197)
1 8km + 1
8kmm + 13
co (R', ())
skm+1 13)) Faam p²
2i=1
Cu (1
[1==111114 11++1 ([)) + 1
hf
201) ) + s skm 1,3
ai
Fkm
(198)
v 5kmm-18ey
8e-1 p² D1=1 ai
am (1
II kl
1== (1- (1-1))+111 -3,)
9--1 Ftom
Frm
8km+1 (3){
7e-1 Tkmm
p2
i=
m
am
(1=-11))++1 (1 -Bi)) 3t-1 Tkm
kl
5km 13
(199)
The schematic can see Figure 5.
35

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have summarized our contributions at the end of the introduction.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See Appendix Section B.2.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: See Appendix Sections D.1.
D.2, E,
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide experimental details in Appendix Section
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The source code is available at this repo
6. Experimental Setting/ettalls
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Appendix Section A. Specifically, the hyperparameters are in Table 6 and Table the Appendix.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: When possible, we report the mean and standard deviation of three random seeds.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: See Appendix Section A.3.
9. Code Of Ethics
36

---

Question:
Does the research conducted in the paper conform, in every respect, with the
NeurIPS Code of Ethics https://neurips: httppp////orrrs...cpppii//////a/
Answer: [Yes]
Justification: This research conducted in the paper conforms with the NeurlPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: This paper presents work whose goal is advance the field of Machine Learning. There are some potential societal consequences of our work, none of which we feel must be specifically highlighted here.
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: This paper poses no such risks.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The materials in this paper are used with permission and properly cited.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The weights of trained models are provided in this repo
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing or research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
37


Reviews:
[{"review_id": "6IgFAH0dtN", "reviewer": "Reviewer_BpDA", "summary": "The paper proposes a diffusion-based pretraining method using subgraphs to learn enhanced molecular representations. Unlike previous methods which normally add noise to every atom, this paper proposes adding noise based on subgraphs. The method is evaluated on various downstream tasks to demonstrate its effectiveness, such as 2D and 3D property prediction tasks.", "questions": "See weaknesses.", "limitations": "Discussed.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper offers an interesting perspective on existing molecular diffusion models and proposes several approaches to address them accordingly, which is inspiring.\n\n- The paper presents many experiments and analysis to illustrate the results, which could provide valuable insights to the community.", "weaknesses": "- Although the proposed method effectively addresses the identified limitations, it does not provide much chemical intuition for the design. The authors claim that existing methods neglect the dependency in substructures, but the proposed method does not seem to integrate or learn the inherent molecular substructure information. The decomposition does not seem to be based on significant chemical knowledge, and the interactions or relations between various substructures are not explored. The motivation is not heavily grounded in domain knowledge.\n\n- The method is a little bit confusing. The entire training process includes many steps and three key training objectives: subgraph prediction, expectation state, and k-step same-subgraph diffusion. The paper mainly describes each component separately, but it is unclear how the three objectives are leveraged during the training. Some details are also not clear. For example, “The mask vector $s_t$ is sampled from a discrete distribution $p\\_{s_t}$ (S|G)”—does $p_{s_t}$ vary for each molecule? How exactly is the distribution obtained? Also, an overall framework could help better understand the process.\n\n- The diffusion steps number is 5000, which is quite large. Molecular property prediction datasets normally contain small molecules. Are there any specific reasons for using such large steps? What are the computational costs?\n\n- The proposed method seems impractical due to its complexity, while the performance improvement is not very significant. It is difficult to justify the trade-off between model complexity and performance. Perhaps the authors could provide a more in-depth discussion to advocate for their method compared to other baselines, beyond just the prediction performance.\n\n- Some ablation studies are not conducted. For example, k seems to be an important h-param. What are the effects of different k values? Also, what about the performance of applying only one training objective?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments and positive feedback!\n\n**[Weakness 1. Chemical intuition]**\nThank you for your insightful comment! We agree that chemical intuition and domain knowledge are crucial in molecular representation learning, and we'd like to address how our method incorporates these elements:\n(1) While our method may not explicitly learn predefined substructures, it does capture atomic dependencies within the substructure through diffusion on subgraphs. This approach allows the model to implicitly learn relevant substructure information from the data, rather than relying on predetermined chemical fragments.\n(2) Further, we sincerely clarify that the subgraph is obtained by decomposing the molecule by breaking the rotational bonds. This torsional-based decomposition method can capture the inherent molecular 3D substructure information. For example, the benzene ring is always preserved in the same subgraph since the π bonds in benzene rings aren't torsional. The downstream tasks on 3D property prediction also demonstrate the effectiveness of our method.\n\n\n**[Weakness 2. Overall methods]**\nWe thank the reviewer for this valuable feedback. We acknowledge that our method involves multiple components, which may have led to some confusion. We apologize for any lack of clarity in our presentation and would like to provide a more comprehensive explanation of how these components work together:\n(1) As we described in section 4.5, in the forward process, we use **expectation state diffusion** to get the state $km$ ($m:=\\lfloor t/k \\rfloor$) from the initial state $0$ (Phase I). Then the $(t-km)$-step **same-subgraph diffusion** is employed to get the state $t$ from state $km$ with the same subgraph $s_{km+1}$ (Phase II). In the training step, we use the denoising loss and the **subgraph prediction** loss to train the model (Equation. 16).\n(2) These components are not independent but rather complement each other:\n- The subgraph prediction task helps the model learn local structural features.\n- The expectation state reduces the complexity of $R^t$ and removes the unstable sampling of the masks $s_{1:km}$.\n- The k-step same-subgraph diffusion accumulates more noise on the same subgraph from $km$ to $t$ for facilitating the convergence of the subgraph prediction loss.\n\n(3) The subgraph is represented by a mask vector $s_t$. The distribution $p_{s_t}(s|G)$ is implicitly defined by the torsional-based decomposition method. One can get two subgraphs by breaking a torsional bond. Thus there are $2n$ subgraphs if the molecule has $n$ torsional bonds. $p_{s_t}(s|G)$ is a random distribution over the $2n$ subgraphs. And it varies for each molecule. \n\n**[Weakness 3. Diffusion steps number]**\nThank you for bringing up this important point about the number of diffusion steps and computational costs. From our experiments, we find that using 5000 steps can obtain better performances, as shown in Table 3 for conformation generation. We also evaluate the downstream 3D property prediction tasks when pretraining with different steps, and the results are shown in **Table #3** of the general response.\nThe large diffusion step will not increase the computational cost for the downstream tasks, as we use the diffusion denoising loss for pertaining. We don't need to do the sampling process that requires high computational overhead when we predict the molecular property.\n\n**[Weakness 4. Complexity]**\nThank you for bringing this important concern to our attention. \n(1) As we have shown in the paper, our method significantly outperforms baselines on 3D molecular property prediction, as shown in Table 1 in the main text and Table 8 in the appendix. In addition to the prediction performance, our method also brings the sampling efficiency and generalization ability. The results from Table 3 show that SubgDiff significantly outperforms the baseline when adopting fewer diffusion steps. As shown in Table 4, SubgDiff consistently outperforms the other baselines in the cross-domain conformation generation tasks.\nThe above discussion will be added to the final version.\n(2) Despite the multiple components of the algorithm, the method is practical to implement. We train the model using the objective in Eq(17), which is similar to the typical diffusion model DDPM. The main difference with DDPM are the noisy version $R^t$ computed from Eq(16), i.e. $R^{t} = \\sqrt{\\frac{\\bar\\gamma_{t}\\bar\\alpha_{m}}{\\bar\\gamma_{km}}}R^0 + (\\frac{\\bar\\gamma_{t}}{\\bar\\gamma_{km}} p^2 \\sum_{l=1}^{m} \\frac{\\bar\\alpha_{m}}{\\bar\\alpha_{l}}(1-\\frac{\\bar\\beta_{kl}}{\\bar\\beta_{(l-1)k}}) + 1-\\frac{\\bar\\gamma_{t}}{\\bar\\gamma_{km}})\\epsilon$, where the weighting coefficients can be obtained simply by scalar calculation. This brings very little complexity. In practice, the training speed is similar to the original diffusion model (GeoDiff). In addition, we also provide the source code [here](https://anonymous.4open.science/r/SubGDiff/README.md), which will make it easy to use our method.\n\n**[Weakness 5. Ablation]** Thank you for the constructive suggestion. (1) We provide the sensitivity analysis of $k$ in $k$-step same-subgraph diffusion. The results are shown in **Table #1** of the general response.\n(2) When only using the denoising loss as the training objective, without subgraph prediction loss, the method cannot be used for sampling due to the inaccessible $s_t$ during inference. Nevertheless, the pure denoising objective can be used for pertaining and the results of the downstream force prediction tasks are shown in **Table #2** of the general response.", "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": "", "extra_scores": {}}, {"review_id": "8mKcvVtaSi", "reviewer": "Reviewer_Djew", "summary": "The paper presents a new denoising diffusion probabilistic model (DDPM) named SubgDiff, designed to enhance molecular representation learning by incorporating substructural information into the diffusion process. SubgDiff introduces a mask operation that selects subgraphs for diffusion, aiming to better capture the dependencies among atoms within substructures. The method includes techniques such as subgraph prediction, expectation state diffusion, and k-step same-subgraph diffusion, which together are intended to improve the learning of molecular properties related to 3D conformation. The paper claims superior performance on various downstream tasks, particularly in molecular force predictions.", "questions": "1.\tHow is the stability of subgraph selection ensured during the implementation of the multi-step subgraph diffusion (k-step same-subgraph diffusion) process? Is there a risk of accumulating long-term errors due to inappropriate subgraph choices? \n2.\tDoes the training strategy for the subgraph prediction model $p_\\partial(s_t|R^t)$ include handling of imbalanced data? Specifically, how does the model avoid bias towards subgraphs that appear more frequently in the training data? \n3.\tCan this method also be applied to broader graph representation learning tasks, such as node classification and clustering?\n4.\tDiffusing considering substructures is an interesting aspect from my perspective; besides the functional group substructures in molecular graphs, community substructures [1] are prevalent in other networks. If considering community structures during the diffusion process, what insights might the authors have? This is an interesting question, as community structures are widely present in biological and social networks.\n\n[1] 'Community detection in graph: An embedding method' in IEEE TNSE", "limitations": "This paper introduces a complex diffusion model that leverages subgraph structures to enhance molecular representation learning. However, more limitations should be discussed, such as overfitting, computational efficiency, and generalization across diverse molecular structures.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tUtilizing subgraphs for diffusion is a novel and intriguing exploration. \n2.\tThe experiments are thorough and demonstrate the effectiveness of SubgDiff across a range of molecular prediction tasks.", "weaknesses": "1.\tThe subgraph prediction model is trained on highly specialized datasets, which might pose a risk of overfitting. \n2.\tThe model integrates multiple complex diffusion stages, which could complicate the training and debugging processes and make them difficult to optimize. Particularly, adjusting hyperparameters and verifying model stability might require additional effort.\n3.\tGiven the complexity involved in the expectation state and k-step diffusion processes, the model may have high demands on computational resources.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments and thoughtful feedback on our work.\n\n**[W1. Highly specialized datasets]** Thanks for highlighting this important concern. The subgraph prediction model shares the molecular encoder with the denoising network and has an additional classification head. The subgraph predictor is trained together with the denoising network on the same dataset. We use the commonly used datasets GEOM and PCQM4Mv2 to ensure fair comparisons with the baselines. The overfitting risk can be avoided by using a larger dataset for pretraining. Our experiment in the cross-domain task shows that our proposed method has good generalization ability and may alleviate the overfitting problem to some extent.\n\n**[W2&W3. Multiple stages and complexity]** We appreciate the reviewer's concern about the complexity of SubgDiff. While we acknowledge our model is more complex than the vanilla diffusion model, we believe this complexity is acceptable and can be justified by performance improvements, especially in molecular force predictions. To address your concerns:\n- Training and debugging: compared to the original diffusion model, our model only adds a classification head for subgraph prediction and one additional loss. So the complication of training and debugging can be under control.\n- Optimization and hyperparameters: Our method brings two more hyperparameters: the weight $\\lambda$ for subgraph prediction loss and $k$ for the $k$-step same-subgraph diffusion. We conduct ablation studies (see the Tables in the general response) and provide recommended hyperparameters based on our experiments.\n- Model stability: We implement gradient clipping, careful learning rate scheduling, and conduct long-term training experiments to ensure stability, which are the common techniques for diffusion model training. \n- Computational resources: In practice, compared with typical diffusion models, the introduced expectation state and $k$-step same-graph diffusion do not bring much computational overhead. We explain it from training and inference respectively. \n(1）Training: The expectation state and k-step diffusion only affect the weights of adding noise. SubgDiff can directly compute the $R^t$ from $R^0$ like the classical diffusion model as shown in Eq(16) of the main paper, i.e. $R^{t} = \\sqrt{\\frac{\\bar\\gamma_{t}\\bar\\alpha_{m}}{\\bar\\gamma_{km}}}R^0 + (\\frac{\\bar\\gamma_{t}}{\\bar\\gamma_{km}} p^2 \\sum_{l=1}^{m} \\frac{\\bar\\alpha_{m}}{\\bar\\alpha_{l}}(1-\\frac{\\bar\\beta_{kl}}{\\bar\\beta_{(l-1)k}}) + 1-\\frac{\\bar\\gamma_{t}}{\\bar\\gamma_{km}})\\epsilon$, where the weight calculations are scalar and do not require much computational overhead. In practice, the training time of one batch is similar to the GeoDiff baseline.\n(2) Inference: The sampling computation also only needs to calculate some scalar coefficients as shown in Eq（18). From Table 3, we can see that SubgDiff needs fewer sampling steps compared to baseline GeoDiff, demonstrating the computational efficiency of the proposed method.\nWe also provide the source code [here](https://anonymous.4open.science/r/SubGDiff/README.md) to make SubgDiff accessible.\n\n**[Q1. Multi-step subgraph diffusion]** Thanks for the insightful question! During training, the subgraph will be randomly sampled from the subgraph set generated from the torsional-based decomposition method, which randomly selects a torsional bond among the given molecule and breaks the molecule into two parts. Each part is a connected subgraph and then randomly chosen as the subgraph. This predefined subgraph set will ensure the appropriateness of subgraph choices. For each time step, the probability of each node being included in the diffusion process is 0.5. We use a relatively large diffusion time step to make uniform diffusion for each node and ensure the avoid the risk of accumulating errors.\n\n**[Q2. Imbalanced data]** Thanks for the valuable question. The current training strategy does not consider this issue, so as to make a fair comparison with the baselines. In our subgraph sampling strategy, each node has the same probability of 0.5 to be included in the diffusion process. So as the diffusion step is large, our method is comparable to the baselines w.r.t the same training dataset. However, the imbalanced data is an important issue. It may be addressed by clustering the dataset and using a stratified sampling strategy during training.\n\n**[Q3. Broader graph representation learning tasks]** Thanks a lot for the insightful comments. In general, the introduced subgraph prediction can be used as an auxiliary loss for representation learning, if there is subgraph information available. According to the applications, our framework can be easily adjusted to incorporate substructure information into the representation learning model.\n\n**[Q4. Community substructures]**\nThanks for raising this interesting question regarding the potential application of the substructure-aware diffusion process to community structures in other networks. We're excited to share our thoughts on this matter:\n- The parallel between functional group substructures in molecular graphs and community structures in other networks is indeed compelling. This analogy suggests that our approach could potentially be adapted to other domains with hierarchical or modular structures.\n- Applying a similar diffusion process to networks with community structures could potentially yield several insights, such as community evolution over time, inter-community interactions, and hierarchical community structure.\n- There exist potential applications in other domains such as biological networks or social networks. In biological networks, a community-aware diffusion model could potentially: a) Predict the impact of perturbations (e.g., gene knockouts) on functional modules; b) Identify key regulatory hubs that influence multiple communities.\n\n**[Limitataions]**\nThanks so much for the advice. We will add those discussions to the revised paper.", "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": "", "extra_scores": {}}, {"review_id": "t66lUMKe4h", "reviewer": "Reviewer_VBzc", "summary": "The paper proposed SubgDiff which is a diffusion model used in self-supervised learning setup to enhance the molecular representation learning. It introduces motif enhancement during the diffusion process to force the model to learn more structure information.", "questions": "Please refer to Weaknesses.", "limitations": "No", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The idea of enhancing motif information in the diffusion process is promising.\n2. The paper is well-organized and easy to follow.", "weaknesses": "1. The authors directly use the baseline results from the MoleculeSDE paper in their table; however, the results for MoleculeSDE are significantly lower than those reported in the original paper.\n2. The proposed method is more like a graph diffusion model than a molecular representation learning model. It is limited to representation learning within a self-supervised learning framework. It would be beneficial to explicitly state this in the abstract or introduction.\n3. The paper lacks an ablation study to evaluate the contribution of each component of SubgDiff to molecular representation learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper and sharing these important points. These comments highlight several areas where we can improve the clarity and thoroughness of our paper. We acknowledge the issues raised and would like to address each point.\n\n> W1: The authors directly use the baseline results from the MoleculeSDE paper in their table; however, the results for MoleculeSDE are significantly lower than those reported in the original paper.\n\n**A1:** Thanks for your careful review! We use the released code provided by the MoleculeSDE and report the reproduced results. Our method applies the same pretraining framework as MoleculeSDE except for the diffusion model. Specifically, we replace the original diffusion model in MoleculeSDE with our SubgDiff. Hence, we used the reproduced results from MolecularSDE instead of the results reported in the paper for a fair comparison.\n\n\n> W2: The proposed method is more like a graph diffusion model than a molecular representation learning model. It is limited to representation learning within a self-supervised learning framework. It would be beneficial to explicitly state this in the abstract or introduction.\n\n**A2:** We agree with the reviewer that our method bears similarities to graph diffusion models and that its primary focus is on representation learning within a self-supervised framework. Indeed, we use the denoising objective in diffusion as the self-supervised task. We acknowledge that we should have been more explicit about this positioning of our work. To address this:\na) We will revise the abstract to clearly state that SubgDiff is a graph diffusion model focused on molecular representation learning in a self-supervised context.\nb) In the introduction, we will add a paragraph explicitly discussing the nature of our model and its place within the broader landscape of molecular self-supervised learning.\nc) We will review the entire manuscript to ensure consistent and clear communication about the scope and nature of our method.\n\n> W3: The paper lacks an ablation study to evaluate the contribution of each component of SubgDiff to molecular representation learning.\n\n**A3:** Thank you for the helpful suggestion. (1) Our method contains three components, where the subgraph prediction and expectation state are designed to work together for the sampling. Without these two techniques, the method cannot be used for sampling due to the inaccessible $s_t$ during inference. For the $k$-step same subgraph component, we provide the results concerning different $k$ to demonstrate its significance, as shown in **Table #1** in the general response. \n(2) The expectation state and subgraph prediction can be evaluated in the self-supervised learning context. The results of the downstream force prediction tasks are shown in **Table #2** in the general response.\nThe results suggest that the subgraph prediction loss plays a more important role in molecular representation learning.", "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": "", "extra_scores": {}}, {"review_id": "tUDiUIgbO9", "reviewer": "Reviewer_aMXv", "summary": "SubgDiff  is introduced to improve molecular representation learning by integrating substructural information into the diffusion model framework. It offers three key technical contributions (subgraph prediction, expectation state, and k-step same subgraph diffusion) to enhance the network's understanding of molecular substructures. Experiments were carried out on several downstream tasks, particularly molecular force predictions.", "questions": "* Does the denoising processing require to start with a R^T that has a clear topology? If not, how to sample s_t?\n* How to select k in the k-step same-subgraph diffusion? Any general guidance?\n* How are the SubgDiff learning results fed into downstream tasks?", "limitations": "None", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "* Incorporate substructural information into diffusion model", "weaknesses": "* The denoising process need better explanations.\n* Diffusion models excel at generating new samples. The application of SubgDiff to molecular property prediction/classification does not show its strengths. The state-of-the-art results are missed in Section 5.1. Section 5.2 doesn't compare generation results with the state-of-the-art.\n* Explain COV-R and MAT-R in section 5.2", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive suggestions and useful feedback!\n\n> W1: The denoising process needs better explanations.\n\n**AW1:** Thanks for the suggestive advice. We will further describe the details of the denoising process in the paper. The only difference from the traditional diffusion backward process is that when $t\\%k=0$, we use the subgraph predictor to update the subgraph to be denoised. Additionally, the whole denoising process can be referred to in Algorithm #1 below.\n\n> W2: Diffusion models excel at generating new samples. The application of SubgDiff to molecular property prediction/classification does not show its strengths. The state-of-the-art results are missing in Section 5.1. Section 5.2 doesn't compare generation results with the state-of-the-art.\n\n**AW2**: Thanks a lot for the insightful comments. \nWe agree that diffusion models are good at generation. As a generation model, the objective of diffusion can be used as a task for self-supervised learning. This is one of our motivations and our primary focus is on improving the diffusion model for representation learning. Our experiments on molecular property prediction show that our approach outperforms various molecular self-supervised learning methods and the original diffusion model. \n\n(1) In Section 5.1, our method uses the same framework as MoleculeSDE, which is the state-of-the-art technique for molecular pretraining. The experiments also demonstrate our method can gain significant improvement over the other self-supervised techniques. Besides, our method significantly outperforms MoleculeSDE, especially in 3D property prediction (Table 1 in the main paper), which indicates the effectiveness of our method by incorporating the information of substructures. \n\n(2) In Section 5.2, we focus on evaluating the substructure's effect on the diffusion model itself and the experiment verifies the success of our design (Table 2 in the main text). Our results show that our approach is able to improve the sampling efficiency and generalization ability, compared to the original diffusion model. We also discuss other conformational generation models in the related work and provide the performance of other SOTA methods in Appendix Table 12.\n \n> W3: Explain COV-R and MAT-R in section 5.2\n \n**AW3:** Thanks for pointing this out. The definition of COV-R and MAT-R can be found in Eq (32) and Eq (33) in Appendix A.5.2. \n\n\"Let $S_g$ and $S_r$ denote the sets of generated and reference conformers respectively, then the Coverage and Matching metrics  can be defined as\n\n$COV-R(S_g, S_r) = \\frac{1}{ |S_r| }  |\n   \\lbrace C\\in S_r |  \\operatorname{RMSD}(C, \\hat{C}) \\le \\delta,  \\hat{C} \\in S_g \\rbrace\n| $\n\n$ {MAT-R}(S_g, S_r) = \\frac{1}{| S_r |}\n    \\sum\\limits_{C \\in S_r}\n    \\min\\limits_{\\hat{C} \\in S_g} \\operatorname{RMSD}(C, \\hat{C})$\nwhere $\\delta$ is a threshold.\"\n\nWe will put it into Section 5.1 in the main paper. \n\n    \n> Q1: Does the denoising processing require starting with a $R^T$ that has a clear topology? If not, how to sample s_t?\n\n**AQ1:** Thanks for the question. Our diffusion model is able to generate the molecule conformation conditioned on the molecule graph. So the graph topology is given. But the initial conformation $R^T$ is sampled from the Gaussian noise. And the $s_t$ is predicted through the subgraph predictor $s_\\vartheta(\\mathcal{G}, R^t,t)$ trained by the objective in Eq (17). Algorithm 2 in the main paper and Algorithm 4 in the appendix also describe the whole process. We also put the algorithm #1 below. \n\n- **Algorithm #1: Sampling from SubgDiff**\n----\nSample $R^T  \\sim \\mathcal N(\\mathbf 0, \\mathbf I)$ \\//Random noise initialization\n\n**For**{t = T **To** 1}\n\n{\n   1.  $\\mathbf z  \\sim \\mathcal N(\\mathbf 0, \\mathbf I)$ if $t>1$, else $\\mathbf z =\\mathbf 0$ \\// Random noise\n\n  2. **If** $t\\%k==0$  or $t==T$: $\\hat{s} \\gets s_\\vartheta(\\mathcal{G},R^{t},t)$ \\//subgraph prediction\n\n   3.  $\\hat\\epsilon \\gets \\epsilon_\\theta (\\mathcal G,R^{t},t)$ \\// Posterior\n\n   4.  $R^{t-1} \\gets$ Equation 18 \\// sampling coordinate\n    \n}\n\n**Return** $R^0$\n\n\n\n>Q2: How to select k in the k-step same-subgraph diffusion? Any general guidance?\n\n**AQ2:** Thanks for this valuable question. $k$ cannot be too large. A large $k$ will make the subgraph prediction task too easy such that the model learns little substructure information. However, a very small $k$ makes it difficult for the subgraph predictor to converge. From our experience, $k$ and the number of the diffusion step $N$ can maintain a certain ratio, e.g. $N:k=20$ in our experiments, and the model can achieve better performance. In practice, $k$ is a hyper-parameter that can be tuned according to the sampling performance.\nWe also provide the sensitivity analysis of $k$ in Table #1 of the general response.\n\n\n>Q3: How are the SubgDiff learning results fed into downstream tasks?\n\n**AQ3:** Thank you for the important question. The denoising network $\\epsilon_\\theta (\\mathcal G, R^{t},t)$ contains a 3D encoder based on SchNet and a 2D encoder based on GIN. After we finish the training of SubgDiff, the encoders can be used as the pretrained models for the downstream tasks. This feeding method follows the conventional generative self-supervised techniques, e.g.MolecularSDE[1] and Denoising[2]. We also provide the details in Appendix A.4.\n\n[1] Shengchao Liu, et al. A group symmetric stochastic differential equation model for molecule multi-modal pretraining. ICML 2023\n\n[2] Sheheryar Zaidi, et al. Pre-training via denoising for molecular property prediction. ICLR 2023.", "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": "", "extra_scores": {}}]

Author rebuttal:
## General Response

Dear reviewers,

Thanks to all the reviewers for your time and effort during the review process and the constructive advice. In addition to the response to each reviewer individually, we conduct the ablation study and sensitivity analysis of the k-step same subgraph and diffusion steps for our method. The results are shown below in Table #1, Table #2, and Table #3.
We also provide curves of subgraph prediction error during training in the attached PDF.

**Table #1. The sensitivity analysis for different $k$ in $k$-step same subgraph diffusion on conformation generation.**
| k | 1(ld sampling*) | 10 | 25| 50 |
|----------|---------|----------|----------|--------------|
|  COV-R(Mean)     ↑    |     89.70    |     88.06             |     **89.78**     |       89.02    |
|   COV-R(Median) ↑       |  93.96      |      93.26      |  **94.17**         |    93.21      |
|   MAT-R(Mean) (Å)   ↓     |     0.5235     |     0.2623     |      **0.2417**          |      0.2706    |
|   MAT-R(Median) (Å)   ↓     |    0.2710      |      0.2597    |       **0.2449**       |       0.2709    |
|  COV-P(Mean)  (%)  ↑      |      49.90    |       47.38     |         **50.03**      |          48.63|
|  COV-P(Median)  (%)  ↑      |   47.00       |      47.06    |           **48.31**        |          46.77|
| MAT-P(Mean) (Å) ↓ |      4.7816    |        4.2922 |                   **0.5571**              |  0.7512       |
| MAT-P(Median) (Å) ↓ |   0.5378       |   0.5615       |             **0.4921**          |    0.4995      |

*The subgraph predictor cannot predict the correct subgraphs so we use ld sampling rather than DDPM when k=1. 

**Table #2 Ablation study for the pretrained model and evaluate the model on MD17 force prediction (MAE) downstream task.**
| Component | Aspirin ↓ | Benzene ↓ | Ethanol ↓ | Malonaldehyde ↓ | Naphthalene ↓ | Salicylic ↓ | Toluene ↓ | Uracil ↓ |
|--------------------|-----------|-----------|-----------|-----------------|---------------|-------------|-----------|----------|
| w/o subgraph prediction loss| 1.193     | 0.305     | 0.321    | 0. 478  | 0.456 | 0.678       | 0.406     | 0.470    |
| w/o k-step same subgraph| 1.011     | 0.278     | 0.289    | 0. 461  | 0.448 | 0.665      | 0.377     | 0.464    |
| w/o expectation state| 0.931     | 0.281     | 0.276    | 0. 465  | 0.421 | 0.601       | 0.380     | 0.446    |
| expectation state + subgraph prediction loss | **0.880** |**0.252** | **0.258** | **0.459**   | **0.325** | **0.572**   | **0.362** | **0.420**|

**Table #3 The sensitivity analysis of diffusion steps for the pretrained model and evaluate the model on MD17 force prediction (MAE) downstream task.**
| #diffusion steps| Aspirin ↓ | Benzene ↓ | Ethanol ↓ | Malonaldehyde ↓ | Naphthalene ↓ | Salicylic ↓ | Toluene ↓ | Uracil ↓ |
|--|--|-|-|--|-|-|-|-|
| 100 | 0.921     | 0.282       | 0.275     | 0.471    | 0.348 | 0.596   | 0.394     | 0.465    |
| 1000 | 0.901     | 0.261     | 0.266    | 0. 463  | 0.336 | 0.590       | 0.385     | 0.438    |
| 5000    | **0.880** |**0.252** | **0.258** | **0.459**       | **0.325**     | **0.572**   | **0.362** | **0.420**|


Best regards,

The authors

Meta-review: This paper proposes a subgraph-based diffusion model for molecular representation learning, where noises are added to subgraphs in the diffusion process. Experimental results in multiple different tasks including molecular representation learning and conformation generation prove the competitiveness of the proposed approach over existing baselines. Overall, the subgraph-based diffusion model is new and novel. The experiments are strong.
Decision: Accept (poster)