RoFt-Mol: Benchmarking Robust Fine-tuning with Molecular Graph Foundation Models

We thank reviewer iuUw for appreciating that our benchmark aligns with practical scenarios and recognizing our contribution in method introduction. We will respond to the questions as follows.

W1 / Q1: More powerful pre-trained model

Regarding the Mole-BERT and GraphMAE, we directly use the checkpoints from existing works as improving pre-training methods is beyond the scope of this work. And selecting more complex backbones presents uncertainty about whether they can be effectively pre-trained using current pre-training methods.

Thank you for suggesting the inclusion of additional pre-trained models. We have added MoleculeSTM [1] as a pre-trained model for our downstream fine-tuning evaluations and have discussed related multi-modal molecular graph models to our revised draft. We are currently benchmarking the results and aim to incorporate the available findings into the new paper draft by the end of the rebuttal period. However, due to time constraints and computational costs, some benchmarking results on the large MUV and HIV datasets may not be ready by then; we plan to include any remaining findings in a subsequent version of the paper.

[1]. Liu, Shengchao, et al. "Multi-modal molecule structure–text model for text-based retrieval and editing." Nature Machine Intelligence 5.12 (2023): 1447-1457.

Q2: How do scaffold and size split impact the robustness of fine-tuning methods in molecular property prediction tasks?

In general, scaffold and size splits introduce larger distribution shifts compared to random splits, making them more receptive to advanced fine-tuning strategies beyond basic vanilla fine-tuning. Between scaffold and size splits, scaffold splitting calls for more specialized fine-tuning techniques, resulting in significant performance gains over baseline methods, particularly in severe few-shot scenarios like few-shot 50. As discussed in Finding 5, size splits are notably vulnerable under linear probing with fixed representations and a tunable prediction head. The prediction head tends to overfit to the specific mapping from representations to output labels for molecules within a certain size range, impairing its ability to generalize to out-of-distribution molecules of differing sizes.

W5: The findings in this paper seem common, which is similar in other public datasets, could you provide more insights on MOLECULAR research?

• Evaluation on Regression Tasks: Previous research on robust fine-tuning has largely focused on classification tasks [1, 2], especially in computer vision and NLP. In NLP, tasks like token generation are akin to classification because they involve selecting the correct token from a discrete set, rather than focusing on numerical precision. However, in molecular research, regression tasks, particularly for molecular property predictions, are prevalent and critical. Understanding robust fine-tuning for regression datasets is essential due to their numerical sensitivity. As noted in the second takeaway of our revised introduction, we found that regression tasks generally have a lower risk of overfitting compared to classification, owing to the need for precise numerical labels and detailed molecular modeling. These insights are vital for choosing suitable fine-tuning methods based on task type.
• Consideration of Diverse Pre-training Models: Previous research predominantly examines self-supervised pre-trained models like CLIP [1, 2]. However, there's increasing interest in supervised pre-trained molecular models [3, 4]. By evaluating both self-supervised and supervised models, our study reflects the diverse molecular models in practice and provides insights into how pre-training affects downstream fine-tuning choices. In our revised draft introduction, we compare these models and outline fine-tuning methods suited to each in the respective takeaways.
• Analysis unique for Molecular Tasks: We provide analysis specifically within the context of molecular tasks. First, in Q1 and Finding 1, we highlight the distinct needs of regression versus classification tasks: classification tasks often require coarser features, like functional groups, while regression tasks demand finer details, such as charge distribution and hydrogen bond patterns. Second, we identify the misalignment between general-purpose self-supervised pre-training in molecular settings (which focuses on substructure discovery [5]) and the task-specific requirements of downstream applications. This misalignment explains why linear probing underperforms in molecular tasks and underscores the importance of weight-based methods, which effectively integrate broad pre-trained patterns with task-specific fine-tuning to enhance performance.

These molecular-specific insights have also directly inspired the development of DWiSE-FT, which leverages weight-based methods to address the unique needs of both regression and classification tasks effectively. For more detailed discussions on these findings and further insights relevant to molecular research, please refer to sections 4.1 and 4.2 of our revised paper.

[1]. Kumar, Ananya, et al. "Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution." International Conference on Learning Representations.

[2]. Wortsman, Mitchell, et al. "Robust fine-tuning of zero-shot models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[3]. Shoghi, Nima, et al. "From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction." The Twelfth International Conference on Learning Representations.

[4]. Beaini, Dominique, et al. "Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets." The Twelfth International Conference on Learning Representations.

[5]. Wang, Hanchen, et al. "Evaluating self-supervised learning for molecular graph embeddings." Advances in Neural Information Processing Systems 36 (2024).

Q: Why choose GraphMAE and Mole-BERT as pre-training model

We selected these two pre-training models because they represent leading strategies in both reconstruction-based and contrastive-based pre-training approaches. Additionally, they are recognized as well-known and relatively new state-of-the-art models in the field.

We thank reviewer CyuA for acknowledging the importance of our findings and the introduction of DWiSE-FT. We will address the listed questions below.

W1: The paper does not provide a detailed comparison of DWiSE-FT with other state-of-the-art methods in the field.

Thank you for your interest in our proposed DWiSE-FT method and the feedback. We chose to compare our DWiSE-FT method with WiSE-FT and $L^2$-SP because they are the top-performing algorithms among all SOTA robust fine-tuning methods included in our benchmark, as shown in Figures 2(a) and 2(c). By doing so, we effectively compare our method against all included baselines, since these two represent the best performances. Additionally, we report the "best" scores for each experimental setting in Table 3, which encapsulate the optimal possible results considering all fine-tuning methods together. Therefore, we believe our comparisons can show the position and importance of DWiSE-FT. if the reviewer knows some potentially SOTA methods to compare, we would like to incorporate. If the rebuttal time is limited, we would like to put in the final version.

W2: The abbreviation PT in Line 15 has no explanations

Thanks for pointing out the confusion. To avoid confusion, we omit using the “PT” abbreviation in the revised draft.

W3: Computational efficiency and scalability of DWiSE-FT

DWiSE-FT is significantly more efficient in optimization compared to other fine-tuning methods involving re-training. DWiSE-FT is essentially a post-hoc weight interpolation without any training needed given the interpolation coefficients, but we further introduce an automatic hyperparameter search to find the optimal interpolation coefficients. Even with the hyperparameter search, it only has $l$ number of parameters to optimize where $l$ is the number of layers in the encoder. In contrast, other fine-tuning methods initially require fine-tuning on model weights that scale with $l*d$ with $d$ being the model hidden dimension. Additionally, the tuning on hyperparameters will further lead to multiple rounds of model re-training.

W4: The performance of DWiSE-FT in combination with other pre-training paradigms

Thank you for highlighting this intriguing idea. Indeed, we have conducted some preliminary experiments combining DWiSE-FT with other pre-training approaches, such as graphMAE and supervised pre-training, and observed promising improvements over WiSE-FT and $L^2$-SP. However, due to time constraints, we did not include a complete evaluation in the current paper. We will include these comprehensive evaluations in a future version.

We appreciate reviewer Cxf8 for acknowledging the significance of the direction highlighted in our benchmark, as well as recognizing our efforts in benchmarking and proposing the new method. We will address the weaknesses and questions raised below.

W1: Lack of new datasets

We would like to kindly state that "Datasets and Benchmarks" is an established focus area at ICLR. Although specific details about this direction are not outlined on the ICLR webpage, we reference the NeurIPS Datasets and Benchmarks Track, which defines its scope as encompassing "benchmarks on new or existing datasets, as well as benchmarking tools." This indicates that contributing a new dataset is not necessary; impactful work can also concentrate on developing benchmarks or tools that are applied to existing datasets.

W2: Over-simplified backbones

As improving pre-training methods is beyond the scope of this work, we utilize checkpoints from existing studies. Selecting more complex backbones presents uncertainty about whether they can be effectively pre-trained using current pre-training methods.

W3: Limited pre-training for Graphium

We adopted the pre-trained model on ToyMix for Graphium, as larger pre-training datasets do not currently have available checkpoints. We attempted to reproduce the pre-training on LargeMix using the Graphium library, but encountered computational constraints when processing the datasets. Additionally, using ToyMix allows for a more fair comparison with other self-supervised pre-trained models regarding the scale of the pre-training model and data. In the future, we will attempt again to obtain the LargeMix model and evaluate whether the downstream fine-tuning performances align with the trends observed using the ToyMix model.

Q: Table format and claim revision

Thanks for pointing out the inappropriate wordings and points to improve in the paper presentation, we have updated those in our new draft.

We thank reviewer u2Et for appreciating our paper presentation and insights provided, we will respond to the concerns accordingly below.

W1: Limited understanding about the properties of architectures that leads to the performance difference

Thanks for raising the clarification question. We suppose the properties of architectures you are referring to are the mechanisms of the fine-tuning methods. To better analyze and discuss the fine-tuning performances of different methods, we categorize these 8 different fine-tuning methods to 3 groups based on their fine-tuning mechanisms. With the first takeaway in our revised draft, we point out the key insights that explain the performance of different categories of methods. For instance,

• As discussed in finding 1, since weight-based methods have the property to ensemble weights from pre-trained and fine-tuned models, they are able to extract molecular patterns learned from general-purpose self-supervised pre-training [1] and combine them together with task-specific knowledge from fine-tuning.
• As discussed in finding 4, since regularization-based methods have the mechanism in enforcing close proximity of fine-tuned representation to pre-trained ones, the fine-tuned model successfully preserves domain-relevant representations learned from supervised pre-training given similar pre-training and fine-tuning tasks. More findings that explain the fine-tuning mechanisms can be found in section 4.

[1]. Wang, Hanchen, et al. "Evaluating self-supervised learning for molecular graph embeddings." Advances in Neural Information Processing Systems 36 (2024).

W2: The study’s scope does not extend to multi-task, multi-modality foundation models nor foundation models across diverse scientific domains.

Thank you for suggesting potential directions to extend our benchmark. Looking ahead, one future direction is to investigate robust fine-tuning with larger multi-task and multi-modality models. Currently, our benchmark focuses on robust fine-tuning methods for molecular graph foundation models involving small molecules. We have included diverse fine-tuning methods and representative pre-trained models to support this focus. Furthermore, our range of scope is aligned with previous robust fine-tuning studies, which commonly focus on a single pre-trained model, like CLIP, and primarily target a single application domain, such as image or object classification [2, 3]. As these prior works demonstrate, it is not essential for robust fine-tuning studies to encompass a wide array of pre-trained models across diverse applications.

[2]. Kumar, Ananya, et al. "Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution." International Conference on Learning Representations.

[3]. Wortsman, Mitchell, et al. "Robust fine-tuning of zero-shot models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

W3: Additional fine-tuning methods from other categories

Thank you for providing the relevant works. However, the studies you mentioned primarily focus on pre-training foundation models that span multi-task applications from molecules to materials, and they do not introduce novel downstream fine-tuning techniques. Since our aim is to benchmark fine-tuning methods, we believe that these suggested works fall outside the direct scope of our current focus and are not directly comparable.

Q: Sensitivity analysis on DWISE-FT hyperparameters

This is an insightful question that highlights a key aspect of fine-tuning methods. Our proposed DWiSE-FT includes hyperparameters being the initial values for the mixing coefficients in each encoder layer. Since DWiSE-FT automatically optimizes these coefficients using the validation loss, the initializations are less critical compared to the manually set hyperparameters in other methods that cannot be adjusted. In our experiments, we tested various starting points for the mixing coefficients within the range of [0,1], and found that the variance in final results was small. This is because the coefficients converge to their optimal values, regardless of the initial settings, given adequate training time. We have also clarified this in the revised draft, specifically in Appendix E, where we discuss the hyperparameter tuning.