SubgDiff: A Subgraph Diffusion Model to Improve Molecular Representation Learning

Authors: jiying  zhang
zijing  liu
yu  wang
bin  feng
yu  li
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-06 10:02:23

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.

Main paper: 
 introduction: 
 (Adapted with permission from [26] . Copyright 2018 American Chemical Society.) Molecular representation learning (MRL) has attracted tremendous attention due to its significant role in learning from limited labeled data for applications like AI-based drug discovery [36, 3, 4, 24] and material science [32] . From the perspective of physical chemistry, the 3D molecular conformation can largely determine the properties of molecules and the activities of drugs [6] . Thus, numerous geometric neural network architectures and self-supervised learning strategies have been proposed to explore 3D molecular structures to improve performance on downstream molecular property prediction tasks [34, 55, 22, 57] . demonstrated strong capabilities of molecular 3D conformation generation [51, 15, 58, 42] . The training process of a DPM for conformation generation can be viewed as the reconstruction of the original conformation from a 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 [31] . Inspired by this intuition, several works have used this technique for molecule pretraining [23, 55] . 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 network (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 [10] or atomic coordinates of 3D molecular geometry [51, 55] . However, this approach treats each atom as an individual particle, overlooking the substructure within molecules, which is pivotal in molecular representation learning [54, 45, 27] . As shown in Figure 1 , 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 [26] . As a result, uniformly adding same-scale Gaussian noise to all atoms makes it 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 a 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 a randomly selected subgraph at each time step in the diffusion process (Figure 2 ). In the training phase of SubgDiff, a 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 k-step same-subgraph diffusion, to train the model effectively and enhance its sampling capability.With the ability to 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 SubgDiff that adopts subgraph prediction, expectation state and k-step same-subgraph diffusion to improve its sampling and training; (3) the proposed representation learning method achieves superior performance on various downstream tasks, especially molecular forces prediction. 
preliminaries: 
 Notations. We use I to denote the identity matrix with dimensionality implied by context, ⊙ to represent the element product, and diag(s) to denote the diagonal matrix with diagonal elements of the vector s. If not specified, both ϵ and z represent noise sampled from the standard Gaussian distribution N (0, I). The topological molecular graph can be denoted as G(V, E, X) where V is the set of nodes, E is the set of edges, X is the node feature matrix, and its corresponding 3D Conformational Molecular Graph is represented asG 3D (G, R), where R = [R 1 , • • • , R |V| ] ∈ R |V|×3is the set of 3D coordinates of atoms.DDPM. Denoising diffusion probabilistic models (DDPM) [9] is a typical diffusion model [39] which consists of a diffusion (aka forward) and a reverse process. In the setting of molecular conformation generation, the diffusion model adds noise on the 3D molecular coordinates R [51] .Forward and reverse process. Given the fixed variance schedule β 1 , β 2 , • • • , β T , the posterior distribution q(R 1:T |R 0 ) that is fixed to a Markov chain can be written asEQUATIONTo simplify notation, we consider the diffusion on single atom coordinate R v and omit the subscript v to get the general notion R throughout the paper. Let α t = 1 -β t , ᾱt = t i=1 (1 -β i ), and then the sampling of R t at any time step t has the closed form: q(R t |R 0 ) = N (R t ; √ ᾱt R 0 , (1 -ᾱt )I).The reverse process of DDPM is defined as a Markov chain starting from a Gaussian distribution p(R T ) = N (R T ; 0, I):EQUATIONwhereσ t = 1-ᾱt-1 1-ᾱt β t denote time-dependent constant. In DDPM, µ θ (R t , t) is parameterized as µ θ (R t , t) = 1 ᾱt (R t -βt √ 1-ᾱt ϵ θ (R t , t)) and ϵ θ , i.e., the denoising network, is parameterized by a neural network where the inputs are R t and time step t.Training and sampling. The training objective of DDPM is:EQUATIONAfter training, samples are generated through the reverse process p θ (R 0:T ). Specifically, R T is first sampled from N (0, I), and R t in each step is predicted as follows,R t-1 = 1 √ αt (R t - 1 -αt √ 1 -ᾱt ϵθ(R t , t)) + σtz, z ∼ N (0, I).(4) 
subgdiff: 
 Directly using DDPM on atomic coordinates of 3D molecules means each atom is viewed as an independent single data point. However, the substructures play an important role in molecular generation [14] and representation learning [56] . Ignoring the inherent interactions among atoms within substructures may hinder the denoising network from learning features for molecular properties related to 3D conformation. In this paper, we propose to involve a mask operation in each diffusion step, leading to a new diffusion SubgDiff for molecular representation learning. Each mask corresponds to a 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 same-subgraph diffusion for training to effectively capture the substructure information. In the forward process of DDPM, we have𝑅 𝑡-1 𝑅 𝑡 𝑠 𝑡 = 0 𝑠 𝑡 = 1R t v = √ 1 -β t R t-1 v + √ β t ϵ t-1, ∀v ∈ V, in which the Gaussian noise ϵ t-1 is injected to every atom. Moreover, the training objective in Equation 3shows that the denoising networks would always predict a Gaussian noise for all atoms. Neither the forward nor reverse process of DDPM takes into account the substructure of the molecule. Instead, in SubgDiff, a mask vectors t = [s t1 , • • • , s t |V| ] ⊤ ∈ {0, 1} |V| isintroduced to determine which atoms will be added noise at step t. The mask vector s t is sampled from a discrete distribution p st (S | G) to select a subset of the atoms. In molecular graphs, the discrete mask distribution p st (S | G) is equivalent to the subgraph distribution, defined over a predefined sample space χ = {G i sub } N i=1 , where each sample is a connected subgraph extracted from G. Further, the distribution p st (S | G) should keep the selected connected subgraph to cohere with the molecular substructures. Here, we adopt a Torsional-based decomposition method [15] (Details in Appendix A.2). With the mask vector as latent variables s 1:t , the state transition of the forward process can be formulated as (Figure 3 ):EQUATIONwhich can be rewritten asR t v = 1 -s tv β t R t-1 v + s tv β t ϵ t-1 .The posterior distribution q(R 1:T |R 0 , s 1:T ) can be expressed as matrix form:EQUATION)To simplify the notation, we consider the diffusion on a single node v and omit the subscript v in R t v and s tv to get the notion R t and s t . By definingEQUATIONThe reverse process is decomposed as follows:EQUATIONwhere p θ (R t-1 |R t , s t ) and p θ (s t |R t ) are both learnable models. In the context of molecular learning, the model can be regarded as first predicting which subgraph s t should be denoised and then using the noise prediction network p θ (R t-1 |R t , s t ) to denoise the node position in the subgraph.However, it is tricky to generate a 3D structure by adopting the typical training and sampling method used in Ho et al. [9] . Specifically, following Ho et al. [9] , the reverse process can be optimized by maximizing the variational lower bound (VLB) of log p(R 0 ) as follows,log p(R 0 ) ≥ T t=1 E q(R t ,st|R 0 ) log p ϑ (s t |R t ) q(s t )subgraph prediction term+ E q(R 1 ,s1|R 0 ) log p θ (R 0 |R 1 ) reconstruction term - E q(s1:t) D KL (q(R T |R 0 , s 1:T ) ∥ p(R T ))prior matching termEQUATIONDetails of the derivation are provided in the Appendix D.2. The subgraph predictor p ϑ (s t |R t ) in the first term can be parameterized by a node classifier s ϑ . For the denoising matching term that is closely related to sampling, by Bayes rule, the posterior q(R t-1 |R t , R 0 , s 1:t ) can be written as:EQUATIONµ q (R t , R 0 , s 1:t , ϵ 0 ) = 1 √ 1 -β t s t (R t - β t s t (1 -β t s t )(1 -γt-1 ) + β t s t ϵ 0 ),where σ q (t) is the standard deviation and s 1:t-1 are contained in γt-1 . Following DDPM and parameterizingEQUATIONwhere BCE(s t , s ϑ ) is the binary cross entropy loss, λ is the weight used for the trade-off, and s ϑ is the subgraph predictor implemented as a node classifier with G 3D (G, R t ) as input and shares a molecule encoder with ϵ θ . The BCE loss employed here uses the subgraph selected at time-step t as the target, thereby explicitly compelling the denoising network to capture substructure information from molecules. Eventually, the s ϑ can be used to infer the mask vector ŝt = s ϑ (G, R t , t) during sampling. Thus, the sampling process is:EQUATIONHowever, using Equation 11 and Equation 12 directly for training and sampling faces two issues. First, the inability to access s 1:t-1 in µ q during the sampling process hinders the step-wise denoising procedure, posing a challenge to the utilization of conventional sampling methods in this context. Inferring s 1:t-1 solely from R t using another model p θ (s 1:t-1 |R t , s t ) is also difficult due to the intricate modulation of noise introduced in R t 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 t. However, the noise scale β t is relatively small when t is small, especially if the diffusion step is large (e.g. 1000).As a result, it is difficult to precisely predict the subgraph.Next, to effectively tackle the above issues, we design two techniques: expectation state diffusion and k-step same subgraph diffusion.We first devise a new way to calculate the denoising term. To eliminate the effect of mask series s 1:t and improve the training of subgraph prediction loss, we use a new lower bound of the denoising term as follows:E q(R t ,s1:t|R 0 ) D KL (q(R t-1 |R t , R 0 , s 1:t ) ∥ p θ (R t-1 |R t , s t )) ≤ E q(R t ,s1:t|R 0 ) D KL (q(R t-1 |R t , R 0 , s 1:t ) ∥ p θ (R t-1 |R t , s t )) ,where q(R t-1 |R t , R 0 , s 1:t ) is defined asq(R t |Es 1:t-1 R t-1 ,R 0 ,st)q(Es 1:t-1 R t-1 |R 0 ) q(R t |R 0 ,s1:t). It is an approximated posterior that only relies on the expectation of R t and s t . This lower bound defines a new forward process, in which, state 0 to state t -1 use the E s1:t-1 R t-1 and state t remains as Equation 6. Assume each node v ∈ V, s tv ∼ Bern(p) (i.i.d. w.r.t. t). Formally, we haveEQUATIONEQUATIONwhere 2 and ᾱt := t i=1 α i are general form of α j and ᾱj in DDPM (p = 1), respectively. Intuitively, this process is equivalent to using mean state ER t-1 to replace R t-1 during the forward process. This estimation is reasonable since the expectation E s1:t-1 R t-1 is like a cluster center of R t-1 , which can represent R t-1 properly. Thus, the approximated posterior becomesα i := (p √ 1 -β i + 1 -p)EQUATIONwhereµ q (R t , R 0 , s 1:t , ϵ 0 ) := 1 √ 1 -β t s t (R t - s t β t (s t β t + (1 -s t β t )p 2 t-1 i=1 ᾱt-1 ᾱi β i ϵ 0 ) σ 2 q (t) := s t β t p 2 t-1 i=1 ᾱt-1 ᾱi β i /(s t β t + p 2 (1 -s t β t ) t-1 i=1 ᾱt-1 ᾱi β i ).We parameterize p θ (R t-1 |R t , s t ) as N (R t-1 ; µ q (R t , R 0 , s 1:t-1 , s ϑ (G, R t , t), ϵ θ (R t , t)), σ q (t)I), and adopt the same training objective as Equation 11. By employing the sampling method R t-1 = µ q (R t , R 0 , s 1:t-1 , s ϑ (G, R t , t), ϵ θ (R t , t)) + σ q (t)z, we observe that the proposed expectation state enables a step-by-step execution of the sampling process. Moreover, using expectation is beneficial to reduce the complexity of R t for predicting the mask s t during training. This will improve the denoising network to perceive the substructure when we use the diffusion model for self-supervised learning.To reduce the complexity of the mask series (s 1 , s 2 , • • • , s T ) 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 k 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 > k, k ∈ N):EQUATIONwhereσ k t = (1 - t i=t-k (1 -s t-k+1 β i )I.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, k(i + 1)], the mask vectors{s j } k(i+1)j=ki+2 are equal to s ki+1 . SubgDiff then adopts the expectation state at the split time step {ik|i = 1, 2, • • • } to eliminate the effect of {s ik+1 |i = 1, 2, . . .}, that is, gets the expectation of ER ik at step ik w.r.t. s ik+1 . Overall, the diffusion process of SubgDiff is a two-phase diffusion process. In the first phase, the state 1 to state k⌊(t -1)/k⌋ use the expectation state diffusion, while in the second phase, state k(⌊(t -1)/k⌋) + 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, Phase I: Step 0 → km: E s 1:km R km = √ ᾱm R 0 + p m l=1 ᾱm ᾱl (1 - kl i=(l-1)k+1 (1 -β i ))ϵ 0 , where α j = (p kj i=(j-1)k+1 (1 -β i ) + 1 -p)ො 𝒔 2𝑘+1 = ො 𝒔 3𝑘 𝐑 0 𝐑 𝑘 𝐑 𝑘 𝐑 2𝑘 𝐑 2𝑘 𝐑 𝑇 Reverse Process ො 𝒔 𝑘 ො 𝒔 𝑘+1 = ො 𝒔 2𝑘 𝒔 𝑘𝑚+1 ො 𝒔 𝑘𝑚+1 = ො 𝒔 𝑇{s ik+1 |i = 1, 2, • • • }. Phase II: Step km + 1 → t: The phase is a (t -km)-step same mask diffusion. R t = t i=km+1 (1 -β i s km+1 )E s 1:km R km + 1 - t i=km+1 (1 -β i s km+1 )ϵ km . Let γ i = 1 -β i s km+1 , γt = t i=1 γ i , and βt = t i=1 (1 -β i ).We can drive the single-step state transition:EQUATION)Then we reuse the training objective in Equation 11as the objective of SubgDiff:EQUATIONwhere R t is calculated by Equation 17.Sampling. Although the forward process uses the expectation state w.r.t. s, we can only update the mask ŝt at t = ik, i = 1, 2, • • • because sampling only needs to get a subgraph from the distribution in the k-step interval. Eventually, adopting the δ defined in Equation 17, the sampling process (Figure 5 ) is shown below,EQUATIONwhere z ∼ N (0, I), m = ⌊(t -1)/k⌋ and ŝkm+1 = s ϑ (G, R km+k , km + k). The subgraph selected by ŝkm+1 will be generated in from the steps km + k to km. The mask predictor can be viewed as a discriminator of important subgraphs, indicating the optimal subgraph should be recovered in the next k steps. After one subgraph (substructure) is generated properly, the model can gently fine-tune the other parts of the molecule (c.f. the video in supplementary material). This subgraph diffusion would intuitively increase the robustness and generalization of the generation process, which is also verified by the experiments in Section 5.2. The training and sampling algorithms of SubgDiff are summarized in Algorithm 1 and Algorithm 2.Input: A molecular graph G 3D , k for same mask diffusion, m := ⌊(t -1)/k⌋ Sample t ∼ U(1, ..., T ) , ϵ ∼ N (0, I) Sample s km+1 ∼ p s km+1 (S | G) ▷ Sample a subgraph R t ← q(R t |R 0 ) ▷ Equation 17 L 1 = BCE(s km+1 , s ϑ (G, R t , t)) ▷ Subgraph prediction loss L 2 = ∥diag(s km+1 )(ϵ -ϵ θ (G, R t , t))∥ 2 ▷ Denoising loss optimizer. step(E t,R 0 ,st,ϵ [λL 1 + L 2 ]) ▷ Optimize parameters θ, ϑ Algorithm 2: Sampling from SubgDiff k is the same as training, for k-step same-subgraph diffusion; Sample R T ∼ N (0, I) ▷ Random noise initialization for t = T to 1 do z ∼ N (0, I) if t > 1, else z = 0 ▷ Random noise If t%k == 0 or t == T : ŝ ← s ϑ (G, R t , t) ▷ Subgraph prediction ε ← ϵ θ (G, R t , t) ▷ Posterior R t-1 ← Equation 19▷ sampling end return R 0 
experiments: 
 We conduct experiments to address the following two questions: 1) Can substructures improve the representation ability of the denoising network when using diffusion as 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 a 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 [51] on the task of molecular conformation generation.To verify the introduced substructure in the diffusion can enhance the denoising network for representation learning, we pretrain with SubgDiff objective and finetune on various downstream tasks.Dataset and settings. For pretraining, we follow [22] and use PCQM4Mv2 dataset [11] . It's a subdataset of PubChemQC [28] with 3.4 million molecules with 3D geometric conformations. We use various molecular property prediction datasets as downstream tasks. For tasks with 3D conformations, we consider the dataset MD17 and follow the literature [34, 35, 23] of using 1K for training and 1K 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 [47] .Pretraining framework. To explore the potential of the proposed method for representation learning, we consider MoleculeSDE [22] , 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-supervised methods [Type Prediction, Angle Prediction], and three contrastive methods [InfoNCE [30] and EBM-NCE [21] and 3D InfoGraph [23] ]. Two denoising baselines are also included [GeoSSL [23] , Denoising [55] and MoleculeSDE]. For 2D tasks, the baselines are AttrMask, ContexPred [12] , InfoGraph [44] , MolCLR [46] , 3D InfoMax [43] , GraphMVP [21] and MoleculeSDE.Results. As shown in Table 1 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, SubgDiff achieves SOTA performance compared to all the baselines. This also indicates that the proposed SubgDiff objective is promising for molecular representation learning due to the involvement of the knowledge of substructures during training. More results on the QM9 dataset [34] on the quantum mechanics property prediction can be found in Appendix A.4.3.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 [51] , we utilize the GEOM-QM9 [33] 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 Table 3 . From the results, it is easy to see that SubgDiff significantly outperforms the GeoDiff baseline on both metrics (COV-R and MAT-R) across different sampling steps. It indicates that by training with the substructure information, SubgDiff has a 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 can accelerate the sampling efficiency (10x). 13 ) show that SubgDiff consistently outperforms GeoDiff and other models trained on the in-domain dataset, demonstrating the introduced substructure effectively enhances the robustness and generalization of the diffusion model. 
conclusion: 
 We present a 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. A Experiment details and more results 
 
 Appendix: We conduct an alignment analysis to validate that our method can 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.A scaffold, which is the core structure of a molecule, holds significant importance in the field of chemistry as it serves as a foundation for systematic exploration of molecular cores and building blocks. It is usually represented by a substructure of a molecule and can be regarded as the subgraph in our SubgDiff.In our analysis, we select the nine most prevalent scaffolds from each dataset (BBBP, Sider, ClinTox, and Bace) and assign each molecule to a cluster according to its scaffold. To quantify the molecule embedding, we compute the Silhouette index of the embeddings for each dataset.As shown in Table 5 , SubgDiff enables the generation of more distinctive representations of 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 in Figure 6 . In this paper, we pre-define the mask distribution to be a discrete distribution, with sample space χ = {G i sub } N i=1 , and p t (S = G i sub ) = 1/N, t > 1, where G i sub is the subgraph split by the Torsionalbased decomposition methods [15] . The decomposition approach will cut off one torsional edge in a 3D molecule to make the molecule into two components, each of which contains at least two atoms. The two components are represented as two complementary mask vectors (i.e. s ′ + s = 1). Thus n torsional edges in G i 3D will generate 2n subgraphs. Finally, for each atom v, the s tv ∼ Bern(0.5), i.e. p = 0.5 in SubgDiff.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 . We use the pretraining framework MoleculeSDE proposed by [22] and extend our SubgDiff to multimodality pertaining. The two key components of MoleculeSDE are two SDEs(stochastic differential equations Song et al. [41] ): an SDE from 2D topology to 3D conformation (2D → 3D) and an SDE 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.2D topological molecular graph. A topological molecular graph is denoted asg 2D = G(V, E, X),where X is the atom attribute matrix and X is the bond attribute matrix. The 2D graph representation with graph neural network (GNN) is:EQUATION)where GIN is the a powerful 2D graph neural network [48] andH 2D = [h 0 2D , h 1 2D , . . .],where h i 2D is the i-th node representation.3D conformational molecular graph. The molecular conformation is denoted asg 3D := G 3D (G, R).The conformational representations are obtained by a 3D GNN SchNet [34] :EQUATIONwhereH 3D = [h 0 3D , h 1 3D , . . .],and h i 3D is the i-th node representation.The first objective is the conditional generation from topology to conformation, p(y|x), implemented as SubgDiff. The denoising network we adopt is the SE(3)-equivariance network (S 2D→3D θ) used in MoleculeSDE. The details of the network architecture refer to [22] .Therefore, the training objective from 2D topology graph to 3D confirmation is:EQUATIONwhere s 2D→3D ϑ (x, R t , t) gets the invariant feature from S θ and introduces a 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.e., p(x|y). We also use the SE(3)-invariant score network S 3D→2D θ proposed by MoleculeSDE. The details of the network architecture refer to [22] . For modeling S 3D→2D θ , it needs to satisfy the SE(3)-invariance symmetry property. The inputs are 3D conformational representation y, the noised 2D information x t at time t, and time t. The output of S 3D→2D θ is the Gaussian noise, as (ϵ X , ϵ E ). The diffused 2D information contains two parts: x t = (X t , E t ). For node feature X, the training objective isEQUATIONFor edge feature E, we define a mask matrix S from mask vector s: S ij = 1 if s i = 1 or s j = 1, otherwise, S ij = 0. Eventually, the ojective can be written as:EQUATION)EQUATIONThen the score network S 3D→2D θ is also decomposed into two parts for the atoms and bonds: S Xt θ (x t ) and S Et θ (x t ). Similarly, the mask predictor s 3D→2D ϑ is also decomposed into two parts for the atoms and bonds: s Xt ϑ (x t ) and s Et ϑ (x t ). Similar to the topology to conformation generation procedure, the s 3D→2D ϑ (x, R t , t) gets the invariant feature from S 3D→2D θ and introduces a mask head (MLP) to read out the mask prediction.Learning. Following MoleculeSDE, we incorporate a contrastive loss called EBM-NCE [21] . EBM-NCE provides an alternative approach to estimate the mutual information I(X; Y ) and is anticipated to complement the generative self-supervised learning (SSL) method. As a result, the ultimate objective is:EQUATIONwhere α 1 , α 2 , α 3 are three coefficient hyperparameters.Dataset. For pretraining, following MoleculeSDE, we use PCQM4Mv2 [11] . It's a sub-dataset of PubChemQC [28] with 3.4 million molecules with both the topological graph and geometric conformations. For finetuning, in addition to QM9 [33] , we also include MD17. To be specific, MD17 comprises eight molecular dynamics simulations focused on small organic molecules. These datasets were initially presented by Chmiela et al. [5] for the development of energy-conserving force fields using GDML. Each dataset features the trajectory of an individual molecule, encompassing a broad spectrum of conformations. The objective is to predict energies and forces for each trajectory by employing a single model.We begin by incorporating three coordinate-MI-unaware SSL methods: (1) Type Prediction, which aims to predict the atom type of masked atoms; (2) Angle Prediction, which focuses on 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 [23] , we include two contrastive baselines: (4) GeoSSL-InfoNCE [30] and (5) GeoSSL-EBM-NCE [21] . We also incorporate a generative SSL baseline named (6) GeoSSL-RR (RR for Representation Reconstruction). The above baselines are pre-trained on a subset of 1M molecules with 3D geometries from Molecule3D [52] and we reuse the results reported by [23] with SchNet as backbone.Baselines for 2D topology pretraining. We pick up the most promising ones as follows. Attr-Mask [12, 20] , ContexPred [12] , InfoGraph [44] , and MolCLR [46] .Baselines for 2D and 3D multi-modality pretraining. We include MoleculeSDE [22] (Variance Exploding (VE) and Variance Preserving (VP)) as a crucial baseline to verify the effectiveness of our methods due to the same pertaining framework. We reproduce the results from the released Code.By adopting the pertaining setting in Appendix A.4.2, we also take the QM9 dataset for finetuning and follow the literature [34, 35, 22] , using 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 [34, 23] . 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 GeoDiff from naive denoising pretraining perspective [55] .The results are shown in Table 9 . Following prior works [51] , we utilize the GEOM-QM9 [33] and GEOM-Drugs [2] datasets. The former dataset comprises small molecules of up to 9 heavy atoms, while the latter contains larger druglike compounds. We reuse the data split provided by Xu et al. [51] . 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 [51] as the GNN encoder for extracting the 3D molecular information. In the l-th layer, the GFN receives node embeddings h l ∈ R n×b (where b represents the feature dimension) and corresponding coordinate embeddings x l ∈ R n×3 as input. It then produces the output h l+1 and x l+1 according to the following process:EQUATION)EQUATION)EQUATION)where Φ are implemented as feed-forward networks and d ij denotes interatomic distances. The initial embedding h 0 is composed of atom embedding and time step embedding while x 0 represents atomic coordinates. N (i) is the neighborhood of i th 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 ):EQUATION)where εi is equivalent and ŝi is invariant.Settings. For GeoDiff [51] 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.To compare the generated and ground truth conformer ensembles, we employ the same evaluation metrics as in a prior study [7] : Average Minimum RMSD (AMR) and Coverage. These metrics 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 metrics built upon root-mean-square deviation (RMSD), which is defined as the normalized Frobenius norm of two atomic coordinates matrices, after alignment by Kabsch algorithm [16] . Formally, let S g and S r denote the sets of generated and reference conformers respectively, then the Coverage and Matching metrics [49] can be defined as:EQUATION)EQUATION)where δ is a threshold. The other two metrics COV-P and MAT-P can be defined similarly but with the generated sets S g and reference sets S r exchanged.EQUATIONEQUATIONIn practice, S g is set as twice of the size of S r for each molecule.A.5.3 Comparison with GeoDiff using Langevin Dynamics sampling method.In order to verify that our 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 (LD sampling) [40] :EQUATIONwhere z t ∼ N (0, I) and hσ 2 t . h is the hyper-parameter referring to step size and σ t is the noise schedule in the forward process. We use various time-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 is 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 ). i) Baselines: We compare SubgDiff with 7 state-of-the-art baselines: CVGAE [25] , GraphDG [38] , CGCF [49] , ConfVAE [50] , ConfGF [37] and GeoDiff [51] . For the above baselines, we reuse the experimental results reported by [51] . For GeoDiff [51] , we use the checkpoints released in public GitHub to reproduce the results. ii) Results: The results on the GEOM-QM9 dataset are reported in Table 12 . From the results, we get the following observation: SubgDiff significantly outperforms the baselines on COV-R, indicating the SubgDiff tends to explore more conformations. This implicitly demonstrates the subgraph will help fine-tune the generated conformation to be a potential conformation.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 . 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, 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. Masks on diffusion models. Previous works also share a similar idea of subgraph (mask) diffusion, such as MDM [31] , MDSM [18] 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 t ∼ p st (S) at each time step, hence a suitable discrete distribution p(S) can ensure that almost all nodes can be added noise. ii) Purpose: MDSM and MDM concentrate on self-supervised pre-training, while SubgDiff serves as a potent generative model and self-supervised pre-training algorithm. Notably, when s t = s 0 , ∀t, SubgDiff can recover to MDSM.Graph generation models. D3FG [19] : D3FG adopts three different diffusion models (D3PM, DDPM, and SO(3) Diffusion) to generate three different parts of molecules(linkerr types, center atom position, and functional group orientations), respectively. In general, these three parts can also be viewed as three subgraphs(subset). DiffPACK [59] is an Autoregressive generative method that predicts the torsional angle χ i (i = 1, 2, .., 4) of protein side-chains with the condition χ 1,...,i-1 , where χ i is a predefined subset of atoms. It uses a torsional-based diffusion model to approximate the distribution p(χ i |χ 1,...,i-1 ), in which every subset χ i needs a separate score network to estimate. Essentially, both D3FG and DiffPACK can be viewed as selecting a 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. [17] 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 [53] is proposed to use subgraphs as minimum units to train a latent diffusion model, while our method directly involves the subgraph during the forward process, which is a new type of diffusion model.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 a 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 a requirement for the number of steps in sampling to be the same as that in training, which can be improved with the recent progress of diffusion models such as DDIM.Thus,EQUATIONParameterizingp θ (R t-1 |R t ) in the reverse process as N (R t-1 ; 1 µ1 (R t -This lemma can be easily extended to the conditional version:Lemma C.2 Assume the forward and reverse processes of the diffusion model are both Markov chains. {R t } T t=0 are the states and y 1 , y 2 are the given conditions. Given the forward Gaussian distributionEQUATIONwhere z ∼ N (0, I).D.1 SubgDiff (1-same step and without expectation state)Here, we utilize the binary characteristic of the mask vector to derive the ELBO for SubgDiff, and we also provide a general proof in Section D.2:log p(R 0 ) ≥ E q(R 1:T ,s1:T |R 0 ) log p(R 0:T , s 1:T ) q(R 1:T |R 0 , s 1:T )q(s 1:T ) (69)EQUATION(74) The first term is mask prediction while the second term is similar to the ELBO of the classical diffusion model. The only difference is the s t in q(R t |R t-1 , s t ). According to Bayes rule, we can rewrite each transition as:EQUATIONwhere δ a (x) := δ(x -a) is Dirac delta function, that is, δ a (x) = 0 if x ̸ = a and ∞ -∞ δ a (x)dx = 1. Without loss of generality, assume that s 1 and s T both equal 1. Armed with this new equation, we drive the second term:EQUATION= E q(R 1:T |R 0 )   log p(R T )p θ (R 0 |R 1 ) q(R 1 |R 0 ) + log t∈{t|st=0} p θ (R t-1 |R t ) δ R t-1 (R t ) + log t∈{t|st=1} p θ (R t-1 |R t ) q(R t-1 |R t ,R 0 ) q(R t |R 0 )( ( ( ( (EQUATION+E q(R 1:T |R 0 ) log p(R T ) q(R T |R 0 ) + t∈{t|st=1} E q(R 1:T |R 0 ) log p θ (R t-1 |R t ) q(R t-1 |R t , R 0 ) (85) = t∈{t|st=0} E q(R 1:T |R 0 ) log p θ (R t-1 |R t ) δ R t-1 (R t ) + E q(R 1 |R 0 ) log p θ (R 0 |R 1 ) (86) +E q(R T |R 0 ) log p(R T ) q(R T |R 0 ) + t∈{t|st=1} E q(R t ,R t-1 |R 0 ) log p θ (R t-1 |R t ) q(R t-1 |R t , R 0 ) (87) = t∈{t|st=0} E q(R 1:T |R 0 ) log p θ (R t-1 |R t ) δ R t-1 (R t ) decay term + E q(R 1 |R 0 ) log p θ (R 0 |R 1 )reconstruction term (88)-D KL (q(R T |R 0 ) ∥ p(R T ))prior matching term- t∈{t|st=1} E q(R t |R 0 ) D KL (q(R t-1 |R t , R 0 ) ∥ p θ (R t-1 |R t ))denoising matching term (89)Here, the decay term represents the terms with s t = 0, which are unnecessary to minimize when we set p θ (R t-1 |R t ) := δ R t-1 (R t ). Eventually, the ELOB can be rewritten as follows:log p(R 0 ) ≥ T t=1 E q(R 1:T |R 0 ) log p ϑ (s t |R t ) q(s t ) mask prediction term + E q(R 1 |R 0 ) log p θ (R 0 |R 1 ) reconstruction term -D KL (q(R T |R 0 ) ∥ p(R T ))prior matching term- t∈{t|st=1} E q(R t |R 0 ) D KL (q(R t-1 |R t , R 0 ) ∥ p θ (R t-1 |R t ))denoising matching term (90)The mask prediction term can be implemented by a node classifier and the denoising matching term can be calculated via Lemma C.1. In detail,EQUATIONEQUATIONThus, the training objective of SubgDiff is:EQUATIONTo recover the existing work, we omit the mask prediction term (i.e. Let p θ (s t |R t ) := q(s t )) of SubgDiff in the main text.Here, we can derive the ELBO for SubgDiff:EQUATION= log p(R 0:T , s 1:T )q(R 1:T , s 1:T |R 0 ) q(R 1:T , s 1:T |R 0 ) dR 1:T ds 1:T= log p(R 0:T , s 1:T )q(R 1:T |R 0 , s 1:T )q(s 1:T ) q(R 1:T , s 1:T |R 0 ) dR 1:T ds 1:T(96)= log E q(s1:T ) E q(R 1:T |R 0 ,s1:T ) p(R 0:T , s 1:T )) q(R 1:T , s 1:EQUATION≥ E q(R 1:T |R 0 ,s1:T ) log E q(s1:T ) p(R 0:T , s 1:T ) q(R 1:T |R 0 , s 1:T )q(s 1:T ) (98)EQUATIONAccording to Bayes rule, we can rewrite each transition as:EQUATIONArmed with this new equation, we drive the second term:EQUATIONEQUATION+E q(R T |R 0 ,s1:T )q(s1:T ) log p(R T ) q(R T |R 0 , s 1:T ) +T t=2 E q(R t ,R t-1 ,s1:t|R 0 ) log p θ (R t-1 |R t , s t ) q(R t-1 |R t , R 0 , s 1:t ) (116) = E q(R 1 ,s1|R 0 ) log p θ (R 0 |R 1 , s 1 ) reconstruction term (117) -E q(s1:t) D KL (q(R T |R 0 , s 1:T ) ∥ p(R T )) prior matching term - T t=2 E q(R t ,s1:t|R 0 ) D KL (q(R t-1 |R t , R 0 , s 1:t ) ∥ p θ (R t-1 |R t , s t ))denoising matching term (118) Eventually, the ELOB can be rewritten as follows:EQUATION) -E q(s1:t) D KL (q(R T |R 0 , s 1:T ) ∥ p(R T ))prior matching term- T t=2 E q(R t ,s1:t|R 0 ) D KL (q(R t-1 |R t , R 0 , s 1:t ) ∥ p θ (R t-1 |R t , s t ))denoising matching term (120)The mask prediction term can be implemented by a node classifier s ϑ . For the denoising matching term, by Bayes rule, the q(R t-1 |R t , R 0 , s 1:t ) can be written as:EQUATIONFor the naive SubgDiff, we haveEQUATIONEQUATIONThen the denoising matching term can also be calculated via Lemma C.1 (let q(R t |R t-1 , R 0 ) := q(R t |R t-1 , R 0 , s 1:t ) , q(R t-1 |R 0 ) := q(R t-1 |R 0 , s 1:t-1 ) and p θ (R t-1 |R t ) = p θ (R t-1 )). Thus, the training objective of SubgDiff is:EQUATIONThe denoising matching term in (120) can be calculated by only sampling (R t , s t ) instead of (R t , s 1:t ). Specifically, we substitute (121) into the denoising matching term:EQUATION-E q(st)   log E q(s1:t-1) q(R t-1 |R 0 , s 1:t-1 ) :=q(EsR t-1 |R 0 ) + log E q(s1:t-1) q(R t |R t-1 , R 0 , s 1:t ) :=q(R t |EsR t-1 ,R 0 ,st)    (129) = E q(R t ,R t-1 ,|R 0 ,s1:t) E q(s1:t) log p θ (R t-1 |R t , s t ) 1 q(R t |R 0 ,s1:t) -E q(st) log q(E s R t-1 |R 0 ) -E q(st) log q(R t |E s R t-1 , R 0 , s t ) (130) = E q(R t ,R t-1 ,|R 0 ,s1:t)   E q(s1:t) log p θ (R t-1 |R t , s t ) q(R t |EsR t-1 ,R 0 ,st)q(EsR t-1 |R 0 ) q(R t |R 0 ,s1:t)   (131) = E q(R t ,R t-1 ,s1:t|R 0 )   log p θ (R t-1 |R t , s t ) q(R t |EsR t-1 ,R 0 ,st)q(EsR t-1 |R 0 ) q(R t |R 0 ,s1:t)   denoising matching term (132) = E q(R t ,R t-1 ,s1:t|R 0 ) log p θ (R t-1 |R t , s t ) q(R t-1 |R t , R 0 , s 1:t ) (133) = E q(R t ,s1:t|R 0 ) D KL (q(R t-1 |R t , R 0 , s 1:t ) ∥ p θ (R t-1 |R t , s t ))denoising matching term (134) Thus, we should focus on calculating the distributionEQUATIONBy lemma C.1, if we can gain the expression of q(R t |E s R t-1 , R 0 , s t ) and q(E s R t-1 |R 0 ), we can get the training objective and sampling process. Step 0 to Step t -1 (R 0 → R t-1 ): The state space of the mask diffusion should be the mean of the random state.EQUATIONForm (137), we have:EQUATIONEQUATIONEQUATIONEQUATIONwhereEQUATIONII:Step t -1 to Step t (R t-1 → R t ): We build the step t -1 → t is a discrete transition from q(R t-1 |R 0 ), with q(E s R t-1 |R 0 ) = N (R Step 0 to Step t (R 0 → R t ):q(R t |R 0 ) = q(R t |ER t-1 )q(ER t-1 |R 0 )dER where BCE(s t , s ϑ ) = s t log s ϑ (R t , t) + (1 -s t ) log (1 -s ϑ (R t , t)) is Binary Cross Entropy loss. However, training the SubgDiff 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 β t is relatively small when t 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 (t -1)-th step would be mainly affected by the mask sampling from 1 to t -1 step, which heavily increases the difficulty of predicting the noise added between (t -1)-step to t-step.Finally, the sampling can be written as: RThe state space of the mask diffusion should be the mean of the random state.EQUATIONForm Equation 137, we have: (1REQUATION)where α j = (p kj i=(j-1)k+1 (1-β i ) + 1 -p) 2 .When t ∈ N, let m := ⌊(t -1)/k⌋, and we haveR t = t i=km+1(1 -β i s km+1 )ER m×k + 1 - (1EQUATIONwhere (1EQUATIONEQUATIONAccording to the Lemma C.1, we have EQUATIONEQUATIONThe schematic can see Figure 5 . 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.

- 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.

- 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?

- 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.

- 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?
questions: questions: See weaknesses.
limitations: limitations: Discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	The subgraph prediction model is trained on highly specialized datasets, which might pose a risk of overfitting. 
2.	The 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.
3.	Given the complexity involved in the expectation state and k-step diffusion processes, the model may have high demands on computational resources.
questions: questions: 1.	How 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? 
2.	Does 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? 
3.	Can this method also be applied to broader graph representation learning tasks, such as node classification and clustering?
4.	Diffusing 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.

[1] 'Community detection in graph: An embedding method' in IEEE TNSE
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. The paper lacks an ablation study to evaluate the contribution of each component of SubgDiff to molecular representation learning.
questions: questions: Please refer to Weaknesses.
limitations: limitations: No
paper_summary: paper_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.
strength_weakness: strength_weakness: * The denoising process need better explanations.
* 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.
* Explain COV-R and MAT-R in section 5.2
questions: questions: * Does the denoising processing require to start with a R^T that has a clear topology? If not, how to sample s_t?
* How to select k in the k-step same-subgraph diffusion? Any general guidance?
* How are the SubgDiff learning results fed into downstream tasks?
limitations: limitations: None

Comments: title: Kind reminder to the reviewer
Kind reminder to the reviewer
kind reminder to the reviewer
comment: Thank the authors for the responses, which seem sound to me. I am maintaining my score since I am not an expert in diffusion models. I suggest that the authors include these additional information in the revision to make the paper more clear.
Dear reviewer, 

We kindly ask if you could inform us whether your concerns have been adequately addressed in our rebuttal, or if you have any further questions. We are committed to responding promptly to any additional inquiries you may have.

Thank you for your time and valuable feedback.

Best regards,

The Authors
Dear reviewer,

We kindly ask if you could inform us whether your concerns have been adequately addressed in our rebuttal, or if you have any further questions. We are committed to responding promptly to any additional inquiries you may have.

Thank you for your time and valuable feedback.

Best regards,

The Authors
Dear reviewer, 

We kindly ask if you could inform us whether your concerns have been adequately addressed in our rebuttal, or if you have any further questions. We are committed to responding promptly to any additional inquiries you may have.

Thank you for your time and valuable feedback.

Best regards,

The Authors
Dear reviewer Djew,

Thank you so much for the positive rating!
We sincerely appreciate your constructive suggestions and valuable comments for improving our paper. Thank you!

Best regards,
The Authors
Thank you to the authors for their reply, which has resolved most of my issues. I will maintain my positive rating, and I wish you good luck.