MUBen: Benchmarking the Uncertainty of Molecular Representation Models

Dear reviewer,

We highly appreciate your recognition to work and the detailed feedback. We will first address your questions and then provide some discussion on the weakness.

Regarding your question, some explanations are as below:

1. Terminology confusion: For classification, the deterministic method directly predicts the probabilities. So it is considered as a non-UQ baseline. For regression, your understanding is totally correct. But we find that the mean ($\mu$) trained in this way performs almost identical to a non-UQ model with a single output head and trained with MSE loss. Therefore, we can directly consider the RMSE and MAE metrics of this method as the performance of a non-UQ model. Regarding the method name, we refer to it as “deterministic” in contradict to the “distributed” outputs of BNNs. Specifically, “deterministic” methods predict the mean and variance as single values, whereas for BNNs, both mean and variance are sampled from distributions parameterized by the trained network. We acknowledge that this name might be misleading, but we did not find any precise name for such a method in previous works. We’d highly appreciate some suggestions on this topic.

2. How n_ensembles affects variance: In theory, we anticipate that the variance associated with ensemble models initially increases with the number of models, especially when this number is relatively small. This increase tends to stabilize once the number of ensemble models reaches a certain threshold. This threshold is not a fixed value but varies depending on the specific datasets and models under consideration. Based on our experience, this threshold typically ranges from 7 to 15.
   However, we did not explicitly conduct experiments to quantify this threshold in our current work. The primary reason for this is our methodological approach, wherein we calculate the average of logits before applying the output activation functions, i.e., Sigmoid for classification tasks and SoftPlus for regression variance estimation. This approach does not directly yield a measure of ensemble variance.

3. Why choosing ChemBERTa over MoLFormer: In our research, we considered incorporating MoLFormer. However, during our testing and adaptation phase using the provided script, we encountered challenges similar to the documented issue of a missing model checkpoint. This led us to believe that the pre-trained model checkpoint was not made available by the authors, and we subsequently ceased further investigation into this model.
   As an alternative, we identified ChemBERTa as a widely-recognized model also with SMILES input, which also offered ease of implementation. This practical consideration influenced our decision to opt for ChemBERTa over MoLFormer.
   Upon double-checking, we noticed the updates in the MoLFormer GitHub repository that address the checkpoint issue. We plan to include MoLFormer in our future benchmarking efforts and discussions. In fact, while we were investigating the pre-trained models, we noticed that many works did not open-source their models, and we assumed the same situation for MoLFormer.

Regarding the weakness, we find the materials you provided are very helpful in strengthening our paper, and we will include them and related discussions in an updated draft, which will be posted a few days later.
For investigating the effect of pre-training on UQ, we find it could be an interesting future direction, but may also be a bit tricky in practice. The main concern is the model size. To incorporate as much knowledge as possible, pre-trained models are generally large, containing millions of parameters (Table 5). However, such a huge model can severely overfit the training data, even with the aid of UQ methods, when their parameters are initialized randomly instead of from the pre-training, which is exactly what happened in our experiments. But for smaller models, pre-training does not function well: the knowledge gained from pre-training is soon forgotten after a few steps of fine-tuning. This is also the reason why few papers directly compare the performance of pre-trained and non-pre-trained models with the same size on prediction tasks. Nonetheless, we will think about it more carefully and see if we can come up with a solution.

We hope our responses address your concerns and resolve your questions regarding our work.

Best regards,
MUBen authors

Dear reviewer,

We highly appreciate you spending time reviewing our work and providing detailed feedback. We address your concerns below.

1. Limited Contribution: In our introduction, we provided a discussion on the limitations of previous works, the motivation of this research, and the key points of our contribution. We believe that our work goes beyond simply reporting the results from different combinations. It also provides a detailed discussion of the behaviors of different backbone/UQ methods under different situations. Therefore, we consider our work to contribute fairly to the community and meet the requirements of ICLR’s benchmark track.

2. Including more backbones per molecular descriptor: We acknowledge the potential insights gained from testing various backbone models with identical molecular descriptor inputs. However, conducting such expansive experiments is not feasible within the scope of our current research, primarily due to the significant increase in required resources and efforts. Our experimental scale already surpasses that of prior studies [1, 2] by several magnitudes. Incorporating multiple backbone models for each descriptor would necessitate an exponential increase in these resources, making it impractical. Furthermore, it's important to note that while comparing backbone models holds value, our study's primary objective is to explore and evaluate Uncertainty Quantification (UQ) methods and their performances. We believe that the scale and depth of our current experiments are sufficient to substantiate the claims made in our paper. While we are open to integrating additional backbones like MoLFormer, as suggested by other reviewers, a comprehensive investigation of molecular descriptors is beyond our current project's scope. We encourage the scientific community to further this line of inquiry in future research endeavors.

3. Variance among random runs: We presented the variances of multiple runs in the Appendix, from Table 9 to Table 38. We find did not include them in the main article due to space limitations.

4. Unsupported hypothesis: We think our discussions are well-grounded and supported by the figures and tables in both the main article and the Appendix. It would be helpful if you could point out the unrooted hypothesis so that we may have a chance it make it clearer.

5. Conclusion unjustified: In Figure 3, the line indicating temperature scaling is closer to the line y=x, which indicates a better calibration and mitigated overconfidence. In Figure 5, there’s a gap between TorchMD-NET’s and Uni-Mol’s ranks for biophysical properties on ROC-AUC, which is the metric evaluating the property prediction performance.

6. Should use TDC: We acknowledge TDC might be newer, but MoleculeNet is so far the most popular benchmark for modeling molecular representation methods and the first choice of many recent works, such as [3]. We can expand our MUBen to include some TDC datasets in future works, but we still consider MoleculeNet as our primary testbed to present our studies and discoveries.

We hope these clarifications address your concerns effectively. We would be grateful if you could reconsider your evaluation of our manuscript in light of these responses.

Best regards, MUBen authors

[1] Wollschläger, Tom, et al. "Uncertainty Estimation for Molecules: Desiderata and Methods." (2023).
[2] Varivoda, Daniel, et al. "Materials property prediction with uncertainty quantification: A benchmark study." Applied Physics Reviews 10.2 (2023).
[3] Qian, Chen, et al. "Can large language models empower molecular property prediction?." arXiv preprint arXiv:2307.07443 (2023).

Dear Reviewer,

We're appreciative of your constructive feedback on our work. As the rebuttal session draws to a close, we would like to kindly remind you to check out our responses and updates of the draft. We would greatly value your insights on these updates, and welcome any further suggestions or concerns you might have.

Dear Reviewer,

Thank you for your reply! We would like to provide further explanations for some points.

2. More backbones: Yes, you are right about the computation complexity. We apologize for our mistake on that. Still, our study focus is not on how molecular descriptors affect the model performance, so we had to be selective while choosing backbone models to maintain the experiments within a manageable scale. Currently, we have GIN and TorchMD-NET as second backbones for 2D and 3D graphs, and we plan to add MoLFormer for SMILES. We have to seek contributions from the community for more backbone models.

3. Variance: Your observation is accurate, and we did notice such cases. However, as our discussions are mostly based on the ranks averaged on backbones, UQ methods, or datasets, we would assume that the 1-$\sigma$ issue you mentioned is largely mitigated and does not have a significant impact on our conclusions.

4. Unsupported hypothesis: Thanks for the examples, we can address these concerns as follows. 1) The page 6 statement is a summarization of the contributions claimed by the Uni-Mol paper [1]. Such claims are scattered across their Introduction, Methods, Conclusion, and Appendix, so forgive us for not being able to give a precise reference. 2) The page 9 statement compares the performance of Uni-Mol and TorchMD-NET---both trained on 3D graphs but on different data. We believe Figure 5 shows that Uni-Mol largely outperforms TorchMD-NET in property prediction (not uncertainty estimation) for all properties except for Quantum Mechanics (QM) and performs similarly for QM. Combining the results from Figure 2 and Tables 9--22, we believe that our statement is reasonable.

5. Conclusion unjustified: 1) Figure 3. The Temperature Scale (TS) plots are indeed closer to the $y=x$ line for all 4 subfigures, which can be proven by the ECE numbers in the tables. As mentioned in Appendix D1, Paragraph ECE, ECE calculates the average gaps between the calibration plot (Figure 3) and $y=x$. We observe smaller TS ECE values for the models on the presented datasets, as shown in Table 12 and Table 14. Even though Figure 3 may not be extremely obvious, the conclusion stands. 2) NLL, as mentioned in Section 3, serves as a UQ metric rather than an evaluation of the prediction, which is only represented by ROC-AUC for classification and RMSE & MAE for regression. We are uncertain what the sentence "the NLL gap for quantum properties is similar to the ROC gap for biophysical properties" implies. We do not entirely understand how Uni-Mol's better performance in UQ undermines our statement about Uni-Mol's performance being much better than TorchMD-NET in biophysical property prediction. It would be helpful to clarify the statement. 3) The assertion that "there is simply a lot of eyeballing of comparisons rather than a more rigorous approach" challenges our commitment to academic integrity. We kindly ask that you furnish more substantial evidence to support this claim.

About the unaddressed points, we think they are discussed in Section 5 Experiment, Section 6 Conclusion, and Appendix E.

Let us know if you have further questions or concerns.

Best regards,
MUBen Authors

[1] Zhou, Gengmo, et al. "Uni-Mol: a universal 3D molecular representation learning framework." (2023).

Dear reviewer,

We highly appreciate you spending time reviewing our work and providing detailed feedback. We noticed an overlap in your discussion on the weaknesses and questions, and we address them collectively below:

1. Q1&Q2 Computation cost: We indeed included a theoretical analysis of the computation cost of each UQ method. Due to the page limitation of the main article, we presented the results in Table 6 in the Appendix and provided a short discussion in subsection C.2. In addition, if you are also interested in the efficiency of backbone models, we also provided some numbers in Table 5 in the Appendix.
2. Q3 Datasets other than MoleculeNet: From our perspective, we don’t think adopting pre-trained models that have been proven effective on MoleculeNet would have a significant impact on the fairness of our evaluation. The reason is that 1) these models are not pre-trained on the MoleculeNet data, 2) MoleculeNet is still a popular choice to benchmark molecular representation learning methods. We think that the good performance on MoleculeNet is not a result of dataset over-fitting or arduous hyper-parameter tuning, but rather a fair indicator of their general abilities, at least within the data distributions represented by MoleculeNet, which are comprehensive enough to cover most aspects of the material properties. Therefore we don’t reckon that testing UQ methods based on the backbones already having good performance on MoleculeNet would cause biased results.
   In addition, we also provided test results on the randomly split datasets, which deviates from the setting of the previous works and indeed triggered some interesting discoveries, as discussed in Appendix E.3. Therefore, we believe that we have provided a fair and comprehensive analysis of both the backbones and UQ methods in our work. However, we also made it convenient to conduct experiments on any new datasets, as indicated in our Code README. We will continuously improve our benchmark to make it more inclusive.

We hope these clarifications address your concerns effectively. We would be grateful if you could reconsider your evaluation of our manuscript in light of these responses.

Best regards,
MUBen authors

Dear Reviewer,

Thank you for your clarification and apology for our previous misinterpretation of your statement. We would like to address your concern about the pre-training data from several perspectives:

1. To our knowledge, there is no strict evaluation of how MoleculeNet data overlap with the pre-training dataset of each pre-trained model. Therefore, we cannot know concretely whether their good performance comes from the overlapping of pre-training and test data. Still, we cannot directly build the causal relation of “they perform well on this dataset, so their pre-training data probably have already covered the test data distribution”. There are many factors that can affect the performance of molecular representation models, such as the network structure, input data format, pre-training objective, etc. Even if we include new datasets in MUBen, there is no guarantee that the new datasets certainly fall outside of the distribution of the pre-training data.

2. In fact, many backbones are not originally evaluated on the entire MoleculeNet. For example, ChemBERTa is evaluated only on BACE, Lipo, BBBP, and ClinTox [1], and GROVER on BBBP, SIDER, ClinTox, BACE, Tox21 and ToxCast [2]. Under this circumstance, even if we do assume the aforementioned causal relationship is correct, we suppose our evaluation of these models on the Quantum Mechanics properties is unbiased. We can also make a similar claim that TorchMD-NET, which is only pre-trained on Quantum Mechanics data, has an unbiased evaluation of other Biophysics, Physiology, and Physical Chemistry properties. In addition, as mentioned in the previous response, we also include DNN and GIN, which are not pre-trained on any data. Therefore, the test data are guaranteed to be OOD for them, and their evaluations are unbiased. As we observe a similar trend in the performance of UQ methods as the pre-trained backbones, we may conclude that our evaluation is largely accurate.

3. Comparing the results from scaffold split datasets (Table 9 – Table 22) and random split datasets (in-distribution test data, Table 23 – Table 30), we notice that there are performance gaps between the pre-trained models when they are tested on OOD or in-distribution data. This indicates that the models are less familiar with the former dataset, and UQ methods can help improve their evaluation in such situations.

4. We developed this benchmark and codebase for practical purposes, and we have applied MUBen to our own data (apologize for not being able to reveal these data—they are so far still confidential). Our experiments demonstrate the same backbone-wise and UQ-wise performance. Therefore we can claim with certain confidence that our evaluation of MoleculeNet is faithful and able to guide practical backbone and UQ selection.

We hope our response addresses your concern. We would be grateful if you raise your evaluation of our manuscript in light of these responses. Let us know if you have further questions or suggestions!

Best regards,
MUBen authors

[1] Ahmad, Walid, et al. "Chemberta-2: Towards chemical foundation models." arXiv preprint arXiv:2209.01712 (2022).
[2] Rong, Yu, et al. "Self-supervised graph transformer on large-scale molecular data." Advances in Neural Information Processing Systems 33 (2020): 12559-12571.
[3] Zaidi, Sheheryar, et al. "Pre-training via denoising for molecular property prediction." arXiv preprint arXiv:2206.00133 (2022).

Regarding the number M of deep ensembles, we use M=10 on most datasets and only use 3 on 3 very large ones (HIV, MUV, QM9) to reduce the computational cost, as mentioned in Appendix C1, Paragraph “Deep Ensembles” under equation 10. Using smaller M does improve the model performance over the “deterministic” baseline, but it does not necessarily surpass other UQ methods with a much cheaper cost. In addition, using a small M may cause inadequate exploration of the latent space, resulting in less stable performance gain (or loss). A larger M would give a more constant promotion. We did some experiments to investigate the impact of M. The results are shown below:

Table 1, DNN on Lipo dataset

Table 2, DNN on Tox21

Table 3. Uni-Mol on Lipo

Table 4. Uni-Mol on Tox21

We hope our response addresses your concern. We would be grateful if you raise your evaluation of our manuscript in light of these responses. Let us know if you have further questions or suggestions!

Best regards,
MUBen authors

Dear Reviewer,

Thank you for your valuable feedback on our manuscript. We address your concerns as follows:

1. UQ Illustration Update: We acknowledge the need to update the illustration and will include them in the revised version that will be posted in a few days.

2. MolFormer Comparison: Our initial attempt to test and adapt MolFormer using this script encountered the error of missing model checkpoint, similar to this issue. However, we now recognize that these issues have been resolved in the updated GitHub repository. We plan to incorporate MolFormer in our benchmark in future work. In fact, while we were investigating the pre-trained models, we noticed that many works did not open-source their models, and we assumed the same situation for MoLFormer.

3. Discussion on SGLD: We have provided a detailed discussion on this topic, as indicated in the second paragraph below Figure 4, starting with “Although limited in classification efficacy…”

4. Figure 4 Subfigures Selection: We chose these subfigures to compare ChemBERTa with DNN and Uni-Mol with DNN respectively on two datasets to show a general trend. We selected DNN as an anchor since it has the best regression calibration, as shown in Figure 2 (b). As the space for the main paper body is limited, we only selected the most representative figures to support our argument and left other results in the appendix. We have checked all figures during experiments and they show a coherent pattern as displayed and described in our paper.

5. Table 3 Content and Discussion: The results in Table 3 represent a macro average across datasets, backbones, and UQ methods, and are not exclusive to Uni-Mol, as the caption indicates. We spend the second half of the paragraph “Frozen Backbone and Randomly Split Datasets” in section 5 and the whole appendix subsection E3 discussing the behavior of the models on randomly split datasets. So we think the statement “discussion mainly focuses on frozen backbone” is not fully justified. Randomly splitting also performs better on Regression tasks. For regression metrics, the scores are the lower the better, as indicated by the arrows beside each metric.

6. Best Configuration of Method Choices: By “configuration”, if it refers to what backbone/UQ method we should choose under a certain scenario, then we provided some insights in our Conclusion section. If it refers to the model hyper-parameters, we have set them as the default parameters in our code. For the data-splitting strategy, it depends on the problem set-up (whether it is in-distribution or out-of-distribution) and there is no “best configuration” for it.

We respectfully disagree with your point that “novelty is limited”, we would like to argue that we did provide insights in our experiments on the model and UQ selection in different scenarios in experiments (Section 5) and conclusion (Section 6). Despite experimental results suggesting that the ensemble method is a superior choice, we also would like to highlight the contribution of a benchmark paper does not only come with the conclusion and suggestion of which method is the best but also includes:

1. Set a fair comparison and evaluation pipeline and encourage more researchers to study better UQ methods (or specific adaptation or inspiration from the molecular domain)

2. Provide the implementation of the UQ method ready to use for the researchers and practitioners

We hope these clarifications address your concerns effectively. We would be grateful if you could reconsider your evaluation of our manuscript in light of these responses.

Sincerely,
MUBen Authors

Dear Reviewer,

We're appreciative of your constructive feedback on our work. As the rebuttal session draws to a close, we would like to kindly remind you to check out our responses and updates of the draft. We would greatly value your insights on these updates, and welcome any further suggestions or concerns you might have.