Learning Invariant Molecular Representation in Latent Discrete Space

Summary

This paper presents a new graph neural network architecture and objective function that encourages models to identify features that are invariant to distribution shifts in the data. The proposed method, iMoLD performs invariant feature extraction in the latent embedding space and leads to improved performance across an extensive set of molecular property prediction tasks.

Strengths

- The presented idea is novel and leads to improved performance across a variety of datasets and tasks. - Experimentation is extensive with good results. - The ablation analysis is Section 5.3 and sensitivity analysis from Appendix C are useful.

Weaknesses

#### **Incorrect definitions in Section 3.1** - I believe there is some issue in the notation of Section 3.1. Specifically, the definitions of $P_{train}, P_{test}, P_{all}$ as collections of distributions, means that they are not themselves valid probability distributions. I believe some re-normalization would be required here. --- #### **Use of term “Discrete Latent space” is unclear** - Why do the authors claim they have a **“discrete”** latent space? The residual connection between $\mathbf{H}$ and the quantized representation means that embeddings are continuous. Additionally the element-wise gating to create $\mathbf{H}^{\mathrm{Inv}}$ and $\mathbf{H}^{\mathrm{Spu}}$ means the model does not have a discrete latent representation. --- #### **Unclear elements about the learning objective** - The notation in Equation (11) is confusing. Specifically, what is the dimensionality of the $\tilde{\mathbf{z}}_i^{\mathrm{Inv}}$? Are you concatenating multiple batch samples from $\mathbf{z}^{\mathrm{Spu}}$ to $\mathbf{z}_i^{\mathrm{Inv}}$ or just one random one? - There seems to be an inherent tension between the residual connection and the commitment loss $\mathcal{L}_{\mathrm{cmt}}$. That is, if this loss were perfectly minimized, then the residual connection would be negated. - The role of $\gamma$ in the scoring regularization is not well described. --- #### **Baseline presentation is confusing** - It seems that the authors are conflating baselines in terms of loss objectives and in terms of model/architecture designs. It would be good to clarify which baselines rely on the same architecture but have different objectives (e.g. ERM) and which constitute an entirely different modeling scheme (e.g., CIGA). For the baselines that simply differ in objective, it would be good to also make explicit (could go in Appendix) if any model / architecture adjustments were also applied. --- #### **Other minor comments** - In line 147, the notation for edges $\mathcal{E}$ is overloaded, since the same variable is used to denote environments in Section 3.1. - At the end of Section 5.4 (lines 341-345), the authors seem to be mixing the meaning of low/high in terms of whether low = “good” or low = ”bad”. - $D$ and $Score$ should be defined explicitly in Figure 4 caption.

Questions

Q1) It is not clear to me why vector quantization (VQ) is the right “bottleneck” to use here. Other than restricting the model’s expressivity, which can be done in other ways such as weight regularization, why is VQ particularly suited for this setup? Q2) Why is the stop gradient applied in equation 12? Is this simply for computation efficiency / stability? If so, this should be made explicit in the text. Q3) For the GOOD-PCBA experiment, why is average precision (vs. average accuracy, recall, or ROC-AUC) used? Q4) I know that there is an extensive sensitivity analysis in the appendix, but what are the hyperparameter configurations for the reported results in the main text (Tables 1 and 2)? Are the “best” iMoLD models sensitive to hyperparameter choice or do you see a general trend as to which configurations perform best?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

2 fair

Contribution

4 excellent

Limitations

- The current methodology does seem quite intricate with many loss terms that are justified in a somewhat ad hoc manner. - There is no real discussion of limitations / potential pitfalls relative to previous work.

Summary

While significant advances have been made in molecular representation approaches, conventional approaches typically assume that data sources are independent and sampled from the same distribution. However, molecules in real-world drug development often show different characteristics, which might be from a different distribution. This issue is called the out-of-distribution (OOD) problem. OOD challenges the generalization capability of molecular characterization methods and can degrade the performance of downstream tasks. Unlike previous studies' "first-separation-then-encoding" approach, this study proposes a "first-encoding-then-separation" molecular graph representation paradigm. Specifically, the authors first employ a GNN to encode the molecules and then employ a residual vector quantization module to alleviate the overfitting of the training data distribution while preserving the expressiveness of the encoder. Then, they score molecular representations using another GNN that measures the contribution of each dimension to the target in the latent space, thus clearly distinguishing between invariant and spurious representations. Finally, the authors propose a self-supervised learning objective that encourages the recognition of invariant features and effectively preserves label-relevant information while discarding environment-relevant information. The authors conducted experiments on real-world datasets. The experimental results show that the proposed method outperforms the SOTA methods.

Strengths

- Unlike the traditional "first-separation-then-encoding" approach, the authors propose a "first-encoding-then-separation" paradigm that uses an encoding GNN and a scoring GNN to identify invariant features from a graph. The authors use the residual vector quantization module to make a balance between the model's expressivity and generalization. The quantization is used as a bottleneck to strengthen the generalization, and the residual connection complements the model's expressivity. Moreover, the authors design a self-supervised invariant learning objective to facilitate the precise capture of invariant features. This objective is generic, task-independent, and applicable to a variety of tasks. - The model is cleared described. - The authors conducted comprehensive experiments on 18 real-world datasets. The experimental results show that the proposed model achieved stronger generalization against SOTA baselines. - Code has been released and will be valuable for future related research.

Weaknesses

- There are some typos in the paper. For example, a lack of space before references in line 31. - OOD is repeatedly defined in lines 29 and 90. In addition, please use "OOD" for "out-of-distribution" that appears later in the text, such as lines 130 and 353.

Questions

Could the authors show 3D visualization graphs representing the extracted features, as in Fig. 4?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

None.

Summary

This paper proposes a molecular self-supervised learning method for out-of-distribution generalization. The authors introduces a "first-encoding-then-separation" framework to learn invariant features. For doing so, the authors design discrete latent space with VQ-VAE. The experimental results show that their method improves previous baselines in various out-of-distribution downstream tasks.

Strengths

- The paper is well written and easy to understand. - The pre-training objective to separate invariant and spurious features seems to make sense to me.

Weaknesses

- The complexity of the proposed method is high. The loss function contains several tunable parameters and the ablation study (in Figure 5) shows that the performance is quite dependent on the choice of hyperparameters. - It seems vague why discrete latent space is needed.

Questions

- How are the hyperparameters chosen in Table 1, 2? - Is there specific intuition why discrete latent space is useful for out-of-distribution molecular representation learning?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes, the authors addressed the limitations.

Summary

The paper presents an invariant and robust representation learning approach for molecules to improve the out-of-distribution generalization performance of the predictive models. Specifically, they first map the molecule to the latent representation and then do a separation step where they separate the latent word into invariant and spurious representations. They also propose using residual vector quantization on the latent representation to avoid over-fitting while preserving the expressiveness power of the encoder.

Strengths

1. The paper tries to address an interesting problem. 2. The proposed idea is novel. 3. They included a detailed ablation study which helps identify the effectiveness of each component.

Weaknesses

1) The experimental results when compared to the baseline do not have noticeable improvement. 2) Some more details on the experiment section would be helpful, for example in Figure 4.

Questions

1) Intuitively, what is the difference between applying the Frobenius norm directly to the S matrix versus the regularization defined in equation (14)? 2) It is a bit confusing to use "h" in equation (2) and equation (12) to represent different meanings. 3) It is not clear what effect the discrete term Q(h) has on the learned final note representation H', since H' is the sum of the discrete and continuous representations. Is it possible for the model to completely ignore the discretization step and focus only on the continuous representation? 4) In line 191, it is mentioned that "It is worth noting that our separation is not only performed at the node dimension but also takes into account the feature dimension in the latent space." I didn't quite understand this. Could you provide more explanation? 5) Could you explain the intuition behind what S is learning in equation 8? Essentially, it seems S is just reweighing every element in H'. Intuitively, for the parts of H' that are not very important/invariant/main motif, S should be low so that those elements mainly contribute to the spurious representation, and vice versa. But how does the model enforces this? 6) I'm not sure if the learned high-level representation can be seen as the sum of the invariant and spurious representations. In other words, can we really break down the abstract learned representation of such a complex structure into invariant and spurious parts? Would each of these components eventually represent some substructures if decoded? 7) The paper states that the model learns discrete latent representation, but according to equation 6, the continuous representation is added back to the discretized representation. Can we still claim that the final learned representation is discrete? 8) It would be very helpful to provide a brief explanation of how the dataset is split, how the out-of-distribution is represented in the training/test/validation data, and what the terms "covariates" and "concepts" refer to in Table 1. This would provide context, especially for readers who are not familiar with the dataset.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The paper did not discuss the limitations of the work and there is no potential negative societal impact of the work.

Summary

This paper presents a new approach to obtain robust molecular representation through a first-encoding-then-separation method. The proposed method utilizes a graph neural network (GNN) as a molecule encoder and applies a residual vector quantization module to modify the representation. Additionally, a scoring GNN is employed to separate the resulting representations into spurious and invariant categories. The learning process involves contrastive-based self-supervised learning (SSL) loss and task prediction losses. Experimental results on three molecule-related benchmarks demonstrate the superiority of the proposed method over traditional debiasing techniques and recent methods designed specifically for molecule debiasing. Ablation studies and visualization techniques are conducted to provide further insights and analysis.

Strengths

1. The proposed first-encoding-then-separation approach is novel. 2. Experiments on various datasets have shown that the proposed method has the ability to achieve better results.

Weaknesses

1. The motivation is not clear. Why the first-encoding-then-separation approach is reasonable? 2. The reason to combine different components is also not clear, making the technical contribution not strong. The current version is like a straight forward combination without sufficient insight or understanding on the problem. For example, why we need a RVQ module in the molecule representation? 3. It is not clear why the proposed method has the ability to mitigate spurious biases to achieve better OOD results. Is there any theory to support that? 3. The experiments are not sufficient. For example, only improved results have been demonstrated, without sufficient analysis. In the ablation study of different modules are not consistent on different data, making the technique very ad hoc. Though some visualization have been provided, they are not sufficient to support the claim that the proposed method obtain better invariant features. What does it mean by a uniform distribution? Is the uniform distribution equivalent to a good feature?

Questions

See above in the weakness part.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Thank you for the detailed response.

I appreciate the detailed response. I do not have any additional comments or questions at this time.