Enhancing Molecular Property Prediction with Auxiliary Learning and Task-Specific Adaptation

Vishal Dey, Xia Ning

Keywords: Graph Neural Networks, Auxiliary Learning, Molecular Property Prediction, Transfer Learning, Adaptation

Abstract: Pretrained Graph Neural Networks (GNNs) have been widely adopted for various molecular property prediction tasks. Despite their ability to capture rich chemical knowledge, traditional finetuning of such pretrained GNNs on the target task can lead to poor generalization. To address this, we explore the adaptation of pretrained GNNs to the target task by jointly training them with multiple auxiliary tasks. This could enable the GNNs to learn both general and task-specific features, which may benefit the target task. However, an effective adaptation strategy needs to determine the relevance of auxiliary tasks with the target task, which poses a major challenge. In this regard, we investigate multiple strategies to adaptively combine task gradients or learn task weights via bi-level optimization. Our experiments with state-of-the-art pretrained GNNs demonstrate the efficacy of our proposed methods, with improvements of up to 8.45% over finetuning. Overall, this suggests that incorporating auxiliary tasks along with target task fine-tuning can be an effective way to improve the generalizability of pretrained GNNs for molecular property prediction tasks, and thus inspires future research.

Submission Type: Extended abstract (max 4 main pages).

Software: https://github.com/vishaldeyiiest/GraphTA

Submission Number: 199