Subgraph Diffusion for 3D Molecular Representation Learning: Combining Continuous and Discrete

Q3: Can the model be compared against other graph representation learning methods for molecular graphs that do not process 3D coordinates? Such a comparison could more effectively underscore the significance of explicitly incorporating 3D coordinates in practical applications.

A3: Thanks for this constructive suggestion! Following the advice, we compare our method with various graph representation learning methods on 2D molecular property prediction tasks, and the results are presented in Table 3 above. The baselines, including AttrMask[1], ContexPred[1], InfoGraph[2], and MolCLR[3] are all exclusively 2D graph learning methods that don't process 3D coordinates, while other methods all leverage the 3D coordinates information. As indicated in the results in Table 3, methods incorporating 3D coordinates demonstrate a significant performance advantage over those without 3D information. Furthermore, our approach stands out by enabling the simultaneous exploration of 2D and 3D information, leading to state-of-the-art performances.

References:

[1] Hu, W., et al. "Strategies For Pre-training Graph Neural Networks." International Conference on Learning Representations (ICLR). 2020.

[2] Sun, Fan-Yun, et al. "InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization." International Conference on Learning Representations. 2019.

[3]Wang, Yuyang, et al. "Molecular contrastive learning of representations via graph neural networks." Nature Machine Intelligence 4.3 (2022): 279-287.

Q1: Can the diffusion model be adapted to accommodate node features and graph structures? How dependent is the model on the graph structure?

A1: 1) The proposed diffusion model can be extended to node features and graph structures. For node features, similar to the proposed MaskedDiff, we only need to operate the diffusion model on node feature space instead of Euclidean space (3D coordinate). Regarding graph structure, for example in Digress[1], the transition matrices can be designed as $Q_m=diag(Q_n,I_{m-n})$, which means adding noise to the edges of the selected subgraph of size $n$. 2) In our model, the graph structure provides the initial graph topology for subgraph decomposition; in other words, the mask distribution $p(S)$ depends on the graph topology. Besides, the graph structure serves as a condition for the conformational distribution $p(R|\mathcal{G}(A,X))$ and is therefore input to the denoising network to assist in predicting Gaussian noise.

Reference:

[1] DiGress: Discrete Denoising diffusion for graph generation, ICLR 2023

Q2: Could the model be used for other important tasks such as the datasets from MoleculeNet [1] or OGBG [2]?

A2: We conduct additional experiments to evaluate our method on eight classification tasks from MoleculeNet[1]. Utilizing our MaskedDiff to replace the SDE in MoleculeSDE[3], a pretraining framework, we established a MaskedDiff-based pretraining framework. Following the procedure of MoleculeSDE[3], we pretrain the framework on the PCQM4Mv2 dataset and then fine-tune it on the MoleculeNet[1] dataset. The results can be found in Table 3. It is clear that our method consistently outperforms other methods across all tasks.

References:

[1]MoleculeNet: a benchmark for molecular machine learning. Chemical Science.

[2]Open Graph Benchmark: Datasets for Machine Learning on Graphs. NeurIPS 2020.

[3] Liu S, Du W, Ma Z M, et al. A group symmetric stochastic differential equation model for molecule multi-modal pretraining//International Conference on Machine Learning. PMLR, 2023: 21497-21526.

Table 3: Results for molecular property prediction tasks (with 2D topology only). The backbone is GIN.

We thank the reviewer for the positive feedback and constructive comments. We provide pointwise responses below.

Weakness 1: The presentation of the work could be improved. The use of mathematical symbols is burdensome and difficult to follow. Explanatory figures, such as Figure 4, are not self-explanatory due to numerous unexplained symbols. Algorithm 1 is helpful for readers trying to understand the method's training, but it lacks many details. For instance, is the mask predictor different from the noise predictor? Which neural network architectures are used?

Response 1: Thanks so much for your helpful suggestions! We will carefully improve the presentation and clarify the mathematical notations. We have updated Figure 4 to make it easier to follow. For Algorithm 1, 1) The mask predictor and the noise predictor are different. 2) Our network consists of a GNN encoder (e.g. SchNet) for extracting the 3D graph information, a noise predictor (e.g. an MLP), and a mask predictor (e.g. an MLP). Notably, these two predictors share the same GNN encoder. 3) In our experiments, we adopt the graph field network (GFN) from GeoDiff as the GNN encoder to extract 3D molecular information. Specifically, in the $l$-th layer, the GFN receives node embeddings $\mathbf{h}^l \in \mathbb{R}^{n \times b}$ (where $b$ represents the feature dimension) and corresponding coordinate embeddings $\mathbf{x}^l \in \mathbb{R}^{n \times 3}$ as input. It then produces the output $\mathbf{h}^{l+1}$ and $\mathbf{x}^{l+1}$ according to the following process:

Q1: I may missed it but didn't find it: what's the number of diffusion timesteps T for each experiment's config? This is important for audience to have a better understanding of the method's sampling efficiency.

A1: Thanks for the question. In our experiments, the diffusion timestep T is set to 5000, and crucial hyperparameters are provided in Tables 5 and 6 of the Appendix. In addition, we include more results with different numbers of timesteps to further support our method. Please see the Table 4 below. Specifically, we include the experimental results of 200 steps and 500 steps and compare them with GeoDiff. The results suggest that our method significantly outperforms the baseline when the same sampling step is adopted, demonstrating the competitive sampling efficiency of our method.

Table 4: Results on the {GEOM-QM9} dataset. The threshold $\delta=0.5A$

Table 2: Results on 12 quantum mechanics prediction tasks from QM9. We take 110K for training, 10K for validation, and 11K for testing. The evaluation is mean absolute error~(MAE), and the best results are marked in bold. The backbone is SchNet.

Table 3: Results on eight force prediction tasks from MD17. We take 1K for training, 1K for validation, and 48K to 991K molecules for the test concerning different tasks. The evaluation is mean absolute error, and the best results are marked in bold.

Thank you for your endorsement and the positive comments! Below please find the responses to some specific comments.

W1: The main weakness in my mind is about the performance. The generation quality in Tab1 and Tab5 seems even worse than the baseline GeoDiff, and the property prediction results in Tab2 are also not consistently better than baselines and GeoDiff.

Response 1: Thanks for your careful reviews! 1) For a more comprehensive comparison with GeoDiff, we include results using 200 and 500 diffusion steps (Table 4 below). The findings suggest that our method significantly outperforms GeoDiff when reducing the number of diffusion steps, indicating an improvement in sampling efficiency. 2) Additionally, we conduct extra experiments to support property prediction tasks, and the results are shown in Table 1, Table 2, and Table 3 below. In these experiments, we employ another stronger score network (a pretraining framework provided by MoleculeSDE[1]). The new experimental results suggest that our methods can significantly outperform the baselines consistently under the self-supervised learning context, implying our method can be used as an important technique for improving the pretrained model to capture the substructure information during training.

[1] Liu S, Du W, Ma Z M, et al. A group symmetric stochastic differential equation model for molecule multi-modal pretraining//International Conference on Machine Learning. PMLR, 2023: 21497-21526.

Table 1: Results for molecular property prediction tasks (with 2D topology only)

Q6: Incorporating a well-designed prior distribution could also reduce the challenging learning task, e.g., EigenFold. The model should be able to figure out the correlation between data dimensions during training. Would this be more effective and easier to learn than introducing masks for the diffusion process?

A6: Thanks for your valuable advice! 1) A well-designed prior distribution is a potentially effective method to mitigate the independence of Gaussian noise injected into atomic coordinates and simultaneously enhance the score network's ability to perceive substructure. 2) However, the harmonic prior provided in EigenFold cannot be directly used to tackle this learning task. Specifically, sampling from the harmonic prior restricts the diversity of the noisy conformations input to the score networks. Consequently, when using this diffusion method as a pretraining task, it may simplify the learning process for the generation task but fail to enhance the model's capability to perceive substructure. As a result, downstream tasks may not derive significant benefits from the pretraining. 3) Although the harmonic prior contains graph structure information, it does not contribute to enhancing the denoising network $s_\theta(G,R^t,t)$ in capturing substructure. This is because the input $R^t$ already contains sufficient prior information (via adding noise) to help $s_\theta$ predict the noise related to the harmonic prior. In contrast, our method introduces masks, allowing the model to learn prior knowledge by predicting the mask.

Thanks again for your useful insights, which help us a lot to improve our work.

Q4: Diffusion models are known for their slow sampling process. Here the large number of steps (i.e., 5000) would make the sampling process particularly inefficient. For instance, for protein conformation generation, the number of steps is typically 200 or 500. Upon reducing the number of diffusion steps, would the alternating update scheme still work?

A4: Thanks for your valuable question. The proposed method still demonstrates effectiveness when the number of diffusion steps is reduced. Following your instruction, we performed experiments for a diffusion step sensitivity analysis, and the results can be found in Table 5 below. As observed in Table 7, our method significantly outperforms GeoDiff across all evaluation metrics when the number of steps is reduced to 500 and 200. Furthermore, our method even achieves competitive performance against 5000-step ConfGF, highlighting its ability to enhance the denoising efficiency of the denoising network through the introduction of the mask.

Table 7: Results on the {GEOM-QM9} dataset. The threshold $\delta=0.5A$

Q5: Does other sampling techniques, e.g., probability flow ODE, work for the proposed method?

A5: Thanks for your insightful question! 1) Theoretically, ODE-based sampling methods cannot be directly applied to our approach. This is due to our method combining a discrete distribution (mask distribution $p(S)$) into the Gaussian distribution. If we transform the discrete Diffusion process to its corresponding continue version (i.e., increasing the number of noise scales $N\to \infty$, $\Delta t := 1/N \to 0$), the discrete mask distribution wouldn't be converted into a continuous function, resulting in that the Markov chain $\{R^i\}_{i=1}^{N}$ cannot converge to a continuous stochastic process[1] when $\Delta t := 1/N \to 0$. Therefore, obtaining ODE-based sampling methods from rigorous mathematical derivations is not possible. 2) In practice, it is possible to apply ODE-based sampling methods by exclusively utilizing the noise predictor for sampling and disregarding the mask predictor. This is feasible since the noise predictor can be well-trained among all nodes through uniform subgraph sampling.

[1] Song, Yang, et al. "Score-Based Generative Modeling through Stochastic Differential Equations." International Conference on Learning Representations. 2021.

Table 4: Results on eight force prediction tasks from MD17. We take 1K for training, 1K for validation, and 48K to 991K molecules for the test concerning different tasks. The evaluation is mean absolute error, and the best results are marked in bold.

Q3: From Table 5 in the Appendix, the proposed method is inferior to GeoDiff across all evaluation metrics. Minor performance gain on GEOM-QM9 may not be that significant since the target molecules are small in general. That being said, MaskedDiff cannot outperform SOTA models while introducing extra complications during modeling.

A3: Thanks for your careful reviews! It is important to clarify that our method isn't inferior to GeoDiff across all evaluation metrics. We should consider different sampling methods when making comparisons with GeoDiff from Table 5 in the Appendix. 1) Our method outperforms GeoDiff among all the Recall metrics when we use DDPM sampling(ancestral sampling) methods, demonstrating that our method can effectively improve the diversity of the generative samples compared to GeoDiff. The results are shown in Table 5 below, where the values are obtained by evaluating from the released checkpoint. 2) For a more comprehensive comparison, we include results using 500 diffusion steps. The findings suggest that our method significantly outperforms GeoDiff when reducing the number of diffusion steps, indicating an improvement in sampling efficiency. On the other hand, 3) it is true that our methods cannot outperform GeoDiff when we use the same 5000-step Langevin dynamics(LD) sampling(Table 6). This is attributed to the LD sampling method's inability to consider mask prediction when using our model trained by SubGDiff, which could be a potential problem for generating relatively larger molecular graphs in the Drugs dataset with a large number of steps (e.g., 5000). This challenge can be alleviated when reducing the number of steps (e.g. 500), as demonstrated in Table 6.

We will add more descriptions to avoid this misunderstanding. Thanks again for pointing it out!

Table 5: Results on the {GEOM-Drugs} dataset. The threshold $\delta=1.25A$ and the sampling method is DDPM(ancestral sampling)

Table 6: Results on the {GEOM-Drugs} dataset. The threshold $\delta=1.25A$. Both models adopt the same Langevin dynamics sampling method.

Table 2: Results for molecular property prediction tasks (with 2D topology only)

Table 3: Results on 12 quantum mechanics prediction tasks from QM9. We take 110K for training, 10K for validation, and 11K for testing. The evaluation is mean absolute error~(MAE), and the best results are marked in bold. The backbone is SchNet.

We thank the valuable comments from the reviewer! We provide point-wise responses below.

Q1：The number of diffusion steps is the same as that in GeoDiff, i.e., 5000 steps, which is quite a lot. Given this large amount of diffusion steps, the alternating update scheme might work well. However, I am not sure if using more cost-effective sampling methods would still benefit from masked diffusion.

A1: Thanks for your insightful comments. We conduct experiments with fewer sampling steps (500,200) to evaluate our method and the results are shown in Table 1 below. The findings suggest that: 1) With 500 steps, our method can outperform existing methods like ConfGF, which uses 5000 sampling steps; 2) Our method can consistently outperform GeoDiff for both DDPM and LD sampling methods, particularly when the number of diffusion steps is small. This indicates that different sampling methods can benefit from our masked diffusion.

Table 1: Results on the {GEOM-QM9} dataset. The threshold $\delta=0.5A$

Reference:

[1] 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.

Q2: Taking masking into consideration makes the diffusion process more complicated by design. Therefore, one should show that this extra modeling complexity is worthwhile by showing more promising performance on various downstream applications. The marginal performance gain on the GEOM-QM9 dataset (not GEOM-Drugs) is not convincing enough.

A2: Thanks a lot for the helpful advice. We conduct additional experiments on self-supervised learning tasks, evaluating our approach to various downstream tasks. These tasks include 2D molecular property prediction (MoleculeNet) and 3D molecular property predictions (QM9 and MD17). The results are presented in Tables 2, 3, and 4 below. In the experiments, we incorporate the score network used in MoleculeSDE[1] to assess the effectiveness of SubGDiff. The results indicate that our methods surpass the baseline models in the realm of self-supervised learning. This implies that our technique can serve as a valuable tool for enhancing pretrained models in capturing substructure information during training.

Beyond representation learning, we highlight that our method exhibits superior sampling efficiency for conformation generation (refer to Table 1 above).

Reference:

[1] Liu S, Du W, Ma Z M, et al. A group symmetric stochastic differential equation model for molecule multi-modal pretraining//International Conference on Machine Learning. PMLR, 2023: 21497-21526.

We thank the reviewer for the constructive comments. We provide pointwise responses below.

Q1. My major concern is that, the overall motivation is not convincing or valid. The authors think diffusing individual atoms in a molecule may constrain the capacity of the models, as the connections among these atoms can be important. However, the methods tend to be a mask to choose a subset of atoms, which can be anywhere of the original graph and disconnected at all. As a comparison, there exist a bunch of studies [1,2] that consider REAL subgraphs (like motifs, functional groups, etc) in generation by leveraging domain knowledge. The motivation of these existing works is more convincing and valid.

A1: We sincerely thank the reviewer's comment! 1) The reviewer rightly emphasizes the importance of connections among subgraph atoms, and we agree that randomly selecting a subset of atoms could result in disconnected subgraphs. Our proposed method is a versatile framework, allowing us to design specific mask distributions to ensure the connected property of the selected subgraph. We will provide detailed instructions on constraining the mask distribution when introducing maskedDiff. 2) In our setting, we employ the torsional-based decomposition method (refer to sec 5.1 in the main text), in which the torsional edge in a molecule will be randomly selected and cut off to make the molecule into two components. This effectively guarantees the connectivity of the subgraphs. 3) Using the REAl subgraphs proposed by [1,2] in our framework is also a potentially effective approach. One key difference between our work and [1,2] is that in our work the atoms of a subgraph are denoised in the same step, and no 3D information of the subgraph (3D local structure) is directly input to the model. Our model aims to learn these prior domain knowledge so as to enhance the capability of 3D representation. While in [1,2], the 3D structure of the REAL subgraphs is given to the model as a prior and their aim is to make the generation task easier.

Q2. The title and scope of the paper are a little weird. The work is actually not regular 3D molecular representation learning, but for molecular generation and self-supervised based representation learning. I suggest the authors consider this in the revision.

A2: Thanks for the suggestion. One of our motivations is that a generation task can be used for self-supervised representation learning. To better incorporate the 3D information, we choose the 3D molecular conformation generation as our representation learning task and introduce subgraphs. 3D molecular representation learning is our focus. We will try to make this clearer in the paper and maybe use a more appropriate title.

Q4: To me, the implementation of the training and sampling of diffusion models on a subset of a graph is not trivial, but there exist several studies for similar purposes, like [1,2,3]. The authors may want to clearly state the unique technical contribution compared with these studies.

A4: Thanks a lot for suggesting these studies. We highlight the differences between our method with [1,2,3] as follows.

• Compare to [1]: D3FG[1] adopts three different diffusion models (D3PM,DDPM,and SO(3) Diffusion) to generate three different components of molecules(linkerr types, center atom position, and functional group orientations), respectively. In general, these three components can also be considered as three subgraphs(subset). Essentially, D3FG firstly selects the subgraph and only injects noise into the fixed subgraph during the entire diffusion process, while our method has the flexibility to select different subgraphs from a mask distribution at each time step during the forward process. Further, our model is actually a single diffusion model that can enhance the denoising network to perceive substructure information by incorporating a mask variable in it.

• Compare to [2]: DIffPAck[2] is an Autoregressive generative method predicting the torsional angle $\chi_i (i=1,2,..,4)$ of protein side-chains conditioned on $\chi_{1,...,i-1}$, where $\chi_i$ is a predefined subset of atoms. It uses a torsional-based diffusion model to approximate the distribution $p(\chi_i| \chi_{1,...,i-1})$. Each $\chi_i$ requires a separate score network for estimation. Essentially, it can be viewed as selecting a subset initially and then adding noise exclusively to the fixed subset throughout the entire diffusion process. In contrast, our method takes a different approach by randomly sampling a subset from a mask distribution $p(S)$ at each time step during the forward process. This approach offers greater flexibility and cost-effectiveness, requiring only a score network and a subgraph predictor.

• Compare to [3]: [3] introduces an autoregressive diffusion model GraphARM, which absorbs one node in each time-step by masking it along with its connecting edges during the forward process. Unlike GraphARM[3], instead of masking individual nodes, selects a subgraph in each time-step to inject the Gaussian noise. This is equivalent to masking several nodes during the forward process. Moreover, GraphARM mandates that the number of steps must be the same as the number of nodes due to the usage of the absorbing state. In contrast, our method can set any number of diffusion steps theoretically since we use the real-value Gaussian noise.

We will add those comparisons to our revised version. Thanks for your constructive suggestions.

References:

[1] Haitao Lin et al, Functional-Group-Based Diffusion for Pocket-Specific Molecule Generation and Elaboration, NeurIPS 2023.

[2] Yangtian Zhang et al, DiffPack: A Torsional Diffusion Model for Autoregressive Protein Side-Chain Packing, NeurIPS 2023.

[3] Lingkai Kong, Autoregressive Diffusion Model for Graph Generation, ICML 2023.

Table 2: Results on 12 quantum mechanics prediction tasks from QM9. We take 110K for training, 10K for validation, and 11K for testing. The evaluation is mean absolute error~(MAE), and the best results are marked in bold. The backbone is SchNet.

Table 3: Results on eight force prediction tasks from MD17. We take 1K for training, 1K for validation, and 48K to 991K molecules for the test concerning different tasks. The evaluation is mean absolute error, and the best results are marked in bold.

Q3. Experimental results on GEOM-0M9 are very marginal, especially on the Co-P and MATP metrics where GeoDifis much better.I also have concerns about the experiment setting that both generation and representation learning are conducted. However, for each one, the setup (like datasets) is not sufficient or the results are not strong enough compared to baseline methods. A better solution is to only study one problem (like generation) with sufficient and strong results. This is also related to the title and scope of this work.

A3: Thank you for your valuable suggestions. 1) We have updated the results for conformation generation in the revised version. 2) Following your advice, we've shifted our focus to self-supervised representation learning here. Additional experiments were conducted, and the results are presented in Tables 1, 2, and 3 below. In these experiments, we assessed the efficacy of SubGDiff by incorporating the score network from MoleculeSDE[1], a pretraining framework. The results demonstrate that our method outperforms baseline models in the realm of self-supervised learning. This underscores the utility of our technique in enhancing pre-trained models to capture substructure information during the pre-training process.

[1] Liu S, Du W, Ma Z M, et al. A group symmetric stochastic differential equation model for molecule multi-modal pretraining[C]//International Conference on Machine Learning. PMLR, 2023: 21497-21526.

Table 1: Results for molecular property prediction tasks (with 2D topology only)

Thank you so much for your response! We would like to provide the following clarification.

Summary: The paper studies 3D molecular graphs, and performs diffusion on a subgraph sampled by a learnable mask vector. Then the major work is to fulfill the training and sampling of the diffusion model based on the mask strategy.

Reponse: We clarify that our method does not adpot a learnable mask vector during the diffusion process. Actually, a mask vector corresponds to a subgraph, and the mask vector is sampled from a predefined mask distribution. Thus our selected subgraph is actually a REAL subgraph. In our method and experiments, we pre-define the mask distribution to be a discrete distribution with a sample space $\chi =\\{G ^i_{sub}\\} _{i=1}^{N}$, and $p _{s _t} (\mathcal{S} = G _{sub} ^i )$ , $= 1/N,t>1$, where $G ^i _{sub}$ is the subgraph split by the torsional-based decomposition methods [1]. The decomposition approach will cut off one torsional edge in a 3D molecule to make the molecule into two connected components, each containing at least two atoms. The torisional edge here acutally corresponds to the Bridge(or Cut Edge) in Graph Theory. The two components are represented as two complementary mask vectors (i.e. $\mathbf{s}'+\mathbf{s}=\mathbf{1}$). Thus $n$ torsional edges in $G _{3D}$ will generate $2n$ subgraphs. Finally, for each atom $v$, the $s _{t_v} \sim \text{Bern}(0.5)$, ensuring that every atom can be added noise.

We have revised our PDF (highlighted in orange) for clarification.

Reference:

[1] Jing B, Corso G, Chang J, et al. Torsional diffusion for molecular conformer generation[J]. Advances in Neural Information Processing Systems, 2022, 35: 24240-24253.

Further comment 1: The motivation of diffusing a subset is still not convincing to me.

Response 1: The motivation for using subgraph (subset) sampling during diffusion is to enhance the perception of 3D molecular structures in the denoising network. The subgraph (subset) is predefined, not a random subset. The mask distribution is equivalent to the subgraph distribution, defined over a predefined sample space $\chi =\\{G ^i_{sub}\\} _{i=1}^{N}$, where each sample is a connected subgraph extracted from $G$. The mask predictor $s _{\vartheta}$ learns the substructure information by predicting the mask, which serve as an additonal task to denoising. Thus, the denoising network $\epsilon _\theta$ can improve its ability to capture the substructure through sharing the GNN encoder with the mask predictor $s _{\vartheta}$.

The whole training process can be summarized in the following algorithm:

for $i$ in $\{1,2,...,K\}$; do

1. Sample $G_{3D} \sim q(\mathrm X)$ // Sample a 3D molecule graph
2. Sample $t \sim \mathcal{U}(1, ..., T)$, $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I}))$
3. Create sample space $\chi = \\{ G _{sub} ^i \\} _{i=1} ^{N}$; $G ^i _{sub} \in$ Torsional_decomposition$(G _{3D})$ / / For the mask distribution $p _{s _t}(\mathcal{S})$
4. Sample $\mathbf{s}^t \sim p_{s_t}(\mathcal{S})$ //Sample a mask vector (subgraph)
5. $R^t \gets q(R^t | R^{0})$ // Eq. 14
6. $\mathcal{L} = \\|\text{diag}(\mathbf{s} _t) (\epsilon-\epsilon _{\theta}(\mathcal{G}, R^t, t))\\|^2 +\lambda \mathrm{BCE}(\mathbf{s} _t, s _\vartheta(\mathcal{G}, R^t, t))$
7. $\text{optimizer.step}(\mathcal{L})$

end

I thank the authors for their efforts in adding more experiments on SSL. However, the motivation of diffusing a subset is still not convincing to me; the scope of this work has not been changed; and empirical studies on generation are still weak. Given these considerations, I feel a significant revision is needed so that the paper can be resubmitted to another venue. I will maintain my score.