CFD: Learning Generalized Molecular Representation via Concept-Enhanced Feedback Disentanglement

3. Distinction between Our Concept Mining and Residual Vector Quantization in iMoLD?

Our molecule concept learning module is significantly different from Residual Vector Quantization (RVQ). In essence, RVQ is to perform feature replacement. Our method is to build the connection to capture critical molecule substructures.

(1) The motivation is different. RVQ aims to leverage the discretized continuous representation to improve the model generalization. As shown in Fig. 6, our method is to learn a series of concepts that could capture critical molecule structure information involving plentiful task-specific invariant molecule characteristics.

(2) The operations are different. RVQ first introduces a shared learnable codebook as a discrete latent space. For each node representation in molecular graph, RVQ looks up and fetches the nearest neighbor in the codebook and outputs it as the result. Meanwhile, a sum operation is further used to strengthen the representation ability. During training, similar to VQ-VAE, RVQ employs the exponential moving average updates for the codebook.

Differently, our concepts are first sampled from a learnable Gaussian distribution. And we employ KAN operation to promote the learning of concepts. Importantly, as shown in Eq. (6), the weighted residual operation is to align critical molecule substructure features to the corresponding concept, building the connection between molecule substructures and corresponding concepts. During training, we only utilize the task loss, e.g., molecule property prediction, to promote the learned concepts to capture plentiful task-specific invariant molecule characteristics.

(3) Experimental results show the effectiveness of our method. Taking GOOD-ZINC as the example, for iMoLD, we first remove the RVQ operation and keep other operations unchanged. We observe that removing RVQ does not affect the performance. Instead, our concept mining module could be plugged into iMoLD to further improve its performance. For the covariate case of scaffold, the performance is improved by 0.0142. Furthermore, the prediction results of the molecule's ground-state conformation in Fig. 4 and 5 also show that our learned concepts are beneficial for capturing critical substructures.

Thanks for your valuable comments. We have added these contents to the revision and provide a figure to further illustrate the distinction (as shown Sec. G in appendix).

4. Further Analysis of the Loss Function

Eq. (10) is the optimization objective, consisting of three terms, i.e., $L_{pred}$, $L_{fd}$ (as shown in Eq. (8)), and $L_{vd}^{t}$ (as shown in Eq. (5)). Here, we provide two propositions:

Proposition 1. $L_{vd}^{t}$ mainly is a contrastive loss. Minimizing $L_{vd}^{t}$ could enlarge the gap between the disentangled $\tilde{Z}^{Gen}$ and $Z^{Spu}$. Meanwhile, $\tilde{Z}^{Gen}$ is promoted to retain plentiful input-related content.

Proposition 2. Minimizing $L_{fd}$ is instrumental in achieving the invariant representation $Z^{Gen}$.

Thanks for your valuable comments. In Sec. H of appendix, we give corresponding analyses about these two propositions.

5. Why CIGA, DisC, and CAL fail to the adapt on GOOD-ZINC and GOOD-PCBA?

The reason that leads to the weak disentangled ability of CIGA, DisC, and CAL mainly lies in that they only exploit the one-step disentangled mechanism and could not apply well to the complex molecules involving more atoms and substructures. Particularly, CIGA employs the one-step alignment to learn invariant representation. DisC utilizes a parameterized edge mask generator to separate causal and bias subgraphs, whose performance highly relies on the generation quality. Finally, CAL exploits attention modules to estimate the causal and shortcut features of the input graph. Only using attention could not effectively disentangle complex molecules.

This further demonstrates that using iterative feedback disentanglement is an effective mechanism for learning generalized molecule representations.

Dear Reviewer,

Thanks for your valuable comments and recognition for our work. To advance the development of biochemical research, e.g., drug discovery, learning generalized molecule representations is important. Currently, most existing methods utilize graphs to represent molecules and mainly employ a one-step disentangled mechanism to obtain invariant characteristics. However, we observe that this one-step method could not cope with molecules involving complex entangled structures. Thus, we first design a concept-enhanced feedback disentanglement to iteratively conduct molecule decomposition. Extensive experimental results demonstrate the effectiveness of our method.

In the following, for the helpful questions, we provide corresponding replies:

1. About Chemical Inspiration of Our Method

This method is mainly based on the following chemical observations:

(1) A chemical molecule is a complex composed of multiple atoms and bonds. The reason for the OOD problem is the structural change of molecules. For example, changes in molecular chirality can lead to deviations in molecular properties. Thus, learning robust molecule representation is important for molecule-based tasks, e.g., drug discovery.

(2) Lines 082-095 have indicated that existing methods could be classified into two types, i.e., First-Separation-Then-Encoding and First-Encoding-Then-Separation (as shown in Fig. 1 (a) and (b)). As indicated in Lines 272-283 and Fig. 3, these methods for learning generalized molecular representation generally employ a one-step disentangled strategy, i.e., directly separating the input into generalized and spurious parts. However, in practice, we may encounter some unknown biomacromolecules containing more atoms and highly complex structures, e.g., entanglement of multiple substructures. At this time, using the original one-step mechanism could not obtain satisfactory disentangled results. To this end, we exploit the feedback idea to iteratively and progressively separate generalized molecular representation.

(3) In general, the molecule's ground-state conformation belongs to the substructures of molecules and determines their properties. Therefore, it is important to capture substructure-related characteristics. To this end, we define a series of concepts to build connections between molecule substructures and corresponding concepts, which further improves the performance.

Thanks for your valuable comments. We have added these contents to the revision (as shown in Sec. E of appendix).

2. Evaluation on Other Graph-based Out-of-Distribution Generalization Problem

Thanks for your recognition for our work. We follow the work [1] and further evaluate our method on Two-Piece Graph Datasets [1]. Particularly, each dataset is generated from a variation of two-piece graph model, denoted as {a, b}, where a refers to the invariant correlation strength and b refers to the spurious correlation strength. The results are shown as follows:

We can observe that our method indeed improves the performance. Particularly, when the invariant correlation strength is weaker than the spurious correlation strength, the performance of our method outperforms state-of-the-art methods, further demonstrating that our method is indeed beneficial for disentangling invariant characteristics.

Thanks for your valuable comments. We have added these contents to the revision (as shown in Sec. F of appendix).

[1] Chen Y, Bian Y, Zhou K, et al. Does invariant graph learning via environment augmentation learn invariance?[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2] Fan S, Wang X, Mo Y, et al. Debiasing graph neural networks via learning disentangled causal substructure[J]. Advances in Neural Information Processing Systems, 2022, 35: 24934-24946.

[3] Chen Y, Zhang Y, Bian Y, et al. Learning causally invariant representations for out-of-distribution generalization on graphs[J]. Advances in Neural Information Processing Systems, 2022, 35: 22131-22148.

Dear Reviewer:

Thanks for your recognition of our work. We will modify this paper carefully according to your valuable comments.

1. The Reason of Using KAN for Learning Concepts

In general, molecules contain multiple atoms and bonds. For certain biomacromolecules, their structures are entanglement of various substructures, which poses the challenge for learning molecule concepts.

Inspired by the Kolmogorov-Arnold representation theorem [1, 2], KANs allow learning custom activations of the network. In this way, it is possible to analyze the contribution of individual components of the input, then providing a more transparent view of the network's decision-making process.

Thus, using KANs could sufficiently convert the analyzed components to the corresponding concepts that involve plentiful molecule-aware substructure information, which is beneficial for strengthening the generalization of molecule representations.

Lines 474-477 have provided an ablation experiment of KANs. We observe that using KANs could indeed improve the performance, indicating its superiorities.

[1] Arnold V I. On functions of three variables[J]. Collected Works: Representations of Functions, Celestial Mechanics and KAM Theory, 1957–1965, 2009: 5-8.

[2] Kolmogorov A N. On the representation of continuous functions of several variables by superpositions of continuous functions of a smaller number of variables[M]. American Mathematical Society, 1961.

2. Further Analysis of Contrastive Loss

As shown in Eq. (5), by means of variational disentanglement, $\tilde{Z}^{Gen}$ and $Z^{Spu}$ separately represent the corresponding outputs. Particularly, minimizing the contrastive loss could enlarge the gap between $\tilde{Z}^{Gen}$ and $Z^{Spu}$, and promote $\tilde{Z}^{Gen}$ to contain plentiful input-related information. Then, as shown in Eq. (7), $\tilde{Z}^{Gen}$ is used to calculate the generalized output at the current step. Therefore, high-similarity $\tilde{Z}^{Gen}$ constructs the foundation of forming the final generalized representation.

3. Examining Importance of Concepts Across Iterations

Table 5 analyzes the impact of iteration number. We can observe that when the iteration number increases, inserting concepts could improve the performance effectively, which indicates the effectiveness of the learned concepts. To further analyze the impact of iteration number on concepts, we also make a new ablation experiment, i.e., removing the concepts at the 4th, and 8th iteration. The results are shown as follows:

We can observe that when the iteration number is small, inserting concepts could significantly affect the performance. Instead, when the iteration number is large, the impact of concepts is somewhat weak. The reason may be that when the iteration number is small, leveraging concepts is beneficial for capturing critical substructures, improving the performance. As the iteration number increases, concepts progressively pay attention to meaningful substructure characteristics, reducing the iteration impact on performance.

Thanks for your valuable comment. We have added these contents in Sec. I of appendix.

4. More Analysis of Molecule Concepts

We perform an ablation experiment of the decoder that calculates task-related representation. We observe that for classification and regression tasks, using decoders with more layers leads to the performance degradation. The reason may be that this operation introduces more parameters, resulting in the overfitting risk.

In Fig. 7, we show some visualization examples of molecule concepts. Since these concepts are learned from all molecule data, they are not specific to certain kinds of molecules, which is beneficial for improving the generalization of molecule concepts. For different predictions, we can observe that the learned concepts could indeed focus on critical molecule substructures, which strengthens the performance of molecule prediction.

Thanks for your valuable comment. We have added these contents in Sec. J of appendix.

Dear Reviewer:

Thanks for your helpful comments. We will modify this paper carefully.

Learning robust molecule representation is important for advancing the development of biochemical research. Currently, existing methods mainly employ a one-step disentangled mechanism to learn invariant representations. However, when encountering some unknown biomacromolecules containing more atoms and highly complex structures, one-step methods could not sufficiently decompose structure-entangled molecules, affecting the performance of downstream tasks. To this end, we first propose a feedback mechanism to iteratively disentangle invariant representations. In the submitted paper and appendix, we perform extensive experiments and visualization analysis to indicate the effectiveness of our method.

1. More Detail Introduction

Thanks for your comments. We will modify this paper carefully and add more details.

(1) $\mu_{c}$ and $\sigma_{c}$ are taken as the parameters. And we take the final objective $\mathcal{L}_{\rm task}$ (as shown in Eq. (10)) to perform optimization.

(2) The goal of using residual operation is to align the molecule substructures to the corresponding concepts and build the connection between molecule substructures and concepts, which is instrumental in enhancing the generalization of molecule representations.

(3) Through multiple graph networks, each node of the output integrates features of its neighborhood nodes, which could be considered as a representation of a molecule substructure. By means of calculating weight and residual operation (as shown in Eq. (6)), each concept could capture substructure characteristics related to the current concept. In Fig. 7 of appendix, we show the substructures captured by the corresponding concepts. We observe that our method is indeed helpful for leveraging concepts to capture molecule substructures.

(4) The goal of Eq. (1) is mainly to interpret the reason leading to the OOD problem on molecule representation learning. Lines 197-199 have indicated that the reason lies in that the training data only cover very limited environments in $\xi$ while the model is expected to perform well on all environments. Thus, how to learn generalized representations in limited environments is important for strengthening the robustness against distribution shifts.

2. Analysis of KAN

In general, molecules contain multiple atoms and bonds. For certain biomacromolecules, their structures are entanglement of various substructures, which poses the challenge for learning molecule concepts.

Inspired by the Kolmogorov-Arnold representation theorem [1, 2], KAN allows learning custom activations of the network. In this way, it is possible to analyze the contribution of individual components of the input, then providing a more transparent view of the network's decision-making process.

Thus, using KAN could sufficiently convert the analyzed components to the corresponding concepts that involve plentiful molecule-aware substructure information, which is beneficial for strengthening the generalization of molecule representations.

Lines 474-477 have provided ablation experiments of KAN. During mining molecule concepts (Eq. 6 and 7), we employ KAN to improve the concept accuracy. Here, we find that replacing KAN with MLP results in performance degradation. For example, for Cov-Sold case, the performance is decreased by around 0.6%, indicating KAN's effectiveness.

Besides, in Sec. D of appendix, we have analyzed the impact of hyper-parameters. Meanwhile, we also provided more experimental analyses in the modified version, e.g., Sec. I, K, L, and M of appendix. These analyses all indicate the effectiveness of our method.

[1] Arnold V I. On functions of three variables[J]. Collected Works: Representations of Functions, Celestial Mechanics and KAM Theory, 1957–1965, 2009: 5-8.

[2] Kolmogorov A N. On the representation of continuous functions of several variables by superpositions of continuous functions of a smaller number of variables[M]. American Mathematical Society, 1961.

My real concern is that the features after multi-disentanglement still contain spurious components, making them unable to effectively represent different substructures. How does the multi-step feedback ensure that the concatenated feature is truly invariant?

If the desired feature is not obtained by the one-step disentanglement, the final step of feedback-disentanglement should yield the invariant feature instead of concatenating all features, because $Z_0^{\text{Gen}}, Z_1^{\text{Gen}}, \dots$ still contain spurious components. Please conduct a comparison experiment that considers only the features from the final step of disentanglement.

Plus, KAN's rationality was not accepted.

If the author can response, I am willing to raise the score further.

Dear Reviewer:

Thanks for your recognition of our work. We will modify this paper carefully according to your valuable comments.

1. Quantify Contribution of Feedback Mechanism

To quantify the contribution of feedback iterations, we make an ablation experiment of the iteration number. Here, we only change the iteration number. Other operations are kept unchanged. The results are shown as follows:

We can observe that compared with one-step disentanglement, using multiple feedback iterations indeed improves the performance, indicating that our method is beneficial for decomposing invariant molecule representations. Besides, we also observe that when the iteration number is larger than 12, the performance improvement is limited. The reason is that for the current dataset, using 8 iterations is sufficient for disentangling highly complex molecules.

Finally, we observe that complex molecules usually contain more atoms. Therefore, to further demonstrate the effectiveness of our method on highly complex molecules, we select those molecules with more than 12 atoms for additional evaluation. The results are shown as follows:

We can see that for complex molecules, using one-step disentanglement could not decompose invariant characteristics effectively. Using multiple feedback iterations could indeed decouple the representations involving plentiful invariant information, which strengthens the generalization ability.

Thanks for your valuable comments. We have modified Table 5 and added these content in Sec. L of the appendix.

2. Protein-Ligand Binding-Affinity Prediction

For OOD molecular representation learning, in Table 1 and 2, we show the corresponding results. Furthermore, in Table 3, our method is evaluated on an additional task, i.e., molecule's ground-state prediction. These results demonstrate that our method could effectively disentangle invariant molecule characteristics, which strengthens the generalization ability.

Finally, to further demonstrate the superiorities of our method, we evaluate our method on protein-ligand binding-affinity prediction. Binding affinity is an important metric in biology and medicinal chemistry that measures the strength of a reversible association between biological macromolecules, such as proteins or DNA, and small-molecule ligands, such as drugs [1]. Here, we follow the work [1] and employ the same datasets and metrics to verify our feedback disentanglement method. The results are shown as follows:

Here, MAT, DimeNet, CMPNN are three GNN-based methods. RMSE, MAE, and R are three pre-defined metrics. We can observe that for the challenging protein-ligand binding-affinity prediction, employing our method still improves the performance for the given three metrics. This further demonstrates that for biological macromolecules, using feedback disentanglement is meaningful, which improves the generalization of the learned molecule representations.

Thanks for your valuable comments. We have added these content in Sec. M of the appendix.

[1] Boulougouri, Maria, Pierre Vandergheynst, and Daniel Probst. "Molecular set representation learning." Nature Machine Intelligence 6.7 (2024): 754-763.