Venue: NeurIPS 2022 Conference

Reviews:

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

This paper introduces techniques to enhance the robustness of molecule representation learning against distribution shifts. The authors made the observation that bio-chemical properties of molecules are usually invariantly associated with certain substructures across different environments such as scaffolds, sizes, etc. There are three important pieces in their modeling process. One is architecture-wise they use graph-level embeddings to attend a list of substructures. Second they device a new learning objective from mutual information and invariant learning to help select causal substructures. Third they use a latent representation for environment to mitigate the current issue of dealing molecular environments (e.g., often man-made, not always available, too many if using scaffolds). Overall, it's a solid paper.
strengths_and_weaknesses: ## Strengths
- the paper addresses an important topic in molecule representation learning (distribution shift).
- it introduces interesting innovative techniques to help mitigate issues of current methods

## Weaknesses
- some key notations are not clearly defined - care is needed for the writing
questions: ## Questions

- what's the exact definition of $\mathcal{G}^e$, can you add text and math formula to explain?
- it seems that the second term of eqn 7 could be simplified to $\beta \frac{1}{|\mathcal{G}|} \sum_{(G, y) \in \mathcal{G}} -log q_{\theta}(y | G)$, can you confirm?
- BBBP benchmark ROC seems much lower than typical methods would produce, is it some mistake? reference: https://arxiv.org/abs/2111.12951 appendix G table s1.
- Often times bio-chemical properties are affected by interactions between substructures, would be interesting to see if adding self-attention in the molecule encoder would help
limitations: NA
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: This work aims to address the problem of the distributional shift between training data and test data when solving the property prediction task.  The idea is to learn a representation that is invariant with respect to the environment as the environment change results in the distributional shift. The main assumption is that, regarding the property that we want to predict, there are specific structures that determine the property and other substructures that are not relevant (so no matter how that structure changes the property does not change ) and to avoid the distributional shift one needs to learn the representation that is invariant with respect to the change of the environment.
strengths_and_weaknesses: strength: the problem they try to address is interesting and valuable to the community
weakness: The assumptions and proposed model are not really convincing. The experiments section is not clearly described. Hard to follow for readers who are not familiar with the dataset, so at least in appendix explains for each data set what is the property they are predicting, and what is the environment they considering here, and how the dataset is split to train test sets to model the distributional shift. 
                  
questions: 1. The definition of distribution shift here is defined as having a different scaffold of the molecules from training to test set. In this case, is each molecule regarded as coming from a different environment as they all have different structures, it is not clear how we measure  there is a distribution shift (how different the structure of the molecule should be to be considered as there is a distributional shift? is the size of the molecules? the number of the atom, the type of the atom, to which degree) 

2. The main idea is that there is spurious structure and causal structure wrt the property.  First of all, this assumption holds for most of the properties (at least the ones that they use here). Can we really identify substructures that determine the property and are totally disentangled from the rest of the structure (spurious structure)? What if the property is related to the global structure of the molecule? 

If it is the case that the assumption holds,  if we have enough data and enough representation power, isn't the predictive model should learn to predict the property from the casual structure and ignoring the rest?  which means the representation will be naturally pushed to learn such invariant representation.  Unless we are really in a data scarcity case,  and if we are in the data scarcity case then learning to infer the environmental factor (unsupervised) would be also hard. 

3. In equation 7, the first term and last term exactly the same except one is the sum over G^e and the other one is over G , G^e represents the graph G and y sampled from environment e, so if you take the expectation at the last term with respect to e (E_e) as you have there isn't this two-term exact the same?

4. Intuitively, it seems like what the model does is given the graph and property, learn to infer the environment, and a predictive model that learns to predict the property y for the graph given environment e (objective 6). Then use this inference and predictive model to learn another predictive model that is not restricted to the environment. But my question is if the model learned from objective 6 can infer the environment and predict the property given that environment is the problem solved? so even if we have a distributional shift we can use it to predict the property for the graph from the new environment?
limitations: The assumptions and proposed model are not really convincing. 
The experiment section is a bit confusing and does not provide very structured information. For example, in table 2, what are the properties that are being predicted, and what elements are considered as the environment (still scaffold? ). If scaffold, size, and assay are considered as the environment then what is the property where it is being predicted? Also if the environment label is accessible, does the model still do inference or use it directly?
review_title: null

final_score: 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.
review_content: summary: This paper is mainly designing a learning framework to learn causal-invariance molecular representations, improving the generalization of MRL in OOD settings. This paper observes that the molecule usually consists of two parts: one part is causal substructure that determines the behaviors of molecular properties, another part is spurious substructure that constitute the molecular structures but has no impact on molecular behaviors. Based on this observation, this paper builds a structural casual model to learn more robust molecular representation learning. Specifically, new model includes substructure-aware attention-based molecular encoders and VAE-based environment inference model. To enable two-stage optimization, this paper proposes a new learning objective function consisting of the ELBO objective and the ERM objective. 
strengths_and_weaknesses: Strengths:
(1) This reviewer thinks that this paper has good novelty. It proposes an important but overlooked observation in molecular representation learning. In short, the problem addressed in this paper is significant, the proposed solution is sufficiently novelty and the motivation is reasonable; 
(2) The presentation of each section is quite clear and well-organized. It is not hard for this reviewer to understand the novel idea and elegant design of new MRL paradigm;
(3) Sufficient experiments over different public benchmarks under OOD settings and ablation studies demonstrate the effectiveness of each component;
Weaknesses:
(1) Decomposing strategies (BRICS) and scaffold introduction as environment actually implicitly introduce additional knowledge to MRL. The decomposing step and scaffold introduction both follow hand-engineered rules by domain experts and may have non-trivial impacts on model performance. 
questions: (1) Does different decomposing strategy affect the model performance?
(2) What if the model is trained in a simple multi-task learning setting? In other words, setting environment prediction as an auxiliary task? Would this have equivalent performance to the proposed method? Why it is necessary to design environment prediction in a VAE learning way?
limitations: This reviewer thinks the proposed method in this paper has the major limitations that the scaffold and decomposing strategy may not be robust to OOD setting although adding this additional knowledge could improve robustness of MRL.
review_title: null

Meta-review: metareview: All reviewers agreed that this paper should be accepted because of the strong author response during the rebuttal phase. Specifically the reviewers appreciated the motivation of the paper, its clarity, and added explanation and experiments included during the rebuttal. Authors: please carefully revise the manuscript based on the suggestions by the reviewers: they made many careful suggestions to improve the work and stressed that the paper should only be accepted once these changes are implemented. To these suggestions I urge the authors to add another: I strongly suggest removing section 3.2. This data generating process is not validated and is not at all necessary for your approach. The SCM is never referred to again outside of this section. All that is necessary is that one can view molecules as coming from different environments or contexts and predicting this context is useful to improve generalization. Finally, with the added space I suggest expanding figure 1 to add more examples of “environments” and make this clearer in the figure: right now you only mention briefly in the caption that different scaffolds can be thought of as different environments. If you could include more / better examples that align with your experiments this will make the motivation clearer. Once these changes are made the paper will be a nice addition to the conference!
confidence: Less certain
Decision: Accept