Title: SubgDiff: A Subgraph Diffusion Model to Improve Molecular Representation Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Diffusion model

Subgraph

Molecular representation learning
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

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.

1 Introduction

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 [37, 3, 4, 25] and material science [33]. From the perspective of physical chemistry, the 3D molecular conformation can largely determine the properties of molecules and the activities of drugs [6, 7]. 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 [35, 56, 23, 58].

Meanwhile, diffusion probabilistic models (DPMs) have shown remarkable power to generate realistic samples, especially in synthesizing high-quality images and videos [40, 10]. By modeling the generation as a reverse diffusion process, DPMs transform a random noise into a sample in the target distribution. Recently, diffusion models/flow models have also

Solution

0.8

Adsorbed

Surface

0.6

Solid

P

0.4

0.2

0

C1

C2

C4

C5

9

DDPM vs MaskedDiff

𝝐𝑖

𝐑 𝒊

:

Atomic coordinates of 3D molecule

𝝐1

𝐬2

𝝐3

𝝐𝑇 , 𝒔 𝑇

𝐬1

𝝐2

𝐬3

𝐑 0

𝐑 1

𝐑 1

𝐑 2

𝐑 2

𝐑 𝑇

SUBGDIFF

𝐑 1

𝐑 0

𝐑 2

𝝐1

𝝐3

𝝐𝑇

𝝐2

𝐑 𝑇

DDPM

: Subgraph selected by mask vector

𝐬𝑖

百

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 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 [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 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 [11] or atomic coordinates of 3D molecular geometry [52, 56]. However, this approach treats each atom as an individual particle, overlooking the substructure within molecules, which is pivotal in molecular representation learning [55, 46, 28]. 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 [27]. 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.

10

2 Related work

Diffusion models on graphs. The diffusion models on graphs can be mainly divided into two categories: continuous diffusion and discrete diffusion. Continuous diffusion applies a Gaussian noise process on each node or edge [14, 30], 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 CVGAE [26], GraphDG [39], CGCF [50], ConfVAE [51], ConfGF [38] and GeoMol [8]. Recently, diffusion-based methods have shown competitive performance. Torsional Diffusion [16] raises a diffusion process on the hypertorus defined by torsion angles. However, it is not suitable as a representation learning technique due to 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 et al. [56] and MoleculeSDE [23], etc. However, those studies have not considered the molecular substructure in the pertaining. In this paper, we concentrate on 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.1.

3 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 as G 3 D ( G , R ) , where R = [ R 1 , · · · , R |V| ] ∈ R |V|× 3 is the set of 3D coordinates of atoms.

DDPM. Denoising diffusion probabilistic models (DDPM) [10] is a typical diffusion model [40] 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 [52].

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 as

To 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 ) :

where σ 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:

After 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,

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 [15] and representation learning [57]. 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.

4.1 Involving subgraph into diffusion process

In the forward process of DDPM, we have R 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 3 shows 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 vector s t = [ s t 1 , · · · , s t |V| ] ⊤ ∈ { 0 , 1 } |V| is introduced to determine which atoms will be added noise at step t . The mask vector s t is sampled from a discrete distribution p s t ( S | G ) to select a subset of the atoms. In molecular graphs, the discrete

𝑅 𝑡-1

𝑅 𝑡

= 0

= 1

mask distribution p s t ( 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 s t ( S | G ) should keep the selected connected subgraph to cohere with the molecular substructures. Here, we adopt a Torsional-based decomposition method [16] (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):

which can be rewritten as R t v = √ 1 -s t v β t R t -1 v + √ s t v β t ϵ t -1 .

The posterior distribution q ( R 1: T | R 0 , s 1: T ) can be expressed as matrix form:

To simplify the notation, we consider the diffusion on a single node v and omit the subscript v in R t v and s t v to get the notion R t and s t . By defining γ t := 1 -s t β t , ¯ γ t := ∏ t i =1 (1 -s t β t ) , the closed form of sampling R t given R 0 is

4.2 Reverse process learning

The reverse process is decomposed as follows:

where 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. [10]. Specifically, following Ho et al. [10], the reverse process can be optimized by maximizing the variational lower bound (VLB) of log p ( R 0 ) as follows,

Details 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:

where σ q ( t ) is the standard deviation and s 1: t -1 are contained in ¯ γ t -1 . Following DDPM and parameterizing 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 ) , ϵ θ ( G , R t , t )) , σ q ( t ) I ) , the training objective is

where 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 3 D ( 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 ˆ s t = s ϑ ( G , R t , t ) during sampling. Thus, the sampling process is:

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

4.3 Expectation state 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:

where ˆ q ( R t -1 | R t , R 0 , s 1: t ) is defined as q ( R t | E s 1: t -1 R t -1 ,R 0 ,s t ) q ( E s 1: t -1 R t -1 | R 0 ) q ( R t | R 0 ,s 1: 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 s 1: t -1 R t -1 and state t remains as Equation 6. Assume each node v ∈ V , s t v ∼ Bern ( p ) (i.i.d. w.r.t. t ). Formally, we have

where α i := ( p √ 1 -β i + 1 -p ) 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 E R t -1 to replace R t -1 during the forward process. This estimation is reasonable since the expectation E s 1: 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

where

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.

4.4 k -step same-subgraph diffusion.

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 ):

t i = t -k

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, 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 E R 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,

𝐑 0

𝐑 𝑘

𝐑 𝑘

𝐑 2𝑘

𝐑 2𝑘

𝐑 𝑇

𝝐

𝑘-1 𝑚+1:𝑘𝑚

𝐑 0

(

𝑡 - 𝑘𝑚

)-step same-subgraph diffusion

𝝐1:𝑘

𝔼𝐬1:𝑘

𝐑 𝑘

𝐬𝑖

= 𝐬1 𝑖=1

𝑘

𝐬𝑖

= 𝐬𝑘𝑚-𝑚+1 𝑖=𝑘𝑚-𝑚+1

𝑘𝑚

𝔼𝐬1:𝑘𝑚

𝐑 𝑘m

𝝐𝑘+1:2𝑘

𝐬𝑖

= 𝐬𝑘+1 𝑖=𝑘+1

2𝑘

𝐑 𝑡

𝝐𝑘𝑚+1:𝑡

𝐬𝑖

= 𝐬𝑘𝑚+1 𝑖=𝑘𝑚+1

𝑡

Expectation state diffusion

𝐬𝑘𝑚+1

𝔼𝐬1:𝑘𝑚

𝐑 𝑘m

𝐬𝑖

: mask

random variable

;

𝝐𝑖

𝔼𝒔𝐑

:

the expectation of

𝐑

w.r.t

𝐬

: subgraph selected by mask

sample

𝐬𝑘𝑚+1

MaskedDiff with k-same step diffusion

𝝐1:𝑘

𝝐2𝑘+1:3𝑘

𝝐𝑇 , 𝒔 𝑇

𝒔1

𝒔2𝑘+1

Forward Process

15

𝝐1:𝑘

ො 𝒔2𝑘

𝝐2𝑘+1:3𝑘

𝝐𝑇 , 𝒔 𝑇

𝒔1

= ො 𝒔𝑘

𝝐𝑘+1:2𝑘

ො 𝒔2𝑘+1

= ො 𝒔 3𝑘

𝐑 0

𝐑 𝑘

𝐑 𝑘

𝐑 2𝑘

𝐑 2𝑘

𝐑 𝑇

Reverse Process

ො 𝒔𝑘

ො 𝒔𝑘+1

= ො 𝒔 2𝑘

ො

𝒔𝑘𝑚+1 = ො 𝒔 𝑇

case k = 1 ) and ¯ α i = ∏ t j =1 α j . In the rest of the paper, α j denotes the general version without a special statement. Actually, E s 1: km R km only calculate the expectation w.r.t. random variables { s ik +1 | i = 1 , 2 , · · · } .

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: q ( R t | R t -1 ) = N ( R t ; √ γ t R t -1 , (1 -γ t ) I ) and

Then we reuse the training objective in Equation 11 as the objective of SubgDiff:

where 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 ˆ s 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,

where z ∼ N ( 0 , I ) , m = ⌊ ( t -1) /k ⌋ and ˆ s km +1 = s ϑ ( G , R km + k , km + k ) . The subgraph selected by ˆ s 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.

𝒔𝑘+1

𝝐𝑇 , 𝒔 𝑇

12

Algorithm 1: Training SubgDiff

Algorithm 2: Sampling from SubgDiff

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 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 [52] on 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 representation learning, we pretrain with SubgDiff objective and finetune on various downstream tasks.

Dataset and settings. For pretraining, we follow [23] and use PCQM4Mv2 dataset [12]. It's a subdataset of PubChemQC [29] 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 [35, 36, 24] 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 [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 2 D → 3 D model and the mask operation is extended to the node feature and graph adjacency for the 3 D → 2 D 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 [31] and EBM-NCE [22] and 3D InfoGraph [24]]. Two denoising baselines are also included [GeoSSL [24], Denoising [56] and MoleculeSDE]. 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 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 [35] on the quantum mechanics property prediction can 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-QM9 [34] and GEOMDrugs [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).

Domain generalization. To further illustrate the benefits of SubgDiff, we design two crossdomain tasks: (1) Training on QM9 (small molecular with up to 9 heavy atoms) and testing on Drugs (medium-sized organic compounds); (2) Training on Drugs and testing on QM9. The results (Table 4 and Appendix Table 13) show that SubgDiff consistently outperforms GeoDiff and other models trained on the in-domain dataset, demonstrating the introduced sub-

structure effectively enhances the robustness and generalization of the diffusion model.

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

References

Juan Lopez Alcaraz and Nils Strodthoff. Diffusion-based time series imputation and forecasting with structured state space models. Transactions on Machine Learning Research , 2022.

Simon Axelrod and Rafael Gomez-Bombarelli. GEOM, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data , 9(1):185, 2022.

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.16208 , 2023.

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.13193 , 2024.

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 , 3(5):e1603015, 2017.

Aurora J Cruz-Cabeza and Joel Bernstein. Conformational polymorphism. Chemical reviews , 114(4):2170-2191, 2014.

Bin Feng, Zequn Liu, Nanlan Huang, Zhiping Xiao, Haomiao Zhang, Srbuhi Mirzoyan, Hanwen Xu, Jiaran Hao, Yinghui Xu, Ming Zhang, et al. A bioactivity foundation model using pairwise meta-learning. Nature Machine Intelligence , 6(8):962-974, 2024.

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-13769, 2021.

Kilian Konstantin Haefeli, Karolis Martinkus, Nathanaël Perraudin, and Roger Wattenhofer. Diffusion models for graphs benefit from discrete state spaces. In The First Learning on Graphs Conference , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning , pp. 8867-8887, 2022.

Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint arXiv:2005.00687 , 2020.

Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. In International Conference on Learning Representations , 2020.

John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph-based protein design. Advances in neural information processing systems , 32, 2019.

Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Hierarchical generation of molecular graphs using structural motifs. In International conference on machine learning , pp. 4839-4848, 2020.

Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi S Jaakkola. Torsional diffusion for molecular conformer generation. In Advances in Neural Information Processing Systems , 2022.

Wolfgang Kabsch. A solution for the best rotation to relate two sets of vectors. Acta Crystallographica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallography , 32 (5):922-923, 1976.

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 Learning Research , pp. 17391-17408, 23-29 Jul 2023.

Jiachen Lei, Peng Cheng, Zhongjie Ba, and Kui Ren. Masked diffusion models are fast learners. arXiv preprint arXiv:2306.11363 , 2023.

Haitao Lin, Yufei Huang, Odin Zhang, Yunfan Liu, Lirong Wu, Siyuan Li, Zhiyuan Chen, and Stan Z. Li. Functional-group-based diffusion for pocket-specific molecule generation and elaboration. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

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.

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.

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.

Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)invariant denoising distance matching. In The Eleventh International Conference on Learning Representations , 2023.

Xingyu Lu, He Cao, Zijing Liu, Shengyuan Bai, Leqing Chen, Yuan Yao, Hai-Tao Zheng, and Yu Li. MoleculeQA: A dataset to evaluate factual accuracy in molecular comprehension. arXiv preprint arXiv:2403.08192 , 2024.

Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry prediction using a deep generative graph neural network. Scientific reports , 9(1):20381, 2019.

Veselina Marinova, Geoffrey PF Wood, Ivan Marziano, and Matteo Salvalaglio. Dynamics and thermodynamics of ibuprofen conformational isomerism at the crystal/solution interface. Journal of chemical theory and computation , 14(12):6484-6494, 2018.

Siqi Miao, Mia Liu, and Pan Li. Interpretable and generalizable graph learning via stochastic attention mechanism. In International Conference on Machine Learning , pp. 15524-15543, 2022.

Maho Nakata and Tomomi Shimazaki. Pubchemqc project: a large-scale first-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling , 57(6):1300-1308, 2017.

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-4484, 2020.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Zixuan Pan, Jianxu Chen, and Yiyu Shi. Masked diffusion as self-supervised representation learner. arXiv preprint arXiv:2308.05695 , 2023.

Robert Pollice, Gabriel dos Passos Gomes, Matteo Aldeghi, Riley J Hickman, Mario Krenn, Cyrille Lavigne, Michael Lindner-D'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.

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.

Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems , 30, 2017.

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.

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.

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.

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-8958, 2020.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning , pp. 2256-2265, 2015.

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.

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 Learning Representations , 2020.

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.

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-20502, 2022.

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.

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 , 62(11):2713-2725, 2022.

Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence , 4(3):279287, 2022.

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(2):513-530, 2017.

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations , 2018.

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.

Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael Gomez-Bombarelli, and Jian Tang. An end-to-end framework for molecular conformation generation via bilevel programming. In International Conference on Machine Learning , pp. 11537-11547, 2021.

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.

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 arXiv:2110.01717 , 2021.

Kuo Yang, Zhengyang Zhou, Limin Li, Pengkun Wang, Xu Wang, and Yang Wang. Subdiff: Subgraph latent diffusion model, 2024.

Zhaoning Yu and Hongyang Gao. Molecular representation learning via heterogeneous motif graph neural networks. In International Conference on Machine Learning , pp. 25581-25594, 2022.

Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. In The Eleventh International Conference on Learning Representations , 2023.

Xuan Zang, Xianbing Zhao, and Buzhou Tang. Hierarchical molecular graph self-supervised learning for property prediction. Communications Chemistry , 6(1):34, 2023.

Jiying Zhang, Xi Xiao, Long-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-First International Joint Conference on Artificial Intelligence, IJCAI-22 , pp. 37303736. International Joint Conferences on Artificial Intelligence Organization, 7 2022. doi: 10.24963/ijcai.2022/518. Main Track.

Jiying Zhang, Zijing Liu, Shengyuan Bai, He Cao, Yu Li, and Lei Zhang. Efficient antibody structure refinement using energy-guided se (3) flow matching. In IEEE International Conference on Bioinformatics and Biomedicine (BIBM) . IEEE, 2024.

Yangtian Zhang, Zuobai Zhang, Bozitao Zhong, Sanchit Misra, and Jian Tang. DiffPack: A torsional diffusion model for autoregressive protein side-chain packing. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

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.

Contents

A Experiment details and more results

A.1 Visualization of representations

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

Appendix

SubGDiff

Silhouette: 0.2632

Silhouette: 0.1739

SubGDiff

Silhouette: 0.1780

Clintox

Silhouette: 0.2571

BACE

150

100

0

10

>

Y

150

30

20

10

10

20

30

40

100

10

10

Y

80

Q

8

10

10

0Z-

-10

10

J

15

10

10

15

BBBP

Sider

A.2 Mask distribution

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 [16]. 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 3 D will generate 2 n subgraphs. Finally, for each atom v , the s t v ∼ Bern (0 . 5) , i.e. p = 0 . 5 in 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.

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 multimodality pertaining. The two key components of MoleculeSDE are two SDEs(stochastic differential equations Song et al. [42]): 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 as g 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:

where GIN is the a powerful 2D graph neural network [49] and H 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 as g 3D := G 3 D ( G , R ) . The conformational representations are obtained by a 3D GNN SchNet [35]:

where H 3D = [ h 0 3D , h 1 3D , . . . ] , and h i 3D is the i -th node representation.

An SE(3)-Equivariant Conformation Generation 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 [23].

Therefore, the training objective from 2D topology graph to 3D confirmation is:

where 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 [23]. 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 is

For 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:

Then the score network S 3D → 2D θ is also decomposed into two parts for the atoms and bonds: S X t θ ( x t ) and S E t θ ( x t ) . Similarly, the mask predictor s 3D → 2D ϑ is also decomposed into two parts for the atoms and bonds: s X t ϑ ( x t ) and s E t ϑ ( 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 [22]. EBMNCE 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:

where α 1 , α 2 , α 3 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, 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.

Baselines for 3D property prediction 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 [24], we include two contrastive baselines: (4) GeoSSL-InfoNCE [31] and (5) GeoSSL-EBM-NCE [22]. 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 [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. AttrMask [13, 21], ContexPred [13], InfoGraph [45], and MolCLR [47].

Baselines for 2D and 3D multi-modality pretraining. We include MoleculeSDE[23](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.

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 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 GeoDiff from naive denoising pretraining perspective [56]. The results are shown in Table 9.

A.5 Settings and more results on 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 druglike 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 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:

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):

where ˆ ϵ i is equivalent and ˆ s i 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 metrics for conformation generation.

To compare the generated and ground truth conformer ensembles, we employ the same evaluation metrics as in a prior study [8]: Average Minimum RMSD (AMR) and Coverage. These metrics

Ƹ

Ƹ

去噪

Ƹ

Adding

noise

GNN Encoder

: Mask prediction;

𝜖𝑡

: Noise prediction

𝐑 0

𝜖𝑡

Noise head

Mask head

𝐑 𝑡

ℒ1

ℒ2

ℒ1

);

ℒ2 =

diag 𝐬𝑡

𝜖 -

𝜖𝑡

= 𝑠ϑ

𝒢, 𝐑

𝑡

, 𝑡

;

𝜖𝑡

= 𝜖𝜃

𝒢, 𝐑

𝑡

, 𝑡

;

Fine-tuning

下游任务 分子 性质预测 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 [17]. Formally, let S g and S r denote the sets of generated and reference conformers respectively, then the Cov erage and Mat ching metrics [50] can be defined as:

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.

In 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)[41] :

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

.

Ƹ

Ƹ

Ƹ

Ƹ

Ƹ

A.5.4 Comparison with SOTAs.

i) Baselines: We compare SubgDiff with 7 state-of-the-art baselines: CVGAE [26], GraphDG [39], CGCF [50], ConfVAE [51], ConfGF [38] and GeoDiff [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 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 possible conformations. This implicitly demonstrates the subgraph will help fine-tune the generated conformation to be a potential conformation.

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.

A.6 Sentivity analysis of k 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, 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.

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 MDSMis 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 s t ( 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 [20]: 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[60] 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. [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 a 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 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.

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-likelihood -log p θ ( R 0 ) , in which the tricky terms are L t -1 = D KL ( q ( R t -1 | R t , R 0 ) || p θ ( R t -1 | R t ))) , T ≥ t > 1 . Here we provide a lemma that tells us the posterior distribution q ( R t -1 | R t , R 0 ) used in the training and sampling algorithms of the diffusion model can be determined by q ( R t | R t -1 , R 0 ) , q ( R t -1 | R 0 ) . Formally, we have

Lemma C.1 Assume the forward and reverse processes of the diffusion model are both Markov chains. Given the forward Gaussian distribution q ( R t | R t -1 , R 0 ) = N ( R t ; µ 1 R t -1 , σ 2 1 I ) , q ( R t -1 | R 0 ) = N ( R t -1 ; µ 2 R 0 , σ 2 2 I ) and ϵ 0 ∼ N ( 0 , I ) , the distribution q ( R t -1 | R t , R 0 ) is

Parameterizing p θ ( R t -1 | R t ) in the reverse process as

the training objective of the DPM can be written as

and the sampling process is

where z ∼ N ( 0 , I ) .

Once we get the variables ( µ 1 , σ 1 , µ 2 , σ 2 ) , we can directly obtain the training objective and sampling process via lemma C.1, which will help the design of new diffusion models.

Proof: Given the forward Gaussian distribution q ( R t | R t -1 , R 0 ) = N ( R t ; µ 1 R t -1 , σ 2 1 I ) and q ( R t -1 | R 0 ) = N ( R t -1 ; µ 2 R 0 , σ 2 2 I ) , we have

From the DDPM, we know training a diffusion model should optimize the ELBO of the data

To compute the KL divergence D KL ( q ( R t -1 | R t , R 0 ) ∥ p θ ( R t -1 | R t )) , we first rewrite q ( R t -1 | R t , R 0 ) by Bayes rule

We can rewrite our variance equation as Σ q ( t ) = σ 2 q ( t ) I , where:

From (39), we have the relationship between R t and R 0 :

Substituting this into µ q ( R t , R 0 ) , we can get

Thus,

Parameterizing p θ ( R t -1 | R t ) in the reverse process as N ( R t -1 ; 1 µ 1 ( R t -σ 2 1 √ µ 2 1 σ 2 2 + σ 2 1 ϵ θ ( R t , t )) , σ 2 1 σ 2 2 µ 2 1 σ 2 2 + σ 2 1 I ) , and the corresponding optimization problem becomes:

Therefore, the training objective of the DPM can be written as

During the reverse process, we sample R t -1 ∼ p θ ( R t -1 | R t ) . Formally, the sampling (reverse) process is

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 distribution

and ϵ 0 ∼ N ( 0 , I ) , we have the distribution q ( R t | R 0 , y 1 , y 2 ) = N ( R t ; µ 1 µ 2 R 0 , ( σ 2 1 + µ 2 1 σ 2 2 ) I ) . Thus, the posterior distribution q ( R t -1 | R t , R 0 ) is

Parameterizing p θ ( R t -1 | R t , y 1 , y 2 ) in the reverse process as

the training objective of the DPM can be written as

and the sampling process is

where z ∼ N ( 0 , I ) .

D Derivations of training objectives

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:

(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:

̸

where δ 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:

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:

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,

Thus, the training objective of SubgDiff is:

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

D.2 ELBO of SubgDiff

Here, we can derive the ELBO for SubgDiff:

(103)

According to Bayes rule, we can rewrite each transition as:

Armed with this new equation, we drive the second term:

Eventually, the ELOB can be rewritten as follows:

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:

For the naive SubgDiff, we have

Then 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:

D.2.1 Expectation of s 1: T

The 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:

︸

︷︷ denoising matching term

︸

Thus, we should focus on calculating the distribution

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

D.3 Single-step subgraph diffusion

D.3.1 Training

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

Form (137), we have:

where α i := ( p √ 1 -β i +1 -p ) 2 and ¯ α t = ∏ t i =1 α i .

II: 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

Step 0 to Step t ( R 0 → R t ):

Thus, from subsection D.2.1, the training objective of 1-step SubgDiff is:

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.

D.3.2 Sampling

Finally, the sampling can be written as:

where s = s ( R , t ) and

𝐑 𝑘𝑚

𝐑 𝑘𝑚+1

𝐑 𝑚+1 𝑘

𝝐𝑘𝑚+1

𝒔𝑘𝑚+1

𝝐𝑘𝑚+3

𝒔𝑘𝑚+1

𝝐𝑘𝑚+2

𝒔𝑘𝑚+1

𝝐

𝑚+1 𝑘

𝒔𝑘𝑚+1

𝐑 𝑘𝑚+2

E Expectation state distribution

ො

ො

𝝐𝑇, ො 𝒔𝑇

ො

𝝐3, ො 𝒔3

ො

𝝐2, ො 𝒔2

𝝐1 , ො 𝒔𝟏

The state space of the mask diffusion should be the mean of the random state.

Form Equation 137, we have:

where α i := ( p 1 -β i +1 -p ) 2 and ¯ α t = ∏ t i =1 α i .

Finally, the Expectation state distribution is:

F The derivation of SubgDiff

When t is an integer multiple of k ,

where α j = ( p √ ∏ kj i =( j -1) k +1 (1 -β i ) + 1 -p ) 2 .

When t ∈ N , let m := ⌊ ( t -1) /k ⌋ , and we have

where γ i = 1 -β i s km +1 .

F.0.1 Sampling

According to the Lemma C.1, we have

The schematic can see Figure 5.

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 C, D.1, D.2, E, F.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it 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 A.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions 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/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 7 in 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 computer 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

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This research conducted in the paper conforms with the NeurIPS 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 to 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

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

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 as 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
comments: 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.

STRENGTHS:
* Incorporate substructural information into diffusion model

WEAKNESSES:
* 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:
* 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:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-07

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

STRENGTHS:
1. The idea of enhancing motif information in the diffusion process is promising.
2. 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.
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:
Please refer to Weaknesses.

LIMITATIONS:
No
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1.\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.\n\nWEAKNESSES:\n1.\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.\n\nQUESTIONS:\n1.\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\n\nLIMITATIONS:\nThis 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nDiscussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-19"}