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

Authors: Shikun Liu, Deyu Zou, Nima Shoghi, Victor Fung, Kai Liu, Pan Li
Primary area: datasets and benchmarks
Keywords: Molecular representation learning, Fine-tuning

Abstract:
In the era of foundation models, fine-tuning pre-trained models for specific downstream tasks has become crucial. This drives the need for robust fine-tuning methods to address challenges such as model overfitting and sparse labeling. Molecular graph foundation models (MGFMs) face unique difficulties that complicate fine-tuning. These models are limited by smaller pre-training datasets and more severe data scarcity for downstream tasks, both of which require enhanced model generalization. Moreover, MGFMs must accommodate diverse pre-training objectives, including both regression and classification tasks. To better understand and improve fine-tuning techniques under these conditions, we classify eight fine-tuning methods into three mechanisms: weight-based fine-tuning, representation-based fine-tuning, and partial fine-tuning. We benchmark these methods on downstream regression and classification tasks across both supervised and self-supervised pre-trained models in diverse labeling settings. This extensive evaluation provides valuable insights and informs the design of a refined robust fine-tuning method, DWiSE-FT. This approach combines the strengths of simple post-hoc weight interpolation with more complex weight ensemble fine-tuning methods, delivering improved performance across both task types while maintaining the ease of use inherent in post-hoc weight interpolation.

Decision: reject
Reviewer scores: 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1729084035218

thread_type:
official_review

<thread_object>
creator: reviewer_iuuw

<content>
summary:
This paper addresses the challenges and importance of fine-tuning pre-trained molecular graph models. The work explores 8 fine-tuning methods categorized into weight-based, representation-based, and partial FT, benchmarked across 12 datasets with 36 experimental settings. These settings aim to simulate real-world FT scenarios, including OOD generalization and Fewshot settings. The paper highlights the strengths of different FT strategies depending on whether tasks are classification or regression-based. A new method, DWiSE-FT, is proposed, combining aspects of weight-based methods to enhance fine-tuning efficiency.

strengths:
1. The paper benchmarks 8 FT methods across 12 datasets and various experimental settings, providing a thorough evaluation of the fine-tuning methods for molecular graph models.

2. This paper proposes DWiSE-FT based on their findings, and DWiSE-FT performs well on regression tasks.

3. The benchmark is similar to the practical scenarios of molecular representation learning, by including scaffold and size splitting.

weaknesses:
1. The pre-trained molecular models evaluated in this paper, such as Mole-BERT and Graph MAE are small in scale compared to other foundation models. The authors should evaluate more powerful pre-trained models such as [1] and [2].

2. Though DWiSE-FT shows improved performance, it is based on existing FT strategies rather than novel FT methods.

questions:
1.  Recently, many multi-modal molecular graph models have been proposed. These models can gain additional information from texts and thus are more expressive. Can you conduct experiments on these more expressive molecular graph models such as [1] and [2]? 

2. How do scaffold and size splits, which simulate OOD challenges, impact the robustness of fine-tuning methods in molecular property prediction tasks?

[1] MolCA: Molecular Graph-Language Modeling with Cross-Modal Projector and Uni-Modal Adapter
[2] Multi-modal Molecule Structure-text Model for Text-based Retrieval and Editing
</content>
</thread_object>

</thread>
<thread>
creation date:
1730456177619

thread_type:
official_review

<thread_object>
creator: reviewer_cyua

<content>
summary:
This paper presents a benchmarking study on robust fine-tuning (FT) with molecular graph foundation models. The authors evaluate different fine-tuning strategies under supervised and self-supervised pre-training (PT) paradigms, using various molecular datasets. They introduce a refined method called DWiSE-FT, which consolidates the strengths of existing FT methods and demonstrates improved efficiency while maintaining promising results for regression tasks.

strengths:
1. Comprehensive Evaluation:The paper conducts an extensive evaluation of different fine-tuning strategies across multiple datasets, providing a comprehensive understanding of their performance.
2. The introduction of DWiSE-FT represents a novel approach to fine-tuning, combining the strengths of existing methods.
3. The paper provides valuable insights into the design of FT methodologies and practical guidance for molecular representation learning. The findings have implications for both researchers and practitioners working in the field of molecular graph representation learning.

weaknesses:
1. While the paper compares different fine-tuning strategies, it does not provide a detailed comparison of DWiSE-FT with other state-of-the-art methods in the field. This paper presents a benchmarking study, the comparisons and findings are important and novel. I want to know the importance of this proposed method DWiSE-FT , can it be presented in a research paper?
2. The abbreviation PT in Line 15 has no explanations.
3. In sec 2.2, the authors present different fine-tuning methods, the proposed Adaptive post-hoc ensemble method is similar with WiSE-FT. More information about the computational efficiency and scalability of DWiSE-FT would be beneficial. 
4. Explore the performance of DWiSE-FT in combination with other pre-training paradigms.
5. The findings in this paper seems common, which is similar in other public datasets, could you provide more insights on MOLECULAR research?

questions:
1. In Line111-115, why choose GraphMAE and Mole-BERT as the PT model?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730588884067

thread_type:
official_review

<thread_object>
creator: reviewer_cxf8

<content>
summary:
In this work, the authors proposed a novel benchmark to understand the impact of different fine-tuning methods for pre-trained molecular foundation models in downstream tasks. Specifically, 8 fine-tuning methods grouped in 3 categories are benchmarked over 12 datasets and 36 experimental settings. Insights are then derived from the experimental results to understand which methods are best suited for which needs. In addition, a refined method named DWiSE-FT is proposed to enable more efficient fine-tuning with competitive performance.

strengths:
Recently, various foundation models for molecular graphs have been proposed, however, to apply them in downstream tasks, a well-designed fine-tuning procedure is required. This work benchmarks a wide range of fine-tuning methods to provide insight into this aspect. Here are the strengths of the work: 

- **Important direction.** Understanding how to select the best fine-tuning strategy for a downstream task is highly important for molecular foundation models. 

- **Extensive experiments.** In this work, 8 FT methods are examined and compared across 12 datasets and 36 experimental settings, providing an empirical comparison for a variety of settings. 

- **novel method DWiSE-FT.** The authors proposed DWiSE-FT which is a great candidate for regression datasets.

weaknesses:
Here are the weaknesses of this work. 

- **lack of new datasets.** The main contribution of this work is proposing a novel benchmark for fine-tuning. However, there are no novel datasets to provide more benchmark environments outside of existing datasets. 

- **over-simplified backbone.** The authors mention that "high model expressiveness" is needed to capture the semantics of molecular datasets. However, for experiments, only a 5-layer GIN architecture is used as the backbone, which is known to be limited by the 1-WL test in expressiveness. A more powerful architecture can be used such as GraphGPS[1] or Graphormer [2]. Experiments with a more powerful backbone PT architecture can also provide more insight into the impact of the choice of PT architecture. 

- **limited pre-training for Graphium.** In the Graphium paper, three categories of datasets are provided: Toymix, Largemix, and UltraLarge. It would be more interesting to observe results for PT models trained on Largemix or both Toymix and Largemix. 


[1] Rampášek L, Galkin M, Dwivedi VP, Luu AT, Wolf G, Beaini D. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems. 2022 Dec 6;35:14501-15.

[2] Ying C, Cai T, Luo S, Zheng S, Ke G, He D, Shen Y, Liu TY. Do transformers really perform badly for graph representation?. Advances in neural information processing systems. 2021 Dec 6;34:28877-88.

questions:
-  The tables in the paper are too small and very hard to read or draw conclusions from. For example, Table 3 is very small and not readable at all.  The authors should update the format of the Table or move less relevant content to the appendix. 

- I can't entire agree with the claim on line 47 "pre-trained on insufficient amount of PT data
(1M-100M samples) and vocabularies". The authors made several comparisons between the molecular foundation model and foundation models from NLP / vision throughout the paper. However, the reality is that molecular data is significantly harder to curate and needs to be explicitly constructed for effective learning. Thus, I don't believe that the same scale of available data for NLP or vision would ever be available for molecular learning. For example, 100 M molecular graphs are already a large number and the focus should instead be in more to design effective PT and FT strategies. If possible, the authors should revise this claim.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730653068726

thread_type:
official_review

<thread_object>
creator: reviewer_u2et

<content>
summary:
The paper explores optimal Fine Tuning (FT) strategies for molecular representation, addressing challenges such as label scarcity and distribution shifts. It introduces an enhanced FT method, DWiSE-FT, tailored for diverse pre-trained molecular graph models. This method promises efficiency and automation in specific FT scenarios, consistently achieving top-ranking results.

strengths:
- I have found the paper well written and self-contained. I think a non-expert could find most of the information in the paper, and I appreciate this aspect.

 - Figure 1, which demonstrate the problem domain and architecture, are interesting and easy to read. I commend the authors for their explicit effort in making these illustrations clear and informative.

 - The insights are didactical and well communicated. The conclusion given by the experiments looks interesting and valuable for future practitioners, while I think a synthesis would be beneficial for the reader.

weaknesses:
Despite these merits, I have the following concerns about the paper.

1- While there is a careful analysis of the different design decisions/performance tradeoffs, I feel that there is only a limited understanding about what are the properties of the Architecture that lead to these decisions/performance differences.


2-  The study's scope, while broad, does not extend to a multi-task, multi-modality approach that could significantly enhance its applicability and impact. It does not explore foundation models across diverse scientific domains such as RNA,  and proteins nor does it address varied scientific tasks likechemical reactions. Expanding the research to cover these aspects would substantially enrich the study's utility and relevance.

3- While the study incorporates several representative FineTuning (FT) methods from various categories, it does not investigate additional FT methods from other categories that could offer significant insights. I have mentioned some noteworthy models here that deserve further exploration:
“DPA-2: A Large Atomic Model as a Multi-Task Learner [Zhang 2023]”
“Scalable Training of Trustworthy and Energy-Efficient Predictive Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNN [Pasini 2024]”
“MiniMol: A Parameter-Efficient Foundation Model for Molecular Learning [Klaser 2024]”

questions:
is there any sensitivity analysis on key hyperparameters of DWiSE-FT? I would also suggest comparing the hyperparameter sensitivity of DWiSE-FT to that of baseline methods.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732305252066

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

We sincerely appreciate your time and effort in providing us with valuable feedback to enhance our paper. We are grateful for your recognition of the importance of our benchmark, which offers valuable insights to the field, and for your appreciation of our comprehensive evaluation and the novel method we proposed, DWiSE-FT.

Additionally, we thank you for raising insightful questions and suggestions. We will address each of these points in our detailed responses to individual reviewers. In the new paper draft, we have polished the writing and adjusted the presentation for clarity, while maintaining all the original conclusions and results. We are also incorporating MoleculeSTM [1], as suggested by reviewer iuUw as an additional pre-trained model for our downstream fine-tuning evaluations. We are currently benchmarking the results and will incorporate the available findings in the updated paper draft by the end of the rebuttal period. However, due to time constraints and computational costs, some results for the large MUV and HIV datasets may not be finalized in time for the rebuttal period, and we will include them in a future version.

[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.
</content>
</thread_object>

</thread>


Metareview:
metareview:
### Summary
The paper explores fine-tuning (FT) strategies for pre-trained molecular graph models, focusing on improving downstream performance under label scarcity and distribution shifts. The study benchmarks 8 FT methods across 12 datasets and 36 experimental settings, highlighting insights for regression and classification tasks. A new method, DWiSE-FT, is proposed to enhance fine-tuning efficiency, achieving competitive results, particularly for regression tasks.

### Strengths
- The study systematically evaluates 8 FT methods across multiple datasets and experimental settings, offering valuable insights into FT strategies for molecular graph models.
- The benchmark closely simulates real-world scenarios such as OOD generalization and few-shot learning, making the findings useful for practitioners.
- The proposed DWiSE-FT method consolidates existing FT strategies, providing improved efficiency and competitive performance, especially for regression tasks.

### Weaknesses
- Limited Scope of Pre-training and Architectures: The study uses relatively small-scale pre-trained models (e.g., Mole-BERT, Graph MAE) and a limited backbone architecture (5-layer GIN), which lacks expressiveness. The evaluation does not explore more powerful models like GraphGPS, Graphormer, or larger-scale pre-training paradigms.
- Lack of Novelty: While DWiSE-FT performs well, it is primarily a refinement of existing FT strategies rather than a fundamentally novel method. Its computational efficiency and scalability compared to other FT methods are not clearly demonstrated.
- The study focuses solely on molecular graphs and does not extend to multi-task, multi-modal scenarios (e.g., RNA, protein structures, or chemical reactions). The findings lack deeper insights specific to molecular research, with conclusions appearing generalizable to other public datasets.


While the paper provides a valuable benchmark of FT strategies for molecular graph models and introduces an improved method (DWiSE-FT), it suffers from limited novelty, narrow evaluation scope, and lack of scalability tests. The reliance on small-scale pre-trained models and backbones weakens the paper's impact. Extending the study to more powerful architectures, diverse pre-training paradigms, and broader tasks would significantly enhance its contributions.