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

Shikun Liu, Deyu Zou, Nima Shoghi, Victor Fung, Kai Liu, Pan Li

Keywords: Molecular representation learning, Robust 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 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, representation-based, and partial fine-tuning. We
benchmark these methods on downstream regression and classification tasks across
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, ROFT-MOL. 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.

Primary Area: AL/ML Datasets & Benchmarks for life sciences (e.g. climate, health, life sciences, physics, social sciences)

Submission Number: 1552